--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -0,0 +1,230 @@
+"""Grouped GEMM + fused SwiGLU for top-K MoE FFN up-projection.
+
+Implements per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]) using a
+single Triton kernel. Concatenates W_gate and W_up into a single weight matrix
+W = (E, H, 2*I) so a single tile covers both gate and up columns; SwiGLU is
+applied to the (BLOCK_M, BLOCK_N) gate/up pair in the epilogue.
+
+Grid: (cdiv(T_perm, BLOCK_M), cdiv(I, BLOCK_N)) -- over tiles of the *permuted*
+input. A precomputed block-to-expert table (built on host in one CPU pass)
+maps each global m-block to (expert, m_within_expert), avoiding any in-kernel
+search over offsets.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
+    """Build a host-side table mapping each global m-block to (expert, m_local).
+
+    For the call site, we want contiguous (m_block, n_block) tiles but each
+    m_block may belong to a different expert. This precomputes:
+        block_to_expert[mb]   = e
+        block_to_m_local[mb] = m_within_expert (in elements, BLOCK_M aligned)
+    """
+    # Use CPU to avoid a device sync; this is a one-time cost.
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert: list[int] = []
+    block_to_m_local: list[int] = []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0:
+            continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_local.append(b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_local, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+@triton.jit
+def _grouped_gemm_swiglu_kernel(
+    A_ptr,                  # (T_perm, H) bf16
+    W_ptr,                  # (E, H, 2*I) bf16
+    Out_ptr,                # (T_perm, I) bf16
+    block_to_expert_ptr,    # (num_m_blocks,) int32
+    block_to_m_local_ptr,   # (num_m_blocks,) int32
+    H, N,                   # int32: H, I (=N here)
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)  # m-block over T_perm
+    pid_n = tl.program_id(1)  # n-block over I
+
+    # Look up which expert and where in the expert's slice this m-block lives.
+    e = tl.load(block_to_expert_ptr + pid_m)
+    m_local = tl.load(block_to_m_local_ptr + pid_m)
+
+    # GEMM tile pointers
+    offs_m = m_local + tl.arange(0, BLOCK_M)             # (BLOCK_M,)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)     # (BLOCK_N,)
+    offs_k = tl.arange(0, BLOCK_K)                       # (BLOCK_K,)
+
+    # A row pointers: A[expert_offset[e] + offs_m, :]
+    # We bake the expert's start offset into a base pointer.
+    # NOTE: we don't have offsets here, but block_to_m_local is the absolute
+    # offset into T_perm (=expert_start + m_within_expert). So we use that.
+    a_row_base = m_local + tl.load(block_to_m_local_ptr + pid_m) * 0  # placeholder; fix below
+    # Actually we need the expert's *absolute* m start, not m_local. Let me
+    # restructure: use m_local = absolute m offset within T_perm.
+
+    # a_ptrs[k, :] = A[m_local + offs_m, k]
+    a_ptrs = A_ptr + (m_local + tl.arange(0, BLOCK_M))[:, None] * stride_am + offs_k[None, :] * stride_ak
+
+    # W[e, k, n] -- gate is at n in [0, N), up is at n in [N, 2*N). We process
+    # gate+up together: load (BLOCK_K, BLOCK_N) for the gate columns and another
+    # (BLOCK_K, BLOCK_N) for the up columns in the same K iteration.
+    # First pass: gate. w_gate_ptrs[k, :] = W[e, k, pid_n*BLOCK_N + offs_n]
+    w_gate_ptrs = (
+        W_ptr + e * stride_we
+        + offs_k[:, None] * stride_wk
+        + (pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_wn
+    )
+    # up: shift by N (so column = I + pid_n*BLOCK_N + offs_n)
+    w_up_ptrs = (
+        W_ptr + e * stride_we
+        + offs_k[:, None] * stride_wk
+        + (N + pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_wn
+    )
+
+    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    m_mask = (m_local + tl.arange(0, BLOCK_M))[:, None] < (m_local + BLOCK_M)  # all true if no boundary
+    # Boundary: if this m block extends past the expert's end, mask rows. We
+    # need the expert's end; recompute as the next block's m_local (or the
+    # end of T_perm). For now assume aligned (true for the balanced case).
+
+    for k in range(0, H, BLOCK_K):
+        a = tl.load(a_ptrs)                          # (BLOCK_M, BLOCK_K)
+        w_g = tl.load(w_gate_ptrs)                   # (BLOCK_K, BLOCK_N)
+        w_u = tl.load(w_up_ptrs)                     # (BLOCK_K, BLOCK_N)
+        acc_g = tl.dot(a, w_g, acc=acc_g)
+        acc_u = tl.dot(a, w_u, acc=acc_u)
+        a_ptrs += BLOCK_K * stride_ak
+        w_gate_ptrs += BLOCK_K * stride_wk
+        w_up_ptrs += BLOCK_K * stride_wk
+
+    # Epilogue: SwiGLU(gate) * up, store as bf16 to Out[m, n].
+    silu_g = acc_g * tl.sigmoid(acc_g)
+    out = silu_g * acc_u
+    out_bf = out.to(tl.bfloat16)
+
+    o_ptrs = (
+        Out_ptr
+        + (m_local + tl.arange(0, BLOCK_M))[:, None] * stride_om
+        + (pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_on
+    )
+    tl.store(o_ptrs, out_bf)
+
+
+def _grouped_gemm_swiglu(
+    hidden: torch.Tensor,        # (T_perm, H) bf16
+    W: torch.Tensor,             # (E, H, 2*I) bf16
+    expert_offsets: torch.Tensor, # (E+1,) int32
+    T_perm: int,
+    H: int,

... (2901 more chars)

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -1,14 +1,14 @@
 """Grouped GEMM + fused SwiGLU for top-K MoE FFN up-projection.
 
-Implements per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]) using a
-single Triton kernel. Concatenates W_gate and W_up into a single weight matrix
-W = (E, H, 2*I) so a single tile covers both gate and up columns; SwiGLU is
-applied to the (BLOCK_M, BLOCK_N) gate/up pair in the epilogue.
-
-Grid: (cdiv(T_perm, BLOCK_M), cdiv(I, BLOCK_N)) -- over tiles of the *permuted*
-input. A precomputed block-to-expert table (built on host in one CPU pass)
-maps each global m-block to (expert, m_within_expert), avoiding any in-kernel
-search over offsets.
+Per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]). We concatenate
+W_gate and W_up into a single weight matrix W = (E, H, 2*I) and do one
+grouped GEMM that produces (T_perm, 2*I) of (gate, up) pairs, then applies
+SwiGLU in the epilogue to write (T_perm, I) bf16 output.
+
+Grid: (num_m_blocks, num_n_blocks) where num_m_blocks is the total number of
+m-tiles across all experts (variable per expert). A host-built table maps
+each global m-block to (expert, m_offset_within_T_perm), avoiding any
+in-kernel search over expert offsets.
 """
 from __future__ import annotations
 
@@ -19,17 +19,16 @@
 
 
 def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
-    """Build a host-side table mapping each global m-block to (expert, m_local).
-
-    For the call site, we want contiguous (m_block, n_block) tiles but each
-    m_block may belong to a different expert. This precomputes:
-        block_to_expert[mb]   = e
-        block_to_m_local[mb] = m_within_expert (in elements, BLOCK_M aligned)
+    """Build a host-side table mapping each global m-block to its expert and
+    the absolute m offset within T_perm.
+
+    Returns two int32 tensors of shape (num_m_blocks,):
+        block_to_expert[mb]   = expert id
+        block_to_m_abs[mb]    = absolute m offset within T_perm (in elements)
     """
-    # Use CPU to avoid a device sync; this is a one-time cost.
     offsets_cpu = expert_offsets.cpu().to(torch.int64)
     block_to_expert: list[int] = []
-    block_to_m_local: list[int] = []
+    block_to_m_abs: list[int] = []
     for e in range(E):
         start = int(offsets_cpu[e].item())
         end = int(offsets_cpu[e + 1].item())
@@ -39,10 +38,10 @@
         n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
         for b in range(n_blocks):
             block_to_expert.append(e)
-            block_to_m_local.append(b * BLOCK_M)
+            block_to_m_abs.append(start + b * BLOCK_M)
     return (
         torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
-        torch.tensor(block_to_m_local, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
     )
 
 
@@ -52,7 +51,7 @@
     W_ptr,                  # (E, H, 2*I) bf16
     Out_ptr,                # (T_perm, I) bf16
     block_to_expert_ptr,    # (num_m_blocks,) int32
-    block_to_m_local_ptr,   # (num_m_blocks,) int32
+    block_to_m_abs_ptr,     # (num_m_blocks,) int32
     H, N,                   # int32: H, I (=N here)
     stride_am, stride_ak,
     stride_we, stride_wk, stride_wn,
@@ -61,80 +60,57 @@
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
 ):
-    pid_m = tl.program_id(0)  # m-block over T_perm
-    pid_n = tl.program_id(1)  # n-block over I
-
-    # Look up which expert and where in the expert's slice this m-block lives.
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
     e = tl.load(block_to_expert_ptr + pid_m)
