--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -0,0 +1,190 @@
+"""Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
+
+Per-expert dense GEMM with SwiGLU fused in the epilogue. Loops over experts
+in Python (each expert is a dense tile); the Triton kernel computes both
+x @ W_gate and x @ W_up simultaneously and applies silu(gate) * up.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Triton kernel: fused dual GEMM + SwiGLU
+# ---------------------------------------------------------------------------
+# C[m,n] = silu(A[m,k] @ W_gate[k,n]) * (A[m,k] @ W_up[k,n])
+#
+# A is (M, K), W_gate is (K, N), W_up is (K, N), out is (M, N).
+# ---------------------------------------------------------------------------
+@triton.jit
+def _grouped_gemm_swiglu_kernel(
+    A_ptr,          # base ptr to (n_e, H)
+    W_gate_ptr,     # base ptr to (H, I) for this expert
+    W_up_ptr,       # base ptr to (H, I) for this expert
+    Out_ptr,        # base ptr to (n_e, I) for this expert
+    M: tl.constexpr,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    m_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    n_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    m_mask = m_offsets < M
+    n_mask = n_offsets < N
+
+    acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, K, BLOCK_K):
+        k_offsets = k + tl.arange(0, BLOCK_K)
+        k_mask = k_offsets < K
+
+        # Load A block: (BLOCK_M, BLOCK_K), col-major read row-major store
+        # A is row-major with stride K
+        a_ptrs = A_ptr + m_offsets[:, None] * K + k_offsets[None, :]
+        a = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
+
+        # Load W_gate block: (BLOCK_K, BLOCK_N), row-major stride N
+        wg_ptrs = W_gate_ptr + k_offsets[:, None] * N + n_offsets[None, :]
+        wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+
+        # Load W_up block: (BLOCK_K, BLOCK_N), row-major stride N
+        wu_ptrs = W_up_ptr + k_offsets[:, None] * N + n_offsets[None, :]
+        wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+
+        acc_gate = tl.dot(a, wg, acc_gate, out_dtype=tl.float32)
+        acc_up = tl.dot(a, wu, acc_up, out_dtype=tl.float32)
+
+    # SwiGLU: silu(gate) * up
+    # silu(x) = x * sigmoid(x)
+    acc_gate = acc_gate.astype(tl.bfloat16)
+    acc_gate_fp32 = acc_gate.to(tl.float32)
+    sigmoid = 1.0 / (1.0 + tl.exp(-acc_gate_fp32))
+    silu = acc_gate.to(tl.float32) * sigmoid
+    out_fp32 = silu * acc_up
+
+    out_ptr = Out_ptr + m_offsets[:, None] * N + n_offsets[None, :]
+    out = out_fp32.to(tl.bfloat16)
+    tl.store(out_ptr, out, mask=m_mask[:, None] & n_mask[None, :])
+
+
+def grouped_gemm_swiglu_launch(x, w_gate, w_up, n_tokens, out):
+    """Launch fused dual GEMM + SwiGLU for one expert.
+
+    x:        (n_tokens, H) bf16
+    w_gate:   (H, I) bf16
+    w_up:     (H, I) bf16
+    n_tokens: int
+    out:      (n_tokens, I) bf16
+    """
+    H = x.shape[1]
+    I = w_gate.shape[1]
+
+    # Heuristic block sizes
+    # SM120 has many SMs, try larger tiles for larger matrices
+    if n_tokens >= 1024 and I >= 1024:
+        BLOCK_M = 256
+        BLOCK_N = 256
+        BLOCK_K = 128
+    elif n_tokens >= 512:
+        BLOCK_M = 128
+        BLOCK_N = 128
+        BLOCK_K = 128
+    elif n_tokens >= 128:
+        BLOCK_M = 128
+        BLOCK_N = 64
+        BLOCK_K = 64
+    else:
+        BLOCK_M = 64
+        BLOCK_N = 64
+        BLOCK_K = 64
+
+    grid = (
+        triton.cdiv(n_tokens, BLOCK_M),
+        triton.cdiv(I, BLOCK_N),
+    )
+
+    _grouped_gemm_swiglu_kernel[grid](
+        x, w_gate, w_up, out,
+        M=n_tokens,
+        K=H,
+        N=I,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+
+
+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
+        out = torch.empty(T_perm, self.I, dtype=torch.bfloat16, device=hidden_states.device)
+
+        for e in range(self.E):
+            start = int(expert_offsets[e].item())
+            end = int(expert_offsets[e + 1].item())
+            if end == start:
+                continue
+            n_e = end - start
+            x_e = hidden_states[start:end].contiguous()
+            w_gate_e = self.W_gate[e].contiguous()
+            w_up_e = self.W_up[e].contiguous()
+            out_e = out[start:end]
+            grouped_gemm_swiglu_launch(x_e, w_gate_e, w_up_e, n_e, out_e)
+
+        return out
+
+
+# Module-level shape shims rewritten by check.py / benchmark.py per shape.
+T_total = 32768
+H 

... (368 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
@@ -65,11 +65,9 @@
         acc_up = tl.dot(a, wu, acc_up, out_dtype=tl.float32)
 
     # SwiGLU: silu(gate) * up
-    # silu(x) = x * sigmoid(x)
-    acc_gate = acc_gate.astype(tl.bfloat16)
-    acc_gate_fp32 = acc_gate.to(tl.float32)
-    sigmoid = 1.0 / (1.0 + tl.exp(-acc_gate_fp32))
-    silu = acc_gate.to(tl.float32) * sigmoid
+    # silu(x) = x * sigmoid(x); everything in fp32, cast to bf16 on store
+    sigmoid = 1.0 / (1.0 + tl.exp(-acc_gate))
+    silu = acc_gate * sigmoid
     out_fp32 = silu * acc_up
 
     out_ptr = Out_ptr + m_offsets[:, None] * N + n_offsets[None, :]

--- 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
@@ -12,25 +12,27 @@
 import triton.language as tl
 
 
-# ---------------------------------------------------------------------------
-# Triton kernel: fused dual GEMM + SwiGLU
-# ---------------------------------------------------------------------------
-# C[m,n] = silu(A[m,k] @ W_gate[k,n]) * (A[m,k] @ W_up[k,n])
-#
-# A is (M, K), W_gate is (K, N), W_up is (K, N), out is (M, N).
-# ---------------------------------------------------------------------------
 @triton.jit
 def _grouped_gemm_swiglu_kernel(
     A_ptr,          # base ptr to (n_e, H)
     W_gate_ptr,     # base ptr to (H, I) for this expert
     W_up_ptr,       # base ptr to (H, I) for this expert
     Out_ptr,        # base ptr to (n_e, I) for this expert
-    M: tl.constexpr,
-    K: tl.constexpr,
-    N: tl.constexpr,
+    M,
+    K,
+    N,
+    stride_am,
+    stride_ak,
+    stride_wg_k,
+    stride_wg_n,
+    stride_wu_k,
+    stride_wu_n,
+    stride_om,
+    stride_on,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
@@ -44,67 +46,63 @@
     acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    for k in range(0, K, BLOCK_K):
-        k_offsets = k + tl.arange(0, BLOCK_K)
-        k_mask = k_offsets < K
+    for k_start in range(0, K, BLOCK_K * NUM_STAGES):
+        for stage in range(NUM_STAGES):
+            k = k_start + stage * BLOCK_K
+            if k >= K:
+                break
+            k_offsets = k + tl.arange(0, BLOCK_K)
+            k_mask = k_offsets < K
 
-        # Load A block: (BLOCK_M, BLOCK_K), col-major read row-major store
-        # A is row-major with stride K
-        a_ptrs = A_ptr + m_offsets[:, None] * K + k_offsets[None, :]
-        a = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
+            a_ptrs = A_ptr + m_offsets[:, None] * stride_am + k_offsets[None, :] * stride_ak
+            a = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
 
-        # Load W_gate block: (BLOCK_K, BLOCK_N), row-major stride N
-        wg_ptrs = W_gate_ptr + k_offsets[:, None] * N + n_offsets[None, :]
-        wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+            wg_ptrs = W_gate_ptr + k_offsets[:, None] * stride_wg_k + n_offsets[None, :] * stride_wg_n
+            wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
 
-        # Load W_up block: (BLOCK_K, BLOCK_N), row-major stride N
-        wu_ptrs = W_up_ptr + k_offsets[:, None] * N + n_offsets[None, :]
-        wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+            wu_ptrs = W_up_ptr + k_offsets[:, None] * stride_wu_k + n_offsets[None, :] * stride_wu_n
+            wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
 
-        acc_gate = tl.dot(a, wg, acc_gate, out_dtype=tl.float32)
-        acc_up = tl.dot(a, wu, acc_up, out_dtype=tl.float32)
+            acc_gate = tl.dot(a, wg, acc_gate, out_dtype=tl.float32)
+            acc_up = tl.dot(a, wu, acc_up, out_dtype=tl.float32)
 
     # SwiGLU: silu(gate) * up
-    # silu(x) = x * sigmoid(x); everything in fp32, cast to bf16 on store
     sigmoid = 1.0 / (1.0 + tl.exp(-acc_gate))
     silu = acc_gate * sigmoid
     out_fp32 = silu * acc_up
 
-    out_ptr = Out_ptr + m_offsets[:, None] * N + n_offsets[None, :]
+    out_ptrs = Out_ptr + m_offsets[:, None] * stride_om + n_offsets[None, :] * stride_on
     out = out_fp32.to(tl.bfloat16)
-    tl.store(out_ptr, out, mask=m_mask[:, None] & n_mask[None, :])
+    tl.store(out_ptrs, out, mask=m_mask[:, None] & n_mask[None, :])
 
 
 def grouped_gemm_swiglu_launch(x, w_gate, w_up, n_tokens, out):
-    """Launch fused dual GEMM + SwiGLU for one expert.
-
-    x:        (n_tokens, H) bf16
-    w_gate:   (H, I) bf16
-    w_up:     (H, I) bf16
-    n_tokens: int
-    out:      (n_tokens, I) bf16
-    """
+    """Launch fused dual GEMM + SwiGLU for one expert."""
     H = x.shape[1]
     I = w_gate.shape[1]
 
-    # Heuristic block sizes
-    # SM120 has many SMs, try larger tiles for larger matrices
-    if n_tokens >= 1024 and I >= 1024:
-        BLOCK_M = 256
-        BLOCK_N = 256
-        BLOCK_K = 128
-    elif n_tokens >= 512:
+    # Constrain block sizes to fit in shared memory on SM120 (~101KB limit)
+    # Each tile at bf16: 2 bytes per element
+    # Shared memory needed per stage: A(BLOCK_M, BLOCK_K) + W_gate(BLOCK_K, BLOCK_N) + W_up(BLOCK_K, BLOCK_N)
+    # For 128x128x64: 2*(128*64 + 128*128 + 128*128) = 2*(8192 + 16384 + 16384) = 81664 bytes
+    # Still large with num_stages>1. Let's use 64x64x64 or 128x64x64.
+    # Actually, Triton doesn't necessarily load all stages at once for non-pipelined kernels.
+    # Let's use moderate sizes.
+    if n_tokens >= 512 and I >= 512:
         BLOCK_M = 128
         BLOCK_N = 128
-        BLOCK_K = 128
-    elif n_tokens >= 128:
+        BLOCK_K = 64
+        NUM_STAGES = 2
+    elif n_tokens >= 256 and I >= 256:
         BLOCK_M = 128
         BLOCK_N = 64
         BLOCK_K = 64
+        NUM_STAGES = 2
     else:
         BLOCK_M = 64
         BLOCK_N = 64
         BLOCK_K = 64
+        NUM_STAGES = 2
 
     grid = (
         triton.cdiv(n_tokens, BLOCK_M),
@@ -116,22 +114,23 @@
         M=n_tokens,
         K=H,
         N=I,
+        stride_am=x.stride(0),
+        stride_ak=x.stride(1),
+        stride_wg_k=w_gate.stride(0),
+        stride_wg_n=w_gate.stride(1),
+        stride_wu_k=w_up.stride(0),
+        stride_wu_n=w_up.stride(1),
+        stride_om=out.stride(0),
+        stride_on=out.stride(1),
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
         BLOCK_K=BLOCK_K,
+        NUM_STAGES=NUM_

... (1432 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,9 +1,4 @@
-"""Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
-
-Per-expert dense GEMM with SwiGLU fused in the epilogue. Loops over experts
-in Python (each expert is a dense tile); the Triton kernel computes both
-x @ W_gate and x @ W_up simultaneously and applies silu(gate) * up.
-"""
+"""Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection."""
 from __future__ import annotations
 
 import torch
@@ -14,158 +9,94 @@
 
 @triton.jit
 def _grouped_gemm_swiglu_kernel(
-    A_ptr,          # base ptr to (n_e, H)
-    W_gate_ptr,     # base ptr to (H, I) for this expert
-    W_up_ptr,       # base ptr to (H, I) for this expert
-    Out_ptr,        # base ptr to (n_e, I) for this expert
+    A_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    Out_ptr,
     M,
     K,
     N,
-    stride_am,
-    stride_ak,
-    stride_wg_k,
-    stride_wg_n,
-    stride_wu_k,
-    stride_wu_n,
-    stride_om,
-    stride_on,
+    stride_a_row,
+    stride_a_col,
+    stride_w_row,
+    stride_w_col,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
-    NUM_STAGES: tl.constexpr,
 ):
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
-    m_offsets = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    n_offsets = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    m_offs = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    k_offs = tl.arange(0, BLOCK_K)
 
-    m_mask = m_offsets < M
-    n_mask = n_offsets < N
+    a_base = A_ptr + m_offs[:, None] * stride_a_row
+    wg_base = W_gate_ptr + n_offs[None, :] * stride_w_col
+    wu_base = W_up_ptr + n_offs[None, :] * stride_w_col
 
     acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    for k_start in range(0, K, BLOCK_K * NUM_STAGES):
-        for stage in range(NUM_STAGES):
-            k = k_start + stage * BLOCK_K
-            if k >= K:
-                break
-            k_offsets = k + tl.arange(0, BLOCK_K)
-            k_mask = k_offsets < K
+    for _ in range(0, K, BLOCK_K):
+        a_ptrs = a_base + k_offs[None, :] * stride_a_col
+        wg_ptrs = wg_base + k_offs[:, None] * stride_w_row
+        wu_ptrs = wu_base + k_offs[:, None] * stride_w_row
 
-            a_ptrs = A_ptr + m_offsets[:, None] * stride_am + k_offsets[None, :] * stride_ak
-            a = tl.load(a_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
+        a = tl.load(a_ptrs, mask=(m_offs[:, None] < M) & (k_offs[None, :] < K), other=0.0)
+        wg = tl.load(wg_ptrs, mask=(k_offs[:, None] < K) & (n_offs[None, :] < N), other=0.0)
+        wu = tl.load(wu_ptrs, mask=(k_offs[:, None] < K) & (n_offs[None, :] < N), other=0.0)
 
-            wg_ptrs = W_gate_ptr + k_offsets[:, None] * stride_wg_k + n_offsets[None, :] * stride_wg_n
-            wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+        acc_gate += tl.dot(a, wg)
+        acc_up += tl.dot(a, wu)
 
-            wu_ptrs = W_up_ptr + k_offsets[:, None] * stride_wu_k + n_offsets[None, :] * stride_wu_n
-            wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0)
+        a_base += BLOCK_K * stride_a_col
+        wg_base += BLOCK_K * stride_w_row
+        wu_base += BLOCK_K * stride_w_row
 
-            acc_gate = tl.dot(a, wg, acc_gate, out_dtype=tl.float32)
-            acc_up = tl.dot(a, wu, acc_up, out_dtype=tl.float32)
+    # SwiGLU
+    silu = acc_gate * tl.sigmoid(acc_gate)
+    out_fp = silu * acc_up
+    out_bf = out_fp.to(tl.bfloat16)
 
-    # SwiGLU: silu(gate) * up
-    sigmoid = 1.0 / (1.0 + tl.exp(-acc_gate))
-    silu = acc_gate * sigmoid
-    out_fp32 = silu * acc_up
-
-    out_ptrs = Out_ptr + m_offsets[:, None] * stride_om + n_offsets[None, :] * stride_on
-    out = out_fp32.to(tl.bfloat16)
-    tl.store(out_ptrs, out, mask=m_mask[:, None] & n_mask[None, :])
+    out_ptrs = Out_ptr + m_offs[:, None] * N + n_offs[None, :]
+    tl.store(out_ptrs, out_bf, mask=(m_offs[:, None] < M) & (n_offs[None, :] < N))
 
 
-def grouped_gemm_swiglu_launch(x, w_gate, w_up, n_tokens, out):
-    """Launch fused dual GEMM + SwiGLU for one expert."""
-    H = x.shape[1]
-    I = w_gate.shape[1]
-
-    # Constrain block sizes to fit in shared memory on SM120 (~101KB limit)
-    # Each tile at bf16: 2 bytes per element
-    # Shared memory needed per stage: A(BLOCK_M, BLOCK_K) + W_gate(BLOCK_K, BLOCK_N) + W_up(BLOCK_K, BLOCK_N)
-    # For 128x128x64: 2*(128*64 + 128*128 + 128*128) = 2*(8192 + 16384 + 16384) = 81664 bytes
-    # Still large with num_stages>1. Let's use 64x64x64 or 128x64x64.
-    # Actually, Triton doesn't necessarily load all stages at once for non-pipelined kernels.
-    # Let's use moderate sizes.
-    if n_tokens >= 512 and I >= 512:
-        BLOCK_M = 128
-        BLOCK_N = 128
-        BLOCK_K = 64
-        NUM_STAGES = 2
-    elif n_tokens >= 256 and I >= 256:
-        BLOCK_M = 128
-        BLOCK_N = 64
-        BLOCK_K = 64
-        NUM_STAGES = 2
-    else:
-        BLOCK_M = 64
-        BLOCK_N = 64
-        BLOCK_K = 64
-        NUM_STAGES = 2
-
-    grid = (
-        triton.cdiv(n_tokens, BLOCK_M),
-        triton.cdiv(I, BLOCK_N),
-    )
-
+def _launch(x, w_gate_e, w_up_e, n_e, out_e, bm=128, bn=128, bk=128):
+    grid = (triton.cdiv(n_e, bm), triton.cdiv(out_e.shape[1], bn))
     _grouped_gemm_swiglu_kernel[grid](
-        x, w_gate, w_up, out,
-        M=n_tokens,
-        K=H,
-        N=I,
-        stride_am=x.stride(0),
-        stride_ak=x.stride(1),
-        stride_wg_k=w_gate.stride(0),
-        stride_wg_n=w_gate.stride(1),
-        stride_wu_k=w_up.stride(0),
-        stride_wu_n=w_up.stride(1),
-        stride_om=out.stride(0),
-        stride_on=out.stride(1),
-        BLOCK_M=BLOCK_M,
-        BLOC

... (2368 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
@@ -63,14 +63,23 @@
     tl.store(out_ptrs, out_bf, mask=(m_offs[:, None] < M) & (n_offs[None, :] < N))
 
 
-def _launch(x, w_gate_e, w_up_e, n_e, out_e, bm=128, bn=128, bk=128):
-    grid = (triton.cdiv(n_e, bm), triton.cdiv(out_e.shape[1], bn))
+def _launch(x, w_gate_e, w_up_e, n_e, out_e):
+    I = out_e.shape[1]
+    # 64x64x128 fits in shared memory (num_stages=2 -> 98KB, just under 101KB limit)
+    if n_e >= 512 and I >= 512:
+        bm, bn, bk = 64, 128, 128
+    elif n_e >= 128:
+        bm, bn, bk = 64, 64, 128
+    else:
+        bm, bn, bk = 32, 32, 64
+    grid = (triton.cdiv(n_e, bm), triton.cdiv(I, bn))
     _grouped_gemm_swiglu_kernel[grid](
         x, w_gate_e, w_up_e, out_e,
-        M=n_e, K=x.shape[1], N=out_e.shape[1],
+        M=n_e, K=x.shape[1], N=I,
         stride_a_row=x.stride(0), stride_a_col=x.stride(1),
         stride_w_row=w_gate_e.stride(0), stride_w_col=w_gate_e.stride(1),
         BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk,
+        num_stages=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
@@ -1,4 +1,9 @@
-"""Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection."""
+"""Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
+
+Single kernel handles all experts in one launch to avoid Python loop overhead.
+Grid is (E, cdiv(M_expert, BLOCK_M), cdiv(N, BLOCK_N)) where M_expert varies
+per expert (computed from expert_offsets).
+"""
 from __future__ import annotations
 
 import torch
@@ -8,104 +13,146 @@
 
 
 @triton.jit
-def _grouped_gemm_swiglu_kernel(
-    A_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    Out_ptr,
-    M,
-    K,
-    N,
-    stride_a_row,
-    stride_a_col,
-    stride_w_row,
-    stride_w_col,
-    BLOCK_M: tl.constexpr,
+def _grouped_gemm_swiglu_single_kernel(
+    X_ptr,           # (T_perm, H), row-major
+    W_gate_ptr,      # (E, H, I), row-major
+    W_up_ptr,        # (E, H, I), row-major
+    Out_ptr,         # (T_perm, I), row-major
+    Offsets_ptr,     # (E+1,) int32
+    T_perm,
+    H,
+    I,
+    E,
+    stride_x_row,
+    stride_w_H,
+    stride_w_I,
+    stride_o_row,
+    BLOCK_H: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
-    pid_m = tl.program_id(0)
-    pid_n = tl.program_id(1)
+    # Grid: (E, cdiv(max_tokens, BLOCK_M), cdiv(I, BLOCK_N))
+    pid_e = tl.program_id(0)
+    pid_bm = tl.program_id(1)
+    pid_n = tl.program_id(2)
 
-    m_offs = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    # Read expert offset bounds
+    off_start = tl.load(Offsets_ptr + pid_e)
+    off_end = tl.load(Offsets_ptr + pid_e + 1)
+    n_e = off_end - off_start
+
+    # M dimension for this tile
+    m_base = pid_bm * 256  # We fix BLOCK_M = 256 for simplicity
+    m_offs = m_base + tl.arange(0, 256)
+
+    # Global row indices within hidden_states
+    row_global = off_start + m_offs
+    m_valid = m_offs < n_e
+
+    # N (output) dimension
     n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    k_offs = tl.arange(0, BLOCK_K)
+    n_valid = n_offs < I
 
-    a_base = A_ptr + m_offs[:, None] * stride_a_row
-    wg_base = W_gate_ptr + n_offs[None, :] * stride_w_col
-    wu_base = W_up_ptr + n_offs[None, :] * stride_w_col
+    # K (hidden) dimension
+    k_offs = tl.arange(0, BLOCK_H)
+    k_valid = k_offs < H
 
-    acc_gate = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-    acc_up = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    # Pointers
+    # X[row, k] = X_ptr + row * stride_x_row + k
+    # W_gate[e, k, n] = W_gate_ptr + e*stride_w_H*H + k*stride_w_I + n
 
-    for _ in range(0, K, BLOCK_K):
-        a_ptrs = a_base + k_offs[None, :] * stride_a_col
-        wg_ptrs = wg_base + k_offs[:, None] * stride_w_row
-        wu_ptrs = wu_base + k_offs[:, None] * stride_w_row
+    acc_gate = tl.zeros((256, BLOCK_N), dtype=tl.float32)
+    acc_up = tl.zeros((256, BLOCK_N), dtype=tl.float32)
 
-        a = tl.load(a_ptrs, mask=(m_offs[:, None] < M) & (k_offs[None, :] < K), other=0.0)
-        wg = tl.load(wg_ptrs, mask=(k_offs[:, None] < K) & (n_offs[None, :] < N), other=0.0)
-        wu = tl.load(wu_ptrs, mask=(k_offs[:, None] < K) & (n_offs[None, :] < N), other=0.0)
+    # Load base W pointers (add expert offset once)
+    w_e_start = pid_e * stride_w_H * H  # offset for this expert's weights
+    wg_base = W_gate_ptr + w_e_start
+    wu_base = W_up_ptr + w_e_start
 
-        acc_gate += tl.dot(a, wg)
-        acc_up += tl.dot(a, wu)
+    for h_start in range(0, H, BLOCK_H):
+        h_offset = h_start + k_offs
+        h_mask = h_offset < H
 
-        a_base += BLOCK_K * stride_a_col
-        wg_base += BLOCK_K * stride_w_row
-        wu_base += BLOCK_K * stride_w_row
+        # Load X block: (256, BLOCK_H)
+        x_ptrs = X_ptr + row_global[:, None] * stride_x_row + h_offset[None, :]
+        x = tl.load(x_ptrs, mask=m_valid[:, None] & h_mask[None, :], other=0.0)
+
+        # Load W_gate block: (BLOCK_H, BLOCK_N) for expert e
+        wg_ptrs = wg_base + h_offset[:, None] * stride_w_I + n_offs[None, :]
+        wg = tl.load(wg_ptrs, mask=h_mask[:, None] & n_valid[None, :], other=0.0)
+
+        # Load W_up block: (BLOCK_H, BLOCK_N) for expert e
+        wu_ptrs = wu_base + h_offset[:, None] * stride_w_I + n_offs[None, :]
+        wu = tl.load(wu_ptrs, mask=h_mask[:, None] & n_valid[None, :], other=0.0)
+
+        acc_gate += tl.dot(x, wg)
+        acc_up += tl.dot(x, wu)
 
     # SwiGLU
-    silu = acc_gate * tl.sigmoid(acc_gate)
-    out_fp = silu * acc_up
-    out_bf = out_fp.to(tl.bfloat16)
+    gate_fp32 = acc_gate
+    silu = gate_fp32 * tl.sigmoid(gate_fp32)
+    out_fp32 = silu * acc_up
 
-    out_ptrs = Out_ptr + m_offs[:, None] * N + n_offs[None, :]
-    tl.store(out_ptrs, out_bf, mask=(m_offs[:, None] < M) & (n_offs[None, :] < N))
-
-
-def _launch(x, w_gate_e, w_up_e, n_e, out_e):
-    I = out_e.shape[1]
-    # 64x64x128 fits in shared memory (num_stages=2 -> 98KB, just under 101KB limit)
-    if n_e >= 512 and I >= 512:
-        bm, bn, bk = 64, 128, 128
-    elif n_e >= 128:
-        bm, bn, bk = 64, 64, 128
-    else:
-        bm, bn, bk = 32, 32, 64
-    grid = (triton.cdiv(n_e, bm), triton.cdiv(I, bn))
-    _grouped_gemm_swiglu_kernel[grid](
-        x, w_gate_e, w_up_e, out_e,
-        M=n_e, K=x.shape[1], N=I,
-        stride_a_row=x.stride(0), stride_a_col=x.stride(1),
-        stride_w_row=w_gate_e.stride(0), stride_w_col=w_gate_e.stride(1),
-        BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk,
-        num_stages=1,
-    )
+    # Store
+    out_ptrs = Out_ptr + row_global[:, None] * stride_o_row + n_offs[None, :]
+    out = out_fp32.to(tl.bfloat16)
+    tl.store(out_ptrs, out, mask=m_valid[:, None] & n_valid[None, :])
 
 
 class Model(nn.Module):
     def __init__(self, T_total, H, I, E, K):  # noqa: E741
         super().__init__()
+        self.T_total = T_total
+        self.H = H
         self.I = I
+        self.E = E
+   

... (2141 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,8 +1,9 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-Single kernel handles all experts in one launch to avoid Python loop overhead.
-Grid is (E, cdiv(M_expert, BLOCK_M), cdiv(N, BLOCK_N)) where M_expert varies
-per expert (computed from expert_offsets).
+Single kernel handles all experts in one launch. The 3D grid is
+(E, cdiv(M_expert, BM), cdiv(I, BN)), but we use a 2D grid
+(E * cdiv(total_rows/avg, BM), cdiv(I, BN)) and compute M from offsets.
+We keep tiles small enough for shared memory.
 """
 from __future__ import annotations
 
