KernelBench mega · RTX PRO 6000
Kimi-Linear Decode Claude Fable 5
24.61×geomean speedup across shapes
manually audited: clean
RTX PRO 6000 (rented PRO box) fable-5 cell, 24.61x geomean, graded by clean sequential isolated re-run (rebench_check.log PASS, cosine 0.9999/1.0000 on out/S/cache across seeds 0-2 x ctx 2048/8192). Genuine single-launch cooperative megakernel: 33 grid.sync-separated phases in one cudaLaunchCooperativeKernel per step() -- int4 dequant-GEMV (fp16 magic-number nibble expansion, zero-point correction via group sums), KDA short-conv + delta-rule state update, absorbed-latent MLA over the bf16 c_kv cache with online softmax, MoE router + top-8 + expert GEMVs accumulated atomically into residuals, all RMSNorms in-kernel, plus L2 prefetch warming from idle CTAs. No CUDA graphs, no torch.compile, no Triton, no forbidden libraries, no sidecar modules (only imports: os, torch, load_inline). Workspace check.py/benchmark.py/reference.py/shapes.py/problem.yaml/PROMPT.txt are byte-identical to the repo deck. Empirical recompute test on the live PRO-box workspace passed all cases.
harnessor-fable
Kernel source (redacted)
"""Single-kernel megakernel decode for the Kimi-Linear W4A16 hybrid unit.
The whole per-token forward -- 3x(KDA + MoE) + (MLA + MoE), every int4
dequant-GEMV, short conv, KDA state update, MLA latent attention, MoE router +
expert FFNs, all RMSNorms and residuals -- runs as ONE cooperative CUDA kernel
launch per step(). Phases inside the kernel are separated by grid.sync()
(cooperative groups), so sequential dependencies are honored without extra
launches.
Design notes:
* int4 GEMV: weights stay packed in global memory. Each (128-col tile,
k-chunk) unit is one CTA pass: 16 warps each take 4 consecutive packed rows
per 128-row quant group, lanes take 4 consecutive columns (uint32 loads,
fully coalesced). Nibbles are expanded 2-at-a-time with the fp16 magic
trick ((w>>s & 0x000F000F) | 0x64006400 -> half2(1024+n), then -1024) and
accumulated in half2 against a /8-prescaled activation vector staged in
shared memory; per-group fold applies scale, and warp 0 alone applies the
-scale*zero*groupsum(x) correction. Cross-warp reduce through smem, then
store or fp32 atomicAdd (split-K / expert accumulation directly into the
residual-preloaded destination -- no zeroing or extra reduce phases).
* MLA uses the absorbed-latent form: q_nope is folded through kv_b once
(q_abs = q . W_knope per latent channel), attention runs directly on the
bf16 c_kv cache, and the context vector is folded through the W_v half of
kv_b afterwards -- the cache is never expanded to per-head K/V.
* The MLA cache lives in a capacity buffer owned by the model; step() returns
views [:L+1], so cat() disappears. The first step of a fresh state copies
the fed cache into the buffer inside the same kernel (phase Z).
* KDA conv windows double-buffer (ping-pong) between two internal banks so
the shifted window can be written while other CTAs still read the old one.
* MoE: the router runs as two phases -- 16 CTAs compute the 64 expert
logits in parallel (4 dot products each), then one warp does softmax +
top-8; gate/up/down expert GEMVs read the selected ids from scratch.
* Idle CTAs in compute-light phases issue prefetch.global.L2::evict_last
on weights used a phase or two later (shared experts, next attention
layer, the growing c_kv cache), hiding DRAM latency behind grid syncs.
Warming only pays in bandwidth-idle windows; in streaming phases it
competes with demand reads and was measured to hurt.
A slow eager path is NOT kept around: this file's step() always uses the fused
kernel. Debugging happens via the `debug_max_phase` attribute (early-exit after
N grid barriers) used by scratch tests during development.
"""
from __future__ import annotations
import os
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
EPS = 1.0e-6
# ---------------------------------------------------------------------------
# CUDA source: one cooperative kernel, 33 grid.sync-separated phases.
# ---------------------------------------------------------------------------
_CUDA_SRC = r"""
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#include <vector>
namespace cg = cooperative_groups;
#define D_HID 2304
#define C_KDA 4096
#define NH 32
#define DK 128
#define M_INTER 1024
#define HQ_MLA 6144
#define KVL 512
#define QR 64
#define KVA 576
#define NKVB 8192
#define NBLK 512
#define KDA_SCALE 0.08838834764831845f
#define MLA_SCALE 0.07216878364870323f
#define SMEM_BYTES 81920
typedef __nv_bfloat16 bf16;
typedef __nv_bfloat162 bf162;
struct KWeights {
const bf16* attn_norm[4];
const bf16* moe_norm[4];
const uint8_t* kda_wq[3][5]; // q,k,v,g,o
const bf16* kda_sc[3][5];
const bf16* kda_zr[3][5];
const bf16* beta_w[3]; // [32,2304]
const bf16* conv_w[3]; // [3,4096,4]
const uint8_t* mla_wq[4]; // q, kv_a, kv_b, o
const bf16* mla_sc[4];
const bf16* mla_zr[4];
const bf16* router[4]; // [64,2304]
const uint8_t* moe_wq[4][6]; // gate,up,down,s_gate,s_up,s_down
const bf16* moe_sc[4][6];
const bf16* moe_zr[4][6];
};
struct KScratch {
float* x; // [5][D]
float* hattn; // [4][D]
float* qkvg; // [4*C]
float* betaout; // [32]
float* obuf; // [4096]
__half* xn_moe; // [D]
int* ids; // [8]
float* wts; // [8]
float* logits; // [64]
float* hg; // [9*1024]
float* hu; // [9*1024]
float* q6144; // [6144]
float* kva; // [576]
__half* qabs; // [512*32] (c-major: qabs[c*32+h])
__half* qr; // [64*32]
float* scores; // [LCAP*32]
float* mx; // [NCHMAX*32]
float* sm; // [NCHMAX*32]
float* ctx; // [32*512]
bf16* outb; // [D]
bf16* ckv; // [LCAP*512]
bf16* krope; // [LCAP*64]
};
struct KStep {
const bf16* hidden;
float* S[3];
const bf16* winr[9]; // [layer*3+kind] -> [3][4096]
bf16* winw[9];
const bf16* ckv_src;
const bf16* kr_src;
int copy_len;
int L;
int max_phase;
};
struct KArg {
KWeights w;
KScratch s;
KStep a;
float routed;
float theta;
};
// ---------------------------------------------------------------- int4 GEMV
__device__ __forceinline__ __half2 deq2(uint32_t w, int sh) {
uint32_t v = ((w >> sh) & 0x000F000Fu) | 0x64006400u;
__half2 h = *reinterpret_cast<__half2*>(&v);
return __hsub2(h, __float2half2_rn(1024.f));
}
// y[col0+t] (for t < ncols) = 8 * sum_k xs[k] * deq(w[k, col0+t]) over groups
// [g0, g0+ng). xs is the FULL-K activation staged in smem, prescaled by 1/8.
// gsum[g] = sum of xs over group g. smem_red needs 16*128 floats.
__device__ void gemv_i4(const uint8_t* __restrict__ wq,
const bf16* __restrict__ sc,
const bf16* __restrict__ zr,
int N, int g0, int ng, int col0, int ncols,
const __half* __restrict__ xs,
const float* __restrict__ gsum,
float* __restrict__ out, float wmul, bool atomic,
float* smem_red) {
const int lane = threadIdx.x & 31;
const int w = threadIdx.x >> 5;
const int c4 = lane * 4;
const bool active = (c4 < ncols);
float y0 = 0.f, y1 = 0.f, y2 = 0.f, y3 = 0.f;
if (active) {
#pragma unroll 2
for (int g = g0; g < g0 + ng; ++g) {
const uint8_t* rp = wq + (size_t)(g * 64 + w * 4) * N + col0 + c4;
const __half2* xp = reinterpret_cast<const __half2*>(xs + (g * 64 + w * 4) * 2);
__half2 a02 = __float2half2_rn(0.f);
__half2 a13 = __float2half2_rn(0.f);
#pragma unroll
for (int r = 0; r < 4; ++r) {
uint32_t wrd = *reinterpret_cast<const uint32_t*>(rp);
rp += N;
__half2 xpair = xp[r];
__half2 xlo = __half2half2(__low2half(xpair));
__half2 xhi = __half2half2(__high2half(xpair));
a02 = __hfma2(deq2(wrd, 0), xlo, a02);
a13 = __hfma2(deq2(wrd, 8), xlo, a13);
a02 = __hfma2(deq2(wrd, 4), xhi, a02);
a13 = __hfma2(deq2(wrd, 12), xhi, a13);
}
float2 f02 = __half22float2(a02);
float2 f13 = __half22float2(a13);
const bf16* sp = sc + (size_t)g * N + col0 + c4;
const bf16* zp = zr + (size_t)g * N + col0 + c4;
bf162 sA = *reinterpret_cast<const bf162*>(sp);
bf162 sB = *reinterpret_cast<const bf162*>(sp + 2);
bf162 zA = *reinterpret_cast<const bf162*>(zp);
bf162 zB = *reinterpret_cast<const bf162*>(zp + 2);
float xg = (w == 0) ? gsum[g] : 0.f;
y0 += __bfloat162float(sA.x) * (f02.x - __bfloat162float(zA.x) * xg);
y1 += __bfloat162float(sA.y) * (f13.x - __bfloat162float(zA.y) * xg);
y2 += __bfloat162float(sB.x) * (f02.y - __bfloat162float(zB.x) * xg);
y3 += __bfloat162float(sB.y) * (f13.y - __bfloat162float(zB.y) * xg);
}
}
__syncthreads();
if (active) {
smem_red[w * 128 + c4 + 0] = y0;
smem_red[w * 128 + c4 + 1] = y1;
smem_red[w * 128 + c4 + 2] = y2;
smem_red[w * 128 + c4 + 3] = y3;
}
__syncthreads();
const int t = threadIdx.x;
if (t < ncols) {
float acc = 0.f;
#pragma unroll
for (int i = 0; i < 16; ++i) acc += smem_red[i * 128 + t];
acc = acc * 8.f * wmul;
if (atomic) atomicAdd(out + t, acc);
else out[t] = acc;
}
}
// ------------------------------------------------------------- x preparation
__device__ float block_sum(float v, float* red) {
for (int o = 16; o; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
if ((threadIdx.x & 31) == 0) red[threadIdx.x >> 5] = v;
__syncthreads();
if (threadIdx.x == 0) {
float tot = 0.f;
for (int i = 0; i < 16; ++i) tot += red[i];
red[16] = tot;
}
__syncthreads();
return red[16];
}
__device__ void fill_gsum(const __half* xs, int n, float* gsum) {
const int w = threadIdx.x >> 5, lane = threadIdx.x & 31;
for (int g = w; g < n / 128; g += 16) {
float s = 0.f;
for (int j = lane; j < 128; j += 32) s += __half2float(xs[g * 128 + j]);
for (int o = 16; o; o >>= 1) s += __shfl_down_sync(0xffffffffu, s, o);
if (lane == 0) gsum[g] = s;
}
__syncthreads();
}
// rmsnorm(src)*normw -> xs (half, /8), gsum
__device__ void prep_norm(const float* src, const bf16* normw, int n,
__half* xs, float* gsum, float* red) {
float ss = 0.f;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
float v = src[i];
ss += v * v;
}
float tot = block_sum(ss, red);
float scale = rsqrtf(tot / (float)n + 1.0e-6f);
for (int i = threadIdx.x; i < n; i += blockDim.x) {
float v = src[i] * scale * __bfloat162float(normw[i]);
v = __bfloat162float(__float2bfloat16(v)); // reference emits bf16
xs[i] = __float2half(v * 0.125f);
}
__syncthreads();
fill_gsum(xs, n, gsum);
}
// bf16-round(src [+ src2]) -> xs (half, /8), gsum
__device__ void prep_bf16(const float* src, const float* src2, int n,
__half* xs, float* gsum) {
for (int i = threadIdx.x; i < n; i += blockDim.x) {
float v = src[i] + (src2 ? src2[i] : 0.f);
v = __bfloat162float(__float2bfloat16(v));
xs[i] = __float2half(v * 0.125f);
}
__syncthreads();
fill_gsum(xs, n, gsum);
}
// -------------------------------------------------------------------- phases
__device__ void phase_Z(const KArg& A) {
const int T = blockIdx.x * blockDim.x + threadIdx.x;
const int NT = gridDim.x * blockDim.x;
for (int i = T; i < D_HID; i += NT) A.s.x[i] = __bfloat162float(A.a.hidden[i]);
for (int i = T; i < NH * KVL; i += NT) A.s.ctx[i] = 0.f;
if (A.a.copy_len > 0) {
const uint4* s1 = reinterpret_cast<const uint4*>(A.a.ckv_src);
uint4* d1 = reinterpret_cast<uint4*>(A.s.ckv);
int n1 = A.a.copy_len * 64; // 512 bf16 = 64 uint4 per row
for (int i = T; i < n1; i += NT) d1[i] = s1[i];
const uint4* s2 = reinterpret_cast<const uint4*>(A.a.kr_src);
uint4* d2 = reinterpret_cast<uint4*>(A.s.krope);
int n2 = A.a.copy_len * 8;
for (int i = T; i < n2; i += NT) d2[i] = s2[i];
}
}
__device__ void l2_warm(const uint8_t* const* bases, const size_t* sizes,
int nreg, int pu, int np) {
size_t total = 0;
for (int i = 0; i < nreg; ++i) total += sizes[i];
size_t b0 = total * pu / np, b1 = total * (pu + 1) / np;
for (size_t off = b0 + (size_t)threadIdx.x * 128; off < b1;
off += (size_t)blockDim.x * 128) {
size_t o = off;
for (int i = 0; i < nreg; ++i) {
if (o < sizes[i]) {
asm volatile("prefetch.global.L2::evict_last [%0];" :: "l"(bases[i] + o));
break;
}
o -= sizes[i];
}
}
}
__device__ void kda_phaseA(const KArg& A, int layer, char* SM) {
__half* xs = (__half*)SM; // 2304 half
float* gsum = (float*)(SM + 8192); // 18
float* red = (float*)(SM + 8320); // 17
float* ysred = (float*)(SM + 8448); // 16*128
prep_norm(A.s.x + layer * D_HID, A.w.attn_norm[layer], D_HID, xs, gsum, red);
for (int u = blockIdx.x; u < 133; u += gridDim.x) {
if (u < 128) {
int mat = u >> 5, tile = u & 31;
gemv_i4(A.w.kda_wq[layer][mat], A.w.kda_sc[layer][mat], A.w.kda_zr[layer][mat],
C_KDA, 0, 18, tile * 128, 128, xs, gsum,
A.s.qkvg + mat * C_KDA + tile * 128, 1.f, false, ysred);
} else if (u == 128) {
int hh = threadIdx.x >> 4;
int ss = threadIdx.x & 15;
float acc = 0.f;
const bf16* bw = A.w.beta_w[layer] + (size_t)hh * D_HID;
for (int d = ss; d < D_HID; d += 16)
acc += __half2float(xs[d]) * 8.f * __bfloat162float(bw[d]);
for (int o = 8; o; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o, 16);
if (ss == 0) {
acc = __bfloat162float(__float2bfloat16(acc));
A.s.betaout[hh] = 1.f / (1.f + expf(-acc));
}
} else {
int j0 = (u - 129) * 576;
for (int i = threadIdx.x; i < 576; i += blockDim.x)
A.s.hattn[layer * D_HID + j0 + i] = A.s.x[layer * D_HID + j0 + i];
}
}
}
__device__ void kda_phaseB(const KArg& A, int layer, char* SM) {
float* qc = (float*)SM; // 128
float* kc = qc + 128;
float* vc = kc + 128;
float* dec = vc + 128;
float* red = dec + 128; // 16*32
float* predv = red + 512; // 32
for (int u = blockIdx.x; u < 128; u += gridDim.x) {
__syncthreads();
const int h = u >> 2, j0 = (u & 3) * 32;
const int t = threadIdx.x;
if (t < 384) {
const int kind = t >> 7, c = t & 127;
const int ch = h * 128 + c;
const bf16* cw = A.w.conv_w[layer] + (size_t)(kind * C_KDA + ch) * 4;
const bf16* wr = A.a.winr[layer * 3 + kind];
bf16 p1b = wr[1 * C_KDA + ch];
bf16 p2b = wr[2 * C_KDA + ch];
bf16 valb = __float2bfloat16(A.s.qkvg[kind * C_KDA + ch]);
float s = __bfloat162float(wr[ch]) * __bfloat162float(cw[0])
+ __bfloat162float(p1b) * __bfloat162float(cw[1])
+ __bfloat162float(p2b) * __bfloat162float(cw[2])
+ __bfloat162float(valb) * __bfloat162float(cw[3]);
s = s / (1.f + expf(-s));
s = __bfloat162float(__float2bfloat16(s));
if (kind == 0) qc[c] = s * KDA_SCALE;
else if (kind == 1) kc[c] = s;
else vc[c] = s;
if (j0 == 0) {
bf16* ww = A.a.winw[layer * 3 + kind];
ww[ch] = p1b;
ww[C_KDA + ch] = p2b;
ww[2 * C_KDA + ch] = valb;
}
} else {
const int c = t - 384;
float graw = __bfloat162float(__float2bfloat16(A.s.qkvg[3 * C_KDA + h * 128 + c]));
dec[c] = 1.f / (1.f + expf(graw)); // exp(-softplus(g)) == sigmoid(-g)
}
__syncthreads();
const int j = t & 31, ic = t >> 5;
const int jc = j0 + j;
const int i0 = ic * 8;
float* Sg = A.a.S[layer] + (size_t)h * DK * DK;
float sreg[8];
float pp = 0.f;
#pragma unroll
for (int r = 0; r < 8; ++r) {
float sv = Sg[(i0 + r) * DK + jc] * dec[i0 + r];
sreg[r] = sv;
pp += sv * kc[i0 + r];
}
red[ic * 32 + j] = pp;
__syncthreads();
if (t < 32) {
float s = 0.f;
for (int i = 0; i < 16; ++i) s += red[i * 32 + t];
predv[t] = s;
}
__syncthreads();
const float dv = A.s.betaout[h] * (vc[jc] - predv[j]);
float oo = 0.f;
#pragma unroll
for (int r = 0; r < 8; ++r) {
float sv = sreg[r] + kc[i0 + r] * dv;
Sg[(i0 + r) * DK + jc] = sv;
oo += sv * qc[i0 + r];
}
__syncthreads();
red[ic * 32 + j] = oo;
__syncthreads();
if (t < 32) {
float s = 0.f;
for (int i = 0; i < 16; ++i) s += red[i * 32 + t];
A.s.obuf[h * 128 + j0 + t] = s;
}
}
}
__device__ void oproj_phase(const KArg& A, const uint8_t* wq, const bf16* sc,
const bf16* zr, float* dst, const float* src2,
char* SM) {
__half* xs = (__half*)SM; // 4096 half
float* gsum = (float*)(SM + 8192); // 32
float* ysred = (float*)(SM + 8448);
prep_bf16(A.s.obuf, src2, C_KDA, xs, gsum);
for (int u = blockIdx.x; u < 144; u += gridDim.x) {
int tile = u % 18, kchunk = u / 18;
gemv_i4(wq, sc, zr, D_HID, kchunk * 4, 4, tile * 128, 128,
xs, gsum, dst + tile * 128, 1.f, true, ysred);
}
}
// Warm L2 with [b0,b1) of the concatenated regions (fire-and-forget).
// Router phase A: 16 CTAs compute 4 expert logits each (bf16-rounded, into
// global scratch) while four more preload x[l+1] with the attention residual.
__device__ void moe_logits(const KArg& A, int layer, char* SM) {
const int u = blockIdx.x;
if (u < 16) {
float* xh = (float*)SM; // 2304
float* red = xh + 2304; // 17
float* er = red + 32; // 16 partials (4 experts x 4 warps)
const float* hsrc = A.s.hattn + layer * D_HID;
float ss = 0.f;
for (int i = threadIdx.x; i < D_HID; i += blockDim.x) {
float v = __bfloat162float(__float2bfloat16(hsrc[i]));
xh[i] = v;
ss += v * v;
}
float tot = block_sum(ss, red);
float scale = rsqrtf(tot / (float)D_HID + 1.0e-6f);
for (int i = threadIdx.x; i < D_HID; i += blockDim.x) {
float v = xh[i] * scale * __bfloat162float(A.w.moe_norm[layer][i]);
v = __bfloat162float(__float2bfloat16(v));
xh[i] = v;
if (u == 0) A.s.xn_moe[i] = __float2half(v);
}
__syncthreads();
const int el = threadIdx.x >> 7, t7 = threadIdx.x & 127;
const bf16* rw = A.w.router[layer] + (size_t)(u * 4 + el) * D_HID;
float acc = 0.f;
for (int d = t7; d < D_HID; d += 128) acc += xh[d] * __bfloat162float(rw[d]);
for (int o = 16; o; o >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, o);
if ((t7 & 31) == 0) er[el * 4 + (t7 >> 5)] = acc;
__syncthreads();
if (t7 == 0) {
float sacc = er[el * 4] + er[el * 4 + 1] + er[el * 4 + 2] + er[el * 4 + 3];
A.s.logits[u * 4 + el] = __bfloat162float(__float2bfloat16(sacc));
}
} else if (u < 20) {
int j0 = (u - 16) * 576;
for (int i = threadIdx.x; i < 576; i += blockDim.x)
A.s.x[(layer + 1) * D_HID + j0 + i] = A.s.hattn[layer * D_HID + j0 + i];
} else {
// bandwidth-idle: warm L2 with the shared-expert down-proj plus a
// slice of the next attention layer's weight stream.
const uint8_t* bases[3];
size_t sizes[3];
int nreg = 0;
bases[nreg] = A.w.moe_wq[layer][5];
sizes[nreg++] = (size_t)512 * 2304;
const size_t SZ = (size_t)1152 * C_KDA;
if (layer < 2) {
bases[nreg] = A.w.kda_wq[layer + 1][0];
sizes[nreg++] = SZ;
bases[nreg] = A.w.kda_wq[layer + 1][1];
sizes[nreg++] = SZ;
} else if (layer == 2) {
bases[nreg] = A.w.mla_wq[0];
sizes[nreg++] = (size_t)1152 * HQ_MLA;
} else {
bases[nreg] = A.w.kda_wq[0][0];
sizes[nreg++] = SZ;
bases[nreg] = A.w.kda_wq[0][1];
sizes[nreg++] = SZ;
}
l2_warm(bases, sizes, nreg, u - 20, gridDim.x - 20);
}
}
// Router phase B: warp 0 of CTA 0 does softmax + top-8 (min-index tie-break,
// matching torch.topk) over the 64 logits into global ids/wts; every other
// CTA warms L2 with the shared-expert gate/up weights the next phase reads.
__device__ void moe_pick(const KArg& A, int layer) {
if (blockIdx.x != 0) {
const size_t BLK = (size_t)1152 * 1024;
const size_t SZ = (size_t)1152 * C_KDA;
const uint8_t* bases[4] = { A.w.moe_wq[layer][3], A.w.moe_wq[layer][4],
nullptr, nullptr };
size_t sizes[4] = { BLK, BLK, 0, 0 };
int nreg = 2;
if (layer == 2) {
bases[nreg] = A.w.mla_wq[1];
sizes[nreg++] = (size_t)1152 * KVA;
} else {
int nl = (layer + 1) & 3; // layer 3 -> next step's layer 0
bases[nreg] = A.w.kda_wq[nl][2];
sizes[nreg++] = SZ;
bases[nreg] = A.w.kda_wq[nl][3];
sizes[nreg++] = SZ;
}
l2_warm(bases, sizes, nreg, blockIdx.x - 1, (int)gridDim.x - 1);
return;
}
if (threadIdx.x >= 32) return;
const int lane = threadIdx.x;
const float* lg = A.s.logits;
float v0 = lg[lane], v1 = lg[lane + 32];
float mx = fmaxf(v0, v1);
for (int o = 16; o; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
float sum = expf(v0 - mx) + expf(v1 - mx);
for (int o = 16; o; o >>= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
float inv = 1.f / sum;
float wsum = 0.f;
#pragma unroll
for (int j = 0; j < 8; ++j) {
float best = fmaxf(v0, v1);
for (int o = 16; o; o >>= 1) best = fmaxf(best, __shfl_xor_sync(0xffffffffu, best, o));
int win = (v0 == best) ? lane : ((v1 == best) ? lane + 32 : 64);
for (int o = 16; o; o >>= 1) win = min(win, __shfl_xor_sync(0xffffffffu, win, o));
float p = expf(best - mx) * inv;
wsum += p;
if (lane == 0) { A.s.ids[j] = win; A.s.wts[j] = p; }
if (win == lane) v0 = -1e30f;
else if (win == lane + 32) v1 = -1e30f;
}
__syncwarp();
float sc2 = A.routed / (wsum + 1e-9f);
if (lane < 8) A.s.wts[lane] *= sc2;
}
__device__ void moe_gateup(const KArg& A, int layer, char* SM) {
__half* xs = (__half*)SM; // 2304 half
float* gsum = (float*)(SM + 8192); // 18
float* ysred = (float*)(SM + 8448);
for (int i = threadIdx.x; i < D_HID; i += blockDim.x)
xs[i] = __hmul(A.s.xn_moe[i], __float2half(0.125f));
__syncthreads();
fill_gsum(xs, D_HID, gsum);
for (int u = blockIdx.x; u < 144; u += gridDim.x) {
int slot = u >> 4, gu = (u >> 3) & 1, tile = u & 7;
const uint8_t* wq;
const bf16 *sc, *zr;
if (slot < 8) {
int e = A.s.ids[slot];
wq = A.w.moe_wq[layer][gu] + (size_t)e * (1152 * 1024);
sc = A.w.moe_sc[layer][gu] + (size_t)e * (18 * 1024);
zr = A.w.moe_zr[layer][gu] + (size_t)e * (18 * 1024);
} else {
wq = A.w.moe_wq[layer][3 + gu];
sc = A.w.moe_sc[layer][3 + gu];
zr = A.w.moe_zr[layer][3 + gu];
}
gemv_i4(wq, sc, zr, M_INTER, 0, 18, tile * 128, 128, xs, gsum,
(gu ? A.s.hu : A.s.hg) + slot * M_INTER + tile * 128, 1.f, false, ysred);
}
}
__device__ void moe_down(const KArg& A, int layer, char* SM) {
__half* xs = (__half*)SM; // 1024 half
float* gsum = (float*)(SM + 8192); // 8
float* ysred = (float*)(SM + 8448);
for (int u = blockIdx.x; u < 162; u += gridDim.x) {
__syncthreads();
int slot = u / 18, tile = u % 18;
for (int i = threadIdx.x; i < M_INTER; i += blockDim.x) {
float g = A.s.hg[slot * M_INTER + i];
float up = A.s.hu[slot * M_INTER + i];
float act = g / (1.f + expf(-g)) * up;
xs[i] = __float2half(act * 0.125f);
}
__syncthreads();
fill_gsum(xs, M_INTER, gsum);
const uint8_t* wq;
const bf16 *sc, *zr;
float wmul;
if (slot < 8) {
int e = A.s.ids[slot];
wq = A.w.moe_wq[layer][2] + (size_t)e * (512 * 2304);
sc = A.w.moe_sc[layer][2] + (size_t)e * (8 * 2304);
zr = A.w.moe_zr[layer][2] + (size_t)e * (8 * 2304);
wmul = A.s.wts[slot];
} else {
wq = A.w.moe_wq[layer][5];
sc = A.w.moe_sc[layer][5];
zr = A.w.moe_zr[layer][5];
wmul = 1.f;
}
gemv_i4(wq, sc, zr, D_HID, 0, 8, tile * 128, 128, xs, gsum,
A.s.x + (layer + 1) * D_HID + tile * 128, wmul, true, ysred);
}
}
__device__ void mla_phaseA(const KArg& A, char* SM) {
__half* xs = (__half*)SM;
float* gsum = (float*)(SM + 8192);
float* red = (float*)(SM + 8320);
float* ysred = (float*)(SM + 8448);
prep_norm(A.s.x + 3 * D_HID, A.w.attn_norm[3], D_HID, xs, gsum, red);
for (int u = blockIdx.x; u < 57; u += gridDim.x) {
if (u < 48) {
gemv_i4(A.w.mla_wq[0], A.w.mla_sc[0], A.w.mla_zr[0], HQ_MLA,
0, 18, u * 128, 128, xs, gsum, A.s.q6144 + u * 128,
1.f, false, ysred);
} else if (u < 53) {
int tile = u - 48;
int nc = min(128, KVA - tile * 128);
gemv_i4(A.w.mla_wq[1], A.w.mla_sc[1], A.w.mla_zr[1], KVA,
0, 18, tile * 128, nc, xs, gsum, A.s.kva + tile * 128,
1.f, false, ysred);
} else if (u < 57) {
int j0 = (u - 53) * 576;
for (int i = threadIdx.x; i < 576; i += blockDim.x)
A.s.hattn[3 * D_HID + j0 + i] = A.s.x[3 * D_HID + j0 + i];
} else {
// idle slots: warm L2 with the absorbed kv_b matrix (next phase)
// and the ckv/krope cache rows the scores phase streams.
const int L1 = A.a.L + 1;
const uint8_t* bases[3];
size_t sizes[3];
bases[0] = A.w.mla_wq[2];
sizes[0] = (size_t)512 * NKVB;
bases[1] = (const uint8_t*)A.s.ckv;
sizes[1] = (size_t)L1 * KVL * 2;
bases[2] = (const uint8_t*)A.s.krope;
sizes[2] = (size_t)L1 * QR * 2;
l2_warm(bases, sizes, 3, u - 57, gridDim.x - 57);
}
}
}
__device__ void mla_phaseB(const KArg& A, char* SM) {
float* qn = (float*)SM; // [32][128] scaled q_nope
for (int idx = threadIdx.x; idx < NH * 128; idx += blockDim.x) {
int hh = idx >> 7, dd = idx & 127;
float v = __bfloat162float(__float2bfloat16(
A.s.q6144[hh * 192 + dd]));
qn[idx] = v * MLA_SCALE;
}
__syncthreads();
const int L = A.a.L;
for (int u = blockIdx.x; u < 257; u += gridDim.x) {
if (u < 256) {
const int c2 = u, gc = c2 >> 6;
const int w = threadIdx.x >> 5, lane = threadIdx.x & 31;
const uint8_t* wr = A.w.mla_wq[2] + (size_t)c2 * NKVB;
const bf16* scb = A.w.mla_sc[2] + (size_t)gc * NKVB;
const bf16* zrb = A.w.mla_zr[2] + (size_t)gc * NKVB;
for (int hh = w * 2; hh < w * 2 + 2; ++hh) {
const int d0 = lane * 4;
uint32_t wrd = *reinterpret_cast<const uint32_t*>(wr + hh * 256 + d0);
float alo = 0.f, ahi = 0.f;
#pragma unroll
for (int jj = 0; jj < 4; ++jj) {
uint32_t b = (wrd >> (8 * jj)) & 0xFFu;
int n = hh * 256 + d0 + jj;
float s = __bfloat162float(scb[n]);
float z = __bfloat162float(zrb[n]);
float qv = qn[hh * 128 + d0 + jj];
alo += qv * ((float)(b & 15u) - z) * s;
ahi += qv * ((float)(b >> 4) - z) * s;
}
for (int o = 16; o; o >>= 1) {
alo += __shfl_down_sync(0xffffffffu, alo, o);
ahi += __shfl_down_sync(0xffffffffu, ahi, o);
}
if (lane == 0) {
A.s.qabs[(2 * c2) * NH + hh] = __float2half(alo);
A.s.qabs[(2 * c2 + 1) * NH + hh] = __float2half(ahi);
}
}
} else {
for (int i = threadIdx.x; i < KVL; i += blockDim.x)
A.s.ckv[(size_t)L * KVL + i] = __float2bfloat16(A.s.kva[i]);
if (threadIdx.x < 32) {
const int r = threadIdx.x;
float inv = powf(A.theta, -((float)(2 * r)) / 64.f);
float ang = (float)L * inv;
float cc = cosf(ang), sn = sinf(ang);
float ke = __bfloat162float(__float2bfloat16(
A.s.kva[KVL + 2 * r]));
float ko = __bfloat162float(__float2bfloat16(
A.s.kva[KVL + 2 * r + 1]));
A.s.krope[(size_t)L * QR + 2 * r] = __float2bfloat16(ke * cc - ko * sn);
A.s.krope[(size_t)L * QR + 2 * r + 1] = __float2bfloat16(ko * cc + ke * sn);
for (int hh = 0; hh < NH; ++hh) {
int b0 = hh * 192 + 128 + 2 * r;
float qe = __bfloat162float(__float2bfloat16(A.s.q6144[b0]));
float qo = __bfloat162float(__float2bfloat16(A.s.q6144[b0 + 1]));
A.s.qr[(2 * r) * NH + hh] = __float2half((qe * cc - qo * sn) * MLA_SCALE);
A.s.qr[(2 * r + 1) * NH + hh] = __float2half((qo * cc + qe * sn) * MLA_SCALE);
}
}
}
}
}
__device__ void mla_scores(const KArg& A, char* SM) {
const int L1 = A.a.L + 1;
// adaptive chunk rows: smallest of 16/32/64/... so one round covers L1
int crows = 16;
while ((L1 + crows - 1) / crows > (int)gridDim.x) crows <<= 1;
const int NCH = (L1 + crows - 1) / crows;
if (blockIdx.x >= NCH && (int)gridDim.x > NCH) {
// idle chunks: warm L2 with the o-proj weights used two phases later
const uint8_t* bases[1] = { A.w.mla_wq[3] };
size_t sizes[1] = { (size_t)2048 * 2304 };
l2_warm(bases, sizes, 1, blockIdx.x - NCH, (int)gridDim.x - NCH);
}
__half* qa_s = (__half*)SM; // 512*32
__half* qr_s = (__half*)(SM + 32768); // 64*32
__half* ckv_s = (__half*)(SM + 36864); // 32*512
float* red = (float*)(SM + 69632); // 32*32 per-batch dots
__half* kr_s = (__half*)(SM + 77824); // 32*64
if (blockIdx.x >= NCH) return; // idle CTA: skip straight to
// the grid barrier in the caller
{
const uint32_t* qsrc = reinterpret_cast<const uint32_t*>(A.s.qabs);
uint32_t* qdst = reinterpret_cast<uint32_t*>(qa_s);
for (int i = threadIdx.x; i < KVL * 16; i += blockDim.x) qdst[i] = qsrc[i];
const uint32_t* rsrc = reinterpret_cast<const uint32_t*>(A.s.qr);
uint32_t* rdst = reinterpret_cast<uint32_t*>(qr_s);
for (int i = threadIdx.x; i < QR * 16; i += blockDim.x) rdst[i] = rsrc[i];
}
__syncthreads();
for (int u = blockIdx.x; u < NCH; u += gridDim.x) {
__syncthreads();
const int l0 = u * crows;
const int cnt = min(crows, L1 - l0);
float m_run = -1e30f, s_run = 0.f; // online softmax stats (threads<32)
for (int b = 0; b * 32 < cnt; ++b) {
const int bc = min(32, cnt - b * 32);
{
const int lb = l0 + b * 32;
for (int idx = threadIdx.x; idx < 32 * 256; idx += blockDim.x) {
int row = idx >> 8, cp = (idx & 255) * 2;
if (row < bc) {
bf162 v = *reinterpret_cast<const bf162*>(A.s.ckv + (size_t)(lb + row) * KVL + cp);
*reinterpret_cast<__half2*>(ckv_s + row * KVL + cp) =
__floats2half2_rn(__bfloat162float(v.x), __bfloat162float(v.y));
}
}
for (int idx = threadIdx.x; idx < 32 * 32; idx += blockDim.x) {
int row = idx >> 5, cp = (idx & 31) * 2;
if (row < bc) {
bf162 v = *reinterpret_cast<const bf162*>(A.s.krope + (size_t)(lb + row) * QR + cp);
*reinterpret_cast<__half2*>(kr_s + row * QR + cp) =
__floats2half2_rn(__bfloat162float(v.x), __bfloat162float(v.y));
}
}
}
__syncthreads();
{
const int lloc = threadIdx.x >> 4;
const int h2 = threadIdx.x & 15;
if (lloc < bc) {
float sx = 0.f, sy = 0.f;
const __half* crow = ckv_s + lloc * KVL;
const __half2* qa2 = reinterpret_cast<const __half2*>(qa_s) + h2;
#pragma unroll 1
for (int seg = 0; seg < 8; ++seg) {
__half2 a0 = __float2half2_rn(0.f);
__half2 a1 = __float2half2_rn(0.f);
#pragma unroll 8
for (int c = seg * 64; c < seg * 64 + 64; c += 2) {
a0 = __hfma2(__half2half2(crow[c]), qa2[c * 16], a0);
a1 = __hfma2(__half2half2(crow[c + 1]), qa2[(c + 1) * 16], a1);
}
float2 f0 = __half22float2(a0);
float2 f1 = __half22float2(a1);
sx += f0.x + f1.x;
sy += f0.y + f1.y;
}
const __half* krow = kr_s + lloc * QR;
const __half2* qr2 = reinterpret_cast<const __half2*>(qr_s) + h2;
__half2 a0 = __float2half2_rn(0.f);
#pragma unroll 8
for (int r = 0; r < QR; ++r)
a0 = __hfma2(__half2half2(krow[r]), qr2[r * 16], a0);
float2 f0 = __half22float2(a0);
sx += f0.x;
sy += f0.y;
red[lloc * 32 + h2 * 2] = sx;
red[lloc * 32 + h2 * 2 + 1] = sy;
float2 sv = make_float2(sx, sy);
*reinterpret_cast<float2*>(A.s.scores + (size_t)(l0 + b * 32 + lloc) * NH + h2 * 2) = sv;
}
}
__syncthreads();
if (threadIdx.x < 32) {
const int hh = threadIdx.x;
float bm = -1e30f;
for (int l = 0; l < bc; ++l) bm = fmaxf(bm, red[l * 32 + hh]);
float bs = 0.f;
for (int l = 0; l < bc; ++l) bs += expf(red[l * 32 + hh] - bm);
float nm = fmaxf(m_run, bm);
s_run = s_run * expf(m_run - nm) + bs * expf(bm - nm);
m_run = nm;
}
}
if (threadIdx.x < 32) {
A.s.mx[u * 32 + threadIdx.x] = m_run;
A.s.sm[u * 32 + threadIdx.x] = s_run;
}
}
}
__device__ void mla_ctx(const KArg& A, char* SM) {
const int L1 = A.a.L + 1;
int crows = 16; // same adaptive chunking as scores
while ((L1 + crows - 1) / crows > (int)gridDim.x) crows <<= 1;
const int NCH = (L1 + crows - 1) / crows;
const int NCHS = NCH;
float* red = (float*)SM; // 16*32
float* Msm = red + 512; // 32
float* Sinv = Msm + 32; // 32
__half* p_s = (__half*)(SM + 4096); // crows*32 (<=128*32)
if (blockIdx.x >= NCH) return;
const int part = threadIdx.x >> 5, hh = threadIdx.x & 31;
if (part < 16) {
float m = -1e30f;
for (int ch = part; ch < NCHS; ch += 16) m = fmaxf(m, A.s.mx[ch * 32 + hh]);
red[part * 32 + hh] = m;
}
__syncthreads();
if (threadIdx.x < 32) {
float m = -1e30f;
for (int p = 0; p < 16; ++p) m = fmaxf(m, red[p * 32 + threadIdx.x]);
Msm[threadIdx.x] = m;
}
__syncthreads();
if (part < 16) {
float s = 0.f;
for (int ch = part; ch < NCHS; ch += 16)
s += A.s.sm[ch * 32 + hh] * expf(A.s.mx[ch * 32 + hh] - Msm[hh]);
red[part * 32 + hh] = s;
}
__syncthreads();
if (threadIdx.x < 32) {
float s = 0.f;
for (int p = 0; p < 16; ++p) s += red[p * 32 + threadIdx.x];
Sinv[threadIdx.x] = 1.f / s;
}
__syncthreads();
for (int u = blockIdx.x; u < NCH; u += gridDim.x) {
__syncthreads();
const int l0 = u * crows;
const int cnt = min(crows, L1 - l0);
for (int idx = threadIdx.x; idx < crows * 32; idx += blockDim.x) {
int ll = idx >> 5, h5 = idx & 31;
float p = 0.f;
if (ll < cnt)
p = expf(A.s.scores[(size_t)(l0 + ll) * NH + h5] - Msm[h5]) * Sinv[h5];
p_s[ll * 32 + h5] = __float2half(p);
}
__syncthreads();
const int c = threadIdx.x;
__half2 acc[16];
#pragma unroll
for (int i = 0; i < 16; ++i) acc[i] = __float2half2_rn(0.f);
for (int ll = 0; ll < cnt; ++ll) {
__half2 cv = __half2half2(__float2half(
__bfloat162float(A.s.ckv[(size_t)(l0 + ll) * KVL + c])));
const __half2* pr = reinterpret_cast<const __half2*>(p_s + ll * 32);
#pragma unroll
for (int i = 0; i < 16; ++i) acc[i] = __hfma2(pr[i], cv, acc[i]);
}
#pragma unroll
for (int i = 0; i < 16; ++i) {
float2 f = __half22float2(acc[i]);
atomicAdd(A.s.ctx + (size_t)(2 * i) * KVL + c, f.x);
atomicAdd(A.s.ctx + (size_t)(2 * i + 1) * KVL + c, f.y);
}
}
}
__device__ void mla_wv(const KArg& A, char* SM) {
__half* xs = (__half*)SM; // 512
float* gsum = (float*)(SM + 8192); // 4
float* ysred = (float*)(SM + 8448);
for (int u = blockIdx.x; u < NH; u += gridDim.x) {
__syncthreads();
for (int i = threadIdx.x; i < KVL; i += blockDim.x)
xs[i] = __float2half(A.s.ctx[u * KVL + i] * 0.125f);
__syncthreads();
fill_gsum(xs, KVL, gsum);
gemv_i4(A.w.mla_wq[2], A.w.mla_sc[2], A.w.mla_zr[2], NKVB, 0, 4,
u * 256 + 128, 128, xs, gsum, A.s.obuf + u * 128, 1.f, false,
ysred);
}
}
__device__ void tail_phase(const KArg& A) {
const int T = blockIdx.x * blockDim.x + threadIdx.x;
const int NT = gridDim.x * blockDim.x;
for (int i = T; i < D_HID; i += NT)
A.s.outb[i] = __float2bfloat16(A.s.x[4 * D_HID + i]);
}
// ------------------------------------------------------------------- kernel
extern "C" __global__ void __launch_bounds__(NBLK, 1) mega_kernel(KArg A) {
cg::grid_group grid = cg::this_grid();
extern __shared__ char SM[];
int ph = 0;
#define PHASE_END() do { grid.sync(); if (++ph >= A.a.max_phase) return; } while (0)
phase_Z(A);
PHASE_END();
for (int layer = 0; layer < 3; ++layer) {
kda_phaseA(A, layer, SM);
PHASE_END();
kda_phaseB(A, layer, SM);
PHASE_END();
oproj_phase(A, A.w.kda_wq[layer][4], A.w.kda_sc[layer][4],
A.w.kda_zr[layer][4], A.s.hattn + layer * D_HID, nullptr, SM);
PHASE_END();
moe_logits(A, layer, SM);
PHASE_END();
moe_pick(A, layer);
PHASE_END();
moe_gateup(A, layer, SM);
PHASE_END();
moe_down(A, layer, SM);
PHASE_END();
}
mla_phaseA(A, SM);
PHASE_END();
mla_phaseB(A, SM);
PHASE_END();
mla_scores(A, SM);
PHASE_END();
mla_ctx(A, SM);
PHASE_END();
mla_wv(A, SM);
PHASE_END();
oproj_phase(A, A.w.mla_wq[3], A.w.mla_sc[3], A.w.mla_zr[3],
A.s.hattn + 3 * D_HID, nullptr, SM);
PHASE_END();
moe_logits(A, 3, SM);
PHASE_END();
moe_pick(A, 3);
PHASE_END();
moe_gateup(A, 3, SM);
PHASE_END();
moe_down(A, 3, SM);
PHASE_END();
tail_phase(A);
#undef PHASE_END
}
// --------------------------------------------------------------------- host
static KArg g_arg;
static bf16* g_winbank = nullptr; // [2][3][3][3][C]
static int g_grid = 0;
static bool g_attr_set = false;
static void launch(cudaStream_t stream) {
void* p = (void*)&g_arg;
void* args[] = {p};
cudaError_t err = cudaLaunchCooperativeKernel(
(void*)mega_kernel, dim3(g_grid), dim3(NBLK), args, SMEM_BYTES, stream);
TORCH_CHECK(err == cudaSuccess, "coop launch failed: ", cudaGetErrorString(err));
}
void begin(std::vector<torch::Tensor> ws, std::vector<torch::Tensor> scr,
torch::Tensor winbank,
torch::Tensor S0, torch::Tensor S1, torch::Tensor S2,
std::vector<torch::Tensor> winr,
torch::Tensor ckv_src, torch::Tensor kr_src,
torch::Tensor hidden, int64_t L, double routed, double theta,
int64_t max_phase) {
if (!g_attr_set) {
cudaFuncSetAttribute((void*)mega_kernel,
cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_BYTES);
int dev = 0;
cudaGetDevice(&dev);
cudaDeviceProp prop;
cudaGetDeviceProperties(&prop, dev);
int per_sm = 0;
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&per_sm, (void*)mega_kernel, NBLK, SMEM_BYTES);
TORCH_CHECK(per_sm >= 1, "megakernel: zero occupancy");
g_grid = prop.multiProcessorCount * per_sm;
g_attr_set = true;
}
// ---- weights, fixed order built in Python (see _weight_list) ----
size_t i = 0;
for (int l = 0; l < 3; ++l) {
g_arg.w.attn_norm[l] = (const bf16*)ws[i++].data_ptr();
g_arg.w.moe_norm[l] = (const bf16*)ws[i++].data_ptr();
for (int m = 0; m < 5; ++m) {
g_arg.w.kda_wq[l][m] = (const uint8_t*)ws[i++].data_ptr();
g_arg.w.kda_sc[l][m] = (const bf16*)ws[i++].data_ptr();
g_arg.w.kda_zr[l][m] = (const bf16*)ws[i++].data_ptr();
}
g_arg.w.beta_w[l] = (const bf16*)ws[i++].data_ptr();
g_arg.w.conv_w[l] = (const bf16*)ws[i++].data_ptr();
g_arg.w.router[l] = (const bf16*)ws[i++].data_ptr();
for (int m = 0; m < 6; ++m) {
g_arg.w.moe_wq[l][m] = (const uint8_t*)ws[i++].data_ptr();
g_arg.w.moe_sc[l][m] = (const bf16*)ws[i++].data_ptr();
g_arg.w.moe_zr[l][m] = (const bf16*)ws[i++].data_ptr();
}
}
g_arg.w.attn_norm[3] = (const bf16*)ws[i++].data_ptr();
g_arg.w.moe_norm[3] = (const bf16*)ws[i++].data_ptr();
for (int m = 0; m < 4; ++m) {
g_arg.w.mla_wq[m] = (const uint8_t*)ws[i++].data_ptr();
g_arg.w.mla_sc[m] = (const bf16*)ws[i++].data_ptr();
g_arg.w.mla_zr[m] = (const bf16*)ws[i++].data_ptr();
}
g_arg.w.router[3] = (const bf16*)ws[i++].data_ptr();
for (int m = 0; m < 6; ++m) {
g_arg.w.moe_wq[3][m] = (const uint8_t*)ws[i++].data_ptr();
g_arg.w.moe_sc[3][m] = (const bf16*)ws[i++].data_ptr();
g_arg.w.moe_zr[3][m] = (const bf16*)ws[i++].data_ptr();
}
TORCH_CHECK(i == ws.size(), "weight list size mismatch");
// ---- scratch, fixed order (see _scratch_list) ----
size_t j = 0;
g_arg.s.x = (float*)scr[j++].data_ptr();
g_arg.s.hattn = (float*)scr[j++].data_ptr();
g_arg.s.qkvg = (float*)scr[j++].data_ptr();
g_arg.s.betaout = (float*)scr[j++].data_ptr();
g_arg.s.obuf = (float*)scr[j++].data_ptr();
g_arg.s.xn_moe = (__half*)scr[j++].data_ptr();
g_arg.s.ids = (int*)scr[j++].data_ptr();
g_arg.s.wts = (float*)scr[j++].data_ptr();
g_arg.s.logits = (float*)scr[j++].data_ptr();
g_arg.s.hg = (float*)scr[j++].data_ptr();
g_arg.s.hu = (float*)scr[j++].data_ptr();
g_arg.s.q6144 = (float*)scr[j++].data_ptr();
g_arg.s.kva = (float*)scr[j++].data_ptr();
g_arg.s.qabs = (__half*)scr[j++].data_ptr();
g_arg.s.qr = (__half*)scr[j++].data_ptr();
g_arg.s.scores = (float*)scr[j++].data_ptr();
g_arg.s.mx = (float*)scr[j++].data_ptr();
g_arg.s.sm = (float*)scr[j++].data_ptr();
g_arg.s.ctx = (float*)scr[j++].data_ptr();
g_arg.s.outb = (bf16*)scr[j++].data_ptr();
g_arg.s.ckv = (bf16*)scr[j++].data_ptr();
g_arg.s.krope = (bf16*)scr[j++].data_ptr();
TORCH_CHECK(j == scr.size(), "scratch list size mismatch");
g_winbank = (bf16*)winbank.data_ptr();
g_arg.a.S[0] = (float*)S0.data_ptr();
g_arg.a.S[1] = (float*)S1.data_ptr();
g_arg.a.S[2] = (float*)S2.data_ptr();
for (int k = 0; k < 9; ++k) {
g_arg.a.winr[k] = (const bf16*)winr[k].data_ptr();
g_arg.a.winw[k] = g_winbank + (size_t)k * 3 * C_KDA; // bank 0
}
g_arg.a.ckv_src = (const bf16*)ckv_src.data_ptr();
g_arg.a.kr_src = (const bf16*)kr_src.data_ptr();
g_arg.a.hidden = (const bf16*)hidden.data_ptr();
g_arg.a.copy_len = (int)L;
g_arg.a.L = (int)L;
g_arg.a.max_phase = (int)max_phase;
g_arg.routed = (float)routed;
g_arg.theta = (float)theta;
launch(at::cuda::getCurrentCUDAStream());
}
void nextstep(torch::Tensor hidden, int64_t L, int64_t widx_w, int64_t max_phase) {
const size_t bank = (size_t)9 * 3 * C_KDA;
for (int k = 0; k < 9; ++k) {
g_arg.a.winr[k] = g_winbank + (size_t)(1 - widx_w) * bank + (size_t)k * 3 * C_KDA;
g_arg.a.winw[k] = g_winbank + (size_t)widx_w * bank + (size_t)k * 3 * C_KDA;
}
g_arg.a.hidden = (const bf16*)hidden.data_ptr();
g_arg.a.copy_len = 0;
g_arg.a.L = (int)L;
g_arg.a.max_phase = (int)max_phase;
launch(at::cuda::getCurrentCUDAStream());
}
"""
_CPP_SRC = r"""
#include <torch/extension.h>
#include <vector>
void begin(std::vector<torch::Tensor> ws, std::vector<torch::Tensor> scr,
torch::Tensor winbank,
torch::Tensor S0, torch::Tensor S1, torch::Tensor S2,
std::vector<torch::Tensor> winr,
torch::Tensor ckv_src, torch::Tensor kr_src,
torch::Tensor hidden, int64_t L, double routed, double theta,
int64_t max_phase);
void nextstep(torch::Tensor hidden, int64_t L, int64_t widx_w, int64_t max_phase);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("begin", &begin);
m.def("next", &nextstep);
}
"""
_ext = None
def _get_ext():
global _ext
if _ext is None:
_ext = load_inline(
name="kimi_mega_v1",
cpp_sources=[_CPP_SRC],
cuda_sources=[_CUDA_SRC],
extra_cuda_cflags=["-O3", "-gencode=arch=compute_120,code=sm_120"],
verbose=False,
)
return _ext
# ---------------------------------------------------------------------------
# Module tree mirroring reference.py's state_dict layout exactly.
# ---------------------------------------------------------------------------
class _QL(nn.Module):
def __init__(self, in_f, out_f, group=128):
super().__init__()
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))
class _QE(nn.Module):
def __init__(self, n, in_f, out_f, group=128):
super().__init__()
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))
class _KDA(nn.Module):
def __init__(self, cfg):
super().__init__()
H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
self.q_proj = _QL(d, H * Dk)
self.k_proj = _QL(d, H * Dk)
self.v_proj = _QL(d, H * Dk)
self.g_proj = _QL(d, H * Dk)
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 = _QL(H * Dk, d)
class _MLA(nn.Module):
def __init__(self, cfg):
super().__init__()
H, d = cfg.mla_heads, cfg.hidden
self.q_proj = _QL(d, H * (cfg.qk_nope + cfg.qk_rope))
self.kv_a = _QL(d, cfg.kv_lora + cfg.qk_rope)
self.kv_b = _QL(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head))
self.o_proj = _QL(H * cfg.v_head, d)
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 = _QE(E, d, m)
self.up = _QE(E, d, m)
self.down = _QE(E, m, d)
self.s_gate = _QE(cfg.n_shared, d, m)
self.s_up = _QE(cfg.n_shared, d, m)
self.s_down = _QE(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)
class Model(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
assert tuple(cfg.pattern) == ("K", "K", "K", "M"), "kernel is specialized"
assert cfg.n_experts == 64 and cfg.moe_inter == 1024 and cfg.hidden == 2304
self.blocks = nn.ModuleList(_Block(cfg, k) for k in cfg.pattern)
self._tr = None # (expected_len, widx, S0_ptr) for continuation detect
self._cap = 0
self._dev = None
self.debug_max_phase = 1 << 30
# ---- internal buffers (plain attributes: must NOT enter state_dict) ----
def _alloc(self, dev, L):
cap = L + 4096
if self._dev == dev and self._cap >= L + 64:
return
C = 4096
nch = cap // 32 + 2
f32 = dict(device=dev, dtype=torch.float32)
f16 = dict(device=dev, dtype=torch.float16)
bf = dict(device=dev, dtype=torch.bfloat16)
self._ckv = torch.zeros(cap, 512, **bf)
self._kr = torch.zeros(cap, 64, **bf)
self._winbank = torch.zeros(2, 9, 3, C, **bf)
self._scr = [
torch.zeros(5 * 2304, **f32), # x
torch.zeros(4 * 2304, **f32), # hattn
torch.zeros(4 * C, **f32), # qkvg
torch.zeros(32, **f32), # betaout
torch.zeros(C, **f32), # obuf
torch.zeros(2304, **f16), # xn_moe
torch.zeros(8, dtype=torch.int32, device=dev), # ids
torch.zeros(8, **f32), # wts
torch.zeros(64, **f32), # logits
torch.zeros(9 * 1024, **f32), # hg
torch.zeros(9 * 1024, **f32), # hu
torch.zeros(6144, **f32), # q6144
torch.zeros(576, **f32), # kva
torch.zeros(512 * 32, **f16), # qabs
torch.zeros(64 * 32, **f16), # qr
torch.zeros(cap * 32, **f32), # scores
torch.zeros(nch * 32, **f32), # mx
torch.zeros(nch * 32, **f32), # sm
torch.zeros(32 * 512, **f32), # ctx
torch.zeros(2304, **bf), # outb
self._ckv.view(-1),
self._kr.view(-1),
]
self._out = self._scr[19].view(2304)
self._cap = cap
self._dev = dev
def _weight_list(self):
ws = []
for i in range(3):
b = self.blocks[i]
ws += [b.attn_norm, b.moe_norm]
for p in (b.attn.q_proj, b.attn.k_proj, b.attn.v_proj, b.attn.g_proj, b.attn.o_proj):
ws += [p.w_q, p.scales, p.zeros]
ws += [b.attn.beta_proj.weight, b.attn.conv_w, b.moe.router.weight]
for p in (b.moe.gate, b.moe.up, b.moe.down, b.moe.s_gate, b.moe.s_up, b.moe.s_down):
ws += [p.w_q, p.scales, p.zeros]
b = self.blocks[3]
ws += [b.attn_norm, b.moe_norm]
for p in (b.attn.q_proj, b.attn.kv_a, b.attn.kv_b, b.attn.o_proj):
ws += [p.w_q, p.scales, p.zeros]
ws += [b.moe.router.weight]
for p in (b.moe.gate, b.moe.up, b.moe.down, b.moe.s_gate, b.moe.s_up, b.moe.s_down):
ws += [p.w_q, p.scales, p.zeros]
return [w.contiguous() for w in ws]
@torch.no_grad()
def step(self, hidden, state):
ext = _get_ext()
L = state[3]["c_kv"].shape[0]
cont = (
self._tr is not None
and state[3]["c_kv"].data_ptr() == self._ckv.data_ptr()
and L == self._tr[0]
and state[0]["S"].data_ptr() == self._tr[2]
)
if not cont:
dev = hidden.device
self._alloc(dev, L)
winr = []
for i in range(3):
for k in ("cq", "ck", "cv"):
t = state[i][k]
assert t.is_contiguous() and t.dtype == torch.bfloat16
winr.append(t)
assert state[i]["S"].is_contiguous() and state[i]["S"].dtype == torch.float32
ckv_src = state[3]["c_kv"].contiguous()
kr_src = state[3]["k_rope"].contiguous()
assert L + 1 <= self._cap
ext.begin(
self._weight_list(), self._scr, self._winbank.view(-1),
state[0]["S"], state[1]["S"], state[2]["S"], winr,
ckv_src, kr_src, hidden.contiguous(), L,
float(self.cfg.routed_scaling), float(self.cfg.rope_theta),
self.debug_max_phase,
)
widx = 0
else:
widx = 1 - self._tr[1]
ext.next(hidden, L, widx, self.debug_max_phase)
# hand back views into our capacity buffers / window bank
state[3]["c_kv"] = self._ckv[: L + 1]
state[3]["k_rope"] = self._kr[: L + 1]
for i in range(3):
state[i]["cq"] = self._winbank[widx, i * 3 + 0]
state[i]["ck"] = self._winbank[widx, i * 3 + 1]
state[i]["cv"] = self._winbank[widx, i * 3 + 2]
self._tr = (L + 1, widx, state[0]["S"].data_ptr())
return self._out, state
20260719_121747_or-fable_anthropic_claude-fable-5_02_kimi_linear_decode