KernelBench mega · RTX PRO 6000
Kimi-Linear Decode Gemini 4 Argon (High)
27.03×geomean speedup across shapes
manually audited: clean
Gemini 4 Argon wrote one hand-written CUDA kernel that runs the whole four-layer decode step in a single launch: 376 CTAs (two per SM) synchronised by a software grid barrier, int4 GEMVs that fold the zero point into one subtraction per group, KDA short-conv and delta-rule state updates, MLA attention in the 512-dim latent space with W_uk absorbed into the query, a 188-way split-K flash-decoding pass, and router top-8 plus shared expert MoE in every layer. 27.03x over baseline.
harnessagy
Kernel source (redacted)
"""Kimi-Linear W4A16 Hybrid Decode Megakernel for SM120 Blackwell."""
from __future__ import annotations
import os
from dataclasses import dataclass, field
import ninja
import torch
import torch.nn as nn
os.environ["PATH"] = ninja.BIN_DIR + os.pathsep + os.environ.get("PATH", "")
from torch.utils.cpp_extension import load_inline
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)))
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.zeros(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)
# --------------------------------------------------------------------------- #
# Single Fused CUDA Megakernel across 188 SMs (376 Persistent Co-Resident Blocks)
# --------------------------------------------------------------------------- #
_CUDA_SRC = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <stdint.h>
#include <math.h>
constexpr int HIDDEN = 2304;
constexpr int KDA_H = 32;
constexpr int KDA_DK = 128;
constexpr int KDA_C = 4096;
constexpr int MLA_H = 32;
constexpr int KV_LORA = 512;
constexpr int QK_NOPE = 128;
constexpr int QK_ROPE = 64;
constexpr int V_HEAD = 128;
constexpr int Q_DIM = 192;
constexpr int KVB_DIM = 256;
constexpr int MOE_E = 64;
constexpr int MOE_ACTIVE = 8;
constexpr int MOE_TOTAL_ACTIVE = 9;
constexpr int MOE_INTER = 1024;
constexpr int GROUP = 128;
constexpr int NUM_SMS = 376;
constexpr int BLOCK_THREADS = 256;
constexpr int NUM_L_SPLITS = 188;
__device__ uint32_t g_sync_arr[2] = {0u, 0u};
__device__ __forceinline__ void grid_barrier(uint32_t& local_epoch) {
__syncthreads();
if (threadIdx.x == 0) {
uint32_t next_epoch = local_epoch + 1u;
asm volatile("membar.gl;" ::: "memory");
uint32_t arrived = atomicAdd(&g_sync_arr[0], 1u);
if (arrived == NUM_SMS - 1) {
*reinterpret_cast<volatile uint32_t*>(&g_sync_arr[0]) = 0u;
asm volatile("membar.gl;" ::: "memory");
*reinterpret_cast<volatile uint32_t*>(&g_sync_arr[1]) = next_epoch;
} else {
while (*reinterpret_cast<volatile uint32_t*>(&g_sync_arr[1]) != next_epoch) {
__nanosleep(40);
}
asm volatile("membar.gl;" ::: "memory");
}
local_epoch = next_epoch;
}
__syncthreads();
}
struct QLinearPtrs {
const uint8_t* w_q;
const __nv_bfloat16* scales;
const __nv_bfloat16* zeros;
};
struct MoEPtrs {
const __nv_bfloat16* router_w;
QLinearPtrs gate;
QLinearPtrs up;
QLinearPtrs down;
QLinearPtrs s_gate;
QLinearPtrs s_up;
QLinearPtrs s_down;
};
struct KDABlockPtrs {
const __nv_bfloat16* attn_norm;
const __nv_bfloat16* moe_norm;
QLinearPtrs q_proj;
QLinearPtrs k_proj;
QLinearPtrs v_proj;
QLinearPtrs g_proj;
const __nv_bfloat16* beta_w;
const __nv_bfloat16* conv_w;
QLinearPtrs o_proj;
MoEPtrs moe;
float* S;
__nv_bfloat16* cq;
__nv_bfloat16* ck;
__nv_bfloat16* cv;
};
struct MLABlockPtrs {
const __nv_bfloat16* attn_norm;
const __nv_bfloat16* moe_norm;
QLinearPtrs q_proj;
QLinearPtrs kv_a;
QLinearPtrs kv_b;
QLinearPtrs o_proj;
MoEPtrs moe;
const __nv_bfloat16* c_kv_src;
const __nv_bfloat16* k_rope_src;
__nv_bfloat16* c_kv_dst;
__nv_bfloat16* k_rope_dst;
int pos;
};
struct ScratchPtrs {
float* norm_x; // [2304]
float* attn_out_fp32; // [2304]
float* moe_out_fp32; // [2304]
__nv_bfloat16* res_h; // [2304]
float* kda_qkvg_fp32; // [4 * 4096]
float* kda_o; // [4096]
float* kda_beta; // [32]
float* mla_q_kv_fp32; // [6144 + 576]
float* mla_q_abs; // [32 * 512]
float* mla_q_rope; // [32 * 64]
float* mla_att_m; // [188 * 32]
float* mla_att_d; // [188 * 32]
float* mla_att_o; // [188 * 32 * 512]
float* router_logits; // [64]
float* moe_gate_fp32; // [9 * 1024]
float* moe_up_fp32; // [9 * 1024]
};
struct KernelParams {
KDABlockPtrs kda[3];
MLABlockPtrs mla;
ScratchPtrs scratch;
const __nv_bfloat16* hidden_in;
__nv_bfloat16* hidden_out;
};
__device__ __forceinline__ float bf16_to_f32(__nv_bfloat16 v) {
return __bfloat162float(v);
}
__device__ __forceinline__ __nv_bfloat16 f32_to_bf16(float v) {
return __float2bfloat16(v);
}
__device__ __forceinline__ float round_bf16(float v) {
return __bfloat162float(__float2bfloat16(v));
}
__device__ __forceinline__ float warp_reduce_sum(float val) {
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1) {
val += __shfl_xor_sync(0xffffffff, val, offset);
}
return val;
}
__device__ __forceinline__ void accum_u32_fast(uint32_t p, float2 xv, uint32_t magic, float (&g_acc)[4]) {
uint32_t p_ev = p & 0x0F0F0F0Fu;
uint32_t p_od = (p >> 4) & 0x0F0F0F0Fu;
g_acc[0] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7044)), g_acc[0]);
g_acc[1] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7144)), g_acc[1]);
g_acc[2] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7244)), g_acc[2]);
g_acc[3] = fmaf(xv.x, __uint_as_float(__byte_perm(p_ev, magic, 0x7344)), g_acc[3]);
g_acc[0] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7044)), g_acc[0]);
g_acc[1] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7144)), g_acc[1]);
g_acc[2] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7244)), g_acc[2]);
g_acc[3] = fmaf(xv.y, __uint_as_float(__byte_perm(p_od, magic, 0x7344)), g_acc[3]);
}
// Full group (8 kp-rows per warp, 64 kp-rows across 8 warps, 128 contiguous output columns per warp)
__device__ __forceinline__ void w4a16_group_tile_128(
const uint8_t* __restrict__ w_q,
const __nv_bfloat16* __restrict__ scales,
const __nv_bfloat16* __restrict__ zeros,
const float* __restrict__ smem_x,
int N,
int n0,
int g,
int warp_id,
int lane,
float (&acc)[4]
) {
const int kp_base = g * 64 + warp_id * 8;
const int col = n0 + lane * 4;
const uint8_t* w_ptr = w_q + col;
const float2* x_f2 = reinterpret_cast<const float2*>(smem_x + kp_base * 2);
float g_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float x_sum = 0.0f;
#pragma unroll
for (int r = 0; r < 8; ++r) {
uint32_t p = *reinterpret_cast<const uint32_t*>(w_ptr + (kp_base + r) * N);
float2 xv = x_f2[r];
x_sum += xv.x + xv.y;
uint32_t magic = 0x43000000u + (threadIdx.x >> 20);
accum_u32_fast(p, xv, magic, g_acc);
}
const int sz_idx = g * N + col;
uint2 s_u2 = *reinterpret_cast<const uint2*>(scales + sz_idx);
uint2 z_u2 = *reinterpret_cast<const uint2*>(zeros + sz_idx);
const __nv_bfloat16* s_bf = reinterpret_cast<const __nv_bfloat16*>(&s_u2);
const __nv_bfloat16* z_bf = reinterpret_cast<const __nv_bfloat16*>(&z_u2);
#pragma unroll
for (int c = 0; c < 4; ++c) {
float sv = bf16_to_f32(s_bf[c]);
float zv = bf16_to_f32(z_bf[c]) + 128.0f;
acc[c] += sv * (g_acc[c] - zv * x_sum);
}
}
// Fused Gate+Up full-group tile (reads x_f2 and computes x_sum ONCE for both Gate and Up!)
__device__ __forceinline__ void w4a16_gate_up_group_tile_128(
const uint8_t* __restrict__ gw_q,
const __nv_bfloat16* __restrict__ g_scales,
const __nv_bfloat16* __restrict__ g_zeros,
const uint8_t* __restrict__ uw_q,
const __nv_bfloat16* __restrict__ u_scales,
const __nv_bfloat16* __restrict__ u_zeros,
const float* __restrict__ smem_x,
int N,
int n0,
int g,
int warp_id,
int lane,
float (&acc_g)[4],
float (&acc_u)[4]
) {
const int kp_base = g * 64 + warp_id * 8;
const int col = n0 + lane * 4;
const uint8_t* gw_ptr = gw_q + col;
const uint8_t* uw_ptr = uw_q + col;
const float2* x_f2 = reinterpret_cast<const float2*>(smem_x + kp_base * 2);
float gg_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float ug_acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float x_sum = 0.0f;
#pragma unroll
for (int r = 0; r < 8; ++r) {
int row_off = (kp_base + r) * N;
uint32_t gp = *reinterpret_cast<const uint32_t*>(gw_ptr + row_off);
uint32_t up = *reinterpret_cast<const uint32_t*>(uw_ptr + row_off);
float2 xv = x_f2[r];
x_sum += xv.x + xv.y;
uint32_t magic = 0x43000000u + (threadIdx.x >> 20);
accum_u32_fast(gp, xv, magic, gg_acc);
accum_u32_fast(up, xv, magic, ug_acc);
}
const int sz_idx = g * N + col;
uint2 gs_u2 = *reinterpret_cast<const uint2*>(g_scales + sz_idx);
uint2 gz_u2 = *reinterpret_cast<const uint2*>(g_zeros + sz_idx);
uint2 us_u2 = *reinterpret_cast<const uint2*>(u_scales + sz_idx);
uint2 uz_u2 = *reinterpret_cast<const uint2*>(u_zeros + sz_idx);
const __nv_bfloat16* gs_bf = reinterpret_cast<const __nv_bfloat16*>(&gs_u2);
const __nv_bfloat16* gz_bf = reinterpret_cast<const __nv_bfloat16*>(&gz_u2);
const __nv_bfloat16* us_bf = reinterpret_cast<const __nv_bfloat16*>(&us_u2);
const __nv_bfloat16* uz_bf = reinterpret_cast<const __nv_bfloat16*>(&uz_u2);
#pragma unroll
for (int c = 0; c < 4; ++c) {
acc_g[c] += bf16_to_f32(gs_bf[c]) * (gg_acc[c] - (bf16_to_f32(gz_bf[c]) + 128.0f) * x_sum);
acc_u[c] += bf16_to_f32(us_bf[c]) * (ug_acc[c] - (bf16_to_f32(uz_bf[c]) + 128.0f) * x_sum);
}
}
__device__ __forceinline__ void save_tile_128(
float (&acc)[4],
float* smem_red_prev, // [8 * 128]
int warp_id,
int lane
) {
*reinterpret_cast<float4*>(&smem_red_prev[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]);
acc[0] = 0.0f; acc[1] = 0.0f; acc[2] = 0.0f; acc[3] = 0.0f;
}
__device__ __forceinline__ void flush_dual_tile_128(
float* dst_cur,
float* dst_prev,
float (&acc)[4],
float* smem_red_cur, // [8 * 128]
float* smem_red_prev, // [8 * 128]
int tid,
int warp_id,
int lane,
float scale_cur = 1.0f,
float scale_prev = 1.0f
) {
*reinterpret_cast<float4*>(&smem_red_cur[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]);
__syncthreads();
if (tid < 128) {
if (dst_cur != nullptr) {
float sum = smem_red_cur[0 * 128 + tid] + smem_red_cur[1 * 128 + tid]
+ smem_red_cur[2 * 128 + tid] + smem_red_cur[3 * 128 + tid]
+ smem_red_cur[4 * 128 + tid] + smem_red_cur[5 * 128 + tid]
+ smem_red_cur[6 * 128 + tid] + smem_red_cur[7 * 128 + tid];
atomicAdd(&dst_cur[tid], sum * scale_cur);
}
} else if (dst_prev != nullptr) {
int c = tid - 128;
float sum = smem_red_prev[0 * 128 + c] + smem_red_prev[1 * 128 + c]
+ smem_red_prev[2 * 128 + c] + smem_red_prev[3 * 128 + c]
+ smem_red_prev[4 * 128 + c] + smem_red_prev[5 * 128 + c]
+ smem_red_prev[6 * 128 + c] + smem_red_prev[7 * 128 + c];
atomicAdd(&dst_prev[c], sum * scale_prev);
}
}
__device__ __forceinline__ void save_tile_128_pair(
float (&acc_g)[4],
float (&acc_u)[4],
float* smem_red_g_prev,
float* smem_red_u_prev,
int warp_id,
int lane
) {
*reinterpret_cast<float4*>(&smem_red_g_prev[warp_id * 128 + lane * 4]) = make_float4(acc_g[0], acc_g[1], acc_g[2], acc_g[3]);
*reinterpret_cast<float4*>(&smem_red_u_prev[warp_id * 128 + lane * 4]) = make_float4(acc_u[0], acc_u[1], acc_u[2], acc_u[3]);
acc_g[0] = 0.0f; acc_g[1] = 0.0f; acc_g[2] = 0.0f; acc_g[3] = 0.0f;
acc_u[0] = 0.0f; acc_u[1] = 0.0f; acc_u[2] = 0.0f; acc_u[3] = 0.0f;
}
__device__ __forceinline__ void flush_dual_tile_128_pair(
float* dst_g_cur,
float* dst_u_cur,
float* dst_g_prev,
float* dst_u_prev,
float (&acc_g)[4],
float (&acc_u)[4],
float* smem_red_g_cur,
float* smem_red_u_cur,
float* smem_red_g_prev,
float* smem_red_u_prev,
int tid,
int warp_id,
int lane
) {
*reinterpret_cast<float4*>(&smem_red_g_cur[warp_id * 128 + lane * 4]) = make_float4(acc_g[0], acc_g[1], acc_g[2], acc_g[3]);
*reinterpret_cast<float4*>(&smem_red_u_cur[warp_id * 128 + lane * 4]) = make_float4(acc_u[0], acc_u[1], acc_u[2], acc_u[3]);
__syncthreads();
if (tid < 128) {
if (dst_g_cur != nullptr) {
float sum_g = smem_red_g_cur[0 * 128 + tid] + smem_red_g_cur[1 * 128 + tid]
+ smem_red_g_cur[2 * 128 + tid] + smem_red_g_cur[3 * 128 + tid]
+ smem_red_g_cur[4 * 128 + tid] + smem_red_g_cur[5 * 128 + tid]
+ smem_red_g_cur[6 * 128 + tid] + smem_red_g_cur[7 * 128 + tid];
atomicAdd(&dst_g_cur[tid], sum_g);
}
if (dst_g_prev != nullptr) {
float sum_gp = smem_red_g_prev[0 * 128 + tid] + smem_red_g_prev[1 * 128 + tid]
+ smem_red_g_prev[2 * 128 + tid] + smem_red_g_prev[3 * 128 + tid]
+ smem_red_g_prev[4 * 128 + tid] + smem_red_g_prev[5 * 128 + tid]
+ smem_red_g_prev[6 * 128 + tid] + smem_red_g_prev[7 * 128 + tid];
atomicAdd(&dst_g_prev[tid], sum_gp);
}
} else {
int c = tid - 128;
if (dst_u_cur != nullptr) {
float sum_u = smem_red_u_cur[0 * 128 + c] + smem_red_u_cur[1 * 128 + c]
+ smem_red_u_cur[2 * 128 + c] + smem_red_u_cur[3 * 128 + c]
+ smem_red_u_cur[4 * 128 + c] + smem_red_u_cur[5 * 128 + c]
+ smem_red_u_cur[6 * 128 + c] + smem_red_u_cur[7 * 128 + c];
atomicAdd(&dst_u_cur[c], sum_u);
}
if (dst_u_prev != nullptr) {
float sum_up = smem_red_u_prev[0 * 128 + c] + smem_red_u_prev[1 * 128 + c]
+ smem_red_u_prev[2 * 128 + c] + smem_red_u_prev[3 * 128 + c]
+ smem_red_u_prev[4 * 128 + c] + smem_red_u_prev[5 * 128 + c]
+ smem_red_u_prev[6 * 128 + c] + smem_red_u_prev[7 * 128 + c];
atomicAdd(&dst_u_prev[c], sum_up);
}
}
}
__device__ __forceinline__ void flush_tile_128(
float* dst_fp32,
float (&acc)[4],
float* smem_red, // [8 * 128]
int tid,
int warp_id,
int lane,
float scale_factor = 1.0f
) {
*reinterpret_cast<float4*>(&smem_red[warp_id * 128 + lane * 4]) = make_float4(acc[0], acc[1], acc[2], acc[3]);
__syncthreads();
if (tid < 128) {
float sum = smem_red[0 * 128 + tid] + smem_red[1 * 128 + tid]
+ smem_red[2 * 128 + tid] + smem_red[3 * 128 + tid]
+ smem_red[4 * 128 + tid] + smem_red[5 * 128 + tid]
+ smem_red[6 * 128 + tid] + smem_red[7 * 128 + tid];
atomicAdd(&dst_fp32[tid], sum * scale_factor);
}
__syncthreads();
acc[0] = 0.0f; acc[1] = 0.0f; acc[2] = 0.0f; acc[3] = 0.0f;
}
__device__ __forceinline__ void compute_moe_routing_smem(
const float* __restrict__ router_logits,
int* smem_top_idx,
float* smem_top_w,
int tid
) {
if (tid < 32) {
float l0 = router_logits[tid];
float l1 = router_logits[tid + 32];
float m = fmaxf(l0, l1);
#pragma unroll
for (int off = 16; off > 0; off >>= 1) {
m = fmaxf(m, __shfl_xor_sync(0xffffffff, m, off));
}
float p0 = expf(l0 - m);
float p1 = expf(l1 - m);
float s = p0 + p1;
s = warp_reduce_sum(s);
p0 /= s;
p1 /= s;
float w_sum = 0.0f;
#pragma unroll
for (int j = 0; j < MOE_ACTIVE; ++j) {
float best_v = (p0 >= p1) ? p0 : p1;
int best_i = (p0 >= p1) ? tid : (tid + 32);
#pragma unroll
for (int off = 16; off > 0; off >>= 1) {
float other_v = __shfl_xor_sync(0xffffffff, best_v, off);
int other_i = __shfl_xor_sync(0xffffffff, best_i, off);
if (other_v > best_v || (other_v == best_v && other_i < best_i)) {
best_v = other_v;
best_i = other_i;
}
}
if (tid == 0) {
smem_top_idx[j] = best_i;
smem_top_w[j] = best_v;
}
w_sum += best_v;
if (best_i == tid) p0 = -1.0f;
if (best_i == tid + 32) p1 = -1.0f;
}
if (tid < MOE_ACTIVE) {
smem_top_w[tid] = smem_top_w[tid] / (w_sum + 1.0e-9f) * 2.446f;
}
if (tid == 0) {
smem_top_idx[MOE_ACTIVE] = 0;
smem_top_w[MOE_ACTIVE] = 1.0f;
}
}
}
__device__ void run_moe_phases(
uint32_t& local_epoch,
const __nv_bfloat16* __restrict__ moe_norm,
const MoEPtrs& moe,
const ScratchPtrs& sc,
float* smem_x, // [6144] floats
float* smem_red, // [1024] floats
int* smem_top_idx,
float* smem_top_w
) {
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int warp_id = tid >> 5;
const int lane = tid & 31;
// ---------------- Phase E1: Attn Residual + MoE RMSNorm + Router + Zeroing ----------------
if (bid < 64) {
const __nv_bfloat16* __restrict__ res_in = sc.res_h + HIDDEN;
const float* __restrict__ a_in = sc.attn_out_fp32;
const __nv_bfloat16* __restrict__ rw_row = moe.router_w + bid * HIDDEN;
float hv[9];
__nv_bfloat16 h_bf[9];
float m_norm[9];
float rw_val[9];
float local_sq = 0.0f;
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
float h_old = bf16_to_f32(res_in[i]);
float a_out = round_bf16(a_in[i]);
h_bf[m] = f32_to_bf16(h_old + a_out);
m_norm[m] = bf16_to_f32(moe_norm[i]);
rw_val[m] = bf16_to_f32(rw_row[i]);
hv[m] = bf16_to_f32(h_bf[m]);
local_sq += hv[m] * hv[m];
}
if (bid == 0) {
#pragma unroll
for (int m = 0; m < 9; ++m) {
sc.res_h[tid + m * BLOCK_THREADS] = h_bf[m];
}
}
local_sq = warp_reduce_sum(local_sq);
if (lane == 0) smem_red[warp_id] = local_sq;
__syncthreads();
if (tid < 32) {
float b_sq = (tid < 8) ? smem_red[tid] : 0.0f;
b_sq = warp_reduce_sum(b_sq);
if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f);
}
__syncthreads();
float inv_rms = smem_red[0];
float r_dot = 0.0f;
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
float xn = round_bf16(hv[m] * inv_rms * m_norm[m]);
smem_x[i] = xn;
if (bid == 0) {
sc.norm_x[i] = xn;
}
r_dot += xn * rw_val[m];
}
r_dot = warp_reduce_sum(r_dot);
if (lane == 0) smem_red[warp_id] = r_dot;
__syncthreads();
if (tid < 32) {
float sum = (tid < 8) ? smem_red[tid] : 0.0f;
sum = warp_reduce_sum(sum);
if (tid == 0) sc.router_logits[bid] = round_bf16(sum);
}
} else {
const int zero_tid = (bid - 64) * BLOCK_THREADS + tid;
const int zero_stride = (NUM_SMS - 64) * BLOCK_THREADS;
for (int i = zero_tid; i < MOE_TOTAL_ACTIVE * MOE_INTER; i += zero_stride) {
sc.moe_gate_fp32[i] = 0.0f;
sc.moe_up_fp32[i] = 0.0f;
}
for (int i = zero_tid; i < HIDDEN; i += zero_stride) {
sc.moe_out_fp32[i] = 0.0f;
}
}
grid_barrier(local_epoch);
// ---------------- Phase E2: Top-8 Selection + Fused Gate & Up W4A16 GEMV ----------------
compute_moe_routing_smem(sc.router_logits, smem_top_idx, smem_top_w, tid);
if (bid >= 64) {
for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) {
smem_x[i] = sc.norm_x[i];
}
}
__syncthreads();
float* smem_red_u = smem_x + 2304;
float* smem_red_g_prev = smem_x + 3328;
float* smem_red_u_prev = smem_x + 4352;
const int eff_bid = (bid < 188) ? (bid << 1) : (((bid - 188) << 1) + 1);
// 9 active experts * 8 col_tiles (1024/128) * 18 full_groups (2304/128) = 1296 units
constexpr int TOTAL_GU_GROUPS = MOE_TOTAL_ACTIVE * 8 * 18;
const int gu_start = (eff_bid * TOTAL_GU_GROUPS) / NUM_SMS;
const int gu_end = ((eff_bid + 1) * TOTAL_GU_GROUPS) / NUM_SMS;
float acc_g[4] = {0.0f, 0.0f, 0.0f, 0.0f};
float acc_u[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_sc = -1;
int prev_sc = -1;
for (int u = gu_start; u < gu_end; ++u) {
int sc_idx = u / 18;
int g = u % 18;
if (sc_idx != cur_sc) {
if (cur_sc >= 0) {
save_tile_128_pair(acc_g, acc_u, smem_red_g_prev, smem_red_u_prev, warp_id, lane);
prev_sc = cur_sc;
}
cur_sc = sc_idx;
}
int s = sc_idx >> 3;
int c_tile = sc_idx & 7;
int e = smem_top_idx[s];
const QLinearPtrs& g_ptrs = (s < MOE_ACTIVE) ? moe.gate : moe.s_gate;
const QLinearPtrs& u_ptrs = (s < MOE_ACTIVE) ? moe.up : moe.s_up;
const size_t wq_off = size_t(e) * (HIDDEN / 2) * MOE_INTER;
const size_t sz_off = size_t(e) * (HIDDEN / GROUP) * MOE_INTER;
w4a16_gate_up_group_tile_128(
g_ptrs.w_q + wq_off, g_ptrs.scales + sz_off, g_ptrs.zeros + sz_off,
u_ptrs.w_q + wq_off, u_ptrs.scales + sz_off, u_ptrs.zeros + sz_off,
smem_x, MOE_INTER, c_tile * 128, g, warp_id, lane, acc_g, acc_u
);
}
flush_dual_tile_128_pair(
(cur_sc >= 0) ? (sc.moe_gate_fp32 + cur_sc * 128) : nullptr,
(cur_sc >= 0) ? (sc.moe_up_fp32 + cur_sc * 128) : nullptr,
(prev_sc >= 0) ? (sc.moe_gate_fp32 + prev_sc * 128) : nullptr,
(prev_sc >= 0) ? (sc.moe_up_fp32 + prev_sc * 128) : nullptr,
acc_g, acc_u, smem_red, smem_red_u, smem_red_g_prev, smem_red_u_prev, tid, warp_id, lane
);
grid_barrier(local_epoch);
// ---------------- Phase E3: MoE Down W4A16 GEMV (Per-Expert Full-Group Mapped, 41 blocks/expert) ----------------
if (eff_bid < 369) {
const int s_my = eff_bid / 41;
const int b_in_s = eff_bid % 41;
const int u_start = (b_in_s * 144) / 41;
const int u_end = ((b_in_s + 1) * 144) / 41;
const float4 g4 = reinterpret_cast<const float4*>(sc.moe_gate_fp32 + s_my * MOE_INTER)[tid];
const float4 u4 = reinterpret_cast<const float4*>(sc.moe_up_fp32 + s_my * MOE_INTER)[tid];
reinterpret_cast<float4*>(smem_x)[tid] = make_float4(
(g4.x / (1.0f + expf(-g4.x))) * u4.x,
(g4.y / (1.0f + expf(-g4.y))) * u4.y,
(g4.z / (1.0f + expf(-g4.z))) * u4.z,
(g4.w / (1.0f + expf(-g4.w))) * u4.w
);
__syncthreads();
const int e = smem_top_idx[s_my];
const float w_scale = smem_top_w[s_my];
const QLinearPtrs& d_ptrs = (s_my < MOE_ACTIVE) ? moe.down : moe.s_down;
const uint8_t* dw_q = d_ptrs.w_q + size_t(e) * (MOE_INTER / 2) * HIDDEN;
const __nv_bfloat16* d_scales = d_ptrs.scales + size_t(e) * (MOE_INTER / GROUP) * HIDDEN;
const __nv_bfloat16* d_zeros = d_ptrs.zeros + size_t(e) * (MOE_INTER / GROUP) * HIDDEN;
float acc_d[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_ct = -1;
int prev_ct = -1;
for (int u = u_start; u < u_end; ++u) {
int c_tile = u >> 3;
int g = u & 7;
if (c_tile != cur_ct) {
if (cur_ct >= 0) {
save_tile_128(acc_d, smem_red_g_prev, warp_id, lane);
prev_ct = cur_ct;
}
cur_ct = c_tile;
}
w4a16_group_tile_128(
dw_q, d_scales, d_zeros,
smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_d
);
}
flush_dual_tile_128(
(cur_ct >= 0) ? (sc.moe_out_fp32 + cur_ct * 128) : nullptr,
(prev_ct >= 0) ? (sc.moe_out_fp32 + prev_ct * 128) : nullptr,
acc_d, smem_red, smem_red_g_prev, tid, warp_id, lane, w_scale, w_scale
);
}
grid_barrier(local_epoch);
}
__global__ __launch_bounds__(BLOCK_THREADS, 2)
void kimi_decode_megakernel(KernelParams params) {
uint32_t local_epoch = (threadIdx.x == 0) ? *reinterpret_cast<volatile uint32_t*>(&g_sync_arr[1]) : 0u;
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int warp_id = tid >> 5;
const int lane = tid & 31;
const int eff_bid = (bid < 188) ? (bid << 1) : (((bid - 188) << 1) + 1);
__shared__ float smem_x[6144];
__shared__ float smem_red[8 * 128];
__shared__ int smem_top_idx[16];
__shared__ float smem_top_w[16];
ScratchPtrs sc = params.scratch;
float* smem_red_prev = smem_x + 4096;
// =========================================================================
// BLOCKS 0, 1, 2: KDA + MoE
// =========================================================================
for (int b = 0; b < 3; ++b) {
const KDABlockPtrs kb = params.kda[b];
// ---------------- Phase K1: Residual Add + Pre-Attn RMSNorm + Beta_proj (blocks 0..31) + Zeroing ----------------
if (bid < 32) {
const __nv_bfloat16* __restrict__ res_in = sc.res_h;
const float* __restrict__ m_in = sc.moe_out_fp32;
const __nv_bfloat16* __restrict__ a_norm = kb.attn_norm;
const __nv_bfloat16* __restrict__ bw = kb.beta_w + bid * HIDDEN;
float hv[9], an_val[9], bw_val[9];
__nv_bfloat16 h_bf[9];
float local_sq = 0.0f;
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
if (b == 0) {
h_bf[m] = params.hidden_in[i];
} else {
float h_old = bf16_to_f32(res_in[i]);
float m_out = round_bf16(m_in[i]);
h_bf[m] = f32_to_bf16(h_old + m_out);
}
an_val[m] = bf16_to_f32(a_norm[i]);
bw_val[m] = bf16_to_f32(bw[i]);
hv[m] = bf16_to_f32(h_bf[m]);
local_sq += hv[m] * hv[m];
}
if (bid == 0) {
#pragma unroll
for (int m = 0; m < 9; ++m) {
(sc.res_h + HIDDEN)[tid + m * BLOCK_THREADS] = h_bf[m];
}
}
local_sq = warp_reduce_sum(local_sq);
if (lane == 0) smem_red[warp_id] = local_sq;
__syncthreads();
if (tid < 32) {
float b_sq = (tid < 8) ? smem_red[tid] : 0.0f;
b_sq = warp_reduce_sum(b_sq);
if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f);
}
__syncthreads();
float inv_rms = smem_red[0];
float b_dot = 0.0f;
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
float xn = round_bf16(hv[m] * inv_rms * an_val[m]);
smem_x[i] = xn;
if (bid == 0) {
sc.norm_x[i] = xn;
}
b_dot += xn * bw_val[m];
}
b_dot = warp_reduce_sum(b_dot);
if (lane == 0) smem_red[warp_id] = b_dot;
__syncthreads();
if (tid < 32) {
float sum = (tid < 8) ? smem_red[tid] : 0.0f;
sum = warp_reduce_sum(sum);
if (tid == 0) {
float b_bf = round_bf16(sum);
sc.kda_beta[bid] = 1.0f / (1.0f + expf(-b_bf));
}
}
} else {
const int z_tid = (bid - 32) * BLOCK_THREADS + tid;
const int z_stride = (NUM_SMS - 32) * BLOCK_THREADS;
for (int i = z_tid; i < 4 * KDA_C; i += z_stride) {
sc.kda_qkvg_fp32[i] = 0.0f;
}
for (int i = z_tid; i < HIDDEN; i += z_stride) {
sc.attn_out_fp32[i] = 0.0f;
}
}
grid_barrier(local_epoch);
// ---------------- Phase K2: KDA q/k/v/g W4A16 GEMV ----------------
if (bid >= 32) {
for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) {
smem_x[i] = sc.norm_x[i];
}
__syncthreads();
}
constexpr int TOTAL_QKVG_UNITS = 4 * 32 * 18;
const int qk_start = (eff_bid * TOTAL_QKVG_UNITS) / NUM_SMS;
const int qk_end = ((eff_bid + 1) * TOTAL_QKVG_UNITS) / NUM_SMS;
float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_pc = -1;
int prev_pc = -1;
for (int u = qk_start; u < qk_end; ++u) {
int pc_idx = u / 18;
int g = u % 18;
if (pc_idx != cur_pc) {
if (cur_pc >= 0) {
save_tile_128(acc, smem_red_prev, warp_id, lane);
prev_pc = cur_pc;
}
cur_pc = pc_idx;
}
int p_id = pc_idx >> 5;
int c_tile = pc_idx & 31;
const QLinearPtrs& ql = (p_id == 0) ? kb.q_proj :
(p_id == 1) ? kb.k_proj :
(p_id == 2) ? kb.v_proj : kb.g_proj;
w4a16_group_tile_128(
ql.w_q, ql.scales, ql.zeros,
smem_x, KDA_C, c_tile * 128, g, warp_id, lane, acc
);
}
flush_dual_tile_128(
(cur_pc >= 0) ? (sc.kda_qkvg_fp32 + cur_pc * 128) : nullptr,
(prev_pc >= 0) ? (sc.kda_qkvg_fp32 + prev_pc * 128) : nullptr,
acc, smem_red, smem_red_prev, tid, warp_id, lane
);
grid_barrier(local_epoch);
// ---------------- Phase K3: Short Causal Conv + KDA Recurrent State S Update (128 blocks) ----------------
if (bid < 128) {
const int h = bid >> 2;
const int bj = bid & 3;
float* s_q = smem_x + 0;
float* s_k = smem_x + 128;
float* s_v = smem_x + 256;
float* s_g = smem_x + 384;
if (tid < 128) {
const int c = h * 128 + tid;
float q_raw = round_bf16(sc.kda_qkvg_fp32[0 * KDA_C + c]);
const __nv_bfloat16* cw_q = kb.conv_w + (0 * KDA_C + c) * 4;
float q_c = bf16_to_f32(kb.cq[0 * KDA_C + c]) * bf16_to_f32(cw_q[0])
+ bf16_to_f32(kb.cq[1 * KDA_C + c]) * bf16_to_f32(cw_q[1])
+ bf16_to_f32(kb.cq[2 * KDA_C + c]) * bf16_to_f32(cw_q[2])
+ q_raw * bf16_to_f32(cw_q[3]);
float q_silu = round_bf16(q_c / (1.0f + expf(-q_c)));
s_q[tid] = q_silu * 0.08838834764831845f;
float k_raw = round_bf16(sc.kda_qkvg_fp32[1 * KDA_C + c]);
const __nv_bfloat16* cw_k = kb.conv_w + (1 * KDA_C + c) * 4;
float k_c = bf16_to_f32(kb.ck[0 * KDA_C + c]) * bf16_to_f32(cw_k[0])
+ bf16_to_f32(kb.ck[1 * KDA_C + c]) * bf16_to_f32(cw_k[1])
+ bf16_to_f32(kb.ck[2 * KDA_C + c]) * bf16_to_f32(cw_k[2])
+ k_raw * bf16_to_f32(cw_k[3]);
s_k[tid] = round_bf16(k_c / (1.0f + expf(-k_c)));
float v_raw = round_bf16(sc.kda_qkvg_fp32[2 * KDA_C + c]);
const __nv_bfloat16* cw_v = kb.conv_w + (2 * KDA_C + c) * 4;
float v_c = bf16_to_f32(kb.cv[0 * KDA_C + c]) * bf16_to_f32(cw_v[0])
+ bf16_to_f32(kb.cv[1 * KDA_C + c]) * bf16_to_f32(cw_v[1])
+ bf16_to_f32(kb.cv[2 * KDA_C + c]) * bf16_to_f32(cw_v[2])
+ v_raw * bf16_to_f32(cw_v[3]);
s_v[tid] = round_bf16(v_c / (1.0f + expf(-v_c)));
float g_raw = round_bf16(sc.kda_qkvg_fp32[3 * KDA_C + c]);
float sp = (g_raw > 20.0f) ? g_raw : log1pf(expf(g_raw));
s_g[tid] = expf(-sp);
}
__syncthreads();
const int tid_i = tid >> 5;
const int tid_j = tid & 31;
const int j = bj * 32 + tid_j;
float* S_head = kb.S + h * (128 * 128);
float s_reg[16];
float pred_part = 0.0f;
#pragma unroll
for (int m = 0; m < 16; ++m) {
int i = tid_i * 16 + m;
float sv = S_head[i * 128 + j] * s_g[i];
s_reg[m] = sv;
pred_part += sv * s_k[i];
}
smem_red[tid_i * 32 + tid_j] = pred_part;
__syncthreads();
float pred = smem_red[0 * 32 + tid_j] + smem_red[1 * 32 + tid_j]
+ smem_red[2 * 32 + tid_j] + smem_red[3 * 32 + tid_j]
+ smem_red[4 * 32 + tid_j] + smem_red[5 * 32 + tid_j]
+ smem_red[6 * 32 + tid_j] + smem_red[7 * 32 + tid_j];
float diff_j = sc.kda_beta[h] * (s_v[j] - pred);
float o_part = 0.0f;
#pragma unroll
for (int m = 0; m < 16; ++m) {
int i = tid_i * 16 + m;
float sv = s_reg[m] + s_k[i] * diff_j;
S_head[i * 128 + j] = sv;
o_part += sv * s_q[i];
}
smem_red[tid_i * 32 + tid_j] = o_part;
__syncthreads();
if (tid_i == 0) {
float o_val = smem_red[0 * 32 + tid_j] + smem_red[1 * 32 + tid_j]
+ smem_red[2 * 32 + tid_j] + smem_red[3 * 32 + tid_j]
+ smem_red[4 * 32 + tid_j] + smem_red[5 * 32 + tid_j]
+ smem_red[6 * 32 + tid_j] + smem_red[7 * 32 + tid_j];
sc.kda_o[h * 128 + j] = round_bf16(o_val);
}
}
grid_barrier(local_epoch);
// ---------------- Phase K4: KDA o_proj W4A16 GEMV + Conv Window Update ----------------
if (bid < 16) {
int c = bid * BLOCK_THREADS + tid;
kb.cq[0 * KDA_C + c] = kb.cq[1 * KDA_C + c];
kb.cq[1 * KDA_C + c] = kb.cq[2 * KDA_C + c];
kb.cq[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[0 * KDA_C + c]);
kb.ck[0 * KDA_C + c] = kb.ck[1 * KDA_C + c];
kb.ck[1 * KDA_C + c] = kb.ck[2 * KDA_C + c];
kb.ck[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[1 * KDA_C + c]);
kb.cv[0 * KDA_C + c] = kb.cv[1 * KDA_C + c];
kb.cv[1 * KDA_C + c] = kb.cv[2 * KDA_C + c];
kb.cv[2 * KDA_C + c] = f32_to_bf16(sc.kda_qkvg_fp32[2 * KDA_C + c]);
}
for (int i = tid; i < KDA_C; i += BLOCK_THREADS) {
smem_x[i] = sc.kda_o[i];
}
__syncthreads();
constexpr int TOTAL_O_GROUPS = 18 * 32;
const int o_start = (eff_bid * TOTAL_O_GROUPS) / NUM_SMS;
const int o_end = ((eff_bid + 1) * TOTAL_O_GROUPS) / NUM_SMS;
float acc_o[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_oc = -1;
int prev_oc = -1;
for (int u = o_start; u < o_end; ++u) {
int c_tile = u >> 5;
int g = u & 31;
if (c_tile != cur_oc) {
if (cur_oc >= 0) {
save_tile_128(acc_o, smem_red_prev, warp_id, lane);
prev_oc = cur_oc;
}
cur_oc = c_tile;
}
w4a16_group_tile_128(
kb.o_proj.w_q, kb.o_proj.scales, kb.o_proj.zeros,
smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_o
);
}
flush_dual_tile_128(
(cur_oc >= 0) ? (sc.attn_out_fp32 + cur_oc * 128) : nullptr,
(prev_oc >= 0) ? (sc.attn_out_fp32 + prev_oc * 128) : nullptr,
acc_o, smem_red, smem_red_prev, tid, warp_id, lane
);
grid_barrier(local_epoch);
run_moe_phases(local_epoch, kb.moe_norm, kb.moe, sc, smem_x, smem_red, smem_top_idx, smem_top_w);
}
// =========================================================================
// BLOCK 3: MLA + MoE
// =========================================================================
const MLABlockPtrs mb = params.mla;
// ---------------- Phase M1: Residual Add + Pre-Attn RMSNorm + Cache Copy ----------------
if (bid == 0) {
const __nv_bfloat16* __restrict__ res_in = sc.res_h;
const float* __restrict__ m_in = sc.moe_out_fp32;
const __nv_bfloat16* __restrict__ a_norm = mb.attn_norm;
float hv[9], an_val[9];
__nv_bfloat16 h_bf[9];
float local_sq = 0.0f;
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
float h_old = bf16_to_f32(res_in[i]);
float m_out = round_bf16(m_in[i]);
h_bf[m] = f32_to_bf16(h_old + m_out);
an_val[m] = bf16_to_f32(a_norm[i]);
hv[m] = bf16_to_f32(h_bf[m]);
local_sq += hv[m] * hv[m];
}
#pragma unroll
for (int m = 0; m < 9; ++m) {
(sc.res_h + HIDDEN)[tid + m * BLOCK_THREADS] = h_bf[m];
}
local_sq = warp_reduce_sum(local_sq);
if (lane == 0) smem_red[warp_id] = local_sq;
__syncthreads();
if (tid < 32) {
float b_sq = (tid < 8) ? smem_red[tid] : 0.0f;
b_sq = warp_reduce_sum(b_sq);
if (tid == 0) smem_red[0] = rsqrtf(b_sq / float(HIDDEN) + 1.0e-6f);
}
__syncthreads();
float inv_rms = smem_red[0];
#pragma unroll
for (int m = 0; m < 9; ++m) {
int i = tid + m * BLOCK_THREADS;
float xn = round_bf16(hv[m] * inv_rms * an_val[m]);
smem_x[i] = xn;
sc.norm_x[i] = xn;
}
} else {
const int z_tid = (bid - 1) * BLOCK_THREADS + tid;
const int z_stride = (NUM_SMS - 1) * BLOCK_THREADS;
for (int i = z_tid; i < 53 * 128; i += z_stride) {
sc.mla_q_kv_fp32[i] = 0.0f;
}
for (int i = z_tid; i < HIDDEN; i += z_stride) {
sc.attn_out_fp32[i] = 0.0f;
}
if (mb.c_kv_src != mb.c_kv_dst) {
const size_t total_ckv_u4 = (size_t(mb.pos) * KV_LORA) / 8;
const uint4* src_ckv = reinterpret_cast<const uint4*>(mb.c_kv_src);
uint4* dst_ckv = reinterpret_cast<uint4*>(mb.c_kv_dst);
for (size_t i = z_tid; i < total_ckv_u4; i += z_stride) {
dst_ckv[i] = src_ckv[i];
}
const size_t total_kr_u4 = (size_t(mb.pos) * QK_ROPE) / 8;
const uint4* src_kr = reinterpret_cast<const uint4*>(mb.k_rope_src);
uint4* dst_kr = reinterpret_cast<uint4*>(mb.k_rope_dst);
for (size_t i = z_tid; i < total_kr_u4; i += z_stride) {
dst_kr[i] = src_kr[i];
}
}
}
grid_barrier(local_epoch);
// ---------------- Phase M2: MLA q_proj + kv_a W4A16 GEMV ----------------
if (bid > 0) {
for (int i = tid; i < HIDDEN; i += BLOCK_THREADS) {
smem_x[i] = sc.norm_x[i];
}
__syncthreads();
}
{
constexpr int TOTAL_QKV_GROUPS = 53 * 18;
const int mq_start = (eff_bid * TOTAL_QKV_GROUPS) / NUM_SMS;
const int mq_end = ((eff_bid + 1) * TOTAL_QKV_GROUPS) / NUM_SMS;
float acc_mq[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_mqc = -1;
int prev_mqc = -1;
for (int u = mq_start; u < mq_end; ++u) {
int c_tile = u / 18;
int g = u % 18;
if (c_tile != cur_mqc) {
if (cur_mqc >= 0) {
save_tile_128(acc_mq, smem_red_prev, warp_id, lane);
prev_mqc = cur_mqc;
}
cur_mqc = c_tile;
}
if (c_tile < 48) {
w4a16_group_tile_128(
mb.q_proj.w_q, mb.q_proj.scales, mb.q_proj.zeros,
smem_x, 6144, c_tile * 128, g, warp_id, lane, acc_mq
);
} else {
int eff_lane = (c_tile < 52) ? lane : (lane & 15);
w4a16_group_tile_128(
mb.kv_a.w_q, mb.kv_a.scales, mb.kv_a.zeros,
smem_x, 576, (c_tile - 48) * 128, g, warp_id, eff_lane, acc_mq
);
}
}
flush_dual_tile_128(
(cur_mqc >= 0) ? (sc.mla_q_kv_fp32 + cur_mqc * 128) : nullptr,
(prev_mqc >= 0) ? (sc.mla_q_kv_fp32 + prev_mqc * 128) : nullptr,
acc_mq, smem_red, smem_red_prev, tid, warp_id, lane
);
}
grid_barrier(local_epoch);
// ---------------- Phase M3: Absorbed Query + RoPE + KV Cache Append + Zeroing kda_o (256 blocks) ----------------
if (bid < 256) {
const int h = bid >> 3;
const int g_half = bid & 7;
const int g = g_half >> 1;
float* s_q_scaled = smem_x;
if (tid < 128) {
float qn = round_bf16(sc.mla_q_kv_fp32[h * Q_DIM + tid]);
int col = h * KVB_DIM + tid;
float sv = bf16_to_f32(mb.kv_b.scales[g * 8192 + col]);
float zv = bf16_to_f32(mb.kv_b.zeros[g * 8192 + col]);
float qs = qn * sv;
s_q_scaled[tid] = qs;
smem_red[tid] = qs * zv;
}
__syncthreads();
if (tid < 32) {
float zsum = smem_red[tid] + smem_red[tid + 32] + smem_red[tid + 64] + smem_red[tid + 96];
zsum = warp_reduce_sum(zsum);
if (tid == 0) smem_red[0] = zsum;
}
__syncthreads();
const float q_zero_sum = smem_red[0];
const float4 qs4 = *reinterpret_cast<const float4*>(&s_q_scaled[lane * 4]);
const int col_base = h * KVB_DIM + lane * 4;
#pragma unroll
for (int m = 0; m < 4; ++m) {
int rp = g_half * 32 + warp_id * 4 + m;
uint32_t p = *reinterpret_cast<const uint32_t*>(&mb.kv_b.w_q[rp * 8192 + col_base]);
float dot_ev = qs4.x * float(p & 0x0Fu) + qs4.y * float((p >> 8) & 0x0Fu)
+ qs4.z * float((p >> 16) & 0x0Fu) + qs4.w * float((p >> 24) & 0x0Fu);
float dot_od = qs4.x * float((p >> 4) & 0x0Fu) + qs4.y * float((p >> 12) & 0x0Fu)
+ qs4.z * float((p >> 20) & 0x0Fu) + qs4.w * float((p >> 28) & 0x0Fu);
dot_ev = warp_reduce_sum(dot_ev);
dot_od = warp_reduce_sum(dot_od);
if (lane == 0) {
sc.mla_q_abs[h * KV_LORA + rp * 2 + 0] = dot_ev - q_zero_sum;
sc.mla_q_abs[h * KV_LORA + rp * 2 + 1] = dot_od - q_zero_sum;
}
}
} else if (bid == 256) {
for (int idx = tid; idx < MLA_H * (QK_ROPE / 2); idx += BLOCK_THREADS) {
int h = idx >> 5;
int pair = idx & 31;
float inv_freq = powf(10000.0f, -float(2 * pair) / float(QK_ROPE));
float ang = float(mb.pos) * inv_freq;
float c_v = cosf(ang);
float s_v = sinf(ang);
int base = h * Q_DIM + QK_NOPE + 2 * pair;
float ev = round_bf16(sc.mla_q_kv_fp32[base + 0]);
float od = round_bf16(sc.mla_q_kv_fp32[base + 1]);
sc.mla_q_rope[h * QK_ROPE + 2 * pair + 0] = round_bf16(ev * c_v - od * s_v);
sc.mla_q_rope[h * QK_ROPE + 2 * pair + 1] = round_bf16(od * c_v + ev * s_v);
}
} else if (bid == 257) {
__nv_bfloat16* dst_ckv = mb.c_kv_dst + size_t(mb.pos) * KV_LORA;
for (int r = tid; r < KV_LORA; r += BLOCK_THREADS) {
dst_ckv[r] = f32_to_bf16(sc.mla_q_kv_fp32[6144 + r]);
}
if (tid < (QK_ROPE / 2)) {
int pair = tid;
float inv_freq = powf(10000.0f, -float(2 * pair) / float(QK_ROPE));
float ang = float(mb.pos) * inv_freq;
float c_v = cosf(ang);
float s_v = sinf(ang);
int base = 6144 + KV_LORA + 2 * pair;
float ev = round_bf16(sc.mla_q_kv_fp32[base + 0]);
float od = round_bf16(sc.mla_q_kv_fp32[base + 1]);
__nv_bfloat16* dst_kr = mb.k_rope_dst + size_t(mb.pos) * QK_ROPE;
dst_kr[2 * pair + 0] = f32_to_bf16(ev * c_v - od * s_v);
dst_kr[2 * pair + 1] = f32_to_bf16(od * c_v + ev * s_v);
}
} else {
const int z_tid = (bid - 258) * BLOCK_THREADS + tid;
const int z_stride = (NUM_SMS - 258) * BLOCK_THREADS;
for (int i = z_tid; i < 4096; i += z_stride) {
sc.kda_o[i] = 0.0f;
}
}
grid_barrier(local_epoch);
// ---------------- Phase M4: 188-Split 2-Head/Warp (376 blocks) Cooperative-Smem FlashDecoding ----------------
{
const int split_id = bid >> 1;
const int h_group = bid & 1;
const int h0 = h_group * 16 + warp_id * 2 + 0;
const int h1 = h_group * 16 + warp_id * 2 + 1;
const int total_len = mb.pos + 1;
const int l_start = (split_id * total_len) / NUM_L_SPLITS;
const int l_end = ((split_id + 1) * total_len) / NUM_L_SPLITS;
float q_abs_0[16], q_abs_1[16];
const float* qa0 = sc.mla_q_abs + h0 * KV_LORA + lane * 4;
const float* qa1 = sc.mla_q_abs + h1 * KV_LORA + lane * 4;
#pragma unroll
for (int c = 0; c < 4; ++c) {
float4 v0 = *reinterpret_cast<const float4*>(qa0 + c * 128);
float4 v1 = *reinterpret_cast<const float4*>(qa1 + c * 128);
q_abs_0[c * 4 + 0] = v0.x; q_abs_0[c * 4 + 1] = v0.y; q_abs_0[c * 4 + 2] = v0.z; q_abs_0[c * 4 + 3] = v0.w;
q_abs_1[c * 4 + 0] = v1.x; q_abs_1[c * 4 + 1] = v1.y; q_abs_1[c * 4 + 2] = v1.z; q_abs_1[c * 4 + 3] = v1.w;
}
float2 qr0 = *reinterpret_cast<const float2*>(sc.mla_q_rope + h0 * QK_ROPE + lane * 2);
float2 qr1 = *reinterpret_cast<const float2*>(sc.mla_q_rope + h1 * QK_ROPE + lane * 2);
float o_acc_0[16], o_acc_1[16];
#pragma unroll
for (int m = 0; m < 16; ++m) {
o_acc_0[m] = 0.0f;
o_acc_1[m] = 0.0f;
}
float m_i0 = -1.0e30f, d_i0 = 0.0f;
float m_i1 = -1.0e30f, d_i1 = 0.0f;
constexpr float mla_scale = 0.07216878364870322f;
uint4* smem_ckv_u4 = reinterpret_cast<uint4*>(smem_x);
uint32_t* smem_kr = reinterpret_cast<uint32_t*>(smem_ckv_u4 + 512);
for (int l_base = l_start; l_base < l_end; l_base += 8) {
int l_tok = l_base + warp_id;
if (l_tok < l_end) {
const uint4* g_ckv = reinterpret_cast<const uint4*>(mb.c_kv_dst + size_t(l_tok) * KV_LORA);
smem_ckv_u4[warp_id * 64 + lane] = g_ckv[lane];
smem_ckv_u4[warp_id * 64 + 32 + lane] = g_ckv[32 + lane];
const uint32_t* g_kr = reinterpret_cast<const uint32_t*>(mb.k_rope_dst + size_t(l_tok) * QK_ROPE);
smem_kr[warp_id * 32 + lane] = g_kr[lane];
}
__syncthreads();
int rem = l_end - l_base;
if (rem > 8) rem = 8;
for (int w = 0; w < rem; ++w) {
const __nv_bfloat16* s_ckv = reinterpret_cast<const __nv_bfloat16*>(smem_ckv_u4 + w * 64) + lane * 4;
uint32_t kr_u32 = smem_kr[w * 32 + lane];
const __nv_bfloat16* kr_bf = reinterpret_cast<const __nv_bfloat16*>(&kr_u32);
float kr0 = bf16_to_f32(kr_bf[0]);
float kr1 = bf16_to_f32(kr_bf[1]);
float c_val[16];
#pragma unroll
for (int c = 0; c < 4; ++c) {
uint2 u2 = *reinterpret_cast<const uint2*>(s_ckv + c * 128);
const __nv_bfloat16* bf = reinterpret_cast<const __nv_bfloat16*>(&u2);
c_val[c * 4 + 0] = bf16_to_f32(bf[0]);
c_val[c * 4 + 1] = bf16_to_f32(bf[1]);
c_val[c * 4 + 2] = bf16_to_f32(bf[2]);
c_val[c * 4 + 3] = bf16_to_f32(bf[3]);
}
float dot0 = qr0.x * kr0 + qr0.y * kr1;
float dot1 = qr1.x * kr0 + qr1.y * kr1;
#pragma unroll
for (int m = 0; m < 16; ++m) {
dot0 += q_abs_0[m] * c_val[m];
dot1 += q_abs_1[m] * c_val[m];
}
float s0 = warp_reduce_sum(dot0) * mla_scale;
float s1 = warp_reduce_sum(dot1) * mla_scale;
float mn0 = fmaxf(m_i0, s0);
float al0 = expf(m_i0 - mn0);
float p0 = expf(s0 - mn0);
d_i0 = d_i0 * al0 + p0;
m_i0 = mn0;
float mn1 = fmaxf(m_i1, s1);
float al1 = expf(m_i1 - mn1);
float p1 = expf(s1 - mn1);
d_i1 = d_i1 * al1 + p1;
m_i1 = mn1;
#pragma unroll
for (int m = 0; m < 16; ++m) {
o_acc_0[m] = o_acc_0[m] * al0 + p0 * c_val[m];
o_acc_1[m] = o_acc_1[m] * al1 + p1 * c_val[m];
}
}
__syncthreads();
}
const int sh0 = split_id * MLA_H + h0;
const int sh1 = split_id * MLA_H + h1;
if (lane == 0) {
sc.mla_att_m[sh0] = m_i0;
sc.mla_att_d[sh0] = d_i0;
sc.mla_att_m[sh1] = m_i1;
sc.mla_att_d[sh1] = d_i1;
}
if (l_start < l_end) {
float* od0 = sc.mla_att_o + size_t(sh0) * KV_LORA + lane * 4;
float* od1 = sc.mla_att_o + size_t(sh1) * KV_LORA + lane * 4;
#pragma unroll
for (int c = 0; c < 4; ++c) {
*reinterpret_cast<float4*>(od0 + c * 128) = make_float4(
o_acc_0[c * 4 + 0], o_acc_0[c * 4 + 1], o_acc_0[c * 4 + 2], o_acc_0[c * 4 + 3]
);
*reinterpret_cast<float4*>(od1 + c * 128) = make_float4(
o_acc_1[c * 4 + 0], o_acc_1[c * 4 + 1], o_acc_1[c * 4 + 2], o_acc_1[c * 4 + 3]
);
}
}
}
grid_barrier(local_epoch);
// ---------------- Phase M5: 128-Block Parallel Split-K Combine + Un-absorption (c_out @ W_uv) ----------------
if (bid < 128) {
const int h = bid >> 2;
const int g = bid & 3;
float my_m = (tid < NUM_L_SPLITS) ? sc.mla_att_m[tid * MLA_H + h] : -1.0e30f;
float my_d = (tid < NUM_L_SPLITS) ? sc.mla_att_d[tid * MLA_H + h] : 0.0f;
float w_max = my_m;
#pragma unroll
for (int off = 16; off > 0; off >>= 1) {
w_max = fmaxf(w_max, __shfl_xor_sync(0xffffffff, w_max, off));
}
if (lane == 0) smem_red[warp_id] = w_max;
__syncthreads();
if (tid < 8) {
float gm = smem_red[tid];
#pragma unroll
for (int off = 4; off > 0; off >>= 1) {
gm = fmaxf(gm, __shfl_xor_sync(0xff, gm, off));
}
if (tid == 0) smem_red[0] = gm;
}
__syncthreads();
float m_max = smem_red[0];
float my_w = (my_d > 0.0f) ? expf(my_m - m_max) : 0.0f;
float wd_sum = warp_reduce_sum(my_d * my_w);
if (lane == 0) smem_red[warp_id] = wd_sum;
__syncthreads();
if (tid < 32) {
float gd = (tid < 8) ? smem_red[tid] : 0.0f;
gd = warp_reduce_sum(gd);
if (tid == 0) smem_red[0] = 1.0f / gd;
}
__syncthreads();
float inv_d = smem_red[0];
if (tid < NUM_L_SPLITS) {
smem_red[16 + tid] = my_w * inv_d;
}
__syncthreads();
float4 o_sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
const int g_off = g * 128 + lane * 4;
for (int s = warp_id; s < NUM_L_SPLITS; s += 8) {
float ws = smem_red[16 + s];
if (ws > 0.0f) {
float4 v4 = *reinterpret_cast<const float4*>(sc.mla_att_o + size_t(s * MLA_H + h) * KV_LORA + g_off);
o_sum.x += v4.x * ws;
o_sum.y += v4.y * ws;
o_sum.z += v4.z * ws;
o_sum.w += v4.w * ws;
}
}
*reinterpret_cast<float4*>(&smem_x[warp_id * 128 + lane * 4]) = o_sum;
__syncthreads();
if (tid < 128) {
float sum = smem_x[0 * 128 + tid] + smem_x[1 * 128 + tid]
+ smem_x[2 * 128 + tid] + smem_x[3 * 128 + tid]
+ smem_x[4 * 128 + tid] + smem_x[5 * 128 + tid]
+ smem_x[6 * 128 + tid] + smem_x[7 * 128 + tid];
smem_x[g * 128 + tid] = sum;
}
__syncthreads();
float acc_uv[4] = {0.0f, 0.0f, 0.0f, 0.0f};
const int n0 = h * KVB_DIM + QK_NOPE;
w4a16_group_tile_128(
mb.kv_b.w_q, mb.kv_b.scales, mb.kv_b.zeros,
smem_x, 8192, n0, g, warp_id, lane, acc_uv
);
flush_tile_128(sc.kda_o + h * V_HEAD, acc_uv, smem_red, tid, warp_id, lane);
}
grid_barrier(local_epoch);
// ---------------- Phase M6: MLA o_proj W4A16 GEMV ----------------
for (int i = tid; i < 4096; i += BLOCK_THREADS) {
smem_x[i] = round_bf16(sc.kda_o[i]);
}
__syncthreads();
{
constexpr int TOTAL_O_GROUPS = 18 * 32;
const int o_start = (eff_bid * TOTAL_O_GROUPS) / NUM_SMS;
const int o_end = ((eff_bid + 1) * TOTAL_O_GROUPS) / NUM_SMS;
float acc_o[4] = {0.0f, 0.0f, 0.0f, 0.0f};
int cur_oc = -1;
int prev_oc = -1;
for (int u = o_start; u < o_end; ++u) {
int c_tile = u >> 5;
int g = u & 31;
if (c_tile != cur_oc) {
if (cur_oc >= 0) {
save_tile_128(acc_o, smem_red_prev, warp_id, lane);
prev_oc = cur_oc;
}
cur_oc = c_tile;
}
w4a16_group_tile_128(
mb.o_proj.w_q, mb.o_proj.scales, mb.o_proj.zeros,
smem_x, HIDDEN, c_tile * 128, g, warp_id, lane, acc_o
);
}
flush_dual_tile_128(
(cur_oc >= 0) ? (sc.attn_out_fp32 + cur_oc * 128) : nullptr,
(prev_oc >= 0) ? (sc.attn_out_fp32 + prev_oc * 128) : nullptr,
acc_o, smem_red, smem_red_prev, tid, warp_id, lane
);
}
grid_barrier(local_epoch);
// ---------------- Phases M7, M8, M9: MLA MoE ----------------
run_moe_phases(local_epoch, mb.moe_norm, mb.moe, sc, smem_x, smem_red, smem_top_idx, smem_top_w);
// ---------------- Final Epilogue: Write Final Hidden Output ----------------
if (bid < 9) {
int i = bid * BLOCK_THREADS + tid;
float h_old = bf16_to_f32(sc.res_h[i]);
float m_out = round_bf16(sc.moe_out_fp32[i]);
params.hidden_out[i] = f32_to_bf16(h_old + m_out);
}
}
extern "C" void launch_kimi_megakernel(const void* params_ptr, cudaStream_t stream) {
const KernelParams* p = reinterpret_cast<const KernelParams*>(params_ptr);
kimi_decode_megakernel<<<NUM_SMS, BLOCK_THREADS, 0, stream>>>(*p);
}
"""
_CPP_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <stdint.h>
extern "C" void launch_kimi_megakernel(const void* params_ptr, cudaStream_t stream);
struct QLinearPtrs {
const void* w_q;
const void* scales;
const void* zeros;
};
struct MoEPtrs {
const void* router_w;
QLinearPtrs gate;
QLinearPtrs up;
QLinearPtrs down;
QLinearPtrs s_gate;
QLinearPtrs s_up;
QLinearPtrs s_down;
};
struct KDABlockPtrs {
const void* attn_norm;
const void* moe_norm;
QLinearPtrs q_proj;
QLinearPtrs k_proj;
QLinearPtrs v_proj;
QLinearPtrs g_proj;
const void* beta_w;
const void* conv_w;
QLinearPtrs o_proj;
MoEPtrs moe;
void* S;
void* cq;
void* ck;
void* cv;
};
struct MLABlockPtrs {
const void* attn_norm;
const void* moe_norm;
QLinearPtrs q_proj;
QLinearPtrs kv_a;
QLinearPtrs kv_b;
QLinearPtrs o_proj;
MoEPtrs moe;
const void* c_kv_src;
const void* k_rope_src;
void* c_kv_dst;
void* k_rope_dst;
int pos;
};
struct ScratchPtrs {
void* norm_x;
void* attn_out_fp32;
void* moe_out_fp32;
void* res_h;
void* kda_qkvg_fp32;
void* kda_o;
void* kda_beta;
void* mla_q_kv_fp32;
void* mla_q_abs;
void* mla_q_rope;
void* mla_att_m;
void* mla_att_d;
void* mla_att_o;
void* router_logits;
void* moe_gate_fp32;
void* moe_up_fp32;
};
struct KernelParams {
KDABlockPtrs kda[3];
MLABlockPtrs mla;
ScratchPtrs scratch;
const void* hidden_in;
void* hidden_out;
};
static KernelParams g_params;
void init_static_params(const std::vector<int64_t>& ptrs) {
int idx = 0;
auto next_ptr = [&]() -> void* { return reinterpret_cast<void*>(ptrs[idx++]); };
auto read_ql = [&]() -> QLinearPtrs {
QLinearPtrs q;
q.w_q = next_ptr();
q.scales = next_ptr();
q.zeros = next_ptr();
return q;
};
auto read_moe = [&]() -> MoEPtrs {
MoEPtrs m;
m.router_w = next_ptr();
m.gate = read_ql();
m.up = read_ql();
m.down = read_ql();
m.s_gate = read_ql();
m.s_up = read_ql();
m.s_down = read_ql();
return m;
};
for (int b = 0; b < 3; ++b) {
g_params.kda[b].attn_norm = next_ptr();
g_params.kda[b].moe_norm = next_ptr();
g_params.kda[b].q_proj = read_ql();
g_params.kda[b].k_proj = read_ql();
g_params.kda[b].v_proj = read_ql();
g_params.kda[b].g_proj = read_ql();
g_params.kda[b].beta_w = next_ptr();
g_params.kda[b].conv_w = next_ptr();
g_params.kda[b].o_proj = read_ql();
g_params.kda[b].moe = read_moe();
}
g_params.mla.attn_norm = next_ptr();
g_params.mla.moe_norm = next_ptr();
g_params.mla.q_proj = read_ql();
g_params.mla.kv_a = read_ql();
g_params.mla.kv_b = read_ql();
g_params.mla.o_proj = read_ql();
g_params.mla.moe = read_moe();
g_params.scratch.norm_x = next_ptr();
g_params.scratch.attn_out_fp32 = next_ptr();
g_params.scratch.moe_out_fp32 = next_ptr();
g_params.scratch.res_h = next_ptr();
g_params.scratch.kda_qkvg_fp32 = next_ptr();
g_params.scratch.kda_o = next_ptr();
g_params.scratch.kda_beta = next_ptr();
g_params.scratch.mla_q_kv_fp32 = next_ptr();
g_params.scratch.mla_q_abs = next_ptr();
g_params.scratch.mla_q_rope = next_ptr();
g_params.scratch.mla_att_m = next_ptr();
g_params.scratch.mla_att_d = next_ptr();
g_params.scratch.mla_att_o = next_ptr();
g_params.scratch.router_logits = next_ptr();
g_params.scratch.moe_gate_fp32 = next_ptr();
g_params.scratch.moe_up_fp32 = next_ptr();
}
void run_step(
torch::Tensor hidden_in,
torch::Tensor hidden_out,
torch::Tensor S0, torch::Tensor cq0, torch::Tensor ck0, torch::Tensor cv0,
torch::Tensor S1, torch::Tensor cq1, torch::Tensor ck1, torch::Tensor cv1,
torch::Tensor S2, torch::Tensor cq2, torch::Tensor ck2, torch::Tensor cv2,
torch::Tensor c_kv_src, torch::Tensor k_rope_src,
torch::Tensor c_kv_dst, torch::Tensor k_rope_dst,
int64_t pos
) {
g_params.hidden_in = hidden_in.data_ptr();
g_params.hidden_out = hidden_out.data_ptr();
g_params.kda[0].S = S0.data_ptr();
g_params.kda[0].cq = cq0.data_ptr();
g_params.kda[0].ck = ck0.data_ptr();
g_params.kda[0].cv = cv0.data_ptr();
g_params.kda[1].S = S1.data_ptr();
g_params.kda[1].cq = cq1.data_ptr();
g_params.kda[1].ck = ck1.data_ptr();
g_params.kda[1].cv = cv1.data_ptr();
g_params.kda[2].S = S2.data_ptr();
g_params.kda[2].cq = cq2.data_ptr();
g_params.kda[2].ck = ck2.data_ptr();
g_params.kda[2].cv = cv2.data_ptr();
g_params.mla.c_kv_src = c_kv_src.data_ptr();
g_params.mla.k_rope_src = k_rope_src.data_ptr();
g_params.mla.c_kv_dst = c_kv_dst.data_ptr();
g_params.mla.k_rope_dst = k_rope_dst.data_ptr();
g_params.mla.pos = static_cast<int>(pos);
cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
launch_kimi_megakernel(&g_params, stream);
}
"""
_EXT = load_inline(
name="kimi_megakernel_v5",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=["init_static_params", "run_step"],
no_implicit_headers=True,
extra_cuda_cflags=["-O3", "--use_fast_math"],
)
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._prepared_dev = 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)
self._prepared_dev = None
def _prepare_buffers(self, dev: torch.device):
self._norm_x = torch.empty(2304, dtype=torch.float32, device=dev)
self._attn_out_fp32 = torch.empty(2304, dtype=torch.float32, device=dev)
self._moe_out_fp32 = torch.empty(2304, dtype=torch.float32, device=dev)
self._res_h = torch.empty(2 * 2304, dtype=torch.bfloat16, device=dev)
self._kda_qkvg_fp32 = torch.empty(4 * 4096, dtype=torch.float32, device=dev)
self._kda_o = torch.empty(4096, dtype=torch.float32, device=dev)
self._kda_beta = torch.empty(32, dtype=torch.float32, device=dev)
self._mla_q_kv_fp32 = torch.empty(53 * 128, dtype=torch.float32, device=dev)
self._mla_q_abs = torch.empty(32 * 512, dtype=torch.float32, device=dev)
self._mla_q_rope = torch.empty(32 * 64, dtype=torch.float32, device=dev)
self._mla_att_m = torch.empty(188 * 32, dtype=torch.float32, device=dev)
self._mla_att_d = torch.empty(188 * 32, dtype=torch.float32, device=dev)
self._mla_att_o = torch.empty(188 * 32 * 512, dtype=torch.float32, device=dev)
self._router_logits = torch.empty(64, dtype=torch.float32, device=dev)
self._moe_gate_fp32 = torch.empty(9 * 1024, dtype=torch.float32, device=dev)
self._moe_up_fp32 = torch.empty(9 * 1024, dtype=torch.float32, device=dev)
self._c_kv_cap = torch.empty((20480, 512), dtype=torch.bfloat16, device=dev)
self._k_rope_cap = torch.empty((20480, 64), dtype=torch.bfloat16, device=dev)
self._hidden_out = torch.empty(2304, dtype=torch.bfloat16, device=dev)
ptrs = []
def add_ql(ql):
ptrs.extend([ql.w_q.data_ptr(), ql.scales.data_ptr(), ql.zeros.data_ptr()])
def add_moe(moe):
ptrs.append(moe.router.weight.data_ptr())
add_ql(moe.gate)
add_ql(moe.up)
add_ql(moe.down)
add_ql(moe.s_gate)
add_ql(moe.s_up)
add_ql(moe.s_down)
for b in range(3):
blk = self.blocks[b]
ptrs.append(blk.attn_norm.data_ptr())
ptrs.append(blk.moe_norm.data_ptr())
add_ql(blk.attn.q_proj)
add_ql(blk.attn.k_proj)
add_ql(blk.attn.v_proj)
add_ql(blk.attn.g_proj)
ptrs.append(blk.attn.beta_proj.weight.data_ptr())
ptrs.append(blk.attn.conv_w.data_ptr())
add_ql(blk.attn.o_proj)
add_moe(blk.moe)
blk3 = self.blocks[3]
ptrs.append(blk3.attn_norm.data_ptr())
ptrs.append(blk3.moe_norm.data_ptr())
add_ql(blk3.attn.q_proj)
add_ql(blk3.attn.kv_a)
add_ql(blk3.attn.kv_b)
add_ql(blk3.attn.o_proj)
add_moe(blk3.moe)
ptrs.extend([
self._norm_x.data_ptr(),
self._attn_out_fp32.data_ptr(),
self._moe_out_fp32.data_ptr(),
self._res_h.data_ptr(),
self._kda_qkvg_fp32.data_ptr(),
self._kda_o.data_ptr(),
self._kda_beta.data_ptr(),
self._mla_q_kv_fp32.data_ptr(),
self._mla_q_abs.data_ptr(),
self._mla_q_rope.data_ptr(),
self._mla_att_m.data_ptr(),
self._mla_att_d.data_ptr(),
self._mla_att_o.data_ptr(),
self._router_logits.data_ptr(),
self._moe_gate_fp32.data_ptr(),
self._moe_up_fp32.data_ptr(),
])
_EXT.init_static_params(ptrs)
self._prepared_dev = dev
def _apply(self, fn):
ret = super()._apply(fn)
dev = self.blocks[0].attn_norm.device
if dev.type == "cuda":
self._prepare_buffers(dev)
return ret
def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
ret = super().load_state_dict(state_dict, strict=strict, assign=assign)
dev = self.blocks[0].attn_norm.device
if dev.type == "cuda":
self._prepare_buffers(dev)
return ret
def step(self, hidden: torch.Tensor, state: list):
if self._prepared_dev != hidden.device:
self._prepare_buffers(hidden.device)
s0, s1, s2, s3 = state[0], state[1], state[2], state[3]
c_kv_src = s3["c_kv"]
k_rope_src = s3["k_rope"]
pos = c_kv_src.shape[0]
_EXT.run_step(
hidden,
self._hidden_out,
s0["S"], s0["cq"], s0["ck"], s0["cv"],
s1["S"], s1["cq"], s1["ck"], s1["cv"],
s2["S"], s2["cq"], s2["ck"], s2["cv"],
c_kv_src, k_rope_src,
self._c_kv_cap, self._k_rope_cap,
pos,
)
s3["c_kv"] = self._c_kv_cap[: pos + 1]
s3["k_rope"] = self._k_rope_cap[: pos + 1]
return self._hidden_out, state
def init_state(cfg: Config, context_len: int, seed: int = 0, device: str = "cuda"):
g = torch.Generator(device=device).manual_seed(seed)
state = []
for k in cfg.pattern:
if k == "K":
H, Dk, W = cfg.kda_heads, cfg.kda_head_dim, cfg.short_conv
S = torch.randn(H, Dk, Dk, generator=g, device=device, dtype=torch.float32) * 0.05
cq = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1
ck = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1
cv = torch.randn(W - 1, H * Dk, generator=g, device=device, dtype=cfg.dtype) * 0.1
state.append({"S": S, "cq": cq, "ck": ck, "cv": cv})
else:
c_kv = torch.randn(context_len, cfg.kv_lora, generator=g, device=device, dtype=cfg.dtype) * 0.1
k_rope = torch.randn(context_len, cfg.qk_rope, generator=g, device=device, dtype=cfg.dtype) * 0.1
state.append({"c_kv": c_kv, "k_rope": k_rope})
return state
def init_token(cfg: Config, seed: int = 0, device: str = "cuda") -> torch.Tensor:
g = torch.Generator(device=device).manual_seed(seed)
return torch.randn(cfg.hidden, generator=g, device=device, dtype=cfg.dtype) * 0.2
20261013_175145_agy_gemini-4-argon-high_02_kimi_linear_decode