# SPDX-License-Identifier: Apache-2.0 """Best-known LeapQuant flush kernel for Kimi Delta Attention (per-key-channel gate; see README.md). Warp-specialised TileLang kernel for sm_100 (B200): one persistent CTA per SM, a TMA producer warp and three 128-thread consumer groups, products on tensor cores (mma.sync m16n8k16, bf16 hi/lo operands), the subspace iteration on an fp16 copy of the rebuilt state in shared memory with the old rank-4 part carried algebraically, CholeskyQR2 with a relative pivot floor for the new Compensator Tokens, the residual accumulated in fp32 from the checkpoint, the buffered updates and both rank-4 parts, the due programs balanced across CTAs by a rank scan. run(...) functional form with the input order / outputs of definition.json (packs the inputs into a slot pool, calls the kernel, returns the new checkpoint as fresh tensors) -- used for correctness Pool, flush_inplace deployment form: caller-owned strided slot pool updated in place -- what benchmark.py times Requires: torch (CUDA), tilelang 0.1.12, CUDA toolkit 13 (nvcc), NVIDIA B200. """ import torch, tilelang, tilelang.language as T from tilelang.layout import make_swizzled_layout F32, BF16, I32, I8 = "float32", "bfloat16", "int32", "int8" QMAX = 127.0 EPS = 1e-12 HV, K, V, L, R = 32, 128, 128, 16, 4 # heads per sequence, key dim, value dim, window length, rank SS, OFF = 33 * V * K, 12288 # pool geometry in fp32 words: slot stride (one spare head per slot), base offset @tilelang.jit(pass_configs={"tl.disable_thread_storage_sync": True}) def make_flush_kernel(NS, HV, K, V, L, SS, R=4, ITERS=1, SMS=148, NCG=3, MAXIT=128, PROF=False, PAD=8, KDA=False, GB16=True, PROFIT=False, DROPLO=False, MMT="bf16", ABL="", PNMIN=1e-20, GRAM16=False, NSTAGE=0, NMAX=0, KEXACT=True, **_ignored): """TileLang generator of the flush kernel (KDA=True: the per-channel gate of this request; KDA=False is the scalar-gate variant of requests/conless-leapquant-gdn-flush). The rebuilt state B (state minus the old rank-4 part) lives in shared memory as fp16 for the subspace iteration; the old rank-4 part is carried algebraically; the residual is accumulated in fp32 from the stage (see KEXACT below) and quantised row-wise. Keyword switches other than the shapes are ablations (defaults are the deployed configuration).""" # KEXACT (KDA): the residual E is accumulated from fp32 terms only -- the stage (codes, scales, decayed keys, values) is still resident, # so B is rebuilt in registers instead of read back from the fp16 St_s, and the old Compensator Tokens stay unscaled in UQ_s rows # 0..R-1: E = (deq(S0) + Q_old U_old^T) pc + ring buffer - Q U^T with pc = exp(Pbuf) in fp32. The subspace iteration still runs on # the fp16 St_s and an fp16 pc U_old^T copy in the zero padding rows R..2R-1 (only the new factors' quality depends on those). # E then goes to St_s as fp16 scaled by a power of two per (warp, half), so tiny heads do not fall into fp16 subnormals. # Without KEXACT (pc U_old^T rounded in place, E from the fp16 B, unscaled) the reconstruction error is 1-8 % above the fp32 # reference's; with it 0.1-0.3 % (+20 % flush time: the second dequantisation). KEX = bool(KDA and KEXACT) assert not KEX or 2 * R <= 8 # Gbuf holds fp16 factors; it is declared bf16 so the TMA target (Kbl_s, later the lo key tile) matches -- the bits are reinterpreted on read assert 1 <= R <= 4 and K == 128 and V == 128 and (not KDA or GB16) and PAD == 8 F16 = "float16" N = T.symbolic("N") KW = K // 4 NPROD = 32 NCONS = 128 * NCG NT = NPROD + NCONS NWARPS = NT // 32 NDUE = min(NS - 8, NMAX) if NMAX else NS - 8 # rows that can be due at once: one call's rows (NMAX = the engine's batch cap) or the whole table MAXIT = max(MAXIT, (NDUE * HV + SMS - 1) // SMS + 1) # dlist capacity: a CTA's share of all programs when every such row is due # (with NMAX the list is ~350 B instead of 5.5 KB on an 8155-slot engine: on sm_120 that is the difference between a 2-stage flush fitting in 99 KB or not) KMAX = min((NS + NWARPS - 1) // NWARPS, 40) # n-steps per warp kept in registers by the compaction (3 x 40 registers; N <= 40 * NWARPS = # 520 sequences at NCG = 3, vLLM's max_num_seqs 512); larger batches take a serial tail # with the flags recomputed in the rank pass (64 steps spilled: +3 us none-due) NSTAGE = NSTAGE or NCG # NSTAGE > NCG: each group double-buffers its stage (the TMA warp prefetches the group's assert NSTAGE % NCG == 0 # next due program while it works) -- worth it where one group is all that fits (sm_120) SPG = NSTAGE // NCG KP = K + PAD # St_s row stride: conflict-free fragment stores, ldmatrix rows and 16 B row reads; the pad holds the fp16 Q tile KV8 = K + V + 8 # UQn_s row stride (halfwords): 528 B -> conflict-free ldmatrix rows P16OFF = 4 * KV8 # fp16 P tile (K x 8) at halfword offset 4 rows into UQn_s (after the fp32 Q plane) assert 4 * KV8 * 2 >= (V + 4) * 4 * 4 and 4 * KV8 >= K * 8 @T.prim_func def flush(Sq: T.StridedTensor((NS, HV, V, KW), (SS, V * K, KW, 1), I32), Sk: T.StridedTensor((NS, HV, K), (SS, V * K, 1), F32), Sv: T.StridedTensor((NS, HV, V), (SS, V * K, 1), F32), Skv: T.StridedTensor((NS, HV, K + V), (SS, V * K, 1), F32), Kbuf: T.Tensor((NS, HV, L, K), BF16), Ubuf: T.Tensor((NS, HV, L, V), BF16), Wbuf: T.Tensor((NS, HV, L), F32), Pbuf: T.Tensor((NS, HV, K if KDA else 1), F32), Gbuf: T.Tensor((NS, HV, L, K) if KDA else (1, 1, 1, 1), BF16), hcnt: T.Tensor((NS, HV), I32), Sq16: T.StridedTensor((NS, HV, V, K // 2), (2 * SS, 2 * V * K, K // 2, 1), "int16"), # unused (kept for the v2 call signature) UQ: T.StridedTensor((NS, HV, R, K + V), (2 * SS, 2 * V * K, K + V, 1), F16), # fp16 factors U^T | Q^T after the scales Q0: T.Tensor((V, R), F32), # deterministic start of the subspace iteration idx: T.Tensor((N,), I32)): with T.Kernel(SMS, threads=NPROD + NCONS) as bid: Qs = T.alloc_shared((NSTAGE, V, KW), I32) skv_s = T.alloc_shared((NSTAGE, K + V), F32) Kb_s = T.alloc_shared((NSTAGE, 2, L, K // 2), BF16) # two 128 B-wide halves, 128 B swizzle (conflict-free ldmatrix) Ub_s = T.alloc_shared((NSTAGE, 2, L, V // 2), BF16) Wb_s = T.alloc_shared((NSTAGE, 32), F32) # 128 B per stage (TMA destination alignment) UQ_s = T.alloc_shared((NSTAGE, 4, 8, 64), F16) # old factors U^T | Q^T in four 128 B-wide column blocks (swizzled: conflict-free ldmatrix rows), k-rows padded to 8 (rows R.. zero) St_s = T.alloc_shared((NCG, V, KP), F16) # B (fp16) for the iteration, then the residual E in place; columns K.. = the fp16 Q tile (V x 8, columns R.. zero) UQn_s = T.alloc_shared((NCG, 8, KV8), F16) # new -U^T (rows R.. zero) for the residual mma; aliases: PQ_s (fp32 Q rows), P16 (fp16 P tile), cpart (column partials) PQ_s = T.view(UQn_s, (NCG, 2, KV8 // 2, 4), F32) # plane 0 = Q (V x R) fp32 for the CholeskyQR rows UQn1 = T.view(UQn_s, (NCG, 8 * KV8), F16) # flat view: the fp16 P tile P16[c][r] at P16OFF + 8 c + r cpart = T.view(UQn_s, (NCG, 8, KV8 // 2), F32) # per-warp column partials of |E| (rows 0..3) rsk = T.view(UQn_s, (NCG, 8, KV8 // 2), F32) # reciprocal column scales: row 6 of the cpart view (free once the partials are read) dlist = T.alloc_shared((MAXIT,), I32) # this CTA's due programs p (global rank == bid mod SMS), in rank order ndue_s = T.alloc_shared((1,), I32) wsum_s = T.alloc_shared((NWARPS,), I32) # per-warp due counts of the block-wide compaction if KDA: Kbl_s = T.alloc_shared((NSTAGE, 2, L, K // 2), BF16) # per-entry decay factors (TMA target, fp16 bits) -> lo part of the decayed keys pb_s = T.alloc_shared((NSTAGE, K), F32) pc_s = T.alloc_shared((NCG, K), F32) loaded = T.alloc_barrier([32] * NSTAGE) consumed = T.alloc_barrier([128] * NSTAGE) if KDA: T.annotate_layout({Qs: make_swizzled_layout(Qs), Kb_s: make_swizzled_layout(Kb_s), Ub_s: make_swizzled_layout(Ub_s), Kbl_s: make_swizzled_layout(Kbl_s), UQ_s: make_swizzled_layout(UQ_s)}) else: T.annotate_layout({Qs: make_swizzled_layout(Qs), Kb_s: make_swizzled_layout(Kb_s), Ub_s: make_swizzled_layout(Ub_s), UQ_s: make_swizzled_layout(UQ_s)}) tx = T.get_thread_binding() # ---- block-wide balanced compaction (no atomics): every CTA scans all N * HV due flags -- warp w takes the sequences # n = w, w + NWARPS, ..., lane hv the head (one broadcast idx load + one coalesced 128 B hcnt load per n; an uncoalesced # gather costs 32 L1 lookups per warp-load and was 8x slower) -- ranks the due programs with ballot/popc in the fixed order # (warp, n-step, hv) and keeps the ranks == bid (mod SMS). The flush is latency-bound at 1/16 duty (1-3 programs per # group): the static assignment's binomial imbalance (max 12 of mean 3.6 per CTA at bs 256) cost ~1.4x; this gives +-1. # dlist holds q = slot * HV + hv. ---- tc = T.alloc_local((4,), "int64") if PROF: tc[0] = T.call_extern("int64", "clock64") wlane = tx % 32 wid = tx // 32 kn = T.alloc_local((1,), I32) kn[0] = T.ceildiv(N, NWARPS) # n-steps per warp cnt = T.alloc_local((1,), I32) cnt[0] = 0 slk = T.alloc_local((KMAX,), I32) hck = T.alloc_local((KMAX,), I32) mk = T.alloc_local((KMAX,), "uint32") for k in T.unroll(KMAX): # all loads in flight at once (explicit clamps: no guard branches) slk[k] = idx[T.max(T.min(wid + NWARPS * k, N - 1), 0)] for k in T.unroll(KMAX): hck[k] = hcnt[T.min(T.max(slk[k], 0), NS - 1), wlane] for k in T.unroll(KMAX): f = T.if_then_else((k < kn[0]) & (wid + NWARPS * k < N) & (slk[k] > 0) & (hck[k] == L), 1, 0) mk[k] = T.call_extern("uint32", "__ballot_sync", T.uint32(0xFFFFFFFF), f) cnt[0] += T.call_extern(I32, "__popc", mk[k]) sl8 = T.alloc_local((8,), I32) hc8 = T.alloc_local((8,), I32) for ck in T.serial(T.ceildiv(kn[0] - KMAX, 8)): # tail (N > KMAX * NWARPS only): serial chunks, flags recomputed in the rank pass for j in T.unroll(8): sl8[j] = idx[T.max(T.min(wid + NWARPS * (KMAX + ck * 8 + j), N - 1), 0)] for j in T.unroll(8): hc8[j] = hcnt[T.min(T.max(sl8[j], 0), NS - 1), wlane] for j in T.unroll(8): k = KMAX + ck * 8 + j f = T.if_then_else((k < kn[0]) & (wid + NWARPS * k < N) & (sl8[j] > 0) & (hc8[j] == L), 1, 0) cnt[0] += T.call_extern(I32, "__popc", T.call_extern("uint32", "__ballot_sync", T.uint32(0xFFFFFFFF), f)) if PROF: tc[1] = T.call_extern("int64", "clock64") if wlane == 0: wsum_s[wid] = cnt[0] T.sync_threads(14, NPROD + NCONS) base = T.alloc_local((1,), I32) base[0] = 0 for ww in T.serial(NWARPS): base[0] += T.if_then_else(ww < wid, wsum_s[ww], 0) tot = T.alloc_local((1,), I32) tot[0] = 0 for ww in T.serial(NWARPS): tot[0] += wsum_s[ww] T.sync_threads(14, NPROD + NCONS) # wsum_s reads done before dlist (which may share its smem) is written lmask = T.cast((T.int64(1) << T.cast(wlane, "int64")) - T.int64(1), "uint32") if cnt[0] > 0: # warps without due programs skip the pass (masks past kn are zero) rst = T.alloc_local((2,), I32) # next rank == bid (mod SMS) at or above base, and its dlist index (one division per thread) rst[0] = base[0] + (bid - base[0] % SMS + SMS) % SMS rst[1] = rst[0] // SMS for k in T.unroll(KMAX): # static indexing of the mask registers; one predicate per step, no division m = mk[k] rank = base[0] + T.call_extern(I32, "__popc", m & lmask) if (((m >> T.cast(wlane, "uint32")) & T.uint32(1)) == T.uint32(1)) & (rank == rst[0]): dlist[rst[1]] = idx[wid + NWARPS * k] * HV + wlane base[0] += T.call_extern(I32, "__popc", m) if rst[0] < base[0]: # the warp's rank sequence passed the target (warp-uniform; at most one target per step) rst[0] += SMS rst[1] += 1 for ck in T.serial(T.ceildiv(kn[0] - KMAX, 8)): for j in T.unroll(8): sl8[j] = idx[T.max(T.min(wid + NWARPS * (KMAX + ck * 8 + j), N - 1), 0)] for j in T.unroll(8): hc8[j] = hcnt[T.min(T.max(sl8[j], 0), NS - 1), wlane] for j in T.unroll(8): k = KMAX + ck * 8 + j f = T.if_then_else((k < kn[0]) & (wid + NWARPS * k < N) & (sl8[j] > 0) & (hc8[j] == L), 1, 0) m = T.call_extern("uint32", "__ballot_sync", T.uint32(0xFFFFFFFF), f) if f == 1: rank = base[0] + T.call_extern(I32, "__popc", m & lmask) if rank % SMS == bid: dlist[rank // SMS] = sl8[j] * HV + wlane base[0] += T.call_extern(I32, "__popc", m) if tx == 0: ndue_s[0] = T.max(T.ceildiv(tot[0] - bid, SMS), 0) if PROF: tc[2] = T.call_extern("int64", "clock64") if tx < NPROD: for i in T.serial(NSTAGE * 4 * (8 - R) * 64 // NPROD): # zero the padding k-rows of UQ_s once (only rows < R are ever written) e = tx + NPROD * i UQ_s[e // (4 * (8 - R) * 64), (e // ((8 - R) * 64)) % 4, R + (e % ((8 - R) * 64)) // 64, e % 64] = T.cast(0.0, F16) T.sync_threads(14, NPROD + NCONS) if tx < NPROD: T.fence_proxy_async() for k in T.serial(ndue_s[0]): q = dlist[k] gc = k % NCG ig = k // NCG s = gc * SPG + ig % SPG hv = q % HV slot = q // HV T.mbarrier_wait_parity(consumed[s], ((ig // SPG) & 1) ^ 1) T.tma_copy(Sq[slot, hv, :, :], Qs[s, :, :], barrier=loaded[s]) T.tma_copy(Skv[slot, hv, :], skv_s[s, :], barrier=loaded[s]) T.tma_copy(Kbuf[slot, hv, :, 0:K // 2], Kb_s[s, 0, :, :], barrier=loaded[s]) T.tma_copy(Kbuf[slot, hv, :, K // 2:K], Kb_s[s, 1, :, :], barrier=loaded[s]) T.tma_copy(Ubuf[slot, hv, :, 0:V // 2], Ub_s[s, 0, :, :], barrier=loaded[s]) T.tma_copy(Ubuf[slot, hv, :, V // 2:V], Ub_s[s, 1, :, :], barrier=loaded[s]) T.tma_copy(Wbuf[slot, hv, :], Wb_s[s, 0:L], barrier=loaded[s]) for cb in T.unroll(4): T.tma_copy(UQ[slot, hv, :, cb * 64:(cb + 1) * 64], UQ_s[s, cb, 0:R, :], barrier=loaded[s]) if KDA: T.tma_copy(Gbuf[slot, hv, :, 0:K // 2], Kbl_s[s, 0, :, :], barrier=loaded[s]) T.tma_copy(Gbuf[slot, hv, :, K // 2:K], Kbl_s[s, 1, :, :], barrier=loaded[s]) T.tma_copy(Pbuf[slot, hv, :], pb_s[s, :], barrier=loaded[s]) T.mbarrier_arrive(loaded[s]) else: g = (tx - NPROD) // 128 t = (tx - NPROD) % 128 w = t // 32 lane = t % 32 q4 = lane % 4 r0 = lane // 4 bar = 8 + g qw = T.alloc_local((4,), I32) svr_l = T.alloc_local((4,), F32) rmax = T.alloc_local((4,), F32) afrag = T.alloc_local((32 if not KDA else 16,), BF16) araw = T.alloc_local((8,), BF16) wv = T.alloc_local((4,), F32) bfrag = T.alloc_local((32 if KDA else 16,), BF16) # 2 x (ring-buffer B tile, KDA: hi + lo) -- alternating per unrolled step afrag2 = T.alloc_local((8,), F16) # rank-R A fragments (m16n8k8) per m-tile: old Q^T (residual) afrag3 = T.alloc_local((8,), F16) # rank-R A fragments of the new factors (residual) cfrag = T.alloc_local((64,), F32) # half of B / E in fragment layout: index (ntl * 4 + (mt * 2 + r)) * 2 + c1, ntl = local n-tile 0..7 efrag = T.alloc_local((8,), F16) # ldmatrix.x4 result: B values of (ntl, top), (ntl, bottom), (ntl+1, top), (ntl+1, bottom) colp = T.alloc_local((32,), F32) # column partials of |E| (16 n-tiles x 2 columns of this lane) skl = T.alloc_local((16,), F32) # this lane's 16 column scales of the current half (dequant) ah = T.alloc_local((32,), F16) # m16n8k16 A fragments (ldmatrix.x4 from the fp16 St_s): 2 k-parities x 2 m-tiles bh2 = T.alloc_local((8,), F16) # B fragments (Q or P column pairs, fp16, ldmatrix.x2.trans): 2 k-parities ao = T.alloc_local((16,), F16) # A fragments of the old factor (rows 0..7 of a 16-row m-tile; rows 8..15 zero): 2 k-parities bfr = T.alloc_local((8,), F16) # 2 x rank-R B fragment (residual) bhi = T.alloc_local((4,), I32) # |E| B fragments (sign bits cleared) of one n-tile, 2 k-steps aone = T.alloc_local((8,), F16) # A fragment of ones (column sums on the tensor cores) c32 = T.alloc_local((16,), F32) # 2 m-tiles x 4, two accumulator chains (even/odd k-steps) mc = T.alloc_local((4,), F32) # R x R correction matrix accumulator (M = Q_old^T Q, M2 = U_old^T P), full in every warp mm = T.alloc_local((8,), F32) # M[a][2 q4 + i] for a = 0..3 (shuffled from lanes 4 a + q4) uo = T.alloc_local((16,), F32) # old factor entries U_old[a][c] / Q_old[a][v] of this lane's 4 rows pr = T.alloc_local((4,), F32) # this thread's row of Q gg = T.alloc_local((10,), F32) # Gram entries (0,0) (1,0) (1,1) (2,0) (2,1) (2,2) (3,0) (3,1) (3,2) (3,3) ll = T.alloc_local((10,), F32) # Cholesky factor, same packing uh = T.alloc_local((4,), F16) qh = T.alloc_local((4,), F16) el = T.alloc_local((L,), F32) # KDA: this column's per-entry decay factors pcl = T.alloc_local((16 if KEX else 1,), F32) # KEX: pc of this lane's residual columns emx = T.alloc_local((1,), F32) # KEX: block max of |E|, then the power-of-two scale esc = T.alloc_local((2,), F32) # KEX: per-half inverse E scale of this warp's rows dcur = T.alloc_local((1,), F32) eh = T.alloc_local((32,), F16) # quantisation: 2 x 2 x 8 halfwords of E (alternating per row pass) ph = T.alloc_local((8,), F16) # fp16 pairs for stmatrix (vectorised cast -> one F2FP.PACK per pair) ex = T.alloc_local((16,), F32) # quantisation: this lane's 16 scaled residual values of the row rk = T.alloc_local((16,), F32) # quantisation: this lane's 16 reciprocal column scales tk = T.alloc_local((10,), "int64") tk2 = T.alloc_local((6,), "int64") tacc = T.alloc_local((10,), F32) for i in T.unroll(10): tacc[i] = 0.0 if PROF: tacc[7] = T.cast(tc[1] - tc[0], F32) # compaction: loads + ballots tacc[9] = T.cast(tc[2] - tc[1], F32) # compaction: scan + rank pass for i in T.unroll(8): aone[i] = T.cast(1.0, F16) for i in T.unroll(2): # the old-factor A fragments: register pairs 1 and 3 (m-rows 8..15) stay zero ao[2 + i] = T.cast(0.0, F16) ao[6 + i] = T.cast(0.0, F16) ao[10 + i] = T.cast(0.0, F16) ao[14 + i] = T.cast(0.0, F16) for r in T.unroll(4): # the fp16 Q tile's padding columns R..7 stay zero (this thread's row) St_s[g, t, K + 4 + r] = T.cast(0.0, F16) for r in T.unroll(4 - R): St_s[g, t, K + R + r] = T.cast(0.0, F16) for it0 in T.serial(T.ceildiv(ndue_s[0] - g, NCG)): k = g + it0 * NCG q = dlist[k] hv = q % HV s = g * SPG + it0 % SPG if PROF: tk[0] = T.call_extern("int64", "clock64") T.mbarrier_wait_parity(loaded[s], (it0 // SPG) & 1) if PROF: tk[1] = T.call_extern("int64", "clock64") slot = q // HV if KDA: # pre-pass (thread per column c = t): pc = exp(G); decayed keys w_j k_j D_j as bf16 hi (in place) + lo (over the factors); # U^T pre-multiplied by pc (the checkpoint's column decay) pcc = T.exp(pb_s[s, t]) pc_s[g, t] = pcc for jj in T.unroll(L): el[jj] = T.cast(T.reinterpret(F16, Kbl_s[s, t // (K // 2), jj, t % (K // 2)]), F32) T.sync_threads(bar, 128) # every lane holds its factors before the lo tile overwrites them dcur[0] = 1.0 for jj in T.unroll(L): # backward: D_{L-1} = 1, D_j = D_{j+1} e_{j+1} j = L - 1 - jj kd = T.cast(Kb_s[s, t // (K // 2), j, t % (K // 2)], F32) * (Wb_s[s, j] * dcur[0]) khi = T.cast(kd, BF16) Kb_s[s, t // (K // 2), j, t % (K // 2)] = khi Kbl_s[s, t // (K // 2), j, t % (K // 2)] = T.cast(kd - T.cast(khi, F32), BF16) dcur[0] = dcur[0] * el[j] for r in T.unroll(R): # pc U^T in place (fp16: 2^-11 of the rank-R term, measured irrelevant next to the int8 residual) if KEX: UQ_s[s, t // 64, R + r, t % 64] = T.cast(T.cast(UQ_s[s, t // 64, r, t % 64], F32) * pcc, F16) else: UQ_s[s, t // 64, r, t % 64] = T.cast(T.cast(UQ_s[s, t // 64, r, t % 64], F32) * pcc, F16) T.sync_threads(bar, 128) for i in T.unroll(4): svr_l[i] = skv_s[s, K + w * 32 + (i // 2) * 16 + (i % 2) * 8 + lane // 4] * (1.0 / QMAX) else: pn = T.max(Pbuf[slot, hv, 0], PNMIN) # floor: a fully decayed head has a denormal product (2e-41 seen on 9B layer 11) and 1/pn = inf rpn = 1.0 / pn # poisoned the residual (Sk = inf -> NaN outputs); the floor changes E by <= 1e-20 |Q_old U_old^T| if KDA: pn = 1.0 if not KDA: for i in T.unroll(4): svr_l[i] = skv_s[s, K + w * 32 + (i // 2) * 16 + (i % 2) * 8 + lane // 4] * (pn / QMAX) # (1) A fragments of the ring buffer u_j (GDN: w_j-scaled, bf16 hi + lo; KDA: exact bf16) if KDA: for mt in T.unroll(2): T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Ub_s[s, (w * 32 + mt * 16) // (V // 2), (lane % 8) + 8 * (lane // 16), (w * 32 + mt * 16) % (V // 2) + 8 * ((lane // 8) % 2)], "r", extent=8), T.access_ptr(afrag[mt * 8], "w", extent=8)) else: for i in T.unroll(4): wv[i] = Wb_s[s, 2 * q4 + (i % 2) + 8 * (i // 2)] for mt in T.unroll(2): T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Ub_s[s, (w * 32 + mt * 16) // (V // 2), (lane % 8) + 8 * (lane // 16), (w * 32 + mt * 16) % (V // 2) + 8 * ((lane // 8) % 2)], "r", extent=8), T.access_ptr(araw[0], "w", extent=8)) for i in T.unroll(8): wu = wv[(i % 2) + 2 * (i // 4)] * T.cast(araw[i], F32) afrag[mt * 16 + i] = T.cast(wu, BF16) afrag[mt * 16 + 8 + i] = T.cast(wu - T.cast(afrag[mt * 16 + i], F32), BF16) # (2) two 64-column halves: B = ring buffer (mma) + pn * deq(S0), out to St_s (fp16, stmatrix); the old rank-R part stays algebraic for h in T.unroll(2): # dequantise first (independent FFMA chains, no accumulator dependency), then the ring-buffer mma accumulates on top for ntl in T.unroll(8): for c1 in T.vectorized(2): if KDA: skl[ntl * 2 + c1] = skv_s[s, (h * 8 + ntl) * 8 + q4 * 2 + c1] * pc_s[g, (h * 8 + ntl) * 8 + q4 * 2 + c1] else: skl[ntl * 2 + c1] = skv_s[s, (h * 8 + ntl) * 8 + q4 * 2 + c1] for rr in (T.unroll(0) if "nodeq" in ABL else T.unroll(4)): v0 = w * 32 + (rr // 2) * 16 + (rr % 2) * 8 + r0 for jcl in T.unroll(4): # 16-column chunk = n-tiles 2jc, 2jc+1; lane q4 needs word q4 // 2 of each n-tile jc = h * 4 + jcl for nn in T.unroll(2): qw[nn] = Qs[s, v0, jc * 4 + nn * 2 + q4 // 2] for nn in T.unroll(2): ntl = jcl * 2 + nn pkw = qw[nn] ^ T.int32(-2139062144) for c1 in T.unroll(2): # the magic-number subtraction must stay a separate exact fp32 op: folding -K*skl into an FFMA x = T.reinterpret(F32, T.call_extern(I32, "__byte_perm", pkw, T.int32(0x4B000000), T.int32(0x7650) + (q4 % 2) * 2 + c1)) - 8388736.0 # costs ulp(2^23 skl) = half a code cfrag[(ntl * 4 + rr) * 2 + c1] = x * svr_l[rr] * skl[ntl * 2 + c1] # stmatrix.x4: matrices (ntl, rows 0-7), (ntl, rows 8-15), (ntl+1, rows 0-7), (ntl+1, rows 8-15) of m-tile mt for npl in T.unroll(4): n0 = (h * 4 + npl) * 16 bo = (npl % 2) * (16 if KDA else 8) T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Kb_s[s, n0 // (K // 2), (lane % 8) + 8 * ((lane // 8) % 2), n0 % (K // 2) + 8 * (lane // 16)], "r", extent=8), T.access_ptr(bfrag[bo], "w", extent=8)) if KDA: T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Kbl_s[s, n0 // (K // 2), (lane % 8) + 8 * ((lane // 8) % 2), n0 % (K // 2) + 8 * (lane // 16)], "r", extent=8), T.access_ptr(bfrag[bo + 8], "w", extent=8)) for mt in T.unroll(2): for nn in T.unroll(2): if KDA: T.ptx_mma("float32", "m16n8k16", "row", "col", "bf16", "bf16", "fp32", afrag.data, mt * 8, bfrag.data, bo + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) T.ptx_mma("float32", "m16n8k16", "row", "col", "bf16", "bf16", "fp32", afrag.data, mt * 8, bfrag.data, bo + 8 + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) else: T.ptx_mma("float32", "m16n8k16", "row", "col", MMT, MMT, "fp32", afrag.data, mt * 16, bfrag.data, bo + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) if not DROPLO: T.ptx_mma("float32", "m16n8k16", "row", "col", MMT, MMT, "fp32", afrag.data, mt * 16 + 8, bfrag.data, bo + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) for mt in (T.unroll(0) if "nostb" in ABL else T.unroll(2)): for npl in T.unroll(4): for jj in T.unroll(4): # 4 fp16 pairs: (ntl, rows 0-7), (ntl, rows 8-15), (ntl+1, ...) -- one F2FP.PACK each for c1 in T.vectorized(2): ph[jj * 2 + c1] = T.cast(cfrag[((npl * 2 + jj // 2) * 4 + mt * 2 + jj % 2) * 2 + c1], F16) T.call_extern("handle", "tl::ptx_stmatrix_m8n8_x4", T.access_ptr(St_s[g, w * 32 + mt * 16 + 8 * ((lane // 8) % 2) + lane % 8, (h * 8 + npl * 2 + lane // 16) * 8], "w", extent=8), T.reinterpret(I32, ph[0:2]), T.reinterpret(I32, ph[2:4]), T.reinterpret(I32, ph[4:6]), T.reinterpret(I32, ph[6:8])) # the iteration's start Q0 (fp16 tile in the St_s pad columns; this thread's row) -- before the barrier that publishes B for r in T.unroll(R): St_s[g, t, K + r] = T.cast(Q0[t, r], F16) T.sync_threads(bar, 128) if PROF: tk[2] = T.call_extern("int64", "clock64") # ================= subspace iteration on B (fp16 mma) with the algebraic rank-R corrections ================= for itr in T.serial(ITERS + 1): if PROFIT: tk2[0] = T.call_extern("int64", "clock64") # ---- P = B^T Q (K x R): A[m=c][k=v] = St_s[v][c] via ldmatrix.x4.trans; B[k=v][n=r] = Q16[v][r] via ldmatrix.x2.trans; # alongside M = Q_old^T Q (R x R): A[m=r][k=v] = Q_old^T (rows 8..15 zero), same B fragments ---- for i in T.unroll(16): c32[i] = 0.0 for i in T.unroll(4): mc[i] = 0.0 for k0 in T.unroll(V // 16): kp = k0 % 2 # alternating fragment registers: the next step's ldmatrix can overlap this step's mma T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(St_s[g, k0 * 16 + 8 * ((lane // 8) % 2) + lane % 8, K], "r", extent=4), T.access_ptr(bh2[kp * 4], "w", extent=4)) for mt in T.unroll(2): m0 = w * 32 + mt * 16 T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(St_s[g, k0 * 16 + 8 * (lane // 16) + lane % 8, m0 + 8 * ((lane // 8) % 2)], "r", extent=8), T.access_ptr(ah[(kp * 2 + mt) * 8], "w", extent=8)) T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", ah.data, (kp * 2 + mt) * 8, bh2.data, kp * 4, c32.data, (k0 % 2) * 8 + mt * 4, T.bool(False)) T.ptx_ldmatrix(T.bool(False), 1, T.access_ptr(UQ_s[s, (K + k0 * 16) // 64, lane % 8, (K + k0 * 16) % 64], "r", extent=2), T.access_ptr(ao[kp * 8], "w", extent=2)) T.ptx_ldmatrix(T.bool(False), 1, T.access_ptr(UQ_s[s, (K + k0 * 16 + 8) // 64, lane % 8, (K + k0 * 16 + 8) % 64], "r", extent=2), T.access_ptr(ao[kp * 8 + 4], "w", extent=2)) T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", ao.data, kp * 8, bh2.data, kp * 4, mc.data, 0, T.bool(False)) for i in T.unroll(8): c32[i] += c32[8 + i] if PROFIT: tk2[1] = T.call_extern("int64", "clock64") # ---- P += pn U_old M in the fragments: rows c = m0 + r0 (+8), columns n = 2 q4 + i; M[a][n] sits in lane 4 a + q4 ---- for a in T.unroll(4): mm[a * 2] = T.shfl_sync(mc[0], a * 4 + q4) mm[a * 2 + 1] = T.shfl_sync(mc[1], a * 4 + q4) for mt in T.unroll(2): for hh in T.unroll(2): cc = w * 32 + mt * 16 + hh * 8 + r0 for a in T.unroll(4): uo[(mt * 2 + hh) * 4 + a] = T.if_then_else(a < R, T.cast(UQ_s[s, cc // 64, (R + a) if KEX else a, cc % 64], F32), 0.0) for i in T.unroll(2): c32[mt * 4 + hh * 2 + i] += pn * (uo[(mt * 2 + hh) * 4] * mm[i] + uo[(mt * 2 + hh) * 4 + 1] * mm[2 + i] + uo[(mt * 2 + hh) * 4 + 2] * mm[4 + i] + uo[(mt * 2 + hh) * 4 + 3] * mm[6 + i]) # ---- the fp16 P tile P16[c][r] (columns >= R zero: lanes q4 >= R/2 write zeros) via stmatrix.x2 (rows 0-7, rows 8-15 of each m-tile) ---- for mt in T.unroll(2): for c1 in T.vectorized(2): ph[c1] = T.cast(T.if_then_else(2 * q4 + c1 < R, c32[mt * 4 + c1], 0.0), F16) ph[2 + c1] = T.cast(T.if_then_else(2 * q4 + c1 < R, c32[mt * 4 + 2 + c1], 0.0), F16) T.call_extern("handle", "tl::ptx_stmatrix_m8n8_x2", T.access_ptr(UQn1[g, P16OFF + (w * 32 + mt * 16 + 8 * ((lane // 8) % 2) + lane % 8) * 8], "w", extent=8), T.reinterpret(I32, ph[0:2]), T.reinterpret(I32, ph[2:4])) T.sync_threads(bar, 128) if PROFIT: tk2[2] = T.call_extern("int64", "clock64") tacc[0] += T.cast(tk2[1] - tk2[0], F32) tacc[1] += T.cast(tk2[2] - tk2[1], F32) if itr < ITERS: # ---- Q = B P (V x R): A[m=v][k=c] = St_s[v][c] via ldmatrix.x4; B[k=c][n=r] = P16[c][r] via ldmatrix.x2.trans; alongside M2 = U_old^T P ---- for i in T.unroll(16): c32[i] = 0.0 for i in T.unroll(4): mc[i] = 0.0 for k0 in T.unroll(K // 16): kp = k0 % 2 T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(UQn1[g, P16OFF + (k0 * 16 + 8 * ((lane // 8) % 2) + lane % 8) * 8], "r", extent=4), T.access_ptr(bh2[kp * 4], "w", extent=4)) for mt in T.unroll(2): m0 = w * 32 + mt * 16 T.ptx_ldmatrix(T.bool(False), 4, T.access_ptr(St_s[g, m0 + 8 * ((lane // 8) % 2) + lane % 8, k0 * 16 + 8 * (lane // 16)], "r", extent=8), T.access_ptr(ah[(kp * 2 + mt) * 8], "w", extent=8)) T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", ah.data, (kp * 2 + mt) * 8, bh2.data, kp * 4, c32.data, (k0 % 2) * 8 + mt * 4, T.bool(False)) T.ptx_ldmatrix(T.bool(False), 1, T.access_ptr(UQ_s[s, (k0 * 16) // 64, lane % 8, (k0 * 16) % 64], "r", extent=2), T.access_ptr(ao[kp * 8], "w", extent=2)) T.ptx_ldmatrix(T.bool(False), 1, T.access_ptr(UQ_s[s, (k0 * 16 + 8) // 64, lane % 8, (k0 * 16 + 8) % 64], "r", extent=2), T.access_ptr(ao[kp * 8 + 4], "w", extent=2)) T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", ao.data, kp * 8, bh2.data, kp * 4, mc.data, 0, T.bool(False)) for i in T.unroll(8): c32[i] += c32[8 + i] # ---- Q += pn Q_old M2 in the fragments, then the fp32 rows to PQ_s plane 0 for the CholeskyQR ---- for a in T.unroll(4): mm[a * 2] = T.shfl_sync(mc[0], ((R + a) if KEX else a) * 4 + q4) mm[a * 2 + 1] = T.shfl_sync(mc[1], ((R + a) if KEX else a) * 4 + q4) for mt in T.unroll(2): for hh in T.unroll(2): vv = w * 32 + mt * 16 + hh * 8 + r0 for a in T.unroll(4): uo[(mt * 2 + hh) * 4 + a] = T.if_then_else(a < R, T.cast(UQ_s[s, (K + vv) // 64, a, (K + vv) % 64], F32), 0.0) for i in T.unroll(2): c32[mt * 4 + hh * 2 + i] += pn * (uo[(mt * 2 + hh) * 4] * mm[i] + uo[(mt * 2 + hh) * 4 + 1] * mm[2 + i] + uo[(mt * 2 + hh) * 4 + 2] * mm[4 + i] + uo[(mt * 2 + hh) * 4 + 3] * mm[6 + i]) if q4 * 2 < R: for mt in T.unroll(2): PQ_s[g, 0, w * 32 + mt * 16 + r0, 2 * q4] = c32[mt * 4] PQ_s[g, 0, w * 32 + mt * 16 + r0, 2 * q4 + 1] = c32[mt * 4 + 1] PQ_s[g, 0, w * 32 + mt * 16 + r0 + 8, 2 * q4] = c32[mt * 4 + 2] PQ_s[g, 0, w * 32 + mt * 16 + r0 + 8, 2 * q4 + 1] = c32[mt * 4 + 3] T.sync_threads(bar, 128) if PROFIT: tk2[3] = T.call_extern("int64", "clock64") tacc[2] += T.cast(tk2[3] - tk2[2], F32) # ---- CholeskyQR2 of Q (V rows = 128 threads), relative pivot floor for rank-deficient heads. Pass 1 in fp32 (shuffle tree); # pass 2's Gram on the tensor cores in fp16 (8 mma; the factors are stored in fp16 anyway and the residual uses the same # fp16 factors, so a 2^-11 non-orthonormality costs nothing). A 3xTF32 Gram for pass 1 was slower (48 dependent mma). ---- for r in T.unroll(4): pr[r] = T.if_then_else(r < R, PQ_s[g, 0, t, r], 0.0) for cq in range(2): if cq == 0 or not GRAM16: # pass 1 (kappa up to 1e3): fp32 Gram by the warp-shuffle tree + smem partials. GRAM16=False: pass 2 too -- # after a floored pass-1 pivot (rank-deficient head, e.g. svals 4.8 / 0.2 / 8e-3 / 1e-3) the 4th column has norm 1e-3..1e-5, # its fp16 Gram entry is subnormal/zero, rsqrt blows the column up by 1e3..1e15 and U_4 = B^T q_4 with it for r in T.unroll(4): for c in T.unroll(r + 1): gg[r * (r + 1) // 2 + c] = pr[r] * pr[c] for off in (T.unroll(0) if "nogram" in ABL else T.unroll(5)): for i in T.unroll(10): gg[i] += T.shfl_xor(gg[i], 16 >> off) if lane == 0: for i in T.unroll(10): cpart[g, 4, w * 16 + i] = gg[i] # per-warp Gram partials (cpart row 4: free during the iteration) T.sync_threads(bar, 128) for i in T.unroll(10): gg[i] = cpart[g, 4, i] + cpart[g, 4, 16 + i] + cpart[g, 4, 32 + i] + cpart[g, 4, 48 + i] else: # pass 2 (Q nearly orthonormal): Gram on the tensor cores from the fp16 Q tile, for i in T.unroll(4): # A = Q^T (rows R.. zero) and B = Q are the same ldmatrix.x2.trans fragment mc[i] = 0.0 for k0 in T.unroll(V // 16): kp = k0 % 2 T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(St_s[g, k0 * 16 + 8 * ((lane // 8) % 2) + lane % 8, K], "r", extent=4), T.access_ptr(bh2[kp * 4], "w", extent=4)) for i in T.unroll(2): ao[kp * 8 + i] = bh2[kp * 4 + i] ao[kp * 8 + 4 + i] = bh2[kp * 4 + 2 + i] T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", ao.data, kp * 8, bh2.data, kp * 4, mc.data, 0, T.bool(False)) gg[0] = T.shfl_sync(mc[0], 0) # G[a][c] sits in lane 4 a + c // 2, register c % 2 gg[1] = T.shfl_sync(mc[0], 4) gg[2] = T.shfl_sync(mc[1], 4) gg[3] = T.shfl_sync(mc[0], 8) gg[4] = T.shfl_sync(mc[1], 8) gg[5] = T.shfl_sync(mc[0], 9) gg[6] = T.shfl_sync(mc[0], 12) gg[7] = T.shfl_sync(mc[1], 12) gg[8] = T.shfl_sync(mc[0], 13) gg[9] = T.shfl_sync(mc[1], 13) ll[0] = T.rsqrt(T.max(gg[0], 1e-30)) # reciprocals of the diagonal (1/L_jj) ll[1] = gg[1] * ll[0] ll[2] = T.rsqrt(T.max(T.max(gg[2] - ll[1] * ll[1], 1e-6 * gg[2]), 1e-30)) ll[3] = gg[3] * ll[0] ll[4] = (gg[4] - ll[3] * ll[1]) * ll[2] ll[5] = T.rsqrt(T.max(T.max(gg[5] - ll[3] * ll[3] - ll[4] * ll[4], 1e-6 * gg[5]), 1e-30)) ll[6] = gg[6] * ll[0] ll[7] = (gg[7] - ll[6] * ll[1]) * ll[2] ll[8] = (gg[8] - ll[6] * ll[3] - ll[7] * ll[4]) * ll[5] ll[9] = T.rsqrt(T.max(T.max(gg[9] - ll[6] * ll[6] - ll[7] * ll[7] - ll[8] * ll[8], 1e-6 * gg[9]), 1e-30)) pr[0] = pr[0] * ll[0] pr[1] = (pr[1] - pr[0] * ll[1]) * ll[2] pr[2] = (pr[2] - pr[0] * ll[3] - pr[1] * ll[4]) * ll[5] pr[3] = (pr[3] - pr[0] * ll[6] - pr[1] * ll[7] - pr[2] * ll[8]) * ll[9] for r in T.unroll(R): St_s[g, t, K + r] = T.cast(pr[r], F16) # the fp16 Q tile (pass 2's Gram operand, then the next product) T.sync_threads(bar, 128) if PROFIT: tk2[4] = T.call_extern("int64", "clock64") tacc[3] += T.cast(tk2[4] - tk2[3], F32) if PROF: tk[3] = T.call_extern("int64", "clock64") # ---- the fp16 factors (U^T[r][c] = P16[c][r], Q^T[r][v] = Q16[v][r]) to global and, after a barrier, -U^T over the P/Q planes for the residual mma ---- for r in T.unroll(R): uh[r] = UQn1[g, P16OFF + t * 8 + r] qh[r] = St_s[g, t, K + r] UQ[slot, hv, r, t] = uh[r] UQ[slot, hv, r, K + t] = qh[r] T.sync_threads(bar, 128) for r in T.unroll(R): UQn_s[g, r, t] = -uh[r] # -U^T: the residual mma accumulates the subtraction for r in T.unroll(8 - R): UQn_s[g, R + r, t] = T.cast(0.0, F16) T.sync_threads(bar, 128) # ---- residual E = fp16(B) + pn Q_old U_old^T - Q U^T in two halves (fp16 mma, fp32 accumulate), |E| column partials, E over B in St_s ---- for mt in T.unroll(2): # A[m=v][k=r] = Q_old^T[r][v]: x2.trans -> (m 0-7), (m 8-15) of k-rows 0-7 T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(UQ_s[s, (K + w * 32 + mt * 16 + 8 * ((lane // 8) % 2)) // 64, lane % 8, (K + w * 32 + mt * 16 + 8 * ((lane // 8) % 2)) % 64], "r", extent=4), T.access_ptr(afrag2[mt * 4], "w", extent=4)) # A[m=v][k=r] = Q16[v][r] (row-major [m][k], k 0-7): x2 (non-trans) -> rows 0-7, rows 8-15 T.ptx_ldmatrix(T.bool(False), 2, T.access_ptr(St_s[g, w * 32 + mt * 16 + 8 * ((lane // 8) % 2) + lane % 8, K], "r", extent=4), T.access_ptr(afrag3[mt * 4], "w", extent=4)) if KEX: # the ring buffer's A fragments again (the rebuild's registers are long dead) for i in T.unroll(4): svr_l[i] = skv_s[s, K + w * 32 + (i // 2) * 16 + (i % 2) * 8 + lane // 4] * (1.0 / QMAX) for mt in T.unroll(2): T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Ub_s[s, (w * 32 + mt * 16) // (V // 2), (lane % 8) + 8 * (lane // 16), (w * 32 + mt * 16) % (V // 2) + 8 * ((lane // 8) % 2)], "r", extent=8), T.access_ptr(afrag[mt * 8], "w", extent=8)) for h in T.unroll(2): if KEX: # KEX: B is rebuilt here in fp32 from the stage instead of read back from the fp16 St_s (its 2^-11 rounding cost +1 % # reconstruction error): E = (deq(S0) + Q_old U_old^T) pc + ring buffer - Q U^T, every term accumulated in fp32 for ntl in T.unroll(8): for c1 in T.vectorized(2): skl[ntl * 2 + c1] = skv_s[s, (h * 8 + ntl) * 8 + q4 * 2 + c1] for rr in T.unroll(4): v0 = w * 32 + (rr // 2) * 16 + (rr % 2) * 8 + r0 for jcl in T.unroll(4): jc = h * 4 + jcl for nn in T.unroll(2): qw[nn] = Qs[s, v0, jc * 4 + nn * 2 + q4 // 2] for nn in T.unroll(2): ntl = jcl * 2 + nn pkw = qw[nn] ^ T.int32(-2139062144) for c1 in T.unroll(2): x = T.reinterpret(F32, T.call_extern(I32, "__byte_perm", pkw, T.int32(0x4B000000), T.int32(0x7650) + (q4 % 2) * 2 + c1)) - 8388736.0 cfrag[(ntl * 4 + rr) * 2 + c1] = x * svr_l[rr] * skl[ntl * 2 + c1] # E = pn (B / pn + Q_old U_old^T) - Q U^T: the accumulator starts from B (scaled by 1 / pn on GDN, loaded with ldmatrix.x4 in the # fragment layout), so both rank-R mma chains accumulate straight into it and there is no elementwise add afterwards for mt in (T.unroll(0) if KEX else T.unroll(2)): for npl in T.unroll(4): T.ptx_ldmatrix(T.bool(False), 4, T.access_ptr(St_s[g, w * 32 + mt * 16 + 8 * ((lane // 8) % 2) + lane % 8, (h * 8 + npl * 2 + lane // 16) * 8], "r", extent=8), T.access_ptr(efrag[0], "w", extent=8)) for jj in T.unroll(4): for c1 in T.unroll(2): if KDA: cfrag[((npl * 2 + jj // 2) * 4 + mt * 2 + jj % 2) * 2 + c1] = T.cast(efrag[jj * 2 + c1], F32) else: cfrag[((npl * 2 + jj // 2) * 4 + mt * 2 + jj % 2) * 2 + c1] = T.cast(efrag[jj * 2 + c1], F32) * rpn for npl in T.unroll(4): n0 = (h * 4 + npl) * 16 T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(UQ_s[s, n0 // 64, lane % 8, n0 % 64 + 8 * ((lane // 8) % 2)], "r", extent=4), T.access_ptr(bfr[(npl % 2) * 4], "w", extent=4)) for mt in T.unroll(2): for nn in T.unroll(2): T.ptx_mma("float32", "m16n8k8", "row", "col", "fp16", "fp16", "fp32", afrag2.data, mt * 4, bfr.data, (npl % 2) * 4 + nn * 2, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) if not KDA: for i in T.unroll(64): cfrag[i] = cfrag[i] * pn if KEX: for nt in T.unroll(8): for c1 in T.unroll(2): pcl[nt * 2 + c1] = pc_s[g, h * 64 + nt * 8 + 2 * q4 + c1] for mr in T.unroll(4): for c1 in T.unroll(2): cfrag[(nt * 4 + mr) * 2 + c1] = cfrag[(nt * 4 + mr) * 2 + c1] * pcl[nt * 2 + c1] for npl in T.unroll(4): # + the ring buffer (bf16 hi + lo decayed keys, exact bf16 values) n0 = (h * 4 + npl) * 16 bo = (npl % 2) * 16 T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Kb_s[s, n0 // (K // 2), (lane % 8) + 8 * ((lane // 8) % 2), n0 % (K // 2) + 8 * (lane // 16)], "r", extent=8), T.access_ptr(bfrag[bo], "w", extent=8)) T.ptx_ldmatrix(T.bool(True), 4, T.access_ptr(Kbl_s[s, n0 // (K // 2), (lane % 8) + 8 * ((lane // 8) % 2), n0 % (K // 2) + 8 * (lane // 16)], "r", extent=8), T.access_ptr(bfrag[bo + 8], "w", extent=8)) for mt in T.unroll(2): for nn in T.unroll(2): T.ptx_mma("float32", "m16n8k16", "row", "col", "bf16", "bf16", "fp32", afrag.data, mt * 8, bfrag.data, bo + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) T.ptx_mma("float32", "m16n8k16", "row", "col", "bf16", "bf16", "fp32", afrag.data, mt * 8, bfrag.data, bo + 8 + nn * 4, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) for npl in T.unroll(4): n0 = (h * 4 + npl) * 16 T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(UQn_s[g, lane % 8, n0 + 8 * ((lane // 8) % 2)], "r", extent=4), T.access_ptr(bfr[(npl % 2) * 4], "w", extent=4)) for mt in T.unroll(2): for nn in T.unroll(2): T.ptx_mma("float32", "m16n8k8", "row", "col", "fp16", "fp16", "fp32", afrag3.data, mt * 4, bfr.data, (npl % 2) * 4 + nn * 2, cfrag.data, ((npl * 2 + nn) * 4 + mt * 2) * 2, T.bool(False)) if KEX: # E goes to the fp16 St_s scaled by a power of two per (warp, half): the block maximum lands in [2^14, 2^15), so a tiny # head keeps 11 significant bits instead of sinking into fp16 subnormals (the column sums and the quantisation undo it exactly) emx[0] = 0.0 for i in T.unroll(64): emx[0] = T.max(emx[0], T.abs(cfrag[i])) for i in T.unroll(5): emx[0] = T.max(emx[0], T.shfl_xor(emx[0], 1 << i)) ebx = T.min(T.max(268 - ((T.reinterpret(I32, emx[0]) >> 23) & 255), 2), 252) esc[h] = T.reinterpret(F32, (254 - ebx) << 23) # 2^-e: undoes the scaling emx[0] = T.reinterpret(F32, ebx << 23) # 2^e for i in T.unroll(64): cfrag[i] = cfrag[i] * emx[0] for mt in T.unroll(2): for npl in T.unroll(4): if "noste" not in ABL: for jj in T.unroll(4): for c1 in T.vectorized(2): ph[jj * 2 + c1] = T.cast(cfrag[((npl * 2 + jj // 2) * 4 + mt * 2 + jj % 2) * 2 + c1], F16) T.call_extern("handle", "tl::ptx_stmatrix_m8n8_x4", T.access_ptr(St_s[g, w * 32 + mt * 16 + 8 * ((lane // 8) % 2) + lane % 8, (h * 8 + npl * 2 + lane // 16) * 8], "w", extent=8), T.reinterpret(I32, ph[0:2]), T.reinterpret(I32, ph[2:4]), T.reinterpret(I32, ph[4:6]), T.reinterpret(I32, ph[6:8])) # ---- |E| column partials of this warp's 32 rows on the tensor cores (this half's 8 n-tiles): C[m][c] = sum_v 1 * |E[v][c]| with # A = ones and B = |E| read back from the fp16 E tile (ldmatrix.x2.trans + sign-bit clear), accumulated into the dead cfrag; # every m-row of C holds the sums -- lane rows r0 == 0 are used. No shuffle reduction. ---- T.sync_warp() for i in T.unroll(32): cfrag[i] = 0.0 for ntl in T.unroll(8): for ks in T.unroll(2): T.ptx_ldmatrix(T.bool(True), 2, T.access_ptr(St_s[g, w * 32 + ks * 16 + 8 * ((lane // 8) % 2) + lane % 8, (h * 8 + ntl) * 8], "r", extent=4), T.access_ptr(bh2[ks * 4], "w", extent=4)) for i in T.unroll(4): bhi[i] = T.reinterpret(I32, bh2[i * 2:i * 2 + 2]) & T.int32(0x7FFF7FFF) for ks in T.unroll(2): T.ptx_mma("float32", "m16n8k16", "row", "col", "fp16", "fp16", "fp32", aone.data, 0, bhi.data, ks * 2, cfrag.data, ntl * 4, T.bool(False)) for ntl in T.unroll(8): for c1 in T.unroll(2): if KEX: colp[(h * 8 + ntl) * 2 + c1] = cfrag[ntl * 4 + c1] * esc[h] else: colp[(h * 8 + ntl) * 2 + c1] = cfrag[ntl * 4 + c1] T.sync_threads(bar, 128) # the residual mma operands (UQn_s, UQ_s) are consumed: the partials may take the buffer, the stage is free T.fence_proxy_async() T.mbarrier_arrive(consumed[s]) if r0 == 0: for nt in T.unroll(16): for c1 in T.vectorized(2): cpart[g, w, nt * 8 + q4 * 2 + c1] = colp[nt * 2 + c1] T.sync_threads(bar, 128) if PROF: tk[4] = T.call_extern("int64", "clock64") # (4) column scale s_k (thread per column) from the 4 warp partials, reciprocal to smem msum = (cpart[g, 0, t] + cpart[g, 1, t] + cpart[g, 2, t] + cpart[g, 3, t]) * (1.0 / V) rskn = T.min(T.rsqrt(msum), 1.0 / EPS) # 1/max(sqrt(m), EPS) without the IEEE division slow path Sk[slot, hv, t] = T.max(T.sqrt(msum), EPS) rsk[g, 6, t] = rskn T.sync_threads(bar, 128) if PROF: tk[5] = T.call_extern("int64", "clock64") # (5) row-wise quantisation: warp w rows 32w.., 8 lanes per row (16 columns each), 4 rows per pass; # row max over the 8 lanes, int8 codes packed in registers, one coalesced 16 B store per lane and row for i in T.unroll(4): for c1 in T.vectorized(4): if KEX: # lanes 0-3 hold columns of half 0, lanes 4-7 of half 1 rk[i * 4 + c1] = rsk[g, 6, (lane % 8) * 16 + i * 4 + c1] * T.if_then_else(lane % 8 < 4, esc[0], esc[1]) else: rk[i * 4 + c1] = rsk[g, 6, (lane % 8) * 16 + i * 4 + c1] for i in T.unroll(8): v0 = w * 32 + i * 4 + lane // 8 rmax[0] = 0.0 for hh in T.unroll(2): for c1 in T.vectorized(8): eh[((i % 2) * 2 + hh) * 8 + c1] = St_s[g, v0, (lane % 8) * 16 + hh * 8 + c1] for hh in T.unroll(2): for c1 in T.unroll(8): ex[hh * 8 + c1] = T.cast(eh[((i % 2) * 2 + hh) * 8 + c1], F32) * rk[hh * 8 + c1] rmax[0] = T.max(rmax[0], T.abs(ex[hh * 8 + c1])) rmax[0] = T.max(rmax[0], T.shfl_xor(rmax[0], 1)) rmax[0] = T.max(rmax[0], T.shfl_xor(rmax[0], 2)) rmax[0] = T.max(rmax[0], T.shfl_xor(rmax[0], 4)) rmax[0] = T.max(rmax[0], EPS) rmax[1] = QMAX * T.call_extern(F32, "__frcp_rn", rmax[0]) if lane % 8 == 0: Sv[slot, hv, v0] = rmax[0] if "nocode" not in ABL: for c1 in T.unroll(4): qw[c1] = T.call_extern(I32, "__byte_perm", T.call_extern(I32, "__byte_perm", T.reinterpret(I32, ex[c1 * 4] * rmax[1] + 12582912.0), T.reinterpret(I32, ex[c1 * 4 + 1] * rmax[1] + 12582912.0), T.int32(0x0040)), T.call_extern(I32, "__byte_perm", T.reinterpret(I32, ex[c1 * 4 + 2] * rmax[1] + 12582912.0), T.reinterpret(I32, ex[c1 * 4 + 3] * rmax[1] + 12582912.0), T.int32(0x0040)), T.int32(0x5410)) for c1 in T.vectorized(4): Sq[slot, hv, v0, (lane % 8) * 4 + c1] = qw[c1] if PROF: tk[6] = T.call_extern("int64", "clock64") if t < L: Wbuf[slot, hv, t] = 0.0 if KDA: Pbuf[slot, hv, t] = 0.0 for jj in T.unroll(L): Gbuf[slot, hv, jj, t] = T.reinterpret(BF16, T.cast(1.0, F16)) # entry factors -> fp16 1.0 else: if t == 0: Pbuf[slot, hv, 0] = 1.0 # hcnt stays at L: 'flushed, ring empty' for the step kernels (a reset here would change the due flags other CTAs are still scanning) T.sync_threads(bar, 128) # St_s / UQn_s reuse by the next program of this group if PROF: tk[7] = T.call_extern("int64", "clock64") if not PROFIT: for i in T.unroll(7): tacc[i] += T.cast(tk[i + 1] - tk[i], F32) tacc[8] += 1.0 if PROF and g == 0 and t == 0 and bid < HV: for i in T.unroll(10): Wbuf[0, bid, i] = tacc[i] return flush class Pool: """Caller-owned state pool, the layout the serving engine uses. Every (slot, head) owns a V x K fp32-sized region (64 KB) of one flat buffer; the checkpoint lives in its first 19 KB: bytes [0, 16384) int8 codes, row-major [V, K] bytes [16384, 16896) fp32 column scales s_k [K] bytes [16896, 17408) fp32 row scales s_v [V] bytes [17408, 19456) fp16 Compensator Tokens [R, K + V]: row r = (u[r, :K] | q[r, :V]) Slots are SS fp32 words apart; slot 0 is reserved (idx 0 = "no sequence"). The update buffer (kbuf, ubuf, gbuf, w), the cumulative log-gate pcum and the per-(slot, head) fill counter hcnt are separate dense tensors. A program (slot, head) is due when hcnt == L.""" def __init__(self, num_slots, device="cuda"): NS = self.NS = num_slots self.raw = torch.zeros(NS * SS + OFF, dtype=torch.float32, device=device); st = self.raw.untyped_storage() def view(dtype, shape, strides, off): return torch.empty(0, dtype=dtype, device=device).set_(st).as_strided(shape, strides, off) self.codes = view(torch.int8, (NS, HV, V, K), (4 * SS, 4 * V * K, K, 1), 4 * OFF) self.codes32 = view(torch.int32, (NS, HV, V, K // 4), (SS, V * K, K // 4, 1), OFF) self.codes16 = view(torch.int16, (NS, HV, V, K // 2), (2 * SS, 2 * V * K, K // 2, 1), 2 * OFF) self.s_k = view(torch.float32, (NS, HV, K), (SS, V * K, 1), OFF + V * K // 4) self.s_v = view(torch.float32, (NS, HV, V), (SS, V * K, 1), OFF + V * K // 4 + K) self.s_kv = view(torch.float32, (NS, HV, K + V), (SS, V * K, 1), OFF + V * K // 4) self.uq = view(torch.float16, (NS, HV, R, K + V), (2 * SS, 2 * V * K, K + V, 1), 2 * (OFF + V * K // 4 + K + V)) self.kbuf = torch.zeros(NS, HV, L, K, dtype=torch.bfloat16, device=device) self.ubuf = torch.zeros(NS, HV, L, V, dtype=torch.bfloat16, device=device) self.w = torch.zeros(NS, HV, L, device=device) self.pcum = torch.zeros(NS, HV, K, device=device) self.gbuf = torch.ones(NS, HV, L, K, dtype=torch.float16, device=device) self.hcnt = torch.zeros(NS, HV, dtype=torch.int32, device=device) self.kernel = make_flush_kernel(NS, HV, K, V, L, SS, R=R, ITERS=1, NCG=3, KDA=True, GB16=True, SMS=torch.cuda.get_device_properties(self.raw.device).multi_processor_count) def flush_inplace(pool, idx, q0): """Deployment form. idx: int32 [N] slot per sequence (0 = skip). For every (idx[n], head) with hcnt == L: rebuild the state from the checkpoint and the update buffer, write the new checkpoint over the old one, reset w = 0, pcum = 0, gbuf = 1 (hcnt stays at L: "flushed, empty").""" pool.kernel(pool.codes32, pool.s_k, pool.s_v, pool.s_kv, pool.kbuf, pool.ubuf, pool.w, pool.pcum, pool.gbuf.view(torch.bfloat16), pool.hcnt, pool.codes16, pool.uq, q0, idx) _POOLS = {} def run(codes, s_k, s_v, u, q, kbuf, ubuf, gbuf, w, pcum, q0): B = codes.shape[0] pool = _POOLS.get((B, codes.device)) if pool is None: pool = _POOLS[(B, codes.device)] = Pool(B + 8, codes.device) sl = slice(1, B + 1) pool.codes[sl] = codes; pool.s_k[sl] = s_k; pool.s_v[sl] = s_v; pool.uq[sl, :, :, :K] = u; pool.uq[sl, :, :, K:] = q pool.kbuf[sl] = kbuf; pool.ubuf[sl] = ubuf; pool.gbuf[sl] = gbuf; pool.w[sl] = w; pool.pcum[sl] = pcum pool.hcnt.zero_(); pool.hcnt[sl] = L flush_inplace(pool, torch.arange(1, B + 1, dtype=torch.int32, device=codes.device), q0.contiguous()) return (pool.codes[sl].clone(), pool.s_k[sl].clone(), pool.s_v[sl].clone(), pool.uq[sl, :, :, :K].clone(), pool.uq[sl, :, :, K:].clone())