"""Kimi-Linear W4A8 hybrid decode fused into ONE persistent Triton megakernel. One `step(hidden, state)` == one `_mega[(P,)](...)` launch. The whole per-token forward lives inside that launch: 3x(KDA + MoE) + 1x(MLA + MoE), every int4 dequant-GEMV, the short causal conv, the KDA recurrent state update, the MLA latent-cache attention (weight-absorbed), the MoE router / top-8 / expert GEMVs, all 8 RMSNorms and all 8 residual adds. The grid is persistent (two CTAs per SM) and the 33 dataflow dependencies between "stages" are resolved by software grid-wide barriers (release atomic arrival + volatile poll of a monotonically increasing target), so no kernel boundary is ever needed. Decode at batch 1 is pure weight streaming: 230 MB of int4 projections and routed experts per token (plus 2-19 MB of MLA cache over the scored context range), which is a 127 us floor at the 1818 GB/s this kernel's GEMV loop sustains. Every design choice below is about keeping that stream saturated and about not paying for anything else. Design notes that matter for reading the kernel: * GEMV, W4A8 with dp4a. Weights are repacked once (at first step) into a task-major int32 arena W32[tile][kchunk][word][col]: the 8 nibbles of one int32 are 8 consecutive contraction elements of ONE output column, ordered so that `w & 0x0F0F0F0F` and `(w >> 4) & 0x0F0F0F0F` are each a 4-lane uint8 vector aligned with 4 *consecutive* int8 activations. A task is 512 output columns x 128 contraction elements == 16 int32 words x 512 columns == 32 KB (x3 for the two stages with split-k to spare, see KM below), and each word is retired by two `dp4a.u32.s32` against a *uniform* int32 of packed activations -- 4 weights per instruction, exact in int32. The scales come out of the sum: sum_k (v-z)*s*x == s*xs*(sum_k v*xq - z*sum_k xq) so the whole epilogue is two FFMAs per output column, once per 32 KB. A standalone replica of this loop runs at 1818 GB/s against a 1861 GB/s load-only ceiling for the same access pattern. * Two CTAs per SM. The same replica reaches only 1533 GB/s with one CTA per SM: a single 4-warp CTA cannot keep enough loads in flight to cover DRAM latency. * Residual adds ride the GEMV epilogue. The four attention o_projs and the four MoE down_projs write their split-k partials with `red.global.add.f32` directly onto a residual stream that an earlier stage seeded, with the MoE router weight folded in. That removes 8 reduce-and-add stages, 8 grid barriers and ~8 MB/token of partial traffic. Because L2 atomics do not invalidate L1, the four addresses involved are read back with `.cg`. * MLA absorption. kv_b is folded into the query (qa[h,c] = sum_d q_nope[h,d] Wb[c,h,d]) and into the output (o[h,d] = sum_c pc[h,c] Wb[c,h,128+d]), so the cache pass costs O(L*512) instead of the reference's O(L*512*8192). Softmax is flash-decode style: 114 l-chunks x 2 head-groups == 228 programs, each with a running max, then a rescaling combine. * Cross-CTA memory model. Apart from the four `.cg` residual addresses, every shared scratch address is written at most once per launch, every write is >=32 B (sector) aligned, and no program reads an address before the barrier that follows its write -- so plain L1-cached loads and stores are safe (L1 dirty state is per-32 B-sector, and the release fence in the barrier writes dirty sectors back to L2). """ from __future__ import annotations import torch import torch.nn as nn import triton import triton.language as tl # --------------------------------------------------------------------------- # # config mirror (names/shapes must match reference.py for load_state_dict) # # Every constant the kernel reads is a tl.constexpr: triton requires that of # jit-visible globals, and constexpr still behaves like an int on the host # (__index__), so the packing code below shares one set of definitions. # --------------------------------------------------------------------------- # _c = tl.constexpr D = _c(2304) # hidden H = _c(32) # heads (kda and mla) DK = _c(128) # kda head dim C4 = _c(H * DK) # 4096 KVL = _c(512) # kv lora rank QKN = _c(128) # qk nope QKR = _c(64) # qk rope VH = _c(128) # v head QH = _c(QKN + QKR) # 192 KVBW = _c(QKN + VH) # 256 MI = _c(1024) # moe intermediate NEXP = _c(64) NBLK = _c(4) NSM = torch.cuda.get_device_properties(0).multi_processor_count if torch.cuda.is_available() else 114 # Two CTAs per SM, not one. A standalone replica of the _gemv inner loop tops # out at 1533 GB/s with 114 CTAs and 1818 GB/s with 228 (the pure weight stream # reaches 1761 and 1861), so the second resident CTA buys ~19% on every int4 # GEMV. It costs ~9 us on the grid barriers, which have twice as many arrivals # to gather; the net at ctx 2048 is -13 us. P = 2 * NSM # persistent grid: two CTAs per SM (host-side int) NWARP = 4 NBAR = 33 # barriers per launch NT = _c(512) # GEMV tile width (output columns) NTB = _c(128) # small tile width (kv_b) RU = _c(4) # word loads per pipelined iteration NSTG = _c(4) # pipeline depth GA = _c(2) # mla head groups NCH = _c(P // GA) # mla l-chunks NCHP = _c(triton.next_power_of_2(int(NCH))) BL = _c(16) # mla rows per dot QXW = _c(C4 // 4) # per-program int8-activation words (max 4096 elements) QSW = _c(8 * (C4 // 128)) # per-program (scale, sum-of-int8) pairs, stride 8 # per-stage GEMV task shapes: NTILE tiles of NT cols x KS k-chunks of 128. # # DNT is the tile width of the two stages whose output width is D=2304, which # tiles badly: 2304 is 4.5 tiles of 512, so 11% of every weight byte loaded is # padding and the final round runs on 46 of the 114 CTAs. 256 divides 2304 # exactly and removes both -- and is 8 us *slower*, because 256 columns over # 128 threads is an 8-byte-per-thread load and the weight stream needs 16 to # reach full bandwidth. Keeping it as a named knob so the answer stays visible. DNT = _c(512) # o_proj / down width # # KM is how many 128-groups one task contracts over, so KS*KM == KF == the true # k-chunk count of the weight arena and KM>1 is a pure reduction of split-k # partial traffic (and of the reduce depth in the consumer stage). It is free # only while NTILE*KS*nslot still covers the grid: qkvg goes 576 tasks in 3 # rounds -> 192 in 1, both 84.2% efficient, and gate/up 648/3 -> 216/1, both # 94.7%. qkva (252 -> 126) and the two o_projs would lose occupancy, so they # stay at KM=1 and KS == KF. QKVG_NTILE, QKVG_KF, QKVG_KM = _c(4 * C4 // NT), _c(D // 128), _c(3) QKVG_KS = _c(QKVG_KF // QKVG_KM) # 32, 18, 3 -> 6 OP_NTILE, OP_KS, OP_KM = _c((D + DNT - 1) // DNT), _c(C4 // 128), _c(1) QKVA_NTILE, QKVA_KS, QKVA_KM = _c(14), _c(D // 128), _c(1) # 14, 18 GU_NTILE, GU_KF, GU_KM = _c(2 * MI // NT), _c(D // 128), _c(3) GU_KS = _c(GU_KF // GU_KM) # 4, 18, 3 -> 6 DN_NTILE, DN_KS, DN_KM = _c((D + DNT - 1) // DNT), _c(MI // 128), _c(1) KVBN_NTILE = _c(H * (KVL // 128)) # 128 KVBV_NTILE, KVBV_KS = _c(H), _c(KVL // 128) # 32, 4 QKVA_NPAD = _c(QKVA_NTILE * NT) # 7168 # ---- int32 weight arena ---- W_QKVG = _c(QKVG_NTILE * QKVG_KF * 16 * NT) W_OP = _c(OP_NTILE * OP_KS * 16 * DNT) W_KDA = _c(W_QKVG + W_OP) W_QKVA = _c(QKVA_NTILE * QKVA_KS * 16 * NT) W_KVBN = _c(KVBN_NTILE * 16 * NTB) W_KVBV = _c(KVBV_NTILE * KVBV_KS * 16 * NTB) W_MLA = _c(W_QKVA + W_KVBN + W_KVBV + W_OP) W_GU = _c(GU_NTILE * GU_KF * 16 * NT) W_DN = _c(DN_NTILE * DN_KS * 16 * DNT) W_EXP = _c(W_GU + W_DN) W_MOE = _c((NEXP + 1) * W_EXP) A_KDA = _c(0) A_MLA = _c(3 * W_KDA) A_MOE = _c(A_MLA + W_MLA) W_TOT = _c(A_MOE + NBLK * W_MOE) # ---- parallel scale arena (bf16, holds s/8) and zero arena (int8) ---- S_QKVG = _c(QKVG_NTILE * QKVG_KF * NT) S_OP = _c(OP_NTILE * OP_KS * DNT) S_KDA = _c(S_QKVG + S_OP) S_QKVA = _c(QKVA_NTILE * QKVA_KS * NT) S_KVBV = _c(KVBV_NTILE * KVBV_KS * NTB) S_MLA = _c(S_QKVA + S_KVBV + S_OP) S_GU = _c(GU_NTILE * GU_KF * NT) S_DN = _c(DN_NTILE * DN_KS * DNT) S_EXP = _c(S_GU + S_DN) S_MOE = _c((NEXP + 1) * S_EXP) B_KDA = _c(0) B_MLA = _c(3 * S_KDA) B_MOE = _c(B_MLA + S_MLA) S_TOT = _c(B_MOE + NBLK * S_MOE) # ---- bf16 side arena for the small dense weights ---- WB_CW = _c(0) # [3][3][4][4096] conv_w tap-major WB_BW = _c(WB_CW + 3 * 3 * 4 * C4) # [3][32][2304] beta_proj WB_RW = _c(WB_BW + 3 * H * D) # [4][64][2304] router WB_NRM = _c(WB_RW + NBLK * NEXP * D) # [4][2][2304] norm weights WB_TOT = _c(WB_NRM + NBLK * 2 * D) # ---- bf16 scratch ---- CB_QA = _c(0) # [32][512] absorbed query CB_QR = _c(CB_QA + H * KVL) # [32][64] rope'd query CB_APC = _c(CB_QR + H * QKR) # [NCH][32][512] pc partials CB_TOT = _c(CB_APC + NCH * H * KVL) # --------------------------------------------------------------------------- # # fp32 scratch layout # --------------------------------------------------------------------------- # _off = 0 def _a(n): global _off o = _off _off += (int(n) + 63) // 64 * 64 # keep every region 256 B aligned return _c(o) SC_INVF = _a(QKR // 2) # rope inverse frequencies SC_EZ = _a(8) # stays zero: EPTR for non-MoE gemvs SC_SN8 = _a(H * 4 * QKN) # kv_b nope scales, s [h][g][d] SC_SNZ = _a(H * 4 * QKN) # kv_b nope scales, s*z SC_XNP = _a(P * D) # private normed activation SC_TOP = _a(P * 32) # private [9 expert ids | 8 probs] SC_XQP = _a(P * KVL) # private scalar operands (kv_b gemvs) SC_QX = _a(P * QXW) # private int8 activation, 4 per word SC_QS = _a(P * QSW) # private (scale, sum-of-int8) pairs SC_QHH = _a(NBLK * 72 * 32) # shared int8 moe hidden [slot*8+g][32] SC_QHS = _a(NBLK * 72 * 8) # shared (scale, sum) pairs (stride 8) SC_PQ = _a(3 * QKVG_KS * 4 * C4) # q/k/v/g partials SC_QKV = _a(3 * 4 * C4) # q, k, v, exp(g) SC_BETA = _a(3 * H * 8) # beta (stride 8) SC_O1 = _a(3 * C4) # kda attention output SC_HA = _a(NBLK * D) # hidden after attn sublayer SC_HM = _a(NBLK * D) # hidden after moe sublayer SC_RLOG = _a(NBLK * NEXP * 8) # router logits (stride 8) SC_PGU = _a(NBLK * 9 * GU_KS * 2 * MI) # gate/up partials SC_HH = _a(NBLK * 9 * MI) # silu(gate)*up SC_PQA = _a(QKVA_KS * QKVA_NPAD) # mla q_proj|kv_a partials SC_AM = _a(NCH * H) # per-chunk running max SC_AL = _a(NCH * H) # per-chunk sumexp SC_O1M = _a(C4) # mla attention output SC_TOT = _c(_off) # strides that appear inside the kernel as ` * ` PQ_SLOT = _c(QKVG_KS * 4 * C4) # per-kda-block q/k/v/g partial slot PGU_SLOT = _c(9 * GU_KS * 2 * MI) # per-moe-block gate/up partial slot S_SLOT = _c(H * DK * DK) # per-kda-block recurrent state CV_SLOT = _c(3 * C4) # one conv state (3 taps x 4096) MLA_SCALE = _c(float(QH) ** -0.5) KDA_SCALE = _c(float(DK) ** -0.5) ROUTED_SCALE = _c(2.446) # --------------------------------------------------------------------------- # # module skeletons: identical parameter/buffer names to reference.py # --------------------------------------------------------------------------- # class QuantLinear(nn.Module): def __init__(self, in_f, out_f, group=128): super().__init__() self.in_f, self.out_f, self.group = in_f, out_f, group self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16)) class QuantExperts(nn.Module): def __init__(self, n, in_f, out_f, group=128): super().__init__() self.n, self.in_f, self.out_f, self.group = n, in_f, out_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, in_f // group, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16)) class KDA(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg d, c = cfg.hidden, cfg.kda_heads * cfg.kda_head_dim self.q_proj = QuantLinear(d, c, cfg.group) self.k_proj = QuantLinear(d, c, cfg.group) self.v_proj = QuantLinear(d, c, cfg.group) self.g_proj = QuantLinear(d, c, cfg.group) self.beta_proj = nn.Linear(d, cfg.kda_heads, bias=False, dtype=cfg.dtype) self.conv_w = nn.Parameter(torch.empty(3, c, cfg.short_conv, dtype=cfg.dtype)) self.o_proj = QuantLinear(c, d, cfg.group) class MLA(nn.Module): def __init__(self, cfg): 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) 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): 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) # --------------------------------------------------------------------------- # # device code # --------------------------------------------------------------------------- # @triton.jit def _dp4a(a, b, c): """c + dot(4 uint8 lanes of a, 4 int8 lanes of b), exact in int32.""" return tl.inline_asm_elementwise("dp4a.u32.s32 $0, $1, $2, $3;", "=r,r,r,r", [a, b, c], dtype=tl.int32, is_pure=True, pack=1) @triton.jit def _rq(y): """round-to-nearest-even for |y| <= 127, without a libdevice call.""" return (y + 12582912.0) - 12582912.0 @triton.jit def _quant(SCF, src, dq, ds, NG: tl.constexpr, QG: tl.constexpr, CG: tl.constexpr = 0): """int8-quantize NG groups of 128 fp32 activations for the dp4a GEMVs. Per group: 32 packed int32 words (4 int8 lanes each, bitcast into the fp32 arena so no extra kernel argument is needed) plus the pair (scale, sum-of-int8) that the GEMV epilogue needs for the zero-point correction. QG groups are reduced per pass. A group is 128 elements == one element per thread, so its max and its sum are both *whole-CTA* reductions (shared memory, two barriers); doing one group at a time cost ~0.7 us each, and at 272 groups per token that was a quarter of the whole step. Widening the tile amortises the two reductions over QG groups at no extra traffic. QG must be a power of two and is chosen per call site: rounding NG up to a multiple of QG costs masked-off lanes, so a caller with NG=1 passes QG=1. CG=1 bypasses L1 on the source, for the one caller whose input arrived by L2 atomics (see the mla o_proj). """ r32 = tl.arange(0, 32)[None, :, None] c4 = tl.arange(0, 4)[None, None, :] for g0 in tl.static_range(0, NG, QG): g = g0 + tl.arange(0, QG) m = g < NG if CG != 0: xt = tl.load(SCF + src + g[:, None, None] * 128 + r32 * 4 + c4, mask=m[:, None, None], other=0.0, cache_modifier=".cg") else: xt = tl.load(SCF + src + g[:, None, None] * 128 + r32 * 4 + c4, mask=m[:, None, None], other=0.0) mx = tl.max(tl.max(tl.abs(xt), axis=2), axis=1) q = _rq(xt * (127.0 / (mx + 1.0e-30))[:, None, None]).to(tl.int32) tl.store(SCF + dq + g[:, None] * 32 + tl.arange(0, 32)[None, :], tl.sum((q & 0xFF) << (c4 * 8), axis=2).to(tl.float32, bitcast=True), mask=m[:, None]) tl.store(SCF + ds + 8 * g, mx * (1.0 / 127.0), mask=m) tl.store(SCF + ds + 8 * g + 1, tl.sum(tl.sum(q, axis=2), axis=1).to(tl.float32), mask=m) tl.debug_barrier() @triton.jit def _bf(x): """round to bf16 precision, keep the fp32 container (matches reference casts).""" return x.to(tl.bfloat16).to(tl.float32) @triton.jit def _bar(BAR, tgt, pid): """grid-wide barrier on a single monotonically increasing arrival counter. `tgt` is the cumulative arrival count over the whole launch (and across launches), so the counter is never reset and no sense reversal is needed. Two variants measured slower and are worth not re-inventing: * spreading arrivals over 8 counters 64 B apart and polling their sum (+0.25 us/barrier: 8x the poll traffic, and if the sum is taken with `tl.sum` over a tile it also puts a shared-memory reduction in the spin loop); * having the last arriver publish a release flag on a separate line so waiters spin somewhere that is not being atomically updated (+0.4 us/barrier: the publisher's own atomic has to complete a full L2 round trip before the flag store can even be issued). Cross-CTA data is left in L1: the release on the arrival makes it visible, and no shared address is written twice within a launch, so no waiter can be holding a stale line for it. """ tl.debug_barrier() tl.atomic_add(BAR, 1, sem="release", scope="gpu") while tl.load(BAR, volatile=True) < tgt: pass tl.debug_barrier() @triton.jit def _cpy(dst, src, n, pid, P: tl.constexpr, W: tl.constexpr): for i in range(pid * W, n, P * W): tl.store(dst + i + tl.arange(0, W), tl.load(src + i + tl.arange(0, W))) # A whole-hidden pass is walked in DNC chunks of DCH. Two constraints pull in # opposite directions: a 128-element chunk is one element per thread, so *every* # `tl.sum` over it is a full shared-memory CTA reduction (~0.7 us) -- but a # single 4096-wide tile is 32 live registers per thread and pushes this kernel, # already at 255 registers, into heavy spilling. DCH=1024 keeps 8 elements per # thread while still reducing only once per pass, at the cost of 3 iterations. DCH = _c(1024) DNC = _c(3) @triton.jit def _dotd(SCF, aoff, WB, woff): """ over the hidden vector: x fp32 in scratch, w bf16 in the weight blob. Partials stay in a DCH-wide vector so the cross-thread reduction happens once at the end instead of once per chunk. """ acc = tl.zeros((DCH,), tl.float32) for i in tl.static_range(0, DNC): o = i * DCH + tl.arange(0, DCH) m = o < D acc += (tl.load(SCF + aoff + o, mask=m, other=0.0) * tl.load(WB + woff + o, mask=m, other=0.0).to(tl.float32)) return tl.sum(acc) @triton.jit def _nld(SCF, HB, hoff, use_bf, o): """load a slice of the residual stream, bf16-rounded like the reference.""" if use_bf != 0: v = tl.load(HB + o).to(tl.float32) else: v = _bf(tl.load(SCF + hoff + o, cache_modifier=".cg")) return v @triton.jit def _norm(WB, SCF, HB, hoff, woff, use_bf, pid): """rmsnorm(h)*w -> the program-private activation slot XNP[pid], bf16 rounded. Deliberately redundant across programs: the result is program-private, so the GEMV that consumes it needs no barrier and hits in its own L1. The hidden is read exactly once and held in registers across the CTA reduction (18 floats per thread). Re-reading it for the rescale pass was cheaper while it came back out of L1, but the residual stream is now written by L2 atomics and has to be read `.cg`, so a second pass would be a second round trip -- 4 MB of it, all P programs pulling the same 9 KB. 2304 == 2048 + 256, so two exact power-of-two tiles cover it with no mask. """ oa = tl.arange(0, 2048) ob = 2048 + tl.arange(0, 256) ha = _nld(SCF, HB, hoff, use_bf, oa) hb = _nld(SCF, HB, hoff, use_bf, ob) rs = tl.rsqrt((tl.sum(ha * ha) + tl.sum(hb * hb)) / 2304.0 + 1.0e-6) tl.store(SCF + SC_XNP + pid * D + oa, _bf(ha * rs * tl.load(WB + woff + oa).to(tl.float32))) tl.store(SCF + SC_XNP + pid * D + ob, _bf(hb * rs * tl.load(WB + woff + ob).to(tl.float32))) tl.debug_barrier() @triton.jit def _gemv(WA, SA, ZA, SCF, EPTR, wbase, sbase, pbase, xqb, xsb, xqstr, xsstr, estw, ests, ntile, ks, nslot, npad, atom, wptr, wnorm, pid, KM: tl.constexpr, NTW: tl.constexpr, P: tl.constexpr): """fused int4 dequant-GEMV stage, W4A8. nslot*ntile*ks tasks; one task is NTW output columns x KM*128 contraction elements (== KM*16 int32 weight words per column, one scale group each). EPTR[slot] picks the expert (a zero float for the non-MoE stages). With atom=0 each task drops its split-k partial in its own scratch slot and a later stage reduces them. With atom!=0 the partial is instead red.add-ed straight into the residual stream at `pbase` (pre-seeded with the residual by an earlier stage, `npad` output columns wide, `wptr[slot]*wnorm` weighting the first nslot-1 slots). That deletes the reduce+residual stage that would otherwise follow -- and, worth far more, the grid barrier in front of it. Both operands are integers: `w & 0x0F0F0F0F` is a 4-lane uint8 vector of the nibbles for 4 consecutive contraction indices (see _pack), so one dp4a against one *uniform* int32 of packed activations retires 4 weights, and the 128-element dot product is exact in int32. The group scale, the activation scale and the zero-point correction are applied once per task: sum_k (v-z)*s*x == s*xs*(sum_k v*xq - z*sum_k xq) That is ~4.5x fewer instructions per weight byte than dequantising to fp32, which is what this loop is actually bound by. The task shape is passed as *runtime* values on purpose: the five GEMV stages then share two compiled copies of this heavily unrolled inner loop (one per distinct KM), not five, which cuts both compile time and instruction footprint. The extra integer math is amortised over 32 KB of weight loads per task. """ nl = tl.arange(0, NTW) nk = ntile * ks for t in range(pid, nslot * nk, P): s = t // nk r = t - s * nk k = r // ntile i = r - k * ntile e = tl.load(EPTR + s).to(tl.int32) po = pbase + (s * ks + k) * npad + i * NTW g0 = (i * ks + k) * KM # first of this task's KM 128-groups val = tl.zeros((NTW,), tl.float32) for j in tl.static_range(KM): wo = wbase + e * estw + (g0 + j) * (16 * NTW) so = sbase + e * ests + (g0 + j) * NTW qo = xqb + s * xqstr + (k * KM + j) * 32 zo = xsb + (s * xsstr + k * KM + j) * 8 # Issued before the accumulation loop, not after it: the epilogue # operands are independent of `a`, so hoisting them lets their # ~500-cycle latency hide under the dp4a chain instead of extending # it. They are kept in their narrow load types (bf16 scale, int8 # zero) so the live set across the loop is 3 registers per thread # rather than 8. sc = tl.load(SA + so + nl) zv = tl.load(ZA + so + nl) xs = tl.load(SCF + zo) xz = tl.load(SCF + zo + 1) a = tl.zeros((NTW,), tl.int32) for m in tl.range(0, 16, RU, num_stages=NSTG): for u in tl.static_range(RU): w = tl.load(WA + wo + (m + u) * NTW + nl) xl = tl.load(SCF + qo + 2 * (m + u)).to(tl.int32, bitcast=True) xh = tl.load(SCF + qo + 2 * (m + u) + 1).to(tl.int32, bitcast=True) a = _dp4a(w & 0x0F0F0F0F, xl, a) a = _dp4a((w >> 4) & 0x0F0F0F0F, xh, a) val += ((sc.to(tl.float32) * (8.0 * xs)) * (a.to(tl.float32) - zv.to(tl.float32) * xz)) if atom == 0: tl.store(SCF + po + nl, val) else: if s < nslot - 1: wj = tl.load(SCF + wptr + s) * wnorm else: wj = 1.0 # red.global.add.f32 at L2. The mask drops the tail of the last # NTW-tile (D is not a multiple of NTW). Readers of pbase must not # have touched it earlier in the launch -- atomics land in L2 and do # not invalidate a stale L1 line -- which holds here: the stage that # seeds pbase with the residual only stores to it. oi = i * NTW + nl tl.atomic_add(SCF + pbase + oi, val * wj, mask=oi < npad, sem="relaxed") @triton.jit def _redp(SCF, pbase, off, KS: tl.constexpr, NPAD: tl.constexpr, n: tl.constexpr): """sum the KS split-k partials of n contiguous output columns.""" acc = tl.zeros((n,), tl.float32) for k in tl.static_range(KS): acc += tl.load(SCF + pbase + k * NPAD + off + tl.arange(0, n)) return acc @triton.jit def _gemv_small(WA, SCF, wo, qo, NTW: tl.constexpr): """dp4a GEMV over exactly 128 contraction elements; scales applied outside.""" acc = tl.zeros((NTW,), tl.int32) nl = tl.arange(0, NTW) for m in tl.range(0, 16, RU, num_stages=NSTG): for u in tl.static_range(RU): w = tl.load(WA + wo + (m + u) * NTW + nl) acc = _dp4a(w & 0x0F0F0F0F, tl.load(SCF + qo + 2 * (m + u)).to(tl.int32, bitcast=True), acc) acc = _dp4a((w >> 4) & 0x0F0F0F0F, tl.load(SCF + qo + 2 * (m + u) + 1).to(tl.int32, bitcast=True), acc) return acc @triton.jit def _moe(WB, WA, SA, ZA, SCF, BAR, b, nb, pid, sid: tl.constexpr, SSTOP: tl.constexpr, P: tl.constexpr): """4 stages: norm+router | top8+gate/up | silu*up+seed | down (accumulating).""" if sid + 1 <= SSTOP: _norm(WB, SCF, WB, SC_HA + b * D, WB_NRM + b * 2 * D + D, 0, pid) _quant(SCF, SC_XNP + pid * D, SC_QX + pid * QXW, SC_QS + pid * QSW, 18, 16) for e in range(pid, NEXP, P): tl.store(SCF + SC_RLOG + (b * NEXP + e) * 8, _bf(_dotd(SCF, SC_XNP + pid * D, WB, WB_RW + (b * NEXP + e) * D))) nb += P _bar(BAR, nb, pid) if sid + 2 <= SSTOP: ee = tl.arange(0, NEXP) lg = tl.load(SCF + SC_RLOG + b * NEXP * 8 + ee * 8) mx = tl.max(lg) den = tl.sum(tl.exp(lg - mx)) rk = tl.sum(((lg[None, :] > lg[:, None]) | ((lg[None, :] == lg[:, None]) & (ee[None, :] < ee[:, None]))).to(tl.int32), axis=1) tl.store(SCF + SC_TOP + pid * 32 + rk, ee.to(tl.float32), mask=rk < 8) tl.store(SCF + SC_TOP + pid * 32 + 16 + rk, tl.exp(lg - mx) / den, mask=rk < 8) tl.store(SCF + SC_TOP + pid * 32 + 8, NEXP + 0.0) tl.debug_barrier() _gemv(WA, SA, ZA, SCF, SCF + SC_TOP + pid * 32, A_MOE + b * W_MOE, B_MOE + b * S_MOE, SC_PGU + b * PGU_SLOT, SC_QX + pid * QXW, SC_QS + pid * QSW, 0, 0, W_EXP, S_EXP, GU_NTILE, GU_KS, 9, 2 * MI, 0, 0, 0.0, pid, GU_KM, NT, P) nb += P _bar(BAR, nb, pid) if sid + 3 <= SSTOP: # Seed the block output with the residual; the down_proj GEMV of the next # stage red.adds the expert contributions on top of it. A separate # buffer from SC_HA is required, not merely tidy: SC_HA was *read* by the # rmsnorm two stages ago, so the reading CTAs hold it in L1, and L2 # atomics would not invalidate those lines. for c in range(pid, 72, P): o = c * 32 + tl.arange(0, 32) tl.store(SCF + SC_HM + b * D + o, _bf(tl.load(SCF + SC_HA + b * D + o, cache_modifier=".cg"))) pg = SC_PGU + b * PGU_SLOT for t in range(pid, 72, P): s = t // 8 g = t % 8 base = pg + s * (GU_KS * 2 * MI) gv = _redp(SCF, base, g * 128, GU_KS, 2 * MI, 128) uv = _redp(SCF, base, MI + g * 128, GU_KS, 2 * MI, 128) hv = (gv * tl.sigmoid(gv)) * uv ho = SC_HH + b * 9 * MI + s * MI + g * 128 tl.store(SCF + ho + tl.arange(0, 128), hv) # one task == exactly one 128-element quantisation group, so the # int8 form of the down_proj input is produced here (once) instead # of redundantly in all P programs. Per-block slot: like every # other shared region this must be written at most once per launch, # or a reader's L1 can still hold the previous block's line. tl.debug_barrier() _quant(SCF, ho, SC_QHH + (b * 72 + t) * 32, SC_QHS + (b * 72 + t) * 8, 1, 1) nb += P _bar(BAR, nb, pid) if sid + 4 <= SSTOP: # The router weight of each of the 8 routed slots (the 9th is the shared # expert, weight 1) folds into the GEMV epilogue, so the 9 x 8 split-k # partials never reach memory: they red.add straight into SC_HM. wsum = tl.sum(tl.load(SCF + SC_TOP + pid * 32 + 16 + tl.arange(0, 8))) _gemv(WA, SA, ZA, SCF, SCF + SC_TOP + pid * 32, A_MOE + b * W_MOE + W_GU, B_MOE + b * S_MOE + S_GU, SC_HM + b * D, SC_QHH + b * 72 * 32, SC_QHS + b * 72 * 8, MI // 4, MI // 128, W_EXP, S_EXP, DN_NTILE, DN_KS, 9, D, 1, SC_TOP + pid * 32 + 16, ROUTED_SCALE / (wsum + 1.0e-9), pid, DN_KM, DNT, P) nb += P _bar(BAR, nb, pid) # HM is consumed by the next block's rmsnorm @triton.jit(do_not_specialize=["bar0", "L"]) def _mega(HBF, HOUT, WA, SA, ZA, WB, SCF, SCB, BAR, bar0, SB, CVI, CVO, CKV, KRP, S0, S1, S2, CQ0, CK0, CV0, CQ1, CK1, CV1, CQ2, CK2, CV2, CKVS, KRPS, L, DOCOPY, SSTOP: tl.constexpr, P: tl.constexpr): pid = tl.program_id(0) nb = bar0 # ---------------- adopt a freshly fed state into our own buffers --------- # statically unrolled, so no later stage needs a runtime-selected pointer if DOCOPY != 0: _cpy(SB + 0 * S_SLOT, S0, S_SLOT, pid, P, 512) _cpy(SB + 1 * S_SLOT, S1, S_SLOT, pid, P, 512) _cpy(SB + 2 * S_SLOT, S2, S_SLOT, pid, P, 512) _cpy(CVI + 0 * CV_SLOT, CQ0, CV_SLOT, pid, P, 128) _cpy(CVI + 1 * CV_SLOT, CK0, CV_SLOT, pid, P, 128) _cpy(CVI + 2 * CV_SLOT, CV0, CV_SLOT, pid, P, 128) _cpy(CVI + 3 * CV_SLOT, CQ1, CV_SLOT, pid, P, 128) _cpy(CVI + 4 * CV_SLOT, CK1, CV_SLOT, pid, P, 128) _cpy(CVI + 5 * CV_SLOT, CV1, CV_SLOT, pid, P, 128) _cpy(CVI + 6 * CV_SLOT, CQ2, CV_SLOT, pid, P, 128) _cpy(CVI + 7 * CV_SLOT, CK2, CV_SLOT, pid, P, 128) _cpy(CVI + 8 * CV_SLOT, CV2, CV_SLOT, pid, P, 128) _cpy(CKV, CKVS, L * KVL, pid, P, 512) _cpy(KRP, KRPS, L * QKR, pid, P, 64) # ========================= 3 x (KDA + MoE) ============================= for b in range(0, 3): # ---- norm + q/k/v/g projection ---- if 1 <= SSTOP: _norm(WB, SCF, HBF, SC_HM + tl.maximum(b - 1, 0) * D, WB_NRM + b * 2 * D, tl.where(b == 0, 1, 0), pid) _quant(SCF, SC_XNP + pid * D, SC_QX + pid * QXW, SC_QS + pid * QSW, 18, 16) _gemv(WA, SA, ZA, SCF, SCF + SC_EZ, A_KDA + b * W_KDA, B_KDA + b * S_KDA, SC_PQ + b * PQ_SLOT, SC_QX + pid * QXW, SC_QS + pid * QSW, 0, 0, 0, 0, QKVG_NTILE, QKVG_KS, 1, 4 * C4, 0, 0, 0.0, pid, QKVG_KM, NT, P) nb += P _bar(BAR, nb, pid) # ---- short causal conv + silu, gate decay, beta ---- if 2 <= SSTOP: # seed the attention residual stream for the o_proj two stages down for c in range(pid, 72, P): o = c * 32 + tl.arange(0, 32) if b == 0: hv = tl.load(HBF + o).to(tl.float32) else: hv = _bf(tl.load(SCF + SC_HM + (b - 1) * D + o, cache_modifier=".cg")) tl.store(SCF + SC_HA + b * D + o, hv) pq = SC_PQ + b * PQ_SLOT for c in range(pid, 128, P): mi = c // 32 col = (c % 32) * 128 nlq = col + tl.arange(0, 128) acc = _redp(SCF, pq + mi * C4, col, QKVG_KS, 4 * C4, 128) if mi == 3: # exp(-softplus(g)) == sigmoid(-g) tl.store(SCF + SC_QKV + b * 4 * C4 + 3 * C4 + nlq, tl.sigmoid(-_bf(acc))) else: val = _bf(acc) ci = CVI + (b * 3 + mi) * CV_SLOT + nlq p0 = tl.load(ci).to(tl.float32) p1 = tl.load(ci + C4).to(tl.float32) p2 = tl.load(ci + 2 * C4).to(tl.float32) cw = WB + WB_CW + (b * 3 + mi) * 4 * C4 + nlq o = (p0 * tl.load(cw).to(tl.float32) + p1 * tl.load(cw + C4).to(tl.float32) + p2 * tl.load(cw + 2 * C4).to(tl.float32) + val * tl.load(cw + 3 * C4).to(tl.float32)) o = _bf(o * tl.sigmoid(o)) o = tl.where(mi == 0, o * KDA_SCALE, o) tl.store(SCF + SC_QKV + b * 4 * C4 + mi * C4 + nlq, o) co = CVO + (b * 3 + mi) * CV_SLOT + nlq tl.store(co, p1.to(tl.bfloat16)) tl.store(co + C4, p2.to(tl.bfloat16)) tl.store(co + 2 * C4, val.to(tl.bfloat16)) for h in range(pid, H, P): tl.store(SCF + SC_BETA + b * H * 8 + h * 8, tl.sigmoid(_bf(_dotd(SCF, SC_XNP + pid * D, WB, WB_BW + (b * H + h) * D)))) nb += P _bar(BAR, nb, pid) # ---- gated delta rule: decay, predict, correct, read out ---- if 3 <= SSTOP: qb = SC_QKV + b * 4 * C4 dk = tl.arange(0, 128) dv = tl.arange(0, 8) for t in range(pid, H * 16, P): h = t // 16 d0 = (t % 16) * 8 ge = tl.load(SCF + qb + 3 * C4 + h * 128 + dk) kk = tl.load(SCF + qb + C4 + h * 128 + dk) qq = tl.load(SCF + qb + h * 128 + dk) vv = tl.load(SCF + qb + 2 * C4 + h * 128 + d0 + dv) bt = tl.load(SCF + SC_BETA + b * H * 8 + h * 8) sadr = SB + b * S_SLOT + h * (DK * DK) + dk[:, None] * DK + d0 + dv[None, :] sg = tl.load(sadr) * ge[:, None] pred = tl.sum(sg * kk[:, None], axis=0) aa = tl.sum(sg * qq[:, None], axis=0) dvv = bt * (vv - pred) tl.store(sadr, sg + kk[:, None] * dvv[None, :]) tl.store(SCF + SC_O1 + b * C4 + h * 128 + d0 + dv, _bf(aa + dvv * tl.sum(kk * qq))) nb += P _bar(BAR, nb, pid) # ---- kda o_proj, accumulated onto the residual ---- if 4 <= SSTOP: _quant(SCF, SC_O1 + b * C4, SC_QX + pid * QXW, SC_QS + pid * QSW, OP_KS, 16) _gemv(WA, SA, ZA, SCF, SCF + SC_EZ, A_KDA + b * W_KDA + W_QKVG, B_KDA + b * S_KDA + S_QKVG, SC_HA + b * D, SC_QX + pid * QXW, SC_QS + pid * QSW, 0, 0, 0, 0, OP_NTILE, OP_KS, 1, D, 1, 0, 0.0, pid, OP_KM, DNT, P) nb += P _bar(BAR, nb, pid) _moe(WB, WA, SA, ZA, SCF, BAR, b, nb, pid, 4, SSTOP, P) nb += 4 * P # ============================ MLA + MoE ================================ # ---- norm + q_proj | kv_a ---- if 31 <= SSTOP: _norm(WB, SCF, WB, SC_HM + 2 * D, WB_NRM + 3 * 2 * D, 0, pid) _quant(SCF, SC_XNP + pid * D, SC_QX + pid * QXW, SC_QS + pid * QSW, 18, 16) _gemv(WA, SA, ZA, SCF, SCF + SC_EZ, A_MLA, B_MLA, SC_PQA, SC_QX + pid * QXW, SC_QS + pid * QSW, 0, 0, 0, 0, QKVA_NTILE, QKVA_KS, 1, QKVA_NPAD, 0, 0, 0.0, pid, QKVA_KM, NT, P) nb += P _bar(BAR, nb, pid) # ---- absorbed query (kv_b nope half), rope, cache append ---- if 32 <= SSTOP: d128 = tl.arange(0, 128) for t in range(pid, KVBN_NTILE, P): h = t // 4 g = t % 4 qn = _bf(_redp(SCF, SC_PQA, h * QH, QKVA_KS, QKVA_NPAD, 128)) * MLA_SCALE sn8 = tl.load(SCF + SC_SN8 + (h * 4 + g) * 128 + d128) snz = tl.load(SCF + SC_SNZ + (h * 4 + g) * 128 + d128) cst = tl.sum(qn * snz) # exact zero-point term, from fp32 q tl.store(SCF + SC_XQP + pid * KVL + d128, qn * sn8) tl.debug_barrier() _quant(SCF, SC_XQP + pid * KVL, SC_QX + pid * QXW, SC_QS + pid * QSW, 1, 1) acc = _gemv_small(WA, SCF, (A_MLA + W_QKVA) + t.to(tl.int64) * (16 * NTB), SC_QX + pid * QXW, NTB) xs = tl.load(SCF + SC_QS + pid * QSW) tl.store(SCB + CB_QA + h * KVL + g * 128 + tl.arange(0, NTB), (xs * acc.to(tl.float32) - cst).to(tl.bfloat16)) tl.debug_barrier() i32 = tl.arange(0, 32) ang = L.to(tl.float32) * tl.load(SCF + SC_INVF + i32) cs = tl.cos(ang) sn = tl.sin(ang) ev = tl.arange(0, 64)[None, :] == (2 * i32)[:, None] od = tl.arange(0, 64)[None, :] == (2 * i32 + 1)[:, None] if pid == 0: ck = _bf(_redp(SCF, SC_PQA, H * QH, QKVA_KS, QKVA_NPAD, KVL)) tl.store(CKV + L.to(tl.int64) * KVL + tl.arange(0, KVL), ck.to(tl.bfloat16)) if pid == 1: kr = _bf(_redp(SCF, SC_PQA, H * QH + KVL, QKVA_KS, QKVA_NPAD, 64)) ke = tl.sum(tl.where(ev, kr[None, :], 0.0), axis=1) ko = tl.sum(tl.where(od, kr[None, :], 0.0), axis=1) tl.store(KRP + L.to(tl.int64) * QKR + 2 * i32, _bf(ke * cs - ko * sn).to(tl.bfloat16)) tl.store(KRP + L.to(tl.int64) * QKR + 2 * i32 + 1, _bf(ko * cs + ke * sn).to(tl.bfloat16)) for zc in range(pid, 32, P): # zero the stage-34 accumulator tl.store(SCF + SC_O1M + zc * 128 + tl.arange(0, 128), tl.zeros((128,), tl.float32)) for hh in range(pid - 2, H, P): if hh >= 0: qr = _bf(_redp(SCF, SC_PQA, hh * QH + QKN, QKVA_KS, QKVA_NPAD, 64)) qe = tl.sum(tl.where(ev, qr[None, :], 0.0), axis=1) qo = tl.sum(tl.where(od, qr[None, :], 0.0), axis=1) tl.store(SCB + CB_QR + hh * QKR + 2 * i32, _bf(qe * cs - qo * sn).to(tl.bfloat16)) tl.store(SCB + CB_QR + hh * QKR + 2 * i32 + 1, _bf(qo * cs + qe * sn).to(tl.bfloat16)) nb += P _bar(BAR, nb, pid) # ---- latent attention: flash-decode over (l-chunk, head-group) ---- if 33 <= SSTOP: hg = pid % GA ci = pid // GA if ci < NCH: h0 = hg * 16 hr = tl.arange(0, 16) cl = tl.arange(0, KVL) qa = tl.trans(tl.load(SCB + CB_QA + (h0 + hr)[:, None] * KVL + cl[None, :])) qr = tl.trans(tl.load(SCB + CB_QR + (h0 + hr)[:, None] * QKR + tl.arange(0, QKR)[None, :])) cw = (L + NCH) // NCH l0 = ci * cw l1 = tl.minimum(l0 + cw, L + 1) mx = tl.full((16,), -1.0e30, tl.float32) ls = tl.zeros((16,), tl.float32) pc = tl.zeros((16, KVL), tl.float32) for lb in range(l0, l1, BL): rows = lb + tl.arange(0, BL) msk = rows < l1 ck = tl.load(CKV + rows[:, None].to(tl.int64) * KVL + cl[None, :], mask=msk[:, None], other=0.0) kr = tl.load(KRP + rows[:, None].to(tl.int64) * QKR + tl.arange(0, QKR)[None, :], mask=msk[:, None], other=0.0) s = tl.dot(ck, qa) + tl.dot(kr, qr) * MLA_SCALE s = tl.where(msk[:, None], s, -1.0e30) mn = tl.maximum(mx, tl.max(s, axis=0)) al = tl.exp(mx - mn) pe = tl.exp(s - mn[None, :]) ls = ls * al + tl.sum(pe, axis=0) pc = pc * al[:, None] + tl.dot(tl.trans(pe).to(tl.bfloat16), ck) mx = mn tl.store(SCF + SC_AM + ci * H + h0 + hr, mx) tl.store(SCF + SC_AL + ci * H + h0 + hr, ls) tl.store(SCB + CB_APC + (ci * H + h0 + hr)[:, None] * KVL + cl[None, :], pc.to(tl.bfloat16)) nb += P _bar(BAR, nb, pid) # ---- softmax combine + absorbed output projection (kv_b v half) ---- if 34 <= SSTOP: # One task per (head, 128-wide slice of the latent), not per head: with # 32 heads only 32 of the P programs had anything to do here and the # stage ran at 14% occupancy while pulling 3.7 MB of partial contexts. # The four slices of a head red.add into the same 128 outputs. nl = tl.arange(0, NTB) cc = tl.arange(0, NCHP) cm = cc < NCH for t in range(pid, H * KVBV_KS, P): h = t // KVBV_KS g = t % KVBV_KS am = tl.load(SCF + SC_AM + cc * H + h, mask=cm, other=-1.0e30) al = tl.load(SCF + SC_AL + cc * H + h, mask=cm, other=0.0) sc = tl.exp(am - tl.max(am)) pcv = tl.sum(sc[:, None] * tl.load( SCB + CB_APC + (cc * H + h)[:, None] * KVL + g * 128 + nl[None, :], mask=cm[:, None], other=0.0).to(tl.float32), axis=0) / tl.sum(sc * al) tl.store(SCF + SC_XQP + pid * KVL + nl, pcv) tl.debug_barrier() _quant(SCF, SC_XQP + pid * KVL, SC_QX + pid * QXW, SC_QS + pid * QSW, 1, 1) a = _gemv_small(WA, SCF, (A_MLA + W_QKVA + W_KVBN) + t.to(tl.int64) * (16 * NTB), SC_QX + pid * QXW, NTB) so = (B_MLA + S_QKVA) + t.to(tl.int64) * NTB qs = SC_QS + pid * QSW f = tl.load(SA + so + nl).to(tl.float32) * (8.0 * tl.load(SCF + qs)) z = tl.load(ZA + so + nl).to(tl.float32) tl.atomic_add(SCF + SC_O1M + h * VH + nl, f * (a.to(tl.float32) - z * tl.load(SCF + qs + 1)), sem="relaxed") tl.debug_barrier() # seed the attention residual stream for the o_proj of the next stage for c in range(pid, 72, P): o = c * 32 + tl.arange(0, 32) tl.store(SCF + SC_HA + 3 * D + o, _bf(tl.load(SCF + SC_HM + 2 * D + o, cache_modifier=".cg"))) nb += P _bar(BAR, nb, pid) # ---- mla o_proj, accumulated onto the residual ---- if 35 <= SSTOP: _quant(SCF, SC_O1M, SC_QX + pid * QXW, SC_QS + pid * QSW, OP_KS, 16, 1) _gemv(WA, SA, ZA, SCF, SCF + SC_EZ, A_MLA + W_QKVA + W_KVBN + W_KVBV, B_MLA + S_QKVA + S_KVBV, SC_HA + 3 * D, SC_QX + pid * QXW, SC_QS + pid * QSW, 0, 0, 0, 0, OP_NTILE, OP_KS, 1, D, 1, 0, 0.0, pid, OP_KM, DNT, P) nb += P _bar(BAR, nb, pid) _moe(WB, WA, SA, ZA, SCF, BAR, 3, nb, pid, 35, SSTOP, P) # ---- emit the bf16 next-token hidden (same program that wrote HM) ---- for c in range(pid, 72, P): o = c * 32 + tl.arange(0, 32) tl.store(HOUT + o, tl.load(SCF + SC_HM + 3 * D + o, cache_modifier=".cg").to(tl.bfloat16)) # --------------------------------------------------------------------------- # # host-side weight repacking # --------------------------------------------------------------------------- # NIB = (0, 2, 4, 6, 1, 3, 5, 7) """nibble slot of contraction element j inside a packed word. Element j lands in byte j%4, low nibble for j<4 and high nibble for j>=4, so `w & 0x0F0F0F0F` and `(w >> 4) & 0x0F0F0F0F` are each a 4-lane uint8 vector lined up with 4 *consecutive* int8 activations -- one dp4a each. """ def _pack(wq, scales, zeros, K, N, ntw, ntile): """(.., K//2, N) uint8 -> task-major int32 arena, plus (s/8) bf16 and z int8.""" lead = tuple(wq.shape[:-2]) ng = K // 128 npad = ntile * ntw dev = wq.device v = torch.zeros(lead + (K, npad), dtype=torch.int32, device=dev) v[..., 0::2, :N] = (wq & 0xF).to(torch.int32) v[..., 1::2, :N] = ((wq >> 4) & 0xF).to(torch.int32) v = v.view(lead + (ng, 16, 8, ntile, ntw)) w32 = torch.zeros(lead + (ng, 16, ntile, ntw), dtype=torch.int32, device=dev) for j in range(8): w32 |= v.select(-3, j) << (4 * NIB[j]) w32 = w32.movedim(-2, -4).contiguous() # [.., ntile, ng, 16, ntw] sp = torch.zeros(lead + (ng, npad), dtype=torch.float32, device=dev) sp[..., :N] = scales.float() / 8.0 s8 = sp.view(lead + (ng, ntile, ntw)).movedim(-2, -3).contiguous().to(torch.bfloat16) zp = torch.zeros(lead + (ng, npad), dtype=torch.float32, device=dev) zp[..., :N] = zeros.float() z8 = zp.view(lead + (ng, ntile, ntw)).movedim(-2, -3).contiguous().to(torch.int8) return w32.reshape(-1), s8.reshape(-1), z8.reshape(-1) def _unpack_kvb(wq): v = torch.zeros((KVL, H * KVBW), dtype=torch.int32, device=wq.device) v[0::2] = (wq & 0xF).to(torch.int32) v[1::2] = ((wq >> 4) & 0xF).to(torch.int32) return v.view(KVL, H, KVBW) def _pack_kvbn(wq): """kv_b nope half, transposed: [h][g][word][c]; 8 nibbles == 8 consecutive d.""" vn = _unpack_kvb(wq)[:, :, :QKN].permute(1, 0, 2).contiguous() # (h, c, d) vn = vn.view(H, 4, NTB, 16, 8) w32 = torch.zeros((H, 4, NTB, 16), dtype=torch.int32, device=wq.device) for j in range(8): w32 |= vn[..., j] << (4 * NIB[j]) return w32.permute(0, 1, 3, 2).contiguous().reshape(-1) def _pack_kvbv(wq, scales, zeros): """kv_b v half, natural orientation: [h][g][word][dv], contraction over c.""" vv = _unpack_kvb(wq)[:, :, QKN:].reshape(4, 16, 8, H, VH) w32 = torch.zeros((4, 16, H, VH), dtype=torch.int32, device=wq.device) for j in range(8): w32 |= vv[:, :, j] << (4 * NIB[j]) w32 = w32.permute(2, 0, 1, 3).contiguous() s = (scales.float() / 8.0).view(4, H, KVBW)[:, :, QKN:].permute(1, 0, 2).contiguous() z = zeros.float().view(4, H, KVBW)[:, :, QKN:].permute(1, 0, 2).contiguous() return w32.reshape(-1), s.reshape(-1).to(torch.bfloat16), z.reshape(-1).to(torch.int8) # --------------------------------------------------------------------------- # 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._ready = False self._sstop = 39 # ------------------------------------------------------------------ # def _build(self): dev = self.blocks[0].attn_norm.device cfg = self.cfg wa = torch.zeros(W_TOT, dtype=torch.int32, device=dev) sa = torch.zeros(S_TOT, dtype=torch.bfloat16, device=dev) za = torch.zeros(S_TOT, dtype=torch.int8, device=dev) wb = torch.zeros(WB_TOT, dtype=torch.bfloat16, device=dev) scf = torch.zeros(SC_TOT, dtype=torch.float32, device=dev) def put(off, soff, packed): w32, s8, z8 = packed wa[off:off + w32.numel()] = w32 if s8 is not None: sa[soff:soff + s8.numel()] = s8 za[soff:soff + z8.numel()] = z8 for b in range(3): k = self.blocks[b].attn put(A_KDA + b * W_KDA, B_KDA + b * S_KDA, _pack(torch.cat([k.q_proj.w_q, k.k_proj.w_q, k.v_proj.w_q, k.g_proj.w_q], 1), torch.cat([k.q_proj.scales, k.k_proj.scales, k.v_proj.scales, k.g_proj.scales], 1), torch.cat([k.q_proj.zeros, k.k_proj.zeros, k.v_proj.zeros, k.g_proj.zeros], 1), D, 4 * C4, NT, QKVG_NTILE)) put(A_KDA + b * W_KDA + W_QKVG, B_KDA + b * S_KDA + S_QKVG, _pack(k.o_proj.w_q, k.o_proj.scales, k.o_proj.zeros, C4, D, DNT, OP_NTILE)) wb[WB_CW + b * 12 * C4: WB_CW + (b + 1) * 12 * C4] = \ k.conv_w.detach().permute(0, 2, 1).reshape(-1) wb[WB_BW + b * H * D: WB_BW + (b + 1) * H * D] = k.beta_proj.weight.detach().reshape(-1) m = self.blocks[3].attn wq = torch.zeros(D // 2, QKVA_NPAD, dtype=torch.uint8, device=dev) sc = torch.zeros(D // 128, QKVA_NPAD, dtype=torch.bfloat16, device=dev) ze = torch.zeros(D // 128, QKVA_NPAD, dtype=torch.bfloat16, device=dev) wq[:, :H * QH] = m.q_proj.w_q sc[:, :H * QH] = m.q_proj.scales ze[:, :H * QH] = m.q_proj.zeros wq[:, H * QH:H * QH + KVL + QKR] = m.kv_a.w_q sc[:, H * QH:H * QH + KVL + QKR] = m.kv_a.scales ze[:, H * QH:H * QH + KVL + QKR] = m.kv_a.zeros put(A_MLA, B_MLA, _pack(wq, sc, ze, D, QKVA_NPAD, NT, QKVA_NTILE)) put(A_MLA + W_QKVA, 0, (_pack_kvbn(m.kv_b.w_q), None, None)) put(A_MLA + W_QKVA + W_KVBN, B_MLA + S_QKVA, _pack_kvbv(m.kv_b.w_q, m.kv_b.scales, m.kv_b.zeros)) put(A_MLA + W_QKVA + W_KVBN + W_KVBV, B_MLA + S_QKVA + S_KVBV, _pack(m.o_proj.w_q, m.o_proj.scales, m.o_proj.zeros, C4, D, DNT, OP_NTILE)) s = m.kv_b.scales.float().view(4, H, KVBW)[:, :, :QKN].permute(1, 0, 2).contiguous() z = m.kv_b.zeros.float().view(4, H, KVBW)[:, :, :QKN].permute(1, 0, 2).contiguous() scf[SC_SN8:SC_SN8 + H * 4 * QKN] = s.reshape(-1) scf[SC_SNZ:SC_SNZ + H * 4 * QKN] = (s * z).reshape(-1) qkr = int(QKR) # torch needs a real int, not a constexpr scf[SC_INVF:SC_INVF + qkr // 2] = 1.0 / ( cfg.rope_theta ** (torch.arange(0, qkr, 2, device=dev, dtype=torch.float32) / qkr)) for b in range(NBLK): mo = self.blocks[b].moe for e in range(NEXP + 1): if e < NEXP: g, u, dn, i = mo.gate, mo.up, mo.down, e else: g, u, dn, i = mo.s_gate, mo.s_up, mo.s_down, 0 off = A_MOE + b * W_MOE + e * W_EXP soff = B_MOE + b * S_MOE + e * S_EXP put(off, soff, _pack(torch.cat([g.w_q[i], u.w_q[i]], 1), torch.cat([g.scales[i], u.scales[i]], 1), torch.cat([g.zeros[i], u.zeros[i]], 1), D, 2 * MI, NT, GU_NTILE)) put(off + W_GU, soff + S_GU, _pack(dn.w_q[i], dn.scales[i], dn.zeros[i], MI, D, DNT, DN_NTILE)) wb[WB_RW + b * NEXP * D: WB_RW + (b + 1) * NEXP * D] = \ mo.router.weight.detach().reshape(-1) wb[WB_NRM + b * 2 * D: WB_NRM + b * 2 * D + D] = self.blocks[b].attn_norm.detach() wb[WB_NRM + b * 2 * D + D: WB_NRM + (b + 1) * 2 * D] = self.blocks[b].moe_norm.detach() self._wa, self._sa, self._za, self._wb, self._scf = wa, sa, za, wb, scf self._scb = torch.zeros(CB_TOT, dtype=torch.bfloat16, device=dev) self._bar = torch.zeros(8 * 64, dtype=torch.int32, device=dev) self._barbase = 0 self._hout = torch.zeros(D, dtype=torch.bfloat16, device=dev) self._sbuf = torch.zeros(3, H, DK, DK, dtype=torch.float32, device=dev) self._cv = torch.zeros(2, 9, 3, C4, dtype=torch.bfloat16, device=dev) self._par = 0 self._ckv = None self._krp = None self._cap = 0 self._ready = True def _cache(self, ckv): n = ckv.shape[0] if self._ckv is not None and ckv.data_ptr() == self._ckv.data_ptr() and n < self._cap: return 0 self._cap = n + 512 self._ckv = torch.empty(self._cap, KVL, dtype=torch.bfloat16, device=ckv.device) self._krp = torch.empty(self._cap, QKR, dtype=torch.bfloat16, device=ckv.device) return 1 # ------------------------------------------------------------------ # @torch.no_grad() def step(self, hidden, state): if not self._ready: self._build() ki = [i for i, k in enumerate(self.cfg.pattern) if k == "K"] mi = self.cfg.pattern.index("M") docopy = self._cache(state[mi]["c_kv"]) if state[ki[0]]["S"].data_ptr() != self._sbuf.data_ptr(): docopy = 1 L = int(state[mi]["c_kv"].shape[0]) cs = [state[i][n] for i in ki for n in ("cq", "ck", "cv")] _mega[(P,)]( hidden, self._hout, self._wa, self._sa, self._za, self._wb, self._scf, self._scb, self._bar, self._barbase, self._sbuf, self._cv[self._par], self._cv[1 - self._par], self._ckv, self._krp, state[ki[0]]["S"], state[ki[1]]["S"], state[ki[2]]["S"], cs[0], cs[1], cs[2], cs[3], cs[4], cs[5], cs[6], cs[7], cs[8], state[mi]["c_kv"], state[mi]["k_rope"], L, docopy, self._sstop, P, num_warps=NWARP, ) self._barbase += NBAR * P out = self._cv[1 - self._par] self._par = 1 - self._par for j, i in enumerate(ki): state[i]["S"] = self._sbuf[j] state[i]["cq"] = out[j * 3 + 0] state[i]["ck"] = out[j * 3 + 1] state[i]["cv"] = out[j * 3 + 2] state[mi]["c_kv"] = self._ckv[:L + 1] state[mi]["k_rope"] = self._krp[:L + 1] return self._hout, state