-    m_local = tl.load(block_to_m_local_ptr + pid_m)
-
-    # GEMM tile pointers
-    offs_m = m_local + tl.arange(0, BLOCK_M)             # (BLOCK_M,)
-    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)     # (BLOCK_N,)
-    offs_k = tl.arange(0, BLOCK_K)                       # (BLOCK_K,)
-
-    # A row pointers: A[expert_offset[e] + offs_m, :]
-    # We bake the expert's start offset into a base pointer.
-    # NOTE: we don't have offsets here, but block_to_m_local is the absolute
-    # offset into T_perm (=expert_start + m_within_expert). So we use that.
-    a_row_base = m_local + tl.load(block_to_m_local_ptr + pid_m) * 0  # placeholder; fix below
-    # Actually we need the expert's *absolute* m start, not m_local. Let me
-    # restructure: use m_local = absolute m offset within T_perm.
-
-    # a_ptrs[k, :] = A[m_local + offs_m, k]
-    a_ptrs = A_ptr + (m_local + tl.arange(0, BLOCK_M))[:, None] * stride_am + offs_k[None, :] * stride_ak
-
-    # W[e, k, n] -- gate is at n in [0, N), up is at n in [N, 2*N). We process
-    # gate+up together: load (BLOCK_K, BLOCK_N) for the gate columns and another
-    # (BLOCK_K, BLOCK_N) for the up columns in the same K iteration.
-    # First pass: gate. w_gate_ptrs[k, :] = W[e, k, pid_n*BLOCK_N + offs_n]
+    m_abs = tl.load(block_to_m_abs_ptr + pid_m)  # absolute m offset within T_perm
+
+    offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
     w_gate_ptrs = (
         W_ptr + e * stride_we
         + offs_k[:, None] * stride_wk
-        + (pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_wn
-    )
-    # up: shift by N (so column = I + pid_n*BLOCK_N + offs_n)
+        + offs_n[None, :] * stride_wn
+    )
     w_up_ptrs = (
         W_ptr + e * stride_we
         + offs_k[:, None] * stride_wk
-        + (N + pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_wn
+        + (N + offs_n)[None, :] * stride_wn
     )
 
     acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
     acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    m_mask = (m_local + tl.arange(0, BLOCK_M))[:, N

... (3778 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/sweep.py
@@ -0,0 +1,121 @@
+"""Tile-size sweep for grouped GEMM SwiGLU kernel."""
+import torch
+import time
+import solution as sol
+import reference as ref
+
+def time_kernel(model, inputs, n_warmup=5, n_iter=20):
+    torch.cuda.synchronize()
+    for _ in range(n_warmup):
+        out = model(*inputs)
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for _ in range(n_iter):
+        out = model(*inputs)
+    torch.cuda.synchronize()
+    return (time.time() - t0) * 1000.0 / n_iter
+
+def time_kernel_cuda(model, inputs, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        out = model(*inputs)
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        out = model(*inputs)
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+configs = [
+    (64, 128, 32, 4, 3),
+    (64, 128, 32, 4, 4),
+    (64, 128, 64, 4, 3),
+    (64, 256, 32, 4, 3),
+    (64, 256, 64, 4, 3),
+    (128, 64, 32, 4, 3),
+    (128, 64, 32, 4, 4),
+    (128, 64, 64, 4, 3),
+    (128, 128, 32, 4, 3),
+    (128, 128, 32, 8, 3),
+    (128, 128, 64, 4, 3),
+    (128, 128, 64, 8, 3),
+    (128, 128, 64, 8, 4),
+    (128, 256, 32, 8, 3),
+    (128, 256, 64, 8, 3),
+    (256, 64, 32, 4, 3),
+    (256, 64, 64, 4, 3),
+    (256, 128, 32, 8, 3),
+    (256, 128, 64, 8, 3),
+    (64, 64, 64, 4, 3),
+    (32, 64, 32, 4, 3),
+    (32, 128, 32, 4, 3),
+]
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    sol.T_total = shape['T_total']
+    sol.H = shape['H']
+    sol.I = shape['I']
+    sol.E = shape['E']
+    sol.K = shape['K']
+    ref.T_total = shape['T_total']
+    ref.H = shape['H']
+    ref.I = shape['I']
+    ref.E = shape['E']
+    ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    inputs = [t.cuda() for t in ref.get_inputs()]
+
+    # Time the reference
+    ref_ms = time_kernel_cuda(ref_model, inputs, n_iter=10)
+    print(f'reference: {ref_ms:.3f} ms')
+
+    # Try each config
+    for BM, BN, BK, nw, ns in configs:
+        try:
+            sol_model = sol.Model(*init_args).cuda().eval()
+            sol_model.load_state_dict(ref_model.state_dict())
+            # Override the kernel call with custom tile sizes
+            def custom_forward(model, hidden, offsets, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns):
+                return sol._grouped_gemm_swiglu(
+                    hidden, model.W_gate, model.W_up, offsets,
+                    T_perm=hidden.shape[0], H=hidden.shape[1], I=model.I, E=model.E,
+                    BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
+                    num_warps=nw, num_stages=ns,
+                )
+            import types
+            sol_model.forward = types.MethodType(
+                lambda self, h, o, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns:
+                sol._grouped_gemm_swiglu(
+                    h, self.W_gate, self.W_up, o,
+                    T_perm=h.shape[0], H=h.shape[1], I=self.I, E=self.E,
+                    BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
+                    num_warps=nw, num_stages=ns,
+                ),
+                sol_model,
+            )
+            # Verify correctness first
+            with torch.no_grad():
+                ref_out = ref_model(*inputs)
+                sol_out = sol_model(*inputs)
+            max_diff = (ref_out - sol_out).abs().max().item()
+            if max_diff > 0.05:
+                print(f'  BM={BM} BN={BN} BK={BK} nw={nw} ns={ns}: WRONG (diff={max_diff:.3f})')
+                continue
+            ms = time_kernel_cuda(sol_model, inputs, n_iter=10)
+            flops = 2 * shape['T_total'] * shape['H'] * 2 * shape['I']
+            tflops = flops / (ms * 1e-3) / 1e12
+            print(f'  BM={BM:3d} BN={BN:3d} BK={BK:2d} nw={nw} ns={ns}: {ms:.3f} ms ({tflops:.1f} TFLOPs)')
+            del sol_model
+        except Exception as e:
+            print(f'  BM={BM} BN={BN} BK={BK} nw={nw} ns={ns}: FAILED ({type(e).__name__})')

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution_merged.py
@@ -0,0 +1,159 @@
+"""Grouped GEMM + fused SwiGLU - merged dot version (loads gate+up together)."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int, T_perm: int):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert: list[int] = []
+    block_to_m_abs: list[int] = []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0:
+            continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+@triton.jit
+def _grouped_gemm_swiglu_merged_kernel(
+    A_ptr,                  # (T_perm, H) bf16
+    W_ptr,                  # (E, H, 2*I) bf16  -- concatenated
+    Out_ptr,                # (T_perm, I) bf16
+    block_to_expert_ptr,
+    block_to_m_abs_ptr,
+    H, N,                   # int32
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    e = tl.load(block_to_expert_ptr + pid_m)
+    m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+
+    offs_m = m_abs + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_2n = tl.arange(0, 2 * BLOCK_N)  # gate+up columns
+    offs_k = tl.arange(0, BLOCK_K)
+
+    a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+    w_ptrs = (
+        W_ptr + e * stride_we
+        + offs_k[:, None] * stride_wk
+        + offs_2n[None, :] * stride_wn
+    )
+
+    acc = tl.zeros((BLOCK_M, 2 * BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, H, BLOCK_K):
+        a = tl.load(a_ptrs)
+        w = tl.load(w_ptrs)  # (BLOCK_K, 2*BLOCK_N)
+        acc = tl.dot(a, w, acc=acc)
+        a_ptrs += BLOCK_K * stride_ak
+        w_ptrs += BLOCK_K * stride_wk
+
+    # Split: first BLOCK_N is gate, last BLOCK_N is up.
+    # Use static reshape to extract halves.
+    acc_2d = tl.reshape(acc, (BLOCK_M, 2, BLOCK_N))
+    gate = tl.reshape(acc_2d[:, 0, :], (BLOCK_M, BLOCK_N))
+    up = tl.reshape(acc_2d[:, 1, :], (BLOCK_M, BLOCK_N))
+
+    silu_g = gate * tl.sigmoid(gate)
+    out = (silu_g * up).to(tl.bfloat16)
+
+    o_ptrs = (
+        Out_ptr
+        + offs_m[:, None] * stride_om
+        + offs_n[None, :] * stride_on
+    )
+    tl.store(o_ptrs, out)
+
+
+def _grouped_gemm_swiglu(
+    hidden, W_gate, W_up, expert_offsets,
+    T_perm, H, I, E, BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages,
+):
+    W = torch.cat([W_gate, W_up], dim=-1).contiguous()
+    block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M, T_perm)
+    num_m_blocks = block_to_expert.shape[0]
+    num_n_blocks = (I + BLOCK_N - 1) // BLOCK_N
+
+    out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden.device)
+
+    grid = (num_m_blocks, num_n_blocks)
+    _grouped_gemm_swiglu_merged_kernel[grid](
+        hidden, W, out,
+        block_to_expert, block_to_m_abs,
+        H, I,
+        hidden.stride(0), hidden.stride(1),
+        W.stride(0), W.stride(1), W.stride(2),
+        out.stride(0), out.stride(1),
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        num_warps=num_warps, num_stages=num_stages,
+    )
+    return out
+
+
+class Model(nn.Module):
+    def __init__(self, T_total, H, I, E, K):
+        super().__init__()
+        self.T_total, self.H, self.I, self.E, self.K = T_total, H, I, E, K
+        self.W_gate = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
+        self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
+        nn.init.normal_(self.W_gate, std=0.02)
+        nn.init.normal_(self.W_up, std=0.02)
+
+    def forward(self, hidden_states, expert_offsets):
+        T_perm, H = hidden_states.shape
+        return _grouped_gemm_swiglu(
+            hidden_states, self.W_gate, self.W_up, expert_offsets,
+            T_perm=T_perm, H=H, I=self.I, E=self.E,
+            BLOCK_M=64, BLOCK_N=128, BLOCK_K=64,
+            num_warps=4, num_stages=3,
+        )
+
+
+T_total = 32768
+H = 4096
+I = 1536
+E = 128
+K = 8
+
+
+def _build_routing(T_total, E, K, device="cpu"):
+    T_perm = T_total * K
+    base = T_perm // E
+    rem = T_perm - base * E
+    counts = torch.full((E,), base, dtype=torch.int32, device=device)
+    counts[:rem] += 1
+    offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
+    offsets[1:] = torch.cumsum(counts, dim=0)
+    return offsets
+
+
+def get_inputs():
+    T_perm = T_total * K
+    hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
+    expert_offsets = _build_routing(T_total, E, K)
+    return [hidden_states, expert_offsets]
+
+
+def get_init_inputs():
+    return [T_total, H, I, E, K]

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution_persistent.py
@@ -0,0 +1,179 @@
+"""Persistent kernel version: launch fixed number of programs, each handles many tiles.
+
+Improves L2 cache reuse by ensuring consecutive work for the same expert is
+done on the same SM, keeping W[e] hot in L2.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int, T_perm: int):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert: list[int] = []
+    block_to_m_abs: list[int] = []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0:
+            continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+@triton.jit
+def _persistent_grouped_kernel(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    # Each program processes multiple tiles in strided order.
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        # Decode tile_id into (m_block, n_block) with swizzle for L2 locality.
+        # Group_m swizzle: ensure consecutive tile_ids share the same expert.
+        GROUP_M: tl.constexpr = 8
+        num_pid_m = num_m_blocks
+        num_pid_n = num_n_blocks
+        num_pid_in_group = GROUP_M * num_pid_n
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+
+        offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+        w_gate_ptrs = (
+            W_gate_ptr + e * stride_we
+            + offs_k[:, None] * stride_wk
+            + offs_n[None, :] * stride_wn
+        )
+        w_up_ptrs = (
+            W_up_ptr + e * stride_we
+            + offs_k[:, None] * stride_wk
+            + offs_n[None, :] * stride_wn
+        )
+
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w_g = tl.load(w_gate_ptrs)
+            w_u = tl.load(w_up_ptrs)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_ptrs += BLOCK_K * stride_ak
+            w_gate_ptrs += BLOCK_K * stride_wk
+            w_up_ptrs += BLOCK_K * stride_wk
+
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+
+        o_ptrs = (
+            Out_ptr
+            + offs_m[:, None] * stride_om
+            + offs_n[None, :] * stride_on
+        )
+        tl.store(o_ptrs, out)
+
+
+def _grouped_gemm_swiglu(
+    hidden, W_gate, W_up, expert_offsets,
+    T_perm, H, I, E, BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, num_sms=188,
+):
+    block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M, T_perm)
+    num_m_blocks = block_to_expert.shape[0]
+    num_n_blocks = (I + BLOCK_N - 1) // BLOCK_N
+    num_tiles = num_m_blocks * num_n_blocks
+
+    out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden.device)
+
+    grid = (num_sms,)
+    _persistent_grouped_kernel[grid](
+        hidden, W_gate, W_up, out,
+        block_to_expert, block_to_m_abs,
+        H, I,
+        num_m_blocks, num_n_blocks, num_tiles,
+        hidden.stride(0), hidden.stride(1),
+        W_gate.stride(0), W_gate.stride(1), W_gate.stride(2),
+        out.stride(0), out.stride(1),
+        NUM_SMS=num_sms,
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        num_warps=num_warps, num_stages=num_stages,
+    )
+    return out
+
+
+class Model(nn.Module):
+    def __init__(self, T_total, H, I, E, K):
+        super().__init__()
+        self.T_total, self.H, self.I, self.E, self.K = T_total, H, I, E, K
+        self.W_gate = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
+        self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
+        nn.init.normal_(self.W_gate, std=0.02)
+        nn.init.normal_(self.W_up, std=0.02)
+
+    def forward(self, hidden_states, expert_offsets):
+        T_perm, H = hidden_states.shape
+        return _grouped_gemm_swiglu(
+            hidden_states, self.W_gate, self.W_up, expert_offsets,
+            T_perm=T_perm, H=H, I=self.I, E=self.E,
+            BLOCK_M=64, BLOCK_N=128, BLOCK_K=64,
+            num_warps=4, num_stages=3,
+        )
+
+
+T_total = 32768
+H = 4096
+I = 1536
+E = 128
+K = 8
+
+
+def _build_routing(T_total, E, K, device="cpu"):
+    T_perm = T_total * K
+    base = T_perm // E
+    rem = T_perm - base * E
+    counts = torch.full((E,), base, dtype=torch.int32, device=device)
+    counts[:rem] += 1
+    offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
+    offsets[1:] = torch.cumsum(counts, dim=0)
+    retu

... (290 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/sweep_persistent.py
@@ -0,0 +1,146 @@
+"""Sweep over persistent kernel configs."""
+import torch
+import time
+import importlib
+
+def time_kernel_cuda(fn, inputs, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn(*inputs)
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn(*inputs)
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+# Different group sizes
+configs = [
+    # BM, BN, BK, nw, ns, num_sms, group_m
+    (64, 128, 64, 4, 3, 188, 4),
+    (64, 128, 64, 4, 3, 188, 8),
+    (64, 128, 64, 4, 3, 188, 16),
+    (64, 128, 64, 4, 3, 188, 32),
+    (64, 128, 64, 4, 3, 188, 1),  # no swizzle
+    (64, 128, 64, 4, 3, 376, 8),  # 2 waves
+    (64, 128, 64, 4, 3, 752, 8),  # 4 waves
+    (64, 128, 64, 4, 3, 94, 8),   # half SMs
+    (64, 128, 64, 4, 2, 188, 8),
+    (64, 128, 32, 4, 3, 188, 8),
+    (64, 64, 64, 4, 3, 188, 8),
+    (64, 64, 64, 4, 4, 188, 8),
+    (32, 64, 64, 4, 3, 188, 8),
+    (32, 64, 64, 4, 4, 188, 8),
+    (32, 128, 32, 4, 3, 188, 8),
+    (128, 64, 64, 4, 3, 188, 8),
+    (128, 64, 64, 4, 2, 188, 8),
+    (128, 64, 32, 4, 3, 188, 8),
+    (128, 128, 32, 8, 3, 188, 8),
+    (128, 128, 64, 8, 2, 188, 8),
+    (64, 128, 128, 4, 2, 188, 8),
+]
+
+import solution_persistent as sol
+import reference as ref
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    sol.T_total = shape['T_total']; sol.H = shape['H']; sol.I = shape['I']; sol.E = shape['E']; sol.K = shape['K']
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    inputs = [t.cuda() for t in ref.get_inputs()]
+
+    for BM, BN, BK, nw, ns, ns_sms, gm in configs:
+        try:
+            sol_model = sol.Model(*init_args).cuda().eval()
+            sol_model.load_state_dict(ref_model.state_dict())
+            import types
+            def make_forward(BM, BN, BK, nw, ns, ns_sms, gm):
+                def fwd(self, h, o):
+                    return sol._grouped_gemm_swiglu(
+                        h, self.W_gate, self.W_up, o,
+                        T_perm=h.shape[0], H=h.shape[1], I=self.I, E=self.E,
+                        BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
+                        num_warps=nw, num_stages=ns, num_sms=ns_sms,
+                    )
+                return fwd
+            sol_model.forward = types.MethodType(make_forward(BM, BN, BK, nw, ns, ns_sms, gm), sol_model)
+
+            # Monkey-patch the kernel swizzle
+            import triton
+            @triton.jit
+            def _custom_persistent(
+                A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+                block_to_expert_ptr, block_to_m_abs_ptr,
+                H, N,
+                num_m_blocks, num_n_blocks, num_tiles,
+                stride_am, stride_ak,
+                stride_we, stride_wk, stride_wn,
+                stride_om, stride_on,
+                NUM_SMS: tl.constexpr,
+                BLOCK_M: tl.constexpr,
+                BLOCK_N: tl.constexpr,
+                BLOCK_K: tl.constexpr,
+                GROUP_M: tl.constexpr,
+            ):
+                from triton.language import constexpr
+                pid = tl.program_id(0)
+                for tile_id in range(pid, num_tiles, NUM_SMS):
+                    num_pid_in_group = GROUP_M * num_n_blocks
+                    group_id = tile_id // num_pid_in_group
+                    first_pid_m = group_id * GROUP_M
+                    group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+                    pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+                    pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+                    e = tl.load(block_to_expert_ptr + pid_m)
+                    m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+                    offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+                    w_gate_ptrs = W_gate_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+                    w_up_ptrs = W_up_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+                    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+                    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+                    for k in range(0, H, BLOCK_K):
+                        a = tl.load(a_ptrs)
+                        w_g = tl.load(w_gate_ptrs)
+                        w_u = tl.load(w_up_ptrs)
+                        acc_g = tl.dot(a, w_g, acc=acc_g)
+                        acc_u = tl.dot(a, w_u, acc=acc_u)
+                        a_ptrs += BLOCK_K * stride_ak
+                        w_gate_ptrs += BLOCK_K * stride_wk
+                        w_up_ptrs += BLOCK_K * stride_wk
+                    silu_g = acc_g * tl.sigmoid(acc_g)
+                    out = (silu_g * acc_u).to(tl.bfloat16)
+                    o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+                    tl.store(o_ptrs, out)
+
+            import triton.language as tl
+            sol._persistent_grouped_kernel = _custom_persistent
+
+            with torch.no_grad():
+                ref_out = ref_model(*inputs)
+                sol_out = sol_model(*inputs)
+            max_diff = (ref_out - s

... (706 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/sweep_group_m.py
@@ -0,0 +1,176 @@
+"""Sweep GROUP_M for persistent kernel."""
+import torch
+import time
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+        offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+        w_gate_ptrs = W_gate_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        w_up_ptrs = W_up_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w_g = tl.load(w_gate_ptrs)
+            w_u = tl.load(w_up_ptrs)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_ptrs += BLOCK_K * stride_ak
+            w_gate_ptrs += BLOCK_K * stride_wk
+            w_up_ptrs += BLOCK_K * stride_wk
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+        o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, inputs, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn(*inputs)
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn(*inputs)
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+import solution as sol
+import reference as ref
+
+configs = [
+    # BM, BN, BK, nw, ns, num_sms, group_m
+    (64, 128, 64, 4, 3, 188, 1),
+    (64, 128, 64, 4, 3, 188, 4),
+    (64, 128, 64, 4, 3, 188, 8),
+    (64, 128, 64, 4, 3, 188, 16),
+    (64, 128, 64, 4, 3, 188, 32),
+    (64, 128, 64, 4, 3, 188, 64),
+    (64, 128, 64, 4, 3, 188, 128),
+    (64, 128, 64, 4, 3, 376, 8),
+    (64, 128, 64, 4, 3, 752, 8),
+    (64, 128, 32, 4, 3, 188, 8),
+    (64, 128, 32, 4, 4, 188, 8),
+    (64, 128, 32, 4, 3, 188, 16),
+    (64, 128, 32, 4, 3, 188, 32),
+    (64, 64, 64, 4, 4, 188, 8),
+    (64, 64, 64, 4, 5, 188, 8),
+    (32, 128, 32, 4, 4, 188, 8),
+    (32, 128, 64, 4, 4, 188, 8),
+    (128, 64, 64, 4, 3, 188, 8),
+    (128, 64, 32, 4, 3, 188, 8),
+    (128, 128, 32, 8, 3, 188, 8),
+    (128, 128, 64, 8, 2, 188, 8),
+    (128, 64, 64, 4, 2, 188, 8),
+]
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    sol.T_total = shape['T_total']; sol.H = shape['H']; sol.I = shape['I']; sol.E = shape['E']; sol.K = shape['K']
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    W_gate = ref_model.W_gate
+    W_up = ref_model.W_up
+    hidden = torch.randn(shape['T_total']*shape['K'], shape['H'], dtype=torch.bfloat16, device='cuda') * 0.1
+    T_perm = shape['T_total'] * shape['K']
+    base = T_perm // shape['E']
+    rem = T_perm - base * shape['E']
+    counts = torch.full((shape['E'],), base, dtype=torch.int32, device='cuda')
+    counts[:rem] += 1
+    offsets = torch.zeros(shape['E']+1, dtype=torch.int32, device='cuda')
+    offsets[1:] = torch.cumsum(counts, dim=0)
+    inputs = [hidden, offsets]
+
+    for BM, BN, BK, nw, ns, ns_sms, gm in configs:
+        try:
+            def fn(h=hidden, o=offsets, wg=W_gate, wu=W_up, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns, ns_sms=ns_sms, gm=gm):
+                bte, btma = _build_block_table(o, shape['E'], BM)
+                nmb = bte.shape[0]
+                nnb = (shape['I'] + BN - 1) // BN
+                nt = nmb * nnb
+                out = torch.empty(T_perm, shape['I'], dtype=torch.bfloat16, device='cuda')
+                _kernel[(ns_sms,)](
+                    h, wg, wu, out,
+                    bte, btma,
+     

... (1231 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/sweep2.py
@@ -0,0 +1,183 @@
+"""Sweep GROUP_M for persistent kernel - 2nd round."""
+import torch
+import time
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+        offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+        w_gate_ptrs = W_gate_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        w_up_ptrs = W_up_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w_g = tl.load(w_gate_ptrs)
+            w_u = tl.load(w_up_ptrs)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_ptrs += BLOCK_K * stride_ak
+            w_gate_ptrs += BLOCK_K * stride_wk
+            w_up_ptrs += BLOCK_K * stride_wk
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+        o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, inputs, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn(*inputs)
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn(*inputs)
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+import solution as sol
+import reference as ref
+
+configs = [
+    # (BM, BN, BK, nw, ns, num_sms, group_m)
+    (128, 64, 64, 4, 3, 188, 8),
+    (128, 64, 64, 4, 3, 188, 4),
+    (128, 64, 64, 4, 3, 188, 16),
+    (128, 64, 64, 4, 3, 188, 32),
+    (128, 64, 64, 4, 3, 188, 64),
+    (128, 64, 64, 4, 4, 188, 8),
+    (128, 64, 64, 4, 5, 188, 8),
+    (128, 64, 64, 8, 3, 188, 8),
+    (128, 64, 128, 4, 2, 188, 8),
+    (128, 64, 128, 4, 3, 188, 8),
+    (128, 64, 32, 4, 4, 188, 8),
+    (128, 64, 32, 4, 5, 188, 8),
+    (128, 64, 32, 4, 6, 188, 8),
+    (128, 64, 32, 8, 3, 188, 8),
+    (128, 64, 32, 8, 4, 188, 8),
+    (128, 32, 64, 4, 3, 188, 8),
+    (128, 32, 32, 4, 4, 188, 8),
+    (128, 32, 32, 4, 5, 188, 8),
+    (128, 32, 32, 4, 6, 188, 8),
+    (64, 64, 64, 4, 4, 188, 8),
+    (64, 64, 64, 4, 5, 188, 8),
+    (64, 64, 32, 4, 5, 188, 8),
+    (64, 64, 32, 4, 6, 188, 8),
+    (64, 32, 64, 4, 5, 188, 8),
+    (64, 32, 32, 4, 5, 188, 8),
+    (64, 32, 32, 4, 6, 188, 8),
+    (32, 64, 64, 4, 5, 188, 8),
+    (32, 64, 32, 4, 5, 188, 8),
+    (32, 32, 64, 4, 5, 188, 8),
+]
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    sol.T_total = shape['T_total']; sol.H = shape['H']; sol.I = shape['I']; sol.E = shape['E']; sol.K = shape['K']
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    W_gate = ref_model.W_gate
+    W_up = ref_model.W_up
+    hidden = torch.randn(shape['T_total']*shape['K'], shape['H'], dtype=torch.bfloat16, device='cuda') * 0.1
+    T_perm = shape['T_total'] * shape['K']
+    base = T_perm // shape['E']
+    rem = T_perm - base * shape['E']
+    counts = torch.full((shape['E'],), base, dtype=torch.int32, device='cuda')
+    counts[:rem] += 1
+    offsets = torch.zeros(shape['E']+1, dtype=torch.int32, device='cuda')
+    offsets[1:] = torch.cumsum(counts, dim=0)
+    inputs = [hidden, offsets]
+
+    for BM, BN, BK, nw, ns, ns_sms, gm in configs:
+        try:
+            def fn(h=hidden, o=offsets, wg=W_gate, wu=W_up, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns, ns_sms=ns_sms, gm=gm):
+                bte, btma = _build_block_table(o, shape['E'], BM)
+                nmb = bte.shape[0]
+                nnb = (shape['I'] + BN - 1) // BN
+  

... (1422 more chars)

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -3,6 +3,9 @@
 Per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]). Persistent
 Triton kernel with fused SwiGLU epilogue and a host-built block table
 to skip the in-kernel search for the expert owning each m-tile.
+
+Autotune sweeps BLOCK_N, BLOCK_K, GROUP_M, num_warps, num_stages per
+(H, N, BLOCK_M) combination to find the best tile for each shape.
 """
 from __future__ import annotations
 
@@ -33,20 +36,23 @@
     )
 
 
-# Autotune over a few configs that worked in the sweep. One kernel per
-# autotune key keeps the cached compilation focused on the actual shape.
+# Autotune over BLOCK_N, BLOCK_K, GROUP_M, num_warps, num_stages.
+# BLOCK_M is fixed per call (passed as a constexpr), so the autotune
+# key includes it for shape-specialized configs.
 @triton.autotune(
     configs=[
-        triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 16}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=4),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 256, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 16}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=4),
+        triton.Config({"BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=5),
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=8, num_stages=3),
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 32, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_N": 32, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=4),
     ],
-    key=["H", "N", "BLOCK_M_max"],
+    key=["H", "N", "BLOCK_M"],
 )
 @triton.jit
 def _persistent_grouped_kernel(
@@ -58,7 +64,6 @@
     stride_we, stride_wk, stride_wn,
     stride_om, stride_on,
     NUM_SMS: tl.constexpr,
-    BLOCK_M_max: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
@@ -139,23 +144,25 @@
 ):
     block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M)
     num_m_blocks = block_to_expert.shape[0]
-    num_n_blocks = (I + 64 - 1) // 64  # Will be overridden by autotune's BLOCK_N
-    num_tiles = num_m_blocks * num_n_blocks
+    num_sms = _get_num_sms()
 
     out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden.device)
-    num_sms = _get_num_sms()
+
+    # Use max possible num_n_blocks to give autotune headroom; the actual
+    # num_n_blocks depends on BLOCK_N, which is autotuned.
+    max_n = (I + 31) // 32
+    num_tiles = num_m_blocks * max_n
 
     grid = (num_sms,)
     _persistent_grouped_kernel[grid](
         hidden, W_gate, W_up, out,
         block_to_expert, block_to_m_abs,
         H, I,
-        num_m_blocks, num_n_blocks, num_tiles,
+        num_m_blocks, max_n, num_tiles,
         hidden.stride(0), hidden.stride(1),
         W_gate.stride(0), W_gate.stride(1), W_gate.stride(2),
         out.stride(0), out.stride(1),
         NUM_SMS=num_sms,
-        BLOCK_M_max=BLOCK_M,
         BLOCK_M=BLOCK_M,
     )
     return out
@@ -182,12 +189,10 @@
         expert_offsets: torch.Tensor,
     ) -> torch.Tensor:
         T_perm, H = hidden_states.shape
-        # BLOCK_M is fixed; autotune picks the best tile sizes.
-        BLOCK_M = 128
         return _grouped_gemm_swiglu(
             hidden_states, self.W_gate, self.W_up, expert_offsets,
             T_perm=T_perm, H=H, I=self.I, E=self.E,
-            BLOCK_M=BLOCK_M,
+            BLOCK_M=128,
         )
 
 

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -78,6 +78,10 @@
         pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
         pid_n = (tile_id % num_pid_in_group) // group_size_m
 
+        # Bounds check: skip if pid_n is out of range for the autotuned BLOCK_N.
+        if pid_n >= num_n_blocks:
+            continue
+
         e = tl.load(block_to_expert_ptr + pid_m)
         m_abs = tl.load(block_to_m_abs_ptr + pid_m)
 

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -4,8 +4,8 @@
 Triton kernel with fused SwiGLU epilogue and a host-built block table
 to skip the in-kernel search for the expert owning each m-tile.
 
-Autotune sweeps BLOCK_N, BLOCK_K, GROUP_M, num_warps, num_stages per
-(H, N, BLOCK_M) combination to find the best tile for each shape.
+Per-shape configs (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M)
+picked from a benchmark sweep; see comment at _PICK_CONFIG.
 """
 from __future__ import annotations
 
@@ -36,24 +36,6 @@
     )
 
 
-# Autotune over BLOCK_N, BLOCK_K, GROUP_M, num_warps, num_stages.
-# BLOCK_M is fixed per call (passed as a constexpr), so the autotune
-# key includes it for shape-specialized configs.
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_N": 64, "BLOCK_K": 64, "GROUP_M": 16}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=4),
-        triton.Config({"BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=5),
-        triton.Config({"BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=8, num_stages=3),
-        triton.Config({"BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_N": 128, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_N": 32, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_N": 32, "BLOCK_K": 64, "GROUP_M": 8}, num_warps=4, num_stages=4),
-    ],
-    key=["H", "N", "BLOCK_M"],
-)
 @triton.jit
 def _persistent_grouped_kernel(
     A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
@@ -78,10 +60,6 @@
         pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
         pid_n = (tile_id % num_pid_in_group) // group_size_m
 
-        # Bounds check: skip if pid_n is out of range for the autotuned BLOCK_N.
-        if pid_n >= num_n_blocks:
-            continue
-
         e = tl.load(block_to_expert_ptr + pid_m)
         m_abs = tl.load(block_to_m_abs_ptr + pid_m)
 
@@ -145,31 +123,45 @@
     I: int,
     E: int,
     BLOCK_M: int,
+    BLOCK_N: int,
+    BLOCK_K: int,
+    num_warps: int,
+    num_stages: int,
+    group_m: int,
 ):
     block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M)
     num_m_blocks = block_to_expert.shape[0]
+    num_n_blocks = (I + BLOCK_N - 1) // BLOCK_N
+    num_tiles = num_m_blocks * num_n_blocks
     num_sms = _get_num_sms()
 
     out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden.device)
-
-    # Use max possible num_n_blocks to give autotune headroom; the actual
-    # num_n_blocks depends on BLOCK_N, which is autotuned.
-    max_n = (I + 31) // 32
-    num_tiles = num_m_blocks * max_n
 
     grid = (num_sms,)
     _persistent_grouped_kernel[grid](
         hidden, W_gate, W_up, out,
         block_to_expert, block_to_m_abs,
         H, I,
-        num_m_blocks, max_n, num_tiles,
+        num_m_blocks, num_n_blocks, num_tiles,
         hidden.stride(0), hidden.stride(1),
         W_gate.stride(0), W_gate.stride(1), W_gate.stride(2),
         out.stride(0), out.stride(1),
         NUM_SMS=num_sms,
-        BLOCK_M=BLOCK_M,
+        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_M=group_m,
+        num_warps=num_warps, num_stages=num_stages,
     )
     return out
+
+
+# Per-shape best configs from sweep.
+# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M)
+def _pick_config(T_total: int, H: int, I: int, E: int, K: int):
+    """Heuristic per-shape config selection based on sweep data."""
+    T_perm = T_total * K
+    avg_tokens = T_perm // E
+    # For very small per-expert (shape 1), use BM=128 BN=64 BK=64 nw=4 ns=3.
+    # For larger shapes, also BM=128 BN=64 BK=64 nw=4 ns=3 was best.
+    return (128, 64, 64, 4, 3, 8)
 
 
 class Model(nn.Module):
@@ -193,10 +185,12 @@
         expert_offsets: torch.Tensor,
     ) -> torch.Tensor:
         T_perm, H = hidden_states.shape
+        BM, BN, BK, nw, ns, gm = _pick_config(self.T_total, H, self.I, self.E, self.K)
         return _grouped_gemm_swiglu(
             hidden_states, self.W_gate, self.W_up, expert_offsets,
             T_perm=T_perm, H=H, I=self.I, E=self.E,
-            BLOCK_M=128,
+            BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK,
+            num_warps=nw, num_stages=ns, group_m=gm,
         )
 
 

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -1,11 +1,8 @@
 """Grouped GEMM + fused SwiGLU for top-K MoE FFN up-projection.
 
 Per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]). Persistent
-Triton kernel with fused SwiGLU epilogue and a host-built block table
+Triton kernel with fused SwiGLU epilogue and a device-side block table
 to skip the in-kernel search for the expert owning each m-tile.
-
-Per-shape configs (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M)
-picked from a benchmark sweep; see comment at _PICK_CONFIG.
 """
 from __future__ import annotations
 
@@ -15,25 +12,24 @@
 import triton.language as tl
 
 
-def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
-    """Map each global m-block to (expert, m_abs). One CPU sync at forward time."""
-    offsets_cpu = expert_offsets.cpu().to(torch.int64)
-    block_to_expert: list[int] = []
-    block_to_m_abs: list[int] = []
-    for e in range(E):
-        start = int(offsets_cpu[e].item())
-        end = int(offsets_cpu[e + 1].item())
-        n_e = end - start
-        if n_e == 0:
-            continue
-        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
-        for b in range(n_blocks):
-            block_to_expert.append(e)
-            block_to_m_abs.append(start + b * BLOCK_M)
-    return (
-        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
-        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
-    )
+@triton.jit
+def _build_block_table_kernel(
+    counts_ptr,                # (E,) int32
+    expert_offsets_ptr,        # (E+1,) int32
+    block_to_expert_ptr,       # (num_m_blocks,) int32
+    block_to_m_abs_ptr,        # (num_m_blocks,) int32
+    expert_starts_ptr,         # (E+1,) int32 -- cumsum of block counts
+    E, BLOCK_M: tl.constexpr,
+):
+    """For each expert e, fill its block range in the table."""
+    e = tl.program_id(0)
+    count = tl.load(counts_ptr + e)
+    expert_start = tl.load(expert_offsets_ptr + e)
+    block_start = tl.load(expert_starts_ptr + e)
+
+    for b in range(count):
+        tl.store(block_to_expert_ptr + block_start + b, e)
+        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
 
 
 @triton.jit
@@ -113,6 +109,33 @@
     return _NUM_SMS
 
 
+def _build_block_table_gpu(expert_offsets: torch.Tensor, E: int, BLOCK_M: int, T_perm: int):
+    """Build the block table on device, avoiding a CPU sync.
+
+    Steps (all on GPU):
+      1. counts[e] = offsets[e+1] - offsets[e]
+      2. blocks_per_expert[e] = (counts[e] + BM - 1) // BM
+      3. expert_block_starts (cumsum of blocks_per_expert, with 0 prepend)
+      4. fill block_to_expert and block_to_m_abs
+    """
+    counts = expert_offsets[1:] - expert_offsets[:-1]   # (E,) int32
+    blocks_per_expert = (counts + BLOCK_M - 1) // BLOCK_M  # (E,) int32
+
+    # cumsum of blocks_per_expert, prepended with 0
+    expert_block_starts = torch.zeros(E + 1, dtype=torch.int32, device=expert_offsets.device)
+    expert_block_starts[1:] = torch.cumsum(blocks_per_expert, dim=0)
+
+    total_blocks = expert_block_starts[E].item()  # single sync
+    block_to_expert = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
+    block_to_m_abs = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
+
+    _build_block_table_kernel[(E,)](
+        counts, expert_offsets, block_to_expert, block_to_m_abs, expert_block_starts,
+        E, BLOCK_M=BLOCK_M,
+    )
+    return block_to_expert, block_to_m_abs
+
+
 def _grouped_gemm_swiglu(
     hidden: torch.Tensor,
     W_gate: torch.Tensor,
@@ -129,7 +152,7 @@
     num_stages: int,
     group_m: int,
 ):
-    block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M)
+    block_to_expert, block_to_m_abs = _build_block_table_gpu(expert_offsets, E, BLOCK_M, T_perm)
     num_m_blocks = block_to_expert.shape[0]
     num_n_blocks = (I + BLOCK_N - 1) // BLOCK_N
     num_tiles = num_m_blocks * num_n_blocks
@@ -153,14 +176,8 @@
     return out
 
 
-# Per-shape best configs from sweep.
 # (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M)
 def _pick_config(T_total: int, H: int, I: int, E: int, K: int):
-    """Heuristic per-shape config selection based on sweep data."""
-    T_perm = T_total * K
-    avg_tokens = T_perm // E
-    # For very small per-expert (shape 1), use BM=128 BN=64 BK=64 nw=4 ns=3.
-    # For larger shapes, also BM=128 BN=64 BK=64 nw=4 ns=3 was best.
     return (128, 64, 64, 4, 3, 8)
 
 

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/test_block_table.py
@@ -0,0 +1,52 @@
+"""Debug GPU block table."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _build_block_table_kernel(
+    counts_ptr, expert_offsets_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr, expert_starts_ptr,
+    BLOCK_M: tl.constexpr,
+):
+    e = tl.program_id(0)
+    count = tl.load(counts_ptr + e)
+    expert_start = tl.load(expert_offsets_ptr + e)
+    block_start = tl.load(expert_starts_ptr + e)
+
+    for b in range(count):
+        tl.store(block_to_expert_ptr + block_start + b, e)
+        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
+
+
+E, BLOCK_M = 4, 8
+counts = torch.tensor([16, 8, 24, 16], dtype=torch.int32, device='cuda')
+expert_offsets = torch.tensor([0, 16, 24, 48, 64], dtype=torch.int32, device='cuda')
+blocks_per_expert = (counts + BLOCK_M - 1) // BLOCK_M
+print("blocks_per_expert:", blocks_per_expert)
+expert_starts = torch.zeros(E+1, dtype=torch.int32, device='cuda')
+expert_starts[1:] = torch.cumsum(blocks_per_expert, dim=0)
+print("expert_starts:", expert_starts)
+total_blocks = expert_starts[E].item()
+print("total_blocks:", total_blocks)
+
+block_to_expert = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
+block_to_m_abs = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
+_build_block_table_kernel[(E,)](
+    counts, expert_offsets, block_to_expert, block_to_m_abs, expert_starts,
+    BLOCK_M=BLOCK_M,
+)
+print("block_to_expert:", block_to_expert)
+print("block_to_m_abs:  ", block_to_m_abs)
+
+# Expected
+expected_expert = []
+expected_m_abs = []
+for e in range(E):
+    count = blocks_per_expert[e].item()
+    for b in range(count):
+        expected_expert.append(e)
+        expected_m_abs.append(b * BLOCK_M)
+print("expected_expert:", expected_expert)
+print("expected_m_abs: ", expected_m_abs)

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -16,20 +16,26 @@
 def _build_block_table_kernel(
     counts_ptr,                # (E,) int32
     expert_offsets_ptr,        # (E+1,) int32
+    expert_starts_ptr,         # (E+1,) int32 -- cumsum of block counts (with 0 at [0])
     block_to_expert_ptr,       # (num_m_blocks,) int32
     block_to_m_abs_ptr,        # (num_m_blocks,) int32
-    expert_starts_ptr,         # (E+1,) int32 -- cumsum of block counts
-    E, BLOCK_M: tl.constexpr,
+    E,
+    BLOCK_M: tl.constexpr,
 ):
-    """For each expert e, fill its block range in the table."""
+    """One program per (e, b) pair: total programs = sum of blocks_per_expert.
+
+    We launch with a fixed upper-bound grid and exit early if pid exceeds
+    the total. Use a 2D grid (e, b) where b has a max bound.
+    """
     e = tl.program_id(0)
+    b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
+    if b >= count:
+        return
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
-
-    for b in range(count):
-        tl.store(block_to_expert_ptr + block_start + b, e)
-        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
+    tl.store(block_to_expert_ptr + block_start + b, e)
+    tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -116,7 +116,7 @@
 
 
 def _build_block_table_gpu(expert_offsets: torch.Tensor, E: int, BLOCK_M: int, T_perm: int):
-    """Build the block table on device, avoiding a CPU sync.
+    """Build the block table on device, avoiding a CPU sync (except for total_blocks).
 
     Steps (all on GPU):
       1. counts[e] = offsets[e+1] - offsets[e]
@@ -135,8 +135,12 @@
     block_to_expert = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
     block_to_m_abs = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
 
-    _build_block_table_kernel[(E,)](
-        counts, expert_offsets, block_to_expert, block_to_m_abs, expert_block_starts,
+    # max possible blocks per expert
+    max_b = ((T_perm // E + BLOCK_M - 1) // BLOCK_M) + 1
+    grid = (E, max_b)
+    _build_block_table_kernel[grid](
+        counts, expert_offsets, expert_block_starts,
+        block_to_expert, block_to_m_abs,
         E, BLOCK_M=BLOCK_M,
     )
     return block_to_expert, block_to_m_abs

--- a//workspace/problems/06_sonic_moe_swiglu/test_block_table.py
+++ b//workspace/problems/06_sonic_moe_swiglu/test_block_table.py
@@ -6,18 +6,20 @@
 
 @triton.jit
 def _build_block_table_kernel(
-    counts_ptr, expert_offsets_ptr,
-    block_to_expert_ptr, block_to_m_abs_ptr, expert_starts_ptr,
+    counts_ptr, expert_offsets_ptr, expert_starts_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    E,
     BLOCK_M: tl.constexpr,
 ):
     e = tl.program_id(0)
+    b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
+    if b >= count:
+        return
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
-
-    for b in range(count):
-        tl.store(block_to_expert_ptr + block_start + b, e)
-        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
+    tl.store(block_to_expert_ptr + block_start + b, e)
+    tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
 
 
 E, BLOCK_M = 4, 8
@@ -33,9 +35,13 @@
 
 block_to_expert = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
 block_to_m_abs = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
-_build_block_table_kernel[(E,)](
-    counts, expert_offsets, block_to_expert, block_to_m_abs, expert_starts,
-    BLOCK_M=BLOCK_M,
+max_b = (counts.max().item() // BLOCK_M) + 1
+grid = (E, max_b)
+print("grid:", grid)
+_build_block_table_kernel[grid](
+    counts, expert_offsets, expert_starts,
+    block_to_expert, block_to_m_abs,
+    E, BLOCK_M=BLOCK_M,
 )
 print("block_to_expert:", block_to_expert)
 print("block_to_m_abs:  ", block_to_m_abs)
@@ -50,3 +56,4 @@
         expected_m_abs.append(b * BLOCK_M)
 print("expected_expert:", expected_expert)
 print("expected_m_abs: ", expected_m_abs)
+

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/test_block_table2.py
@@ -0,0 +1,48 @@
+"""Test simpler version of block table kernel."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(
+    counts_ptr, expert_offsets_ptr, expert_starts_ptr,
+    out_e_ptr, out_m_ptr, out_idx_ptr,
+    BLOCK_M: tl.constexpr,
+):
+    e = tl.program_id(0)
+    b = tl.program_id(1)
+    count = tl.load(counts_ptr + e)
+    if b >= count:
+        return
+    expert_start = tl.load(expert_offsets_ptr + e)
+    block_start = tl.load(expert_starts_ptr + e)
+    # Store e, m_abs, and the final write position
+    tl.store(out_e_ptr + e * 4 + b, e)
+    tl.store(out_m_ptr + e * 4 + b, expert_start + b * BLOCK_M)
+    tl.store(out_idx_ptr + e * 4 + b, block_start + b)
+
+
+E, BLOCK_M = 4, 8
+counts = torch.tensor([16, 8, 24, 16], dtype=torch.int32, device='cuda')
+expert_offsets = torch.tensor([0, 16, 24, 48, 64], dtype=torch.int32, device='cuda')
+blocks_per_expert = (counts + BLOCK_M - 1) // BLOCK_M
+print("blocks_per_expert:", blocks_per_expert)
+expert_starts = torch.zeros(E+1, dtype=torch.int32, device='cuda')
+expert_starts[1:] = torch.cumsum(blocks_per_expert, dim=0)
+print("expert_starts:", expert_starts)
+
+out_e = torch.zeros(E * 4, dtype=torch.int32, device='cuda')
+out_m = torch.zeros(E * 4, dtype=torch.int32, device='cuda')
+out_idx = torch.zeros(E * 4, dtype=torch.int32, device='cuda')
+max_b = 4
+grid = (E, max_b)
+print("grid:", grid)
+_kernel[grid](
+    counts, expert_offsets, expert_starts,
+    out_e, out_m, out_idx,
+    BLOCK_M=BLOCK_M,
+)
+print("out_e:", out_e)
+print("out_m:", out_m)
+print("out_idx:", out_idx)

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -22,20 +22,17 @@
     E,
     BLOCK_M: tl.constexpr,
 ):
-    """One program per (e, b) pair: total programs = sum of blocks_per_expert.
-
-    We launch with a fixed upper-bound grid and exit early if pid exceeds
-    the total. Use a 2D grid (e, b) where b has a max bound.
-    """
+    """One program per (e, b) pair. Skips writes for b >= count using masks."""
     e = tl.program_id(0)
     b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
-    if b >= count:
-        return
+    # Use mask instead of early return (Triton doesn't always honor return).
+    valid = b < count
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
-    tl.store(block_to_expert_ptr + block_start + b, e)
-    tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
+    if valid:
+        tl.store(block_to_expert_ptr + block_start + b, e)
+        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -22,17 +22,16 @@
     E,
     BLOCK_M: tl.constexpr,
 ):
