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_intracta_combine_partial

def fa4_ws_intracta_combine_partial[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_norm: Array[Float32, (ov_depth // m_pack)]) -> Tuple[Float32, Float32]

T==1 cross-CTA (split-K) partial: the 4-way intra-CTA combine emitting the NORMALIZED per-CTA band into o_band_norm (registers) + returning (M_cta, L_cta), for the cross-CTA reduce-scatter third level, instead of packing to o_smem.

Implemented as fa4_ws_level1_combine (the unnormalized m_pack-way merge, which writes warp g's band [g*own_cols, +own_cols) and returns (m_wg, l_wg)) followed by an explicit recip(l_wg) normalize -- an identical result to fa4_ws_intracta_combine but emitted to registers. For T==1 only WG0 is active, so its (m_wg, l_wg) is the full-CTA (M_cta, L_cta) and its per-warp band is the normalized O_cta band the third level scatters.

Returns:

Tuple[Float32, Float32]