KernelBench cuda · RTX PRO 6000
GLM-5.2 Fused MoE DeepSeek V4 Flash (0731)
4.11%geomean peak fraction across shapes
manually audited: clean
Manual static audit covered all 186 lines of solution.py, all 397 substantive lines of scratch/fused_moe.cu, all 280369 transcript events and 306 agent tool calls, result.json, check.log, benchmark.log, the CUDA-language sidecar, and the frozen grader/template files. solution.py:113-160 derives a fresh expert sort, token gather, weights, group offsets, and tile maps from the current x, expert_ids, and expert_weights on every forward. It then launches run_moe on the live routed and model-weight tensors at solution.py:162-172 and computes the always-on shared expert from the current x and shared weights at solution.py:174-182. scratch/fused_moe.cu:61-179 performs the routed gate/up tensor-core GEMMs and SwiGLU, lines 190-308 perform the routed down GEMMs and weighted atomic accumulation, and lines 358-397 allocate fresh intermediate and zeroed fp32 output tensors before launching both kernels. The only retained Python object is the compiled extension handle; no output, input identity/data_ptr key, CUDA graph, result table, constant answer, stack/caller or check.py sniff, reference import, or forbidden library exists in the final computation. Therefore no empirical same-buffer overwrite cache test is required. The trace writes only this cell's solution.py and fused_moe.cu plus disposable /tmp kernel experiments. Its foreign run IDs appear only in passive ps output while diagnosing GPU/process contention; no tool input names a foreign run and no foreign artifact is opened. All seven archived repo grader files compare byte-for-byte equal to template_files, consistent with result.json template_mutated=false. The agent read check.py and eval sources but did not edit them, set KBH_NUMERIC_STRESS, alter tolerances, or branch the submitted code on grader behavior. check.log:6-7 records the PTX/CUDA language gate and PASS from that unmodified checker; check.py:75-92 runs nominal, small_hidden, and large_hidden numeric-stress cases for seeds 42/123/456, but check.log gives no per-case magnitudes, so none are claimed. cuda_language.json reports framework=ptx, triton_cheat=false, dsl_cheat=false, forbidden_hits=[], and ok=true. The archived pre-regrade benchmark.log records per-shape fractions 0.0599, 0.0604, 0.0001, 0.0920, 0.0080, and 0.0159, peak_fraction 0.0114, and RESULT: OK; its result.json records correct=true and zero check/benchmark exit codes. Publication metrics come from the parent-run isolated regrade on 2026-08-03 using NVIDIA RTX PRO 6000 Blackwell Server Edition: correct=true and peak_fraction=0.0411. The archived 0.0114 value is retained here only as pre-regrade provenance and is superseded by that isolated publish-grade result.
harnessor-fableagent session6h 12mtotal wall6h 25mcheck6mbenchmark5moutput tokens—gpu-lock wait41mgpu-lock held12mregimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
shape 021.112 ms17.6%0.62 TB/s · 34% of 1.8 TB/s HBM · also 88 TFLOPS (18% of compute)
shape 120.883 ms17.9%0.62 TB/s · 35% of 1.8 TB/s HBM · also 90 TFLOPS (18% of compute)
shape 21.671 ms0.1%7.74 TB/s · 430% of 1.8 TB/s HBM · also 0 TFLOPS (0% of compute)
shape 329.305 ms25.3%127 TFLOPS · 25% of 500 TF bf16 peak · also 0.45 TB/s (25% of HBM)
shape 419.943 ms2.3%0.65 TB/s · 36% of 1.8 TB/s HBM · also 12 TFLOPS (2% of compute)
shape 518.786 ms4.8%0.69 TB/s · 38% of 1.8 TB/s HBM · also 24 TFLOPS (5% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(17.6% · 17.9% · 0.1% · 25.3% · 2.3% · 4.8%) = 4.1%
Kernel source (redacted)
"""GLM-5.2-class fused MoE layer (CUDA, SM120 / Blackwell, bf16).
Fused vLLM-style weight layout with GLM routing (E=256 routed, top_k=8,
1 always-on shared expert). Reference math per expert:
h = silu(x @ w1_gate.T) * (x @ w1_up.T); y = h @ w2.T
out = sum_shared(y_s) + sum_k weight[t,k] * y_routed[e_k]
Two grouped-GEMM CUDA kernels (fused_moe.cu, WMMA m16n16k16 tensor cores):
A) gate/up -> h buffer
B) down -> weighted fp32 atomic accumulate into out
Routing is given; tokens are grouped (sorted) by expert on the host.
"""
from __future__ import annotations
import os
from pathlib import Path
_HERE = Path(__file__).resolve().parent
def _find_cuda_home():
"""Locate a real CUDA toolkit (nvcc), bypassing harness PATH shims."""
for c in (os.environ.get("CUDA_HOME"), os.environ.get("TORCH_CUDA_HOME")):
if c and os.path.exists(os.path.join(c, "bin", "nvcc")):
return c
for c in ("/usr/local/cuda-12.8", "/usr/local/cuda", "/usr/local/cuda-13",
"/usr/local/cuda-13.0"):
if os.path.exists(os.path.join(c, "bin", "nvcc")):
return c
return None
# Must run BEFORE importing torch.utils.cpp_extension (which caches CUDA_HOME
# at module-import time from the environment).
_ch = _find_cuda_home()
if _ch is not None:
os.environ["CUDA_HOME"] = _ch
import torch # noqa: E402
import torch.nn as nn # noqa: E402
from torch.utils import cpp_extension # noqa: E402
_EXT = None
# Kernel tile sizes (must match fused_moe.cu).
BM_A, BI_A, BK_A = 128, 32, 32
BM_B, BN_B, BK_B = 128, 64, 32
THREADS = 256
def _load_ext():
global _EXT
if _EXT is not None:
return _EXT
cuda_src = (_HERE / "fused_moe.cu").read_text()
cpp_src = r"""
#include <torch/extension.h>
torch::Tensor run_moe(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w2_r,
torch::Tensor w1_s, torch::Tensor w2_s,
torch::Tensor toks_all, torch::Tensor wts_all,
torch::Tensor g_start,
torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
int64_t E, int64_t H, int64_t I, int64_t T);
torch::Tensor gate_up_only(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w1_s,
torch::Tensor toks_all, torch::Tensor g_start,
torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
int64_t E, int64_t H, int64_t I);
torch::Tensor down_only(torch::Tensor h_all, torch::Tensor w2_r, torch::Tensor w2_s,
torch::Tensor wts, torch::Tensor toks_all, torch::Tensor g_start,
torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
int64_t E, int64_t H, int64_t I, int64_t T);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("run_moe", &run_moe, "GLM-5.2 fused MoE (gate/up + down grouped GEMMs)");
m.def("gate_up_only", &gate_up_only, "debug: gate/up -> h");
m.def("down_only", &down_only, "debug: down -> out");
}
"""
_EXT = cpp_extension.load_inline(
name="fused_moe_ext",
cpp_sources=cpp_src,
cuda_sources=cuda_src,
extra_cuda_cflags=["-arch=sm_120a", "-O3", "--expt-relaxed-constexpr"],
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)
def forward(
self,
x: torch.Tensor,
expert_ids: torch.Tensor,
expert_weights: torch.Tensor,
) -> torch.Tensor:
T, H = x.shape
E, top_k, I = self.E, self.top_k, self.I
dev = x.device
ext = _load_ext()
M_r = T * top_k
# ---- group (sort) routed assignments by expert ---------------------
ids32 = expert_ids.reshape(-1).to(torch.int32)
toks_flat = (
torch.arange(T, device=dev)
.view(T, 1)
.expand(T, top_k)
.reshape(-1)
.to(torch.int32)
)
order = torch.argsort(ids32, stable=True)
toks_sorted = toks_flat[order]
wts_sorted = expert_weights.reshape(-1)[order]
cnts = torch.bincount(ids32[order], minlength=E)
g_start = torch.cat(
[
torch.zeros(1, dtype=cnts.dtype, device=dev),
torch.cumsum(cnts, 0),
]
).to(torch.int32) # (E+1): routed groups only
toks_all = toks_sorted
wts_all = wts_sorted
# gathered activations: contiguous rows so kernel A streams xg sequentially
xg = x[toks_all.to(torch.int64)].contiguous()
# ---- build tile maps for kernel A and B (vectorized) ---------------
gcnts = (g_start[1:] - g_start[:-1]).to(torch.int64)
def tiles(BM):
nt = ((gcnts + BM - 1) // BM)
total = int(nt.sum())
tg = torch.repeat_interleave(
torch.arange(E, device=dev, dtype=torch.int64), nt
)
gstarts = torch.cumsum(
torch.cat([torch.zeros(1, dtype=torch.int64, device=dev), nt]), 0
)
pos = torch.arange(total, device=dev, dtype=torch.int64) - gstarts[tg]
tr0 = g_start.to(torch.int64)[tg] + pos * BM
return tg.to(torch.int32), tr0.to(torch.int32)
tgA, tr0A = tiles(BM_A)
tgB, tr0B = tiles(BM_B)
out_f32 = ext.run_moe(xg,
self.w1_routed.detach(),
self.w2_routed.detach(),
self.w1_shared.detach(),
self.w2_shared.detach(),
toks_all,
wts_all,
g_start,
tgA, tr0A, tgB, tr0B,
E, H, I, T,
)
# Shared expert: dense bf16 GEMMs (read once, vs the grouped kernel's
# per-tile re-read), accumulated in fp32.
w1s = self.w1_shared[0]
w2s = self.w2_shared[0]
gate = x @ w1s[:I].T
up = x @ w1s[I:].T
h_s = torch.nn.functional.silu(gate) * up
y_s = h_s @ w2s.T
return (out_f32 + y_s.float()).to(torch.bfloat16)
def get_init_inputs():
return [4096, 256, 8, 1, 4096, 2048]
# ==================================================================
# ===== sidecar: fused_moe.cu (17626 bytes, loaded by solution.py) =====
# ==================================================================
// GLM-5.2 fused MoE — grouped-GEMM CUDA kernels (SM120 / Blackwell, bf16).
//
// Layout (vLLM-style fused weights):
// w1_routed: (E, 2I, H) bf16 — [0:I)=gate, [I:2I)=up
// w2_routed: (E, H, I) bf16
// w1_shared / w2_shared: (n_shared, ...) — the always-on shared expert.
//
// Two kernels, unified "row" space:
// rows [0, T*top_k) -> routed assignments, sorted by expert.
// rows [T*top_k, +T) -> shared-expert tokens (all T), weight 1.0.
// g_start[g]..g_start[g+1] -> rows belonging to group g (g in [0,E) routed, g=E shared).
//
// Kernel A (gate_up): h[row, :] = silu(x@w1_gate.T) * (x@w1_up.T) -> h_all (M_total, I)
// Kernel B (down): out[token] += wts[row] * (h[row] @ w2.T) -> atomic fp32 accumulate
//
// WMMA (nvcuda::wmma) m16n16k16 bf16 tensor-core path with cp.async double
// buffering over the contraction dimension.
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <mma.h>
using namespace nvcuda;
using bf16 = __nv_bfloat16;
__device__ __forceinline__ float silu_f(float x) {
return x / (1.0f + expf(-x));
}
// ---- cp.async helpers (16-byte copies, SM80+) ----------------------------
__device__ __forceinline__ 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));
}
__device__ __forceinline__ void cp_commit() {
asm volatile("cp.async.commit_group;\n");
}
template <int N>
__device__ __forceinline__ void cp_wait() {
asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
}
// wmma m16n16k16 f32 accumulator fragment -> (row, col) within 16x16 tile.
__device__ __forceinline__ void frag_row_col(int lane, int i, int& r, int& c) {
r = (i & 2) ? (8 + (lane >> 2)) : (lane >> 2);
c = ((lane & 3) << 1) | (i & 1) | ((i & 4) << 1);
}
// ---------------------------------------------------------------------------
// Kernel A: grouped gate/up GEMM -> h_all (M_total, I)
// ---------------------------------------------------------------------------
#define BM_A 128
#define BI_A 32
#define BK_A 32
#define THREADS_A 256
#define GS_A 4 // I-slices processed per block (xg rows reused from L2 across them)
// warps per block = THREADS_A / 32 = 16; rows per warp = BM_A / 16 = 16
__global__ void __launch_bounds__(THREADS_A)
gate_up_kernel(
const bf16* __restrict__ x, // (M_total, H) gathered xg
const bf16* __restrict__ w1_r, // (E, 2I, H)
const bf16* __restrict__ w1_s, // (n_shared, 2I, H)
const int* __restrict__ toks, // (M_total)
const int* __restrict__ g_start, // (G+1)
const int* __restrict__ tile_group, // (num_tiles)
const int* __restrict__ tile_row0, // (num_tiles)
bf16* __restrict__ h_all, // (M_total, I)
const int H, const int I, const int E, const int num_sgroups)
{
// grid = (num_sgroups, num_tiles): a block handles GS_A consecutive I-slices
// of one tile, re-reading its xg rows from L2 across the group.
const int sgroup = blockIdx.x;
const int tile = blockIdx.y;
const int g = tile_group[tile];
const int row0 = tile_row0[tile];
const int g0 = g_start[g];
const int g1 = g_start[g + 1];
const int cnt = g1 - g0;
const int nrow = min(BM_A, g1 - row0);
const bf16* w1 = (g < E) ? (w1_r + (size_t)g * (size_t)2 * I * H) : w1_s;
__shared__ __align__(16) bf16 xs[2][BM_A][BK_A];
__shared__ __align__(16) bf16 wg[2][BI_A][BK_A];
__shared__ __align__(16) bf16 wu[2][BI_A][BK_A];
const int tid = threadIdx.x;
const int warp = tid >> 5;
const int lane = tid & 31;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc_g[BI_A / 16], acc_u[BI_A / 16];
const int XBLOCKS = (BM_A * BK_A) / 8;
const int WBLOCKS = (BI_A * BK_A) / 8;
for (int gs_i = 0; gs_i < GS_A; gs_i++) {
const int slice = sgroup * GS_A + gs_i;
const int ibase = slice * BI_A;
#pragma unroll
for (int i = 0; i < BI_A / 16; i++) {
wmma::fill_fragment(acc_g[i], 0.f);
wmma::fill_fragment(acc_u[i], 0.f);
}
// Prefetch chunk 0 into buffer 0.
for (int b = tid; b < XBLOCKS; b += THREADS_A) {
int r = b / (BK_A / 8);
int k8 = b % (BK_A / 8);
const bf16* src = (r < nrow) ? (x + (size_t)(row0 + r) * H + k8 * 8)
: (x + k8 * 8);
cp_async16(&xs[0][r][k8 * 8], src);
}
for (int b = tid; b < WBLOCKS; b += THREADS_A) {
int n = b / (BK_A / 8);
int k8 = b % (BK_A / 8);
cp_async16(&wg[0][n][k8 * 8], w1 + (size_t)(ibase + n) * H + k8 * 8);
cp_async16(&wu[0][n][k8 * 8], w1 + (size_t)(I + ibase + n) * H + k8 * 8);
}
cp_commit();
for (int k0 = 0; k0 < H; k0 += BK_A) {
const int cur = (k0 / BK_A) & 1;
const int nxt = cur ^ 1;
if (k0 + BK_A < H) {
for (int b = tid; b < XBLOCKS; b += THREADS_A) {
int r = b / (BK_A / 8);
int k8 = b % (BK_A / 8);
const bf16* src = (r < nrow) ? (x + (size_t)(row0 + r) * H + k0 + BK_A + k8 * 8)
: (x + k0 + BK_A + k8 * 8);
cp_async16(&xs[nxt][r][k8 * 8], src);
}
for (int b = tid; b < WBLOCKS; b += THREADS_A) {
int n = b / (BK_A / 8);
int k8 = b % (BK_A / 8);
cp_async16(&wg[nxt][n][k8 * 8], w1 + (size_t)(ibase + n) * H + k0 + BK_A + k8 * 8);
cp_async16(&wu[nxt][n][k8 * 8], w1 + (size_t)(I + ibase + n) * H + k0 + BK_A + k8 * 8);
}
cp_commit();
}
if (k0 + BK_A < H) cp_wait<1>();
else cp_wait<0>();
__syncthreads();
const int mbase = warp * 16;
#pragma unroll
for (int ks = 0; ks < BK_A / 16; ks++) {
wmma::fragment<wmma::matrix_a, 16, 16, 16, bf16, wmma::row_major> af;
wmma::load_matrix_sync(af, &xs[cur][mbase][ks * 16], BK_A);
#pragma unroll
for (int nt = 0; nt < BI_A / 16; nt++) {
wmma::fragment<wmma::matrix_b, 16, 16, 16, bf16, wmma::col_major> bg, bu;
wmma::load_matrix_sync(bg, &wg[cur][nt * 16][ks * 16], BK_A);
wmma::load_matrix_sync(bu, &wu[cur][nt * 16][ks * 16], BK_A);
wmma::mma_sync(acc_g[nt], af, bg, acc_g[nt]);
wmma::mma_sync(acc_u[nt], af, bu, acc_u[nt]);
}
}
__syncthreads();
}
// Write h = silu(gate) * up directly from accumulator fragments.
#pragma unroll
for (int nt = 0; nt < BI_A / 16; nt++) {
#pragma unroll
for (int i = 0; i < 8; i++) {
int rr, cc;
frag_row_col(lane, i, rr, cc);
const int grow = warp * 16 + rr;
if (grow < nrow) {
float hv = silu_f(acc_g[nt].x[i]) * acc_u[nt].x[i];
h_all[(size_t)(row0 + grow) * I + ibase + nt * 16 + cc] = __float2bfloat16(hv);
}
}
}
}
}
// ---------------------------------------------------------------------------
// Kernel B: grouped down GEMM -> weighted atomic accumulate into out (T, H) fp32
// ---------------------------------------------------------------------------
#define BM_B 128
#define BN_B 64
#define BK_B 32
#define THREADS_B 256
#define HS_B 4 // H-slices processed per block (h rows reused from L2 across them)
__global__ void __launch_bounds__(THREADS_B)
down_kernel(
const bf16* __restrict__ h_all, // (M_total, I)
const bf16* __restrict__ w2_r, // (E, H, I)
const bf16* __restrict__ w2_s, // (n_shared, H, I)
const bf16* __restrict__ wts, // (M_total)
const int* __restrict__ toks, // (M_total)
const int* __restrict__ g_start,
const int* __restrict__ tile_group,
const int* __restrict__ tile_row0,
float* __restrict__ out, // (T, H) fp32
const int H, const int I, const int T, const int E, const int num_hgroups)
{
// grid = (num_hgroups, num_tiles): a block processes HS_B consecutive H-slices
// of one tile, re-reading its h rows from L2 across the group.
const int hslice0 = blockIdx.x * HS_B;
const int tile = blockIdx.y;
const int g = tile_group[tile];
const int row0 = tile_row0[tile];
const int g0 = g_start[g];
const int g1 = g_start[g + 1];
const int cnt = g1 - g0;
const int nrow = min(BM_B, g1 - row0);
const bf16* w2 = (g < E) ? (w2_r + (size_t)g * H * I) : w2_s;
__shared__ __align__(16) bf16 hs[2][BM_B][BK_B];
__shared__ __align__(16) bf16 w2t[2][BN_B][BK_B];
const int tid = threadIdx.x;
const int warp = tid >> 5;
const int lane = tid & 31;
// 2D warp tiling: warps_m x warps_n over the (BM_B, BN_B) output tile.
const int warps_m = BM_B / 16;
const int warps_n = (THREADS_B / 32) / warps_m;
const int wm = warp % warps_m;
const int wn = warp / warps_m;
const int mbase = wm * 16;
const int n_ntiles_per_warp = BN_B / (16 * warps_n);
const int n_base = wn * (BN_B / warps_n);
wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc[n_ntiles_per_warp];
const int HBLOCKS = (BM_B * BK_B) / 8;
const int WBLOCKS = (BN_B * BK_B) / 8;
for (int hs_i = 0; hs_i < HS_B; hs_i++) {
const int hb = (hslice0 + hs_i) * BN_B;
#pragma unroll
for (int i = 0; i < n_ntiles_per_warp; i++) wmma::fill_fragment(acc[i], 0.f);
// prefetch chunk 0
for (int b = tid; b < HBLOCKS; b += THREADS_B) {
int r = b / (BK_B / 8);
int k8 = b % (BK_B / 8);
const bf16* src = (r < nrow) ? (h_all + (size_t)(row0 + r) * I + k8 * 8)
: (h_all + k8 * 8);
cp_async16(&hs[0][r][k8 * 8], src);
}
for (int b = tid; b < WBLOCKS; b += THREADS_B) {
int n = b / (BK_B / 8);
int k8 = b % (BK_B / 8);
cp_async16(&w2t[0][n][k8 * 8], w2 + (size_t)(hb + n) * I + k8 * 8);
}
cp_commit();
for (int k0 = 0; k0 < I; k0 += BK_B) {
const int cur = (k0 / BK_B) & 1;
const int nxt = cur ^ 1;
if (k0 + BK_B < I) {
for (int b = tid; b < HBLOCKS; b += THREADS_B) {
int r = b / (BK_B / 8);
int k8 = b % (BK_B / 8);
const bf16* src = (r < nrow) ? (h_all + (size_t)(row0 + r) * I + k0 + BK_B + k8 * 8)
: (h_all + k0 + BK_B + k8 * 8);
cp_async16(&hs[nxt][r][k8 * 8], src);
}
for (int b = tid; b < WBLOCKS; b += THREADS_B) {
int n = b / (BK_B / 8);
int k8 = b % (BK_B / 8);
cp_async16(&w2t[nxt][n][k8 * 8], w2 + (size_t)(hb + n) * I + k0 + BK_B + k8 * 8);
}
cp_commit();
}
if (k0 + BK_B < I) cp_wait<1>();
else cp_wait<0>();
__syncthreads();
#pragma unroll
for (int ks = 0; ks < BK_B / 16; ks++) {
wmma::fragment<wmma::matrix_a, 16, 16, 16, bf16, wmma::row_major> af;
wmma::load_matrix_sync(af, &hs[cur][mbase][ks * 16], BK_B);
#pragma unroll
for (int nt = 0; nt < n_ntiles_per_warp; nt++) {
wmma::fragment<wmma::matrix_b, 16, 16, 16, bf16, wmma::col_major> bf;
wmma::load_matrix_sync(bf, &w2t[cur][n_base + nt * 16][ks * 16], BK_B);
wmma::mma_sync(acc[nt], af, bf, acc[nt]);
}
}
__syncthreads();
}
// weighted atomic accumulate into out.
#pragma unroll
for (int nt = 0; nt < n_ntiles_per_warp; nt++) {
#pragma unroll
for (int i = 0; i < 8; i++) {
int rr, cc;
frag_row_col(lane, i, rr, cc);
const int grow = mbase + rr;
if (grow < nrow) {
int tok = toks[row0 + grow];
float wt = __bfloat162float(wts[row0 + grow]);
atomicAdd(&out[(size_t)tok * H + hb + n_base + nt * 16 + cc], wt * acc[nt].x[i]);
}
}
}
}
}
// ---------------------------------------------------------------------------
// Host entry points
// ---------------------------------------------------------------------------
torch::Tensor gate_up_only(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w1_s,
torch::Tensor toks_all, torch::Tensor g_start,
torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
int64_t E, int64_t H, int64_t I) {
int64_t M_total = toks_all.size(0);
auto h_all = torch::empty({M_total, I}, x.options().dtype(torch::kBFloat16));
auto stream = at::cuda::getCurrentCUDAStream();
const bf16* xp = reinterpret_cast<const bf16*>(x.data_ptr<at::BFloat16>());
const bf16* w1r = reinterpret_cast<const bf16*>(w1_r.data_ptr<at::BFloat16>());
const bf16* w1s = reinterpret_cast<const bf16*>(w1_s.data_ptr<at::BFloat16>());
const int* toksp = toks_all.data_ptr<int>();
const int* gsp = g_start.data_ptr<int>();
const int* tgA = tile_group_A.data_ptr<int>();
const int* tr0A = tile_row0_A.data_ptr<int>();
bf16* hp = reinterpret_cast<bf16*>(h_all.data_ptr<at::BFloat16>());
int num_tiles_A = (int)tile_group_A.numel();
dim3 gridA(I / (BI_A * GS_A), num_tiles_A);
gate_up_kernel<<<gridA, THREADS_A, 0, stream>>>(
xp, w1r, w1s, toksp, gsp, tgA, tr0A, hp, (int)H, (int)I, (int)E, gridA.y);
return h_all;
}
torch::Tensor down_only(torch::Tensor h_all, torch::Tensor w2_r, torch::Tensor w2_s,
torch::Tensor wts, torch::Tensor toks_all, torch::Tensor g_start,
torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
int64_t E, int64_t H, int64_t I, int64_t T) {
int64_t M_total = toks_all.size(0);
auto out = torch::zeros({T, H}, h_all.options().dtype(torch::kFloat32));
auto stream = at::cuda::getCurrentCUDAStream();
const bf16* hp = reinterpret_cast<const bf16*>(h_all.data_ptr<at::BFloat16>());
const bf16* w2r = reinterpret_cast<const bf16*>(w2_r.data_ptr<at::BFloat16>());
const bf16* w2s = reinterpret_cast<const bf16*>(w2_s.data_ptr<at::BFloat16>());
const bf16* wtsp = reinterpret_cast<const bf16*>(wts.data_ptr<at::BFloat16>());
const int* toksp = toks_all.data_ptr<int>();
const int* gsp = g_start.data_ptr<int>();
const int* tgB = tile_group_B.data_ptr<int>();
const int* tr0B = tile_row0_B.data_ptr<int>();
float* op = out.data_ptr<float>();
int num_tiles_B = (int)tile_group_B.numel();
dim3 gridB(H / (BN_B * HS_B), num_tiles_B);
down_kernel<<<gridB, THREADS_B, 0, stream>>>(
hp, w2r, w2s, wtsp, toksp, gsp, tgB, tr0B, op, (int)H, (int)I, (int)T, (int)E, gridB.y);
return out;
}
torch::Tensor run_moe(torch::Tensor x, torch::Tensor w1_r, torch::Tensor w2_r,
torch::Tensor w1_s, torch::Tensor w2_s,
torch::Tensor toks_all, torch::Tensor wts_all,
torch::Tensor g_start,
torch::Tensor tile_group_A, torch::Tensor tile_row0_A,
torch::Tensor tile_group_B, torch::Tensor tile_row0_B,
int64_t E, int64_t H, int64_t I, int64_t T) {
int64_t M_total = toks_all.size(0);
auto h_all = torch::empty({M_total, I}, x.options().dtype(torch::kBFloat16));
auto out = torch::zeros({T, H}, x.options().dtype(torch::kFloat32));
auto stream = at::cuda::getCurrentCUDAStream();
const bf16* xp = reinterpret_cast<const bf16*>(x.data_ptr<at::BFloat16>());
const bf16* w1r = reinterpret_cast<const bf16*>(w1_r.data_ptr<at::BFloat16>());
const bf16* w2r = reinterpret_cast<const bf16*>(w2_r.data_ptr<at::BFloat16>());
const bf16* w1s = reinterpret_cast<const bf16*>(w1_s.data_ptr<at::BFloat16>());
const bf16* w2s = reinterpret_cast<const bf16*>(w2_s.data_ptr<at::BFloat16>());
const int* toksp = toks_all.data_ptr<int>();
const bf16* wtsp = reinterpret_cast<const bf16*>(wts_all.data_ptr<at::BFloat16>());
const int* gsp = g_start.data_ptr<int>();
const int* tgA = tile_group_A.data_ptr<int>();
const int* tr0A = tile_row0_A.data_ptr<int>();
const int* tgB = tile_group_B.data_ptr<int>();
const int* tr0B = tile_row0_B.data_ptr<int>();
bf16* hp = reinterpret_cast<bf16*>(h_all.data_ptr<at::BFloat16>());
float* op = out.data_ptr<float>();
int num_tiles_A = (int)tile_group_A.numel();
int num_tiles_B = (int)tile_group_B.numel();
dim3 gridA(I / (BI_A * GS_A), num_tiles_A);
dim3 gridB(H / (BN_B * HS_B), num_tiles_B);
gate_up_kernel<<<gridA, THREADS_A, 0, stream>>>(
xp, w1r, w1s, toksp, gsp, tgA, tr0A, hp, (int)H, (int)I, (int)E, gridA.y);
down_kernel<<<gridB, THREADS_B, 0, stream>>>(
hp, w2r, w2s, wtsp, toksp, gsp, tgB, tr0B, op, (int)H, (int)I, (int)T, (int)E, gridB.y);
return out;
}
20260802_203348_or-fable_deepseek_deepseek-v4-flash-0731_01_glm52_fused_moe