-    """One program per (e, b) pair. Skips writes for b >= count using masks."""
+    """One program per (e, b) pair. Uses store mask to skip invalid writes."""
     e = tl.program_id(0)
     b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
-    # Use mask instead of early return (Triton doesn't always honor return).
     valid = b < count
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
-    if valid:
-        tl.store(block_to_expert_ptr + block_start + b, e)
-        tl.store(block_to_m_abs_ptr + block_start + b, expert_start + b * BLOCK_M)
+    write_pos = block_start + b
+    tl.store(block_to_expert_ptr + write_pos, e, mask=valid)
+    tl.store(block_to_m_abs_ptr + write_pos, expert_start + b * BLOCK_M, mask=valid)
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/test_block_table.py
+++ b//workspace/problems/06_sonic_moe_swiglu/test_block_table.py
@@ -33,8 +33,8 @@
 total_blocks = expert_starts[E].item()
 print("total_blocks:", total_blocks)
 
-block_to_expert = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
-block_to_m_abs = torch.zeros(total_blocks, dtype=torch.int32, device='cuda')
+block_to_expert = torch.full((total_blocks,), -1, dtype=torch.int32, device='cuda')
+block_to_m_abs = torch.full((total_blocks,), -1, dtype=torch.int32, device='cuda')
 max_b = (counts.max().item() // BLOCK_M) + 1
 grid = (E, max_b)
 print("grid:", grid)

--- a//workspace/problems/06_sonic_moe_swiglu/test_correctness.py
+++ b//workspace/problems/06_sonic_moe_swiglu/test_correctness.py
@@ -17,6 +17,11 @@
 torch.manual_seed(42)
 inputs = [t.cuda() for t in reference.get_inputs()]
 
+# Print first few offsets
+print("expert_offsets[:5]:", inputs[1][:5])
+print("expert_offsets[-5:]:", inputs[1][-5:])
+print("hidden[0, :5]:", inputs[0][0, :5])
+
 with torch.no_grad():
     ref_out = ref_model(*inputs)
     sol_out = sol_model(*inputs)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/test_correctness2.py
@@ -0,0 +1,47 @@
+"""Debug the GPU block table on real data."""
+import torch
+import solution
+import reference
+
+shape = {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8}
+T_total, H, I, E, K = shape['T_total'], shape['H'], shape['I'], shape['E'], shape['K']
+solution.T_total = T_total; solution.H = H; solution.I = I; solution.E = E; solution.K = K
+reference.T_total = T_total; reference.H = H; reference.I = I; reference.E = E; reference.K = K
+
+init_args = reference.get_init_inputs()
+ref_model = reference.Model(*init_args).cuda().eval()
+sol_model = solution.Model(*init_args).cuda().eval()
+sol_model.load_state_dict(ref_model.state_dict())
+
+torch.manual_seed(42)
+inputs = [t.cuda() for t in reference.get_inputs()]
+hidden, expert_offsets = inputs
+T_perm = hidden.shape[0]
+
+# Manually build the block table the way the host version does
+import time
+t0 = time.time()
+block_to_expert_host, block_to_m_abs_host = solution._build_block_table(expert_offsets, E, 128)
+t1 = time.time()
+print(f"host build: {(t1-t0)*1000:.3f} ms")
+print(f"host block_to_expert[:20]: {block_to_expert_host[:20]}")
+print(f"host block_to_m_abs[:20]: {block_to_m_abs_host[:20]}")
+print(f"host total blocks: {block_to_expert_host.shape[0]}")
+
+# Manually build the GPU block table
+t0 = time.time()
+block_to_expert_gpu, block_to_m_abs_gpu = solution._build_block_table_gpu(expert_offsets, E, 128, T_perm)
+t1 = time.time()
+print(f"GPU build: {(t1-t0)*1000:.3f} ms")
+print(f"GPU block_to_expert[:20]: {block_to_expert_gpu[:20]}")
+print(f"GPU block_to_m_abs[:20]: {block_to_m_abs_gpu[:20]}")
+print(f"GPU total blocks: {block_to_expert_gpu.shape[0]}")
+
+# Compare
+print(f"\nexpert equal: {torch.equal(block_to_expert_host, block_to_expert_gpu)}")
+print(f"m_abs equal: {torch.equal(block_to_m_abs_host, block_to_m_abs_gpu)}")
+if not torch.equal(block_to_expert_host, block_to_expert_gpu):
+    diff_idx = (block_to_expert_host != block_to_expert_gpu).nonzero()
+    print(f"first expert diff at: {diff_idx[:5].flatten()}")
+    for i in diff_idx[:5].flatten():
+        print(f"  [{i}]: host={block_to_expert_host[i].item()}, gpu={block_to_expert_gpu[i].item()}")

--- a//workspace/problems/06_sonic_moe_swiglu/test_correctness2.py
+++ b//workspace/problems/06_sonic_moe_swiglu/test_correctness2.py
@@ -20,8 +20,25 @@
 
 # Manually build the block table the way the host version does
 import time
+def _build_block_table_host(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
 t0 = time.time()
-block_to_expert_host, block_to_m_abs_host = solution._build_block_table(expert_offsets, E, 128)
+block_to_expert_host, block_to_m_abs_host = _build_block_table_host(expert_offsets, E, 128)
 t1 = time.time()
 print(f"host build: {(t1-t0)*1000:.3f} ms")
 print(f"host block_to_expert[:20]: {block_to_expert_host[:20]}")

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -22,16 +22,21 @@
     E,
     BLOCK_M: tl.constexpr,
 ):
-    """One program per (e, b) pair. Uses store mask to skip invalid writes."""
+    """One program per (e, b) pair. Uses tl.where to gate writes (Triton's
+    `if ... return` doesn't always short-circuit)."""
     e = tl.program_id(0)
     b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
-    valid = b < count
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
     write_pos = block_start + b
-    tl.store(block_to_expert_ptr + write_pos, e, mask=valid)
-    tl.store(block_to_m_abs_ptr + write_pos, expert_start + b * BLOCK_M, mask=valid)
+    # Always write, but write 0 or "no-op" for invalid positions.
+    # Use a sentinel: if invalid, write 0 which we'll filter out later.
+    # Better: only write to the right position by computing a conditional.
+    expert_val = tl.where(b < count, e, 0)
+    m_abs_val = tl.where(b < count, expert_start + b * BLOCK_M, 0)
+    tl.store(block_to_expert_ptr + write_pos, expert_val)
+    tl.store(block_to_m_abs_ptr + write_pos, m_abs_val)
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -22,21 +22,17 @@
     E,
     BLOCK_M: tl.constexpr,
 ):
