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_level2_combine_partial

def fa4_ws_level2_combine_partial[m_pack: Int, rows: Int, ov_depth: Int, use_fma: Bool = True](local_row: UInt32, band_g: UInt32, warp_group_idx: UInt32, own_max: Float32, own_sum: Float32, scale_log2e: Float32, o_band: Array[Float32, (ov_depth // m_pack)], l2_stage_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], l2_maxsum_smem: Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False], mut o_band_norm: Array[Float32, (ov_depth // m_pack)]) -> Tuple[Float32, Float32]

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.

Same WG0-writes-all transport as the base (WG1 scatters its o_band + (m_1, l_1) into l2_stage_smem/l2_maxsum_smem, single named_barrier[2*WARPGROUP_SIZE](3), WG0 gathers + 2-way combines + normalizes), only the sink differs (registers, not o_smem). WG1 returns a placeholder (own_max, own_sum) -- it does not feed the third level (WG0 owns every band's per-CTA partial). See §5c / Phase D of the living doc.

Returns:

Tuple[Float32, Float32]