--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -0,0 +1,198 @@
+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"]
+
+
+def _get_configs(H: int, I: int, M_avg: int) -> list[triton.Config]:
+    """Return heuristic configs based on problem dimensions."""
+    configs = []
+    # M tile: pick based on average tokens-per-expert
+    for bm in [32, 64, 128]:
+        if bm > M_avg and bm != 32:
+            continue
+        for bn in [64, 128, 256]:
+            if bn >= I:
+                bn = I
+            for bk in [32, 64]:
+                for nw in [4, 8]:
+                    configs.append(
+                        triton.Config(
+                            {"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk},
+                            num_warps=nw,
+                            num_stages=3,
+                        )
+                    )
+    return configs
+
+
+@triton.autotune(
+    configs=[],
+    key=["H", "I"],
+    use_cuda_graph=True,
+)
+@triton.jit
+def grouped_gemm_swiglu_kernel(
+    hidden_states_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    expert_offsets_ptr,
+    output_ptr,
+    H: tl.constexpr,
+    I: tl.constexpr,
+    stride_hidden_h: tl.constexpr,
+    stride_gate_e: tl.constexpr,
+    stride_gate_h: tl.constexpr,
+    stride_gate_i: tl.constexpr,
+    stride_up_e: tl.constexpr,
+    stride_up_h: tl.constexpr,
+    stride_up_i: tl.constexpr,
+    stride_out_i: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    e = pid // num_n_tiles
+    n_tile_idx = pid % num_n_tiles
+
+    start = tl.load(expert_offsets_ptr + e)
+    end = tl.load(expert_offsets_ptr + e + 1)
+    n_tokens = end - start
+    if n_tokens <= 0:
+        return
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    w_gate_base = W_gate_ptr + e * stride_gate_e
+    w_up_base = W_up_ptr + e * stride_up_e
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+
+        x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
+
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            x = tl.load(
+                x_ptr + m_off[:, None] * stride_hidden_h + k_off[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            w_gate = tl.load(
+                w_gate_base
+                + k_off[:, None] * stride_gate_h
+                + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up = tl.load(
+                w_up_base
+                + k_off[:, None] * stride_up_h
+                + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+
+            acc_gate = tl.dot(x, w_gate, acc_gate)
+            acc_up = tl.dot(x, w_up, acc_up)
+
+        gate_silu = tl.sigmoid(acc_gate) * acc_gate
+        result = (gate_silu * acc_up).to(tl.bfloat16)
+
+        out_ptr = output_ptr + (start + m_start) * stride_out_i
+        tl.store(
+            out_ptr + m_off[:, None] * stride_out_i + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, T_total: int, H: int, I: int, E: int, K: int):
+        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,
+        expert_offsets: torch.Tensor,
+    ) -> torch.Tensor:
+        T_perm, _ = hidden_states.shape
+        out = torch.empty(T_perm, self.I, dtype=torch.bfloat16, device=hidden_states.device)
+
+        M_avg = T_perm // self.E
+        grid = (self.E * triton.cdiv(self.I, 128),)
+
+        grouped_gemm_swiglu_kernel[grid](
+            hidden_states,
+            self.W_gate,
+            self.W_up,
+            expert_offsets,
+            out,
+            self.H,
+            self.I,
+            hidden_states.stride(0),
+            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),
+            BLOCK_M=64,
+            BLOCK_N=128,
+            BLOCK_K=32,
+            num_warps=4,
+            num_stages=3,
+        )
+        return out
+
+
+T_total = 32768
+H = 4096
+I = 1536
+E = 128
+K = 8
+
+
+def _build_routing(T_total: int, E: int, K: int, device: str = "cpu") -> torch.Tensor:
+    T_perm = T_total * K
+    base = T_perm // E
+    rem = T_perm - base * E
+    counts = torch.full((E,), base, dtype=torch.int32, device=device)
+    counts[:rem] += 1
+    offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
+    offsets[1:] = torch.cumsum(counts, dim=0)
+    return offsets
+
+
+def get_inputs():
+    T_perm = T_total * K
+    hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
+    expert_offsets = _build_routing(T_total, E, K)
+    return [hidden_states, expert_offsets]
+
+
+def get_init_in

... (42 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
@@ -10,33 +10,6 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-def _get_configs(H: int, I: int, M_avg: int) -> list[triton.Config]:
-    """Return heuristic configs based on problem dimensions."""
-    configs = []
-    # M tile: pick based on average tokens-per-expert
-    for bm in [32, 64, 128]:
-        if bm > M_avg and bm != 32:
-            continue
-        for bn in [64, 128, 256]:
-            if bn >= I:
-                bn = I
-            for bk in [32, 64]:
-                for nw in [4, 8]:
-                    configs.append(
-                        triton.Config(
-                            {"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk},
-                            num_warps=nw,
-                            num_stages=3,
-                        )
-                    )
-    return configs
-
-
-@triton.autotune(
-    configs=[],
-    key=["H", "I"],
-    use_cuda_graph=True,
-)
 @triton.jit
 def grouped_gemm_swiglu_kernel(
     hidden_states_ptr,

--- 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
@@ -115,7 +115,12 @@
         out = torch.empty(T_perm, self.I, dtype=torch.bfloat16, device=hidden_states.device)
 
         M_avg = T_perm // self.E
-        grid = (self.E * triton.cdiv(self.I, 128),)
+        bn = 128 if self.I >= 128 else 32
+        bm = 64
+        bk = 32
+        nw = 4
+        ns = 3
+        grid = (self.E * triton.cdiv(self.I, bn),)
 
         grouped_gemm_swiglu_kernel[grid](
             hidden_states,
@@ -133,11 +138,11 @@
             self.W_up.stride(1),
             self.W_up.stride(2),
             out.stride(0),
-            BLOCK_M=64,
-            BLOCK_N=128,
-            BLOCK_K=32,
-            num_warps=4,
-            num_stages=3,
+            BLOCK_M=bm,
+            BLOCK_N=bn,
+            BLOCK_K=bk,
+            num_warps=nw,
+            num_stages=ns,
         )
         return out
 

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug.py
@@ -0,0 +1,158 @@
+"""Debug script: test the triton kernel on a small case and compare with reference."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+# Small case: 2 experts, 8 tokens each, small dims
+T_total = 4
+H = 32
+I = 16
+E = 2
+K = 2
+T_perm = T_total * K  # 8
+
+# Round-robin routing
+base = T_perm // E
+rem = T_perm - base * E
+counts = torch.full((E,), base, dtype=torch.int32)
+counts[:rem] += 1
+expert_offsets = torch.zeros(E + 1, dtype=torch.int32)
+expert_offsets[1:] = torch.cumsum(counts, dim=0)
+print(f"expert_offsets: {expert_offsets}")
+
+hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
+W_gate = torch.randn(E, H, I, dtype=torch.bfloat16) * 0.02
+W_up = torch.randn(E, H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference
+ref_out = torch.empty(T_perm, I, dtype=torch.bfloat16)
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    x_e = hidden_states[start:end]
+    gate = x_e @ W_gate[e]
+    up = x_e @ W_up[e]
+    ref_out[start:end] = torch.nn.functional.silu(gate) * up
+
+
+# Triton kernel (copy from solution)
+@triton.jit
+def grouped_gemm_swiglu_kernel(
+    hidden_states_ptr,
+    W_gate_ptr,
+    W_up_ptr,
+    expert_offsets_ptr,
+    output_ptr,
+    H: tl.constexpr,
+    I: tl.constexpr,
+    stride_hidden_h: tl.constexpr,
+    stride_gate_e: tl.constexpr,
+    stride_gate_h: tl.constexpr,
+    stride_gate_i: tl.constexpr,
+    stride_up_e: tl.constexpr,
+    stride_up_h: tl.constexpr,
+    stride_up_i: tl.constexpr,
+    stride_out_i: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    e = pid // num_n_tiles
+    n_tile_idx = pid % num_n_tiles
+
+    start = tl.load(expert_offsets_ptr + e)
+    end = tl.load(expert_offsets_ptr + e + 1)
+    n_tokens = end - start
+    if n_tokens <= 0:
+        return
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    w_gate_base = W_gate_ptr + e * stride_gate_e
+    w_up_base = W_up_ptr + e * stride_up_e
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+
+        x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
+
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            x = tl.load(
+                x_ptr + m_off[:, None] * stride_hidden_h + k_off[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            w_gate = tl.load(
+                w_gate_base
+                + k_off[:, None] * stride_gate_h
+                + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up = tl.load(
+                w_up_base
+                + k_off[:, None] * stride_up_h
+                + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+
+            acc_gate = tl.dot(x, w_gate, acc_gate)
+            acc_up = tl.dot(x, w_up, acc_up)
+
+        gate_silu = tl.sigmoid(acc_gate) * acc_gate
+        result = (gate_silu * acc_up).to(tl.bfloat16)
+
+        out_ptr = output_ptr + (start + m_start) * stride_out_i
+        tl.store(
+            out_ptr + m_off[:, None] * stride_out_i + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+
+triton_out = torch.empty(T_perm, I, dtype=torch.bfloat16, device='cuda')
+hs_cuda = hidden_states.cuda()
+wg_cuda = W_gate.cuda()
+wu_cuda = W_up.cuda()
+eo_cuda = expert_offsets.cuda()
+
+grid = (E * triton.cdiv(I, 16),)
+grouped_gemm_swiglu_kernel[grid](
+    hs_cuda, wg_cuda, wu_cuda, eo_cuda, triton_out,
+    H, I,
+    hs_cuda.stride(0),
+    wg_cuda.stride(0), wg_cuda.stride(1), wg_cuda.stride(2),
+    wu_cuda.stride(0), wu_cuda.stride(1), wu_cuda.stride(2),
+    triton_out.stride(0),
+    BLOCK_M=4, BLOCK_N=16, BLOCK_K=8, num_warps=4, num_stages=2,
+)
+
+print(f"Reference:\n{ref_out}")
+print(f"Triton:\n{triton_out.cpu()}")
+diff = (ref_out.float() - triton_out.cpu().float()).abs()
+print(f"Max diff: {diff.max().item()}")
+print(f"Mean diff: {diff.mean().item()}")
+
+if torch.allclose(ref_out.float(), triton_out.cpu().float(), atol=0.02, rtol=0.02):
+    print("PASS: results match")
+else:
+    print("FAIL: results differ")
+    # Print per-expert comparison
+    for e in range(E):
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        ref_slice = ref_out[start:end].float()
+        sol_slice = triton_out[start:end].cpu().float()
+        d = (ref_slice - sol_slice).abs().max().item()
+        print(f"  Expert {e}: tokens {start}-{end}, max_diff={d:.6f}")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug.py
@@ -135,7 +135,7 @@
     wg_cuda.stride(0), wg_cuda.stride(1), wg_cuda.stride(2),
     wu_cuda.stride(0), wu_cuda.stride(1), wu_cuda.stride(2),
     triton_out.stride(0),
-    BLOCK_M=4, BLOCK_N=16, BLOCK_K=8, num_warps=4, num_stages=2,
+    BLOCK_M=4, BLOCK_N=16, BLOCK_K=16, num_warps=4, num_stages=2,
 )
 
 print(f"Reference:\n{ref_out}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug2.py
@@ -0,0 +1,141 @@
+"""Test with actual problem shapes."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+T_total = 2048  # small enough to debug easily
+H = 4096
+I = 1536
+E = 4
+K = 8
+T_perm = T_total * K
+
+# Round-robin routing
+base = T_perm // E
+rem = T_perm - base * E
+counts = torch.full((E,), base, dtype=torch.int32)
+counts[:rem] += 1
+expert_offsets = torch.zeros(E + 1, dtype=torch.int32)
+expert_offsets[1:] = torch.cumsum(counts, dim=0)
+print(f"expert_offsets: {expert_offsets}")
+print(f"Per-expert token counts: {[int(expert_offsets[i+1]-expert_offsets[i]) for i in range(E)]}")
+
+hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
+W_gate = torch.randn(E, H, I, dtype=torch.bfloat16) * 0.02
+W_up = torch.randn(E, H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference
+ref_out = torch.empty(T_perm, I, dtype=torch.bfloat16)
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    x_e = hidden_states[start:end]
+    gate = x_e @ W_gate[e]
+    up = x_e @ W_up[e]
+    ref_out[start:end] = torch.nn.functional.silu(gate) * up
+
+# Per-expert reference for comparison
+for e in range(E):
+    s = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    n = end - s
+    x_e = hidden_states[s:end]
+    gate_ref = x_e @ W_gate[e]
+    up_ref = x_e @ W_up[e]
+    out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+    
+    # Now test Triton for this expert
+    BLOCK_M = 64
+    BLOCK_N = 128
+    BLOCK_K = 32
+    
+    @triton.jit
+    def single_expert_kernel(
+        x_ptr, w_gate_ptr, w_up_ptr, out_ptr,
+        n_tokens: tl.constexpr,
+        H: tl.constexpr, I: tl.constexpr,
+        x_stride: tl.constexpr,
+        w_gate_stride_h: tl.constexpr, w_gate_stride_i: tl.constexpr,
+        w_up_stride_h: tl.constexpr, w_up_stride_i: tl.constexpr,
+        out_stride: tl.constexpr,
+        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    ):
+        num_n_tiles = tl.cdiv(I, BLOCK_N)
+        pid = tl.program_id(0)
+        n_tile_idx = pid
+        
+        n_start = n_tile_idx * BLOCK_N
+        n_off = tl.arange(0, BLOCK_N)
+        n_mask = n_off < I - n_start
+        
+        for m_start in range(0, n_tokens, BLOCK_M):
+            m_off = tl.arange(0, BLOCK_M)
+            m_mask = m_start + m_off < n_tokens
+            
+            x_base = x_ptr + m_start * x_stride
+            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, H, BLOCK_K):
+                k_off = tl.arange(0, BLOCK_K)
+                k_mask = k_start + k_off < H
+                
+                x = tl.load(
+                    x_base + m_off[:, None] * x_stride + k_off[None, :],
+                    mask=m_mask[:, None] & k_mask[None, :],
+                )
+                w_gate = tl.load(
+                    w_gate_ptr + k_off[:, None] * w_gate_stride_h + (n_start + n_off)[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+                w_up = tl.load(
+                    w_up_ptr + k_off[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+                acc_gate = tl.dot(x, w_gate, acc_gate)
+                acc_up = tl.dot(x, w_up, acc_up)
+            
+            gate_silu = tl.sigmoid(acc_gate) * acc_gate
+            result = (gate_silu * acc_up).to(tl.bfloat16)
+            
+            out_base = out_ptr + m_start * out_stride
+            tl.store(
+                out_base + m_off[:, None] * out_stride + n_off[None, :],
+                result,
+                mask=m_mask[:, None] & n_mask[None, :],
+            )
+    
+    num_n_tiles = triton.cdiv(I, BLOCK_N)
+    triton_out_e = torch.zeros(n, I, dtype=torch.bfloat16, device='cuda')
+    
+    x_cuda = x_e.unsqueeze(0).cuda()  # add batch dim... no, keep as (n, H)
+    x_cuda = x_e.cuda()
+    wg_cuda = W_gate[e:e+1].cuda().squeeze(0)
+    wu_cuda = W_up[e:e+1].cuda().squeeze(0)
+    
+    # Use parameters without the E dimension
+    single_expert_kernel[(num_n_tiles,)](
+        x_cuda, wg_cuda, wu_cuda, triton_out_e,
+        n, H, I,
+        x_cuda.stride(0),
+        wg_cuda.stride(0), wg_cuda.stride(1),
+        wu_cuda.stride(0), wu_cuda.stride(1),
+        triton_out_e.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+    )
+    
+    diff = (out_ref.cuda().float() - triton_out_e.float()).abs()
+    print(f"  Expert {e} (n={n}): max_diff={diff.max().item():.6f}")
+    
+    if diff.max().item() > 1.0:
+        print(f"    WARNING: large error!")
+        # Check individual tile
+        for ni in range(num_n_tiles):
+            ns = ni * BLOCK_N
+            ne = min(ns + BLOCK_N, I)
+            ref_tile = out_ref[:, ns:ne].float()
+            sol_tile = triton_out_e[:, ns:ne].cpu().float()
+            tile_diff = (ref_tile - sol_tile).abs()
+            print(f"    N-tile {ni} ({ns}:{ne}): max_diff={tile_diff.max().item():.6f}")

--- 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,18 +63,18 @@
             k_mask = k_start + k_off < H
 
             x = tl.load(
-                x_ptr + m_off[:, None] * stride_hidden_h + k_off[None, :],
+                x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
                 mask=m_mask[:, None] & k_mask[None, :],
             )
             w_gate = tl.load(
                 w_gate_base
-                + k_off[:, None] * stride_gate_h
+                + (k_start + k_off)[:, None] * stride_gate_h
                 + (n_start + n_off)[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )
             w_up = tl.load(
                 w_up_base
-                + k_off[:, None] * stride_up_h
+                + (k_start + k_off)[:, None] * stride_up_h
                 + (n_start + n_off)[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
@@ -0,0 +1,149 @@
+"""Debug remaining error: test individual components."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H = 4096
+I = 1536
+E = 128
+K = 8
+T_total = 4096
+T_perm = T_total * K
+
+# Single expert test
+n_tokens = 512  # tokens for this expert
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_up = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference (full precision)
+gate_ref = x.float() @ w_gate.float()
+up_ref = x.float() @ w_up.float()
+out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+
+# Reference (bf16 matmul)
+gate_bf16 = x @ w_gate  # bf16 matmul
+up_bf16 = x @ w_up
+out_bf16 = torch.nn.functional.silu(gate_bf16) * up_bf16
+
+# Compare
+diff_fp32 = (out_ref - out_bf16.float()).abs()
+print(f"bf16 reference vs fp32 reference: max_diff={diff_fp32.max().item():.6f}")
+
+# Now test with a single M-block
+BLOCK_M = 64
+BLOCK_N = 128
+BLOCK_K = 32
+
+@triton.jit
+def single_expert_kernel(
+    x_ptr, w_gate_ptr, w_up_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_gate_stride_h, w_gate_stride_i,
+    w_up_stride_h, w_up_stride_i,
+    out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    USE_MANUAL_SIGMOID: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+    
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+    
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        
+        x_base = x_ptr + m_start * x_stride
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+            
+            x_tile = tl.load(
+                x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            w_gate_tile = tl.load(
+                w_gate_ptr + (k_start + k_off)[:, None] * w_gate_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up_tile = tl.load(
+                w_up_ptr + (k_start + k_off)[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate)
+            acc_up = tl.dot(x_tile, w_up_tile, acc_up)
+        
+        if USE_MANUAL_SIGMOID:
+            gate_silu = (1.0 / (1.0 + tl.exp(-acc_gate))) * acc_gate
+        else:
+            gate_silu = tl.sigmoid(acc_gate) * acc_gate
+        result = (gate_silu * acc_up).to(tl.bfloat16)
+        
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+num_n_tiles = triton.cdiv(I, BLOCK_N)
+
+for use_manual in [False, True]:
+    triton_out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    single_expert_kernel[(num_n_tiles,)](
+        x.cuda(), w_gate.cuda(), w_up.cuda(), triton_out,
+        n_tokens, H, I,
+        x.stride(0), w_gate.stride(0), w_gate.stride(1),
+        w_up.stride(0), w_up.stride(1),
+        triton_out.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        USE_MANUAL_SIGMOID=use_manual,
+    )
+    
+    diff_vs_bf16 = (out_bf16.cuda().float() - triton_out.float()).abs()
+    diff_vs_fp32 = (out_ref.cuda().float() - triton_out.float()).abs()
+    
+    method = "manual sigmoid" if use_manual else "tl.sigmoid"
+    print(f"\n{method}:")
+    print(f"  vs bf16 ref: max_diff={diff_vs_bf16.max().item():.6f}")
+    print(f"  vs fp32 ref: max_diff={diff_vs_fp32.max().item():.6f}")
+    
+    # Per-tile breakdown
+    for ni in range(num_n_tiles):
+        ns = ni * BLOCK_N
+        ne = min(ns + BLOCK_N, I)
+        tile_diff = diff_vs_bf16[0, ns:ne].abs().max().item() if n_tokens > 0 else 0
+        print(f"  Tile {ni} ({ns}:{ne}): diff_vs_bf16_max={tile_diff:.6f}")
+
+# Also test: does accumulating twice give the same result?
+print("\n--- Reproducibility test ---")
+out1 = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+single_expert_kernel[(num_n_tiles,)](
+    x.cuda(), w_gate.cuda(), w_up.cuda(), out1,
+    n_tokens, H, I,
+    x.stride(0), w_gate.stride(0), w_gate.stride(1),
+    w_up.stride(0), w_up.stride(1),
+    out1.stride(0),
+    BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+    USE_MANUAL_SIGMOID=False,
+)
+out2 = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+single_expert_kernel[(num_n_tiles,)](
+    x.cuda(), w_gate.cuda(), w_up.cuda(), out2,
+    n_tokens, H, I,
+    x.stride(0), w_gate.stride(0), w_gate.stride(1),
+    w_up.stride(0), w_up.stride(1),
+    out2.stride(0),
+    BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+    USE_MANUAL_SIGMOID=False,
+)
+print(f"Run1 vs Run2: max_diff={(out1.float()-out2.float()).abs().max().item():.6f}")

--- 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
@@ -79,8 +79,8 @@
                 mask=k_mask[:, None] & n_mask[None, :],
             )
 
-            acc_gate = tl.dot(x, w_gate, acc_gate)
-            acc_up = tl.dot(x, w_up, acc_up)
+            acc_gate = tl.dot(x, w_gate, acc_gate, input_precision='ieee')
+            acc_up = tl.dot(x, w_up, acc_up, input_precision='ieee')
 
         gate_silu = tl.sigmoid(acc_gate) * acc_gate
         result = (gate_silu * acc_up).to(tl.bfloat16)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
@@ -79,8 +79,8 @@
                 w_up_ptr + (k_start + k_off)[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )
-            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate)
-            acc_up = tl.dot(x_tile, w_up_tile, acc_up)
+            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate, input_precision='ieee')
+            acc_up = tl.dot(x_tile, w_up_tile, acc_up, input_precision='ieee')
         
         if USE_MANUAL_SIGMOID:
             gate_silu = (1.0 / (1.0 + tl.exp(-acc_gate))) * acc_gate

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
@@ -79,8 +79,8 @@
                 w_up_ptr + (k_start + k_off)[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )
-            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate, input_precision='ieee')
-            acc_up = tl.dot(x_tile, w_up_tile, acc_up, input_precision='ieee')
+            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate, input_precision='ieee', max_num_imprecise_acc=0)
+            acc_up = tl.dot(x_tile, w_up_tile, acc_up, input_precision='ieee', max_num_imprecise_acc=0)
         
         if USE_MANUAL_SIGMOID:
             gate_silu = (1.0 / (1.0 + tl.exp(-acc_gate))) * acc_gate

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug3.py
@@ -98,31 +98,23 @@
 num_n_tiles = triton.cdiv(I, BLOCK_N)
 
 for use_manual in [False, True]:
-    triton_out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
-    single_expert_kernel[(num_n_tiles,)](
-        x.cuda(), w_gate.cuda(), w_up.cuda(), triton_out,
-        n_tokens, H, I,
-        x.stride(0), w_gate.stride(0), w_gate.stride(1),
-        w_up.stride(0), w_up.stride(1),
-        triton_out.stride(0),
-        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
-        USE_MANUAL_SIGMOID=use_manual,
-    )
-    
-    diff_vs_bf16 = (out_bf16.cuda().float() - triton_out.float()).abs()
-    diff_vs_fp32 = (out_ref.cuda().float() - triton_out.float()).abs()
-    
-    method = "manual sigmoid" if use_manual else "tl.sigmoid"
-    print(f"\n{method}:")
-    print(f"  vs bf16 ref: max_diff={diff_vs_bf16.max().item():.6f}")
-    print(f"  vs fp32 ref: max_diff={diff_vs_fp32.max().item():.6f}")
-    
-    # Per-tile breakdown
-    for ni in range(num_n_tiles):
-        ns = ni * BLOCK_N
-        ne = min(ns + BLOCK_N, I)
-        tile_diff = diff_vs_bf16[0, ns:ne].abs().max().item() if n_tokens > 0 else 0
-        print(f"  Tile {ni} ({ns}:{ne}): diff_vs_bf16_max={tile_diff:.6f}")
+    for bk in [32, 64, 128]:
+        triton_out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+        num_nt = triton.cdiv(I, BLOCK_N)
+        single_expert_kernel[(num_nt,)](
+            x.cuda(), w_gate.cuda(), w_up.cuda(), triton_out,
+            n_tokens, H, I,
+            x.stride(0), w_gate.stride(0), w_gate.stride(1),
+            w_up.stride(0), w_up.stride(1),
+            triton_out.stride(0),
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=bk,
+            USE_MANUAL_SIGMOID=use_manual,
+        )
+        
+        diff_vs_bf16 = (out_bf16.cuda().float() - triton_out.float()).abs()
+        method = "manual sigmoid" if use_manual else "tl.sigmoid"
+        print(f"\n{method} BK={bk}:")
+        print(f"  vs bf16 ref: max_diff={diff_vs_bf16.max().item():.6f}")
 
 # Also test: does accumulating twice give the same result?
 print("\n--- Reproducibility test ---")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug4.py
@@ -0,0 +1,123 @@
+"""Investigate non-determinism and compare against reference."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_up = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference (bf16 precision)
+gate_bf16 = x @ w_gate
+up_bf16 = x @ w_up
+out_bf16 = torch.nn.functional.silu(gate_bf16) * up_bf16
+
+# Reference (full fp32)
+gate_fp32 = x.float() @ w_gate.float()
+up_fp32 = x.float() @ w_up.float()
+out_fp32 = torch.nn.functional.silu(gate_fp32) * up_fp32
+
+print(f"bf16 ref vs fp32 ref: max={((out_bf16.float()-out_fp32).abs().max().item()):.6f}")
+
+# Triton kernel
+@triton.jit
+def test_kernel(
+    x_ptr, w_gate_ptr, w_up_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_gate_stride_h, w_gate_stride_i,
+    w_up_stride_h, w_up_stride_i,
+    out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    PREC: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            if PREC == 0:
+                # Standard bf16 dot
+                x_tile = tl.load(
+                    x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                    mask=m_mask[:, None] & k_mask[None, :],
+                )
+            else:
+                # Load as fp32
+                x_tile = tl.load(
+                    x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                    mask=m_mask[:, None] & k_mask[None, :],
+                ).to(tl.float32)
+
+            w_gate_tile = tl.load(
+                w_gate_ptr + (k_start + k_off)[:, None] * w_gate_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up_tile = tl.load(
+                w_up_ptr + (k_start + k_off)[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+
+            if PREC == 0:
+                acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate)
+                acc_up = tl.dot(x_tile, w_up_tile, acc_up)
+            elif PREC == 1:
+                acc_gate = tl.dot(x_tile, w_gate_tile.to(tl.float32), acc_gate)
+                acc_up = tl.dot(x_tile, w_up_tile.to(tl.float32), acc_up)
+            elif PREC == 2:
+                x_f32 = x_tile.to(tl.float32)
+                acc_gate = tl.dot(x_f32, w_gate_tile.to(tl.float32), acc_gate)
+                acc_up = tl.dot(x_f32, w_up_tile.to(tl.float32), acc_up)
+
+        gate_silu = tl.sigmoid(acc_gate) * acc_gate
+        result = (gate_silu * acc_up).to(tl.bfloat16)
+
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+num_nt = triton.cdiv(I, 128)
+
+for prec, name in [(0, "bf16 dot"), (1, "bf16@fp32 dot"), (2, "fp32 dot")]:
+    print(f"\n=== {name} ===")
+    diffs_vs_bf16 = []
+    diffs_vs_fp32 = []
+    for run in range(5):
+        out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+        test_kernel[(num_nt,)](
+            x.cuda(), w_gate.cuda(), w_up.cuda(), out,
+            n_tokens, H, I,
+            x.stride(0), w_gate.stride(0), w_gate.stride(1),
+            w_up.stride(0), w_up.stride(1),
+            out.stride(0),
+            BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+            PREC=prec,
+        )
+        dvb = (out_bf16.cuda().float() - out.float()).abs().max().item()
+        dvf = (out_fp32.cuda().float() - out.float()).abs().max().item()
+        diffs_vs_bf16.append(dvb)
+        diffs_vs_fp32.append(dvf)
+        print(f"  Run {run}: vs_bf16={dvb:.6f} vs_fp32={dvf:.6f}")
+
+    print(f"  --- vs bf16: min={min(diffs_vs_bf16):.6f} max={max(diffs_vs_bf16):.6f} range={max(diffs_vs_bf16)-min(diffs_vs_bf16):.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug5.py
@@ -0,0 +1,136 @@
+"""Compare GEMM intermediate values between Triton and reference."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_up = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference GEMMs (bf16 matmul)
+gate_ref = x @ w_gate   # (n_tokens, I) bf16
+up_ref = x @ w_up       # (n_tokens, I) bf16
+out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+
+# Reference GEMM (fp32, to see if bf16 truncation is the issue)
+gate_fp32 = x.float() @ w_gate.float()
+up_fp32 = x.float() @ w_up.float()
+
+print("GEMM differences:")
+print(f"  gate bf16 vs fp32: max={((gate_ref.float() - gate_fp32).abs().max().item()):.6f}")
+print(f"  up   bf16 vs fp32: max={((up_ref.float() - up_fp32).abs().max().item()):.6f}")
+
+# Triton kernel that computes JUST the GEMM and outputs bf16 gate and up
+@triton.jit
+def triton_gemm_kernel(
+    x_ptr, w_gate_ptr, w_up_ptr, gate_out_ptr, up_out_ptr,
+    n_tokens, H, I,
+    x_stride, w_gate_stride_h, w_gate_stride_i,
+    w_up_stride_h, w_up_stride_i,
+    gate_stride, up_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            x_tile = tl.load(
+                x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            w_gate_tile = tl.load(
+                w_gate_ptr + (k_start + k_off)[:, None] * w_gate_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up_tile = tl.load(
+                w_up_ptr + (k_start + k_off)[:, None] * w_up_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            acc_gate = tl.dot(x_tile, w_gate_tile, acc_gate)
+            acc_up = tl.dot(x_tile, w_up_tile, acc_up)
+
+        # Store gate and up separately as bf16
+        gate_base = gate_out_ptr + m_start * gate_stride
+        up_base = up_out_ptr + m_start * up_stride
+        tl.store(
+            gate_base + m_off[:, None] * gate_stride + n_off[None, :],
+            acc_gate.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+        tl.store(
+            up_base + m_off[:, None] * up_stride + n_off[None, :],
+            acc_up.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+num_nt = triton.cdiv(I, 128)
+gate_triton = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+up_triton = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+
+triton_gemm_kernel[(num_nt,)](
+    x.cuda(), w_gate.cuda(), w_up.cuda(), gate_triton, up_triton,
+    n_tokens, H, I,
+    x.stride(0), w_gate.stride(0), w_gate.stride(1),
+    w_up.stride(0), w_up.stride(1),
+    gate_triton.stride(0), up_triton.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+)
+
+print("\nTriton GEMM vs reference:")
+gate_diff = (gate_ref.cuda().float() - gate_triton.float()).abs()
+up_diff = (up_ref.cuda().float() - up_triton.float()).abs()
+print(f"  gate: max_diff={gate_diff.max().item():.6f} mean={gate_diff.mean().item():.6f}")
+print(f"  up:   max_diff={up_diff.max().item():.6f} mean={up_diff.mean().item():.6f}")
+
+# Now check the silu * up stage
+# Reference SwiGLU
+out_ref_swiglu = torch.nn.functional.silu(gate_ref) * up_ref
+
+# Triton SwiGLU (using Triton-computed gate and up with torch silu)
+out_triton_gemm_torch_silu = torch.nn.functional.silu(gate_triton.cpu()) * up_triton.cpu()
+diff_gemm = (out_ref_swiglu.float() - out_triton_gemm_torch_silu.float()).abs()
+print(f"\nSwiGLU: triton GEMM + torch silu vs ref:")
+print(f"  max_diff={diff_gemm.max().item():.6f}")
+
+# Compute silu and multiply in fp32 from triton's gate/up
+gate_triton_f32 = gate_triton.float()
+up_triton_f32 = up_triton.float()
+out_triton_fp32 = torch.nn.functional.silu(gate_triton_f32) * up_triton_f32
+diff_fp32_swiglu = (out_ref_swiglu.float().cuda() - out_triton_fp32.cuda()).abs()
+print(f"\nSwiGLU: triton GEMM fp32 silu vs ref:")
+print(f"  max_diff={diff_fp32_swiglu.max().item():.6f}")
+
+# The full pipeline: triton GEMM -> bf16 cutoff -> silu -> multiply -> bf16
+# vs reference
+print(f"\nFull pipeline triton vs ref:")
+out_triton_full = torch.nn.functional.silu(gate_triton.cpu()) * up_triton.cpu()
+diff_full = (out_ref_swiglu.float() - out_triton_full.float()).abs()
+print(f"  max_diff={diff_full.max().item():.6f}")
+
+# Also check: what if gate_triton is bf16 and up_triton is bf16, then silu and mul in fp32 on GPU?
+gate_triton_gpu = gate_triton.cuda()
+up_triton_gpu = up_triton.cuda()
+out_triton_gpu_silu = torch.nn.functional.silu(gate_triton_gpu.float()) * up_triton_gpu.float()
+diff_gpu_fp32 = (out_ref_swiglu.float().cuda() - out_triton_gpu_silu).abs()
+print(f"  (GPU fp32 silu): max_diff={diff_gpu_fp32.max().item():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug6.py
@@ -0,0 +1,111 @@
+"""Tile size sweep for GEMM accuracy."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Reference
+gate_ref = x @ w_gate  # cuBLAS bf16 GEMM
+
+@triton.jit
+def gemm_kernel(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k_start in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            x_tile = tl.load(
+                x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            w_tile = tl.load(
+                w_ptr + (k_start + k_off)[:, None] * w_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            acc = tl.dot(x_tile, w_tile, acc)
+
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+import itertools
+
+configs = []
+for bm, bn, bk, nw in itertools.product(
+    [32, 64, 128],
+    [64, 128, 256],
+    [16, 32, 64],
+    [4, 8],
+):
+    if bk < 16: continue
+    if bn > I: continue
+    configs.append((bm, bn, bk, nw))
+
+# Also test single-k-tile (load all of H at once if possible)
+for bm, bn, nw in itertools.product([16, 32], [64, 128], [4]):
+    bk = H  # full K dimension
+    configs.append((bm, bn, bk, nw))
+
+results = []
+for bm, bn, bk, nw in configs:
+    if bk == H:
+        try:
+            out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+            num_nt = triton.cdiv(I, bn)
+            gemm_kernel[(num_nt,)](
+                x.cuda(), w_gate.cuda(), out,
+                n_tokens, H, I,
+                x.stride(0), w_gate.stride(0), out.stride(0),
+                BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk,
+                num_warps=nw, num_stages=2,
+            )
+            diff = (gate_ref.cuda().float() - out.float()).abs().max().item()
+            results.append((bm, bn, bk, nw, diff, "OK"))
+            print(f"BM={bm:3d} BN={bn:3d} BK={bk:4d} NW={nw}: diff={diff:.6f} OK")
+        except Exception as e:
+            print(f"BM={bm:3d} BN={bn:3d} BK={bk:4d} NW={nw}: FAILED - {e}")
+    else:
+        out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+        num_nt = triton.cdiv(I, bn)
+        gemm_kernel[(num_nt,)](
+            x.cuda(), w_gate.cuda(), out,
+            n_tokens, H, I,
+            x.stride(0), w_gate.stride(0), out.stride(0),
+            BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk,
+            num_warps=nw, num_stages=3,
+        )
+        diff = (gate_ref.cuda().float() - out.float()).abs().max().item()
+        results.append((bm, bn, bk, nw, diff, "OK"))
+        print(f"BM={bm:3d} BN={bn:3d} BK={bk:4d} NW={nw}: diff={diff:.6f} OK")
+
+print("\n=== Best configs ===")
+for r in sorted(results, key=lambda x: x[4])[:10]:
+    print(f"BM={r[0]:3d} BN={r[1]:3d} BK={r[2]:4d} NW={r[3]}: diff={r[4]:.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug7.py
@@ -0,0 +1,158 @@
+"""Test if Triton bf16 GEMM differs fundamentally from cuBLAS."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# cuBLAS bf16 matmul
+ref = x @ w
+
+# Test tl.dot vs manual accumulation in fp32
+x_f32 = x.float().cuda()
+w_f32 = w.float().cuda()
+ref_f32 = x_f32 @ w_f32
+
+print(f"cuBLAS bf16 vs fp32 ref: max={((ref.float()-ref_f32.cpu()).abs().max().item()):.6f}")
+
+# Now: use Triton to load tiles and compute dot product manually (no tl.dot)
+# We'll compute the dot product by multiplying and summing in fp32
+@triton.jit
+def manual_dot_kernel(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k_start in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            # Convert to fp32 manually, multiply, accumulate
+            x_tile = tl.load(
+                x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            ).to(tl.float32)
+            w_tile = tl.load(
+                w_ptr + (k_start + k_off)[:, None] * w_stride_h + (n_start + n_off)[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            ).to(tl.float32)
+
+            # Manual dot product: element-wise multiply, then sum along K
+            prod = x_tile[:, :, None] * w_tile[None, :, :]
+            # This creates (BLOCK_M, BLOCK_K, BLOCK_N) and sums along K
+            # But this is insanely expensive in registers
+            
+            # Better: compute via reduction
+            # acc[m,n] += sum_k x[m,k] * w[k,n]
+            for k in range(BLOCK_K):
+                acc += x_tile[:, k:k+1] * w_tile[k:k+1, :]
+
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+# Actually, that manual approach is terrible. Let me test something simpler.
+# Test: what if we just do ONE K-tile and verify the partial sum is correct?
+
+@triton.jit
+def single_tile_gemm(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    TILE_K_START: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        k_start = TILE_K_START
+        k_off = tl.arange(0, BLOCK_K)
+        k_mask = k_start + k_off < H
+
+        x_tile = tl.load(
+            x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+            mask=m_mask[:, None] & k_mask[None, :],
+        )
+        w_tile = tl.load(
+            w_ptr + (k_start + k_off)[:, None] * w_stride_h + (n_start + n_off)[None, :],
+            mask=k_mask[:, None] & n_mask[None, :],
+        )
+        acc = tl.dot(x_tile, w_tile, acc)
+
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+# Test: compute partial sums for each K-tile and compare with cuBLAS partials
+BLOCK_K = 32
+BLOCK_M = 64
+BLOCK_N = 128
+num_k_tiles = H // BLOCK_K
+
+for k_tile in range(num_k_tiles):
+    k_start = k_tile * BLOCK_K
+    k_end = k_start + BLOCK_K
+    
+    # cuBLAS reference for this K-tile only
+    ref_partial = x[:, k_start:k_end] @ w[k_start:k_end, :]
+    
+    # Triton partial
+    out_triton = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    num_nt = triton.cdiv(I, BLOCK_N)
+    single_tile_gemm[(num_nt,)](
+        x.cuda(), w.cuda(), out_triton,
+        n_tokens, H, I,
+        x.stride(0), w.stride(0), out_triton.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        TILE_K_START=k_start,
+        num_warps=4, num_stages=2,
+    )
+    
+    diff = (ref_partial.cuda().float() - out_triton.float()).abs()
+    maxd = diff.max().item()
+    
+    if maxd > 0.01:
+        print(f"K-tile {k_tile} ({k_start}:{k_end}): max_diff={maxd:.6f} (LARGE!)")
+    else:
+        pass  # suppress OK tiles for brevity
+
+# Summarize: how many tiles have large errors?
+print(f"\n(n_tokens=512, checking all {num_k_tiles} K-tiles)")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug8.py
@@ -0,0 +1,93 @@
+"""Magnitude of Triton bf16 GEMM error: is it relative or absolute?"""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+ref = x @ w  # cuBLAS
+ref_vals = ref.float()
+print(f"Reference stats: mean={ref_vals.mean():.6f} std={ref_vals.std():.6f} min={ref_vals.min():.6f} max={ref_vals.max():.6f}")
+
+# Test: does tl.dot with a single K-tile give different result from cuBLAS?
+@triton.jit
+def single_tile_gemm(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, H, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    TILE_K_START: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        k_start = TILE_K_START
+        k_off = tl.arange(0, BLOCK_K)
+        k_mask = k_start + k_off < H
+        x_tile = tl.load(
+            x_base + m_off[:, None] * x_stride + (k_start + k_off)[None, :],
+            mask=m_mask[:, None] & k_mask[None, :],
+        )
+        w_tile = tl.load(
+            w_ptr + (k_start + k_off)[:, None] * w_stride_h + (n_start + n_off)[None, :],
+            mask=k_mask[:, None] & n_mask[None, :],
+        )
+        acc = tl.dot(x_tile, w_tile, acc)
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+BLOCK_M = 64
+BLOCK_N = 128
+BLOCK_K = 32
+num_nt = triton.cdiv(I, BLOCK_N)
+
+# Compare individual K-tiles
+for k_start in list(range(0, H, BLOCK_K))[:5]:  # first 5 tiles
+    ref_partial = x[:, k_start:k_start+BLOCK_K] @ w[k_start:k_start+BLOCK_K, :]
+    out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    single_tile_gemm[(num_nt,)](
+        x.cuda(), w.cuda(), out,
+        n_tokens, H, I,
+        x.stride(0), w.stride(0), out.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        TILE_K_START=k_start,
+        num_warps=4, num_stages=2,
+    )
+    diff = (ref_partial.cuda().float() - out.float()).abs()
+    ref_partial_stats = ref_partial.float()
+    print(f"K-tile {k_start:4d}: ref mean={ref_partial_stats.mean():.4f} std={ref_partial_stats.std():.4f}  diff max={diff.max().item():.6f}")
+
+# Test: does single-tile accumulate correctly?
+# Sum over ALL K-tiles using Triton
+out_all = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+for k_start in range(0, H, BLOCK_K):
+    partial = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    single_tile_gemm[(num_nt,)](
+        x.cuda(), w.cuda(), partial,
+        n_tokens, H, I,
+        x.stride(0), w.stride(0), partial.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        TILE_K_START=k_start,
+        num_warps=4, num_stages=2,
+    )
+    out_all = out_all + partial
+diff = (ref.cuda().float() - out_all.float()).abs()
+print(f"\nAccumulated all K-tiles separately: max_diff={diff.max().item():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug9.py
@@ -0,0 +1,101 @@
+"""Find which elements have large errors in single K-tile GEMM."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+H = 4096
+I = 1536
+n_tokens = 512
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+BLOCK_K = 32
+k_start = 0
+k_end = k_start + BLOCK_K
+
+# cuBLAS reference for first K-tile
+ref_partial = x[:, k_start:k_end] @ w[k_start:k_end, :]  # (512, 1536)
+
+# Manual fp32 computation
+x_f32 = x[:, k_start:k_end].float()  # (512, 32)
+w_f32 = w[k_start:k_end, :].float()  # (32, 1536)
+manual_fp32 = x_f32 @ w_f32  # (512, 1536)
+
+print("Ref (cuBLAS bf16) vs manual fp32:")
+diff_bf16_vs_fp32 = (ref_partial.float() - manual_fp32.cpu()).abs()
+print(f"  max={diff_bf16_vs_fp32.max().item():.6f} mean={diff_bf16_vs_fp32.mean().item():.6f}")
+
+# Now test Triton's tl.dot for the same tile
+@triton.jit
+def single_tile_gemm(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        k_off = tl.arange(0, BLOCK_K)
+        x_tile = tl.load(
+            x_base + m_off[:, None] * x_stride + k_off[None, :],
+            mask=m_mask[:, None],
+        )
+        w_tile = tl.load(
+            w_ptr + k_off[:, None] * w_stride_h + (n_start + n_off)[None, :],
+            mask=n_mask[None, :],
+        )
+        acc = tl.dot(x_tile, w_tile, acc)
+        out_base = out_ptr + m_start * out_stride
+        tl.store(
+            out_base + m_off[:, None] * out_stride + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+# Pass weights without E dimension
+single_tile_gemm[(triton.cdiv(I, 128),)](
+    x.cuda(), w[:BLOCK_K].contiguous().cuda(), out,
+    n_tokens, I,
+    x.stride(0), w[:BLOCK_K].stride(0), out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=2,
+)
+
+diff = (ref_partial.cuda().float() - out.float()).abs()
+print(f"\nTriton single tile vs cuBLAS:")
+print(f"  max={diff.max().item():.6f} mean={diff.mean().item():.6f}")
+
+# Find the worst elements
+flat_diff = diff.flatten()
+worst_idx = flat_diff.argsort(descending=True)[:5]
+for idx in worst_idx:
+    m = idx // I
+    n = idx % I
+    print(f"  Element ({m},{n}): ref={ref_partial[m,n].item():.4f} triton={out[m,n].item():.4f} diff={flat_diff[idx].item():.4f}")
+
+# Check: are certain columns systematically wrong?
+print(f"\nColumn-wise max diff:")
+col_means = diff.max(dim=0).values
+best_cols = col_means.argsort()[:5]
+worst_cols = col_means.argsort(descending=True)[:5]
+print(f"  Best columns: {[(c.item(), col_means[c].item()) for c in best_cols]}")
+print(f"  Worst columns: {[(c.item(), col_means[c].item()) for c in worst_cols]}")
+
+# Check max diff per row
+row_means = diff.max(dim=1).values
+best_rows = row_means.argsort()[:5]
+worst_rows = row_means.argsort(descending=True)[:5]
+print(f"  Best rows: {[(r.item(), row_means[r].item()) for r in best_rows]}")
+print(f"  Worst rows: {[(r.item(), row_means[r].item()) for r in worst_rows]}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug10.py
@@ -0,0 +1,89 @@
+"""Verify Triton dot product with exact element comparison."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+H, I, n_tokens = 4096, 1536, 512  
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Test specific problematic element
+m, n = 110, 70
+
+# Manual dot product for the first K-tile (k=0..31)
+ref_partial = x[:, :32] @ w[:32, :]
+manual = sum(x[m, k].float().item() * w[k, n].float().item() for k in range(32))
+print(f"Manual dot: sum_k x[110,k] * w[k,70] over k=0..31 = {manual:.6f}")
+print(f"cuBLAS: ref_partial[110,70] = {ref_partial[110,70].item():.6f}")
+print(f"Are they equal? {abs(manual - ref_partial[110,70].item()) < 1e-6}")
+
+# Now check what Triton loads
+BLOCK_M = 64
+BLOCK_N = 128
+BLOCK_K = 32
+
+@triton.jit
+def debug_load_kernel(
+    x_ptr, w_ptr, out_ptr,
+    n_tokens, I,
+    x_stride, w_stride_h, out_stride,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    num_n_tiles = tl.cdiv(I, BLOCK_N)
+    pid = tl.program_id(0)
+    n_tile_idx = pid
+    n_start = n_tile_idx * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+        x_base = x_ptr + m_start * x_stride
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        k_off = tl.arange(0, BLOCK_K)
+        
+        # Instead of tl.dot, manually compute one element
+        x_tile = tl.load(
+            x_base + m_off[:, None] * x_stride + k_off[None, :],
+            mask=m_mask[:, None],
+        )
+        w_tile = tl.load(
+            w_ptr + k_off[:, None] * w_stride_h + (n_start + n_off)[None, :],
+            mask=n_mask[None, :],
+        )
+        
+        # Store raw loaded values for debugging
+        out0 = out_ptr + m_start * out_stride
+        
+        # Store first row of x_tile and first column of w_tile
+        x_row0 = x_tile[0, :]
+        w_col0 = w_tile[:, 0]
+        
+        # Store x_row0 and w_col0 to separate output locations
+        for kk in range(BLOCK_K):
+            tl.store(out0 + kk * 2, x_row0[kk], mask=m_mask[0:1])  # store x values
+            tl.store(out0 + kk * 2 + 128, w_col0[kk], mask=m_mask[0:1])  # store w values
+
+out_debug = torch.zeros(256, dtype=torch.bfloat16, device='cuda')
+debug_load_kernel[(1,)](
+    x.cuda(), w[:32].contiguous().cuda(), out_debug,
+    n_tokens, I,
+    x.stride(0), w[:32].stride(0), 1,
+    BLOCK_M=BLOCK_M, BLOCK_N=128, BLOCK_K=BLOCK_K,
+    num_warps=4, num_stages=2,
+)
+
+# Check loaded values
+print(f"\nFirst 32 x values for row 0: {x[0, :32]}")
+print(f"Triton loaded x row 0 vals: {out_debug[0:64:2]}")
+
+print(f"\nFirst 32 w values for col 0: {w[:32, 0]}")
+print(f"Triton loaded w col 0 vals: {out_debug[1:65:2]}")
+
+# Check: do Triton's loaded values match Python?
+x_match = torch.allclose(x[0, :32].cuda(), out_debug[0:64:2], atol=1e-6, rtol=1e-6)
+w_match = torch.allclose(w[:32, 0].cuda(), out_debug[1:65:2], atol=1e-6, rtol=1e-6)
+print(f"x values match: {x_match}")
+print(f"w values match: {w_match}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug11.py
@@ -0,0 +1,109 @@
+"""Minimal reproducer: does tl.dot bf16 give wrong results on SM120?"""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+# Tiny test: 16x32 @ 32x16
+M, K, N = 16, 32, 16
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+b = torch.randn(K, N, dtype=torch.bfloat16) * 0.02
+
+ref = a @ b  # cuBLAS
+ref_fp32 = a.float() @ b.float()  # fp32 reference
+
+# Also check fp16
+a_fp16 = a.to(torch.float16)
+b_fp16 = b.to(torch.float16)
+ref_fp16 = a_fp16 @ b_fp16
+
+print(f"cuBLAS bf16 vs fp32: max={((ref.float()-ref_fp32).abs().max().item()):.6f}")
+print(f"cuBLAS fp16 vs fp32: max={((ref_fp16.float()-ref_fp32).abs().max().item()):.6f}")
+
+@triton.jit
+def tiny_gemm(a_ptr, b_ptr, out_ptr,
+              BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    off_m = pid * BLOCK_M
+    off_n = 0
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    
+    a = tl.load(a_ptr + m_off[:, None] * K + k_off[None, :])
+    b = tl.load(b_ptr + k_off[:, None] * N + n_off[None, :])
+    c = tl.dot(a, b)
+    tl.store(out_ptr + m_off[:, None] * N + n_off[None, :], c.to(tl.bfloat16))
+
+# bf16 test
+out_bf16 = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+tiny_gemm[(1,)](
+    a.cuda(), b.cuda(), out_bf16,
+    BLOCK_M=M, BLOCK_N=N, BLOCK_K=K,
+    num_warps=4
+)
+
+print(f"\nTriton bf16 vs cuBLAS bf16: max={((out_bf16.cpu().float() - ref.float()).abs().max().item()):.6f}")
+
+# Test with fp16 input to tl.dot
+@triton.jit
+def tiny_gemm_fp16(a_ptr, b_ptr, out_ptr,
+                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    
+    a = tl.load(a_ptr + m_off[:, None] * K + k_off[None, :])
+    b = tl.load(b_ptr + k_off[:, None] * N + n_off[None, :])
+    c = tl.dot(a.to(tl.float16), b.to(tl.float16))
+    tl.store(out_ptr + m_off[:, None] * N + n_off[None, :], c.to(tl.bfloat16))
+
+out_fp16 = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+tiny_gemm_fp16[(1,)](
+    a.cuda().to(torch.float16), b.cuda().to(torch.float16), out_fp16,
+    BLOCK_M=M, BLOCK_N=N, BLOCK_K=K,
+    num_warps=4
+)
+
+print(f"Triton fp16 vs cuBLAS fp16: max={((out_fp16.cpu().float() - ref_fp16.float()).abs().max().item()):.6f}")
+
+# Also try split into smaller K tiles
+@triton.jit
+def tiny_gemm_splitk(a_ptr, b_ptr, out_ptr,
+                     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k_start in range(0, K, BLOCK_K):
+        k_off = tl.arange(0, BLOCK_K)
+        a = tl.load(a_ptr + m_off[:, None] * K + (k_start + k_off)[None, :])
+        b = tl.load(b_ptr + (k_start + k_off)[:, None] * N + n_off[None, :])
+        acc = tl.dot(a, b, acc)
+    
+    tl.store(out_ptr + m_off[:, None] * N + n_off[None, :], acc.to(tl.bfloat16))
+
+out_split = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+tiny_gemm_splitk[(1,)](
+    a.cuda(), b.cuda(), out_split,
+    BLOCK_M=M, BLOCK_N=N, BLOCK_K=16,
+    num_warps=4
+)
+print(f"Triton bf16 splitk vs cuBLAS bf16: max={((out_split.cpu().float() - ref.float()).abs().max().item()):.6f}")
+
+# Check if the issue is with larger K
+K2 = 256
+a2 = torch.randn(M, K2, dtype=torch.bfloat16) * 0.1
+b2 = torch.randn(K2, N, dtype=torch.bfloat16) * 0.02
+ref2 = a2 @ b2
+
+out2 = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+tiny_gemm[(1,)](
+    a2.cuda(), b2.cuda(), out2,
+    BLOCK_M=M, BLOCK_N=N, BLOCK_K=K2,
+    num_warps=4
+)
+print(f"\nK=256 triton vs cuBLAS: max={((out2.cpu().float() - ref2.float()).abs().max().item()):.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug12.py
@@ -0,0 +1,67 @@
+"""Try different approaches to get correct bf16 matmul on SM120."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K, N = 64, 32, 64
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+b = torch.randn(K, N, dtype=torch.bfloat16) * 0.02
+
+ref = a @ b
+ref_fp32 = a.float() @ b.float()
+
+print(f"cuBLAS bf16 vs fp32: max={((ref.float()-ref_fp32).abs().max().item()):.6f}")
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr,
+                M, K, N, a_stride, b_stride, out_stride,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+                DTYPE: tl.constexpr):
+    pid = tl.program_id(0)
+    n_tile = pid
+    n_start = n_tile * BLOCK_N
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    
+    for m_start in range(0, M, BLOCK_M):
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k_start in range(0, K, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            a_load = tl.load(a_ptr + (m_start + m_off)[:, None] * a_stride + (k_start + k_off)[None, :],
+                           mask=(m_start + m_off)[:, None] < M)
+            b_load = tl.load(b_ptr + (k_start + k_off)[:, None] * b_stride + (n_start + n_off)[None, :],
+                           mask=(n_start + n_off)[None, :] < N)
+            
+            if DTYPE == 0:
+                # bf16 dot
+                acc = tl.dot(a_load, b_load, acc)
+            elif DTYPE == 1:
+                # fp16 dot
+                acc = tl.dot(a_load.to(tl.float16), b_load.to(tl.float16), acc)
+            elif DTYPE == 2:
+                # fp32 dot (tf32 on Blackwell)
+                acc = tl.dot(a_load.to(tl.float32), b_load.to(tl.float32), acc)
+        
+        tl.store(out_ptr + (m_start + m_off)[:, None] * out_stride + n_off[None, :],
+                acc.to(tl.bfloat16),
+                mask=(m_start + m_off)[:, None] < M)
+
+num_nt = triton.cdiv(N, 64)
+
+for dtype_name, dt in [("bf16", 0), ("fp16", 1), ("fp32", 2)]:
+    out = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(num_nt,)](
+        a.cuda(), b.cuda(), out,
+        M, K, N, a.stride(0), b.stride(0), out.stride(0),
+        BLOCK_M=M, BLOCK_N=64, BLOCK_K=K,
+        DTYPE=dt,
+        num_warps=4, num_stages=2,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    maxd = diff.max().item()
+    meand = diff.mean().item()
+    # Check sign flips
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    print(f"{dtype_name}: max_diff={maxd:.6f} mean={meand:.6f} sign_flips={sign_flips}/{M*N}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug13.py
@@ -0,0 +1,34 @@
+"""Check if error is specific to N=128 tile size."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K = 64, 32
+
+for N in [16, 32, 64, 128, 256]:
+    a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+    b = torch.randn(K, N, dtype=torch.bfloat16) * 0.02
+    ref = a @ b
+
+    @triton.jit
+    def gemm_kernel(a_ptr, b_ptr, out_ptr,
+                    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+        m_off = tl.arange(0, BLOCK_M)
+        n_off = tl.arange(0, BLOCK_N)
+        k_off = tl.arange(0, BLOCK_K)
+        a = tl.load(a_ptr + m_off[:, None] * K + k_off[None, :])
+        b = tl.load(b_ptr + k_off[:, None] * N + n_off[None, :])
+        c = tl.dot(a, b)
+        tl.store(out_ptr + m_off[:, None] * N + n_off[None, :], c.to(tl.bfloat16))
+
+    out = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(1,)](
+        a.cuda(), b.cuda(), out,
+        BLOCK_M=M, BLOCK_N=N, BLOCK_K=K,
+        num_warps=4,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    print(f"N={N:3d}: max_diff={diff.max().item():.6f} mean={diff.mean().item():.6f} sign_flips={sign_flips}/{M*N}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug14.py
@@ -0,0 +1,35 @@
+"""Test single K-tile with different N sizes. Key: is N=128 buggy?"""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K = 64, 32
+
+for N in [64, 128, 256]:
+    a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+    b = torch.randn(K, N, dtype=torch.bfloat16) * 0.02
+    ref = a @ b
+
+    @triton.jit
+    def gemm_kernel(a_ptr, b_ptr, out_ptr, M, K, N, a_stride, b_stride, out_stride,
+                    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+        m_off = tl.arange(0, BLOCK_M)
+        n_off = tl.arange(0, BLOCK_N)
+        k_off = tl.arange(0, BLOCK_K)
+        a = tl.load(a_ptr + m_off[:, None] * a_stride + k_off[None, :])
+        b = tl.load(b_ptr + k_off[:, None] * b_stride + n_off[None, :])
+        c = tl.dot(a, b)
+        tl.store(out_ptr + m_off[:, None] * out_stride + n_off[None, :], c.to(tl.bfloat16))
+
+    out = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(1,)](
+        a.cuda(), b.cuda(), out, M, K, N,
+        a.stride(0), b.stride(0), out.stride(0),
+        BLOCK_M=M, BLOCK_N=N, BLOCK_K=K,
+        num_warps=4,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    print(f"N={N:3d}: max_diff={diff.max().item():.6f} mean={diff.mean().item():.6f} sign_flips={sign_flips}/{M*N}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug15.py
@@ -0,0 +1,68 @@
+"""Focused: does tl.dot with BLOCK_N=128 give wrong results on SM120?"""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+# Test 1: (64, 32) @ (32, 128) - the single K-tile case
+M, K, N1 = 64, 32, 64
+N2 = 128
+
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+
+b64 = torch.randn(K, N1, dtype=torch.bfloat16) * 0.02
+b128 = torch.randn(K, N2, dtype=torch.bfloat16) * 0.02
+
+ref64 = a @ b64
+ref128 = a @ b128
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b = tl.load(b_ptr + k_off[:, None] * BLOCK_N + n_off[None, :])
+    c = tl.dot(a, b)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+for name, N in [("N=64", N1), ("N=128", N2)]:
+    b = b64 if N == N1 else b128
+    ref = ref64 if N == N1 else ref128
+    out = torch.zeros(M, N, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(1,)](
+        a.cuda(), b.cuda(), out,
+        BLOCK_M=M, BLOCK_N=N, BLOCK_K=K,
+        num_warps=4,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    print(f"{name}: max_diff={diff.max().item():.6f} mean={diff.mean().item():.6f} sign_flips={sign_flips}/{M*N}")
+
+# Now test: what about (4, 32) @ (32, 128)?
+M2 = 4
+a_small = a[:M2]
+
+@triton.jit
+def gemm_kernel_small(a_ptr, b_ptr, out_ptr,
+                      BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    c = tl.dot(
+        tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :]),
+        tl.load(b_ptr + k_off[:, None] * BLOCK_N + n_off[None, :]),
+    )
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+ref_small = a_small @ b128
+out_small = torch.zeros(M2, N2, dtype=torch.bfloat16, device='cuda')
+gemm_kernel_small[(1,)](
+    a_small.cuda(), b128.cuda(), out_small,
+    BLOCK_M=M2, BLOCK_N=N2, BLOCK_K=K,
+    num_warps=4,
+)
+diff_small = (ref_small.cuda().float() - out_small.float()).abs()
+print(f"M=4,N=128: max_diff={diff_small.max().item():.6f} mean={diff_small.mean().item():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
@@ -0,0 +1,136 @@
+"""CLEAN test: one K-tile correctly, compare tl.dot vs cuBLAS."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I = 128, 4096, 1536
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 32
+
+# Test ONE K-tile: columns k_start..k_start+BLOCK_K
+k_start = 0
+x_slice = x[:, k_start:k_start+BLOCK_K].contiguous()  # (n_tokens, 32)
+w_slice = w[k_start:k_start+BLOCK_K, :].contiguous()   # (32, 1536)
+
+# cuBLAS reference
+ref_partial = x_slice @ w_slice
+
+# Clean Triton kernel, explicit strides
+@triton.jit
+def clean_gemm(
+    x_ptr, w_ptr, out_ptr,
+    x_stride_0, w_stride_0,
+    out_stride_0,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    n_tile = pid
+    n_start = n_tile * BLOCK_N
+
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+
+    # Load x[m, k] where k is relative to the tile start
+    # x shape: (n_tokens, BLOCK_K) with stride (BLOCK_K, 1)
+    # So offset = m * BLOCK_K + k
+    for m_base in range(0, n_tokens, BLOCK_M):
+        a = tl.load(
+            x_ptr + (m_base + m_off)[:, None] * x_stride_0 + k_off[None, :],
+            mask=(m_base + m_off)[:, None] < n_tokens,
+        )
+
+        # Load w[k, n] where k is relative to the tile start
+        # w shape: (BLOCK_K, I) with stride (I, 1)
+        # So offset = k * I + n
+        b = tl.load(
+            w_ptr + k_off[:, None] * w_stride_0 + (n_start + n_off)[None, :],
+            mask=(n_start + n_off)[None, :] < I,
+        )
+
+        c = tl.dot(a, b)
+
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_stride_0 + n_off[None, :],
+            c.to(tl.bfloat16),
+            mask=(m_base + m_off)[:, None] < n_tokens,
+        )
+
+num_nt = triton.cdiv(I, BLOCK_N)
+out_triton = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+clean_gemm[(num_nt,)](
+    x_slice.cuda(), w_slice.cuda(), out_triton,
+    x_slice.stride(0), w_slice.stride(0),
+    out_triton.stride(0),
+    BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+    num_warps=4, num_stages=2,
+)
+
+diff = (ref_partial.cuda().float() - out_triton.float()).abs()
+print(f"Single K-tile (k_start=0):")
+print(f"  max_diff={diff.max().item():.6f}  mean={diff.mean().item():.6f}")
+sign_flips = ((ref_partial.cuda() * out_triton) < 0).sum().item()
+print(f"  sign_flips={sign_flips}/{n_tokens*I}")
+
+# Also test: full GEMM with all K tiles, CORRECTLY advancing k_start
+@triton.jit
+def full_gemm(
+    x_ptr, w_ptr, out_ptr,
+    H, I,
+    x_stride_0, w_stride_0, w_stride_1,
+    out_stride_0,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    n_tile = pid
+    n_start = n_tile * BLOCK_N
+
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+
+    for m_base in range(0, n_tokens, BLOCK_M):
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_base + k_off < H
+
+            a = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_stride_0 + (k_base + k_off)[None, :],
+                mask=(m_base + m_off)[:, None] < n_tokens,
+            )
+            b = tl.load(
+                w_ptr + (k_base + k_off)[:, None] * w_stride_0 + (n_start + n_off)[None, :],
+                mask=(k_base + k_off)[:, None] < H,
+            )
+
+            acc = tl.dot(a, b, acc)
+
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_stride_0 + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=(m_base + m_off)[:, None] < n_tokens,
+        )
+
+# Full GEMM reference
+ref_full = x @ w
+
+out_full = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+full_gemm[(num_nt,)](
+    x.cuda(), w.cuda(), out_full,
+    H, I,
+    x.stride(0), w.stride(0), w.stride(1),
+    out_full.stride(0),
+    BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+    num_warps=4, num_stages=3,
+)
+
+diff_full = (ref_full.cuda().float() - out_full.float()).abs()
+print(f"\nFull GEMM (all K tiles, correct offsets):")
+print(f"  max_diff={diff_full.max().item():.6f}  mean={diff_full.mean().item():.6f}")
+sign_flips_full = ((ref_full.cuda() * out_full) < 0).sum().item()
+print(f"  sign_flips={sign_flips_full}/{n_tokens*I}")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
@@ -23,6 +23,7 @@
 @triton.jit
 def clean_gemm(
     x_ptr, w_ptr, out_ptr,
+    n_tokens, I,
     x_stride_0, w_stride_0,
     out_stride_0,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
@@ -35,25 +36,16 @@
     n_off = tl.arange(0, BLOCK_N)
     k_off = tl.arange(0, BLOCK_K)
 
-    # Load x[m, k] where k is relative to the tile start
-    # x shape: (n_tokens, BLOCK_K) with stride (BLOCK_K, 1)
-    # So offset = m * BLOCK_K + k
     for m_base in range(0, n_tokens, BLOCK_M):
         a = tl.load(
             x_ptr + (m_base + m_off)[:, None] * x_stride_0 + k_off[None, :],
             mask=(m_base + m_off)[:, None] < n_tokens,
         )
-
-        # Load w[k, n] where k is relative to the tile start
-        # w shape: (BLOCK_K, I) with stride (I, 1)
-        # So offset = k * I + n
         b = tl.load(
             w_ptr + k_off[:, None] * w_stride_0 + (n_start + n_off)[None, :],
             mask=(n_start + n_off)[None, :] < I,
         )
-
         c = tl.dot(a, b)
-
         tl.store(
             out_ptr + (m_base + m_off)[:, None] * out_stride_0 + n_off[None, :],
             c.to(tl.bfloat16),

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
@@ -72,8 +72,8 @@
 @triton.jit
 def full_gemm(
     x_ptr, w_ptr, out_ptr,
-    H, I,
-    x_stride_0, w_stride_0, w_stride_1,
+    n_tokens, H, I,
+    x_stride_0, w_stride_0,
     out_stride_0,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
@@ -89,7 +89,6 @@
 
         for k_base in range(0, H, BLOCK_K):
             k_off = tl.arange(0, BLOCK_K)
-            k_mask = k_base + k_off < H
 
             a = tl.load(
                 x_ptr + (m_base + m_off)[:, None] * x_stride_0 + (k_base + k_off)[None, :],

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
@@ -56,6 +56,7 @@
 out_triton = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
 clean_gemm[(num_nt,)](
     x_slice.cuda(), w_slice.cuda(), out_triton,
+    n_tokens, I,
     x_slice.stride(0), w_slice.stride(0),
     out_triton.stride(0),
     BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_clean.py
@@ -114,8 +114,8 @@
 out_full = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
 full_gemm[(num_nt,)](
     x.cuda(), w.cuda(), out_full,
-    H, I,
-    x.stride(0), w.stride(0), w.stride(1),
+    n_tokens, H, I,
+    x.stride(0), w.stride(0),
     out_full.stride(0),
     BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
     num_warps=4, num_stages=3,

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_workaround.py
@@ -0,0 +1,85 @@
+"""Test workarounds for bf16 tl.dot bug on SM120."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I = 128, 4096, 1536
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_up = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Full reference
+gate_ref = x @ w_gate
+up_ref = x @ w_up
+out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+
+BLOCK_M, BLOCK_N, BLOCK_K = 64, 128, 32
+num_nt = triton.cdiv(I, BLOCK_N)
+
+@triton.jit
+def full_kernel(
+    x_ptr, wg_ptr, wu_ptr, out_ptr,
+    n_tokens, H, I,
+    x_s0, wg_s0, wu_s0, out_s0,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    USE_FP16: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    nt = pid
+    n_start = nt * BLOCK_N
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+
+    for m_base in range(0, n_tokens, BLOCK_M):
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+
+            a_raw = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=(m_base + m_off)[:, None] < n_tokens,
+            )
+            wg_raw = tl.load(
+                wg_ptr + (k_base + k_off)[:, None] * wg_s0 + (n_start + n_off)[None, :],
+                mask=(k_base + k_off)[:, None] < H,
+            )
+            wu_raw = tl.load(
+                wu_ptr + (k_base + k_off)[:, None] * wu_s0 + (n_start + n_off)[None, :],
+                mask=(k_base + k_off)[:, None] < H,
+            )
+
+            if USE_FP16:
+                acc_g = tl.dot(a_raw.to(tl.float16), wg_raw.to(tl.float16), acc_g)
+                acc_u = tl.dot(a_raw.to(tl.float16), wu_raw.to(tl.float16), acc_u)
+            else:
+                a_f32 = a_raw.to(tl.float32)
+                acc_g = tl.dot(a_f32, wg_raw.to(tl.float32), acc_g)
+                acc_u = tl.dot(a_f32, wu_raw.to(tl.float32), acc_u)
+
+        gate_silu = tl.sigmoid(acc_g) * acc_g
+        result = (gate_silu * acc_u).to(tl.bfloat16)
+
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            result,
+            mask=(m_base + m_off)[:, None] < n_tokens,
+        )
+
+for name, use_fp16 in [("fp16_gemm", True), ("fp32_gemm", False)]:
+    out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    full_kernel[(num_nt,)](
+        x.cuda(), w_gate.cuda(), w_up.cuda(), out,
+        n_tokens, H, I,
+        x.stride(0), w_gate.stride(0), w_up.stride(0), out.stride(0),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        USE_FP16=use_fp16,
+        num_warps=4, num_stages=3,
+    )
+    diff = (out_ref.cuda().float() - out.float()).abs()
+    print(f"{name}: max_diff={diff.max().item():.6f}  mean={diff.mean().item():.6f}")
+    sign_flips = ((out_ref.cuda() * out) < 0).sum().item()
+    print(f"  sign_flips={sign_flips}/{n_tokens*I}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_load.py
@@ -0,0 +1,80 @@
+"""Test: does Triton load/store bf16 data correctly?"""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+@triton.jit
+def load_store_kernel(in_ptr, out_ptr, N):
+    pid = tl.program_id(0)
+    off = pid * 128 + tl.arange(0, 128)
+    mask = off < N
+    x = tl.load(in_ptr + off, mask=mask)
+    tl.store(out_ptr + off, x, mask=mask)
+
+x = torch.randn(1024, dtype=torch.bfloat16) * 0.1
+out = torch.zeros(1024, dtype=torch.bfloat16, device='cuda')
+load_store_kernel[(8,)](
+    x.cuda(), out, 1024,
+    num_warps=4,
+)
+print(f"Load/store test: max_diff={(x.cuda() - out).abs().max().item():.6f}")
+
+# Now test: does tl.dot with tiny amounts work?
+@triton.jit
+def tiny_dot(a_ptr, b_ptr, out_ptr,
+             BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b = tl.load(b_ptr + k_off[:, None] * BLOCK_N + n_off[None, :])
+    c = tl.dot(a, b)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+# Test with very small values where bf16 has limited precision
+for scale in [0.1, 1.0, 10.0]:
+    a = torch.randn(16, 16, dtype=torch.bfloat16) * scale
+    b = torch.randn(16, 16, dtype=torch.bfloat16) * scale * 0.1
+    ref = a @ b
+    
+    out = torch.zeros(16, 16, dtype=torch.bfloat16, device='cuda')
+    tiny_dot[(1,)](
+        a.cuda(), b.cuda(), out,
+        BLOCK_M=16, BLOCK_N=16, BLOCK_K=16,
+        num_warps=4,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    print(f"16x16 dot (scale={scale}): max_diff={diff.max().item():.6f} sign_flips={sign_flips}/256")
+    
+# Now test: does the issue happen with any specific row count?
+for M in [1, 2, 4, 8, 16, 32, 64]:
+    a = torch.randn(M, 32, dtype=torch.bfloat16) * 0.1
+    b = torch.randn(32, 128, dtype=torch.bfloat16) * 0.02
+    ref = a @ b
+    
+    @triton.jit
+    def gemm_kernel(a_ptr, b_ptr, out_ptr, M, N,
+                    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+        m_off = tl.arange(0, BLOCK_M)
+        n_off = tl.arange(0, BLOCK_N)
+        k_off = tl.arange(0, BLOCK_K)
+        a = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :], mask=m_off[:, None] < M)
+        b = tl.load(b_ptr + k_off[:, None] * N + n_off[None, :])
+        c = tl.dot(a, b)
+        tl.store(out_ptr + m_off[:, None] * N + n_off[None, :], c.to(tl.bfloat16), mask=m_off[:, None] < M)
+    
+    out = torch.zeros(M, 128, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(1,)](
+        a.cuda(), b.cuda(), out, M, 128,
+        BLOCK_M=M, BLOCK_N=128, BLOCK_K=32,
+        num_warps=4,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sign_flips = ((ref.cuda() * out) < 0).sum().item()
+    if sign_flips > 0 or diff.max() > 0.001:
+        print(f"M={M:2d} N=128: max_diff={diff.max():.6f} sign_flips={sign_flips}/{M*128}")
+    else:
+        print(f"M={M:2d} N=128: OK")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_scale.py
@@ -0,0 +1,66 @@
+"""Test GEMM at different scales to find where errors start."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, N, a_s0, b_s0, out_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    n_start = nt * BLOCK_N
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+
+    a = tl.load(a_ptr + m_off[:, None] * a_s0 + k_off[None, :])
+    b = tl.load(b_ptr + k_off[:, None] * b_s0 + (n_start + n_off)[None, :])
+    c = tl.dot(a, b)
+    tl.store(out_ptr + m_off[:, None] * out_s0 + n_off[None, :], c.to(tl.bfloat16))
+
+# Test 1: fixed K=32, N=128, vary M up to 64
+print("Fixed (64,32) x (32,128):")
+a = torch.randn(64, 32, dtype=torch.bfloat16) * 0.1
+b = torch.randn(32, 128, dtype=torch.bfloat16) * 0.02
+ref = a @ b
+out = torch.zeros(64, 128, dtype=torch.bfloat16, device='cuda')
+gemm_kernel[(1,)](
+    a.cuda(), b.cuda(), out, 128, a.stride(0), b.stride(0), out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32, num_warps=4,
+)
+diff = (ref.cuda().float() - out.float()).abs()
+sf = ((ref.cuda() * out) < 0).sum().item()
+print(f"  max_diff={diff.max():.6f} sign_flips={sf}")
+
+# Test 2: vary K, keep M,N fixed
+for K in [32, 64, 128, 256, 512, 1024]:
+    a = torch.randn(64, K, dtype=torch.bfloat16) * 0.1
+    b = torch.randn(K, 128, dtype=torch.bfloat16) * 0.02
+    ref = a @ b
+    
+    @triton.jit
+    def gemm_split(a_ptr, b_ptr, out_ptr, a_s0, b_s0, out_s0,
+                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+        pid = tl.program_id(0)
+        nt = pid
+        n_start = nt * BLOCK_N
+        m_off = tl.arange(0, BLOCK_M)
+        n_off = tl.arange(0, BLOCK_N)
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k_base in range(0, K, BLOCK_K):
+            k_off2 = tl.arange(0, BLOCK_K)
+            a2 = tl.load(a_ptr + m_off[:, None] * a_s0 + (k_base + k_off2)[None, :])
+            b2 = tl.load(b_ptr + (k_base + k_off2)[:, None] * b_s0 + (n_start + n_off)[None, :])
+            acc = tl.dot(a2, b2, acc)
+        tl.store(out_ptr + m_off[:, None] * out_s0 + n_off[None, :], acc.to(tl.bfloat16))
+    
+    out = torch.zeros(64, 128, dtype=torch.bfloat16, device='cuda')
+    gemm_split[(1,)](
+        a.cuda(), b.cuda(), out, a.stride(0), b.stride(0), out.stride(0),
+        BLOCK_M=64, BLOCK_N=128, BLOCK_K=32, num_warps=4, num_stages=3,
+    )
+    diff = (ref.cuda().float() - out.float()).abs()
+    sf = ((ref.cuda() * out) < 0).sum().item()
+    print(f"K={K:4d} ({(K+31)//32} tiles): max_diff={diff.max():.6f} sign_flips={sf}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_stride.py
@@ -0,0 +1,89 @@
+"""Test if stride width in the N dimension causes tl.dot errors."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K = 64, 32
+N_tile = 128
+
+# Case 1: w has shape (32, 128) - contiguous, stride(0)=128
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+b1 = torch.randn(K, N_tile, dtype=torch.bfloat16) * 0.02
+ref1 = a @ b1
+
+# Case 2: w has shape (32, 1536) - stride(0)=1536, we access first 128 columns
+w_big = torch.randn(K, 1536, dtype=torch.bfloat16) * 0.02
+b2 = w_big[:, :128]  # stride(0) = 1536
+ref2 = a @ b2  # cuBLAS computes the same dot product
+
+print(f"Ref1 vs Ref2 (should be same): max={((ref1-ref2).abs().max().item()):.6f}")
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, b_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + n_off[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+for name, b, b_s0 in [("stride=128", b1, b1.stride(0)), ("stride=1536", b2, b2.stride(0))]:
+    out = torch.zeros(M, N_tile, dtype=torch.bfloat16, device='cuda')
+    gemm_kernel[(1,)](
+        a.cuda(), b.cuda(), out, b_s0,
+        BLOCK_M=M, BLOCK_N=N_tile, BLOCK_K=K,
+        num_warps=4,
+    )
+    ref = a @ b  # cuBLAS with whatever shape
+    diff = (ref.cuda().float() - out.float()).abs()
+    sf = ((ref.cuda() * out) < 0).sum().item()
+    print(f"{name}: max_diff={diff.max():.6f} sign_flips={sf}/{M*N_tile}")
+
+# Test: what if the B tensor has a large stride in the K dimension too?
+# Use a big tensor with K rows but skip rows
+w_wide = torch.randn(4096, 1536, dtype=torch.bfloat16) * 0.02
+b3 = w_wide[:K, :N_tile]  # stride(0)=1536, stride(1)=1
+ref3 = a @ b3
+out3 = torch.zeros(M, N_tile, dtype=torch.bfloat16, device='cuda')
+gemm_kernel[(1,)](
+    a.cuda(), b3.cuda(), out3, b3.stride(0),
+    BLOCK_M=M, BLOCK_N=N_tile, BLOCK_K=K,
+    num_warps=4,
+)
+diff3 = (ref3.cuda().float() - out3.float()).abs()
+sf3 = ((ref3.cuda() * out3) < 0).sum().item()
+print(f"big_view (K=32,N=128 from (4096,1536)): max_diff={diff3.max():.6f} sign_flips={sf3}/{M*N_tile}")
+
+# Test: do sign flips appear with multiple N-tiles?
+out_multi = torch.zeros(M, 1536, dtype=torch.bfloat16, device='cuda')
+num_nt = triton.cdiv(1536, N_tile)
+
+@triton.jit
+def gemm_multi(a_ptr, b_ptr, out_ptr, b_s0, I,
+               BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    n_start = pid * BLOCK_N
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + (n_start + n_off)[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * I + n_off[None, :], c.to(tl.bfloat16))
+
+a_big = torch.randn(64, 32, dtype=torch.bfloat16) * 0.1
+w_big2 = torch.randn(32, 1536, dtype=torch.bfloat16) * 0.02
+ref_big = a_big @ w_big2
+
+gemm_multi[(num_nt,)](
+    a_big.cuda(), w_big2.cuda(), out_multi, w_big2.stride(0), 1536,
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4,
+)
+diff_big = (ref_big.cuda().float() - out_multi.float()).abs()
+sf_big = ((ref_big.cuda() * out_multi) < 0).sum().item()
+print(f"\nMulti N-tile (64,32) x (32,1536): max_diff={diff_big.max():.6f} sign_flips={sf_big}/{64*1536}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_tmafix.py
@@ -0,0 +1,41 @@
+"""Test TMA workaround: try small num_stages or different approach."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K, N_tile = 64, 32, 128
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+w_big = torch.randn(K, 1536, dtype=torch.bfloat16) * 0.02
+b = w_big[:, :128]  # stride(0)=1536
+ref = a @ b
+
+print(f"b shape={b.shape}, stride={b.stride()}")
+print(f"cuBLAS ref: mean={ref.float().mean():.6f}")
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, b_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + n_off[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+# Test with different num_stages
+for ns in [0, 1, 2, 3, 4]:
+    out = torch.zeros(M, N_tile, dtype=torch.bfloat16, device='cuda')
+    try:
+        gemm_kernel[(1,)](
+            a.cuda(), b.cuda(), out, b.stride(0),
+            BLOCK_M=M, BLOCK_N=N_tile, BLOCK_K=K,
+            num_warps=4, num_stages=ns,
+        )
+        diff = (ref.cuda().float() - out.float()).abs()
+        sf = ((ref.cuda() * out) < 0).sum().item()
+        print(f"num_stages={ns}: max_diff={diff.max():.6f} sign_flips={sf}/{M*N_tile}")
+    except Exception as e:
+        print(f"num_stages={ns}: FAILED - {e}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_stride2.py
@@ -0,0 +1,37 @@
+"""Find which strides work and which don't for tl.dot."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+M, K, BN = 64, 32, 128
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, b_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + n_off[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+for stride in [128, 129, 130, 132, 144, 160, 192, 256, 512, 1024, 1536, 2048]:
+    b = torch.randn(K, stride, dtype=torch.bfloat16) * 0.02
+    # We use first BN=128 columns
+    ref = a @ b[:, :BN]
+    out = torch.zeros(M, BN, dtype=torch.bfloat16, device='cuda')
+    try:
+        gemm_kernel[(1,)](
+            a.cuda(), b.cuda(), out, b.stride(0),
+            BLOCK_M=M, BLOCK_N=BN, BLOCK_K=K,
+            num_warps=4,
+        )
+        diff = (ref.cuda().float() - out.float()).abs()
+        sf = ((ref.cuda() * out) < 0).sum().item()
+        flag = "SIGNFLIP" if sf > 0 else "OK"
+        print(f"stride={stride:4d}: max_diff={diff.max():.6f} {flag} sf={sf}")
+    except Exception as e:
+        print(f"stride={stride:4d}: FAILED - {e}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_notma.py
@@ -0,0 +1,45 @@
+"""Test if disabling TMA fixes the stride issue."""
+import os
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+M, K, BN = 64, 32, 128
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+w_big = torch.randn(K, 1536, dtype=torch.bfloat16) * 0.02
+b = w_big[:, :128]  # view with stride 1536
+ref = a @ b
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, b_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + n_off[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+# Test with TRITON_DISABLE_TMA
+for disable_tma in [None, "1"]:
+    if disable_tma:
+        os.environ["TRITON_DISABLE_TMA"] = "1"
+    else:
+        os.environ.pop("TRITON_DISABLE_TMA", None)
+    
+    out = torch.zeros(M, BN, dtype=torch.bfloat16, device='cuda')
+    try:
+        gemm_kernel[(1,)](
+            a.cuda(), b.cuda(), out, b.stride(0),
+            BLOCK_M=M, BLOCK_N=BN, BLOCK_K=K,
+            num_warps=4, num_stages=2,
+        )
+        diff = (ref.cuda().float() - out.float()).abs()
+        sf = ((ref.cuda() * out) < 0).sum().item()
+        flag = "SIGNFLIP" if sf > 0 else "OK"
+        print(f"DISABLE_TMA={disable_tma}: max_diff={diff.max():.6f} {flag} sf={sf}")
+    except Exception as e:
+        print(f"DISABLE_TMA={disable_tma}: FAILED {e}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_notma2.py
@@ -0,0 +1,33 @@
+"""Test if disabling TMA fixes the stride issue. Run with env var."""
+import os
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+M, K, BN = 64, 32, 128
+a = torch.randn(M, K, dtype=torch.bfloat16) * 0.1
+w_big = torch.randn(K, 1536, dtype=torch.bfloat16) * 0.02
+b = w_big[:, :128]  # view with stride 1536
+ref = a @ b
+
+@triton.jit
+def gemm_kernel(a_ptr, b_ptr, out_ptr, b_s0,
+                BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    b_t = tl.load(b_ptr + k_off[:, None] * b_s0 + n_off[None, :])
+    c = tl.dot(a_t, b_t)
+    tl.store(out_ptr + m_off[:, None] * BLOCK_N + n_off[None, :], c.to(tl.bfloat16))
+
+out = torch.zeros(M, BN, dtype=torch.bfloat16, device='cuda')
+gemm_kernel[(1,)](
+    a.cuda(), b.cuda(), out, b.stride(0),
+    BLOCK_M=M, BLOCK_N=BN, BLOCK_K=K,
+    num_warps=4, num_stages=2,
+)
+diff = (ref.cuda().float() - out.float()).abs()
+sf = ((ref.cuda() * out) < 0).sum().item()
+print(f"DISABLE_TMA={os.environ.get('TRITON_DISABLE_TMA', '0')}: max_diff={diff.max():.6f} sign_flips={sf}/{M*BN}")

--- 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
@@ -13,22 +13,24 @@
 @triton.jit
 def grouped_gemm_swiglu_kernel(
     hidden_states_ptr,
-    W_gate_ptr,
-    W_up_ptr,
+    W_gate_tiled_ptr,
+    W_up_tiled_ptr,
     expert_offsets_ptr,
     output_ptr,
     H: tl.constexpr,
     I: tl.constexpr,
+    BLOCK_N: tl.constexpr,
     stride_hidden_h: tl.constexpr,
-    stride_gate_e: tl.constexpr,
-    stride_gate_h: tl.constexpr,
-    stride_gate_i: tl.constexpr,
-    stride_up_e: tl.constexpr,
-    stride_up_h: tl.constexpr,
-    stride_up_i: tl.constexpr,
+    stride_gate_tiled_e: tl.constexpr,
+    stride_gate_tiled_t: tl.constexpr,
+    stride_gate_tiled_h: tl.constexpr,
+    stride_gate_tiled_n: tl.constexpr,
+    stride_up_tiled_e: tl.constexpr,
+    stride_up_tiled_t: tl.constexpr,
+    stride_up_tiled_h: tl.constexpr,
+    stride_up_tiled_n: tl.constexpr,
     stride_out_i: tl.constexpr,
     BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(0)
@@ -46,8 +48,8 @@
     n_off = tl.arange(0, BLOCK_N)
     n_mask = n_off < I - n_start
 
-    w_gate_base = W_gate_ptr + e * stride_gate_e
-    w_up_base = W_up_ptr + e * stride_up_e
+    w_gate_tile_ptr = W_gate_tiled_ptr + e * stride_gate_tiled_e + n_tile_idx * stride_gate_tiled_t
+    w_up_tile_ptr = W_up_tiled_ptr + e * stride_up_tiled_e + n_tile_idx * stride_up_tiled_t
 
     for m_start in range(0, n_tokens, BLOCK_M):
         m_off = tl.arange(0, BLOCK_M)
@@ -67,20 +69,20 @@
                 mask=m_mask[:, None] & k_mask[None, :],
             )
             w_gate = tl.load(
-                w_gate_base
-                + (k_start + k_off)[:, None] * stride_gate_h
-                + (n_start + n_off)[None, :],
+                w_gate_tile_ptr
+                + (k_start + k_off)[:, None] * stride_gate_tiled_h
+                + n_off[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )
             w_up = tl.load(
-                w_up_base
-                + (k_start + k_off)[:, None] * stride_up_h
-                + (n_start + n_off)[None, :],
+                w_up_tile_ptr
+                + (k_start + k_off)[:, None] * stride_up_tiled_h
+                + n_off[None, :],
                 mask=k_mask[:, None] & n_mask[None, :],
             )
 
-            acc_gate = tl.dot(x, w_gate, acc_gate, input_precision='ieee')
-            acc_up = tl.dot(x, w_up, acc_up, input_precision='ieee')
+            acc_gate = tl.dot(x, w_gate, acc_gate)
+            acc_up = tl.dot(x, w_up, acc_up)
 
         gate_silu = tl.sigmoid(acc_gate) * acc_gate
         result = (gate_silu * acc_up).to(tl.bfloat16)
@@ -105,6 +107,28 @@
         self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
         nn.init.normal_(self.W_gate, std=0.02)
         nn.init.normal_(self.W_up, std=0.02)
+        self._gate_tiled: torch.Tensor | None = None
+        self._up_tiled: torch.Tensor | None = None
+
+    def _make_tiled(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reorganize (E, H, I) weights to (E, num_nt, H, BLOCK_N) for TMA-compatible access."""
+        if self._gate_tiled is not None:
+            return self._gate_tiled, self._up_tiled
+        E, H, I = self.E, self.H, self.I
+        bn = 128 if I >= 128 else 32
+        num_nt = I // bn if I % bn == 0 else triton.cdiv(I, bn)
+        pad = num_nt * bn - I
+        if pad > 0:
+            wg = torch.nn.functional.pad(self.W_gate.view(E, H, -1), (0, pad))
+            wu = torch.nn.functional.pad(self.W_up.view(E, H, -1), (0, pad))
+        else:
+            wg = self.W_gate
+            wu = self.W_up
+        gt = wg.view(E, H, num_nt, bn).permute(0, 2, 1, 3).contiguous()
+        ut = wu.view(E, H, num_nt, bn).permute(0, 2, 1, 3).contiguous()
+        self._gate_tiled = gt
+        self._up_tiled = ut
+        return gt, ut
 
     def forward(
         self,
@@ -114,29 +138,33 @@
         T_perm, _ = hidden_states.shape
         out = torch.empty(T_perm, self.I, dtype=torch.bfloat16, device=hidden_states.device)
 
-        M_avg = T_perm // self.E
         bn = 128 if self.I >= 128 else 32
         bm = 64
         bk = 32
         nw = 4
         ns = 3
-        grid = (self.E * triton.cdiv(self.I, bn),)
-
+        num_nt = triton.cdiv(self.I, bn)
+        grid = (self.E * num_nt,)
+
+        gt, ut = self._make_tiled()
         grouped_gemm_swiglu_kernel[grid](
             hidden_states,
-            self.W_gate,
-            self.W_up,
+            gt,
+            ut,
             expert_offsets,
             out,
             self.H,
             self.I,
+            bn,
             hidden_states.stride(0),
-            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),
+            gt.stride(0),
+            gt.stride(1),
+            gt.stride(2),
+            gt.stride(3),
+            ut.stride(0),
+            ut.stride(1),
+            ut.stride(2),
+            ut.stride(3),
             out.stride(0),
             BLOCK_M=bm,
             BLOCK_N=bn,

--- 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
@@ -19,7 +19,6 @@
     output_ptr,
     H: tl.constexpr,
     I: tl.constexpr,
-    BLOCK_N: tl.constexpr,
     stride_hidden_h: tl.constexpr,
     stride_gate_tiled_e: tl.constexpr,
     stride_gate_tiled_t: tl.constexpr,

--- 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
@@ -19,6 +19,7 @@
     output_ptr,
     H: tl.constexpr,
     I: tl.constexpr,
+    BLOCK_N: tl.constexpr,
     stride_hidden_h: tl.constexpr,
     stride_gate_tiled_e: tl.constexpr,
     stride_gate_tiled_t: tl.constexpr,

--- 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
@@ -167,7 +167,6 @@
             ut.stride(3),
             out.stride(0),
             BLOCK_M=bm,
-            BLOCK_N=bn,
             BLOCK_K=bk,
             num_warps=nw,
             num_stages=ns,

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_tiled.py
@@ -0,0 +1,90 @@
+"""Verify tiled kernel correctness."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I = 128, 4096, 1536
+E, num_nt = 1, I // 128
+BN = 128
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w_gate = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_up = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Tiled weights: (E, num_nt, H, BN)
+wg_tiled = w_gate.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+wu_tiled = w_up.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+
+print(f"Tiled shape: {wg_tiled.shape}")
+print(f"Tiled strides: {wg_tiled.stride()}")
+
+# Reference
+gate_ref = x @ w_gate
+up_ref = x @ w_up
+out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+
+@triton.jit
+def tiled_kernel(x_ptr, wg_ptr, wu_ptr, out_ptr, H, I, BN,
+                 x_s0, wg_s0, wg_s1, wg_s2,
+                 wu_s0, wu_s1, wu_s2, out_s0,
+                 BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt_cnt = tl.cdiv(I, BLOCK_N)
+    e = pid // nt_cnt
+    nt = pid % nt_cnt
+    n_start = nt * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+    
+    wg_base = wg_ptr + e * wg_s0 + nt * wg_s1
+    wu_base = wu_ptr + e * wu_s0 + nt * wu_s1
+    
+    for m_base in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_base + m_off < n_tokens
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_base + k_off < H
+            x_t = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
+            )
+            wg_t = tl.load(
+                wg_base + (k_base + k_off)[:, None] * wg_s2 + n_off[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            wu_t = tl.load(
+                wu_base + (k_base + k_off)[:, None] * wu_s2 + n_off[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            acc_g = tl.dot(x_t, wg_t, acc_g)
+            acc_u = tl.dot(x_t, wu_t, acc_u)
+        
+        result = (tl.sigmoid(acc_g) * acc_g * acc_u).to(tl.bfloat16)
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+tiled_kernel[(E * num_nt,)](
+    x.cuda(), wg_tiled.cuda(), wu_tiled.cuda(), out,
+    n_tokens, H, I, BN,
+    x.stride(0),
+    wg_tiled.stride(0), wg_tiled.stride(1), wg_tiled.stride(2),
+    wu_tiled.stride(0), wu_tiled.stride(1), wu_tiled.stride(2),
+    out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=3,
+)
+
+diff = (out_ref.cuda().float() - out.float()).abs()
+print(f"Tiled kernel vs ref: max_diff={diff.max().item():.6f} mean={diff.mean().item():.6f}")
+sf = ((out_ref.cuda() * out) < 0).sum().item()
+print(f"Sign flips: {sf}/{n_tokens*I}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_tiled2.py
@@ -0,0 +1,120 @@
+"""Verify tiled kernel correctness - all args passed properly."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I, E = 128, 4096, 1536, 1
+BN, num_nt = 128, I // 128
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+wg = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+wu = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+wg_t = wg.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+wu_t = wu.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+
+gate_ref = x @ wg
+up_ref = x @ wu
+out_ref = torch.nn.functional.silu(gate_ref) * up_ref
+
+@triton.jit
+def tiled_kernel(
+    x_ptr, wg_ptr, wu_ptr, out_ptr,
+    n_tokens, H, I, BN,
+    x_s0, wg_s0, wg_s1, wg_s2,
+    wu_s0, wu_s1, wu_s2, out_s0,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    nt_cnt = tl.cdiv(I, BLOCK_N)
+    e = pid // nt_cnt
+    nt = pid % nt_cnt
+
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - nt * BLOCK_N
+
+    wg_base = wg_ptr + e * wg_s0 + nt * wg_s1
+    wu_base = wu_ptr + e * wu_s0 + nt * wu_s1
+
+    for m_base in range(0, n_tokens, BLOCK_M):
+        m_mask = m_base + m_off < n_tokens
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            x_t = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=m_mask[:, None],
+            )
+            wg_t = tl.load(
+                wg_base + (k_base + k_off)[:, None] * wg_s2 + n_off[None, :],
+                mask=n_mask[None, :],
+            )
+            wu_t = tl.load(
+                wu_base + (k_base + k_off)[:, None] * wu_s2 + n_off[None, :],
+                mask=n_mask[None, :],
+            )
+            acc_g = tl.dot(x_t, wg_t, acc_g)
+            acc_u = tl.dot(x_t, wu_t, acc_u)
+
+        result = (tl.sigmoid(acc_g) * acc_g * acc_u).to(tl.bfloat16)
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+tiled_kernel[(E * num_nt,)](
+    x.cuda(), wg_t.cuda(), wu_t.cuda(), out,
+    n_tokens, H, I, BN,
+    x.stride(0),
+    wg_t.stride(0), wg_t.stride(1), wg_t.stride(2),
+    wu_t.stride(0), wu_t.stride(1), wu_t.stride(2),
+    out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=3,
+)
+
+diff = (out_ref.cuda().float() - out.float()).abs()
+print(f"Tiled kernel: max_diff={diff.max():.6f} mean={diff.mean():.6f}")
+sf = ((out_ref.cuda() * out) < 0).sum().item()
+print(f"Sign flips: {sf}/{n_tokens*I}")
+
+# Also test just the GEMM (no silu)
+@triton.jit
+def gemm_only(x_ptr, w_ptr, out_ptr, n_tokens, H, I, BN,
+              x_s0, w_s0, w_s1, w_s2, out_s0,
+              BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    w_base = w_ptr + nt * w_s1
+    for m_base in range(0, n_tokens, BLOCK_M):
+        m_mask = m_base + m_off < n_tokens
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            x_t = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=m_mask[:, None],
+            )
+            w_t = tl.load(
+                w_base + (k_base + k_off)[:, None] * w_s2 + n_off[None, :],
+            )
+            acc = tl.dot(x_t, w_t, acc)
+        tl.store(out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+                acc.to(tl.bfloat16), mask=m_mask[:, None])
+
+out_g = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+gemm_only[(num_nt,)](x.cuda(), wg_t.cuda(), out_g, n_tokens, H, I, BN,
+    x.stride(0), wg_t.stride(0), wg_t.stride(1), wg_t.stride(2), out_g.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32, num_warps=4, num_stages=3)
+dg = (gate_ref.cuda().float() - out_g.float()).abs()
+sfg = ((gate_ref.cuda() * out_g) < 0).sum().item()
+print(f"\nGEMM only (gate): max_diff={dg.max():.6f} sign_flips={sfg}/{n_tokens*I}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_dataload.py
@@ -0,0 +1,123 @@
+"""Verify data loading from tiled weights is correct."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+H, I, BN = 4096, 1536, 128
+num_nt = I // BN
+
+# Create data
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_copy = w.clone()
+
+# Tiled version
+w_t = w.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+# Verify: w_t[0, t, :, :] should equal w[:, t*BN:(t+1)*BN]
+for t in range(min(3, num_nt)):
+    print(f"Tile {t}: match={torch.allclose(w_t[0,t,:,:], w[:, t*BN:(t+1)*BN], atol=1e-10)}")
+
+# Now test: does tl.load with tiled pointer give the right data?
+@triton.jit
+def load_test(w_ptr, out_ptr, H, BN,
+              w_s0, w_s1, w_s2, out_s0,
+              BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    
+    w_base = w_ptr + nt * w_s1
+    
+    for m_base in range(0, H, BLOCK_M):
+        # Load from tiled weights: shape (H, BN), stride (BN, 1)
+        tile = tl.load(
+            w_base + (m_base + m_off)[:, None] * w_s2 + n_off[None, :],
+            mask=(m_base + m_off)[:, None] < H,
+        )
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            tile.to(tl.bfloat16),
+            mask=(m_base + m_off)[:, None] < H,
+        )
+
+# Load tile 0 using Triton
+out_t = torch.zeros(H, BN, dtype=torch.bfloat16, device='cuda')
+load_test[(1,)](
+    w_t.cuda(), out_t, H, BN,
+    w_t.stride(0), w_t.stride(1), w_t.stride(2), out_t.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=3,
+)
+
+ref_tile = w[:, 0:BN]
+diff = (ref_tile.cuda().float() - out_t.float()).abs()
+print(f"\nLoad test tile 0: max_diff={diff.max():.6f}")
+
+# Now test with a K-tile load (BLOCK_K=32, BLOCK_N=128)
+@triton.jit
+def ktile_load(w_ptr, out_ptr, H, BN,
+               w_s0, w_s1, w_s2, out_s0,
+               BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    
+    w_base = w_ptr + nt * w_s1
+    
+    for m_base in range(0, H, BLOCK_M):
+        tile = tl.load(
+            w_base + (m_base + m_off)[:, None] * w_s2 + n_off[None, :],
+            mask=(m_base + m_off)[:, None] < H,
+        )
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            tile.to(tl.bfloat16),
+            mask=(m_base + m_off)[:, None] < H,
+        )
+
+out_k = torch.zeros(H, BN, dtype=torch.bfloat16, device='cuda')
+ktile_load[(1,)](
+    w_t.cuda(), out_k, H, BN,
+    w_t.stride(0), w_t.stride(1), w_t.stride(2), out_k.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=3,
+)
+
+diff_k = (ref_tile.cuda().float() - out_k.float()).abs()
+print(f"K-tile load test: max_diff={diff_k.max():.6f}")
+
+# Test: tl.load + tl.dot from tiled weights
+a = torch.randn(64, 32, dtype=torch.bfloat16) * 0.1
+# For K-tile at k=0: B = w[0:32, 0:128]
+ref_gemm = a @ w[:32, :BN]
+
+@triton.jit
+def gemm_tiled(a_ptr, w_ptr, out_ptr, w_s0, w_s1, w_s2, out_s0,
+               BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    m_off = tl.arange(0, BLOCK_M)
+    n_off = tl.arange(0, BLOCK_N)
+    k_off = tl.arange(0, BLOCK_K)
+    
+    w_base = w_ptr + nt * w_s1
+    
+    a_t = tl.load(a_ptr + m_off[:, None] * BLOCK_K + k_off[None, :])
+    w_t = tl.load(w_base + k_off[:, None] * w_s2 + n_off[None, :])
+    c = tl.dot(a_t, w_t)
+    tl.store(out_ptr + m_off[:, None] * out_s0 + n_off[None, :], c.to(tl.bfloat16))
+
+out_gemm = torch.zeros(64, BN, dtype=torch.bfloat16, device='cuda')
+gemm_tiled[(1,)](
+    a.cuda(), w_t.cuda(), out_gemm,
+    w_t.stride(0), w_t.stride(1), w_t.stride(2), out_gemm.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4,
+)
+
+diff_g = (ref_gemm.cuda().float() - out_gemm.float()).abs()
+sf_g = ((ref_gemm.cuda() * out_gemm) < 0).sum().item()
+print(f"\nGEMM from tiled: max_diff={diff_g.max():.6f} sign_flips={sf_g}/{64*BN}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_multitile.py
@@ -0,0 +1,75 @@
+"""Test full multi-tile GEMM with tiled weights."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I, E = 128, 4096, 1536, 1
+BN, num_nt = 128, I // 128
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+wg = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+
+# Tiled
+wg_t = wg.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+
+# Reference
+gate_ref = x @ wg
+
+@triton.jit
+def gemm_multi(x_ptr, w_ptr, out_ptr, n_tokens, H, I, BN,
+               x_s0, w_s0, w_s1, w_s2, out_s0,
+               BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid = tl.program_id(0)
+    nt_cnt = tl.cdiv(I, BLOCK_N)
+    e = pid // nt_cnt
+    nt = pid % nt_cnt
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - nt * BLOCK_N
+    
+    w_base = w_ptr + e * w_s0 + nt * w_s1
+    
+    for m_base in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_base + m_off < n_tokens
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            x_t = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=m_mask[:, None],
+            )
+            w_t = tl.load(
+                w_base + (k_base + k_off)[:, None] * w_s2 + n_off[None, :],
+                mask=n_mask[None, :],
+            )
+            acc = tl.dot(x_t, w_t, acc)
+        
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+gemm_multi[(E * num_nt,)](
+    x.cuda(), wg_t.cuda(), out, n_tokens, H, I, BN,
+    x.stride(0),
+    wg_t.stride(0), wg_t.stride(1), wg_t.stride(2), out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    num_warps=4, num_stages=3,
+)
+
+diff = (gate_ref.cuda().float() - out.float()).abs()
+print(f"Multi-N-tile GEMM: max_diff={diff.max():.6f} mean={diff.mean():.6f}")
+sf = ((gate_ref.cuda() * out) < 0).sum().item()
+print(f"Sign flips: {sf}/{n_tokens*I}")
+
+# Per-N-tile breakdown
+for nt in range(num_nt):
+    ns = nt * BN
+    ne = min(ns + BN, I)
+    tile_diff = diff[:, ns:ne]
+    print(f"  NTile {nt} ({ns}:{ne}): max_diff={tile_diff.max():.6f} mean={tile_diff.mean():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_multitile2.py
@@ -0,0 +1,70 @@
+"""Debug multi-tile GEMM - pass NUM_NT as constexpr."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I, E, BN = 128, 4096, 1536, 1, 128
+num_nt = I // BN
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_t = w.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+
+ref = x @ w
+
+@triton.jit
+def gm(x_ptr, w_ptr, out_ptr, n_tokens, H, I,
+       x_s0, w_s0, w_s1, w_s2, out_s0,
+       BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+       NUM_NT: tl.constexpr):
+    pid = tl.program_id(0)
+    e = pid // NUM_NT
+    nt = pid % NUM_NT
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - nt * BLOCK_N
+    
+    w_base = w_ptr + e * w_s0 + nt * w_s1
+    
+    for m_base in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_base + m_off < n_tokens
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        
+        for k_base in range(0, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            x_t = tl.load(
+                x_ptr + (m_base + m_off)[:, None] * x_s0 + (k_base + k_off)[None, :],
+                mask=m_mask[:, None],
+            )
+            w_t = tl.load(
+                w_base + (k_base + k_off)[:, None] * w_s2 + n_off[None, :],
+            )
+            acc = tl.dot(x_t, w_t, acc)
+        
+        tl.store(
+            out_ptr + (m_base + m_off)[:, None] * out_s0 + n_off[None, :],
+            acc.to(tl.bfloat16),
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
+
+out = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+gm[(E * num_nt,)](
+    x.cuda(), w_t.cuda(), out, n_tokens, H, I,
+    x.stride(0), w_t.stride(0), w_t.stride(1), w_t.stride(2), out.stride(0),
+    BLOCK_M=64, BLOCK_N=128, BLOCK_K=32,
+    NUM_NT=num_nt,
+    num_warps=4, num_stages=3,
+)
+
+diff = (ref.cuda().float() - out.float()).abs()
+print(f"Multi-tile (NUM_NT constexpr): max_diff={diff.max():.6f} mean={diff.mean():.6f}")
+sf = ((ref.cuda() * out) < 0).sum().item()
+print(f"Sign flips: {sf}/{n_tokens*I}")
+
+for nt in range(num_nt):
+    ns = nt * BN
+    ne = min(ns + BN, I)
+    td = diff[:, ns:ne]
+    print(f"  Tile {nt} ({ns}:{ne}): max={td.max():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_compare.py
@@ -0,0 +1,79 @@
+"""Compare single-launch multi-tile vs individual single-tile launches."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I, E, BN = 128, 4096, 1536, 1, 128
+num_nt = I // BN
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_t = w.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+ref = x @ w
+
+@triton.jit
+def gm_single(xp, wp, op, N, H, xs, ws2, os0,
+              BM: tl.constexpr, BN2: tl.constexpr, BK: tl.constexpr):
+    mo = tl.arange(0, BM)
+    no = tl.arange(0, BN2)
+    for mb in range(0, N, BM):
+        mm = mb + mo < N
+        ag = tl.zeros((BM, BN2), dtype=tl.float32)
+        for kb in range(0, H, BK):
+            ko = tl.arange(0, BK)
+            ag = tl.dot(tl.load(xp + (mb+mo)[:,None]*xs + (kb+ko)[None,:], mask=mm[:,None]),
+                       tl.load(wp + (kb+ko)[:,None]*ws2 + no[None,:]), ag)
+        tl.store(op + (mb+mo)[:,None]*os0 + no[None,:], ag.to(tl.bfloat16), mask=mm[:,None])
+
+# Approach 1: single launch for all tiles
+@triton.jit
+def gm_multi(xp, wp, op, N, H, I, xs, ws0, ws1, ws2, os0,
+             BM: tl.constexpr, BN2: tl.constexpr, BK: tl.constexpr, NNT: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    no = tl.arange(0, BN2)
+    nm = no < I - nt * BN2
+    wb = wp + nt * ws1
+    for mb in range(0, N, BM):
+        mo = tl.arange(0, BM)
+        mm = mb + mo < N
+        ag = tl.zeros((BM, BN2), dtype=tl.float32)
+        for kb in range(0, H, BK):
+            ko = tl.arange(0, BK)
+            ag = tl.dot(tl.load(xp + (mb+mo)[:,None]*xs + (kb+ko)[None,:], mask=mm[:,None]),
+                       tl.load(wb + (kb+ko)[:,None]*ws2 + no[None,:]), ag)
+        tl.store(op + (mb+mo)[:,None]*os0 + no[None,:], ag.to(tl.bfloat16), mask=mm[:,None] & nm[None,:])
+
+# Multi-launch: launch each tile separately
+out_separate = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+for nt in range(num_nt):
+    w_slice = w_t[:, nt, :, :].contiguous()
+    temp = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+    gm_single[(1,)](
+        x.cuda(), w_slice.cuda(), temp, n_tokens, H,
+        x.stride(0), w_slice.stride(0), temp.stride(0),
+        BM=64, BN2=128, BK=32, num_warps=4, num_stages=3,
+    )
+    out_separate[:, nt*BN:(nt+1)*BN] = temp[:, nt*BN:(nt+1)*BN]
+
+diff_sep = (ref.cuda().float() - out_separate.float()).abs()
+sf_sep = ((ref.cuda()*out_separate)<0).sum().item()
+print(f"Separate launches: max_diff={diff_sep.max():.6f} sign_flips={sf_sep}/{n_tokens*I}")
+
+# Single launch
+out_single = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+gm_multi[(num_nt,)](
+    x.cuda(), w_t.cuda(), out_single, n_tokens, H, I,
+    x.stride(0), w_t.stride(0), w_t.stride(1), w_t.stride(2), out_single.stride(0),
+    BM=64, BN2=128, BK=32, NNT=num_nt, num_warps=4, num_stages=3,
+)
+
+diff_single = (ref.cuda().float() - out_single.float()).abs()
+sf_single = ((ref.cuda()*out_single)<0).sum().item()
+print(f"Single launch: max_diff={diff_single.max():.6f} sign_flips={sf_single}/{n_tokens*I}")
+
+# Compare separate vs single
+diff_ss = (out_separate.cuda().float() - out_single.float()).abs()
+print(f"Separate vs Single: max_diff={diff_ss.max():.6f}")

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/debug_compare2.py
@@ -0,0 +1,80 @@
+"""Compare single-launch multi-tile vs reference, check for consistent errors."""
+import torch
+import triton
+import triton.language as tl
+
+torch.manual_seed(42)
+
+n_tokens, H, I, E, BN = 128, 4096, 1536, 1, 128
+num_nt = I // BN
+
+x = torch.randn(n_tokens, H, dtype=torch.bfloat16) * 0.1
+w = torch.randn(H, I, dtype=torch.bfloat16) * 0.02
+w_t = w.view(1, H, num_nt, BN).permute(0, 2, 1, 3).contiguous()
+ref = x @ w
+
+# Approach: directly tile the output using a single-tile kernel per tile
+# Instead of passing temp tensor with I cols, pass one with BN cols per tile
+out_correct = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+for nt in range(num_nt):
+    w_slice = w_t[:, nt, :, :].contiguous().squeeze(0)
+    
+    @triton.jit
+    def gm(xp, wp, op, N, H, xs, ws2, os0,
+           BM: tl.constexpr, BN2: tl.constexpr, BK: tl.constexpr):
+        mo = tl.arange(0, BM)
+        no = tl.arange(0, BN2)
+        for mb in range(0, N, BM):
+            mm = mb + mo < N
+            ag = tl.zeros((BM, BN2), dtype=tl.float32)
+            for kb in range(0, H, BK):
+                ko = tl.arange(0, BK)
+                ag = tl.dot(tl.load(xp + (mb+mo)[:,None]*xs + (kb+ko)[None,:], mask=mm[:,None]),
+                           tl.load(wp + (kb+ko)[:,None]*ws2 + no[None,:]), ag)
+            tl.store(op + (mb+mo)[:,None]*os0 + no[None,:], ag.to(tl.bfloat16), mask=mm[:,None])
+    
+    out_nt = torch.zeros(n_tokens, BN, dtype=torch.bfloat16, device='cuda')
+    gm[(1,)](
+        x.cuda(), w_slice.unsqueeze(0).cuda(), out_nt, n_tokens, H,
+        x.stride(0), w_slice.stride(0), out_nt.stride(0),
+        BM=64, BN2=128, BK=32, num_warps=4, num_stages=3,
+    )
+    out_correct[:, nt*BN:(nt+1)*BN] = out_nt
+
+diff_correct = (ref.cuda().float() - out_correct.float()).abs()
+sf_correct = ((ref.cuda()*out_correct)<0).sum().item()
+print(f"Per-tile contiguous: max_diff={diff_correct.max():.6f} sign_flips={sf_correct}/{n_tokens*I}")
+
+# Now try with the multi-tile kernel
+@triton.jit
+def gm_multi(xp, wp, op, N, H, I, xs, ws0, ws1, ws2, os0,
+             BM: tl.constexpr, BN2: tl.constexpr, BK: tl.constexpr, NNT: tl.constexpr):
+    pid = tl.program_id(0)
+    nt = pid
+    no = tl.arange(0, BN2)
+    nm = no < I - nt * BN2
+    wb = wp + nt * ws1
+    for mb in range(0, N, BM):
+        mo = tl.arange(0, BM)
+        mm = mb + mo < N
+        ag = tl.zeros((BM, BN2), dtype=tl.float32)
+        for kb in range(0, H, BK):
+            ko = tl.arange(0, BK)
+            ag = tl.dot(tl.load(xp + (mb+mo)[:,None]*xs + (kb+ko)[None,:], mask=mm[:,None]),
+                       tl.load(wb + (kb+ko)[:,None]*ws2 + no[None,:]), ag)
+        tl.store(op + (mb+mo)[:,None]*os0 + no[None,:], ag.to(tl.bfloat16), mask=mm[:,None] & nm[None,:])
+
+out_multi = torch.zeros(n_tokens, I, dtype=torch.bfloat16, device='cuda')
+gm_multi[(num_nt,)](
+    x.cuda(), w_t.cuda(), out_multi, n_tokens, H, I,
+    x.stride(0), w_t.stride(0), w_t.stride(1), w_t.stride(2), out_multi.stride(0),
+    BM=64, BN2=128, BK=32, NNT=num_nt, num_warps=4, num_stages=3,
+)
+
+diff_multi = (ref.cuda().float() - out_multi.float()).abs()
+sf_multi = ((ref.cuda()*out_multi)<0).sum().item()
+print(f"Multi-tile: max_diff={diff_multi.max():.6f} sign_flips={sf_multi}/{n_tokens*I}")
+
+# Compare correct vs multi
+diff_cm = (out_correct.cuda().float() - out_multi.float()).abs()
+print(f"Correct vs Multi: max_diff={diff_cm.max():.6f}")

--- 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
@@ -11,7 +11,7 @@
 
 
 @triton.jit
-def grouped_gemm_swiglu_kernel(
+def expert_kernel(
     hidden_states_ptr,
     W_gate_tiled_ptr,
     W_up_tiled_ptr,
@@ -19,24 +19,23 @@
     output_ptr,
     H: tl.constexpr,
     I: tl.constexpr,
-    BLOCK_N: tl.constexpr,
     stride_hidden_h: tl.constexpr,
-    stride_gate_tiled_e: tl.constexpr,
-    stride_gate_tiled_t: tl.constexpr,
-    stride_gate_tiled_h: tl.constexpr,
-    stride_gate_tiled_n: tl.constexpr,
-    stride_up_tiled_e: tl.constexpr,
-    stride_up_tiled_t: tl.constexpr,
-    stride_up_tiled_h: tl.constexpr,
-    stride_up_tiled_n: tl.constexpr,
+    tiled_stride_e: tl.constexpr,
+    tiled_stride_t: tl.constexpr,
+    tiled_stride_h: tl.constexpr,
+    tiled_stride_n: tl.constexpr,
+    tiled_up_stride_e: tl.constexpr,
+    tiled_up_stride_t: tl.constexpr,
+    tiled_up_stride_h: tl.constexpr,
+    tiled_up_stride_n: tl.constexpr,
     stride_out_i: tl.constexpr,
     BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    NUM_NT: tl.constexpr,
 ):
     pid = tl.program_id(0)
-    num_n_tiles = tl.cdiv(I, BLOCK_N)
-    e = pid // num_n_tiles
-    n_tile_idx = pid % num_n_tiles
+    e = pid
 
     start = tl.load(expert_offsets_ptr + e)
     end = tl.load(expert_offsets_ptr + e + 1)
@@ -44,55 +43,59 @@
     if n_tokens <= 0:
         return
 
-    n_start = n_tile_idx * BLOCK_N
-    n_off = tl.arange(0, BLOCK_N)
-    n_mask = n_off < I - n_start
-
-    w_gate_tile_ptr = W_gate_tiled_ptr + e * stride_gate_tiled_e + n_tile_idx * stride_gate_tiled_t
-    w_up_tile_ptr = W_up_tiled_ptr + e * stride_up_tiled_e + n_tile_idx * stride_up_tiled_t
-
-    for m_start in range(0, n_tokens, BLOCK_M):
-        m_off = tl.arange(0, BLOCK_M)
-        m_mask = m_start + m_off < n_tokens
-
-        x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
-
-        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, H, BLOCK_K):
-            k_off = tl.arange(0, BLOCK_K)
-            k_mask = k_start + k_off < H
-
-            x = tl.load(
-                x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
-                mask=m_mask[:, None] & k_mask[None, :],
+    w_gate_base = W_gate_tiled_ptr + e * tiled_stride_e
+    w_up_base = W_up_tiled_ptr + e * tiled_up_stride_e
+
+    for nt in range(NUM_NT):
+        n_start = nt * BLOCK_N
+        n_off = tl.arange(0, BLOCK_N)
+        n_mask = n_off < I - n_start
+
+        w_gate_nt_ptr = w_gate_base + nt * tiled_stride_t
+        w_up_nt_ptr = w_up_base + nt * tiled_up_stride_t
+
+        for m_start in range(0, n_tokens, BLOCK_M):
+            m_off = tl.arange(0, BLOCK_M)
+            m_mask = m_start + m_off < n_tokens
+
+            x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
+
+            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, H, BLOCK_K):
+                k_off = tl.arange(0, BLOCK_K)
+                k_mask = k_start + k_off < H
+
+                x = tl.load(
+                    x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
+                    mask=m_mask[:, None] & k_mask[None, :],
+                )
+                w_gate = tl.load(
+                    w_gate_nt_ptr
+                    + (k_start + k_off)[:, None] * tiled_stride_h
+                    + n_off[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+                w_up = tl.load(
+                    w_up_nt_ptr
+                    + (k_start + k_off)[:, None] * tiled_up_stride_h
+                    + n_off[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+
+                acc_gate = tl.dot(x, w_gate, acc_gate)
+                acc_up = tl.dot(x, w_up, acc_up)
+
+            gate_silu = tl.sigmoid(acc_gate) * acc_gate
+            result = (gate_silu * acc_up).to(tl.bfloat16)
+
+            out_ptr = output_ptr + (start + m_start) * stride_out_i
+            tl.store(
+                out_ptr + m_off[:, None] * stride_out_i + (n_start + n_off)[None, :],
+                result,
+                mask=m_mask[:, None] & n_mask[None, :],
             )
-            w_gate = tl.load(
-                w_gate_tile_ptr
-                + (k_start + k_off)[:, None] * stride_gate_tiled_h
-                + n_off[None, :],
-                mask=k_mask[:, None] & n_mask[None, :],
-            )
-            w_up = tl.load(
-                w_up_tile_ptr
-                + (k_start + k_off)[:, None] * stride_up_tiled_h
-                + n_off[None, :],
-                mask=k_mask[:, None] & n_mask[None, :],
-            )
-
-            acc_gate = tl.dot(x, w_gate, acc_gate)
-            acc_up = tl.dot(x, w_up, acc_up)
-
-        gate_silu = tl.sigmoid(acc_gate) * acc_gate
-        result = (gate_silu * acc_up).to(tl.bfloat16)
-
-        out_ptr = output_ptr + (start + m_start) * stride_out_i
-        tl.store(
-            out_ptr + m_off[:, None] * stride_out_i + n_off[None, :],
-            result,
-            mask=m_mask[:, None] & n_mask[None, :],
-        )
 
 
 class Model(nn.Module):
@@ -110,13 +113,11 @@
         self._gate_tiled: torch.Tensor | None = None
         self._up_tiled: torch.Tensor | None = None
 
-    def _make_tiled(self) -> tuple[torch.Tensor, torch.Tensor]:
-        """Reorganize (E, H, I) weights to (E, num_nt, H, BLOCK_N) for TMA-compatible access."""
+    def _make_tiled(self, bn: int) -> tuple[torch.Tensor, torch.Tensor]:
         if self._gate_tiled is not None:
             return self._gate_tiled, self._up_tiled
       

... (1090 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
@@ -10,6 +10,20 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
+# Autotune configs
+_configs = [
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+]
+
+
+@triton.autotune(_configs, key=[])
 @triton.jit
 def expert_kernel(
     hidden_states_ptr,
@@ -32,10 +46,16 @@
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
-    NUM_NT: tl.constexpr,
 ):
     pid = tl.program_id(0)
-    e = pid
+
+    # Precompute N-tile count from I and BLOCK_N
+    # Triton can't use tl.cdiv in constexpr context, use Python
+    # Actually tl.cdiv is available as a regular function
+    num_nt = tl.cdiv(I, BLOCK_N)
+
+    e = pid // num_nt
+    nt = pid % num_nt
 
     start = tl.load(expert_offsets_ptr + e)
     end = tl.load(expert_offsets_ptr + e + 1)
@@ -43,59 +63,55 @@
     if n_tokens <= 0:
         return
 
-    w_gate_base = W_gate_tiled_ptr + e * tiled_stride_e
-    w_up_base = W_up_tiled_ptr + e * tiled_up_stride_e
-
-    for nt in range(NUM_NT):
-        n_start = nt * BLOCK_N
-        n_off = tl.arange(0, BLOCK_N)
-        n_mask = n_off < I - n_start
-
-        w_gate_nt_ptr = w_gate_base + nt * tiled_stride_t
-        w_up_nt_ptr = w_up_base + nt * tiled_up_stride_t
-
-        for m_start in range(0, n_tokens, BLOCK_M):
-            m_off = tl.arange(0, BLOCK_M)
-            m_mask = m_start + m_off < n_tokens
-
-            x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
-
-            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, H, BLOCK_K):
-                k_off = tl.arange(0, BLOCK_K)
-                k_mask = k_start + k_off < H
-
-                x = tl.load(
-                    x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
-                    mask=m_mask[:, None] & k_mask[None, :],
-                )
-                w_gate = tl.load(
-                    w_gate_nt_ptr
-                    + (k_start + k_off)[:, None] * tiled_stride_h
-                    + n_off[None, :],
-                    mask=k_mask[:, None] & n_mask[None, :],
-                )
-                w_up = tl.load(
-                    w_up_nt_ptr
-                    + (k_start + k_off)[:, None] * tiled_up_stride_h
-                    + n_off[None, :],
-                    mask=k_mask[:, None] & n_mask[None, :],
-                )
-
-                acc_gate = tl.dot(x, w_gate, acc_gate)
-                acc_up = tl.dot(x, w_up, acc_up)
-
-            gate_silu = tl.sigmoid(acc_gate) * acc_gate
-            result = (gate_silu * acc_up).to(tl.bfloat16)
-
-            out_ptr = output_ptr + (start + m_start) * stride_out_i
-            tl.store(
-                out_ptr + m_off[:, None] * stride_out_i + (n_start + n_off)[None, :],
-                result,
-                mask=m_mask[:, None] & n_mask[None, :],
+    n_start = nt * BLOCK_N
+    n_off = tl.arange(0, BLOCK_N)
+    n_mask = n_off < I - n_start
+
+    w_gate_ptr = W_gate_tiled_ptr + e * tiled_stride_e + nt * tiled_stride_t
+    w_up_ptr = W_up_tiled_ptr + e * tiled_up_stride_e + nt * tiled_up_stride_t
+
+    for m_start in range(0, n_tokens, BLOCK_M):
+        m_off = tl.arange(0, BLOCK_M)
+        m_mask = m_start + m_off < n_tokens
+
+        x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
+
+        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, H, BLOCK_K):
+            k_off = tl.arange(0, BLOCK_K)
+            k_mask = k_start + k_off < H
+
+            x = tl.load(
+                x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
+                mask=m_mask[:, None] & k_mask[None, :],
             )
+            w_gate = tl.load(
+                w_gate_ptr
+                + (k_start + k_off)[:, None] * tiled_stride_h
+                + n_off[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+            w_up = tl.load(
+                w_up_ptr
+                + (k_start + k_off)[:, None] * tiled_up_stride_h
+                + n_off[None, :],
+                mask=k_mask[:, None] & n_mask[None, :],
+            )
+
+            acc_gate = tl.dot(x, w_gate, acc_gate)
+            acc_up = tl.dot(x, w_up, acc_up)
+
+        gate_silu = tl.sigmoid(acc_gate) * acc_gate
+        result = (gate_silu * acc_up).to(tl.bfloat16)
+
+        out_ptr = output_ptr + (start + m_start) * stride_out_i
+        tl.store(
+            out_ptr + m_off[:, None] * stride_out_i + (n_start + n_off)[None, :],
+            result,
+            mask=m_mask[:, None] & n_mask[None, :],
+        )
 
 
 class Model(nn.Module):
@@ -117,11 +133,11 @@
         if self._gate_tiled is not None:
             return self._gate_tiled, self._up_tiled
         E, H, I = self.E, self.H, self.I
-        num_nt = triton.cdiv(I, bn)
+        num_nt = (I + bn - 1) // bn
         pad = num_nt * bn - I
         if pad > 0:
-            wg = torch

... (1203 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,3 +1,4 @@
+"""Per-expert kernel with N-tile iteration - correct + optimized."""
 from __future__ import annotations
 
 import torch
@@ -10,20 +11,20 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-# Autotune configs
-_configs = [
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
-    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=3),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-]
-
-
-@triton.autotune(_configs, key=[])
+@triton.autotune(
+    [
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=3),
+    ],
+    key=[],
+    warmup=5,
+    rep=10,
+)
 @triton.jit
 def expert_kernel(
     hidden_states_ptr,
@@ -37,25 +38,16 @@
     tiled_stride_e: tl.constexpr,
     tiled_stride_t: tl.constexpr,
     tiled_stride_h: tl.constexpr,
-    tiled_stride_n: tl.constexpr,
     tiled_up_stride_e: tl.constexpr,
     tiled_up_stride_t: tl.constexpr,
     tiled_up_stride_h: tl.constexpr,
-    tiled_up_stride_n: tl.constexpr,
     stride_out_i: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(0)
-
-    # Precompute N-tile count from I and BLOCK_N
-    # Triton can't use tl.cdiv in constexpr context, use Python
-    # Actually tl.cdiv is available as a regular function
-    num_nt = tl.cdiv(I, BLOCK_N)
-
-    e = pid // num_nt
-    nt = pid % num_nt
+    e = pid
 
     start = tl.load(expert_offsets_ptr + e)
     end = tl.load(expert_offsets_ptr + e + 1)
@@ -63,55 +55,61 @@
     if n_tokens <= 0:
         return
 
-    n_start = nt * BLOCK_N
-    n_off = tl.arange(0, BLOCK_N)
-    n_mask = n_off < I - n_start
-
-    w_gate_ptr = W_gate_tiled_ptr + e * tiled_stride_e + nt * tiled_stride_t
-    w_up_ptr = W_up_tiled_ptr + e * tiled_up_stride_e + nt * tiled_up_stride_t
-
-    for m_start in range(0, n_tokens, BLOCK_M):
-        m_off = tl.arange(0, BLOCK_M)
-        m_mask = m_start + m_off < n_tokens
-
-        x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
-
-        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, H, BLOCK_K):
-            k_off = tl.arange(0, BLOCK_K)
-            k_mask = k_start + k_off < H
-
-            x = tl.load(
-                x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
-                mask=m_mask[:, None] & k_mask[None, :],
+    w_gate_base = W_gate_tiled_ptr + e * tiled_stride_e
+    w_up_base = W_up_tiled_ptr + e * tiled_up_stride_e
+
+    num_nt = tl.cdiv(I, BLOCK_N)
+
+    for nt in range(num_nt):
+        n_start = nt * BLOCK_N
+        n_off = tl.arange(0, BLOCK_N)
+        n_mask = n_off < I - n_start
+
+        w_gate_ptr = w_gate_base + nt * tiled_stride_t
+        w_up_ptr = w_up_base + nt * tiled_up_stride_t
+
+        for m_start in range(0, n_tokens, BLOCK_M):
+            m_off = tl.arange(0, BLOCK_M)
+            m_mask = m_start + m_off < n_tokens
+
+            x_ptr = hidden_states_ptr + (start + m_start) * stride_hidden_h
+
+            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, H, BLOCK_K):
+                k_off = tl.arange(0, BLOCK_K)
+                k_mask = k_start + k_off < H
+
+                x = tl.load(
+                    x_ptr + m_off[:, None] * stride_hidden_h + (k_start + k_off)[None, :],
+                    mask=m_mask[:, None] & k_mask[None, :],
+                )
+                w_gate = tl.load(
+                    w_gate_ptr
+                    + (k_start + k_off)[:, None] * tiled_stride_h
+                    + n_off[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+                w_up = tl.load(
+                    w_up_ptr
+                    + (k_start + k_off)[:, None] * tiled_up_stride_h
+                    + n_off[None, :],
+                    mask=k_mask[:, None] & n_mask[None, :],
+                )
+
+                acc_gate = tl.dot(x, w_gate, acc_gate)
+                acc_up = tl.dot(x, w_up, acc_up)
+
+            gate_silu = tl.sigmoid(acc_gate) * acc_gate
+            result = (gate_silu * acc_up).to(tl.bfloat16)
+
+            out_ptr = output_ptr + (start + m_start) * stride_out_i
+            tl.store(
+                out_ptr + m_off[:, None] * stride_out_i + (n_start + n_off)[None, :],
+                result,
+                mask=m_mask[:, None] & n_mask[None, :],
             )
-            w_gate = tl

... (1604 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,4 +1,3 @@
-"""Per-expert kernel with N-tile iteration - correct + optimized."""
 from __future__ import annotations
 
 import torch
@@ -11,20 +10,6 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 
-@triton.autotune(
-    [
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=8, num_stages=3),
-    ],
-    key=[],
-    warmup=5,
-    rep=10,
-)
 @triton.jit
 def expert_kernel(
     hidden_states_ptr,
@@ -38,9 +23,11 @@
     tiled_stride_e: tl.constexpr,
     tiled_stride_t: tl.constexpr,
     tiled_stride_h: tl.constexpr,
+    tiled_stride_n: tl.constexpr,
     tiled_up_stride_e: tl.constexpr,
     tiled_up_stride_t: tl.constexpr,
     tiled_up_stride_h: tl.constexpr,
+    tiled_up_stride_n: tl.constexpr,
     stride_out_i: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -154,6 +141,11 @@
         out = torch.empty(T_perm, self.I, dtype=torch.bfloat16, device=hidden_states.device)
 
         bn = 128
+        bm = 64
+        bk = 32
+        nw = 4
+        ns = 3
+
         gt, ut = self._make_tiled(bn)
 
         expert_kernel[(self.E,)](
@@ -168,10 +160,17 @@
             gt.stride(0),
             gt.stride(1),
             gt.stride(2),
+            gt.stride(3),
             ut.stride(0),
             ut.stride(1),
             ut.stride(2),
+            ut.stride(3),
             out.stride(0),
+            BLOCK_M=bm,
+            BLOCK_N=bn,
+            BLOCK_K=bk,
+            num_warps=nw,
+            num_stages=ns,
         )
         return out