-    """One program per (e, b) pair. Uses tl.where to gate writes (Triton's
-    `if ... return` doesn't always short-circuit)."""
+    """One program per (e, b) pair. Use scalar mask with tl.store."""
     e = tl.program_id(0)
     b = tl.program_id(1)
     count = tl.load(counts_ptr + e)
     expert_start = tl.load(expert_offsets_ptr + e)
     block_start = tl.load(expert_starts_ptr + e)
     write_pos = block_start + b
-    # Always write, but write 0 or "no-op" for invalid positions.
-    # Use a sentinel: if invalid, write 0 which we'll filter out later.
-    # Better: only write to the right position by computing a conditional.
-    expert_val = tl.where(b < count, e, 0)
-    m_abs_val = tl.where(b < count, expert_start + b * BLOCK_M, 0)
-    tl.store(block_to_expert_ptr + write_pos, expert_val)
-    tl.store(block_to_m_abs_ptr + write_pos, m_abs_val)
+    # Scalar mask should work for tl.store.
+    if b < count:
+        tl.store(block_to_expert_ptr + write_pos, e)
+        tl.store(block_to_m_abs_ptr + write_pos, expert_start + b * BLOCK_M)
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -1,7 +1,7 @@
 """Grouped GEMM + fused SwiGLU for top-K MoE FFN up-projection.
 
 Per-expert h_e = silu(x_e @ W_gate[e]) * (x_e @ W_up[e]). Persistent
-Triton kernel with fused SwiGLU epilogue and a device-side block table
+Triton kernel with fused SwiGLU epilogue and a host-built block table
 to skip the in-kernel search for the expert owning each m-tile.
 """
 from __future__ import annotations
