"""Sonic-MoE up-projection: grouped GEMM + fused SwiGLU on H100 (SM90). Per expert e, with x_e = hidden_states[start:end] (n_e rows): h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]) Custom Triton kernel. The grid tiles the (M, N) output plane of a *single* expert's GEMM and reduces over K=H in the inner loop; the SwiGLU epilogue (silu(gate) * up) is applied in-register on the fp32 accumulators before the bf16 store. Experts are processed one at a time in a Python loop: each expert's two weight matrices (H*I*2 bf16 ~= 25 MB) and its activation strip (n_e*H*2 bf16 ~= 17 MB) both fit comfortably in the 52 MB L2, so every byte is read from HBM roughly once -- the canonical GEMM data-reuse schedule -- and the kernel runs at the HBM bandwidth ceiling rather than thrashing L2 with all 128 experts' weights in flight at once. This is a from-scratch indexed grouped-GEMM that dispatches two expert-weight matrices per block and fuses the SwiGLU nonlinear; no vendor matmul path is used. """ from __future__ import annotations import torch import triton import triton.language as tl OP_TYPE = "grouped_gemm_swiglu" SUPPORTED_PRECISIONS = ["bf16"] HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"] @triton.autotune( configs=[ # BLOCK_N is held at 128 across every config because the launch grid's # N-tile count (computed in Python from BLOCK_N before autotune runs) # must agree with the block size the chosen config actually uses. Only # BLOCK_K and num_stages are swept. triton.Config( {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128}, num_warps=8, num_stages=2, ), triton.Config( {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=3, ), triton.Config( {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=6, ), triton.Config( {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=4, ), ], key=["M_max", "N", "K"], ) @triton.jit def _grouped_swiglu_kernel( hidden_ptr, # (T_perm, H) bf16 wgate_ptr, # (H, I) bf16 -- single expert, already sliced wup_ptr, # (H, I) bf16 -- single expert, already sliced out_ptr, # (T_perm, I) bf16 start_row: int, # this expert's start row in hidden / output n_e: int, # this expert's token count H: int, I: int, stride_hm, stride_hk, stride_wgh, stride_wgi, stride_wuh, stride_wui, stride_om, stride_on, M_max, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, ): # Grid: (m-tile, n-tile) for ONE expert (caller launches per expert). pm = tl.program_id(0) pn = tl.program_id(1) m_start = pm * BLOCK_M n_start = pn * BLOCK_N if m_start >= n_e: return a_ptr = hidden_ptr + start_row * stride_hm o_ptr = out_ptr + start_row * stride_om a = tl.make_block_ptr( a_ptr, (n_e, H), (stride_hm, stride_hk), (m_start, 0), (BLOCK_M, BLOCK_K), (1, 0), ) bg = tl.make_block_ptr( wgate_ptr, (H, I), (stride_wgh, stride_wgi), (0, n_start), (BLOCK_K, BLOCK_N), (1, 0), ) bu = tl.make_block_ptr( wup_ptr, (H, I), (stride_wuh, stride_wui), (0, n_start), (BLOCK_K, BLOCK_N), (1, 0), ) gate_acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) up_acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) for k in range(0, H, BLOCK_K): a_tile = tl.load(a, boundary_check=(0, 1), padding_option="zero") bg_tile = tl.load(bg, boundary_check=(0, 1), padding_option="zero") bu_tile = tl.load(bu, boundary_check=(0, 1), padding_option="zero") gate_acc = tl.dot(a_tile, bg_tile, gate_acc) up_acc = tl.dot(a_tile, bu_tile, up_acc) a = tl.advance(a, (0, BLOCK_K)) bg = tl.advance(bg, (BLOCK_K, 0)) bu = tl.advance(bu, (BLOCK_K, 0)) sig = 1.0 / (1.0 + tl.math.exp(-gate_acc)) silu = gate_acc * sig out_val = (silu * up_acc).to(tl.bfloat16) o = tl.make_block_ptr( o_ptr, (n_e, I), (stride_om, stride_on), (m_start, n_start), (BLOCK_M, BLOCK_N), (1, 0), ) tl.store(o, out_val, boundary_check=(0, 1)) class Model(torch.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 = torch.nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16)) self.W_up = torch.nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16)) torch.nn.init.normal_(self.W_gate, std=0.02) torch.nn.init.normal_(self.W_up, std=0.02) def forward( self, hidden_states: torch.Tensor, # (T_perm, H) bf16 expert_offsets: torch.Tensor, # (E+1,) int32 ) -> torch.Tensor: T_perm, H = hidden_states.shape E = self.E I = self.I out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device) # Pull per-expert (start, n_e) to plain Python ints once, up front, so # the per-expert loop issues pure GPU kernel launches with no CPU sync # and no per-launch host-side tensor slicing. off_np = expert_offsets.to(torch.int32).tolist() splits = [(off_np[e], off_np[e + 1]) for e in range(E)] n_tiles_I = triton.cdiv(I, 128) for e in range(E): start, end = splits[e] n_e = end - start if n_e == 0: continue m_tiles = triton.cdiv(n_e, 128) wgate_e = self.W_gate[e] # (H, I) wup_e = self.W_up[e] # (H, I) _grouped_swiglu_kernel[(m_tiles, n_tiles_I)]( hidden_states, wgate_e, wup_e, out, start, n_e, H, I, *hidden_states.stride(), *wgate_e.stride(), *wup_e.stride(), *out.stride(), n_e, ) 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: 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]