IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, ImplicitlyCopyable, ImplicitlyDeletable, 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

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact