"""Kimi-Linear W4A16 hybrid decode — single-launch CUDA megakernel. The timed path is one cooperative __global__ kernel that fuses every int4 dequant-GEMV, KDA/MLA attention, MoE (top-8 + shared), RMSNorm and residuals for all four layers in the [K,K,K,M] motif. No CUDA graphs, no torch.compile, no per-op Python kernel loop. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.cpp_extension import load_inline # --------------------------------------------------------------------------- # # Config / quant (same layout as reference so state_dict keys match) # --------------------------------------------------------------------------- # EPS = 1.0e-6 GROUP_SIZE = 128 @dataclass(frozen=True) class Config: hidden: int = 2304 kda_heads: int = 32 kda_head_dim: int = 128 short_conv: int = 4 mla_heads: int = 32 kv_lora: int = 512 qk_nope: int = 128 qk_rope: int = 64 v_head: int = 128 rope_theta: float = 10000.0 n_experts: int = 64 n_active: int = 8 n_shared: int = 1 moe_inter: int = 1024 routed_scaling: float = 2.446 group: int = 128 pattern: tuple = ("K", "K", "K", "M") dtype: torch.dtype = field(default=torch.bfloat16) def build_config(shape: dict) -> Config: return Config(n_experts=int(shape.get("n_experts", 64))) def _pack_int4(w_q: torch.Tensor) -> torch.Tensor: lo = w_q[0::2] & 0xF hi = w_q[1::2] & 0xF return (lo | (hi << 4)).contiguous() def quantize(w_io: torch.Tensor, group: int = GROUP_SIZE): K, N = w_io.shape ng = K // group wg = w_io.view(ng, group, N).float() wmin = wg.min(dim=1, keepdim=True).values wmax = wg.max(dim=1, keepdim=True).values scales = (wmax - wmin).clamp_min(1e-8) / 15.0 zeros = (-wmin / scales).round().clamp(0, 15) w_q = ((wg / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N) return _pack_int4(w_q), scales.squeeze(1).to(torch.bfloat16), zeros.squeeze(1).to(torch.bfloat16) class QuantLinear(nn.Module): def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() assert in_f % group == 0 and in_f % 2 == 0 self.in_f, self.out_f, self.group = in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16)) def init_random(self, gen: torch.Generator, std: float = 0.02) -> None: w = torch.randn(self.in_f, self.out_f, generator=gen) * std wq, s, z = quantize(w, self.group) self.w_q.copy_(wq) self.scales.copy_(s) self.zeros.copy_(z) class QuantExperts(nn.Module): def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE): super().__init__() self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group ng = in_f // group self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8)) self.register_buffer("scales", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16)) def init_random(self, gen: torch.Generator, std: float = 0.02) -> None: for e in range(self.n): w = torch.randn(self.in_f, self.out_f, generator=gen) * std wq, s, z = quantize(w, self.group) self.w_q[e].copy_(wq) self.scales[e].copy_(s) self.zeros[e].copy_(z) class KDA(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg 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: Config): 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: Config): super().__init__() self.cfg = cfg d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype) self.gate = QuantExperts(E, d, m, cfg.group) self.up = QuantExperts(E, d, m, cfg.group) self.down = QuantExperts(E, m, d, cfg.group) self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group) self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group) class Block(nn.Module): def __init__(self, cfg: Config, kind: str): super().__init__() self.kind = kind self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype)) self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype)) self.attn = KDA(cfg) if kind == "K" else MLA(cfg) self.moe = MoE(cfg) # --------------------------------------------------------------------------- # # CUDA megakernel (single cooperative launch for the full step) # --------------------------------------------------------------------------- # _CUDA_SRC = r""" #include #include #include #include #include #include #include namespace cg = cooperative_groups; static constexpr int BLOCK = 256; static constexpr int MAX_K = 4096; // o_proj K = H*Dk static constexpr float EPS = 1e-6f; __device__ __forceinline__ float b2f(__nv_bfloat16 x) { return __bfloat162float(x); } __device__ __forceinline__ __nv_bfloat16 f2b(float x) { return __float2bfloat16(x); } __device__ __forceinline__ float silu(float x) { return x / (1.f + expf(-x)); } __device__ __forceinline__ float softplus(float x) { if (x > 20.f) return x; if (x < -20.f) return expf(x); return log1pf(expf(x)); } // Grid-stride W4A16 GEMV. x read from global (L2-resident for K~2-4k). // Avoid large static smem so the megakernel stays under the 48KB default. // Does NOT grid-barrier — caller grid.sync() when consumers need y. __device__ void gemv_int4_body( const float* __restrict__ x, const uint8_t* __restrict__ w_q, const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ zeros, float* __restrict__ y, int K, int N, int group ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; const int ng = K / group; for (int n = tid; n < N; n += stride) { float acc = 0.f; #pragma unroll 1 for (int g = 0; g < ng; ++g) { float s = b2f(scales[g * N + n]); float z = b2f(zeros[g * N + n]); const int row0 = (g * group) >> 1; #pragma unroll 8 for (int r = 0; r < (group >> 1); ++r) { int row = row0 + r; uint8_t packed = w_q[row * N + n]; int k = row << 1; acc = fmaf(x[k], ((float)(packed & 0xF) - z) * s, acc); acc = fmaf(x[k + 1], ((float)((packed >> 4) & 0xF) - z) * s, acc); } } y[n] = acc; } } __device__ void gemv_int4( const float* __restrict__ x, const uint8_t* __restrict__ w_q, const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ zeros, float* __restrict__ y, int K, int N, int group ) { gemv_int4_body(x, w_q, scales, zeros, y, K, N, group); } // BF16 weight GEMV (tiny: beta [H,D], router [E,D]) __device__ void gemv_bf16( const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, float* __restrict__ y, int in_f, int out_f ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; for (int n = tid; n < out_f; n += stride) { float acc = 0.f; const __nv_bfloat16* row = w + n * in_f; #pragma unroll 4 for (int k = 0; k < in_f; ++k) acc = fmaf(x[k], b2f(row[k]), acc); y[n] = acc; } } __device__ void rmsnorm( const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, float* __restrict__ y, int D ) { // mean of squares via block reduce then grid reduce into y[0] temp? use shared // Simple two-pass: first compute sumsq with atomics into a single float* // We use y as output and a provided sumsq pointer. } // RMSNorm using a provided 1-element workspace float* sumsq (zeroed by thread 0) __device__ void rmsnorm_into( cg::grid_group &grid, const float* __restrict__ x, const __nv_bfloat16* __restrict__ w, float* __restrict__ y, float* __restrict__ sumsq_ws, int D ) { if (blockIdx.x == 0 && threadIdx.x == 0) sumsq_ws[0] = 0.f; grid.sync(); const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; float local = 0.f; for (int i = tid; i < D; i += stride) { float v = x[i]; local += v * v; } __shared__ float sh[BLOCK]; sh[threadIdx.x] = local; __syncthreads(); for (int s = BLOCK/2; s > 0; s >>= 1) { if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s]; __syncthreads(); } if (threadIdx.x == 0) atomicAdd(sumsq_ws, sh[0]); grid.sync(); float inv = rsqrtf(sumsq_ws[0] / (float)D + EPS); for (int i = tid; i < D; i += stride) { y[i] = x[i] * inv * b2f(w[i]); } grid.sync(); } __device__ void vec_add( cg::grid_group &grid, float* __restrict__ a, const float* __restrict__ b, int n ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; for (int i = tid; i < n; i += stride) a[i] += b[i]; grid.sync(); } __device__ void bf16_to_float( const __nv_bfloat16* __restrict__ x, float* __restrict__ y, int n ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; for (int i = tid; i < n; i += stride) y[i] = b2f(x[i]); __syncthreads(); } __device__ void float_to_bf16( const float* __restrict__ x, __nv_bfloat16* __restrict__ y, int n ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; for (int i = tid; i < n; i += stride) y[i] = f2b(x[i]); __syncthreads(); } // Short depthwise conv kernel=4 causal + SiLU. win is [3, C] previous, val is [C]. // out[C], new_win[3,C] = cat(win[1:], val) __device__ void short_conv( cg::grid_group &grid, const float* __restrict__ val, const __nv_bfloat16* __restrict__ prev, // [KW-1, C] const __nv_bfloat16* __restrict__ conv_w, // [C, KW] for this channel group float* __restrict__ out, __nv_bfloat16* __restrict__ new_prev, int C, int KW ) { const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; for (int c = tid; c < C; c += stride) { float acc = 0.f; // prev rows 0..KW-2, then val for (int t = 0; t < KW - 1; ++t) { acc += b2f(prev[t * C + c]) * b2f(conv_w[c * KW + t]); } acc += val[c] * b2f(conv_w[c * KW + (KW - 1)]); out[c] = silu(acc); // shift window for (int t = 0; t < KW - 2; ++t) { new_prev[t * C + c] = prev[(t + 1) * C + c]; } new_prev[(KW - 2) * C + c] = f2b(val[c]); } // no grid.sync — caller batches multiple short_convs then syncs } // KDA recurrent update for all heads. S[H,Dk,Dk] float, q/k/v/g [H,Dk], beta[H] __device__ void kda_update( cg::grid_group &grid, float* __restrict__ S, const float* __restrict__ q, const float* __restrict__ k, const float* __restrict__ v, const float* __restrict__ g, const float* __restrict__ beta, float* __restrict__ o, int H, int Dk ) { // One head per (block-ish); threads cooperate within head on Dk*Dk const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; // First: S *= exp(g)[:,None]; compute pred; update S; compute o // Pass 1: scale S by exp(g) for (int idx = tid; idx < H * Dk * Dk; idx += stride) { int h = idx / (Dk * Dk); int rem = idx % (Dk * Dk); int r = rem / Dk; // k-dim // int c = rem % Dk; float eg = expf(g[h * Dk + r]); S[idx] *= eg; } grid.sync(); // Pass 2: pred[h, d] = sum_r S[h,r,d] * k[h,r] // Use o as temporary for pred then overwrite for (int idx = tid; idx < H * Dk; idx += stride) { int h = idx / Dk; int d = idx % Dk; float pred = 0.f; for (int r = 0; r < Dk; ++r) { pred += S[(h * Dk + r) * Dk + d] * k[h * Dk + r]; } o[idx] = pred; // temp } grid.sync(); // Pass 3: S += beta * k[:,None] * (v - pred)[None,:] for (int idx = tid; idx < H * Dk * Dk; idx += stride) { int h = idx / (Dk * Dk); int rem = idx % (Dk * Dk); int r = rem / Dk; int d = rem % Dk; float b = beta[h]; float kk = k[h * Dk + r]; float delta = v[h * Dk + d] - o[h * Dk + d]; S[idx] += b * kk * delta; } grid.sync(); // Pass 4: o[h,d] = sum_r S[h,r,d] * q[h,r] for (int idx = tid; idx < H * Dk; idx += stride) { int h = idx / Dk; int d = idx % Dk; float acc = 0.f; for (int r = 0; r < Dk; ++r) { acc += S[(h * Dk + r) * Dk + d] * q[h * Dk + r]; } o[idx] = acc; } grid.sync(); } // Apply RoPE to a vector of dim (even/odd pairs). in/out float [dim] __device__ void apply_rope_vec( float* __restrict__ x, int pos, int dim, float theta ) { // only thread 0 path for small dim — caller handles parallel } // MLA with absorb: score via c_kv @ absorbed_q, online-ish softmax + v absorb // Simplified correct path: materialize scores [L,H], then softmax, then o. struct LayerPtrs { // norms const __nv_bfloat16* attn_norm; const __nv_bfloat16* moe_norm; // kind: 0=KDA, 1=MLA int kind; // KDA const uint8_t *q_wq, *k_wq, *v_wq, *g_wq, *o_wq; const __nv_bfloat16 *q_s, *k_s, *v_s, *g_s, *o_s; const __nv_bfloat16 *q_z, *k_z, *v_z, *g_z, *o_z; const __nv_bfloat16* beta_w; // [H, D] const __nv_bfloat16* conv_w; // [3, C, 4] float* S; // [H,Dk,Dk] __nv_bfloat16 *cq, *ck, *cv; // [3, C] each (KW-1=3) // MLA const uint8_t *mq_wq, *kva_wq, *kvb_wq, *mo_wq; const __nv_bfloat16 *mq_s, *kva_s, *kvb_s, *mo_s; const __nv_bfloat16 *mq_z, *kva_z, *kvb_z, *mo_z; __nv_bfloat16* c_kv; // [L+1, kv_lora] output buffer (write full) __nv_bfloat16* k_rope; // [L+1, qk_rope] const __nv_bfloat16* c_kv_in; // [L, kv_lora] const __nv_bfloat16* k_rope_in; // [L, qk_rope] // MoE const __nv_bfloat16* router_w; // [E, D] const uint8_t *gate_wq, *up_wq, *down_wq; const __nv_bfloat16 *gate_s, *up_s, *down_s; const __nv_bfloat16 *gate_z, *up_z, *down_z; const uint8_t *sgate_wq, *sup_wq, *sdown_wq; const __nv_bfloat16 *sgate_s, *sup_s, *sdown_s; const __nv_bfloat16 *sgate_z, *sup_z, *sdown_z; }; struct MegaParams { int D, H, Dk, C, group; int kv_lora, qk_nope, qk_rope, v_head; int n_experts, n_active, n_shared, moe_inter; int L; // context len (cache size before append) int n_layers; float scale_kda, scale_mla, rope_theta, routed_scaling; float* hidden; // [D] float working float* ws; // large workspace int ws_elems; LayerPtrs layers[4]; }; __device__ void gemv_int4_no_sync( const float* __restrict__ x, const uint8_t* __restrict__ w_q, const __nv_bfloat16* __restrict__ scales, const __nv_bfloat16* __restrict__ zeros, float* __restrict__ y, int K, int N, int group, int tid, int stride ) { (void)tid; (void)stride; gemv_int4_body(x, w_q, scales, zeros, y, K, N, group); } __global__ void megakernel(MegaParams* __restrict__ Pp) { cg::grid_group grid = cg::this_grid(); // Pointer arg keeps launch ABI small (better occupancy than 2KB by-value). MegaParams &P = *Pp; const int tid = blockIdx.x * blockDim.x + threadIdx.x; const int stride = gridDim.x * blockDim.x; const int D = P.D; const int H = P.H; const int Dk = P.Dk; const int C = P.C; const int group = P.group; const int m = P.moe_inter; const int E = P.n_experts; const int Ka = P.n_active; const int L = P.L; // Workspace layout (floats) // 0: sumsq (1) // 1: x_norm [D] // 1+D: tmp0 [max(C, H*(nope+rope), kv_lora+rope, m*Ka, D*2 ...)] float* sumsq = P.ws; float* x_norm = P.ws + 1; float* tmp0 = P.ws + 1 + D; // 16k region float* tmp1 = tmp0 + 16384; float* tmp2 = tmp1 + 16384; float* tmp3 = tmp2 + 16384; // Bigger region for MoE / MLA (scores need up to ~524k floats) float* big0 = P.ws + 1 + D + 65536; float* big1 = big0 + 65536; // scores [L1*H] lives here; rest of ws after for (int li = 0; li < P.n_layers; ++li) { LayerPtrs &lp = P.layers[li]; // ----- attn RMSNorm ----- rmsnorm_into(grid, P.hidden, lp.attn_norm, x_norm, sumsq, D); if (lp.kind == 0) { // ===== KDA ===== // q,k,v,g projections fused: one x load, 4 weight streams, Nflat=4*C float* q = tmp0; float* k = tmp1; float* v = tmp2; float* g = tmp3; { const int Nflat = 4 * C; const int ng = D / group; const uint8_t* wqs[4] = {lp.q_wq, lp.k_wq, lp.v_wq, lp.g_wq}; const __nv_bfloat16* scs[4] = {lp.q_s, lp.k_s, lp.v_s, lp.g_s}; const __nv_bfloat16* zss[4] = {lp.q_z, lp.k_z, lp.v_z, lp.g_z}; float* outs[4] = {q, k, v, g}; for (int n0 = blockIdx.x * blockDim.x + threadIdx.x; n0 < Nflat; n0 += gridDim.x * blockDim.x) { int pi = n0 / C; int col = n0 % C; float acc = 0.f; const uint8_t* wq = wqs[pi]; const __nv_bfloat16* sc = scs[pi]; const __nv_bfloat16* zz = zss[pi]; for (int gg = 0; gg < ng; ++gg) { float s = b2f(sc[gg * C + col]); float z = b2f(zz[gg * C + col]); int row0 = (gg * group) >> 1; #pragma unroll 4 for (int r = 0; r < (group >> 1); ++r) { int row = row0 + r; uint8_t packed = wq[row * C + col]; int kk = row << 1; acc = fmaf(x_norm[kk], ((float)(packed & 0xF) - z) * s, acc); acc = fmaf(x_norm[kk + 1], ((float)((packed >> 4) & 0xF) - z) * s, acc); } } outs[pi][col] = acc; } } grid.sync(); // short conv on q,k,v — conv_w is [3, C, 4] float* q2 = big0; float* k2 = big0 + C; float* v2 = big0 + 2 * C; // temp new prev windows in big1 __nv_bfloat16* new_cq = reinterpret_cast<__nv_bfloat16*>(big1); __nv_bfloat16* new_ck = new_cq + 3 * C; __nv_bfloat16* new_cv = new_ck + 3 * C; short_conv(grid, q, lp.cq, lp.conv_w + 0 * C * 4, q2, new_cq, C, 4); short_conv(grid, k, lp.ck, lp.conv_w + 1 * C * 4, k2, new_ck, C, 4); short_conv(grid, v, lp.cv, lp.conv_w + 2 * C * 4, v2, new_cv, C, 4); grid.sync(); // write back conv windows for (int i = tid; i < 3 * C; i += stride) { lp.cq[i] = new_cq[i]; lp.ck[i] = new_ck[i]; lp.cv[i] = new_cv[i]; } grid.sync(); // reshape to [H,Dk], scale q, softplus gate, beta float* qh = tmp0; float* kh = tmp1; float* vh = tmp2; float* gh = tmp3; for (int i = tid; i < C; i += stride) { qh[i] = q2[i] * P.scale_kda; kh[i] = k2[i]; vh[i] = v2[i]; gh[i] = -softplus(g[i]); } grid.sync(); float* beta = big0; gemv_bf16(x_norm, lp.beta_w, beta, D, H); grid.sync(); for (int i = tid; i < H; i += stride) beta[i] = 1.f / (1.f + expf(-beta[i])); grid.sync(); float* o = big0 + H; // [C] kda_update(grid, lp.S, qh, kh, vh, gh, beta, o, H, Dk); // o_proj float* attn_out = tmp0; gemv_int4_body(o, lp.o_wq, lp.o_s, lp.o_z, attn_out, C, D, group); grid.sync(); vec_add(grid, P.hidden, attn_out, D); } else { // ===== MLA ===== const int q_dim = H * (P.qk_nope + P.qk_rope); const int kv_dim = P.kv_lora + P.qk_rope; const int nope = P.qk_nope; const int rope = P.qk_rope; const int vh = P.v_head; const int r = P.kv_lora; float* q_full = tmp0; // [q_dim] float* kv = tmp1; // [kv_dim] gemv_int4_body(x_norm, lp.mq_wq, lp.mq_s, lp.mq_z, q_full, D, q_dim, group); gemv_int4_body(x_norm, lp.kva_wq, lp.kva_s, lp.kva_z, kv, D, kv_dim, group); grid.sync(); // split q into nope/rope, apply rope at position L // q layout per head: [nope | rope] float* q_nope = big0; // [H, nope] float* q_rope = big0 + H*nope; // [H, rope] float* c_new = big0 + H*(nope+rope); // [r] float* k_rope_new = c_new + r; // [rope] // RoPE inv freq for this position for (int i = tid; i < H * nope; i += stride) { int h = i / nope; int d = i % nope; q_nope[i] = q_full[h * (nope + rope) + d]; } for (int i = tid; i < H * rope; i += stride) { int h = i / rope; int d = i % rope; float val = q_full[h * (nope + rope) + nope + d]; // will rope below q_rope[i] = val; } for (int i = tid; i < r; i += stride) c_new[i] = kv[i]; for (int i = tid; i < rope; i += stride) k_rope_new[i] = kv[r + i]; grid.sync(); // Apply RoPE to q_rope [H,rope] and k_rope_new [rope] for (int i = tid; i < H * (rope / 2); i += stride) { int h = i / (rope / 2); int j = i % (rope / 2); float inv = 1.f / powf(P.rope_theta, (2.f * j) / (float)rope); float ang = (float)L * inv; float co = cosf(ang), si = sinf(ang); int base = h * rope + 2 * j; float even = q_rope[h * rope + 2 * j]; float odd = q_rope[h * rope + 2 * j + 1]; q_rope[h * rope + 2 * j] = even * co - odd * si; q_rope[h * rope + 2 * j + 1] = odd * co + even * si; } for (int j = tid; j < rope / 2; j += stride) { float inv = 1.f / powf(P.rope_theta, (2.f * j) / (float)rope); float ang = (float)L * inv; float co = cosf(ang), si = sinf(ang); float even = k_rope_new[2 * j]; float odd = k_rope_new[2 * j + 1]; k_rope_new[2 * j] = even * co - odd * si; k_rope_new[2 * j + 1] = odd * co + even * si; } grid.sync(); const int L1 = L + 1; // Absorb q into kv_b rank: q_abs[h, rr] float* q_abs = tmp0; // [H, r] { const int Nkv = H * (nope + vh); for (int idx = tid; idx < H * r; idx += stride) { int h = idx / r; int rr = idx % r; int g = rr / group; float acc = 0.f; int n0 = h * (nope + vh); #pragma unroll 4 for (int d = 0; d < nope; ++d) { int n = n0 + d; uint8_t packed = lp.kvb_wq[(rr >> 1) * Nkv + n]; float nibble = (rr & 1) ? (float)((packed >> 4) & 0xF) : (float)(packed & 0xF); float s = b2f(lp.kvb_s[g * Nkv + n]); float z = b2f(lp.kvb_z[g * Nkv + n]); acc = fmaf((nibble - z) * s, q_nope[h * nope + d], acc); } q_abs[idx] = acc; } } grid.sync(); // scores layout [H, L1] for contiguous softmax over L // Read past tokens from c_kv_in / k_rope_in (avoid full copy first); // new token from c_new / k_rope_new. Write full cache later. float* scores = big1; // [H * L1] for (int idx = tid; idx < H * L1; idx += stride) { int h = idx / L1; int l = idx % L1; float s = 0.f; const float* qa = q_abs + h * r; if (l < L) { const __nv_bfloat16* ck = lp.c_kv_in + l * r; #pragma unroll 4 for (int rr = 0; rr < r; rr += 4) { s = fmaf(b2f(ck[rr]), qa[rr], s); s = fmaf(b2f(ck[rr + 1]), qa[rr + 1], s); s = fmaf(b2f(ck[rr + 2]), qa[rr + 2], s); s = fmaf(b2f(ck[rr + 3]), qa[rr + 3], s); } const __nv_bfloat16* kr = lp.k_rope_in + l * rope; const float* qr = q_rope + h * rope; #pragma unroll 4 for (int d = 0; d < rope; d += 4) { s = fmaf(b2f(kr[d]), qr[d], s); s = fmaf(b2f(kr[d + 1]), qr[d + 1], s); s = fmaf(b2f(kr[d + 2]), qr[d + 2], s); s = fmaf(b2f(kr[d + 3]), qr[d + 3], s); } } else { for (int rr = 0; rr < r; ++rr) s = fmaf(c_new[rr], qa[rr], s); const float* qr = q_rope + h * rope; for (int d = 0; d < rope; ++d) s = fmaf(k_rope_new[d], qr[d], s); } scores[idx] = s * P.scale_mla; } grid.sync(); // Softmax over L (contiguous in scores[h*L1 + l]) for (int h = tid; h < H; h += stride) { float* sh = scores + h * L1; float mmax = -1e30f; for (int l = 0; l < L1; ++l) if (sh[l] > mmax) mmax = sh[l]; float sum = 0.f; for (int l = 0; l < L1; ++l) { float e = expf(sh[l] - mmax); sh[l] = e; sum += e; } float inv = 1.f / sum; for (int l = 0; l < L1; ++l) sh[l] *= inv; } grid.sync(); // ctxt[h, rr] = sum_l p[h,l] * c_kv[l, rr] // Parallel over (h, rr); p is contiguous in l for fixed h. float* ctxt = tmp2; // [H, r] for (int idx = tid; idx < H * r; idx += stride) { int h = idx / r; int rr = idx % r; float acc = 0.f; const float* sh = scores + h * L1; #pragma unroll 4 for (int l = 0; l < L; ++l) acc = fmaf(sh[l], b2f(lp.c_kv_in[l * r + rr]), acc); acc = fmaf(sh[L], c_new[rr], acc); ctxt[idx] = acc; } grid.sync(); // o[h,d] via absorbed V float* o_mla = tmp1; { const int Nkv = H * (nope + vh); for (int idx = tid; idx < H * vh; idx += stride) { int h = idx / vh; int d = idx % vh; int n = h * (nope + vh) + nope + d; float acc = 0.f; const float* ct = ctxt + h * r; for (int rr = 0; rr < r; ++rr) { int g = rr / group; uint8_t packed = lp.kvb_wq[(rr >> 1) * Nkv + n]; float nibble = (rr & 1) ? (float)((packed >> 4) & 0xF) : (float)(packed & 0xF); float s = b2f(lp.kvb_s[g * Nkv + n]); float z = b2f(lp.kvb_z[g * Nkv + n]); acc = fmaf(ct[rr], (nibble - z) * s, acc); } o_mla[idx] = acc; } } // Write expanded cache for state (vectorized-ish copy) for (int i = tid; i < L * r; i += stride) lp.c_kv[i] = lp.c_kv_in[i]; for (int i = tid; i < L * rope; i += stride) lp.k_rope[i] = lp.k_rope_in[i]; for (int i = tid; i < r; i += stride) lp.c_kv[L * r + i] = f2b(c_new[i]); for (int i = tid; i < rope; i += stride) lp.k_rope[L * rope + i] = f2b(k_rope_new[i]); grid.sync(); float* attn_out = tmp2; gemv_int4_body(o_mla, lp.mo_wq, lp.mo_s, lp.mo_z, attn_out, H * vh, D, group); grid.sync(); vec_add(grid, P.hidden, attn_out, D); } // ----- MoE ----- rmsnorm_into(grid, P.hidden, lp.moe_norm, x_norm, sumsq, D); // router logits float* logits = tmp0; // [E] gemv_bf16(x_norm, lp.router_w, logits, D, E); grid.sync(); // softmax + topk (serial on thread 0 — E=64 tiny) float* top_w = tmp1; // [Ka] int* top_idx = reinterpret_cast(tmp1 + Ka); if (blockIdx.x == 0 && threadIdx.x == 0) { float mmax = -1e30f; for (int e = 0; e < E; ++e) if (logits[e] > mmax) mmax = logits[e]; float sum = 0.f; for (int e = 0; e < E; ++e) { logits[e] = expf(logits[e] - mmax); sum += logits[e]; } float inv = 1.f / sum; for (int e = 0; e < E; ++e) logits[e] *= inv; // top-k for (int k = 0; k < Ka; ++k) { int best = -1; float bv = -1.f; for (int e = 0; e < E; ++e) { // skip already chosen bool used = false; for (int j = 0; j < k; ++j) if (top_idx[j] == e) { used = true; break; } if (!used && logits[e] > bv) { bv = logits[e]; best = e; } } top_idx[k] = best; top_w[k] = bv; } float wsum = 0.f; for (int k = 0; k < Ka; ++k) wsum += top_w[k]; float scale = P.routed_scaling / (wsum + 1e-9f); for (int k = 0; k < Ka; ++k) top_w[k] *= scale; } grid.sync(); // Batched MoE: all experts' gate/up/down with minimal grid.sync. const int nexp = Ka + P.n_shared; float* moe_out = big0; // [D] // big1 holds gate_all[nexp*m], up_all[nexp*m], down_all[nexp*D] float* gate_all = big1; float* up_all = gate_all + nexp * m; float* down_all = up_all + nexp * m; const int ng_d = D / group; const int ng_m = m / group; const int e_stride_gu = (D / 2) * m; const int e_stride_gu_sz = ng_d * m; const int e_stride_dn = (m / 2) * D; const int e_stride_dn_sz = ng_m * D; // Resolve expert ids into shared mem (per-block) __shared__ int sh_idx[16]; __shared__ float sh_wt[16]; if (threadIdx.x < nexp) { if (threadIdx.x < Ka) { sh_idx[threadIdx.x] = top_idx[threadIdx.x]; sh_wt[threadIdx.x] = top_w[threadIdx.x]; } else { sh_idx[threadIdx.x] = threadIdx.x - Ka; sh_wt[threadIdx.x] = 1.f; } } __syncthreads(); // Fused multi-expert gate+up GEMV (SiLU fused). Flattened N = nexp*m. { const int Nflat = nexp * m; const int ng = D / group; for (int n0 = (blockIdx.x * blockDim.x + threadIdx.x); n0 < Nflat; n0 += gridDim.x * blockDim.x) { int ki = n0 / m; int col = n0 % m; int e = sh_idx[ki]; bool routed = (ki < Ka); const uint8_t* gw = routed ? lp.gate_wq : lp.sgate_wq; const __nv_bfloat16* gs = routed ? lp.gate_s : lp.sgate_s; const __nv_bfloat16* gz = routed ? lp.gate_z : lp.sgate_z; const uint8_t* uw = routed ? lp.up_wq : lp.sup_wq; const __nv_bfloat16* us = routed ? lp.up_s : lp.sup_s; const __nv_bfloat16* uz = routed ? lp.up_z : lp.sup_z; gw += e * e_stride_gu; gs += e * e_stride_gu_sz; gz += e * e_stride_gu_sz; uw += e * e_stride_gu; us += e * e_stride_gu_sz; uz += e * e_stride_gu_sz; float acc_g = 0.f, acc_u = 0.f; for (int g = 0; g < ng; ++g) { float sg = b2f(gs[g * m + col]), zg = b2f(gz[g * m + col]); float su = b2f(us[g * m + col]), zu = b2f(uz[g * m + col]); int row0 = (g * group) >> 1; #pragma unroll 4 for (int r = 0; r < (group >> 1); ++r) { int row = row0 + r; int k = row << 1; uint8_t pg = gw[row * m + col]; uint8_t pu = uw[row * m + col]; acc_g = fmaf(x_norm[k], ((float)(pg & 0xF) - zg) * sg, acc_g); acc_g = fmaf(x_norm[k + 1], ((float)((pg >> 4) & 0xF) - zg) * sg, acc_g); acc_u = fmaf(x_norm[k], ((float)(pu & 0xF) - zu) * su, acc_u); acc_u = fmaf(x_norm[k + 1], ((float)((pu >> 4) & 0xF) - zu) * su, acc_u); } } gate_all[n0] = silu(acc_g) * acc_u; } } grid.sync(); // Fused multi-expert down GEMV into down_all[nexp, D], then weighted sum. { const int Nflat = nexp * D; const int ng = m / group; for (int n0 = tid; n0 < Nflat; n0 += stride) { int ki = n0 / D; int col = n0 % D; int e = sh_idx[ki]; bool routed = (ki < Ka); const uint8_t* dw = (routed ? lp.down_wq : lp.sdown_wq) + e * e_stride_dn; const __nv_bfloat16* ds = (routed ? lp.down_s : lp.sdown_s) + e * e_stride_dn_sz; const __nv_bfloat16* dz = (routed ? lp.down_z : lp.sdown_z) + e * e_stride_dn_sz; const float* hh = gate_all + ki * m; float acc = 0.f; for (int g = 0; g < ng; ++g) { float s = b2f(ds[g * D + col]); float z = b2f(dz[g * D + col]); int row0 = (g * group) >> 1; #pragma unroll 4 for (int r = 0; r < (group >> 1); ++r) { int row = row0 + r; uint8_t packed = dw[row * D + col]; int k = row << 1; acc = fmaf(hh[k], ((float)(packed & 0xF) - z) * s, acc); acc = fmaf(hh[k + 1], ((float)((packed >> 4) & 0xF) - z) * s, acc); } } down_all[n0] = acc; } grid.sync(); for (int i = tid; i < D; i += stride) { float acc = 0.f; for (int ki = 0; ki < nexp; ++ki) acc += sh_wt[ki] * down_all[ki * D + i]; moe_out[i] = acc; } grid.sync(); } vec_add(grid, P.hidden, moe_out, D); } } // Host launcher static MegaParams* d_params = nullptr; void ensure_params() { if (!d_params) { cudaMalloc(&d_params, sizeof(MegaParams)); } } // We pass everything via a host MegaParams copied to device each call. // Tensors are raw pointers. #define PTR_BF(t) (reinterpret_cast((t).data_ptr())) #define PTR_BF_M(t) (reinterpret_cast<__nv_bfloat16*>((t).data_ptr())) #define PTR_U8(t) (reinterpret_cast((t).data_ptr())) #define PTR_F(t) (reinterpret_cast((t).data_ptr())) #define PTR_CF(t) (reinterpret_cast((t).data_ptr())) void launch_megakernel( torch::Tensor hidden_f, // [D] float, in/out torch::Tensor workspace, // float workspace // flat list built in python — we'll use a simpler packed approach int64_t D, int64_t H, int64_t Dk, int64_t group, int64_t kv_lora, int64_t qk_nope, int64_t qk_rope, int64_t v_head, int64_t n_experts, int64_t n_active, int64_t n_shared, int64_t moe_inter, int64_t L, int64_t n_layers, double scale_kda, double scale_mla, double rope_theta, double routed_scaling, // layers as lists of tensors — use c10::List const std::vector& t_attn_norm, const std::vector& t_moe_norm, const std::vector& kinds, // KDA quant bundles: for each layer, 5 linears * 3 + beta + conv + S + cq/ck/cv // Too many args — pack into a single byte blob of pointers from Python. torch::Tensor ptr_blob // int64 tensor of packed pointer bits + meta ) { // ptr_blob layout defined by Python packer TORCH_CHECK(hidden_f.is_cuda() && hidden_f.dtype() == torch::kFloat32); TORCH_CHECK(workspace.is_cuda() && workspace.dtype() == torch::kFloat32); MegaParams hp; memset(&hp, 0, sizeof(hp)); hp.D = (int)D; hp.H = (int)H; hp.Dk = (int)Dk; hp.C = (int)(H*Dk); hp.group = (int)group; hp.kv_lora = (int)kv_lora; hp.qk_nope = (int)qk_nope; hp.qk_rope = (int)qk_rope; hp.v_head = (int)v_head; hp.n_experts = (int)n_experts; hp.n_active = (int)n_active; hp.n_shared = (int)n_shared; hp.moe_inter = (int)moe_inter; hp.L = (int)L; hp.n_layers = (int)n_layers; hp.scale_kda = (float)scale_kda; hp.scale_mla = (float)scale_mla; hp.rope_theta = (float)rope_theta; hp.routed_scaling = (float)routed_scaling; hp.hidden = PTR_F(hidden_f); hp.ws = PTR_F(workspace); hp.ws_elems = (int)workspace.numel(); // Unpack pointer blob: Python packs for each layer a fixed sequence of int64 ptrs // See Python for order. const int64_t* p = ptr_blob.data_ptr(); int off = 0; for (int li = 0; li < n_layers; ++li) { LayerPtrs &lp = hp.layers[li]; lp.kind = (int)kinds[li]; lp.attn_norm = reinterpret_cast(p[off++]); lp.moe_norm = reinterpret_cast(p[off++]); if (lp.kind == 0) { lp.q_wq = reinterpret_cast(p[off++]); lp.q_s = reinterpret_cast(p[off++]); lp.q_z = reinterpret_cast(p[off++]); lp.k_wq = reinterpret_cast(p[off++]); lp.k_s = reinterpret_cast(p[off++]); lp.k_z = reinterpret_cast(p[off++]); lp.v_wq = reinterpret_cast(p[off++]); lp.v_s = reinterpret_cast(p[off++]); lp.v_z = reinterpret_cast(p[off++]); lp.g_wq = reinterpret_cast(p[off++]); lp.g_s = reinterpret_cast(p[off++]); lp.g_z = reinterpret_cast(p[off++]); lp.o_wq = reinterpret_cast(p[off++]); lp.o_s = reinterpret_cast(p[off++]); lp.o_z = reinterpret_cast(p[off++]); lp.beta_w = reinterpret_cast(p[off++]); lp.conv_w = reinterpret_cast(p[off++]); lp.S = reinterpret_cast(p[off++]); lp.cq = reinterpret_cast<__nv_bfloat16*>(p[off++]); lp.ck = reinterpret_cast<__nv_bfloat16*>(p[off++]); lp.cv = reinterpret_cast<__nv_bfloat16*>(p[off++]); } else { lp.mq_wq = reinterpret_cast(p[off++]); lp.mq_s = reinterpret_cast(p[off++]); lp.mq_z = reinterpret_cast(p[off++]); lp.kva_wq = reinterpret_cast(p[off++]); lp.kva_s = reinterpret_cast(p[off++]); lp.kva_z = reinterpret_cast(p[off++]); lp.kvb_wq = reinterpret_cast(p[off++]); lp.kvb_s = reinterpret_cast(p[off++]); lp.kvb_z = reinterpret_cast(p[off++]); lp.mo_wq = reinterpret_cast(p[off++]); lp.mo_s = reinterpret_cast(p[off++]); lp.mo_z = reinterpret_cast(p[off++]); lp.c_kv_in = reinterpret_cast(p[off++]); lp.k_rope_in = reinterpret_cast(p[off++]); lp.c_kv = reinterpret_cast<__nv_bfloat16*>(p[off++]); lp.k_rope = reinterpret_cast<__nv_bfloat16*>(p[off++]); } // MoE (all layers) lp.router_w = reinterpret_cast(p[off++]); lp.gate_wq = reinterpret_cast(p[off++]); lp.gate_s = reinterpret_cast(p[off++]); lp.gate_z = reinterpret_cast(p[off++]); lp.up_wq = reinterpret_cast(p[off++]); lp.up_s = reinterpret_cast(p[off++]); lp.up_z = reinterpret_cast(p[off++]); lp.down_wq = reinterpret_cast(p[off++]); lp.down_s = reinterpret_cast(p[off++]); lp.down_z = reinterpret_cast(p[off++]); lp.sgate_wq = reinterpret_cast(p[off++]); lp.sgate_s = reinterpret_cast(p[off++]); lp.sgate_z = reinterpret_cast(p[off++]); lp.sup_wq = reinterpret_cast(p[off++]); lp.sup_s = reinterpret_cast(p[off++]); lp.sup_z = reinterpret_cast(p[off++]); lp.sdown_wq = reinterpret_cast(p[off++]); lp.sdown_s = reinterpret_cast(p[off++]); lp.sdown_z = reinterpret_cast(p[off++]); } // Best measured on this box: ~128 CTAs (GEMV BW vs grid.sync tradeoff). int need = 128; if (L > 8192) need = 160; int dev = hidden_f.get_device(); cudaDeviceProp prop; cudaGetDeviceProperties(&prop, dev); int nb_per_sm = 0; cudaOccupancyMaxActiveBlocksPerMultiprocessor(&nb_per_sm, megakernel, BLOCK, 0); int max_coop = max(1, nb_per_sm * prop.multiProcessorCount); int num_blocks = need < max_coop ? need : max_coop; // Device-side params — keeps kernel arg as a single pointer. static MegaParams* d_hp = nullptr; if (!d_hp) { cudaMalloc(&d_hp, sizeof(MegaParams)); } cudaMemcpyAsync(d_hp, &hp, sizeof(MegaParams), cudaMemcpyHostToDevice, at::cuda::getCurrentCUDAStream().stream()); auto stream = at::cuda::getCurrentCUDAStream(); void* args[] = { &d_hp }; cudaError_t err = cudaLaunchCooperativeKernel( (void*)megakernel, dim3(num_blocks), dim3(BLOCK), args, 0, stream.stream()); TORCH_CHECK(err == cudaSuccess, "cooperative launch failed: ", cudaGetErrorString(err), " blocks=", num_blocks, " sizeof(MegaParams)=", sizeof(MegaParams)); } """ _mod = None _CPP_SRC = r""" #include #include void launch_megakernel( torch::Tensor hidden_f, torch::Tensor workspace, int64_t D, int64_t H, int64_t Dk, int64_t group, int64_t kv_lora, int64_t qk_nope, int64_t qk_rope, int64_t v_head, int64_t n_experts, int64_t n_active, int64_t n_shared, int64_t moe_inter, int64_t L, int64_t n_layers, double scale_kda, double scale_mla, double rope_theta, double routed_scaling, const std::vector& t_attn_norm, const std::vector& t_moe_norm, const std::vector& kinds, torch::Tensor ptr_blob ); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("launch_megakernel", &launch_megakernel, "Kimi-Linear W4A16 decode megakernel"); } """ def _get_mod(): global _mod if _mod is not None: return _mod import os os.environ.setdefault("CUDA_HOME", "/usr/local/cuda-13") os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0") build_dir = Path(__file__).resolve().parent / "_mega_build" build_dir.mkdir(exist_ok=True) _mod = load_inline( name="kimi_linear_mega_v26", cpp_sources=[_CPP_SRC], cuda_sources=[_CUDA_SRC], functions=None, extra_cuda_cflags=[ "-O3", "--use_fast_math", "-lineinfo", "-std=c++17", ], extra_ldflags=["-lcuda"], with_cuda=True, verbose=True, build_directory=str(build_dir), ) return _mod def _ip(t: torch.Tensor) -> int: return t.data_ptr() def _pack_layer_ptrs(block: Block, st: dict, c_kv_out=None, k_rope_out=None) -> list: """Pack int64 pointers for one layer in the order expected by the CUDA launcher.""" p: list[int] = [] p += [_ip(block.attn_norm), _ip(block.moe_norm)] if block.kind == "K": attn: KDA = block.attn # type: ignore for lin in (attn.q_proj, attn.k_proj, attn.v_proj, attn.g_proj, attn.o_proj): p += [_ip(lin.w_q), _ip(lin.scales), _ip(lin.zeros)] p += [_ip(attn.beta_proj.weight), _ip(attn.conv_w)] p += [_ip(st["S"]), _ip(st["cq"]), _ip(st["ck"]), _ip(st["cv"])] else: attn: MLA = block.attn # type: ignore for lin in (attn.q_proj, attn.kv_a, attn.kv_b, attn.o_proj): p += [_ip(lin.w_q), _ip(lin.scales), _ip(lin.zeros)] p += [_ip(st["c_kv"]), _ip(st["k_rope"]), _ip(c_kv_out), _ip(k_rope_out)] moe: MoE = block.moe p.append(_ip(moe.router.weight)) for ex in (moe.gate, moe.up, moe.down, moe.s_gate, moe.s_up, moe.s_down): p += [_ip(ex.w_q), _ip(ex.scales), _ip(ex.zeros)] return p class Model(nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern) self.reset_parameters() self._ws = None self._hidden_f = None self._hidden_out = None self._mod = None self._weight_ptrs = None # cached static weight/buffer pointers self._kinds = [0 if k == "K" else 1 for k in cfg.pattern] self._mla_cap = 0 self._mla_c = None self._mla_k = None def reset_parameters(self): g = torch.Generator(device="cpu").manual_seed(1234) for mod in self.modules(): if isinstance(mod, (QuantLinear, QuantExperts)): mod.init_random(g) elif isinstance(mod, nn.Linear): nn.init.normal_(mod.weight, 0.0, 0.02, generator=g) elif isinstance(mod, KDA): nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g) def _ensure_ws(self, device): n = 4 * 1024 * 1024 if self._ws is None or self._ws.device != device or self._ws.numel() < n: self._ws = torch.empty(n, device=device, dtype=torch.float32) self._hidden_f = torch.empty(self.cfg.hidden, device=device, dtype=torch.float32) self._hidden_out = torch.empty(self.cfg.hidden, device=device, dtype=torch.bfloat16) self._weight_ptrs = None self._mla_cap = 0 def _cache_weight_ptrs(self): """Pack weight/param pointers that are stable across steps (not state).""" ptrs = [] for blk in self.blocks: ptrs += [_ip(blk.attn_norm), _ip(blk.moe_norm)] if blk.kind == "K": attn: KDA = blk.attn # type: ignore for lin in (attn.q_proj, attn.k_proj, attn.v_proj, attn.g_proj, attn.o_proj): ptrs += [_ip(lin.w_q), _ip(lin.scales), _ip(lin.zeros)] ptrs += [_ip(attn.beta_proj.weight), _ip(attn.conv_w)] # state slots filled later: S, cq, ck, cv ptrs += [0, 0, 0, 0] else: attn: MLA = blk.attn # type: ignore for lin in (attn.q_proj, attn.kv_a, attn.kv_b, attn.o_proj): ptrs += [_ip(lin.w_q), _ip(lin.scales), _ip(lin.zeros)] # state: c_kv_in, k_rope_in, c_kv_out, k_rope_out ptrs += [0, 0, 0, 0] moe: MoE = blk.moe ptrs.append(_ip(moe.router.weight)) for ex in (moe.gate, moe.up, moe.down, moe.s_gate, moe.s_up, moe.s_down): ptrs += [_ip(ex.w_q), _ip(ex.scales), _ip(ex.zeros)] self._weight_ptrs = ptrs @torch.no_grad() def step(self, hidden, state): cfg = self.cfg device = hidden.device self._ensure_ws(device) if self._mod is None: self._mod = _get_mod() if self._weight_ptrs is None: self._cache_weight_ptrs() self._hidden_f.copy_(hidden.float()) # Grow MLA buffers if needed L = 0 mla_i = None for i, blk in enumerate(self.blocks): if blk.kind == "M": mla_i = i L = state[i]["c_kv"].shape[0] need = L + 1 if self._mla_cap < need or self._mla_c is None or self._mla_c.device != device: cap = max(need, self._mla_cap * 2 if self._mla_cap else need) self._mla_c = torch.empty(cap, cfg.kv_lora, device=device, dtype=cfg.dtype) self._mla_k = torch.empty(cap, cfg.qk_rope, device=device, dtype=cfg.dtype) self._mla_cap = cap break # Fill state pointers into a copy of the cached weight ptr list ptrs = list(self._weight_ptrs) off = 0 for i, blk in enumerate(self.blocks): off += 2 # norms if blk.kind == "K": off += 5 * 3 + 2 # 5 linears * 3 + beta + conv st = state[i] if st["S"].dtype != torch.float32: st["S"] = st["S"].float() st["S"] = st["S"].contiguous() st["cq"] = st["cq"].contiguous() st["ck"] = st["ck"].contiguous() st["cv"] = st["cv"].contiguous() ptrs[off] = _ip(st["S"]) ptrs[off + 1] = _ip(st["cq"]) ptrs[off + 2] = _ip(st["ck"]) ptrs[off + 3] = _ip(st["cv"]) off += 4 else: off += 4 * 3 # 4 linears st = state[i] st["c_kv"] = st["c_kv"].contiguous() st["k_rope"] = st["k_rope"].contiguous() c_out = self._mla_c[: L + 1] k_out = self._mla_k[: L + 1] ptrs[off] = _ip(st["c_kv"]) ptrs[off + 1] = _ip(st["k_rope"]) ptrs[off + 2] = _ip(c_out) ptrs[off + 3] = _ip(k_out) off += 4 # MoE: router + 6 experts * 3 off += 1 + 6 * 3 ptr_blob = torch.tensor(ptrs, dtype=torch.int64) # CPU H, Dk = cfg.kda_heads, cfg.kda_head_dim self._mod.launch_megakernel( self._hidden_f, self._ws, cfg.hidden, H, Dk, cfg.group, cfg.kv_lora, cfg.qk_nope, cfg.qk_rope, cfg.v_head, cfg.n_experts, cfg.n_active, cfg.n_shared, cfg.moe_inter, L, len(self.blocks), float(Dk ** -0.5), float((cfg.qk_nope + cfg.qk_rope) ** -0.5), float(cfg.rope_theta), float(cfg.routed_scaling), [b.attn_norm for b in self.blocks], [b.moe_norm for b in self.blocks], self._kinds, ptr_blob, ) out = self._hidden_f.to(dtype=torch.bfloat16) if mla_i is not None: state[mla_i]["c_kv"] = self._mla_c[: L + 1] state[mla_i]["k_rope"] = self._mla_k[: L + 1] return out, state # Optional helpers (not required by check, but useful) def init_state(cfg: Config, context_len: int, seed: int) -> list: dev = torch.device("cuda:0") g = torch.Generator(device=dev).manual_seed(seed) H, Dk = cfg.kda_heads, cfg.kda_head_dim C = H * Dk state = [] for kind in cfg.pattern: if kind == "K": state.append({ "S": torch.randn(H, Dk, Dk, device=dev, generator=g) * 0.05, "cq": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1, "ck": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1, "cv": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1, }) else: state.append({ "c_kv": torch.randn(context_len, cfg.kv_lora, device=dev, generator=g, dtype=cfg.dtype) * 0.1, "k_rope": torch.randn(context_len, cfg.qk_rope, device=dev, generator=g, dtype=cfg.dtype) * 0.1, }) return state def init_token(cfg: Config, seed: int) -> torch.Tensor: dev = torch.device("cuda:0") g = torch.Generator(device=dev).manual_seed(seed + 1) return torch.randn(cfg.hidden, device=dev, generator=g, dtype=cfg.dtype) * 0.25