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

fa4_splitk_combine_write

def fa4_splitk_combine_write[output_type: DType, band_cols: Int, //, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], P: Int, wg_j_offset: Int, iters_per_wg: Int, use_fma: Bool, output_swizzle_mode: TensorMapSwizzle = config.swizzle_mode, tma_bpo: Int = Int(0)](local_row: UInt32, local_warp_idx: UInt32, warp_group_idx: UInt32, own_rank: UInt32, own_max: Float32, own_sum: Float32, scale_log2e: Float32, stage_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], maxsum_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], o_smem_arg: Pointer[Scalar[output_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], publish_mbar: Pointer[SharedMemBarrier, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], ragged_tma_store: RaggedTMA3DTile[output_type, output_swizzle_mode, BM=(config // config), 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=tma_bpo], num_output_rows: Int32, out_head_idx: UInt32, out_row_idx: UInt32, mut o_final: Array[Float32, band_cols])

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.

o_final arrives PRE-SEEDED by fa4_splitk_stage_partial with THIS partition's own-rank b NORMALIZED O_cta band (UNSCALED). The own band is read by no peer (each peer reads only ITS band from us), so it is handed off in registers rather than round-tripped through stage_smem. This pass scales that seed in place by w[b], then accumulates the other np-1 ranks via DSMEM reads of their stage_smem bands.

Layout is the SWIZZLE_NONE row-per-thread pattern (mirrors fa4_scale_write_output): the f32 stage is 16 B (F=4) block-major, read in bulk v4 DSMEM loads; the bf16 write is 16 B (o_sw_K=8) block-major.

NO round-2 cluster barrier: the bf16 pack is retargeted into THIS partition's OWN-band f32 stage slice (which no peer reads -- each peer reads only ITS band from us), so the write clobbers nothing peers are reading. The peer-read PEER bands of stage_smem are kept alive until every partition finishes reading by the kernel's terminal cluster_sync() (kernel.mojo) -- the same trailing-sync discipline the DSMEM spike validated. Only round-1 (publish_mbar phase 0, in the caller) is needed, to make peers' staged O_cta + (max,sum) visible before the DSMEM reads. (An earlier reduce-scatter revision aliased the bf16 write over peer-read f32 bytes and so needed a phase-1 barrier; the own-slice retarget removes that hazard. Writer-combines-all needed only a WG-local barrier because peers never wrote.)