--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/_test_triton.py
@@ -0,0 +1,26 @@
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _test_dot(A, B, C, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    rm = pid_m * M + tl.arange(0, M)
+    rn = pid_n * N + tl.arange(0, N)
+    rk = tl.arange(0, K)
+    a = tl.load(A + rm[:, None] * K + rk[None, :])
+    b = tl.load(B + rk[:, None] * N + rn[None, :])
+    c = tl.dot(a.to(tl.bfloat16), b.to(tl.bfloat16))
+    tl.store(C + rm[:, None] * N + rn[None, :], c.to(tl.float32))
+
+M, N, K = 64, 64, 32
+A = torch.randn(M, K, dtype=torch.bfloat16, device='cuda')
+B = torch.randn(K, N, dtype=torch.bfloat16, device='cuda')
+C = torch.empty(M, N, dtype=torch.float32, device='cuda')
+_test_dot[(1, 1)](A, B, C, M, N, K)
+ref = (A.float() @ B.float())
+print('Dot test max diff:', (C - ref).abs().max().item())
+print('Dot test matches:', torch.allclose(C, ref, atol=0.01, rtol=0.01))
+print('SM:', torch.cuda.get_device_capability(0))
+print('Triton OK on SM120')

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -0,0 +1,211 @@
+"""Grouped GEMM + fused SwiGLU for MoE up-projection.
+
+Triton kernel: per expert e, computes 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.
+
+Each kernel instance handles one (expert, m_tile, n_tile) triple, iterating
+over the reduction dimension K (hidden dim H) in tiles.  Two accumulators are
+maintained (gate + up) and SwiGLU is fused in the epilogue.
+"""
+from __future__ import annotations
+
+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"]
+
+
+# ---------------------------------------------------------------------------
+# Triton kernel
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def _grouped_swiglu_kernel(
+    # Pointers
+    hidden_states_ptr,   # (T_perm, H) bf16
+    W_gate_ptr,          # (E, H, I) bf16
+    W_up_ptr,            # (E, H, I) bf16
+    out_ptr,             # (T_perm, I) bf16
+    expert_offsets_ptr,  # (E+1,) int32
+    # Strides
+    stride_xh, stride_xw,         # hidden_states row/col stride
+    stride_wg_e, stride_wg_h, stride_wg_i,  # W_gate
+    stride_wu_e, stride_wu_h, stride_wu_i,  # W_up
+    stride_oh, stride_ow,          # output row/col stride
+    # Dimensions
+    H: tl.constexpr,
+    I: tl.constexpr,
+    # Tile sizes
+    BM: tl.constexpr,   # tile over token (M) dim
+    BN: tl.constexpr,   # tile over output (N) dim
+    BK: tl.constexpr,   # tile over reduction (K) dim
+):
+    # 3-D grid: (num_m_tiles, num_n_tiles, E)
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    expert_id = tl.program_id(2)
+
+    # Expert row range
+    start = tl.load(expert_offsets_ptr + expert_id)
+    end = tl.load(expert_offsets_ptr + expert_id + 1)
+    n_rows = end - start
+
+    # This tile's M range (relative to expert start)
+    row_off = pid_m * BM
+    # Early exit if entire tile is beyond this expert's rows
+    if row_off >= n_rows:
+        return
+
+    # Accumulators
+    acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
+
+    # N range for this tile
+    col_off = pid_n * BN
+
+    # Base pointers for this expert
+    x_base  = hidden_states_ptr + (start + row_off) * stride_xh
+    wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
+    wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
+
+    # Offsets for loads
+    offs_m = tl.arange(0, BM)
+    offs_n = tl.arange(0, BN)
+    m_mask = (row_off + offs_m) < n_rows
+    n_mask = (col_off + offs_n) < I
+
+    for k in range(0, tl.cdiv(H, BK)):
+        k_off = k * BK
+        offs_k = tl.arange(0, BK)
+        k_mask = (k_off + offs_k) < H
+
+        # Load x tile (BM, BK) - shared by gate and up
+        x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+        x_mask_2d = m_mask[:, None] & k_mask[None, :]
+        x_tile = tl.load(x_ptrs, mask=x_mask_2d, other=0.0).to(tl.bfloat16)
+
+        # Load W_gate tile (BK, BN)
+        wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
+        wg_mask_2d = k_mask[:, None] & n_mask[None, :]
+        wg_tile = tl.load(wg_ptrs, mask=wg_mask_2d, other=0.0).to(tl.bfloat16)
+
+        # Load W_up tile (BK, BN)
+        wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
+        wu_tile = tl.load(wu_ptrs, mask=wg_mask_2d, other=0.0).to(tl.bfloat16)
+
+        # Accumulate both GEMMs
+        acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
+        acc_up   = tl.dot(x_tile, wu_tile, acc_up)
+
+    # Epilogue: SwiGLU = silu(gate) * up
+    silu_gate = acc_gate * tl.sigmoid(acc_gate)
+    result = silu_gate * acc_up
+    result_bf16 = result.to(tl.bfloat16)
+
+    # Store output
+    out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
+    out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
+    out_mask_2d = m_mask[:, None] & n_mask[None, :]
+    tl.store(out_ptrs, result_bf16, mask=out_mask_2d)
+
+
+# ---------------------------------------------------------------------------
+# Python wrapper
+# ---------------------------------------------------------------------------
+
+class Model(nn.Module):
+    """Up-projection of a top-K MoE FFN with fused SwiGLU.
+
+    Inputs at call time:
+      hidden_states:    (T_perm, H)  bf16, already permuted to expert order
+      expert_offsets:   (E+1,)       int32, prefix sums of token counts per expert
+
+    Output:
+      gated_up:         (T_perm, I)  bf16
+    """
+
+    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)
+
+    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
+        I = self.I  # noqa: E741
+        E = self.E
+
+        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
+
+        # Tile sizes
+        BM = 64
+        BN = 64
+        BK = 32
+
+        # Compute per-expert row counts to determine grid dims
+        offsets_cpu = expert_offsets.cpu()
+        max_rows = int((offsets_cpu[1:] - offsets_cpu[:-1]).max().item())
+        num_m_tile

