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

fa4_load

def fa4_load[KVLUTType: MHAOperand, MaxSeqLenType: OptionallyStaticInt, MaskType: MHAMask, //, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], *, ValidLengthType: OptionalPointer, _is_cache_length_accurate: Bool, is_leader: Bool](smem: SM100AttentionSMem[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 // config), 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 // config), 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.k_rows_per_cta(), 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.k_rows_per_cta(), 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(config.v_tma_box_rows(KVLUTType.page_size), Int(1), config.v_tma_box_cols(), __list_literal__=NoneType(None)), config.swizzle_mode](), _ragged_shape[Int(3), KVLUTType.dtype, IndexList(config.v_tma_box_rows(KVLUTType.page_size), Int(1), config.v_tma_box_cols(), __list_literal__=NoneType(None)), config.swizzle_mode]()], kv_lut: KVLUTType)

Issues the TMA loads for the Q, K, and V tiles consumed by one FA4 attention warp group on SM100.

Drives the producer side of the K/V pipeline across both fused-KV and split-KV modes, handling pair-CTA half-tile offsets, split-K windowing, partial-page bounds for sub-page KV tiles, and mask-driven iteration. Q is issued on the K barriers (fused with the first K stage) and on a separate Q1 barrier in two-Q mode. K and V TMA coordinates are produced from a shared kv_lut.populate so V reuses K's paged-row indices without a second LUT lookup.

Parameters:

KVLUTType (MHAOperand): Paged KV-cache lookup table type supplying the element dtype, page_size, and row-to-page mappings used by kv_lut (inferred).
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 (inferred).
MaskType (MHAMask): Attention mask type driving per-tile skip and load decisions via start_column, last_masked_set_end, and status queries (inferred).
config (FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_]): SM100 FA4 attention configuration providing tile sizes (BM, BN, BK0), stage counts, swizzle mode, GQA grouping, and split-K 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.
is_leader (Bool): Whether this CTA is the leader (even-ranked) CTA in a pair-CTA cluster, or always True in single-CTA mode; the leader issues expect_bytes and selects the first half of K/V rows.

Args:

smem (SM100AttentionSMem[config]): Shared-memory allocator for Q, K, V, and barrier storage.
score_row (UInt32): Row index of this tile within the score matrix.
num_keys (UInt32): Number of valid KV keys in the sequence.
seq_info (SeqInfo): Per-sequence metadata (prompt index, head index, etc.).
max_seq_len (MaxSeqLenType): Maximum sequence length, optionally a compile-time constant.
mask (MaskType): Attention mask governing which KV tiles are skipped.
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 // config), 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 // config), 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 the Q tile load.
k_tma_op (TMATensorTile[KVLUTType.dtype, Int(3), _padded_shape[Int(3), KVLUTType.dtype, IndexList(kv_sub_tile_rows(config.k_rows_per_cta(), 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.k_rows_per_cta(), KVLUTType.page_size), Int(1), config, __list_literal__=NoneType(None)), config.swizzle_mode]()]): TMA descriptor for the K tile load.
v_tma_op (TMATensorTile[KVLUTType.dtype, Int(3), _padded_shape[Int(3), KVLUTType.dtype, IndexList(config.v_tma_box_rows(KVLUTType.page_size), Int(1), config.v_tma_box_cols(), __list_literal__=NoneType(None)), config.swizzle_mode](), _ragged_shape[Int(3), KVLUTType.dtype, IndexList(config.v_tma_box_rows(KVLUTType.page_size), Int(1), config.v_tma_box_cols(), __list_literal__=NoneType(None)), config.swizzle_mode]()]): TMA descriptor for the V tile load.
kv_lut (KVLUTType): Paged KV-cache lookup table producing per-tile row indices.