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 function

depth512_load

def depth512_load[KVLUTType: MHAOperand, MaskType: MHAMask, qkv_dtype: DType, config: Depth512SM100Config[qkv_dtype], ValidLengthType: OptionalPointer, _is_cache_length_accurate: Bool, MaxSeqLenType: OptionallyStaticInt, is_leader: Bool](smem: Depth512AttentionSMem[config], score_row: UInt32, num_keys: UInt32, seq_info: SeqInfo, max_seq_len: MaxSeqLenType, mask: MaskType, q_tma_op: TMATensorTile[KVLUTType.dtype, Int(4) if config.fuse_gqa else Int(3), _padded_shape[Int(4) if config.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=config.fuse_gqa, num_qk_stages=config.num_qk_stages](), config.swizzle_mode](), _ragged_shape[Int(4) if config.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=config.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]()], kv_lut: KVLUTType)

Issues TMA loads for one pair-CTA attention tile's Q, K, and V data.

Drives the producer side of the staged KV pipeline for a single attention tile, loading Q (per-CTA), K (per-CTA half of BN), and V (depth-split sub-tiles) into SMEM via non-multicast TMA copies. The peeled-first iteration co-arrives Q with the first K depth stage on a shared barrier so the MMA can start immediately; the main loop issues full-tile K/V loads; the peeled-last iteration handles partial pages when page_size < BN. Mask checks use PairBM so both CTAs make identical skip decisions and pipeline barriers stay synchronized.

Parameters:

KVLUTType (MHAOperand): Paged key-value lookup table type supplying the element dtype, page_size, and row-to-page mappings used by kv_lut.
MaskType (MHAMask): Attention mask type driving per-tile skip and load decisions via start_column, last_masked_set_end, and status queries.
qkv_dtype (DType): Element DType of the Q, K, and V tensors; must match KVLUTType.dtype and config.qkv_dtype.
config (Depth512SM100Config[qkv_dtype]): Depth-512 SM100 pair-CTA attention configuration providing tile sizes (BM, BN, BK0, BK1), stage counts, swizzle mode, GQA grouping, and split_o controls.
ValidLengthType (OptionalPointer): Optional pointer type for per-sequence valid lengths; is_null is False for ragged variable-length sequences.
_is_cache_length_accurate (Bool): Whether the reported KV cache length exactly matches the count of valid tokens.
MaxSeqLenType (OptionallyStaticInt): Type of the maximum sequence length, either a comptime static Int or a runtime value; selects the decoding vs. prefill path via _is_decoding.
is_leader (Bool): Whether this CTA is the leader (even-ranked) CTA in the pair-CTA cluster; the leader issues expect_bytes and selects the first half of K/V rows.

Args:

smem (Depth512AttentionSMem[config]): Shared-memory allocator and stage buffers for the kernel.
score_row (UInt32): Starting row of the Q tile within the current sequence.
num_keys (UInt32): Number of valid KV keys for the current sequence.
seq_info (SeqInfo): Per-sequence metadata (prompt index, head index, etc.).
max_seq_len (MaxSeqLenType): Maximum sequence length (static or dynamic).
mask (MaskType): Attention mask driving per-tile skip/load decisions.
q_tma_op (TMATensorTile[KVLUTType.dtype, Int(4) if config.fuse_gqa else Int(3), _padded_shape[Int(4) if config.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=config.fuse_gqa, num_qk_stages=config.num_qk_stages](), config.swizzle_mode](), _ragged_shape[Int(4) if config.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=config.fuse_gqa, num_qk_stages=config.num_qk_stages](), config.swizzle_mode]()]): TMA descriptor for Q tile loads.
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]()]): TMA descriptor for K tile loads (per-CTA half).
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]()]): TMA descriptor for V tile loads (per-sub-tile).
kv_lut (KVLUTType): Paged key-value lookup table supplying row-to-page mappings.