For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo struct

SM100MHADepth512

struct SM100MHADepth512[KVLUTType: MHAOperand, output_type: DType, MaskType: MHAMask, SchedulerType: MHATileScheduler, config: Depth512SM100Config[KVLUTType.dtype], ValidLengthType: OptionalPointer, KVRowOffsetsType: OptionalPointer, _is_cache_length_accurate: Bool, MaxSeqLenType: OptionallyStaticInt, PartitionType: MHAPartitionScheme]

Implements pair-CTA SM100 (Blackwell) multi-head attention prefill for depth=512.

Dispatches the 12-warp (3 warp-group) schedule across two cooperating CTAs: softmax, correction, MMA, load, and spare warps each run a specialized sub-kernel that shares a single Depth512AttentionSMem allocation.

Parameters

• ​KVLUTType (MHAOperand): MHA operand describing the KV cache lookup table; its dtype is the Q, K, V element type and its page_size drives the KV sub-tile row counts.
• ​output_type (DType): Output DType of the attention result (O store).
• ​MaskType (MHAMask): Mask type applied to the attention score tiles, such as a causal mask.
• ​SchedulerType (MHATileScheduler): Tile scheduler that assigns work tiles to CTAs and advances the per-CTA iteration state.
• ​config (Depth512SM100Config[KVLUTType.dtype]): Depth-512 SM100 tile and pipeline configuration (Depth512SM100Config), parameterized by the KV dtype. Holds BM, BN, head counts, depths, staging, and the TMEM/SMEM budget.
• ​ValidLengthType (OptionalPointer): Optional pointer type for the per-batch valid sequence lengths; when non-null the kernel runs in ragged mode.
• ​KVRowOffsetsType (OptionalPointer): Optional pointer type for the KV input row offsets; when non-null used to compute the per-batch KV sequence length.
• ​_is_cache_length_accurate (Bool): Whether the cache length is known to be accurate, letting the start position be zeroed in ragged mode.
• ​MaxSeqLenType (OptionallyStaticInt): Optionally-static maximum sequence length type; a static value of 1 selects decoding mode (asserted false for depth-512).
• ​PartitionType (MHAPartitionScheme): KV cache partition scheme; partitioning is asserted unsupported for the depth-512 pair-CTA kernel.

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable, RegisterPassable, TrivialRegisterPassable

comptime members

accum_type

comptime accum_type = DType.float32

BM

comptime BM = config.BM

BM_eff

comptime BM_eff = config.BM_eff()

BN

comptime BN = config.BN

cta_group

comptime cta_group = 2

fuse_gqa

comptime fuse_gqa = config.fuse_gqa

group

comptime group = config.group

k_sub_BN

comptime k_sub_BN = kv_sub_tile_rows((config // Int(2)), KVLUTType.page_size)

num_q_heads

comptime num_q_heads = config.num_q_heads

page_size

comptime page_size = KVLUTType.page_size

PairBM

comptime PairBM = (config * Int(2))

PairBM_mask

comptime PairBM_mask = (config.BM_eff() * Int(2))

PositionType

comptime PositionType = MHAPosition[Int((mul config.BM, 2)), config.BN, config.qk_depth, config.qk_depth, config.num_q_heads, config.group, _is_decoding[MaxSeqLenType]()]

qkv_type

comptime qkv_type = KVLUTType.dtype

ragged

comptime ragged = not ValidLengthType.is_null.__bool__()

SmemType

comptime SmemType = Depth512AttentionSMem[config]

v_sub_BN

comptime v_sub_BN = kv_sub_tile_rows(config.BK1, KVLUTType.page_size)

Methods

kernel

static def kernel(q_tma_op: TMATensorTile[KVLUTType.dtype, Int(4) if Self.fuse_gqa else Int(3), _padded_shape[Int(4) if Self.fuse_gqa else Int(3), KVLUTType.dtype, q_smem_shape[KVLUTType.dtype, config.swizzle_mode, BM=config.BM, group=config.group, depth=config.qk_depth, decoding=False, fuse_gqa=Self.fuse_gqa, num_qk_stages=config.num_qk_stages](), config.swizzle_mode](), _ragged_shape[Int(4) if Self.fuse_gqa else Int(3), KVLUTType.dtype, q_smem_shape[KVLUTType.dtype, config.swizzle_mode, BM=config.BM, group=config.group, depth=config.qk_depth, decoding=False, fuse_gqa=Self.fuse_gqa, num_qk_stages=config.num_qk_stages](), config.swizzle_mode]()], k_tma_op: TMATensorTile[KVLUTType.dtype, Int(3), _padded_shape[Int(3), KVLUTType.dtype, IndexList(kv_sub_tile_rows((config // Int(2)), KVLUTType.page_size), Int(1), config, __list_literal__=NoneType(None)), config.swizzle_mode](), _ragged_shape[Int(3), KVLUTType.dtype, IndexList(kv_sub_tile_rows((config // Int(2)), KVLUTType.page_size), Int(1), config, __list_literal__=NoneType(None)), config.swizzle_mode]()], v_tma_op: TMATensorTile[KVLUTType.dtype, Int(3), _padded_shape[Int(3), KVLUTType.dtype, IndexList(kv_sub_tile_rows(config.BK1, KVLUTType.page_size), Int(1), config, __list_literal__=NoneType(None)), config.swizzle_mode](), _ragged_shape[Int(3), KVLUTType.dtype, IndexList(kv_sub_tile_rows(config.BK1, KVLUTType.page_size), Int(1), config, __list_literal__=NoneType(None)), config.swizzle_mode]()], ragged_tma_store: RaggedTMA3DTile[output_type, TensorMapSwizzle.SWIZZLE_NONE, BM=config.BM, BN=config.ov_depth, middle_dim=config.num_kv_heads if config.fuse_gqa else config.num_q_heads, group=config.group if config.fuse_gqa else Int(1), tma_blocks_per_op=o_store_tma_blocks_per_op[output_type, TensorMapSwizzle(Int32(0)), config.ov_depth, config.group if config.fuse_gqa else Int(1), Int(1)]()], kv_lut: KVLUTType, scale: Float32, batch_size: UInt32, num_keys_arg: UInt32, pack: Pack[MaskType, SchedulerType, ValidLengthType, NullPointer[DType.float32], KVRowOffsetsType, MaxSeqLenType, PartitionType])

mask_status

static def mask_status(mask: MaskType, seq_id: UInt32, score_row: UInt32, kv_row: UInt32) -> TileMaskStatus

Returns:

TileMaskStatus