KernelBench mega · RTX PRO 6000
Kimi-Linear Decode Gemini 3.8 Flash (High)
2.74×geomean speedup across shapes
manually audited: clean
RTX PRO 6000 Gemini 3.8 Flash (high) cell (2.74x isolated regrade; in-run was 2.75x). Genuine single-launch W4A16 decode megakernel, honestly parallelised in its GEMVs and serialised in its attention: one block per SM at 256 threads (about 12.5% occupancy), the KDA state update spills a 128-float per-thread array, and the whole MLA phase runs under blockIdx.x < 32 while each head streams the latent cache twice, so speedup decays 3.35x -> 2.73x -> 2.27x with context. Same-buffer overwrite on the box 2026-09-03, sole GPU owner: primed step cos(ref,sol)=0.9999; continuation step with the hidden buffer overwritten in place (same data_ptr) cos(out1,out2)=0.0335 and cos(ref,sol)=0.9999; fresh state with a third token cos(out2,out3)=-0.0339 and cos(ref,sol)=0.9999. Transcript: every outputs/runs string is this run's own id, no results/, leaderboard, DEVLOG, src/eval or bin/python reads, no env or key printing; all eight template files byte-identical. The agent self-terminated at 25 minutes on an unlimited budget.
harnessagy
Kernel source (redacted)
"""W4A16 Kimi-Linear Hybrid Decode Single Fused Megakernel.
Single GPU kernel launch per decode step.
Fuses the entire 4-layer motif:
- 3 KDA layers (RMSNorm, q/k/v/g/beta GEMVs, short causal conv, recurrence on S, o_proj)
- 1 MLA layer (RMSNorm, q/kv_a GEMVs, RoPE, cache update, absorbed Q/O, attention, o_proj)
- 4 MoE blocks (RMSNorm, router, top-8 + 1 shared expert gate/up/down GEMVs)
- All residual additions
into a single cooperative CUDA __global__ kernel launch.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
EPS = 1.0e-6
GROUP_SIZE = 128
CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <c10/cuda/CUDAStream.h>
namespace cg = cooperative_groups;
#define HIDDEN 2304
#define KDA_HEADS 32
#define KDA_HEAD_DIM 128
#define KDA_CHANNELS 4096
#define CONV_LEN 4
#define MLA_HEADS 32
#define KV_LORA 512
#define QK_NOPE 128
#define QK_ROPE 64
#define V_HEAD 128
#define MOE_INTER 1024
#define N_EXPERTS 64
#define N_ACTIVE 8
#define N_SHARED 1
#define ROUTED_SCALING 2.446f
#define EPS 1e-6f
struct QuantWeight {
const uint8_t* w_q;
const __nv_bfloat16* scales;
const __nv_bfloat16* zeros;
int in_f;
int out_f;
};
struct KDAWeights {
QuantWeight q_proj;
QuantWeight k_proj;
QuantWeight v_proj;
QuantWeight g_proj;
const __nv_bfloat16* beta_proj;
const __nv_bfloat16* conv_w;
QuantWeight o_proj;
};
struct MLAWeights {
QuantWeight q_proj;
QuantWeight kv_a;
QuantWeight kv_b;
QuantWeight o_proj;
};
struct MoEWeights {
const __nv_bfloat16* router;
QuantWeight gate;
QuantWeight up;
QuantWeight down;
QuantWeight s_gate;
QuantWeight s_up;
QuantWeight s_down;
};
struct BlockWeights {
const __nv_bfloat16* attn_norm;
const __nv_bfloat16* moe_norm;
KDAWeights kda;
MLAWeights mla;
MoEWeights moe;
};
struct MegakernelParams {
BlockWeights blocks[4];
};
struct KDAState {
float* S;
__nv_bfloat16* cq;
__nv_bfloat16* ck;
__nv_bfloat16* cv;
};
struct MLAState {
const __nv_bfloat16* old_c_kv;
const __nv_bfloat16* old_k_rope;
__nv_bfloat16* new_c_kv;
__nv_bfloat16* new_k_rope;
int pos;
};
struct Scratchpad {
float x[HIDDEN];
float x_norm[HIDDEN];
float x_group_sum[32];
float gemv_out[18432];
float beta[32];
float kda_q[KDA_CHANNELS];
float kda_k[KDA_CHANNELS];
float kda_v[KDA_CHANNELS];
float kda_g[KDA_CHANNELS];
float kda_o[KDA_CHANNELS];
float attn_out[HIDDEN];
float router_logits[N_EXPERTS];
float router_probs[N_EXPERTS];
int topk_idx[N_ACTIVE];
float topk_w[N_ACTIVE];
float moe_inter[9 * MOE_INTER];
float moe_out[HIDDEN];
float mla_q[6144];
float mla_kv[576];
float q_abs[MLA_HEADS * KV_LORA];
float q_rope[MLA_HEADS * QK_ROPE];
float c_att[MLA_HEADS * KV_LORA];
float scores[MLA_HEADS * 16384];
};
__device__ void device_rmsnorm(
const float* __restrict__ x,
const __nv_bfloat16* __restrict__ weight,
float* __restrict__ out,
cg::grid_group& grid
) {
__shared__ float sh_val[32];
int tid = threadIdx.x;
if (blockIdx.x == 0) {
float sum_sq = 0.0f;
for (int i = tid; i < HIDDEN; i += blockDim.x) {
float val = x[i];
sum_sq += val * val;
}
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
sum_sq += __shfl_down_sync(0xffffffff, sum_sq, mask);
}
if ((tid % 32) == 0) {
sh_val[tid / 32] = sum_sq;
}
__syncthreads();
if (tid == 0) {
float total = 0.0f;
int num_warps = blockDim.x / 32;
for (int w = 0; w < num_warps; ++w) total += sh_val[w];
float rsqrt_val = rsqrtf(total / (float)HIDDEN + EPS);
sh_val[0] = rsqrt_val;
}
__syncthreads();
float r = sh_val[0];
for (int i = tid; i < HIDDEN; i += blockDim.x) {
float w = __bfloat162float(weight[i]);
out[i] = x[i] * r * w;
}
}
grid.sync();
}
__device__ void compute_x_group_sums(
const float* __restrict__ x,
float* __restrict__ x_group_sums,
int in_f,
cg::grid_group& grid
) {
if (blockIdx.x == 0) {
int num_groups = in_f / 128;
int tid = threadIdx.x;
if (tid < num_groups) {
float s = 0.0f;
int base = tid * 128;
#pragma unroll 8
for (int i = 0; i < 128; ++i) {
s += x[base + i];
}
x_group_sums[tid] = s;
}
}
grid.sync();
}
__device__ inline float compute_col_slice(
const float* __restrict__ sh_x,
const float* __restrict__ x_group_sums,
const uint8_t* __restrict__ w_q,
const __nv_bfloat16* __restrict__ scales,
const __nv_bfloat16* __restrict__ zeros,
int in_f, int out_f, int col,
int g_start, int g_end
) {
float acc = 0.0f;
for (int g = g_start; g < g_end; ++g) {
float scale = __bfloat162float(scales[g * out_f + col]);
float zero = __bfloat162float(zeros[g * out_f + col]);
float x_sum = x_group_sums[g];
int k_start = g * 128;
const float2* x_pairs = (const float2*)&sh_x[k_start];
const uint8_t* w_ptr = &w_q[(k_start / 2) * out_f + col];
float group_acc = 0.0f;
#pragma unroll 8
for (int p = 0; p < 64; ++p) {
uint8_t byte_val = *w_ptr;
w_ptr += out_f;
float w0 = (float)(byte_val & 0x0F);
float w1 = (float)(byte_val >> 4);
float2 x_val = x_pairs[p];
group_acc += x_val.x * w0 + x_val.y * w1;
}
acc += (group_acc - zero * x_sum) * scale;
}
return acc;
}
__global__ void kimi_decode_megakernel_opt(
const __nv_bfloat16* __restrict__ hidden_in,
__nv_bfloat16* __restrict__ hidden_out,
KDAState st_kda0,
KDAState st_kda1,
KDAState st_kda2,
MLAState st_mla,
const MegakernelParams* __restrict__ params,
Scratchpad* __restrict__ scratch
) {
cg::grid_group grid = cg::this_grid();
int global_tid = blockIdx.x * blockDim.x + threadIdx.x;
int num_threads = gridDim.x * blockDim.x;
int wid = global_tid / 32;
int lane = threadIdx.x % 32;
int total_warps = num_threads / 32;
__shared__ float sh_x[HIDDEN];
for (int i = global_tid; i < HIDDEN; i += num_threads) {
scratch->x[i] = __bfloat162float(hidden_in[i]);
}
grid.sync();
#pragma unroll 1
for (int b = 0; b < 4; ++b) {
const BlockWeights& blk = params->blocks[b];
// 1. Attention RMSNorm
device_rmsnorm(scratch->x, blk.attn_norm, scratch->x_norm, grid);
for (int i = threadIdx.x; i < HIDDEN; i += blockDim.x) {
sh_x[i] = scratch->x_norm[i];
}
__syncthreads();
if (b < 3) {
// ---------------------------------------------------------------
// KDA Block
// ---------------------------------------------------------------
KDAState& st_kda = (b == 0) ? st_kda0 : ((b == 1) ? st_kda1 : st_kda2);
compute_x_group_sums(scratch->x_norm, scratch->x_group_sum, HIDDEN, grid);
for (int i = global_tid; i < 16384; i += num_threads) {
scratch->gemv_out[i] = 0.0f;
}
grid.sync();
for (int task = wid; task < 1024; task += total_warps) {
int tile = task / 2;
int k_slice = task % 2;
int col = tile * 32 + lane;
int proj = col / KDA_CHANNELS;
int c = col % KDA_CHANNELS;
const QuantWeight& qw = (proj == 0) ? blk.kda.q_proj :
((proj == 1) ? blk.kda.k_proj :
((proj == 2) ? blk.kda.v_proj : blk.kda.g_proj));
int g_start = k_slice * 9;
int g_end = g_start + 9;
float val = compute_col_slice(sh_x, scratch->x_group_sum, qw.w_q, qw.scales, qw.zeros, HIDDEN, KDA_CHANNELS, c, g_start, g_end);
atomicAdd(&scratch->gemv_out[col], val);
}
if (wid == 0 && lane < 32) {
int c = lane;
float b_acc = 0.0f;
#pragma unroll 8
for (int k = 0; k < HIDDEN; ++k) {
b_acc += sh_x[k] * __bfloat162float(blk.kda.beta_proj[c * HIDDEN + k]);
}
scratch->beta[c] = b_acc;
}
grid.sync();
// Conv & Act
for (int c = global_tid; c < KDA_CHANNELS; c += num_threads) {
float val_q = scratch->gemv_out[c];
float val_k = scratch->gemv_out[KDA_CHANNELS + c];
float val_v = scratch->gemv_out[2 * KDA_CHANNELS + c];
float val_g = scratch->gemv_out[3 * KDA_CHANNELS + c];
float pq0 = __bfloat162float(st_kda.cq[0 * KDA_CHANNELS + c]);
float pq1 = __bfloat162float(st_kda.cq[1 * KDA_CHANNELS + c]);
float pq2 = __bfloat162float(st_kda.cq[2 * KDA_CHANNELS + c]);
float wq0 = __bfloat162float(blk.kda.conv_w[(0 * KDA_CHANNELS + c) * 4 + 0]);
float wq1 = __bfloat162float(blk.kda.conv_w[(0 * KDA_CHANNELS + c) * 4 + 1]);
float wq2 = __bfloat162float(blk.kda.conv_w[(0 * KDA_CHANNELS + c) * 4 + 2]);
float wq3 = __bfloat162float(blk.kda.conv_w[(0 * KDA_CHANNELS + c) * 4 + 3]);
float out_q = pq0 * wq0 + pq1 * wq1 + pq2 * wq2 + val_q * wq3;
out_q = (out_q / (1.0f + expf(-out_q))) * 0.08838834764831845f;
st_kda.cq[0 * KDA_CHANNELS + c] = __float2bfloat16(pq1);
st_kda.cq[1 * KDA_CHANNELS + c] = __float2bfloat16(pq2);
st_kda.cq[2 * KDA_CHANNELS + c] = __float2bfloat16(val_q);
scratch->kda_q[c] = out_q;
float pk0 = __bfloat162float(st_kda.ck[0 * KDA_CHANNELS + c]);
float pk1 = __bfloat162float(st_kda.ck[1 * KDA_CHANNELS + c]);
float pk2 = __bfloat162float(st_kda.ck[2 * KDA_CHANNELS + c]);
float wk0 = __bfloat162float(blk.kda.conv_w[(1 * KDA_CHANNELS + c) * 4 + 0]);
float wk1 = __bfloat162float(blk.kda.conv_w[(1 * KDA_CHANNELS + c) * 4 + 1]);
float wk2 = __bfloat162float(blk.kda.conv_w[(1 * KDA_CHANNELS + c) * 4 + 2]);
float wk3 = __bfloat162float(blk.kda.conv_w[(1 * KDA_CHANNELS + c) * 4 + 3]);
float out_k = pk0 * wk0 + pk1 * wk1 + pk2 * wk2 + val_k * wk3;
out_k = out_k / (1.0f + expf(-out_k));
st_kda.ck[0 * KDA_CHANNELS + c] = __float2bfloat16(pk1);
st_kda.ck[1 * KDA_CHANNELS + c] = __float2bfloat16(pk2);
st_kda.ck[2 * KDA_CHANNELS + c] = __float2bfloat16(val_k);
scratch->kda_k[c] = out_k;
float pv0 = __bfloat162float(st_kda.cv[0 * KDA_CHANNELS + c]);
float pv1 = __bfloat162float(st_kda.cv[1 * KDA_CHANNELS + c]);
float pv2 = __bfloat162float(st_kda.cv[2 * KDA_CHANNELS + c]);
float wv0 = __bfloat162float(blk.kda.conv_w[(2 * KDA_CHANNELS + c) * 4 + 0]);
float wv1 = __bfloat162float(blk.kda.conv_w[(2 * KDA_CHANNELS + c) * 4 + 1]);
float wv2 = __bfloat162float(blk.kda.conv_w[(2 * KDA_CHANNELS + c) * 4 + 2]);
float wv3 = __bfloat162float(blk.kda.conv_w[(2 * KDA_CHANNELS + c) * 4 + 3]);
float out_v = pv0 * wv0 + pv1 * wv1 + pv2 * wv2 + val_v * wv3;
out_v = out_v / (1.0f + expf(-out_v));
st_kda.cv[0 * KDA_CHANNELS + c] = __float2bfloat16(pv1);
st_kda.cv[1 * KDA_CHANNELS + c] = __float2bfloat16(pv2);
st_kda.cv[2 * KDA_CHANNELS + c] = __float2bfloat16(val_v);
scratch->kda_v[c] = out_v;
float g_val = (val_g > 20.0f) ? -val_g : -log1pf(expf(val_g));
scratch->kda_g[c] = g_val;
if (c < 32) {
scratch->beta[c] = 1.0f / (1.0f + expf(-scratch->beta[c]));
}
}
grid.sync();
// KDA Recurrence on S
if (blockIdx.x < KDA_HEADS) {
int h = blockIdx.x;
int tid = threadIdx.x;
__shared__ float sh_k[128];
__shared__ float sh_q[128];
__shared__ float sh_exp_g[128];
__shared__ float sh_b;
if (tid < 128) {
sh_k[tid] = scratch->kda_k[h * 128 + tid];
sh_q[tid] = scratch->kda_q[h * 128 + tid];
sh_exp_g[tid] = expf(scratch->kda_g[h * 128 + tid]);
}
if (tid == 0) {
sh_b = scratch->beta[h];
}
__syncthreads();
if (tid < 128) {
int j = tid;
float* col_S_ptr = &st_kda.S[h * 16384 + j];
float pred_j = 0.0f;
float col_S[128];
#pragma unroll 4
for (int i = 0; i < 128; ++i) {
float s_val = col_S_ptr[i * 128] * sh_exp_g[i];
col_S[i] = s_val;
pred_j += s_val * sh_k[i];
}
float v_j = scratch->kda_v[h * 128 + j];
float diff_j = v_j - pred_j;
float beta_h = sh_b;
float o_j = 0.0f;
#pragma unroll 4
for (int i = 0; i < 128; ++i) {
float s_up = col_S[i] + beta_h * sh_k[i] * diff_j;
col_S_ptr[i * 128] = s_up;
o_j += s_up * sh_q[i];
}
scratch->kda_o[h * 128 + j] = o_j;
}
}
grid.sync();
// KDA o_proj
compute_x_group_sums(scratch->kda_o, scratch->x_group_sum, KDA_CHANNELS, grid);
for (int i = global_tid; i < HIDDEN; i += num_threads) {
scratch->attn_out[i] = 0.0f;
}
grid.sync();
for (int task = wid; task < 576; task += total_warps) {
int tile = task / 8;
int k_slice = task % 8;
int col = tile * 32 + lane;
if (col < HIDDEN) {
int g_start = k_slice * 4;
int g_end = g_start + 4;
float val = compute_col_slice(scratch->kda_o, scratch->x_group_sum, blk.kda.o_proj.w_q, blk.kda.o_proj.scales, blk.kda.o_proj.zeros, KDA_CHANNELS, HIDDEN, col, g_start, g_end);
atomicAdd(&scratch->attn_out[col], val);
}
}
grid.sync();
} else {
// ---------------------------------------------------------------
// MLA Block
// ---------------------------------------------------------------
compute_x_group_sums(scratch->x_norm, scratch->x_group_sum, HIDDEN, grid);
for (int i = global_tid; i < 6144; i += num_threads) scratch->mla_q[i] = 0.0f;
grid.sync();
for (int task = wid; task < 768; task += total_warps) {
int tile = task / 4;
int k_slice = task % 4;
int col = tile * 32 + lane;
int g_start = (k_slice < 3) ? (k_slice * 4) : 12;
int g_end = (k_slice < 3) ? (g_start + 4) : 18;
float val = compute_col_slice(sh_x, scratch->x_group_sum, blk.mla.q_proj.w_q, blk.mla.q_proj.scales, blk.mla.q_proj.zeros, HIDDEN, 6144, col, g_start, g_end);
atomicAdd(&scratch->mla_q[col], val);
}
for (int i = global_tid; i < 576; i += num_threads) scratch->mla_kv[i] = 0.0f;
grid.sync();
for (int task = wid; task < 144; task += total_warps) {
int tile = task / 8;
int k_slice = task % 8;
int col = tile * 32 + lane;
if (col < 576) {
int g_start = (k_slice < 7) ? (k_slice * 2) : 14;
int g_end = (k_slice < 7) ? (g_start + 2) : 18;
float val = compute_col_slice(sh_x, scratch->x_group_sum, blk.mla.kv_a.w_q, blk.mla.kv_a.scales, blk.mla.kv_a.zeros, HIDDEN, 576, col, g_start, g_end);
atomicAdd(&scratch->mla_kv[col], val);
}
}
grid.sync();
// MLA RoPE & Cache Update
int pos = st_mla.pos;
int L = pos + 1;
int total_c_kv_elements = pos * KV_LORA;
for (int i = global_tid; i < total_c_kv_elements; i += num_threads) {
st_mla.new_c_kv[i] = st_mla.old_c_kv[i];
}
int total_k_rope_elements = pos * QK_ROPE;
for (int i = global_tid; i < total_k_rope_elements; i += num_threads) {
st_mla.new_k_rope[i] = st_mla.old_k_rope[i];
}
if (blockIdx.x == 0 && threadIdx.x < 32) {
int j = threadIdx.x;
float inv = 1.0f / powf(10000.0f, (2.0f * j) / 64.0f);
float ang = pos * inv;
float cos_val = cosf(ang);
float sin_val = sinf(ang);
float k_even = scratch->mla_kv[512 + 2 * j];
float k_odd = scratch->mla_kv[512 + 2 * j + 1];
st_mla.new_k_rope[pos * 64 + 2 * j] = __float2bfloat16(k_even * cos_val - k_odd * sin_val);
st_mla.new_k_rope[pos * 64 + 2 * j + 1] = __float2bfloat16(k_odd * cos_val + k_even * sin_val);
#pragma unroll
for (int m = 0; m < 16; ++m) {
int r = j * 16 + m;
st_mla.new_c_kv[pos * 512 + r] = __float2bfloat16(scratch->mla_kv[r]);
}
}
for (int h = wid; h < MLA_HEADS; h += total_warps) {
int j = lane;
float inv = 1.0f / powf(10000.0f, (2.0f * j) / 64.0f);
float ang = pos * inv;
float cos_val = cosf(ang);
float sin_val = sinf(ang);
int q_base = h * (QK_NOPE + QK_ROPE) + QK_NOPE;
float q_even = scratch->mla_q[q_base + 2 * j];
float q_odd = scratch->mla_q[q_base + 2 * j + 1];
scratch->q_rope[h * 64 + 2 * j] = q_even * cos_val - q_odd * sin_val;
scratch->q_rope[h * 64 + 2 * j + 1] = q_odd * cos_val + q_even * sin_val;
}
grid.sync();
// MLA Absorbed Q: q_abs[h, r]
for (int idx = wid; idx < 16384; idx += total_warps) {
int h = idx / 512;
int r = idx % 512;
int q_base = h * (QK_NOPE + QK_ROPE);
int row_idx = r / 2;
bool is_even = (r % 2 == 0);
int g = r / 128;
int col_base = h * 256;
float dot = 0.0f;
#pragma unroll
for (int m = 0; m < 4; ++m) {
int d = lane * 4 + m;
int col = col_base + d;
uint8_t b = blk.mla.kv_b.w_q[row_idx * 8192 + col];
float w_val = is_even ? (float)(b & 0x0F) : (float)(b >> 4);
float zero = __bfloat162float(blk.mla.kv_b.zeros[g * 8192 + col]);
float scale = __bfloat162float(blk.mla.kv_b.scales[g * 8192 + col]);
dot += scratch->mla_q[q_base + d] * ((w_val - zero) * scale);
}
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
dot += __shfl_down_sync(0xffffffff, dot, mask);
}
if (lane == 0) {
scratch->q_abs[h * 512 + r] = dot;
}
}
grid.sync();
// MLA Attention Scores, Softmax & Weighted Sum
if (blockIdx.x < MLA_HEADS) {
int h = blockIdx.x;
int tid = threadIdx.x;
int warp_id = tid / 32;
int lane_id = tid % 32;
__shared__ float sh_q_abs[512];
__shared__ float sh_q_rope[64];
__shared__ float sh_max;
__shared__ float sh_sum;
for (int i = tid; i < 512; i += blockDim.x) sh_q_abs[i] = scratch->q_abs[h * 512 + i];
for (int i = tid; i < 64; i += blockDim.x) sh_q_rope[i] = scratch->q_rope[h * 64 + i];
__syncthreads();
float local_max = -1e30f;
float scale_attn = 0.07216878364870322f;
for (int l = warp_id; l < L; l += 8) {
float dot = 0.0f;
#pragma unroll
for (int i = 0; i < 16; ++i) {
int r = lane_id * 16 + i;
float c_val = __bfloat162float(st_mla.new_c_kv[l * 512 + r]);
dot += c_val * sh_q_abs[r];
}
#pragma unroll
for (int i = 0; i < 2; ++i) {
int d = lane_id * 2 + i;
float k_val = __bfloat162float(st_mla.new_k_rope[l * 64 + d]);
dot += k_val * sh_q_rope[d];
}
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
dot += __shfl_down_sync(0xffffffff, dot, mask);
}
if (lane_id == 0) {
float sc = dot * scale_attn;
scratch->scores[h * 16384 + l] = sc;
if (sc > local_max) local_max = sc;
}
}
__shared__ float warp_max[8];
if (lane_id == 0) warp_max[warp_id] = local_max;
__syncthreads();
if (tid == 0) {
float m = warp_max[0];
for (int w = 1; w < 8; ++w) if (warp_max[w] > m) m = warp_max[w];
sh_max = m;
}
__syncthreads();
float global_max = sh_max;
float local_sum = 0.0f;
for (int l = tid; l < L; l += blockDim.x) {
float p = expf(scratch->scores[h * 16384 + l] - global_max);
scratch->scores[h * 16384 + l] = p;
local_sum += p;
}
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
local_sum += __shfl_down_sync(0xffffffff, local_sum, mask);
}
__shared__ float warp_sum[8];
if (lane_id == 0) warp_sum[warp_id] = local_sum;
__syncthreads();
if (tid == 0) {
float s = 0.0f;
for (int w = 0; w < 8; ++w) s += warp_sum[w];
sh_sum = (s > 0.0f) ? (1.0f / s) : 0.0f;
}
__syncthreads();
float inv_sum = sh_sum;
float acc0 = 0.0f;
float acc1 = 0.0f;
int r0 = tid;
int r1 = tid + 256;
#pragma unroll 4
for (int l = 0; l < L; ++l) {
float p = scratch->scores[h * 16384 + l] * inv_sum;
const __nv_bfloat16* c_l = &st_mla.new_c_kv[l * 512];
acc0 += p * __bfloat162float(c_l[r0]);
acc1 += p * __bfloat162float(c_l[r1]);
}
scratch->c_att[h * 512 + r0] = acc0;
scratch->c_att[h * 512 + r1] = acc1;
}
grid.sync();
// MLA Absorb O
for (int i = global_tid; i < 4096; i += num_threads) scratch->kda_o[i] = 0.0f;
grid.sync();
for (int task = wid; task < 256; task += total_warps) {
int tile = task / 2;
int k_slice = task % 2;
int col_idx = tile * 32 + lane;
int h = col_idx / 128;
int d = col_idx % 128;
int col = h * 256 + 128 + d;
int g_start = k_slice * 2;
int g_end = g_start + 2;
float acc = 0.0f;
for (int g = g_start; g < g_end; ++g) {
float scale = __bfloat162float(blk.mla.kv_b.scales[g * 8192 + col]);
float zero = __bfloat162float(blk.mla.kv_b.zeros[g * 8192 + col]);
int r_start = g * 128;
int row_start = r_start / 2;
#pragma unroll 4
for (int p = 0; p < 64; ++p) {
uint8_t b = blk.mla.kv_b.w_q[(row_start + p) * 8192 + col];
float w0 = ((float)(b & 0x0F) - zero) * scale;
float w1 = ((float)(b >> 4) - zero) * scale;
acc += scratch->c_att[h * 512 + r_start + 2 * p] * w0 + scratch->c_att[h * 512 + r_start + 2 * p + 1] * w1;
}
}
atomicAdd(&scratch->kda_o[col_idx], acc);
}
grid.sync();
// MLA o_proj
compute_x_group_sums(scratch->kda_o, scratch->x_group_sum, 4096, grid);
for (int i = global_tid; i < HIDDEN; i += num_threads) scratch->attn_out[i] = 0.0f;
grid.sync();
for (int task = wid; task < 576; task += total_warps) {
int tile = task / 8;
int k_slice = task % 8;
int col = tile * 32 + lane;
if (col < HIDDEN) {
int g_start = k_slice * 4;
int g_end = g_start + 4;
float val = compute_col_slice(scratch->kda_o, scratch->x_group_sum, blk.mla.o_proj.w_q, blk.mla.o_proj.scales, blk.mla.o_proj.zeros, 4096, HIDDEN, col, g_start, g_end);
atomicAdd(&scratch->attn_out[col], val);
}
}
grid.sync();
}
// Residual Add 1: x = x + attn_out
for (int i = global_tid; i < HIDDEN; i += num_threads) {
scratch->x[i] += scratch->attn_out[i];
}
grid.sync();
// -------------------------------------------------------------------
// MoE
// -------------------------------------------------------------------
device_rmsnorm(scratch->x, blk.moe_norm, scratch->x_norm, grid);
for (int i = threadIdx.x; i < HIDDEN; i += blockDim.x) {
sh_x[i] = scratch->x_norm[i];
}
__syncthreads();
if (blockIdx.x == 0) {
int tid = threadIdx.x;
if (tid < N_EXPERTS) {
float logit = 0.0f;
#pragma unroll 8
for (int k = 0; k < HIDDEN; ++k) {
logit += sh_x[k] * __bfloat162float(blk.moe.router[tid * HIDDEN + k]);
}
scratch->router_logits[tid] = logit;
}
__syncthreads();
if (tid == 0) {
float m = scratch->router_logits[0];
for (int e = 1; e < N_EXPERTS; ++e) {
if (scratch->router_logits[e] > m) m = scratch->router_logits[e];
}
float s = 0.0f;
for (int e = 0; e < N_EXPERTS; ++e) {
float p = expf(scratch->router_logits[e] - m);
scratch->router_probs[e] = p;
s += p;
}
float inv_s = (s > 0.0f) ? (1.0f / s) : 0.0f;
for (int e = 0; e < N_EXPERTS; ++e) {
scratch->router_probs[e] *= inv_s;
}
for (int j = 0; j < N_ACTIVE; ++j) {
float best_val = -1.0f;
int best_idx = -1;
for (int e = 0; e < N_EXPERTS; ++e) {
float p = scratch->router_probs[e];
bool selected = false;
for (int prev = 0; prev < j; ++prev) {
if (scratch->topk_idx[prev] == e) {
selected = true;
break;
}
}
if (!selected && p > best_val) {
best_val = p;
best_idx = e;
}
}
scratch->topk_idx[j] = best_idx;
scratch->topk_w[j] = best_val;
}
float sum_topk = 0.0f;
for (int j = 0; j < N_ACTIVE; ++j) sum_topk += scratch->topk_w[j];
float norm_factor = ROUTED_SCALING / (sum_topk + 1e-9f);
for (int j = 0; j < N_ACTIVE; ++j) scratch->topk_w[j] *= norm_factor;
}
}
grid.sync();
// MoE Gate & Up
compute_x_group_sums(scratch->x_norm, scratch->x_group_sum, HIDDEN, grid);
for (int i = global_tid; i < 18432; i += num_threads) scratch->gemv_out[i] = 0.0f;
grid.sync();
for (int task = wid; task < 1152; task += total_warps) {
int tile = task / 2;
int k_slice = task % 2;
int col_id = tile * 32 + lane;
int e = col_id / 2048;
int rem = col_id % 2048;
bool is_gate = (rem < 1024);
int c = is_gate ? rem : (rem - 1024);
int expert_id = (e < N_ACTIVE) ? scratch->topk_idx[e] : 0;
const QuantWeight& qw = (e < N_ACTIVE) ? (is_gate ? blk.moe.gate : blk.moe.up)
: (is_gate ? blk.moe.s_gate : blk.moe.s_up);
int w_offset = expert_id * (HIDDEN / 2) * MOE_INTER;
int s_offset = expert_id * (HIDDEN / 128) * MOE_INTER;
int g_start = k_slice * 9;
int g_end = g_start + 9;
float val = compute_col_slice(sh_x, scratch->x_group_sum, qw.w_q + w_offset, qw.scales + s_offset, qw.zeros + s_offset, HIDDEN, MOE_INTER, c, g_start, g_end);
atomicAdd(&scratch->gemv_out[col_id], val);
}
grid.sync();
for (int i = global_tid; i < 9 * MOE_INTER; i += num_threads) {
int e = i / MOE_INTER;
int c = i % MOE_INTER;
float g_val = scratch->gemv_out[e * 2048 + c];
float u_val = scratch->gemv_out[e * 2048 + 1024 + c];
float silu_g = g_val / (1.0f + expf(-g_val));
scratch->moe_inter[e * MOE_INTER + c] = silu_g * u_val;
}
for (int i = global_tid; i < HIDDEN; i += num_threads) {
scratch->moe_out[i] = 0.0f;
}
grid.sync();
// MoE Down
for (int task = wid; task < 1296; task += total_warps) {
int tile = task / 2;
int k_slice = task % 2;
int col_id = tile * 32 + lane;
int e = col_id / HIDDEN;
int d = col_id % HIDDEN;
int expert_id = (e < N_ACTIVE) ? scratch->topk_idx[e] : 0;
float w_e = (e < N_ACTIVE) ? scratch->topk_w[e] : 1.0f;
const QuantWeight& qw = (e < N_ACTIVE) ? blk.moe.down : blk.moe.s_down;
int w_offset = expert_id * (MOE_INTER / 2) * HIDDEN;
int s_offset = expert_id * (MOE_INTER / 128) * HIDDEN;
float acc = 0.0f;
const float* expert_inter = &scratch->moe_inter[e * MOE_INTER];
const float2* inter_pairs = (const float2*)expert_inter;
int g_start = k_slice * 4;
int g_end = g_start + 4;
for (int g = g_start; g < g_end; ++g) {
float scale = __bfloat162float(qw.scales[s_offset + g * HIDDEN + d]);
float zero = __bfloat162float(qw.zeros[s_offset + g * HIDDEN + d]);
int k_start = g * 128;
const uint8_t* w_ptr = &qw.w_q[w_offset + (k_start / 2) * HIDDEN + d];
#pragma unroll 8
for (int p = 0; p < 64; ++p) {
uint8_t byte_val = *w_ptr;
w_ptr += HIDDEN;
float w0 = ((float)(byte_val & 0x0F) - zero) * scale;
float w1 = ((float)(byte_val >> 4) - zero) * scale;
float2 val = inter_pairs[(k_start / 2) + p];
acc += val.x * w0 + val.y * w1;
}
}
atomicAdd(&scratch->moe_out[d], w_e * acc);
}
grid.sync();
for (int i = global_tid; i < HIDDEN; i += num_threads) {
scratch->x[i] += scratch->moe_out[i];
}
grid.sync();
}
for (int i = global_tid; i < HIDDEN; i += num_threads) {
hidden_out[i] = __float2bfloat16(scratch->x[i]);
}
}
void launch_kimi_decode_megakernel_opt(
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 old_c_kv, torch::Tensor old_k_rope,
torch::Tensor new_c_kv, torch::Tensor new_k_rope,
int pos,
torch::Tensor params_tensor,
torch::Tensor scratch_tensor
) {
KDAState st0 = { s0.data_ptr<float>(), (__nv_bfloat16*)cq0.data_ptr<at::BFloat16>(), (__nv_bfloat16*)ck0.data_ptr<at::BFloat16>(), (__nv_bfloat16*)cv0.data_ptr<at::BFloat16>() };
KDAState st1 = { s1.data_ptr<float>(), (__nv_bfloat16*)cq1.data_ptr<at::BFloat16>(), (__nv_bfloat16*)ck1.data_ptr<at::BFloat16>(), (__nv_bfloat16*)cv1.data_ptr<at::BFloat16>() };
KDAState st2 = { s2.data_ptr<float>(), (__nv_bfloat16*)cq2.data_ptr<at::BFloat16>(), (__nv_bfloat16*)ck2.data_ptr<at::BFloat16>(), (__nv_bfloat16*)cv2.data_ptr<at::BFloat16>() };
MLAState st_mla = {
(const __nv_bfloat16*)old_c_kv.data_ptr<at::BFloat16>(),
(const __nv_bfloat16*)old_k_rope.data_ptr<at::BFloat16>(),
(__nv_bfloat16*)new_c_kv.data_ptr<at::BFloat16>(),
(__nv_bfloat16*)new_k_rope.data_ptr<at::BFloat16>(),
pos
};
const MegakernelParams* params_ptr = (const MegakernelParams*)params_tensor.data_ptr<uint8_t>();
Scratchpad* scratch_ptr = (Scratchpad*)scratch_tensor.data_ptr<uint8_t>();
const __nv_bfloat16* in_ptr = (const __nv_bfloat16*)hidden_in.data_ptr<at::BFloat16>();
__nv_bfloat16* out_ptr = (__nv_bfloat16*)hidden_out.data_ptr<at::BFloat16>();
void* args[] = {
&in_ptr, &out_ptr,
&st0, &st1, &st2, &st_mla,
¶ms_ptr, &scratch_ptr
};
cudaLaunchCooperativeKernel(
(void*)kimi_decode_megakernel_opt,
dim3(188),
dim3(256),
args,
0,
c10::cuda::getCurrentCUDAStream(hidden_in.get_device())
);
}
void setup_kda_block(
torch::Tensor params_tensor,
int b,
torch::Tensor attn_norm, torch::Tensor moe_norm,
torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
torch::Tensor k_wq, torch::Tensor k_s, torch::Tensor k_z,
torch::Tensor v_wq, torch::Tensor v_s, torch::Tensor v_z,
torch::Tensor g_wq, torch::Tensor g_s, torch::Tensor g_z,
torch::Tensor beta_w, torch::Tensor conv_w,
torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
torch::Tensor router_w,
torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
) {
MegakernelParams* p = (MegakernelParams*)params_tensor.data_ptr<uint8_t>();
BlockWeights& blk = p->blocks[b];
blk.attn_norm = (const __nv_bfloat16*)attn_norm.data_ptr<at::BFloat16>();
blk.moe_norm = (const __nv_bfloat16*)moe_norm.data_ptr<at::BFloat16>();
blk.kda.q_proj = { q_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)q_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)q_z.data_ptr<at::BFloat16>(), HIDDEN, KDA_CHANNELS };
blk.kda.k_proj = { k_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)k_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)k_z.data_ptr<at::BFloat16>(), HIDDEN, KDA_CHANNELS };
blk.kda.v_proj = { v_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)v_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)v_z.data_ptr<at::BFloat16>(), HIDDEN, KDA_CHANNELS };
blk.kda.g_proj = { g_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)g_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)g_z.data_ptr<at::BFloat16>(), HIDDEN, KDA_CHANNELS };
blk.kda.beta_proj = (const __nv_bfloat16*)beta_w.data_ptr<at::BFloat16>();
blk.kda.conv_w = (const __nv_bfloat16*)conv_w.data_ptr<at::BFloat16>();
blk.kda.o_proj = { o_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)o_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)o_z.data_ptr<at::BFloat16>(), KDA_CHANNELS, HIDDEN };
blk.moe.router = (const __nv_bfloat16*)router_w.data_ptr<at::BFloat16>();
blk.moe.gate = { gate_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)gate_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)gate_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.up = { up_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)up_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)up_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.down = { down_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)down_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)down_z.data_ptr<at::BFloat16>(), MOE_INTER, HIDDEN };
blk.moe.s_gate = { s_gate_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_gate_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_gate_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.s_up = { s_up_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_up_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_up_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.s_down = { s_down_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_down_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_down_z.data_ptr<at::BFloat16>(), MOE_INTER, HIDDEN };
}
void setup_mla_block(
torch::Tensor params_tensor,
int b,
torch::Tensor attn_norm, torch::Tensor moe_norm,
torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
torch::Tensor kva_wq, torch::Tensor kva_s, torch::Tensor kva_z,
torch::Tensor kvb_wq, torch::Tensor kvb_s, torch::Tensor kvb_z,
torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
torch::Tensor router_w,
torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
) {
MegakernelParams* p = (MegakernelParams*)params_tensor.data_ptr<uint8_t>();
BlockWeights& blk = p->blocks[b];
blk.attn_norm = (const __nv_bfloat16*)attn_norm.data_ptr<at::BFloat16>();
blk.moe_norm = (const __nv_bfloat16*)moe_norm.data_ptr<at::BFloat16>();
blk.mla.q_proj = { q_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)q_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)q_z.data_ptr<at::BFloat16>(), HIDDEN, 6144 };
blk.mla.kv_a = { kva_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)kva_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)kva_z.data_ptr<at::BFloat16>(), HIDDEN, 576 };
blk.mla.kv_b = { kvb_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)kvb_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)kvb_z.data_ptr<at::BFloat16>(), KV_LORA, 8192 };
blk.mla.o_proj = { o_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)o_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)o_z.data_ptr<at::BFloat16>(), KDA_CHANNELS, HIDDEN };
blk.moe.router = (const __nv_bfloat16*)router_w.data_ptr<at::BFloat16>();
blk.moe.gate = { gate_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)gate_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)gate_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.up = { up_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)up_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)up_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.down = { down_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)down_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)down_z.data_ptr<at::BFloat16>(), MOE_INTER, HIDDEN };
blk.moe.s_gate = { s_gate_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_gate_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_gate_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.s_up = { s_up_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_up_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_up_z.data_ptr<at::BFloat16>(), HIDDEN, MOE_INTER };
blk.moe.s_down = { s_down_wq.data_ptr<uint8_t>(), (const __nv_bfloat16*)s_down_s.data_ptr<at::BFloat16>(), (const __nv_bfloat16*)s_down_z.data_ptr<at::BFloat16>(), MOE_INTER, HIDDEN };
}
int get_params_size() {
return sizeof(MegakernelParams);
}
int get_scratch_size() {
return sizeof(Scratchpad);
}
"""
CPP_SRC = """
#include <torch/extension.h>
void launch_kimi_decode_megakernel_opt(
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 old_c_kv, torch::Tensor old_k_rope,
torch::Tensor new_c_kv, torch::Tensor new_k_rope,
int pos,
torch::Tensor params_tensor,
torch::Tensor scratch_tensor
);
void setup_kda_block(
torch::Tensor params_tensor,
int b,
torch::Tensor attn_norm, torch::Tensor moe_norm,
torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
torch::Tensor k_wq, torch::Tensor k_s, torch::Tensor k_z,
torch::Tensor v_wq, torch::Tensor v_s, torch::Tensor v_z,
torch::Tensor g_wq, torch::Tensor g_s, torch::Tensor g_z,
torch::Tensor beta_w, torch::Tensor conv_w,
torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
torch::Tensor router_w,
torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
);
void setup_mla_block(
torch::Tensor params_tensor,
int b,
torch::Tensor attn_norm, torch::Tensor moe_norm,
torch::Tensor q_wq, torch::Tensor q_s, torch::Tensor q_z,
torch::Tensor kva_wq, torch::Tensor kva_s, torch::Tensor kva_z,
torch::Tensor kvb_wq, torch::Tensor kvb_s, torch::Tensor kvb_z,
torch::Tensor o_wq, torch::Tensor o_s, torch::Tensor o_z,
torch::Tensor router_w,
torch::Tensor gate_wq, torch::Tensor gate_s, torch::Tensor gate_z,
torch::Tensor up_wq, torch::Tensor up_s, torch::Tensor up_z,
torch::Tensor down_wq, torch::Tensor down_s, torch::Tensor down_z,
torch::Tensor s_gate_wq, torch::Tensor s_gate_s, torch::Tensor s_gate_z,
torch::Tensor s_up_wq, torch::Tensor s_up_s, torch::Tensor s_up_z,
torch::Tensor s_down_wq, torch::Tensor s_down_s, torch::Tensor s_down_z
);
int get_params_size();
int get_scratch_size();
"""
_MODULE = None
def _get_cuda_module():
global _MODULE
if _MODULE is None:
_MODULE = load_inline(
name="kimi_megakernel_lib",
cpp_sources=CPP_SRC,
cuda_sources=CUDA_SRC,
functions=[
"launch_kimi_decode_megakernel_opt",
"setup_kda_block",
"setup_mla_block",
"get_params_size",
"get_scratch_size"
],
extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_120,code=sm_120"]
)
return _MODULE
@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[[REDACTED: IP]] & 0xF
hi = w_q[[REDACTED: IP]] & 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[[REDACTED: IP]] = w_packed & 0xF
out[[REDACTED: IP]] = (w_packed >> 4) & 0xF
return out
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)
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 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)
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 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)
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)
class KDA(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
self.q_proj = QuantLinear(d, H * Dk, cfg.group)
self.k_proj = QuantLinear(d, H * Dk, cfg.group)
self.v_proj = QuantLinear(d, H * Dk, cfg.group)
self.g_proj = QuantLinear(d, H * Dk, cfg.group)
self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype)
self.conv_w = nn.Parameter(torch.empty(3, H * Dk, cfg.short_conv, dtype=cfg.dtype))
self.o_proj = QuantLinear(H * Dk, d, cfg.group)
self.scale = Dk ** -0.5
class MLA(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
H, d = cfg.mla_heads, cfg.hidden
self.q_proj = QuantLinear(d, H * (cfg.qk_nope + cfg.qk_rope), cfg.group)
self.kv_a = QuantLinear(d, cfg.kv_lora + cfg.qk_rope, cfg.group)
self.kv_b = QuantLinear(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head), cfg.group)
self.o_proj = QuantLinear(H * cfg.v_head, d, cfg.group)
self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5
class MoE(nn.Module):
def __init__(self, cfg: Config):
super().__init__()
self.cfg = cfg
d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts
self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype)
self.gate = QuantExperts(E, d, m, cfg.group)
self.up = QuantExperts(E, d, m, cfg.group)
self.down = QuantExperts(E, m, d, cfg.group)
self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group)
self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group)
self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group)
class Block(nn.Module):
def __init__(self, cfg: Config, kind: str):
super().__init__()
self.kind = kind
self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
self.moe = MoE(cfg)
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.mod = _get_cuda_module()
self.params_size = self.mod.get_params_size()
self.scratch_size = self.mod.get_scratch_size()
self.params_tensor_cpu = torch.zeros(self.params_size, dtype=torch.uint8, device="cpu")
self.params_tensor_cuda = None
self.scratch_tensor = None
self._synced_weights = False
self.reset_parameters()
def reset_parameters(self):
g = torch.Generator(device="cpu").manual_seed(1234)
for mod in self.modules():
if isinstance(mod, (QuantLinear, QuantExperts)):
mod.init_random(g)
elif isinstance(mod, nn.Linear):
nn.init.normal_(mod.weight, 0.0, 0.02, generator=g)
elif isinstance(mod, KDA):
nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g)
def _sync_weights(self, device):
for b in range(4):
blk = self.blocks[b]
if b < 3:
self.mod.setup_kda_block(
self.params_tensor_cpu, b,
blk.attn_norm, blk.moe_norm,
blk.attn.q_proj.w_q, blk.attn.q_proj.scales, blk.attn.q_proj.zeros,
blk.attn.k_proj.w_q, blk.attn.k_proj.scales, blk.attn.k_proj.zeros,
blk.attn.v_proj.w_q, blk.attn.v_proj.scales, blk.attn.v_proj.zeros,
blk.attn.g_proj.w_q, blk.attn.g_proj.scales, blk.attn.g_proj.zeros,
blk.attn.beta_proj.weight, blk.attn.conv_w,
blk.attn.o_proj.w_q, blk.attn.o_proj.scales, blk.attn.o_proj.zeros,
blk.moe.router.weight,
blk.moe.gate.w_q, blk.moe.gate.scales, blk.moe.gate.zeros,
blk.moe.up.w_q, blk.moe.up.scales, blk.moe.up.zeros,
blk.moe.down.w_q, blk.moe.down.scales, blk.moe.down.zeros,
blk.moe.s_gate.w_q, blk.moe.s_gate.scales, blk.moe.s_gate.zeros,
blk.moe.s_up.w_q, blk.moe.s_up.scales, blk.moe.s_up.zeros,
blk.moe.s_down.w_q, blk.moe.s_down.scales, blk.moe.s_down.zeros,
)
else:
self.mod.setup_mla_block(
self.params_tensor_cpu, b,
blk.attn_norm, blk.moe_norm,
blk.attn.q_proj.w_q, blk.attn.q_proj.scales, blk.attn.q_proj.zeros,
blk.attn.kv_a.w_q, blk.attn.kv_a.scales, blk.attn.kv_a.zeros,
blk.attn.kv_b.w_q, blk.attn.kv_b.scales, blk.attn.kv_b.zeros,
blk.attn.o_proj.w_q, blk.attn.o_proj.scales, blk.attn.o_proj.zeros,
blk.moe.router.weight,
blk.moe.gate.w_q, blk.moe.gate.scales, blk.moe.gate.zeros,
blk.moe.up.w_q, blk.moe.up.scales, blk.moe.up.zeros,
blk.moe.down.w_q, blk.moe.down.scales, blk.moe.down.zeros,
blk.moe.s_gate.w_q, blk.moe.s_gate.scales, blk.moe.s_gate.zeros,
blk.moe.s_up.w_q, blk.moe.s_up.scales, blk.moe.s_up.zeros,
blk.moe.s_down.w_q, blk.moe.s_down.scales, blk.moe.s_down.zeros,
)
if self.params_tensor_cuda is None or self.params_tensor_cuda.device != device:
self.params_tensor_cuda = self.params_tensor_cpu.to(device)
else:
self.params_tensor_cuda.copy_(self.params_tensor_cpu)
if self.scratch_tensor is None or self.scratch_tensor.device != device:
self.scratch_tensor = torch.zeros(self.scratch_size, dtype=torch.uint8, device=device)
self._synced_weights = True
def load_state_dict(self, state_dict, strict=True):
res = super().load_state_dict(state_dict, strict=strict)
self._synced_weights = False
return res
def step(self, hidden: torch.Tensor, state: list) -> tuple[torch.Tensor, list]:
if not self._synced_weights:
self._sync_weights(hidden.device)
pos = state[3]["c_kv"].shape[0]
new_c_kv = torch.empty((pos + 1, 512), dtype=torch.bfloat16, device=hidden.device)
new_k_rope = torch.empty((pos + 1, 64), dtype=torch.bfloat16, device=hidden.device)
out_hidden = torch.empty_like(hidden)
# EXACTLY ONE GPU KERNEL LAUNCH
self.mod.launch_kimi_decode_megakernel_opt(
hidden, out_hidden,
state[0]["S"], state[0]["cq"], state[0]["ck"], state[0]["cv"],
state[1]["S"], state[1]["cq"], state[1]["ck"], state[1]["cv"],
state[2]["S"], state[2]["cq"], state[2]["ck"], state[2]["cv"],
state[3]["c_kv"], state[3]["k_rope"],
new_c_kv, new_k_rope,
pos,
self.params_tensor_cuda, self.scratch_tensor
)
state[3]["c_kv"] = new_c_kv
state[3]["k_rope"] = new_k_rope
return out_hidden, state
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
20260902_231753_agy_gemini-3.8-flash-high_02_kimi_linear_decode