"""First draft: Triton grouped GEMM with fused SwiGLU. The op: per expert e, compute h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]) where x_e is the slice of permuted hidden states routed to expert e. Strategy: - Pre-concatenate W_gate and W_up into W = (E, H, 2*I) bf16 - Pre-build a per-(expert, m_tile) work schedule on device - Launch a persistent Triton kernel that does the grouped GEMM with SwiGLU fused - Each CTA processes one (expert, m_tile) and computes BLOCK_M rows of output I """ from __future__ import annotations import torch import torch.nn as nn from typing import Tuple def _silu(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.silu(x) class Model(nn.Module): """Reference-style Model; the actual kernel is launched in forward().""" 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) # Combined weight is computed lazily in forward() (and cached) so state_dict # roundtrip works. self._W_cached = None def _get_W(self) -> torch.Tensor: # Concatenate along the N dim: (E, H, 2*I) where the first I are gate, last I are up. if self._W_cached is not None and self._W_cached.shape[0] == self.E: return self._W_cached W = torch.cat([self.W_gate, self.W_up], dim=-1).contiguous() self._W_cached = W return W def forward(self, hidden_states, expert_offsets): W = self._get_W() return _triton_grouped_swiglu(hidden_states, W, expert_offsets, self.H, self.I) # Import the kernel here so the actual compute runs in Triton. import triton import triton.language as tl @triton.jit def _grouped_swiglu_kernel( x_ptr, w_ptr, out_ptr, expert_offsets_ptr, expert_for_tile_ptr, # (total_tiles,) expert id for each tile m_start_for_tile_ptr, # (total_tiles,) starting row in permuted X for each tile m_size_for_tile_ptr, # (total_tiles,) actual m size of this tile H, I, # hidden dim, intermediate dim stride_xm, stride_xk, stride_we, stride_wk, stride_wn, stride_om, stride_on, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, # block size in N direction (for I, may be < I) BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr, ): pid = tl.program_id(0) # Look up tile info e = tl.load(expert_for_tile_ptr + pid) m_start = tl.load(m_start_for_tile_ptr + pid) m_size = tl.load(m_size_for_tile_ptr + pid) # Skip empty tiles (e == -1) if e == -1: return # Compute two passes: gate (n_offset=0) and up (n_offset=I) # Each pass iterates over the N dimension in chunks of BLOCK_N # We process all gate and up N tiles in two separate accumulators n_iters = tl.cdiv(I, BLOCK_N) # For each N tile we need separate gate, up accumulators # Allocate BLOCK_M x BLOCK_N accumulators # But we can also process them in pairs (gate_i, up_i) for the same i # This lets us use one N loop iteration # Per N tile: compute gate = X @ W[:,:I_block] and up = X @ W[:,I_block+I:I_block+2I_block] # Then SwiGLU and write offs_m = tl.arange(0, BLOCK_M) offs_k = tl.arange(0, BLOCK_K) m_mask = offs_m < m_size for n_idx in range(n_iters): n_offset_gate = n_idx * BLOCK_N n_offset_up = I + n_idx * BLOCK_N # for up pass, W[:, I+n_idx*BLOCK_N : I+(n_idx+1)*BLOCK_N] # First pass: gate a_ptrs = x_ptr + (m_start + offs_m)[:, None] * stride_xm + offs_k[None, :] * stride_xk offs_n = n_offset_gate + tl.arange(0, BLOCK_N) b_ptrs = w_ptr + e * stride_we + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(H, BLOCK_K)): a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) b = tl.load(b_ptrs) acc_gate += tl.dot(a, b) a_ptrs += BLOCK_K * stride_xk b_ptrs += BLOCK_K * stride_wk # Second pass: up a_ptrs = x_ptr + (m_start + offs_m)[:, None] * stride_xm + offs_k[None, :] * stride_xk offs_n = n_offset_up + tl.arange(0, BLOCK_N) b_ptrs = w_ptr + e * stride_we + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(H, BLOCK_K)): a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) b = tl.load(b_ptrs) acc_up += tl.dot(a, b) a_ptrs += BLOCK_K * stride_xk b_ptrs += BLOCK_K * stride_wk # SwiGLU silu_gate = acc_gate * tl.sigmoid(acc_gate) out_val = silu_gate * acc_up # Write to output n_mask = (n_offset_gate + tl.arange(0, BLOCK_N)) < I out_n = n_offset_gate + tl.arange(0, BLOCK_N) out_ptrs = out_ptr + (m_start + offs_m)[:, None] * stride_om + out_n[None, :] * stride_on tl.store(out_ptrs, out_val.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]) def _triton_grouped_swiglu( hidden_states: torch.Tensor, # (T_perm, H) bf16 W: torch.Tensor, # (E, H, 2*I) bf16 expert_offsets: torch.Tensor, # (E+1,) int32 H: int, I: int, # noqa: E741 ) -> torch.Tensor: T_perm, H_ = hidden_states.shape E = W.shape[0] I_ = I assert W.shape[1] == H assert W.shape[2] == 2 * I out = torch.empty(T_perm, I_, dtype=torch.bfloat16, device=hidden_states.device) # Build per-tile work schedule # For each expert, ceil(n_e / BLOCK_M) tiles # First, on GPU: compute n_e = offsets[e+1] - offsets[e] # Then for each e, n_tiles = cdiv(n_e, BLOCK_M) # Total tiles = sum of n_tiles # We'll do this on the CPU/GPU and pass to the kernel BLOCK_M = 64 BLOCK_N = 64 # one tile covers part of I BLOCK_K = 64 GROUP_M = 8 # Build schedule on device using torch ops offs_cpu = expert_offsets.cpu().to(torch.int64) n_per_expert = (offs_cpu[1:] - offs_cpu[:-1]).to(torch.int32) n_tiles_per_expert = (n_per_expert + (BLOCK_M - 1)) // BLOCK_M # m_offsets_in_permuted: where each expert starts # For expert e with n_e tokens: # tile t (0 <= t < n_tiles) covers rows [m_start + t*BLOCK_M, m_start + min((t+1)*BLOCK_M, n_e)) # where m_start = expert_offsets[e] total_tiles = int(n_tiles_per_expert.sum().item()) expert_for_tile = torch.empty(total_tiles, dtype=torch.int32, device='cpu') m_start_for_tile = torch.empty(total_tiles, dtype=torch.int32, device='cpu') m_size_for_tile = torch.empty(total_tiles, dtype=torch.int32, device='cpu') e_idx = 0 for e in range(E): m_start_e = int(offs_cpu[e].item()) n_e = int(offs_cpu[e + 1].item()) - m_start_e n_tiles_e = (n_e + BLOCK_M - 1) // BLOCK_M for t in range(n_tiles_e): expert_for_tile[e_idx] = e m_start_for_tile[e_idx] = m_start_e + t * BLOCK_M m_size_for_tile[e_idx] = min(BLOCK_M, n_e - t * BLOCK_M) e_idx += 1 expert_for_tile = expert_for_tile.to(hidden_states.device, non_blocking=True) m_start_for_tile = m_start_for_tile.to(hidden_states.device, non_blocking=True) m_size_for_tile = m_size_for_tile.to(hidden_states.device, non_blocking=True) grid = (total_tiles,) _grouped_swiglu_kernel[grid]( hidden_states, W, out, expert_offsets, expert_for_tile, m_start_for_tile, m_size_for_tile, H, I_, hidden_states.stride(0), hidden_states.stride(1), W.stride(0), W.stride(1), W.stride(2), out.stride(0), out.stride(1), BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_M=GROUP_M, num_warps=4, num_stages=2, ) return out # Test if __name__ == "__main__": import sys sys.path.insert(0, "/workspace/problems/06_sonic_moe_swiglu") import reference torch.manual_seed(0) T_total, H, I, E, K = 4096, 2048, 1024, 64, 4 reference.T_total = T_total reference.H = H reference.I = I reference.E = E reference.K = K device = torch.device("cuda:0") ref_model = reference.Model(T_total, H, I, E, K).to(device).eval() sol_model = Model(T_total, H, I, E, K).to(device).eval() sol_model.load_state_dict(ref_model.state_dict()) inputs = [t.to(device) for t in reference.get_inputs()] with torch.no_grad(): ref_out = ref_model(*inputs) sol_out = sol_model(*inputs) diff = (ref_out.float() - sol_out.float()).abs() print("max abs diff:", diff.max().item()) print("max rel diff:", (diff / (ref_out.float().abs() + 1e-9)).max().item()) print("PASS" if diff.max().item() < 0.02 else "FAIL")