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

def fa4_ws_intracta_combine[output_type: DType, //, 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], o_smem: Pointer[Scalar[output_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False])

Intra-CTA m_pack-way LSE combine for the .ws MMA_M=32 datapath.

The m_pack packed-TMEM datapath quarters are m_pack independent key-partitions over the SAME rows query rows: warp g = partition_g holds partition g's unnormalized O_g (rows x ov_depth) in its C-TMEM band (all warps share the base o_tmem; the hardware subpartition routes each warp its own quarter, so there is no per-warp column offset), plus its per-row (m_g, l_g) = (own_max, own_sum). Because the partitions live in DIFFERENT warps, the merge is a cross-warp reduction and must transit SMEM (a warp shuffle cannot reach across warps).

Merges, per query row r:

m       = max_g m_g
scale_g = exp2((m_g - m) * scale_log2e)   # *scale_log2e gated on use_fma
O       = sum_g scale_g * O_g
l       = sum_g scale_g * l_g
out     = O * recip(l)

Identical LSE math to fa4_lse_combine_write / fa4_splitk_combine_write; only the transport is intra-CTA across warp quarters. Assumes >= 1 non-empty partition (a launched tile always does >= 1 partition of real work). An individual empty partition supplies m_g = -inf, l_g = 0, and a FINITE zero-filled O_g (so scale_g = 0 and 0 * 0 = 0, never 0 * inf = NaN); there is no all-empty guard, mirroring fa4_splitk_combine_write's

=1-valid-key assumption.

The caller must have made O_g visible in TMEM (wait the O-producer barrier and issue tcgen05_fence_after()) before this call, exactly as the fa4_lse_combine_write caller does. Output is written to o_smem in the SWIZZLE_NONE 16 B block-major layout ([block, rows, o_sw_K], o_sw_K = 16 // size_of[output_type]()) that a SWIZZLE_NONE O-TMA store consumes; the caller performs the egress (the O TMA store in the kernel).

Parameters:

output_type (DType): The o_smem element dtype (bf16 in-kernel; f32 for tests).
m_pack (Int): Number of key-partitions == datapath quarters (4 for MMA_M=32).
rows (Int): Query rows per tile (config.BM, 32 for MMA_M=32).
ov_depth (Int): Output value depth.
use_fma (Bool): Gate the * scale_log2e step (matches fa4_splitk_combine_write).

Args:

local_row (UInt32): This lane's query row (tid & 31).
partition_g (UInt32): This warp's partition / datapath quarter (tid >> 5).
warp_group_idx (UInt32): WG-scoped named_barrier id.
own_max (Float32): This partition's per-row local max m_g[local_row] (raw units).
own_sum (Float32): This partition's per-row local sum l_g[local_row].
scale_log2e (Float32): The softmax log2(e) scale.
o_tmem (UInt32): Base C-TMEM column address of the shared O band.
stage_smem (Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]): SMEM scratch, >= m_pack*rows*ov_depth f32 (raw O staging, SWIZZLE_NONE 16 B block-major [m_pack, stage_blocks, rows, 4]).
maxsum_smem (Pointer[Float32, MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]): SMEM scratch, >= m_pack*rows*2 f32 ((m, l) staging).
o_smem (Pointer[Scalar[output_type], MutAnyOrigin, address_space=AddressSpace.SHARED, _safe=False]): Output sink, SWIZZLE_NONE 16 B block-major [block, rows, o_sw_K] of output_type.