"""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 #include #include #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(out + (long)gr * I + gc) = *reinterpret_cast(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<<>>( reinterpret_cast(x), reinterpret_cast(w), reinterpret_cast<__nv_bfloat16*>(out), offs, d_tmA, d_tmB, T_perm, H, N2, I, E, NN); }