@@ -12,27 +12,25 @@
 import triton.language as tl
 
 
-@triton.jit
-def _build_block_table_kernel(
-    counts_ptr,                # (E,) int32
-    expert_offsets_ptr,        # (E+1,) int32
-    expert_starts_ptr,         # (E+1,) int32 -- cumsum of block counts (with 0 at [0])
-    block_to_expert_ptr,       # (num_m_blocks,) int32
-    block_to_m_abs_ptr,        # (num_m_blocks,) int32
-    E,
-    BLOCK_M: tl.constexpr,
-):
-    """One program per (e, b) pair. Use scalar mask with tl.store."""
-    e = tl.program_id(0)
-    b = tl.program_id(1)
-    count = tl.load(counts_ptr + e)
-    expert_start = tl.load(expert_offsets_ptr + e)
-    block_start = tl.load(expert_starts_ptr + e)
-    write_pos = block_start + b
-    # Scalar mask should work for tl.store.
-    if b < count:
-        tl.store(block_to_expert_ptr + write_pos, e)
-        tl.store(block_to_m_abs_ptr + write_pos, expert_start + b * BLOCK_M)
+def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
+    """Map each global m-block to (expert, m_abs). One CPU sync at forward time."""
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert: list[int] = []
+    block_to_m_abs: list[int] = []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0:
+            continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
 
 
 @triton.jit
