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_reduce_scatter_write

def fa4_splitk_reduce_scatter_write[output_type: DType, //, config: FA4Config[config.qkv_dtype, rope_dtype_=config.rope_dtype_, scale_dtype_=config.scale_dtype_], P: Int, use_fma: Bool, single_source: Bool = False, zero_fill: Bool = False, output_swizzle_mode: TensorMapSwizzle = config.swizzle_mode, tma_bpo: Int = Int(0)](local_row: UInt32, local_warp_idx: UInt32, warp_group_idx: UInt32, partition_idx: UInt32, own_max: Float32, own_sum: Float32, scale_log2e: Float32, final_scale_local: Float32, final_scale_peer: 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], own_o_tmem: TMemTile[DType.float32, config.BM, config.padded_ov_depth], peer_o_tmem: TMemTile[DType.float32, config.BM, config.padded_ov_depth], 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)

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.

Per band (WG0), staging split so the publish lands between peers and own:

fa4_splitk_stage_partial[do_own=False] stages only the PEER-visible blocks to stage_smem (peers read those for THEIR bands).
Phase-0 publish: WG0 arrive_clusters every peer (peers read only our peer bands + maxsum, never our own band, so publish now).
fa4_splitk_stage_partial[do_peers=False] stages the OWN band into the register array o_final (UNSCALED) LAST -- its TMEM load overlaps the cluster sync, and o_final stays out of the peer-staging live range.
ALL entering WGs wait(0) (peers' staged O_cta/(max,sum) visible), then fa4_splitk_combine_write (WG0) scales the o_final seed by w[b], DSMEM-accumulates the other ranks' bands, and TMA-stores its band -- packing the bf16 into its OWN-band dead f32 slice (no peer reads it), so no round-2 cluster barrier is needed (the terminal cluster_sync() keeps the peer-read bands alive through peers' reads).

Called by BOTH warp groups (WG1 only participates in the phase-0 wait; WG0 does the staging/combine). EVERY partition -- empty or not -- joins round-1 (publish_mbar phase 0), so the publish_mbar count always completes. The own-no-band partition (non-pow2 depth at high P) still stages its FULL O_cta to smem (its complement is the full range) so peers can read it for their bands; it writes no bf16 output, so it just falls through to the kernel's terminal cluster_sync() after round-1.