KernelBench mega · RTX PRO 6000
Kimi-Linear Decode Kimi K3 (1M)
9.79×geomean speedup across shapes
manually audited: clean
Clean, genuine single-launch persistent Triton megakernel at 9.7846x geomean decode speedup (0.747/0.852/0.975 ms/tok vs baseline 7.615/8.304/9.199 at ctx 2048/8192/16384). The entire per-token forward -- 3 KDA blocks (fused int4-dequant split-K GEMVs for q/k/v/g sharing one x load, beta projection, short conv with a double-buffered conv-window ping-pong, gated-delta S recurrence with per-head fuzzy start flags, o_proj+residual, router/top-8 MoE with shared expert), 1 MLA block (q/kv_a projections, absorb-identity attention over the compressed bf16 latent cache via head-grouped online-softmax flash decode with adaptive chunking and in-kernel cache append/RoPE, kv_b value projection, o_proj, MoE), plus both RMSNorms per block -- executes as ONE Triton kernel launch per step. Stage ordering uses atomic item counters with acquire/release semantics and nanosleep-throttled polling; a warmup-based occupancy probe picks the CTA count once. int4 weights are repacked once into flat uint8/f32 regions and dequantized in-register inside every GEMV (no bf16 weight materialization). Weight flattening is cached across steps but invalidated by a register_load_state_dict_post_hook, so fresh check.py seeds re-materialize; no output memoization, no input-identity keying, no data_ptr dispatch in the final solution. A KIMI_DEBUG_EAGER env flag selects a slow eager oracle used only for the agent's own debugging; grading ran the megakernel path.
harnesskinetic-claude
Kernel source (redacted)
"""Kimi-Linear W4A16 hybrid decode — single-launch persistent megakernel.
The entire per-token forward (3 KDA + 1 MLA attention layers, 4 MoE FFNs,
every int4 dequant-GEMV, the short causal convs, the KDA recurrent-state
update, the MLA latent-cache attention via the absorb identity, the MoE
router + top-8 + shared expert GEMVs, both RMSNorms per block, residuals, and
all state/cache updates) is fused into ONE Triton kernel launch per step().
No CUDA graphs, no torch.compile, no per-op kernel loop — one persistent
co-resident grid executes ~20 sequential stages per token.
Design (profiling-driven):
* int4 weights are repacked once into one flat uint8 region (group-128
asymmetric scales/zeros + dense weights in a flat fp32 region), streamed
exactly once per token. Dequant is fused into every GEMV; no bf16 weight
ever lives in DRAM (the whole bandwidth advantage of W4A16 kept).
* Stages synchronize via item counters (release/acquire, gpu scope): each
CTA spins on the producer counter, processes grid-strided items, bumps
its own counter. Chain-ordered DAG => deadlock-free. High-value stages
also use per-head fuzzy flags so a head's KDA recurrence starts as soon
as its panels land rather than waiting for the full grid.
* GEMVs are split-K into ~3-group-deep items (batch-1 decode is
latency-bound, not bandwidth-bound: long dependent load chains hurt),
writing fp32 partials; nonlinearities (bf16 rounds/silu/softmax) are
applied by consumers AFTER summing partials, preserving ref rounding.
q/k/v and gate/up GEMVs are fused into shared-x items for 2-3x load
issue density per latency window.
* RMSNorms precompute once per norm site into a bf16 xn scratch; MoE
h=silu(g)*u precomputes once per block; the MoE tail is one weighted
reduce.
* MLA decode uses the absorb identity (scores = (W_k^T q_nope).c_kv +
q_rope.k_rope; out = (sum_l p_l c_kv_l) W_v^T), so the growing cache
stays compressed bf16 latents, read ~once per token via a head-grouped
(4 heads/tile) online-softmax flash pass with adaptive chunking
(64/128/256-row chunks by context length).
"""
from __future__ import annotations
import os
# Triton >=3.2 rejects plain-int module globals referenced inside @triton.jit;
# explicitly allow reading frozen layout constants (scratch offsets, strides)
# at trace time (documented escape hatch; avoids threading dozens of offsets
# through every kernel argument list).
os.environ.setdefault("TRITON_ALLOW_NON_CONSTEXPR_GLOBALS", "1")
# Triton's on-disk cache does not content-hash globals read under the flag
# above; stale kernels from older sources can be reused. Pin the cache root
# to this file's content hash so behavior is deterministic per source.
import hashlib as _hl
import tempfile as _tf
_fp = _hl.sha256(__file__.encode()).hexdigest()[:12]
os.environ.setdefault("TRITON_CACHE_DIR", os.path.join(
_tf.gettempdir(), f"triton_kimi_cache_{_fp}"))
import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl
OP_TYPE = "kimi_linear_w4a16_decode"
EPS = 1.0e-6
GROUP_SIZE = 128
_NUM_WARPS = int(os.environ.get("KIMI_NW", "8"))
_NCTA_REQ = int(os.environ.get("KIMI_NCTA", "188"))
_DEBUG_EAGER = os.environ.get("KIMI_DEBUG_EAGER", "") == "1"
_TIMING = int(os.environ.get("KIMI_TIMING", "0"))
_SLEEPNS = int(os.environ.get("KIMI_SLEEPNS", "400"))
_ACHUNK = 256 # cache rows per attention work item
_ACH_MAX = 68 # max attention chunks (~17.4k ctx)
_COPY_R = 512
_NSM = 188
# --------------------------------------------------------------------------- #
# scratch (fp32) element offsets
# --------------------------------------------------------------------------- #
SC_XA = 0
SC_XB = 2304
SC_QP = 4608 # q/k/v/g partials [kind 4][split 3][4096]
SC_BETA = SC_QP + 6 * 4 * 4096 # 53760 (9, 32)
SC_OPRP = SC_BETA + 288 # 54048 (4, 2304)
SC_LOGIT = SC_OPRP + 8 * 2304 # 63264 (9, 64)
SC_GUP = SC_LOGIT + 576 # 63840 [kind 2][slot 9][split 3][1024]
SC_DNP = SC_GUP + 2 * 9 * 6 * 1024 # 119136 (2, 9, 2304)
SC_HX = SC_DNP + 2 * 9 * 2304 # h = silu(g)*u (9, 1024)
SC_ORA = SC_HX + 9 * 1024 # 169824 (4096)
SC_QABS = SC_ORA + 4096 # 164704 (16384)
SC_PVF = SC_QABS + 16384 # 181088 (16384)
SC_QROPE = SC_PVF + 16384 # 197472 (2048)
SC_KVP = SC_QROPE + 2048 # 199520 (3, 576)
SC_MQP = SC_KVP + 3 * 576 # 201248 (3, 6144)
SC_ATT = SC_MQP + 3 * 6144 # 219680
ATT_STRIDE = 514
SC_RIDX = SC_ATT + _ACH_MAX * 32 * ATT_STRIDE # 8 i32 expert ids
SC_RW = SC_RIDX + 16 # 8 f32 route weights
SC_DXN = SC_RW + 16
SC_SHAD = SC_DXN + 2304
SC_SHADL = SC_SHAD + 4 * 2304
SC_LG0 = SC_SHADL + 4 * 576
SC_GSUM = SC_LG0 + 576
SC_XN = SC_GSUM + 18 * 1024
SC_TOTAL = SC_XN + 2304
WIN_PP_STRIDE = 3 * 3 * 4096
WIN_LAYER_STRIDE = 2 * WIN_PP_STRIDE
WIN_TOTAL = 3 * WIN_LAYER_STRIDE
_CT_COPY = 46
N_SLOTS = 48
HQF_BASE = 64 # per-launch per-head flags (3 blocks x 32)
BETA_FLAG = 96 # beta-done slots (3 blocks)
XR_MAX = 192
# per-block weight field ids (par index = b*40 + fid; f32 mirror at +256)
TF_QW, TF_QS, TF_QZ = 0, 1, 2
TF_KW, TF_KS, TF_KZ = 3, 4, 5
TF_VW, TF_VS, TF_VZ = 6, 7, 8
TF_GW, TF_GS, TF_GZ = 9, 10, 11
TF_OW, TF_OS, TF_OZ = 12, 13, 14
TF_BETA, TF_CONV, TF_AN, TF_MN, TF_ROUT = 15, 16, 17, 18, 19
TF_EXW, TF_EXSZ = 20, 21
TF_SGW, TF_SUW, TF_SDW = 22, 23, 24
TF_SGS, TF_SUS, TF_SDS = 25, 26, 27
TF_SGZ, TF_SUZ, TF_SDZ = 28, 29, 30
MF_QW, MF_QS, MF_QZ = 0, 1, 2
MF_KVAW, MF_KVAS, MF_KVAZ = 3, 4, 5
MF_KVBW, MF_KVBS, MF_KVBZ = 6, 7, 8
MF_OW, MF_OS, MF_OZ = 9, 10, 11
MF_AN, MF_MN, MF_ROUT = 12, 13, 14
MF_EXW, MF_EXSZ = 15, 16
MF_SGW, MF_SUW, MF_SDW = 17, 18, 19
MF_SGS, MF_SUS, MF_SDS = 20, 21, 22
MF_SGZ, MF_SUZ, MF_SDZ = 23, 24, 25
PAR_LEN = 512
_40 = 40
GW_BYTES = 2304 * 1024 // 2
GS_ELEMS = 18 * 1024
EXB_U8 = 3 * GW_BYTES
EXB_F32 = 6 * GS_ELEMS
# --------------------------------------------------------------------------- #
# module tree — identical structure/names to the reference for state_dict
# --------------------------------------------------------------------------- #
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__()
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
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
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)
class Block(nn.Module):
def __init__(self, cfg, kind: str):
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 _align256(x):
return (x + 255) // 256 * 256
# --------------------------------------------------------------------------- #
# megakernel helpers
# --------------------------------------------------------------------------- #
@triton.jit
def _wait(ctr, pos, target, dbg, cb, TIMING: tl.constexpr):
one = tl.arange(0, 1)
cur = tl.sum(tl.atomic_add(ctr + pos + one, 0, sem="acquire", scope="gpu"))
if cur < target:
while tl.sum(tl.load(ctr + pos + one, volatile=True)) < target:
# throttle polling to keep L2 atomic traffic off the workers
if _SLEEPNS > 0:
tl.inline_asm_elementwise("nanosleep.u32 2000; mov.u32 $0, $1;", "=r,r", [one],
dtype=tl.int32, is_pure=False, pack=1)
tl.atomic_add(ctr + pos + one, 0, sem="acquire", scope="gpu")
if TIMING:
tl.atomic_min(dbg + 2 * (pos - cb) + one, tl.inline_asm_elementwise(
"mov.u64 $0, %globaltimer;", "=l,l", [one], dtype=tl.int64, is_pure=False, pack=1))
tl.debug_barrier()
@triton.jit
def _bump(ctr, pos, n, pid, NCTA, dbg, cb, TIMING: tl.constexpr):
tl.debug_barrier()
if TIMING:
one = tl.arange(0, 1)
tl.atomic_max(dbg + 2 * (pos - cb) + 1 + one, tl.inline_asm_elementwise(
"mov.u64 $0, %globaltimer;", "=l,l", [one], dtype=tl.int64, is_pure=False, pack=1))
if pid < n:
cnt = (n - pid + NCTA - 1) // NCTA
tl.atomic_add(ctr + pos + tl.arange(0, 1), cnt, sem="release", scope="gpu")
@triton.jit
def _xn_make(scr, nx_ptr, x_ptr, nrm, ctr, ctr_base, sflag, pid, ncta, dbg, TIMING: tl.constexpr,
skip_wait: tl.constexpr):
# produce normed+bf16-rounded x into SC_XN from raw fp32 x_ptr (2304)
r16 = tl.arange(0, 16)
c128 = tl.arange(0, 128)
rstd = _rstd(x_ptr, 2304)
for it in range(pid, 9, ncta):
xk = it * 256 + tl.arange(0, 256)
xv = tl.load(x_ptr + xk).to(tl.float32)
nv = tl.load(nrm + xk).to(tl.float32)
xv = ((xv * rstd) * nv).to(tl.bfloat16).to(tl.float32)
tl.store(scr + SC_XN + xk, xv)
if not skip_wait:
_bump(ctr, ctr_base + sflag, 9, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + sflag, 9, dbg, ctr_base, TIMING)
@triton.jit
def _rstd(x_ptr, K):
"""Wide parallel sum-of-squares for RMSNorm (one or two shot loads)."""
ssq = 0.0
r16 = tl.arange(0, 16)
c128 = tl.arange(0, 128)
for k0 in range(0, K, 2048):
idx = k0 + r16[:, None] * 128 + c128[None, :]
xv = tl.load(x_ptr + idx, mask=idx < K, other=0.0).to(tl.float32)
ssq += tl.sum(xv * xv)
return tl.rsqrt(ssq / K + 1e-6)
@triton.jit
def _route(logits_ptr):
"""Rebuild router logits from 9 split partials, softmax, top-8."""
j = tl.arange(0, 64)
lg = tl.zeros((64,), dtype=tl.float32)
for s in tl.static_range(9):
lg += tl.load(logits_ptr + s * 64 + j)
lg = lg.to(tl.bfloat16).to(tl.float32)
m = tl.max(lg)
e = tl.exp(lg - m)
p = e / tl.sum(e)
idxs = tl.zeros((8,), dtype=tl.int32)
ws = tl.zeros((8,), dtype=tl.float32)
pw = p
acc_w = 0.0
for t in tl.static_range(8):
i = tl.argmax(pw, 0)
v = tl.max(pw, 0)
acc_w += v
id8 = tl.arange(0, 8)
idxs = tl.where(id8 == t, i.to(tl.int32), idxs)
ws = tl.where(id8 == t, v, ws)
pw = tl.where(j == i, -1.0, pw)
return idxs, ws * (2.446 / (acc_w + 1e-9))
@triton.jit
def _moe_in(base, par, slot, EXWF, EXSZF, SGF, SUF, SDF, SGSF, SUSF, SDSF, SGZF, SUZF, SDZF,
kind, e):
"""Address offsets of expert weight for (slot, kind): kind 0=gate 1=up 2=down."""
if slot < 8:
wb = tl.load(par + base + EXWF) + e * (3 * GW_BYTES)
sb = tl.load(par + 256 + base + EXSZF) + e * (6 * GS_ELEMS)
if kind == 0:
wq = wb
sp = sb
zp = sb + GS_ELEMS
elif kind == 1:
wq = wb + GW_BYTES
sp = sb + 2 * GS_ELEMS
zp = sb + 3 * GS_ELEMS
else:
wq = wb + 2 * GW_BYTES
sp = sb + 4 * GS_ELEMS
zp = sb + 5 * GS_ELEMS
else:
if kind == 0:
wq = tl.load(par + base + SGF)
sp = tl.load(par + 256 + base + SGSF)
zp = tl.load(par + 256 + base + SGZF)
elif kind == 1:
wq = tl.load(par + base + SUF)
sp = tl.load(par + 256 + base + SUSF)
zp = tl.load(par + 256 + base + SUZF)
else:
wq = tl.load(par + base + SDF)
sp = tl.load(par + 256 + base + SDSF)
zp = tl.load(par + 256 + base + SDZF)
return wq, sp, zp
@triton.jit
def _gemv_split(scr, x_ptr, w, s, z, NSTR, N, rstd, nrm, ns, sp, NG0,
NORMED: tl.constexpr, BN: tl.constexpr):
"""Split-k fused int4 GEMV piece: groups [6*sp, 6*sp+6) of the (K,NSTR) weight.
Returns fp32 partial accumulator over that group range (no rounding; the
consumer applies rounding/nonlinearities after summing partials).
"""
acc = tl.zeros((BN,), dtype=tl.float32)
for g in range(NG0 * sp, NG0 * sp + NG0):
sv = tl.load(s + g * N + ns, mask=ns < N, other=0.0).to(tl.float32)
zv = tl.load(z + g * N + ns, mask=ns < N, other=0.0).to(tl.float32)
dot = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wp = tl.load(w + (k0 // 2 + r32)[:, None] * N + ns[None, :],
mask=(ns < N)[None, :], other=0)
lo = (wp & 0xF).to(tl.float32)
hi = ((wp >> 4) & 0xF).to(tl.float32)
xe = tl.load(x_ptr + k0 + 2 * tl.arange(0, 32)).to(tl.float32)
xo = tl.load(x_ptr + k0 + 1 + 2 * tl.arange(0, 32)).to(tl.float32)
dot += tl.sum(lo * xe[:, None], 0) + tl.sum(hi * xo[:, None], 0)
sx += tl.sum(xe) + tl.sum(xo)
acc += sv * (dot - zv * sx)
return acc
@triton.jit
def _megakernel(
wq_g, wf_g, par, scr, win,
s0_ptr, s1_ptr, s2_ptr,
cq0_ptr, ck0_ptr, cv0_ptr,
cq1_ptr, ck1_ptr, cv1_ptr,
cq2_ptr, ck2_ptr, cv2_ptr,
ckv_in, kr_in, kvbig, krbig,
inh_ptr, outh_ptr, ctr, dbg,
ctr_base, L0, nc, flags, ppflags,
BN: tl.constexpr, BR: tl.constexpr, TIMING: tl.constexpr,
):
pid = tl.program_id(0)
ncta = tl.num_programs(0)
cold_cache = flags & 1
wqb = wq_g
wfb = wf_g
pb = par
kda_slo0 = (2 * ((0) * 5 + (0)))
# ================= KDA blocks =================
for b in tl.static_range(3):
XB_IN = SC_XA if b % 2 == 0 else SC_XB
XB_OUT = SC_XB if b % 2 == 0 else SC_XA
BASE = b * _40
x_ptr = scr + XB_IN
# ---- QKVG (split-k int4 partials + beta) ; block 0 prepends EXT
if b == 0:
for it in range(pid, 9, ncta):
i = it * 256 + tl.arange(0, 256)
v = tl.load(inh_ptr + i).to(tl.float32)
tl.store(scr + SC_XA + i, v)
_bump(ctr, ctr_base + 1, 9, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 1, 9, dbg, ctr_base, TIMING)
_xn_make(scr, SC_XN, scr + SC_XA, wfb + tl.load(pb + 256 + BASE + TF_AN), ctr, ctr_base, 1, pid, ncta, dbg, TIMING, False)
else:
_wait(ctr, ctr_base + (2 * ((b - 1) * 5 + (4))), 414, dbg, ctr_base, TIMING)
_xn_make(scr, SC_XN, scr + XB_IN, wfb + tl.load(pb + 256 + BASE + TF_AN), ctr, ctr_base, (2 * ((b) * 5 + (0))) + 1, pid, ncta, dbg, TIMING, False)
rstd = _rstd(x_ptr, 2304)
nrm = wfb + tl.load(pb + 256 + BASE + TF_AN)
NQ = 4096 // BN
wq0 = wqb + tl.load(pb + BASE + TF_QW)
sq0 = wfb + tl.load(pb + 256 + BASE + TF_QS)
zq0 = wfb + tl.load(pb + 256 + BASE + TF_QZ)
wk0 = wqb + tl.load(pb + BASE + TF_KW)
sk0 = wfb + tl.load(pb + 256 + BASE + TF_KS)
zk0 = wfb + tl.load(pb + 256 + BASE + TF_KZ)
wv0 = wqb + tl.load(pb + BASE + TF_VW)
sv0 = wfb + tl.load(pb + 256 + BASE + TF_VS)
zv0 = wfb + tl.load(pb + 256 + BASE + TF_VZ)
for it in range(pid, NQ * 6, ncta):
n0 = (it // 6) * BN
sp = it % 6
ns = n0 + tl.arange(0, BN)
aq = tl.zeros((BN,), dtype=tl.float32)
ak = tl.zeros((BN,), dtype=tl.float32)
av = tl.zeros((BN,), dtype=tl.float32)
for g in range(sp * 3, sp * 3 + 3):
svg = tl.load(sq0 + g * 4096 + ns).to(tl.float32)
zvg = tl.load(zq0 + g * 4096 + ns).to(tl.float32)
skg = tl.load(sk0 + g * 4096 + ns).to(tl.float32)
zkg = tl.load(zk0 + g * 4096 + ns).to(tl.float32)
svc = tl.load(sv0 + g * 4096 + ns).to(tl.float32)
zvc = tl.load(zv0 + g * 4096 + ns).to(tl.float32)
dotq = tl.zeros((BN,), dtype=tl.float32)
dotk = tl.zeros((BN,), dtype=tl.float32)
dotv = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wq_ = tl.load(wq0 + (k0 // 2 + r32)[:, None] * 4096 + ns[None, :])
wk_ = tl.load(wk0 + (k0 // 2 + r32)[:, None] * 4096 + ns[None, :])
wv_ = tl.load(wv0 + (k0 // 2 + r32)[:, None] * 4096 + ns[None, :])
xe = tl.load(scr + SC_XN + k0 + 2 * tl.arange(0, 32))
xo = tl.load(scr + SC_XN + k0 + 1 + 2 * tl.arange(0, 32))
loq = (wq_ & 0xF).to(tl.float32)
hiq = ((wq_ >> 4) & 0xF).to(tl.float32)
lok = (wk_ & 0xF).to(tl.float32)
hik = ((wk_ >> 4) & 0xF).to(tl.float32)
lov = (wv_ & 0xF).to(tl.float32)
hiv = ((wv_ >> 4) & 0xF).to(tl.float32)
dotq += tl.sum(loq * xe[:, None], 0) + tl.sum(hiq * xo[:, None], 0)
dotk += tl.sum(lok * xe[:, None], 0) + tl.sum(hik * xo[:, None], 0)
dotv += tl.sum(lov * xe[:, None], 0) + tl.sum(hiv * xo[:, None], 0)
sx += tl.sum(xe) + tl.sum(xo)
aq += svg * (dotq - zvg * sx)
ak += skg * (dotk - zkg * sx)
av += svc * (dotv - zvc * sx)
tl.store(scr + SC_QP + (0 * 6 + sp) * 4096 + ns, aq)
tl.store(scr + SC_QP + (1 * 6 + sp) * 4096 + ns, ak)
tl.store(scr + SC_QP + (2 * 6 + sp) * 4096 + ns, av)
tl.debug_barrier()
tl.atomic_add(ctr + ctr_base + HQF_BASE + b * 32 + (it // 6) + tl.arange(0, 1), 1, sem="release", scope="gpu")
for it in range(pid, NQ * 6, ncta):
n0 = (it // 6) * BN
sp = it % 6
ns = n0 + tl.arange(0, BN)
wg0 = wqb + tl.load(pb + BASE + TF_GW)
sg0 = wfb + tl.load(pb + 256 + BASE + TF_GS)
zg0 = wfb + tl.load(pb + 256 + BASE + TF_GZ)
acc = _gemv_split(scr, scr + SC_XN, wg0, sg0, zg0, 4096, 4096, rstd, nrm, ns, sp, 3, True, BN)
tl.store(scr + SC_QP + (3 * 6 + sp) * 4096 + ns, acc)
tl.debug_barrier()
tl.atomic_add(ctr + ctr_base + HQF_BASE + b * 32 + (it // 6) + tl.arange(0, 1), 1, sem="release", scope="gpu")
NGEMV = 4 * NQ * 6
for it in range(pid, 9, ncta):
c = it
bp = wfb + tl.load(pb + 256 + BASE + TF_BETA)
xk = tl.arange(0, 256)
xe = tl.load(scr + SC_XN + c * 256 + xk)
for hb in tl.static_range(4):
rows = hb * 8 + tl.arange(0, 8)
wt = tl.load(bp + rows[:, None] * 2304 + c * 256 + xk[None, :]).to(tl.float32)
part = tl.sum(wt * xe[None, :], 1)
tl.store(scr + SC_BETA + c * 32 + rows, part)
tl.debug_barrier()
tl.atomic_add(ctr + ctr_base + BETA_FLAG + b + tl.arange(0, 1), 1, sem="release", scope="gpu")
_bump(ctr, ctr_base + (2 * ((b) * 5 + (0))), NGEMV + 9, pid, ncta, dbg, ctr_base, TIMING)
# ---- REC: conv + recurrence (fuzzy per-head gated)
warm = (ppflags >> b) & 1
ppq = (ppflags >> (3 + b)) & 1
oneff = tl.arange(0, 1)
S_ptr = s0_ptr if b == 0 else (s1_ptr if b == 1 else s2_ptr)
cq_in = cq0_ptr if b == 0 else (cq1_ptr if b == 1 else cq2_ptr)
ck_in = ck0_ptr if b == 0 else (ck1_ptr if b == 1 else ck2_ptr)
cv_in = cv0_ptr if b == 0 else (cv1_ptr if b == 1 else cv2_ptr)
convp = wfb + tl.load(pb + 256 + BASE + TF_CONV)
wsrc = win + b * WIN_LAYER_STRIDE + ppq * WIN_PP_STRIDE
wdst = win + b * WIN_LAYER_STRIDE + (1 - ppq) * WIN_PP_STRIDE
TOTR = 128
for it in range(pid, TOTR, ncta):
h = it // 4
tt = it % 4
curf = tl.sum(tl.atomic_add(ctr + ctr_base + HQF_BASE + b * 32 + h + oneff, 0, sem="acquire", scope="gpu"))
while curf < 12:
tl.inline_asm_elementwise("nanosleep.u32 400; mov.u32 $0, $1;", "=r,r", [oneff],
dtype=tl.int32, is_pure=False, pack=1)
curf = tl.sum(tl.atomic_add(ctr + ctr_base + HQF_BASE + b * 32 + h + oneff, 0, sem="acquire", scope="gpu"))
curb = tl.sum(tl.atomic_add(ctr + ctr_base + BETA_FLAG + b + oneff, 0, sem="acquire", scope="gpu"))
while curb < 9:
tl.inline_asm_elementwise("nanosleep.u32 400; mov.u32 $0, $1;", "=r,r", [oneff],
dtype=tl.int32, is_pure=False, pack=1)
curb = tl.sum(tl.atomic_add(ctr + ctr_base + BETA_FLAG + b + oneff, 0, sem="acquire", scope="gpu"))
tl.debug_barrier()
dv0 = tt * 32
ch0 = h * 128
dk = tl.arange(0, 128)
dr = tl.arange(0, 32)
cwq0 = tl.load(convp + 0 * 16384 + (ch0 + dk) * 4 + 0)
cwq1 = tl.load(convp + 0 * 16384 + (ch0 + dk) * 4 + 1)
cwq2 = tl.load(convp + 0 * 16384 + (ch0 + dk) * 4 + 2)
cwq3 = tl.load(convp + 0 * 16384 + (ch0 + dk) * 4 + 3)
cwk0 = tl.load(convp + 1 * 16384 + (ch0 + dk) * 4 + 0)
cwk1 = tl.load(convp + 1 * 16384 + (ch0 + dk) * 4 + 1)
cwk2 = tl.load(convp + 1 * 16384 + (ch0 + dk) * 4 + 2)
cwk3 = tl.load(convp + 1 * 16384 + (ch0 + dk) * 4 + 3)
cwv0 = tl.load(convp + 2 * 16384 + (ch0 + dv0 + dr) * 4 + 0)
cwv1 = tl.load(convp + 2 * 16384 + (ch0 + dv0 + dr) * 4 + 1)
cwv2 = tl.load(convp + 2 * 16384 + (ch0 + dv0 + dr) * 4 + 2)
cwv3 = tl.load(convp + 2 * 16384 + (ch0 + dv0 + dr) * 4 + 3)
# raw projections = 3 split partial sums, bf16 rounded (ref order)
qraw = (tl.load(scr + SC_QP + (0 * 6 + 0) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (0 * 6 + 1) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (0 * 6 + 2) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (0 * 6 + 3) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (0 * 6 + 4) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (0 * 6 + 5) * 4096 + ch0 + dk)).to(tl.bfloat16).to(tl.float32)
kraw = (tl.load(scr + SC_QP + (1 * 6 + 0) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (1 * 6 + 1) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (1 * 6 + 2) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (1 * 6 + 3) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (1 * 6 + 4) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (1 * 6 + 5) * 4096 + ch0 + dk)).to(tl.bfloat16).to(tl.float32)
graw = (tl.load(scr + SC_QP + (3 * 6 + 0) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (3 * 6 + 1) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (3 * 6 + 2) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (3 * 6 + 3) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (3 * 6 + 4) * 4096 + ch0 + dk)
+ tl.load(scr + SC_QP + (3 * 6 + 5) * 4096 + ch0 + dk)).to(tl.bfloat16).to(tl.float32)
vraw = (tl.load(scr + SC_QP + (2 * 6 + 0) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (2 * 6 + 1) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (2 * 6 + 2) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (2 * 6 + 3) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (2 * 6 + 4) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (2 * 6 + 5) * 4096 + ch0 + dv0 + dr)).to(tl.bfloat16).to(tl.float32)
if warm == 1:
qw0 = tl.load(wsrc + 0 * 12288 + 0 * 4096 + ch0 + dk)
qw1 = tl.load(wsrc + 0 * 12288 + 1 * 4096 + ch0 + dk)
qw2 = tl.load(wsrc + 0 * 12288 + 2 * 4096 + ch0 + dk)
kw0 = tl.load(wsrc + 1 * 12288 + 0 * 4096 + ch0 + dk)
kw1 = tl.load(wsrc + 1 * 12288 + 1 * 4096 + ch0 + dk)
kw2 = tl.load(wsrc + 1 * 12288 + 2 * 4096 + ch0 + dk)
vw0 = tl.load(wsrc + 2 * 12288 + 0 * 4096 + ch0 + dv0 + dr)
vw1 = tl.load(wsrc + 2 * 12288 + 1 * 4096 + ch0 + dv0 + dr)
vw2 = tl.load(wsrc + 2 * 12288 + 2 * 4096 + ch0 + dv0 + dr)
kqw1 = tl.load(wsrc + 0 * 12288 + 1 * 4096 + ch0 + dv0 + dr)
kqw2 = tl.load(wsrc + 0 * 12288 + 2 * 4096 + ch0 + dv0 + dr)
kkw1 = tl.load(wsrc + 1 * 12288 + 1 * 4096 + ch0 + dv0 + dr)
kkw2 = tl.load(wsrc + 1 * 12288 + 2 * 4096 + ch0 + dv0 + dr)
kvw1 = vw1
kvw2 = vw2
else:
qw0 = tl.load(cq_in + 0 * 4096 + ch0 + dk).to(tl.float32)
qw1 = tl.load(cq_in + 1 * 4096 + ch0 + dk).to(tl.float32)
qw2 = tl.load(cq_in + 2 * 4096 + ch0 + dk).to(tl.float32)
kw0 = tl.load(ck_in + 0 * 4096 + ch0 + dk).to(tl.float32)
kw1 = tl.load(ck_in + 1 * 4096 + ch0 + dk).to(tl.float32)
kw2 = tl.load(ck_in + 2 * 4096 + ch0 + dk).to(tl.float32)
vw0 = tl.load(cv_in + 0 * 4096 + ch0 + dv0 + dr).to(tl.float32)
vw1 = tl.load(cv_in + 1 * 4096 + ch0 + dv0 + dr).to(tl.float32)
vw2 = tl.load(cv_in + 2 * 4096 + ch0 + dv0 + dr).to(tl.float32)
kqw1 = tl.load(cq_in + 1 * 4096 + ch0 + dv0 + dr).to(tl.float32)
kqw2 = tl.load(cq_in + 2 * 4096 + ch0 + dv0 + dr).to(tl.float32)
kkw1 = tl.load(ck_in + 1 * 4096 + ch0 + dv0 + dr).to(tl.float32)
kkw2 = tl.load(ck_in + 2 * 4096 + ch0 + dv0 + dr).to(tl.float32)
kvw1 = vw1
kvw2 = vw2
# window dst shift (own 32 channels only, all 3 windows)
tl.store(wdst + 0 * 12288 + 0 * 4096 + ch0 + dv0 + dr, kqw1)
tl.store(wdst + 0 * 12288 + 1 * 4096 + ch0 + dv0 + dr, kqw2)
tl.store(wdst + 0 * 12288 + 2 * 4096 + ch0 + dv0 + dr, vraw)
tl.store(wdst + 1 * 12288 + 0 * 4096 + ch0 + dv0 + dr, kkw1)
tl.store(wdst + 1 * 12288 + 1 * 4096 + ch0 + dv0 + dr, kkw2)
tl.store(wdst + 1 * 12288 + 2 * 4096 + ch0 + dv0 + dr,
(tl.load(scr + SC_QP + (1 * 6 + 0) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (1 * 6 + 1) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (1 * 6 + 2) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (1 * 6 + 3) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (1 * 6 + 4) * 4096 + ch0 + dv0 + dr)
+ tl.load(scr + SC_QP + (1 * 6 + 5) * 4096 + ch0 + dv0 + dr)).to(tl.bfloat16).to(tl.float32))
tl.store(wdst + 2 * 12288 + 0 * 4096 + ch0 + dv0 + dr, kvw1)
tl.store(wdst + 2 * 12288 + 1 * 4096 + ch0 + dv0 + dr, kvw2)
tl.store(wdst + 2 * 12288 + 2 * 4096 + ch0 + dv0 + dr, vraw)
# conv = sum over 4 rows, silu, bf16 round (ref rounds conv outputs)
qacc = qw0 * cwq0 + qw1 * cwq1 + qw2 * cwq2 + qraw * cwq3
kacc = kw0 * cwk0 + kw1 * cwk1 + kw2 * cwk2 + kraw * cwk3
vacc = vw0 * cwv0 + vw1 * cwv1 + vw2 * cwv2 + vraw * cwv3
qs = (qacc * tl.sigmoid(qacc)).to(tl.bfloat16).to(tl.float32) * 0.08838834764831845
ks = (kacc * tl.sigmoid(kacc)).to(tl.bfloat16).to(tl.float32)
vs = (vacc * tl.sigmoid(vacc)).to(tl.bfloat16).to(tl.float32)
decay = tl.sigmoid(-graw)
beta = tl.load(scr + SC_BETA + h)
for jb in tl.static_range(1, 9):
beta += tl.load(scr + SC_BETA + jb * 32 + h)
beta = tl.sigmoid(beta.to(tl.bfloat16).to(tl.float32))
sptr = S_ptr + h * 128 * 128 + dk[:, None] * 128 + dv0 + dr[None, :]
Sb = tl.load(sptr)
Sb = Sb * decay[:, None]
pred = tl.sum(Sb * ks[:, None], 0)
err = vs - pred
Sb = Sb + beta * ks[:, None] * err[None, :]
tl.store(sptr, Sb)
o = tl.sum(Sb * qs[:, None], 0)
tl.store(scr + SC_ORA + ch0 + dv0 + dr, o.to(tl.bfloat16).to(tl.float32))
_bump(ctr, ctr_base + (2 * ((b) * 5 + (1))), TOTR, pid, ncta, dbg, ctr_base, TIMING)
# ---- OPROJ: split-k partials + reduce(residual)
_wait(ctr, ctr_base + (2 * ((b) * 5 + (1))), TOTR, dbg, ctr_base, TIMING)
w_o = wqb + tl.load(pb + BASE + TF_OW)
s_o = wfb + tl.load(pb + 256 + BASE + TF_OS)
z_o = wfb + tl.load(pb + 256 + BASE + TF_OZ)
NOP = (2304 // BN) * 4
for it in range(pid, NOP, ncta):
n0 = (it // 4) * BN
sp = it % 4
ns = n0 + tl.arange(0, BN)
acc = _gemv_split(scr, scr + SC_ORA, w_o, s_o, z_o, 2304, 2304, 0.0, wfb, ns, sp, 8, False, BN)
tl.store(scr + SC_OPRP + sp * 2304 + ns, acc)
_bump(ctr, ctr_base + (2 * ((b) * 5 + (2))) + 1, NOP, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + (2 * ((b) * 5 + (2))) + 1, NOP, dbg, ctr_base, TIMING)
NQO = 2304 // BN
for it in range(pid, NQO, ncta):
ns = it * BN + tl.arange(0, BN)
a4 = tl.load(scr + SC_OPRP + 0 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 1 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 2 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 3 * 2304 + ns)
cur = tl.load(scr + XB_IN + ns)
tl.store(scr + XB_OUT + ns, (cur + a4).to(tl.bfloat16).to(tl.float32))
_bump(ctr, ctr_base + (2 * ((b) * 5 + (2))) + 2, NQO, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + (2 * ((b) * 5 + (2))) + 2, NQO, dbg, ctr_base, TIMING)
# ---- router partials (folded into OPROJ stage)
_xn_make(scr, SC_XN, scr + XB_OUT, wfb + tl.load(pb + 256 + BASE + TF_MN), ctr, ctr_base, (2 * ((b) * 5 + (2))) + 2, pid, ncta, dbg, TIMING, False)
mptr = scr + XB_OUT
rstd_m = _rstd(mptr, 2304)
nrmm = wfb + tl.load(pb + 256 + BASE + TF_MN)
rwp = wfb + tl.load(pb + 256 + BASE + TF_ROUT)
if b == 0:
for it in range(pid, 9, ncta):
kc9 = it * 256
xk9 = kc9 + tl.arange(0, 256)
x9 = tl.load(mptr + xk9).to(tl.float32)
nv9 = tl.load(nrmm + xk9).to(tl.float32)
x9 = ((x9 * rstd_m) * nv9).to(tl.bfloat16).to(tl.float32)
TOTRO = 72
for it in range(pid, TOTRO, ncta):
kc = it // 8
jc = it % 8
rows = jc * 8 + tl.arange(0, 8)
k0 = kc * 256
xk = tl.arange(0, 256)
xe = tl.load(scr + SC_XN + k0 + xk)
wt = tl.load(rwp + rows[:, None] * 2304 + k0 + xk[None, :]).to(tl.float32)
part = tl.sum(wt * xe[None, :], 1)
tl.store(scr + SC_LOGIT + kc * 64 + rows, part)
_bump(ctr, ctr_base + (2 * ((b) * 5 + (2))), NOP + NQO + TOTRO, pid, ncta, dbg, ctr_base, TIMING)
# ---- RGU: gate/up fused gemv (waits full OPROJ incl. router)
_wait(ctr, ctr_base + (2 * ((b) * 5 + (2))), NOP + NQO + TOTRO, dbg, ctr_base, TIMING)
idxs, wsel = _route(scr + SC_LOGIT)
NGCH = 1024 // BN
TOTG = 9 * NGCH * 6
for it in range(pid, TOTG, ncta):
slot = it // (NGCH * 6)
n0 = ((it // 6) % NGCH) * BN
sp = it % 6
ns = n0 + tl.arange(0, BN)
e = tl.sum(tl.where(tl.arange(0, 8) == slot, idxs, 0))
wg, sg, zg = _moe_in(BASE, pb, slot,
TF_EXW, TF_EXSZ, TF_SGW, TF_SUW, TF_SDW,
TF_SGS, TF_SUS, TF_SDS, TF_SGZ, TF_SUZ, TF_SDZ, 0, e)
wu, su, zu = _moe_in(BASE, pb, slot,
TF_EXW, TF_EXSZ, TF_SGW, TF_SUW, TF_SDW,
TF_SGS, TF_SUS, TF_SDS, TF_SGZ, TF_SUZ, TF_SDZ, 1, e)
ag = tl.zeros((BN,), dtype=tl.float32)
au = tl.zeros((BN,), dtype=tl.float32)
for g in range(sp * 3, sp * 3 + 3):
svg = tl.load(wfb + sg + g * 1024 + ns).to(tl.float32)
zvg = tl.load(wfb + zg + g * 1024 + ns).to(tl.float32)
sug = tl.load(wfb + su + g * 1024 + ns).to(tl.float32)
zug = tl.load(wfb + zu + g * 1024 + ns).to(tl.float32)
dg = tl.zeros((BN,), dtype=tl.float32)
du = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wg_ = tl.load(wqb + wg + (k0 // 2 + r32)[:, None] * 1024 + ns[None, :])
wu_ = tl.load(wqb + wu + (k0 // 2 + r32)[:, None] * 1024 + ns[None, :])
log = (wg_ & 0xF).to(tl.float32)
hig = ((wg_ >> 4) & 0xF).to(tl.float32)
lou = (wu_ & 0xF).to(tl.float32)
hiu = ((wu_ >> 4) & 0xF).to(tl.float32)
xe = tl.load(scr + SC_XN + k0 + 2 * tl.arange(0, 32))
xo = tl.load(scr + SC_XN + k0 + 1 + 2 * tl.arange(0, 32))
dg += tl.sum(log * xe[:, None], 0) + tl.sum(hig * xo[:, None], 0)
du += tl.sum(lou * xe[:, None], 0) + tl.sum(hiu * xo[:, None], 0)
sx += tl.sum(xe) + tl.sum(xo)
ag += svg * (dg - zvg * sx)
au += sug * (du - zug * sx)
tl.store(scr + SC_GUP + ((0 * 9 + slot) * 6 + sp) * 1024 + ns, ag)
tl.store(scr + SC_GUP + ((1 * 9 + slot) * 6 + sp) * 1024 + ns, au)
_bump(ctr, ctr_base + (2 * ((b) * 5 + (3))), TOTG, pid, ncta, dbg, ctr_base, TIMING)
# ---- DOWN: h-precompute (own slot) -> split gemv -> reduce
_wait(ctr, ctr_base + (2 * ((b) * 5 + (3))), TOTG, dbg, ctr_base, TIMING)
NHX = 9 * (1024 // 128)
for it in range(pid, NHX, ncta):
slot = it // (1024 // 128)
n0 = (it % (1024 // 128)) * 128
nsr = n0 + tl.arange(0, 128)
g0 = (tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 0) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 1) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 2) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 3) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 4) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 5) * 1024 + nsr))
u0 = (tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 0) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 1) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 2) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 3) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 4) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 5) * 1024 + nsr))
tl.store(scr + SC_HX + slot * 1024 + nsr, g0 * tl.sigmoid(g0) * u0)
_bump(ctr, ctr_base + 42 + b, NHX, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 42 + b, NHX, dbg, ctr_base, TIMING)
NDN = 2304 // BN
TOTDG = 9 * NDN * 2
for it in range(pid, TOTDG, ncta):
slot = it // (NDN * 2)
n0 = ((it // 2) % NDN) * BN
sp = it % 2
ns = n0 + tl.arange(0, BN)
e = tl.sum(tl.where(tl.arange(0, 8) == slot, idxs, 0))
w, s, z = _moe_in(BASE, pb, slot,
TF_EXW, TF_EXSZ, TF_SGW, TF_SUW, TF_SDW,
TF_SGS, TF_SUS, TF_SDS, TF_SGZ, TF_SUZ, TF_SDZ, 2, e)
acc = tl.zeros((BN,), dtype=tl.float32)
for g in range(sp * 4, sp * 4 + 4):
sv = tl.load(wfb + s + g * 2304 + ns).to(tl.float32)
zv = tl.load(wfb + z + g * 2304 + ns).to(tl.float32)
dot = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wp = tl.load(wqb + w + (k0 // 2 + r32)[:, None] * 2304 + ns[None, :])
lo = (wp & 0xF).to(tl.float32)
hi = ((wp >> 4) & 0xF).to(tl.float32)
# h = silu(g)*u precomputed into SC_HX (see HPRE stage)
he = tl.load(scr + SC_HX + slot * 1024 + k0 + 2 * tl.arange(0, 32))
ho = tl.load(scr + SC_HX + slot * 1024 + k0 + 1 + 2 * tl.arange(0, 32))
dot += tl.sum(lo * he[:, None], 0) + tl.sum(hi * ho[:, None], 0)
sx += tl.sum(he) + tl.sum(ho)
acc += sv * (dot - zv * sx)
tl.store(scr + SC_DNP + (sp * 9 + slot) * 2304 + ns, acc)
_bump(ctr, ctr_base + (2 * ((b) * 5 + (4))) + 1, TOTDG, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + (2 * ((b) * 5 + (4))) + 1, TOTDG, dbg, ctr_base, TIMING)
for it in range(pid, NDN, ncta):
ns = it * BN + tl.arange(0, BN)
outv = tl.load(scr + XB_OUT + ns)
for slot in tl.static_range(9):
wsel_s = tl.sum(tl.where(tl.arange(0, 8) == slot, wsel, 0.0))
if slot == 8:
wsel_s = 1.0
pv = tl.load(scr + SC_DNP + (0 * 9 + slot) * 2304 + ns)
pv += tl.load(scr + SC_DNP + (1 * 9 + slot) * 2304 + ns)
outv += wsel_s * pv
tl.store(scr + XB_OUT + ns, outv)
_bump(ctr, ctr_base + (2 * ((b) * 5 + (4))), NHX + TOTDG + NDN, pid, ncta, dbg, ctr_base, TIMING)
# ================= MLA block ============
BASE = 3 * _40
mla_x = scr + SC_XB
# ---- QABS: q_proj/kv_a split partials -> absorb + rope + append
_wait(ctr, ctr_base + (2 * ((2) * 5 + (4))), 414, dbg, ctr_base, TIMING)
_xn_make(scr, SC_XN, scr + SC_XB, wfb + tl.load(pb + 256 + BASE + MF_AN), ctr, ctr_base, 44, pid, ncta, dbg, TIMING, False)
rstd = _rstd(mla_x, 2304)
nrm = wfb + tl.load(pb + 256 + BASE + MF_AN)
NQQ = 6144 // BN
NKV = (576 + BN - 1) // BN
NQG = NQQ * 3 + NKV * 3
for it in range(pid, NQG, ncta):
if it < NQQ * 3:
n0 = (it // 3) * BN
sp = it % 6
ns = n0 + tl.arange(0, BN)
w = wqb + tl.load(pb + BASE + MF_QW)
s = wfb + tl.load(pb + 256 + BASE + MF_QS)
z = wfb + tl.load(pb + 256 + BASE + MF_QZ)
acc = _gemv_split(scr, scr + SC_XN, w, s, z, 6144, 6144, rstd, nrm, ns, sp, 6, True, BN)
tl.store(scr + SC_MQP + sp * 6144 + ns, acc)
else:
it2 = it - NQQ * 3
n0 = (it2 // 3) * BN
sp = it2 % 3
ns = n0 + tl.arange(0, BN)
nmk = ns < 576
w = wqb + tl.load(pb + BASE + MF_KVAW)
s = wfb + tl.load(pb + 256 + BASE + MF_KVAS)
z = wfb + tl.load(pb + 256 + BASE + MF_KVAZ)
acc = _gemv_split(scr, scr + SC_XN, w, s, z, 576, 576, rstd, nrm, ns, sp, 6, True, BN)
tl.store(scr + SC_KVP + sp * 576 + ns, acc, mask=nmk)
_bump(ctr, ctr_base + 31, NQG, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 31, NQG, dbg, ctr_base, TIMING)
kvbw = wqb + tl.load(pb + BASE + MF_KVBW)
kvbs = wfb + tl.load(pb + 256 + BASE + MF_KVBS)
kvbz = wfb + tl.load(pb + 256 + BASE + MF_KVBZ)
TOTA = 512 + 1
log_th = -9.210340371976184 / 32.0
for it in range(pid, TOTA, ncta):
if it < 512:
h = it // 16
cc = it % 16
c0 = cc * 32
d = tl.arange(0, 128)
qn = (tl.load(scr + SC_MQP + 0 * 6144 + h * 192 + d)
+ tl.load(scr + SC_MQP + 1 * 6144 + h * 192 + d)
+ tl.load(scr + SC_MQP + 2 * 6144 + h * 192 + d)).to(tl.bfloat16).to(tl.float32)
g0 = c0 // 128
sv = tl.load(kvbs + g0 * 8192 + h * 256 + d).to(tl.float32)
zv = tl.load(kvbz + g0 * 8192 + h * 256 + d).to(tl.float32)
sq = sv * qn
zsq = tl.sum(zv * sq)
r16 = tl.arange(0, 16)
wp = tl.load(kvbw + (c0 // 2 + r16)[:, None] * 8192 + (h * 256 + d)[None, :])
lo = (wp & 0xF).to(tl.float32)
hi = ((wp >> 4) & 0xF).to(tl.float32)
de = tl.sum(lo * sq[None, :], 1) - zsq
do = tl.sum(hi * sq[None, :], 1) - zsq
tl.store(scr + SC_QABS + h * 512 + c0 + 2 * tl.arange(0, 16), de)
tl.store(scr + SC_QABS + h * 512 + c0 + 2 * tl.arange(0, 16) + 1, do)
if cc == 0:
# rope q_rope[h] from summed partials (round before rope, ref order)
freq = tl.exp(log_th * tl.arange(0, 32).to(tl.float32))
ang = L0 * freq
cv = tl.cos(ang)
svv = tl.sin(ang)
p0 = h * 192 + 128 + 2 * tl.arange(0, 32)
p1 = h * 192 + 129 + 2 * tl.arange(0, 32)
ee = (tl.load(scr + SC_MQP + 0 * 6144 + p0)
+ tl.load(scr + SC_MQP + 1 * 6144 + p0)
+ tl.load(scr + SC_MQP + 2 * 6144 + p0)).to(tl.bfloat16).to(tl.float32)
eo = (tl.load(scr + SC_MQP + 0 * 6144 + p1)
+ tl.load(scr + SC_MQP + 1 * 6144 + p1)
+ tl.load(scr + SC_MQP + 2 * 6144 + p1)).to(tl.bfloat16).to(tl.float32)
re = (ee * cv - eo * svv).to(tl.bfloat16).to(tl.float32)
ro = (eo * cv + ee * svv).to(tl.bfloat16).to(tl.float32)
tl.store(scr + SC_QROPE + h * 64 + 2 * tl.arange(0, 32), re)
tl.store(scr + SC_QROPE + h * 64 + 2 * tl.arange(0, 32) + 1, ro)
else:
i512 = tl.arange(0, 512)
ckv = (tl.load(scr + SC_KVP + 0 * 576 + i512)
+ tl.load(scr + SC_KVP + 1 * 576 + i512)
+ tl.load(scr + SC_KVP + 2 * 576 + i512)).to(tl.bfloat16)
tl.store(kvbig + L0 * 512 + i512, ckv)
freq = tl.exp(log_th * tl.arange(0, 32).to(tl.float32))
ang = L0 * freq
cv = tl.cos(ang)
svv = tl.sin(ang)
ee = (tl.load(scr + SC_KVP + 0 * 576 + 512 + 2 * tl.arange(0, 32))
+ tl.load(scr + SC_KVP + 1 * 576 + 512 + 2 * tl.arange(0, 32))
+ tl.load(scr + SC_KVP + 2 * 576 + 512 + 2 * tl.arange(0, 32)))
eo = (tl.load(scr + SC_KVP + 0 * 576 + 512 + 1 + 2 * tl.arange(0, 32))
+ tl.load(scr + SC_KVP + 1 * 576 + 512 + 1 + 2 * tl.arange(0, 32))
+ tl.load(scr + SC_KVP + 2 * 576 + 512 + 1 + 2 * tl.arange(0, 32)))
re = (ee * cv - eo * svv).to(tl.bfloat16)
ro = (eo * cv + ee * svv).to(tl.bfloat16)
tl.store(krbig + L0 * 64 + 2 * tl.arange(0, 32), re)
tl.store(krbig + L0 * 64 + 2 * tl.arange(0, 32) + 1, ro)
_bump(ctr, ctr_base + 30, TOTA, pid, ncta, dbg, ctr_base, TIMING)
# ---- ATT: online softmax over latent cache (256-row chunks, prefetched)
_wait(ctr, ctr_base + 30, TOTA, dbg, ctr_base, TIMING)
# ATT: head-grouped online softmax (4 heads per item share each cache tile)
TOTATT = 8 * nc
i512a = tl.arange(0, 512)
i64a = tl.arange(0, 64)
HPI: tl.constexpr = 4
for it in range(pid, TOTATT, ncta):
hg = it // nc
ch = it % nc
r0 = ch * _ACHUNK
r1 = tl.minimum(r0 + _ACHUNK, L0 + 1)
lim = tl.minimum(r1, L0)
hh = hg * HPI + tl.arange(0, HPI)
qa = tl.load(scr + SC_QABS + hh[:, None] * 512 + i512a[None, :]) # (HPI, 512)
qr = tl.load(scr + SC_QROPE + hh[:, None] * 64 + i64a[None, :]) # (HPI, 64)
m_run = tl.full((HPI,), -1e30, dtype=tl.float32)
l_run = tl.zeros((HPI,), dtype=tl.float32)
pv = tl.zeros((HPI, 512), dtype=tl.float32)
# double-buffered prefetch pipeline over BR-row tiles
mk0 = r0 + tl.arange(0, BR) < lim
rlast = r0
ct0 = tl.load(ckv_in + (r0 + tl.arange(0, BR))[:, None] * 512 + i512a[None, :],
mask=mk0[:, None], other=0.0).to(tl.float32)
krt0 = tl.load(kr_in + (r0 + tl.arange(0, BR))[:, None] * 64 + i64a[None, :],
mask=mk0[:, None], other=0.0).to(tl.float32)
for rr in range(r0 + BR, lim, BR):
mk1 = rr + tl.arange(0, BR) < lim
ct1 = tl.load(ckv_in + (rr + tl.arange(0, BR))[:, None] * 512 + i512a[None, :],
mask=mk1[:, None], other=0.0).to(tl.float32)
krt1 = tl.load(kr_in + (rr + tl.arange(0, BR))[:, None] * 64 + i64a[None, :],
mask=mk1[:, None], other=0.0).to(tl.float32)
s = (tl.sum(ct0[None, :, :] * qa[:, None, :], 2)
+ tl.sum(krt0[None, :, :] * qr[:, None, :], 2)) * 0.07216878364870323
s = tl.where((rlast + tl.arange(0, BR) < lim)[None, :], s, -1e30)
m_new = tl.maximum(m_run, tl.max(s, 1))
alpha = tl.exp(m_run - m_new)
wj = tl.exp(s - m_new[:, None])
l_run = l_run * alpha + tl.sum(wj, 1)
pv = pv * alpha[:, None] + tl.sum(wj[:, :, None] * ct0[None, :, :], 1)
m_run = m_new
ct0 = ct1
krt0 = krt1
rlast = rr
s = (tl.sum(ct0[None, :, :] * qa[:, None, :], 2)
+ tl.sum(krt0[None, :, :] * qr[:, None, :], 2)) * 0.07216878364870323
s = tl.where((rlast + tl.arange(0, BR) < lim)[None, :], s, -1e30)
m_new = tl.maximum(m_run, tl.max(s, 1))
alpha = tl.exp(m_run - m_new)
wj = tl.exp(s - m_new[:, None])
l_run = l_run * alpha + tl.sum(wj, 1)
pv = pv * alpha[:, None] + tl.sum(wj[:, :, None] * ct0[None, :, :], 1)
m_run = m_new
if r1 > L0:
ct = tl.load(kvbig + L0 * 512 + i512a).to(tl.float32)
krt = tl.load(krbig + L0 * 64 + i64a).to(tl.float32)
sb = (tl.sum(qa * ct[None, :], 1) + tl.sum(qr * krt[None, :], 1)) * 0.07216878364870323
m_new = tl.maximum(m_run, sb)
alpha = tl.exp(m_run - m_new)
wjb = tl.exp(sb - m_new)
l_run = l_run * alpha + wjb
pv = pv * alpha[:, None] + wjb[:, None] * ct[None, :]
m_run = m_new
for hi in tl.static_range(HPI):
base = SC_ATT + ((ch * 32 + hg * HPI + hi) * ATT_STRIDE)
tl.store(scr + base + 0, tl.sum(tl.where(tl.arange(0, HPI) == hi, m_run, 0.0)))
tl.store(scr + base + 1, tl.sum(tl.where(tl.arange(0, HPI) == hi, l_run, 0.0)))
tl.store(scr + base + 2 + i512a,
tl.sum(tl.where(tl.arange(0, HPI)[:, None] == hi, pv, 0.0), 0))
# seed copy (cold only)
NCOPY = 0 if cold_cache == 0 else (L0 + _COPY_R - 1) // _COPY_R
for it in range(pid + TOTATT, TOTATT + NCOPY, ncta):
it0 = it - TOTATT
r0c = it0 * _COPY_R
for rr in range(r0c, tl.minimum(r0c + _COPY_R, L0), 32):
i32 = tl.arange(0, 32)
mk = rr + i32 < L0
v = tl.load(ckv_in + (rr + i32)[:, None] * 512 + i512a[None, :], mask=mk[:, None], other=0.0)
tl.store(kvbig + (rr + i32)[:, None] * 512 + i512a[None, :], v, mask=mk[:, None])
v2 = tl.load(kr_in + (rr + i32)[:, None] * 64 + i64a[None, :], mask=mk[:, None], other=0.0)
tl.store(krbig + (rr + i32)[:, None] * 64 + i64a[None, :], v2, mask=mk[:, None])
_bump(ctr, ctr_base + 32, TOTATT + NCOPY, pid, ncta, dbg, ctr_base, TIMING)
# ---- CAO: combine partials (intra) -> AOUT through W_v
_wait(ctr, ctr_base + 32, TOTATT + NCOPY, dbg, ctr_base, TIMING)
for it in range(pid, 32, ncta):
h = it
m_star = -1e30
for c in range(0, nc):
m_c = tl.load(scr + SC_ATT + (c * 32 + h) * ATT_STRIDE + 0)
m_star = tl.maximum(m_star, m_c)
l_star = 0.0
pvf = tl.zeros((512,), dtype=tl.float32)
for c in range(0, nc):
base = SC_ATT + (c * 32 + h) * ATT_STRIDE
m_c = tl.load(scr + base + 0)
l_c = tl.load(scr + base + 1)
pv_c = tl.load(scr + base + 2 + i512a)
wv = tl.exp(m_c - m_star)
l_star += l_c * wv
pvf += pv_c * wv
tl.store(scr + SC_PVF + h * 512 + i512a, pvf / l_star)
_bump(ctr, ctr_base + 35, 32, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 35, 32, dbg, ctr_base, TIMING)
TOTAO = 128
for it in range(pid, TOTAO, ncta):
h = it // 4
dc = it % 4
d0 = dc * 32
ns32 = h * 256 + 128 + d0 + tl.arange(0, 32)
acc = tl.zeros((32,), dtype=tl.float32)
for g in range(0, 4):
sv = tl.load(kvbs + g * 8192 + ns32).to(tl.float32)
zv = tl.load(kvbz + g * 8192 + ns32).to(tl.float32)
dot = tl.zeros((32,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wp = tl.load(kvbw + (k0 // 2 + r32)[:, None] * 8192 + ns32[None, :])
lo = (wp & 0xF).to(tl.float32)
hi = ((wp >> 4) & 0xF).to(tl.float32)
xe = tl.load(scr + SC_PVF + h * 512 + k0 + 2 * tl.arange(0, 32))
xo = tl.load(scr + SC_PVF + h * 512 + k0 + 1 + 2 * tl.arange(0, 32))
dot += tl.sum(lo * xe[:, None], 0) + tl.sum(hi * xo[:, None], 0)
sx += tl.sum(xe) + tl.sum(xo)
acc += sv * (dot - zv * sx)
tl.store(scr + SC_ORA + h * 128 + d0 + tl.arange(0, 32), acc.to(tl.bfloat16).to(tl.float32))
_bump(ctr, ctr_base + 34, TOTAO + 32, pid, ncta, dbg, ctr_base, TIMING)
# ---- MLA OPROJ (split + reduce)
_wait(ctr, ctr_base + 34, TOTAO + 32, dbg, ctr_base, TIMING)
w_o = wqb + tl.load(pb + BASE + MF_OW)
s_o = wfb + tl.load(pb + 256 + BASE + MF_OS)
z_o = wfb + tl.load(pb + 256 + BASE + MF_OZ)
NOP = (2304 // BN) * 4
for it in range(pid, NOP, ncta):
n0 = (it // 4) * BN
sp = it % 4
ns = n0 + tl.arange(0, BN)
acc = _gemv_split(scr, scr + SC_ORA, w_o, s_o, z_o, 2304, 2304, 0.0, wfb, ns, sp, 8, False, BN)
tl.store(scr + SC_OPRP + sp * 2304 + ns, acc)
_bump(ctr, ctr_base + 37, NOP, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 37, NOP, dbg, ctr_base, TIMING)
NQO = 2304 // BN
for it in range(pid, NQO, ncta):
ns = it * BN + tl.arange(0, BN)
a4 = tl.load(scr + SC_OPRP + 0 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 1 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 2 * 2304 + ns)
a4 += tl.load(scr + SC_OPRP + 3 * 2304 + ns)
cur = tl.load(scr + SC_XB + ns)
tl.store(scr + SC_XA + ns, (cur + a4).to(tl.bfloat16).to(tl.float32))
_bump(ctr, ctr_base + 47, NQO, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 47, NQO, dbg, ctr_base, TIMING)
# ---- MLA router partials (folded into MOPR stage)
_xn_make(scr, SC_XN, scr + SC_XA, wfb + tl.load(pb + 256 + BASE + MF_MN), ctr, ctr_base, 47, pid, ncta, dbg, TIMING, False)
mptr = scr + SC_XA
rstd_m = _rstd(mptr, 2304)
nrmm = wfb + tl.load(pb + 256 + BASE + MF_MN)
rwp = wfb + tl.load(pb + 256 + BASE + MF_ROUT)
for it in range(pid, 9, ncta):
kc9 = it * 256
xk9 = kc9 + tl.arange(0, 256)
x9 = tl.load(mptr + xk9).to(tl.float32)
nv9 = tl.load(nrmm + xk9).to(tl.float32)
x9 = ((x9 * rstd_m) * nv9).to(tl.bfloat16).to(tl.float32)
TOTRO = 72
for it in range(pid, TOTRO, ncta):
kc = it // 8
jc = it % 8
rows = jc * 8 + tl.arange(0, 8)
k0 = kc * 256
xk = tl.arange(0, 256)
xe = tl.load(scr + SC_XN + k0 + xk)
wt = tl.load(rwp + rows[:, None] * 2304 + k0 + xk[None, :]).to(tl.float32)
part = tl.sum(wt * xe[None, :], 1)
tl.store(scr + SC_LOGIT + kc * 64 + rows, part)
_bump(ctr, ctr_base + 36, NOP + NQO + TOTRO, pid, ncta, dbg, ctr_base, TIMING)
# ---- MLA RGU: gate/up fused gemv (waits full MOPR incl. router)
_wait(ctr, ctr_base + 36, NOP + NQO + TOTRO, dbg, ctr_base, TIMING)
idxs, wsel = _route(scr + SC_LOGIT)
NGCH = 1024 // BN
TOTG = 9 * NGCH * 6
for it in range(pid, TOTG, ncta):
slot = it // (NGCH * 6)
n0 = ((it // 6) % NGCH) * BN
sp = it % 6
ns = n0 + tl.arange(0, BN)
e = tl.sum(tl.where(tl.arange(0, 8) == slot, idxs, 0))
wg, sg, zg = _moe_in(BASE, pb, slot,
MF_EXW, MF_EXSZ, MF_SGW, MF_SUW, MF_SDW,
MF_SGS, MF_SUS, MF_SDS, MF_SGZ, MF_SUZ, MF_SDZ, 0, e)
wu, su, zu = _moe_in(BASE, pb, slot,
MF_EXW, MF_EXSZ, MF_SGW, MF_SUW, MF_SDW,
MF_SGS, MF_SUS, MF_SDS, MF_SGZ, MF_SUZ, MF_SDZ, 1, e)
ag = tl.zeros((BN,), dtype=tl.float32)
au = tl.zeros((BN,), dtype=tl.float32)
for g in range(sp * 3, sp * 3 + 3):
svg = tl.load(wfb + sg + g * 1024 + ns).to(tl.float32)
zvg = tl.load(wfb + zg + g * 1024 + ns).to(tl.float32)
sug = tl.load(wfb + su + g * 1024 + ns).to(tl.float32)
zug = tl.load(wfb + zu + g * 1024 + ns).to(tl.float32)
dg = tl.zeros((BN,), dtype=tl.float32)
du = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wg_ = tl.load(wqb + wg + (k0 // 2 + r32)[:, None] * 1024 + ns[None, :])
wu_ = tl.load(wqb + wu + (k0 // 2 + r32)[:, None] * 1024 + ns[None, :])
log = (wg_ & 0xF).to(tl.float32)
hig = ((wg_ >> 4) & 0xF).to(tl.float32)
lou = (wu_ & 0xF).to(tl.float32)
hiu = ((wu_ >> 4) & 0xF).to(tl.float32)
xe = tl.load(scr + SC_XN + k0 + 2 * tl.arange(0, 32))
xo = tl.load(scr + SC_XN + k0 + 1 + 2 * tl.arange(0, 32))
dg += tl.sum(log * xe[:, None], 0) + tl.sum(hig * xo[:, None], 0)
du += tl.sum(lou * xe[:, None], 0) + tl.sum(hiu * xo[:, None], 0)
sx += tl.sum(xe) + tl.sum(xo)
ag += svg * (dg - zvg * sx)
au += sug * (du - zug * sx)
tl.store(scr + SC_GUP + ((0 * 9 + slot) * 6 + sp) * 1024 + ns, ag)
tl.store(scr + SC_GUP + ((1 * 9 + slot) * 6 + sp) * 1024 + ns, au)
_bump(ctr, ctr_base + 38, TOTG, pid, ncta, dbg, ctr_base, TIMING)
# ---- MLA DOWN + FINAL
_wait(ctr, ctr_base + 38, TOTG, dbg, ctr_base, TIMING)
NHX = 9 * (1024 // 128)
for it in range(pid, NHX, ncta):
slot = it // (1024 // 128)
n0 = (it % (1024 // 128)) * 128
nsr = n0 + tl.arange(0, 128)
g0 = (tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 0) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 1) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 2) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 3) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 4) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((0 * 9 + slot) * 6 + 5) * 1024 + nsr))
u0 = (tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 0) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 1) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 2) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 3) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 4) * 1024 + nsr)
+ tl.load(scr + SC_GUP + ((1 * 9 + slot) * 6 + 5) * 1024 + nsr))
tl.store(scr + SC_HX + slot * 1024 + nsr, g0 * tl.sigmoid(g0) * u0)
_bump(ctr, ctr_base + 45, NHX, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 45, NHX, dbg, ctr_base, TIMING)
NDN = 2304 // BN
TOTDG = 9 * NDN * 2
for it in range(pid, TOTDG, ncta):
slot = it // (NDN * 2)
n0 = ((it // 2) % NDN) * BN
sp = it % 2
ns = n0 + tl.arange(0, BN)
e = tl.sum(tl.where(tl.arange(0, 8) == slot, idxs, 0))
w, s, z = _moe_in(BASE, pb, slot,
MF_EXW, MF_EXSZ, MF_SGW, MF_SUW, MF_SDW,
MF_SGS, MF_SUS, MF_SDS, MF_SGZ, MF_SUZ, MF_SDZ, 2, e)
acc = tl.zeros((BN,), dtype=tl.float32)
for g in range(sp * 4, sp * 4 + 4):
sv = tl.load(wfb + s + g * 2304 + ns).to(tl.float32)
zv = tl.load(wfb + z + g * 2304 + ns).to(tl.float32)
dot = tl.zeros((BN,), dtype=tl.float32)
sx = 0.0
for kk in tl.static_range(2):
k0 = g * 128 + kk * 64
r32 = tl.arange(0, 32)
wp = tl.load(wqb + w + (k0 // 2 + r32)[:, None] * 2304 + ns[None, :])
lo = (wp & 0xF).to(tl.float32)
hi = ((wp >> 4) & 0xF).to(tl.float32)
he = tl.load(scr + SC_HX + slot * 1024 + k0 + 2 * tl.arange(0, 32))
ho = tl.load(scr + SC_HX + slot * 1024 + k0 + 1 + 2 * tl.arange(0, 32))
dot += tl.sum(lo * he[:, None], 0) + tl.sum(hi * ho[:, None], 0)
sx += tl.sum(he) + tl.sum(ho)
acc += sv * (dot - zv * sx)
tl.store(scr + SC_DNP + (sp * 9 + slot) * 2304 + ns, acc)
_bump(ctr, ctr_base + 41, TOTDG, pid, ncta, dbg, ctr_base, TIMING)
_wait(ctr, ctr_base + 41, TOTDG, dbg, ctr_base, TIMING)
for it in range(pid, NDN, ncta):
ns = it * BN + tl.arange(0, BN)
outv = tl.load(scr + SC_XA + ns)
for slot in tl.static_range(9):
wsel_s = tl.sum(tl.where(tl.arange(0, 8) == slot, wsel, 0.0))
if slot == 8:
wsel_s = 1.0
pv = tl.load(scr + SC_DNP + (0 * 9 + slot) * 2304 + ns)
pv += tl.load(scr + SC_DNP + (1 * 9 + slot) * 2304 + ns)
outv += wsel_s * pv
tl.store(scr + SC_XA + ns, outv)
tl.store(outh_ptr + ns, outv.to(tl.bfloat16))
_bump(ctr, ctr_base + 40, NHX + TOTDG + NDN, pid, ncta, dbg, ctr_base, TIMING)
# --------------------------------------------------------------------------- #
# flat weight materialization
# --------------------------------------------------------------------------- #
def _materialize(model: "Model", dev):
plans = []
def add(fid, region, numel, getter):
plans.append((fid, region, numel, getter))
def qline(fidw, fids, fidz, m):
add(fidw, "u8", m.w_q.numel(), lambda: m.w_q)
add(fids, "f32", m.scales.numel(), lambda: m.scales)
add(fidz, "f32", m.zeros.numel(), lambda: m.zeros)
for b in range(3):
base = b * _40
blk = model.blocks[b]
at = blk.attn
qline(base + TF_QW, base + TF_QS, base + TF_QZ, at.q_proj)
qline(base + TF_KW, base + TF_KS, base + TF_KZ, at.k_proj)
qline(base + TF_VW, base + TF_VS, base + TF_VZ, at.v_proj)
qline(base + TF_GW, base + TF_GS, base + TF_GZ, at.g_proj)
qline(base + TF_OW, base + TF_OS, base + TF_OZ, at.o_proj)
add(base + TF_BETA, "f32", at.beta_proj.weight.numel(), lambda at=at: at.beta_proj.weight)
add(base + TF_CONV, "f32", at.conv_w.numel(), lambda blk=blk: blk.attn.conv_w)
add(base + TF_AN, "f32", 2304, lambda blk=blk: blk.attn_norm)
add(base + TF_MN, "f32", 2304, lambda blk=blk: blk.moe_norm)
add(base + TF_ROUT, "f32", 64 * 2304, lambda blk=blk: blk.moe.router.weight)
add(base + TF_EXW, "u8", 64 * EXB_U8, None)
add(base + TF_EXSZ, "f32", 64 * EXB_F32, None)
for fidw, fids, fidz, mod in (
(TF_SGW, TF_SGS, TF_SGZ, blk.moe.s_gate),
(TF_SUW, TF_SUS, TF_SUZ, blk.moe.s_up),
(TF_SDW, TF_SDS, TF_SDZ, blk.moe.s_down),
):
add(base + fidw, "u8", mod.w_q[0].numel(), lambda mod=mod: mod.w_q[0])
add(base + fids, "f32", mod.scales[0].numel(), lambda mod=mod: mod.scales[0])
add(base + fidz, "f32", mod.zeros[0].numel(), lambda mod=mod: mod.zeros[0])
base = 3 * _40
blk = model.blocks[3]
at = blk.attn
qline(base + MF_QW, base + MF_QS, base + MF_QZ, at.q_proj)
qline(base + MF_KVAW, base + MF_KVAS, base + MF_KVAZ, at.kv_a)
qline(base + MF_KVBW, base + MF_KVBS, base + MF_KVBZ, at.kv_b)
qline(base + MF_OW, base + MF_OS, base + MF_OZ, at.o_proj)
add(base + MF_AN, "f32", 2304, lambda blk=blk: blk.attn_norm)
add(base + MF_MN, "f32", 2304, lambda blk=blk: blk.moe_norm)
add(base + MF_ROUT, "f32", 64 * 2304, lambda blk=blk: blk.moe.router.weight)
add(base + MF_EXW, "u8", 64 * EXB_U8, None)
add(base + MF_EXSZ, "f32", 64 * EXB_F32, None)
for fidw, fids, fidz, mod in (
(MF_SGW, MF_SGS, MF_SGZ, blk.moe.s_gate),
(MF_SUW, MF_SUS, MF_SUZ, blk.moe.s_up),
(MF_SDW, MF_SDS, MF_SDZ, blk.moe.s_down),
):
add(base + fidw, "u8", mod.w_q[0].numel(), lambda mod=mod: mod.w_q[0])
add(base + fids, "f32", mod.scales[0].numel(), lambda mod=mod: mod.scales[0])
add(base + fidz, "f32", mod.zeros[0].numel(), lambda mod=mod: mod.zeros[0])
par = torch.zeros(PAR_LEN, dtype=torch.int32)
off8 = 0
off32 = 0
offs = {}
for fid, region, numel, getter in plans:
if region == "u8":
off8 = _align256(off8)
par[fid] = off8
offs[fid] = off8
off8 += numel
else:
off32 = _align256(off32)
par[256 + fid] = off32 // 4
offs[256 + fid] = off32 // 4
off32 += numel * 4
wq = torch.zeros(off8 + 256, dtype=torch.uint8, device=dev)
wf = torch.zeros(off32 // 4 + 256, dtype=torch.float32, device=dev)
# copies
for fid, region, numel, getter in plans:
if getter is None:
continue
if region == "u8":
o8 = offs[fid]
src = getter().contiguous().view(-1)
wq[o8:o8 + src.numel()].copy_(src)
else:
o32 = offs[256 + fid]
src = getter().float().contiguous().view(-1)
wf[o32:o32 + src.numel()].copy_(src)
# expert tensors
for b in range(3):
base = b * _40
moe = model.blocks[b].moe
o8 = offs[base + TF_EXW]
o32 = offs[256 + base + TF_EXSZ]
for e in range(64):
du8 = o8 + e * EXB_U8
df32 = o32 + e * EXB_F32
for i, mod in enumerate((moe.gate, moe.up, moe.down)):
src = mod.w_q[e].contiguous().view(-1)
wq[du8 + i * GW_BYTES: du8 + i * GW_BYTES + src.numel()].copy_(src)
ss = mod.scales[e].float().contiguous().view(-1)
zz = mod.zeros[e].float().contiguous().view(-1)
wf[df32 + (2 * i) * GS_ELEMS: df32 + (2 * i) * GS_ELEMS + ss.numel()].copy_(ss)
wf[df32 + (2 * i + 1) * GS_ELEMS: df32 + (2 * i + 1) * GS_ELEMS + zz.numel()].copy_(zz)
base = 3 * _40
moe = model.blocks[3].moe
o8 = offs[base + MF_EXW]
o32 = offs[256 + base + MF_EXSZ]
for e in range(64):
du8 = o8 + e * EXB_U8
df32 = o32 + e * EXB_F32
for i, mod in enumerate((moe.gate, moe.up, moe.down)):
src = mod.w_q[e].contiguous().view(-1)
wq[du8 + i * GW_BYTES: du8 + i * GW_BYTES + src.numel()].copy_(src)
ss = mod.scales[e].float().contiguous().view(-1)
zz = mod.zeros[e].float().contiguous().view(-1)
wf[df32 + (2 * i) * GS_ELEMS: df32 + (2 * i) * GS_ELEMS + ss.numel()].copy_(ss)
wf[df32 + (2 * i + 1) * GS_ELEMS: df32 + (2 * i + 1) * GS_ELEMS + zz.numel()].copy_(zz)
model._flat = (wq, wf, par.to(dev))
# --------------------------------------------------------------------------- #
# Model
# --------------------------------------------------------------------------- #
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._flat = None
self._scratch = None
self._win = None
self._ctr = None
self._ctr_base = 0
self._outh = None
self._ncta = None
self.register_load_state_dict_post_hook(_invalidate)
def _win_view(self, layer, pp, qkv):
base = layer * WIN_LAYER_STRIDE + pp * WIN_PP_STRIDE + qkv * 3 * 4096
return self._win[base: base + 3 * 4096].view(3, 4096)
def step(self, hidden, state):
if _DEBUG_EAGER:
return _step_eager(self, hidden, state)
dev = hidden.device
if self._flat is None:
_materialize(self, dev)
self._scratch = torch.zeros(SC_TOTAL, dtype=torch.float32, device=dev)
self._win = torch.zeros(WIN_TOTAL, dtype=torch.float32, device=dev)
self._ctr = torch.zeros(65536, dtype=torch.int32, device=dev)
self._ctr_base = 0
self._outh = torch.zeros(2304, dtype=torch.bfloat16, device=dev)
self._dbg = torch.zeros(128, dtype=torch.int64, device=dev)
wq, wf, par = self._flat
scr = self._scratch
win = self._win
ppflags = 0
for i in range(3):
st = state[i]
if "_pp" in st:
pp = int(st["_pp"]) & 1
ppflags |= (1 << i) | (pp << (3 + i))
mst = state[3]
if "_kv" in mst:
cold_cache = 0
kvbig, krbig = mst["_kv"]
L0 = int(mst["_L"])
ckvin, krin = kvbig, krbig
else:
cold_cache = 1
L0 = int(mst["c_kv"].shape[0])
cap = L0 + 4096
kvbig = torch.empty((cap, 512), dtype=torch.bfloat16, device=dev)
krbig = torch.empty((cap, 64), dtype=torch.bfloat16, device=dev)
ckvin, krin = mst["c_kv"], mst["k_rope"]
achunk = 64 if L0 + 1 <= 4096 else (128 if L0 + 1 <= 8192 else 256)
nc = (L0 + 1 + achunk - 1) // achunk
if self._ctr_base + XR_MAX > self._ctr.numel():
self._ctr = torch.zeros(65536, dtype=torch.int32, device=dev)
self._ctr_base = 0
ctr_base = self._ctr_base
self._ctr_base += XR_MAX
flags = cold_cache
if self._ncta is None:
self._ncta = _pick_ncta(wq, wf, par, scr, win, state, ckvin, krin,
kvbig, krbig, hidden, self._outh, self._ctr,
ctr_base, L0, nc, flags, ppflags, achunk)
ncta = self._ncta
_megakernel[(ncta,)](
wq, wf, par, scr, win,
state[0]["S"], state[1]["S"], state[2]["S"],
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"],
ckvin, krin, kvbig, krbig,
hidden, self._outh, self._ctr, self._dbg,
ctr_base, L0, nc, flags, ppflags,
BN=int(os.environ.get("KIMI_BN", "128")), BR=int(os.environ.get("KIMI_BR", "16")), TIMING=_TIMING,
num_warps=_NUM_WARPS,
)
new_state = list(state)
for i in range(3):
st = state[i]
if "_pp" in st:
npp = 1 - (int(st["_pp"]) & 1)
else:
npp = 0
new_state[i] = {
"S": st["S"],
"cq": self._win_view(i, npp, 0),
"ck": self._win_view(i, npp, 1),
"cv": self._win_view(i, npp, 2),
"_pp": npp,
}
new_state[3] = {
"c_kv": kvbig[: L0 + 1],
"k_rope": krbig[: L0 + 1],
"_kv": (kvbig, krbig),
"_L": L0 + 1,
}
return self._outh, new_state
# ------------------------------------------------------------------ #
# eager debug path (mirrors the reference math exactly)
# ------------------------------------------------------------------ #
def _eager(self, hidden, state):
return _step_eager(self, hidden, state)
def _invalidate(module, incompatible_keys=None, *args, **kwargs):
module._flat = None
module._ncta = None
def _pick_ncta(wq, wf, par, scr, win, state, ckvin, krin, kvbig, krbig,
inh, outh, ctr, ctr_base, L0, nc, flags, ppflags, achunk):
try:
kc = _megakernel.warmup(
wq, wf, par, scr, win,
state[0]["S"], state[1]["S"], state[2]["S"],
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"],
ckvin, krin, kvbig, krbig,
inh, outh, ctr, torch.zeros(128, dtype=torch.int64, device=inh.device),
ctr_base, L0, nc, flags, ppflags, achunk,
BN=128, BR=16, TIMING=0,
grid=(1,), num_warps=_NUM_WARPS,
)
n_regs = getattr(kc, "n_regs", None)
shared = getattr(getattr(kc, "metadata", None), "shared", 0) or 0
if n_regs is None:
return _NCTA_REQ
per_cta = n_regs * _NUM_WARPS * 32
by_regs = max(1, 65536 // max(1, per_cta))
by_smem = max(1, 232448 // max(1, shared)) if shared else 32
by_warps = 64 // _NUM_WARPS
per_sm = max(1, min(by_regs, by_smem, by_warps, 32))
return max(_NSM, min(_NSM * per_sm, _NCTA_REQ))
except Exception:
return _NSM
# --------------------------------------------------------------------------- #
# eager reference (debug only)
# --------------------------------------------------------------------------- #
def _dequant(w_q, scales, zeros, K, group):
wu = torch.empty((K, w_q.shape[1]), dtype=torch.uint8, device=w_q.device)
wu[[REDACTED: IP]] = w_q & 0xF
wu[[REDACTED: IP]] = (w_q >> 4) & 0xF
s = scales.repeat_interleave(group, dim=0)
z = zeros.repeat_interleave(group, dim=0)
return (wu.to(torch.bfloat16) - z) * s
def _rmsnorm(x, w):
xf = x.float()
xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + EPS)
return (xf * w.float()).to(x.dtype)
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)
def _qmm(x, ql):
w = _dequant(ql.w_q, ql.scales, ql.zeros, ql.in_f, ql.group)
return (x.float() @ w.float()).to(torch.bfloat16)
def _kda_eager(at, x, st, cfg):
H, Dk = cfg.kda_heads, cfg.kda_head_dim
q, k, v = _qmm(x, at.q_proj), _qmm(x, at.k_proj), _qmm(x, at.v_proj)
nst = {}
outs = []
for idx, (val, prev) in enumerate(((q, st["cq"]), (k, st["ck"]), (v, st["cv"]))):
win = torch.cat([prev, val[None]], dim=0)
w = at.conv_w[idx].float().transpose(0, 1)
out = F.silu((win.float() * w).sum(0)).to(val.dtype)
outs.append(out)
nst[["cq", "ck", "cv"][idx]] = win[1:]
q, k, v = outs
q = q.view(H, Dk).float() * at.scale
k = k.view(H, Dk).float()
v = v.view(H, Dk).float()
g = (-F.softplus(_qmm(x, at.g_proj).float())).view(H, Dk)
beta = torch.sigmoid(at.beta_proj(x).float())
S = st["S"] * g.exp()[:, :, None]
pred = (S * k[:, :, None]).sum(1)
S = S + beta[:, None, None] * k[:, :, None] * (v - pred)[:, None, :]
o = (S * q[:, :, None]).sum(1)
nst["S"] = S
return _qmm(o.reshape(H * Dk).to(torch.bfloat16), at.o_proj), nst
def _mla_eager(at, x, st, cfg):
H = cfg.mla_heads
pos = st["c_kv"].shape[0]
q = _qmm(x, at.q_proj).view(H, cfg.qk_nope + cfg.qk_rope)
q_nope = q[:, : cfg.qk_nope].float()
q_rope = q[:, cfg.qk_nope:]
kv = _qmm(x, at.kv_a)
c_kv = kv[: cfg.kv_lora]
k_rope = kv[cfg.kv_lora:]
cos, sin = _rope_cossin(pos, cfg.qk_rope, cfg.rope_theta, x.device)
q_rope = _apply_rope(q_rope, cos, sin).float()
k_rope = _apply_rope(k_rope, cos, sin)
nst = {
"c_kv": torch.cat([st["c_kv"], c_kv[None]], 0),
"k_rope": torch.cat([st["k_rope"], k_rope[None]], 0),
}
kvb = _qmm(nst["c_kv"], at.kv_b).view(-1, H, cfg.qk_nope + cfg.v_head).float()
k_nope = kvb[..., : cfg.qk_nope]
v = kvb[..., cfg.qk_nope:]
scores = (torch.einsum("hd,lhd->lh", q_nope, k_nope)
+ torch.einsum("hd,ld->lh", q_rope, nst["k_rope"].float())) * at.scale
p = torch.softmax(scores, dim=0)
o = torch.einsum("lh,lhd->hd", p, v)
return _qmm(o.reshape(H * cfg.v_head).to(torch.bfloat16), at.o_proj), nst
def _moe_eager(moe, x, cfg):
probs = torch.softmax(moe.router(x).float(), dim=-1)
w, idx = torch.topk(probs, cfg.n_active)
w = w / (w.sum() + 1e-9) * cfg.routed_scaling
out = x.new_zeros(cfg.hidden, dtype=torch.float32)
xf = x.float()
for j in range(cfg.n_active):
e = int(idx[j])
h = F.silu(_qmm(x, _ExpertView2(moe.gate, e)).float()) * _qmm(x, _ExpertView2(moe.up, e)).float()
out = out + w[j] * (h @ _dequant(moe.down.w_q[e], moe.down.scales[e], moe.down.zeros[e], moe.down.in_f, moe.down.group).float())
for s in range(cfg.n_shared):
h = F.silu(_qmm(x, _ExpertView2(moe.s_gate, s)).float()) * _qmm(x, _ExpertView2(moe.s_up, s)).float()
out = out + h @ _dequant(moe.s_down.w_q[s], moe.s_down.scales[s], moe.s_down.zeros[s], moe.s_down.in_f, moe.s_down.group).float()
return out.to(torch.bfloat16)
class _ExpertView2:
def __init__(self, qe, e):
self.w_q = qe.w_q[e]
self.scales = qe.scales[e]
self.zeros = qe.zeros[e]
self.in_f = qe.in_f
self.group = qe.group
def _step_eager(model, hidden, state):
with torch.no_grad():
x = hidden.clone()
new_state = []
for i, blk in enumerate(model.blocks):
st = state[i]
xn = _rmsnorm(x, blk.attn_norm)
if blk.kind == "K":
h, nst = _kda_eager(blk.attn, xn, st, model.cfg)
else:
h, nst = _mla_eager(blk.attn, xn, st, model.cfg)
x = x + h
xn = _rmsnorm(x, blk.moe_norm)
x = x + _moe_eager(blk.moe, xn, model.cfg)
new_state.append(nst)
return x, new_state
20260716_233348_kinetic-claude_kinetic-0715_1m__02_kimi_linear_decode