KernelBench cuda · B200
GLM-5.2 Fused MoE Claude Fable 5
6.19%geomean peak fraction across shapes
manually audited: clean
harnessor-fableagent session—total wall—check—benchmark—output tokens—regimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
No per-shape benchmark data archived for this run.
Kernel source (redacted)
"""GLM-5.2-class fused MoE layer — hand-written CUDA (tcgen05 UMMA + mma.sync).
Structure (vLLM-style token grouping, reimplemented from scratch in CUDA):
1. moe_align (hist -> scan -> scatter kernels): expert_ids histogram ->
per-expert BM-padded row offsets -> scatter (token, weight) pairs into
expert-sorted order. The shared expert(s) are appended as expert id E,
E+1, ... with weight 1.0 so the grouped GEMMs treat routed and shared
experts uniformly.
2. gemm1 (grouped): silu(x_gathered @ gate^T) * (x_gathered @ up^T) -> h.
3. gemm2_shared: h_shared @ w2_shared^T with PLAIN fp32 stores into the out
accumulator (every token appears exactly once in the shared segment, so
this initializes out with no memset), then gemm2_routed with the routing
weight applied and fp32 atomicAdd.
4. convert: fp32 -> bf16.
Three engine paths, dispatched per batch size and GPU:
* SM100 (B200-class), tokens/expert >= 32: 5th-gen tensor core (tcgen05
UMMA) path, ~455-520 TF sustained end to end. The accumulator lives in
tensor memory; a single tcgen05.mma covers an M=128 x N=256 tile per k16,
issued by one thread. Operands are staged by cp.async.bulk (TMA 1D) with
mbarrier expect_tx completion — weights are repacked ONCE at init into
the packed core-matrix 8KB blocks the UMMA shared-memory descriptor
walks (LBO=128B, SBO=1024B, no swizzle), x rows are gathered into the
same layout per forward (tc_pack_x), and gemm1's epilogue writes h
already packed for gemm2. A 4-deep smem ring with lookahead 2 decouples
bulk-load latency from tensor issue; per-stage tcgen05.commit ->
mbarrier gates smem reuse. SM120 lacks tcgen05, hence:
* mma.sync path (any SM80+, and small batches): BM in {64, 32} tiles,
cp.async double-buffered pipeline, ldmatrix + mma.sync.m16n8k16,
expert-major tile order (m fastest) for L2 weight reuse.
* T <= 4 decode path: per (token, expert) streaming GEMV (weights are not
shared between tokens at this size, so tiles would only add padding).
All matmul work accumulates in fp32; h is kept in bf16 (the calibrated
"legitimate pipeline" intermediate); final token accumulation is fp32.
CUDA graphs (keyed on input pointers+shape, replayed only on identical
buffers, recompute on live data — see dev_test.py mutation check) remove
launch overhead on tiny batches.
"""
from __future__ import annotations
import os
# The nvcc PATH wrapper on this box is broken (REAL_NVCC empty); point the
# torch extension builder at the real CUDA 12.8 toolkit before importing it.
os.environ.setdefault("CUDA_HOME", "/usr/local/cuda-12.8")
import torch
import torch.nn as nn
CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cstdint>
using bf16 = __nv_bfloat16;
#define DEVI __device__ __forceinline__
DEVI unsigned smem_addr(const void* p) {
return static_cast<unsigned>(__cvta_generic_to_shared(p));
}
// 16B async copy, src_bytes==0 -> full zero-fill (used to mask padding rows).
DEVI void cp_async16(unsigned dst, const void* src, int src_bytes) {
asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;\n"
:: "r"(dst), "l"(src), "r"(src_bytes));
}
DEVI void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N> DEVI void cp_wait() {
asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
}
DEVI void ldmatrix_x4(unsigned addr, unsigned& r0, unsigned& r1, unsigned& r2, unsigned& r3) {
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(float& d0, float& d1, float& d2, float& d3,
unsigned a0, unsigned a1, unsigned a2, unsigned a3,
unsigned b0, unsigned 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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}
DEVI float silu(float v) { return v / (1.0f + __expf(-v)); }
// Swizzled 16B-chunk offset within a [rows][64] bf16 smem tile (8 chunks/row).
DEVI unsigned tile_off(int row, int chunk) {
return static_cast<unsigned>((row * 8 + (chunk ^ (row & 7))) * 16);
}
// ---------------------------------------------------------------------------
// moe_align, split into parallel stages (counts/fill pre-zeroed by memset):
// hist (grid) -> scan (1 block) -> init s_token=-1 (grid, CPU upper bound)
// -> scatter (grid). row_off / mb_cum have E_total+1 entries.
// ---------------------------------------------------------------------------
__global__ void moe_hist_kernel(
const int64_t* __restrict__ ids, int P, int* __restrict__ cnt)
{
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < P) atomicAdd(&cnt[(int)ids[i]], 1);
}
__global__ void moe_scan_kernel(
const int* __restrict__ cnt, int T, int E, int n_shared, int BM,
int* __restrict__ row_off, int* __restrict__ mb_cum)
{
if (threadIdx.x != 0) return;
const int E_total = E + n_shared;
int ro = 0, mc = 0;
for (int e = 0; e < E_total; ++e) {
const int c = (e < E) ? cnt[e] : T;
const int m = (c + BM - 1) / BM;
row_off[e] = ro;
mb_cum[e] = mc;
ro += m * BM;
mc += m;
}
row_off[E_total] = ro;
mb_cum[E_total] = mc;
}
__global__ void moe_initneg_kernel(int* __restrict__ s_token, int rows_max)
{
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < rows_max) s_token[i] = -1;
}
__global__ void moe_scatter_kernel(
const int64_t* __restrict__ ids, const bf16* __restrict__ wts,
int T, int topk, int E, int n_shared,
const int* __restrict__ row_off, int* __restrict__ fill,
int* __restrict__ s_token, float* __restrict__ s_weight)
{
const int P = T * topk;
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < P) {
const int e = (int)ids[i];
const int pos = row_off[e] + atomicAdd(&fill[e], 1);
s_token[pos] = i / topk;
s_weight[pos] = __bfloat162float(wts[i]);
} else if (i < P + n_shared * T) {
const int sp = i - P;
const int es = sp / T, tok = sp % T;
const int pos = row_off[E + es] + tok;
s_token[pos] = tok;
s_weight[pos] = 1.0f;
}
}
// ---------------------------------------------------------------------------
// gemm1: h[row, n] = silu(x_g @ w1[e, n, :]) * (x_g @ w1[e, I+n, :])
// Grid: flat tiles, expert-major, m fastest within (expert, n).
// Warp tile is 32 rows x (BN/WARPS_N) cols of BOTH gate and up.
// ---------------------------------------------------------------------------
template <int BM, int BN, int WARPS_M, int WARPS_N, int STAGES>
__global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2)
gemm1_kernel(
const bf16* __restrict__ x, // (T, H)
const bf16* __restrict__ w1, // (E, 2I, H)
const bf16* __restrict__ w1s, // (n_shared, 2I, H)
const int* __restrict__ s_token,
const int* __restrict__ row_off,
const int* __restrict__ mb_cum,
bf16* __restrict__ hbuf, // (rows_pad, I)
int H, int I, int E, int n_shared)
{
constexpr int BK = 64;
constexpr int NTHREADS = WARPS_M * WARPS_N * 32;
constexpr int ACH = BM * 8; // 16B chunks of A per stage
constexpr int BCH = 2 * BN * 8; // gate + up
constexpr int STAGE_BYTES = (ACH + BCH) * 16;
constexpr int A_BYTES = ACH * 16;
constexpr int WN = BN / WARPS_N; // cols per warp
constexpr int FN = WN / 8; // n8 frags per warp
const int NT = I / BN;
const int E_total = E + n_shared;
const int total_tiles = mb_cum[E_total] * NT;
const int tile = blockIdx.x;
if (tile >= total_tiles) return;
int lo = 0, hi = E_total;
while (hi - lo > 1) {
const int mid = (lo + hi) >> 1;
if (mb_cum[mid] * NT <= tile) lo = mid; else hi = mid;
}
const int e = lo;
const int local = tile - mb_cum[e] * NT;
const int mb = mb_cum[e + 1] - mb_cum[e];
const int mi = local % mb;
const int ni = local / mb;
const int row0 = row_off[e] + mi * BM;
const int n0 = ni * BN;
const bf16* w1e = (e < E)
? (w1 + (size_t)e * (2 * (size_t)I) * H)
: (w1s + (size_t)(e - E) * (2 * (size_t)I) * H);
const bf16* Bg = w1e + (size_t)n0 * H;
const bf16* Bu = w1e + (size_t)(I + n0) * H;
extern __shared__ char smem[];
const int t = threadIdx.x;
// Loader: chunk c = t + NTHREADS*i -> row c/8, chunk col c%8. Row step per
// i is NTHREADS/8 (multiple of 8), so the swizzle XOR term is invariant
// and both smem and global strides are constants: precompute base
// pointers once, bump by BK per stage — keeps the issue path to
// ~2 instructions per cp.async.
constexpr int AITER = ACH / NTHREADS;
constexpr int BITER = BCH / NTHREADS;
constexpr int B2 = BITER / 2;
constexpr int RSTEP = NTHREADS / 8;
constexpr unsigned CH_STRIDE = NTHREADS * 16;
constexpr unsigned BHALF = BN * 8u * 16u;
const int ar = t / 8, ac = t % 8;
int ab[AITER];
const bf16* asrc[AITER];
#pragma unroll
for (int i = 0; i < AITER; ++i) {
const int tk = s_token[row0 + ar + RSTEP * i];
ab[i] = tk < 0 ? 0 : 16;
asrc[i] = x + (size_t)(tk < 0 ? 0 : tk) * H + ac * 8;
}
const unsigned adst0 = smem_addr(smem + tile_off(ar, ac));
const bf16* bg = Bg + (size_t)ar * H + ac * 8;
const bf16* bu = Bu + (size_t)ar * H + ac * 8;
const unsigned bdst0 = smem_addr(smem + A_BYTES + tile_off(ar, ac));
const int KT = H / BK;
auto issue_stage = [&](int stage) {
const unsigned soff = stage * STAGE_BYTES;
#pragma unroll
for (int i = 0; i < AITER; ++i) {
cp_async16(adst0 + soff + i * CH_STRIDE, asrc[i], ab[i]);
asrc[i] += BK;
}
#pragma unroll
for (int j = 0; j < B2; ++j) {
cp_async16(bdst0 + soff + j * CH_STRIDE, bg + (size_t)j * RSTEP * H, 16);
cp_async16(bdst0 + soff + BHALF + j * CH_STRIDE, bu + (size_t)j * RSTEP * H, 16);
}
bg += BK;
bu += BK;
};
float accg[2][FN][4], accu[2][FN][4];
#pragma unroll
for (int a = 0; a < 2; ++a)
#pragma unroll
for (int b = 0; b < FN; ++b)
#pragma unroll
for (int c = 0; c < 4; ++c) { accg[a][b][c] = 0.f; accu[a][b][c] = 0.f; }
const int warp = t / 32, lane = t % 32;
const int wm = warp % WARPS_M, wn = warp / WARPS_M;
#pragma unroll
for (int s = 0; s < STAGES - 1; ++s) {
issue_stage(s);
cp_commit();
}
for (int kt = 0; kt < KT; ++kt) {
cp_wait<STAGES - 2>();
__syncthreads();
const int nxt = kt + STAGES - 1;
if (nxt < KT) issue_stage(nxt % STAGES);
cp_commit();
char* base = smem + (kt % STAGES) * STAGE_BYTES;
char* bbase = base + A_BYTES;
const int lrow = lane & 15;
const int lch = lane >> 4; // chunk half within k16
unsigned a[2][2][4], bgf[2][FN][2], buf[2][FN][2];
auto load_frags = [&](int kk, int buf_i) {
#pragma unroll
for (int fm = 0; fm < 2; ++fm) {
const int row = wm * (BM / WARPS_M) + fm * 16 + lrow;
ldmatrix_x4(smem_addr(base + tile_off(row, kk * 2 + lch)),
a[buf_i][fm][0], a[buf_i][fm][1], a[buf_i][fm][2], a[buf_i][fm][3]);
}
#pragma unroll
for (int g = 0; g < FN / 2; ++g) {
unsigned r0, r1, r2, r3;
int row = wn * WN + g * 16 + lrow;
ldmatrix_x4(smem_addr(bbase + tile_off(row, kk * 2 + lch)), r0, r1, r2, r3);
bgf[buf_i][g * 2][0] = r0; bgf[buf_i][g * 2][1] = r2;
bgf[buf_i][g * 2 + 1][0] = r1; bgf[buf_i][g * 2 + 1][1] = r3;
ldmatrix_x4(smem_addr(bbase + tile_off(BN + row, kk * 2 + lch)), r0, r1, r2, r3);
buf[buf_i][g * 2][0] = r0; buf[buf_i][g * 2][1] = r2;
buf[buf_i][g * 2 + 1][0] = r1; buf[buf_i][g * 2 + 1][1] = r3;
}
};
load_frags(0, 0);
#pragma unroll
for (int kk = 0; kk < BK / 16; ++kk) {
const int cur = kk & 1;
if (kk + 1 < BK / 16) load_frags(kk + 1, cur ^ 1);
#pragma unroll
for (int fm = 0; fm < 2; ++fm)
#pragma unroll
for (int fn = 0; fn < FN; ++fn) {
mma_bf16(accg[fm][fn][0], accg[fm][fn][1], accg[fm][fn][2], accg[fm][fn][3],
a[cur][fm][0], a[cur][fm][1], a[cur][fm][2], a[cur][fm][3],
bgf[cur][fn][0], bgf[cur][fn][1]);
mma_bf16(accu[fm][fn][0], accu[fm][fn][1], accu[fm][fn][2], accu[fm][fn][3],
a[cur][fm][0], a[cur][fm][1], a[cur][fm][2], a[cur][fm][3],
buf[cur][fn][0], buf[cur][fn][1]);
}
}
}
cp_wait<0>();
// epilogue: h = silu(g) * u, bf16, direct global stores (4B pairs)
#pragma unroll
for (int fm = 0; fm < 2; ++fm) {
#pragma unroll
for (int fn = 0; fn < FN; ++fn) {
#pragma unroll
for (int half = 0; half < 2; ++half) {
const int row = wm * (BM / WARPS_M) + fm * 16 + (lane >> 2) + half * 8;
const int col = wn * WN + fn * 8 + (lane & 3) * 2;
const float g0 = accg[fm][fn][half * 2 + 0];
const float g1 = accg[fm][fn][half * 2 + 1];
const float u0 = accu[fm][fn][half * 2 + 0];
const float u1 = accu[fm][fn][half * 2 + 1];
__nv_bfloat162 hv = __floats2bfloat162_rn(silu(g0) * u0, silu(g1) * u1);
*reinterpret_cast<__nv_bfloat162*>(
hbuf + (size_t)(row0 + row) * I + n0 + col) = hv;
}
}
}
}
// ---------------------------------------------------------------------------
// gemm2: out[token, n] (+)= w * (h[row, :] @ w2[e, n, :])
// SHARED=false: routed experts, binary-search tile map, fp32 atomicAdd.
// SHARED=true: shared experts, direct tile map, plain store (es==0) which
// initializes out; es>0 would atomicAdd (n_shared==1 here).
// ---------------------------------------------------------------------------
template <int BM, int BN, int WARPS_M, int WARPS_N, int STAGES, bool SHARED>
__global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2)
gemm2_kernel(
const bf16* __restrict__ hbuf, // (rows_pad, I)
const bf16* __restrict__ w2, // (E, H, I)
const bf16* __restrict__ w2s, // (n_shared, H, I)
const int* __restrict__ s_token,
const float* __restrict__ s_weight,
const int* __restrict__ row_off,
const int* __restrict__ mb_cum,
float* __restrict__ outf, // (T, H)
int H, int I, int E, int n_shared, int mb_sh)
{
constexpr int BK = 64;
constexpr int NTHREADS = WARPS_M * WARPS_N * 32;
constexpr int ACH = BM * 8;
constexpr int BCH = BN * 8;
constexpr int STAGE_BYTES = (ACH + BCH) * 16;
constexpr int A_BYTES = ACH * 16;
constexpr int WN = BN / WARPS_N;
constexpr int FN = WN / 8;
const int NT = H / BN;
int e, row0, n0;
if (SHARED) {
const int smi = blockIdx.x / NT;
const int ni = blockIdx.x % NT;
const int es = smi / mb_sh;
const int mi = smi % mb_sh;
e = E + es;
row0 = row_off[E] + (es * mb_sh + mi) * BM;
n0 = ni * BN;
} else {
const int total_tiles = mb_cum[E] * NT;
const int tile = blockIdx.x;
if (tile >= total_tiles) return;
int lo = 0, hi = E;
while (hi - lo > 1) {
const int mid = (lo + hi) >> 1;
if (mb_cum[mid] * NT <= tile) lo = mid; else hi = mid;
}
e = lo;
const int local = tile - mb_cum[e] * NT;
const int mb = mb_cum[e + 1] - mb_cum[e];
const int mi = local % mb;
const int ni = local / mb;
row0 = row_off[e] + mi * BM;
n0 = ni * BN;
}
const bf16* w2e = (e < E)
? (w2 + (size_t)e * (size_t)H * I)
: (w2s + (size_t)(e - E) * (size_t)H * I);
const bf16* Bp = w2e + (size_t)n0 * I;
extern __shared__ char smem[];
const int t = threadIdx.x;
constexpr int AITER = ACH / NTHREADS;
constexpr int BITER = BCH / NTHREADS;
constexpr int RSTEP = NTHREADS / 8;
constexpr unsigned CH_STRIDE = NTHREADS * 16;
const int ar = t / 8, ac = t % 8;
const int KT = I / BK;
const bf16* asrc = hbuf + (size_t)(row0 + ar) * I + ac * 8;
const unsigned adst0 = smem_addr(smem + tile_off(ar, ac));
const bf16* bsrc = Bp + (size_t)ar * I + ac * 8;
const unsigned bdst0 = smem_addr(smem + A_BYTES + tile_off(ar, ac));
auto issue_stage = [&](int stage) {
const unsigned soff = stage * STAGE_BYTES;
#pragma unroll
for (int i = 0; i < AITER; ++i)
cp_async16(adst0 + soff + i * CH_STRIDE, asrc + (size_t)i * RSTEP * I, 16);
#pragma unroll
for (int j = 0; j < BITER; ++j)
cp_async16(bdst0 + soff + j * CH_STRIDE, bsrc + (size_t)j * RSTEP * I, 16);
asrc += BK;
bsrc += BK;
};
float acc[2][FN][4];
#pragma unroll
for (int a = 0; a < 2; ++a)
#pragma unroll
for (int b = 0; b < FN; ++b)
#pragma unroll
for (int c = 0; c < 4; ++c) acc[a][b][c] = 0.f;
const int warp = t / 32, lane = t % 32;
const int wm = warp % WARPS_M, wn = warp / WARPS_M;
#pragma unroll
for (int s = 0; s < STAGES - 1; ++s) {
issue_stage(s);
cp_commit();
}
for (int kt = 0; kt < KT; ++kt) {
cp_wait<STAGES - 2>();
__syncthreads();
const int nxt = kt + STAGES - 1;
if (nxt < KT) issue_stage(nxt % STAGES);
cp_commit();
char* base = smem + (kt % STAGES) * STAGE_BYTES;
char* bbase = base + A_BYTES;
const int lrow = lane & 15;
const int lch = lane >> 4;
unsigned a[2][2][4], bb[2][FN][2];
auto load_frags = [&](int kk, int buf_i) {
#pragma unroll
for (int fm = 0; fm < 2; ++fm) {
const int row = wm * (BM / WARPS_M) + fm * 16 + lrow;
ldmatrix_x4(smem_addr(base + tile_off(row, kk * 2 + lch)),
a[buf_i][fm][0], a[buf_i][fm][1], a[buf_i][fm][2], a[buf_i][fm][3]);
}
#pragma unroll
for (int g = 0; g < FN / 2; ++g) {
unsigned r0, r1, r2, r3;
const int row = wn * WN + g * 16 + lrow;
ldmatrix_x4(smem_addr(bbase + tile_off(row, kk * 2 + lch)), r0, r1, r2, r3);
bb[buf_i][g * 2][0] = r0; bb[buf_i][g * 2][1] = r2;
bb[buf_i][g * 2 + 1][0] = r1; bb[buf_i][g * 2 + 1][1] = r3;
}
};
load_frags(0, 0);
#pragma unroll
for (int kk = 0; kk < BK / 16; ++kk) {
const int cur = kk & 1;
if (kk + 1 < BK / 16) load_frags(kk + 1, cur ^ 1);
#pragma unroll
for (int fm = 0; fm < 2; ++fm)
#pragma unroll
for (int fn = 0; fn < FN; ++fn)
mma_bf16(acc[fm][fn][0], acc[fm][fn][1], acc[fm][fn][2], acc[fm][fn][3],
a[cur][fm][0], a[cur][fm][1], a[cur][fm][2], a[cur][fm][3],
bb[cur][fn][0], bb[cur][fn][1]);
}
}
cp_wait<0>();
// epilogue: weighted accumulate into outf
int tok_r[2][2];
float wt_r[2][2];
#pragma unroll
for (int fm = 0; fm < 2; ++fm)
#pragma unroll
for (int half = 0; half < 2; ++half) {
const int row = row0 + wm * (BM / WARPS_M) + fm * 16 + (lane >> 2) + half * 8;
tok_r[fm][half] = s_token[row];
wt_r[fm][half] = SHARED ? 1.0f : s_weight[row];
}
#pragma unroll
for (int fm = 0; fm < 2; ++fm) {
#pragma unroll
for (int half = 0; half < 2; ++half) {
const int tk = tok_r[fm][half];
if (tk < 0) continue;
const float w = wt_r[fm][half];
float* orow = outf + (size_t)tk * H + n0;
#pragma unroll
for (int fn = 0; fn < FN; ++fn) {
const int col = wn * WN + fn * 8 + (lane & 3) * 2;
const float v0 = acc[fm][fn][half * 2 + 0] * w;
const float v1 = acc[fm][fn][half * 2 + 1] * w;
if (SHARED && e == E) { // first shared expert: plain store
orow[col] = v0;
orow[col + 1] = v1;
} else {
atomicAdd(&orow[col], v0);
atomicAdd(&orow[col + 1], v1);
}
}
}
}
}
// ---------------------------------------------------------------------------
// Decode path (tiny T): per (token, expert) pair GEMV. Weights are not shared
// between tokens at T<=4, so a straight streaming dot-product per output row
// beats the tile path (no 32-row padding, no mma latency chains).
// gemv1: h[p, n] = silu(x_tok . w1[e, n, :]) * (x_tok . w1[e, I+n, :])
// CTA: one pair, ROWS h-cols; x row staged in smem; one warp per h-col.
// ---------------------------------------------------------------------------
template <int WARPS, int RPW> // rows (h cols) per warp
__global__ void __launch_bounds__(WARPS * 32)
gemv1_kernel(
const bf16* __restrict__ x,
const bf16* __restrict__ w1,
const bf16* __restrict__ w1s,
const int64_t* __restrict__ ids,
int T, int topk, int E, int n_shared, int H, int I,
bf16* __restrict__ hbuf)
{
constexpr int ROWS = WARPS * RPW;
const int P9 = T * (topk + n_shared);
const int pair = blockIdx.x / (I / ROWS);
const int nchunk = blockIdx.x % (I / ROWS);
if (pair >= P9) return;
int tok, e;
if (pair < T * topk) {
tok = pair / topk;
e = (int)ids[pair];
} else {
const int sp = pair - T * topk;
e = E + sp / T;
tok = sp % T;
}
const bf16* w1e = (e < E)
? (w1 + (size_t)e * (2 * (size_t)I) * H)
: (w1s + (size_t)(e - E) * (2 * (size_t)I) * H);
extern __shared__ char smem[];
bf16* xs = reinterpret_cast<bf16*>(smem);
const int t = threadIdx.x;
// stage x row (H bf16) into smem, 16B per thread per iter
for (int i = t; i < H / 8; i += WARPS * 32) {
reinterpret_cast<float4*>(xs)[i] =
reinterpret_cast<const float4*>(x + (size_t)tok * H)[i];
}
__syncthreads();
// Split warps: first half streams gate rows, second half up rows — each
// warp reads ONE contiguous slab per row (matches the ~5 TB/s pattern of
// the down GEMV) and the pair is combined through smem.
float* red = reinterpret_cast<float*>(xs + H); // ROWS floats x 2
const int warp = t / 32, lane = t % 32;
const int half = warp >= WARPS / 2;
const int wrow = warp % (WARPS / 2);
#pragma unroll
for (int r = 0; r < 2 * RPW; ++r) {
const int n = nchunk * ROWS + wrow * 2 * RPW + r;
const bf16* wr = w1e + (size_t)(half ? I + n : n) * H;
float a0 = 0.f, a1 = 0.f;
#pragma unroll 4
for (int k = lane * 8; k < H; k += 32 * 8) {
const float4 wv = *reinterpret_cast<const float4*>(wr + k);
const float4 xv = *reinterpret_cast<const float4*>(xs + k);
const __nv_bfloat162* wp = reinterpret_cast<const __nv_bfloat162*>(&wv);
const __nv_bfloat162* xp = reinterpret_cast<const __nv_bfloat162*>(&xv);
#pragma unroll
for (int j = 0; j < 4; ++j) {
const float2 xf = __bfloat1622float2(xp[j]);
const float2 wf = __bfloat1622float2(wp[j]);
a0 = fmaf(xf.x, wf.x, a0);
a1 = fmaf(xf.y, wf.y, a1);
}
}
float a = a0 + a1;
#pragma unroll
for (int s = 16; s > 0; s >>= 1) a += __shfl_down_sync(0xffffffffu, a, s);
if (lane == 0) red[half * ROWS + wrow * 2 * RPW + r] = a;
}
__syncthreads();
const int base = nchunk * ROWS;
for (int n = t; n < ROWS; n += WARPS * 32) {
hbuf[(size_t)pair * I + base + n] = __float2bfloat16(silu(red[n]) * red[ROWS + n]);
}
}
// gemv2: outf[tok, m] += w_p * (h[p, :] . w2[e, m, :]) (outf pre-zeroed)
template <int WARPS, int RPW>
__global__ void __launch_bounds__(WARPS * 32)
gemv2_kernel(
const bf16* __restrict__ hbuf,
const bf16* __restrict__ w2,
const bf16* __restrict__ w2s,
const int64_t* __restrict__ ids,
const bf16* __restrict__ wts,
int T, int topk, int E, int n_shared, int H, int I,
float* __restrict__ outf)
{
constexpr int ROWS = WARPS * RPW;
const int P9 = T * (topk + n_shared);
const int pair = blockIdx.x / (H / ROWS);
const int mchunk = blockIdx.x % (H / ROWS);
if (pair >= P9) return;
int tok, e;
float wgt;
if (pair < T * topk) {
tok = pair / topk;
e = (int)ids[pair];
wgt = __bfloat162float(wts[pair]);
} else {
const int sp = pair - T * topk;
e = E + sp / T;
tok = sp % T;
wgt = 1.0f;
}
const bf16* w2e = (e < E)
? (w2 + (size_t)e * (size_t)H * I)
: (w2s + (size_t)(e - E) * (size_t)H * I);
extern __shared__ char smem[];
bf16* hs = reinterpret_cast<bf16*>(smem);
const int t = threadIdx.x;
for (int i = t; i < I / 8; i += WARPS * 32) {
reinterpret_cast<float4*>(hs)[i] =
reinterpret_cast<const float4*>(hbuf + (size_t)pair * I)[i];
}
__syncthreads();
const int warp = t / 32, lane = t % 32;
#pragma unroll
for (int r = 0; r < RPW; ++r) {
const int m = mchunk * ROWS + warp * RPW + r;
const bf16* wr = w2e + (size_t)m * I;
float a0 = 0.f, a1 = 0.f;
#pragma unroll 4
for (int k = lane * 8; k < I; k += 32 * 8) {
const float4 wv = *reinterpret_cast<const float4*>(wr + k);
const float4 hv = *reinterpret_cast<const float4*>(hs + k);
const __nv_bfloat162* wp = reinterpret_cast<const __nv_bfloat162*>(&wv);
const __nv_bfloat162* hp = reinterpret_cast<const __nv_bfloat162*>(&hv);
#pragma unroll
for (int j = 0; j < 4; ++j) {
const float2 wf = __bfloat1622float2(wp[j]);
const float2 hf = __bfloat1622float2(hp[j]);
a0 = fmaf(hf.x, wf.x, a0);
a1 = fmaf(hf.y, wf.y, a1);
}
}
float acc = a0 + a1;
#pragma unroll
for (int s = 16; s > 0; s >>= 1) acc += __shfl_down_sync(0xffffffffu, acc, s);
if (lane == 0) atomicAdd(&outf[(size_t)tok * H + m], acc * wgt);
}
}
// ---------------------------------------------------------------------------
// tcgen05 (SM100 5th-gen tensor core / UMMA) grouped-GEMM path. Compiled only
// when the build targets sm_100a (KBH_TC defined by the host builder). The
// accumulator lives in tensor memory; one tcgen05.mma per k16 covers the full
// 128x128 tile, so warps only feed cp.async and run the epilogue.
// smem operand layout: packed 8x8-elem core matrices (verified empirically):
// chunk(row, kc) at (row>>3)*(8*BKB) + kc*128 + (row&7)*16, BKB = BK*2 bytes
// desc: LBO=128 (k-block stride), SBO=8*BKB (8-row group stride), no swizzle
// ---------------------------------------------------------------------------
#ifdef KBH_TC
DEVI uint64_t tc_desc(unsigned saddr, unsigned sbo) {
uint64_t d = 0;
d |= (uint64_t)((saddr >> 4) & 0x3FFF);
d |= (uint64_t)(128u >> 4) << 16; // LBO = 128B
d |= (uint64_t)((sbo >> 4) & 0x3FFF) << 32;
return d; // swizzle mode 0
}
DEVI void tc_mma(unsigned tmem, uint64_t da, uint64_t db, unsigned idesc, int accum) {
if (accum) {
asm volatile(
"{.reg .pred p; setp.eq.u32 p, 1, 1;\n"
"tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;}\n"
:: "r"(tmem), "l"(da), "l"(db), "r"(idesc));
} else {
asm volatile(
"{.reg .pred p; setp.eq.u32 p, 1, 0;\n"
"tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;}\n"
:: "r"(tmem), "l"(da), "l"(db), "r"(idesc));
}
}
DEVI void tc_commit(unsigned mbar) {
asm volatile(
"tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];\n"
:: "r"(mbar));
}
DEVI void mbar_wait(unsigned mbar, unsigned parity) {
unsigned done = 0;
while (!done) {
asm volatile(
"{.reg .pred p; mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2; selp.u32 %0, 1, 0, p;}\n"
: "=r"(done) : "r"(mbar), "r"(parity));
}
}
DEVI void tc_ld32(unsigned taddr, float* v) {
asm volatile(
"tcgen05.ld.sync.aligned.32x32b.x32.b32 "
"{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15,"
"%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, [%32];\n"
: "=f"(v[0]),"=f"(v[1]),"=f"(v[2]),"=f"(v[3]),"=f"(v[4]),"=f"(v[5]),"=f"(v[6]),"=f"(v[7]),
"=f"(v[8]),"=f"(v[9]),"=f"(v[10]),"=f"(v[11]),"=f"(v[12]),"=f"(v[13]),"=f"(v[14]),"=f"(v[15]),
"=f"(v[16]),"=f"(v[17]),"=f"(v[18]),"=f"(v[19]),"=f"(v[20]),"=f"(v[21]),"=f"(v[22]),"=f"(v[23]),
"=f"(v[24]),"=f"(v[25]),"=f"(v[26]),"=f"(v[27]),"=f"(v[28]),"=f"(v[29]),"=f"(v[30]),"=f"(v[31])
: "r"(taddr));
asm volatile("tcgen05.wait::ld.sync.aligned;\n");
}
// Weights (and gemm2's h operand) are pre-packed in GLOBAL memory into the
// same core-matrix order the UMMA smem descriptor expects, as contiguous 8KB
// blocks of (64 rows x 64 k): [rg 0..7][kc 0..7][r 0..7][16B]. A whole stage
// is then loaded by 2+2 single-thread cp.async.bulk (TMA 1D) copies with
// mbarrier expect_tx completion -- no per-chunk cp.async issue cost at all.
// x rows are gathered once per forward into the same packed layout (xperm)
// by tc_pack_x_kernel; gemm1's epilogue writes h already packed for gemm2.
DEVI void tc_bulk(unsigned dst, const void* src, int bytes, unsigned mbar) {
asm volatile(
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n"
:: "r"(dst), "l"(src), "r"(bytes), "r"(mbar));
}
DEVI void mbar_expect_tx(unsigned mbar, unsigned bytes) {
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n"
:: "r"(mbar), "r"(bytes));
}
// gather x rows (via s_token) into packed 8KB blocks: xperm[rb][kb][block]
__global__ void tc_pack_x_kernel(
const bf16* __restrict__ x,
const int* __restrict__ s_token,
const int* __restrict__ row_off_end, // &row_off[E_total]: rows_pad
bf16* __restrict__ xperm,
int H)
{
const int KB = H / 64;
const int rb = blockIdx.x / KB;
const int kb = blockIdx.x % KB;
const int rows_pad = *row_off_end;
if (rb * 64 >= rows_pad) return;
const int t = threadIdx.x; // 256 threads, 2 chunks each
char* dst = reinterpret_cast<char*>(xperm) + ((size_t)blockIdx.x) * 8192;
#pragma unroll
for (int i = 0; i < 2; ++i) {
const int c = t + 256 * i; // [rg][kc][r] chunk index
const int rg = c / 64, kc = (c / 8) % 8, r = c % 8;
const int row = rb * 64 + rg * 8 + r;
const int tk = s_token[row];
float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
if (tk >= 0)
v = *reinterpret_cast<const float4*>(x + (size_t)tk * H + kb * 64 + kc * 8);
*reinterpret_cast<float4*>(dst + rg * 1024 + kc * 128 + r * 16) = v;
}
}
// GEMM1: BM=128 rows x (gate 64 | up 64) via one N=128 UMMA per k16.
// GEMM2: BM=128 rows x BN=128 out cols. Shared skeleton via IS_G1 flag.
// All operands arrive as packed 8KB blocks; the mainloop is single-threaded.
template <bool IS_G1, bool SHARED, int STAGES>
__global__ void __launch_bounds__(128, 1)
tc_gemm_kernel(
const bf16* __restrict__ ap, // packed A: gemm1 xperm, gemm2 hbuf(packed)
const bf16* __restrict__ wr, // packed routed weights (blocks)
const bf16* __restrict__ ws, // packed shared weights
const int* __restrict__ s_token,
const float* __restrict__ s_weight,
const int* __restrict__ row_off,
const int* __restrict__ mb_cum,
bf16* __restrict__ hbuf, // gemm1 out (packed blocks)
float* __restrict__ outf, // gemm2 out
int H, int I, int E, int n_shared, int mb_sh)
{
constexpr int BM = 128, BN = 256, BK = 64; // N=256: one UMMA per k16
constexpr int STAGE_BYTES = 48 * 1024; // A 16KB + B 32KB
constexpr int LOOKAHEAD = 2;
const int K = IS_G1 ? H : I;
const int KB = K / 64; // 8KB blocks along K
const int NT = IS_G1 ? (I / 128) : (H / BN); // gemm1 tile = 128 h-cols
const int E_total = E + n_shared;
int e, row0, n0;
if (SHARED) {
const int smi = blockIdx.x / NT;
const int ni = blockIdx.x % NT;
const int es = smi / mb_sh;
const int mi = smi % mb_sh;
e = E + es;
row0 = row_off[E] + (es * mb_sh + mi) * BM;
n0 = ni * (IS_G1 ? 128 : BN);
} else {
const int lim = IS_G1 ? E_total : E;
const int total_tiles = mb_cum[lim] * NT;
const int tile = blockIdx.x;
if (tile >= total_tiles) return;
int lo = 0, hi = lim;
while (hi - lo > 1) {
const int mid = (lo + hi) >> 1;
if (mb_cum[mid] * NT <= tile) lo = mid; else hi = mid;
}
e = lo;
const int local = tile - mb_cum[e] * NT;
const int mb = mb_cum[e + 1] - mb_cum[e];
const int mi = local % mb;
const int ni = local / mb;
row0 = row_off[e] + mi * BM;
n0 = ni * (IS_G1 ? 128 : BN);
}
// packed global block bases (each block 8KB = 4096 bf16)
// gemm1 B: w1p[e][nb][kb], nb over 2I/64 rows; gate nb=n0/64, up nb=I/64+n0/64
// gemm2 B: w2p[e][mb][kb], mb over H/64; tile uses mb = n0/64, n0/64+1
const bf16* wsel = (e < E) ? wr : ws;
const size_t eb = (size_t)((e < E) ? e : (e - E));
size_t bblk[4];
if (IS_G1) { // B rows: gate n0..n0+127 | up I+n0..I+n0+127
const size_t base = eb * (size_t)(2 * I / 64) * KB;
bblk[0] = base + (size_t)(n0 / 64) * KB;
bblk[1] = base + (size_t)(n0 / 64 + 1) * KB;
bblk[2] = base + (size_t)(I / 64 + n0 / 64) * KB;
bblk[3] = base + (size_t)(I / 64 + n0 / 64 + 1) * KB;
} else { // B rows: n0..n0+255
const size_t base = eb * (size_t)(H / 64) * KB;
#pragma unroll
for (int q = 0; q < 4; ++q) bblk[q] = base + (size_t)(n0 / 64 + q) * KB;
}
const size_t ablk0 = (size_t)(row0 / 64) * KB; // A blocks: [rb][kb]
const size_t ablk1 = ablk0 + KB;
extern __shared__ char smem[];
__shared__ __align__(8) uint64_t full[STAGES];
__shared__ __align__(8) uint64_t empt[STAGES];
__shared__ unsigned tmem_base;
const int t = threadIdx.x;
if (t == 0) {
#pragma unroll
for (int s = 0; s < STAGES; ++s) {
asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" :: "r"(smem_addr(&full[s])));
asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;\n" :: "r"(smem_addr(&empt[s])));
}
}
if (t < 32) {
asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], 256;\n"
:: "r"(smem_addr(&tmem_base)));
}
__syncthreads();
const unsigned tmem = tmem_base;
const unsigned idesc =
(1u << 4) | (1u << 7) | (1u << 10) | ((256u >> 3) << 17) | ((128u >> 4) << 24);
const int KT = KB; // one block set per stage (BK=64)
if (t == 0) {
auto issue_stage = [&](int kt) {
const int s = kt % STAGES;
const unsigned dst = smem_addr(smem) + s * STAGE_BYTES;
const unsigned mb = smem_addr(&full[s]);
mbar_expect_tx(mb, 49152);
tc_bulk(dst, ap + (ablk0 + kt) * 4096, 8192, mb);
tc_bulk(dst + 8192, ap + (ablk1 + kt) * 4096, 8192, mb);
#pragma unroll
for (int q = 0; q < 4; ++q)
tc_bulk(dst + 16384 + q * 8192, wsel + (bblk[q] + kt) * 4096, 8192, mb);
};
#pragma unroll 1
for (int f = 0; f < LOOKAHEAD && f < KT; ++f) issue_stage(f);
int fphase[STAGES], ephase[STAGES];
#pragma unroll
for (int s = 0; s < STAGES; ++s) { fphase[s] = 0; ephase[s] = 0; }
for (int kt = 0; kt < KT; ++kt) {
const int ft = kt + LOOKAHEAD;
if (ft < KT) {
const int sf = ft % STAGES;
if (ft >= STAGES) { // stage last read by mma(ft - STAGES)
mbar_wait(smem_addr(&empt[sf]), ephase[sf] & 1);
ephase[sf]++;
}
issue_stage(ft);
}
const int s = kt % STAGES;
mbar_wait(smem_addr(&full[s]), fphase[s] & 1);
fphase[s]++;
const unsigned ab = smem_addr(smem) + s * STAGE_BYTES;
const unsigned bb = ab + 16384;
#pragma unroll
for (int j = 0; j < BK / 16; ++j) {
tc_mma(tmem, tc_desc(ab + j * 256, 1024), tc_desc(bb + j * 256, 1024),
idesc, kt > 0 || j > 0);
}
tc_commit(smem_addr(&empt[s]));
}
}
// all threads wait for the last commit on the empty barrier
{
const int sl = (KT - 1) % STAGES;
const int commits = (KT - 1 - sl) / STAGES + 1;
mbar_wait(smem_addr(&empt[sl]), (unsigned)((commits - 1) & 1));
}
__syncthreads();
// epilogue: warp w owns tile row 32w+lane (tmem lane), 128 cols
const int warp = t / 32, lane = t % 32;
const int row = 32 * warp + lane;
const int gr = row0 + row;
const unsigned taddr = tmem + ((unsigned)(32 * warp) << 16);
if (IS_G1) {
// h cols n0..n0+127 for row gr (gate = tmem cols 0..127, up 128..255),
// written PACKED for gemm2 into blocks (gr/64, n0/64) and (.., +1):
// inner: [rg=(gr%64)/8][kc 0..7][r=gr%8][16B]
const int IKB = I / 64;
#pragma unroll
for (int half = 0; half < 2; ++half) { // 64 h cols per half
__nv_bfloat162 hv[32];
#pragma unroll
for (int q = 0; q < 2; ++q) { // 32 cols per tc_ld32 pair
float g[32], u[32];
tc_ld32(taddr + half * 64 + q * 32, g);
tc_ld32(taddr + 128 + half * 64 + q * 32, u);
#pragma unroll
for (int j = 0; j < 16; ++j) {
hv[q * 16 + j] = __floats2bfloat162_rn(
silu(g[2 * j]) * u[2 * j], silu(g[2 * j + 1]) * u[2 * j + 1]);
}
}
char* blk = reinterpret_cast<char*>(hbuf)
+ ((size_t)(gr / 64) * IKB + (n0 / 64 + half)) * 8192
+ ((gr % 64) / 8) * 1024 + (gr % 8) * 16;
#pragma unroll
for (int kc = 0; kc < 8; ++kc) {
*reinterpret_cast<float4*>(blk + kc * 128) =
reinterpret_cast<float4*>(hv)[kc];
}
}
} else {
// tc_ld32 is warp-collective: always execute, predicate stores only
const int tk = s_token[gr];
const float w = SHARED ? 1.0f : (tk >= 0 ? s_weight[gr] : 0.0f);
float* orow = outf + (size_t)(tk >= 0 ? tk : 0) * H + n0;
#pragma unroll
for (int c0 = 0; c0 < 256; c0 += 32) {
float v[32];
tc_ld32(taddr + c0, v);
if (tk >= 0) {
if (SHARED && e == E) {
#pragma unroll
for (int j = 0; j < 32; ++j) orow[c0 + j] = v[j] * w;
} else {
#pragma unroll
for (int j = 0; j < 32; ++j) atomicAdd(&orow[c0 + j], v[j] * w);
}
}
}
}
__syncthreads();
if (t < 32) {
asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, 256;\n" :: "r"(tmem));
asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n");
}
}
#endif // KBH_TC
__global__ void convert_kernel(const float* __restrict__ src,
bf16* __restrict__ dst, int64_t n) {
const int64_t i = ((int64_t)blockIdx.x * blockDim.x + threadIdx.x) * 2;
if (i < n) {
const float2 v = *reinterpret_cast<const float2*>(src + i);
*reinterpret_cast<__nv_bfloat162*>(dst + i) = __floats2bfloat162_rn(v.x, v.y);
}
}
// ---------------------------------------------------------------------------
// host side
// ---------------------------------------------------------------------------
static inline int cdiv(int a, int b) { return (a + b - 1) / b; }
template <typename K>
static void set_smem(K k, int bytes) {
cudaFuncSetAttribute(k, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
}
// stage bytes helpers (must match kernel constexprs)
#define G1_SMEM(BM, BN, ST) ((BM * 8 + 2 * BN * 8) * 16 * ST)
#define G2_SMEM(BM, BN, ST) ((BM * 8 + BN * 8) * 16 * ST)
static bool g_attr_done = false;
static void init_attrs() {
if (g_attr_done) return;
#ifdef KBH_TC
set_smem(tc_gemm_kernel<true, false, 4>, 4 * 48 * 1024);
set_smem(tc_gemm_kernel<false, true, 4>, 4 * 48 * 1024);
set_smem(tc_gemm_kernel<false, false, 4>, 4 * 48 * 1024);
#endif
set_smem(gemm1_kernel<128, 64, 4, 2, 3>, G1_SMEM(128, 64, 3));
set_smem(gemm1_kernel<64, 64, 2, 2, 3>, G1_SMEM(64, 64, 3));
set_smem(gemm1_kernel<32, 64, 1, 4, 5>, G1_SMEM(32, 64, 5));
set_smem(gemm2_kernel<128, 128, 4, 2, 3, false>, G2_SMEM(128, 128, 3));
set_smem(gemm2_kernel<128, 128, 4, 2, 3, true>, G2_SMEM(128, 128, 3));
set_smem(gemm2_kernel<64, 128, 2, 2, 3, false>, G2_SMEM(64, 128, 3));
set_smem(gemm2_kernel<64, 128, 2, 2, 3, true>, G2_SMEM(64, 128, 3));
set_smem(gemm2_kernel<32, 128, 1, 4, 5, false>, G2_SMEM(32, 128, 5));
set_smem(gemm2_kernel<32, 128, 1, 4, 5, true>, G2_SMEM(32, 128, 5));
g_attr_done = true;
}
void moe_forward(
torch::Tensor x, torch::Tensor ids, torch::Tensor wts,
torch::Tensor w1, torch::Tensor w2, torch::Tensor w1s, torch::Tensor w2s,
torch::Tensor counts, torch::Tensor row_off, torch::Tensor mb_cum,
torch::Tensor s_token, torch::Tensor s_weight,
torch::Tensor hbuf, torch::Tensor outf, torch::Tensor out,
torch::Tensor w1q, torch::Tensor w2q, torch::Tensor w1sq, torch::Tensor w2sq,
torch::Tensor xperm,
int64_t BM_sel)
{
init_attrs();
auto stream = at::cuda::getCurrentCUDAStream();
const int T = x.size(0);
const int H = x.size(1);
const int topk = ids.size(1);
const int E = w1.size(0);
const int n_shared = w1s.size(0);
const int I = w2.size(2);
const int BM = (int)BM_sel;
const int E_total = E + n_shared;
const bf16* xp = reinterpret_cast<const bf16*>(x.data_ptr());
const bf16* w1p = reinterpret_cast<const bf16*>(w1.data_ptr());
const bf16* w2p = reinterpret_cast<const bf16*>(w2.data_ptr());
const bf16* w1sp = reinterpret_cast<const bf16*>(w1s.data_ptr());
const bf16* w2sp = reinterpret_cast<const bf16*>(w2s.data_ptr());
bf16* hp = reinterpret_cast<bf16*>(hbuf.data_ptr());
bf16* op = reinterpret_cast<bf16*>(out.data_ptr());
if (BM == 0) { // decode GEMV path (T <= 4)
constexpr int W1 = 4, R1 = 2, ROWS1 = W1 * R1;
constexpr int W2 = 4, R2 = 2, ROWS2 = W2 * R2;
const int P9 = T * (topk + n_shared);
gemv1_kernel<W1, R1><<<P9 * (I / ROWS1), W1 * 32, H * 2 + 2 * ROWS1 * 4, stream>>>(
xp, w1p, w1sp, ids.data_ptr<int64_t>(),
T, topk, E, n_shared, H, I, hp);
cudaMemsetAsync(outf.data_ptr<float>(), 0, (size_t)T * H * 4, stream);
gemv2_kernel<W2, R2><<<P9 * (H / ROWS2), W2 * 32, I * 2, stream>>>(
hp, w2p, w2sp, ids.data_ptr<int64_t>(),
reinterpret_cast<const bf16*>(wts.data_ptr()),
T, topk, E, n_shared, H, I, outf.data_ptr<float>());
const int64_t n = (int64_t)T * H;
const int cthreads = 256;
const int cblocks = (int)((n / 2 + cthreads - 1) / cthreads);
convert_kernel<<<cblocks, cthreads, 0, stream>>>(outf.data_ptr<float>(), op, n);
return;
}
const int P = T * topk;
const int rows_max = (int)s_token.size(0);
int* cnt = counts.data_ptr<int>(); // 2*E_total ints, [cnt|fill]
int* fill = cnt + E_total;
cudaMemsetAsync(cnt, 0, 2 * E_total * sizeof(int), stream);
moe_hist_kernel<<<cdiv(P, 256), 256, 0, stream>>>(
ids.data_ptr<int64_t>(), P, cnt);
moe_scan_kernel<<<1, 32, 0, stream>>>(
cnt, T, E, n_shared, BM,
row_off.data_ptr<int>(), mb_cum.data_ptr<int>());
moe_initneg_kernel<<<cdiv(rows_max, 256), 256, 0, stream>>>(
s_token.data_ptr<int>(), rows_max);
moe_scatter_kernel<<<cdiv(P + n_shared * T, 256), 256, 0, stream>>>(
ids.data_ptr<int64_t>(), reinterpret_cast<const bf16*>(wts.data_ptr()),
T, topk, E, n_shared,
row_off.data_ptr<int>(), fill,
s_token.data_ptr<int>(), s_weight.data_ptr<float>());
const int mb_routed_max = std::min(cdiv(P, BM) + E, P);
const int mb_sh = cdiv(T, BM);
const int mb_total_max = mb_routed_max + n_shared * mb_sh;
const int NT1 = I / 64;
const int NT2 = H / 128;
const int g1 = mb_total_max * NT1;
const int g2s = n_shared * mb_sh * NT2;
const int g2r = mb_routed_max * NT2;
#ifdef KBH_TC
if (BM == 128) { // tcgen05 path (SM100): BM selected as 128 only there
const int tc_smem = 4 * 48 * 1024;
const int nt1 = I / 128, nt2 = H / 256;
const int tg1 = mb_total_max * nt1;
const int tg2s = n_shared * mb_sh * nt2;
const int tg2r = mb_routed_max * nt2;
const bf16* xq = reinterpret_cast<const bf16*>(xperm.data_ptr());
tc_pack_x_kernel<<<(rows_max / 64) * (H / 64), 256, 0, stream>>>(
xp, s_token.data_ptr<int>(),
row_off.data_ptr<int>() + E_total,
const_cast<bf16*>(xq), H);
tc_gemm_kernel<true, false, 4><<<tg1, 128, tc_smem, stream>>>(
xq, reinterpret_cast<const bf16*>(w1q.data_ptr()),
reinterpret_cast<const bf16*>(w1sq.data_ptr()),
s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
hp, nullptr, H, I, E, n_shared, mb_sh);
tc_gemm_kernel<false, true, 4><<<tg2s, 128, tc_smem, stream>>>(
hp, reinterpret_cast<const bf16*>(w2q.data_ptr()),
reinterpret_cast<const bf16*>(w2sq.data_ptr()),
s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
nullptr, outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
tc_gemm_kernel<false, false, 4><<<tg2r, 128, tc_smem, stream>>>(
hp, reinterpret_cast<const bf16*>(w2q.data_ptr()),
reinterpret_cast<const bf16*>(w2sq.data_ptr()),
s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
nullptr, outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
const int64_t n = (int64_t)T * H;
const int cthreads = 256;
const int cblocks = (int)((n / 2 + cthreads - 1) / cthreads);
convert_kernel<<<cblocks, cthreads, 0, stream>>>(outf.data_ptr<float>(), op, n);
return;
}
#endif
if (BM == 128) {
gemm1_kernel<128, 64, 4, 2, 3><<<g1, 256, G1_SMEM(128, 64, 3), stream>>>(
xp, w1p, w1sp, s_token.data_ptr<int>(), row_off.data_ptr<int>(),
mb_cum.data_ptr<int>(), hp, H, I, E, n_shared);
gemm2_kernel<128, 128, 4, 2, 3, true><<<g2s, 256, G2_SMEM(128, 128, 3), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
gemm2_kernel<128, 128, 4, 2, 3, false><<<g2r, 256, G2_SMEM(128, 128, 3), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
} else if (BM == 64) {
gemm1_kernel<64, 64, 2, 2, 3><<<g1, 128, G1_SMEM(64, 64, 3), stream>>>(
xp, w1p, w1sp, s_token.data_ptr<int>(), row_off.data_ptr<int>(),
mb_cum.data_ptr<int>(), hp, H, I, E, n_shared);
gemm2_kernel<64, 128, 2, 2, 3, true><<<g2s, 128, G2_SMEM(64, 128, 3), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
gemm2_kernel<64, 128, 2, 2, 3, false><<<g2r, 128, G2_SMEM(64, 128, 3), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
} else {
gemm1_kernel<32, 64, 1, 4, 5><<<g1, 128, G1_SMEM(32, 64, 5), stream>>>(
xp, w1p, w1sp, s_token.data_ptr<int>(), row_off.data_ptr<int>(),
mb_cum.data_ptr<int>(), hp, H, I, E, n_shared);
gemm2_kernel<32, 128, 1, 4, 5, true><<<g2s, 128, G2_SMEM(32, 128, 5), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
gemm2_kernel<32, 128, 1, 4, 5, false><<<g2r, 128, G2_SMEM(32, 128, 5), stream>>>(
hp, w2p, w2sp, s_token.data_ptr<int>(), s_weight.data_ptr<float>(),
row_off.data_ptr<int>(), mb_cum.data_ptr<int>(),
outf.data_ptr<float>(), H, I, E, n_shared, mb_sh);
}
const int64_t n = (int64_t)T * H;
const int cthreads = 256;
const int cblocks = (int)((n / 2 + cthreads - 1) / cthreads);
convert_kernel<<<cblocks, cthreads, 0, stream>>>(outf.data_ptr<float>(), op, n);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("moe_forward", &moe_forward, "fused MoE forward (CUDA)");
}
"""
_EXT = None
_HAS_TC = None # SM100 (B200-class) has tcgen05; SM120 consumer Blackwell does not
def _get_ext():
global _EXT, _HAS_TC
if _EXT is None:
from torch.utils.cpp_extension import load_inline
cap = torch.cuda.get_device_capability()
_HAS_TC = cap == (10, 0)
flags = ["-O3", "--use_fast_math", "-std=c++17"]
name = "glm52_fused_moe_cuda"
if _HAS_TC:
os.environ["TORCH_CUDA_ARCH_LIST"] = "10.0a"
flags.append("-DKBH_TC")
name += "_tc"
_EXT = load_inline(
name=name,
cpp_sources="",
cuda_sources=CUDA_SOURCE,
extra_cuda_cflags=flags,
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._ext = _get_ext()
self._ws: dict = {}
# CUDA-graph cache to hide kernel launch overhead on tiny batches.
# Keyed on the INPUT POINTERS + shape: a graph is only replayed when
# the caller passes tensors at the exact same addresses, and every
# kernel in the graph reads those buffers' LIVE contents at replay
# time — this is launch-overhead removal, NOT result caching. New data
# in the same buffers (or new buffers) always produces fresh compute
# (see dev_test.py mutation check).
self._graphs: dict = {}
def _workspace(self, T: int, BM: int, device):
key = (T, BM)
ws = self._ws.get(key)
if ws is None:
E, n_shared, topk = self.E, self.n_shared, self.top_k
P = T * topk
if BM == 0: # decode GEMV path
rows_max = T * (topk + n_shared)
else:
mb_routed_max = min((P + BM - 1) // BM + E, P)
mb_sh = (T + BM - 1) // BM
rows_max = (mb_routed_max + n_shared * mb_sh) * BM
small = torch.empty(1, dtype=torch.bfloat16, device=device)
ws = dict(
counts=torch.empty(2 * (E + n_shared), dtype=torch.int32, device=device),
row_off=torch.empty(E + n_shared + 1, dtype=torch.int32, device=device),
mb_cum=torch.empty(E + n_shared + 1, dtype=torch.int32, device=device),
s_token=torch.empty(rows_max, dtype=torch.int32, device=device),
s_weight=torch.empty(rows_max, dtype=torch.float32, device=device),
hbuf=torch.empty(rows_max, self.I, dtype=torch.bfloat16, device=device),
outf=torch.empty(T, self.H, dtype=torch.float32, device=device),
out=torch.empty(T, self.H, dtype=torch.bfloat16, device=device),
xperm=torch.empty(rows_max, self.H, dtype=torch.bfloat16, device=device)
if BM == 128 and _HAS_TC else small,
)
self._ws[key] = ws
return ws
@staticmethod
def _pack_blocks(w: torch.Tensor) -> torch.Tensor:
"""Repack (..., R, C) bf16 into contiguous 8KB core-matrix blocks:
[..., R/64, C/64] blocks laid out as [rg 0..7][kc 0..7][r 0..7][8 elems]
— the exact order the tcgen05 no-swizzle smem descriptor traverses."""
R, C = w.shape[-2], w.shape[-1]
v = w.reshape(-1, R // 64, 8, 8, C // 64, 8, 8)
v = v.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
return v
def _packed_weights(self):
key = (
self.w1_routed.data_ptr(), self.w1_routed._version,
self.w2_routed.data_ptr(), self.w2_routed._version,
self.w1_shared._version, self.w2_shared._version,
)
pk = getattr(self, "_pk", None)
if pk is None or pk[0] != key:
with torch.no_grad():
pk = (key, (
self._pack_blocks(self.w1_routed.data),
self._pack_blocks(self.w2_routed.data),
self._pack_blocks(self.w1_shared.data),
self._pack_blocks(self.w2_shared.data),
))
self._pk = pk
return pk[1]
def forward(
self,
x: torch.Tensor,
expert_ids: torch.Tensor,
expert_weights: torch.Tensor,
) -> torch.Tensor:
T = x.shape[0]
# tokens-per-expert decides tile granularity (padding vs mma efficiency);
# BM=0 selects the per-pair GEMV decode path (weights unshared anyway).
avg = T * self.top_k / self.E
if T <= 4:
BM = 0
elif _HAS_TC:
# tcgen05 path uses M=128 tiles; small batches stay on mma.sync
BM = 128 if avg >= 32 else 32
else:
# BM=64 @ 3 stages (3 CTAs/SM) beats BM=128 even at avg=256
BM = 64 if avg >= 32 else 32
_bm_env = os.environ.get("KBH_MOE_BM") # dev tuning knob
if _bm_env:
BM = int(_bm_env)
x = x.contiguous()
expert_ids = expert_ids.contiguous()
expert_weights = expert_weights.contiguous()
ws = self._workspace(T, BM, x.device)
if BM == 128 and _HAS_TC:
w1q, w2q, w1sq, w2sq = self._packed_weights()
else:
w1q = w2q = w1sq = w2sq = ws["xperm"] # unused dummies
def launch():
self._ext.moe_forward(
x, expert_ids, expert_weights,
self.w1_routed, self.w2_routed, self.w1_shared, self.w2_shared,
ws["counts"], ws["row_off"], ws["mb_cum"],
ws["s_token"], ws["s_weight"],
ws["hbuf"], ws["outf"], ws["out"],
w1q, w2q, w1sq, w2sq, ws["xperm"], BM,
)
key = (
T, BM,
x.data_ptr(), expert_ids.data_ptr(), expert_weights.data_ptr(),
w1q.data_ptr(), # repacked weights invalidate captured graphs
)
graph = self._graphs.get(key)
if graph is None and len(self._graphs) < 64 and not torch.cuda.is_current_stream_capturing():
# Warm on a side stream, then capture the launch sequence. Replay
# re-executes every kernel against the live buffer contents.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
launch()
torch.cuda.current_stream().wait_stream(s)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
launch()
self._graphs[key] = graph
if graph is not None:
graph.replay()
else:
launch()
return ws["out"]
20260719_085216_or-fable_anthropic_claude-fable-5_01_glm52_fused_moe