--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/solution.py
@@ -0,0 +1,631 @@
+"""Grouped GEMM + fused SwiGLU for MoE up-projection on Blackwell (SM120).
+
+Implements: h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e])
+using custom CUDA kernel with tensor cores and load_inline.
+"""
+from __future__ import annotations
+
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+cuda_source = r"""
+#include <cuda/bfloat16.h>
+#include <cuda_pipeline.h>
+#include <mma.h>
+
+using namespace nvcuda;
+using cutlass::pipe;
+
+constexpr int WARP_SIZE = 32;
+constexpr int MMA_M = 16;
+constexpr int MMA_N = 16;
+constexpr int MMA_K = 16;
+
+// SwiGLU activation: silu(x) * x = x / (1 + exp(-x)) * x
+__device__ __forceinline__ __nv_bfloat16 sigmoid_swish(__nv_bfloat16 x) {
+    float fx = __bfloat16_to_float(x);
+    float sig = fx / (1.0f + expf(-fx));
+    return __float_to_bfloat16(sig * fx);
+}
+
+template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
+__global__ void grouped_gemm_swiglu_kernel(
+    const __nv_bfloat16* __restrict__ hidden_states,  // (T_perm, H)
+    const __nv_bfloat16* __restrict__ W_gate,         // (E, H, I)
+    const __nv_bfloat16* __restrict__ W_up,          // (E, H, I)
+    const int* __restrict__ expert_offsets,           // (E+1,)
+    __nv_bfloat16* __restrict__ output,              // (T_perm, I)
+    int T_perm, int H, int I, int E
+) {
+    extern __shared__ char smem_buf[];
+
+    int expert_idx = blockIdx.x;
+    if (expert_idx >= E) return;
+
+    int start_row = expert_offsets[expert_idx];
+    int end_row = expert_offsets[expert_idx + 1];
+    int N_e = end_row - start_row;
+    if (N_e <= 0) return;
+
+    const __nv_bfloat16* Wg = W_gate + expert_idx * H * I;
+    const __nv_bfloat16* Wu = W_up + expert_idx * H * I;
+    const __nv_bfloat16* X = hidden_states + start_row * H;
+    __nv_bfloat16* Y = output + start_row * I;
+
+    int tid = threadIdx.x;
+    int warp_id = tid / WARP_SIZE;
+    int lane_id = tid % WARP_SIZE;
+
+    int total_blocks = gridDim.x;
+    int block_idx_in_expert = blockIdx.y;
+    int num_blocks_per_expert = gridDim.y;
+
+    int row_start = block_idx_in_expert * BLOCK_M;
+    int row_end = min(row_start + BLOCK_M, N_e);
+
+    if (row_start >= N_e) return;
+
+    int num_rows_this_block = row_end - row_start;
+    int num_tiles_k = (H + BLOCK_K - 1) / BLOCK_K;
+
+    __nv_bfloat16* smem = (__nv_bfloat16*)smem_buf;
+
+    for (int row = row_start + lane_id; row < row_end; row += WARP_SIZE) {
+        for (int ki = 0; ki < num_tiles_k; ++ki) {
+            if (ki > 0) {
+                __syncthreads();
+            }
+            int k_start = ki * BLOCK_K;
+            int k_len = min(BLOCK_K, H - k_start);
+
+            if (lane_id < k_len) {
+                for (int jj = 0; jj < BLOCK_N; ++jj) {
+                    int col = jj;
+                    int k_idx = k_start + lane_id;
+                    smem[jj * BLOCK_K + lane_id] = X[row * H + k_idx];
+                }
+            }
+            if (lane_id < BLOCK_K && lane_id >= k_len) {
+                for (int jj = 0; jj < BLOCK_N; ++jj) {
+                    smem[jj * BLOCK_K + lane_id] = __float_to_bfloat16(0.0f);
+                }
+            }
+
+            for (int jj = 0; jj < BLOCK_N; ++jj) {
+                int w_idx = k_start * I + jj * BLOCK_K + lane_id;
+                if (k_start + lane_id < H && jj < I) {
+                    smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + lane_id] = Wg[w_idx];
+                } else {
+                    smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + lane_id] = __float_to_bfloat16(0.0f);
+                }
+            }
+
+            __syncthreads();
+
+            for (int ii = 0; ii < BLOCK_M; ++ii) {
+                if (row_start + ii < row_end) {
+                    float acc = 0.0f;
+                    for (int kk = 0; kk < BLOCK_K; ++kk) {
+                        __nv_bfloat16 a_val = smem[ii * BLOCK_K + kk];
+                        __nv_bfloat16 b_val = smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + kk];
+                        acc += __bfloat16_to_float(a_val) * __bfloat16_to_float(b_val);
+                    }
+                    int out_idx = (row_start + ii) * I + jj;
+                    if (ki == 0) {
+                        Y[out_idx] = __float_to_bfloat16(acc);
+                    } else {
+                        Y[out_idx] = __float_to_bfloat16(__bfloat16_to_float(Y[out_idx]) + acc);
+                    }
+                }
+            }
+            __syncthreads();
+        }
+    }
+}
+
+template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
+__global__ void grouped_gemm_swiglu_kernel_v2(
+    const __nv_bfloat16* __restrict__ hidden_states,
+    const __nv_bfloat16* __restrict__ W_gate,
+    const __nv_bfloat16* __restrict__ W_up,
+    const int* __restrict__ expert_offsets,
+    __nv_bfloat16* __restrict__ output,
+    int T_perm, int H, int I, int E
+) {
+    extern __shared__ char smem_buf[];
+
+    int expert_idx = blockIdx.x;
+    if (expert_idx >= E) return;
+
+    int start_row = expert_offsets[expert_idx];
+    int end_row = expert_offsets[expert_idx + 1];
+    int N_e = end_row - start_row;
+    if (N_e <= 0) return;
+
+    const __nv_bfloat16* Wg = W_gate + expert_idx * H * I;
+    const __nv_bfloat16* Wu = W_up + expert_idx * H * I;
+    const __nv_bfloat16* X = hidden_states + start_row * H;
+    __nv_bfloat16* Y = output + start_row * I;
+
+    int tid = threadIdx.x;
+    int num_threads = blockDim.x;
+
+    int total_blocks = gridDim.x;
+    int block_idx = blockIdx.y;
+    int blocks_per_expert = gridDim.y;
+
+    constexpr int TILES_M = 4;
+    constexpr int TILES_N = 4;
+    constexpr int THREADS_PER_ROW = 8;
+
+    int row_tile = block_idx / TILES_N;
+    int col_tile = block_idx % TILES_N;
+
+    int row_start = row_tile * BLOCK_M;
+    int row_end = min(row_start + BLOCK_M, N_e);
+
+  

... (16508 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
@@ -5,44 +5,43 @@
 """
 from __future__ import annotations
 
-import math
-from typing import List
-
 import torch
 import torch.nn as nn
 from torch.utils.cpp_extension import load_inline
 
 cuda_source = r"""
-#include <cuda/bfloat16.h>
-#include <cuda_pipeline.h>
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
 #include <mma.h>
 
 using namespace nvcuda;
-using cutlass::pipe;
 
 constexpr int WARP_SIZE = 32;
-constexpr int MMA_M = 16;
-constexpr int MMA_N = 16;
-constexpr int MMA_K = 16;
-
-// SwiGLU activation: silu(x) * x = x / (1 + exp(-x)) * x
-__device__ __forceinline__ __nv_bfloat16 sigmoid_swish(__nv_bfloat16 x) {
-    float fx = __bfloat16_to_float(x);
-    float sig = fx / (1.0f + expf(-fx));
-    return __float_to_bfloat16(sig * fx);
-}
-
+
+// SwiGLU: silu(x) * y = x / (1 + exp(-x)) * y
+__device__ __forceinline__ float sigmoid_swish(float x) {
+    return x / (1.0f + expf(-x));
+}
+
+__device__ __forceinline__ __nv_bfloat162 sigmoid_swish(__nv_bfloat162 x) {
+    float2 fx = __bfloat1622float2(x);
+    fx.x = sigmoid_swish(fx.x);
+    fx.y = sigmoid_swish(fx.y);
+    return __float22bfloat162_rn(fx);
+}
+
+// Kernel for small experts (when N_e <= 256)
 template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
-__global__ void grouped_gemm_swiglu_kernel(
-    const __nv_bfloat16* __restrict__ hidden_states,  // (T_perm, H)
-    const __nv_bfloat16* __restrict__ W_gate,         // (E, H, I)
-    const __nv_bfloat16* __restrict__ W_up,          // (E, H, I)
-    const int* __restrict__ expert_offsets,           // (E+1,)
-    __nv_bfloat16* __restrict__ output,              // (T_perm, I)
+__global__ void grouped_gemm_swiglu_kernel_small(
+    const __nv_bfloat16* __restrict__ hidden_states,
+    const __nv_bfloat16* __restrict__ W_gate,
+    const __nv_bfloat16* __restrict__ W_up,
+    const int* __restrict__ expert_offsets,
+    __nv_bfloat16* __restrict__ gate_out,
+    __nv_bfloat16* __restrict__ up_out,
     int T_perm, int H, int I, int E
 ) {
-    extern __shared__ char smem_buf[];
-
     int expert_idx = blockIdx.x;
     if (expert_idx >= E) return;
 
@@ -54,480 +53,153 @@
     const __nv_bfloat16* Wg = W_gate + expert_idx * H * I;
     const __nv_bfloat16* Wu = W_up + expert_idx * H * I;
     const __nv_bfloat16* X = hidden_states + start_row * H;
-    __nv_bfloat16* Y = output + start_row * I;
-
-    int tid = threadIdx.x;
-    int warp_id = tid / WARP_SIZE;
-    int lane_id = tid % WARP_SIZE;
-
-    int total_blocks = gridDim.x;
-    int block_idx_in_expert = blockIdx.y;
-    int num_blocks_per_expert = gridDim.y;
-
-    int row_start = block_idx_in_expert * BLOCK_M;
-    int row_end = min(row_start + BLOCK_M, N_e);
-
-    if (row_start >= N_e) return;
-
-    int num_rows_this_block = row_end - row_start;
-    int num_tiles_k = (H + BLOCK_K - 1) / BLOCK_K;
-
-    __nv_bfloat16* smem = (__nv_bfloat16*)smem_buf;
-
-    for (int row = row_start + lane_id; row < row_end; row += WARP_SIZE) {
-        for (int ki = 0; ki < num_tiles_k; ++ki) {
-            if (ki > 0) {
-                __syncthreads();
-            }
-            int k_start = ki * BLOCK_K;
-            int k_len = min(BLOCK_K, H - k_start);
-
-            if (lane_id < k_len) {
-                for (int jj = 0; jj < BLOCK_N; ++jj) {
-                    int col = jj;
-                    int k_idx = k_start + lane_id;
-                    smem[jj * BLOCK_K + lane_id] = X[row * H + k_idx];
-                }
-            }
-            if (lane_id < BLOCK_K && lane_id >= k_len) {
-                for (int jj = 0; jj < BLOCK_N; ++jj) {
-                    smem[jj * BLOCK_K + lane_id] = __float_to_bfloat16(0.0f);
-                }
-            }
-
-            for (int jj = 0; jj < BLOCK_N; ++jj) {
-                int w_idx = k_start * I + jj * BLOCK_K + lane_id;
-                if (k_start + lane_id < H && jj < I) {
-                    smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + lane_id] = Wg[w_idx];
-                } else {
-                    smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + lane_id] = __float_to_bfloat16(0.0f);
-                }
-            }
-
-            __syncthreads();
-
-            for (int ii = 0; ii < BLOCK_M; ++ii) {
-                if (row_start + ii < row_end) {
-                    float acc = 0.0f;
-                    for (int kk = 0; kk < BLOCK_K; ++kk) {
-                        __nv_bfloat16 a_val = smem[ii * BLOCK_K + kk];
-                        __nv_bfloat16 b_val = smem[BLOCK_N * BLOCK_K + jj * BLOCK_K + kk];
-                        acc += __bfloat16_to_float(a_val) * __bfloat16_to_float(b_val);
-                    }
-                    int out_idx = (row_start + ii) * I + jj;
-                    if (ki == 0) {
-                        Y[out_idx] = __float_to_bfloat16(acc);
-                    } else {
-                        Y[out_idx] = __float_to_bfloat16(__bfloat16_to_float(Y[out_idx]) + acc);
-                    }
-                }
-            }
-            __syncthreads();
-        }
-    }
-}
-
-template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
-__global__ void grouped_gemm_swiglu_kernel_v2(
-    const __nv_bfloat16* __restrict__ hidden_states,
-    const __nv_bfloat16* __restrict__ W_gate,
-    const __nv_bfloat16* __restrict__ W_up,
-    const int* __restrict__ expert_offsets,
-    __nv_bfloat16* __restrict__ output,
-    int T_perm, int H, int I, int E
-) {
-    extern __shared__ char smem_buf[];
-
-    int expert_idx = blockIdx.x;
-    if (expert_idx >= E) return;
-
-    int start_row = expert_offsets[expert_idx];
-    int end_row = expert_offsets[expert_idx + 1];
-    int N_e = end_row - start_row;
-    if (N_e <= 0) return;
-
-    const __nv_bfloat16* Wg = W_gate + expert_idx * H * I;
-    const __nv_bfloat16* Wu = W_up + expert_idx * H * I;
-    const __nv_b

... (20382 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,229 +1,184 @@
 """Grouped GEMM + fused SwiGLU for MoE up-projection on Blackwell (SM120).
 
 Implements: h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e])
