KernelBench mega · RTX PRO 6000
Kimi-Linear Decode GLM-5.3 Flash
13.64×geomean speedup across shapes
manually audited: clean
GLM-5.3 Flash built a genuine persistent single-launch Triton megakernel for the complete Kimi-Linear decode motif. The isolated regrade passed and measured a 13.6418x geometric-mean speedup over the optimized PyTorch baseline. It streams packed int4 weights through fused dequant-GEMVs and uses software device barriers between the KDA, MLA, and MoE phases.
harnessor-fable
Kernel source (redacted)
"""Single-launch W4A16 megakernel for the Kimi-Linear hybrid decode unit (batch 1).
The entire per-token forward -- 3 KDA (gated-delta) layers + 1 MLA layer, each
followed by a 64-expert MoE (top-8 + shared), RMSNorms, residuals, the short
causal depthwise conv, the KDA recurrent-state update, and the MLA latent-cache
attention -- runs as ONE Triton kernel launch (`Model.step` fires exactly one).
How a single launch is possible: the kernel is persistent. The grid has one CTA
per SM (co-resident by construction) and the CTAs walk through ~23 phases
separated by software global barriers (release/acquire spin counters in L2).
Every int4 dequant-GEMV unpacks nibbles and folds the per-group scales/zeros
into the dot product, so the int4 weight stream is read exactly once and a
dequantized bf16 weight matrix is never materialized. MLA uses the absorb path:
kv_b is never evaluated against the context; per-head query latents (one small
GEMV against kv_b's nope slice) dot the cached latents for scores, and the
softmax-weighted latent sum is folded back through the value slice afterwards.
State handling: the KDA recurrent state S is updated in place; short-conv
windows ping-pong between two ring slots (a foreign fed window is read directly
on first touch); the MLA latent cache lives in a preallocated capacity buffer --
a foreign cache is bulk-copied once inside the kernel, later steps append one
row. All source-pointer selection (fed tensor vs ring slot) happens on the host,
so the kernel itself never branches on pointers.
"""
from __future__ import annotations
import os
from dataclasses import dataclass, field
import torch
import torch.nn as nn
import triton
import triton.language as tl
OP_TYPE = "kimi_linear_w4a16_decode"
HARDWARE_REQUIRED = ["RTX_PRO_6000"]
EPS = 1.0e-6
GROUP_SIZE = 128
# ------------------------------------------------------------------ #
# geometry: _PY_* plain ints for host code, tl.constexpr twins for kernels
# ------------------------------------------------------------------ #
_PY_D = 2304
_PY_H = 32
_PY_DK = 128
_PY_HKD = _PY_H * _PY_DK # 4096
_PY_QK_NOPE = 128
_PY_QK_ROPE = 64
_PY_QOUT = _PY_H * (_PY_QK_NOPE + _PY_QK_ROPE) # 6144
_PY_KVL = 512
_PY_KVA = _PY_KVL + _PY_QK_ROPE # 576
_PY_VHD = 128
_PY_KVB_N = _PY_H * (_PY_QK_NOPE + _PY_VHD) # 8192
_PY_M = 1024
_PY_NE = 64
_PY_KP_QKVG = _PY_D // 2 # 1152 packed rows
_PY_NS_QKVG = 4 * _PY_HKD # 16384 columns
_PY_NS_QA = _PY_QOUT + _PY_KVA # 6720 columns
_PY_KP_KVB = _PY_KVL // 2 # 256
_PY_KP_O = _PY_HKD // 2 # o_proj packed rows (K=H*DK)
_PY_GP = GROUP_SIZE // 2 # packed rows per group
_PY_ATT_SCALE = (_PY_QK_NOPE + _PY_QK_ROPE) ** -0.5
# window strides are shared by host views AND kernel address math -> constexpr
_PY_WIN_LSTRIDE = tl.constexpr(2 * 3 * 3 * _PY_HKD) # win elems per KDA layer
_PY_WIN_KSTRIDE = tl.constexpr(3 * 3 * _PY_HKD) # win elems per slot
_PY_WIN_KSTRIDEC = tl.constexpr(3 * _PY_HKD) # win elems per kind
D: tl.constexpr = tl.constexpr(_PY_D)
HKD: tl.constexpr = tl.constexpr(_PY_HKD)
DKC: tl.constexpr = tl.constexpr(_PY_DK)
QOUT: tl.constexpr = tl.constexpr(_PY_QOUT)
KVL: tl.constexpr = tl.constexpr(_PY_KVL)
KP_KVB: tl.constexpr = tl.constexpr(_PY_KP_KVB)
KVB_N: tl.constexpr = tl.constexpr(_PY_KVB_N)
MOEI: tl.constexpr = tl.constexpr(_PY_M)
NE: tl.constexpr = tl.constexpr(_PY_NE)
NS_QKVG: tl.constexpr = tl.constexpr(_PY_NS_QKVG)
KP_QKVG: tl.constexpr = tl.constexpr(_PY_KP_QKVG)
NS_QA: tl.constexpr = tl.constexpr(_PY_NS_QA)
KP_O: tl.constexpr = tl.constexpr(_PY_KP_O)
GP: tl.constexpr = tl.constexpr(_PY_GP)
# fp32 workspace element offsets (constexpr: shared by host allocs and kernels)
# CNT slot holding the router-slice arrival counter; never reset -- each phase
# adds exactly _RS, so the last slice of a phase is the one whose returned old
# value satisfies (old + 1) % _RS == 0
_CNT_RLOG: tl.constexpr = tl.constexpr(22)
_RS: tl.constexpr = tl.constexpr(8) # router slices (8 experts)
O_BETA: tl.constexpr = tl.constexpr(0) # 32 betas
O_ROUTE: tl.constexpr = tl.constexpr(32) # 8 idx + 8 weights
O_QLAT: tl.constexpr = tl.constexpr(O_ROUTE + 16) # 32*512 query latents
O_OLAT: tl.constexpr = tl.constexpr(O_QLAT + 16384) # merged latent summary
O_QR: tl.constexpr = tl.constexpr(O_OLAT + 16384) # 32*64 roped q_rope
O_HH: tl.constexpr = tl.constexpr(O_QR + 2048) # 9*1024 expert intermed
O_ATT: tl.constexpr = tl.constexpr(O_HH + 9216) # partials 128*16*512
O_ATTS: tl.constexpr = tl.constexpr(O_ATT + 128 * 16 * 512) # stats 128*16*2
O_MOACC: tl.constexpr = tl.constexpr(O_ATTS + 128 * 16 * 2) # 2304 moe-down acc
O_RLOG: tl.constexpr = tl.constexpr(O_MOACC + D) # 64 folded router logits
O_RLOGP: tl.constexpr = tl.constexpr(O_RLOG + NE) # _RS per-slice partials
WSF_SIZE = O_RLOGP + _RS * NE
# bf16 workspace element offsets
OB_CONV: tl.constexpr = tl.constexpr(0) # 4*4096 conv'd q,k,v + g
OB_H: tl.constexpr = tl.constexpr(OB_CONV + 16384) # 2304
OB_XP = OB_H + 2304 # spare 2*2304
OB_KO: tl.constexpr = tl.constexpr(OB_XP + 4608) # 4096 kda attn out
OB_QA: tl.constexpr = tl.constexpr(OB_KO + 4096) # 6720 q_all | kv_a
OB_MO: tl.constexpr = tl.constexpr(OB_QA + 6720) # 4096 mla attn out
OB_HID = OB_MO + 4096 # 2304 (spare)
WSB_SIZE = OB_HID + 2304
ATT_LC: tl.constexpr = tl.constexpr(64) # context chunks
ATT_HG: tl.constexpr = tl.constexpr(2) # head groups
ATT_BT: tl.constexpr = tl.constexpr(64) # token tile
ATT_NU: tl.constexpr = tl.constexpr(ATT_LC * ATT_HG)
# per-block offset-table entries (int32; constexpr so kernels can use them)
TB_QW: tl.constexpr = tl.constexpr(0) # q|k|v|g (KDA) or q|kv_a (MLA)
TB_QS: tl.constexpr = tl.constexpr(1)
TB_QZ: tl.constexpr = tl.constexpr(2)
TB_OPW: tl.constexpr = tl.constexpr(3) # o_proj
TB_OPS: tl.constexpr = tl.constexpr(4)
TB_OPZ: tl.constexpr = tl.constexpr(5)
TB_KBW: tl.constexpr = tl.constexpr(6) # kv_b (MLA only)
TB_KBS: tl.constexpr = tl.constexpr(7)
TB_KBZ: tl.constexpr = tl.constexpr(8)
TB_GW: tl.constexpr = tl.constexpr(9) # routed gate experts (region base)
TB_GS: tl.constexpr = tl.constexpr(10)
TB_GZ: tl.constexpr = tl.constexpr(11)
TB_UW: tl.constexpr = tl.constexpr(12) # routed up
TB_US: tl.constexpr = tl.constexpr(13)
TB_UZ: tl.constexpr = tl.constexpr(14)
TB_DW: tl.constexpr = tl.constexpr(15) # routed down
TB_DS: tl.constexpr = tl.constexpr(16)
TB_DZ: tl.constexpr = tl.constexpr(17)
TB_SGW: tl.constexpr = tl.constexpr(18) # shared gate
TB_SGS: tl.constexpr = tl.constexpr(19)
TB_SGZ: tl.constexpr = tl.constexpr(20)
TB_SUW: tl.constexpr = tl.constexpr(21) # shared up
TB_SUS: tl.constexpr = tl.constexpr(22)
TB_SUZ: tl.constexpr = tl.constexpr(23)
TB_SDW: tl.constexpr = tl.constexpr(24) # shared down
TB_SDS: tl.constexpr = tl.constexpr(25)
TB_SDZ: tl.constexpr = tl.constexpr(26)
TB_ROUTER: tl.constexpr = tl.constexpr(27)
TB_ANORM: tl.constexpr = tl.constexpr(28)
TB_MNORM: tl.constexpr = tl.constexpr(29)
TB_CONV: tl.constexpr = tl.constexpr(30) # KDA only
TB_BETA: tl.constexpr = tl.constexpr(31)
TB_STRIDE_W: tl.constexpr = tl.constexpr(32) # per-expert arena strides
TB_STRIDE_SZ: tl.constexpr = tl.constexpr(33)
TB_SIZE = 36
# ==================================================================== #
# device helpers
# ==================================================================== #
@triton.jit
def _fence_gpu():
"""Hard PTX memory fence (not removable by the compiler)."""
tl.inline_asm_elementwise(
"fence.acq_rel.gpu; mov.u32 $0, 0;", "=r", [],
dtype=tl.int32, is_pure=False, pack=1)
@triton.jit
def _nap():
"""Backoff between barrier polls: keeps idle warps from flooding the
memory pipe with volatile requests while peers are still working."""
tl.inline_asm_elementwise(
"nanosleep.u32 512; mov.u32 $0, 0;", "=r", [],
dtype=tl.int32, is_pure=False, pack=1)
@triton.jit
def _bar(cnt_ptr, idx, target):
"""Global software barrier: arrive (release), backed-off spin, then a
hard acquire fence (the fence alone gives the acquire edge -- no extra
contended RMW needed)."""
tl.atomic_add(cnt_ptr + idx, 1, sem="release", scope="gpu")
n = tl.load(cnt_ptr + idx, volatile=True)
while n < target:
_nap()
n = tl.load(cnt_ptr + idx, volatile=True)
_fence_gpu()
@triton.jit
def _tick(pt, off, PROF):
"""KIMI_PROF instrumentation: stash %globaltimer at PT[off] (ns)."""
if PROF:
t = tl.inline_asm_elementwise(
"mov.u64 $0, %globaltimer;", "=l", [],
dtype=tl.int64, is_pure=False, pack=1)
tl.store(pt + off, t)
@triton.jit
def _pf2(ptrs):
"""Fire-and-forget L2 prefetch of a vector of u32 addresses."""
tl.inline_asm_elementwise(
"prefetch.global.L2 [$1]; mov.u32 $0, 0;", "=r,l", [ptrs],
dtype=tl.int32, is_pure=False, pack=1)
@triton.jit
def _softplus(x):
# stable softplus: log(1+e^x) = max(x,0) + log(1+e^{-|x|})
return tl.where(x > 20.0, x, tl.maximum(x, 0.0) + tl.log(1.0 + tl.exp(-tl.abs(x))))
@triton.jit
def _rrms(X):
ssq = 0.0
for d0 in range(0, D, 256):
offs = d0 + tl.arange(0, 256)
v = tl.load(X + offs).to(tl.float32)
ssq += tl.sum(v * v)
return 1.0 / tl.sqrt(ssq / D + 1e-6)
@triton.jit
def _gv(WA, SA, ZA, woff, soff, zoff, XP, NPR, rrms, n_base,
KPU: tl.constexpr, NS: tl.constexpr, BN: tl.constexpr,
USE_NORM: tl.constexpr, XF32: tl.constexpr):
"""Fused int4 dequant-GEMV over a column-major u32-packed matrix.
WT[n, t] packs the eight weights k=8t..8t+7 of output column n into one
32-bit word (nibble j -> weight k=8t+j). A [BN, 16] word tile covers
exactly one scale group (128 k values) and streams as coalesced 32-bit
loads; no dequantized weight is ever materialized.
"""
offs = tl.arange(0, BN)
cols = n_base + offs
tt = tl.arange(0, 32)
acc = tl.zeros([BN], tl.float32)
for t0 in range(0, KPU, 32):
# one whole scale group of words per round: maximal bytes in flight,
# which is what keeps the cold-DRAM stream near bandwidth-bound
V = tl.load(WA + woff + cols[:, None] * KPU + (t0 + tt)[None, :])
acc2 = tl.zeros([BN, 32], tl.float32)
sxp = tl.zeros([32], tl.float32)
for j in tl.static_range(8):
nib = ((V >> (4 * j)) & 0xF).to(tl.float32)
kj = 8 * (t0 + tt) + j
if USE_NORM:
xj = ((tl.load(XP + kj).to(tl.float32) * rrms
* tl.load(NPR + kj).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
elif XF32:
xj = tl.load(XP + kj)
else:
xj = tl.load(XP + kj).to(tl.float32)
acc2 += nib * xj[None, :]
sxp += xj
gsel = (t0 // 16) + tl.arange(0, 2)
sv = tl.load(SA + soff + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
zv = tl.load(ZA + zoff + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
ss = tl.sum(tl.reshape(sxp, [2, 16]), axis=1)
part = tl.sum(tl.reshape(acc2, [BN, 2, 16]), axis=2) - tl.trans(zv) * ss[None, :]
acc += tl.sum(tl.trans(sv) * part, axis=1)
return acc
@triton.jit
def _gv2(WA, SA, ZA, wa, sa, za, wb, sb, zb, XP, NPR, rrms, n_base,
KPU: tl.constexpr, NS: tl.constexpr, BN: tl.constexpr,
USE_NORM: tl.constexpr):
"""_gv for TWO matrices at once: one whole-group word tile per matrix per
round, so four scale groups are folded together per iteration."""
offs = tl.arange(0, BN)
cols = n_base + offs
tt = tl.arange(0, 32)
acca = tl.zeros([BN], tl.float32)
accb = tl.zeros([BN], tl.float32)
for t0 in range(0, KPU, 32):
Va = tl.load(WA + wa + cols[:, None] * KPU + (t0 + tt)[None, :])
Vb = tl.load(WA + wb + cols[:, None] * KPU + (t0 + tt)[None, :])
acc2a = tl.zeros([BN, 32], tl.float32)
acc2b = tl.zeros([BN, 32], tl.float32)
sxp = tl.zeros([32], tl.float32)
for j in tl.static_range(8):
nja = ((Va >> (4 * j)) & 0xF).to(tl.float32)
njb = ((Vb >> (4 * j)) & 0xF).to(tl.float32)
kj = 8 * (t0 + tt) + j
if USE_NORM:
xj = ((tl.load(XP + kj).to(tl.float32) * rrms
* tl.load(NPR + kj).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
else:
xj = tl.load(XP + kj).to(tl.float32)
acc2a += nja * xj[None, :]
acc2b += njb * xj[None, :]
sxp += xj
gsel = (t0 // 16) + tl.arange(0, 2)
sva = tl.load(SA + sa + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
zva = tl.load(ZA + za + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
svb = tl.load(SA + sb + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
zvb = tl.load(ZA + zb + gsel[:, None] * NS + cols[None, :]).to(tl.float32)
ss = tl.sum(tl.reshape(sxp, [2, 16]), axis=1)
parta = tl.sum(tl.reshape(acc2a, [BN, 2, 16]), axis=2) - tl.trans(zva) * ss[None, :]
partb = tl.sum(tl.reshape(acc2b, [BN, 2, 16]), axis=2) - tl.trans(zvb) * ss[None, :]
acca += tl.sum(tl.trans(sva) * parta, axis=1)
accb += tl.sum(tl.trans(svb) * partb, axis=1)
return acca, accb
@triton.jit
def _conv_ep(src_win, dst_win, cw_ptr, cur, ch0):
"""silu(depthwise causal conv over the 4-row window) + window shift."""
offs = ch0 + tl.arange(0, 128)
p0 = tl.load(src_win + 0 * HKD + offs).to(tl.float32)
p1 = tl.load(src_win + 1 * HKD + offs).to(tl.float32)
p2 = tl.load(src_win + 2 * HKD + offs).to(tl.float32)
c0 = tl.load(cw_ptr + 0 * HKD + offs).to(tl.float32)
c1 = tl.load(cw_ptr + 1 * HKD + offs).to(tl.float32)
c2 = tl.load(cw_ptr + 2 * HKD + offs).to(tl.float32)
c3 = tl.load(cw_ptr + 3 * HKD + offs).to(tl.float32)
o = p0 * c0 + p1 * c1 + p2 * c2 + cur.to(tl.float32) * c3
o = o * tl.sigmoid(o)
tl.store(dst_win + 0 * HKD + offs, p1.to(tl.bfloat16))
tl.store(dst_win + 1 * HKD + offs, p2.to(tl.bfloat16))
tl.store(dst_win + 2 * HKD + offs, cur)
return o.to(tl.bfloat16)
# -------------------------------------------------------------------- #
# KDA block: PA gemv+conv, PB recurrence, PC o_proj, PD moe gate/up, PE down
# -------------------------------------------------------------------- #
@triton.jit
def _kda_block(pid, cnt0, cnt_ptr, tgt, flags_off,
WA, SA, ZA, PB, TB, WSF, WSB, WINL,
WQ, WK, WV, Sptr, XIN, XOUT, epoch,
pbase=0, pt=None, PROF=0):
_tick(pt, pid * 64 + pbase + 0, PROF)
# ---- PA: fused q|k|v|g dequant-GEMV + conv epilogue (+ beta on CTA0) ----
_tick(pt, pid * 64 + pbase + 6, PROF)
tqw = tl.load(TB + TB_QW)
tqs = tl.load(TB + TB_QS)
tqz = tl.load(TB + TB_QZ)
tcw = tl.load(TB + TB_CONV)
tan_ = tl.load(TB + TB_ANORM)
rrms = _rrms(XIN)
if pid < NS_QKVG // 128:
n0 = pid * 128
acc = _gv(WA, SA, ZA, tqw, tqs, tqz, XIN, PB + tan_, rrms, n0,
KPU=KP_QKVG // 4, NS=NS_QKVG, BN=128, USE_NORM=True, XF32=False)
ten = n0 // HKD
ch0 = n0 % HKD
if ten == 3:
# g is consumed raw (no conv, no window state)
tl.store(WSB + OB_CONV + 3 * HKD + ch0 + tl.arange(0, 128), acc.to(tl.bfloat16))
else:
if ten == 0:
sw = WQ
cw = PB + tcw
elif ten == 1:
sw = WK
cw = PB + tcw + HKD * 4
else:
sw = WV
cw = PB + tcw + 2 * HKD * 4
cd = _conv_ep(sw, WINL + ten * _PY_WIN_KSTRIDEC, cw, acc.to(tl.bfloat16), ch0)
tl.store(WSB + OB_CONV + ten * HKD + ch0 + tl.arange(0, 128), cd)
if (pid >= 128) and (pid < 128 + D // 72):
# one idle CTA per head folds beta[h] = sigmoid(w_beta[h] . x_norm);
# pid 0 must not serialize this behind its own GEMV tile
tbeta = tl.load(TB + TB_BETA)
hh = pid - 128
bb = 0.0
for d0 in range(0, D, 256):
offs = d0 + tl.arange(0, 256)
xv = ((tl.load(XIN + offs).to(tl.float32) * rrms
* tl.load(PB + tan_ + offs).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
wt = tl.load(PB + tbeta + hh * D + offs).to(tl.float32)
bb += tl.sum(wt * xv)
bb = tl.sigmoid(bb.to(tl.bfloat16).to(tl.float32))
tl.store(WSF + O_BETA + hh, bb)
_tick(pt, pid * 64 + pbase + 7, PROF)
_bar(cnt_ptr, cnt0 + 0, tgt)
_tick(pt, pid * 64 + pbase + 1, PROF)
# ---- PB: per-head gated-delta recurrence (S updated in place) ----------
if pid < 32:
h = pid
hb = WSB + OB_CONV + h * DKC
arv = tl.arange(0, DKC)
q = tl.load(hb + arv).to(tl.float32) * 0.08838834764831845
v = tl.load(hb + 2 * HKD + arv).to(tl.float32)
bh = tl.load(WSF + O_BETA + h)
Sp = Sptr + h * DKC * DKC
pred = tl.zeros([DKC], tl.float32)
for d0 in range(0, DKC, 32):
rr = d0 + tl.arange(0, 32)
gt = tl.exp(-_softplus(tl.load(hb + 3 * HKD + rr).to(tl.float32)))
kt = tl.load(hb + HKD + rr).to(tl.float32)
Sp_t = Sp + rr[:, None] * DKC + arv[None, :]
St = tl.load(Sp_t)
St = St * gt[:, None]
tl.store(Sp_t, St)
pred += tl.sum(St * kt[:, None], axis=0)
upd = v - pred
oacc = tl.zeros([DKC], tl.float32)
for d0 in range(0, DKC, 32):
rr = d0 + tl.arange(0, 32)
qt = tl.load(hb + rr).to(tl.float32) * 0.08838834764831845
kt = tl.load(hb + HKD + rr).to(tl.float32)
Sp_t = Sp + rr[:, None] * DKC + arv[None, :]
St = tl.load(Sp_t)
St = St + (bh * kt)[:, None] * upd[None, :]
tl.store(Sp_t, St)
oacc += tl.sum(St * qt[:, None], axis=0)
tl.store(WSB + OB_KO + h * DKC + arv, oacc.to(tl.bfloat16))
_bar(cnt_ptr, cnt0 + 1, tgt)
_tick(pt, pid * 64 + pbase + 2, PROF)
# ---- PC: o_proj dequant-GEMV (bf16-rounded) + residual -> h ------------
topw = tl.load(TB + TB_OPW)
tops = tl.load(TB + TB_OPS)
topz = tl.load(TB + TB_OPZ)
if pid < D // 64:
n0 = pid * 64
acc = _gv(WA, SA, ZA, topw, tops, topz, WSB + OB_KO, PB, rrms, n0,
KPU=KP_O // 4, NS=D, BN=64, USE_NORM=False, XF32=False)
offs = n0 + tl.arange(0, 64)
op = acc.to(tl.bfloat16)
hv = (op.to(tl.float32) + tl.load(XIN + offs).to(tl.float32)).to(tl.bfloat16)
tl.store(WSB + OB_H + offs, hv)
if (pid >= 64) and (pid < 64 + NE):
# idle during o_proj: L2-prefetch this block's router rows so the PD
# routing slices don't pay cold DRAM for their weight reads
trot = tl.load(TB + TB_ROUTER)
r = pid - 64
vw = tl.arange(0, 256)
for i in range(0, D // 256):
_pf2(PB + trot + r * D + i * 256 + vw)
_bar(cnt_ptr, cnt0 + 2, tgt)
_tick(pt, pid * 64 + pbase + 3, PROF)
# ---- PD: parallel routing slices + routed/shared gate|up GEMVs ----------
tmn = tl.load(TB + TB_MNORM)
if (pid >= 1) and (pid < 1 + _RS):
# slice s folds logits for experts [8s, 8s+8); the last slice to
# arrive (atomic counter, no barrier) softmaxes, picks top-8 and
# publishes the routing flags the gate/up workers spin on
sl = pid - 1
_tick(pt, pid * 64 + pbase + 13, PROF)
rr2 = _rrms(WSB + OB_H)
trot = tl.load(TB + TB_ROUTER)
ex = 8 * sl + tl.arange(0, 8)
# hoist all nine weight chunks above the fold so the DRAM round
# trips overlap instead of serializing behind the accumulator
o0 = 0 + tl.arange(0, 256)
w0 = tl.load(PB + trot + ex[:, None] * D + o0[None, :])
o1 = 256 + tl.arange(0, 256)
w1 = tl.load(PB + trot + ex[:, None] * D + o1[None, :])
o2 = 512 + tl.arange(0, 256)
w2 = tl.load(PB + trot + ex[:, None] * D + o2[None, :])
o3 = 768 + tl.arange(0, 256)
w3 = tl.load(PB + trot + ex[:, None] * D + o3[None, :])
o4 = 1024 + tl.arange(0, 256)
w4 = tl.load(PB + trot + ex[:, None] * D + o4[None, :])
o5 = 1280 + tl.arange(0, 256)
w5 = tl.load(PB + trot + ex[:, None] * D + o5[None, :])
o6 = 1536 + tl.arange(0, 256)
w6 = tl.load(PB + trot + ex[:, None] * D + o6[None, :])
o7 = 1792 + tl.arange(0, 256)
w7 = tl.load(PB + trot + ex[:, None] * D + o7[None, :])
o8 = 2048 + tl.arange(0, 256)
w8 = tl.load(PB + trot + ex[:, None] * D + o8[None, :])
x0 = ((tl.load(WSB + OB_H + o0).to(tl.float32) * rr2
* tl.load(PB + tmn + o0).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x1 = ((tl.load(WSB + OB_H + o1).to(tl.float32) * rr2
* tl.load(PB + tmn + o1).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x2 = ((tl.load(WSB + OB_H + o2).to(tl.float32) * rr2
* tl.load(PB + tmn + o2).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x3 = ((tl.load(WSB + OB_H + o3).to(tl.float32) * rr2
* tl.load(PB + tmn + o3).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x4 = ((tl.load(WSB + OB_H + o4).to(tl.float32) * rr2
* tl.load(PB + tmn + o4).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x5 = ((tl.load(WSB + OB_H + o5).to(tl.float32) * rr2
* tl.load(PB + tmn + o5).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x6 = ((tl.load(WSB + OB_H + o6).to(tl.float32) * rr2
* tl.load(PB + tmn + o6).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x7 = ((tl.load(WSB + OB_H + o7).to(tl.float32) * rr2
* tl.load(PB + tmn + o7).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x8 = ((tl.load(WSB + OB_H + o8).to(tl.float32) * rr2
* tl.load(PB + tmn + o8).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
logits = (
tl.sum(w0.to(tl.float32) * x0[None, :], 1) + tl.sum(w1.to(tl.float32) * x1[None, :], 1) + tl.sum(w2.to(tl.float32) * x2[None, :], 1) + tl.sum(w3.to(tl.float32) * x3[None, :], 1) + tl.sum(w4.to(tl.float32) * x4[None, :], 1) + tl.sum(w5.to(tl.float32) * x5[None, :], 1) + tl.sum(w6.to(tl.float32) * x6[None, :], 1) + tl.sum(w7.to(tl.float32) * x7[None, :], 1) + tl.sum(w8.to(tl.float32) * x8[None, :], 1))
# private partial slot + release-counted arrival; no atomic folds so
# the publisher's fixed-order sum below is deterministic
tl.store(WSF + O_RLOGP + sl * NE + ex, logits)
arrived = tl.atomic_add(cnt_ptr + _CNT_RLOG, 1, sem="acq_rel", scope="gpu")
if arrived % _RS == _RS - 1:
full = tl.zeros([NE], tl.float32)
for s3 in tl.static_range(_RS):
full += tl.load(WSF + O_RLOGP + s3 * NE + tl.arange(0, NE),
volatile=True)
# top-8 selection only needs the ORDER; the full-softmax
# denominator cancels in the top-8 renormalization below
p = full.to(tl.bfloat16).to(tl.float32)
sel = tl.zeros([8], tl.float32)
val = tl.zeros([8], tl.float32)
for j in range(8):
s = tl.argmax(p, axis=0)
mv = tl.max(p, axis=0)
sel = tl.where(tl.arange(0, 8) == j, s.to(tl.float32), sel)
val = tl.where(tl.arange(0, 8) == j, mv, val)
p = tl.where(tl.arange(0, NE) == s, -1.0, p)
val = tl.exp(val - tl.max(val, axis=0))
wn = val / (tl.sum(val, axis=0) + 1e-9) * 2.446
tl.store(WSF + O_ROUTE + tl.arange(0, 8), sel)
tl.store(WSF + O_ROUTE + 8 + tl.arange(0, 8), wn)
tl.atomic_xchg(cnt_ptr + flags_off, epoch * 4 + flags_off + 1,
sem="release", scope="gpu")
_tick(pt, pid * 64 + pbase + 15, PROF)
if (pid >= 1 + _RS) and (pid < 1 + _RS + 9 * (MOEI // 64)):
uid = pid - 1 - _RS
if uid < 8 * (MOEI // 64):
# routed units wait for the flags; the shared expert's weights are
# input-independent, so its units start streaming immediately
ready = tl.load(cnt_ptr + flags_off, volatile=True)
want = epoch * 4 + flags_off + 1
while ready < want:
_nap()
ready = tl.load(cnt_ptr + flags_off, volatile=True)
_fence_gpu()
_tick(pt, pid * 64 + pbase + 10, PROF)
e9 = uid // (MOEI // 64)
jc = (uid % (MOEI // 64)) * 64
if e9 == 8: # shared expert lives in its own single-slot region
gw = tl.load(TB + TB_SGW)
gs = tl.load(TB + TB_SGS)
gz = tl.load(TB + TB_SGZ)
uw = tl.load(TB + TB_SUW)
us = tl.load(TB + TB_SUS)
uz = tl.load(TB + TB_SUZ)
else:
es = tl.load(WSF + O_ROUTE + e9).to(tl.int32)
stw = es * tl.load(TB + TB_STRIDE_W)
ssz = es * tl.load(TB + TB_STRIDE_SZ)
gw = tl.load(TB + TB_GW) + stw
gs = tl.load(TB + TB_GS) + ssz
gz = tl.load(TB + TB_GZ) + ssz
uw = tl.load(TB + TB_UW) + stw
us = tl.load(TB + TB_US) + ssz
uz = tl.load(TB + TB_UZ) + ssz
rr2 = _rrms(WSB + OB_H)
ag, au = _gv2(WA, SA, ZA, gw, gs, gz, uw, us, uz, WSB + OB_H, PB + tmn,
rr2, jc, KPU=KP_QKVG // 4, NS=MOEI, BN=64, USE_NORM=True)
hg = ag * tl.sigmoid(ag)
tl.store(WSF + O_HH + e9 * MOEI + jc + tl.arange(0, 64), hg * au)
_tick(pt, pid * 64 + pbase + 12, PROF)
if pid == 153:
# idle during gate/up: pre-zero the fp32 accumulator PE will atomically
# fold into (ordered before PE by the barrier below)
for z0 in range(0, D, 128):
tl.store(WSF + O_MOACC + z0 + tl.arange(0, 128),
tl.zeros([128], tl.float32))
_bar(cnt_ptr, cnt0 + 3, tgt)
_tick(pt, pid * 64 + pbase + 4, PROF)
# ---- PE: down projection, one (expert, col-tile) unit per CTA ----------
# each of the 9 expert slots x 18 col-tiles runs as its own GEMV over the
# full intermediate axis and atomically folds into the shared accumulator;
# a final pass adds the residual and casts.
sdw = tl.load(TB + TB_SDW)
sds = tl.load(TB + TB_SDS)
sdz = tl.load(TB + TB_SDZ)
if pid < 9 * (D // 128):
e9 = pid // (D // 128)
j = pid % (D // 128)
n0 = j * 128
accp = WSF + O_MOACC + n0 + tl.arange(0, 128)
if e9 == 8:
ad = _gv(WA, SA, ZA, sdw, sds, sdz, WSF + O_HH + e9 * MOEI, PB, rrms,
n0, KPU=(MOEI // 2) // 4, NS=D, BN=128, USE_NORM=False,
XF32=True)
tl.atomic_add(accp, ad)
else:
wv = tl.load(WSF + O_ROUTE + 8 + e9)
es = tl.load(WSF + O_ROUTE + e9).to(tl.int32)
dw = tl.load(TB + TB_DW) + es * tl.load(TB + TB_STRIDE_W)
ds = tl.load(TB + TB_DS) + es * tl.load(TB + TB_STRIDE_SZ)
dz = tl.load(TB + TB_DZ) + es * tl.load(TB + TB_STRIDE_SZ)
ad = _gv(WA, SA, ZA, dw, ds, dz, WSF + O_HH + e9 * MOEI, PB, rrms,
n0, KPU=(MOEI // 2) // 4, NS=D, BN=128, USE_NORM=False,
XF32=True)
tl.atomic_add(accp, wv * ad)
_bar(cnt_ptr, 28 + cnt0 // 5, tgt)
if pid < D // 64:
n0 = pid * 64
offs = n0 + tl.arange(0, 64)
mo = tl.load(WSF + O_MOACC + offs)
ov = (mo.to(tl.bfloat16).to(tl.float32)
+ tl.load(WSB + OB_H + offs).to(tl.float32)).to(tl.bfloat16)
tl.store(XOUT + offs, ov)
_tick(pt, pid * 64 + pbase + 5, PROF)
# -------------------------------------------------------------------- #
# MLA block: PA gemv, PB rope/cache/qlat, PC flash-decode, MRG, PE1 value
# fold, PE2 o_proj, PGU moe gate/up, PE3 moe down
# -------------------------------------------------------------------- #
@triton.jit
def _mla_block(pid, cnt0, cnt_ptr, tgt, flags_off,
WA, SA, ZA, PB, TM, WSF, WSB,
CKVCAP, KRCAP, CKVSRC, KRSRC, XIN, HO,
epoch, FTM, POS, CTXF,
pbase=0, pt=None, PROF=0):
_tick(pt, pid * 64 + pbase + 0, PROF)
# ---- PA: fused q_proj | kv_a dequant-GEMV ------------------------------
tqw = tl.load(TM + TB_QW)
tqs = tl.load(TM + TB_QS)
tqz = tl.load(TM + TB_QZ)
tan_ = tl.load(TM + TB_ANORM)
rrms = _rrms(XIN)
if pid < NS_QA // 64:
n0 = pid * 64
acc = _gv(WA, SA, ZA, tqw, tqs, tqz, XIN, PB + tan_, rrms, n0,
KPU=KP_QKVG // 4, NS=NS_QA, BN=64, USE_NORM=True, XF32=False)
tl.store(WSB + OB_QA + n0 + tl.arange(0, 64), acc.to(tl.bfloat16))
_bar(cnt_ptr, cnt0 + 0, tgt)
_tick(pt, pid * 64 + pbase + 1, PROF)
# ---- PB: rope + cache append (+ first-touch bulk copy) + q latents -----
tkw = tl.load(TM + TB_KBW)
tks = tl.load(TM + TB_KBS)
tkz = tl.load(TM + TB_KBZ)
if pid == 0:
j = tl.arange(0, 32)
inv = tl.exp((-(2.0 * j.to(tl.float32)) / 64.0) * 9.210340371976184)
ang = POS.to(tl.float32) * inv
cs = tl.cos(ang)
sn = tl.sin(ang)
hb = tl.arange(0, 32)[:, None] * 192 + 128
ev = tl.load(WSB + OB_QA + hb + 2 * j[None, :]).to(tl.float32)
od = tl.load(WSB + OB_QA + hb + 2 * j[None, :] + 1).to(tl.float32)
ob = tl.arange(0, 32)[:, None] * 64
tl.store(WSF + O_QR + ob + 2 * j[None, :], ev * cs[None, :] - od * sn[None, :])
tl.store(WSF + O_QR + ob + 2 * j[None, :] + 1, od * cs[None, :] + ev * sn[None, :])
kev = tl.load(WSB + OB_QA + QOUT + KVL + 2 * j).to(tl.float32)
kod = tl.load(WSB + OB_QA + QOUT + KVL + 2 * j + 1).to(tl.float32)
tl.store(KRCAP + POS * 64 + 2 * j, (kev * cs - kod * sn).to(tl.bfloat16))
tl.store(KRCAP + POS * 64 + 2 * j + 1, (kod * cs + kev * sn).to(tl.bfloat16))
cv = tl.load(WSB + OB_QA + QOUT + tl.arange(0, 512))
tl.store(CKVCAP + POS * 512 + tl.arange(0, 512), cv)
if (pid >= 128) and (FTM != 0):
r = pid - 128
while r < CTXF:
cv = tl.load(CKVSRC + r * 512 + tl.arange(0, 512))
tl.store(CKVCAP + r * 512 + tl.arange(0, 512), cv)
kv = tl.load(KRSRC + r * 64 + tl.arange(0, 64))
tl.store(KRCAP + r * 64 + tl.arange(0, 64), kv)
r += 60
if pid < 128:
h = pid // 4
n0 = h * 256 + (pid % 4) * 128
acc = _gv(WA, SA, ZA, tkw, tks, tkz, WSB + OB_QA + h * 192, PB, rrms, n0,
KPU=KP_KVB // 4, NS=KVB_N, BN=128, USE_NORM=False, XF32=False)
tl.store(WSF + O_QLAT + h * KVL + (pid % 4) * 128 + tl.arange(0, 128), acc)
_bar(cnt_ptr, cnt0 + 1, tgt)
_tick(pt, pid * 64 + pbase + 2, PROF)
# ---- PC: flash-decode absorb attention over the latent cache -----------
if pid < ATT_NU:
T = POS + 1
cl = ((T + ATT_LC * ATT_BT - 1) // (ATT_LC * ATT_BT)) * ATT_BT
lci = pid // ATT_HG
gh = pid % ATT_HG
hl = tl.arange(0, 16)
heads = gh * 16 + hl
lo0 = lci * cl
hi1 = tl.minimum(lo0 + cl, T)
qrb = tl.load(WSF + O_QR + heads[:, None] * 64 + tl.arange(0, 64)[None, :]).to(tl.bfloat16)
qlb = tl.load(WSF + O_QLAT + heads[:, None] * KVL + tl.arange(0, KVL)[None, :]).to(tl.bfloat16)
qlT = tl.trans(qlb)
m_i = tl.full([16], -1.0e30, tl.float32)
l_i = tl.zeros([16], tl.float32)
acc = tl.zeros([16, KVL], tl.float32)
for lo in range(lo0, hi1, ATT_BT):
ls = lo + tl.arange(0, ATT_BT)
msk = ls < hi1
ckv = tl.load(CKVCAP + ls[:, None] * 512 + tl.arange(0, 512)[None, :],
mask=msk[:, None], other=0.0)
kr = tl.load(KRCAP + ls[:, None] * 64 + tl.arange(0, 64)[None, :],
mask=msk[:, None], other=0.0)
sc = (tl.dot(ckv, qlT) + tl.dot(kr, tl.trans(qrb))) * 0.07216878364870323
sc = tl.where(msk[:, None], sc, -1.0e30)
m_new = tl.maximum(m_i, tl.max(sc, axis=0))
alpha = tl.exp(m_i - m_new)
pmat = tl.exp(sc - m_new[None, :])
l_i = l_i * alpha + tl.sum(pmat, axis=0)
acc = acc * alpha[:, None] + tl.dot(tl.trans(pmat.to(tl.bfloat16)), ckv)
m_i = m_new
tl.store(WSF + O_ATT + pid * 16 * KVL + hl[:, None] * KVL + tl.arange(0, KVL)[None, :], acc)
tl.store(WSF + O_ATTS + pid * 32 + hl * 2, m_i)
tl.store(WSF + O_ATTS + pid * 32 + hl * 2 + 1, l_i)
_bar(cnt_ptr, cnt0 + 2, tgt)
_tick(pt, pid * 64 + pbase + 3, PROF)
# ---- merge partials -> per-head latent summary --------------------------
if pid < 128:
hh_ = pid // 4
gh = hh_ // 16
hl = hh_ % 16
offs = (pid % 4) * 128 + tl.arange(0, 128)
M = -1.0e30
for c in range(ATT_LC):
M = tl.maximum(M, tl.load(WSF + O_ATTS + (c * ATT_HG + gh) * 32 + hl * 2))
num = tl.zeros([128], tl.float32)
den = 0.0
for c in range(ATT_LC):
sl = c * ATT_HG + gh
m_c = tl.load(WSF + O_ATTS + sl * 32 + hl * 2)
l_c = tl.load(WSF + O_ATTS + sl * 32 + hl * 2 + 1)
a = tl.exp(m_c - M)
ac = tl.load(WSF + O_ATT + sl * 16 * KVL + hl * KVL + offs)
num += a * ac
den += a * l_c
tl.store(WSF + O_OLAT + hh_ * KVL + offs, num / den)
_bar(cnt_ptr, cnt0 + 3, tgt)
_tick(pt, pid * 64 + pbase + 4, PROF)
# ---- fold latent back through kv_b value slice -> attn out --------------
if pid < 32:
acc = _gv(WA, SA, ZA, tkw, tks, tkz, WSF + O_OLAT + pid * KVL, PB, rrms,
pid * 256 + 128, KPU=KP_KVB // 4, NS=KVB_N, BN=128,
USE_NORM=False, XF32=True)
tl.store(WSB + OB_MO + pid * 128 + tl.arange(0, 128), acc.to(tl.bfloat16))
_bar(cnt_ptr, cnt0 + 4, tgt)
_tick(pt, pid * 64 + pbase + 5, PROF)
# ---- o_proj (bf16-rounded) + residual -> h ------------------------------
topw = tl.load(TM + TB_OPW)
tops = tl.load(TM + TB_OPS)
topz = tl.load(TM + TB_OPZ)
if pid < D // 64:
n0 = pid * 64
acc = _gv(WA, SA, ZA, topw, tops, topz, WSB + OB_MO, PB, rrms, n0,
KPU=KP_O // 4, NS=D, BN=64, USE_NORM=False, XF32=False)
offs = n0 + tl.arange(0, 64)
op = acc.to(tl.bfloat16)
hv = (op.to(tl.float32) + tl.load(XIN + offs).to(tl.float32)).to(tl.bfloat16)
tl.store(WSB + OB_H + offs, hv)
if (pid >= 64) and (pid < 64 + NE):
# L2-prefetch the MLA block's router rows ahead of the routing slices
trot = tl.load(TM + TB_ROUTER)
r = pid - 64
vw = tl.arange(0, 256)
for i in range(0, D // 256):
_pf2(PB + trot + r * D + i * 256 + vw)
_bar(cnt_ptr, cnt0 + 5, tgt)
_tick(pt, pid * 64 + pbase + 6, PROF)
# ---- parallel routing slices + gate/up ----------------------------------
tmn = tl.load(TM + TB_MNORM)
if (pid >= 1) and (pid < 1 + _RS):
# router slice: fold 8 experts' logits; last arrival publishes flags
sl = pid - 1
_tick(pt, pid * 64 + pbase + 13, PROF)
rr2 = _rrms(WSB + OB_H)
trot = tl.load(TM + TB_ROUTER)
ex = 8 * sl + tl.arange(0, 8)
# hoist all nine weight chunks above the fold so the DRAM round
# trips overlap instead of serializing behind the accumulator
o0 = 0 + tl.arange(0, 256)
w0 = tl.load(PB + trot + ex[:, None] * D + o0[None, :])
o1 = 256 + tl.arange(0, 256)
w1 = tl.load(PB + trot + ex[:, None] * D + o1[None, :])
o2 = 512 + tl.arange(0, 256)
w2 = tl.load(PB + trot + ex[:, None] * D + o2[None, :])
o3 = 768 + tl.arange(0, 256)
w3 = tl.load(PB + trot + ex[:, None] * D + o3[None, :])
o4 = 1024 + tl.arange(0, 256)
w4 = tl.load(PB + trot + ex[:, None] * D + o4[None, :])
o5 = 1280 + tl.arange(0, 256)
w5 = tl.load(PB + trot + ex[:, None] * D + o5[None, :])
o6 = 1536 + tl.arange(0, 256)
w6 = tl.load(PB + trot + ex[:, None] * D + o6[None, :])
o7 = 1792 + tl.arange(0, 256)
w7 = tl.load(PB + trot + ex[:, None] * D + o7[None, :])
o8 = 2048 + tl.arange(0, 256)
w8 = tl.load(PB + trot + ex[:, None] * D + o8[None, :])
x0 = ((tl.load(WSB + OB_H + o0).to(tl.float32) * rr2
* tl.load(PB + tmn + o0).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x1 = ((tl.load(WSB + OB_H + o1).to(tl.float32) * rr2
* tl.load(PB + tmn + o1).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x2 = ((tl.load(WSB + OB_H + o2).to(tl.float32) * rr2
* tl.load(PB + tmn + o2).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x3 = ((tl.load(WSB + OB_H + o3).to(tl.float32) * rr2
* tl.load(PB + tmn + o3).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x4 = ((tl.load(WSB + OB_H + o4).to(tl.float32) * rr2
* tl.load(PB + tmn + o4).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x5 = ((tl.load(WSB + OB_H + o5).to(tl.float32) * rr2
* tl.load(PB + tmn + o5).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x6 = ((tl.load(WSB + OB_H + o6).to(tl.float32) * rr2
* tl.load(PB + tmn + o6).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x7 = ((tl.load(WSB + OB_H + o7).to(tl.float32) * rr2
* tl.load(PB + tmn + o7).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
x8 = ((tl.load(WSB + OB_H + o8).to(tl.float32) * rr2
* tl.load(PB + tmn + o8).to(tl.float32)).to(tl.bfloat16)).to(tl.float32)
logits = (
tl.sum(w0.to(tl.float32) * x0[None, :], 1) + tl.sum(w1.to(tl.float32) * x1[None, :], 1) + tl.sum(w2.to(tl.float32) * x2[None, :], 1) + tl.sum(w3.to(tl.float32) * x3[None, :], 1) + tl.sum(w4.to(tl.float32) * x4[None, :], 1) + tl.sum(w5.to(tl.float32) * x5[None, :], 1) + tl.sum(w6.to(tl.float32) * x6[None, :], 1) + tl.sum(w7.to(tl.float32) * x7[None, :], 1) + tl.sum(w8.to(tl.float32) * x8[None, :], 1))
# private partial slot + release-counted arrival; no atomic folds so
# the publisher's fixed-order sum below is deterministic
tl.store(WSF + O_RLOGP + sl * NE + ex, logits)
arrived = tl.atomic_add(cnt_ptr + _CNT_RLOG, 1, sem="acq_rel", scope="gpu")
if arrived % _RS == _RS - 1:
full = tl.zeros([NE], tl.float32)
for s3 in tl.static_range(_RS):
full += tl.load(WSF + O_RLOGP + s3 * NE + tl.arange(0, NE),
volatile=True)
# top-8 selection only needs the ORDER; the full-softmax
# denominator cancels in the top-8 renormalization below
p = full.to(tl.bfloat16).to(tl.float32)
sel = tl.zeros([8], tl.float32)
val = tl.zeros([8], tl.float32)
for j in range(8):
s = tl.argmax(p, axis=0)
mv = tl.max(p, axis=0)
sel = tl.where(tl.arange(0, 8) == j, s.to(tl.float32), sel)
val = tl.where(tl.arange(0, 8) == j, mv, val)
p = tl.where(tl.arange(0, NE) == s, -1.0, p)
val = tl.exp(val - tl.max(val, axis=0))
wn = val / (tl.sum(val, axis=0) + 1e-9) * 2.446
tl.store(WSF + O_ROUTE + tl.arange(0, 8), sel)
tl.store(WSF + O_ROUTE + 8 + tl.arange(0, 8), wn)
tl.atomic_xchg(cnt_ptr + flags_off, epoch * 4 + flags_off + 1,
sem="release", scope="gpu")
_tick(pt, pid * 64 + pbase + 15, PROF)
if (pid >= 1 + _RS) and (pid < 1 + _RS + 9 * (MOEI // 64)):
uid = pid - 1 - _RS
if uid < 8 * (MOEI // 64):
# routed units wait for the flags; the shared expert's weights are
# input-independent, so its units start streaming immediately
ready = tl.load(cnt_ptr + flags_off, volatile=True)
want = epoch * 4 + flags_off + 1
while ready < want:
_nap()
ready = tl.load(cnt_ptr + flags_off, volatile=True)
_fence_gpu()
e9 = uid // (MOEI // 64)
jc = (uid % (MOEI // 64)) * 64
if e9 == 8:
gw = tl.load(TM + TB_SGW)
gs = tl.load(TM + TB_SGS)
gz = tl.load(TM + TB_SGZ)
uw = tl.load(TM + TB_SUW)
us = tl.load(TM + TB_SUS)
uz = tl.load(TM + TB_SUZ)
else:
es = tl.load(WSF + O_ROUTE + e9).to(tl.int32)
stw = es * tl.load(TM + TB_STRIDE_W)
ssz = es * tl.load(TM + TB_STRIDE_SZ)
gw = tl.load(TM + TB_GW) + stw
gs = tl.load(TM + TB_GS) + ssz
gz = tl.load(TM + TB_GZ) + ssz
uw = tl.load(TM + TB_UW) + stw
us = tl.load(TM + TB_US) + ssz
uz = tl.load(TM + TB_UZ) + ssz
rr2 = _rrms(WSB + OB_H)
ag, au = _gv2(WA, SA, ZA, gw, gs, gz, uw, us, uz, WSB + OB_H, PB + tmn,
rr2, jc, KPU=KP_QKVG // 4, NS=MOEI, BN=64, USE_NORM=True)
hg = ag * tl.sigmoid(ag)
tl.store(WSF + O_HH + e9 * MOEI + jc + tl.arange(0, 64), hg * au)
if (pid >= 153) and (pid < 153 + D // 128):
# idle during gate/up: pre-zero the fp32 accumulator for the down phase
z0 = (pid - 153) * 128
tl.store(WSF + O_MOACC + z0 + tl.arange(0, 128),
tl.zeros([128], tl.float32))
_bar(cnt_ptr, cnt0 + 6, tgt)
_tick(pt, pid * 64 + pbase + 7, PROF)
# ---- down + residual -> final hidden (one unit per CTA, atomic fold) ----
sdw = tl.load(TM + TB_SDW)
sds = tl.load(TM + TB_SDS)
sdz = tl.load(TM + TB_SDZ)
if pid < 9 * (D // 128):
e9 = pid // (D // 128)
j = pid % (D // 128)
n0 = j * 128
accp = WSF + O_MOACC + n0 + tl.arange(0, 128)
if e9 == 8:
ad = _gv(WA, SA, ZA, sdw, sds, sdz, WSF + O_HH + e9 * MOEI, PB, rrms,
n0, KPU=(MOEI // 2) // 4, NS=D, BN=128, USE_NORM=False,
XF32=True)
tl.atomic_add(accp, ad)
else:
wv = tl.load(WSF + O_ROUTE + 8 + e9)
es = tl.load(WSF + O_ROUTE + e9).to(tl.int32)
dw = tl.load(TM + TB_DW) + es * tl.load(TM + TB_STRIDE_W)
ds = tl.load(TM + TB_DS) + es * tl.load(TM + TB_STRIDE_SZ)
dz = tl.load(TM + TB_DZ) + es * tl.load(TM + TB_STRIDE_SZ)
ad = _gv(WA, SA, ZA, dw, ds, dz, WSF + O_HH + e9 * MOEI, PB, rrms,
n0, KPU=(MOEI // 2) // 4, NS=D, BN=128, USE_NORM=False,
XF32=True)
tl.atomic_add(accp, wv * ad)
_bar(cnt_ptr, 31, tgt)
if pid < D // 64:
n0 = pid * 64
offs = n0 + tl.arange(0, 64)
mo = tl.load(WSF + O_MOACC + offs)
ov = (mo.to(tl.bfloat16).to(tl.float32)
+ tl.load(WSB + OB_H + offs).to(tl.float32)).to(tl.bfloat16)
tl.store(HO + offs, ov)
@triton.jit(do_not_specialize=["epoch", "G", "FTM", "POS", "CTXF"])
def _mega(WA, SA, ZA, PB, T0, T1, T2, T3, WSF, WSB, CNT,
HI, HO, XP0, XP1, WIN,
W0Q, W0K, W0V, W1Q, W1K, W1V, W2Q, W2K, W2V,
S0, S1, S2,
CKVCAP, KRCAP, CKVSRC, KRSRC,
epoch, G, FTM, POS, CTXF, pt=None, PROF=0):
pid = tl.program_id(0)
tgt = (epoch + 1) * G
_tick(pt, pid * 64 + 0, PROF)
_kda_block(pid, 0, CNT, tgt, 24, WA, SA, ZA, PB, T0, WSF, WSB,
WIN, W0Q, W0K, W0V, S0, HI, XP0, epoch, 1, pt, PROF)
_bar(CNT, 4, tgt)
_kda_block(pid, 5, CNT, tgt, 25, WA, SA, ZA, PB, T1, WSF, WSB,
WIN + _PY_WIN_LSTRIDE, W1Q, W1K, W1V, S1, XP0, XP1, epoch,
17, pt, PROF)
_bar(CNT, 9, tgt)
_kda_block(pid, 10, CNT, tgt, 26, WA, SA, ZA, PB, T2, WSF, WSB,
WIN + 2 * _PY_WIN_LSTRIDE, W2Q, W2K, W2V, S2, XP1, XP0, epoch,
33, pt, PROF)
_bar(CNT, 14, tgt)
_mla_block(pid, 15, CNT, tgt, 27, WA, SA, ZA, PB, T3, WSF, WSB,
CKVCAP, KRCAP, CKVSRC, KRSRC, XP0, HO, epoch, FTM, POS, CTXF,
49, pt, PROF)
_tick(pt, pid * 64 + 58, PROF)
# ==================================================================== #
# host side
# ==================================================================== #
@dataclass(frozen=True)
class Config:
hidden: int = 2304
kda_heads: int = 32
kda_head_dim: int = 128
short_conv: int = 4
mla_heads: int = 32
kv_lora: int = 512
qk_nope: int = 128
qk_rope: int = 64
v_head: int = 128
rope_theta: float = 10000.0
n_experts: int = 64
n_active: int = 8
n_shared: int = 1
moe_inter: int = 1024
routed_scaling: float = 2.446
group: int = 128
pattern: tuple = ("K", "K", "K", "M")
dtype: torch.dtype = field(default=torch.bfloat16)
def build_config(shape: dict) -> Config:
return Config(n_experts=int(shape.get("n_experts", 64)))
def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
lo = w_q[0::2] & 0xF
hi = w_q[1::2] & 0xF
return (lo | (hi << 4)).contiguous()
def quantize(w_io: torch.Tensor, group: int = GROUP_SIZE):
K, N = w_io.shape
ng = K // group
wg = w_io.view(ng, group, N).float()
wmin = wg.min(dim=1, keepdim=True).values
wmax = wg.max(dim=1, keepdim=True).values
scales = (wmax - wmin).clamp_min(1e-8) / 15.0
zeros = (-wmin / scales).round().clamp(0, 15)
w_q = ((wg / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
return _pack_int4(w_q), scales.squeeze(1).to(torch.bfloat16), zeros.squeeze(1).to(torch.bfloat16)
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))
def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
w = torch.randn(self.in_f, self.out_f, generator=gen) * std
wq, s, z = quantize(w, self.group)
self.w_q.copy_(wq)
self.scales.copy_(s)
self.zeros.copy_(z)
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))
def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
for e in range(self.n):
w = torch.randn(self.in_f, self.out_f, generator=gen) * std
wq, s, z = quantize(w, self.group)
self.w_q[e].copy_(wq)
self.scales[e].copy_(s)
self.zeros[e].copy_(z)
class KDA(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
d, Hd = cfg.hidden, cfg.kda_heads * cfg.kda_head_dim
self.q_proj = QuantLinear(d, Hd, cfg.group)
self.k_proj = QuantLinear(d, Hd, cfg.group)
self.v_proj = QuantLinear(d, Hd, cfg.group)
self.g_proj = QuantLinear(d, Hd, cfg.group)
self.beta_proj = nn.Linear(d, cfg.kda_heads, bias=False, dtype=cfg.dtype)
self.conv_w = nn.Parameter(torch.empty(3, Hd, cfg.short_conv, dtype=cfg.dtype))
self.o_proj = QuantLinear(Hd, d, cfg.group)
self.scale = cfg.kda_head_dim ** -0.5
class MLA(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
d = cfg.hidden
self.q_proj = QuantLinear(d, cfg.mla_heads * (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, cfg.mla_heads * (cfg.qk_nope + cfg.v_head), cfg.group)
self.o_proj = QuantLinear(cfg.mla_heads * cfg.v_head, d, cfg.group)
self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5
class MoE(nn.Module):
def __init__(self, cfg: Config):
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: Config, 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)
class Model(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
self.reset_parameters()
self._rt = None
self._epoch = 0
self._wpar = [0, 0, 0]
self._clen = 0
def reset_parameters(self):
g = torch.Generator(device="cpu").manual_seed(1234)
for mod in self.modules():
if isinstance(mod, (QuantLinear, QuantExperts)):
mod.init_random(g)
elif isinstance(mod, nn.Linear):
nn.init.normal_(mod.weight, 0.0, 0.02, generator=g)
elif isinstance(mod, KDA):
nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g)
def load_state_dict(self, *a, **kw): # arenas must be rebuilt from new weights
self._rt = None
self._epoch = 0
self._wpar = [0, 0, 0]
self._clen = 0
return super().load_state_dict(*a, **kw)
# ---- arena/table construction ------------------------------------------ #
def _ensure_runtime(self, dev):
if self._rt is not None:
return self._rt
assert tuple(self.cfg.pattern) == ("K", "K", "K", "M")
blk = self.blocks
wa, sa, za, pb = [], [], [], []
def push_w(ws, ss, zs):
off = (sum(x.numel() for x in wa),
sum(x.numel() for x in sa),
sum(x.numel() for x in za))
# repack to column-major u32: WT[n, t] holds the 8 nibbles
# k=8t..8t+7 of column n; same total bytes as the byte-packed form
k2, n = ws.shape[-2], ws.shape[-1]
wtn = ws.reshape(-1, k2, n).transpose(-2, -1).contiguous()
bb = wtn.view(-1, k2 // 4, 4).int()
wt = (bb[..., 0] | (bb[..., 1] << 8)
| (bb[..., 2] << 16) | (bb[..., 3] << 24))
wa.append(wt.reshape(-1))
sa.append(ss.reshape(-1))
za.append(zs.reshape(-1))
return off
def lin(mods, tbl, i):
ws = torch.cat([m.w_q for m in mods], dim=1).contiguous()
ss = torch.cat([m.scales for m in mods], dim=1).contiguous()
zs = torch.cat([m.zeros for m in mods], dim=1).contiguous()
tbl[i], tbl[i + 1], tbl[i + 2] = push_w(ws, ss, zs)
def exp_(mod, tbl, i):
tbl[i], tbl[i + 1], tbl[i + 2] = push_w(mod.w_q, mod.scales, mod.zeros)
def pb_push(t):
off = sum(x.numel() for x in pb)
pb.append(t.detach().reshape(-1).contiguous())
return off
def moe_tbl(moe, tbl):
exp_(moe.gate, tbl, TB_GW)
exp_(moe.up, tbl, TB_UW)
exp_(moe.down, tbl, TB_DW)
exp_(moe.s_gate, tbl, TB_SGW)
exp_(moe.s_up, tbl, TB_SUW)
exp_(moe.s_down, tbl, TB_SDW)
tbl[TB_ROUTER] = pb_push(moe.router.weight)
tbl[TB_STRIDE_W] = (_PY_D // 2) * _PY_M // 4 # u32 words
tbl[TB_STRIDE_SZ] = (_PY_D // GROUP_SIZE) * _PY_M
# one table (+ own arena regions) per block -- weights differ per block
def kda_tbl(b):
tbl = [0] * TB_SIZE
ak = blk[b].attn
lin([ak.q_proj, ak.k_proj, ak.v_proj, ak.g_proj], tbl, TB_QW)
lin([ak.o_proj], tbl, TB_OPW)
moe_tbl(blk[b].moe, tbl)
tbl[TB_CONV] = pb_push(ak.conv_w.permute(0, 2, 1))
tbl[TB_BETA] = pb_push(ak.beta_proj.weight)
tbl[TB_ANORM] = pb_push(blk[b].attn_norm)
tbl[TB_MNORM] = pb_push(blk[b].moe_norm)
return tbl
def mla_tbl():
tbl = [0] * TB_SIZE
bm = blk[3].attn
lin([bm.q_proj, bm.kv_a], tbl, TB_QW)
lin([bm.kv_b], tbl, TB_KBW)
lin([bm.o_proj], tbl, TB_OPW)
moe_tbl(blk[3].moe, tbl)
tbl[TB_ANORM] = pb_push(blk[3].attn_norm)
tbl[TB_MNORM] = pb_push(blk[3].moe_norm)
return tbl
T0, T1, T2 = kda_tbl(0), kda_tbl(1), kda_tbl(2)
T3 = mla_tbl()
rt = {
"WA": torch.cat(wa).contiguous().to(dev),
"SA": torch.cat(sa).contiguous().to(dev),
"ZA": torch.cat(za).contiguous().to(dev),
"PB": torch.cat(pb).contiguous().to(dev),
"T0": torch.tensor(T0, dtype=torch.int32, device=dev),
"T1": torch.tensor(T1, dtype=torch.int32, device=dev),
"T2": torch.tensor(T2, dtype=torch.int32, device=dev),
"T3": torch.tensor(T3, dtype=torch.int32, device=dev),
"WSF": torch.zeros(WSF_SIZE, dtype=torch.float32, device=dev),
"WSB": torch.zeros(WSB_SIZE, dtype=torch.bfloat16, device=dev),
"CNT": torch.zeros(32, dtype=torch.int32, device=dev),
"PROF": 1 if os.getenv("KIMI_PROF") else 0,
"PT": (torch.zeros(int(os.getenv("KIMI_GMUL", "1"))
* torch.cuda.get_device_properties(dev.index or 0)
.multi_processor_count * 64,
dtype=torch.int64, device=dev)
if os.getenv("KIMI_PROF") else None),
"HO": torch.empty(_PY_D, dtype=torch.bfloat16, device=dev),
"XP0": torch.empty(_PY_D, dtype=torch.bfloat16, device=dev),
"XP1": torch.empty(_PY_D, dtype=torch.bfloat16, device=dev),
"capc": None, "capk": None, "capn": 0,
"G": torch.cuda.get_device_properties(dev.index or 0).multi_processor_count
* int(os.getenv("KIMI_GMUL", "1")),
}
WIN = torch.zeros(3 * _PY_WIN_LSTRIDE, dtype=torch.bfloat16, device=dev)
rt["WINL"] = [WIN[l * _PY_WIN_LSTRIDE:(l + 1) * _PY_WIN_LSTRIDE] for l in range(3)]
wv = []
for l in range(3):
per_layer = []
for slot in range(2):
per_slot = []
for k in range(3):
o = l * _PY_WIN_LSTRIDE + slot * _PY_WIN_KSTRIDE + k * _PY_WIN_KSTRIDEC
per_slot.append(WIN[o:o + _PY_WIN_KSTRIDEC])
per_layer.append(per_slot)
wv.append(per_layer)
rt["WV"] = wv
rt["WIN"] = WIN
self._rt = rt
return rt
def step(self, hidden, state):
rt = self._rt
if rt is None:
rt = self._ensure_runtime(hidden.device)
# resolve KDA conv-window sources: our ring slot or the foreign tensor
wsrc = []
dsts = []
for l in range(3):
stl = state[l]
cur = rt["WV"][l][self._wpar[l]]
if stl["cq"].data_ptr() == cur[0].data_ptr():
d = 1 - self._wpar[l]
wsrc.extend(cur)
else:
d = 0
wsrc.extend((stl["cq"], stl["ck"], stl["cv"]))
self._wpar[l] = d
dsts.append(d)
# resolve MLA cache: our capacity buffer or a foreign cache
stm = state[3]
fc, fk = stm["c_kv"], stm["k_rope"]
capc, capk = rt["capc"], rt["capk"]
ctxf = fc.shape[0]
if capc is not None and fc.data_ptr() == capc.data_ptr() and fk.data_ptr() == capk.data_ptr():
ftm = 0
pos = self._clen
else:
ftm = 1
pos = ctxf
if rt["capn"] < ctxf + 64:
need = ctxf + 4096
rt["capc"] = torch.empty(need, _PY_KVL, dtype=torch.bfloat16, device=fc.device)
rt["capk"] = torch.empty(need, _PY_QK_ROPE, dtype=torch.bfloat16, device=fk.device)
rt["capn"] = need
capc, capk = rt["capc"], rt["capk"]
self._clen = pos + 1
epoch = self._epoch
_mega[(rt["G"],)](
rt["WA"], rt["SA"], rt["ZA"], rt["PB"],
rt["T0"], rt["T1"], rt["T2"], rt["T3"],
rt["WSF"], rt["WSB"], rt["CNT"],
hidden, rt["HO"], rt["XP0"], rt["XP1"],
rt["WIN"],
wsrc[0], wsrc[1], wsrc[2], wsrc[3], wsrc[4], wsrc[5],
wsrc[6], wsrc[7], wsrc[8],
state[0]["S"], state[1]["S"], state[2]["S"],
capc, capk,
fc if ftm else capc, fk if ftm else capk,
epoch, rt["G"], ftm, pos, ctxf,
pt=rt["PT"] if rt["PROF"] else rt["CNT"], PROF=rt["PROF"],
num_warps=int(os.getenv("KIMI_WARPS", "8")), num_stages=int(os.getenv("KIMI_STAGES", "1")))
self._epoch = epoch + 1
for l in range(3):
stl = state[l]
d = dsts[l]
stl["cq"] = rt["WV"][l][d][0]
stl["ck"] = rt["WV"][l][d][1]
stl["cv"] = rt["WV"][l][d][2]
stm["c_kv"] = capc[:pos + 1]
stm["k_rope"] = capk[:pos + 1]
return rt["HO"], state
# -------------------------------------------------------------------- #
# standalone state helpers (mirror the harness generators; no imports)
# -------------------------------------------------------------------- #
def init_state(cfg: Config, context_len: int, seed: int) -> list:
dev = torch.device("cuda:0")
g = torch.Generator(device=dev).manual_seed(seed)
H, Dk = cfg.kda_heads, cfg.kda_head_dim
C = H * Dk
state = []
for kind in cfg.pattern:
if kind == "K":
state.append({
"S": torch.randn(H, Dk, Dk, device=dev, generator=g) * 0.05,
"cq": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
"ck": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
"cv": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
})
else:
state.append({
"c_kv": torch.randn(context_len, cfg.kv_lora, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
"k_rope": torch.randn(context_len, cfg.qk_rope, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
})
return state
def init_token(cfg: Config, seed: int) -> torch.Tensor:
dev = torch.device("cuda:0")
g = torch.Generator(device=dev).manual_seed(seed + 1)
return torch.randn(cfg.hidden, device=dev, generator=g, dtype=cfg.dtype) * 0.25
20260825_141528_or-fable_stealth_ox-alpha_02_kimi_linear_decode