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