-using custom CUDA kernel with tensor cores and load_inline.
+using Triton kernel with grouped GEMM.
 """
 from __future__ import annotations
 
 import torch
 import torch.nn as nn
-from torch.utils.cpp_extension import load_inline
-
-cuda_source = r"""
-#include <torch/extension.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <mma.h>
-
-using namespace nvcuda;
-
-constexpr int WARP_SIZE = 32;
-
-// SwiGLU: silu(x) * y = x / (1 + exp(-x)) * y
-__device__ __forceinline__ float sigmoid_swish(float x) {
-    return x / (1.0f + expf(-x));
-}
-
-__device__ __forceinline__ __nv_bfloat162 sigmoid_swish(__nv_bfloat162 x) {
-    float2 fx = __bfloat1622float2(x);
-    fx.x = sigmoid_swish(fx.x);
-    fx.y = sigmoid_swish(fx.y);
-    return __float22bfloat162_rn(fx);
-}
-
-// Kernel for small experts (when N_e <= 256)
-template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
-__global__ void grouped_gemm_swiglu_kernel_small(
-    const __nv_bfloat16* __restrict__ hidden_states,
-    const __nv_bfloat16* __restrict__ W_gate,
-    const __nv_bfloat16* __restrict__ W_up,
-    const int* __restrict__ expert_offsets,
-    __nv_bfloat16* __restrict__ gate_out,
-    __nv_bfloat16* __restrict__ up_out,
-    int T_perm, int H, int I, int E
-) {
-    int expert_idx = blockIdx.x;
-    if (expert_idx >= E) return;
-
-    int start_row = expert_offsets[expert_idx];
-    int end_row = expert_offsets[expert_idx + 1];
-    int N_e = end_row - start_row;
-    if (N_e <= 0) return;
-
-    const __nv_bfloat16* Wg = W_gate + expert_idx * H * I;
-    const __nv_bfloat16* Wu = W_up + expert_idx * H * I;
-    const __nv_bfloat16* X = hidden_states + start_row * H;
-
-    int row_block = blockIdx.y;
-    int col_block = blockIdx.z;
-
-    int row_start = row_block * BLOCK_M;
-    int col_start = col_block * BLOCK_N;
-
-    if (row_start >= N_e || col_start >= I) return;
-
-    int lane_id = threadIdx.x % WARP_SIZE;
-    int warp_id = threadIdx.x / WARP_SIZE;
-    int num_warps = blockDim.x / WARP_SIZE;
-
-    int num_k_tiles = (H + BLOCK_K - 1) / BLOCK_K;
-
-    float gate_acc[2] = {0.0f, 0.0f};
-    float up_acc[2] = {0.0f, 0.0f};
-
-    __nv_bfloat16 a_reg[2];
-    __nv_bfloat16 b_gate_reg[2];
-    __nv_bfloat16 b_up_reg[2];
-
-    for (int kt = 0; kt < num_k_tiles; ++kt) {
-        for (int mi = 0; mi < 2; ++mi) {
-            int row = row_start + mi * 16 + warp_id;
-            int k = kt * BLOCK_K + lane_id;
-            if (row < N_e && k < H) {
-                a_reg[mi] = X[row * H + k];
-            } else {
-                a_reg[mi] = __float_to_bfloat16(0.0f);
-            }
-        }
-
-        for (int ni = 0; ni < 2; ++ni) {
-            int col = col_start + ni * 16 + (lane_id % 16);
-            int k = kt * BLOCK_K + warp_id;
-            if (col < I && k < H) {
-                b_gate_reg[ni] = Wg[k * I + col];
-                b_up_reg[ni] = Wu[k * I + col];
-            } else {
-                b_gate_reg[ni] = __float_to_bfloat16(0.0f);
-                b_up_reg[ni] = __float_to_bfloat16(0.0f);
-            }
-        }
-
-        for (int mi = 0; mi < 2; ++mi) {
-            for (int ni = 0; ni < 2; ++ni) {
-                gate_acc[mi] += __bfloat16_to_float(a_reg[mi]) * __bfloat16_to_float(b_gate_reg[ni]);
-                up_acc[mi] += __bfloat16_to_float(a_reg[mi]) * __bfloat16_to_float(b_up_reg[ni]);
-            }
-        }
-    }
-
-    for (int mi = 0; mi < 2; ++mi) {
-        int row = row_start + mi * 16 + warp_id;
-        for (int ni = 0; ni < 2; ++ni) {
-            int col = col_start + ni * 16 + (lane_id % 16);
-            if (row < N_e && col < I) {
-                gate_out[start_row * I + row * I + col] = __float_to_bfloat16(gate_acc[mi]);
-                up_out[start_row * I + row * I + col] = __float_to_bfloat16(up_acc[mi]);
-            }
-        }
-    }
-}
-
-// Kernel for applying SwiGLU and combining gate+up outputs
-template <int BLOCK_M, int BLOCK_N>
-__global__ void swiglu_combine_kernel(
-    __nv_bfloat16* gate_out,
-    __nv_bfloat16* up_out,
-    __nv_bfloat16* output,
-    int T_perm, int I
-) {
-    int row = blockIdx.x * BLOCK_M + threadIdx.x;
-    int col = blockIdx.y * BLOCK_N + threadIdx.y;
-
-    if (row < T_perm && col < I) {
-        float gate_val = __bfloat16_to_float(gate_out[row * I + col]);
-        float up_val = __bfloat16_to_float(up_out[row * I + col]);
-        float result = sigmoid_swish(gate_val) * up_val;
-        output[row * I + col] = __float_to_bfloat16(result);
-    }
-}
-
-// Main launcher
-void grouped_gemm_swiglu_cuda(
-    const torch::Tensor& hidden_states,
-    const torch::Tensor& W_gate,
-    const torch::Tensor& W_up,
-    const torch::Tensor& expert_offsets,
-    torch::Tensor& output,
-    int T_total, int H, int I, int E, int K
-) {
-    int T_perm = T_total * K;
-
-    constexpr int BLOCK_M = 32;
-    constexpr int BLOCK_N = 32;
-    constexpr int BLOCK_K = 64;
-
-    dim3 gemm_grid(E, 8, (I + BLOCK_N - 1) / BLOCK_N);
-
-    auto gate_out = torch::empty({T_perm, I}, hidden_states.options());
-    auto up_out = torch::empty({T_perm, I}, hidden_states.options());
-
-    grouped_gemm_swiglu_kernel_small<BLOCK_M, BLOCK_N, BLOCK_K><<<gemm_grid, 256>>>(
-        reinterpret_cast<const __nv_bfloat16*>(hidden_states.data_ptr()),
-        reinterpret_cast<const __nv_bfloat16*>(W_gate.data_ptr()),
-        reinterpret_cast<const __nv_bfloat16*>(W_up.data_ptr()),
-        reinterpret_cast<const int*>(expert_offsets.data_ptr()),
-        reinterpret_cast<__nv_bfloat16*>(gate_out.data_ptr()),
-        reinterpret_cast<__nv_bfloat16*>(up_out.data_ptr()),
-        T_perm, H, I, E
-    );
-
-    dim3 swiglu_grid

... (8761 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
@@ -11,24 +11,16 @@
 import triton.language as tl
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64}, num_stages=3, num_warps=4),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=3, num_warps=4),
-    ],
-    key=["N_e", "I", "H"],
-)
 @triton.jit
-def triton_gemm_kernel_gate(
+def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
     a_stride, b_stride, c_stride,
     N_e, I, H,
-    block_m: tl.constexpr, block_n: tl.constexpr, block_k: tl.constexpr,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(N_e, block_m)
-    num_pid_n = tl.cdiv(I, block_n)
+    num_pid_m = tl.cdiv(N_e, BLOCK_M)
+    num_pid_n = tl.cdiv(I, BLOCK_N)
     num_pid_in_group = num_pid_m * num_pid_n
     group_id = pid // num_pid_in_group
     first_pid_m = group_id * num_pid_m
@@ -37,15 +29,15 @@
     pid_m = first_pid_m + (pid % group_size_m)
     pid_n = (pid % num_pid_in_group) // group_size_m
 
-    offs_m = (pid_m * block_m + tl.arange(0, block_m)) % N_e
-    offs_n = (pid_n * block_n + tl.arange(0, block_n)) % I
-    offs_k = tl.arange(0, block_k)
+    offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % N_e
+    offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % I
+    offs_k = tl.arange(0, BLOCK_K)
 
     a_ptrs = a_ptr + (offs_m[:, None] * a_stride + offs_k[None, :] * H)
     b_ptrs = b_ptr + (offs_k[:, None] * I + offs_n[None, :])
 
-    accumulator = tl.zeros((block_m, block_n), dtype=tl.float32)
-    for k in range(0, tl.cdiv(H, block_k)):
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(H, BLOCK_K)):
         a_mask = (offs_m[:, None] < N_e) & (offs_k[None, :] < H)
         b_mask = (offs_k[:, None] < H) & (offs_n[None, :] < I)
 
@@ -54,82 +46,18 @@
 
         accumulator += tl.dot(a, b)
 
-        a_ptrs += block_k * H
-        b_ptrs += block_k * I
-        offs_k += block_k
+        a_ptrs += BLOCK_K * H
+        b_ptrs += BLOCK_K * I
+        offs_k += BLOCK_K
 
     c_ptrs = c_ptr + (offs_m[:, None] * I + offs_n[None, :])
     c_mask = (offs_m[:, None] < N_e) & (offs_n[None, :] < I)
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 64}, num_stages=3, num_warps=4),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 32, "BLOCK_K": 64}, num_stages=3, num_warps=4),
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32}, num_stages=3, num_warps=4),
-    ],
-    key=["N_e", "I", "H"],
-)
-@triton.jit
-def triton_gemm_kernel_up(
-    a_ptr, b_ptr, c_ptr,
-    a_stride, b_stride, c_stride,
-    N_e, I, H,
-    block_m: tl.constexpr, block_n: tl.constexpr, block_k: tl.constexpr,
-):
-    pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(N_e, block_m)
-    num_pid_n = tl.cdiv(I, block_n)
-    num_pid_in_group = num_pid_m * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * num_pid_m
-    group_size_m = min(num_pid_m, N_e - first_pid_m)
-
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m
-
-    offs_m = (pid_m * block_m + tl.arange(0, block_m)) % N_e
-    offs_n = (pid_n * block_n + tl.arange(0, block_n)) % I
-    offs_k = tl.arange(0, block_k)
-
-    a_ptrs = a_ptr + (offs_m[:, None] * a_stride + offs_k[None, :] * H)
-    b_ptrs = b_ptr + (offs_k[:, None] * I + offs_n[None, :])
-
-    accumulator = tl.zeros((block_m, block_n), dtype=tl.float32)
-    for k in range(0, tl.cdiv(H, block_k)):
-        a_mask = (offs_m[:, None] < N_e) & (offs_k[None, :] < H)
-        b_mask = (offs_k[:, None] < H) & (offs_n[None, :] < I)
-
-        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
-        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
-
-        accumulator += tl.dot(a, b)
-
-        a_ptrs += block_k * H
-        b_ptrs += block_k * I
-        offs_k += block_k
-
-    c_ptrs = c_ptr + (offs_m[:, None] * I + offs_n[None, :])
-    c_mask = (offs_m[:, None] < N_e) & (offs_n[None, :] < I)
-    tl.store(c_ptrs, accumulator, mask=c_mask)
-
-
-@triton.jit
-def swiglu_epilogue(output, gate, up, stride, N, I):
-    pid = tl.program_id(axis=0)
-    offs_m = pid
-    offs_n = tl.arange(0, I)
-    gate_val = tl.load(gate + offs_m * stride + offs_n)
-    up_val = tl.load(up + offs_m * stride + offs_n)
-    silu_gate = gate_val / (1.0 + tl.exp(-gate_val))
-    result = silu_gate * up_val
-    tl.store(output + offs_m * stride + offs_n, result)
-
-
-def triton_grouped_gemm_gate(hidden_states, W_gate, expert_offsets, E):
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E, is_gate=True):
     T_perm, H = hidden_states.shape
