K2 in CuTe [FA3]

This is the in-depth companion to the K2 section of the FA3 worklog. Refer back there for the FA2 algorithm, the three-phase structure, the SMEM tile layout, and the in-place rewriting tricks. This post covers the CuTe mechanics: how the kernel is written in CuTe Python DSL, the per-phase code with variable tables and unwrapping, and the new primitives K2 introduces.

The kernel file is kernels/k2.py in the project repo. One CuTe class K2FMHA plus a compile_kernel / run_kernel dispatcher pair.

Block sizes and the threading model

Working constants:

Br = 64
Bc = 64
THREADS = 256
THREADS_PER_ROW = 4

Why 4 threads per row. 256 threads, 64 output rows; each row gets \(256 / 64 = 4\) threads. Those 4 threads form a “row group” and are jointly responsible for one row of \(O\) across all KV iterations. Inside that row group, the work splits three ways:

1. Storing one row of \(O\) in registers. The row is \(d\) wide. The 4 threads split the \(d\) axis evenly: each thread holds \(d / 4\) FP32 values of \(O\) in registers. For \(d = 128\) that is 32 values per thread; for \(d = 64\), 16.

2. Reducing one row of \(S\). For the row max and row sum, the row is \(B_c = 64\) wide. The 4 threads split the \(B_c\) axis: each thread scans \(B_c / 4 = 16\) elements of the row, producing a local partial. The four partials are then combined with a sub-warp shuffle (covered later in this post).

3. Computing one row of \(P V_j\). Same \(d\)-axis sharding as \(O\): each thread accumulates \(d / 4\) output columns by walking the \(B_c\) axis.

The output accumulator \(O_\text{acc}\) is the only state sharded across the 4 threads (along \(d\)). The scalar softmax state \((m_i, \ell_i)\) lives in registers replicated across all 4 threads: they all do the same reductions, so they all see the same values. This replication is what avoids any SMEM round-trip for \((m_i, \ell_i)\).

The thread-to-(row, slice) mapping.