@@ -13,100 +14,131 @@
 
 
 @triton.jit
-def _grouped_gemm_swiglu_single_kernel(
-    X_ptr,           # (T_perm, H), row-major
-    W_gate_ptr,      # (E, H, I), row-major
-    W_up_ptr,        # (E, H, I), row-major
-    Out_ptr,         # (T_perm, I), row-major
-    Offsets_ptr,     # (E+1,) int32
-    T_perm,
+def _grouped_gemm_swiglu_kernel(
+    X_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    Out_ptr,
+    Offsets_ptr,
     H,
     I,
-    E,
     stride_x_row,
-    stride_w_H,
-    stride_w_I,
+    stride_w_row,
+    stride_w_col,
     stride_o_row,
-    BLOCK_H: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    BLOCK_K: tl.constexpr,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
 ):
-    # Grid: (E, cdiv(max_tokens, BLOCK_M), cdiv(I, BLOCK_N))
+    pid_mn = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    # We encode expert id in the upper bits of pid_mn for small E,
+    # OR we restructure as 3D grid: (E, BN, BK). Let me think differently.
+    # Actually the standard approach: grid (total_tiles_M, cdiv(I, BN)).
+    # We need to know which expert each tile belongs to and the intra-expert row.
+    pass
+
+
+# Better approach: use 3D grid (E, cdiv(max_M, BM), cdiv(I, BN)).
+# The kernel computes row = off_start[pid_e] + pid_bm * BM + local_m.
+# Shared memory: BM * BK + 2 * BK * BN bf16 values.
+# For BM=128, BN=128, BK=64: 128*64 + 2*64*128 = 8192 + 16384 = 24576 * 2 = 49152 (fits!)
+# For BM=64, BN=64, BK=128: 64*128 + 2*128*64 = 8192 + 16384 = 24576 * 2 = 49152 (fits!)
+
+@triton.jit
+def _gemma_gemm_kernel(
+    X_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    Out_ptr,
+    Offsets_ptr,
+    H,
+    I,
+    stride_x_row,
+    stride_w_row,
+    stride_w_col,
+    stride_o_row,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
     pid_e = tl.program_id(0)
     pid_bm = tl.program_id(1)
     pid_n = tl.program_id(2)
 
-    # Read expert offset bounds
-    off_start = tl.load(Offsets_ptr + pid_e)
-    off_end = tl.load(Offsets_ptr + pid_e + 1)
-    n_e = off_end - off_start
+    # Expert bounds
+    os = tl.load(Offsets_ptr + pid_e).to(tl.int64)
+    oe = tl.load(Offsets_ptr + pid_e + 1).to(tl.int64)
+    n_e = oe - os
 
-    # M dimension for this tile
-    m_base = pid_bm * 256  # We fix BLOCK_M = 256 for simplicity
-    m_offs = m_base + tl.arange(0, 256)
+    # Local M offset
+    m_local = pid_bm * BM + tl.arange(0, BM)
+    m_valid = m_local < n_e
 
-    # Global row indices within hidden_states
-    row_global = off_start + m_offs
-    m_valid = m_offs < n_e
+    # Global row
+    row_g = os + m_local
 
-    # N (output) dimension
-    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    n_valid = n_offs < I
+    # N offsets
+    n_local = pid_n * BN + tl.arange(0, BN)
+    n_valid = n_local < I
 
-    # K (hidden) dimension
-    k_offs = tl.arange(0, BLOCK_H)
-    k_valid = k_offs < H
+    # K offsets
+    k_local = tl.arange(0, BK)
 
-    # Pointers
-    # X[row, k] = X_ptr + row * stride_x_row + k
-    # W_gate[e, k, n] = W_gate_ptr + e*stride_w_H*H + k*stride_w_I + n
+    # Accumulators
+    acc_g = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_u = tl.zeros((BM, BN), dtype=tl.float32)
 
-    acc_gate = tl.zeros((256, BLOCK_N), dtype=tl.float32)
-    acc_up = tl.zeros((256, BLOCK_N), dtype=tl.float32)
+    # Weight base for this expert
+    wg_base = W_gate_ptr + pid_e * stride_w_row * H
+    wu_base = W_up_ptr + pid_e * stride_w_row * H
 
-    # Load base W pointers (add expert offset once)
-    w_e_start = pid_e * stride_w_H * H  # offset for this expert's weights
-    wg_base = W_gate_ptr + w_e_start
-    wu_base = W_up_ptr + w_e_start
+    for k_start in range(0, H, BK):
+        k_offs = k_start + k_local
+        k_mask = k_offs < H
 
-    for h_start in range(0, H, BLOCK_H):
-        h_offset = h_start + k_offs
-        h_mask = h_offset < H
+        # A[m, k]
+        a_ptrs = X_ptr + row_g[:, None] * stride_x_row + k_offs[None, :]
+        a = tl.load(a_ptrs, mask=m_valid[:, None] & k_mask[None, :], other=0.0)
 
-        # Load X block: (256, BLOCK_H)
-        x_ptrs = X_ptr + row_global[:, None] * stride_x_row + h_offset[None, :]
-        x = tl.load(x_ptrs, mask=m_valid[:, None] & h_mask[None, :], other=0.0)
+        # W_gate[e][k, n]
+        wg_ptrs = wg_base + k_offs[:, None] * stride_w_row + n_local[None, :] * stride_w_col
+        wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        # Load W_gate block: (BLOCK_H, BLOCK_N) for expert e
-        wg_ptrs = wg_base + h_offset[:, None] * stride_w_I + n_offs[None, :]
-        wg = tl.load(wg_ptrs, mask=h_mask[:, None] & n_valid[None, :], other=0.0)
+        # W_up[e][k, n]
+        wu_ptrs = wu_base + k_offs[:, None] * stride_w_row + n_local[None, :] * stride_w_col
+        wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        # Load W_up block: (BLOCK_H, BLOCK_N) for expert e
-        wu_ptrs = wu_base + h_offset[:, None] * stride_w_I + n_offs[None, :]
-        wu = tl.load(wu_ptrs, mask=h_mask[:, None] & n_valid[None, :], other=0.0)
-
-        acc_gate += tl.dot(x, wg)
-        acc_up += tl.dot(x, wu)
+        acc_g += tl.dot(a, wg)
+        acc_u += tl.dot(a, wu)
 
     #

... (3445 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,9 +1,7 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-Single kernel handles all experts in one launch. The 3D grid is
-(E, cdiv(M_expert, BM), cdiv(I, BN)), but we use a 2D grid
-(E * cdiv(total_rows/avg, BM), cdiv(I, BN)) and compute M from offsets.
-We keep tiles small enough for shared memory.
+Flattened grid approach: grid dimension 0 covers all row-blocks across all experts.
+A per-block expert-lookup uses binary search on offsets (or a cumulative tile-count array).
 """
 from __future__ import annotations
 
@@ -11,6 +9,23 @@
 import torch.nn as nn
 import triton
 import triton.language as tl
+
+
+# We use a simple approach: grid is (total_m_blocks, cdiv(I, BN), E)
+# Actually, let's keep it as (E, cdiv(n_tokens_for_expert_e, BM), cdiv(I, BN))
+# but compute the grid size from actual per-expert tile counts.
+
+# The issue with 3D grid: the grid size in dimension 1 is not uniform across experts.
+# Standard solution: flatten dimension 0+1 into a single dimension.
+#
+# For each expert e, num_tiles_e = cdiv(n_e, BM)
+# cumulative_tiles[e] = sum(num_tiles_e' for e' < e)
+# total_tiles = cumulative_tiles[E]
+#
+# Grid: (total_tiles, cdiv(I, BN))
+# Given pid_0 in range [cumulative_tiles[e], cumulative_tiles[e+1]),
+#   expert_id = e
+#   pid_bm = pid_0 - cumulative_tiles[e]
 
 
 @triton.jit
@@ -20,94 +35,91 @@
     W_up_ptr,
     Out_ptr,
     Offsets_ptr,
+    TileOffs_ptr,  # (E+1,) int32, cumulative start tile indices
     H,
     I,
+    BM,
+    BN,
     stride_x_row,
     stride_w_row,
     stride_w_col,
     stride_o_row,
-    BM: tl.constexpr,
-    BN: tl.constexpr,
-    BK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
 ):
-    pid_mn = tl.program_id(0)
+    pid_tile = tl.program_id(0)
     pid_n = tl.program_id(1)
 
-    # We encode expert id in the upper bits of pid_mn for small E,
-    # OR we restructure as 3D grid: (E, BN, BK). Let me think differently.
-    # Actually the standard approach: grid (total_tiles_M, cdiv(I, BN)).
-    # We need to know which expert each tile belongs to and the intra-expert row.
-    pass
-
-
-# Better approach: use 3D grid (E, cdiv(max_M, BM), cdiv(I, BN)).
-# The kernel computes row = off_start[pid_e] + pid_bm * BM + local_m.
-# Shared memory: BM * BK + 2 * BK * BN bf16 values.
-# For BM=128, BN=128, BK=64: 128*64 + 2*64*128 = 8192 + 16384 = 24576 * 2 = 49152 (fits!)
-# For BM=64, BN=64, BK=128: 64*128 + 2*128*64 = 8192 + 16384 = 24576 * 2 = 49152 (fits!)
-
-@triton.jit
-def _gemma_gemm_kernel(
-    X_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    Out_ptr,
-    Offsets_ptr,
-    H,
-    I,
-    stride_x_row,
-    stride_w_row,
-    stride_w_col,
-    stride_o_row,
-    BM: tl.constexpr,
-    BN: tl.constexpr,
-    BK: tl.constexpr,
-):
-    pid_e = tl.program_id(0)
-    pid_bm = tl.program_id(1)
-    pid_n = tl.program_id(2)
-
-    # Expert bounds
-    os = tl.load(Offsets_ptr + pid_e).to(tl.int64)
-    oe = tl.load(Offsets_ptr + pid_e + 1).to(tl.int64)
-    n_e = oe - os
-
-    # Local M offset
-    m_local = pid_bm * BM + tl.arange(0, BM)
+    # Find which expert this tile belongs to via binary search on TileOffs.
+    # TileOffs has E+1 entries; we need the largest e such that TileOffs[e] <= pid_tile.
+    # Since E is small (64 or 128), linear search is fine.
+    lo = 0
+    hi = Offsets_ptr.numel() - 2  # E-1, but we don't know E at compile time
+    # We'll just do a simple binary search
+    expert = 0
+    
+    # Binary search: find largest lo such that TileOffs[lo] <= pid_tile
+    left = 0
+    right = 128  # upper bound on E
+    while left < right:
+        mid = (left + right + 1) // 2
+        val = tl.load(TileOffs_ptr + mid).to(tl.int32)
+        if val <= pid_tile:
+            left = mid
+        else:
+            right = mid - 1
+    expert = left
+
+    tile_base = tl.load(TileOffs_ptr + expert).to(tl.int32)
+    pid_bm = pid_tile - tile_base
+
+    # Expert row bounds
+    row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
+    row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
+    n_e = row_end - row_start
+
+    # Local row offsets within expert
+    m_local = pid_bm * BLOCK_M + tl.arange(0, BLOCK_M)
     m_valid = m_local < n_e
 
-    # Global row
-    row_g = os + m_local
-
-    # N offsets
-    n_local = pid_n * BN + tl.arange(0, BN)
+    # Global rows
+    row_g = row_start + m_local
+
+    # Column offsets
+    n_local = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     n_valid = n_local < I
 
     # K offsets
-    k_local = tl.arange(0, BK)
+    k_local = tl.arange(0, BLOCK_K)
 
     # Accumulators
-    acc_g = tl.zeros((BM, BN), dtype=tl.float32)
-    acc_u = tl.zeros((BM, BN), dtype=tl.float32)
-
-    # Weight base for this expert
-    wg_base = W_gate_ptr + pid_e * stride_w_row * H
-    wu_base = W_up_ptr + pid_e * stride_w_row * H
-
-    for k_start in range(0, H, BK):
+    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    # Weight base for this expert: W is (E, H, I), row-major
+    # stride_w_row = W.stride(1) = I, stride_w_col = W.stride(2) = 1
+    # W[e, k, n] = W_ptr + e*H*I + k*I + n
+    wg_base = W_gate_ptr + expert * H * I.to(tl.int64)
+    wu_base = W_up_ptr + expert * H * I.to(tl.int64)
+
+    for k_start in range(0, H, BLOCK_K):
         k_offs = k_start + k_local
         k_mask = k_offs < H
 
-        # A[m, k]
-        a_ptrs = X_ptr + row_g[:, None] * stride_x_row + k_offs[None, :]
+        # A[m, k]: (BLOCK_M, BLOCK_K)
+        a_ptrs = X_ptr + row_g.to(tl.int64)[:, None] * stride_x_row + k_offs[None, :] * stride_x_row // BLOCK_M
+        # Actually, stride_x_row might not be H if the tensor isn't contiguous...
+        # Let me just use stride_x_row directly
+       

... (4233 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,7 +1,7 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-Flattened grid approach: grid dimension 0 covers all row-blocks across all experts.
-A per-block expert-lookup uses binary search on offsets (or a cumulative tile-count array).
+Flattened 2D grid: (cdiv(T_perm, BM), cdiv(I, BN)) with expert lookup
+via binary search on offsets.
 """
 from __future__ import annotations
 
@@ -11,23 +11,6 @@
 import triton.language as tl
 
 
-# We use a simple approach: grid is (total_m_blocks, cdiv(I, BN), E)
-# Actually, let's keep it as (E, cdiv(n_tokens_for_expert_e, BM), cdiv(I, BN))
-# but compute the grid size from actual per-expert tile counts.
-
-# The issue with 3D grid: the grid size in dimension 1 is not uniform across experts.
-# Standard solution: flatten dimension 0+1 into a single dimension.
-#
-# For each expert e, num_tiles_e = cdiv(n_e, BM)
-# cumulative_tiles[e] = sum(num_tiles_e' for e' < e)
-# total_tiles = cumulative_tiles[E]
-#
-# Grid: (total_tiles, cdiv(I, BN))
-# Given pid_0 in range [cumulative_tiles[e], cumulative_tiles[e+1]),
-#   expert_id = e
-#   pid_bm = pid_0 - cumulative_tiles[e]
-
-
 @triton.jit
 def _grouped_gemm_swiglu_kernel(
     X_ptr,
@@ -35,132 +18,97 @@
     W_up_ptr,
     Out_ptr,
     Offsets_ptr,
-    TileOffs_ptr,  # (E+1,) int32, cumulative start tile indices
+    T_perm,
     H,
-    I,
-    BM,
-    BN,
-    stride_x_row,
-    stride_w_row,
-    stride_w_col,
-    stride_o_row,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    BLOCK_K: tl.constexpr,
+    I_val,
+    E,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
 ):
-    pid_tile = tl.program_id(0)
+    pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
-    # Find which expert this tile belongs to via binary search on TileOffs.
-    # TileOffs has E+1 entries; we need the largest e such that TileOffs[e] <= pid_tile.
-    # Since E is small (64 or 128), linear search is fine.
+    # Find expert for the first row in this block
+    row_0 = pid_m * BM
+    # Binary search for expert: find smallest e such that Offsets[e+1] > row_0
     lo = 0
-    hi = Offsets_ptr.numel() - 2  # E-1, but we don't know E at compile time
-    # We'll just do a simple binary search
-    expert = 0
-    
-    # Binary search: find largest lo such that TileOffs[lo] <= pid_tile
-    left = 0
-    right = 128  # upper bound on E
-    while left < right:
-        mid = (left + right + 1) // 2
-        val = tl.load(TileOffs_ptr + mid).to(tl.int32)
-        if val <= pid_tile:
-            left = mid
+    hi = E - 1
+    while lo < hi:
+        mid = (lo + hi) // 2
+        off = tl.load(Offsets_ptr + mid + 1).to(tl.int64)
+        if off <= row_0:
+            lo = mid + 1
         else:
-            right = mid - 1
-    expert = left
+            hi = mid
+    expert = lo
 
-    tile_base = tl.load(TileOffs_ptr + expert).to(tl.int32)
-    pid_bm = pid_tile - tile_base
-
-    # Expert row bounds
     row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
     row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
-    n_e = row_end - row_start
 
-    # Local row offsets within expert
-    m_local = pid_bm * BLOCK_M + tl.arange(0, BLOCK_M)
-    m_valid = m_local < n_e
+    # M offset (intra-expert rows)
+    m_local = pid_m * BM + tl.arange(0, BM) - row_start
+    m_valid = (pid_m * BM + tl.arange(0, BM)) < row_end
 
     # Global rows
-    row_g = row_start + m_local
+    row_g = pid_m * BM + tl.arange(0, BM)
+    row_valid = row_g < T_perm & m_valid
 
-    # Column offsets
-    n_local = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    n_valid = n_local < I
+    # N offset
+    n_offs = pid_n * BN + tl.arange(0, BN)
+    n_valid = n_offs < I_val
 
-    # K offsets
-    k_local = tl.arange(0, BLOCK_K)
+    # K offset
+    k_offs = tl.arange(0, BK)
 
-    # Accumulators
-    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc_g = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_u = tl.zeros((BM, BN), dtype=tl.float32)
 
-    # Weight base for this expert: W is (E, H, I), row-major
-    # stride_w_row = W.stride(1) = I, stride_w_col = W.stride(2) = 1
-    # W[e, k, n] = W_ptr + e*H*I + k*I + n
-    wg_base = W_gate_ptr + expert * H * I.to(tl.int64)
-    wu_base = W_up_ptr + expert * H * I.to(tl.int64)
+    # Weight base: W is (E, H, I), row major
+    # W[e, k, n] offset = e * H * I + k * I + n
+    w_base = expert * H * I_val
 
-    for k_start in range(0, H, BLOCK_K):
-        k_offs = k_start + k_local
-        k_mask = k_offs < H
+    for k_start in range(0, H, BK):
+        k_mask = k_start + k_offs < H
+        k = k_start + k_offs
 
-        # A[m, k]: (BLOCK_M, BLOCK_K)
-        a_ptrs = X_ptr + row_g.to(tl.int64)[:, None] * stride_x_row + k_offs[None, :] * stride_x_row // BLOCK_M
-        # Actually, stride_x_row might not be H if the tensor isn't contiguous...
-        # Let me just use stride_x_row directly
-        a_ptrs = X_ptr + row_g.to(tl.int64)[:, None] * stride_x_row + k_offs[None, :]
-        a = tl.load(a_ptrs, mask=m_valid[:, None] & k_mask[None, :], other=0.0)
+        # A[row, k]
+        a_ptrs = X_ptr + row_g[:, None] * H + k[None, :]
+        a = tl.load(a_ptrs, mask=row_valid[:, None] & k_mask[None, :], other=0.0)
 
-        # W_gate: (BLOCK_K, BLOCK_N) -> W[e, k, n]
-        wg_ptrs = wg_base + k_offs[:, None].to(tl.int64) * stride_w_row + n_local[None, :].to(tl.int64) * stride_w_col
+        # W_gate[e, k, n]
+        wg_ptrs = W_gate_ptr + w_base + k[:, None] * I_val + n_offs[None, :]
         wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        # W_up
-        wu_ptrs = wu_base + k_offs[:, None].to(tl.int64) * stride_w_row + n_local[None, :].to(tl.int64) * stride_

... (2914 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,7 +1,7 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-Flattened 2D grid: (cdiv(T_perm, BM), cdiv(I, BN)) with expert lookup
-via binary search on offsets.
+Single kernel launch, flattened 2D grid (cdiv(T_perm, BM), cdiv(I, BN)).
+Each kernel block determines its owning expert by binary-searching offsets.
 """
 from __future__ import annotations
 
@@ -21,7 +21,6 @@
     T_perm,
     H,
     I_val,
-    E,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
@@ -29,67 +28,69 @@
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
-    # Find expert for the first row in this block
-    row_0 = pid_m * BM
-    # Binary search for expert: find smallest e such that Offsets[e+1] > row_0
+    # First row this block covers
+    row_start_block = pid_m * BM
+
+    # Binary search for expert: find smallest e where Offsets[e+1] > row_start_block
     lo = 0
-    hi = E - 1
+    hi = 127  # max E is 128
     while lo < hi:
-        mid = (lo + hi) // 2
+        mid = (lo + hi) >> 1
         off = tl.load(Offsets_ptr + mid + 1).to(tl.int64)
-        if off <= row_0:
+        if off <= row_start_block:
             lo = mid + 1
         else:
             hi = mid
     expert = lo
 
+    # Read expert boundaries
     row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
     row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
 
-    # M offset (intra-expert rows)
-    m_local = pid_m * BM + tl.arange(0, BM) - row_start
-    m_valid = (pid_m * BM + tl.arange(0, BM)) < row_end
+    # Local row offsets within this expert
+    m_local = pid_m * BM - row_start + tl.arange(0, BM)
+    m_valid = m_local >= 0 and m_local < (row_end - row_start)
+    row_valid = row_start_block + tl.arange(0, BM) < T_perm
 
-    # Global rows
     row_g = pid_m * BM + tl.arange(0, BM)
-    row_valid = row_g < T_perm & m_valid
 
     # N offset
     n_offs = pid_n * BN + tl.arange(0, BN)
     n_valid = n_offs < I_val
 
-    # K offset
+    # K
     k_offs = tl.arange(0, BK)
 
+    # Accumulators
     acc_g = tl.zeros((BM, BN), dtype=tl.float32)
     acc_u = tl.zeros((BM, BN), dtype=tl.float32)
 
-    # Weight base: W is (E, H, I), row major
-    # W[e, k, n] offset = e * H * I + k * I + n
-    w_base = expert * H * I_val
+    # Weight base: expert * H * I
+    w_base = expert.to(tl.int64) * H * I_val
 
+    # Main GEMM loop
     for k_start in range(0, H, BK):
-        k_mask = k_start + k_offs < H
         k = k_start + k_offs
+        k_mask = k < H
 
-        # A[row, k]
-        a_ptrs = X_ptr + row_g[:, None] * H + k[None, :]
+        # X[row, k]: stride is H (row-major contiguous)
+        a_ptrs = X_ptr + row_g[:, None].to(tl.int64) * H + k[None, :]
         a = tl.load(a_ptrs, mask=row_valid[:, None] & k_mask[None, :], other=0.0)
 
-        # W_gate[e, k, n]
+        # W_gate[expert, k, n]: (H, I) slice
         wg_ptrs = W_gate_ptr + w_base + k[:, None] * I_val + n_offs[None, :]
         wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        # W_up[e, k, n]
+        # W_up[expert, k, n]
         wu_ptrs = W_up_ptr + w_base + k[:, None] * I_val + n_offs[None, :]
         wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        acc_g = tl.dot(a, wg, acc_g, out_dtype=tl.float32)
-        acc_u = tl.dot(a, wu, acc_u, out_dtype=tl.float32)
+        acc_g += tl.dot(a, wg)
+        acc_u += tl.dot(a, wu)
 
     # SwiGLU
     out_fp = acc_g * tl.sigmoid(acc_g) * acc_u
-    out_ptrs = Out_ptr + row_g[:, None] * I_val + n_offs[None, :]
+    out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * I_val + n_offs[None, :]
     tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=row_valid[:, None] & n_valid[None, :])
 
 
@@ -97,17 +98,16 @@
     T_perm = X.size(0)
     H = X.size(1)
     I_val = out.size(1)
-    E_val = expert_offsets.size(0) - 1
 
-    BM = 256
-    BN = 128
+    BM = 128
+    BN = 64
     BK = 128
 
     grid = (triton.cdiv(T_perm, BM), triton.cdiv(I_val, BN))
 
     _grouped_gemm_swiglu_kernel[grid](
         X, W_gate, W_up, out, expert_offsets,
-        T_perm, H, I_val, E_val,
+        T_perm, H, I_val,
         BM=BM, BN=BN, BK=BK,
     )
 

--- 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
@@ -105,6 +105,7 @@
 
     grid = (triton.cdiv(T_perm, BM), triton.cdiv(I_val, BN))
 
+    # num_stages=1 to stay within shared memory limit on SM120
     _grouped_gemm_swiglu_kernel[grid](
         X, W_gate, W_up, out, expert_offsets,
         T_perm, H, I_val,

--- 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
@@ -105,11 +105,11 @@
 
     grid = (triton.cdiv(T_perm, BM), triton.cdiv(I_val, BN))
 
-    # num_stages=1 to stay within shared memory limit on SM120
     _grouped_gemm_swiglu_kernel[grid](
         X, W_gate, W_up, out, expert_offsets,
         T_perm, H, I_val,
         BM=BM, BN=BN, BK=BK,
+        num_stages=1, num_warps=8,
     )
 
 

--- 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
@@ -49,7 +49,7 @@
 
     # Local row offsets within this expert
     m_local = pid_m * BM - row_start + tl.arange(0, BM)
-    m_valid = m_local >= 0 and m_local < (row_end - row_start)
+    m_valid = (m_local >= 0) & (m_local < (row_end - row_start))
     row_valid = row_start_block + tl.arange(0, BM) < T_perm
 
     row_g = pid_m * BM + tl.arange(0, BM)

--- 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,7 +1,7 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-Single kernel launch, flattened 2D grid (cdiv(T_perm, BM), cdiv(I, BN)).
-Each kernel block determines its owning expert by binary-searching offsets.
+3D grid (E, cdiv(MAX_M, BM), cdiv(I, BN)) - no binary search needed.
+Blocks whose range falls outside the expert's row range are no-ops.
 """
 from __future__ import annotations
 
@@ -18,71 +18,68 @@
     W_up_ptr,
     Out_ptr,
     Offsets_ptr,
-    T_perm,
+    MAX_M,
     H,
     I_val,
+    stride_x_row,
+    stride_o_row,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
 ):
-    pid_m = tl.program_id(0)
-    pid_n = tl.program_id(1)
+    pid_e = tl.program_id(0)
+    pid_bm = tl.program_id(1)
+    pid_n = tl.program_id(2)
 
-    # First row this block covers
-    row_start_block = pid_m * BM
+    # Expert row bounds
+    row_start = tl.load(Offsets_ptr + pid_e).to(tl.int64)
+    row_end = tl.load(Offsets_ptr + pid_e + 1).to(tl.int64)
+    n_e = row_end - row_start
 
-    # Binary search for expert: find smallest e where Offsets[e+1] > row_start_block
-    lo = 0
-    hi = 127  # max E is 128
-    while lo < hi:
-        mid = (lo + hi) >> 1
-        off = tl.load(Offsets_ptr + mid + 1).to(tl.int64)
-        if off <= row_start_block:
-            lo = mid + 1
-        else:
-            hi = mid
-    expert = lo
+    # First row of this block
+    block_first = pid_bm * BM
 
-    # Read expert boundaries
-    row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
-    row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
+    # If this block is entirely outside the expert's range, skip
+    if block_first >= n_e:
+        return
 
-    # Local row offsets within this expert
-    m_local = pid_m * BM - row_start + tl.arange(0, BM)
-    m_valid = (m_local >= 0) & (m_local < (row_end - row_start))
-    row_valid = row_start_block + tl.arange(0, BM) < T_perm
+    # Local row offsets within expert
+    m_local = block_first + tl.arange(0, BM)
+    m_valid = m_local < n_e
 
-    row_g = pid_m * BM + tl.arange(0, BM)
+    # Global row indices
+    row_g = row_start + m_local
 
-    # N offset
+    # Column offsets
     n_offs = pid_n * BN + tl.arange(0, BN)
     n_valid = n_offs < I_val
 
-    # K
+    # Reduction dimension offsets
     k_offs = tl.arange(0, BK)
 
     # Accumulators
     acc_g = tl.zeros((BM, BN), dtype=tl.float32)
     acc_u = tl.zeros((BM, BN), dtype=tl.float32)
 
-    # Weight base: expert * H * I
-    w_base = expert.to(tl.int64) * H * I_val
+    # Weight base for this expert: W is (E, H, I), row-major
+    wg_base = W_gate_ptr + pid_e * H * I_val
+    wu_base = W_up_ptr + pid_e * H * I_val
 
-    # Main GEMM loop
+    # Main GEMM loop over K=H dimension
     for k_start in range(0, H, BK):
         k = k_start + k_offs
         k_mask = k < H
 
-        # X[row, k]: stride is H (row-major contiguous)
-        a_ptrs = X_ptr + row_g[:, None].to(tl.int64) * H + k[None, :]
-        a = tl.load(a_ptrs, mask=row_valid[:, None] & k_mask[None, :], other=0.0)
+        # Load A: (BM, BK)
+        a_ptrs = X_ptr + row_g[:, None].to(tl.int64) * stride_x_row + k[None, :]
+        a = tl.load(a_ptrs, mask=m_valid[:, None] & k_mask[None, :], other=0.0)
 
-        # W_gate[expert, k, n]: (H, I) slice
-        wg_ptrs = W_gate_ptr + w_base + k[:, None] * I_val + n_offs[None, :]
+        # Load W_gate[e, k, n]: (BK, BN)
+        wg_ptrs = wg_base + k[:, None] * I_val + n_offs[None, :]
         wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
-        # W_up[expert, k, n]
-        wu_ptrs = W_up_ptr + w_base + k[:, None] * I_val + n_offs[None, :]
+        # Load W_up[e, k, n]: (BK, BN)
+        wu_ptrs = wu_base + k[:, None] * I_val + n_offs[None, :]
         wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
 
         acc_g += tl.dot(a, wg)
@@ -90,26 +87,52 @@
 
     # SwiGLU
     out_fp = acc_g * tl.sigmoid(acc_g) * acc_u
-    out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * I_val + n_offs[None, :]
-    tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=row_valid[:, None] & n_valid[None, :])
+    out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * stride_o_row + n_offs[None, :]
+    tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=m_valid[:, None] & n_valid[None, :])
 
 
 def _launch(X, W_gate, W_up, out, expert_offsets):
     T_perm = X.size(0)
     H = X.size(1)
     I_val = out.size(1)
+    E = expert_offsets.size(0) - 1
 
-    BM = 128
-    BN = 64
-    BK = 128
+    # Heuristic: pick block sizes that maximize FLOPS/tile while fitting in shared memory
+    # Shared memory: (BM * BK + 2 * BK * BN) * 2 bytes (bf16, num_stages=1)
+    # For BM=256, BN=256, BK=64: (256*64 + 2*64*256)*2 = (16384+32768)*2 = 98304 < 101k OK
+    # For BM=256, BN=128, BK=128: (256*128 + 2*128*128)*2 = (32768+32768)*2 = 131072 - too big
+    # For BM=128, BN=128, BK=128: (128*128 + 2*128*128)*2 = (16384+32768)*2 = 98304 < 101k OK
+    if H >= 2048:
+        BM = 128
+        BN = 128
+        BK = 128
+    else:
+        BM = 64
+        BN = 128
+        BK = 64
 
-    grid = (triton.cdiv(T_perm, BM), triton.cdiv(I_val, BN))
+    # For SM120, try larger blocks when possible
+    # 128x256x64: (128*64 + 2*64*256)*2 = (8192+32768)*2 = 81920 - OK
+    # 256x128x64: (256*64 + 2*64*128)*2 = (16384+16384)*2 = 65536 - OK
+    if I_val >= 1024 and H >= 4096:
+        BM = 256
+        BN = 128
+        BK = 64
+    elif I_val >= 4096:
+        BM = 128
+        BN = 128
+        BK = 128
+
+    grid = (E, triton.cdiv(T_perm, BM), triton.cdiv(I_val, BN))
 
     _grouped_gemm_swiglu_kernel[grid](
         X, W_gate, W_up, out, expert_offsets,
         T_perm, H, I_val,
+        stride_x_row=X.stride(0),
+        stride_o_row=out.stride(0),
 

... (495 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,7 +1,7 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-3D grid (E, cdiv(MAX_M, BM), cdiv(I, BN)) - no binary search needed.
-Blocks whose range falls outside the expert's row range are no-ops.
+3D grid (E, cdiv(n_tokens_for_e, BM), cdiv(I, BN)) - grid size is
+computed from actual expert token counts, no wasted blocks.
 """
 from __future__ import annotations
 
@@ -18,7 +18,6 @@
     W_up_ptr,
     Out_ptr,
     Offsets_ptr,
-    MAX_M,
     H,
     I_val,
     stride_x_row,
@@ -36,15 +35,8 @@
     row_end = tl.load(Offsets_ptr + pid_e + 1).to(tl.int64)
     n_e = row_end - row_start
 
-    # First row of this block
-    block_first = pid_bm * BM
-
-    # If this block is entirely outside the expert's range, skip
-    if block_first >= n_e:
-        return
-
     # Local row offsets within expert
-    m_local = block_first + tl.arange(0, BM)
+    m_local = pid_bm * BM + tl.arange(0, BM)
     m_valid = m_local < n_e
 
     # Global row indices
@@ -88,6 +80,129 @@
     # SwiGLU
     out_fp = acc_g * tl.sigmoid(acc_g) * acc_u
     out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * stride_o_row + n_offs[None, :]
+    tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=m_valid[:, None] & n_valid[None, :])
+
+
+def _expert_grid(expert_offsets, BM, BN, I):
+    """Build 3D grid from expert_offsets: (E_list, M_list, N_list) where
+    E_list is just [0,1,...,E-1] repeated, and M_list/N_list are per-expert tile counts."""
+    E = len(expert_offsets) - 1
+    counts = expert_offsets[1:] - expert_offsets[:-1]  # tokens per expert
+    m_tiles = (counts + BM - 1) // BM  # cdiv per expert
+    n_tiles = (I + BN - 1) // BN
+
+    # Create grid as tensor of shape (total_work_items, 3)
+    # Actually, Triton doesn't support tensor grids. We need to use a flat list approach.
+    # Instead, let's use a 1D grid and decode (e, m, n) from the flat index.
+    total_work = int(m_tiles.sum()) * n_tiles
+    return E, int(m_tiles.sum()), n_tiles, m_tiles, total_work
+
+
+@triton.jit
+def _grouped_gemm_1d_kernel(
+    X_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    Out_ptr,
+    Offsets_ptr,
+    MTilesPrefix_ptr,  # cumulative sum of m_tiles: (E+1,) int32, prefix sum of m_tiles
+    H,
+    I_val,
+    total_m_tiles,
+    BN: tl.constexpr,
+    stride_x_row,
+    stride_o_row,
+    BM: tl.constexpr,
+    BLOCK_BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid_work = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    pid_bm_in_expert = pid_work  # each pid_work handles one M-tile for one expert
+    # But we need to know WHICH expert AND which M-tile.
+    # Actually, let's restructure: grid = (total_m_tiles, cdiv(I, BN))
+    # Then for each pid_m in total_m_tiles, find which expert it belongs to.
+    pass
+
+
+# Simplest approach: 2D flattened grid (total_tiles_M, cdiv(I, BN))
+# where total_tiles_M = sum_e cdiv(n_tokens_e, BM)
+# Use binary search on cumulative_M_tiles to find expert index.
+
+
+@triton.jit
+def _grouped_gemm_flat_kernel(
+    X_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    Out_ptr,
+    Offsets_ptr,
+    CumMTiles_ptr,  # (E+1,) int32, cumulative M-tiles: CumMTiles[e] = sum of cdiv(n_tokens_e', BM) for e' < e
+    H,
+    I_val,
+    n_experts,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+    BK: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    # Find expert for this M-tile: largest e such that CumMTiles[e] <= pid_m
+    lo = 0
+    hi = n_experts - 1
+    while lo < hi:
+        mid = (lo + hi + 1) >> 1
+        val = tl.load(CumMTiles_ptr + mid).to(tl.int32)
+        if val <= pid_m:
+            lo = mid
+        else:
+            hi = mid - 1
+    expert = lo
+
+    # M-tile index within this expert
+    tile_off = tl.load(CumMTiles_ptr + expert).to(tl.int32)
+    pid_bm = pid_m - tile_off
+
+    # Expert row bounds
+    row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
+    row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
+    n_e = row_end - row_start
+
+    m_local = pid_bm * BM + tl.arange(0, BM)
+    m_valid = m_local < n_e
+
+    row_g = row_start + m_local
+
+    n_offs = pid_n * BN + tl.arange(0, BN)
+    n_valid = n_offs < I_val
+
+    k_offs = tl.arange(0, BK)
+
+    acc_g = tl.zeros((BM, BN), dtype=tl.float32)
+    acc_u = tl.zeros((BM, BN), dtype=tl.float32)
+
+    wg_base = W_gate_ptr + expert * H * I_val
+    wu_base = W_up_ptr + expert * H * I_val
+
+    for k_start in range(0, H, BK):
+        k = k_start + k_offs
+        k_mask = k < H
+
+        a_ptrs = X_ptr + row_g[:, None].to(tl.int64) * H + k[None, :]
+        a = tl.load(a_ptrs, mask=m_valid[:, None] & k_mask[None, :], other=0.0)
+
+        wg_ptrs = wg_base + k[:, None] * I_val + n_offs[None, :]
+        wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
+
+        wu_ptrs = wu_base + k[:, None] * I_val + n_offs[None, :]
+        wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
+
+        acc_g += tl.dot(a, wg)
+        acc_u += tl.dot(a, wu)
+
+    out_fp = acc_g * tl.sigmoid(acc_g) * acc_u
+    out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * I_val + n_offs[None, :]
     tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=m_valid[:, None] & n_valid[None, :])
 
 
@@ -96,40 +211,42 @@
     H = X.size(1)
     I_val = out.size(1)
     E = expert_offsets.size(0) - 1
-
-    # Heuristic: pick block sizes that maximize FLOPS/tile while fitting in shared memory
-    # Shared memory: (BM * BK + 2 * BK * BN) * 2 bytes (bf16, num_stages=1)
-    # For BM=256, BN=256, BK=64: (256*64 + 2*64*256)*2 = (16384+32768)*2 = 98304 < 101k OK
-    # For BM=256, BN=128, BK=128: (256*128 + 2*128*128)*2 = (32768+32768)*2 = 131072 - too big
-    # For BM=128, BN=128, BK=128: (128*128 + 2*128*128)*2 = (16384+32768)*2 = 98304 < 101k OK
-    if H >= 

... (1858 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,7 +1,9 @@
 """Triton grouped GEMM + fused SwiGLU for Sonic-MoE up-projection.
 
-3D grid (E, cdiv(n_tokens_for_e, BM), cdiv(I, BN)) - grid size is
-computed from actual expert token counts, no wasted blocks.
+2D grid: (sum_e cdiv(n_e, BM), cdiv(I, BN)) with per-block expert lookup
+via binary search on cumulative tile counts.
+
+Block sizes tuned for SM120 shared memory budget (~98KB usable).
 """
 from __future__ import annotations
 
@@ -13,177 +15,61 @@
 
 @triton.jit
 def _grouped_gemm_swiglu_kernel(
-    X_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    Out_ptr,
-    Offsets_ptr,
+    X_ptr,             # (T_perm, H), row-major
+    W_gate_ptr,        # (E, H, I), row-major
+    W_up_ptr,          # (E, H, I), row-major
+    Out_ptr,           # (T_perm, I), row-major
+    Offsets_ptr,       # (E+1,) int32
+    CumMTiles_ptr,     # (E+1,) int32
     H,
     I_val,
-    stride_x_row,
-    stride_o_row,
     BM: tl.constexpr,
     BN: tl.constexpr,
     BK: tl.constexpr,
+    NUM_SMS: tl.constexpr,
 ):
-    pid_e = tl.program_id(0)
-    pid_bm = tl.program_id(1)
-    pid_n = tl.program_id(2)
+    pid = tl.program_id(0)
+    pid_n = tl.program_id(1)
 
-    # Expert row bounds
-    row_start = tl.load(Offsets_ptr + pid_e).to(tl.int64)
-    row_end = tl.load(Offsets_ptr + pid_e + 1).to(tl.int64)
+    # 1D -> (pid_m, ...) mapping
+    pid_m = pid
+
+    # Find expert: largest e such that CumMTiles[e] <= pid_m
+    # E is at most 128; binary search
+    lo = 0
+    hi = 127
+    while lo < hi:
+        mid = (lo + hi + 1) >> 1
+        if tl.load(CumMTiles_ptr + mid).to(tl.int32) <= pid_m:
+            lo = mid
+        else:
+            hi = mid - 1
+    expert = lo
+
+    # Intra-expert M-tile index
+    m_tile_idx = pid_m - tl.load(CumMTiles_ptr + expert).to(tl.int32)
+
+    # Expert row range
+    row_start = tl.load(Offsets_ptr + expert).to(tl.int64)
+    row_end = tl.load(Offsets_ptr + expert + 1).to(tl.int64)
     n_e = row_end - row_start
 
-    # Local row offsets within expert
-    m_local = pid_bm * BM + tl.arange(0, BM)
-    m_valid = m_local < n_e
-
-    # Global row indices
-    row_g = row_start + m_local
+    # Row offsets
+    m_offs = m_tile_idx * BM + tl.arange(0, BM)
+    m_valid = m_offs < n_e
+    row_g = row_start + m_offs
 
     # Column offsets
     n_offs = pid_n * BN + tl.arange(0, BN)
     n_valid = n_offs < I_val
 
-    # Reduction dimension offsets
-    k_offs = tl.arange(0, BK)
-
-    # Accumulators
-    acc_g = tl.zeros((BM, BN), dtype=tl.float32)
-    acc_u = tl.zeros((BM, BN), dtype=tl.float32)
-
-    # Weight base for this expert: W is (E, H, I), row-major
-    wg_base = W_gate_ptr + pid_e * H * I_val
-    wu_base = W_up_ptr + pid_e * H * I_val
-
-    # Main GEMM loop over K=H dimension
-    for k_start in range(0, H, BK):
-        k = k_start + k_offs
-        k_mask = k < H
-
-        # Load A: (BM, BK)
-        a_ptrs = X_ptr + row_g[:, None].to(tl.int64) * stride_x_row + k[None, :]
-        a = tl.load(a_ptrs, mask=m_valid[:, None] & k_mask[None, :], other=0.0)
-
-        # Load W_gate[e, k, n]: (BK, BN)
-        wg_ptrs = wg_base + k[:, None] * I_val + n_offs[None, :]
-        wg = tl.load(wg_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
-
-        # Load W_up[e, k, n]: (BK, BN)
-        wu_ptrs = wu_base + k[:, None] * I_val + n_offs[None, :]
-        wu = tl.load(wu_ptrs, mask=k_mask[:, None] & n_valid[None, :], other=0.0)
-
-        acc_g += tl.dot(a, wg)
-        acc_u += tl.dot(a, wu)
-
-    # SwiGLU
-    out_fp = acc_g * tl.sigmoid(acc_g) * acc_u
-    out_ptrs = Out_ptr + row_g[:, None].to(tl.int64) * stride_o_row + n_offs[None, :]
-    tl.store(out_ptrs, out_fp.to(tl.bfloat16), mask=m_valid[:, None] & n_valid[None, :])
-
-
-def _expert_grid(expert_offsets, BM, BN, I):
-    """Build 3D grid from expert_offsets: (E_list, M_list, N_list) where
-    E_list is just [0,1,...,E-1] repeated, and M_list/N_list are per-expert tile counts."""
-    E = len(expert_offsets) - 1
-    counts = expert_offsets[1:] - expert_offsets[:-1]  # tokens per expert
-    m_tiles = (counts + BM - 1) // BM  # cdiv per expert
-    n_tiles = (I + BN - 1) // BN
-
-    # Create grid as tensor of shape (total_work_items, 3)
-    # Actually, Triton doesn't support tensor grids. We need to use a flat list approach.
-    # Instead, let's use a 1D grid and decode (e, m, n) from the flat index.
-    total_work = int(m_tiles.sum()) * n_tiles
-    return E, int(m_tiles.sum()), n_tiles, m_tiles, total_work
-
-
-@triton.jit
-def _grouped_gemm_1d_kernel(
-    X_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    Out_ptr,
-    Offsets_ptr,
-    MTilesPrefix_ptr,  # cumulative sum of m_tiles: (E+1,) int32, prefix sum of m_tiles
-    H,
-    I_val,
-    total_m_tiles,
-    BN: tl.constexpr,
-    stride_x_row,
-    stride_o_row,
-    BM: tl.constexpr,
-    BLOCK_BN: tl.constexpr,
-    BK: tl.constexpr,
-):
-    pid_work = tl.program_id(0)
-    pid_n = tl.program_id(1)
-    pid_bm_in_expert = pid_work  # each pid_work handles one M-tile for one expert
-    # But we need to know WHICH expert AND which M-tile.
-    # Actually, let's restructure: grid = (total_m_tiles, cdiv(I, BN))
-    # Then for each pid_m in total_m_tiles, find which expert it belongs to.
-    pass
-
-
-# Simplest approach: 2D flattened grid (total_tiles_M, cdiv(I, BN))
-# where total_tiles_M = sum_e cdiv(n_tokens_e, BM)
-# Use binary search on cumulative_M_tiles to find expert index.
-
-
-@triton.jit
-def _grouped_gemm_flat_kernel(
-    X_ptr,
-    W_gate_ptr,
-    W_up_ptr,
-    Out_ptr,
-    Offsets_ptr,
-    CumMTiles_ptr,  # (E+1,) int32, cumulative M-tiles: CumMTiles[e] = sum of cdiv(n_tokens_e', BM) for e' < e
-    H,
-    I_val,
-    n_experts,
-    BM: tl.constexpr,
-    BN: tl.constexpr,
-    BK: tl.constexpr,
-):
-    pid_m = tl.program_id(0)
-  

... (3921 more chars)