-    gate_out = torch.empty_like(hidden_states)
+    output = torch.empty_like(hidden_states)
 
     for e in range(E):
         start = expert_offsets[e].item()
@@ -137,48 +65,26 @@
         if end == start:
             continue
         N_e = end - start
-        I = W_gate.shape[2]
+        I = W.shape[2]
 
         x_e = hidden_states[start:end]
-        w_e = W_gate[e]
+        w_e = W[e]
 
-        grid_m = triton.cdiv(N_e, 32)
+        grid_m = triton.cdiv(N_e, 64)
         grid_n = triton.cdiv(I, 64)
         grid = (grid_m * grid_n,)
 
-        triton_gemm_kernel_gate[grid](
-            x_e, w_e, gate_out[start:end],
-            x_e.stride(0), w_e.stride(0), gate_out[start:end].stride(0),
+        BLOCK_M = 64
+        BLOCK_N = 64
+        BLOCK_K = 64
+
+        triton_gemm_kernel[grid](
+            x_e, w_e, output[start:end],
+            x_e.stride(0), w_e.stride(0), output[start:end].stride(0),
             N_e, I

... (1538 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
@@ -14,7 +14,6 @@
 @triton.jit
 def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
-    a_stride, b_stride, c_stride,
     N_e, I, H,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
@@ -29,20 +28,22 @@
     pid_m = first_pid_m + (pid % group_size_m)
     pid_n = (pid % num_pid_in_group) // group_size_m
 
-    offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % N_e
-    offs_n = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % I
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
 
-    a_ptrs = a_ptr + (offs_m[:, None] * a_stride + offs_k[None, :] * H)
-    b_ptrs = b_ptr + (offs_k[:, None] * I + offs_n[None, :])
+    a_ptrs = a_ptr + offs_m[:, None] * H + offs_k[None, :] * 1
+    b_ptrs = b_ptr + offs_k[:, None] * I + offs_n[None, :] * 1
 
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(H, BLOCK_K)):
-        a_mask = (offs_m[:, None] < N_e) & (offs_k[None, :] < H)
-        b_mask = (offs_k[:, None] < H) & (offs_n[None, :] < I)
+        mask_m = offs_m[:, None] < N_e
+        mask_k = offs_k[None, :] < H
+        a = tl.load(a_ptrs, mask=mask_m & mask_k, other=0.0)
 
-        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
-        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+        mask_k = offs_k[:, None] < H
+        mask_n = offs_n[None, :] < I
+        b = tl.load(b_ptrs, mask=mask_k & mask_n, other=0.0)
 
         accumulator += tl.dot(a, b)
 
@@ -50,13 +51,17 @@
         b_ptrs += BLOCK_K * I
         offs_k += BLOCK_K
 
-    c_ptrs = c_ptr + (offs_m[:, None] * I + offs_n[None, :])
-    c_mask = (offs_m[:, None] < N_e) & (offs_n[None, :] < I)
-    tl.store(c_ptrs, accumulator, mask=c_mask)
+    offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    c_ptrs = c_ptr + offs_m_out[:, None] * I + offs_n_out[None, :] * 1
+    mask_m_out = offs_m_out[:, None] < N_e
+    mask_n_out = offs_n_out[None, :] < I
+    tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E, is_gate=True):
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=32, BLOCK_N=32, BLOCK_K=64):
     T_perm, H = hidden_states.shape
+    I = W.shape[2]
     output = torch.empty_like(hidden_states)
 
     for e in range(E):
@@ -65,22 +70,17 @@
         if end == start:
             continue
         N_e = end - start
-        I = W.shape[2]
 
-        x_e = hidden_states[start:end]
-        w_e = W[e]
+        x_e_ptr = hidden_states + start * H
+        w_e_ptr = W + e * H * I
+        out_e_ptr = output + start * I
 
-        grid_m = triton.cdiv(N_e, 64)
-        grid_n = triton.cdiv(I, 64)
+        grid_m = triton.cdiv(N_e, BLOCK_M)
+        grid_n = triton.cdiv(I, BLOCK_N)
         grid = (grid_m * grid_n,)
 
-        BLOCK_M = 64
-        BLOCK_N = 64
-        BLOCK_K = 64
-
         triton_gemm_kernel[grid](
-            x_e, w_e, output[start:end],
-            x_e.stride(0), w_e.stride(0), output[start:end].stride(0),
+            x_e_ptr, w_e_ptr, out_e_ptr,
             N_e, I, H,
             BLOCK_M, BLOCK_N, BLOCK_K,
         )
@@ -101,8 +101,8 @@
         nn.init.normal_(self.W_up, std=0.02)
 
     def forward(self, hidden_states: torch.Tensor, expert_offsets: torch.Tensor) -> torch.Tensor:
-        gate_out = triton_grouped_gemm(hidden_states, self.W_gate, expert_offsets, self.E, is_gate=True)
-        up_out = triton_grouped_gemm(hidden_states, self.W_up, expert_offsets, self.E, is_gate=False)
+        gate_out = triton_grouped_gemm(hidden_states, self.W_gate, expert_offsets, self.E)
+        up_out = triton_grouped_gemm(hidden_states, self.W_up, expert_offsets, self.E)
         silu_gate = torch.nn.functional.silu(gate_out)
         return silu_gate * up_out
 
@@ -133,4 +133,4 @@
 
 
 def get_init_inputs():
-    return [T_total, H, I, E, K]
+    return [T_total, H, I, E, K]

--- 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
@@ -15,6 +15,7 @@
 def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
     N_e, I, H,
+    stride_a, stride_b, stride_c,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
@@ -32,8 +33,8 @@
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
 
-    a_ptrs = a_ptr + offs_m[:, None] * H + offs_k[None, :] * 1
-    b_ptrs = b_ptr + offs_k[:, None] * I + offs_n[None, :] * 1
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * stride_b
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * stride_c
 
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(H, BLOCK_K)):
@@ -47,41 +48,42 @@
 
         accumulator += tl.dot(a, b)
 
-        a_ptrs += BLOCK_K * H
-        b_ptrs += BLOCK_K * I
+        a_ptrs += BLOCK_K * stride_b
+        b_ptrs += BLOCK_K * stride_b
         offs_k += BLOCK_K
 
     offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    c_ptrs = c_ptr + offs_m_out[:, None] * I + offs_n_out[None, :] * 1
+    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c + offs_n_out[None, :] * 1
     mask_m_out = offs_m_out[:, None] < N_e
     mask_n_out = offs_n_out[None, :] < I
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=32, BLOCK_N=32, BLOCK_K=64):
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=32, BLOCK_N=64, BLOCK_K=64):
     T_perm, H = hidden_states.shape
     I = W.shape[2]
     output = torch.empty_like(hidden_states)
 
     for e in range(E):
-        start = expert_offsets[e].item()
-        end = expert_offsets[e + 1].item()
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
         if end == start:
             continue
         N_e = end - start
 
-        x_e_ptr = hidden_states + start * H
-        w_e_ptr = W + e * H * I
-        out_e_ptr = output + start * I
+        x_e = hidden_states[start:end]
+        w_e = W[e]
+        out_e = output[start:end]
 
         grid_m = triton.cdiv(N_e, BLOCK_M)
         grid_n = triton.cdiv(I, BLOCK_N)
         grid = (grid_m * grid_n,)
 
         triton_gemm_kernel[grid](
-            x_e_ptr, w_e_ptr, out_e_ptr,
+            x_e, w_e, out_e,
             N_e, I, H,
+            x_e.stride(0), w_e.stride(0), out_e.stride(0),
             BLOCK_M, BLOCK_N, BLOCK_K,
         )
     return output

--- 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
@@ -33,8 +33,8 @@
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
 
-    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * stride_b
-    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * stride_c
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * 1
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * 1
 
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(H, BLOCK_K)):
@@ -48,8 +48,8 @@
 
         accumulator += tl.dot(a, b)
 
-        a_ptrs += BLOCK_K * stride_b
-        b_ptrs += BLOCK_K * stride_b
+        a_ptrs += BLOCK_K * stride_a
+        b_ptrs += BLOCK_K * 1
         offs_k += BLOCK_K
 
     offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
@@ -60,7 +60,7 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=32, BLOCK_N=64, BLOCK_K=64):
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=64, BLOCK_N=64, BLOCK_K=64):
     T_perm, H = hidden_states.shape
     I = W.shape[2]
     output = torch.empty_like(hidden_states)
@@ -83,7 +83,7 @@
         triton_gemm_kernel[grid](
             x_e, w_e, out_e,
             N_e, I, H,
-            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            x_e.stride(0), w_e.stride(1), out_e.stride(0),
             BLOCK_M, BLOCK_N, BLOCK_K,
         )
     return output

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -0,0 +1,20 @@
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def simple_kernel(a_ptr, c_ptr, N, BLOCK_M: tl.constexpr):
+    pid = tl.program_id(axis=0)
+    offs = pid * BLOCK_M + tl.arange(0, BLOCK_M)
+    mask = offs < N
+    vals = tl.load(a_ptr + offs * 1, mask=mask, other=0.0)
+    tl.store(c_ptr + offs * 1, vals, mask=mask)
+
+N = 100
+BLOCK_M = 32
+a = torch.randn(N, device='cuda')
+c = torch.zeros_like(a)
+
+grid = (triton.cdiv(N, BLOCK_M),)
+simple_kernel[grid](a, c, N, BLOCK_M)
+print('Result matches:', torch.allclose(a, c))

--- 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
@@ -60,7 +60,7 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E, BLOCK_M=64, BLOCK_N=64, BLOCK_K=64):
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
     T_perm, H = hidden_states.shape
     I = W.shape[2]
     output = torch.empty_like(hidden_states)
@@ -76,15 +76,15 @@
         w_e = W[e]
         out_e = output[start:end]
 
-        grid_m = triton.cdiv(N_e, BLOCK_M)
-        grid_n = triton.cdiv(I, BLOCK_N)
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
         grid = (grid_m * grid_n,)
 
         triton_gemm_kernel[grid](
             x_e, w_e, out_e,
             N_e, I, H,
-            x_e.stride(0), w_e.stride(1), out_e.stride(0),
-            BLOCK_M, BLOCK_N, BLOCK_K,
+            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            64, 64, 64,
         )
     return output
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,20 +1,127 @@
+"""Test Triton grouped GEMM"""
 import torch
 import triton
 import triton.language as tl
+import solution
 
 @triton.jit
-def simple_kernel(a_ptr, c_ptr, N, BLOCK_M: tl.constexpr):
+def triton_gemm_kernel(
+    a_ptr, b_ptr, c_ptr,
+    N_e, I, H,
+    stride_a, stride_b, stride_c,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
     pid = tl.program_id(axis=0)
-    offs = pid * BLOCK_M + tl.arange(0, BLOCK_M)
-    mask = offs < N
-    vals = tl.load(a_ptr + offs * 1, mask=mask, other=0.0)
-    tl.store(c_ptr + offs * 1, vals, mask=mask)
+    num_pid_m = tl.cdiv(N_e, BLOCK_M)
+    num_pid_n = tl.cdiv(I, BLOCK_N)
+    num_pid_in_group = num_pid_m * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * num_pid_m
+    group_size_m = min(num_pid_m, N_e - first_pid_m)
 
-N = 100
-BLOCK_M = 32
-a = torch.randn(N, device='cuda')
-c = torch.zeros_like(a)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
 
-grid = (triton.cdiv(N, BLOCK_M),)
-simple_kernel[grid](a, c, N, BLOCK_M)
-print('Result matches:', torch.allclose(a, c))+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * 1
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * 1
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(H, BLOCK_K)):
+        mask_m = offs_m[:, None] < N_e
+        mask_k = offs_k[None, :] < H
+        a = tl.load(a_ptrs, mask=mask_m & mask_k, other=0.0)
+
+        mask_k = offs_k[:, None] < H
+        mask_n = offs_n[None, :] < I
+        b = tl.load(b_ptrs, mask=mask_k & mask_n, other=0.0)
+
+        accumulator += tl.dot(a, b)
+
+        a_ptrs += BLOCK_K * stride_a
+        b_ptrs += BLOCK_K * 1
+        offs_k += BLOCK_K
+
+    offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c + offs_n_out[None, :] * 1
+    mask_m_out = offs_m_out[:, None] < N_e
+    mask_n_out = offs_n_out[None, :] < I
+    tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
+
+
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
+    T_perm, H = hidden_states.shape
+    I = W.shape[2]
+    output = torch.empty_like(hidden_states)
+
+    for e in range(E):
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        if end == start:
+            continue
+        N_e = end - start
+
+        x_e = hidden_states[start:end]
+        w_e = W[e]
+        out_e = output[start:end]
+
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
+        grid = (grid_m * grid_n,)
+
+        triton_gemm_kernel[grid](
+            x_e, w_e, out_e,
+            N_e, I, H,
+            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            64, 64, 64,
+        )
+    return output
+
+
+# Test with small shape
+solution.T_total = 4096
+solution.H = 2048
+solution.I = 1024
+solution.E = 64
+solution.K = 4
+
+model = solution.Model(*solution.get_init_inputs())
+model = model.cuda()
+inputs = [t.cuda() for t in solution.get_inputs()]
+print('Input shapes ok')
+print('W_gate shape:', model.W_gate.shape)
+print('expert_offsets[:5]:', inputs[1][:5])
+
+# Try first expert only
+e = 0
+start = int(inputs[1][e].item())
+end = int(inputs[1][e+1].item())
+print(f'Expert 0: rows {start}:{end}')
+print(f'x_e shape: {inputs[0][start:end].shape}')
+print(f'w_e shape: {model.W_gate[e].shape}')
+
+# Try gemm for just expert 0
+N_e = end - start
+I = model.W_gate.shape[2]
+H = model.H
+
+x_e = inputs[0][start:end]
+w_e = model.W_gate[e]
+out_e = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+
+grid_m = triton.cdiv(N_e, 64)
+grid_n = triton.cdiv(I, 64)
+grid = (grid_m * grid_n,)
+print(f'Grid: {grid}, N_e={N_e}, I={I}, H={H}')
+
+triton_gemm_kernel[grid](
+    x_e, w_e, out_e,
+    N_e, I, H,
+    x_e.stride(0), w_e.stride(0), out_e.stride(0),
+    64, 64, 64,
+)
+print('Expert 0 GEMM succeeded, out_e:', out_e.shape)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -53,35 +53,6 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
-    T_perm, H = hidden_states.shape
-    I = W.shape[2]
-    output = torch.empty_like(hidden_states)
-
-    for e in range(E):
-        start = int(expert_offsets[e].item())
-        end = int(expert_offsets[e + 1].item())
-        if end == start:
-            continue
-        N_e = end - start
-
-        x_e = hidden_states[start:end]
-        w_e = W[e]
-        out_e = output[start:end]
-
-        grid_m = triton.cdiv(N_e, 64)
-        grid_n = triton.cdiv(I, 64)
-        grid = (grid_m * grid_n,)
-
-        triton_gemm_kernel[grid](
-            x_e, w_e, out_e,
-            N_e, I, H,
-            x_e.stride(0), w_e.stride(0), out_e.stride(0),
-            64, 64, 64,
-        )
-    return output
-
-
 # Test with small shape
 solution.T_total = 4096
 solution.H = 2048