... (1548 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -3,9 +3,9 @@
 Triton kernel: per expert e, computes 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.
 
-Each kernel instance handles one (expert, m_tile, n_tile) triple, iterating
-over the reduction dimension K (hidden dim H) in tiles.  Two accumulators are
-maintained (gate + up) and SwiGLU is fused in the epilogue.
+Each kernel instance handles one (expert, m_tile) pair and iterates over ALL
+output-column (N) tiles, fusing both gate and up GEMMs with SwiGLU epilogue.
+This way x_tile is loaded once from DRAM and reused across all N tiles.
 """
 from __future__ import annotations
 
@@ -19,115 +19,77 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-# ---------------------------------------------------------------------------
-# Triton kernel
-# ---------------------------------------------------------------------------
-
 @triton.jit
 def _grouped_swiglu_kernel(
-    # Pointers
-    hidden_states_ptr,   # (T_perm, H) bf16
-    W_gate_ptr,          # (E, H, I) bf16
-    W_up_ptr,            # (E, H, I) bf16
-    out_ptr,             # (T_perm, I) bf16
-    expert_offsets_ptr,  # (E+1,) int32
-    # Strides
-    stride_xh, stride_xw,         # hidden_states row/col stride
-    stride_wg_e, stride_wg_h, stride_wg_i,  # W_gate
-    stride_wu_e, stride_wu_h, stride_wu_i,  # W_up
-    stride_oh, stride_ow,          # output row/col stride
-    # Dimensions
-    H: tl.constexpr,
-    I: tl.constexpr,
-    # Tile sizes
-    BM: tl.constexpr,   # tile over token (M) dim
-    BN: tl.constexpr,   # tile over output (N) dim
-    BK: tl.constexpr,   # tile over reduction (K) dim
+    hidden_states_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    out_ptr,
+    expert_offsets_ptr,
+    stride_xh, stride_xw,
+    stride_wg_e, stride_wg_h, stride_wg_i,
+    stride_wu_e, stride_wu_h, stride_wu_i,
+    stride_oh, stride_ow,
+    H,
+    I,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
 ):
-    # 3-D grid: (num_m_tiles, num_n_tiles, E)
     pid_m = tl.program_id(0)
-    pid_n = tl.program_id(1)
-    expert_id = tl.program_id(2)
+    expert_id = tl.program_id(1)
 
-    # Expert row range
     start = tl.load(expert_offsets_ptr + expert_id)
     end = tl.load(expert_offsets_ptr + expert_id + 1)
     n_rows = end - start
 
-    # This tile's M range (relative to expert start)
     row_off = pid_m * BM
-    # Early exit if entire tile is beyond this expert's rows
     if row_off >= n_rows:
         return
 
-    # Accumulators
-    acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
-    acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
-
-    # N range for this tile
-    col_off = pid_n * BN
-
-    # Base pointers for this expert
-    x_base  = hidden_states_ptr + (start + row_off) * stride_xh
-    wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
-    wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
-
-    # Offsets for loads
     offs_m = tl.arange(0, BM)
     offs_n = tl.arange(0, BN)
+    offs_k = tl.arange(0, BK)
     m_mask = (row_off + offs_m) < n_rows
-    n_mask = (col_off + offs_n) < I
 
-    for k in range(0, tl.cdiv(H, BK)):
-        k_off = k * BK
-        offs_k = tl.arange(0, BK)
-        k_mask = (k_off + offs_k) < H
+    # Loop over output-column tiles (N dimension)
+    n_tiles = tl.cdiv(I, BN)
+    for pid_n in range(n_tiles):
+        col_off = pid_n * BN
+        n_mask = (col_off + offs_n) < I
 
-        # Load x tile (BM, BK) - shared by gate and up
-        x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-        x_mask_2d = m_mask[:, None] & k_mask[None, :]
-        x_tile = tl.load(x_ptrs, mask=x_mask_2d, other=0.0).to(tl.bfloat16)
+        acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
+        acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
 
-        # Load W_gate tile (BK, BN)
-        wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-        wg_mask_2d = k_mask[:, None] & n_mask[None, :]
-        wg_tile = tl.load(wg_ptrs, mask=wg_mask_2d, other=0.0).to(tl.bfloat16)
+        x_base  = hidden_states_ptr + (start + row_off) * stride_xh
+        wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
+        wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
-        # Load W_up tile (BK, BN)
-        wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-        wu_tile = tl.load(wu_ptrs, mask=wg_mask_2d, other=0.0).to(tl.bfloat16)
+        for k in range(0, tl.cdiv(H, BK)):
+            k_off = k * BK
+            k_mask = (k_off + offs_k) < H
 
-        # Accumulate both GEMMs
-        acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
-        acc_up   = tl.dot(x_tile, wu_tile, acc_up)
+            x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+            x_tile = tl.load(x_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0).to(tl.bfloat16)
 
-    # Epilogue: SwiGLU = silu(gate) * up
-    silu_gate = acc_gate * tl.sigmoid(acc_gate)
-    result = silu_gate * acc_up
-    result_bf16 = result.to(tl.bfloat16)
+            wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
+            wg_tile = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
 
-    # Store output
-    out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
-    out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
-    out_mask_2d = m_mask[:, None] & n_mask[None, :]
-    tl.store(out_ptrs, result_bf16, mask=out_mask_2d)
+            wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[N

... (3392 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -3,9 +3,8 @@
 Triton kernel: per expert e, computes 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.
 
-Each kernel instance handles one (expert, m_tile) pair and iterates over ALL
-output-column (N) tiles, fusing both gate and up GEMMs with SwiGLU epilogue.
-This way x_tile is loaded once from DRAM and reused across all N tiles.
+Autotuned tile sizes. Each kernel instance handles one (expert, m_tile, n_tile)
+triple with fused gate+up GEMMs and SwiGLU epilogue.
 """
 from __future__ import annotations
 
@@ -19,6 +18,18 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
+@triton.autotune(
+    configs=[
+        triton.Config({'BM': 64, 'BN': 64, 'BK': 32}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 64, 'BN': 128, 'BK': 32}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 128, 'BN': 64, 'BK': 32}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 128, 'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
+        triton.Config({'BM': 64, 'BN': 64, 'BK': 64}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 64, 'BN': 128, 'BK': 64}, num_warps=8, num_stages=2),
+        triton.Config({'BM': 128, 'BN': 64, 'BK': 64}, num_warps=8, num_stages=2),
+    ],
+    key=['H', 'I', 'max_rows'],
+)
 @triton.jit
 def _grouped_swiglu_kernel(
     hidden_states_ptr,
@@ -32,12 +43,14 @@
     stride_oh, stride_ow,
     H,
     I,
+    max_rows,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
-    expert_id = tl.program_id(1)
+    pid_n = tl.program_id(1)
+    expert_id = tl.program_id(2)
 
     start = tl.load(expert_offsets_ptr + expert_id)
     end = tl.load(expert_offsets_ptr + expert_id + 1)
@@ -47,46 +60,44 @@
     if row_off >= n_rows:
         return
 
+    col_off = pid_n * BN
+
+    acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
+
     offs_m = tl.arange(0, BM)
     offs_n = tl.arange(0, BN)
     offs_k = tl.arange(0, BK)
     m_mask = (row_off + offs_m) < n_rows
+    n_mask = (col_off + offs_n) < I
 
-    # Loop over output-column tiles (N dimension)
-    n_tiles = tl.cdiv(I, BN)
-    for pid_n in range(n_tiles):
-        col_off = pid_n * BN
-        n_mask = (col_off + offs_n) < I
+    x_base  = hidden_states_ptr + (start + row_off) * stride_xh
+    wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
+    wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
-        acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
-        acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
+    for k in range(0, tl.cdiv(H, BK)):
+        k_off = k * BK
+        k_mask = (k_off + offs_k) < H
+        mask_2d = m_mask[:, None] & k_mask[None, :]
 
-        x_base  = hidden_states_ptr + (start + row_off) * stride_xh
-        wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
-        wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
+        x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+        x_tile = tl.load(x_ptrs, mask=mask_2d, other=0.0).to(tl.bfloat16)
 
-        for k in range(0, tl.cdiv(H, BK)):
-            k_off = k * BK
-            k_mask = (k_off + offs_k) < H
+        wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
+        wg_tile = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
 
-            x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-            x_tile = tl.load(x_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0).to(tl.bfloat16)
+        wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
+        wu_tile = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
 
-            wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-            wg_tile = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
+        acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
+        acc_up   = tl.dot(x_tile, wu_tile, acc_up)
 
-            wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-            wu_tile = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
+    silu_gate = acc_gate * tl.sigmoid(acc_gate)
+    result = (silu_gate * acc_up).to(tl.bfloat16)
 
-            acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
-            acc_up   = tl.dot(x_tile, wu_tile, acc_up)
-
-        silu_gate = acc_gate * tl.sigmoid(acc_gate)
-        result = (silu_gate * acc_up).to(tl.bfloat16)
-
-        out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
-        out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
-        tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
+    out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
+    out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
+    tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
 
 
 class Model(nn.Module):
@@ -113,23 +124,22 @@
 
         out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
 
-        BM = 128
-        BN = 128
-        BK = 32
-
-        # Compute max rows on GPU to avoid sync
         diffs = expert_offsets[1:] - expert_offsets[:-1]
         max_rows = int(diffs.max().item())
-        num_m_tiles = triton.cdiv(max_rows, BM)
+        # Use a fixed upper bound to avoid sync -- num_m_tiles will be large
+        # enough; extra tiles early-exit.  T_perm

... (760 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -3,8 +3,9 @@
 Triton kernel: per expert e, computes 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.
 
-Autotuned tile sizes. Each kernel instance handles one (expert, m_tile, n_tile)
-triple with fused gate+up GEMMs and SwiGLU epilogue.
+Each kernel instance handles one (expert, n_tile, m_tile) triple.  Grid is
+ordered (E, num_n_tiles, num_m_tiles) so that consecutive programs share the
+same weight tile in L2 (adjacent m_tiles load the same (BK, BN) weight slice).
 """
 from __future__ import annotations
 
@@ -18,18 +19,6 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({'BM': 64, 'BN': 64, 'BK': 32}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 64, 'BN': 128, 'BK': 32}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 128, 'BN': 64, 'BK': 32}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 128, 'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
-        triton.Config({'BM': 64, 'BN': 64, 'BK': 64}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 64, 'BN': 128, 'BK': 64}, num_warps=8, num_stages=2),
-        triton.Config({'BM': 128, 'BN': 64, 'BK': 64}, num_warps=8, num_stages=2),
-    ],
-    key=['H', 'I', 'max_rows'],
-)
 @triton.jit
 def _grouped_swiglu_kernel(
     hidden_states_ptr,
@@ -43,17 +32,17 @@
     stride_oh, stride_ow,
     H,
     I,
-    max_rows,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
-    pid_m = tl.program_id(0)
-    pid_n = tl.program_id(1)
-    expert_id = tl.program_id(2)
+    # Grid: (E, num_n_tiles, num_m_tiles)  -- consecutive m_tiles share weights
+    expert_id = tl.program_id(0)
+    pid_n     = tl.program_id(1)
+    pid_m     = tl.program_id(2)
 
     start = tl.load(expert_offsets_ptr + expert_id)
-    end = tl.load(expert_offsets_ptr + expert_id + 1)
+    end   = tl.load(expert_offsets_ptr + expert_id + 1)
     n_rows = end - start
 
     row_off = pid_m * BM
@@ -78,16 +67,17 @@
     for k in range(0, tl.cdiv(H, BK)):
         k_off = k * BK
         k_mask = (k_off + offs_k) < H
-        mask_2d = m_mask[:, None] & k_mask[None, :]
+        mk_mask = m_mask[:, None] & k_mask[None, :]
+        kn_mask = k_mask[:, None] & n_mask[None, :]
 
-        x_ptrs = x_base + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-        x_tile = tl.load(x_ptrs, mask=mask_2d, other=0.0).to(tl.bfloat16)
+        x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+        x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
 
         wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-        wg_tile = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
+        wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
         wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-        wu_tile = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
+        wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
         acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
         acc_up   = tl.dot(x_tile, wu_tile, acc_up)
@@ -124,22 +114,24 @@
 
         out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
 
+        BM = 64
+        BN = 64
+        BK = 32
+
         diffs = expert_offsets[1:] - expert_offsets[:-1]
         max_rows = int(diffs.max().item())
-        # Use a fixed upper bound to avoid sync -- num_m_tiles will be large
-        # enough; extra tiles early-exit.  T_perm = T_total * K, balanced
-        # routing means max_rows ~ T_perm / E.
-        num_m_tiles = triton.cdiv(max_rows, 64)  # min BM
-        num_n_tiles = triton.cdiv(I, 64)  # min BN
+        num_m_tiles = triton.cdiv(max_rows, BM)
+        num_n_tiles = triton.cdiv(I, BN)
 
-        grid = (num_m_tiles, num_n_tiles, E)
+        # Grid: (E, n_tiles, m_tiles) for L2 weight reuse across m_tiles
+        grid = (E, num_n_tiles, num_m_tiles)
         _grouped_swiglu_kernel[grid](
             hidden_states, self.W_gate, self.W_up, out, expert_offsets,
             hidden_states.stride(0), hidden_states.stride(1),
             self.W_gate.stride(0), self.W_gate.stride(1), self.W_gate.stride(2),
             self.W_up.stride(0), self.W_up.stride(1), self.W_up.stride(2),
             out.stride(0), out.stride(1),
-            H=H, I=I, max_rows=max_rows,
+            H=H, I=I, BM=BM, BN=BN, BK=BK,
         )
         return out
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -3,9 +3,8 @@
 Triton kernel: per expert e, computes 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.
 
-Each kernel instance handles one (expert, n_tile, m_tile) triple.  Grid is
-ordered (E, num_n_tiles, num_m_tiles) so that consecutive programs share the
-same weight tile in L2 (adjacent m_tiles load the same (BK, BN) weight slice).
+Grid: (num_m_tiles, num_n_tiles, E). Each kernel instance handles one
+(expert, m_tile, n_tile) triple with fused gate+up GEMMs and SwiGLU epilogue.
 """
 from __future__ import annotations
 
@@ -30,16 +29,15 @@
     stride_wg_e, stride_wg_h, stride_wg_i,
     stride_wu_e, stride_wu_h, stride_wu_i,
     stride_oh, stride_ow,
-    H,
-    I,
+    H: tl.constexpr,
+    I: tl.constexpr,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
-    # Grid: (E, num_n_tiles, num_m_tiles)  -- consecutive m_tiles share weights
-    expert_id = tl.program_id(0)
-    pid_n     = tl.program_id(1)
-    pid_m     = tl.program_id(2)
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    expert_id = tl.program_id(2)
 
     start = tl.load(expert_offsets_ptr + expert_id)
     end   = tl.load(expert_offsets_ptr + expert_id + 1)
@@ -123,8 +121,7 @@
         num_m_tiles = triton.cdiv(max_rows, BM)
         num_n_tiles = triton.cdiv(I, BN)
 
-        # Grid: (E, n_tiles, m_tiles) for L2 weight reuse across m_tiles
-        grid = (E, num_n_tiles, num_m_tiles)
+        grid = (num_m_tiles, num_n_tiles, E)
         _grouped_swiglu_kernel[grid](
             hidden_states, self.W_gate, self.W_up, out, expert_offsets,
             hidden_states.stride(0), hidden_states.stride(1),

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -5,6 +5,7 @@
 
 Grid: (num_m_tiles, num_n_tiles, E). Each kernel instance handles one
 (expert, m_tile, n_tile) triple with fused gate+up GEMMs and SwiGLU epilogue.
+No GPU->CPU sync: grid sized for worst-case expert size with early-exit.
 """
 from __future__ import annotations
 
@@ -62,20 +63,18 @@
     wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
     wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
-    for k in range(0, tl.cdiv(H, BK)):
+    num_k_tiles = H // BK  # H is always divisible by BK for our shapes
+    for k in range(num_k_tiles):
         k_off = k * BK
-        k_mask = (k_off + offs_k) < H
-        mk_mask = m_mask[:, None] & k_mask[None, :]
-        kn_mask = k_mask[:, None] & n_mask[None, :]
 
         x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-        x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
+        x_tile  = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0).to(tl.bfloat16)
 
         wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-        wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+        wg_tile = tl.load(wg_ptrs, mask=n_mask[None, :], other=0.0).to(tl.bfloat16)
 
         wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-        wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+        wu_tile = tl.load(wu_ptrs, mask=n_mask[None, :], other=0.0).to(tl.bfloat16)
 
         acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
         acc_up   = tl.dot(x_tile, wu_tile, acc_up)
@@ -116,9 +115,9 @@
         BN = 64
         BK = 32
 
-        diffs = expert_offsets[1:] - expert_offsets[:-1]
-        max_rows = int(diffs.max().item())
-        num_m_tiles = triton.cdiv(max_rows, BM)
+        # No GPU->CPU sync: use worst-case grid (every expert could have T_perm tokens)
+        # Extra tiles early-exit when row_off >= n_rows.
+        num_m_tiles = triton.cdiv(T_perm, BM)
         num_n_tiles = triton.cdiv(I, BN)
 
         grid = (num_m_tiles, num_n_tiles, E)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -5,7 +5,8 @@
 
 Grid: (num_m_tiles, num_n_tiles, E). Each kernel instance handles one
 (expert, m_tile, n_tile) triple with fused gate+up GEMMs and SwiGLU epilogue.
-No GPU->CPU sync: grid sized for worst-case expert size with early-exit.
+
+Zero GPU->CPU syncs: grid is sized based on known routing structure.
 """
 from __future__ import annotations
 
@@ -63,18 +64,20 @@
     wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
     wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
-    num_k_tiles = H // BK  # H is always divisible by BK for our shapes
-    for k in range(num_k_tiles):
+    for k in range(0, tl.cdiv(H, BK)):
         k_off = k * BK
+        k_mask = (k_off + offs_k) < H
+        mk_mask = m_mask[:, None] & k_mask[None, :]
+        kn_mask = k_mask[:, None] & n_mask[None, :]
 
         x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-        x_tile  = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0).to(tl.bfloat16)
+        x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
 
         wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-        wg_tile = tl.load(wg_ptrs, mask=n_mask[None, :], other=0.0).to(tl.bfloat16)
+        wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
         wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-        wu_tile = tl.load(wu_ptrs, mask=n_mask[None, :], other=0.0).to(tl.bfloat16)
+        wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
         acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
         acc_up   = tl.dot(x_tile, wu_tile, acc_up)
@@ -115,9 +118,10 @@
         BN = 64
         BK = 32
 
-        # No GPU->CPU sync: use worst-case grid (every expert could have T_perm tokens)
-        # Extra tiles early-exit when row_off >= n_rows.
-        num_m_tiles = triton.cdiv(T_perm, BM)
+        # Zero GPU->CPU sync: for balanced routing, max_rows = ceil(T_perm / E).
+        # Grid is sized for worst case; extra tiles early-exit immediately.
+        max_rows_per_expert = (T_perm + E - 1) // E
+        num_m_tiles = triton.cdiv(max_rows_per_expert, BM)
         num_n_tiles = triton.cdiv(I, BN)
 
         grid = (num_m_tiles, num_n_tiles, E)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -5,8 +5,7 @@
 
 Grid: (num_m_tiles, num_n_tiles, E). Each kernel instance handles one
 (expert, m_tile, n_tile) triple with fused gate+up GEMMs and SwiGLU epilogue.
-
-Zero GPU->CPU syncs: grid is sized based on known routing structure.
+Zero GPU->CPU syncs.
 """
 from __future__ import annotations
 
@@ -20,6 +19,17 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
+@triton.autotune(
+    configs=[
+        triton.Config({'BM': 64,  'BN': 64,  'BK': 32}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 64,  'BN': 64,  'BK': 64}, num_warps=4, num_stages=2),
+        triton.Config({'BM': 64,  'BN': 64,  'BK': 32}, num_warps=4, num_stages=3),
+        triton.Config({'BM': 64,  'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
+        triton.Config({'BM': 128, 'BN': 64,  'BK': 32}, num_warps=8, num_stages=2),
+        triton.Config({'BM': 128, 'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
+    ],
+    key=['H', 'I', 'max_rows_per_expert'],
+)
 @triton.jit
 def _grouped_swiglu_kernel(
     hidden_states_ptr,
@@ -31,8 +41,9 @@
     stride_wg_e, stride_wg_h, stride_wg_i,
     stride_wu_e, stride_wu_h, stride_wu_i,
     stride_oh, stride_ow,
-    H: tl.constexpr,
-    I: tl.constexpr,
+    H,
+    I,
+    max_rows_per_expert,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
@@ -114,15 +125,11 @@
 
         out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
 
-        BM = 64
-        BN = 64
-        BK = 32
-
-        # Zero GPU->CPU sync: for balanced routing, max_rows = ceil(T_perm / E).
-        # Grid is sized for worst case; extra tiles early-exit immediately.
         max_rows_per_expert = (T_perm + E - 1) // E
-        num_m_tiles = triton.cdiv(max_rows_per_expert, BM)
-        num_n_tiles = triton.cdiv(I, BN)
+        # Grid based on worst-case (balanced routing).  Autotune picks BM/BN.
+        # Use minimum tile size for grid computation so all configs fit.
+        num_m_tiles = triton.cdiv(max_rows_per_expert, 64)
+        num_n_tiles = triton.cdiv(I, 64)
 
         grid = (num_m_tiles, num_n_tiles, E)
         _grouped_swiglu_kernel[grid](
@@ -131,7 +138,7 @@
             self.W_gate.stride(0), self.W_gate.stride(1), self.W_gate.stride(2),
             self.W_up.stride(0), self.W_up.stride(1), self.W_up.stride(2),
             out.stride(0), out.stride(1),
-            H=H, I=I, BM=BM, BN=BN, BK=BK,
+            H=H, I=I, max_rows_per_expert=max_rows_per_expert,
         )
         return out
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -3,9 +3,9 @@
 Triton kernel: per expert e, computes 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.
 
-Grid: (num_m_tiles, num_n_tiles, E). Each kernel instance handles one
-(expert, m_tile, n_tile) triple with fused gate+up GEMMs and SwiGLU epilogue.
-Zero GPU->CPU syncs.
+Each kernel instance handles one (expert, m_tile) pair and iterates over ALL
+N tiles.  x_tile is loaded once per K-step and reused across all N columns.
+This halves memory traffic vs per-(m,n) tile dispatch.
 """
 from __future__ import annotations
 
@@ -19,17 +19,6 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({'BM': 64,  'BN': 64,  'BK': 32}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 64,  'BN': 64,  'BK': 64}, num_warps=4, num_stages=2),
-        triton.Config({'BM': 64,  'BN': 64,  'BK': 32}, num_warps=4, num_stages=3),
-        triton.Config({'BM': 64,  'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
-        triton.Config({'BM': 128, 'BN': 64,  'BK': 32}, num_warps=8, num_stages=2),
-        triton.Config({'BM': 128, 'BN': 128, 'BK': 32}, num_warps=8, num_stages=2),
-    ],
-    key=['H', 'I', 'max_rows_per_expert'],
-)
 @triton.jit
 def _grouped_swiglu_kernel(
     hidden_states_ptr,
@@ -43,14 +32,12 @@
     stride_oh, stride_ow,
     H,
     I,
-    max_rows_per_expert,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
-    pid_n = tl.program_id(1)
-    expert_id = tl.program_id(2)
+    expert_id = tl.program_id(1)
 
     start = tl.load(expert_offsets_ptr + expert_id)
     end   = tl.load(expert_offsets_ptr + expert_id + 1)
@@ -60,45 +47,48 @@
     if row_off >= n_rows:
         return
 
-    col_off = pid_n * BN
-
-    acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
-    acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
-
     offs_m = tl.arange(0, BM)
     offs_n = tl.arange(0, BN)
     offs_k = tl.arange(0, BK)
     m_mask = (row_off + offs_m) < n_rows
-    n_mask = (col_off + offs_n) < I
 
-    x_base  = hidden_states_ptr + (start + row_off) * stride_xh
-    wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
-    wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
+    x_base = hidden_states_ptr + (start + row_off) * stride_xh
 
-    for k in range(0, tl.cdiv(H, BK)):
-        k_off = k * BK
-        k_mask = (k_off + offs_k) < H
-        mk_mask = m_mask[:, None] & k_mask[None, :]
-        kn_mask = k_mask[:, None] & n_mask[None, :]
+    n_tiles = tl.cdiv(I, BN)
+    for pid_n in range(n_tiles):
+        col_off = pid_n * BN
+        n_mask = (col_off + offs_n) < I
 
-        x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-        x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
+        acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
+        acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
 
-        wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-        wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+        wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
+        wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
-        wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-        wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+        for k in range(0, tl.cdiv(H, BK)):
+            k_off = k * BK
+            k_mask = (k_off + offs_k) < H
+            mk_mask = m_mask[:, None] & k_mask[None, :]
+            kn_mask = k_mask[:, None] & n_mask[None, :]
 
-        acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
-        acc_up   = tl.dot(x_tile, wu_tile, acc_up)
+            x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+            x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
 
-    silu_gate = acc_gate * tl.sigmoid(acc_gate)
-    result = (silu_gate * acc_up).to(tl.bfloat16)
+            wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
+            wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
-    out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
-    out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
-    tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
+            wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
+            wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+
+            acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
+            acc_up   = tl.dot(x_tile, wu_tile, acc_up)
+
+        silu_gate = acc_gate * tl.sigmoid(acc_gate)
+        result = (silu_gate * acc_up).to(tl.bfloat16)
+
+        out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
+        out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
+        tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
 
 
 class Model(nn.Module):
@@ -125,20 +115,21 @@
 
         out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
 
+        BM = 64
+        BN = 64
+        BK = 32
+
         max_rows_per_expert = (T_perm + E - 1) // E
-        # Grid based on worst-case (balanced routing).  Autotune picks BM/BN.
-        # Use minimum tile size for grid computation so all configs fit.
-        num_m_tiles = triton.cdiv(max_rows_per_expert, 64)
-        num_n_tiles = triton.cdiv(I, 64)
+        num_m_tiles = triton.cdiv(max_rows_per_expert, BM)
 
-    

... (591 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -1,11 +1,12 @@
 """Grouped GEMM + fused SwiGLU for MoE up-projection.
 
-Triton kernel: per expert e, computes 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.
+Two separate Triton kernels:
+  1. Grouped GEMM for gate: gate_out = x @ W_gate (per expert)
+  2. Grouped GEMM for up:   up_out   = x @ W_up   (per expert)
+  3. Fused SwiGLU epilogue: out = silu(gate_out) * up_out
 
-Each kernel instance handles one (expert, m_tile) pair and iterates over ALL
-N tiles.  x_tile is loaded once per K-step and reused across all N columns.
-This halves memory traffic vs per-(m,n) tile dispatch.
+This reduces register pressure per GEMM (only one accumulator) and allows
+the scheduler to overlap across the two GEMM launches.
 """
 from __future__ import annotations
 
@@ -20,75 +21,84 @@
 
 
 @triton.jit
-def _grouped_swiglu_kernel(
-    hidden_states_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    out_ptr,
-    expert_offsets_ptr,
-    stride_xh, stride_xw,
-    stride_wg_e, stride_wg_h, stride_wg_i,
-    stride_wu_e, stride_wu_h, stride_wu_i,
-    stride_oh, stride_ow,
-    H,
-    I,
+def _grouped_gemm_kernel(
+    A_ptr,       # (T_perm, H) bf16
+    B_ptr,       # (E, H, I) bf16
+    C_ptr,       # (T_perm, I) bf16
+    offsets_ptr, # (E+1,) int32
+    stride_ah, stride_aw,
+    stride_be, stride_bh, stride_bi,
+    stride_ch, stride_cw,
+    H: tl.constexpr,
+    I: tl.constexpr,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
-    expert_id = tl.program_id(1)
+    pid_n = tl.program_id(1)
+    expert_id = tl.program_id(2)
 
-    start = tl.load(expert_offsets_ptr + expert_id)
-    end   = tl.load(expert_offsets_ptr + expert_id + 1)
+    start = tl.load(offsets_ptr + expert_id)
+    end   = tl.load(offsets_ptr + expert_id + 1)
     n_rows = end - start
 
     row_off = pid_m * BM
     if row_off >= n_rows:
         return
 
+    col_off = pid_n * BN
+
+    acc = tl.zeros((BM, BN), dtype=tl.float32)
+
     offs_m = tl.arange(0, BM)
     offs_n = tl.arange(0, BN)
     offs_k = tl.arange(0, BK)
     m_mask = (row_off + offs_m) < n_rows
+    n_mask = (col_off + offs_n) < I
 
-    x_base = hidden_states_ptr + (start + row_off) * stride_xh
+    a_base = A_ptr + (start + row_off) * stride_ah
+    b_base = B_ptr + expert_id * stride_be + col_off * stride_bi
 
-    n_tiles = tl.cdiv(I, BN)
-    for pid_n in range(n_tiles):
-        col_off = pid_n * BN
-        n_mask = (col_off + offs_n) < I
+    for k in range(0, tl.cdiv(H, BK)):
+        k_off = k * BK
+        k_mask = (k_off + offs_k) < H
 
-        acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
-        acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
+        a_ptrs = a_base + offs_m[:, None] * stride_ah + (k_off + offs_k)[None, :] * stride_aw
+        a_tile = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0).to(tl.bfloat16)
 
-        wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
-        wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
+        b_ptrs = b_base + (k_off + offs_k)[:, None] * stride_bh + offs_n[None, :] * stride_bi
+        b_tile = tl.load(b_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
 
-        for k in range(0, tl.cdiv(H, BK)):
-            k_off = k * BK
-            k_mask = (k_off + offs_k) < H
-            mk_mask = m_mask[:, None] & k_mask[None, :]
-            kn_mask = k_mask[:, None] & n_mask[None, :]
+        acc = tl.dot(a_tile, b_tile, acc)
 
-            x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
-            x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
+    result = acc.to(tl.bfloat16)
 
-            wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
-            wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
+    c_base = C_ptr + (start + row_off) * stride_ch + col_off * stride_cw
+    c_ptrs = c_base + offs_m[:, None] * stride_ch + offs_n[None, :] * stride_cw
+    tl.store(c_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
 
-            wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
-            wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
-            acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
-            acc_up   = tl.dot(x_tile, wu_tile, acc_up)
+@triton.jit
+def _swiglu_epilogue_kernel(
+    gate_ptr, up_ptr, out_ptr,
+    stride_gh, stride_gw,
+    stride_uh, stride_uw,
+    stride_oh, stride_ow,
+    N: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < N
 
-        silu_gate = acc_gate * tl.sigmoid(acc_gate)
-        result = (silu_gate * acc_up).to(tl.bfloat16)
+    g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32)
 
-        out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
-        out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
-        tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
+    silu_g = g * tl.sigmoid(g)
+    result = (silu_g * u).to(tl.bfloat16)
+
+    tl.store(out_ptr + offs, result, mask=mask)
 
 
 class Model(nn.Module):
@@ -113,24 +123,47 @@
         I = self.I  # noqa: E741
         E = self.E
 
-        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
+        gate_out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
+        up_out   = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)


... (1809 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -1,12 +1,7 @@
 """Grouped GEMM + fused SwiGLU for MoE up-projection.
 
-Two separate Triton kernels:
-  1. Grouped GEMM for gate: gate_out = x @ W_gate (per expert)
-  2. Grouped GEMM for up:   up_out   = x @ W_up   (per expert)
-  3. Fused SwiGLU epilogue: out = silu(gate_out) * up_out
-
-This reduces register pressure per GEMM (only one accumulator) and allows
-the scheduler to overlap across the two GEMM launches.
+Triton kernel with fused gate+up GEMMs and SwiGLU epilogue.
+Grid: (num_m_tiles, num_n_tiles, E). Zero GPU->CPU syncs.
 """
 from __future__ import annotations
 
@@ -21,14 +16,16 @@
 
 
 @triton.jit
-def _grouped_gemm_kernel(
-    A_ptr,       # (T_perm, H) bf16
-    B_ptr,       # (E, H, I) bf16
-    C_ptr,       # (T_perm, I) bf16
-    offsets_ptr, # (E+1,) int32
-    stride_ah, stride_aw,
-    stride_be, stride_bh, stride_bi,
-    stride_ch, stride_cw,
+def _grouped_swiglu_kernel(
+    hidden_states_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    out_ptr,
+    expert_offsets_ptr,
+    stride_xh, stride_xw,
+    stride_wg_e, stride_wg_h, stride_wg_i,
+    stride_wu_e, stride_wu_h, stride_wu_i,
+    stride_oh, stride_ow,
     H: tl.constexpr,
     I: tl.constexpr,
     BM: tl.constexpr,
@@ -39,8 +36,8 @@
     pid_n = tl.program_id(1)
     expert_id = tl.program_id(2)
 
-    start = tl.load(offsets_ptr + expert_id)
-    end   = tl.load(offsets_ptr + expert_id + 1)
+    start = tl.load(expert_offsets_ptr + expert_id)
+    end   = tl.load(expert_offsets_ptr + expert_id + 1)
     n_rows = end - start
 
     row_off = pid_m * BM
@@ -49,7 +46,8 @@
 
     col_off = pid_n * BN
 
-    acc = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_gate = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_up   = tl.zeros((BM, BN), dtype=tl.float32)
 
     offs_m = tl.arange(0, BM)
     offs_n = tl.arange(0, BN)
@@ -57,48 +55,40 @@
     m_mask = (row_off + offs_m) < n_rows
     n_mask = (col_off + offs_n) < I
 
-    a_base = A_ptr + (start + row_off) * stride_ah
-    b_base = B_ptr + expert_id * stride_be + col_off * stride_bi
+    x_base  = hidden_states_ptr + (start + row_off) * stride_xh
+    wg_base = W_gate_ptr + expert_id * stride_wg_e + col_off * stride_wg_i
+    wu_base = W_up_ptr   + expert_id * stride_wu_e + col_off * stride_wu_i
 
     for k in range(0, tl.cdiv(H, BK)):
         k_off = k * BK
         k_mask = (k_off + offs_k) < H
+        mk_mask = m_mask[:, None] & k_mask[None, :]
+        kn_mask = k_mask[:, None] & n_mask[None, :]
 
-        a_ptrs = a_base + offs_m[:, None] * stride_ah + (k_off + offs_k)[None, :] * stride_aw
-        a_tile = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0).to(tl.bfloat16)
+        x_ptrs  = x_base  + offs_m[:, None] * stride_xh + (k_off + offs_k)[None, :] * stride_xw
+        x_tile  = tl.load(x_ptrs, mask=mk_mask, other=0.0).to(tl.bfloat16)
 
-        b_ptrs = b_base + (k_off + offs_k)[:, None] * stride_bh + offs_n[None, :] * stride_bi
-        b_tile = tl.load(b_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0).to(tl.bfloat16)
+        wg_ptrs = wg_base + (k_off + offs_k)[:, None] * stride_wg_h + offs_n[None, :] * stride_wg_i
+        wg_tile = tl.load(wg_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
-        acc = tl.dot(a_tile, b_tile, acc)
+        wu_ptrs = wu_base + (k_off + offs_k)[:, None] * stride_wu_h + offs_n[None, :] * stride_wu_i
+        wu_tile = tl.load(wu_ptrs, mask=kn_mask, other=0.0).to(tl.bfloat16)
 
-    result = acc.to(tl.bfloat16)
+        acc_gate = tl.dot(x_tile, wg_tile, acc_gate)
+        acc_up   = tl.dot(x_tile, wu_tile, acc_up)
 
-    c_base = C_ptr + (start + row_off) * stride_ch + col_off * stride_cw
-    c_ptrs = c_base + offs_m[:, None] * stride_ch + offs_n[None, :] * stride_cw
-    tl.store(c_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
+    silu_gate = acc_gate * tl.sigmoid(acc_gate)
+    result = (silu_gate * acc_up).to(tl.bfloat16)
+
+    out_base = out_ptr + (start + row_off) * stride_oh + col_off * stride_ow
+    out_ptrs = out_base + offs_m[:, None] * stride_oh + offs_n[None, :] * stride_ow
+    tl.store(out_ptrs, result, mask=m_mask[:, None] & n_mask[None, :])
 
 
-@triton.jit
-def _swiglu_epilogue_kernel(
-    gate_ptr, up_ptr, out_ptr,
-    stride_gh, stride_gw,
-    stride_uh, stride_uw,
-    stride_oh, stride_ow,
-    N: tl.constexpr,
-    BLOCK: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    offs = pid * BLOCK + tl.arange(0, BLOCK)
-    mask = offs < N
-
-    g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32)
-    u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32)
-
-    silu_g = g * tl.sigmoid(g)
-    result = (silu_g * u).to(tl.bfloat16)
-
-    tl.store(out_ptr + offs, result, mask=mask)
+# Optimal tile config discovered via benchmarking: BM=64, BN=64, BK=32, 4 warps, 2 stages
+_BM = 64
+_BN = 64
+_BK = 32
 
 
 class Model(nn.Module):
@@ -123,47 +113,21 @@
         I = self.I  # noqa: E741
         E = self.E
 
-        gate_out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
-        up_out   = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
-        out      = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
-
-        BM = 64
-        BN = 64
-        BK = 32
+        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
 
         max_rows = (T_perm + E - 1) // E
-        num_m = triton.cdiv(max_rows, BM)
-        num_n = triton.cdiv(I, BN)
+        num_m = triton.cdiv(max_rows, _BM)
+        num_n = triton.cdiv(I, _BN)
 
         grid = (num_m, num_n, E)
-
-        _grouped_gemm_kernel[grid](
-            hidden_states, self.W_gate, gate_out, expert_offsets,
+        _grouped_swiglu_kernel[grid](
+            hidden_states, self.W_gate, self.W_up, out, expert_o

... (1194 more chars)