"""Kimi-Linear W4A16 hybrid decode unit -- single-megakernel solution. The entire per-token decode step (3x KDA + 1x MLA attention, 4x MoE FFN, all int4 dequant GEMVs, convs, state updates, RMSNorms and residuals) runs as ONE CUDA kernel launch inside step(). Grid-wide phase barriers are done with a custom atomic barrier inside a cooperatively-launched persistent kernel (188 blocks x 1024 threads on the RTX PRO 6000). W4A16 dequant is fused into every GEMV: int4 nibbles are unpacked and scaled per-128-group on the fly; no bf16 weight matrix is ever written out. KIMI_EAGER=1 forces the slow-but-faithful eager path (debugging aid). """ from __future__ import annotations import math import os from dataclasses import dataclass, field import torch import torch.nn as nn import torch.nn.functional as F OP_TYPE = "kimi_linear_w4a16_decode" HARDWARE_REQUIRED = ["RTX_PRO_6000"] 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))) # --------------------------------------------------------------------------- # # W4A16 quantization (identical format to reference.py) # --------------------------------------------------------------------------- # 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 _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor: out = torch.empty((K, w_packed.shape[1]), dtype=torch.uint8, device=w_packed.device) out[0::2] = w_packed & 0xF out[1::2] = (w_packed >> 4) & 0xF return out def dequant(w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor, K: int, group: int) -> torch.Tensor: wu = _unpack_int4(w_q, K).to(torch.bfloat16) s = scales.repeat_interleave(group, dim=0) z = zeros.repeat_interleave(group, dim=0) return (wu - z) * s 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 weight_bf(self) -> torch.Tensor: return dequant(self.w_q, self.scales, self.zeros, self.in_f, self.group) def forward(self, x: torch.Tensor) -> torch.Tensor: return (x.float() @ self.weight_bf().float()).to(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)) def weight_bf(self, e: int) -> torch.Tensor: return dequant(self.w_q[e], self.scales[e], self.zeros[e], self.in_f, self.group) # --------------------------------------------------------------------------- # # helpers (eager oracle path) # --------------------------------------------------------------------------- # def _rmsnorm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: 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: int, dim: int, theta: float, 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: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: xf = x.float() even, odd = xf[..., 0::2], xf[..., 1::2] out = torch.empty_like(xf) out[..., 0::2] = even * cos - odd * sin out[..., 1::2] = odd * cos + even * sin return out.to(x.dtype) # --------------------------------------------------------------------------- # # layers -- eager oracle path (used for debugging / fallback) # --------------------------------------------------------------------------- # 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 def _short_conv(self, val, prev, idx): win = torch.cat([prev, val[None]], dim=0) w = self.conv_w[idx].float().transpose(0, 1) out = (win.float() * w).sum(0) return F.silu(out).to(val.dtype), win[1:] def step(self, x, st): H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim q, k, v = self.q_proj(x), self.k_proj(x), self.v_proj(x) q, st["cq"] = self._short_conv(q, st["cq"], 0) k, st["ck"] = self._short_conv(k, st["ck"], 1) v, st["cv"] = self._short_conv(v, st["cv"], 2) q = q.view(H, Dk).float() * self.scale k = k.view(H, Dk).float() v = v.view(H, Dk).float() g = (-F.softplus(self.g_proj(x).float())).view(H, Dk) beta = torch.sigmoid(self.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) st["S"] = S return self.o_proj(o.reshape(H * Dk).to(torch.bfloat16)) 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 def step(self, x, st): cfg = self.cfg H = cfg.mla_heads pos = st["c_kv"].shape[0] q = self.q_proj(x).view(H, cfg.qk_nope + cfg.qk_rope) q_nope = q[:, : cfg.qk_nope].float() q_rope = q[:, cfg.qk_nope :] kv = self.kv_a(x) 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) st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0) st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0) kvb = self.kv_b(st["c_kv"]).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, st["k_rope"].float())) * self.scale p = torch.softmax(scores, dim=0) o = torch.einsum("lh,lhd->hd", p, v) return self.o_proj(o.reshape(H * cfg.v_head).to(torch.bfloat16)) 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) def _ffn(self, x, experts_g, experts_u, experts_d, e): h = F.silu(x.float() @ experts_g.weight_bf(e).float()) * (x.float() @ experts_u.weight_bf(e).float()) return h @ experts_d.weight_bf(e).float() def step(self, x): cfg = self.cfg probs = torch.softmax(self.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) for j in range(cfg.n_active): out = out + w[j] * self._ffn(x, self.gate, self.up, self.down, int(idx[j])) for s in range(cfg.n_shared): out = out + self._ffn(x, self.s_gate, self.s_up, self.s_down, s) return out.to(torch.bfloat16) 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) def step(self, x, st): h = x + self.attn.step(_rmsnorm(x, self.attn_norm), st) return h + self.moe.step(_rmsnorm(h, self.moe_norm)) # --------------------------------------------------------------------------- # # Megakernel integration # --------------------------------------------------------------------------- # from kimi_mega import MegaRunner # local module in this directory 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._runner = None self._eager_only = os.environ.get("KIMI_EAGER", "0") == "1" def step(self, hidden, state): if self._eager_only: return self._step_eager(hidden, state) if self._runner is None: dev = hidden.device self.to(dev) # make sure all buffers are on the right device self._runner = MegaRunner(self, dev) return self._runner.step(hidden, state) def _step_eager(self, hidden, state): with torch.no_grad(): for i, blk in enumerate(self.blocks): hidden = blk.step(hidden, state[i]) return hidden, state # --------------------------------------------------------------------------- # # state / inputs (mirror reference helpers) # --------------------------------------------------------------------------- # 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 # ================================================================== # ===== sidecar: kimi_mega.py (49178 bytes, loaded by solution.py) ===== # ================================================================== """Megakernel runner for the Kimi-Linear W4A16 decode unit. Builds one CUDA kernel (via torch.utils.cpp_extension.load_inline) that executes the ENTIRE per-token decode step as a single cooperative launch: 3x [KDA attn + MoE] + 1x [MLA attn + MoE], all int4 dequant GEMVs fused. Phase ordering inside the kernel is enforced by a custom atomic grid barrier (the kernel is launched cooperatively, so all blocks are resident). """ from __future__ import annotations import os import torch CPP_DECL = r""" #include #include int64_t kimi_mega_setup(std::vector sp_host); int64_t kimi_mega_step(int64_t sp, std::vector dyn, int64_t L, int64_t gen_base, int64_t NCH, int64_t stop_after); """ CUDA_SRC = r""" #include #include #include #include #include using bf16 = __nv_bfloat16; // ------------------------------------------------------------------ consts #define DD 2304 #define HC 4096 #define NHEADS 32 #define HDIM 128 #define GRP 128 #define MOE_M 1024 #define MLB_Q 6144 #define MLB_KVA 576 #define MLB_KVB 8192 #define MLB_QD 192 #define KVL 512 #define QK_NOPE 128 #define QK_ROPE 64 #define V_HEAD 128 #define KDA_SCALE 0.08838834764831845f #define MLA_SCALE 0.07216878364870323f #define RSCALE 2.446f #define NEG_INF (-__int_as_float(0x7f800000)) // static param layout (must match python builder EXACTLY) #define SPK(b) ((b)*19) #define K_NORMA 0 #define K_NORMM 1 #define K_QW 2 #define K_QS 3 #define K_QZ 4 #define K_KW 5 #define K_KS 6 #define K_KZ 7 #define K_VW 8 #define K_VS 9 #define K_VZ 10 #define K_GW 11 #define K_GS 12 #define K_GZ 13 #define K_OW 14 #define K_OS 15 #define K_OZ 16 #define K_CONV 17 #define K_BETA 18 #define SPE(b) (57 + (b)*19) #define E_ROUT 0 #define E_GW 1 #define E_GS 2 #define E_GZ 3 #define E_UW 4 #define E_US 5 #define E_UZ 6 #define E_DW 7 #define E_DS 8 #define E_DZ 9 #define E_SGW 10 #define E_SGS 11 #define E_SGZ 12 #define E_SUW 13 #define E_SUS 14 #define E_SUZ 15 #define E_SDW 16 #define E_SDS 17 #define E_SDZ 18 #define SPM 133 #define M_NORMA 0 #define M_NORMM 1 #define M_QW 2 #define M_QS 3 #define M_QZ 4 #define M_KVAW 5 #define M_KVAS 6 #define M_KVAZ 7 #define M_KVBW 8 #define M_KVBS 9 #define M_KVBZ 10 #define M_OW 11 #define M_OS 12 #define M_OZ 13 #define SPW 147 #define W_Q 0 #define W_K 1 #define W_V 2 #define W_G 3 #define W_O 4 #define W_XN 5 #define W_BX 6 #define W_A 7 #define W_D 8 #define W_HH 9 #define W_TKI 10 #define W_TKW 11 #define W_QM 12 #define W_QF 13 #define W_KRR 14 #define W_AP 15 #define W_CK 16 #define W_KR 17 #define W_CTR 18 #define W_INV 19 #define SP_TOTAL 167 struct DynArgs { long long x_in; long long h_out; long long S[3]; long long cq[3]; long long ck[3]; long long cv[3]; int L; int NCH; int stop_after; unsigned gen_base; }; // ------------------------------------------------------------------ helpers __device__ __forceinline__ float bf2f(bf16 v) { return __bfloat162float(v); } __device__ __forceinline__ bf16 f2bf(float v) { return __float2bfloat16_rn(v); } union B8 { uint4 v; bf16 f[8]; }; union B4 { uint2 v; bf16 f[4]; }; __device__ __forceinline__ float siluf(float x) { return x / (1.f + expf(-x)); } __device__ __forceinline__ float softplusf(float x) { return x > 20.f ? x : log1pf(expf(x)); } // tensor-core mma for the MLA attention phase (bf16, fp32 accumulate) __device__ __forceinline__ void mma_bf16(float* c, const unsigned* a, const unsigned* b) { asm volatile( "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); } #define M3_RT 48 #define M3_STP 584 #define M3_QP 584 __device__ __forceinline__ float block_sum(float v, float* scratch) { int w = threadIdx.x >> 5, l = threadIdx.x & 31; #pragma unroll for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(~0u, v, o); if (l == 0) scratch[w] = v; __syncthreads(); if (w == 0) { float a = (l < (int)(blockDim.x >> 5)) ? scratch[l] : 0.f; #pragma unroll for (int o = 16; o > 0; o >>= 1) a += __shfl_xor_sync(~0u, a, o); if (l == 0) scratch[0] = a; } __syncthreads(); float r = scratch[0]; __syncthreads(); return r; } __device__ __forceinline__ void compute_smx(const float* sx, float* smx, int ng) { int w = threadIdx.x >> 5, l = threadIdx.x & 31; __syncthreads(); if (w < ng) { const float* g = sx + w * 128; float a = g[l] + g[l + 32] + g[l + 64] + g[l + 96]; #pragma unroll for (int o = 16; o > 0; o >>= 1) a += __shfl_xor_sync(~0u, a, o); if (l == 0) smx[w] = a; } __syncthreads(); } __device__ __forceinline__ void load_sx_bf16(const bf16* __restrict__ x, float* sx, int K) { int nv = K >> 3; for (int i = threadIdx.x; i < nv; i += blockDim.x) { B8 u; u.v = __ldcg((const uint4*)(x + 8 * i)); #pragma unroll for (int j = 0; j < 8; ++j) sx[8 * i + j] = bf2f(u.f[j]); } } __device__ __forceinline__ void load_sx_f32(const float* __restrict__ x, float* sx, int K) { int nv = K >> 2; float4* sx4 = (float4*)sx; for (int i = threadIdx.x; i < nv; i += blockDim.x) { sx4[i] = __ldcg((const float4*)(x + 4 * i)); } } __device__ __forceinline__ void gbar(unsigned* cnt, unsigned* gen, unsigned target) { __syncthreads(); if (threadIdx.x == 0) { __threadfence(); unsigned a = atomicAdd(cnt, 1u); if (a == gridDim.x - 1) { atomicExch(cnt, 0u); __threadfence(); atomicAdd(gen, 1u); } while (atomicAdd(gen, 0u) < target) __nanosleep(32); } __syncthreads(); } __device__ __forceinline__ int grab(unsigned* ctr_slot, int ntasks, int* ismall) { __syncthreads(); if (threadIdx.x == 0) ismall[0] = (int)atomicAdd(ctr_slot, 1u); __syncthreads(); return ismall[0]; } // ------------------------------------------------------------------ gemv64 __device__ __forceinline__ void gemv64( const unsigned char* __restrict__ wq, const bf16* __restrict__ sc, const bf16* __restrict__ zr, int N, int g0, int ng, int col0, const float* __restrict__ sx, const float* __restrict__ smx, float* team_part, float* out64) { const int team = threadIdx.x >> 7; const int lane = threadIdx.x & 31; const int wrp = (threadIdx.x >> 5) & 3; const int base = ng >> 3, rem = ng & 7; const int gstart = team * base + min(team, rem); const int gcount = base + (team < rem ? 1 : 0); const int cb = (lane & 7) * 8; const int lsub = lane >> 3; const int ldb = N; float accf[8]; #pragma unroll for (int j = 0; j < 8; ++j) accf[j] = 0.f; const unsigned char* wbase = wq + (size_t)(g0 + gstart) * 64 * (size_t)ldb + col0 + cb; for (int gl = 0; gl < gcount; ++gl) { float dotq[8]; #pragma unroll for (int j = 0; j < 8; ++j) dotq[j] = 0.f; const unsigned char* wr = wbase + (size_t)(gl * 64 + wrp * 4) * (size_t)ldb; uint2 wv[4]; #pragma unroll for (int it = 0; it < 4; ++it) wv[it] = __ldg((const uint2*)(wr + (size_t)(it * 16 + lsub) * ldb)); #pragma unroll for (int it = 0; it < 4; ++it) { int rloc = (gstart + gl) * 64 + wrp * 4 + it * 16 + lsub; float x0 = sx[2 * rloc], x1 = sx[2 * rloc + 1]; const unsigned char* pb = (const unsigned char*)&wv[it]; #pragma unroll for (int j = 0; j < 8; ++j) { int qv = pb[j]; dotq[j] += x0 * (float)(qv & 15) + x1 * (float)(qv >> 4); } } int ga = g0 + gstart + gl; uint4 sv = __ldg((const uint4*)(sc + (size_t)ga * N + col0 + cb)); uint4 zv = __ldg((const uint4*)(zr + (size_t)ga * N + col0 + cb)); const bf16* sp = (const bf16*)&sv; const bf16* zp = (const bf16*)&zv; float sxg = smx[gstart + gl]; // zero-point correction must be applied exactly once per group/column: // only one lane of the team's k-split applies it. bool zleader = (wrp == 0 && lsub == 0); #pragma unroll for (int j = 0; j < 8; ++j) { float sf = bf2f(sp[j]); accf[j] += sf * dotq[j] - (zleader ? sf * bf2f(zp[j]) * sxg : 0.f); } } // reduce the 4 lsub lanes' k-partials (lanes l, l+8, l+16, l+24 share cols) #pragma unroll for (int j = 0; j < 8; ++j) { accf[j] += __shfl_xor_sync(~0u, accf[j], 8); accf[j] += __shfl_xor_sync(~0u, accf[j], 16); } if (lsub == 0) { float* warp_part = team_part + (threadIdx.x >> 5) * 64 + cb; #pragma unroll for (int j = 0; j < 8; ++j) warp_part[j] = accf[j]; } __syncthreads(); if (threadIdx.x < 64) { float a = 0.f; #pragma unroll 8 for (int t = 0; t < 32; ++t) a += team_part[t * 64 + threadIdx.x]; out64[threadIdx.x] = a; } __syncthreads(); } __device__ __forceinline__ void rmsn_apply(float* xf, const bf16* __restrict__ w, float* sx, float* scratch, bf16* ws_xn) { float part = 0.f; for (int i = threadIdx.x; i < DD; i += blockDim.x) part += xf[i] * xf[i]; float tot = block_sum(part, scratch); float rs = 1.f / sqrtf(tot * (1.f / (float)DD) + 1e-6f); for (int i = threadIdx.x; i < DD; i += blockDim.x) { bf16 nb = f2bf(xf[i] * rs * bf2f(__ldg(w + i))); sx[i] = bf2f(nb); if (ws_xn) ws_xn[i] = nb; } __syncthreads(); } // recombine x_b (block b>0 input) from previous-block partials, into h1x __device__ __forceinline__ void recomb_xb(const bf16* xp, const float* A0, const float* Dm, float* h1x, int tid, int bdim) { for (int i = tid; i < DD; i += bdim) { float xb_ = bf2f(__ldcg(xp + i)); float att = bf2f(f2bf(__ldcg(A0 + i) + __ldcg(A0 + DD + i))); float h1 = bf2f(f2bf(xb_ + att)); float macc = __ldcg(Dm + 8 * DD + i); #pragma unroll for (int s = 0; s < 8; ++s) macc += __ldcg(Dm + s * DD + i); h1x[i] = bf2f(f2bf(h1 + bf2f(f2bf(macc)))); } __syncthreads(); } // router + top8 on smem xn vector sx; results in ismall[80..88), w in small[96..104) __device__ __forceinline__ void router_topk(const bf16* rw, const float* sx, float* small, int* ismall, int tid) { int w = tid >> 5, l = tid & 31; if (w < 32) { float a0 = 0.f, a1 = 0.f; const bf16* r0 = rw + w * DD; const bf16* r1 = rw + (w + 32) * DD; for (int i = 0; i < DD / 32; i += 4) { B4 u0, u1; u0.v = __ldg((const uint2*)(r0 + l * 4 + i * 32)); u1.v = __ldg((const uint2*)(r1 + l * 4 + i * 32)); #pragma unroll for (int j = 0; j < 4; ++j) { a0 += bf2f(u0.f[j]) * sx[l * 4 + i * 32 + j]; a1 += bf2f(u1.f[j]) * sx[l * 4 + i * 32 + j]; } } #pragma unroll for (int o = 16; o > 0; o >>= 1) { a0 += __shfl_xor_sync(~0u, a0, o); a1 += __shfl_xor_sync(~0u, a1, o); } if (l == 0) { small[16 + w] = bf2f(f2bf(a0)); small[16 + w + 32] = bf2f(f2bf(a1)); } } __syncthreads(); if (w == 0) { float v0 = small[16 + l], v1 = small[16 + l + 32]; int m0 = l, m1 = l + 32; float gmx = fmaxf(v0, v1); #pragma unroll for (int o = 16; o > 0; o >>= 1) gmx = fmaxf(gmx, __shfl_xor_sync(~0u, gmx, o)); float den = expf(v0 - gmx) + expf(v1 - gmx); #pragma unroll for (int o = 16; o > 0; o >>= 1) den += __shfl_xor_sync(~0u, den, o); for (int i = 0; i < 8; ++i) { float bv = fmaxf(v0, v1); int bi = (v0 >= v1) ? m0 : m1; #pragma unroll for (int o = 16; o > 0; o >>= 1) { float ov = __shfl_xor_sync(~0u, bv, o); int oi = __shfl_xor_sync(~0u, bi, o); if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; } } float p = expf(small[16 + bi] - gmx) / den; if (l == 0) { ismall[80 + i] = bi; small[96 + i] = p; } if (m0 == bi) { v0 = NEG_INF; } if (m1 == bi) { v1 = NEG_INF; } } __syncwarp(); float wsum = small[96 + (l & 7)]; #pragma unroll for (int o = 16; o > 0; o >>= 1) wsum += __shfl_xor_sync(~0u, wsum, o); wsum = wsum * 0.25f; if (l < 8) small[96 + l] = small[96 + l] / (wsum + 1e-9f) * RSCALE; } __syncthreads(); } // gate+up task loop body (shared by K4/M6) __device__ __forceinline__ void moe_gate_up(const long long* SE, unsigned* ctrs, int* ismall, const int* topk8, const float* sx, const float* smx, float* tp0, float* tp1, float* out64, float* small, float* hh) { for (;;) { int t = grab(ctrs, 144, ismall); if (t >= 144) break; int slot_e = t >> 4, tile = t & 15; const unsigned char* wqg; const bf16* scg; const bf16* zrg; const unsigned char* wqu; const bf16* scu; const bf16* zru; if (slot_e < 8) { long long e = topk8[slot_e]; wqg = (const unsigned char*)(const void*)SE[E_GW] + e * (1152 * 1024); scg = (const bf16*)(const void*)SE[E_GS] + e * (18 * 1024); zrg = (const bf16*)(const void*)SE[E_GZ] + e * (18 * 1024); wqu = (const unsigned char*)(const void*)SE[E_UW] + e * (1152 * 1024); scu = (const bf16*)(const void*)SE[E_US] + e * (18 * 1024); zru = (const bf16*)(const void*)SE[E_UZ] + e * (18 * 1024); } else { wqg = (const unsigned char*)(const void*)SE[E_SGW]; scg = (const bf16*)(const void*)SE[E_SGS]; zrg = (const bf16*)(const void*)SE[E_SGZ]; wqu = (const unsigned char*)(const void*)SE[E_SUW]; scu = (const bf16*)(const void*)SE[E_SUS]; zru = (const bf16*)(const void*)SE[E_SUZ]; } gemv64(wqg, scg, zrg, 1024, 0, 18, tile * 64, sx, smx, tp0, out64); if (threadIdx.x < 64) tp1[threadIdx.x] = out64[threadIdx.x]; __syncthreads(); gemv64(wqu, scu, zru, 1024, 0, 18, tile * 64, sx, smx, tp0, out64); if (threadIdx.x < 64) hh[slot_e * MOE_M + tile * 64 + threadIdx.x] = siluf(tp1[threadIdx.x]) * out64[threadIdx.x]; } } // down task loop body (shared by K5/M7) __device__ __forceinline__ void moe_down(const long long* SE, unsigned* ctrs, int* ismall, const int* tki, const float* tkw, const float* hh, float* Dm, float* sx, float* smx, float* tp0, float* out64, float* red) { for (;;) { int t = grab(ctrs, 324, ismall); if (t >= 324) break; int slot_e = t / 36, tile = t % 36; float wgt = (slot_e < 8) ? __ldcg(tkw + slot_e) : 1.f; const unsigned char* wqd; const bf16* scd; const bf16* zrd; if (slot_e < 8) { long long e = __ldcg(tki + slot_e); wqd = (const unsigned char*)(const void*)SE[E_DW] + e * (512 * 2304); scd = (const bf16*)(const void*)SE[E_DS] + e * (8 * 2304); zrd = (const bf16*)(const void*)SE[E_DZ] + e * (8 * 2304); } else { wqd = (const unsigned char*)(const void*)SE[E_SDW]; scd = (const bf16*)(const void*)SE[E_SDS]; zrd = (const bf16*)(const void*)SE[E_SDZ]; } load_sx_f32(hh + slot_e * MOE_M, sx, MOE_M); compute_smx(sx, smx, 8); gemv64(wqd, scd, zrd, 2304, 0, 8, tile * 64, sx, smx, tp0, out64); if (threadIdx.x < 64) Dm[slot_e * DD + tile * 64 + threadIdx.x] = wgt * out64[threadIdx.x]; } } // ------------------------------------------------------------------ kernel extern "C" __global__ void __launch_bounds__(1024, 1) kimi_mega(const long long* __restrict__ SP, const DynArgs dyn) { extern __shared__ char smem_raw[]; float* h1x = (float*)smem_raw; // 2304 float* sx = h1x + 2304; // 2304 float* smx = sx + 2304; // 32 float* tp0 = smx + 32; // 2048 (32 warps x 64 partials) float* tp1 = tp0 + 2048; // 2048 float* out64 = tp1 + 2048; // 128 float* red = out64 + 128; // 64 float* small = red + 64; // 256 int* ismall = (int*)small; unsigned* ctr = (unsigned*)(const void*)SP[SPW + W_CTR]; unsigned* barc = ctr + 30; unsigned* gen = ctr + 31; const int POS = dyn.L; const int Lp1 = dyn.L + 1; const int NCH = dyn.NCH; const int stopa = dyn.stop_after; unsigned target = dyn.gen_base; int slot = 0; const int tid = threadIdx.x; // =================================================================== for (int b = 0; b < 3; ++b) { const long long* SK = SP + SPK(b); const long long* SE = SP + SPE(b); // ---------------- K1 -------------------------------------------- if (slot <= stopa) { const bf16* xin = (const bf16*)(const void*)dyn.x_in; bf16* xbuf = (bf16*)(const void*)SP[SPW + W_BX]; if (b == 0) { for (int i = tid; i < DD; i += blockDim.x) h1x[i] = bf2f(__ldcg(xin + i)); __syncthreads(); } else { const bf16* xp = (b == 1) ? xin : (xbuf + ((b - 1) & 1) * DD); recomb_xb(xp, (const float*)(const void*)SP[SPW + W_A], (const float*)(const void*)SP[SPW + W_D], h1x, tid, blockDim.x); if (blockIdx.x == 0) { bf16* ob = xbuf + (b & 1) * DD; for (int i = tid; i < DD; i += blockDim.x) ob[i] = f2bf(h1x[i]); } } bf16* ws_xn = (blockIdx.x == 0) ? (bf16*)(const void*)SP[SPW + W_XN] : nullptr; rmsn_apply(h1x, (const bf16*)(const void*)SK[K_NORMA], sx, red, ws_xn); compute_smx(sx, smx, 18); unsigned* ctrs = ctr + slot; for (;;) { int t = grab(ctrs, 256, ismall); if (t >= 256) break; int c0 = t * 64; int m = c0 >> 12, n0 = c0 & 4095; const unsigned char* wq; const bf16* sc; const bf16* zr; if (m == 0) { wq = (const unsigned char*)(const void*)SK[K_QW]; sc = (const bf16*)(const void*)SK[K_QS]; zr = (const bf16*)(const void*)SK[K_QZ]; } else if (m == 1) { wq = (const unsigned char*)(const void*)SK[K_KW]; sc = (const bf16*)(const void*)SK[K_KS]; zr = (const bf16*)(const void*)SK[K_KZ]; } else if (m == 2) { wq = (const unsigned char*)(const void*)SK[K_VW]; sc = (const bf16*)(const void*)SK[K_VS]; zr = (const bf16*)(const void*)SK[K_VZ]; } else { wq = (const unsigned char*)(const void*)SK[K_GW]; sc = (const bf16*)(const void*)SK[K_GS]; zr = (const bf16*)(const void*)SK[K_GZ]; } gemv64(wq, sc, zr, 4096, 0, 18, n0, sx, smx, tp0, out64); if (tid < 64) { int cc = n0 + tid; float raw = out64[tid]; if (m < 3) { long long stptr = (m == 0) ? dyn.cq[b] : (m == 1) ? dyn.ck[b] : dyn.cv[b]; bf16* win = (bf16*)(const void*)stptr; bf16 p0 = __ldcg(win + cc); bf16 p1 = __ldcg(win + HC + cc); bf16 p2 = __ldcg(win + 2 * HC + cc); bf16 rawb = f2bf(raw); const bf16* cw = (const bf16*)(const void*)SK[K_CONV] + m * (HC * 4) + cc * 4; float acc = bf2f(p0) * bf2f(__ldg(cw + 0)) + bf2f(p1) * bf2f(__ldg(cw + 1)) + bf2f(p2) * bf2f(__ldg(cw + 2)) + bf2f(rawb) * bf2f(__ldg(cw + 3)); long long wsptr = (m == 0) ? SP[SPW + W_Q] : (m == 1) ? SP[SPW + W_K] : SP[SPW + W_V]; ((bf16*)(const void*)wsptr)[cc] = f2bf(siluf(acc)); win[cc] = p1; win[HC + cc] = p2; win[2 * HC + cc] = rawb; } else { ((float*)(const void*)SP[SPW + W_G])[cc] = -softplusf(bf2f(f2bf(raw))); } } } } gbar(barc, gen, ++target); slot++; // ---------------- K2 -------------------------------------------- if (slot <= stopa) { const bf16* ws_xn = (const bf16*)(const void*)SP[SPW + W_XN]; float* fk = h1x; float* fq = h1x + 128; float* fg = h1x + 256; float* fv = sx; float* prd = sx + 64; float* oscr = tp0; // 1024 floats: tp0+tp1 unsigned* ctrs = ctr + slot; for (;;) { int t = grab(ctrs, 64, ismall); if (t >= 64) break; int h = t >> 1, dv0 = (t & 1) * 64; float part = 0.f; const bf16* bw = (const bf16*)(const void*)SK[K_BETA] + h * DD; for (int i = tid; i < DD; i += blockDim.x) part += bf2f(__ldcg(ws_xn + i)) * bf2f(__ldg(bw + i)); float tot = block_sum(part, red); if (tid == 0) small[104] = 1.f / (1.f + expf(-bf2f(f2bf(tot)))); if (tid < 128) { fk[tid] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_K] + h * HDIM + tid)); fg[tid] = __ldcg((const float*)(const void*)SP[SPW + W_G] + h * HDIM + tid); fq[tid] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_Q] + h * HDIM + tid)) * KDA_SCALE; } else if (tid < 192) { fv[tid - 128] = bf2f(__ldcg((const bf16*)(const void*)SP[SPW + W_V] + h * HDIM + dv0 + tid - 128)); } __syncthreads(); int dv = tid & 63, dkg = tid >> 6; float* Sg = (float*)(const void*)dyn.S[b]; size_t soff = ((size_t)h * HDIM + dkg * 8) * HDIM + dv0 + dv; float sreg[8]; float pp = 0.f; #pragma unroll for (int j = 0; j < 8; ++j) { sreg[j] = __ldcg(Sg + soff + (size_t)j * HDIM); pp += sreg[j] * fk[dkg * 8 + j]; } oscr[dkg * 64 + dv] = pp; __syncthreads(); if (tid < 64) { float a = 0.f; #pragma unroll for (int tt = 0; tt < 16; ++tt) a += oscr[tt * 64 + tid]; prd[tid] = a; } __syncthreads(); float bb = small[104]; float fvd = fv[dv], pd = prd[dv]; float op = 0.f; #pragma unroll for (int j = 0; j < 8; ++j) { float sn = sreg[j] * expf(fg[dkg * 8 + j]) + bb * fk[dkg * 8 + j] * (fvd - pd); sreg[j] = sn; op += sn * fq[dkg * 8 + j]; } #pragma unroll for (int j = 0; j < 8; ++j) *(Sg + soff + (size_t)j * HDIM) = sreg[j]; __syncthreads(); oscr[dkg * 64 + dv] = op; __syncthreads(); if (tid < 64) { float a = 0.f; #pragma unroll for (int tt = 0; tt < 16; ++tt) a += oscr[tt * 64 + tid]; ((bf16*)(const void*)SP[SPW + W_O])[h * HDIM + dv0 + tid] = f2bf(a); } } } gbar(barc, gen, ++target); slot++; // ---------------- K3 -------------------------------------------- if (slot <= stopa) { unsigned* ctrs = ctr + slot; float* A = (float*)(const void*)SP[SPW + W_A]; load_sx_bf16((const bf16*)(const void*)SP[SPW + W_O], sx, 4096); compute_smx(sx, smx, 32); for (;;) { int t = grab(ctrs, 72, ismall); if (t >= 72) break; int kh = t / 36, tile = t % 36; gemv64((const unsigned char*)(const void*)SK[K_OW], (const bf16*)(const void*)SK[K_OS], (const bf16*)(const void*)SK[K_OZ], 2304, kh * 16, 16, tile * 64, sx + kh * 2048, smx + kh * 16, tp0, out64); if (tid < 64) A[kh * DD + tile * 64 + tid] = out64[tid]; } } gbar(barc, gen, ++target); slot++; // ---------------- K4 -------------------------------------------- int topk_i_sh[8]; if (slot <= stopa) { const bf16* xin = (const bf16*)(const void*)dyn.x_in; const bf16* xbuf = (const bf16*)(const void*)SP[SPW + W_BX]; const bf16* xb = (b == 0) ? xin : (xbuf + (b & 1) * DD); const float* A = (const float*)(const void*)SP[SPW + W_A]; for (int i = tid; i < DD; i += blockDim.x) { float att = bf2f(f2bf(__ldcg(A + i) + __ldcg(A + DD + i))); h1x[i] = bf2f(f2bf(bf2f(__ldcg(xb + i)) + att)); } __syncthreads(); rmsn_apply(h1x, (const bf16*)(const void*)SK[K_NORMM], sx, red, nullptr); compute_smx(sx, smx, 18); router_topk((const bf16*)(const void*)SE[E_ROUT], sx, small, ismall, tid); if (blockIdx.x == 0 && tid < 8) { ((int*)(const void*)SP[SPW + W_TKI])[tid] = ismall[80 + tid]; ((float*)(const void*)SP[SPW + W_TKW])[tid] = small[96 + tid]; } #pragma unroll for (int i = 0; i < 8; ++i) topk_i_sh[i] = ismall[80 + i]; moe_gate_up(SE, ctr + slot, ismall, topk_i_sh, sx, smx, tp0, tp1, out64, small, (float*)(const void*)SP[SPW + W_HH]); } gbar(barc, gen, ++target); slot++; // ---------------- K5 -------------------------------------------- if (slot <= stopa) { moe_down(SE, ctr + slot, ismall, (const int*)(const void*)SP[SPW + W_TKI], (const float*)(const void*)SP[SPW + W_TKW], (const float*)(const void*)SP[SPW + W_HH], (float*)(const void*)SP[SPW + W_D], sx, smx, tp0, out64, red); } gbar(barc, gen, ++target); slot++; } // =================================================================== { const int b = 3; const long long* SM = SP + SPM; const long long* SE = SP + SPE(3); // ---------------- M1 -------------------------------------------- if (slot <= stopa) { bf16* xbuf = (bf16*)(const void*)SP[SPW + W_BX]; recomb_xb(xbuf + ((b - 1) & 1) * DD, (const float*)(const void*)SP[SPW + W_A], (const float*)(const void*)SP[SPW + W_D], h1x, tid, blockDim.x); if (blockIdx.x == 0) { bf16* ob = xbuf + (b & 1) * DD; for (int i = tid; i < DD; i += blockDim.x) ob[i] = f2bf(h1x[i]); } __syncthreads(); rmsn_apply(h1x, (const bf16*)(const void*)SM[M_NORMA], sx, red, nullptr); compute_smx(sx, smx, 18); unsigned* ctrs = ctr + slot; for (;;) { int t = grab(ctrs, 105, ismall); if (t >= 105) break; if (t < 96) { gemv64((const unsigned char*)(const void*)SM[M_QW], (const bf16*)(const void*)SM[M_QS], (const bf16*)(const void*)SM[M_QZ], MLB_Q, 0, 18, t * 64, sx, smx, tp0, out64); if (tid < 64) ((bf16*)(const void*)SP[SPW + W_QM])[t * 64 + tid] = f2bf(out64[tid]); } else { int c0 = (t - 96) * 64; gemv64((const unsigned char*)(const void*)SM[M_KVAW], (const bf16*)(const void*)SM[M_KVAS], (const bf16*)(const void*)SM[M_KVAZ], MLB_KVA, 0, 18, c0, sx, smx, tp0, out64); if (tid < 64) { int c = c0 + tid; if (c < KVL) ((bf16*)(const void*)SP[SPW + W_CK])[(size_t)POS * KVL + c] = f2bf(out64[tid]); else ((bf16*)(const void*)SP[SPW + W_KRR])[c - KVL] = f2bf(out64[tid]); } } } } gbar(barc, gen, ++target); slot++; // ---------------- M2 -------------------------------------------- if (slot <= stopa) { const bf16* ws_qm = (const bf16*)(const void*)SP[SPW + W_QM]; bf16* ws_qf = (bf16*)(const void*)SP[SPW + W_QF]; const float* invf = (const float*)(const void*)SP[SPW + W_INV]; const unsigned char* wq = (const unsigned char*)(const void*)SM[M_KVBW]; const bf16* sc = (const bf16*)(const void*)SM[M_KVBS]; const bf16* zr = (const bf16*)(const void*)SM[M_KVBZ]; float* qn_eff = h1x; // 128 float* zred = tp1; // 128 unsigned* ctrs = ctr + slot; int w = tid >> 5, l = tid & 31; for (;;) { int t = grab(ctrs, 128, ismall); if (t >= 128) break; int h = t >> 2, st = t & 3; if (tid < 128) { float qn = bf2f(__ldcg(ws_qm + h * MLB_QD + tid)); float sf = bf2f(__ldg(sc + (size_t)st * MLB_KVB + h * (QK_NOPE + V_HEAD) + tid)); float zf = bf2f(__ldg(zr + (size_t)st * MLB_KVB + h * (QK_NOPE + V_HEAD) + tid)); qn_eff[tid] = qn * sf; zred[tid] = qn * sf * zf; } __syncthreads(); if (tid == 0) { float a = 0.f; for (int i = 0; i < 128; ++i) a += zred[i]; small[108] = a; } __syncthreads(); float zterm = small[108]; for (int it = 0; it < 2; ++it) { int r = st * 64 + w * 2 + it; unsigned u = __ldg((const unsigned*)(wq + (size_t)r * MLB_KVB + h * (QK_NOPE + V_HEAD) + l * 4)); float a_lo = 0.f, a_hi = 0.f; const unsigned char* pb = (const unsigned char*)&u; #pragma unroll for (int j = 0; j < 4; ++j) { a_lo += qn_eff[l * 4 + j] * (float)(pb[j] & 15); a_hi += qn_eff[l * 4 + j] * (float)(pb[j] >> 4); } #pragma unroll for (int o = 16; o > 0; o >>= 1) { a_lo += __shfl_xor_sync(~0u, a_lo, o); a_hi += __shfl_xor_sync(~0u, a_hi, o); } if (l == 0) { ws_qf[h * M3_QP + 2 * r] = f2bf(a_lo - zterm); ws_qf[h * M3_QP + 2 * r + 1] = f2bf(a_hi - zterm); } } __syncthreads(); if (st == 0 && tid < 64) { int i = tid & 31; float ang = (float)POS * __ldcg(invf + i); float cs = cosf(ang), sn = sinf(ang); if (tid < 32) { float x0 = bf2f(__ldcg(ws_qm + h * MLB_QD + QK_NOPE + 2 * i)); float x1 = bf2f(__ldcg(ws_qm + h * MLB_QD + QK_NOPE + 2 * i + 1)); ws_qf[h * M3_QP + 512 + 2 * i] = f2bf(x0 * cs - x1 * sn); ws_qf[h * M3_QP + 512 + 2 * i + 1] = f2bf(x1 * cs + x0 * sn); } else if (h == 0) { const bf16* kr = (const bf16*)(const void*)SP[SPW + W_KRR]; float x0 = bf2f(__ldcg(kr + 2 * i)); float x1 = bf2f(__ldcg(kr + 2 * i + 1)); bf16* kbf = (bf16*)(const void*)SP[SPW + W_KR] + (size_t)POS * QK_ROPE; kbf[2 * i] = f2bf(x0 * cs - x1 * sn); kbf[2 * i + 1] = f2bf(x1 * cs + x0 * sn); } } __syncthreads(); } } gbar(barc, gen, ++target); slot++; // ---------------- M3 (tensor-core flash-decode over latent cache) ---- if (slot <= stopa) { const bf16* ckv = (const bf16*)(const void*)SP[SPW + W_CK]; const bf16* krb = (const bf16*)(const void*)SP[SPW + W_KR]; float* ap = (float*)(const void*)SP[SPW + W_AP]; const bf16* ws_qf = (const bf16*)(const void*)SP[SPW + W_QF]; unsigned* ctrs = ctr + slot; const int rch = (Lp1 + NCH - 1) / NCH; const int w = tid >> 5, l = tid & 31; bf16* stage_s = (bf16*)smem_raw; // [M3_RT][M3_STP] float* sbuf_s = (float*)(smem_raw + M3_RT * M3_STP * 2); // 2 x [32][M3_RT] fp32 __shared__ int stask; __shared__ float sh_rs[32]; for (;;) { __syncthreads(); if (tid == 0) stask = (int)atomicAdd(ctrs, 1u); __syncthreads(); int t = (int)stask; if (t >= NCH) break; int r0 = t * rch, r1 = min(Lp1, r0 + rch); int nrows = r1 - r0; float m = NEG_INF, lsum = 0.f; int mm = w >> 4, ng = w & 15; float accf[4][4]; #pragma unroll for (int i = 0; i < 4; ++i) for (int j = 0; j < 4; ++j) accf[i][j] = 0.f; __syncthreads(); int nb = (nrows + M3_RT - 1) / M3_RT; for (int bb = 0; bb < nb; ++bb) { int rows_b = min(M3_RT, nrows - bb * M3_RT); int rbase = r0 + bb * M3_RT; for (int rr = w; rr < M3_RT; rr += 32) { if (rr < rows_b) { const bf16* crow = ckv + (size_t)(rbase + rr) * 512; const bf16* krow = krb + (size_t)(rbase + rr) * 64; bf16* dst = stage_s + rr * M3_STP; *(uint4*)(dst + l * 16) = __ldcg((const uint4*)crow + l * 2); *(uint4*)(dst + l * 16 + 8) = __ldcg((const uint4*)crow + l * 2 + 1); *(unsigned short*)(dst + 512 + 2 * l) = __ldcg((const unsigned short*)(krow + 2 * l)); *(unsigned short*)(dst + 512 + 2 * l + 1) = __ldcg((const unsigned short*)(krow + 2 * l + 1)); } else { bf16* dst = stage_s + rr * M3_STP; for (int i = l * 8; i < 576; i += 32 * 8) *(uint4*)(dst + i) = make_uint4(0, 0, 0, 0); } } __syncthreads(); // mma1: S[32,RT] = q' @ K^T (24 warps: 2m x 6n x 2 khalf) if (w < 24) { int kg = w / 12, m1 = (w % 12) & 1, nt = (w % 12) >> 1; float c[4] = {0, 0, 0, 0}; const bf16* qa = ws_qf + (m1 * 16) * M3_QP; int rr2 = l >> 2, cc = (l & 3) * 2; int kend = kg ? 36 : 18; for (int ks = kg * 18; ks < kend; ++ks) { unsigned a[4], b[2]; const bf16* ap_ = qa + ks * 16; a[0] = __ldg((const unsigned*)&ap_[rr2 * M3_QP + cc]); a[1] = __ldg((const unsigned*)&ap_[(rr2 + 8) * M3_QP + cc]); a[2] = __ldg((const unsigned*)&ap_[rr2 * M3_QP + cc + 8]); a[3] = __ldg((const unsigned*)&ap_[(rr2 + 8) * M3_QP + cc + 8]); const bf16* bk = stage_s + (nt * 8) * M3_STP + ks * 16; b[0] = *(unsigned*)&bk[rr2 * M3_STP + cc]; b[1] = *(unsigned*)&bk[rr2 * M3_STP + cc + 8]; mma_bf16(c, a, b); } float* sb = sbuf_s + kg * (32 * M3_RT); sb[(m1 * 16 + rr2) * M3_RT + nt * 8 + cc] = c[0]; sb[(m1 * 16 + rr2) * M3_RT + nt * 8 + cc + 1] = c[1]; sb[(m1 * 16 + rr2 + 8) * M3_RT + nt * 8 + cc] = c[2]; sb[(m1 * 16 + rr2 + 8) * M3_RT + nt * 8 + cc + 1] = c[3]; } __syncthreads(); // softmax (warp = head), reads both sbuf layers { int h = w; float vv[2]; vv[0] = sbuf_s[h * M3_RT + l] + sbuf_s[32 * M3_RT + h * M3_RT + l]; vv[1] = (l + 32 < M3_RT) ? (sbuf_s[h * M3_RT + l + 32] + sbuf_s[32 * M3_RT + h * M3_RT + l + 32]) : NEG_INF; vv[0] *= MLA_SCALE; vv[1] *= MLA_SCALE; if (l >= rows_b) vv[0] = NEG_INF; if (l + 32 >= rows_b) vv[1] = NEG_INF; float tmax = fmaxf(vv[0], vv[1]); #pragma unroll for (int o = 16; o > 0; o >>= 1) tmax = fmaxf(tmax, __shfl_xor_sync(~0u, tmax, o)); float mnew = fmaxf(m, tmax); float rs = (m == NEG_INF) ? 0.f : __expf(m - mnew); float p0 = __expf(vv[0] - mnew); float p1 = __expf(vv[1] - mnew); float wsum = p0 + p1; #pragma unroll for (int o = 16; o > 0; o >>= 1) wsum += __shfl_xor_sync(~0u, wsum, o); lsum = lsum * rs + wsum; m = mnew; __syncwarp(); bf16* prow = (bf16*)(sbuf_s + 32 * M3_RT); prow[h * 96 + l] = __float2bfloat16_rn(p0); if (l + 32 < M3_RT) prow[h * 96 + l + 32] = __float2bfloat16_rn(p1); if (l == 0) sh_rs[h] = rs; } __syncthreads(); // acc rescale + mma2: acc[32,512] += P @ V { #pragma unroll for (int i = 0; i < 4; ++i) { #pragma unroll for (int j = 0; j < 4; ++j) accf[i][j] *= sh_rs[mm * 16 + ((j >= 2) ? ((l >> 2) + 8) : (l >> 2))]; } const bf16* pa16 = (const bf16*)(sbuf_s + 32 * M3_RT); int rr2 = l >> 2, cc = (l & 3) * 2; for (int nt2 = 0; nt2 < 4; ++nt2) { for (int ks = 0; ks < 3; ++ks) { unsigned a[4], b[2]; const bf16* ap_ = pa16 + (mm * 16) * 96 + ks * 16; a[0] = *(unsigned*)&ap_[rr2 * 96 + cc]; a[1] = *(unsigned*)&ap_[(rr2 + 8) * 96 + cc]; a[2] = *(unsigned*)&ap_[rr2 * 96 + cc + 8]; a[3] = *(unsigned*)&ap_[(rr2 + 8) * 96 + cc + 8]; unsigned lo = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc) * M3_STP + ng * 32 + nt2 * 8 + rr2]; unsigned hi = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 1) * M3_STP + ng * 32 + nt2 * 8 + rr2]; b[0] = lo | (hi << 16); lo = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 8) * M3_STP + ng * 32 + nt2 * 8 + rr2]; hi = (unsigned)*(const unsigned short*)&stage_s[(ks * 16 + cc + 9) * M3_STP + ng * 32 + nt2 * 8 + rr2]; b[1] = lo | (hi << 16); mma_bf16(accf[nt2], a, b); } } } __syncthreads(); } // partial write if (w < 32 && l == 0) { float* op = ap + ((size_t)t * 32 + w) * 514; op[512] = m; op[513] = lsum; } { int rr2 = l >> 2, cc2 = (l & 3) * 2; float* base = ap + ((size_t)t * 32) * 514; #pragma unroll for (int nt2 = 0; nt2 < 4; ++nt2) { int hrow = mm * 16 + rr2; int col = ng * 32 + nt2 * 8 + cc2; base[hrow * 514 + col] = accf[nt2][0]; base[hrow * 514 + col + 1] = accf[nt2][1]; base[(hrow + 8) * 514 + col] = accf[nt2][2]; base[(hrow + 8) * 514 + col + 1] = accf[nt2][3]; } } } } gbar(barc, gen, ++target); slot++; // ---------------- M4 -------------------------------------------- if (slot <= stopa) { const float* ap = (const float*)(const void*)SP[SPW + W_AP]; unsigned* ctrs = ctr + slot; float* mc = red; // 48 float* wc = tp1; // 48 float* lc = tp1 + 48; // 48 for (;;) { int t = grab(ctrs, 64, ismall); if (t >= 64) break; int h = t >> 1, dv0 = (t & 1) * 64; size_t pbase = (size_t)h * 514; size_t chstride = (size_t)32 * 514; if (tid < 32) { for (int c = tid; c < NCH; c += 32) { mc[c] = __ldcg(ap + (size_t)c * chstride + pbase + 512); lc[c] = __ldcg(ap + (size_t)c * chstride + pbase + 513); } } __syncthreads(); if (tid == 0) { float M = NEG_INF; for (int c = 0; c < NCH; ++c) M = fmaxf(M, mc[c]); float lacc = 0.f; for (int c = 0; c < NCH; ++c) { float wgt = (M == NEG_INF) ? 0.f : expf(mc[c] - M); wc[c] = wgt; lacc += wgt * lc[c]; } small[108] = lacc; } __syncthreads(); float ltot = small[108]; if (tid < 512) { float a = 0.f; for (int c = 0; c < NCH; ++c) a += wc[c] * __ldcg(ap + (size_t)c * chstride + pbase + tid); sx[tid] = a / ltot; } __syncthreads(); compute_smx(sx, smx, 4); gemv64((const unsigned char*)(const void*)SM[M_KVBW], (const bf16*)(const void*)SM[M_KVBS], (const bf16*)(const void*)SM[M_KVBZ], MLB_KVB, 0, 4, h * (QK_NOPE + V_HEAD) + QK_NOPE + dv0, sx, smx, tp0, out64); if (tid < 64) ((bf16*)(const void*)SP[SPW + W_O])[h * V_HEAD + dv0 + tid] = f2bf(out64[tid]); } } gbar(barc, gen, ++target); slot++; // ---------------- M5 -------------------------------------------- if (slot <= stopa) { unsigned* ctrs = ctr + slot; float* A = (float*)(const void*)SP[SPW + W_A]; load_sx_bf16((const bf16*)(const void*)SP[SPW + W_O], sx, 4096); compute_smx(sx, smx, 32); for (;;) { int t = grab(ctrs, 72, ismall); if (t >= 72) break; int kh = t / 36, tile = t % 36; gemv64((const unsigned char*)(const void*)SM[M_OW], (const bf16*)(const void*)SM[M_OS], (const bf16*)(const void*)SM[M_OZ], 2304, kh * 16, 16, tile * 64, sx + kh * 2048, smx + kh * 16, tp0, out64); if (tid < 64) A[kh * DD + tile * 64 + tid] = out64[tid]; } } gbar(barc, gen, ++target); slot++; // ---------------- M6 -------------------------------------------- int topk_i_sh[8]; if (slot <= stopa) { const bf16* xbuf = (const bf16*)(const void*)SP[SPW + W_BX]; const bf16* xb = xbuf + (b & 1) * DD; const float* A = (const float*)(const void*)SP[SPW + W_A]; for (int i = tid; i < DD; i += blockDim.x) { float att = bf2f(f2bf(__ldcg(A + i) + __ldcg(A + DD + i))); h1x[i] = bf2f(f2bf(bf2f(__ldcg(xb + i)) + att)); } __syncthreads(); rmsn_apply(h1x, (const bf16*)(const void*)SM[M_NORMM], sx, red, nullptr); compute_smx(sx, smx, 18); router_topk((const bf16*)(const void*)SE[E_ROUT], sx, small, ismall, tid); if (blockIdx.x == 0 && tid < 8) { ((int*)(const void*)SP[SPW + W_TKI])[tid] = ismall[80 + tid]; ((float*)(const void*)SP[SPW + W_TKW])[tid] = small[96 + tid]; } #pragma unroll for (int i = 0; i < 8; ++i) topk_i_sh[i] = ismall[80 + i]; moe_gate_up(SE, ctr + slot, ismall, topk_i_sh, sx, smx, tp0, tp1, out64, small, (float*)(const void*)SP[SPW + W_HH]); } gbar(barc, gen, ++target); slot++; // ---------------- M7 -------------------------------------------- if (slot <= stopa) { moe_down(SE, ctr + slot, ismall, (const int*)(const void*)SP[SPW + W_TKI], (const float*)(const void*)SP[SPW + W_TKW], (const float*)(const void*)SP[SPW + W_HH], (float*)(const void*)SP[SPW + W_D], sx, smx, tp0, out64, red); } gbar(barc, gen, ++target); slot++; } // =================================================================== // epilogue (no barrier after; resets task counters for the next step) // =================================================================== if (slot <= stopa) { const bf16* xp = (const bf16*)(const void*)SP[SPW + W_BX] + DD; // x_3 const float* A0 = (const float*)(const void*)SP[SPW + W_A]; const float* Dm = (const float*)(const void*)SP[SPW + W_D]; bf16* out = (bf16*)(const void*)dyn.h_out; int cpb = (DD + gridDim.x - 1) / gridDim.x; int c0 = blockIdx.x * cpb, c1 = min(DD, c0 + cpb); for (int i = c0 + tid; i < c1; i += blockDim.x) { float xb_ = bf2f(__ldcg(xp + i)); float att = bf2f(f2bf(__ldcg(A0 + i) + __ldcg(A0 + DD + i))); float h1 = bf2f(f2bf(xb_ + att)); float macc = __ldcg(Dm + 8 * DD + i); #pragma unroll for (int s = 0; s < 8; ++s) macc += __ldcg(Dm + s * DD + i); out[i] = f2bf(h1 + bf2f(f2bf(macc))); } if (blockIdx.x == 0 && tid < 30) ctr[tid] = 0; } else if (blockIdx.x == 0 && tid < 30) { ctr[tid] = 0; } } // ------------------------------------------------------------------ host static bool g_smem_set = false; int64_t kimi_mega_setup(std::vector sp_host) { if (sp_host.size() != SP_TOTAL) { printf("kimi_mega_setup: bad SP size %zu want %d\n", sp_host.size(), SP_TOTAL); return 0; } void* dev = nullptr; cudaError_t e = cudaMalloc(&dev, SP_TOTAL * sizeof(int64_t)); if (e != cudaSuccess) { printf("cudaMalloc: %s\n", cudaGetErrorString(e)); return 0; } e = cudaMemcpy(dev, sp_host.data(), SP_TOTAL * sizeof(int64_t), cudaMemcpyHostToDevice); if (e != cudaSuccess) { printf("cudaMemcpy: %s\n", cudaGetErrorString(e)); return 0; } if (!g_smem_set) { e = cudaFuncSetAttribute((const void*)kimi_mega, cudaFuncAttributeMaxDynamicSharedMemorySize, 99328); if (e != cudaSuccess) { printf("cudaFuncSetAttribute: %s\n", cudaGetErrorString(e)); return 0; } g_smem_set = true; } int coop = 0; cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, 0); if (!coop) { printf("cooperative launch unsupported!\n"); return 0; } int occ = 0; e = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&occ, (const void*)kimi_mega, 1024, 99328); if (e != cudaSuccess || occ < 1) { printf("occupancy: %d %s\n", occ, cudaGetErrorString(e)); return 0; } return (int64_t)dev; } int64_t kimi_mega_step(int64_t sp, std::vector dyn, int64_t L, int64_t gen_base, int64_t NCH, int64_t stop_after) { if (dyn.size() != 14) { printf("bad dyn size\n"); return 1; } DynArgs a; a.x_in = dyn[0]; a.h_out = dyn[1]; for (int i = 0; i < 3; ++i) { a.S[i] = dyn[2 + i]; a.cq[i] = dyn[5 + i]; a.ck[i] = dyn[8 + i]; a.cv[i] = dyn[11 + i]; } a.L = (int)L; a.NCH = (int)NCH; a.stop_after = (int)stop_after; a.gen_base = (unsigned)gen_base; int nsm = 188; cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, 0); const long long* SPp = (const long long*)sp; void* params[2] = { (void*)&SPp, (void*)&a }; cudaError_t e = cudaLaunchCooperativeKernel((const void*)kimi_mega, dim3(nsm), dim3(1024), params, 99328, 0); if (e != cudaSuccess) { printf("launch: %s\n", cudaGetErrorString(e)); return 2; } return 0; } """ _EXT = None SP_TOTAL = 167 def _get_ext(): global _EXT if _EXT is None: from torch.utils.cpp_extension import load_inline import shutil here = os.path.dirname(os.path.abspath(__file__)) os.environ.setdefault("TORCH_EXTENSIONS_DIR", os.path.join(here, ".torch_extensions")) flags = [ "-O3", "-std=c++17", "--generate-code=arch=compute_120,code=sm_120", ] gpp = shutil.which("g++") if gpp: flags.append(f"--compiler-bindir={gpp}") _EXT = load_inline( name="kimi_mega_v1", cpp_sources=[CPP_DECL], cuda_sources=[CUDA_SRC], functions=["kimi_mega_setup", "kimi_mega_step"], extra_cuda_cflags=flags, verbose=False, ) return _EXT class MegaRunner: RINGL = 32768 def __init__(self, model: nn.Module, dev: torch.device): cfg = model.cfg assert cfg.pattern[3] == "M" and all(k == "K" for k in cfg.pattern[:3]), \ "megakernel fast path assumes pattern K,K,K,M" self.model = model self.dev = dev d = cfg.hidden bf = torch.bfloat16 t = lambda *shape, dtype: torch.zeros(*shape, dtype=dtype, device=dev) C = cfg.kda_heads * cfg.kda_head_dim ws = {} ws["ws_q"] = t(C, dtype=bf) ws["ws_k"] = t(C, dtype=bf) ws["ws_v"] = t(C, dtype=bf) ws["ws_g"] = t(C, dtype=torch.float32) ws["ws_o"] = t(C, dtype=bf) ws["ws_xn"] = t(d, dtype=bf) ws["buf_x"] = t(2, d, dtype=bf) ws["A"] = t(2, d, dtype=torch.float32) ws["D"] = t(9, d, dtype=torch.float32) ws["hh"] = t(9, cfg.moe_inter, dtype=torch.float32) ws["topk_i"] = t(8, dtype=torch.int32) ws["topk_w"] = t(8, dtype=torch.float32) ws["ws_qm"] = t(cfg.mla_heads * (cfg.qk_nope + cfg.qk_rope), dtype=bf) ws["ws_qf"] = t(cfg.mla_heads, 584, dtype=bf) ws["kr_raw"] = t(cfg.qk_rope, dtype=bf) ws["attn_part"] = t(128, 32, 514, dtype=torch.float32) ws["c_buf"] = t(self.RINGL, cfg.kv_lora, dtype=bf) ws["k_buf"] = t(self.RINGL, cfg.qk_rope, dtype=bf) ws["ctr"] = t(32, dtype=torch.int32) invf = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.qk_rope, 2, dtype=torch.float32) / cfg.qk_rope)) ws["invf"] = invf.to(dev) ws["h_out"] = t(d, dtype=bf) self.ws = ws self.h_out = ws["h_out"] self.c_buf = ws["c_buf"] self.k_buf = ws["k_buf"] sp = [] def add(x): sp.append(x.data_ptr()) for b in range(3): blk = model.blocks[b]; at = blk.attn add(blk.attn_norm); add(blk.moe_norm) for pr in (at.q_proj, at.k_proj, at.v_proj, at.g_proj, at.o_proj): add(pr.w_q); add(pr.scales); add(pr.zeros) add(at.conv_w); add(at.beta_proj.weight) for b in range(4): moe = model.blocks[b].moe add(moe.router.weight) for qe in (moe.gate, moe.up, moe.down, moe.s_gate, moe.s_up, moe.s_down): add(qe.w_q); add(qe.scales); add(qe.zeros) mla = model.blocks[3]; at = mla.attn add(mla.attn_norm); add(mla.moe_norm) for pr in (at.q_proj, at.kv_a, at.kv_b, at.o_proj): add(pr.w_q); add(pr.scales); add(pr.zeros) add(ws["ws_q"]); add(ws["ws_k"]); add(ws["ws_v"]); add(ws["ws_g"]) add(ws["ws_o"]); add(ws["ws_xn"]); add(ws["buf_x"]); add(ws["A"]) add(ws["D"]); add(ws["hh"]); add(ws["topk_i"]); add(ws["topk_w"]) add(ws["ws_qm"]); add(ws["ws_qf"]); add(ws["kr_raw"]) add(ws["attn_part"]); add(ws["c_buf"]); add(ws["k_buf"]) add(ws["ctr"]); add(ws["invf"]) assert len(sp) == SP_TOTAL, len(sp) self.ext = _get_ext() self.sp = self.ext.kimi_mega_setup(sp) assert self.sp != 0, "megakernel setup failed" self.gen_base = 0 self.mla_idx = 3 self.stop_after = int(os.environ.get("KIMI_STOP_AFTER", "999")) def step(self, hidden, state): st = state[self.mla_idx] c_kv = st["c_kv"] L = c_kv.shape[0] l1 = L + 1 if c_kv.data_ptr() != self.c_buf.data_ptr(): assert l1 <= self.RINGL self.c_buf[:L].copy_(c_kv, non_blocking=True) self.k_buf[:L].copy_(st["k_rope"], non_blocking=True) st["c_kv"] = self.c_buf[:l1] st["k_rope"] = self.k_buf[:l1] dyn = [hidden.data_ptr(), self.h_out.data_ptr()] for kd in ("S", "cq", "ck", "cv"): dyn += [state[b][kd].data_ptr() for b in range(3)] NCH = min(128, max(1, (l1 + 111) // 112)) rc = self.ext.kimi_mega_step(self.sp, dyn, L, self.gen_base, NCH, self.stop_after) self.gen_base += 22 assert rc == 0, f"megakernel step failed rc={rc}" return self.h_out, state