@@ -93,35 +64,37 @@
 model = model.cuda()
 inputs = [t.cuda() for t in solution.get_inputs()]
 print('Input shapes ok')
-print('W_gate shape:', model.W_gate.shape)
-print('expert_offsets[:5]:', inputs[1][:5])
 
-# Try first expert only
-e = 0
-start = int(inputs[1][e].item())
-end = int(inputs[1][e+1].item())
-print(f'Expert 0: rows {start}:{end}')
-print(f'x_e shape: {inputs[0][start:end].shape}')
-print(f'w_e shape: {model.W_gate[e].shape}')
+T_perm, H = inputs[0].shape
+I = model.W_gate.shape[2]
+output = torch.empty_like(inputs[0])
 
-# Try gemm for just expert 0
-N_e = end - start
-I = model.W_gate.shape[2]
-H = model.H
+expert_offsets = inputs[1]
+E = solution.E
 
-x_e = inputs[0][start:end]
-w_e = model.W_gate[e]
-out_e = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    if end == start:
+        continue
+    N_e = end - start
 
-grid_m = triton.cdiv(N_e, 64)
-grid_n = triton.cdiv(I, 64)
-grid = (grid_m * grid_n,)
-print(f'Grid: {grid}, N_e={N_e}, I={I}, H={H}')
+    x_e = inputs[0][start:end]
+    w_e = model.W_gate[e]
+    out_e = output[start:end]
 
-triton_gemm_kernel[grid](
-    x_e, w_e, out_e,
-    N_e, I, H,
-    x_e.stride(0), w_e.stride(0), out_e.stride(0),
-    64, 64, 64,
-)
-print('Expert 0 GEMM succeeded, out_e:', out_e.shape)+    grid_m = triton.cdiv(N_e, 64)
+    grid_n = triton.cdiv(I, 64)
+    grid = (grid_m * grid_n,)
+
+    triton_gemm_kernel[grid](
+        x_e, w_e, out_e,
+        N_e, I, H,
+        x_e.stride(0), w_e.stride(0), out_e.stride(0),
+        64, 64, 64,
+    )
+    if e % 10 == 0:
+        print(f'Expert {e} done, N_e={N_e}, grid={grid}')
+        torch.cuda.synchronize()
+
+print('All experts done')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -73,28 +73,35 @@
 E = solution.E
 
 for e in range(E):
-    start = int(expert_offsets[e].item())
-    end = int(expert_offsets[e + 1].item())
-    if end == start:
-        continue
-    N_e = end - start
+    try:
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        if end == start:
+            continue
+        N_e = end - start
 
-    x_e = inputs[0][start:end]
-    w_e = model.W_gate[e]
-    out_e = output[start:end]
+        x_e = inputs[0][start:end]
+        w_e = model.W_gate[e]
+        out_e = output[start:end]
 
-    grid_m = triton.cdiv(N_e, 64)
-    grid_n = triton.cdiv(I, 64)
-    grid = (grid_m * grid_n,)
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
+        grid = (grid_m * grid_n,)
 
-    triton_gemm_kernel[grid](
-        x_e, w_e, out_e,
-        N_e, I, H,
-        x_e.stride(0), w_e.stride(0), out_e.stride(0),
-        64, 64, 64,
-    )
-    if e % 10 == 0:
-        print(f'Expert {e} done, N_e={N_e}, grid={grid}')
+        triton_gemm_kernel[grid](
+            x_e, w_e, out_e,
+            N_e, I, H,
+            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            64, 64, 64,
+        )
         torch.cuda.synchronize()
+        if e % 10 == 0:
+            print(f'Expert {e} done, N_e={N_e}, grid={grid}')
+    except Exception as ex:
+        print(f'Expert {e} failed: {ex}')
+        print(f'  start={start}, end={end}, N_e={N_e}')
+        print(f'  x_e shape: {x_e.shape}')
+        print(f'  out_e shape: {out_e.shape}')
+        break
 
 print('All experts done')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -8,7 +8,9 @@
 def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
     N_e, I, H,
-    stride_a, stride_b, stride_c,
+    stride_a_row, stride_a_col,  # Row stride and col stride for A
+    stride_b_row, stride_b_col,  # Row stride and col stride for B
+    stride_c_row, stride_c_col,  # Row stride and col stride for C
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
@@ -26,8 +28,8 @@
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
 
-    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * 1
-    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * 1
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a_row + offs_k[None, :] * stride_a_col
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b_row + offs_n[None, :] * stride_b_col
 
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(H, BLOCK_K)):
@@ -41,13 +43,13 @@
 
         accumulator += tl.dot(a, b)
 
-        a_ptrs += BLOCK_K * stride_a
-        b_ptrs += BLOCK_K * 1
+        a_ptrs += BLOCK_K * stride_a_col
+        b_ptrs += BLOCK_K * stride_b_row
         offs_k += BLOCK_K
 
     offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c + offs_n_out[None, :] * 1
+    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c_row + offs_n_out[None, :] * stride_c_col
     mask_m_out = offs_m_out[:, None] < N_e
     mask_n_out = offs_n_out[None, :] < I
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
@@ -67,7 +69,9 @@
 
 T_perm, H = inputs[0].shape
 I = model.W_gate.shape[2]
-output = torch.empty_like(inputs[0])
+
+# Output is (T_perm, I), not (T_perm, H)
+output = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
 
 expert_offsets = inputs[1]
 E = solution.E
@@ -82,16 +86,18 @@
 
         x_e = inputs[0][start:end]
         w_e = model.W_gate[e]
-        out_e = output[start:end]
+        out_e_base = output + start * I  # Pointer to start position in output
 
         grid_m = triton.cdiv(N_e, 64)
         grid_n = triton.cdiv(I, 64)
         grid = (grid_m * grid_n,)
 
         triton_gemm_kernel[grid](
-            x_e, w_e, out_e,
+            x_e, w_e, out_e_base,
             N_e, I, H,
-            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            x_e.stride(0), 1,      # A: row stride, col stride
+            w_e.stride(0), w_e.stride(1),  # B: row stride, col stride
+            I, 1,                   # C: row stride, col stride
             64, 64, 64,
         )
         torch.cuda.synchronize()
@@ -100,8 +106,7 @@
     except Exception as ex:
         print(f'Expert {e} failed: {ex}')
         print(f'  start={start}, end={end}, N_e={N_e}')
-        print(f'  x_e shape: {x_e.shape}')
-        print(f'  out_e shape: {out_e.shape}')
         break
 
-print('All experts done')+print('All experts done')
+print('Output shape:', output.shape)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,4 +1,4 @@
-"""Test Triton grouped GEMM"""
+"""Verify correctness of Triton grouped GEMM + SwiGLU"""
 import torch
 import triton
 import triton.language as tl
@@ -8,9 +8,9 @@
 def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
     N_e, I, H,
-    stride_a_row, stride_a_col,  # Row stride and col stride for A
-    stride_b_row, stride_b_col,  # Row stride and col stride for B
-    stride_c_row, stride_c_col,  # Row stride and col stride for C
+    stride_a_row, stride_a_col,
+    stride_b_row, stride_b_col,
+    stride_c_row, stride_c_col,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
@@ -55,6 +55,37 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
+    T_perm, H = hidden_states.shape
+    I = W.shape[2]
+    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
+
+    for e in range(E):
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        if end == start:
+            continue
+        N_e = end - start
+
+        x_e = hidden_states[start:end]
+        w_e = W[e]
+        out_e_base = output + start * I
+
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
+        grid = (grid_m * grid_n,)
+
+        triton_gemm_kernel[grid](
+            x_e, w_e, out_e_base,
+            N_e, I, H,
+            x_e.stride(0), 1,
+            w_e.stride(0), w_e.stride(1),
+            I, 1,
+            64, 64, 64,
+        )
+    return output
+
+
 # Test with small shape
 solution.T_total = 4096
 solution.H = 2048
@@ -64,49 +95,25 @@
 
 model = solution.Model(*solution.get_init_inputs())
 model = model.cuda()
+
+torch.manual_seed(42)
 inputs = [t.cuda() for t in solution.get_inputs()]
-print('Input shapes ok')
 
-T_perm, H = inputs[0].shape
-I = model.W_gate.shape[2]
+# Reference
+ref_out = solution.reference.Model(*solution.get_init_inputs()).cuda()
+ref_out.load_state_dict(model.state_dict())
+ref_out.eval()
+with torch.no_grad():
+    ref_result = ref_out(*inputs)
 
-# Output is (T_perm, I), not (T_perm, H)
-output = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+# Triton GEMM + SwiGLU
+gate_out = triton_grouped_gemm(inputs[0], model.W_gate, inputs[1], solution.E)
+up_out = triton_grouped_gemm(inputs[0], model.W_up, inputs[1], solution.E)
+silu_gate = torch.nn.functional.silu(gate_out)
+triton_result = silu_gate * up_out
 
-expert_offsets = inputs[1]
-E = solution.E
-
-for e in range(E):
-    try:
-        start = int(expert_offsets[e].item())
-        end = int(expert_offsets[e + 1].item())
-        if end == start:
-            continue
-        N_e = end - start
-
-        x_e = inputs[0][start:end]
-        w_e = model.W_gate[e]
-        out_e_base = output + start * I  # Pointer to start position in output
-
-        grid_m = triton.cdiv(N_e, 64)
-        grid_n = triton.cdiv(I, 64)
-        grid = (grid_m * grid_n,)
-
-        triton_gemm_kernel[grid](
-            x_e, w_e, out_e_base,
-            N_e, I, H,
-            x_e.stride(0), 1,      # A: row stride, col stride
-            w_e.stride(0), w_e.stride(1),  # B: row stride, col stride
-            I, 1,                   # C: row stride, col stride
-            64, 64, 64,
-        )
-        torch.cuda.synchronize()
-        if e % 10 == 0:
-            print(f'Expert {e} done, N_e={N_e}, grid={grid}')
-    except Exception as ex:
-        print(f'Expert {e} failed: {ex}')
-        print(f'  start={start}, end={end}, N_e={N_e}')
-        break
-
-print('All experts done')
-print('Output shape:', output.shape)+print('Reference shape:', ref_result.shape)
+print('Triton shape:', triton_result.shape)
+print('Max diff:', (ref_result - triton_result).abs().max().item())
+print('Mean diff:', (ref_result - triton_result).abs().mean().item())
+print('Close?', torch.allclose(ref_result, triton_result, atol=0.02, rtol=0.02))

