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 module

softmax_warp

Softmax warp group logic for FA4 (SM100 Flash Attention).

comptime values

f32x2

comptime f32x2 = SIMD[DType.float32, SIMDLength(2)]

Functions

• ​fa4_lse_combine_write: LSE-combine two TMEM_O fragments and TMA-store a depth-column slice.
• ​fa4_scale_write_output:
• ​fa4_softmax:
• ​fa4_splitk_combine_write: Split-K pass 2 (reduce-scatter): EVERY partition's WG0 owns the depth-column band [wg_j_offset, wg_j_offset+iters_per_wg) (output blocks). It reads every partition's O_cta FOR ITS BAND + every partition's (max,sum), forms the cluster LSE Gmax/Gsum and per-row weights w_p = sum_p*exp2((max_p-Gmax)*c)/Gsum (per-thread scalars -- this thread owns its row, NO warp.shuffle_idx), weight-sums O_final = Σ_p O_cta_p*w_p (the sum_p cancels the O_cta normalization), casts to bf16, and TMA-stores its band columns to gmem.
• ​fa4_splitk_reduce_scatter_write: Split-K pass 2 dispatch (reduce-scatter): pick THIS partition's depth-column band from the STATIC partition count P and the runtime partition_idx, then STAGE + publish + combine over it. P is config.splitk_partitions (each split-K kernel is compiled once per static P), so the band's wg_j_offset/iters_per_wg are already comptime; only partition_idx is lifted via a comptime-for over the P partitions -- one runtime branch taken, the band params comptime per branch.
• ​fa4_splitk_stage_partial: Split-K pass 1: LSE-combine the two WG O fragments into this partition's NORMALIZED O_cta (f32) over the FULL [0, iters) depth range, routing each block by destination. The OWN band [own_off, own_off+own_iters) is written into the o_final register array UNSCALED (the combine scales it by w[b] after the publish barrier); every other (PEER-visible) block is stored to stage_smem in the SWIZZLE_NONE block-major, bank-conflict-free layout fa4_scale_write_output uses, at f32 width.
• ​fa4_tma_store_o_smem: TMA egress of a fully-written o_smem to gmem: the per-block (tma_blocks_per_op == 0) tail factored out of fa4_scale_write_output's else branch.
• ​fa4_ws_intracta_combine: Intra-CTA m_pack-way LSE combine for the .ws MMA_M=32 datapath.
• ​fa4_ws_intracta_combine_partial: T==1 cross-CTA (split-K) partial: the 4-way intra-CTA combine emitting the NORMALIZED per-CTA band into o_band_norm (registers) + returning (M_cta, L_cta), for the cross-CTA reduce-scatter third level, instead of packing to o_smem.
• ​fa4_ws_level1_combine: Level 1 of the two-level WS combine: per-warpgroup m_pack-way merge of the packed-TMEM datapath quarters into an UNNORMALIZED partial.
• ​fa4_ws_level2_combine_partial: T>=2 cross-CTA (split-K) partial: a fork of fa4_ws_level2_reduce_scatter_write that emits the NORMALIZED per-CTA band into o_band_norm (registers) + returns (M_cta, L_cta) = (gmax, l_acc) for the cross-CTA reduce-scatter third level, instead of packing to o_smem.
• ​fa4_ws_level2_reduce_scatter_write: Level 2 of the two-level WS combine: intra-CTA cross-warpgroup reduce-scatter of the two per-WG Level-1 partials into the normalized output. WG0-writes-all (§5c step 3 / §9.8).
• ​fa4_ws_splitk_reduce_scatter_write: Third level (cross-CTA cluster split-K) reduce-scatter combine for the .ws MMA_M=32 datapath: merges the P partition CTAs' per-CTA normalized partials into the final output, each CTA writing only its OWN depth band.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit