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 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.