KernelBench cuda · RTX PRO 6000
GLM-5.2 Fused MoE Kimi K3 (256k)
5.95%geomean peak fraction across shapes
manually audited: clean
Genuine hand-written SM120 CUDA/PTX fused MoE v2: device-side expert histogram/scan/scatter with the shared expert folded in as expert E, a tile-table build, two grouped TN bf16 mma.sync GEMMs with cp.async multi-stage pipelines (GEMM1 fuses silu(gate)*up, GEMM2 applies routing weights via fp32 atomic accumulation), and an fp32-to-bf16 finalize. A one-time weight repack into contiguous TN*TK blocks is cached on the Model, but it is keyed to model parameters (state), not inputs — every forward recomputes routing and both GEMMs from the live x/expert_ids/expert_weights. No Triton/DSL, forbidden op, output cache, grader tampering, or cross-run contamination. The 0.0595 geomean is a real solution timing and is this model's best clean cell for the problem (beats its round-2 0.0446).
harnesskinetic-claudeagent session2h 4mtotal wall2h 48mcheck7mbenchmark37moutput tokens—gpu-lock wait54mgpu-lock held49mregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
shape 023.266 ms16.0%0.56 TB/s · 31% of 1.8 TB/s HBM · also 80 TFLOPS (16% of compute)
shape 123.101 ms16.2%0.56 TB/s · 31% of 1.8 TB/s HBM · also 81 TFLOPS (16% of compute)
shape 20.531 ms0.2%24.34 TB/s · 1352% of 1.8 TB/s HBM · also 1 TFLOPS (0% of compute)
shape 333.543 ms22.1%111 TFLOPS · 22% of 500 TF bf16 peak · also 0.39 TB/s (22% of HBM)
shape 49.834 ms4.7%1.32 TB/s · 73% of 1.8 TB/s HBM · also 24 TFLOPS (5% of compute)
shape 59.345 ms9.7%1.39 TB/s · 77% of 1.8 TB/s HBM · also 48 TFLOPS (10% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(16.0% · 16.2% · 0.2% · 22.1% · 4.7% · 9.7%) = 6.0%
Kernel source (redacted)
"""GLM-5.2-class fused MoE for RTX PRO 6000 (SM120), hand-written CUDA.
Pipeline (all custom CUDA kernels, bf16 mma.sync tensor cores):
1. routing prep: expert histogram, prefix scan, token scatter (vLLM-style
sorted-token layout with the shared expert folded in as expert E),
tile-table build.
2. GEMM1: grouped TN GEMM (x rows gathered per tile against w1 gate|up
pair) with a cp.async multi-stage pipeline; fused silu(gate)*up epilogue
writes the bf16 intermediate h.
3. GEMM2: grouped TN GEMM (h @ w2^T), epilogue scales by routing weight and
atomically accumulates fp32 into the output.
4. finalize: fp32 -> bf16.
Weights are repacked ONCE per model (first forward; weights are inputs-invariant)
into contiguous TN x TK block layout so the GEMM kernels stream B tiles as
sequential 8 KB DRAM runs instead of 64 B strides of an 8 KB row.
The shared expert is handled as expert index E with weight 1.0, so every
token sees exactly top_k + n_shared rows.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_CPP_DECL = """
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wt, torch::Tensor w1r, torch::Tensor w2r, torch::Tensor w1s, torch::Tensor w2s, long topk, long cfg, long repack);
torch::Tensor moe_repack(torch::Tensor w, long TN, long TK, long halved_rows);
"""
_CUDA_SRC = r"""// GLM-5.2 fused MoE — CUDA kernels v2 (templated configs + optional weight repack).
// Pipeline: routing prep -> GEMM1 (gate+up TN GEMM, silu*mul -> h) ->
// GEMM2 (down TN GEMM, weighted fp32 atomic add) -> finalize (fp32->bf16).
//
// Repack layout (REPACK=true): B streamed as contiguous TN*TK blocks.
// w1p: per (e, band nb over I): [gate KB blocks of TN*TK][up KB blocks]
// w2p: per (e, band nb over H): [KB blocks of TN*TK]
#include <cuda_bf16.h>
#include <cstdint>
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
using bf16 = __nv_bfloat16;
#define DEVI __device__ __forceinline__
DEVI uint32_t smem_u32(const void* p) { return (uint32_t)__cvta_generic_to_shared(p); }
DEVI void cp_async_16(uint32_t dst, const void* src) {
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst), "l"(src));
}
DEVI void cp_async_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N> DEVI void cp_async_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); }
DEVI void ldmatrix_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, uint32_t addr) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(addr));
}
DEVI void mma_bf16_16816(float& c0, float& c1, float& c2, float& c3,
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
uint32_t b0, uint32_t b1) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
: "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
DEVI float silu_f(float x) { return x / (1.0f + __expf(-x)); }
// ---------------------------------------------------------------------------
// Prep kernels (unchanged from v1)
// ---------------------------------------------------------------------------
__global__ void moe_count_kernel(const long* __restrict__ ids, int nslots, int* __restrict__ counts) {
int stride = gridDim.x * blockDim.x;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < nslots; i += stride)
atomicAdd(&counts[(int)ids[i]], 1);
}
__global__ void moe_scan_kernel(int* __restrict__ counts, int* __restrict__ offs, int* __restrict__ next,
int* __restrict__ t1o, int* __restrict__ t2o,
int ETOT, int NS, int T, int TM, int NT1, int NT2) {
if (threadIdx.x != 0) return;
int run = 0, t1 = 0, t2 = 0;
for (int e = 0; e < ETOT; ++e) {
int c = counts[e] + (e >= ETOT - NS ? T : 0);
offs[e] = run; next[e] = run; t1o[e] = t1; t2o[e] = t2;
int mt = (c + TM - 1) / TM;
run += c; t1 += mt * NT1; t2 += mt * NT2;
}
offs[ETOT] = run; t1o[ETOT] = t1; t2o[ETOT] = t2;
}
__global__ void moe_scatter_kernel(const long* __restrict__ ids, const bf16* __restrict__ wt,
int T, int topk, int E, int ETOT,
int* __restrict__ next, int* __restrict__ sorted_token,
float* __restrict__ sorted_wt) {
int stride = gridDim.x * blockDim.x;
int nslots = T * topk;
for (int s = blockIdx.x * blockDim.x + threadIdx.x; s < nslots; s += stride) {
int e = (int)ids[s];
int p = atomicAdd(&next[e], 1);
sorted_token[p] = s / topk;
sorted_wt[p] = __bfloat162float(wt[s]);
}
for (int t = blockIdx.x * blockDim.x + threadIdx.x; t < T; t += stride) {
for (int s = 0; s < ETOT - E; ++s) {
int p = atomicAdd(&next[E + s], 1);
sorted_token[p] = t;
sorted_wt[p] = 1.0f;
}
}
}
__global__ void moe_tile_kernel(const int* __restrict__ offs, const int* __restrict__ t1o,
const int* __restrict__ t2o, int ETOT, int TM, int NT1, int NT2,
int* __restrict__ t1_e, int* __restrict__ t1_r, int* __restrict__ t1_c,
int* __restrict__ t2_e, int* __restrict__ t2_r, int* __restrict__ t2_c) {
int e = blockIdx.x * blockDim.x + threadIdx.x;
if (e >= ETOT) return;
int c = offs[e + 1] - offs[e];
int mt = (c + TM - 1) / TM;
for (int n = 0; n < NT1; ++n)
for (int m = 0; m < mt; ++m) {
int i = t1o[e] + n * mt + m;
t1_e[i] = e; t1_r[i] = offs[e] + m * TM; t1_c[i] = n;
}
for (int n = 0; n < NT2; ++n)
for (int m = 0; m < mt; ++m) {
int i = t2o[e] + n * mt + m;
t2_e[i] = e; t2_r[i] = offs[e] + m * TM; t2_c[i] = n;
}
}
// ---------------------------------------------------------------------------
// Repack kernel: w (ET ot rows R x K) -> per-(e,band) contiguous TNxTK blocks.
// For w1-like tensors, halves (gate/up) are interleaved per band:
// out[(((e*NT + nb)*HALVES + half)*KB + kb)*TN*TK + r*TK + c]
// = w[e][half*HALFROWS + nb*TN + r][kb*TK + c]
// For w2-like (HALVES=1): same with half=0.
// ---------------------------------------------------------------------------
template <int TN, int TK>
__global__ void moe_repack_kernel(const bf16* __restrict__ w, bf16* __restrict__ out,
int ET, long rows, long K, int HALFROWS, int halves) {
// one chunk (8 elems) per thread
const int NT = HALFROWS / TN; // bands per half
const int KB = (int)(K / TK);
const long chunks_per_band = (long)halves * KB * TN * (TK / 8);
long total = (long)ET * NT * chunks_per_band;
for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total; i += (long)gridDim.x * blockDim.x) {
long band = i / chunks_per_band;
long rem = i % chunks_per_band;
int e = (int)(band / NT);
int nb = (int)(band % NT);
int within = (int)(TK / 8); // chunks per row
long rowchunk = rem / within; // (half*KB + kb)*TN + r
int c8 = (int)(rem % within);
long kkblock = rowchunk / TN; // half*KB + kb
int r = (int)(rowchunk % TN);
int half = (int)(kkblock / KB);
int kb = (int)(kkblock % KB);
int n = half * HALFROWS + nb * TN + r;
long k = (long)kb * TK + c8 * 8;
long src = ((long)e * rows + n) * K + k;
reinterpret_cast<uint4*>(out)[i] = reinterpret_cast<const uint4*>(w)[src / 8];
}
}
torch::Tensor moe_repack(torch::Tensor w, long TN, long TK, long halved_rows) {
// w: (ET, rows, K); halved_rows = I for w1 (gate/up halves), = rows for w2
int ET = w.size(0);
long rows = w.size(1);
long K = w.size(2);
auto out = torch::empty_like(w);
int halves = (int)(rows / halved_rows);
long total_chunks = w.numel() / 8;
int thr = 256;
long blk = std::min<long>(2048, (total_chunks + thr - 1) / thr);
auto stream = at::cuda::getCurrentCUDAStream();
auto kfn = (TN == 128 && TK == 32) ? moe_repack_kernel<128, 32>
: (TN == 64 && TK == 32) ? moe_repack_kernel<64, 32>
: (TN == 128 && TK == 64) ? moe_repack_kernel<128, 64>
: moe_repack_kernel<64, 64>;
kfn<<<(int)blk, thr, 0, stream>>>(reinterpret_cast<const bf16*>(w.data_ptr()),
reinterpret_cast<bf16*>(out.data_ptr()),
ET, rows, K, (int)halved_rows, halves);
return out;
}
// ---------------------------------------------------------------------------
// Templated grouped TN GEMM kernels.
// Warp grid: WMS x WNS warps; warp tile WM x WN; threads = 32*WMS*WNS.
// ---------------------------------------------------------------------------
template <int TM, int TN, int TK, int STG, int WM, int WN, int WMS, int WNS, bool REPACK>
struct G {
static constexpr int THREADS = 32 * WMS * WNS;
static constexpr int A_ELEMS = TM * TK;
static constexpr int CH_ROW = TK * 2 / 16; // 16B chunks per row
static constexpr int SW = TK * 2 / 16; // chunks per row (swizzle period)
static constexpr int SW_SHIFT = (TK == 32) ? 1 : 0;
static constexpr int MF = WM / 16; // m16 frags per warp
static constexpr int NF = WN / 8; // n8 frags per warp
DEVI static int swz(int c, int r) { return c ^ ((r >> SW_SHIFT) & (SW - 1)); }
};
// GEMM1: A = gathered x rows (TM x TK), B = gate+up (2*TN x TK) -> h bf16.
template <int TM, int TN, int TK, int STG, int WM, int WN, int WMS, int WNS, bool REPACK>
__global__ void __launch_bounds__(G<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>::THREADS, 1) moe_gemm1_kernel(
const bf16* __restrict__ x, const bf16* __restrict__ wr, const bf16* __restrict__ ws,
const int* __restrict__ offs, const int* __restrict__ sorted_token,
const int* __restrict__ t_e, const int* __restrict__ t_r, const int* __restrict__ t_c,
const int* __restrict__ ntiles_p, bf16* __restrict__ hbuf,
int E, int K, int I) {
using GG = G<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>;
constexpr int THREADS = GG::THREADS;
constexpr int A_ELEMS = GG::A_ELEMS;
constexpr int B_ELEMS = 2 * TN * TK;
constexpr int MF = GG::MF, NF = GG::NF;
extern __shared__ char smem_raw[];
bf16* As = reinterpret_cast<bf16*>(smem_raw);
bf16* Bs = As + STG * A_ELEMS;
int* toks = reinterpret_cast<int*>(Bs + STG * B_ELEMS);
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int wm = (warp / WNS) * WM;
const int wn = (warp % WNS) * WN;
const int ksteps = K / TK;
const int ntiles = *ntiles_p;
constexpr int KB = 4096 / TK; // placeholder; K passed at runtime, KB re-derived below if REPACK
for (int t = blockIdx.x; t < ntiles; t += gridDim.x) {
const int e = t_e[t];
const int row0 = t_r[t];
const int nb = t_c[t];
const int nc = nb * TN;
const int cnt = min(TM, offs[e + 1] - row0);
const int NT1 = I / TN;
const long kb_cnt = K / TK;
const bf16* we;
if (REPACK) {
const long ei = (e < E) ? e : (e - E);
we = (e < E ? wr : ws) + (ei * NT1 + nb) * (2 * kb_cnt) * (TN * TK);
} else {
we = (e < E) ? (wr + (size_t)e * (2 * (size_t)I) * K) : (ws + (size_t)(e - E) * (2 * (size_t)I) * K);
}
(void)KB;
__syncthreads();
if (threadIdx.x < TM) toks[threadIdx.x] = sorted_token[row0 + min((int)threadIdx.x, cnt - 1)];
__syncthreads();
float cg[MF][NF][4], cu[MF][NF][4];
#pragma unroll
for (int i = 0; i < MF; ++i)
#pragma unroll
for (int j = 0; j < NF; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) { cg[i][j][q] = 0.f; cu[i][j][q] = 0.f; }
auto load_stage = [&](int s, int ks) {
const int k0 = ks * TK;
#pragma unroll
for (int i = 0; i < A_ELEMS * 2 / 16 / THREADS; ++i) {
int cidx = threadIdx.x + i * THREADS;
int r = cidx / GG::CH_ROW, c = cidx % GG::CH_ROW;
cp_async_16(smem_u32(As + s * A_ELEMS + r * TK + GG::swz(c, r) * 8),
x + (size_t)toks[r] * K + k0 + c * 8);
}
#pragma unroll
for (int i = 0; i < B_ELEMS * 2 / 16 / THREADS; ++i) {
int cidx = threadIdx.x + i * THREADS;
int r = cidx / GG::CH_ROW, c = cidx % GG::CH_ROW; // r in [0, 2*TN)
bool up = r >= TN;
const bf16* src;
if (REPACK) {
long blk = ((long)(ks) + (up ? kb_cnt : 0)) * (TN * TK) + (r - (up ? TN : 0)) * TK + c * 8;
src = we + blk;
} else {
size_t wrow = up ? ((size_t)I + nc + (r - TN)) : ((size_t)nc + r);
src = we + wrow * K + k0 + c * 8;
}
cp_async_16(smem_u32(Bs + s * B_ELEMS + r * TK + GG::swz(c, r) * 8), src);
}
};
#pragma unroll
for (int s = 0; s < STG - 1; ++s) {
if (s < ksteps) load_stage(s, s);
cp_async_commit();
}
for (int ks = 0; ks < ksteps; ++ks) {
cp_async_wait<STG - 2>();
__syncthreads();
const bf16* At = As + (ks % STG) * A_ELEMS;
const bf16* Bt = Bs + (ks % STG) * B_ELEMS;
#pragma unroll
for (int hh = 0; hh < TK / 16; ++hh) {
uint32_t af[MF][4];
#pragma unroll
for (int mi = 0; mi < MF; ++mi) {
int r = wm + mi * 16 + (lane & 15);
int ch = hh * 2 + (lane >> 4);
ldmatrix_x4(af[mi][0], af[mi][1], af[mi][2], af[mi][3],
smem_u32(At + r * TK + GG::swz(ch, r) * 8));
}
uint32_t bg[NF][2], bu[NF][2];
#pragma unroll
for (int q = 0; q < NF / 2; ++q) {
int g = lane >> 3;
int n = wn + q * 16 + ((g & 2) ? 8 : 0) + (lane & 7);
int ch = hh * 2 + (g & 1);
uint32_t a4[4];
ldmatrix_x4(a4[0], a4[1], a4[2], a4[3], smem_u32(Bt + n * TK + GG::swz(ch, n) * 8));
bg[q * 2 + 0][0] = a4[0]; bg[q * 2 + 0][1] = a4[1];
bg[q * 2 + 1][0] = a4[2]; bg[q * 2 + 1][1] = a4[3];
ldmatrix_x4(a4[0], a4[1], a4[2], a4[3], smem_u32(Bt + TN * TK + n * TK + GG::swz(ch, n) * 8));
bu[q * 2 + 0][0] = a4[0]; bu[q * 2 + 0][1] = a4[1];
bu[q * 2 + 1][0] = a4[2]; bu[q * 2 + 1][1] = a4[3];
}
#pragma unroll
for (int mi = 0; mi < MF; ++mi)
#pragma unroll
for (int nj = 0; nj < NF; ++nj) {
mma_bf16_16816(cg[mi][nj][0], cg[mi][nj][1], cg[mi][nj][2], cg[mi][nj][3],
af[mi][0], af[mi][1], af[mi][2], af[mi][3], bg[nj][0], bg[nj][1]);
mma_bf16_16816(cu[mi][nj][0], cu[mi][nj][1], cu[mi][nj][2], cu[mi][nj][3],
af[mi][0], af[mi][1], af[mi][2], af[mi][3], bu[nj][0], bu[nj][1]);
}
}
if (ks + STG - 1 < ksteps) load_stage((ks + STG - 1) % STG, ks + STG - 1);
cp_async_commit();
}
const int rbase = lane >> 2;
const int cb = (lane & 3) * 2;
#pragma unroll
for (int mi = 0; mi < MF; ++mi)
#pragma unroll
for (int half = 0; half < 2; ++half) {
int rl = wm + mi * 16 + rbase + half * 8;
if (rl < cnt) {
bf16* hrow = hbuf + (size_t)(row0 + rl) * I + nc + wn + cb;
#pragma unroll
for (int nj = 0; nj < NF; ++nj) {
__nv_bfloat162 v;
v.x = __float2bfloat16(silu_f(cg[mi][nj][half * 2 + 0]) * cu[mi][nj][half * 2 + 0]);
v.y = __float2bfloat16(silu_f(cg[mi][nj][half * 2 + 1]) * cu[mi][nj][half * 2 + 1]);
*reinterpret_cast<__nv_bfloat162*>(hrow + nj * 8) = v;
}
}
}
}
}
// GEMM2: A = h rows (TM x TK), B = down (TN x TK) -> out32 atomic add (wt-scaled).
template <int TM, int TN, int TK, int STG, int WM, int WN, int WMS, int WNS, bool REPACK>
__global__ void __launch_bounds__(G<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>::THREADS, 1) moe_gemm2_kernel(
const bf16* __restrict__ hbuf, const bf16* __restrict__ wr, const bf16* __restrict__ ws,
const int* __restrict__ offs, const int* __restrict__ sorted_token, const float* __restrict__ sorted_wt,
const int* __restrict__ t_e, const int* __restrict__ t_r, const int* __restrict__ t_c,
const int* __restrict__ ntiles_p, float* __restrict__ out32,
int E, int K, int H) {
using GG = G<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>;
constexpr int THREADS = GG::THREADS;
constexpr int A_ELEMS = GG::A_ELEMS;
constexpr int B_ELEMS = TN * TK;
constexpr int MF = GG::MF, NF = GG::NF;
extern __shared__ char smem_raw[];
bf16* As = reinterpret_cast<bf16*>(smem_raw);
bf16* Bs = As + STG * A_ELEMS;
const int warp = threadIdx.x >> 5;
const int lane = threadIdx.x & 31;
const int wm = (warp / WNS) * WM;
const int wn = (warp % WNS) * WN;
const int ksteps = K / TK;
const int ntiles = *ntiles_p;
for (int t = blockIdx.x; t < ntiles; t += gridDim.x) {
const int e = t_e[t];
const int row0 = t_r[t];
const int nb = t_c[t];
const int nc = nb * TN;
const int cnt = min(TM, offs[e + 1] - row0);
const int NT2 = H / TN;
const long kb_cnt = K / TK;
const bf16* we;
if (REPACK) {
const long ei = (e < E) ? e : (e - E);
we = (e < E ? wr : ws) + (ei * NT2 + nb) * kb_cnt * (TN * TK);
} else {
we = (e < E) ? (wr + (size_t)e * (size_t)H * K) : (ws + (size_t)(e - E) * (size_t)H * K);
}
__syncthreads();
float cf[MF][NF][4];
#pragma unroll
for (int i = 0; i < MF; ++i)
#pragma unroll
for (int j = 0; j < NF; ++j)
#pragma unroll
for (int q = 0; q < 4; ++q) cf[i][j][q] = 0.f;
auto load_stage = [&](int s, int ks) {
const int k0 = ks * TK;
#pragma unroll
for (int i = 0; i < A_ELEMS * 2 / 16 / THREADS; ++i) {
int cidx = threadIdx.x + i * THREADS;
int r = cidx / GG::CH_ROW, c = cidx % GG::CH_ROW;
cp_async_16(smem_u32(As + s * A_ELEMS + r * TK + GG::swz(c, r) * 8),
hbuf + (size_t)(row0 + min(r, cnt - 1)) * (size_t)K + k0 + c * 8);
}
#pragma unroll
for (int i = 0; i < B_ELEMS * 2 / 16 / THREADS; ++i) {
int cidx = threadIdx.x + i * THREADS;
int r = cidx / GG::CH_ROW, c = cidx % GG::CH_ROW;
const bf16* src;
if (REPACK) {
src = we + (long)ks * (TN * TK) + r * TK + c * 8;
} else {
src = we + (size_t)(nc + r) * (size_t)K + k0 + c * 8;
}
cp_async_16(smem_u32(Bs + s * B_ELEMS + r * TK + GG::swz(c, r) * 8), src);
}
};
#pragma unroll
for (int s = 0; s < STG - 1; ++s) {
if (s < ksteps) load_stage(s, s);
cp_async_commit();
}
for (int ks = 0; ks < ksteps; ++ks) {
cp_async_wait<STG - 2>();
__syncthreads();
const bf16* At = As + (ks % STG) * A_ELEMS;
const bf16* Bt = Bs + (ks % STG) * B_ELEMS;
#pragma unroll
for (int hh = 0; hh < TK / 16; ++hh) {
uint32_t af[MF][4];
#pragma unroll
for (int mi = 0; mi < MF; ++mi) {
int r = wm + mi * 16 + (lane & 15);
int ch = hh * 2 + (lane >> 4);
ldmatrix_x4(af[mi][0], af[mi][1], af[mi][2], af[mi][3],
smem_u32(At + r * TK + GG::swz(ch, r) * 8));
}
uint32_t bg[NF][2];
#pragma unroll
for (int q = 0; q < NF / 2; ++q) {
int g = lane >> 3;
int n = wn + q * 16 + ((g & 2) ? 8 : 0) + (lane & 7);
int ch = hh * 2 + (g & 1);
uint32_t a4[4];
ldmatrix_x4(a4[0], a4[1], a4[2], a4[3], smem_u32(Bt + n * TK + GG::swz(ch, n) * 8));
bg[q * 2 + 0][0] = a4[0]; bg[q * 2 + 0][1] = a4[1];
bg[q * 2 + 1][0] = a4[2]; bg[q * 2 + 1][1] = a4[3];
}
#pragma unroll
for (int mi = 0; mi < MF; ++mi)
#pragma unroll
for (int nj = 0; nj < NF; ++nj)
mma_bf16_16816(cf[mi][nj][0], cf[mi][nj][1], cf[mi][nj][2], cf[mi][nj][3],
af[mi][0], af[mi][1], af[mi][2], af[mi][3], bg[nj][0], bg[nj][1]);
}
if (ks + STG - 1 < ksteps) load_stage((ks + STG - 1) % STG, ks + STG - 1);
cp_async_commit();
}
const int rbase = lane >> 2;
const int cb = (lane & 3) * 2;
#pragma unroll
for (int mi = 0; mi < MF; ++mi)
#pragma unroll
for (int half = 0; half < 2; ++half) {
int rl = wm + mi * 16 + rbase + half * 8;
if (rl < cnt) {
int tok = sorted_token[row0 + rl];
float w = sorted_wt[row0 + rl];
float* orow = out32 + (size_t)tok * H + nc + wn + cb;
#pragma unroll
for (int nj = 0; nj < NF; ++nj) {
atomicAdd(orow + nj * 8 + 0, cf[mi][nj][half * 2 + 0] * w);
atomicAdd(orow + nj * 8 + 1, cf[mi][nj][half * 2 + 1] * w);
}
}
}
}
}
__global__ void moe_finalize_kernel(const float* __restrict__ out32, bf16* __restrict__ out, long n4) {
long i = (long)(blockIdx.x)*blockDim.x + threadIdx.x;
if (i >= n4) return;
const float4 v = reinterpret_cast<const float4*>(out32)[i];
uint2 packed;
__nv_bfloat162 lo, hi;
lo.x = __float2bfloat16(v.x); lo.y = __float2bfloat16(v.y);
hi.x = __float2bfloat16(v.z); hi.y = __float2bfloat16(v.w);
packed.x = *reinterpret_cast<uint32_t*>(&lo);
packed.y = *reinterpret_cast<uint32_t*>(&hi);
reinterpret_cast<uint2*>(out)[i] = packed;
}
// ---------------------------------------------------------------------------
// Config dispatch
// ---------------------------------------------------------------------------
// cfg: 0=A(128,128,32,4 grid2x4) 1=B(128,64,32,6 grid4x2) 2=C(128,128,64,2 grid2x4)
// 3=D(128,64,32,3 grid4x2 2cta) 4=E(64,128,32,4 grid2x4)
template <int TM, int TN, int TK, int STG, int WM, int WN, int WMS, int WNS, bool REPACK>
void launch_gemms(int which, const void* A, const void* wr, const void* ws,
const int* offs, const int* stok, const float* swt,
const int* te, const int* tr, const int* tc, const int* ntp,
void* out, const void* outw, int E, int K, int N, int grid, cudaStream_t stream) {
using GG = G<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>;
if (which == 1) {
constexpr int SMEM = STG * (GG::A_ELEMS * 2 + 2 * TN * TK * 2) + TM * 4 + 16;
auto kfn = moe_gemm1_kernel<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>;
static bool init = false;
if (!init) { cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); init = true; }
kfn<<<grid, GG::THREADS, SMEM, stream>>>((const bf16*)A, (const bf16*)wr, (const bf16*)ws, offs, stok,
te, tr, tc, ntp, (bf16*)out, E, K, N);
} else {
constexpr int SMEM = STG * (GG::A_ELEMS * 2 + TN * TK * 2) + 16;
auto kfn = moe_gemm2_kernel<TM, TN, TK, STG, WM, WN, WMS, WNS, REPACK>;
static bool init = false;
if (!init) { cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM); init = true; }
kfn<<<grid, GG::THREADS, SMEM, stream>>>((const bf16*)A, (const bf16*)wr, (const bf16*)ws, offs, stok, swt,
te, tr, tc, ntp, (float*)outw, E, K, N);
}
}
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor ids, torch::Tensor wt,
torch::Tensor w1r, torch::Tensor w2r, torch::Tensor w1s, torch::Tensor w2s,
long topk_l, long cfg, long repack) {
const int T = x.size(0);
const int H = x.size(1);
const int E = w1r.size(0);
const int I = w1r.size(1) / 2;
const int NS = w1s.size(0);
const int ETOT = E + NS;
const int topk = (int)topk_l;
const int R = T * topk + T * NS;
auto stream = at::cuda::getCurrentCUDAStream();
auto opts_i = torch::TensorOptions().dtype(torch::kInt32).device(x.device());
auto opts_f = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
struct Cfg { int TM, TN, TK, STG, grid, occ; };
static const Cfg CFGS[5] = {
{128, 128, 32, 4, 188, 1},
{128, 64, 32, 6, 188, 1},
{128, 128, 64, 2, 188, 1},
{128, 64, 32, 3, 376, 2},
{64, 128, 32, 4, 188, 1},
};
Cfg C = CFGS[cfg];
const int NT1 = I / C.TN, NT2 = H / C.TN;
const int SEG = 768;
auto ws = torch::empty({6 * SEG}, opts_i);
int* counts = ws.data_ptr<int>();
int* offs = counts + SEG;
int* next = offs + SEG;
int* t1o = next + SEG;
int* t2o = t1o + SEG;
cudaMemsetAsync(counts, 0, ETOT * sizeof(int), stream);
auto sorted_token = torch::empty({R}, opts_i);
auto sorted_wt = torch::empty({R}, opts_f);
const int mt_bound = R / C.TM + ETOT;
auto t1_e = torch::empty({NT1 * mt_bound}, opts_i);
auto t1_r = torch::empty({NT1 * mt_bound}, opts_i);
auto t1_c = torch::empty({NT1 * mt_bound}, opts_i);
auto t2_e = torch::empty({NT2 * mt_bound}, opts_i);
auto t2_r = torch::empty({NT2 * mt_bound}, opts_i);
auto t2_c = torch::empty({NT2 * mt_bound}, opts_i);
auto hbuf = torch::empty({(long)R * I}, torch::TensorOptions().dtype(torch::kBFloat16).device(x.device()));
auto out32 = torch::empty({(long)T * H}, opts_f);
auto out = torch::empty({T, H}, torch::TensorOptions().dtype(torch::kBFloat16).device(x.device()));
const bf16* xp = reinterpret_cast<const bf16*>(x.data_ptr());
const long* idsp = ids.data_ptr<long>();
const bf16* wtp = reinterpret_cast<const bf16*>(wt.data_ptr());
const bf16* w1rp = reinterpret_cast<const bf16*>(w1r.data_ptr());
const bf16* w2rp = reinterpret_cast<const bf16*>(w2r.data_ptr());
const bf16* w1sp = reinterpret_cast<const bf16*>(w1s.data_ptr());
const bf16* w2sp = reinterpret_cast<const bf16*>(w2s.data_ptr());
cudaMemsetAsync(out32.data_ptr<float>(), 0, (size_t)T * H * sizeof(float), stream);
{
int nslots = T * topk;
int thr = 256;
int blk = min(1024, (nslots + thr - 1) / thr + 1);
moe_count_kernel<<<blk, thr, 0, stream>>>(idsp, nslots, counts);
}
moe_scan_kernel<<<1, 32, 0, stream>>>(counts, offs, next, t1o, t2o, ETOT, NS, T, C.TM, NT1, NT2);
{
int thr = 256;
int n = max(T * topk, T);
int blk = min(1024, (n + thr - 1) / thr + 1);
moe_scatter_kernel<<<blk, thr, 0, stream>>>(idsp, wtp, T, topk, E, ETOT, next,
sorted_token.data_ptr<int>(), sorted_wt.data_ptr<float>());
}
{
int thr = 128;
int blk = (ETOT + thr - 1) / thr;
moe_tile_kernel<<<blk, thr, 0, stream>>>(offs, t1o, t2o, ETOT, C.TM, NT1, NT2,
t1_e.data_ptr<int>(), t1_r.data_ptr<int>(), t1_c.data_ptr<int>(),
t2_e.data_ptr<int>(), t2_r.data_ptr<int>(), t2_c.data_ptr<int>());
}
const bool rp = repack != 0;
#define DISPATCH_G(CFGV, TMV, TNV, TKV, STGV, WMV, WNV, WMSV, WNSV) \
if (cfg == CFGV) { \
if (rp) { \
launch_gemms<TMV, TNV, TKV, STGV, WMV, WNV, WMSV, WNSV, true>( \
1, xp, w1rp, w1sp, offs, sorted_token.data_ptr<int>(), nullptr, \
t1_e.data_ptr<int>(), t1_r.data_ptr<int>(), t1_c.data_ptr<int>(), t1o + ETOT, \
hbuf.data_ptr(), nullptr, E, H, I, C.grid, stream); \
launch_gemms<TMV, TNV, TKV, STGV, WMV, WNV, WMSV, WNSV, true>( \
2, hbuf.data_ptr(), w2rp, w2sp, offs, sorted_token.data_ptr<int>(), \
sorted_wt.data_ptr<float>(), t2_e.data_ptr<int>(), t2_r.data_ptr<int>(), t2_c.data_ptr<int>(),\
t2o + ETOT, nullptr, out32.data_ptr(), E, I, H, C.grid, stream); \
} else { \
launch_gemms<TMV, TNV, TKV, STGV, WMV, WNV, WMSV, WNSV, false>( \
1, xp, w1rp, w1sp, offs, sorted_token.data_ptr<int>(), nullptr, \
t1_e.data_ptr<int>(), t1_r.data_ptr<int>(), t1_c.data_ptr<int>(), t1o + ETOT, \
hbuf.data_ptr(), nullptr, E, H, I, C.grid, stream); \
launch_gemms<TMV, TNV, TKV, STGV, WMV, WNV, WMSV, WNSV, false>( \
2, hbuf.data_ptr(), w2rp, w2sp, offs, sorted_token.data_ptr<int>(), \
sorted_wt.data_ptr<float>(), t2_e.data_ptr<int>(), t2_r.data_ptr<int>(), t2_c.data_ptr<int>(),\
t2o + ETOT, nullptr, out32.data_ptr(), E, I, H, C.grid, stream); \
} \
}
DISPATCH_G(0, 128, 128, 32, 4, 64, 32, 2, 4)
DISPATCH_G(1, 128, 64, 32, 6, 32, 32, 4, 2)
DISPATCH_G(2, 128, 128, 64, 2, 64, 32, 2, 4)
DISPATCH_G(3, 128, 64, 32, 3, 32, 32, 4, 2)
DISPATCH_G(4, 64, 128, 32, 4, 32, 32, 2, 4)
#undef DISPATCH_G
{
long n4 = (long)T * H / 4;
int thr = 256;
long blk = (n4 + thr - 1) / thr;
moe_finalize_kernel<<<(int)blk, thr, 0, stream>>>(out32.data_ptr<float>(),
reinterpret_cast<bf16*>(out.data_ptr()), n4);
}
return out;
}
"""
def _build_ext():
import os
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
os.environ.setdefault("MAX_JOBS", "16")
return load_inline(
name="glm52_fused_moe_v2",
cpp_sources=[_CPP_DECL],
cuda_sources=[_CUDA_SRC],
functions=["moe_forward", "moe_repack"],
extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
verbose=False,
)
_EXT = _build_ext()
OP_TYPE = "glm52_fused_moe"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000"]
# cfg id in the CUDA dispatch table (see moe2 CFGS): {TM, TN, TK, STG, grid, occ}
_CFG_ID = 0
_TNS = [128, 64, 128, 64, 128]
_TKS = [32, 32, 64, 32, 32]
class Model(nn.Module):
def __init__(self, T: int, E: int, top_k: int, n_shared: int, H: int, I: int):
super().__init__()
self.T, self.E, self.top_k = T, E, top_k
self.n_shared, self.H, self.I = n_shared, H, I
self.w1_routed = nn.Parameter(torch.empty(E, 2 * I, H, dtype=torch.bfloat16))
self.w2_routed = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
self.w1_shared = nn.Parameter(torch.empty(n_shared, 2 * I, H, dtype=torch.bfloat16))
self.w2_shared = nn.Parameter(torch.empty(n_shared, H, I, dtype=torch.bfloat16))
for p in self.parameters():
nn.init.normal_(p, std=0.02)
self._packed = None
def _ensure_packed(self):
cfg = _CFG_ID
if self._packed is None or self._packed[0] != cfg:
self._packed = (
cfg,
_EXT.moe_repack(self.w1_routed.detach(), _TNS[cfg], _TKS[cfg], self.I),
_EXT.moe_repack(self.w2_routed.detach(), _TNS[cfg], _TKS[cfg], self.H),
_EXT.moe_repack(self.w1_shared.detach(), _TNS[cfg], _TKS[cfg], self.I),
_EXT.moe_repack(self.w2_shared.detach(), _TNS[cfg], _TKS[cfg], self.H),
)
return self._packed
def forward(self, x: torch.Tensor, expert_ids: torch.Tensor, expert_weights: torch.Tensor) -> torch.Tensor:
cfg, w1p, w2p, w1sp, w2sp = self._ensure_packed()
return _EXT.moe_forward(
x.contiguous(),
expert_ids.contiguous(),
expert_weights.contiguous(),
w1p,
w2p,
w1sp,
w2sp,
self.top_k,
cfg,
1,
)
20260716_090533_kinetic-claude_kinetic-0715_01_glm52_fused_moe