KernelBench mega · B200
Kimi-Linear Decode Claude Opus 4.8
19.35×geomean speedup across shapes
manually audited: clean
harnessclaude
Kernel source (redacted)
"""Fused W4A16 Kimi-Linear hybrid decode (batch=1).
Beats baseline.py by fusing the int4 unpack + per-group asymmetric dequant
directly into the GEMV (the weights are streamed once as int4, never
materialized to bf16), plus the MLA absorb reformulation so the latent KV cache
is never re-expanded through kv_b.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl
EPS = 1.0e-6
GROUP = 128
# --------------------------------------------------------------------------- #
# Fused int4 dequant GEMV (batch=1), split-K, optional expert gather.
# y[e,n] = sum_k x[e_or_0,k] * (unpack(wq)[e,k,n] - zeros[e,k//G,n]) * scales[..]
# Using the group identity:
# y = sum_g scales_g * ( sum_{k in g} x_k*w_k - zeros_g * sum_{k in g} x_k )
# --------------------------------------------------------------------------- #
def _gemv_configs():
cfgs = []
for BN in (32, 64, 128):
for SK in (4, 8, 16):
for nw in (1, 2):
for ns in (2, 3, 4):
cfgs.append(triton.Config(
{"BLOCK_N": BN, "SPLIT_K": SK}, num_warps=nw, num_stages=ns))
return cfgs
@triton.autotune(configs=_gemv_configs(), key=["N", "NG", "Eact"], reset_to_zero=["y_ptr"])
@triton.jit
def _gemv_kernel(
x_ptr, wq_ptr, s_ptr, z_ptr, idx_ptr, y_ptr,
K, N, NG, Eact,
stride_xe, stride_xk,
stride_we, stride_wk, stride_wn,
stride_se, stride_sg, stride_sn,
stride_ye, stride_yn,
GROUP: tl.constexpr, USE_IDX: tl.constexpr,
BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
):
pid_n = tl.program_id(0)
pid_e = tl.program_id(1)
pid_k = tl.program_id(2)
if USE_IDX:
e = tl.load(idx_ptr + pid_e).to(tl.int64)
else:
e = pid_e
n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
nmask = n < N
wq_e = wq_ptr + e * stride_we
s_e = s_ptr + e * stride_se
z_e = z_ptr + e * stride_se
x_e = x_ptr + pid_e * stride_xe
acc = tl.zeros([BLOCK_N], dtype=tl.float32)
HG: tl.constexpr = GROUP // 2
rows = tl.arange(0, HG)
for g in range(pid_k, NG, SPLIT_K):
kp = g * HG + rows
wq = tl.load(wq_e + kp[:, None] * stride_wk + n[None, :] * stride_wn,
mask=nmask[None, :], other=0)
lo = (wq & 0xF).to(tl.float32)
hi = ((wq >> 4) & 0xF).to(tl.float32)
xe_ = tl.load(x_e + (2 * kp) * stride_xk)
xo_ = tl.load(x_e + (2 * kp + 1) * stride_xk)
raw = tl.sum(xe_[:, None] * lo + xo_[:, None] * hi, axis=0)
xs = tl.sum(xe_ + xo_)
s = tl.load(s_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
z = tl.load(z_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
acc += s * (raw - z * xs)
if SPLIT_K == 1:
tl.store(y_ptr + pid_e * stride_ye + n * stride_yn, acc, mask=nmask)
else:
tl.atomic_add(y_ptr + pid_e * stride_ye + n * stride_yn, acc, mask=nmask)
def gemv(x, wq, scales, zeros, idx=None, out=None):
"""x:[K] (shared) or [Eact,K]; wq:[E,K//2,N] or [K//2,N]; -> [Eact,N] fp32."""
if wq.dim() == 2:
wq = wq.unsqueeze(0)
scales = scales.unsqueeze(0)
zeros = zeros.unsqueeze(0)
E, Kp, N = wq.shape
K = Kp * 2
NG = scales.shape[1]
use_idx = idx is not None
Eact = idx.numel() if use_idx else E
if x.dim() == 1:
stride_xe, stride_xk = 0, x.stride(0)
else:
stride_xe, stride_xk = x.stride(0), x.stride(1)
if out is None:
out = torch.empty((Eact, N), device=x.device, dtype=torch.float32)
out.zero_()
grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), Eact, meta["SPLIT_K"])
_gemv_kernel[grid](
x, wq, scales, zeros, idx if use_idx else x, out,
K, N, NG, Eact,
stride_xe, stride_xk,
wq.stride(0), wq.stride(1), wq.stride(2),
scales.stride(0), scales.stride(1), scales.stride(2),
out.stride(0), out.stride(1),
GROUP=GROUP, USE_IDX=use_idx,
)
return out
@triton.autotune(configs=_gemv_configs(), key=["N", "NG", "Eact"], reset_to_zero=["y_ptr"])
@triton.jit
def _gemv_acc_kernel(
x_ptr, wq_ptr, s_ptr, z_ptr, idx_ptr, sc_ptr, y_ptr,
K, N, NG, Eact,
stride_xe, stride_xk,
stride_we, stride_wk, stride_wn,
stride_se, stride_sg, stride_sn,
GROUP: tl.constexpr, BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
):
"""Weighted accumulate GEMV: y[n] += scale[e] * sum_k x[e,k]*w[e,k,n], over all e."""
pid_n = tl.program_id(0)
pid_e = tl.program_id(1)
pid_k = tl.program_id(2)
e = tl.load(idx_ptr + pid_e).to(tl.int64)
sc = tl.load(sc_ptr + pid_e)
n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
nmask = n < N
wq_e = wq_ptr + e * stride_we
s_e = s_ptr + e * stride_se
z_e = z_ptr + e * stride_se
x_e = x_ptr + pid_e * stride_xe
acc = tl.zeros([BLOCK_N], dtype=tl.float32)
HG: tl.constexpr = GROUP // 2
rows = tl.arange(0, HG)
for g in range(pid_k, NG, SPLIT_K):
kp = g * HG + rows
wq = tl.load(wq_e + kp[:, None] * stride_wk + n[None, :] * stride_wn,
mask=nmask[None, :], other=0)
lo = (wq & 0xF).to(tl.float32)
hi = ((wq >> 4) & 0xF).to(tl.float32)
xe_ = tl.load(x_e + (2 * kp) * stride_xk)
xo_ = tl.load(x_e + (2 * kp + 1) * stride_xk)
raw = tl.sum(xe_[:, None] * lo + xo_[:, None] * hi, axis=0)
xs = tl.sum(xe_ + xo_)
s = tl.load(s_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
z = tl.load(z_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
acc += s * (raw - z * xs)
tl.atomic_add(y_ptr + n, sc * acc, mask=nmask)
def gemv_acc(x, wq, scales, zeros, idx, scale, out):
"""x:[Eact,K]; accumulate scale[e]*(x[e]@dequant(wq[idx[e]])) into out[N]."""
E, Kp, N = wq.shape
K = Kp * 2
NG = scales.shape[1]
Eact = idx.numel()
out.zero_()
grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), Eact, meta["SPLIT_K"])
_gemv_acc_kernel[grid](
x, wq, scales, zeros, idx, scale, out,
K, N, NG, Eact,
x.stride(0), x.stride(1),
wq.stride(0), wq.stride(1), wq.stride(2),
scales.stride(0), scales.stride(1), scales.stride(2),
GROUP=GROUP,
)
return out
@triton.jit
def _router_logits_kernel(x_ptr, w_ptr, out_ptr, D, BLOCK_D: tl.constexpr):
"""One program per expert: logit[e] = sum_d x[d] * W[e,d].
x rounded to bf16 to match the reference's bf16 router (keeps top-k selection identical)."""
e = tl.program_id(0)
acc = tl.zeros([BLOCK_D], dtype=tl.float32)
for d0 in range(0, D, BLOCK_D):
dd = d0 + tl.arange(0, BLOCK_D)
dm = dd < D
xv = tl.load(x_ptr + dd, mask=dm, other=0.0).to(tl.bfloat16).to(tl.float32)
acc += xv * tl.load(w_ptr + e * D + dd, mask=dm, other=0.0).to(tl.float32)
tl.store(out_ptr + e, tl.sum(acc, 0))
@triton.jit
def _router_topk_kernel(
logits_ptr, idxg_ptr, idxgu_ptr, scale_ptr,
E, scaling, GATE_OFF,
NACT: tl.constexpr, NS: tl.constexpr,
BLOCK_E: tl.constexpr, BLOCK_J: tl.constexpr,
):
"""top-k(logits) + softmax + gather-index build (one program).
Emits idx_g[NACT+NS], idx_gu[2*(NACT+NS)] (gate then up), scale[NACT+NS]."""
e = tl.arange(0, BLOCK_E)
emask = e < E
logits = tl.where(emask, tl.load(logits_ptr + e, mask=emask, other=0.0), -float("inf"))
jr = tl.arange(0, BLOCK_J)
idx_vec = tl.zeros([BLOCK_J], dtype=tl.int32)
val_vec = tl.full([BLOCK_J], -float("inf"), dtype=tl.float32)
for j in range(NACT):
mx = tl.max(logits, 0)
am = tl.argmax(logits, 0).to(tl.int32)
idx_vec = tl.where(jr == j, am, idx_vec)
val_vec = tl.where(jr == j, mx, val_vec)
logits = tl.where(e == am, -float("inf"), logits)
idx_vec = tl.where(jr == NACT, E, idx_vec) # shared expert at slot NACT
routed = jr < NACT
vmax = tl.max(tl.where(routed, val_vec, -float("inf")), 0)
ev = tl.where(routed, tl.exp(val_vec - vmax), 0.0)
sm = ev / tl.sum(ev, 0) * scaling
scale_vec = tl.where(routed, sm, tl.where(jr == NACT, 1.0, 0.0))
ns_mask = jr < (NACT + NS)
tl.store(idxg_ptr + jr, idx_vec, mask=ns_mask)
tl.store(scale_ptr + jr, scale_vec, mask=ns_mask)
tl.store(idxgu_ptr + jr, idx_vec, mask=ns_mask)
tl.store(idxgu_ptr + (NACT + NS) + jr, idx_vec + GATE_OFF, mask=ns_mask)
def router_topk(x, w, n_active, n_shared, gate_off, scaling, E, logits_buf):
dev = x.device
ns = n_active + n_shared
idx_g = torch.empty(ns, device=dev, dtype=torch.int32)
idx_gu = torch.empty(2 * ns, device=dev, dtype=torch.int32)
scale = torch.empty(ns, device=dev, dtype=torch.float32)
D = x.shape[-1]
_router_logits_kernel[(E,)](x, w, logits_buf, D, BLOCK_D=triton.next_power_of_2(D))
_router_topk_kernel[(1,)](
logits_buf, idx_g, idx_gu, scale, E, scaling, gate_off,
NACT=n_active, NS=n_shared,
BLOCK_E=triton.next_power_of_2(E), BLOCK_J=triton.next_power_of_2(ns),
)
return idx_g, idx_gu, scale
@triton.jit
def _silu_mul_kernel(g_ptr, u_ptr, o_ptr, n_elem, BLOCK: tl.constexpr):
pid = tl.program_id(0)
off = pid * BLOCK + tl.arange(0, BLOCK)
m = off < n_elem
g = tl.load(g_ptr + off, mask=m, other=0.0)
u = tl.load(u_ptr + off, mask=m, other=0.0)
tl.store(o_ptr + off, (g * tl.sigmoid(g)) * u, mask=m)
def silu_mul(g, u):
o = torch.empty_like(g)
n = g.numel()
_silu_mul_kernel[(triton.cdiv(n, 1024),)](g, u, o, n, BLOCK=1024)
return o
@triton.jit
def _kda_conv_kernel(
qkv_ptr, cw_ptr, st_ptr, out_ptr, C,
scw_t, scw_n, scw_k, BLOCK: tl.constexpr,
):
"""Depthwise causal conv (kernel 4) + SiLU on q/k/v, with in-place ring shift.
qkv_ptr: [>=3, C] (rows q,k,v); st_ptr: [3,3,C] window state; out_ptr: [3,C]."""
pid_t = tl.program_id(0)
pid_n = tl.program_id(1)
n = pid_n * BLOCK + tl.arange(0, BLOCK)
m = n < C
val = tl.load(qkv_ptr + pid_t * C + n, mask=m, other=0.0).to(tl.float32)
cw = cw_ptr + pid_t * scw_t + n * scw_n
w0 = tl.load(cw + 0 * scw_k, mask=m, other=0.0).to(tl.float32)
w1 = tl.load(cw + 1 * scw_k, mask=m, other=0.0).to(tl.float32)
w2 = tl.load(cw + 2 * scw_k, mask=m, other=0.0).to(tl.float32)
w3 = tl.load(cw + 3 * scw_k, mask=m, other=0.0).to(tl.float32)
sb = st_ptr + pid_t * 3 * C + n
s0 = tl.load(sb + 0 * C, mask=m, other=0.0).to(tl.float32)
s1 = tl.load(sb + 1 * C, mask=m, other=0.0).to(tl.float32)
s2 = tl.load(sb + 2 * C, mask=m, other=0.0).to(tl.float32)
acc = s0 * w0 + s1 * w1 + s2 * w2 + val * w3
tl.store(out_ptr + pid_t * C + n, acc * tl.sigmoid(acc), mask=m)
tl.store(sb + 0 * C, s1.to(tl.bfloat16), mask=m)
tl.store(sb + 1 * C, s2.to(tl.bfloat16), mask=m)
tl.store(sb + 2 * C, val.to(tl.bfloat16), mask=m)
def kda_conv(qkv, conv_w, state):
"""qkv:[>=3,C]; conv_w:[3,C,4]; state:[3,3,C] bf16 (in place). -> out:[3,C] fp32."""
C = qkv.shape[1]
out = torch.empty((3, C), device=qkv.device, dtype=torch.float32)
grid = (3, triton.cdiv(C, 1024))
_kda_conv_kernel[grid](
qkv, conv_w, state, out, C,
conv_w.stride(0), conv_w.stride(1), conv_w.stride(2), BLOCK=1024,
)
return out
@triton.jit
def _kda_recur_kernel(
S_ptr, q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, o_ptr,
stride_sh, stride_sk, stride_sv, scale,
DK: tl.constexpr, DV: tl.constexpr, BLOCK_V: tl.constexpr,
):
pid_h = tl.program_id(0)
pid_v = tl.program_id(1)
koff = tl.arange(0, DK)
voff = pid_v * BLOCK_V + tl.arange(0, BLOCK_V)
key = tl.load(k_ptr + pid_h * DK + koff)
qry = tl.load(q_ptr + pid_h * DK + koff) * scale
graw = tl.load(g_ptr + pid_h * DK + koff)
decay = tl.sigmoid(-graw) # = exp(-softplus(graw))
beta = tl.load(beta_ptr + pid_h)
val = tl.load(v_ptr + pid_h * DV + voff)
sbase = S_ptr + pid_h * stride_sh + koff[:, None] * stride_sk + voff[None, :] * stride_sv
S = tl.load(sbase)
Sd = S * decay[:, None]
pred = tl.sum(Sd * key[:, None], axis=0)
Snew = Sd + (beta * key)[:, None] * (val - pred)[None, :]
o = tl.sum(Snew * qry[:, None], axis=0)
tl.store(sbase, Snew)
tl.store(o_ptr + pid_h * DV + voff, o)
def kda_recur(S, q, k, v, g_raw, beta, scale, BLOCK_V=64):
"""S:[H,DK,DV] fp32 in place; q,k,v,g_raw:[H,DK]; beta:[H]. q scaled, decay=sigmoid(-g_raw)."""
H, DK, DV = S.shape
o = torch.empty((H, DV), device=S.device, dtype=torch.float32)
grid = (H, triton.cdiv(DV, BLOCK_V))
_kda_recur_kernel[grid](
S, q, k, v, g_raw, beta, o,
S.stride(0), S.stride(1), S.stride(2), scale,
DK=DK, DV=DV, BLOCK_V=BLOCK_V,
)
return o
def _unpack_dequant(wq, scales, zeros, group=GROUP):
"""Dequant a single [K//2,N] int4 weight to [K,N] bf16 (only used for kv_b)."""
Kp, N = wq.shape
wu = torch.empty((2 * Kp, N), dtype=torch.uint8, device=wq.device)
wu[[REDACTED: IP]] = wq & 0xF
wu[[REDACTED: IP]] = (wq >> 4) & 0xF
s = scales.repeat_interleave(group, dim=0)
z = zeros.repeat_interleave(group, dim=0)
return (wu.to(torch.bfloat16) - z) * s
# --------------------------------------------------------------------------- #
# modules (buffer/param names identical to reference for state_dict load)
# --------------------------------------------------------------------------- #
class QuantLinear(nn.Module):
def __init__(self, in_f, out_f, group=GROUP):
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, in_f, out_f, group=GROUP):
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))
@triton.jit
def _rmsnorm_kernel(x_ptr, w_ptr, o_ptr, N: tl.constexpr, eps, BLOCK: tl.constexpr):
off = tl.arange(0, BLOCK)
m = off < N
x = tl.load(x_ptr + off, mask=m, other=0.0).to(tl.float32)
w = tl.load(w_ptr + off, mask=m, other=0.0).to(tl.float32)
r = tl.rsqrt(tl.sum(x * x) / N + eps)
tl.store(o_ptr + off, x * r * w, mask=m)
def _rmsnorm(x, w):
"""Fused RMSNorm; returns fp32 (ready for the dequant-GEMV)."""
N = x.shape[-1]
o = torch.empty(N, device=x.device, dtype=torch.float32)
BLOCK = triton.next_power_of_2(N)
_rmsnorm_kernel[(1,)](x, w, o, N, EPS, BLOCK=BLOCK)
return o
@triton.jit
def _rmsnorm_add_kernel(x_ptr, y_ptr, w_ptr, resid_ptr, o_ptr, N: tl.constexpr, eps, BLOCK: tl.constexpr):
off = tl.arange(0, BLOCK)
m = off < N
# round y to bf16 before the add to match reference's bf16 residual exactly
yb = tl.load(y_ptr + off, mask=m, other=0.0).to(tl.bfloat16).to(tl.float32)
s = tl.load(x_ptr + off, mask=m, other=0.0).to(tl.float32) + yb
sb = s.to(tl.bfloat16)
tl.store(resid_ptr + off, sb, mask=m)
sf = sb.to(tl.float32) # match reference: norm the bf16 residual
r = tl.rsqrt(tl.sum(sf * sf) / N + eps)
w = tl.load(w_ptr + off, mask=m, other=0.0).to(tl.float32)
tl.store(o_ptr + off, sf * r * w, mask=m)
def _rmsnorm_add(x, y, w):
"""resid = x + y (bf16); returns (resid, rmsnorm(resid) fp32)."""
N = x.shape[-1]
resid = torch.empty(N, device=x.device, dtype=torch.bfloat16)
o = torch.empty(N, device=x.device, dtype=torch.float32)
_rmsnorm_add_kernel[(1,)](x, y, w, resid, o, N, EPS, BLOCK=triton.next_power_of_2(N))
return resid, o
def _rope_cossin(pos, dim, theta, device):
inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
ang = pos * inv
return torch.cos(ang), torch.sin(ang)
def _apply_rope(x, cos, sin):
xf = x.float()
even, odd = xf[..., [REDACTED: IP]], xf[..., [REDACTED: IP]]
out = torch.empty_like(xf)
out[..., [REDACTED: IP]] = even * cos - odd * sin
out[..., [REDACTED: IP]] = odd * cos + even * sin
return out.to(x.dtype)
class KDA(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
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)
self.scale = Dk ** -0.5
def _stack(self):
self._wq = torch.stack([self.q_proj.w_q, self.k_proj.w_q,
self.v_proj.w_q, self.g_proj.w_q])
self._s = torch.stack([self.q_proj.scales, self.k_proj.scales,
self.v_proj.scales, self.g_proj.scales])
self._z = torch.stack([self.q_proj.zeros, self.k_proj.zeros,
self.v_proj.zeros, self.g_proj.zeros])
def _conv(self, val, prev, idx):
# val:[C] fp32, prev:[3,C] bf16, conv_w[idx]:[C,4]
win = torch.cat([prev.float(), val[None]], dim=0) # [4,C]
w = self.conv_w[idx].float().transpose(0, 1) # [4,C]
out = F.silu((win * w).sum(0)) # [C]
return out, win[1:].to(prev.dtype)
def step(self, x, st):
H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
if not hasattr(self, "_wq"):
self._stack()
xf = x.float()
qkvg = gemv(xf, self._wq, self._s, self._z) # [4, H*Dk]
q, st["cq"] = self._conv(qkvg[0], st["cq"], 0)
k, st["ck"] = self._conv(qkvg[1], st["ck"], 1)
v, st["cv"] = self._conv(qkvg[2], st["cv"], 2)
q = q.view(H, Dk).contiguous()
k = k.view(H, Dk).contiguous()
v = v.view(H, Dk).contiguous()
g_raw = qkvg[3].view(H, Dk).contiguous()
beta = torch.sigmoid(self.beta_proj(x.to(torch.bfloat16)).float()).contiguous()
o = kda_recur(st["S"], q, k, v, g_raw, beta, self.scale) # updates S in place
of = o.reshape(H * Dk)
return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)
def step_g(self, x, S_buf, conv_state):
H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
if not hasattr(self, "_wq"):
self._stack()
xf = x.float()
qkvg = gemv(xf, self._wq, self._s, self._z) # [4, C]
conv_out = kda_conv(qkvg, self.conv_w, conv_state) # [3, C], state updated
beta = torch.sigmoid(self.beta_proj(x.to(torch.bfloat16)).float())
o = kda_recur(S_buf, conv_out[0], conv_out[1], conv_out[2], qkvg[3], beta, self.scale)
of = o.reshape(H * Dk)
return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)
class MLA(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
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)
self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5
def _prep(self):
cfg = self.cfg
H, nope, vh, lora = cfg.mla_heads, cfg.qk_nope, cfg.v_head, cfg.kv_lora
W = _unpack_dequant(self.kv_b.w_q, self.kv_b.scales, self.kv_b.zeros).view(lora, H, nope + vh)
# per-head absorb matrices, bf16 for tensor-core matmuls
self._Wkn = W[:, :, :nope].permute(1, 0, 2).contiguous() # [H,512,nope]
self._Wvn = W[:, :, nope:].permute(1, 0, 2).contiguous() # [H,512,vh]
dev = self.kv_b.w_q.device
self._inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.qk_rope, 2, device=dev, dtype=torch.float32) / cfg.qk_rope))
# combined q_proj || kv_a (same input x); split output at qsplit
self._qkva_wq = torch.cat([self.q_proj.w_q, self.kv_a.w_q], dim=1)
self._qkva_s = torch.cat([self.q_proj.scales, self.kv_a.scales], dim=1)
self._qkva_z = torch.cat([self.q_proj.zeros, self.kv_a.zeros], dim=1)
self._qsplit = H * (nope + cfg.qk_rope)
def step(self, x, st):
cfg = self.cfg
H = cfg.mla_heads
nope, rope, vh, lora = cfg.qk_nope, cfg.qk_rope, cfg.v_head, cfg.kv_lora
if not hasattr(self, "_Wkn"):
self._prep()
pos = st["c_kv"].shape[0]
xf = x.float()
qkva = gemv(xf, self._qkva_wq, self._qkva_s, self._qkva_z).squeeze(0)
q = qkva[:self._qsplit].view(H, nope + rope)
q_nope = q[:, :nope]
q_rope = q[:, nope:]
kv = qkva[self._qsplit:]
c_kv = kv[:lora].to(torch.bfloat16)
k_rope = kv[lora:]
cos, sin = _rope_cossin(pos, rope, cfg.rope_theta, x.device)
q_rope = _apply_rope(q_rope, cos, sin).to(torch.bfloat16)
k_rope = _apply_rope(k_rope, cos, sin).to(torch.bfloat16)
st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0)
st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)
ckv = st["c_kv"] # [L,512] bf16
krope = st["k_rope"] # [L,64] bf16
# absorb q into latent space (per-head [512,nope] @ [nope] -> [512])
q_absorb = torch.einsum("hdc,hc->hd", self._Wkn, q_nope.to(torch.bfloat16)) # [H,512] bf16
scores = (q_absorb @ ckv.t() + q_rope @ krope.t()).float() * self.scale # [H,L]
p = torch.softmax(scores, dim=-1).to(torch.bfloat16) # [H,L]
pc = p @ ckv # [H,512] bf16
o = torch.einsum("hd,hdc->hc", pc, self._Wvn) # [H,vh] bf16
of = o.reshape(H * vh).float()
return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)
def step_g(self, x, c_kv_buf, k_rope_buf, pos, maxlen):
"""Graph variant: pos is a [1] int64 GPU scalar; cache written in place at pos."""
cfg = self.cfg
H = cfg.mla_heads
nope, rope, vh, lora = cfg.qk_nope, cfg.qk_rope, cfg.v_head, cfg.kv_lora
if not hasattr(self, "_Wkn"):
self._prep()
xf = x.float()
qkva = gemv(xf, self._qkva_wq, self._qkva_s, self._qkva_z).squeeze(0)
q = qkva[:self._qsplit].view(H, nope + rope)
q_nope = q[:, :nope]
q_rope = q[:, nope:]
kv = qkva[self._qsplit:]
c_kv_new = kv[:lora].to(torch.bfloat16)
k_rope_new = kv[lora:]
inv = self._inv
ang = pos.float() * inv # [rope/2]
cos, sin = torch.cos(ang), torch.sin(ang)
q_rope = _apply_rope(q_rope, cos, sin).to(torch.bfloat16)
k_rope_new = _apply_rope(k_rope_new, cos, sin).to(torch.bfloat16)
c_kv_buf.index_copy_(0, pos, c_kv_new[None])
k_rope_buf.index_copy_(0, pos, k_rope_new[None])
q_absorb = torch.einsum("hdc,hc->hd", self._Wkn, q_nope.to(torch.bfloat16)) # [H,512]
scores = (q_absorb @ c_kv_buf.t() + q_rope @ k_rope_buf.t()).float() * self.scale # [H,maxlen]
mask = torch.arange(maxlen, device=x.device) > pos
scores = scores.masked_fill(mask[None, :], float("-inf"))
p = torch.softmax(scores, dim=-1).to(torch.bfloat16)
pc = p @ c_kv_buf
o = torch.einsum("hd,hdc->hc", pc, self._Wvn)
of = o.reshape(H * vh).float()
return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)
class MoE(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
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)
def _prep(self):
dev = self.gate.w_q.device
# combined routed+shared expert weights: index E (=n_experts) is the shared expert
E = self.cfg.n_experts
# combined gate||up weights: rows [gate(E), s_gate(1), up(E), s_up(1)]
self._gu_wq = torch.cat([self.gate.w_q, self.s_gate.w_q, self.up.w_q, self.s_up.w_q], 0)
self._gu_s = torch.cat([self.gate.scales, self.s_gate.scales, self.up.scales, self.s_up.scales], 0)
self._gu_z = torch.cat([self.gate.zeros, self.s_gate.zeros, self.up.zeros, self.s_up.zeros], 0)
self._d_wq = torch.cat([self.down.w_q, self.s_down.w_q], 0)
self._d_s = torch.cat([self.down.scales, self.s_down.scales], 0)
self._d_z = torch.cat([self.down.zeros, self.s_down.zeros], 0)
self._gate_off = E + 1 # up block offset in gu
self._out = torch.empty(self.cfg.hidden, device=dev, dtype=torch.float32)
self._logits = torch.empty(E, device=dev, dtype=torch.float32)
def step(self, x):
cfg = self.cfg
if not hasattr(self, "_gu_wq"):
self._prep()
ns = cfg.n_active + cfg.n_shared
xf = x.float()
idx_g, idx_gu, scale = router_topk(
xf, self.router.weight, cfg.n_active, cfg.n_shared,
self._gate_off, cfg.routed_scaling, cfg.n_experts, self._logits)
gu = gemv(xf, self._gu_wq, self._gu_s, self._gu_z, idx=idx_gu) # [2*ns, m]
h = silu_mul(gu[:ns], gu[ns:]) # [ns, m]
gemv_acc(h, self._d_wq, self._d_s, self._d_z, idx_g, scale, self._out)
return self._out
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 step(self, x, st):
h = x + self.attn.step(_rmsnorm(x, self.attn_norm), st)
return h + self.moe.step(_rmsnorm(h, self.moe_norm))
HEADROOM = 256
class _Bundle:
"""Persistent buffers + captured CUDA graph for one context length."""
__slots__ = ("maxlen", "h_in", "h_out", "pos", "cpos", "S", "conv",
"c_kv", "k_rope", "graph", "mla_idx")
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._graphs = {}
self._use_graph = True
# ---- eager reference path (always correct) ---------------------------- #
def _step_eager(self, hidden, state):
for i, blk in enumerate(self.blocks):
hidden = blk.step(hidden, state[i])
return hidden, state
# ---- graph machinery -------------------------------------------------- #
def _forward_graph(self, b: _Bundle):
ki = 0
resid, add = b.h_in, None # deferred residual: actual stream = resid + add
for i, blk in enumerate(self.blocks):
if add is None:
xn = _rmsnorm(resid, blk.attn_norm)
else:
resid, xn = _rmsnorm_add(resid, add, blk.attn_norm) # fold prev moe-residual add
if blk.kind == "K":
a = blk.attn.step_g(xn, b.S[ki], b.conv[ki])
ki += 1
else:
a = blk.attn.step_g(xn, b.c_kv, b.k_rope, b.pos, b.maxlen)
resid, xn2 = _rmsnorm_add(resid, a, blk.moe_norm) # fold attn-residual add
add = blk.moe.step(xn2)
b.h_out.copy_(resid + add.to(torch.bfloat16))
b.pos.add_(1)
def _make_bundle(self, ctx, device):
cfg = self.cfg
H, Dk = cfg.kda_heads, cfg.kda_head_dim
C = H * Dk
b = _Bundle()
b.maxlen = ctx + HEADROOM
b.h_in = torch.zeros(cfg.hidden, device=device, dtype=torch.bfloat16)
b.h_out = torch.zeros(cfg.hidden, device=device, dtype=torch.bfloat16)
b.pos = torch.zeros(1, device=device, dtype=torch.int64)
nk = sum(1 for k in cfg.pattern if k == "K")
b.S = [torch.zeros(H, Dk, Dk, device=device, dtype=torch.float32) for _ in range(nk)]
b.conv = [torch.zeros(3, 3, C, device=device, dtype=torch.bfloat16) for _ in range(nk)]
b.c_kv = torch.zeros(b.maxlen, cfg.kv_lora, device=device, dtype=torch.bfloat16)
b.k_rope = torch.zeros(b.maxlen, cfg.qk_rope, device=device, dtype=torch.bfloat16)
b.mla_idx = cfg.pattern.index("M")
return b
def _capture(self, ctx, device):
b = self._make_bundle(ctx, device)
b.pos.fill_(ctx)
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
for _ in range(3):
self._forward_graph(b)
torch.cuda.current_stream().wait_stream(s)
torch.cuda.synchronize()
b.graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(b.graph):
self._forward_graph(b)
self._graphs[ctx] = b
return b
def _load_state(self, b, state, ctx):
"""Copy a fresh dict state into bundle buffers; rebind dict to buffer views."""
ki = 0
for i, blk in enumerate(self.blocks):
st = state[i]
if blk.kind == "K":
b.S[ki].copy_(st["S"])
b.conv[ki][0].copy_(st["cq"]); b.conv[ki][1].copy_(st["ck"]); b.conv[ki][2].copy_(st["cv"])
st["S"] = b.S[ki]
st["cq"] = b.conv[ki][0]; st["ck"] = b.conv[ki][1]; st["cv"] = b.conv[ki][2]
ki += 1
else:
b.c_kv[:ctx].copy_(st["c_kv"]); b.k_rope[:ctx].copy_(st["k_rope"])
st["c_kv"] = b.c_kv[:ctx]; st["k_rope"] = b.k_rope[:ctx]
b.pos.fill_(ctx)
b.cpos = ctx
def step(self, hidden, state):
if not self._use_graph or not hidden.is_cuda:
return self._step_eager(hidden, state)
mla_idx = self.cfg.pattern.index("M")
ckv = state[mla_idx]["c_kv"]
try:
b = self._cur if getattr(self, "_cur", None) is not None else None
is_cont = (b is not None and ckv.data_ptr() == b.c_kv.data_ptr())
if is_cont and b.cpos + 1 >= b.maxlen: # cache overflow: drop to eager
self._use_graph = False
return self._step_eager(hidden, state)
if not is_cont:
ctx = ckv.shape[0]
b = self._graphs.get(ctx) or self._capture(ctx, hidden.device)
self._load_state(b, state, ctx)
self._cur = b
b.h_in.copy_(hidden)
b.graph.replay()
b.cpos += 1
p = b.cpos
state[mla_idx]["c_kv"] = b.c_kv[:p]
state[mla_idx]["k_rope"] = b.k_rope[:p]
return b.h_out, state
except Exception:
self._use_graph = False
self._cur = None
return self._step_eager(hidden, state)
20260618_170335_claude_claude-opus-4-8_02_kimi_linear_decode