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_softmax

def depth512_softmax[MaskType: MHAMask, qkv_dtype: DType, output_type: 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, scale: Float32, ragged_tma_store: RaggedTMA3DTile[output_type, TensorMapSwizzle.SWIZZLE_NONE, BM=config.BM, BN=config.ov_depth, middle_dim=ragged_tma_store.middle_dim, group=config.group if config.fuse_gqa else Int(1), tma_blocks_per_op=ragged_tma_store.tma_blocks_per_op], num_output_rows: Int32, out_head_idx: UInt32, out_row_idx: UInt32)

Runs the online softmax warp group for pair-CTA SM100 attention.

Loads Q@K' scores (S) from TMEM in pipelined batches, applies the causal mask, computes a running row max and exponentiated probabilities P, writes P to SMEM for the P@V SS MMA, accumulates the row sum, and finally scales the O accumulator by the inverse row sum and TMA-stores the result to global memory. For split_o (d512) configs, cross-thread partial max and sum values are combined via correction SMEM; for d256 each thread owns a unique M row and no exchange is needed.

Parameters:

MaskType (MHAMask): Compile-time mask type for causal/attention masking.
qkv_dtype (DType): DType of the Q/K/V inputs; specializes the config.
output_type (DType): DType of the output store to global memory.
config (Depth512SM100Config[qkv_dtype]): Depth512 SM100 kernel config struct (tile sizes, split_o).
page_size (Int): KV cache page size in tokens, used for paged-attention masking.

Args:

smem (Depth512AttentionSMem[config]): Shared-memory allocator holding S/P/O buffers and barriers.
tmem_addr (UInt32): Base TMEM address for S and O tiles (read once post cluster_sync).
seq_id (UInt32): Sequence index for mask evaluation.
score_row (UInt32): Row offset of the query tile within the sequence.
num_keys (UInt32): Number of valid key columns for masking.
mask (MaskType): Causal/attention mask applied to score batches.
scale (Float32): Softmax scale factor (applied in the log2 domain).
ragged_tma_store (RaggedTMA3DTile[output_type, TensorMapSwizzle.SWIZZLE_NONE, BM=config.BM, BN=config.ov_depth, middle_dim=ragged_tma_store.middle_dim, group=config.group if config.fuse_gqa else Int(1), tma_blocks_per_op=ragged_tma_store.tma_blocks_per_op]): Ragged TMA descriptor for the output store.
num_output_rows (Int32): Dynamic output row count for the TMA store.
out_head_idx (UInt32): Output head index for the TMA store.
out_row_idx (UInt32): Output row index for the TMA store.