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_ws_level1_combine

def fa4_ws_level1_combine[m_pack: Int, rows: Int, ov_depth: Int, use_fma: Bool = True](local_row: UInt32, partition_g: UInt32, warp_group_idx: UInt32, own_max: Float32, own_sum: Float32, scale_log2e: Float32, o_tmem: UInt32, stage_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], maxsum_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], mut o_band: Array[Float32, (ov_depth // m_pack)]) -> Tuple[Float32, Float32]

Level 1 of the two-level WS combine: per-warpgroup m_pack-way merge of the packed-TMEM datapath quarters into an UNNORMALIZED partial.

A fork of fa4_ws_intracta_combine that stops before normalization. It reuses the raw-O staging, the single WG-local barrier, the cross-warp m = max_g m_g, and the scale/l_acc accumulation, but DROPS the recip and, instead of dividing by l and packing to o_smem, returns the per-warpgroup unnormalized partial (per this lane's query row r):

m_wg      = max_g m_g
l_wg      = Σ_g exp2((m_g - m_wg) * scale_log2e) * l_g
o_band[i] = Σ_g exp2((m_g - m_wg) * scale_log2e) * O_g[r, g*own_cols + i]

Warp g = partition_g owns depth band g (own_cols = ov_depth // m_pack columns [g*own_cols, (g+1)*own_cols)) and returns THAT band, unnormalized, in the o_band register array. (m_wg, l_wg) are per-row and identical across the WG's m_pack warps (every warp reduces the same maxsum_smem). Level 2 (fa4_ws_level2_reduce_scatter_write) merges the two warpgroups' (m_wg, l_wg, o_band) and normalizes. See §5c of docs/plans/sm100-fa4-ws-mma32-two-pipeline-splitk.md.

Same caller contract, comptime params, and empty-partition neutral element (m_g = -inf, l_g = 0, finite-0 O_g) as fa4_ws_intracta_combine.

Returns:

Tuple[Float32, Float32]: (m_wg, l_wg) for this lane's row (raw-unit max, unnormalized sum).