"""Grouped GEMM + fused SwiGLU up-projection for MoE FFN (RTX PRO 6000 / SM120). Custom persistent Triton kernel for the variable-length grouped GEMM out[start_e:end_e] = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]) Design notes: - The gate/up weights are repacked once (per weight version) into a single column-interleaved tensor W_comb[e] = [g_0, u_0, g_1, u_1, ...], so each output tile needs only ONE wide dot into a (BM, 2*BN) accumulator; the epilogue de-interleaves and applies SwiGLU. This halves the MMA issue stream versus a dual-accumulator kernel. - Tiles are enumerated expert-major; a persistent grid walks the flattened tile space with an amortized O(E) cursor, so forward() never synchronizes with the host. - SM120 (Blackwell workstation) has no tcgen05/TMEM; the kernel targets the warp-level mma.sync bf16 tensor-core path with cp.async pipelining. - A fully masked generic kernel handles any shape where H % BK or I % BN != 0. """ from __future__ import annotations import torch, torch.nn as nn, triton, triton.language as tl NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count # Per-shape tile configs (tuned on RTX PRO 6000; keyed by (H, I)). CFG_TABLE = { (4096, 1536): dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None), (2048, 1024): dict(BM=128, BN=64, BK=64, warps=8, stages=3, grid_mult=1, maxnreg=232), (2048, 4096): dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None), } CFG_DEFAULT = dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None) def _cfg_for(H, I): # noqa: E741 return CFG_TABLE.get((H, I), CFG_DEFAULT) @triton.jit def _kernel_v5( x_ptr, w_ptr, out_ptr, offs_ptr, stride_xm, stride_wn, stride_om, MAX_ITER, E: tl.constexpr, GRID, H: tl.constexpr, I: tl.constexpr, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, ): pid = tl.program_id(0) NB: tl.constexpr = I // BN total = 0 for ee in range(E): n_e = tl.load(offs_ptr + ee + 1) - tl.load(offs_ptr + ee) total += ((n_e + BM - 1) // BM) * NB e_cur = 0 base_cur = 0 start_cur = tl.load(offs_ptr + 0) end_cur = tl.load(offs_ptr + 1) nm_cur = ((end_cur - start_cur) + BM - 1) // BM rn = tl.arange(0, 2 * BN) rm = tl.arange(0, BM) rk = tl.arange(0, BK) for i in range(0, MAX_ITER): t = pid + i * GRID if t < total: span_cur = nm_cur * NB while t >= base_cur + span_cur: base_cur += span_cur e_cur += 1 start_cur = end_cur end_cur = tl.load(offs_ptr + e_cur + 1) nm_cur = ((end_cur - start_cur) + BM - 1) // BM span_cur = nm_cur * NB off_t = t - base_cur mb = off_t // NB nb = off_t % NB row0 = start_cur + mb * BM wrow0 = e_cur * H x_ptrs = x_ptr + (row0 + rm)[:, None] * stride_xm + rk[None, :] w_ptrs = w_ptr + wrow0.to(tl.int64) * (2 * I) + rk[:, None] * (2 * I) + (nb * 2 * BN + rn)[None, :] acc = tl.zeros((BM, 2 * BN), dtype=tl.float32) for _k in range(0, H // BK): x = tl.load(x_ptrs) w = tl.load(w_ptrs) acc = tl.dot(x, w, acc) x_ptrs += BK w_ptrs += BK * (2 * I) g, u = tl.split(tl.reshape(acc, (BM, BN, 2))) res = g * tl.sigmoid(g) * u out = res.to(tl.bfloat16) rno = tl.arange(0, BN) out_ptrs = out_ptr + (row0 + rm)[:, None] * stride_om + (nb * BN + rno)[None, :] if (row0 + BM) <= end_cur: tl.store(out_ptrs, out) else: row_mask = (row0 + rm) < end_cur tl.store(out_ptrs, out, mask=row_mask[:, None]) @triton.jit def _kernel_generic( x_ptr, wg_ptr, wu_ptr, out_ptr, offs_ptr, stride_xm, stride_om, H, I, MAX_ITER, E: tl.constexpr, GRID, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, ): """Masked dual-accumulator variant for shapes with H % BK or I % BN != 0.""" pid = tl.program_id(0) NB = tl.cdiv(I, BN) total = 0 for ee in range(E): n_e = tl.load(offs_ptr + ee + 1) - tl.load(offs_ptr + ee) total += ((n_e + BM - 1) // BM) * NB e_cur = 0 base_cur = 0 start_cur = tl.load(offs_ptr + 0) end_cur = tl.load(offs_ptr + 1) nm_cur = ((end_cur - start_cur) + BM - 1) // BM rn = tl.arange(0, BN) rm = tl.arange(0, BM) rk = tl.arange(0, BK) for i in range(0, MAX_ITER): t = pid + i * GRID if t < total: span_cur = nm_cur * NB while t >= base_cur + span_cur: base_cur += span_cur e_cur += 1 start_cur = end_cur end_cur = tl.load(offs_ptr + e_cur + 1) nm_cur = ((end_cur - start_cur) + BM - 1) // BM span_cur = nm_cur * NB off_t = t - base_cur mb = off_t // NB nb = off_t % NB row0 = start_cur + mb * BM wrow0 = e_cur.to(tl.int64) * H cols = nb * BN + rn x_ptrs = x_ptr + (row0 + rm)[:, None] * stride_xm + rk[None, :] w_ptrs_g = wg_ptr + wrow0 * I + rk[:, None] * I + cols[None, :] w_ptrs_u = wu_ptr + wrow0 * I + rk[:, None] * I + cols[None, :] acc_g = tl.zeros((BM, BN), dtype=tl.float32) acc_u = tl.zeros((BM, BN), dtype=tl.float32) col_mask = cols < I for kk in range(0, tl.cdiv(H, BK)): k_mask = (kk * BK + rk) < H x = tl.load(x_ptrs, mask=k_mask[None, :], other=0.0) wg = tl.load(w_ptrs_g, mask=k_mask[:, None] & col_mask[None, :], other=0.0) wu = tl.load(w_ptrs_u, mask=k_mask[:, None] & col_mask[None, :], other=0.0) acc_g = tl.dot(x, wg, acc_g) acc_u = tl.dot(x, wu, acc_u) x_ptrs += BK w_ptrs_g += BK * I w_ptrs_u += BK * I res = acc_g * tl.sigmoid(acc_g) * acc_u out = res.to(tl.bfloat16) out_ptrs = out_ptr + (row0 + rm)[:, None] * stride_om + cols[None, :] row_mask = (row0 + rm) < end_cur tl.store(out_ptrs, out, mask=row_mask[:, None] & col_mask[None, :]) class Model(nn.Module): def __init__(self, T_total, H, I, E, K): super().__init__() self.T_total, self.H, self.I, self.E, self.K = T_total, H, I, E, 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._w_comb = None self._w_key = None def _get_w_comb(self): key = (self.W_gate.data_ptr(), self.W_up.data_ptr(), self.W_gate._version, self.W_up._version) if self._w_key != key: E, H, I = self.E, self.H, self.I w = torch.empty(E, H, 2 * I, dtype=torch.bfloat16, device=self.W_gate.device) w[:, :, 0::2] = self.W_gate w[:, :, 1::2] = self.W_up self._w_comb = w self._w_key = key return self._w_comb def forward(self, hidden_states, expert_offsets): T_perm, H = hidden_states.shape I, E = self.I, self.E out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device) cfg = _cfg_for(H, self.I) BM, BN, BK = cfg["BM"], cfg["BN"], cfg["BK"] NB = triton.cdiv(I, BN) upper = (triton.cdiv(T_perm, BM) + E) * NB grid = (min(NUM_SMS * cfg["grid_mult"], upper),) max_iter = triton.cdiv(upper, grid[0]) kw = dict(num_warps=cfg["warps"], num_stages=cfg["stages"]) if cfg.get("maxnreg"): kw["maxnreg"] = cfg["maxnreg"] if H % BK == 0 and I % BN == 0: w_comb = self._get_w_comb() _kernel_v5[grid]( hidden_states, w_comb, out, expert_offsets, hidden_states.stride(0), w_comb.stride(1), out.stride(0), max_iter, E=E, GRID=grid[0], H=H, I=I, BM=BM, BN=BN, BK=BK, **kw) else: _kernel_generic[grid]( hidden_states, self.W_gate, self.W_up, out, expert_offsets, hidden_states.stride(0), out.stride(0), H, I, max_iter, E=E, GRID=grid[0], BM=BM, BN=BN, BK=BK, **kw) return out 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: 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]