KernelBench cuda · RTX PRO 6000
GLM-5.2 Fused MoE Claude Fable 5
8.04%geomean peak fraction across shapes
manually audited: clean
Genuine hand-written CUDA fused MoE via a single load_inline extension with inline-PTX tensor-core kernels (framework=ptx). Design: GPU-side routing (histogram -> one-block exclusive scan + m-block map -> scatter of (token, expert, weight) slots sorted by expert, shared expert folded in as expert index E with weight 1.0), then two grouped GEMMs written from scratch with mma.m16n8k16 bf16->f32, cp.async multi-stage double buffering, ldmatrix, and a fused silu(gate)*up epilogue; per-slot outputs are weight-scaled in the GEMM2 epilogue and summed per token by a vectorized fp32-accumulate reduce. T<=8 decode takes a slot-parallel GEMV path. Weights are re-packed into a k-tiled layout guarded by a (data_ptr, _version) key - a layout transform of live parameters, never a cached output, and empirically proven to invalidate on in-place weight mutation. No forbidden imports (no triton/vllm/flashinfer/_grouped_mm/ reference), no output caching, no CUDA graphs. 0.0804 geomean is an honest mid-tier score: prefill shapes 0.27-0.34, T=1 decode 0.0022 (launch-overhead-bound like every model on this deck).
harnessor-fableagent session4h 44mtotal wall4h 56mcheck5mbenchmark6moutput tokens—gpu-lock wait2mgpu-lock held45mregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
shape 013.551 ms27.4%0.96 TB/s · 53% of 1.8 TB/s HBM · also 137 TFLOPS (27% of compute)
shape 113.639 ms27.4%0.95 TB/s · 53% of 1.8 TB/s HBM · also 137 TFLOPS (27% of compute)
shape 20.414 ms0.2%31.28 TB/s · 1738% of 1.8 TB/s HBM · also 1 TFLOPS (0% of compute)
shape 321.739 ms34.1%171 TFLOPS · 34% of 500 TF bf16 peak · also 0.60 TB/s (33% of HBM)
shape 49.225 ms5.0%1.40 TB/s · 78% of 1.8 TB/s HBM · also 25 TFLOPS (5% of compute)
shape 59.487 ms9.6%1.37 TB/s · 76% 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(27.4% · 27.4% · 0.2% · 34.1% · 5.0% · 9.6%) = 8.0%
Kernel source (redacted)
"""GLM-5.2-class fused MoE — hand-written CUDA (SM120 / RTX PRO 6000).
Strategy
--------
The deck is weight-bandwidth-bound almost everywhere: every touched expert's
weights (w1 33.5 MB + w2 16.8 MB per expert; all 257 experts ~= 12.9 GB at
T>=512) must stream from GDDR7 once. So the design reads each expert's weights
exactly once per forward:
1. Routing build (GPU, no host sync): histogram token-slots per expert ->
exclusive offsets -> scatter slot rows sorted by expert. The shared expert
is folded in as expert index E with routing weight 1.0, giving a single
uniform grouped-GEMM problem with top_k + n_shared slots per token.
2. k_gemm1: grouped GEMM over expert-sorted slot rows. Gathers x rows via
token indices (no permuted copy of x), computes gate and up projections in
one 256-thread block (warps 0-3 gate, warps 4-7 up, sharing the A tile),
bf16 tensor-core mma.m16n8k16 with cp.async double buffering, and fuses
silu(gate)*up in the epilogue -> bf16 workspace h (N, I).
3. k_gemm2: grouped GEMM h @ w2^T, epilogue scales by the slot routing weight
and accumulates with fp32 atomicAdd into out (T, H) fp32; final bf16 cast.
All accumulation is fp32 (mma f32 accumulators + fp32 output buffer); the only
bf16 intermediate is h, matching the blessed real-kernel pipeline.
"""
from __future__ import annotations
import os
import subprocess
import sys
from pathlib import Path
import torch
import torch.nn as nn
# --------------------------------------------------------------------------
# Toolchain bootstrap: torch is cu130 but the system only has CUDA 12.8, so we
# compile with the pip-provided CUDA 13.0 toolchain wheels inside the venv.
# --------------------------------------------------------------------------
_CUDA_WHEELS = [
"nvidia-cuda-nvcc==13.0.88",
"nvidia-cuda-crt==13.0.88",
"nvidia-nvvm==13.0.88",
"nvidia-cuda-cccl==13.0.85",
"nvidia-cuda-runtime==13.0.88",
]
def _find_venv_cuda() -> Path | None:
import site
candidates = list(site.getsitepackages())
try:
candidates.append(str(Path(site.getusersitepackages())))
except Exception:
pass
for sp in candidates:
p = Path(sp) / "nvidia" / "cu13"
if (p / "bin" / "nvcc").exists():
return p
return None
def _bootstrap_cuda_home() -> None:
venv_cuda = _find_venv_cuda()
if venv_cuda is None:
# Fresh environment: install the pinned CUDA 13.0 toolchain wheels.
subprocess.run(
["uv", "pip", "install", *_CUDA_WHEELS],
check=True,
cwd=str(Path(__file__).resolve().parents[2]),
)
venv_cuda = _find_venv_cuda()
if venv_cuda is None:
raise RuntimeError("CUDA 13 toolchain wheels not found; cannot compile")
# nvcc links -lcudart; the wheel only ships libcudart.so.13.
lib = venv_cuda / "lib"
so = lib / "libcudart.so"
if not so.exists():
try:
so.symlink_to("libcudart.so.13")
except FileExistsError:
pass
os.environ["CUDA_HOME"] = str(venv_cuda)
os.environ["PATH"] = f"{venv_cuda}/bin:" + os.environ.get("PATH", "")
_bootstrap_cuda_home()
from torch.utils.cpp_extension import load_inline # noqa: E402
_CPP_SRC = r"""
#include <torch/extension.h>
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor expert_ids,
torch::Tensor expert_weights, torch::Tensor w1_routed,
torch::Tensor w2_routed, torch::Tensor w1_shared,
torch::Tensor w2_shared, torch::Tensor w1p_routed,
torch::Tensor w2p_routed, torch::Tensor w1p_shared,
torch::Tensor w2p_shared);
"""
_CUDA_SRC = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_bf16.h>
#include <cstdint>
using bf16 = __nv_bfloat16;
#define DEVINL __device__ __forceinline__
// ---------------------------------------------------------------- cp.async
DEVINL void cp_async16(void* smem, const void* gmem) {
unsigned s = (unsigned)__cvta_generic_to_shared(smem);
asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(s), "l"(gmem)
: "memory");
}
DEVINL void cp_commit() { asm volatile("cp.async.commit_group;\n" ::: "memory"); }
template <int N>
DEVINL void cp_wait() { asm volatile("cp.async.wait_group %0;\n" ::"n"(N) : "memory"); }
// ---------------------------------------------------------------- ldmatrix
DEVINL void ldsm_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3,
const void* p) {
unsigned s = (unsigned)__cvta_generic_to_shared(p);
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
: "r"(s)
: "memory");
}
DEVINL void ldsm_x4_t(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3,
const void* p) {
unsigned s = (unsigned)__cvta_generic_to_shared(p);
asm volatile(
"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
: "r"(s)
: "memory");
}
// mma.m16n8k16 row.col bf16 -> f32
DEVINL void mma16816(float c[4], 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"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
// ---------------------------------------------------------------- routing
__global__ void k_count(const int64_t* __restrict__ ids, int* __restrict__ counts,
int total) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < total) atomicAdd(&counts[(int)ids[i]], 1);
}
// One-block kernel: offsets + M-block mapping. NE <= ~320, nb <= ~1500.
// Parallel global loads into smem, cheap serial smem scan, parallel blk fill.
__global__ void k_scan(const int* __restrict__ counts, int T, int NE, int E,
int BM, int* __restrict__ offsets,
int* __restrict__ blk_expert, int* __restrict__ blk_row,
int* __restrict__ total_nb, int max_nb) {
__shared__ int sc[520], so[520], sb[520];
int tid = threadIdx.x;
for (int e = tid; e < NE; e += blockDim.x)
sc[e] = (e < E) ? counts[e] : T; // shared experts serve every token
__syncthreads();
if (tid == 0) {
int off = 0, nb = 0;
for (int e = 0; e < NE; e++) {
so[e] = off;
sb[e] = nb;
off += sc[e];
nb += (sc[e] + BM - 1) / BM;
}
so[NE] = off;
offsets[NE] = off;
*total_nb = nb;
}
__syncthreads();
for (int e = tid; e < NE; e += blockDim.x) {
offsets[e] = so[e];
int b = (sc[e] + BM - 1) / BM;
int base = sb[e];
for (int i = 0; i < b && base + i < max_nb; i++) {
blk_expert[base + i] = e;
blk_row[base + i] = so[e] + i * BM;
}
}
}
__global__ void k_scatter(const int64_t* __restrict__ ids,
const bf16* __restrict__ wts,
const int* __restrict__ offsets,
int* __restrict__ cursor, int T, int K, int S, int E,
int* __restrict__ stok, float* __restrict__ sw,
int* __restrict__ tok2slot, int* __restrict__ sexp) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= T * S) return;
int t = i / S, k = i % S;
if (k < K) {
int e = (int)ids[t * K + k];
int pos = offsets[e] + atomicAdd(&cursor[e], 1);
stok[pos] = t;
sw[pos] = __bfloat162float(wts[t * K + k]);
tok2slot[i] = pos;
sexp[pos] = e;
} else {
int s = k - K;
int pos = offsets[E + s] + t;
stok[pos] = t;
sw[pos] = 1.0f;
tok2slot[i] = pos;
sexp[pos] = E + s;
}
}
// ---------------------------------------------------------------- GEMM1
// Grouped GEMM: h[j, n] = silu(x[tok_j] . wg_n) * (x[tok_j] . wu_n)
// Block: BM=64 slot rows x BN=64 intermediate cols. 8 warps: 0-3 gate, 4-7 up.
#define BM 64
#define BN1 64
#define BK 32
#define PAD 8
#define AST (BK + PAD) // smem row stride in elems (80B: ldmatrix conflict-free)
#define NSTG1 3 // gemm1 pipeline depth (46KB smem)
#define ST1 ((BM + 2 * BN1) * AST) // gemm1 per-stage smem elems
#define NSTG2 4 // gemm2 pipeline depth (40KB smem)
#define ST2 ((BM + 64) * AST) // gemm2 per-stage smem elems (BN2=64)
// Templated on block-M: BMT=64 (2 blocks/SM, small N) or BMT=128 (halves
// weight re-reads across m-blocks at large N; 1 block/SM, deeper pipeline).
template <int BMT>
__global__ __launch_bounds__(256) void k_gemm1(
const bf16* __restrict__ x, const bf16* __restrict__ w1r,
const bf16* __restrict__ w1s, const int* __restrict__ stok,
const int* __restrict__ offsets, const int* __restrict__ blk_expert,
const int* __restrict__ blk_row, const int* __restrict__ total_nb,
bf16* __restrict__ hws, int E, int H, int I, int swz) {
constexpr int MT = BMT / 32; // m-tiles of 16 per warp
constexpr int NST = (BMT == 128) ? 4 : NSTG1;
constexpr int STG = (BMT + 2 * BN1) * AST; // per-stage smem elems
// Swizzle (large N only): consecutive scheduled slots cover a PAIR of
// adjacent m-blocks sweeping all n-slices together, so both A tiles stay
// L2-resident AND same-expert pairs stream identical weight tiles
// concurrently (L2 catches the re-read). At small N weights are read once
// either way and the swizzle only hurts scheduling.
int bid, n0;
if (swz) {
int ny = gridDim.y;
int G = swz; // m-blocks per super-tile group
int nxg = (gridDim.x / G) * G;
int lin = blockIdx.y * gridDim.x + blockIdx.x;
int body = nxg * ny;
if (lin < body) {
bid = (lin / (G * ny)) * G + (lin % G);
n0 = ((lin / G) % ny) * BN1;
} else {
bid = nxg + (lin - body) / ny;
n0 = ((lin - body) % ny) * BN1;
}
} else {
bid = blockIdx.x;
n0 = blockIdx.y * BN1;
}
if (bid >= *total_nb) return;
int e = blk_expert[bid];
int row0 = blk_row[bid];
int rows = min(BMT, offsets[e + 1] - row0);
// Weights arrive pre-packed k-tiled: per (expert, n-block j = n0/64) the B
// stream is one contiguous 2MB run of [128 x BK] 8KB tiles (rows 0-63
// gate, 64-127 up), so DRAM sees sequential reads, not 64B/row hops.
const bf16* w1 = (e < E) ? w1r + (size_t)e * (2 * (size_t)I * H)
: w1s + (size_t)(e - E) * (2 * (size_t)I * H);
const bf16* bp = w1 + (size_t)(n0 >> 6) * ((size_t)H * (2 * BN1));
extern __shared__ bf16 smem[];
// NST stages, each [BMT][AST] A then [2*BN1][AST] B (gate|up)
int tid = threadIdx.x;
// A loads: BMT rows x 4 chunks of 16B, gathered by token (BMT/64 passes)
int ar = tid >> 2, ac = (tid & 3) * 8;
const bf16* agp[BMT / 64];
#pragma unroll
for (int p = 0; p < BMT / 64; p++)
agp[p] = x + (size_t)stok[row0 + min(ar + p * 64, rows - 1)] * H + ac;
auto load_stage = [&](int stage, int k0) {
bf16* A = smem + stage * STG;
bf16* B = A + BMT * AST;
#pragma unroll
for (int p = 0; p < BMT / 64; p++)
cp_async16(A + (ar + p * 64) * AST + ac, agp[p] + k0);
// B tile: 128 x BK contiguous at bp + k0*128 -> 512 16B chunks.
const bf16* src = bp + (size_t)k0 * (2 * BN1);
#pragma unroll
for (int p = 0; p < 2; p++) {
int q = tid + p * 256;
cp_async16(B + (q >> 2) * AST + (q & 3) * 8, src + q * 8);
}
cp_commit();
};
#pragma unroll
for (int s = 0; s < NST - 1; s++) load_stage(s, s * BK);
int wid = tid >> 5, lane = tid & 31;
bool is_up = wid >= 4;
int w4 = wid & 3;
int wm = (w4 & 1) * (BMT / 2); // 2x2 warp grid over BMT x 64
int wn = (w4 >> 1) * 32;
float acc[MT][4][4];
#pragma unroll
for (int mt = 0; mt < MT; mt++)
#pragma unroll
for (int nt = 0; nt < 4; nt++)
#pragma unroll
for (int r = 0; r < 4; r++) acc[mt][nt][r] = 0.f;
int nk = H / BK;
for (int kb = 0; kb < nk; kb++) {
// NST-1 groups pending before wait: retire exactly tile kb's group.
// Issuing the next load AFTER the barrier makes a tail barrier
// unnecessary: stage (kb-1)%NST is not read this iteration, and the
// issue that targets stage kb%NST sits behind the next top barrier.
cp_wait<NST - 2>();
__syncthreads();
int kn = kb + NST - 1;
load_stage(kn % NST, min(kn, nk - 1) * BK); // clamped dummy at tail
const bf16* A = smem + (kb % NST) * STG;
const bf16* B = A + BMT * AST + (is_up ? BN1 * AST : 0);
#pragma unroll
for (int ks = 0; ks < 2; ks++) {
uint32_t a[MT][4];
#pragma unroll
for (int mt = 0; mt < MT; mt++) {
const bf16* p =
A + (wm + mt * 16 + (lane & 15)) * AST + ks * 16 + (lane >> 4) * 8;
ldsm_x4(a[mt][0], a[mt][1], a[mt][2], a[mt][3], p);
}
#pragma unroll
for (int nh = 0; nh < 2; nh++) {
// x4.trans: covers ntiles 2nh, 2nh+1 for this k16 step
const bf16* p = B +
(wn + nh * 16 + (lane & 7) + ((lane & 8) ? 8 : 0)) * AST +
ks * 16 + ((lane & 16) ? 8 : 0);
uint32_t b0, b1, b2, b3;
ldsm_x4(b0, b1, b2, b3, p);
#pragma unroll
for (int mt = 0; mt < MT; mt++) {
mma16816(acc[mt][2 * nh + 0], a[mt][0], a[mt][1], a[mt][2],
a[mt][3], b0, b2);
mma16816(acc[mt][2 * nh + 1], a[mt][0], a[mt][1], a[mt][2],
a[mt][3], b1, b3);
}
}
}
// No tail barrier: the only write into stage kb%NST is issued at
// iter kb+1, after that iteration's top barrier.
}
cp_wait<0>();
// Epilogue: gate warps stash fp32 accums in smem; up warps do silu*mul.
__syncthreads();
const int GST = BN1 + 4;
float* sG = (float*)smem; // BMT x 68 f32 <= 34816B, fits dyn smem
int gid = lane >> 2, tq = lane & 3;
if (!is_up) {
#pragma unroll
for (int mt = 0; mt < MT; mt++)
#pragma unroll
for (int nt = 0; nt < 4; nt++) {
int r0m = wm + mt * 16 + gid;
int c0 = wn + nt * 8 + tq * 2;
sG[r0m * GST + c0] = acc[mt][nt][0];
sG[r0m * GST + c0 + 1] = acc[mt][nt][1];
sG[(r0m + 8) * GST + c0] = acc[mt][nt][2];
sG[(r0m + 8) * GST + c0 + 1] = acc[mt][nt][3];
}
}
__syncthreads();
if (is_up) {
#pragma unroll
for (int mt = 0; mt < MT; mt++)
#pragma unroll
for (int nt = 0; nt < 4; nt++)
#pragma unroll
for (int half = 0; half < 2; half++) {
int m = wm + mt * 16 + gid + half * 8;
if (m >= rows) continue;
int c = wn + nt * 8 + tq * 2;
float g0 = sG[m * GST + c], g1 = sG[m * GST + c + 1];
float u0 = acc[mt][nt][half * 2], u1 = acc[mt][nt][half * 2 + 1];
float h0 = g0 / (1.f + __expf(-g0)) * u0;
float h1 = g1 / (1.f + __expf(-g1)) * u1;
__nv_bfloat162 hv = __floats2bfloat162_rn(h0, h1);
*(__nv_bfloat162*)(hws + (size_t)(row0 + m) * I + n0 + c) = hv;
}
}
}
// ---------------------------------------------------------------- GEMM2
// ys[j, n] = w_j * (h[j] . w2_n); per-slot rows, summed later by k_reduce.
// Block: BM=64 rows x BN2=64 out cols; 8 warps 2x4, warp tile 32x16.
#define BN2 64
template <int BMT>
__global__ __launch_bounds__(256) void k_gemm2(
const bf16* __restrict__ hws, const bf16* __restrict__ w2r,
const bf16* __restrict__ w2s, const int* __restrict__ stok,
const float* __restrict__ sw, const int* __restrict__ offsets,
const int* __restrict__ blk_expert, const int* __restrict__ blk_row,
const int* __restrict__ total_nb, bf16* __restrict__ ys, int E, int H,
int I, int swz) {
constexpr int MT = BMT / 32;
constexpr int NST = (BMT == 128) ? 6 : NSTG2;
constexpr int STG = (BMT + BN2) * AST;
int bid, n0;
if (swz) {
// Grouped-m super-tile swizzle; see k_gemm1.
int ny = gridDim.y;
int G = swz;
int nxg = (gridDim.x / G) * G;
int lin = blockIdx.y * gridDim.x + blockIdx.x;
int body = nxg * ny;
if (lin < body) {
bid = (lin / (G * ny)) * G + (lin % G);
n0 = ((lin / G) % ny) * BN2;
} else {
bid = nxg + (lin - body) / ny;
n0 = ((lin - body) % ny) * BN2;
}
} else {
bid = blockIdx.x;
n0 = blockIdx.y * BN2;
}
if (bid >= *total_nb) return;
int e = blk_expert[bid];
int row0 = blk_row[bid];
int rows = min(BMT, offsets[e + 1] - row0);
// Pre-packed k-tiled (see k_gemm1): per (expert, n-block) contiguous
// [64 x BK] tiles.
const bf16* w2 = (e < E) ? w2r + (size_t)e * ((size_t)H * I)
: w2s + (size_t)(e - E) * ((size_t)H * I);
const bf16* bp = w2 + (size_t)(n0 >> 6) * ((size_t)I * BN2);
extern __shared__ bf16 smem[];
// NST stages, each [BMT][AST] A then [BN2][AST] B
int tid = threadIdx.x;
int ar = tid >> 2, ac = (tid & 3) * 8;
const bf16* agp[BMT / 64];
#pragma unroll
for (int p = 0; p < BMT / 64; p++)
agp[p] = hws + (size_t)(row0 + min(ar + p * 64, rows - 1)) * I + ac;
auto load_stage = [&](int stage, int k0) {
bf16* A = smem + stage * STG;
bf16* B = A + BMT * AST;
#pragma unroll
for (int p = 0; p < BMT / 64; p++)
cp_async16(A + (ar + p * 64) * AST + ac, agp[p] + k0);
// B tile: 64 x BK contiguous at bp + k0*64 -> 256 16B chunks.
cp_async16(B + ar * AST + ac, bp + (size_t)k0 * BN2 + tid * 8);
cp_commit();
};
#pragma unroll
for (int s = 0; s < NST - 1; s++) load_stage(s, s * BK);
int wid = tid >> 5, lane = tid & 31;
int wm = (wid & 1) * (BMT / 2); // 2x4 warp grid over BMT x 64
int wn = (wid >> 1) * 16;
float acc[MT][2][4];
#pragma unroll
for (int mt = 0; mt < MT; mt++)
#pragma unroll
for (int nt = 0; nt < 2; nt++)
#pragma unroll
for (int r = 0; r < 4; r++) acc[mt][nt][r] = 0.f;
int nk = I / BK;
for (int kb = 0; kb < nk; kb++) {
cp_wait<NST - 2>();
__syncthreads();
int kn = kb + NST - 1;
load_stage(kn % NST, min(kn, nk - 1) * BK);
const bf16* A = smem + (kb % NST) * STG;
const bf16* B = A + BMT * AST;
#pragma unroll
for (int ks = 0; ks < 2; ks++) {
uint32_t a[MT][4];
#pragma unroll
for (int mt = 0; mt < MT; mt++) {
const bf16* p =
A + (wm + mt * 16 + (lane & 15)) * AST + ks * 16 + (lane >> 4) * 8;
ldsm_x4(a[mt][0], a[mt][1], a[mt][2], a[mt][3], p);
}
const bf16* p = B +
(wn + (lane & 7) + ((lane & 8) ? 8 : 0)) * AST + ks * 16 +
((lane & 16) ? 8 : 0);
uint32_t b0, b1, b2, b3;
ldsm_x4(b0, b1, b2, b3, p);
#pragma unroll
for (int mt = 0; mt < MT; mt++) {
mma16816(acc[mt][0], a[mt][0], a[mt][1], a[mt][2], a[mt][3], b0, b2);
mma16816(acc[mt][1], a[mt][0], a[mt][1], a[mt][2], a[mt][3], b1, b3);
}
}
// No tail barrier (see k_gemm1).
}
cp_wait<0>();
int gid = lane >> 2, tq = lane & 3;
#pragma unroll
for (int mt = 0; mt < MT; mt++)
#pragma unroll
for (int half = 0; half < 2; half++) {
int m = wm + mt * 16 + gid + half * 8;
if (m >= rows) continue;
float w = sw[row0 + m];
bf16* op = ys + (size_t)(row0 + m) * H + n0;
#pragma unroll
for (int nt = 0; nt < 2; nt++) {
int c = wn + nt * 8 + tq * 2;
__nv_bfloat162 v = __floats2bfloat162_rn(
acc[mt][nt][half * 2] * w, acc[mt][nt][half * 2 + 1] * w);
*(__nv_bfloat162*)(op + c) = v;
}
}
}
// ---------------------------------------------------------------- GEMV
// Decode path (tiny N): the tiled GEMMs launch too few blocks to saturate
// DRAM, so stream weights row-parallel instead: one block per (slot, 32
// output cols), one warp row-dotting 4 cols. Activations live in smem.
__device__ inline float dot8(uint4 a, uint4 b) {
const __nv_bfloat162* pa = (const __nv_bfloat162*)&a;
const __nv_bfloat162* pb = (const __nv_bfloat162*)&b;
float s = 0.f;
#pragma unroll
for (int j = 0; j < 4; j++) {
float2 fa = __bfloat1622float2(pa[j]);
float2 fb = __bfloat1622float2(pb[j]);
s = fmaf(fa.x, fb.x, s);
s = fmaf(fa.y, fb.y, s);
}
return s;
}
__device__ inline float warp_sum(float v) {
#pragma unroll
for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(~0u, v, o);
return v;
}
__global__ __launch_bounds__(256) void k_gemv1(
const bf16* __restrict__ x, const bf16* __restrict__ w1r,
const bf16* __restrict__ w1s, const int* __restrict__ stok,
const int* __restrict__ sexp, bf16* __restrict__ hws, int E, int H,
int I) {
int s = blockIdx.x;
int i0 = blockIdx.y * 32;
int e = sexp[s];
const bf16* w1 = (e < E) ? w1r + (size_t)e * (2 * (size_t)I * H)
: w1s + (size_t)(e - E) * (2 * (size_t)I * H);
extern __shared__ bf16 sx[]; // H elems
const uint4* xg = (const uint4*)(x + (size_t)stok[s] * H);
uint4* xs = (uint4*)sx;
for (int c = (int)threadIdx.x; c < H / 8; c += 256) xs[c] = xg[c];
__syncthreads();
int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
int nk = H / 8;
#pragma unroll
for (int ii = 0; ii < 4; ii++) {
int i = i0 + warp * 4 + ii;
const uint4* g = (const uint4*)(w1 + (size_t)i * H);
const uint4* u = (const uint4*)(w1 + (size_t)(I + i) * H);
float ag = 0.f, au = 0.f;
for (int c = lane; c < nk; c += 32) {
uint4 xv = xs[c];
ag += dot8(g[c], xv);
au += dot8(u[c], xv);
}
ag = warp_sum(ag);
au = warp_sum(au);
if (lane == 0) {
float h = ag / (1.f + __expf(-ag)) * au;
hws[(size_t)s * I + i] = __float2bfloat16(h);
}
}
}
__global__ __launch_bounds__(256) void k_gemv2(
const bf16* __restrict__ hws, const bf16* __restrict__ w2r,
const bf16* __restrict__ w2s, const float* __restrict__ sw,
const int* __restrict__ sexp, bf16* __restrict__ ys, int E, int H,
int I) {
int s = blockIdx.x;
int n0 = blockIdx.y * 32;
int e = sexp[s];
const bf16* w2 = (e < E) ? w2r + (size_t)e * ((size_t)H * I)
: w2s + (size_t)(e - E) * ((size_t)H * I);
extern __shared__ bf16 sx[]; // I elems
const uint4* hg = (const uint4*)(hws + (size_t)s * I);
uint4* xs = (uint4*)sx;
for (int c = (int)threadIdx.x; c < I / 8; c += 256) xs[c] = hg[c];
__syncthreads();
float w = sw[s];
int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
int nk = I / 8;
#pragma unroll
for (int ii = 0; ii < 4; ii++) {
int n = n0 + warp * 4 + ii;
const uint4* r = (const uint4*)(w2 + (size_t)n * I);
float a = 0.f;
for (int c = lane; c < nk; c += 32) a += dot8(r[c], xs[c]);
a = warp_sum(a);
if (lane == 0) ys[(size_t)s * H + n] = __float2bfloat16(a * w);
}
}
// ---------------------------------------------------------------- reduce
// out[t, :] = sum_s ys[tok2slot[t*S+s], :]; 16B vectorized, fp32 accum.
__global__ __launch_bounds__(256) void k_reduce(
const bf16* __restrict__ ys, const int* __restrict__ tok2slot, int S,
int H, bf16* __restrict__ out) {
int t = blockIdx.x;
__shared__ int slots[16];
if (threadIdx.x < S) slots[threadIdx.x] = tok2slot[t * S + threadIdx.x];
__syncthreads();
// Each thread owns consecutive 8-elem (16B) chunks.
for (int c0 = threadIdx.x * 8; c0 < H; c0 += blockDim.x * 8) {
float acc[8];
#pragma unroll
for (int r = 0; r < 8; r++) acc[r] = 0.f;
for (int s = 0; s < S; s++) {
uint4 v = *(const uint4*)(ys + (size_t)slots[s] * H + c0);
const __nv_bfloat162* p2 = (const __nv_bfloat162*)&v;
#pragma unroll
for (int j = 0; j < 4; j++) {
float2 f = __bfloat1622float2(p2[j]);
acc[2 * j] += f.x;
acc[2 * j + 1] += f.y;
}
}
uint4 o;
__nv_bfloat162* o2 = (__nv_bfloat162*)&o;
#pragma unroll
for (int j = 0; j < 4; j++)
o2[j] = __floats2bfloat162_rn(acc[2 * j], acc[2 * j + 1]);
*(uint4*)(out + (size_t)t * H + c0) = o;
}
}
// ---------------------------------------------------------------- driver
torch::Tensor moe_forward(torch::Tensor x, torch::Tensor expert_ids,
torch::Tensor expert_weights, torch::Tensor w1_routed,
torch::Tensor w2_routed, torch::Tensor w1_shared,
torch::Tensor w2_shared, torch::Tensor w1p_routed,
torch::Tensor w2p_routed, torch::Tensor w1p_shared,
torch::Tensor w2p_shared) {
TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kBFloat16);
TORCH_CHECK(x.is_contiguous() && w1_routed.is_contiguous() &&
w2_routed.is_contiguous() && w1_shared.is_contiguous() &&
w2_shared.is_contiguous());
auto ids = expert_ids.contiguous();
auto wts = expert_weights.to(torch::kBFloat16).contiguous();
const int T = x.size(0);
const int H = x.size(1);
const int K = ids.size(1);
const int E = w1_routed.size(0);
const int ns = w1_shared.size(0);
const int I = w1_routed.size(1) / 2;
TORCH_CHECK(H % BK == 0 && I % BN1 == 0 && H % BN2 == 0 && I % BK == 0);
const int S = K + ns;
const int64_t N = (int64_t)T * S;
const int NE = E + ns;
const int max_nb = NE + (int)((N + BM - 1) / BM);
auto dev = x.device();
auto i32 = torch::TensorOptions().dtype(torch::kInt32).device(dev);
auto f32 = torch::TensorOptions().dtype(torch::kFloat32).device(dev);
auto bfo = torch::TensorOptions().dtype(torch::kBFloat16).device(dev);
auto counts = torch::zeros({2 * NE}, i32); // [counts | cursor]
auto offsets = torch::empty({NE + 1}, i32);
auto blk = torch::empty({2 * max_nb + 1}, i32); // [expert | row | total]
auto stok = torch::empty({N}, i32);
auto sw = torch::empty({N}, f32);
auto tok2slot = torch::empty({N}, i32);
auto sexp = torch::empty({N}, i32);
auto hws = torch::empty({N, (int64_t)I}, bfo);
auto ys = torch::empty({N, (int64_t)H}, bfo);
auto out = torch::empty({(int64_t)T, (int64_t)H}, bfo);
int* counts_p = counts.data_ptr<int>();
int* cursor_p = counts_p + NE;
int* offsets_p = offsets.data_ptr<int>();
int* blk_e = blk.data_ptr<int>();
int* blk_r = blk_e + max_nb;
int* total_nb = blk_e + 2 * max_nb;
cudaStream_t st = c10::cuda::getCurrentCUDAStream();
{
int total = T * K;
int nb = (total + 255) / 256;
k_count<<<nb, 256, 0, st>>>(ids.data_ptr<int64_t>(), counts_p, total);
}
// Large N: BM=128 halves weight re-reads across m-blocks. Its 1-block/SM
// occupancy only pays off combined with the grouped-m swizzle below
// (measured: 8192 tokens 24.9 -> 21.1 ms). Small N keeps BM=64 for
// block count / occupancy.
const int bm = (N >= 16384) ? 128 : 64;
k_scan<<<1, 32, 0, st>>>(counts_p, T, NE, E, bm, offsets_p, blk_e, blk_r,
total_nb, max_nb);
{
int total = T * S;
int nb = (total + 255) / 256;
k_scatter<<<nb, 256, 0, st>>>(
ids.data_ptr<int64_t>(), (const bf16*)wts.data_ptr(), offsets_p,
cursor_p, T, K, S, E, stok.data_ptr<int>(), sw.data_ptr<float>(),
tok2slot.data_ptr<int>(), sexp.data_ptr<int>());
}
if (T <= 8) {
// Decode path: slot-parallel GEMV saturates DRAM where the tiled
// GEMMs cannot (too few blocks at N = T*S).
dim3 v1((unsigned)N, I / 32);
k_gemv1<<<v1, 256, H * (int)sizeof(bf16), st>>>(
(const bf16*)x.data_ptr(), (const bf16*)w1_routed.data_ptr(),
(const bf16*)w1_shared.data_ptr(), stok.data_ptr<int>(),
sexp.data_ptr<int>(), (bf16*)hws.data_ptr(), E, H, I);
dim3 v2((unsigned)N, H / 32);
k_gemv2<<<v2, 256, I * (int)sizeof(bf16), st>>>(
(const bf16*)hws.data_ptr(), (const bf16*)w2_routed.data_ptr(),
(const bf16*)w2_shared.data_ptr(), sw.data_ptr<float>(),
sexp.data_ptr<int>(), (bf16*)ys.data_ptr(), E, H, I);
k_reduce<<<T, 256, 0, st>>>((const bf16*)ys.data_ptr(),
tok2slot.data_ptr<int>(), S, H,
(bf16*)out.data_ptr());
return out;
}
// Per-variant dynamic smem: gemm1<64>=3*(64+128)*40, gemm1<128>=4*(128+128)*40,
// gemm2<64>=4*(64+64)*40, gemm2<128>=6*(128+64)*40 (all elems, x2 bytes).
const int smem1_64 = NSTG1 * (64 + 2 * BN1) * AST * (int)sizeof(bf16);
const int smem1_128 = 4 * (128 + 2 * BN1) * AST * (int)sizeof(bf16);
const int smem2_64 = NSTG2 * (64 + BN2) * AST * (int)sizeof(bf16);
const int smem2_128 = 6 * (128 + BN2) * AST * (int)sizeof(bf16);
static bool attr_set = false;
if (!attr_set) {
cudaFuncSetAttribute(k_gemm1<64>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem1_64);
cudaFuncSetAttribute(k_gemm1<128>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem1_128);
cudaFuncSetAttribute(k_gemm2<64>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem2_64);
cudaFuncSetAttribute(k_gemm2<128>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem2_128);
attr_set = true;
}
// Swizzle pays off only when m-blocks per expert are plural (weights get
// re-read across m); at small N each weight tile is read once either way.
// Large N also selects the BM=128 variants (bm above).
// swz = m-group size G of the super-tile swizzle (0 = off). Pick G near
// the number of m-blocks per expert so same-expert blocks stream the same
// weight tiles concurrently.
// G=96 measured best at BM=128 across T=4096/8192; small and very large
// G both degrade toward the naive launch orders.
const int swz = (N >= 16384) ? 96 : 0;
dim3 g1(max_nb, I / BN1);
dim3 g2(max_nb, H / BN2);
if (bm == 128) {
k_gemm1<128><<<g1, 256, smem1_128, st>>>(
(const bf16*)x.data_ptr(), (const bf16*)w1p_routed.data_ptr(),
(const bf16*)w1p_shared.data_ptr(), stok.data_ptr<int>(), offsets_p,
blk_e, blk_r, total_nb, (bf16*)hws.data_ptr(), E, H, I, swz);
k_gemm2<128><<<g2, 256, smem2_128, st>>>(
(const bf16*)hws.data_ptr(), (const bf16*)w2p_routed.data_ptr(),
(const bf16*)w2p_shared.data_ptr(), stok.data_ptr<int>(),
sw.data_ptr<float>(), offsets_p, blk_e, blk_r, total_nb,
(bf16*)ys.data_ptr(), E, H, I, swz);
} else {
k_gemm1<64><<<g1, 256, smem1_64, st>>>(
(const bf16*)x.data_ptr(), (const bf16*)w1p_routed.data_ptr(),
(const bf16*)w1p_shared.data_ptr(), stok.data_ptr<int>(), offsets_p,
blk_e, blk_r, total_nb, (bf16*)hws.data_ptr(), E, H, I, swz);
k_gemm2<64><<<g2, 256, smem2_64, st>>>(
(const bf16*)hws.data_ptr(), (const bf16*)w2p_routed.data_ptr(),
(const bf16*)w2p_shared.data_ptr(), stok.data_ptr<int>(),
sw.data_ptr<float>(), offsets_p, blk_e, blk_r, total_nb,
(bf16*)ys.data_ptr(), E, H, I, swz);
}
k_reduce<<<T, 256, 0, st>>>((const bf16*)ys.data_ptr(),
tok2slot.data_ptr<int>(), S, H,
(bf16*)out.data_ptr());
return out;
}
"""
_ext = None
def _get_ext():
global _ext
if _ext is None:
_ext = load_inline(
name="glm52_moe_v2",
cpp_sources=_CPP_SRC,
cuda_sources=_CUDA_SRC,
functions=["moe_forward"],
extra_cuda_cflags=[
"-O3",
"--use_fast_math",
"-gencode=arch=compute_120a,code=sm_120a",
"-std=c++17",
],
verbose=False,
)
return _ext
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._pk = None # (data_ptr, _version) key guarding the packed copies
_get_ext()
def _packed(self):
# k-tiled weight layout so each GEMM pipeline stage load is one
# contiguous 8KB tile (see k_gemm1). Pure layout transform of the
# live parameters, recomputed whenever they are swapped or mutated
# in place (data_ptr / _version change) -- never a cached output.
ws = (self.w1_routed, self.w2_routed, self.w1_shared, self.w2_shared)
key = tuple((w.data_ptr(), w._version) for w in ws)
if self._pk != key:
H, I = self.H, self.I
def pack1(w):
e = w.shape[0]
return (
w.view(e, 2, I // 64, 64, H // 32, 32)
.permute(0, 2, 4, 1, 3, 5)
.contiguous()
)
def pack2(w):
e = w.shape[0]
return (
w.view(e, H // 64, 64, I // 32, 32)
.permute(0, 1, 3, 2, 4)
.contiguous()
)
self._w1p = pack1(self.w1_routed.data)
self._w2p = pack2(self.w2_routed.data)
self._w1sp = pack1(self.w1_shared.data)
self._w2sp = pack2(self.w2_shared.data)
self._pk = key
return self._w1p, self._w2p, self._w1sp, self._w2sp
def forward(
self,
x: torch.Tensor,
expert_ids: torch.Tensor,
expert_weights: torch.Tensor,
) -> torch.Tensor:
w1p, w2p, w1sp, w2sp = self._packed()
return _get_ext().moe_forward(
x.contiguous(),
expert_ids,
expert_weights,
self.w1_routed,
self.w2_routed,
self.w1_shared,
self.w2_shared,
w1p,
w2p,
w1sp,
w2sp,
)
if __name__ == "__main__":
torch.manual_seed(0)
dev = torch.device("cuda:0")
T, E, K, ns, H, I = 64, 256, 8, 1, 4096, 2048
m = Model(T, E, K, ns, H, I).to(dev).eval()
x = torch.randn(T, H, dtype=torch.bfloat16, device=dev)
ids = torch.randint(0, E, (T, K), dtype=torch.int64, device=dev)
w = torch.softmax(torch.randn(T, K, device=dev), -1).to(torch.bfloat16)
with torch.no_grad():
y = m(x, ids, w)
print("out", y.shape, y.dtype, float(y.float().abs().mean()))
20260719_081348_or-fable_anthropic_claude-fable-5_01_glm52_fused_moe