KernelBench hard · RTX PRO 6000
Sonic MoE GLM-5.3 Flash
10.2%geomean peak fraction across shapes
manually audited: clean
Isolated regrade 0.1021 (in-run 0.1039). Packed-weight cache is Parameter identity; TMA descriptors cache activation/weight pointers. Fresh out every call. Same-buffer overwrite of hidden_states reads live bytes. Lint CLEAN. Numeric stress on.
harnessor-fableagent session6h 8mtotal wall6h 11mcheck2mbenchmark88soutput tokens—regimecompute
Per-shape vs governing ceilingeach shape graded against whichever binds — bf16 compute or HBM bandwidth
32768×4096×1536×128×819.071 ms8.6%0.32 TB/s · 18% of 1.8 TB/s HBM · also 43 TFLOPS (9% of compute)
4096×2048×1024×64×40.480 ms14.3%1.33 TB/s · 74% of 1.8 TB/s HBM · also 72 TFLOPS (14% of compute)
16384×2048×4096×64×812.790 ms8.6%0.29 TB/s · 16% of 1.8 TB/s HBM · also 43 TFLOPS (9% of compute)
compute-bound memory-bound · bar + right column = official fraction of the ceiling (the geomean input)
geomean(8.6% · 14.3% · 8.6%) = 10.2%
Kernel source (redacted)
"""Grouped GEMM + fused SwiGLU up-projection for MoE (sm_120 Blackwell).
Strategy
--------
The op per expert e is h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]).
Two GEMMs per expert with a shared activation operand. We fuse them into a
single grouped GEMM over a repacked weight tensor
W_cat[e] = interleave(W_gate[e], W_up[e]) along the last dim -> (H, 2*I)
so column pairs (2i, 2i+1) of the product hold (gate_i, up_i) adjacently. A
single accumulator tile then contains both operands of every SwiGLU lane and
the epilogue writes silu(g)*u directly at (T_perm, I) granularity -- no
intermediate (T_perm, 2I) tensor ever exists.
Implementation: a hand-written CUDA kernel shipped as an embedded fatbin.
Per CTA it computes a 128x256 output tile of the grouped GEMM with
mma.sync.m16n8k16 bf16 tensor-core instructions, streams operands through a
double-buffered 96 KB shared-memory pipeline fed by TMA
(cp.async.bulk.tensor, 128 B swizzle), and fuses silu(g)*u in the epilogue.
Grouping needs no host-side row counts: each CTA derives its
(expert, m-tile, n-tile) assignment from expert_offsets on-device, so there
are no D2H syncs anywhere in the forward pass.
The mma A-fragment register order is the one verified on sm_120 hardware:
[(g,klo), (g+8,klo), (g,khi), (g+8,khi)] -- note this differs from the
PTX ISA doc ordering for m16n8k16.
A Triton implementation of the same fused op is kept as a fallback for
shapes the CUDA fast path does not cover (H % 64 != 0 or N2 % 256 != 0).
"""
from __future__ import annotations
import ctypes
import os
import tempfile
import torch
import torch.nn as nn
import triton
import triton.language as tl
OP_TYPE = "grouped_gemm_swiglu"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
# ---------------------------------------------------------------------------
# Embedded CUDA kernel (compiled for sm_120; loaded once per process).
# ---------------------------------------------------------------------------
_CUDA_SO_B64 = "" # public copy omits embedded ELF fatbin; CUDA source is below
_LIB = None
_LIB_PATH = None
def _load_cuda_lib():
global _LIB, _LIB_PATH
if _LIB is not None or not torch.cuda.is_available():
return _LIB
try:
from base64 import b64decode
path = os.path.join(tempfile.gettempdir(), "moe_swiglu_tma_sm120.so")
blob = b64decode(_CUDA_SO_B64)
if not (os.path.exists(path) and os.path.getsize(path) == len(blob)):
with open(path, "wb") as f:
f.write(blob)
lib = ctypes.CDLL(path)
lib.moe_launch.argtypes = [
ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p,
ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int,
ctypes.c_int, ctypes.c_int, ctypes.c_uint64,
]
lib.moe_launch.restype = None
_LIB, _LIB_PATH = lib, path
except Exception:
_LIB = None
return _LIB
def _next_pow2(n: int) -> int:
p = 1
while p < n:
p *= 2
return p
@triton.autotune(
configs=[
triton.Config({"BM": 128, "BN": 128, "BK": 64}, num_warps=8, num_stages=3),
triton.Config({"BM": 128, "BN": 256, "BK": 64}, num_warps=8, num_stages=2),
triton.Config({"BM": 128, "BN": 256, "BK": 64}, num_warps=8, num_stages=3),
triton.Config({"BM": 256, "BN": 128, "BK": 64}, num_warps=8, num_stages=2),
triton.Config({"BM": 256, "BN": 128, "BK": 64}, num_warps=8, num_stages=3),
triton.Config({"BM": 64, "BN": 256, "BK": 64}, num_warps=8, num_stages=2),
triton.Config({"BM": 64, "BN": 256, "BK": 64}, num_warps=8, num_stages=3),
triton.Config({"BM": 64, "BN": 128, "BK": 64}, num_warps=4, num_stages=4),
triton.Config({"BM": 128, "BN": 64, "BK": 64}, num_warps=4, num_stages=4),
triton.Config({"BM": 128, "BN": 128, "BK": 128}, num_warps=8, num_stages=2),
triton.Config({"BM": 128, "BN": 256, "BK": 128}, num_warps=8, num_stages=2),
],
key=["H", "N2", "I"],
)
@triton.jit
def _grouped_swiglu_kernel(
x_ptr, w_ptr, o_ptr, off_ptr,
T_perm, H, N2, I, E,
BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
EP2: tl.constexpr, EVEN_N: tl.constexpr, EVEN_K: tl.constexpr,
):
pid = tl.program_id(0)
NN = tl.cdiv(N2, BN) # n-tiles per expert
# ---- on-device tile -> (expert, m_tile, n_tile) map --------------------
ea = tl.arange(0, EP2)
emask = ea < E
off0 = tl.load(off_ptr + ea, mask=emask, other=0).to(tl.int32)
off1 = tl.load(off_ptr + ea + 1, mask=emask, other=0).to(tl.int32)
m_cnt = tl.where(emask, (off1 - off0 + BM - 1) // BM, 0)
tpe = m_cnt * NN # tiles per expert
pref = tl.cumsum(tpe, axis=0) # inclusive prefix
if pid >= tl.max(pref, axis=0):
return
expert = tl.sum((pref <= pid).to(tl.int32), axis=0)
sel = ea == expert
before = tl.sum(tl.where(sel, pref - tpe, 0), axis=0)
mstart = tl.sum(tl.where(sel, off0, 0), axis=0)
mend = tl.sum(tl.where(sel, off1, 0), axis=0)
local = pid - before
m_loc = local // NN
n_loc = local % NN
# ---- accumulator tile ---------------------------------------------------
# Mainloop loads are UNMASKED (masked loads defeat the cp.async pipeline);
# out-of-range rows/columns are clamped to valid memory instead and the
# garbage accumulators they produce are discarded by the store mask.
rows = mstart + m_loc * BM + tl.arange(0, BM)
rows = tl.minimum(rows, T_perm - 1)
n2 = n_loc * BN + tl.arange(0, BN) # columns of W_cat (2I wide)
if not EVEN_N:
n2 = tl.minimum(n2, N2 - 1)
i_cols = n_loc * (BN // 2) + tl.arange(0, BN // 2) # fused output cols
rmask = (mstart + m_loc * BM + tl.arange(0, BM)) < mend
omask = i_cols < I
acc = tl.zeros((BM, BN), dtype=tl.float32)
w_row = expert.to(tl.int64) * H * N2
k_main = (H // BK) * BK
for k0 in range(0, k_main, BK):
ks = k0 + tl.arange(0, BK)
a = tl.load(x_ptr + rows[:, None].to(tl.int64) * H + ks[None, :])
b = tl.load(w_ptr + w_row + ks[:, None] * N2 + n2[None, :])
acc = tl.dot(a, b, acc)
if not EVEN_K:
ks = k_main + tl.arange(0, BK)
kmask = ks < H
a = tl.load(
x_ptr + rows[:, None].to(tl.int64) * H + ks[None, :],
mask=kmask[None, :], other=0.0,
)
b = tl.load(
w_ptr + w_row + ks[:, None] * N2 + n2[None, :],
mask=kmask[:, None], other=0.0,
)
acc = tl.dot(a, b, acc)
# ---- fused SwiGLU epilogue: pairs (2i, 2i+1) are adjacent --------------
g, u = tl.split(tl.reshape(acc, (BM, BN // 2, 2)))
o = g * tl.sigmoid(g) * u
tl.store(
o_ptr + rows[:, None].to(tl.int64) * I + i_cols[None, :],
o.to(tl.bfloat16),
mask=rmask[:, None] & omask[None, :],
)
class Model(nn.Module):
"""Up-projection of a top-K MoE FFN with fused SwiGLU."""
def __init__(self, T_total: int, H: int, I: int, E: int, K: int): # noqa: E741
super().__init__()
self.T_total = T_total
self.H = H
self.I = I
self.E = E
self.K = K
self.W_gate = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
nn.init.normal_(self.W_gate, std=0.02)
nn.init.normal_(self.W_up, std=0.02)
self._wcat: torch.Tensor | None = None
self._wkey: tuple | None = None
self._ep2 = _next_pow2(E + 1)
def _packed_weights(self) -> torch.Tensor:
"""(E, H, 2I) with gate/up interleaved along the last dim."""
key = (
self.W_gate._version, self.W_up._version,
self.W_gate.data_ptr(), self.W_up.data_ptr(),
)
if self._wcat is None or self._wkey != key:
with torch.no_grad():
self._wcat = torch.stack(
[self.W_gate.detach(), self.W_up.detach()], dim=-1
).reshape(self.E, self.H, 2 * self.I).contiguous()
self._wkey = key
return self._wcat
def forward(
self,
hidden_states: torch.Tensor, # (T_perm, H) bf16
expert_offsets: torch.Tensor, # (E+1,) int32
) -> torch.Tensor:
x = hidden_states.contiguous()
w = self._packed_weights()
T_perm, H = x.shape
I, N2, E = self.I, 2 * self.I, self.E
out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=x.device)
lib = _load_cuda_lib()
if (
lib is not None
and x.dtype == torch.bfloat16
and H % 64 == 0
and N2 % 256 == 0
and expert_offsets.dtype == torch.int32
and expert_offsets.is_cuda
):
offs = expert_offsets.contiguous()
lib.moe_launch(
x.data_ptr(), w.data_ptr(), out.data_ptr(), offs.data_ptr(),
T_perm, H, N2, I, E, N2 // 256,
torch.cuda.current_stream().cuda_stream,
)
return out
# ---- Triton fallback ----
def grid(meta):
bm = meta["BM"]
nn_ = triton.cdiv(N2, meta["BN"])
mtiles = triton.cdiv(T_perm, bm) + E
return (mtiles * nn_,)
_grouped_swiglu_kernel[grid](
x, w, out, expert_offsets,
T_perm, H, N2, I, E,
EP2=self._ep2, EVEN_N=(N2 % 256 == 0), EVEN_K=(H % 128 == 0),
)
return out
# Module-level shape shims rewritten by check.py / benchmark.py per shape.
T_total = 32768
H = 4096
I = 1536 # noqa: E741
E = 128
K = 8
def _build_routing(T_total: int, E: int, K: int, device: str = "cpu") -> torch.Tensor:
"""Round-robin-ish routing metadata: balanced offsets summing to T_total*K."""
T_perm = T_total * K
base = T_perm // E
rem = T_perm - base * E
counts = torch.full((E,), base, dtype=torch.int32, device=device)
counts[:rem] += 1
offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
offsets[1:] = torch.cumsum(counts, dim=0)
return offsets
def get_inputs():
T_perm = T_total * K
hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
expert_offsets = _build_routing(T_total, E, K)
return [hidden_states, expert_offsets]
def get_init_inputs():
return [T_total, H, I, E, K]
# ==================================================================
# ===== sidecar: moe_tma.cu (11467 bytes, CUDA source for stripped fatbin) =====
# ==================================================================
// sm_120 grouped GEMM + fused SwiGLU, TMA + mbarrier pipeline variant.
#include <cuda.h>
#include <cuda_bf16.h>
#include <cstdint>
#define BM 128
#define BN 256
#define BK 64
#define KSTEPS (BK / 16)
#ifndef NTHREAD
#define NTHREAD 256
#endif
__device__ __forceinline__ unsigned smem_u32(const void* p) {
return (unsigned)__cvta_generic_to_shared(p);
}
__device__ __forceinline__ void ldm_x4(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, unsigned a) {
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) : "r"(a));
}
__device__ __forceinline__ void ldm_x4t(uint32_t& r0, uint32_t& r1, uint32_t& r2, uint32_t& r3, unsigned a) {
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"(a));
}
__device__ __forceinline__ void mma16816(float (&d)[4], const uint32_t (&a)[4],
const uint32_t b0, const 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"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
: "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1));
}
__device__ __forceinline__ void mbar_init(unsigned bar, unsigned count) {
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(bar), "r"(count));
}
__device__ __forceinline__ void mbar_expect_tx(unsigned bar, unsigned bytes) {
asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" ::"r"(bar), "r"(bytes));
}
__device__ __forceinline__ void mbar_wait(unsigned bar, unsigned parity) {
asm volatile(
"{\n\t.reg .pred p;\n"
"L%=:\n\t"
"mbarrier.try_wait.parity.shared::cta.b64 p, [%0], %1;\n\t"
"@!p bra L%=;\n"
"}\n" ::"r"(bar), "r"(parity));
}
__device__ __forceinline__ void tma_2d(void* dst, const void* tmap, int c0, int c1, unsigned bar) {
asm volatile(
"cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes "
"[%0], [%1, {%2, %3}], [%4];\n"
::"r"(smem_u32(dst)), "l"((unsigned long long)tmap), "r"(c0), "r"(c1), "r"(bar)
: "memory");
}
__global__ void __launch_bounds__(NTHREAD, 1) moe_swiglu_kernel(
const __nv_bfloat16* __restrict__ x, // (T_perm, H)
const __nv_bfloat16* __restrict__ w, // (E, H, N2)
__nv_bfloat16* __restrict__ out, // (T_perm, I)
const int* __restrict__ offs, // (E+1,)
const CUtensorMap* __restrict__ tmA, // device-resident tensormaps
const CUtensorMap* __restrict__ tmB,
int T_perm, int H, int N2, int I, int E, int NN) {
extern __shared__ __align__(1024) __nv_bfloat16 smem[];
// stage buffers: A 128x64 (pitch 64, xor-swizzled by TMA 128B mode),
// B: four dense 64(k)x64(col) blocks per stage; block jb == warp column group wn
__nv_bfloat16* sA[2] = {smem, smem + BM * BK};
__nv_bfloat16* sB[2] = {smem + 2 * BM * BK, smem + 2 * BM * BK + BK * BN};
__shared__ __align__(8) unsigned long long bar_full[2];
__shared__ int s_excl[257];
const int tid = threadIdx.x;
const int wid = tid >> 5, lane = tid & 31;
const int wm = wid >> 2, wn = wid & 3;
if (tid == 0) {
mbar_init(smem_u32(&bar_full[0]), 1);
mbar_init(smem_u32(&bar_full[1]), 1);
int run = 0;
for (int e = 0; e < E; ++e) {
s_excl[e] = run;
run += ((offs[e + 1] - offs[e] + BM - 1) / BM) * NN;
}
s_excl[E] = run;
s_excl[256] = run;
}
__syncthreads();
const int total = s_excl[256];
unsigned par[2] = {0, 0};
float acc[4][8][4];
#pragma unroll
for (int mi = 0; mi < 4; ++mi)
#pragma unroll
for (int nj = 0; nj < 8; ++nj)
#pragma unroll
for (int j = 0; j < 4; ++j) acc[mi][nj][j] = 0.f;
// prologue issued inside the main loop on first two iterations via staged logic
int next_issue = 0; // how many stages already issued for current tile
auto issue_stage = [&](int buf, int k0, int mrow0, long ekbase, int ncol0) {
if (tid == 0) {
unsigned bar = smem_u32(&bar_full[buf]);
mbar_expect_tx(bar, (BM * BK + BK * BN) * 2);
tma_2d(sA[buf], tmA, k0, mrow0, bar);
for (int jb = 0; jb < 4; ++jb)
tma_2d(sB[buf] + jb * 64 * 64, tmB, ncol0 + jb * 64, (int)(ekbase + k0), bar);
}
};
for (int t = blockIdx.x; t < total; t += gridDim.x) {
int e = 0;
{
int lo = 0, hi = E - 1;
while (lo < hi) {
int mid = (lo + hi + 1) >> 1;
if (s_excl[mid] <= t) lo = mid;
else hi = mid - 1;
}
e = lo;
}
const int mstart = offs[e];
const int mend = offs[e + 1];
const int local = t - s_excl[e];
const int m_loc = local / NN;
const int n_loc = local % NN;
const int mrow0 = mstart + m_loc * BM;
const int ncol0 = n_loc * BN;
const long ekbase = (long)e * H;
const int nrows_valid = mend - mstart;
next_issue = 0;
for (int ki = 0; ki < H; ki += BK) {
const int buf = (ki / BK) & 1;
// issue up to two stages ahead
while (next_issue <= ki + BK && next_issue < H) {
issue_stage(next_issue / BK & 1, next_issue, mrow0, ekbase, ncol0);
next_issue += BK;
}
mbar_wait(smem_u32(&bar_full[buf]), par[buf]);
par[buf] ^= 1;
#pragma unroll
for (int ks = 0; ks < KSTEPS; ++ks) {
uint32_t af[4][4];
uint32_t bf[4][4];
const int k0 = ks * 16;
#pragma unroll
for (int np = 0; np < 4; ++np) {
// B block jb == wn; rows pitch 64 within block
int kk = k0 + (lane & 15);
int cc = np * 2 + (lane >> 4);
int v = cc ^ (kk & 7);
ldm_x4t(bf[np][0], bf[np][1], bf[np][2], bf[np][3],
smem_u32(sB[buf] + wn * 64 * 64 + kk * 64 + (v << 3)));
}
#pragma unroll
for (int mi = 0; mi < 4; ++mi) {
// HW-verified a-frag reg order: [ (g,klo), (g+8,klo), (g,khi), (g+8,khi) ]
int r = wm * 64 + mi * 16 + (lane & 7) + 8 * ((lane >> 3) & 1);
int c = ks * 2 + ((lane >> 4) & 1);
int v = c ^ (r & 7);
ldm_x4(af[mi][0], af[mi][1], af[mi][2], af[mi][3],
smem_u32(sA[buf] + r * BK + (v << 3)));
}
#pragma unroll
for (int mi = 0; mi < 4; ++mi)
#pragma unroll
for (int np = 0; np < 4; ++np) {
mma16816(acc[mi][np * 2 + 0], af[mi], bf[np][0], bf[np][1]);
mma16816(acc[mi][np * 2 + 1], af[mi], bf[np][2], bf[np][3]);
}
}
__syncthreads(); // everyone done reading buf before refill
}
// ---- fused SwiGLU epilogue ----
__syncthreads();
const int STG_STRIDE = BN / 2 + 8;
__nv_bfloat16* stg = smem;
const int rowp = lane >> 2;
#pragma unroll
for (int mi = 0; mi < 4; ++mi) {
const int r0 = wm * 64 + mi * 16 + rowp;
#pragma unroll
for (int nj = 0; nj < 8; ++nj) {
const int cf = wn * 32 + nj * 4 + ((lane & 3));
float g0 = acc[mi][nj][0], u0 = acc[mi][nj][1];
float g1 = acc[mi][nj][2], u1 = acc[mi][nj][3];
float o0 = (g0 / (1.f + __expf(-g0))) * u0;
float o1 = (g1 / (1.f + __expf(-g1))) * u1;
stg[r0 * STG_STRIDE + cf] = __float2bfloat16(o0);
stg[(r0 + 8) * STG_STRIDE + cf] = __float2bfloat16(o1);
}
}
__syncthreads();
const int ocol0 = ncol0 >> 1;
#pragma unroll
for (int q = 0, ci = tid; q < (BM * (BN / 2) / 8) / NTHREAD; ++q, ci += NTHREAD) {
int r = ci >> 4, c = (ci & 15) << 3;
int gr = mrow0 + r, gc = ocol0 + c;
if (gr < mend && gc < I) {
*reinterpret_cast<uint4*>(out + (long)gr * I + gc) =
*reinterpret_cast<const uint4*>(stg + r * STG_STRIDE + c);
}
}
__syncthreads();
#pragma unroll
for (int mi = 0; mi < 4; ++mi)
#pragma unroll
for (int nj = 0; nj < 8; ++nj)
#pragma unroll
for (int j = 0; j < 4; ++j) acc[mi][nj][j] = 0.f;
}
}
extern "C" void moe_launch(const void* x, const void* w, void* out, const int* offs,
int T_perm, int H, int N2, int I, int E, int NN,
uintptr_t stream) {
static CUtensorMap* d_tmA = nullptr;
static CUtensorMap* d_tmB = nullptr;
static const void* key_x = nullptr;
static const void* key_w = nullptr;
static int kT = -1, kH = -1, kN2 = -1, kE = -1;
static bool attr_set = false;
if (!attr_set) {
size_t smem = 2 * (BM * BK + BK * BN) * 2;
cudaFuncSetAttribute(moe_swiglu_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)smem);
cudaMalloc(&d_tmA, sizeof(CUtensorMap));
cudaMalloc(&d_tmB, sizeof(CUtensorMap));
attr_set = true;
}
if (key_x != x || key_w != w || kT != T_perm || kH != H || kN2 != N2 || kE != E) {
CUtensorMap htm;
cuuint64_t dimA[2] = {(cuuint64_t)H, (cuuint64_t)T_perm};
cuuint64_t strA[1] = {(cuuint64_t)H * 2};
cuuint32_t boxA[2] = {(cuuint32_t)BK, (cuuint32_t)BM};
cuuint32_t est[2] = {1, 1};
cuTensorMapEncodeTiled(&htm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, (void*)x,
dimA, strA, boxA, est,
CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
cudaMemcpyAsync(d_tmA, &htm, sizeof(htm), cudaMemcpyHostToDevice,
(cudaStream_t)stream);
cuuint64_t dimB[2] = {(cuuint64_t)N2, (cuuint64_t)E * H};
cuuint64_t strB[1] = {(cuuint64_t)N2 * 2};
cuuint32_t boxB[2] = {64, (cuuint32_t)BK};
cuTensorMapEncodeTiled(&htm, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, (void*)w,
dimB, strB, boxB, est,
CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
cudaMemcpyAsync(d_tmB, &htm, sizeof(htm), cudaMemcpyHostToDevice,
(cudaStream_t)stream);
key_x = x; key_w = w; kT = T_perm; kH = H; kN2 = N2; kE = E;
}
int nsm = 0;
cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, 0);
moe_swiglu_kernel<<<nsm, NTHREAD, 2 * (BM * BK + BK * BN) * 2, (cudaStream_t)stream>>>(
reinterpret_cast<const __nv_bfloat16*>(x),
reinterpret_cast<const __nv_bfloat16*>(w),
reinterpret_cast<__nv_bfloat16*>(out),
offs, d_tmA, d_tmB, T_perm, H, N2, I, E, NN);
}
20260822_061328_or-fable_stealth_ox-alpha_06_sonic_moe_swiglu