KernelBench mega · RTX PRO 6000
Kimi-Linear Decode GPT-6 Luna
2.48×geomean speedup across shapes
manually audited: clean
Luna wrote one cooperative CUDA launch for the four-layer Kimi-Linear decode. It unpacks int4 weights inside custom GEMVs, parallelizes MLA cache reads across context chunks, and keeps KDA, RoPE, routed and shared MoE, and state updates inside the launch. The final isolated regrade passed the full six-case checker.
harnesscodex
Kernel source (redacted)
from __future__ import annotations
import math
from typing import List
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_CPP = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime_api.h>
#include <vector>
constexpr int NW = 147;
constexpr int NS = 25;
constexpr int NT = 29;
struct Args {
void* w[NW];
void* s[NS];
void* t[NT];
int ctx;
};
extern "C" cudaError_t launch_coop(const Args*, cudaStream_t);
void launch_step(std::vector<torch::Tensor> w,
std::vector<torch::Tensor> s,
std::vector<torch::Tensor> t,
int64_t context) {
TORCH_CHECK(w.size() == NW && s.size() == NS && t.size() == NT,
"internal pointer table size mismatch");
Args a{};
for (int i = 0; i < NW; ++i) a.w[i] = w[i].data_ptr();
for (int i = 0; i < NS; ++i) a.s[i] = s[i].data_ptr();
for (int i = 0; i < NT; ++i) a.t[i] = t[i].data_ptr();
a.ctx = (int)context;
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaError_t err = launch_coop(&a, stream);
TORCH_CHECK(err == cudaSuccess, "cooperative launch failed: ", cudaGetErrorString(err));
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("run", &launch_step);
}
'''
_CUDA = r'''
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <vector>
#include <cmath>
namespace cg = cooperative_groups;
using bf = __nv_bfloat16;
constexpr int NW = 147;
constexpr int NS = 25;
constexpr int NT = 29;
constexpr int MAXOUT = 6144;
constexpr int PST = 8 * MAXOUT;
constexpr int BLOCKS = 188;
constexpr int THREADS = 128;
struct Args {
void* w[NW];
void* s[NS];
void* t[NT];
int ctx;
};
__device__ __forceinline__ float bfload(const bf* p, int i) {
return __bfloat162float(p[i]);
}
__device__ __forceinline__ void bfstore(bf* p, int i, float x) {
p[i] = __float2bfloat16_rn(x);
}
__device__ __forceinline__ float silu(float x) {
return x / (1.0f + __expf(-x));
}
__device__ __forceinline__ float softplus(float x) {
return fmaxf(x, 0.0f) + log1pf(__expf(-fabsf(x)));
}
__device__ __forceinline__ float qweight(const uint8_t* wq, const bf* scales,
const bf* zeros, int K, int N,
int k, int n) {
const uint8_t b = wq[(k >> 1) * N + n];
const int q = (k & 1) ? (b >> 4) : (b & 15);
const int gi = (k >> 7) * N + n;
const float z = bfload(zeros, gi);
const float sc = bfload(scales, gi);
const bf d = __float2bfloat16_rn((float)q - z);
const bf v = __float2bfloat16_rn(__bfloat162float(d) * sc);
return __bfloat162float(v);
}
__device__ __forceinline__ float xload(const void* x, int i, bool fp32) {
return fp32 ? ((const float*)x)[i] : bfload((const bf*)x, i);
}
__device__ __forceinline__ float gemv_one(const uint8_t* wq, const bf* scales,
const bf* zeros, const void* x,
int K, int N, int n, int split,
int split_id, bool xfp32) {
const int chunk = K / split;
const int ks = split_id * chunk;
const int ke = ks + chunk;
float acc = 0.0f;
for (int g = ks >> 7; g < (ke >> 7); ++g) {
const int r0 = g * 64;
const int r1 = r0 + 64;
const float z = bfload(zeros, g * N + n);
const float sc = bfload(scales, g * N + n);
for (int r = r0; r < r1; ++r) {
const uint8_t b = wq[r * N + n];
const bf d0 = __float2bfloat16_rn((float)(b & 15) - z);
const bf d1 = __float2bfloat16_rn((float)(b >> 4) - z);
const float w0 = __bfloat162float(__float2bfloat16_rn(__bfloat162float(d0) * sc));
const float w1 = __bfloat162float(__float2bfloat16_rn(__bfloat162float(d1) * sc));
acc = fmaf(xload(x, 2 * r, xfp32), w0, acc);
acc = fmaf(xload(x, 2 * r + 1, xfp32), w1, acc);
}
}
return acc;
}
__device__ __forceinline__ float warp_sum(float x) {
for (int d = 16; d; d >>= 1) x += __shfl_down_sync(0xffffffff, x, d);
return x;
}
__device__ __forceinline__ float block_sum(float x, float* sh) {
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
x = warp_sum(x);
if (lane == 0) sh[warp] = x;
__syncthreads();
if (warp == 0) {
x = lane < 4 ? sh[lane] : 0.0f;
x = warp_sum(x);
if (lane == 0) sh[0] = x;
}
__syncthreads();
return sh[0];
}
__device__ __forceinline__ float block_sum_root(float x, float* sh) {
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
x = warp_sum(x);
if (lane == 0) sh[warp] = x;
__syncthreads();
if (warp == 0) {
x = lane < 4 ? sh[lane] : 0.0f;
x = warp_sum(x);
if (lane == 0) sh[0] = x;
}
return threadIdx.x == 0 ? sh[0] : 0.0f;
}
__device__ __forceinline__ float block_max(float x, float* sh) {
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
for (int d = 16; d; d >>= 1) x = fmaxf(x, __shfl_down_sync(0xffffffff, x, d));
if (lane == 0) sh[warp] = x;
__syncthreads();
if (warp == 0) {
x = lane < 4 ? sh[lane] : -1.0e30f;
for (int d = 16; d; d >>= 1) x = fmaxf(x, __shfl_down_sync(0xffffffff, x, d));
if (lane == 0) sh[0] = x;
}
__syncthreads();
return sh[0];
}
__device__ void rms(const bf* x, const bf* weight, bf* y, float* sh) {
const int lane = threadIdx.x & 31;
float ss = 0.0f;
for (int i = lane; i < 2304; i += 32) {
const float v = bfload(x, i);
ss = fmaf(v, v, ss);
}
ss = warp_sum(ss);
ss = __shfl_sync(0xffffffff, ss, 0);
const float inv = rsqrtf(ss * (1.0f / 2304.0f) + 1.0e-6f);
for (int i = threadIdx.x; i < 2304; i += THREADS)
bfstore(y, i, bfload(x, i) * inv * bfload(weight, i));
}
__device__ void reduce_gemv(const float* p, int N, int tile, int splits, bf* out) {
const int n0 = tile * 256 + threadIdx.x;
const int n1 = n0 + 128;
if (n0 < N) {
float v = 0.0f;
for (int s = 0; s < splits; ++s) v += p[s * N + n0];
bfstore(out, n0, v);
}
if (n1 < N) {
float v = 0.0f;
for (int s = 0; s < splits; ++s) v += p[s * N + n1];
bfstore(out, n1, v);
}
}
__device__ void write_partial(const Args& a, int wbase, const void* x,
int K, int N, int tile, int splits, int split_id,
bool xfp32, float* partial, int pslot,
int expert, bool expert_bank) {
const int n0 = tile * 256 + threadIdx.x;
const int n1 = n0 + 128;
const int eszq = (K / 2) * N;
const int eszg = (K / 128) * N;
const uint8_t* wq = (const uint8_t*)a.w[wbase];
const bf* sc = (const bf*)a.w[wbase + 1];
const bf* ze = (const bf*)a.w[wbase + 2];
if (expert_bank) {
wq += (size_t)expert * eszq;
sc += (size_t)expert * eszg;
ze += (size_t)expert * eszg;
}
const int po = pslot * PST + split_id * N;
if (n0 < N) partial[po + n0] = gemv_one(wq, sc, ze, x, K, N, n0, splits, split_id, xfp32);
if (n1 < N) partial[po + n1] = gemv_one(wq, sc, ze, x, K, N, n1, splits, split_id, xfp32);
}
__device__ void router_one(const Args& a, int wb, const bf* x, bf* logits,
int e, float* sh) {
const bf* w = (const bf*)a.w[wb];
float sum = 0.0f;
for (int k = threadIdx.x; k < 2304; k += THREADS)
sum = fmaf(bfload(x, k), bfload(w, e * 2304 + k), sum);
sum = block_sum_root(sum, sh);
if (threadIdx.x == 0) bfstore(logits, e, sum);
}
__device__ void beta_one(const Args& a, int wb, const bf* x, bf* beta,
int e, float* sh) {
const bf* w = (const bf*)a.w[wb];
float sum = 0.0f;
for (int k = threadIdx.x; k < 2304; k += THREADS)
sum = fmaf(bfload(x, k), bfload(w, e * 2304 + k), sum);
sum = block_sum_root(sum, sh);
if (threadIdx.x == 0) bfstore(beta, e, sum);
}
__device__ void conv_one(const Args& a, int wbase, const bf* val,
const bf* old, bf* next, bf* out, int which, int tile) {
const bf* cw = (const bf*)a.w[wbase];
for (int z = 0; z < 2; ++z) {
const int c = tile * 256 + threadIdx.x + z * 128;
if (c >= 4096) continue;
const float z0 = bfload(old, 0 * 4096 + c);
const float z1 = bfload(old, 1 * 4096 + c);
const float z2 = bfload(old, 2 * 4096 + c);
const float z3 = bfload(val, c);
float acc = z0 * bfload(cw, which * 4096 * 4 + c * 4 + 0);
acc = fmaf(z1, bfload(cw, which * 4096 * 4 + c * 4 + 1), acc);
acc = fmaf(z2, bfload(cw, which * 4096 * 4 + c * 4 + 2), acc);
acc = fmaf(z3, bfload(cw, which * 4096 * 4 + c * 4 + 3), acc);
bfstore(out, c, silu(acc));
bfstore(next, c, z1);
bfstore(next, 4096 + c, z2);
bfstore(next, 8192 + c, z3);
}
}
__device__ void kda_state(const Args& a, int sb, const bf* q, const bf* k,
const bf* v, const bf* g, const bf* beta,
float* S, bf* out, int head) {
const int vi = threadIdx.x;
const float b = 1.0f / (1.0f + __expf(-bfload(beta, head)));
float pred = 0.0f;
for (int ki = 0; ki < 128; ++ki) {
const int si = (head * 128 + ki) * 128 + vi;
const float decay = __expf(-softplus(bfload(g, head * 128 + ki)));
pred = fmaf(S[si] * decay, bfload(k, head * 128 + ki), pred);
}
const float delta = bfload(v, head * 128 + vi) - pred;
float ov = 0.0f;
for (int ki = 0; ki < 128; ++ki) {
const int si = (head * 128 + ki) * 128 + vi;
const float kval = bfload(k, head * 128 + ki);
const float decay = __expf(-softplus(bfload(g, head * 128 + ki)));
const float sv = S[si] * decay + b * kval * delta;
S[si] = sv;
ov = fmaf(sv, bfload(q, head * 128 + ki) * 0.08838834764831845f, ov);
}
bfstore(out, head * 128 + vi, ov);
}
__device__ void choose_top8(const bf* logits, int32_t* ids, float* weights) {
bool used[64];
float selected[8];
for (int i = 0; i < 64; ++i) used[i] = false;
float mx = -1.0e30f;
for (int j = 0; j < 8; ++j) {
int bi = 0;
float bv = -1.0e30f;
for (int i = 0; i < 64; ++i) {
const float v = bfload(logits, i);
if (!used[i] && v > bv) { bv = v; bi = i; }
}
used[bi] = true;
ids[j] = bi;
selected[j] = bv;
mx = fmaxf(mx, bv);
}
float den = 0.0f;
for (int j = 0; j < 8; ++j) den += __expf(selected[j] - mx);
for (int j = 0; j < 8; ++j) weights[j] = __expf(selected[j] - mx) / den * 2.446f;
}
__device__ void expert_gate_up(const Args& a, int gb, int ub, const bf* x,
const int32_t* ids, float* gp, float* up,
int slot, int tile, int splits, int split) {
const int expert = slot == 8 ? 0 : ids[slot];
const int n0 = tile * 256 + threadIdx.x;
const int n1 = n0 + 128;
const int eszq = (2304 / 2) * 1024;
const int eszg = (2304 / 128) * 1024;
const uint8_t* gw = (const uint8_t*)a.w[gb] + (size_t)expert * eszq;
const bf* gs = (const bf*)a.w[gb + 1] + (size_t)expert * eszg;
const bf* gz = (const bf*)a.w[gb + 2] + (size_t)expert * eszg;
const uint8_t* uw = (const uint8_t*)a.w[ub] + (size_t)expert * eszq;
const bf* us = (const bf*)a.w[ub + 1] + (size_t)expert * eszg;
const bf* uz = (const bf*)a.w[ub + 2] + (size_t)expert * eszg;
const int po = (slot * splits + split) * 1024;
if (n0 < 1024) {
gp[po + n0] = gemv_one(gw, gs, gz, x, 2304, 1024, n0, splits, split, false);
up[po + n0] = gemv_one(uw, us, uz, x, 2304, 1024, n0, splits, split, false);
}
if (n1 < 1024) {
gp[po + n1] = gemv_one(gw, gs, gz, x, 2304, 1024, n1, splits, split, false);
up[po + n1] = gemv_one(uw, us, uz, x, 2304, 1024, n1, splits, split, false);
}
}
__device__ void expert_down(const Args& a, int db, const float* x,
const int32_t* ids, float* p, int slot,
int tile, int splits, int split) {
const int expert = slot == 8 ? 0 : ids[slot];
const int n0 = tile * 256 + threadIdx.x;
const int n1 = n0 + 128;
const int eszq = (1024 / 2) * 2304;
const int eszg = (1024 / 128) * 2304;
const uint8_t* wq = (const uint8_t*)a.w[db] + (size_t)expert * eszq;
const bf* sc = (const bf*)a.w[db + 1] + (size_t)expert * eszg;
const bf* ze = (const bf*)a.w[db + 2] + (size_t)expert * eszg;
const int po = (slot * splits + split) * 2304;
if (n0 < 2304) p[po + n0] = gemv_one(wq, sc, ze, x, 1024, 2304, n0, splits, split, true);
if (n1 < 2304) p[po + n1] = gemv_one(wq, sc, ze, x, 1024, 2304, n1, splits, split, true);
}
__device__ void rope_make(const Args& a, const bf* qfull, const bf* kv,
bf* qrot, bf* krot) {
const int pos = a.ctx;
for (int z = threadIdx.x; z < 32 * 32; z += THREADS) {
const int h = z / 32;
const int i = z & 31;
const float inv = expf(-logf(10000.0f) * ((2.0f * i) / 64.0f));
const float ang = (float)pos * inv;
const float c = cosf(ang), s = sinf(ang);
const float qe = bfload(qfull, h * 192 + 128 + 2 * i);
const float qo = bfload(qfull, h * 192 + 129 + 2 * i);
bfstore(qrot, h * 64 + 2 * i, qe * c - qo * s);
bfstore(qrot, h * 64 + 2 * i + 1, qe * s + qo * c);
}
if (threadIdx.x < 32) {
const int i = threadIdx.x;
const float inv = expf(-logf(10000.0f) * ((2.0f * i) / 64.0f));
const float ang = (float)pos * inv;
const float c = cosf(ang), s = sinf(ang);
const float ke = bfload(kv, 512 + 2 * i);
const float ko = bfload(kv, 513 + 2 * i);
bfstore(krot, 2 * i, ke * c - ko * s);
bfstore(krot, 2 * i + 1, ke * s + ko * c);
}
}
__device__ void rope_cache(const Args& a, const bf* kv, const bf* krot) {
const int oldN = a.ctx * 512;
bf* cn = (bf*)a.s[23];
const bf* co = (const bf*)a.s[21];
if (cn != co) {
for (int i = blockIdx.x * THREADS + threadIdx.x;
i < oldN; i += gridDim.x * THREADS) cn[i] = co[i];
}
if (blockIdx.x == 0) {
for (int i = threadIdx.x; i < 512; i += THREADS) cn[oldN + i] = kv[i];
}
bf* rn = (bf*)a.s[24];
const bf* ro = (const bf*)a.s[22];
if (rn != ro) {
for (int i = blockIdx.x * THREADS + threadIdx.x;
i < a.ctx * 64; i += gridDim.x * THREADS) rn[i] = ro[i];
}
if (blockIdx.x < 64 && threadIdx.x == 0)
rn[a.ctx * 64 + blockIdx.x] = krot[blockIdx.x];
}
__device__ void absorb_q(const Args& a, int wb, const bf* q, float* qlat,
int block) {
if (block >= 160) return;
const int h = block / 5;
const int split = block % 5;
const int begin = split * 104;
const int end = min(begin + 104, 512);
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const uint8_t* wq = (const uint8_t*)a.w[wb];
const bf* sc = (const bf*)a.w[wb + 1];
const bf* ze = (const bf*)a.w[wb + 2];
for (int kb = begin; kb < end; kb += 4) {
const int k = kb + warp;
if (k >= end) continue;
float v = 0.0f;
for (int j = lane; j < 128; j += 32) {
const int n = h * 256 + j;
const float w = qweight(wq, sc, ze, 512, 8192, k, n);
v = fmaf(w, bfload(q, h * 192 + j), v);
}
v = warp_sum(v);
if (lane == 0) qlat[h * 512 + k] = v;
}
}
__device__ void mla_score(const Args& a, const bf* qrot, const bf* kcache,
const bf* ccache, const float* qlat, float* scores,
int block, int parts) {
if (block >= 32 * parts) return;
const int h = block % 32;
const int part = block / 32;
const int L = a.ctx + 1;
const int chunk = (L + parts - 1) / parts;
const int begin = part * chunk;
const int end = min(begin + chunk, L);
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
for (int tb = begin; tb < end; tb += 4) {
const int t = tb + warp;
if (t >= end) continue;
float dot = 0.0f;
for (int k = 0; k < 512; k += 32) {
const float x = bfload(ccache, t * 512 + k + lane);
dot = fmaf(x, qlat[h * 512 + k + lane], dot);
}
dot = warp_sum(dot);
float rope = 0.0f;
for (int j = lane; j < 64; j += 32)
rope = fmaf(bfload(qrot, h * 64 + j), bfload(kcache, t * 64 + j), rope);
rope = warp_sum(rope);
if (lane == 0) scores[h * L + t] = (dot + rope) * 0.07216878364870322f;
}
}
__device__ void mla_softmax(const float* scores, float* probs, int L,
int head, float* sh) {
float mx = -1.0e30f;
for (int t = threadIdx.x; t < L; t += THREADS)
mx = fmaxf(mx, scores[head * L + t]);
mx = block_max(mx, sh);
float den = 0.0f;
for (int t = threadIdx.x; t < L; t += THREADS)
den += __expf(scores[head * L + t] - mx);
den = block_sum(den, sh);
for (int t = threadIdx.x; t < L; t += THREADS)
probs[head * L + t] = __expf(scores[head * L + t] - mx) / den;
__syncthreads();
}
__device__ void latent_reduce_part(const bf* cache, const float* probs,
float* partial, int L, int h,
int part, int parts) {
float acc[4] = {0.f, 0.f, 0.f, 0.f};
const int chunk = (L + parts - 1) / parts;
const int begin = part * chunk;
const int end = min(begin + chunk, L);
for (int t = begin; t < end; ++t) {
const float p = probs[h * L + t];
const int k0 = threadIdx.x;
const int k1 = k0 + 128;
const int k2 = k0 + 256;
const int k3 = k0 + 384;
acc[0] = fmaf(p, bfload(cache, t * 512 + k0), acc[0]);
acc[1] = fmaf(p, bfload(cache, t * 512 + k1), acc[1]);
acc[2] = fmaf(p, bfload(cache, t * 512 + k2), acc[2]);
acc[3] = fmaf(p, bfload(cache, t * 512 + k3), acc[3]);
}
const int base = (h * parts + part) * 512 + threadIdx.x;
partial[base] = acc[0];
partial[base + 128] = acc[1];
partial[base + 256] = acc[2];
partial[base + 384] = acc[3];
}
__device__ void latent_merge(const float* partial, float* latent,
int h, int parts) {
for (int z = 0; z < 4; ++z) {
const int k = threadIdx.x + z * 128;
float sum = 0.0f;
for (int p = 0; p < parts; ++p)
sum += partial[(h * parts + p) * 512 + k];
latent[h * 512 + k] = sum;
}
}
__device__ void v_project(const Args& a, int wb, const float* latent,
bf* out, int h) {
const uint8_t* wq = (const uint8_t*)a.w[wb];
const bf* sc = (const bf*)a.w[wb + 1];
const bf* ze = (const bf*)a.w[wb + 2];
const int n = threadIdx.x;
float acc = 0.0f;
const int col = h * 256 + 128 + n;
for (int k = 0; k < 512; ++k)
acc = fmaf(latent[h * 512 + k], qweight(wq, sc, ze, 512, 8192, k, col), acc);
bfstore(out, h * 128 + n, acc);
}
__global__ void decode_kernel(Args a) {
cg::grid_group grid = cg::this_grid();
__shared__ float red[128];
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int mla_parts = gridDim.x >= 320 ? 10 : 5;
bf* norm = (bf*)a.t[0];
bf* hres = (bf*)a.t[1];
bf* hA = (bf*)a.t[2];
bf* hB = (bf*)a.t[26];
bf* q = (bf*)a.t[3];
bf* kbuf = (bf*)a.t[4];
bf* vbuf = (bf*)a.t[5];
bf* gbuf = (bf*)a.t[6];
bf* beta = (bf*)a.t[7];
bf* attin = (bf*)a.t[8];
bf* attout = (bf*)a.t[9];
float* pp = (float*)a.t[10];
float* pp2 = (float*)a.t[11];
bf* logits = (bf*)a.t[12];
int32_t* ids = (int32_t*)a.t[13];
float* routew = (float*)a.t[14];
float* gatep = (float*)a.t[15];
float* upp = (float*)a.t[16];
float* acts = (float*)a.t[17];
float* downp = (float*)a.t[18];
bf* moeout = (bf*)a.t[19];
float* qlat = (float*)a.t[20];
float* scores = (float*)a.t[21];
float* probs = (float*)a.t[22];
float* latout = (float*)a.t[23];
float* latpart = (float*)a.t[28];
bf* qrot = (bf*)a.t[24];
bf* krot = (bf*)a.t[25];
const bf* cur = (const bf*)a.t[27];
// t[27] aliases the original token and is replaced by the host before launch.
for (int layer = 0; layer < 4; ++layer) {
const int wb = layer < 3 ? layer * 38 : 114;
const int is_kda = layer < 3;
bf* next = (layer & 1) ? hB : hA;
if (bid == 0) rms(cur, (const bf*)a.w[wb + 0], norm, red);
grid.sync();
if (is_kda) {
constexpr int S2304 = 6;
constexpr int S4096 = 8;
constexpr int S1024 = 4;
constexpr int QKV_BLOCKS = 4 * 16 * S2304;
if (bid < QKV_BLOCKS) {
const int p = bid / (16 * S2304);
const int within = bid % (16 * S2304);
const int tile = within / S2304;
const int split = within % S2304;
write_partial(a, wb + 2 + p * 3, norm, 2304, 4096, tile,
S2304, split, false, pp, p, 0, false);
}
if (bid >= QKV_BLOCKS && bid < QKV_BLOCKS + 32)
beta_one(a, wb + 14, norm, beta, bid - QKV_BLOCKS, red);
grid.sync();
if (bid < 64) {
const int p = bid / 16;
const int tile = bid & 15;
bf* out = p == 0 ? q : (p == 1 ? kbuf : (p == 2 ? vbuf : gbuf));
reduce_gemv(pp + p * PST, 4096, tile, S2304, out);
}
grid.sync();
if (bid < 48) {
const int which = bid % 3;
const int tile = bid / 3;
const int layer_s = layer * 7;
const bf* val = which == 0 ? q : (which == 1 ? kbuf : vbuf);
bf* out = which == 0 ? q : (which == 1 ? kbuf : vbuf);
const bf* old = (const bf*)a.s[layer_s + 1 + which];
bf* win = (bf*)a.s[layer_s + 4 + which];
conv_one(a, wb + 15, val, old, win, out, which, tile);
}
grid.sync();
if (bid < 32)
kda_state(a, layer * 7, q, kbuf, vbuf, gbuf, beta,
(float*)a.s[layer * 7], attin, bid);
grid.sync();
if (bid < 9 * S4096) {
const int tile = bid / S4096, split = bid % S4096;
write_partial(a, wb + 16, attin, 4096, 2304, tile,
S4096, split, false, pp2, 0, 0, false);
}
grid.sync();
if (bid < 9) reduce_gemv(pp2, 2304, bid, S4096, attout);
grid.sync();
} else {
constexpr int S2304 = 6;
constexpr int S4096 = 8;
constexpr int Q_BLOCKS = 24 * S2304;
constexpr int KV_BLOCKS = 3 * S2304;
if (bid < Q_BLOCKS + KV_BLOCKS) {
int p, tile, split;
if (bid < Q_BLOCKS) { p = 0; tile = bid / S2304; split = bid % S2304; }
else { int z = bid - Q_BLOCKS; p = 1; tile = z / S2304; split = z % S2304; }
const int N = p == 0 ? 6144 : 576;
write_partial(a, wb + (p == 0 ? 2 : 5), norm, 2304, N,
tile, S2304, split, false, pp, p, 0, false);
}
grid.sync();
if (bid < 24) reduce_gemv(pp, 6144, bid, S2304, q);
if (bid >= 24 && bid < 27) reduce_gemv(pp + PST, 576, bid - 24, S2304, kbuf);
grid.sync();
if (bid == 0) rope_make(a, q, kbuf, qrot, krot);
grid.sync();
rope_cache(a, kbuf, krot);
grid.sync();
if (bid < 160) absorb_q(a, wb + 8, q, qlat, bid);
grid.sync();
if (bid < 32 * mla_parts) mla_score(a, qrot, (const bf*)a.s[24],
(const bf*)a.s[23], qlat, scores, bid, mla_parts);
grid.sync();
if (bid < 32) mla_softmax(scores, probs, a.ctx + 1, bid, red);
grid.sync();
if (bid < 32 * mla_parts)
latent_reduce_part((const bf*)a.s[23], probs, latpart, a.ctx + 1,
bid % 32, bid / 32, mla_parts);
grid.sync();
if (bid < 32) latent_merge(latpart, latout, bid, mla_parts);
grid.sync();
if (bid < 32) v_project(a, wb + 8, latout, attin, bid);
grid.sync();
if (bid < 9 * S4096) {
const int tile = bid / S4096, split = bid % S4096;
write_partial(a, wb + 11, attin, 4096, 2304, tile,
S4096, split, false, pp2, 0, 0, false);
}
grid.sync();
if (bid < 9) reduce_gemv(pp2, 2304, bid, S4096, attout);
grid.sync();
}
if (bid < 9) {
for (int z = 0; z < 2; ++z) {
const int i = bid * 256 + tid + z * 128;
if (i < 2304) bfstore(hres, i, bfload(cur, i) + bfload(attout, i));
}
}
grid.sync();
if (bid == 0) rms(hres, (const bf*)a.w[wb + 1], norm, red);
grid.sync();
if (bid < 64) router_one(a, wb + (is_kda ? 19 : 14), norm, logits, bid, red);
grid.sync();
if (bid == 0 && tid == 0) choose_top8(logits, ids, routew);
grid.sync();
const int gatebase = wb + (is_kda ? 20 : 15);
const int upbase = wb + (is_kda ? 23 : 18);
const int downbase = wb + (is_kda ? 26 : 21);
const int sgatebase = wb + (is_kda ? 29 : 24);
const int supbase = wb + (is_kda ? 32 : 27);
const int sdownbase = wb + (is_kda ? 35 : 30);
constexpr int S2304 = 6;
constexpr int S4096 = 4;
constexpr int S1024 = 4;
constexpr int GATE_BLOCKS = 9 * 4 * S2304;
constexpr int DOWN_BLOCKS = 9 * 9 * S1024;
if (bid < GATE_BLOCKS) {
const int slot = bid / (4 * S2304);
const int rem = bid % (4 * S2304);
const int tile = rem / S2304;
const int split = rem % S2304;
expert_gate_up(a, slot == 8 ? sgatebase : gatebase,
slot == 8 ? supbase : upbase, norm, ids,
gatep, upp, slot, tile, S2304, split);
}
grid.sync();
if (bid < 36) {
const int slot = bid / 4;
const int tile = bid & 3;
for (int z = 0; z < 2; ++z) {
const int n = tile * 256 + tid + z * 128;
if (n >= 1024) continue;
float gv = 0.0f, uv = 0.0f;
for (int s = 0; s < S2304; ++s) {
gv += gatep[(slot * S2304 + s) * 1024 + n];
uv += upp[(slot * S2304 + s) * 1024 + n];
}
acts[slot * 1024 + n] = silu(gv) * uv;
}
}
grid.sync();
if (bid < DOWN_BLOCKS) {
const int slot = bid / (9 * S1024);
const int rem = bid % (9 * S1024);
const int tile = rem / S1024;
const int split = rem % S1024;
const int db = slot == 8 ? sdownbase : downbase;
expert_down(a, db, acts + slot * 1024, ids, downp, slot, tile, S1024, split);
}
grid.sync();
if (bid < 9) {
const int base = bid * 256 + tid;
for (int z = 0; z < 2; ++z) {
const int n = base + z * 128;
if (n >= 2304) continue;
float sum = 0.0f;
for (int e = 0; e < 8; ++e) {
float d = 0.0f;
for (int s = 0; s < S1024; ++s) d += downp[(e * S1024 + s) * 2304 + n];
sum = fmaf(routew[e], d, sum);
}
float shared = 0.0f;
for (int s = 0; s < S1024; ++s) shared += downp[(8 * S1024 + s) * 2304 + n];
bfstore(moeout, n, sum + shared);
}
}
grid.sync();
if (bid < 9) {
for (int z = 0; z < 2; ++z) {
const int i = bid * 256 + tid + z * 128;
if (i < 2304) bfstore(next, i, bfload(hres, i) + bfload(moeout, i));
}
}
grid.sync();
cur = next;
}
}
extern "C" cudaError_t launch_coop(const Args* a, cudaStream_t stream) {
static int cached_grid = 0;
if (cached_grid == 0) {
int active = 0, sms = 0, dev = 0;
cudaError_t e = cudaGetDevice(&dev);
if (e != cudaSuccess) return e;
e = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
if (e != cudaSuccess) return e;
e = cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, decode_kernel, THREADS, 0);
if (e != cudaSuccess) return e;
cached_grid = active * sms;
if (cached_grid > 448) cached_grid = 448;
if (cached_grid < BLOCKS) cached_grid = BLOCKS;
}
void* args[] = {const_cast<Args*>(a)};
return cudaLaunchCooperativeKernel((void*)decode_kernel,
dim3(cached_grid), dim3(THREADS), args, 0, stream);
}
'''
class QuantLinear(nn.Module):
def __init__(self, in_f: int, out_f: int, group: int = 128):
super().__init__()
self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
self.register_buffer("scales", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16))
self.register_buffer("zeros", torch.zeros(in_f // group, out_f, dtype=torch.bfloat16))
class QuantExperts(nn.Module):
def __init__(self, n: int, in_f: int, out_f: int, group: int = 128):
super().__init__()
self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8))
self.register_buffer("scales", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))
self.register_buffer("zeros", torch.zeros(n, in_f // group, out_f, dtype=torch.bfloat16))
class KDA(nn.Module):
def __init__(self, cfg):
super().__init__()
d, c = cfg.hidden, cfg.kda_heads * cfg.kda_head_dim
self.q_proj = QuantLinear(d, c)
self.k_proj = QuantLinear(d, c)
self.v_proj = QuantLinear(d, c)
self.g_proj = QuantLinear(d, c)
self.beta_proj = nn.Linear(d, cfg.kda_heads, bias=False, dtype=cfg.dtype)
self.conv_w = nn.Parameter(torch.empty(3, c, cfg.short_conv, dtype=cfg.dtype))
self.o_proj = QuantLinear(c, d)
class MLA(nn.Module):
def __init__(self, cfg):
super().__init__()
self.q_proj = QuantLinear(cfg.hidden, cfg.mla_heads * (cfg.qk_nope + cfg.qk_rope))
self.kv_a = QuantLinear(cfg.hidden, cfg.kv_lora + cfg.qk_rope)
self.kv_b = QuantLinear(cfg.kv_lora, cfg.mla_heads * (cfg.qk_nope + cfg.v_head))
self.o_proj = QuantLinear(cfg.mla_heads * cfg.v_head, cfg.hidden)
class MoE(nn.Module):
def __init__(self, cfg):
super().__init__()
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)
self.up = QuantExperts(e, d, m)
self.down = QuantExperts(e, m, d)
self.s_gate = QuantExperts(cfg.n_shared, d, m)
self.s_up = QuantExperts(cfg.n_shared, d, m)
self.s_down = QuantExperts(cfg.n_shared, m, d)
class Block(nn.Module):
def __init__(self, cfg, kind):
super().__init__()
self.kind = kind
self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
self.moe = MoE(cfg)
def _push_quant(dst, q):
dst.extend((q.w_q, q.scales, q.zeros))
def _block_weights(block):
out = [block.attn_norm, block.moe_norm]
if block.kind == "K":
a = block.attn
for q in (a.q_proj, a.k_proj, a.v_proj, a.g_proj):
_push_quant(out, q)
out.extend((a.beta_proj.weight, a.conv_w))
_push_quant(out, a.o_proj)
else:
a = block.attn
for q in (a.q_proj, a.kv_a, a.kv_b, a.o_proj):
_push_quant(out, q)
m = block.moe
out.append(m.router.weight)
for q in (m.gate, m.up, m.down, m.s_gate, m.s_up, m.s_down):
_push_quant(out, q)
return out
def _state_buffers(state, device, cache_c, cache_r):
ptrs, outputs = [], []
for layer in range(4):
st = state[layer]
if layer < 3:
ptrs.extend((st["S"], st["cq"], st["ck"], st["cv"]))
news = (st["cq"], st["ck"], st["cv"])
ptrs.extend(news)
outputs.append(news)
else:
ptrs.extend((st["c_kv"], st["k_rope"]))
ptrs.extend((cache_c, cache_r))
length = st["c_kv"].shape[0] + 1
outputs.append((cache_c.narrow(0, 0, length), cache_r.narrow(0, 0, length)))
return ptrs, outputs
def _scratch(device, context, hidden):
b = torch.bfloat16
f = torch.float32
return [
torch.empty((2304,), device=device, dtype=b),
torch.empty((2304,), device=device, dtype=b),
torch.empty((2304,), device=device, dtype=b),
torch.empty((6144,), device=device, dtype=b),
torch.empty((6144,), device=device, dtype=b),
torch.empty((4096,), device=device, dtype=b),
torch.empty((4096,), device=device, dtype=b),
torch.empty((32,), device=device, dtype=b),
torch.empty((4096,), device=device, dtype=b),
torch.empty((2304,), device=device, dtype=b),
torch.empty((4 * 8 * 6144,), device=device, dtype=f),
torch.empty((8 * 6144,), device=device, dtype=f),
torch.empty((64,), device=device, dtype=b),
torch.empty((8,), device=device, dtype=torch.int32),
torch.empty((8,), device=device, dtype=f),
torch.empty((9 * 6 * 1024,), device=device, dtype=f),
torch.empty((9 * 6 * 1024,), device=device, dtype=f),
torch.empty((9 * 1024,), device=device, dtype=f),
torch.empty((9 * 4 * 2304,), device=device, dtype=f),
torch.empty((2304,), device=device, dtype=b),
torch.empty((32 * 512,), device=device, dtype=f),
torch.empty((32 * (context + 1),), device=device, dtype=f),
torch.empty((32 * (context + 1),), device=device, dtype=f),
torch.empty((32 * 512,), device=device, dtype=f),
torch.empty((32 * 64,), device=device, dtype=b),
torch.empty((64,), device=device, dtype=b),
torch.empty((2304,), device=device, dtype=b),
hidden,
torch.empty((32 * 10 * 512,), device=device, dtype=f),
]
class Model(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
self._kernel = None
self._weight_refs = None
self._weight_device = None
self._cache_c = None
self._cache_r = None
def step(self, hidden, state):
if self._kernel is None:
self._kernel = load_inline(
name="kimi_linear_w4a16_megakernel",
cpp_sources=_CPP,
cuda_sources=_CUDA,
functions=None,
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3", "-std=c++17", "-arch=sm_120"],
with_cuda=True,
no_implicit_headers=True,
verbose=False,
)
if self._weight_refs is None or self._weight_device != hidden.device:
weights = []
for block in self.blocks:
weights.extend(_block_weights(block))
self._weight_refs = weights
self._weight_device = hidden.device
oldc = state[3]["c_kv"]
oldr = state[3]["k_rope"]
oldlen = oldc.shape[0]
reusable = (
self._cache_c is not None
and self._cache_c.device == hidden.device
and self._cache_c.data_ptr() == oldc.data_ptr()
and self._cache_r.data_ptr() == oldr.data_ptr()
and self._cache_c.shape[0] >= oldlen + 1
)
if not reusable:
capacity = max(oldlen + 1, 2 * oldlen)
self._cache_c = torch.empty((capacity, 512), device=hidden.device, dtype=oldc.dtype)
self._cache_r = torch.empty((capacity, 64), device=hidden.device, dtype=oldr.dtype)
state_ptrs, state_out = _state_buffers(state, hidden.device, self._cache_c, self._cache_r)
temps = _scratch(hidden.device, oldlen, hidden)
self._kernel.run(self._weight_refs, state_ptrs, temps, int(oldlen))
for i in range(3):
state[i]["cq"], state[i]["ck"], state[i]["cv"] = state_out[i]
state[3]["c_kv"], state[3]["k_rope"] = state_out[3]
return temps[26], state
20260922_123223_codex_gpt-6-luna_02_kimi_linear_decode