"""Single-launch, persistent Triton implementation of the hybrid decode motif. All large matrices stay in the original packed int4 representation. The one resident grid executes the entire motif and synchronizes through an acquire / release device counter. No expanded weight matrices or expanded KV cache exist. """ import torch from torch import nn import triton import triton.language as tl class QuantLinear(nn.Module): def __init__(self, k, n, group=128): super().__init__() self.register_buffer('w_q', torch.empty(k // 2, n, dtype=torch.uint8)) self.register_buffer('scales', torch.empty(k // group, n, dtype=torch.bfloat16)) self.register_buffer('zeros', torch.empty(k // group, n, dtype=torch.bfloat16)) class QuantExperts(nn.Module): def __init__(self, e, k, n, group=128): super().__init__() self.register_buffer('w_q', torch.empty(e, k // 2, n, dtype=torch.uint8)) self.register_buffer('scales', torch.empty(e, k // group, n, dtype=torch.bfloat16)) self.register_buffer('zeros', torch.empty(e, k // group, n, dtype=torch.bfloat16)) class KDA(nn.Module): def __init__(self, c): super().__init__() d, a = c.hidden, c.kda_heads * c.kda_head_dim self.q_proj = QuantLinear(d, a) self.k_proj = QuantLinear(d, a) self.v_proj = QuantLinear(d, a) self.g_proj = QuantLinear(d, a) self.o_proj = QuantLinear(a, d) self.beta_proj = nn.Linear(d, c.kda_heads, bias=False, dtype=c.dtype) self.conv_w = nn.Parameter(torch.empty(3, a, c.short_conv, dtype=c.dtype)) class MLA(nn.Module): def __init__(self, c): super().__init__() self.q_proj = QuantLinear(c.hidden, c.mla_heads * (c.qk_nope + c.qk_rope)) self.kv_a = QuantLinear(c.hidden, c.kv_lora + c.qk_rope) self.kv_b = QuantLinear(c.kv_lora, c.mla_heads * (c.qk_nope + c.v_head)) self.o_proj = QuantLinear(c.mla_heads * c.v_head, c.hidden) class MoE(nn.Module): def __init__(self, c): super().__init__() self.router = nn.Linear(c.hidden, c.n_experts, bias=False, dtype=c.dtype) self.gate = QuantExperts(c.n_experts, c.hidden, c.moe_inter) self.up = QuantExperts(c.n_experts, c.hidden, c.moe_inter) self.down = QuantExperts(c.n_experts, c.moe_inter, c.hidden) self.s_gate = QuantExperts(c.n_shared, c.hidden, c.moe_inter) self.s_up = QuantExperts(c.n_shared, c.hidden, c.moe_inter) self.s_down = QuantExperts(c.n_shared, c.moe_inter, c.hidden) class Block(nn.Module): def __init__(self, c, kind): super().__init__() self.attn_norm = nn.Parameter(torch.ones(c.hidden, dtype=c.dtype)) self.moe_norm = nn.Parameter(torch.ones(c.hidden, dtype=c.dtype)) self.attn = KDA(c) if kind == 'K' else MLA(c) self.moe = MoE(c) @triton.jit def _barrier(Counter, GRID: tl.constexpr): # The grid is smaller than the SM count; every CTA can remain resident. tl.debug_barrier() ticket = tl.atomic_add(Counter, 1, sem='acq_rel', scope='gpu') target = (ticket // GRID + 1) * GRID while tl.atomic_add(Counter, 0, sem='acquire', scope='gpu') < target: pass tl.debug_barrier() @triton.jit def _pointer(W, layer, slot: tl.constexpr, DT: tl.constexpr): return tl.load(W + layer * 40 + slot).to(tl.pointer_type(DT)) @triton.jit def _gemv(X, Q, Scale, Zero, n, K: tl.constexpr, N: tl.constexpr, BN: tl.constexpr): """Stream one group at a time and reduce without materializing any weights.""" kk = tl.arange(0, 128) acc = tl.full((128, BN), 0, tl.float32) for g in range(K // 128): k = g * 128 + kk packed = tl.load(Q + (k[:, None] // 2) * N + n[None, :], n[None, :] < N, 0) iq = (packed >> ((k[:, None] & 1) * 4)) & 15 s = tl.load(Scale + g * N + n, n < N, 0).to(tl.float32) z = tl.load(Zero + g * N + n, n < N, 0).to(tl.float32) wt = ((iq.to(tl.float32) - z[None, :]) * s[None, :]).to(tl.bfloat16).to(tl.float32) x = tl.load(X + k).to(tl.float32) acc += x[:, None] * wt return tl.sum(acc, 0) @triton.jit def _norm(Hidden, X, Weight): d = tl.arange(0, 4096) h = tl.load(Hidden + d, d < 2304, 0).to(tl.float32) w = tl.load(Weight + d, d < 2304, 0).to(tl.float32) inv = tl.rsqrt(tl.sum(h * h, 0) / 2304 + 1.e-6) tl.store(X + d, (h * inv * w).to(tl.bfloat16).to(tl.float32), d < 2304) @triton.jit def _mega(Hin, Hout, W, Counter, Work, S0, CQ0, CK0, CV0, S1, CQ1, CK1, CV1, S2, CQ2, CK2, CV2, Cold, Rold, Cnew, Rnew, L, E: tl.constexpr, GRID: tl.constexpr, BN: tl.constexpr): pid = tl.program_id(0) # Work is fp32 scratch, allocated once per model. Distinct phase lifetimes # allow these arrays to be reused across all four layers. X = Work A = Work + 4096 B = Work + 20480 O = Work + 20544 Idx = Work + 24640 RW = Work + 24672 FF = Work + 24704 FO = Work + 33920 Qlatent = Work + 54656 Parts = Work + 71040 Latent = Work + 137088 if pid == 0: d = tl.arange(0, 4096) tl.store(Hout + d, tl.load(Hin + d, d < 2304, 0), d < 2304) _barrier(Counter, GRID) for layer in range(4): if pid == 0: _norm(Hout, X, _pointer(W, layer, 0, tl.bfloat16)) _barrier(Counter, GRID) if layer < 3: if layer == 0: S, CQ, CK, CV = S0, CQ0, CK0, CV0 elif layer == 1: S, CQ, CK, CV = S1, CQ1, CK1, CV1 else: S, CQ, CK, CV = S2, CQ2, CK2, CV2 for task in range(pid, 4 * (4096 // BN), GRID): which = task // (4096 // BN) n = (task % (4096 // BN)) * BN + tl.arange(0, BN) # Projection slots 2,5,8,11. Runtime selection keeps one body. qp = tl.load(W + layer * 40 + 2 + which * 3).to(tl.pointer_type(tl.uint8)) sp = tl.load(W + layer * 40 + 3 + which * 3).to(tl.pointer_type(tl.bfloat16)) zp = tl.load(W + layer * 40 + 4 + which * 3).to(tl.pointer_type(tl.bfloat16)) val = _gemv(X, qp, sp, zp, n, 2304, 4096, BN).to(tl.bfloat16).to(tl.float32) if which < 3: if which == 0: C = CQ elif which == 1: C = CK else: C = CV cw = _pointer(W, layer, 18, tl.bfloat16) + which * 4096 * 4 c0 = tl.load(C + n).to(tl.float32) c1 = tl.load(C + 4096 + n).to(tl.float32) c2 = tl.load(C + 8192 + n).to(tl.float32) y = c0 * tl.load(cw + n * 4).to(tl.float32) y += c1 * tl.load(cw + n * 4 + 1).to(tl.float32) y += c2 * tl.load(cw + n * 4 + 2).to(tl.float32) y += val * tl.load(cw + n * 4 + 3).to(tl.float32) tl.store(C + n, c1) tl.store(C + 4096 + n, c2) tl.store(C + 8192 + n, val) val = (y / (1 + tl.exp(-y))).to(tl.bfloat16).to(tl.float32) else: val = 1 / (1 + tl.exp(val)) tl.store(A + which * 4096 + n, val) if pid < 32: d = tl.arange(0, 4096) bp = _pointer(W, layer, 17, tl.bfloat16) x = tl.load(X + d, d < 2304, 0) b = tl.load(bp + pid * 2304 + d, d < 2304, 0).to(tl.float32) bv = tl.sum(x * b, 0).to(tl.bfloat16).to(tl.float32) tl.store(B + pid, 1 / (1 + tl.exp(-bv))) _barrier(Counter, GRID) # Independent state columns: each worker owns 32 value channels. for task in range(pid, 128, GRID): head = task // 4 ki = tl.arange(0, 128) vi = (task % 4) * 32 + tl.arange(0, 32) q = tl.load(A + head * 128 + ki) * (128 ** -0.5) k = tl.load(A + 4096 + head * 128 + ki) v = tl.load(A + 8192 + head * 128 + vi) decay = tl.load(A + 12288 + head * 128 + ki) beta = tl.load(B + head) address = S + head * 16384 + ki[:, None] * 128 + vi[None, :] sv = tl.load(address) * decay[:, None] pred = tl.sum(sv * k[:, None], 0) sv += beta * k[:, None] * (v - pred)[None, :] tl.store(address, sv) out = tl.sum(sv * q[:, None], 0).to(tl.bfloat16).to(tl.float32) tl.store(O + head * 128 + vi, out) _barrier(Counter, GRID) else: # Query and current compressed KV projections. for task in range(pid, triton.cdiv(6144, BN) + triton.cdiv(576, BN), GRID): if task < 6144 // BN: n = task * BN + tl.arange(0, BN) val = _gemv(X, _pointer(W, layer, 2, tl.uint8), _pointer(W, layer, 3, tl.bfloat16), _pointer(W, layer, 4, tl.bfloat16), n, 2304, 6144, BN).to(tl.bfloat16).to(tl.float32) tl.store(A + n, val) else: n = (task - 6144 // BN) * BN + tl.arange(0, BN) val = _gemv(X, _pointer(W, layer, 5, tl.uint8), _pointer(W, layer, 6, tl.bfloat16), _pointer(W, layer, 7, tl.bfloat16), n, 2304, 576, BN).to(tl.bfloat16).to(tl.float32) tl.store(B + n, val, n < 576) # Copy the fed caches, including when they came from another model. for t in range(pid * 1024, L * 512, GRID * 1024): ix = t + tl.arange(0, 1024) tl.store(Cnew + ix, tl.load(Cold + ix, ix < L * 512, 0), ix < L * 512) for t in range(pid * 1024, L * 64, GRID * 1024): ix = t + tl.arange(0, 1024) tl.store(Rnew + ix, tl.load(Rold + ix, ix < L * 64, 0), ix < L * 64) _barrier(Counter, GRID) if pid == 0: r = tl.arange(0, 512) tl.store(Cnew + L * 512 + r, tl.load(B + r)) if pid < 32: r = tl.arange(0, 32) angle = L * tl.exp(-tl.log(10000.) * (2 * r / 64.)) co, si = tl.cos(angle), tl.sin(angle) qe = tl.load(A + pid * 192 + 128 + r * 2) qo = tl.load(A + pid * 192 + 129 + r * 2) tl.store(A + pid * 192 + 128 + r * 2, (qe * co - qo * si).to(tl.bfloat16).to(tl.float32)) tl.store(A + pid * 192 + 129 + r * 2, (qo * co + qe * si).to(tl.bfloat16).to(tl.float32)) if pid == 0: ke, ko = tl.load(B + 512 + 2 * r), tl.load(B + 513 + 2 * r) tl.store(Rnew + L * 64 + 2 * r, ke * co - ko * si) tl.store(Rnew + L * 64 + 2 * r + 1, ko * co + ke * si) # Absorb W_k into each head's query, without storing W_k. for task in range(pid, 128, GRID): head = task // 4 k = (task % 4) * 128 + tl.arange(0, 128) n = head * 256 + tl.arange(0, 128) qp = _pointer(W, layer, 8, tl.uint8) sp = _pointer(W, layer, 9, tl.bfloat16) zp = _pointer(W, layer, 10, tl.bfloat16) packed = tl.load(qp + (k[:, None] // 2) * 8192 + n[None, :]) iq = (packed >> ((k[:, None] & 1) * 4)) & 15 sc = tl.load(sp + (k[:, None] // 128) * 8192 + n[None, :]).to(tl.float32) ze = tl.load(zp + (k[:, None] // 128) * 8192 + n[None, :]).to(tl.float32) wt = ((iq.to(tl.float32) - ze) * sc).to(tl.bfloat16).to(tl.float32) q = tl.load(A + head * 192 + tl.arange(0, 128)) ql = tl.sum(wt * q[None, :], 1) tl.store(Qlatent + head * 512 + k, ql) _barrier(Counter, GRID) # Four independent ranges per head, online softmax in latent space. for task in range(pid, 128, GRID): head, split = task // 4, task % 4 d = tl.arange(0, 512) rr = tl.arange(0, 64) ql = tl.load(Qlatent + head * 512 + d) qr = tl.load(A + head * 192 + 128 + rr) m = tl.full((), -float('inf'), tl.float32) den = tl.full((), 0, tl.float32) acc = tl.full((512,), 0, tl.float32) chunks = tl.cdiv(L + 1, 4 * 32) for block in range(chunks): tok = (split * chunks + block) * 32 + tl.arange(0, 32) c = tl.load(Cnew + tok[:, None] * 512 + d[None, :], tok[:, None] <= L, 0).to(tl.float32) kr = tl.load(Rnew + tok[:, None] * 64 + rr[None, :], tok[:, None] <= L, 0).to(tl.float32) score = (tl.sum(c * ql[None, :], 1) + tl.sum(kr * qr[None, :], 1)) * (192 ** -0.5) score = tl.where(tok <= L, score, -float('inf')) nm = tl.maximum(m, tl.max(score, 0)) alpha = tl.exp(m - nm) p = tl.exp(score - nm) acc = acc * alpha + tl.sum(c * p[:, None], 0) den = den * alpha + tl.sum(p, 0) m = nm tl.store(Parts + task * 516 + d, acc) tl.store(Parts + task * 516 + 512, den) tl.store(Parts + task * 516 + 513, m) _barrier(Counter, GRID) if pid < 32: splits = tl.arange(0, 4) d = tl.arange(0, 512) mm = tl.load(Parts + (pid * 4 + splits) * 516 + 513) dd = tl.load(Parts + (pid * 4 + splits) * 516 + 512) factor = tl.exp(mm - tl.max(mm, 0)) av = tl.load(Parts + (pid * 4 + splits[:, None]) * 516 + d[None, :]) out = tl.sum(av * factor[:, None], 0) / tl.sum(dd * factor, 0) tl.store(Latent + pid * 512 + d, out) _barrier(Counter, GRID) for task in range(pid, 4096 // BN, GRID): head = task // (128 // BN) n = head * 256 + 128 + (task % (128 // BN)) * BN + tl.arange(0, BN) val = _gemv(Latent + head * 512, _pointer(W, layer, 8, tl.uint8), _pointer(W, layer, 9, tl.bfloat16), _pointer(W, layer, 10, tl.bfloat16), n, 512, 8192, BN).to(tl.bfloat16).to(tl.float32) tl.store(O + task * BN + tl.arange(0, BN), val) _barrier(Counter, GRID) # Output projection, residual, and second RMS normalization. for task in range(pid, triton.cdiv(2304, BN), GRID): n = task * BN + tl.arange(0, BN) val = _gemv(O, _pointer(W, layer, 14, tl.uint8), _pointer(W, layer, 15, tl.bfloat16), _pointer(W, layer, 16, tl.bfloat16), n, 4096, 2304, BN).to(tl.bfloat16).to(tl.float32) h = tl.load(Hout + n, n < 2304, 0).to(tl.float32) tl.store(Hout + n, h + val, n < 2304) _barrier(Counter, GRID) if pid == 0: _norm(Hout, X, _pointer(W, layer, 1, tl.bfloat16)) _barrier(Counter, GRID) # Tiny bf16 router, rounded like nn.Linear in the oracle. for expert in range(pid, E, GRID): d = tl.arange(0, 4096) x = tl.load(X + d, d < 2304, 0) r = tl.load(_pointer(W, layer, 19, tl.bfloat16) + expert * 2304 + d, d < 2304, 0).to(tl.float32) tl.store(B + expert, tl.sum(x * r, 0).to(tl.bfloat16).to(tl.float32)) _barrier(Counter, GRID) if pid == 0: r = tl.arange(0, triton.next_power_of_2(E)) logits = tl.load(B + r, r < E, -float('inf')) maximum = tl.max(logits, 0) chosen = tl.full((8,), 0, tl.float32) ar = tl.arange(0, 8) for j in tl.static_range(8): best = tl.max(logits, 0) ix = tl.min(tl.where(logits == best, r, 2147483647), 0) tl.store(Idx + j, ix) chosen = tl.where(ar == j, tl.exp(best - maximum), chosen) logits = tl.where(r == ix, -float('inf'), logits) weights = chosen / tl.sum(chosen, 0) * 2.446 tl.store(RW + ar, weights) _barrier(Counter, GRID) for task in range(pid, 9 * (1024 // BN), GRID): selected = task // (1024 // BN) n = (task % (1024 // BN)) * BN + tl.arange(0, BN) if selected < 8: ix = tl.load(Idx + selected).to(tl.int32) qg = _pointer(W, layer, 20, tl.uint8) + ix * (2304 // 2) * 1024 sg = _pointer(W, layer, 21, tl.bfloat16) + ix * 18 * 1024 zg = _pointer(W, layer, 22, tl.bfloat16) + ix * 18 * 1024 qu = _pointer(W, layer, 23, tl.uint8) + ix * (2304 // 2) * 1024 su = _pointer(W, layer, 24, tl.bfloat16) + ix * 18 * 1024 zu = _pointer(W, layer, 25, tl.bfloat16) + ix * 18 * 1024 else: qg, sg, zg = _pointer(W, layer, 29, tl.uint8), _pointer(W, layer, 30, tl.bfloat16), _pointer(W, layer, 31, tl.bfloat16) qu, su, zu = _pointer(W, layer, 32, tl.uint8), _pointer(W, layer, 33, tl.bfloat16), _pointer(W, layer, 34, tl.bfloat16) gv = _gemv(X, qg, sg, zg, n, 2304, 1024, BN) uv = _gemv(X, qu, su, zu, n, 2304, 1024, BN) tl.store(FF + selected * 1024 + n, gv / (1 + tl.exp(-gv)) * uv) _barrier(Counter, GRID) for task in range(pid, 9 * triton.cdiv(2304, BN), GRID): selected = task // triton.cdiv(2304, BN) n = (task % triton.cdiv(2304, BN)) * BN + tl.arange(0, BN) if selected < 8: ix = tl.load(Idx + selected).to(tl.int32) qd = _pointer(W, layer, 26, tl.uint8) + ix * 512 * 2304 sd = _pointer(W, layer, 27, tl.bfloat16) + ix * 8 * 2304 zd = _pointer(W, layer, 28, tl.bfloat16) + ix * 8 * 2304 else: qd, sd, zd = _pointer(W, layer, 35, tl.uint8), _pointer(W, layer, 36, tl.bfloat16), _pointer(W, layer, 37, tl.bfloat16) val = _gemv(FF + selected * 1024, qd, sd, zd, n, 1024, 2304, BN) tl.store(FO + selected * 2304 + n, val, n < 2304) _barrier(Counter, GRID) for task in range(pid, triton.cdiv(2304, 128), GRID): n = task * 128 + tl.arange(0, 128) total = tl.full((128,), 0, tl.float32) for j in tl.static_range(8): total += tl.load(FO + j * 2304 + n, n < 2304, 0) * tl.load(RW + j) total += tl.load(FO + 8 * 2304 + n, n < 2304, 0) total = total.to(tl.bfloat16).to(tl.float32) h = tl.load(Hout + n, n < 2304, 0).to(tl.float32) tl.store(Hout + n, h + total, n < 2304) _barrier(Counter, GRID) 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.register_buffer('_counter', torch.zeros(1, dtype=torch.int32), persistent=False) self.register_buffer('_work', torch.empty(154000, dtype=torch.float32), persistent=False) self.register_buffer('_weights', torch.empty(4, 40, dtype=torch.int64), persistent=False) self._ready = False def _prepare(self): rows = [] for block in self.blocks: a, m = block.attn, block.moe row = [0] * 40 row[0], row[1] = block.attn_norm.data_ptr(), block.moe_norm.data_ptr() def qset(slot, mod): row[slot:slot + 3] = [mod.w_q.data_ptr(), mod.scales.data_ptr(), mod.zeros.data_ptr()] qset(2, a.q_proj) if isinstance(a, KDA): qset(5, a.k_proj) qset(8, a.v_proj) qset(11, a.g_proj) row[17], row[18] = a.beta_proj.weight.data_ptr(), a.conv_w.data_ptr() else: qset(5, a.kv_a) qset(8, a.kv_b) qset(14, a.o_proj) row[19] = m.router.weight.data_ptr() for slot, mod in zip((20, 23, 26, 29, 32, 35), (m.gate, m.up, m.down, m.s_gate, m.s_up, m.s_down)): qset(slot, mod) rows.append(row) self._weights.copy_(torch.tensor(rows, dtype=torch.int64)) self._ready = True def load_state_dict(self, state_dict, strict=True, assign=False): result = super().load_state_dict(state_dict, strict=strict, assign=assign) self._prepare() return result def _apply(self, fn, recurse=True): result = super()._apply(fn, recurse) self._ready = False return result @torch.no_grad() def step(self, hidden, state): if not self._ready: self._prepare() h = torch.empty_like(hidden) oldc, oldr = state[3]['c_kv'], state[3]['k_rope'] length = oldc.shape[0] newc = torch.empty((length + 1, 512), device=hidden.device, dtype=torch.bfloat16) newr = torch.empty((length + 1, 64), device=hidden.device, dtype=torch.bfloat16) _mega[(128,)](hidden, h, self._weights, self._counter, self._work, state[0]['S'], state[0]['cq'], state[0]['ck'], state[0]['cv'], state[1]['S'], state[1]['cq'], state[1]['ck'], state[1]['cv'], state[2]['S'], state[2]['cq'], state[2]['ck'], state[2]['cv'], oldc, oldr, newc, newr, length, E=self.cfg.n_experts, GRID=128, BN=32, num_warps=4, enable_fp_fusion=False) state[3]['c_kv'], state[3]['k_rope'] = newc, newr return h, state