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

fa4_splitk_combine

FA4 SM100 traditional (workspace / unfused) split-K combine kernel.

Merges the P per-partition partial outputs the FA4 1Q attention kernel wrote to a global workspace into the final attention output, using fused Log-Sum-Exp (LSE) weights for numerical stability.

Workspace layout (produced by the attention kernel's do_partition egress):

• o_partial : [P, num_rows_q, num_q_heads, ov_depth] (intermediate dtype), each partition holds its LOCALLY-normalized O_p / l_p.
• lse_partial : [P, num_rows_q, num_q_heads] (f32), fused per-row LSE in log2 domain (lse_p = log2(l_p) + m_p).

Output layout (matches the non-split ragged store): [num_rows_q, num_q_heads, ov_depth].

Per output row-head (token, q_head) (flattened to rh = token*num_q_heads + q_head), reducing over partitions p:

• m* = max_p lse_p
• global_lse = log2(Sum_p exp2(lse_p - m*)) + m*
• scale_p = exp2(lse_p - global_lse) (folds the 1 / denom normalization in: exp2(lse - log2(denom) - m) == exp2(lse - m) / denom)
• O[d] = Sum_p scale_p * o_partial[p, rh, d]

Everything stays in the log2 (base-2) domain to match the FA4 softmax (exp2/log2); crossing into natural-exp would silently corrupt the result.

Perf structure (adapted from mla_decode_combine.mojo's

mla_combine_kernel / mla_combine_kernel_split_parallel)

One warp (WARP_SIZE threads) per output row-head; ov_depth is split across lanes via a vectorized vec_size / elems_per_thread partition (128-bit loads where ov_depth and the intermediate dtype allow it). The LSE reduction (row max, log-sum-exp, per-partition scale_p) is computed ONCE into per-lane registers -- not the 3x redundant global-memory re-read (once for the max, once for the denominator, once per depth-element inside the accumulation loop) the original correctness-first version did. The weighted accumulation loop broadcasts each partition's scale from its owning lane via warp.shuffle_idx and prefetches the NEXT partition's o_partial vector into registers while the CURRENT partition's contribution is folded into the running sum (software pipelining), so the load latency of partition p+1 overlaps the FMA work of partition p.

Two kernel variants share this body, selected by the comptime P_STATIC parameter:

• P_STATIC > 0: num_partitions landed on a rung of the shared splitk_p_ladder that dispatch.mojo's _bucket_ws snaps production P onto. Both the LSE and accumulation loops are comptime for-unrolled over the exact P_STATIC, which also keeps the per-lane local_lse array register-indexed (divmod(p, WARP_SIZE) folds at comptime).
• P_STATIC == 0: an off-rung P. Reachable two ways -- a test force-knob, and _bucket_ws capping its bucketed value at ws_p_ceiling, which is not itself a rung (e.g. B200's sm_count // 4 == 37). The same vectorized
  • prefetched accumulation shape runs over a RUNTIME num_partitions bound, paying a dynamically-indexed local_lse. The array is sized at a fixed comptime ceiling (_P_MAX, covering every current-generation sm_count) so this fallback compiles exactly once, independent of the P it is handed.

Target hardware family: NVIDIA SM100 (B200 / B300).

PDL (Item 4): this kernel is a Programmatic Dependent Launch consumer of the attention producer. It calls wait_on_dependent_grids() (all threads of its one-warp block, before the rh >= RH guard) and launches with pdl_launch_attributes(MHA_PDL_LEVEL). The attention producer does NOT emit a terminal launch_dependent_grids() for the do_partition config (it suppresses its prologue trigger instead, see kernel.mojo); this wait therefore releases on the producer's grid completion, which orders after the workspace egress store. That completion-based release (rather than a post-store terminal trigger) is deliberate: a legal trigger must be issued by ALL threads of the CTA (grid_controls.mojo:105), but the attention kernel's warp-specialized invalid-tile early returns leave no divergence-free post-store convergence point, so a terminal trigger would require restructuring the verified control flow. The PROGRAMMATIC_ STREAM_SERIALIZATION attribute still co-schedules this combine so its launch + prologue overlap attention's tail (measured 1.06-1.13x on P=148 combine-heavy shapes). With -D MHA_PDL=false all of this is comptime-pruned back to plain stream ordering. (An earlier attempt -- a post-store terminal trigger at prologue AND terminal -- double-issued launch_dependent_grids() and tripped CUDA_ERROR_ILLEGAL_INSTRUCTION; the current design issues no producer trigger for this config at all.)

Functions

• ​fa4_splitk_combine: Launches the workspace split-K combine over num_rows_q * num_q_heads output row-heads, one warp (WARP_SIZE threads) each.