|
| 1 | +""" |
| 2 | +Latch-Q / Stream-K flash-attention inner loop — DESIGN SKETCH (WIP, UNTESTED). |
| 3 | +
|
| 4 | +Goal: kill the register spilling that dominates the current kernel. The post-RA |
| 5 | +LLO dump for bq=5888/cin=256 showed ~43k spill ops (85% of all VMEM traffic, |
| 6 | +9,465 spill slots) because the bq=5888-wide score canvas is ~23x the 64-VREG |
| 7 | +file and `bkv_compute_in` only tiles the kv axis — the bq lane-width is never |
| 8 | +tiled, so the allocator spills everything. |
| 9 | +
|
| 10 | +This sketch tiles the *bq* axis into 256-wide strips processed with the low-level |
| 11 | +MXU primitives, keeping the live working set at one [256,256] tile (~64 VREGs) and |
| 12 | +evicting each strip MRB->VPU before the next QK lands. That is the missing axis. |
| 13 | +
|
| 14 | +Status: this is a STRUCTURAL skeleton to reconcile with the working Colab |
| 15 | +baseline (Kunjan's "Latch Q & Stream K"). The index algebra and primitive order |
| 16 | +are worked out; the exact accumulator/staging-register *addresses* and the f32 |
| 17 | +accumulator byte offsets are marked TODO and must be validated on device — these |
| 18 | +primitives have hard rules (no data left in staging/acc on exit) and there is no |
| 19 | +bounds checking. Do NOT wire this into production before it round-trips a |
| 20 | +numerical diff vs _flash_attention_kernel. |
| 21 | +
|
| 22 | +Key constraints (from jax 0.10.0 pltpu primitive docstrings): |
| 23 | + matmul_push_rhs(rhs[256,256], staging_register, mxu_index, transpose=False) |
| 24 | + matmul_acc_lhs(acc_addr, lhs[M,256], mxu_index, load_staged_rhs=None) |
| 25 | + matmul_pop(acc_addr, shape, dtype, mxu_index) -> [M,256] f32, zeroes acc |
| 26 | + - out[m,n] = sum_k lhs[m,k] * staged_rhs[k,n] (contraction over the 256 dim) |
| 27 | + - load_staged_rhs=None REUSES the loaded RHS (this is the "don't re-latch" win) |
| 28 | + - RHS 256x256; LHS M x 256; accumulator f32/i32; nothing left resident on exit |
| 29 | +
|
| 30 | +Layout (kept identical to _flash_attention_kernel so the wrapper/grid is reused): |
| 31 | + TRANSPOSED. scores live as [kv, q]; output o as [head_dim, q]. |
| 32 | + QK (mxu0): lhs=K[kv,hd], rhs=Q_strip pushed transpose=True -> [hd,q] |
| 33 | + => out[kv,q] = sum_hd K[kv,hd]*Q[q,hd] (== current qk) |
| 34 | + PV (mxu1): lhs=V[kv,hd] used as [hd,kv], rhs=p[kv,q] |
| 35 | + => out[hd,q] = sum_kv V[kv,hd]*p[kv,q] (== current o) |
| 36 | + Both contractions are over head_dim => REQUIRES head_dim == 256 to fill the MXU. |
| 37 | +""" |
| 38 | + |
| 39 | + |
| 40 | +import jax.numpy as jnp |
| 41 | +from jax import lax |
| 42 | +from jax.experimental import pallas as pl |
| 43 | +from jax.experimental.pallas import tpu as pltpu |
| 44 | + |
| 45 | +TILE = 256 # MXU-native systolic dimension |
| 46 | +MXU_QK = 0 # dual-MXU: QK on MXU0 ... |
| 47 | +MXU_PV = 1 # ... PV on MXU1, so QK(next) overlaps PV(curr) |
| 48 | +STAGE_Q = 0 # staging register holding the latched Q strip |
| 49 | +STAGE_P = 0 # staging register on MXU1 for the per-tile p |
| 50 | + |
| 51 | +# Accumulator base addresses (f32 slices). TODO(device): confirm these don't |
| 52 | +# overlap and fit the acc file; one [256,256] f32 tile = 256 sublane-rows. |
| 53 | +ACC_QK = 0 |
| 54 | +ACC_PV = 0 |
| 55 | + |
| 56 | + |
| 57 | +def _strip_flash(q_strip, k_ref, v_ref, kv_seq_len, mask_value, exp): |
| 58 | + """One 256-wide bq strip: stream all KV with latched Q, online softmax. |
| 59 | +
|
| 60 | + q_strip: [TILE, head_dim(=256)] — the 256 query rows for this strip. |
| 61 | + k_ref/v_ref: VMEM refs, [kv_seq_len, head_dim]. |
| 62 | + Returns o_strip [head_dim, TILE] (un-normalized) and l [1, TILE] for the caller |
| 63 | + to divide by — matching the current kernel's deferred 1/l epilogue. |
| 64 | + """ |
| 65 | + head_dim = q_strip.shape[1] |
| 66 | + assert head_dim == TILE, "Latch-Q-Stream-K needs head_dim==256 to fill the MXU" |
| 67 | + |
| 68 | + # --- Latch Q ONCE for the whole KV stream (the core idea) ------------------- |
| 69 | + # transpose=True so the staged RHS is Q^T [hd, q]; reused across every kv tile. |
| 70 | + pltpu.matmul_push_rhs(q_strip, STAGE_Q, MXU_QK, transpose=True) |
| 71 | + |
| 72 | + # Running online-softmax state for these 256 queries. Kept tiny + resident: |
| 73 | + # m,l : [1, TILE] o : [head_dim, TILE] |
| 74 | + m = jnp.full((1, TILE), mask_value, jnp.float32) |
| 75 | + l = jnp.zeros((1, TILE), jnp.float32) |
| 76 | + o = jnp.zeros((head_dim, TILE), jnp.float32) |
| 77 | + |
| 78 | + num_kv_tiles = kv_seq_len // TILE # TODO: handle ragged tail like last_compute_body |
| 79 | + first = True |
| 80 | + |
| 81 | + def body(t, carry): |
| 82 | + nonlocal first |
| 83 | + m, l, o = carry |
| 84 | + kv0 = t * TILE |
| 85 | + k_tile = k_ref[pl.ds(kv0, TILE), :] # [kv=256, hd=256] -> lhs for QK |
| 86 | + v_tile = v_ref[pl.ds(kv0, TILE), :] # [kv=256, hd=256] -> lhs for PV |
| 87 | + |
| 88 | + # --- QK on MXU0: load staged Q only on the FIRST tile, then reuse ---------- |
| 89 | + # load_staged_rhs=STAGE_Q on the first acc; None afterwards = no re-latch. |
| 90 | + pltpu.matmul_acc_lhs(ACC_QK, k_tile, MXU_QK, load_staged_rhs=STAGE_Q if first else None) |
| 91 | + qk = pltpu.matmul_pop(ACC_QK, (TILE, TILE), jnp.float32, MXU_QK) # [kv, q] |
| 92 | + first = False |
| 93 | + |
| 94 | + # --- online softmax on the [kv, q] tile (reduce over kv = axis 0) ---------- |
| 95 | + m_curr = qk.max(axis=0, keepdims=True) # [1, q] |
| 96 | + m_next = jnp.maximum(m, m_curr) |
| 97 | + p = exp(qk - m_next) # [kv, q] (EPU) |
| 98 | + alpha = exp(m - m_next) # [1, q] |
| 99 | + l = alpha * l + p.sum(axis=0, keepdims=True) |
| 100 | + |
| 101 | + # --- PV on MXU1: overlaps the next QK on MXU0 ----------------------------- |
| 102 | + # rhs = p [kv, q] (changes every tile -> must push+load each time); |
| 103 | + # lhs = V used as [hd, kv]; out[hd, q] = sum_kv V[kv,hd] p[kv,q]. |
| 104 | + pltpu.matmul_push_rhs(p.astype(v_tile.dtype), STAGE_P, MXU_PV) |
| 105 | + pltpu.matmul_acc_lhs(ACC_PV, v_tile, MXU_PV, load_staged_rhs=STAGE_P) |
| 106 | + o_curr = pltpu.matmul_pop(ACC_PV, (head_dim, TILE), jnp.float32, MXU_PV) |
| 107 | + |
| 108 | + # --- rescale + accumulate o (VPU). The online-softmax rescale is why PV is |
| 109 | + # popped per-tile rather than accumulated in the MXU across all kv tiles — |
| 110 | + # the running-max correction can't be applied inside the MXU accumulator. |
| 111 | + # OPTIMIZATION TARGET: a 2-pass / max-deferred scheme could let PV accumulate |
| 112 | + # in-MXU and drop this per-tile rescale (see assignment: redundant round-trips). |
| 113 | + o = alpha * o + o_curr |
| 114 | + m = m_next |
| 115 | + return m, l, o |
| 116 | + |
| 117 | + m, l, o = lax.fori_loop(0, num_kv_tiles, body, (m, l, o), unroll=True) |
| 118 | + return o, l # caller does o / l (deferred normalize, as today) |
| 119 | + |
| 120 | + |
| 121 | +# --------------------------------------------------------------------------- |
| 122 | +# Open design questions to settle against the working baseline / on device: |
| 123 | +# |
| 124 | +# 1. head_dim==256. With the current head_dim==128 you must either pad the |
| 125 | +# contraction (back to eff=0.5, defeating the point) or pack two heads into |
| 126 | +# the 256 contraction lane. This kernel only pays off WITH the hdim-256 |
| 127 | +# surgery — they compose; sequence them together. |
| 128 | +# |
| 129 | +# 2. Staging/accumulator register budget. STAGE_*/ACC_* are placeholders. Each |
| 130 | +# MXU has its own staging set; confirm one in-flight Q (mxu0) + one in-flight |
| 131 | +# p (mxu1) + the two f32 acc tiles fit, and that nothing is left resident on |
| 132 | +# exit (hard rule, no bounds check). |
| 133 | +# |
| 134 | +# 3. Dual-MXU overlap. QK(t+1) on MXU0 should issue while PV(t) runs on MXU1. |
| 135 | +# With unroll=True the scheduler can interleave; verify in the LLO that the |
| 136 | +# two vmatpush/vmatmul streams target mxu0 vs mxu1 and actually overlap. |
| 137 | +# |
| 138 | +# 4. Redundant-latch audit (the assignment): Q is latched once per strip here |
| 139 | +# (load_staged_rhs only on first tile). If the baseline re-pushes Q per kv |
| 140 | +# tile, that's the first thing to delete. p MUST re-push each tile (it |
| 141 | +# changes) — that one is not redundant. |
| 142 | +# |
| 143 | +# 5. Tail handling: num_kv_tiles assumes kv_seq_len % 256 == 0. Mirror the |
| 144 | +# current kernel's last_compute_body for the ragged tail, OR pad kv to 256. |
| 145 | +# |
| 146 | +# 6. Success metric: re-dump packed-bundles-post-ra and confirm spill stores/ |
| 147 | +# fills (#allocationN_spill) drop from ~43k toward ~0. That, not MXU %, is |
| 148 | +# the number that predicts the speedup. |
| 149 | +# --------------------------------------------------------------------------- |
0 commit comments