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_mma

def depth512_mma[MaskType: MHAMask, qkv_dtype: DType, config: Depth512SM100Config[qkv_dtype], page_size: Int](smem: Depth512AttentionSMem[config], tmem_addr: UInt32, seq_id: UInt32, score_row: UInt32, num_keys: UInt32, mask: MaskType)

Runs the pair-CTA Q@K' and P@V MMAs for depth=256/512 SM100 attention.

Drives the leader CTA of a cta_group=2 pair through the full attention iteration loop: a peeled first iteration seeds S_even and O, then the main loop alternates S_even/S_odd while accumulating into O (or O_lo/O_hi when split_o). Peer CTAs return immediately.

Parameters:

MaskType (MHAMask): The MHAMask specialization governing tile masking behavior.
qkv_dtype (DType): Element dtype of the Q, K, V operands.
config (Depth512SM100Config[qkv_dtype]): Depth-256/512 SM100 attention configuration (tile sizes, pipeline depths, split_o strategy).
page_size (Int): Paged KV cache page size in keys.

Args:

smem (Depth512AttentionSMem[config]): Shared-memory allocator holding Q, K, V, P buffers and mbarriers.
tmem_addr (UInt32): Base TMEM address for this CTA pair's S/O accumulators.
seq_id (UInt32): Sequence index used by the mask.
score_row (UInt32): Starting query row index for this CTA pair.
num_keys (UInt32): Number of valid keys in the KV cache for this sequence.
mask (MaskType): Mask instance used to skip fully-masked tiles.