"""Kimi-Linear W4A16 hybrid decode unit -- single-launch Triton megakernel. One decode step == ONE persistent-grid @triton.jit kernel launch. Inside that launch, a fixed sequence of stages runs the entire per-token forward -- fused int4 dequant-GEMVs for every projection and MoE expert, the short causal conv, the KDA gated delta-rule state update, absorbed-MLA latent-cache attention (including the cache append), the MoE router/top-8 and expert accumulation, both RMSNorms and all residual adds -- synchronized by software global barriers (an atomic counter with a volatile spin + nanosleep backoff) between stages. Weights are streamed exactly once per step in int4; no bf16 weight matrix is ever materialized. Fused dequant-GEMV: per 128-group, the activation chunk is quantized to two int8 levels (value + quantization residual, so the fused product is exact to ~1e-5) once per stage into a private per-program scratch slice; each GEMV task then multiplies the packed int4 nibbles against those int8 rows on tensor cores (imma), and rescales with the group's (scale, zero) on the fly: y[n] += s*(xs1*d1 + xs2*d2) - s*z*(xs1*sum(q1) + xs2*sum(q2)). MLA runs in the absorbed form: q_eff = W_kb_k^T q_nope is computed once (32x512), scores are taken directly against the compressed latent cache with a two-pass softmax (scores+max, then exp+weighted latent sum), and the value projection is folded back through W_kb_v -- so attention touches only the 512-d latents plus 64-d rope keys, never a materialized per-token K/V. Numerics mirror the reference closely: the residual stream is bf16-rounded at every residual add exactly like the reference (h = bf16(x + bf16(attn)); x' = bf16(h + bf16(moe))) via committed buffers XA/XB plus f32 accumulators AACC/MACC (consumers between accumulation and commit derive the rounded value on the fly), and the router logits are bf16-rounded before softmax/top-8 so expert selection matches the reference. The MLA latent cache lives in a preallocated internal buffer; the incoming state's cache is ingested by the kernel itself (stage 0) on the first step of a run, and the returned state is a view of that buffer, so step() never runs any extra kernel or copy. KDA state (S, conv windows) is updated in place. Grid co-residency (required for the in-kernel barriers) is verified at load time by a bounded-spin probe launch; the grid configuration falls back automatically if the target GPU cannot co-schedule the preferred config. """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl OP_TYPE = "kimi_linear_w4a16_decode" HARDWARE_REQUIRED = ["RTX_PRO_6000"] GROUP_SIZE = 128 SCR_SIZE = 83360 PRIV_F = 2592 # per-program f32 scratch stride PRIV_Q = 18432 # per-program int8 scratch stride # ------------------------------------------------------------------------- # # kernel constants (scratch layout offsets in f32 elements, model dims) # ------------------------------------------------------------------------- # EPS = tl.constexpr(1.0e-6) OFF_XA = tl.constexpr(0) # 2304 committed residual (layer input) OFF_QRAW = tl.constexpr(2304) # 16384 raw GEMV accumulators OFF_QKVG = tl.constexpr(18688) # 16384 conv'd q,k,v + decay gate OFF_BETA = tl.constexpr(35072) # 32 beta logits OFF_OATT = tl.constexpr(35104) # 4096 attention output vector OFF_GU = tl.constexpr(39200) # 9*2048 gate/up expert outputs OFF_CTX = tl.constexpr(57632) # 32*512 MLA context accumulator OFF_GMAX = tl.constexpr(74016) # 32 per-head score max OFF_DEN = tl.constexpr(74048) # 32 per-head softmax denominator OFF_XB = tl.constexpr(74080) # 2304 committed residual (post-attn h) OFF_AACC = tl.constexpr(76384) # 2304 attention o_proj accumulator OFF_MACC = tl.constexpr(78688) # 2304 MoE output accumulator OFF_ZERO = tl.constexpr(80992) # 2304 always-zero pad OFF_LOGITS = tl.constexpr(83296) # 64 router logits accumulator # private per-program scratch layout PQ2 = tl.constexpr(9216) # offset of second-level quant (int8) PXS1 = tl.constexpr(0) # [72] group scales, level 1 PXS2 = tl.constexpr(72) # [72] group scales, level 2 PSQ1 = tl.constexpr(144) # [72] group int sums, level 1 PSQ2 = tl.constexpr(216) # [72] group int sums, level 2 PXN = tl.constexpr(288) # [2304] normed activation (f32 of bf16) PSF = tl.constexpr(2592) PSQ = tl.constexpr(18432) D = tl.constexpr(2304) GD = tl.constexpr(18) # D // 128 NQKVG = tl.constexpr(16384) NMQ = tl.constexpr(6720) # MLA q (6144) + kv_a (576) KDA_SCALE = tl.constexpr(0.08838834764831845) # 128 ** -0.5 MLA_SCALE = tl.constexpr(0.07216878364870323) # 192 ** -0.5 SPIN_CAP = tl.constexpr(2000000) # ------------------------------------------------------------------------- # # device helpers # ------------------------------------------------------------------------- # @triton.jit def _bar(BAR, ABORT, TIMES, phase, sid, NPROG: tl.constexpr): tl.debug_barrier() tl.atomic_add(BAR, 1, sem="acq_rel", scope="gpu") phase += NPROG n = 0 done = tl.load(BAR, volatile=True) >= phase while (done == 0) & (n < SPIN_CAP): n += 1 if n > 32: tl.inline_asm_elementwise( "nanosleep.u32 $1; mov.u32 $0, 0;", "=r,r", [tl.full((), 1024, tl.int32)], dtype=tl.int32, is_pure=False, pack=1) done = tl.load(BAR, volatile=True) >= phase if done == 0: tl.atomic_xchg(ABORT, 1) tl.atomic_add(BAR, 0, sem="acquire", scope="gpu") if tl.program_id(0) == 0: tl.store(TIMES + sid, tl.extra.cuda.globaltimer()) tl.debug_barrier() return phase, sid + 1 @triton.jit def _rstd(SCR, offa, offb): ss = 0.0 for part in tl.static_range(2): offs = part * 2048 + tl.arange(0, 2048) m = offs < D a = tl.load(SCR + offa + offs, mask=m, other=0.0, cache_modifier=".cg") b = tl.load(SCR + offb + offs, mask=m, other=0.0, cache_modifier=".cg") v = (a + b.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32) ss += tl.sum(v * v) return tl.rsqrt(ss / D + EPS) @triton.jit def _pick3(a, b, c, i): if i == 0: r = a elif i == 1: r = b else: r = c return r @triton.jit def _prequant(SCR, XQP, XSP, offa, offb, rstd, NRM, MODE: tl.constexpr): """Quantize this stage's activation vector once per program (two-level int8 per 128-group) into private scratch. MODE 0: rmsnorm-ed derived residual (18 groups, also stores the normed f32 vector); MODE 1: bf16 of OATT (32 groups); MODE 2: silu(gate)*up from GU (72 groups).""" pid = tl.program_id(0) qpb = XQP + pid.to(tl.int64) * PSQ psb = XSP + pid.to(tl.int64) * PSF r2048 = tl.arange(0, 2048) r64 = tl.arange(0, 64) if MODE == 0: PARTS: tl.constexpr = 2 NG: tl.constexpr = 18 elif MODE == 1: PARTS: tl.constexpr = 2 NG: tl.constexpr = 32 else: PARTS: tl.constexpr = 5 NG: tl.constexpr = 72 for part in tl.static_range(PARTS): idx = part * 2048 + r2048 if MODE == 0: m = idx < D a = tl.load(SCR + offa + idx, mask=m, other=0.0, cache_modifier=".cg") b = tl.load(SCR + offb + idx, mask=m, other=0.0, cache_modifier=".cg") x2 = (a + b.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32) nw = tl.load(NRM + idx, mask=m, other=0.0) xn = (x2 * rstd * nw).to(tl.bfloat16).to(tl.float32) tl.store(psb + PXN + idx, xn, mask=m) elif MODE == 1: xn = tl.load(SCR + offa + idx, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) else: m = idx < 9216 slot = idx // 1024 mm = idx % 1024 ga = tl.load(SCR + OFF_GU + slot * 2048 + mm, mask=m, other=0.0, cache_modifier=".cg") gb = tl.load(SCR + OFF_GU + slot * 2048 + 1024 + mm, mask=m, other=0.0, cache_modifier=".cg") xn = ga * tl.sigmoid(ga) * gb X = tl.reshape(xn, (16, 128)) xs1 = tl.max(tl.abs(X), 1) / 127.0 + 1e-30 q1 = tl.floor(X / xs1[:, None] + 0.5) rr = X - xs1[:, None] * q1 xs2 = tl.max(tl.abs(rr), 1) / 127.0 + 1e-30 q2 = tl.floor(rr / xs2[:, None] + 0.5) sq1 = tl.sum(q1, 1) sq2 = tl.sum(q2, 1) g16 = part * 16 + tl.arange(0, 16) gm = g16 < NG fm = (part * 2048 + tl.arange(0, 2048)) < NG * 128 # natural order; the consumer separates even/odd k via int16 tricks tl.store(qpb + part * 2048 + tl.arange(0, 2048), tl.reshape(q1.to(tl.int8), (2048,)), mask=fm) tl.store(qpb + PQ2 + part * 2048 + tl.arange(0, 2048), tl.reshape(q2.to(tl.int8), (2048,)), mask=fm) tl.store(psb + PXS1 + g16, xs1, mask=gm) tl.store(psb + PXS2 + g16, xs2, mask=gm) tl.store(psb + PSQ1 + g16, sq1, mask=gm) tl.store(psb + PSQ2 + g16, sq2, mask=gm) tl.debug_barrier() @triton.jit(noinline=True) def _gemv_task(WB, SB, ZB, NW, qg0, gc, nb, QP16, psb, OUTP, oncols, wscale, GPT: tl.constexpr): """Fused int4-dequant GEMV partial task: gpt contiguous k-groups x 256 cols, accumulated atomically into OUTP. Activations come pre-quantized (two int8 levels) from private per-program scratch (read as int16 pairs to split even/odd k); the weight tile for the next group is prefetched while the current one is consumed. Rescale: y += s*(xs1*d1 + xs2*d2) - s*z*(xs1*sum(q1) + xs2*sum(q2)) Shared (noinline) across every projection/expert stage of the kernel.""" r64 = tl.arange(0, 64) r16 = tl.arange(0, 16) cols = nb * 256 + tl.arange(0, 256) cmask = cols < oncols acc = tl.zeros([256], tl.float32) g0 = gc * GPT qpb16 = QP16 + (qg0 + g0) * 64 for gi in tl.range(0, GPT, num_stages=1): g = g0 + gi qg = qg0 + g p1 = tl.load(qpb16 + gi * 64 + r64) p2 = tl.load(qpb16 + PQ2 // 2 + gi * 64 + r64) q1e = ((p1 << 8) >> 8).to(tl.int8) q1o = (p1 >> 8).to(tl.int8) q2e = ((p2 << 8) >> 8).to(tl.int8) q2o = (p2 >> 8).to(tl.int8) xs1 = tl.load(psb + PXS1 + qg) xs2 = tl.load(psb + PXS2 + qg) sq1 = tl.load(psb + PSQ1 + qg) sq2 = tl.load(psb + PSQ2 + qg) Xe = tl.where(r16[:, None] == 0, q1e[None, :], 0) \ + tl.where(r16[:, None] == 1, q2e[None, :], 0) Xo = tl.where(r16[:, None] == 0, q1o[None, :], 0) \ + tl.where(r16[:, None] == 1, q2o[None, :], 0) w = tl.load(WB + (g * 64 + r64)[:, None] * NW + cols[None, :], mask=cmask[None, :], other=0) lo = (w & 15).to(tl.int8) hi = (w >> 4).to(tl.int8) dd = tl.dot(Xe, lo, out_dtype=tl.int32) + tl.dot(Xo, hi, out_dtype=tl.int32) d1 = tl.sum(tl.where(r16[:, None] == 0, dd, 0), 0).to(tl.float32) d2 = tl.sum(tl.where(r16[:, None] == 1, dd, 0), 0).to(tl.float32) sv = tl.load(SB + g * NW + cols, mask=cmask, other=0).to(tl.float32) zv = tl.load(ZB + g * NW + cols, mask=cmask, other=0).to(tl.float32) acc += sv * (xs1 * d1 + xs2 * d2) - sv * zv * (xs1 * sq1 + xs2 * sq2) tl.atomic_add(OUTP + cols, acc * wscale, mask=cmask) @triton.jit def _topk(SCR): """Top-8 experts from the LOGITS accumulator (bf16-rounded like ref).""" r64o = tl.arange(0, 64) lg = tl.load(SCR + OFF_LOGITS + r64o, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) mx = tl.max(lg) p = tl.exp(lg - mx) probs = p / tl.sum(p) rank = tl.full([64], 99, tl.int32) for it in range(8): cur = tl.where(rank < 99, -1.0, probs) mxc = tl.max(cur) cand = tl.where(cur == mxc, r64o, 99999) sel = tl.min(cand) rank = tl.where(r64o == sel, it, rank) wsum = tl.sum(tl.where(rank < 8, probs, 0.0)) return rank, probs, wsum @triton.jit def _oproj(pid, SCR, XQP, XQP16, XSP, OW, OS, OZ, WQL, WKL, WVL, WNT, GPTO: tl.constexpr, NPROG: tl.constexpr): """o_proj GEMV (K=4096) into AACC; zeros GU, MACC, LOGITS; for KDA layers (WNT=12) also shifts the short-conv window state.""" _prequant(SCR, XQP, XSP, OFF_OATT, OFF_ZERO, 0.0, SCR, 1) qp16 = XQP16 + pid.to(tl.int64) * (PSQ // 2) psb = XSP + pid.to(tl.int64) * PSF NGC: tl.constexpr = 32 // GPTO NTO = 9 * NGC NT = NTO + 9 + 3 + WNT for t in range(pid, NT, NPROG): if t < NTO: nb = t // NGC gc = t % NGC _gemv_task(OW, OS, OZ, 2304, 0, gc, nb, qp16, psb, SCR + OFF_AACC, 2304, 1.0, GPTO) elif t < NTO + 9: zt = t - NTO zo = zt * 2048 + tl.arange(0, 2048) tl.store(SCR + OFF_GU + zo, tl.zeros([2048], tl.float32)) elif t < NTO + 11: zt = t - NTO - 9 zo = zt * 2048 + tl.arange(0, 2048) zm = zo < D tl.store(SCR + OFF_MACC + zo, tl.zeros([2048], tl.float32), mask=zm) elif t < NTO + 12: tl.store(SCR + OFF_LOGITS + tl.arange(0, 64), tl.zeros([64], tl.float32)) else: wt = t - NTO - 12 p = wt // 4 c4 = (wt % 4) * 1024 + tl.arange(0, 1024) winp = _pick3(WQL, WKL, WVL, p) raw = tl.load(SCR + OFF_QRAW + p * 4096 + c4, cache_modifier=".cg").to(tl.bfloat16) o1 = tl.load(winp + 4096 + c4, cache_modifier=".cg") o2 = tl.load(winp + 8192 + c4, cache_modifier=".cg") tl.store(winp + c4, o1) tl.store(winp + 4096 + c4, o2) tl.store(winp + 8192 + c4, raw) @triton.jit def _moe_logits(pid, SCR, XQP, XQP16, XSP, NRM, RW, GUW, GUS, GUZ, layer, GPT4A: tl.constexpr, NPROG: tl.constexpr): """Router logits (from private normed x) + shared-expert gate/up + zero QRAW. Runs on h1 derived from (XA, AACC).""" rstd = _rstd(SCR, OFF_XA, OFF_AACC) _prequant(SCR, XQP, XSP, OFF_XA, OFF_AACC, rstd, NRM, 0) qp16 = XQP16 + pid.to(tl.int64) * (PSQ // 2) psb = XSP + pid.to(tl.int64) * PSF r128 = tl.arange(0, 128) NGC: tl.constexpr = GD // GPT4A NTL = 24 NTS = NTL + 8 * NGC NT = NTS + 8 cols0 = tl.arange(0, 256) for t in range(pid, NT, NPROG): if t < NTL: ob = t // 3 kc = t % 3 rows = ob * 8 + tl.arange(0, 8) acc8 = tl.zeros([8], tl.float32) for ch in range(6): cc = kc * 6 + ch xn = tl.load(psb + PXN + cc * 128 + r128) w = tl.load(RW + rows[:, None] * D + cc * 128 + r128[None, :]).to(tl.float32) acc8 += tl.sum(w * xn[None, :], 1) tl.atomic_add(SCR + OFF_LOGITS + rows, acc8) elif t < NTS: ts = t - NTL nb = ts // NGC gc = ts % NGC eb = (layer * 65 + 64) + tl.zeros((), dtype=tl.int64) _gemv_task(GUW + eb * (1152 * 2048), GUS + eb * (GD * 2048), GUZ + eb * (GD * 2048), 2048, 0, gc, nb, qp16, psb, SCR + OFF_GU + 8 * 2048, 2048, 1.0, GPT4A) else: zt = t - NTS zo = zt * 2048 + tl.arange(0, 2048) tl.store(SCR + OFF_QRAW + zo, tl.zeros([2048], tl.float32)) @triton.jit def _moe_gu(pid, SCR, XQP16, XSP, GUW, GUS, GUZ, layer, GPT4: tl.constexpr, NPROG: tl.constexpr): """Routed-expert gate/up (private quant from the logits stage is reused); commits h1 into XB.""" rank, probs, wsum = _topk(SCR) qp16 = XQP16 + pid.to(tl.int64) * (PSQ // 2) psb = XSP + pid.to(tl.int64) * PSF r64o = tl.arange(0, 64) NGC: tl.constexpr = GD // GPT4 NTG = 8 * 8 * NGC NT = NTG + 2 cols0 = tl.arange(0, 256) for t in range(pid, NT, NPROG): if t < NTG: slot = t // (8 * NGC) rem = t % (8 * NGC) nb = rem // NGC gc = rem % NGC e = tl.sum(tl.where(rank == slot, r64o, 0)) eb = (layer * 65 + e).to(tl.int64) _gemv_task(GUW + eb * (1152 * 2048), GUS + eb * (GD * 2048), GUZ + eb * (GD * 2048), 2048, 0, gc, nb, qp16, psb, SCR + OFF_GU + slot * 2048, 2048, 1.0, GPT4) else: ct = t - NTG co = ct * 2048 + tl.arange(0, 2048) cm = co < D av = tl.load(SCR + OFF_XA + co, mask=cm, other=0.0, cache_modifier=".cg") bv = tl.load(SCR + OFF_AACC + co, mask=cm, other=0.0, cache_modifier=".cg") h1 = (av + bv.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32) tl.store(SCR + OFF_XB + co, h1, mask=cm) @triton.jit def _moe_down(pid, SCR, XQP, XQP16, XSP, DWW, DWS, DWZ, layer, GPT5: tl.constexpr, NPROG: tl.constexpr): """MoE down projections into MACC; zeros AACC.""" rank, probs, wsum = _topk(SCR) _prequant(SCR, XQP, XSP, 0, 0, 0.0, SCR, 2) qp16 = XQP16 + pid.to(tl.int64) * (PSQ // 2) psb = XSP + pid.to(tl.int64) * PSF r64o = tl.arange(0, 64) NGC: tl.constexpr = 8 // GPT5 NTD = 9 * 9 * NGC NT = NTD + 2 cols0 = tl.arange(0, 256) for t in range(pid, NT, NPROG): if t < NTD: slot = t // (9 * NGC) rem = t % (9 * NGC) nb = rem // NGC gc = rem % NGC e = tl.where(slot < 8, tl.sum(tl.where(rank == slot, r64o, 0)), 64) ws = tl.where(slot < 8, tl.sum(tl.where(rank == slot, probs, 0.0)) / (wsum + 1e-9) * 2.446, 1.0) eb = (layer * 65 + e).to(tl.int64) _gemv_task(DWW + eb * (512 * 2304), DWS + eb * (8 * 2304), DWZ + eb * (8 * 2304), 2304, slot * 8, gc, nb, qp16, psb, SCR + OFF_MACC, 2304, ws, GPT5) else: zt = t - NTD zo = zt * 2048 + tl.arange(0, 2048) zm = zo < D tl.store(SCR + OFF_AACC + zo, tl.zeros([2048], tl.float32), mask=zm) # ------------------------------------------------------------------------- # # the megakernel # ------------------------------------------------------------------------- # @triton.jit(do_not_specialize=["pos", "copy_flag", "bar_base"]) def _mega( pos, copy_flag, bar_base, BAR, ABORT, TIMES, XIN, XOUT, SCR, XQP, XQP16, XSP, QEFFB, QROPEB, SCORE, CKV, KR, CKV32, KR32, CKVS32, KRS32, S0, S1, S2, CQ0, CK0, CV0, CQ1, CK1, CV1, CQ2, CK2, CV2, KW, KS, KZ, KBETA, KCONV, KOW, KOS, KOZ, ANORM, MNORM, ROUTER, GUW, GUS, GUZ, DWW, DWS, DWZ, MQW, MQS, MQZ, KVBW, KVBS, KVBZ, MOW, MOS, MOZ, RINV, NPROG: tl.constexpr, GPT1: tl.constexpr, GPT4: tl.constexpr, GPT4A: tl.constexpr, GPT5: tl.constexpr, GPTO: tl.constexpr, GPTM1: tl.constexpr, ): pid = tl.program_id(0) phase = bar_base sid = 1 if pid == 0: tl.store(TIMES + 0, tl.extra.cuda.globaltimer()) r128 = tl.arange(0, 128) r64 = tl.arange(0, 64) r32 = tl.arange(0, 32) hh = tl.arange(0, 32) qp16 = XQP16 + pid.to(tl.int64) * (PSQ // 2) psb = XSP + pid.to(tl.int64) * PSF # ---------------- stage 0: copy-in ---------------- if copy_flag == 1: NT0 = 2 + (pos * 288 + 2047) // 2048 else: NT0 = 2 for t in range(pid, NT0, NPROG): if t < 2: o2 = t * 2048 + tl.arange(0, 2048) xm = o2 < D xv = tl.load(XIN + o2, mask=xm, other=0).to(tl.float32) tl.store(SCR + OFF_XB + o2, xv, mask=xm) tl.store(SCR + OFF_MACC + o2, tl.zeros([2048], tl.float32), mask=xm) else: idx = (t - 2) * 2048 + tl.arange(0, 2048) ce = pos * 256 m1 = idx < ce v1 = tl.load(CKVS32 + idx, mask=m1, other=0) tl.store(CKV32 + idx, v1, mask=m1) idx2 = idx - ce m2 = (idx >= ce) & (idx2 < pos * 32) v2 = tl.load(KRS32 + idx2, mask=m2, other=0) tl.store(KR32 + idx2, v2, mask=m2) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # ---------------- KDA layers 0..2 ---------------- for layer in range(3): anrm = ANORM + layer * D mnrm = MNORM + layer * D rw = ROUTER + layer * (64 * D) s_ptr = _pick3(S0, S1, S2, layer) kwb = KW + layer.to(tl.int64) * (1152 * NQKVG) ksb = KS + layer * (GD * NQKVG) kzb = KZ + layer * (GD * NQKVG) # --- K1: qkvg projections (+ beta logits, + commit layer input) --- rstd = _rstd(SCR, OFF_XB, OFF_MACC) _prequant(SCR, XQP, XSP, OFF_XB, OFF_MACC, rstd, anrm, 0) NGC1 = GD // GPT1 NT1G = 64 * NGC1 NT1 = NT1G + 4 + 2 for t in range(pid, NT1, NPROG): if t < NT1G: nb = t // NGC1 gc = t % NGC1 _gemv_task(kwb, ksb, kzb, NQKVG, 0, gc, nb, qp16, psb, SCR + OFF_QRAW, NQKVG, 1.0, GPT1) elif t < NT1G + 4: bt = t - NT1G rows = bt * 8 + tl.arange(0, 8) acc8 = tl.zeros([8], tl.float32) for ch in range(GD): xn = tl.load(psb + PXN + ch * 128 + r128) bw = tl.load(KBETA + layer * (32 * D) + rows[:, None] * D + ch * 128 + r128[None, :]).to(tl.float32) acc8 += tl.sum(bw * xn[None, :], 1) tl.store(SCR + OFF_BETA + rows, acc8.to(tl.bfloat16).to(tl.float32)) else: ct = t - NT1G - 4 co = ct * 2048 + tl.arange(0, 2048) cm = co < D av = tl.load(SCR + OFF_XB + co, mask=cm, other=0.0, cache_modifier=".cg") bv = tl.load(SCR + OFF_MACC + co, mask=cm, other=0.0, cache_modifier=".cg") x2 = (av + bv.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32) tl.store(SCR + OFF_XA + co, x2, mask=cm) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- K2: short conv (inline) + gated delta-rule state update --- wq_l = _pick3(CQ0, CQ1, CQ2, layer) wk_l = _pick3(CK0, CK1, CK2, layer) wv_l = _pick3(CV0, CV1, CV2, layer) for t in range(pid, 256, NPROG): h = t // 8 dvb = t % 8 dv = dvb * 16 + tl.arange(0, 16) hc = h * 128 + r128 cwq = KCONV + (layer * 3) * 16384 + h * 128 rq = tl.load(SCR + OFF_QRAW + hc, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) oq = tl.load(wq_l + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwq + r128) \ + tl.load(wq_l + 4096 + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwq + 4096 + r128) \ + tl.load(wq_l + 8192 + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwq + 8192 + r128) \ + rq * tl.load(cwq + 12288 + r128) qq = (oq * tl.sigmoid(oq)).to(tl.bfloat16).to(tl.float32) * KDA_SCALE cwk = KCONV + (layer * 3 + 1) * 16384 + h * 128 rk = tl.load(SCR + OFF_QRAW + 4096 + hc, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) ok_ = tl.load(wk_l + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwk + r128) \ + tl.load(wk_l + 4096 + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwk + 4096 + r128) \ + tl.load(wk_l + 8192 + hc, cache_modifier=".cg").to(tl.float32) * tl.load(cwk + 8192 + r128) \ + rk * tl.load(cwk + 12288 + r128) kk = (ok_ * tl.sigmoid(ok_)).to(tl.bfloat16).to(tl.float32) vc = h * 128 + dv cwv = KCONV + (layer * 3 + 2) * 16384 + vc rv = tl.load(SCR + OFF_QRAW + 8192 + vc, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) ov = tl.load(wv_l + vc, cache_modifier=".cg").to(tl.float32) * tl.load(cwv) \ + tl.load(wv_l + 4096 + vc, cache_modifier=".cg").to(tl.float32) * tl.load(cwv + 4096) \ + tl.load(wv_l + 8192 + vc, cache_modifier=".cg").to(tl.float32) * tl.load(cwv + 8192) \ + rv * tl.load(cwv + 12288) vv = (ov * tl.sigmoid(ov)).to(tl.bfloat16).to(tl.float32) gr = tl.load(SCR + OFF_QRAW + 12288 + hc, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) dec = tl.sigmoid(-gr) sp = s_ptr + h * 16384 S = tl.load(sp + r128[:, None] * 128 + dv[None, :], cache_modifier=".cg") bl = tl.load(SCR + OFF_BETA + h, cache_modifier=".cg") beta = tl.sigmoid(bl) Sd = S * dec[:, None] pred = tl.sum(Sd * kk[:, None], 0) Sn = Sd + (beta * kk)[:, None] * (vv - pred)[None, :] tl.store(sp + r128[:, None] * 128 + dv[None, :], Sn) oo = tl.sum(Sn * qq[:, None], 0) tl.store(SCR + OFF_OATT + h * 128 + dv, oo) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- K3: o_proj into AACC (+ zero GU/MACC/LOGITS, + shift conv windows) --- _oproj(pid, SCR, XQP, XQP16, XSP, KOW + layer.to(tl.int64) * (2048 * 2304), KOS + layer * (32 * 2304), KOZ + layer * (32 * 2304), wq_l, wk_l, wv_l, 12, GPTO, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- K4a: router logits + shared-expert gate/up (+ zero QRAW) --- _moe_logits(pid, SCR, XQP, XQP16, XSP, mnrm, rw, GUW, GUS, GUZ, layer, GPT4A, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- K4b: routed gate/up (+ commit XB) --- _moe_gu(pid, SCR, XQP16, XSP, GUW, GUS, GUZ, layer, GPT4, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- K5: MoE down into MACC (+ zero AACC) --- _moe_down(pid, SCR, XQP, XQP16, XSP, DWW, DWS, DWZ, layer, GPT5, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # ---------------- MLA layer 3 ---------------- anrm = ANORM + 3 * D mnrm = MNORM + 3 * D rw = ROUTER + 3 * (64 * D) L1 = pos + 1 # --- M1: q + kv_a projections (+ zero CTX/GMAX/DEN, + commit XA) --- rstd = _rstd(SCR, OFF_XB, OFF_MACC) _prequant(SCR, XQP, XSP, OFF_XB, OFF_MACC, rstd, anrm, 0) NGCM = GD // GPTM1 NTMG = 27 * NGCM NTM = NTMG + 9 + 2 for t in range(pid, NTM, NPROG): if t < NTMG: nb = t // NGCM gc = t % NGCM _gemv_task(MQW, MQS, MQZ, NMQ, 0, gc, nb, qp16, psb, SCR + OFF_QRAW, NMQ, 1.0, GPTM1) elif t < NTMG + 9: zt = t - NTMG if zt < 8: zo = zt * 2048 + tl.arange(0, 2048) tl.store(SCR + OFF_CTX + zo, tl.zeros([2048], tl.float32)) else: tl.store(SCR + OFF_GMAX + r32, tl.full([32], -1e30, tl.float32)) tl.store(SCR + OFF_DEN + r32, tl.zeros([32], tl.float32)) else: ct = t - NTMG - 9 co = ct * 2048 + tl.arange(0, 2048) cm = co < D av = tl.load(SCR + OFF_XB + co, mask=cm, other=0.0, cache_modifier=".cg") bv = tl.load(SCR + OFF_MACC + co, mask=cm, other=0.0, cache_modifier=".cg") x2 = (av + bv.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16).to(tl.float32) tl.store(SCR + OFF_XA + co, x2, mask=cm) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M2: absorbed q_eff, rope, cache append --- r32b = tl.arange(0, 32) for t in range(pid, 258, NPROG): if t < 256: h = t // 8 cb = (t % 8) // 2 ch = t % 2 qn = tl.load(SCR + OFF_QRAW + h * 192 + r128, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) s1v = tl.load(KVBS + cb * 8192 + h * 256 + r128).to(tl.float32) z1v = tl.load(KVBZ + cb * 8192 + h * 256 + r128).to(tl.float32) sq = s1v * qn t0 = tl.sum(z1v * sq) w1t = tl.load(KVBW + (cb * 64 + ch * 32 + r32b)[:, None] * 8192 + h * 256 + r128[None, :]) qe1 = tl.sum((w1t & 15).to(tl.float32) * sq[None, :], 1) - t0 qo1 = tl.sum((w1t >> 4).to(tl.float32) * sq[None, :], 1) - t0 tl.store(QEFFB + h * 512 + cb * 128 + ch * 64 + 2 * r32b, qe1.to(tl.bfloat16)) tl.store(QEFFB + h * 512 + cb * 128 + ch * 64 + 2 * r32b + 1, qo1.to(tl.bfloat16)) elif t == 256: pq = SCR + OFF_QRAW + hh[:, None] * 192 + 128 + 2 * r32[None, :] qe2 = tl.load(pq, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) qo2 = tl.load(pq + 1, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) inv2 = tl.load(RINV + r32) ang2 = pos * inv2 cs2 = tl.cos(ang2)[None, :] sn2 = tl.sin(ang2)[None, :] re2 = qe2 * cs2 - qo2 * sn2 ro2 = qo2 * cs2 + qe2 * sn2 tl.store(QROPEB + hh[:, None] * 64 + 2 * r32[None, :], re2.to(tl.bfloat16)) tl.store(QROPEB + hh[:, None] * 64 + 2 * r32[None, :] + 1, ro2.to(tl.bfloat16)) elif t == 257: ke3 = tl.load(SCR + OFF_QRAW + 6656 + 2 * r32, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) ko3 = tl.load(SCR + OFF_QRAW + 6656 + 2 * r32 + 1, cache_modifier=".cg").to(tl.bfloat16).to(tl.float32) inv3 = tl.load(RINV + r32) ang3 = pos * inv3 cs3 = tl.cos(ang3) sn3 = tl.sin(ang3) tl.store(KR + pos.to(tl.int64) * 64 + 2 * r32, (ke3 * cs3 - ko3 * sn3).to(tl.bfloat16)) tl.store(KR + pos.to(tl.int64) * 64 + 2 * r32 + 1, (ko3 * cs3 + ke3 * sn3).to(tl.bfloat16)) o512 = tl.arange(0, 512) cv3 = tl.load(SCR + OFF_QRAW + 6144 + o512, cache_modifier=".cg") tl.store(CKV + pos.to(tl.int64) * 512 + o512, cv3.to(tl.bfloat16)) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M3a: attention scores + per-head max --- NLB = (L1 + 31) // 32 for t in range(pid, NLB, NPROG): rows = t * 32 + r32 rm = rows < L1 acc = tl.zeros([32, 32], tl.float32) for cc in range(4): qe = tl.load(QEFFB + hh[:, None] * 512 + cc * 128 + r128[None, :], cache_modifier=".cg") cvt = tl.load(CKV + rows[:, None].to(tl.int64) * 512 + cc * 128 + r128[None, :], mask=rm[:, None], other=0, cache_modifier=".cg") acc = tl.dot(qe, tl.trans(cvt), acc) qr = tl.load(QROPEB + hh[:, None] * 64 + r64[None, :], cache_modifier=".cg") krt = tl.load(KR + rows[:, None].to(tl.int64) * 64 + r64[None, :], mask=rm[:, None], other=0, cache_modifier=".cg") acc = tl.dot(qr, tl.trans(krt), acc) sc = acc * MLA_SCALE scm = tl.where(rm[None, :], sc, -1e30) tl.atomic_max(SCR + OFF_GMAX + hh, tl.max(scm, 1)) tl.store(SCORE + rows[:, None] * 32 + hh[None, :], tl.trans(scm), mask=rm[:, None]) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M3c: softmax-weighted context accumulation --- NSC = (L1 + 127) // 128 for t in range(pid, 4 * NSC + 2, NPROG): if t >= 4 * NSC: zo5 = (t - 4 * NSC) * 2048 + tl.arange(0, 2048) tl.store(SCR + OFF_OATT + zo5, tl.zeros([2048], tl.float32)) else: lsp = t % NSC cb = t // NSC m = tl.load(SCR + OFF_GMAX + hh, cache_modifier=".cg") acc = tl.zeros([32, 128], tl.float32) dsum = tl.zeros([32], tl.float32) for j in range(4): rows = lsp * 128 + j * 32 + r32 rm = rows < L1 sct = tl.load(SCORE + rows[:, None] * 32 + hh[None, :], mask=rm[:, None], other=-1e30, cache_modifier=".cg") p = tl.exp(sct - m[None, :]) p = tl.where(rm[:, None], p, 0.0) if cb == 0: dsum += tl.sum(p, 0) cvt = tl.load(CKV + rows[:, None].to(tl.int64) * 512 + cb * 128 + r128[None, :], mask=rm[:, None], other=0, cache_modifier=".cg") acc = tl.dot(tl.trans(p.to(tl.bfloat16)), cvt, acc) tl.atomic_add(SCR + OFF_CTX + hh[:, None] * 512 + cb * 128 + r128[None, :], acc) if cb == 0: tl.atomic_add(SCR + OFF_DEN + hh, dsum) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M5: absorbed value projection (partials into zeroed OATT) --- for t in range(pid, 256, NPROG): h = t // 8 jb = (t % 8) // 4 cc = t % 4 den = tl.load(SCR + OFF_DEN + h, cache_modifier=".cg") cols = h * 256 + 128 + jb * 64 + r64 cx = tl.load(SCR + OFF_CTX + h * 512 + cc * 128 + r128, cache_modifier=".cg") / den cxe, cxo = tl.split(tl.reshape(cx, (64, 2))) w = tl.load(KVBW + (cc * 64 + r64)[:, None] * 8192 + cols[None, :]) sv = tl.load(KVBS + cc * 8192 + cols).to(tl.float32) zv = tl.load(KVBZ + cc * 8192 + cols).to(tl.float32) du = tl.sum((w & 15).to(tl.float32) * cxe[:, None], 0) \ + tl.sum((w >> 4).to(tl.float32) * cxo[:, None], 0) acc5 = sv * du - sv * zv * tl.sum(cx) tl.atomic_add(SCR + OFF_OATT + h * 128 + jb * 64 + r64, acc5) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M6: o_proj into AACC (+ zero GU/MACC/LOGITS) --- _oproj(pid, SCR, XQP, XQP16, XSP, MOW, MOS, MOZ, CQ0, CK0, CV0, 0, GPTO, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M7a: router logits + shared gate/up (+ zero QRAW) --- _moe_logits(pid, SCR, XQP, XQP16, XSP, mnrm, rw, GUW, GUS, GUZ, 3, GPT4A, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M7b: routed gate/up (+ commit XB) --- _moe_gu(pid, SCR, XQP16, XSP, GUW, GUS, GUZ, 3, GPT4, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- M8: MoE down into MACC (+ zero AACC) --- _moe_down(pid, SCR, XQP, XQP16, XSP, DWW, DWS, DWZ, 3, GPT5, NPROG) phase, sid = _bar(BAR, ABORT, TIMES, phase, sid, NPROG) # --- F: emit hidden = bf16(XB + bf16(MACC)) --- if pid == 0: for part in tl.static_range(2): o2 = part * 2048 + tl.arange(0, 2048) xm = o2 < D av = tl.load(SCR + OFF_XB + o2, mask=xm, other=0.0, cache_modifier=".cg") bv = tl.load(SCR + OFF_MACC + o2, mask=xm, other=0.0, cache_modifier=".cg") xv = (av + bv.to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16) tl.store(XOUT + o2, xv, mask=xm) tl.store(TIMES + sid, tl.extra.cuda.globaltimer()) # ------------------------------------------------------------------------- # # module skeletons (state_dict layout identical to the reference) # ------------------------------------------------------------------------- # class QuantLinear(nn.Module): def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() self.in_f, self.out_f, self.group = in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16)) class QuantExperts(nn.Module): def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) class KDA(nn.Module): def __init__(self, cfg): super().__init__() H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden self.q_proj = QuantLinear(d, H * Dk, cfg.group) self.k_proj = QuantLinear(d, H * Dk, cfg.group) self.v_proj = QuantLinear(d, H * Dk, cfg.group) self.g_proj = QuantLinear(d, H * Dk, cfg.group) self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype) self.conv_w = nn.Parameter(torch.empty(3, H * Dk, cfg.short_conv, dtype=cfg.dtype)) self.o_proj = QuantLinear(H * Dk, d, cfg.group) class MLA(nn.Module): def __init__(self, cfg): super().__init__() H, d = cfg.mla_heads, cfg.hidden self.q_proj = QuantLinear(d, H * (cfg.qk_nope + cfg.qk_rope), cfg.group) self.kv_a = QuantLinear(d, cfg.kv_lora + cfg.qk_rope, cfg.group) self.kv_b = QuantLinear(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head), cfg.group) self.o_proj = QuantLinear(H * cfg.v_head, d, cfg.group) class MoE(nn.Module): def __init__(self, cfg): super().__init__() d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype) self.gate = QuantExperts(E, d, m, cfg.group) self.up = QuantExperts(E, d, m, cfg.group) self.down = QuantExperts(E, m, d, cfg.group) self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group) class Block(nn.Module): def __init__(self, cfg, kind): super().__init__() self.kind = kind self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype)) self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype)) self.attn = KDA(cfg) if kind == "K" else MLA(cfg) self.moe = MoE(cfg) def _pick_gpt(nb, gk, nprog): """Largest group-chunk (task depth <= 6 groups) keeping >= ~0.75 waves.""" for g in (6, 4, 3, 2): if gk % g == 0 and nb * (gk // g) >= (3 * nprog) // 4: return g return 1 class Model(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern) self._ready = False self.register_load_state_dict_post_hook(Model._invalidate) @staticmethod def _invalidate(module, incompatible_keys): module._ready = False # ------------------------------------------------------------------ # def _prepare(self): cfg = self.cfg assert cfg.hidden == 2304 and cfg.kda_heads == 32 and cfg.kda_head_dim == 128 assert cfg.n_experts == 64 and cfg.moe_inter == 1024 and cfg.n_shared == 1 assert cfg.kv_lora == 512 and cfg.qk_nope == 128 and cfg.qk_rope == 64 dev = self.blocks[0].attn_norm.device assert dev.type == "cuda" self._dev = dev blks = self.blocks kda = [b.attn for b in blks[:3]] mla = blks[3].attn self._kw = torch.stack([torch.cat([a.q_proj.w_q, a.k_proj.w_q, a.v_proj.w_q, a.g_proj.w_q], 1) for a in kda]).contiguous() self._ks = torch.stack([torch.cat([a.q_proj.scales, a.k_proj.scales, a.v_proj.scales, a.g_proj.scales], 1) for a in kda]).contiguous() self._kz = torch.stack([torch.cat([a.q_proj.zeros, a.k_proj.zeros, a.v_proj.zeros, a.g_proj.zeros], 1) for a in kda]).contiguous() self._kbeta = torch.stack([a.beta_proj.weight for a in kda]).contiguous() self._kconv = torch.stack([a.conv_w for a in kda]).float().permute(0, 1, 3, 2).contiguous() self._kow = torch.stack([a.o_proj.w_q for a in kda]).contiguous() self._kos = torch.stack([a.o_proj.scales for a in kda]).contiguous() self._koz = torch.stack([a.o_proj.zeros for a in kda]).contiguous() self._anorm = torch.stack([b.attn_norm for b in blks]).float().contiguous() self._mnorm = torch.stack([b.moe_norm for b in blks]).float().contiguous() self._router = torch.stack([b.moe.router.weight for b in blks]).contiguous() self._guw = torch.stack([torch.cat([torch.cat([b.moe.gate.w_q, b.moe.up.w_q], 2), torch.cat([b.moe.s_gate.w_q, b.moe.s_up.w_q], 2)], 0) for b in blks]).contiguous() self._gus = torch.stack([torch.cat([torch.cat([b.moe.gate.scales, b.moe.up.scales], 2), torch.cat([b.moe.s_gate.scales, b.moe.s_up.scales], 2)], 0) for b in blks]).contiguous() self._guz = torch.stack([torch.cat([torch.cat([b.moe.gate.zeros, b.moe.up.zeros], 2), torch.cat([b.moe.s_gate.zeros, b.moe.s_up.zeros], 2)], 0) for b in blks]).contiguous() self._dww = torch.stack([torch.cat([b.moe.down.w_q, b.moe.s_down.w_q], 0) for b in blks]).contiguous() self._dws = torch.stack([torch.cat([b.moe.down.scales, b.moe.s_down.scales], 0) for b in blks]).contiguous() self._dwz = torch.stack([torch.cat([b.moe.down.zeros, b.moe.s_down.zeros], 0) for b in blks]).contiguous() self._mqw = torch.cat([mla.q_proj.w_q, mla.kv_a.w_q], 1).contiguous() self._mqs = torch.cat([mla.q_proj.scales, mla.kv_a.scales], 1).contiguous() self._mqz = torch.cat([mla.q_proj.zeros, mla.kv_a.zeros], 1).contiguous() self._kvbw = mla.kv_b.w_q.contiguous() self._kvbs = mla.kv_b.scales.contiguous() self._kvbz = mla.kv_b.zeros.contiguous() self._mow = mla.o_proj.w_q.contiguous() self._mos = mla.o_proj.scales.contiguous() self._moz = mla.o_proj.zeros.contiguous() self._rinv = (1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.qk_rope, 2, device=dev, dtype=torch.float32) / cfg.qk_rope))).contiguous() self._scr = torch.zeros(SCR_SIZE, dtype=torch.float32, device=dev) self._qeffb = torch.zeros(32 * 512, dtype=torch.bfloat16, device=dev) self._qropeb = torch.zeros(32 * 64, dtype=torch.bfloat16, device=dev) self._xout = torch.zeros(2304, dtype=torch.bfloat16, device=dev) self._abort = torch.zeros(1, dtype=torch.int32, device=dev) self._times = torch.zeros(64, dtype=torch.int64, device=dev) self._ckv = None self._kr = None self._score = None self._cap = 0 self._len = -1 sm = torch.cuda.get_device_properties(dev).multi_processor_count for mult, nwarp, mreg in ((2, 8, 128), (4, 4, 128), (3, 4, 168), (2, 4, 128), (2, 4, None), (1, 8, None)): nprog = sm * mult self._nprog = nprog self._nwarp = nwarp self._mreg = mreg self._gpt1 = _pick_gpt(64, 18, nprog) self._gpt4 = _pick_gpt(64, 18, nprog) self._gpt4a = _pick_gpt(8, 18, nprog) self._gpt5 = _pick_gpt(81, 8, nprog) self._gpto = _pick_gpt(9, 32, nprog) self._gptm1 = _pick_gpt(27, 18, nprog) self._xqp = torch.zeros(nprog * PRIV_Q, dtype=torch.int8, device=dev) self._xsp = torch.zeros(nprog * PRIV_F, dtype=torch.float32, device=dev) self._bar = torch.zeros(1, dtype=torch.int64, device=dev) self._base = 0 if self._probe(): break else: raise RuntimeError("megakernel: no co-resident grid configuration found") self._ready = True def _probe(self): """Dummy launch on throwaway state; verifies grid co-residency and calibrates the per-step barrier count.""" dev = self._dev dummy_L = 384 s = [torch.zeros(32, 128, 128, dtype=torch.float32, device=dev) for _ in range(3)] cw = [torch.zeros(3, 4096, dtype=torch.bfloat16, device=dev) for _ in range(9)] ckv = torch.zeros(dummy_L + 8, 512, dtype=torch.bfloat16, device=dev) kr = torch.zeros(dummy_L + 8, 64, dtype=torch.bfloat16, device=dev) src_c = torch.zeros(dummy_L, 512, dtype=torch.bfloat16, device=dev) src_k = torch.zeros(dummy_L, 64, dtype=torch.bfloat16, device=dev) score = torch.zeros((dummy_L + 8) * 32, dtype=torch.float32, device=dev) xin = torch.zeros(2304, dtype=torch.bfloat16, device=dev) self._abort.zero_() torch.cuda.synchronize() try: self._launch(xin, dummy_L, 1, s, cw, ckv, kr, src_c, src_k, score) torch.cuda.synchronize() except Exception as e: import os if os.environ.get("MEGA_DEBUG"): import traceback traceback.print_exc() self._last_err = e return False if int(self._abort.item()) != 0: return False cnt = int(self._bar.item()) if cnt <= 0 or cnt % self._nprog != 0: return False self._nbar = cnt // self._nprog self._base = cnt self._scr.zero_() torch.cuda.synchronize() return True def _launch(self, xin, pos, copy_flag, s, cw, ckv, kr, src_c, src_k, score): _mega[(self._nprog,)]( pos, copy_flag, self._base, self._bar, self._abort, self._times, xin, self._xout, self._scr, self._xqp, self._xqp.view(torch.int16), self._xsp, self._qeffb, self._qropeb, score, ckv, kr, ckv.view(torch.int32), kr.view(torch.int32), src_c.view(torch.int32), src_k.view(torch.int32), s[0], s[1], s[2], cw[0], cw[1], cw[2], cw[3], cw[4], cw[5], cw[6], cw[7], cw[8], self._kw, self._ks, self._kz, self._kbeta, self._kconv, self._kow, self._kos, self._koz, self._anorm, self._mnorm, self._router, self._guw, self._gus, self._guz, self._dww, self._dws, self._dwz, self._mqw, self._mqs, self._mqz, self._kvbw, self._kvbs, self._kvbz, self._mow, self._mos, self._moz, self._rinv, NPROG=self._nprog, GPT1=self._gpt1, GPT4=self._gpt4, GPT4A=self._gpt4a, GPT5=self._gpt5, GPTO=self._gpto, GPTM1=self._gptm1, num_warps=self._nwarp, **({"maxnreg": self._mreg} if self._mreg else {}), ) # ------------------------------------------------------------------ # def step(self, hidden, state): if not self._ready: self._prepare() st3 = state[3] L = st3["c_kv"].shape[0] copy_flag = 0 src_c, src_k = st3["c_kv"], st3["k_rope"] if (self._ckv is None or L != self._len or st3["c_kv"].data_ptr() != self._ckv.data_ptr() or st3["k_rope"].data_ptr() != self._kr.data_ptr()): if L + 1 > self._cap: self._cap = L + 1 + 256 self._ckv = torch.empty(self._cap, 512, dtype=torch.bfloat16, device=self._dev) self._kr = torch.empty(self._cap, 64, dtype=torch.bfloat16, device=self._dev) self._score = torch.empty(self._cap * 32, dtype=torch.float32, device=self._dev) copy_flag = 1 s = [state[0]["S"], state[1]["S"], state[2]["S"]] cw = [state[0]["cq"], state[0]["ck"], state[0]["cv"], state[1]["cq"], state[1]["ck"], state[1]["cv"], state[2]["cq"], state[2]["ck"], state[2]["cv"]] self._launch(hidden, L, copy_flag, s, cw, self._ckv, self._kr, src_c, src_k, self._score) self._base += self._nbar * self._nprog st3["c_kv"] = self._ckv[:L + 1] st3["k_rope"] = self._kr[:L + 1] self._len = L + 1 return self._xout, state