@@ -112,37 +110,6 @@
     return _NUM_SMS
 
 
-def _build_block_table_gpu(expert_offsets: torch.Tensor, E: int, BLOCK_M: int, T_perm: int):
-    """Build the block table on device, avoiding a CPU sync (except for total_blocks).
-
-    Steps (all on GPU):
-      1. counts[e] = offsets[e+1] - offsets[e]
-      2. blocks_per_expert[e] = (counts[e] + BM - 1) // BM
-      3. expert_block_starts (cumsum of blocks_per_expert, with 0 prepend)
-      4. fill block_to_expert and block_to_m_abs
-    """
-    counts = expert_offsets[1:] - expert_offsets[:-1]   # (E,) int32
-    blocks_per_expert = (counts + BLOCK_M - 1) // BLOCK_M  # (E,) int32
-
-    # cumsum of blocks_per_expert, prepended with 0
-    expert_block_starts = torch.zeros(E + 1, dtype=torch.int32, device=expert_offsets.device)
-    expert_block_starts[1:] = torch.cumsum(blocks_per_expert, dim=0)
-
-    total_blocks = expert_block_starts[E].item()  # single sync
-    block_to_expert = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
-    block_to_m_abs = torch.empty(total_blocks, dtype=torch.int32, device=expert_offsets.device)
-
-    # max possible blocks per expert
-    max_b = ((T_perm // E + BLOCK_M - 1) // BLOCK_M) + 1
-    grid = (E, max_b)
-    _build_block_table_kernel[grid](
-        counts, expert_offsets, expert_block_starts,
-        block_to_expert, block_to_m_abs,
-        E, BLOCK_M=BLOCK_M,
-    )
-    return block_to_expert, block_to_m_abs
-
-
 def _grouped_gemm_swiglu(
     hidden: torch.Tensor,
     W_gate: torch.Tensor,
@@ -159,7 +126,7 @@
     num_stages: int,
     group_m: int,
 ):
-    block_to_expert, block_to_m_abs = _build_block_table_gpu(expert_offsets, E, BLOCK_M, T_perm)
+    block_to_expert, block_to_m_abs = _build_block_table(expert_offsets, E, BLOCK_M)
     num_m_blocks = block_to_expert.shape[0]
     num_n_blocks = (I + BLOCK_N - 1) // BLOCK_N
     num_tiles = num_m_blocks * num_n_blocks

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -152,7 +152,7 @@
 
 # (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, GROUP_M)
 def _pick_config(T_total: int, H: int, I: int, E: int, K: int):
-    return (128, 64, 64, 4, 3, 8)
+    return (128, 64, 64, 4, 3, 4)
 
 
 class Model(nn.Module):

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution_merged2.py
@@ -0,0 +1,221 @@
+"""Test merged-dot kernel with proper B-tile split."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel_merged(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    """Use single (BLOCK_K, 2*BLOCK_N) B load, then split via reshape+transpose+split."""
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+        offs_m = m_abs + tl.arange(0, BLOCK_M)
+        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        offs_2n = tl.arange(0, 2 * BLOCK_N)
+        offs_k = tl.arange(0, BLOCK_K)
+
+        a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+        # Pointer for the (BLOCK_K, 2*BLOCK_N) tile at W[e]
+        w_combined_ptrs = (
+            W_gate_ptr + e * stride_we
+            + offs_k[:, None] * stride_wk
+            + offs_2n[None, :] * stride_wn
+        )
+
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w = tl.load(w_combined_ptrs)  # (BLOCK_K, 2*BLOCK_N)
+            # Split via reshape + transpose + split.
+            # (BLOCK_K, 2*BLOCK_N) -> (BLOCK_K, 2, BLOCK_N) -> (BLOCK_K, BLOCK_N, 2) -> split
+            w_r = tl.reshape(w, (BLOCK_K, 2, BLOCK_N))
+            w_t = tl.trans(w_r, (0, 2, 1))  # (BLOCK_K, BLOCK_N, 2)
+            w_g, w_u = tl.split(w_t)  # (BLOCK_K, BLOCK_N), (BLOCK_K, BLOCK_N)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_ptrs += BLOCK_K * stride_ak
+            w_combined_ptrs += BLOCK_K * stride_wk
+
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+        o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+@triton.jit
+def _kernel_separate(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+        offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+        w_gate_ptrs = W_gate_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        w_up_ptrs = W_up_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w_g = tl.load(w_gate_ptrs)
+            w_u = tl.load(w_up_ptrs)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_ptrs += BLOCK_K * stride_ak
+            w_gate_ptrs += BLOCK_K * stride_wk
+            w_up_ptrs += BLOCK_K * stride_wk
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+        o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_tot

... (2911 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution_merged3.py
@@ -0,0 +1,174 @@
+"""Test merged-dot with physically concatenated W.
+
+We do a single dot of (BM, K) @ (K, 2*BN) producing (BM, 2*BN), then split
+the result. The split is done via tl.split after a transpose. We pre-concat
+W = cat(W_gate, W_up, dim=-1) so the load is contiguous.
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel_merged(
+    A_ptr, W_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+        offs_m = m_abs + tl.arange(0, BLOCK_M)
+        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        offs_2n = tl.arange(0, 2 * BLOCK_N)
+        offs_k = tl.arange(0, BLOCK_K)
+
+        a_ptrs = A_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+        w_ptrs = W_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_2n[None, :] * stride_wn
+
+        acc = tl.zeros((BLOCK_M, 2 * BLOCK_N), dtype=tl.float32)
+
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_ptrs)
+            w = tl.load(w_ptrs)  # (BLOCK_K, 2*BLOCK_N)
+            acc = tl.dot(a, w, acc=acc)
+            a_ptrs += BLOCK_K * stride_ak
+            w_ptrs += BLOCK_K * stride_wk
+
+        # Split acc (BM, 2*BN) into gate (BM, BN) and up (BM, BN).
+        # Use reshape + transpose + split.
+        acc_r = tl.reshape(acc, (BLOCK_M, 2, BLOCK_N))
+        acc_t = tl.trans(acc_r, (0, 2, 1))  # (BM, BN, 2)
+        gate, up = tl.split(acc_t)  # both (BM, BN)
+
+        silu_g = gate * tl.sigmoid(gate)
+        out = (silu_g * up).to(tl.bfloat16)
+        o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+import reference as ref
+
+# Try different configs
+configs = [
+    (128, 64, 64, 4, 3, 188, 4),
+    (128, 64, 32, 4, 3, 188, 4),
+    (128, 32, 64, 4, 3, 188, 4),
+    (64, 64, 64, 4, 3, 188, 4),
+    (64, 128, 32, 4, 3, 188, 4),
+    (128, 64, 64, 4, 2, 188, 4),
+    (128, 64, 64, 4, 4, 188, 4),
+    (128, 64, 64, 8, 3, 188, 4),
+    (128, 64, 64, 8, 2, 188, 4),
+    (256, 32, 64, 4, 3, 188, 4),
+    (256, 64, 32, 4, 3, 188, 4),
+    (256, 64, 64, 8, 2, 188, 4),
+]
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    W_gate = ref_model.W_gate
+    W_up = ref_model.W_up
+    W_concat = torch.cat([W_gate, W_up], dim=-1).contiguous()  # (E, H, 2*I)
+    hidden = torch.randn(shape['T_total']*shape['K'], shape['H'], dtype=torch.bfloat16, device='cuda') * 0.1
+    T_perm = shape['T_total'] * shape['K']
+    base = T_perm // shape['E']
+    rem = T_perm - base * shape['E']
+    counts = torch.full((shape['E'],), base, dtype=torch.int32, device='cuda')
+    counts[:rem] += 1
+    offsets = torch.zeros(shape['E']+1, dtype=torch.int32, device='cuda')
+    offsets[1:] = torch.cumsum(counts, dim=0)
+
+    for BM, BN, BK, nw, ns, ns_sms, gm in configs:
+        try:
+            bte, btma = _build_block_table(offsets, shape['E'], BM)
+            nmb = bte.shape[0]
+            nnb = (shape['I'] + BN - 1) // BN
+            nt = nmb * nnb
+            out = torch.empty(T_perm, shape['I'], dtype=torch.bfloat16, device='cuda')
+
+            def fn(h=hidden, wg=W_concat, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns, ns_sms=ns_sms, gm=gm, bte=bte, btma=btma, nmb=nmb, nnb=nnb, nt=nt, out=out):
+                _kernel_merged[(ns_sms,)](
+                    h, wg, out,
+                    bte, btma,
+                    shape['H'], shape['I'],
+                    nmb, nnb, nt,
+                    h.stride(0), h.stride(1),
+                    wg.stride(0), wg.stride(1), wg.stride(2),
+                    out.stride(0), 

... (938 more chars)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/solution_nonpersistent.py
@@ -0,0 +1,177 @@
+"""Non-persistent kernel for comparison."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    e = tl.load(block_to_expert_ptr + pid_m)
+    m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+
+    offs_m = m_abs + 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_am + offs_k[None, :] * stride_ak
+    w_gate_ptrs = (
+        W_gate_ptr + e * stride_we
+        + offs_k[:, None] * stride_wk
+        + offs_n[None, :] * stride_wn
+    )
+    w_up_ptrs = (
+        W_up_ptr + e * stride_we
+        + offs_k[:, None] * stride_wk
+        + offs_n[None, :] * stride_wn
+    )
+
+    acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, H, BLOCK_K):
+        a = tl.load(a_ptrs)
+        w_g = tl.load(w_gate_ptrs)
+        w_u = tl.load(w_up_ptrs)
+        acc_g = tl.dot(a, w_g, acc=acc_g)
+        acc_u = tl.dot(a, w_u, acc=acc_u)
+        a_ptrs += BLOCK_K * stride_ak
+        w_gate_ptrs += BLOCK_K * stride_wk
+        w_up_ptrs += BLOCK_K * stride_wk
+
+    silu_g = acc_g * tl.sigmoid(acc_g)
+    out = (silu_g * acc_u).to(tl.bfloat16)
+    o_ptrs = Out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(o_ptrs, out)
+
+
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+    {'T_total': 4096, 'H': 2048, 'I': 1024, 'E': 64, 'K': 4},
+    {'T_total': 16384, 'H': 2048, 'I': 4096, 'E': 64, 'K': 8},
+]
+
+import reference as ref
+
+# Try various configs with non-persistent kernel
+configs = [
+    (128, 64, 64, 4, 3),
+    (64, 64, 64, 4, 3),
+    (64, 128, 64, 4, 3),
+    (64, 128, 32, 4, 3),
+    (64, 128, 32, 4, 4),
+    (32, 64, 64, 4, 4),
+    (32, 64, 64, 4, 5),
+    (32, 128, 32, 4, 4),
+    (32, 128, 32, 4, 5),
+    (32, 128, 32, 4, 6),
+    (16, 64, 64, 4, 4),
+    (16, 64, 64, 4, 5),
+    (16, 64, 64, 4, 6),
+    (16, 128, 32, 4, 4),
+    (16, 128, 32, 4, 5),
+    (16, 128, 32, 4, 6),
+    (16, 128, 32, 4, 8),
+    (16, 64, 128, 4, 3),
+    (16, 64, 128, 4, 4),
+    (16, 32, 128, 4, 3),
+    (16, 32, 128, 4, 4),
+    (16, 32, 128, 4, 5),
+]
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    W_gate = ref_model.W_gate
+    W_up = ref_model.W_up
+    hidden = torch.randn(shape['T_total']*shape['K'], shape['H'], dtype=torch.bfloat16, device='cuda') * 0.1
+    T_perm = shape['T_total'] * shape['K']
+    base = T_perm // shape['E']
+    rem = T_perm - base * shape['E']
+    counts = torch.full((shape['E'],), base, dtype=torch.int32, device='cuda')
+    counts[:rem] += 1
+    offsets = torch.zeros(shape['E']+1, dtype=torch.int32, device='cuda')
+    offsets[1:] = torch.cumsum(counts, dim=0)
+
+    for BM, BN, BK, nw, ns in configs:
+        try:
+            bte, btma = _build_block_table(offsets, shape['E'], BM)
+            nmb = bte.shape[0]
+            nnb = (shape['I'] + BN - 1) // BN
+            out = torch.empty(T_perm, shape['I'], dtype=torch.bfloat16, device='cuda')
+
+            def fn(h=hidden, o=offsets, wg=W_gate, wu=W_up, BM=BM, BN=BN, BK=BK, nw=nw, ns=ns, bte=bte, btma=btma, nmb=nmb, nnb=nnb, out=out):
+                _kernel[(nmb, nnb)](
+                    h, wg, wu, out,
+                    bte, btma,
+                    shape['H'], shape['I'],
+                    nmb, nnb,
+                    h.stride(0), h.stride(1),
+                    wg.stride(0), wg.stride(1), wg.stride(2),
+                    out.stride(0), out.stride(1),
+                    BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK, GROUP_M=1,
+                    num_warps=nw, num_stages=ns,
+                )
+
+            with torch.no_grad():
+                ref_out = ref_model(hidden, offsets)
+                fn()
+                max_diff = (ref_out - out).abs().max().item()
+            if max_diff > 0.05:
+                print(f'  BM={BM} BN={BN} BK={BK} nw={nw} ns={ns}: WRONG (diff={max_diff:.3f})')
+                continue
+            ms = time_kernel_cuda(fn, n_iter=15)
+            flops = 2 * shape['T_total'] * shape['H'] * 2 * shape['I

... (290 more chars)

--- a//workspace/problems/06_sonic_moe_swiglu/test_range.py
+++ b//workspace/problems/06_sonic_moe_swiglu/test_range.py
@@ -40,7 +40,7 @@
         w_up_ptrs = W_up_ptr + e * stride_we + offs_k[:, None] * stride_wk + offs_n[None, :] * stride_wn
         acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
         acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
-        for k in tl.range(0, H, BLOCK_K, num_stages=3, loop_unroll=LOOP_UNROLL):
+        for k in tl.range(0, H, BLOCK_K, num_stages=3):
             a = tl.load(a_ptrs)
             w_g = tl.load(w_gate_ptrs)
             w_u = tl.load(w_up_ptrs)

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/look_ptx.py
@@ -0,0 +1,42 @@
+"""Look at the generated PTX for the kernel."""
+import os
+os.environ['TRITON_KERNEL_DUMP'] = '1'
+os.environ['TRITON_DUMP_DIR'] = '/tmp/triton_dump'
+os.makedirs('/tmp/triton_dump', exist_ok=True)
+
+import torch
+import solution
+import reference
+
+shape = {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8}
+T_total, H, I, E, K = shape['T_total'], shape['H'], shape['I'], shape['E'], shape['K']
+solution.T_total = T_total; solution.H = H; solution.I = I; solution.E = E; solution.K = K
+reference.T_total = T_total; reference.H = H; reference.I = I; reference.E = E; reference.K = K
+
+init_args = reference.get_init_inputs()
+ref_model = reference.Model(*init_args).cuda().eval()
+sol_model = solution.Model(*init_args).cuda().eval()
+sol_model.load_state_dict(ref_model.state_dict())
+
+torch.manual_seed(42)
+inputs = [t.cuda() for t in reference.get_inputs()]
+
+with torch.no_grad():
+    out = sol_model(*inputs)
+
+# Find the generated PTX
+import glob
+ptx_files = glob.glob('/tmp/triton_dump/**/*.ptx', recursive=True)
+print("PTX files:", ptx_files)
+for f in ptx_files:
+    print(f"\n=== {f} ===")
+    with open(f) as fp:
+        content = fp.read()
+        # Print first 2000 chars
+        print(content[:3000])
+        # Count mma instructions
+        mma_count = content.count('mma.sync')
+        print(f"\n\nmma.sync count: {mma_count}")
+        cp_count = content.count('cp.async')
+        print(f"cp.async count: {cp_count}")
+        break

--- /dev/null
+++ b//workspace/problems/06_sonic_moe_swiglu/test_tma.py
@@ -0,0 +1,184 @@
+"""Test kernel using block pointers (TMA on SM90+)."""
+import torch
+import time
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kernel(
+    A_ptr, W_gate_ptr, W_up_ptr, Out_ptr,
+    block_to_expert_ptr, block_to_m_abs_ptr,
+    H, N,
+    num_m_blocks, num_n_blocks, num_tiles,
+    stride_am, stride_ak,
+    stride_we, stride_wk, stride_wn,
+    stride_om, stride_on,
+    NUM_SMS: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_M: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    for tile_id in range(pid, num_tiles, NUM_SMS):
+        num_pid_in_group = GROUP_M * num_n_blocks
+        group_id = tile_id // num_pid_in_group
+        first_pid_m = group_id * GROUP_M
+        group_size_m = min(num_m_blocks - first_pid_m, GROUP_M)
+        pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
+        pid_n = (tile_id % num_pid_in_group) // group_size_m
+
+        e = tl.load(block_to_expert_ptr + pid_m)
+        m_abs = tl.load(block_to_m_abs_ptr + pid_m)
+
+        # Use block pointers for TMA.
+        a_block_ptr = tl.make_block_ptr(
+            base=A_ptr,
+            shape=(num_m_blocks * BLOCK_M + 16, H),  # oversized
+            strides=(stride_am, stride_ak),
+            offsets=(pid_m, 0),
+            block_shape=(BLOCK_M, BLOCK_K),
+            order=(1, 0),
+        )
+        w_gate_block_ptr = tl.make_block_ptr(
+            base=W_gate_ptr,
+            shape=(H, N),  # placeholder
+            strides=(stride_wk, stride_wn),
+            offsets=(0, pid_n * BLOCK_N),
+            block_shape=(BLOCK_K, BLOCK_N),
+            order=(1, 0),
+        )
+        w_up_block_ptr = tl.make_block_ptr(
+            base=W_up_ptr,
+            shape=(H, N),
+            strides=(stride_wk, stride_wn),
+            offsets=(0, pid_n * BLOCK_N),
+            block_shape=(BLOCK_K, BLOCK_N),
+            order=(1, 0),
+        )
+
+        # Add expert offset to base pointers
+        a_block_ptr = tl.advance(a_block_ptr, (m_abs - pid_m * BLOCK_M, 0))
+        # For W, we need to add the expert offset to the base
+        w_gate_block_ptr = tl.make_block_ptr(
+            base=W_gate_ptr + e * stride_we,
+            shape=(H, N),
+            strides=(stride_wk, stride_wn),
+            offsets=(0, pid_n * BLOCK_N),
+            block_shape=(BLOCK_K, BLOCK_N),
+            order=(1, 0),
+        )
+        w_up_block_ptr = tl.make_block_ptr(
+            base=W_up_ptr + e * stride_we,
+            shape=(H, N),
+            strides=(stride_wk, stride_wn),
+            offsets=(0, pid_n * BLOCK_N),
+            block_shape=(BLOCK_K, BLOCK_N),
+            order=(1, 0),
+        )
+
+        acc_g = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        acc_u = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for k in range(0, H, BLOCK_K):
+            a = tl.load(a_block_ptr)  # (BLOCK_M, BLOCK_K)
+            w_g = tl.load(w_gate_block_ptr)
+            w_u = tl.load(w_up_block_ptr)
+            acc_g = tl.dot(a, w_g, acc=acc_g)
+            acc_u = tl.dot(a, w_u, acc=acc_u)
+            a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_K))
+            w_gate_block_ptr = tl.advance(w_gate_block_ptr, (BLOCK_K, 0))
+            w_up_block_ptr = tl.advance(w_up_block_ptr, (BLOCK_K, 0))
+
+        silu_g = acc_g * tl.sigmoid(acc_g)
+        out = (silu_g * acc_u).to(tl.bfloat16)
+        o_ptrs = Out_ptr + (m_abs + tl.arange(0, BLOCK_M))[:, None] * stride_om + (pid_n * BLOCK_N + tl.arange(0, BLOCK_N))[None, :] * stride_on
+        tl.store(o_ptrs, out)
+
+
+# Other imports/funcs same as before
+def _build_block_table(expert_offsets, E, BLOCK_M):
+    offsets_cpu = expert_offsets.cpu().to(torch.int64)
+    block_to_expert, block_to_m_abs = [], []
+    for e in range(E):
+        start = int(offsets_cpu[e].item())
+        end = int(offsets_cpu[e + 1].item())
+        n_e = end - start
+        if n_e == 0: continue
+        n_blocks = (n_e + BLOCK_M - 1) // BLOCK_M
+        for b in range(n_blocks):
+            block_to_expert.append(e)
+            block_to_m_abs.append(start + b * BLOCK_M)
+    return (
+        torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
+        torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
+    )
+
+
+def time_kernel_cuda(fn, n_warmup=5, n_iter=20):
+    for _ in range(n_warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(n_iter):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / n_iter
+
+
+shapes = [
+    {'T_total': 32768, 'H': 4096, 'I': 1536, 'E': 128, 'K': 8},
+]
+
+import reference as ref
+
+for shape in shapes:
+    print(f'\n=== Shape: {shape} ===')
+    ref.T_total = shape['T_total']; ref.H = shape['H']; ref.I = shape['I']; ref.E = shape['E']; ref.K = shape['K']
+
+    init_args = ref.get_init_inputs()
+    ref_model = ref.Model(*init_args).cuda().eval()
+    W_gate = ref_model.W_gate
+    W_up = ref_model.W_up
+    hidden = torch.randn(shape['T_total']*shape['K'], shape['H'], dtype=torch.bfloat16, device='cuda') * 0.1
+    T_perm = shape['T_total'] * shape['K']
+    base = T_perm // shape['E']
+    rem = T_perm - base * shape['E']
+    counts = torch.full((shape['E'],), base, dtype=torch.int32, device='cuda')
+    counts[:rem] += 1
+    offsets = torch.zeros(shape['E']+1, dtype=torch.int32, device='cuda')
+    offsets[1:] = torch.cumsum(counts, dim=0)
+
+    BM, BN, BK, nw, ns, ns_sms, gm = 128, 64, 64, 4, 3, 188, 4
+    bte, btma = _build_block_table(offsets, shape['E'], BM)
+    nmb = bte.shape[0]
+    nnb = (shape['I'] + BN - 1) // BN
+    nt = nmb * nnb
+    out = torch.empty(T_perm, shape['I'], dtype=torch.bfloat16, de

... (878 more chars)

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -12,8 +12,21 @@
 import triton.language as tl
 
 
+_BLOCK_TABLE_CACHE: dict = {}
+
+
 def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
-    """Map each global m-block to (expert, m_abs). One CPU sync at forward time."""
+    """Map each global m-block to (expert, m_abs). Cached by (data_ptr, BLOCK_M).
+
+    Returns:
+        block_to_expert: (num_m_blocks,) int32
+        block_to_m_abs:  (num_m_blocks,) int32
+    """
+    key = (expert_offsets.data_ptr(), int(BLOCK_M), E)
+    cached = _BLOCK_TABLE_CACHE.get(key)
+    if cached is not None:
+        return cached
+
     offsets_cpu = expert_offsets.cpu().to(torch.int64)
     block_to_expert: list[int] = []
     block_to_m_abs: list[int] = []
@@ -27,10 +40,12 @@
         for b in range(n_blocks):
             block_to_expert.append(e)
             block_to_m_abs.append(start + b * BLOCK_M)
-    return (
+    result = (
         torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
         torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
     )
+    _BLOCK_TABLE_CACHE[key] = result
+    return result
 
 
 @triton.jit

--- a//workspace/problems/06_sonic_moe_swiglu/solution.py
+++ b//workspace/problems/06_sonic_moe_swiglu/solution.py
@@ -12,21 +12,8 @@
 import triton.language as tl
 
 
-_BLOCK_TABLE_CACHE: dict = {}
-
-
 def _build_block_table(expert_offsets: torch.Tensor, E: int, BLOCK_M: int):
-    """Map each global m-block to (expert, m_abs). Cached by (data_ptr, BLOCK_M).
-
-    Returns:
-        block_to_expert: (num_m_blocks,) int32
-        block_to_m_abs:  (num_m_blocks,) int32
-    """
-    key = (expert_offsets.data_ptr(), int(BLOCK_M), E)
-    cached = _BLOCK_TABLE_CACHE.get(key)
-    if cached is not None:
-        return cached
-
+    """Map each global m-block to (expert, m_abs). One CPU sync at forward time."""
     offsets_cpu = expert_offsets.cpu().to(torch.int64)
     block_to_expert: list[int] = []
     block_to_m_abs: list[int] = []
@@ -40,12 +27,10 @@
         for b in range(n_blocks):
             block_to_expert.append(e)
             block_to_m_abs.append(start + b * BLOCK_M)
-    result = (
+    return (
         torch.tensor(block_to_expert, dtype=torch.int32, device=expert_offsets.device),
         torch.tensor(block_to_m_abs, dtype=torch.int32, device=expert_offsets.device),
     )
-    _BLOCK_TABLE_CACHE[key] = result
-    return result
 
 
 @triton.jit