--- 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
@@ -15,7 +15,9 @@
 def triton_gemm_kernel(
     a_ptr, b_ptr, c_ptr,
     N_e, I, H,
-    stride_a, stride_b, stride_c,
+    stride_a_row, stride_a_col,
+    stride_b_row, stride_b_col,
+    stride_c_row, stride_c_col,
     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
 ):
     pid = tl.program_id(axis=0)
@@ -33,8 +35,8 @@
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
     offs_k = tl.arange(0, BLOCK_K)
 
-    a_ptrs = a_ptr + offs_m[:, None] * stride_a + offs_k[None, :] * 1
-    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_n[None, :] * 1
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a_row + offs_k[None, :] * stride_a_col
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b_row + offs_n[None, :] * stride_b_col
 
     accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(H, BLOCK_K)):
@@ -48,13 +50,13 @@
 
         accumulator += tl.dot(a, b)
 
-        a_ptrs += BLOCK_K * stride_a
-        b_ptrs += BLOCK_K * 1
+        a_ptrs += BLOCK_K * stride_a_col
+        b_ptrs += BLOCK_K * stride_b_row
         offs_k += BLOCK_K
 
     offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c + offs_n_out[None, :] * 1
+    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c_row + offs_n_out[None, :] * stride_c_col
     mask_m_out = offs_m_out[:, None] < N_e
     mask_n_out = offs_n_out[None, :] < I
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
@@ -63,7 +65,7 @@
 def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
     T_perm, H = hidden_states.shape
     I = W.shape[2]
-    output = torch.empty_like(hidden_states)
+    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
 
     for e in range(E):
         start = int(expert_offsets[e].item())
@@ -74,16 +76,18 @@
 
         x_e = hidden_states[start:end]
         w_e = W[e]
-        out_e = output[start:end]
+        out_e_base = output + start * I
 
         grid_m = triton.cdiv(N_e, 64)
         grid_n = triton.cdiv(I, 64)
         grid = (grid_m * grid_n,)
 
         triton_gemm_kernel[grid](
-            x_e, w_e, out_e,
+            x_e, w_e, out_e_base,
             N_e, I, H,
-            x_e.stride(0), w_e.stride(0), out_e.stride(0),
+            x_e.stride(0), 1,
+            w_e.stride(0), w_e.stride(1),
+            I, 1,
             64, 64, 64,
         )
     return output

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,4 +1,4 @@
-"""Verify correctness of Triton grouped GEMM + SwiGLU"""
+"""Test Triton grouped GEMM with better error handling"""
 import torch
 import triton
 import triton.language as tl
@@ -95,25 +95,9 @@
 
 model = solution.Model(*solution.get_init_inputs())
 model = model.cuda()
-
-torch.manual_seed(42)
 inputs = [t.cuda() for t in solution.get_inputs()]
 
-# Reference
-ref_out = solution.reference.Model(*solution.get_init_inputs()).cuda()
-ref_out.load_state_dict(model.state_dict())
-ref_out.eval()
-with torch.no_grad():
-    ref_result = ref_out(*inputs)
-
-# Triton GEMM + SwiGLU
 gate_out = triton_grouped_gemm(inputs[0], model.W_gate, inputs[1], solution.E)
-up_out = triton_grouped_gemm(inputs[0], model.W_up, inputs[1], solution.E)
-silu_gate = torch.nn.functional.silu(gate_out)
-triton_result = silu_gate * up_out
-
-print('Reference shape:', ref_result.shape)
-print('Triton shape:', triton_result.shape)
-print('Max diff:', (ref_result - triton_result).abs().max().item())
-print('Mean diff:', (ref_result - triton_result).abs().mean().item())
-print('Close?', torch.allclose(ref_result, triton_result, atol=0.02, rtol=0.02))+print('Gate out shape:', gate_out.shape)
+torch.cuda.synchronize()
+print('Done')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,4 +1,4 @@
-"""Test Triton grouped GEMM with better error handling"""
+"""Test Triton grouped GEMM + SwiGLU correctness"""
 import torch
 import triton
 import triton.language as tl
@@ -97,7 +97,21 @@
 model = model.cuda()
 inputs = [t.cuda() for t in solution.get_inputs()]
 
+# Reference
+import reference as ref_module
+ref_model = ref_module.Model(*ref_module.get_init_inputs()).cuda()
+ref_model.load_state_dict(model.state_dict())
+ref_model.eval()
+with torch.no_grad():
+    ref_out = ref_model(*inputs)
+
+# Triton
 gate_out = triton_grouped_gemm(inputs[0], model.W_gate, inputs[1], solution.E)
-print('Gate out shape:', gate_out.shape)
-torch.cuda.synchronize()
-print('Done')+up_out = triton_grouped_gemm(inputs[0], model.W_up, inputs[1], solution.E)
+silu_gate = torch.nn.functional.silu(gate_out)
+triton_result = silu_gate * up_out
+
+print('Reference shape:', ref_out.shape)
+print('Triton shape:', triton_result.shape)
+print('Max diff:', (ref_out - triton_result).abs().max().item())
+print('Close?', torch.allclose(ref_out, triton_result, atol=0.02, rtol=0.02))

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -3,6 +3,7 @@
 import triton
 import triton.language as tl
 import solution