row         = tidx // THREADS_PER_ROW   # 0..63: which Q-row of the tile
col_group   = tidx %  THREADS_PER_ROW   # 0..3:  which 1/4 of that row this thread owns
d_col_start = col_group * (d // THREADS_PER_ROW)

Concretely for \(d = 128\): thread 0 owns row 0, \(d\)-columns 0..31; thread 1 owns row 0, columns 32..63; thread 2 owns row 0, columns 64..95; thread 3 owns row 0, columns 96..127. Thread 4 starts row 1, and so on through thread 255 on row 63.

SMEM layout

sQ_layout  = cute.make_layout((Br, d),  stride=(d, 1))    # BF16
sKV_layout = cute.make_layout((Bc, d),  stride=(d, 1))    # BF16, reused K then V
sS_layout  = cute.make_layout((Br, Bc), stride=(Bc, 1))   # FP32

Size derivation for \(B_r = B_c = 64\), \(d = 128\). Each layout is row-major; total bytes is \(\text{rows} \times \text{cols} \times \text{bytes/element}\).

Total \(48\) KB per CTA. For \(d = 64\) it shrinks to \(8 + 8 + 16 = 32\) KB.

Layouts are plain row-major, no swizzling. K3 and K4 introduce swizzled layouts to fix bank conflicts.

Mapping FA2 to code, phase by phase

The kernel body has three phases: load Q once, KV mainloop, finalize. The mainloop has five sub-phases.

Phase 1: load Q once, init state

m_i = -inf
l_i = 0.0
for c in range_constexpr(cols_per_thread_O):
    O_acc[c] = 0.0

for it in range_constexpr(0, Br * d, THREADS):
    ij = it + tidx
    sQ[ij // d, ij % d] = gQ[q_row_start + ij // d, ij % d]
cute.arch.sync_threads()

Variables.

Unwrapping for \(d = 128\). \(B_r \cdot d = 64 \cdot 128 = 8192\) elements. With 256 threads, each thread handles \(8192 / 256 = 32\) elements, so the outer loop runs 32 iterations (it = 0, 256, 512, ..., 7936).

Look at iteration it = 0. Threads 0..127 have ij = 0..127, so ij // d = 0 and ij % d = 0..127: together they load all of row 0 of the Q-tile. Threads 128..255 have ij = 128..255, so ij // d = 1: they load all of row 1. One “wave” of 256 threads loads 2 rows of Q (because \(256 / d = 2\) when \(d = 128\)).

Iteration it = 256 lands on rows 2 and 3, it = 512 on rows 4 and 5, and so on. After 32 iterations, all 64 rows are loaded. Within each wave, adjacent threads touch adjacent columns of the same row in GMEM, so the loads coalesce.

(For \(d = 64\): \(B_r \cdot d = 4096\), 16 elements per thread, one wave loads \(256 / 64 = 4\) rows, 16 iterations total. Same shape, different numbers.)

The same arithmetic governs Phase 2a (load \(K_j\)) and Phase 2d (load \(V_j\)), with \(B_r \to B_c\).

Phase 2a: load \(K_j\)

for it in range_constexpr(0, Bc * d, THREADS):
    ij = it + tidx
    sKV[ij // d, ij % d] = gK[kv_row_start + ij // d, ij % d]
cute.arch.sync_threads()

Identical shape to the Q load, with \(B_r \to B_c\). kv_row_start = j * Bc is the global KV-row index of this iteration’s tile. For \(B_c = 64\), \(d = 128\): 32 iterations, two rows of \(K_j\) per wave.

Phase 2b: \(S = Q K^\top \cdot \text{scale}\)

for it in range_constexpr(0, Br * Bc, THREADS):
    ij = it + tidx
    si, sj = ij // Bc, ij % Bc
    dot = 0.0
    for k in range_constexpr(d):
        dot += sQ[si, k].to(f32) * sKV[sj, k].to(f32)
    val = dot * scale
    if const_expr(causal):
        if kv_row_start + sj > q_row_start + si:
            val = -inf
    sS[si, sj] = val
cute.arch.sync_threads()

Variables. Each thread now computes one element of \(S\) per inner iteration, not one element of a load.

Note we read \(K_j\) as sKV[sj, k], not sKV[k, sj]: since \(K\) is laid out row-major and \(S = Q K^\top\), the matmul reduces to “row of Q dot row of K”, no transpose in memory.

Unwrapping. \(B_r \cdot B_c = 64 \cdot 64 = 4096\) elements of \(S\), 256 threads, \(4096 / 256 = 16\) elements per thread, 16 outer iterations.

Iteration it = 0: threads 0..63 have ij = 0..63, so si = 0, sj = 0..63. Those 64 threads compute all of row 0 of \(S\). Threads 64..127 produce row 1, 128..191 row 2, 192..255 row 3. One wave fills 4 rows of \(S\); 16 iterations fill all 64 rows.

Causal mask. kv_row_start + sj is the absolute KV position in the sequence; q_row_start + si is the absolute Q position. Any future KV (\(> q\_row\_start + si\)) is set to \(-\infty\) before softmax sees it.

Cost. Each element of \(S\) takes one \(d\)-length FP32 dot product. Per thread: \(16 \cdot d\) FMAs, all scalar, all on the FP32 pipe. K4 replaces this loop nest with one WGMMA call.

Phase 2c: online softmax row update

local_max = -inf
for c_it in range_constexpr(cols_per_thread_S):     # 16 iters when Bc=64
    c = col_group + c_it * THREADS_PER_ROW
    local_max = fmax(local_max, sS[row, c])
row_max = cute.arch.warp_reduction_max(local_max,
                                       threads_in_group=THREADS_PER_ROW)

m_new = fmax(m_i, row_max)
alpha = exp(m_i - m_new)

for c in range_constexpr(cols_per_thread_O):
    O_acc[c] *= alpha

local_sum = 0.0
for c_it in range_constexpr(cols_per_thread_S):
    c = col_group + c_it * THREADS_PER_ROW
    p = exp(sS[row, c] - m_new)
    sS[row, c] = p
    local_sum += p
row_sum = cute.arch.warp_reduction_sum(local_sum,
                                       threads_in_group=THREADS_PER_ROW)

l_i = alpha * l_i + row_sum
m_i = m_new
cute.arch.sync_threads()

Variables and what they map to in FA2 notation.

Row-scan striding. For the row max, each thread scans \(B_c / \text{THREADS\_PER\_ROW} = 16\) columns of sS[row, :]. The pattern c = col_group + c_it * THREADS_PER_ROW gives col_group 0 → columns 0, 4, 8, …, 60; col_group 1 → 1, 5, …, 61; col_group 2 → 2, 6, …, 62; col_group 3 → 3, 7, …, 63. Each of the 4 threads in the row group covers a different residue class mod 4. Together they cover all 64 columns exactly once. The four partials are then combined with warp_reduction_max, after which every thread of the row group holds the full \(\max\) over the row.

In-place exp. After Phase 2c, sS[row, :] no longer holds \(S\); it holds \(P = \exp(S - m_\text{new})\). Elements that were masked to \(-\infty\) exponentiate to \(0\), so they contribute nothing to either row_sum or the subsequent \(PV\) matmul. No branch needed.

Phase 2d: load \(V_j\) (overwrites \(K_j\))

Identical to Phase 2a with \(K \to V\). Same iteration count, same per-thread element count, same coalescing.

Phase 2e: \(O_\text{acc} \mathrel{+}= P V_j\)

for c in range_constexpr(cols_per_thread_O):
    col = d_col_start + c
    pv = 0.0
    for k in range_constexpr(Bc):
        pv += sS[row, k] * sKV[k, col].to(f32)
    O_acc[c] += pv

Variables.

Unwrapping for \(d = 128\). Each thread updates \(d / 4 = 32\) output columns, each by accumulating a \(B_c = 64\)-length dot product. That is \(32 \cdot 64 = 2048\) scalar FMAs per thread per KV iteration. For \(d = 64\) it is \(16 \cdot 64 = 1024\).

Redundant \(P\) reads. The four threads of a row group all read the same sS[row, :], each touching different cols of \(V_j\). That redundancy wastes SMEM bandwidth, but in K2 we are not trying to be efficient; this loop nest, like Phase 2b, is replaced by WGMMA in K4.

Phase 3: finalize

inv_l = cute.arch.rcp_approx(l_i)
if l_i == 0.0 or l_i != l_i:               # fully-masked row
    inv_l = 1.0
for c in range_constexpr(cols_per_thread_O):
    gO[q_row_start + row, d_col_start + c] = (O_acc[c] * inv_l).to(bf16)

rcp_approx is the SFU reciprocal: one instruction. The guard handles rows fully masked by the causal mask (the first row of a Q-tile that begins below the diagonal of a KV-tile sees no valid keys; \(\ell_i = 0\)). Without it, \(0 \cdot \infty\) would propagate NaN.

New CuTe primitives in K2

A handful of primitives that did not appear in K1.

cute.struct and SmemAllocator. A cute.struct declares a fixed-size, aligned SMEM block. Align[..., 1024] enforces 1024-byte alignment; MemRange[dtype, n] reserves \(n\) elements. The struct is materialised through SmemAllocator, and named regions are extracted with get_tensor(layout). This is the standard CuTe pattern for several named SMEM tiles sharing one allocation:

@cute.struct
class SharedStorage:
    sQ:  cute.struct.Align[cute.struct.MemRange[bf16,  cosize(sQ_layout)],  1024]
    sKV: cute.struct.Align[cute.struct.MemRange[bf16,  cosize(sKV_layout)], 1024]
    sS:  cute.struct.Align[cute.struct.MemRange[f32,   cosize(sS_layout)],  1024]

cute.arch.warp_reduction_max / _sum with threads_in_group. Sub-warp shuffle reduction. Each lane brings in a partial scalar; after the call every participating lane holds the same reduced result.

How does it work mechanically? With threads_in_group=4, lanes within each warp are grouped as \(\{0,1,2,3\}, \{4,5,6,7\}, \ldots, \{28,29,30,31\}\). Within each group of 4, the DSL emits a butterfly of two shuffles:

step 1: x = max(x, __shfl_xor_sync(mask, x, 1))   # exchange with lane^1
step 2: x = max(x, __shfl_xor_sync(mask, x, 2))   # exchange with lane^2

After step 1, each pair \((0,1), (2,3), \ldots\) holds the pair max. After step 2, each quad holds the quad max, broadcast to all 4 lanes of the group. Two PTX instructions, no SMEM, no __syncthreads(). warp_reduction_sum is the same butterfly with max replaced by +.

Why this maps cleanly to row groups. Our row-to-thread layout is row = tidx // 4, col_group = tidx % 4. So threads 0..3 are row 0, 4..7 are row 1, …, 28..31 are row 7: 8 row groups per warp, 8 warps per CTA, 64 row groups in total, one per Q-row. That partition lines up exactly with the lane grouping the shuffle uses inside a warp. Each row group’s reduction happens entirely within its own 4 lanes; reductions for different rows are independent and run in parallel.

Why no extra synchronisation. Shuffles execute lockstep across a warp at the SIMT level. Once warp_reduction_max returns, all 4 lanes of every row group hold the same value, and the next instruction (the \(O\) rescale by alpha) can use it directly. No SMEM round-trip, no barrier. This is the structural reason the per-row softmax state \((m_i, \ell_i)\) can live in registers replicated across the row group.

Choice of group size. threads_in_group must divide a warp, so 2, 4, 8, 16, or 32 are the options. A larger group means more parallelism in the reduction (fewer columns to scan per thread, more shuffle steps) but a coarser \(d\)-split for \(O\), so fewer registers per thread to hold the output row. With 4: 16-element partial scan, 2-step shuffle, \(d / 4\) FP32 registers per thread for \(O\) (32 regs for \(d = 128\), comfortably within the per-thread budget at full occupancy on H100).

cutlass.range vs cutlass.range_constexpr. range_constexpr is unrolled at JIT time and needs JIT-time-constant bounds; we use it everywhere the trip count is known. cutlass.range(n_kv, unroll=1) is a runtime loop, used only for the KV mainloop because \(n_\text{kv} = N / B_c\) depends on the runtime sequence length.

cute.math.exp(..., fastmath=True). Lowers to the SFU ex2.approx.f32 (about 2 ULPs). The flag is what makes the softmax fast; without it we would get a much slower software exp.

cute.arch.rcp_approx. SFU reciprocal (about 1 ULP), used in the finalize divide.

The rest (make_rmem_tensor, sync_threads, block_idx, thread_idx, .to(dtype) casts) are direct analogues of CUDA C++ and don’t need separate treatment.

For the next kernel in the worklog, see K3: TMA.