"""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