+import reference as ref_mod
 
 @triton.jit
 def triton_gemm_kernel(
@@ -98,8 +99,7 @@
 inputs = [t.cuda() for t in solution.get_inputs()]
 
 # Reference
-import reference as ref_module
-ref_model = ref_module.Model(*ref_module.get_init_inputs()).cuda()
+ref_model = ref_mod.Model(*ref_mod.get_init_inputs()).cuda()
 ref_model.load_state_dict(model.state_dict())
 ref_model.eval()
 with torch.no_grad():

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -4,6 +4,17 @@
 import triton.language as tl
 import solution
 import reference as ref_mod
+
+# Set shapes for both reference and solution
+for mod, T_total, H, I, E, K in [
+    (solution, 4096, 2048, 1024, 64, 4),
+    (ref_mod, 4096, 2048, 1024, 64, 4),
+]:
+    mod.T_total = T_total
+    mod.H = H
+    mod.I = I
+    mod.E = E
+    mod.K = K
 
 @triton.jit
 def triton_gemm_kernel(
@@ -87,13 +98,6 @@
     return output
 
 
-# Test with small shape
-solution.T_total = 4096
-solution.H = 2048
-solution.I = 1024
-solution.E = 64
-solution.K = 4
-
 model = solution.Model(*solution.get_init_inputs())
 model = model.cuda()
 inputs = [t.cuda() for t in solution.get_inputs()]

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,20 +1,7 @@
-"""Test Triton grouped GEMM + SwiGLU correctness"""
+"""Debug Triton GEMM correctness"""
 import torch
 import triton
 import triton.language as tl
-import solution
-import reference as ref_mod
-
-# Set shapes for both reference and solution
-for mod, T_total, H, I, E, K in [
-    (solution, 4096, 2048, 1024, 64, 4),
-    (ref_mod, 4096, 2048, 1024, 64, 4),
-]:
-    mod.T_total = T_total
-    mod.H = H
-    mod.I = I
-    mod.E = E
-    mod.K = K
 
 @triton.jit
 def triton_gemm_kernel(
@@ -67,55 +54,99 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
-    T_perm, H = hidden_states.shape
-    I = W.shape[2]
-    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
+# Small test case
+N_e = 32
+I = 64
+H = 128
 
-    for e in range(E):
-        start = int(expert_offsets[e].item())
-        end = int(expert_offsets[e + 1].item())
-        if end == start:
-            continue
-        N_e = end - start
+torch.manual_seed(42)
+A = torch.randn(N_e, H, dtype=torch.bfloat16, device='cuda')
+B = torch.randn(H, I, dtype=torch.bfloat16, device='cuda')
+C = torch.zeros(N_e, I, dtype=torch.float32, device='cuda')
 
-        x_e = hidden_states[start:end]
-        w_e = W[e]
-        out_e_base = output + start * I
+# Reference: matmul in fp32
+C_ref = (A.float() @ B.float()).to(torch.bfloat16)
 
-        grid_m = triton.cdiv(N_e, 64)
-        grid_n = triton.cdiv(I, 64)
-        grid = (grid_m * grid_n,)
+grid_m = triton.cdiv(N_e, 64)
+grid_n = triton.cdiv(I, 64)
+grid = (grid_m * grid_n,)
 
-        triton_gemm_kernel[grid](
-            x_e, w_e, out_e_base,
-            N_e, I, H,
-            x_e.stride(0), 1,
-            w_e.stride(0), w_e.stride(1),
-            I, 1,
-            64, 64, 64,
-        )
-    return output
+triton_gemm_kernel[grid](
+    A, B, C,
+    N_e, I, H,
+    A.stride(0), 1,
+    B.stride(0), B.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+
+print('Triton result shape:', C.shape)
+print('Reference result shape:', C_ref.shape)
+print('Max diff (fp32 accum):', (C.float() - C_ref.float()).abs().max().item())
+print('Close (fp32)?', torch.allclose(C.float(), C_ref.float(), atol=0.01, rtol=0.01))
+
+# Try with BLOCK_M=32
+C2 = torch.zeros(N_e, I, dtype=torch.float32, device='cuda')
+grid2 = (triton.cdiv(N_e, 32) * triton.cdiv(I, 64),)
+
+@triton.jit
+def triton_gemm_kernel2(
+    a_ptr, b_ptr, c_ptr,
+    N_e, I, H,
+    stride_a_row, stride_a_col,
+    stride_b_row, stride_b_col,
+    stride_c_row, stride_c_col,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(N_e, BLOCK_M)
+    num_pid_n = tl.cdiv(I, BLOCK_N)
+    num_pid_in_group = num_pid_m * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * num_pid_m
+    group_size_m = min(num_pid_m, N_e - first_pid_m)
+
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    a_ptrs = a_ptr + offs_m[:, None] * stride_a_row + offs_k[None, :] * stride_a_col
+    b_ptrs = b_ptr + offs_k[:, None] * stride_b_row + offs_n[None, :] * stride_b_col
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(H, BLOCK_K)):
+        mask_m = offs_m[:, None] < N_e
+        mask_k = offs_k[None, :] < H
+        a = tl.load(a_ptrs, mask=mask_m & mask_k, other=0.0)
+
+        mask_k = offs_k[:, None] < H
+        mask_n = offs_n[None, :] < I
+        b = tl.load(b_ptrs, mask=mask_k & mask_n, other=0.0)
+
+        accumulator += tl.dot(a, b)
+
+        a_ptrs += BLOCK_K * stride_a_col
+        b_ptrs += BLOCK_K * stride_b_row
+        offs_k += BLOCK_K
+
+    offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c_row + offs_n_out[None, :] * stride_c_col
+    mask_m_out = offs_m_out[:, None] < N_e
+    mask_n_out = offs_n_out[None, :] < I
+    tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-model = solution.Model(*solution.get_init_inputs())
-model = model.cuda()
-inputs = [t.cuda() for t in solution.get_inputs()]
+triton_gemm_kernel2[grid2](
+    A, B, C2,
+    N_e, I, H,
+    A.stride(0), 1,
+    B.stride(0), B.stride(1),
+    I, 1,
+    32, 64, 64,
+)
 
-# Reference
-ref_model = ref_mod.Model(*ref_mod.get_init_inputs()).cuda()
-ref_model.load_state_dict(model.state_dict())
-ref_model.eval()
-with torch.no_grad():
-    ref_out = ref_model(*inputs)
-
-# Triton
-gate_out = triton_grouped_gemm(inputs[0], model.W_gate, inputs[1], solution.E)
-up_out = triton_grouped_gemm(inputs[0], model.W_up, inputs[1], solution.E)
-silu_gate = torch.nn.functional.silu(gate_out)
-triton_result = silu_gate * up_out
-
-print('Reference shape:', ref_out.shape)
-print('Triton shape:', triton_result.shape)
-print('Max diff:', (ref_out - triton_result).abs().max().item())
-print('Close?', torch.allclose(ref_out, triton_result, atol=0.02, rtol=0.02))+print('Max diff (BLOCK_M=32):', (C2.float() - C_ref.float()).abs().max().item())

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,4 +1,4 @@
-"""Debug Triton GEMM correctness"""
+"""Test Triton grouped GEMM correctness"""
 import torch
 import triton
 import triton.language as tl
@@ -54,99 +54,63 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-# Small test case
-N_e = 32
-I = 64
-H = 128
+# Test grouped GEMM on the MoE shape
+T_total, H, I, E, K = 4096, 2048, 1024, 64, 4
+T_perm = T_total * K
 
 torch.manual_seed(42)
-A = torch.randn(N_e, H, dtype=torch.bfloat16, device='cuda')
-B = torch.randn(H, I, dtype=torch.bfloat16, device='cuda')
-C = torch.zeros(N_e, I, dtype=torch.float32, device='cuda')
+hidden = torch.randn(T_perm, H, dtype=torch.bfloat16, device='cuda') * 0.1
+W = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
 
-# Reference: matmul in fp32
-C_ref = (A.float() @ B.float()).to(torch.bfloat16)
+base = T_perm // E
+rem = T_perm - base * E
+counts = torch.full((E,), base, dtype=torch.int32, device='cuda')
+counts[:rem] += 1
+offsets = torch.zeros(E + 1, dtype=torch.int32, device='cuda')
+offsets[1:] = torch.cumsum(counts, dim=0)
 
-grid_m = triton.cdiv(N_e, 64)
-grid_n = triton.cdiv(I, 64)
-grid = (grid_m * grid_n,)
+# Run grouped GEMM
+output = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+for e in range(E):
+    start = int(offsets[e].item())
+    end = int(offsets[e + 1].item())
+    if end == start:
+        continue
+    N_e = end - start
 
-triton_gemm_kernel[grid](
-    A, B, C,
-    N_e, I, H,
-    A.stride(0), 1,
-    B.stride(0), B.stride(1),
-    I, 1,
-    64, 64, 64,
-)
+    x_e = hidden[start:end]
+    w_e = W[e]
+    out_e = output[start:end]
 
-print('Triton result shape:', C.shape)
-print('Reference result shape:', C_ref.shape)
-print('Max diff (fp32 accum):', (C.float() - C_ref.float()).abs().max().item())
-print('Close (fp32)?', torch.allclose(C.float(), C_ref.float(), atol=0.01, rtol=0.01))
+    grid_m = triton.cdiv(N_e, 64)
+    grid_n = triton.cdiv(I, 64)
+    grid = (grid_m * grid_n,)
 
-# Try with BLOCK_M=32
-C2 = torch.zeros(N_e, I, dtype=torch.float32, device='cuda')
-grid2 = (triton.cdiv(N_e, 32) * triton.cdiv(I, 64),)
+    triton_gemm_kernel[grid](
+        x_e, w_e, out_e,
+        N_e, I, H,
+        x_e.stride(0), 1,
+        w_e.stride(0), w_e.stride(1),
+        I, 1,
+        64, 64, 64,
+    )
 
-@triton.jit
-def triton_gemm_kernel2(
-    a_ptr, b_ptr, c_ptr,
-    N_e, I, H,
-    stride_a_row, stride_a_col,
-    stride_b_row, stride_b_col,
-    stride_c_row, stride_c_col,
-    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
-):
-    pid = tl.program_id(axis=0)
-    num_pid_m = tl.cdiv(N_e, BLOCK_M)
-    num_pid_n = tl.cdiv(I, BLOCK_N)
-    num_pid_in_group = num_pid_m * num_pid_n
-    group_id = pid // num_pid_in_group
-    first_pid_m = group_id * num_pid_m
-    group_size_m = min(num_pid_m, N_e - first_pid_m)
+torch.cuda.synchronize()
 
-    pid_m = first_pid_m + (pid % group_size_m)
-    pid_n = (pid % num_pid_in_group) // group_size_m
+# Compare with reference
+ref_output = torch.zeros(T_perm, I, device='cuda', dtype=torch.float32)
+for e in range(E):
+    start = int(offsets[e].item())
+    end = int(offsets[e + 1].item())
+    if end == start:
+        continue
+    x_e = hidden[start:end].float()
+    w_e = W[e].float()
+    ref_output[start:end] = x_e @ w_e
 
-    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    offs_k = tl.arange(0, BLOCK_K)
+ref_output_bf16 = ref_output.to(torch.bfloat16)
 
-    a_ptrs = a_ptr + offs_m[:, None] * stride_a_row + offs_k[None, :] * stride_a_col
-    b_ptrs = b_ptr + offs_k[:, None] * stride_b_row + offs_n[None, :] * stride_b_col
-
-    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(H, BLOCK_K)):
-        mask_m = offs_m[:, None] < N_e
-        mask_k = offs_k[None, :] < H
-        a = tl.load(a_ptrs, mask=mask_m & mask_k, other=0.0)
-
-        mask_k = offs_k[:, None] < H
-        mask_n = offs_n[None, :] < I
-        b = tl.load(b_ptrs, mask=mask_k & mask_n, other=0.0)
-
-        accumulator += tl.dot(a, b)
-
-        a_ptrs += BLOCK_K * stride_a_col
-        b_ptrs += BLOCK_K * stride_b_row
-        offs_k += BLOCK_K
-
-    offs_m_out = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-    offs_n_out = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    c_ptrs = c_ptr + offs_m_out[:, None] * stride_c_row + offs_n_out[None, :] * stride_c_col
-    mask_m_out = offs_m_out[:, None] < N_e
-    mask_n_out = offs_n_out[None, :] < I
-    tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
-
-
-triton_gemm_kernel2[grid2](
-    A, B, C2,
-    N_e, I, H,
-    A.stride(0), 1,
-    B.stride(0), B.stride(1),
-    I, 1,
-    32, 64, 64,
-)
-
-print('Max diff (BLOCK_M=32):', (C2.float() - C_ref.float()).abs().max().item())+max_diff = (output.float() - ref_output).abs().max().item()
+print(f'Max diff (fp32 ref vs bf16 out): {max_diff}')
+print(f'Max diff (bf16 vs bf16): {(output - ref_output_bf16).abs().max().item()}')
+print(f'Close (bf16 atol=0.02)? {torch.allclose(output, ref_output_bf16, atol=0.02)}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,7 +1,23 @@
-"""Test Triton grouped GEMM correctness"""
+"""Debug correctness issue"""
 import torch
 import triton
 import triton.language as tl
+import solution
+import reference as ref_mod
+
+# Test with shape 0
+T_total, H, I, E, K = 32768, 4096, 1536, 128, 8
+T_perm = T_total * K
+
+for mod, t, h, ii, ee, kk in [
+    (solution, T_total, H, I, E, K),
+    (ref_mod, T_total, H, I, E, K),
+]:
+    mod.T_total = t
+    mod.H = h
+    mod.I = ii
+    mod.E = ee
+    mod.K = kk
 
 @triton.jit
 def triton_gemm_kernel(
@@ -54,13 +70,43 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-# Test grouped GEMM on the MoE shape
-T_total, H, I, E, K = 4096, 2048, 1024, 64, 4
-T_perm = T_total * K
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
+    T_perm, H = hidden_states.shape
+    I = W.shape[2]
+    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
 
+    for e in range(E):
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        if end == start:
+            continue
+        N_e = end - start
+
+        x_e = hidden_states[start:end]
+        w_e = W[e]
+        out_e_base = output + start * I
+
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
+        grid = (grid_m * grid_n,)
+
+        triton_gemm_kernel[grid](
+            x_e, w_e, out_e_base,
+            N_e, I, H,
+            x_e.stride(0), 1,
+            w_e.stride(0), w_e.stride(1),
+            I, 1,
+            64, 64, 64,
+        )
+    return output
+
+
+# Just test expert 0
 torch.manual_seed(42)
 hidden = torch.randn(T_perm, H, dtype=torch.bfloat16, device='cuda') * 0.1
-W = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
+
+W_gate = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
+W_up = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
 
 base = T_perm // E
 rem = T_perm - base * E
@@ -69,48 +115,53 @@
 offsets = torch.zeros(E + 1, dtype=torch.int32, device='cuda')
 offsets[1:] = torch.cumsum(counts, dim=0)
 
-# Run grouped GEMM
-output = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
-for e in range(E):
-    start = int(offsets[e].item())
-    end = int(offsets[e + 1].item())
-    if end == start:
-        continue
-    N_e = end - start
+# Expert 0 only
+e = 0
+start = int(offsets[e].item())
+end = int(offsets[e + 1].item())
+N_e = end - start
 
-    x_e = hidden[start:end]
-    w_e = W[e]
-    out_e = output[start:end]
+x_e = hidden[start:end]
+w_g = W_gate[e]
+w_u = W_up[e]
 
-    grid_m = triton.cdiv(N_e, 64)
-    grid_n = triton.cdiv(I, 64)
-    grid = (grid_m * grid_n,)
+print(f'Expert 0: rows {start}:{end} = {N_e} tokens')
 
-    triton_gemm_kernel[grid](
-        x_e, w_e, out_e,
-        N_e, I, H,
-        x_e.stride(0), 1,
-        w_e.stride(0), w_e.stride(1),
-        I, 1,
-        64, 64, 64,
-    )
+# Reference
+gate_ref = x_e.float() @ w_g.float()
+up_ref = x_e.float() @ w_u.float()
+gate_ref_bf16 = gate_ref.to(torch.bfloat16)
+up_ref_bf16 = up_ref.to(torch.bfloat16)
 
+# Triton
+gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+
+grid_m = triton.cdiv(N_e, 64)
+grid_n = triton.cdiv(I, 64)
+grid = (grid_m * grid_n,)
+print(f'Grid: {grid}, N_e={N_e}, I={I}, H={H}')
+
+triton_gemm_kernel[grid](
+    x_e, w_g, gate_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_g.stride(0), w_g.stride(1),
+    I, 1,
+    64, 64, 64,
+)
 torch.cuda.synchronize()
 
-# Compare with reference
-ref_output = torch.zeros(T_perm, I, device='cuda', dtype=torch.float32)
-for e in range(E):
-    start = int(offsets[e].item())
-    end = int(offsets[e + 1].item())
-    if end == start:
-        continue
-    x_e = hidden[start:end].float()
-    w_e = W[e].float()
-    ref_output[start:end] = x_e @ w_e
+triton_gemm_kernel[grid](
+    x_e, w_u, up_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_u.stride(0), w_u.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
 
-ref_output_bf16 = ref_output.to(torch.bfloat16)
-
-max_diff = (output.float() - ref_output).abs().max().item()
-print(f'Max diff (fp32 ref vs bf16 out): {max_diff}')
-print(f'Max diff (bf16 vs bf16): {(output - ref_output_bf16).abs().max().item()}')
-print(f'Close (bf16 atol=0.02)? {torch.allclose(output, ref_output_bf16, atol=0.02)}')+print(f'Gate ref max: {gate_ref.abs().max().item()}, Triton max: {gate_out.abs().max().item()}')
+print(f'Max diff gate: {(gate_out.float() - gate_ref).abs().max().item()}')
+print(f'Max diff up: {(up_out.float() - up_ref).abs().max().item()}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,9 +1,8 @@
-"""Debug correctness issue"""
+"""Debug correctness issue on shape 0"""
 import torch
 import triton
 import triton.language as tl
 import solution
-import reference as ref_mod
 
 # Test with shape 0
 T_total, H, I, E, K = 32768, 4096, 1536, 128, 8
@@ -11,7 +10,6 @@
 
 for mod, t, h, ii, ee, kk in [
     (solution, T_total, H, I, E, K),
-    (ref_mod, T_total, H, I, E, K),
 ]:
     mod.T_total = t
     mod.H = h
@@ -70,98 +68,88 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
-def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
-    T_perm, H = hidden_states.shape
-    I = W.shape[2]
-    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
-
-    for e in range(E):
-        start = int(expert_offsets[e].item())
-        end = int(expert_offsets[e + 1].item())
-        if end == start:
-            continue
-        N_e = end - start
-
-        x_e = hidden_states[start:end]
-        w_e = W[e]
-        out_e_base = output + start * I
-
-        grid_m = triton.cdiv(N_e, 64)
-        grid_n = triton.cdiv(I, 64)
-        grid = (grid_m * grid_n,)
-
-        triton_gemm_kernel[grid](
-            x_e, w_e, out_e_base,
-            N_e, I, H,
-            x_e.stride(0), 1,
-            w_e.stride(0), w_e.stride(1),
-            I, 1,
-            64, 64, 64,
-        )
-    return output
-
-
-# Just test expert 0
+# Build model
+model = solution.Model(*solution.get_init_inputs())
+model = model.cuda()
 torch.manual_seed(42)
-hidden = torch.randn(T_perm, H, dtype=torch.bfloat16, device='cuda') * 0.1
-
-W_gate = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
-W_up = torch.randn(E, H, I, dtype=torch.bfloat16, device='cuda') * 0.02
-
-base = T_perm // E
-rem = T_perm - base * E
-counts = torch.full((E,), base, dtype=torch.int32, device='cuda')
-counts[:rem] += 1
-offsets = torch.zeros(E + 1, dtype=torch.int32, device='cuda')
-offsets[1:] = torch.cumsum(counts, dim=0)
-
-# Expert 0 only
-e = 0
-start = int(offsets[e].item())
-end = int(offsets[e + 1].item())
-N_e = end - start
-
-x_e = hidden[start:end]
-w_g = W_gate[e]
-w_u = W_up[e]
-
-print(f'Expert 0: rows {start}:{end} = {N_e} tokens')
+inputs = [t.cuda() for t in solution.get_inputs()]
 
 # Reference
-gate_ref = x_e.float() @ w_g.float()
-up_ref = x_e.float() @ w_u.float()
-gate_ref_bf16 = gate_ref.to(torch.bfloat16)
-up_ref_bf16 = up_ref.to(torch.bfloat16)
+import reference as ref_mod
+ref_mod.T_total = T_total
+ref_mod.H = H
+ref_mod.I = I
+ref_mod.E = E
+ref_mod.K = K
+ref_model = ref_mod.Model(*ref_mod.get_init_inputs()).cuda()
+ref_model.load_state_dict(model.state_dict())
+ref_model.eval()
+with torch.no_grad():
+    ref_out = ref_model(*inputs)
 
-# Triton
-gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
-up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+# Triton GEMM
+gate_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+up_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+expert_offsets = inputs[1]
 
-grid_m = triton.cdiv(N_e, 64)
-grid_n = triton.cdiv(I, 64)
-grid = (grid_m * grid_n,)
-print(f'Grid: {grid}, N_e={N_e}, I={I}, H={H}')
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    if end == start:
+        continue
+    N_e = end - start
+    I_e = model.I
 
-triton_gemm_kernel[grid](
-    x_e, w_g, gate_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_g.stride(0), w_g.stride(1),
-    I, 1,
-    64, 64, 64,
-)
+    x_e = inputs[0][start:end]
+    w_g = model.W_gate[e]
+    w_u = model.W_up[e]
+    out_g = gate_out[start:start+N_e]
+    out_u = up_out[start:start+N_e]
+
+    grid_m = triton.cdiv(N_e, 64)
+    grid_n = triton.cdiv(I_e, 64)
+    grid = (grid_m * grid_n,)
+
+    triton_gemm_kernel[grid](
+        x_e, w_g, out_g,
+        N_e, I_e, H,
+        x_e.stride(0), 1,
+        w_g.stride(0), w_g.stride(1),
+        I_e, 1,
+        64, 64, 64,
+    )
+    triton_gemm_kernel[grid](
+        x_e, w_u, out_u,
+        N_e, I_e, H,
+        x_e.stride(0), 1,
+        w_u.stride(0), w_u.stride(1),
+        I_e, 1,
+        64, 64, 64,
+    )
+
 torch.cuda.synchronize()
 
-triton_gemm_kernel[grid](
-    x_e, w_u, up_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_u.stride(0), w_u.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
+silu_gate = torch.nn.functional.silu(gate_out)
+triton_out = silu_gate * up_out
 
-print(f'Gate ref max: {gate_ref.abs().max().item()}, Triton max: {gate_out.abs().max().item()}')
-print(f'Max diff gate: {(gate_out.float() - gate_ref).abs().max().item()}')
-print(f'Max diff up: {(up_out.float() - up_ref).abs().max().item()}')+# Find worst expert
+print(f'ref_out shape: {ref_out.shape}, triton_out shape: {triton_out.shape}')
+
+# Find which rows have biggest diff
+diffs = (ref_out.float() - triton_out.float()).abs()
+max_diff, max_idx = diffs.max(), diffs.argmax()
+max_row = max_idx // I
+max_col = max_idx % I
+print(f'Max diff at row {max_row}, col {max_col}: {max_diff.item()}')
+
+# Which expert owns that row?
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e+1].item())
+    if max_row >= start and max_row < end:
+        print(f'Expert {e} owns row {max_row} (offset {start}:{end})')
+        break
+
+# Check expert offsets around that area
+max_expert = max_row // (T_perm // E)  # rough guess
+print(f'Expert offsets around {max_expert}: {expert_offsets[max_expert:max_expert+5]}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -87,7 +87,7 @@
 with torch.no_grad():
     ref_out = ref_model(*inputs)
 
-# Triton GEMM
+# Triton GEMM - use explicit pointer arithmetic
 gate_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
 up_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
 expert_offsets = inputs[1]
@@ -98,32 +98,30 @@
     if end == start:
         continue
     N_e = end - start
-    I_e = model.I
 
     x_e = inputs[0][start:end]
     w_g = model.W_gate[e]
     w_u = model.W_up[e]
-    out_g = gate_out[start:start+N_e]
-    out_u = up_out[start:start+N_e]
 
     grid_m = triton.cdiv(N_e, 64)
-    grid_n = triton.cdiv(I_e, 64)
+    grid_n = triton.cdiv(I, 64)
     grid = (grid_m * grid_n,)
 
+    # Use explicit pointer offset
     triton_gemm_kernel[grid](
-        x_e, w_g, out_g,
-        N_e, I_e, H,
+        x_e, w_g, gate_out + start * I,
+        N_e, I, H,
         x_e.stride(0), 1,
         w_g.stride(0), w_g.stride(1),
-        I_e, 1,
+        I, 1,
         64, 64, 64,
     )
     triton_gemm_kernel[grid](
-        x_e, w_u, out_u,
-        N_e, I_e, H,
+        x_e, w_u, up_out + start * I,
+        N_e, I, H,
         x_e.stride(0), 1,
         w_u.stride(0), w_u.stride(1),
-        I_e, 1,
+        I, 1,
         64, 64, 64,
     )
 
@@ -148,8 +146,4 @@
     end = int(expert_offsets[e+1].item())
     if max_row >= start and max_row < end:
         print(f'Expert {e} owns row {max_row} (offset {start}:{end})')
-        break
-
-# Check expert offsets around that area
-max_expert = max_row // (T_perm // E)  # rough guess
-print(f'Expert offsets around {max_expert}: {expert_offsets[max_expert:max_expert+5]}')+        break

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -87,63 +87,90 @@
 with torch.no_grad():
     ref_out = ref_model(*inputs)
 
-# Triton GEMM - use explicit pointer arithmetic
-gate_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
-up_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+# Test expert 55 specifically
 expert_offsets = inputs[1]
+e = 55
+start = int(expert_offsets[e].item())
+end = int(expert_offsets[e + 1].item())
+N_e = end - start
+print(f'Expert {e}: rows {start}:{end}, N_e={N_e}')
 
-for e in range(E):
-    start = int(expert_offsets[e].item())
-    end = int(expert_offsets[e + 1].item())
-    if end == start:
-        continue
-    N_e = end - start
+x_e = inputs[0][start:end]
+w_g = model.W_gate[e]
+w_u = model.W_up[e]
 
-    x_e = inputs[0][start:end]
-    w_g = model.W_gate[e]
-    w_u = model.W_up[e]
+# Reference computation for this expert
+gate_ref = x_e.float() @ w_g.float()
+up_ref = x_e.float() @ w_u.float()
 
-    grid_m = triton.cdiv(N_e, 64)
-    grid_n = triton.cdiv(I, 64)
-    grid = (grid_m * grid_n,)
+# Triton for this expert
+gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
 
-    # Use explicit pointer offset
-    triton_gemm_kernel[grid](
-        x_e, w_g, gate_out + start * I,
-        N_e, I, H,
-        x_e.stride(0), 1,
-        w_g.stride(0), w_g.stride(1),
-        I, 1,
-        64, 64, 64,
-    )
-    triton_gemm_kernel[grid](
-        x_e, w_u, up_out + start * I,
-        N_e, I, H,
-        x_e.stride(0), 1,
-        w_u.stride(0), w_u.stride(1),
-        I, 1,
-        64, 64, 64,
-    )
+grid_m = triton.cdiv(N_e, 64)
+grid_n = triton.cdiv(I, 64)
+grid = (grid_m * grid_n,)
+print(f'Grid: {grid}')
 
+triton_gemm_kernel[grid](
+    x_e, w_g, gate_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_g.stride(0), w_g.stride(1),
+    I, 1,
+    64, 64, 64,
+)
 torch.cuda.synchronize()
 
-silu_gate = torch.nn.functional.silu(gate_out)
-triton_out = silu_gate * up_out
+triton_gemm_kernel[grid](
+    x_e, w_u, up_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_u.stride(0), w_u.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
 
-# Find worst expert
-print(f'ref_out shape: {ref_out.shape}, triton_out shape: {triton_out.shape}')
+# Check diffs within this expert
+gate_diff = (gate_out.float() - gate_ref).abs()
+up_diff = (up_out.float() - up_ref).abs()
 
-# Find which rows have biggest diff
-diffs = (ref_out.float() - triton_out.float()).abs()
-max_diff, max_idx = diffs.max(), diffs.argmax()
-max_row = max_idx // I
-max_col = max_idx % I
-print(f'Max diff at row {max_row}, col {max_col}: {max_diff.item()}')
+print(f'Gate max diff: {gate_diff.max().item()}')
+print(f'Up max diff: {up_diff.max().item()}')
 
-# Which expert owns that row?
-for e in range(E):
-    start = int(expert_offsets[e].item())
-    end = int(expert_offsets[e+1].item())
-    if max_row >= start and max_row < end:
-        print(f'Expert {e} owns row {max_row} (offset {start}:{end})')
-        break+# Find worst row
+worst_gate_row = gate_diff.max(dim=1).values.argmax()
+print(f'Worst gate row: {worst_gate_row}, diff: {gate_diff.max(dim=1).values[worst_gate_row].item()}')
+
+# Now test with output as a slice
+gate_out2 = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+up_out2 = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+
+# Use slice approach
+triton_gemm_kernel[grid](
+    x_e, w_g, gate_out2[start:end],
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_g.stride(0), w_g.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
+
+triton_gemm_kernel[grid](
+    x_e, w_u, up_out2[start:end],
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_u.stride(0), w_u.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
+
+# Compare
+gate_diff2 = (gate_out2[start:end].float() - gate_ref).abs()
+up_diff2 = (up_out2[start:end].float() - up_ref).abs()
+
+print(f'Gate max diff (slice): {gate_diff2.max().item()}')
+print(f'Up max diff (slice): {up_diff2.max().item()}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -87,90 +87,87 @@
 with torch.no_grad():
     ref_out = ref_model(*inputs)
 
-# Test expert 55 specifically
+# Triton GEMM with slice output
+gate_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
+up_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
 expert_offsets = inputs[1]
+
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e + 1].item())
+    if end == start:
+        continue
+    N_e = end - start
+
+    x_e = inputs[0][start:end]
+    w_g = model.W_gate[e]
+    w_u = model.W_up[e]
+
+    grid_m = triton.cdiv(N_e, 64)
+    grid_n = triton.cdiv(I, 64)
+    grid = (grid_m * grid_n,)
+
+    # Use slice for output
+    triton_gemm_kernel[grid](
+        x_e, w_g, gate_out[start:end],
+        N_e, I, H,
+        x_e.stride(0), 1,
+        w_g.stride(0), w_g.stride(1),
+        I, 1,
+        64, 64, 64,
+    )
+    triton_gemm_kernel[grid](
+        x_e, w_u, up_out[start:end],
+        N_e, I, H,
+        x_e.stride(0), 1,
+        w_u.stride(0), w_u.stride(1),
+        I, 1,
+        64, 64, 64,
+    )
+
+torch.cuda.synchronize()
+
+silu_gate = torch.nn.functional.silu(gate_out)
+triton_out = silu_gate * up_out
+
+# Find worst expert
+diffs = (ref_out.float() - triton_out.float()).abs()
+max_diff, max_idx = diffs.max(), diffs.argmax()
+max_row = max_idx // I
+max_col = max_idx % I
+print(f'Max diff at row {max_row}, col {max_col}: {max_diff.item()}')
+
+# Which expert owns that row?
+for e in range(E):
+    start = int(expert_offsets[e].item())
+    end = int(expert_offsets[e+1].item())
+    if max_row >= start and max_row < end:
+        print(f'Expert {e} owns row {max_row} (offset {start}:{end})')
+        break
+
+# Now check the GEMM for this expert specifically
 e = 55
 start = int(expert_offsets[e].item())
 end = int(expert_offsets[e + 1].item())
 N_e = end - start
-print(f'Expert {e}: rows {start}:{end}, N_e={N_e}')
 
 x_e = inputs[0][start:end]
 w_g = model.W_gate[e]
 w_u = model.W_up[e]
 
-# Reference computation for this expert
 gate_ref = x_e.float() @ w_g.float()
 up_ref = x_e.float() @ w_u.float()
 
-# Triton for this expert
-gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
-up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+# Check if our output buffer for this expert matches reference
+print(f'\nExpert 55 GEMM check:')
+print(f'Gate out buffer rows 112640:114688:')
+print(f'  Max diff: {(gate_out[start:end].float() - gate_ref).abs().max().item()}')
+print(f'  Up out buffer rows 112640:114688:')
+print(f'  Max diff: {(up_out[start:end].float() - up_ref).abs().max().item()}')
 
-grid_m = triton.cdiv(N_e, 64)
-grid_n = triton.cdiv(I, 64)
-grid = (grid_m * grid_n,)
-print(f'Grid: {grid}')
-
-triton_gemm_kernel[grid](
-    x_e, w_g, gate_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_g.stride(0), w_g.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
-
-triton_gemm_kernel[grid](
-    x_e, w_u, up_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_u.stride(0), w_u.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
-
-# Check diffs within this expert
-gate_diff = (gate_out.float() - gate_ref).abs()
-up_diff = (up_out.float() - up_ref).abs()
-
-print(f'Gate max diff: {gate_diff.max().item()}')
-print(f'Up max diff: {up_diff.max().item()}')
-
-# Find worst row
-worst_gate_row = gate_diff.max(dim=1).values.argmax()
-print(f'Worst gate row: {worst_gate_row}, diff: {gate_diff.max(dim=1).values[worst_gate_row].item()}')
-
-# Now test with output as a slice
-gate_out2 = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
-up_out2 = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
-
-# Use slice approach
-triton_gemm_kernel[grid](
-    x_e, w_g, gate_out2[start:end],
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_g.stride(0), w_g.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
-
-triton_gemm_kernel[grid](
-    x_e, w_u, up_out2[start:end],
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_u.stride(0), w_u.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
-
-# Compare
-gate_diff2 = (gate_out2[start:end].float() - gate_ref).abs()
-up_diff2 = (up_out2[start:end].float() - up_ref).abs()
-
-print(f'Gate max diff (slice): {gate_diff2.max().item()}')
-print(f'Up max diff (slice): {up_diff2.max().item()}')+# But what about the overall output?
+print(f'\nOverall output for expert 55 region:')
+print(f'  Gate ref max: {gate_ref.abs().max().item()}')
+print(f'  Gate out max: {gate_out[start:end].abs().max().item()}')
+print(f'  Up ref max: {up_ref.abs().max().item()}')
+print(f'  Up out max: {up_out[start:end].abs().max().item()}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,4 +1,4 @@
-"""Debug correctness issue on shape 0"""
+"""Debug expert 100 in detail"""
 import torch
 import triton
 import triton.language as tl
@@ -74,79 +74,8 @@
 torch.manual_seed(42)
 inputs = [t.cuda() for t in solution.get_inputs()]
 
-# Reference
-import reference as ref_mod
-ref_mod.T_total = T_total
-ref_mod.H = H
-ref_mod.I = I
-ref_mod.E = E
-ref_mod.K = K
-ref_model = ref_mod.Model(*ref_mod.get_init_inputs()).cuda()
-ref_model.load_state_dict(model.state_dict())
-ref_model.eval()
-with torch.no_grad():
-    ref_out = ref_model(*inputs)
-
-# Triton GEMM with slice output
-gate_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
-up_out = torch.empty(T_perm, I, device='cuda', dtype=torch.bfloat16)
 expert_offsets = inputs[1]
-
-for e in range(E):
-    start = int(expert_offsets[e].item())
-    end = int(expert_offsets[e + 1].item())
-    if end == start:
-        continue
-    N_e = end - start
-
-    x_e = inputs[0][start:end]
-    w_g = model.W_gate[e]
-    w_u = model.W_up[e]
-
-    grid_m = triton.cdiv(N_e, 64)
-    grid_n = triton.cdiv(I, 64)
-    grid = (grid_m * grid_n,)
-
-    # Use slice for output
-    triton_gemm_kernel[grid](
-        x_e, w_g, gate_out[start:end],
-        N_e, I, H,
-        x_e.stride(0), 1,
-        w_g.stride(0), w_g.stride(1),
-        I, 1,
-        64, 64, 64,
-    )
-    triton_gemm_kernel[grid](
-        x_e, w_u, up_out[start:end],
-        N_e, I, H,
-        x_e.stride(0), 1,
-        w_u.stride(0), w_u.stride(1),
-        I, 1,
-        64, 64, 64,
-    )
-
-torch.cuda.synchronize()
-
-silu_gate = torch.nn.functional.silu(gate_out)
-triton_out = silu_gate * up_out
-
-# Find worst expert
-diffs = (ref_out.float() - triton_out.float()).abs()
-max_diff, max_idx = diffs.max(), diffs.argmax()
-max_row = max_idx // I
-max_col = max_idx % I
-print(f'Max diff at row {max_row}, col {max_col}: {max_diff.item()}')
-
-# Which expert owns that row?
-for e in range(E):
-    start = int(expert_offsets[e].item())
-    end = int(expert_offsets[e+1].item())
-    if max_row >= start and max_row < end:
-        print(f'Expert {e} owns row {max_row} (offset {start}:{end})')
-        break
-
-# Now check the GEMM for this expert specifically
-e = 55
+e = 100
 start = int(expert_offsets[e].item())
 end = int(expert_offsets[e + 1].item())
 N_e = end - start
@@ -155,19 +84,58 @@
 w_g = model.W_gate[e]
 w_u = model.W_up[e]
 
+print(f'Expert {e}: rows {start}:{end}, N_e={N_e}')
+
+# Reference
 gate_ref = x_e.float() @ w_g.float()
 up_ref = x_e.float() @ w_u.float()
 
-# Check if our output buffer for this expert matches reference
-print(f'\nExpert 55 GEMM check:')
-print(f'Gate out buffer rows 112640:114688:')
-print(f'  Max diff: {(gate_out[start:end].float() - gate_ref).abs().max().item()}')
-print(f'  Up out buffer rows 112640:114688:')
-print(f'  Max diff: {(up_out[start:end].float() - up_ref).abs().max().item()}')
+# Triton output for this expert
+gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
+up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
 
-# But what about the overall output?
-print(f'\nOverall output for expert 55 region:')
-print(f'  Gate ref max: {gate_ref.abs().max().item()}')
-print(f'  Gate out max: {gate_out[start:end].abs().max().item()}')
-print(f'  Up ref max: {up_ref.abs().max().item()}')
-print(f'  Up out max: {up_out[start:end].abs().max().item()}')+grid_m = triton.cdiv(N_e, 64)
+grid_n = triton.cdiv(I, 64)
+grid = (grid_m * grid_n,)
+print(f'Grid: {grid}')
+
+triton_gemm_kernel[grid](
+    x_e, w_g, gate_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_g.stride(0), w_g.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
+
+triton_gemm_kernel[grid](
+    x_e, w_u, up_out,
+    N_e, I, H,
+    x_e.stride(0), 1,
+    w_u.stride(0), w_u.stride(1),
+    I, 1,
+    64, 64, 64,
+)
+torch.cuda.synchronize()
+
+# Find the worst row (row 84, col 30)
+row = 84
+col = 30
+print(f'\nAt local row {row}, col {col}:')
+print(f'  Gate ref: {gate_ref[row, col].item()}')
+print(f'  Gate out: {gate_out[row, col].item()}')
+print(f'  Up ref: {up_ref[row, col].item()}')
+print(f'  Up out: {up_out[row, col].item()}')
+
+# Check all columns for row 84
+row_diff = (gate_out[row].float() - gate_ref[row]).abs()
+print(f'\nRow 84 gate diff stats:')
+print(f'  Max: {row_diff.max().item()}, Mean: {row_diff.mean().item()}')
+print(f'  ArgMax col: {row_diff.argmax().item()}')
+
+# Check all rows for col 30
+col_diff = (gate_out[:, 30].float() - gate_ref[:, 30]).abs()
+print(f'\nCol 30 gate diff stats:')
+print(f'  Max: {col_diff.max().item()}, Mean: {col_diff.mean().item()}')
+print(f'  ArgMax row: {col_diff.argmax().item()}')

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/06_sonic_moe_swiglu/test_triton.py
@@ -1,10 +1,9 @@
-"""Debug expert 100 in detail"""
+"""Test triton_grouped_gemm function directly"""
 import torch
 import triton
 import triton.language as tl
 import solution
 
-# Test with shape 0
 T_total, H, I, E, K = 32768, 4096, 1536, 128, 8
 T_perm = T_total * K
 
@@ -68,74 +67,50 @@
     tl.store(c_ptrs, accumulator, mask=mask_m_out & mask_n_out)
 
 
+def triton_grouped_gemm(hidden_states, W, expert_offsets, E):
+    T_perm, H = hidden_states.shape
+    I = W.shape[2]
+    output = torch.empty(T_perm, I, device=hidden_states.device, dtype=hidden_states.dtype)
+
+    for e in range(E):
+        start = int(expert_offsets[e].item())
+        end = int(expert_offsets[e + 1].item())
+        if end == start:
+            continue
+        N_e = end - start
+
+        x_e = hidden_states[start:end]
+        w_e = W[e]
+        out_e_base = output + start * I
+
+        grid_m = triton.cdiv(N_e, 64)
+        grid_n = triton.cdiv(I, 64)
+        grid = (grid_m * grid_n,)
+
+        triton_gemm_kernel[grid](
+            x_e, w_e, out_e_base,
+            N_e, I, H,
+            x_e.stride(0), 1,
+            w_e.stride(0), w_e.stride(1),
+            I, 1,
+            64, 64, 64,
+        )
+    return output
+
+
 # Build model
 model = solution.Model(*solution.get_init_inputs())
 model = model.cuda()
 torch.manual_seed(42)
 inputs = [t.cuda() for t in solution.get_inputs()]
 
-expert_offsets = inputs[1]
-e = 100
-start = int(expert_offsets[e].item())
-end = int(expert_offsets[e + 1].item())
-N_e = end - start
-
-x_e = inputs[0][start:end]
-w_g = model.W_gate[e]
-w_u = model.W_up[e]
-
-print(f'Expert {e}: rows {start}:{end}, N_e={N_e}')
-
-# Reference
-gate_ref = x_e.float() @ w_g.float()
-up_ref = x_e.float() @ w_u.float()
-
-# Triton output for this expert
-gate_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
-up_out = torch.empty(N_e, I, device='cuda', dtype=torch.bfloat16)
-
-grid_m = triton.cdiv(N_e, 64)
-grid_n = triton.cdiv(I, 64)
-grid = (grid_m * grid_n,)
-print(f'Grid: {grid}')
-
-triton_gemm_kernel[grid](
-    x_e, w_g, gate_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_g.stride(0), w_g.stride(1),
-    I, 1,
-    64, 64, 64,
-)
+# Run triton_grouped_gemm directly
+gate_out = triton_grouped_gemm(inputs[0], model.W_gate, inputs[1], E)
 torch.cuda.synchronize()
 
-triton_gemm_kernel[grid](
-    x_e, w_u, up_out,
-    N_e, I, H,
-    x_e.stride(0), 1,
-    w_u.stride(0), w_u.stride(1),
-    I, 1,
-    64, 64, 64,
-)
-torch.cuda.synchronize()
-
-# Find the worst row (row 84, col 30)
-row = 84
-col = 30
-print(f'\nAt local row {row}, col {col}:')
-print(f'  Gate ref: {gate_ref[row, col].item()}')
-print(f'  Gate out: {gate_out[row, col].item()}')
-print(f'  Up ref: {up_ref[row, col].item()}')
-print(f'  Up out: {up_out[row, col].item()}')
-
-# Check all columns for row 84
-row_diff = (gate_out[row].float() - gate_ref[row]).abs()
-print(f'\nRow 84 gate diff stats:')
-print(f'  Max: {row_diff.max().item()}, Mean: {row_diff.mean().item()}')
-print(f'  ArgMax col: {row_diff.argmax().item()}')
-
-# Check all rows for col 30
-col_diff = (gate_out[:, 30].float() - gate_ref[:, 30]).abs()
-print(f'\nCol 30 gate diff stats:')
-print(f'  Max: {col_diff.max().item()}, Mean: {col_diff.mean().item()}')
-print(f'  ArgMax row: {col_diff.argmax().item()}')+print(f'gate_out shape: {gate_out.shape}')
+print(f'gate_out non-zero count: {(gate_out != 0).sum().item()}')
+print(f'gate_out max: {gate_out.abs().max().item()}')
+print(f'gate_out sample at row 0, col 0: {gate_out[0, 0].item()}')
+print(f'gate_out sample at row 100, col 100: {gate_out[100, 100].item()}')
+print(f'gate_out sample at row 200000, col 100: {gate_out[200000, 100].item()}')

--- 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
@@ -76,18 +76,18 @@
 
         x_e = hidden_states[start:end]
         w_e = W[e]
-        out_e_base = output + start * I
+        out_e = output[start:end]
 
         grid_m = triton.cdiv(N_e, 64)
         grid_n = triton.cdiv(I, 64)
         grid = (grid_m * grid_n,)
 
         triton_gemm_kernel[grid](
-            x_e, w_e, out_e_base,
+            x_e, w_e, out_e,
             N_e, I, H,
             x_e.stride(0), 1,
             w_e.stride(0), w_e.stride(1),
-            I, 1,
+            out_e.stride(0), 1,
             64, 64, 64,
         )
     return output