For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
fa4_mma
def fa4_mma[MaskType: MHAMask, //, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], *, page_size: Int, workspace_split: Bool = False, crossp_effective: Bool = True](smem: SM100AttentionSMem[config], tmem_addr: UInt32, seq_id: UInt32, score_row: UInt32, num_keys: UInt32, mask: MaskType, ws_num_partitions: UInt32 = UInt32(1))
Executes the FA4 MMA warp loop for SM100 Flash Attention.
Computes Q@K' into score TMEM (S0/S1) and P@V into output TMEM (O0/O1)
across the KV tile sequence, coordinating with the load and softmax
warps via barrier pipelines. Supports fused-KV and split-KV modes,
single-Q and two-Q tile configurations, and partial-K contraction for
paged sub-tiles (page_size zero disables paging; sub-BN enables
partial-K).
Args:
- smem (
SM100AttentionSMem[config]): Shared memory allocator holding Q, K, V tiles and barriers. - tmem_addr (
UInt32): Base TMEM address for this CTA's S/O accumulators. - seq_id (
UInt32): Sequence index for the current query batch. - score_row (
UInt32): Row offset of the query tile within the sequence. - num_keys (
UInt32): Number of valid key columns to attend to. - mask (
MaskType): Attention mask controlling tile iteration and masking bounds. - ws_num_partitions (
UInt32): Runtime split-K partition count that windows the KV range whenworkspace_splitisTrue;1disables windowing.