--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -0,0 +1,287 @@
+"""W4A16 weight-only quantized GEMM for RTX PRO 6000 (SM120).
+
+AWQ/GPTQ-style asymmetric int4 with bf16 per-group scales/zeros.
+Fused unpack + GEMM via Triton.
+
+The scheme:
+  w_bf[k, n] = (unpack(w_q)[k, n] - zeros[k // 128, n]) * scales[k // 128, n]
+  out[m, n]  = sum_k x[m, k] * w_bf[k, n]
+
+Key optimizations:
+  - Two-dot pattern: process even/odd K rows with two separate dots. Each
+    group_size=128 K rows become two 64-wide dots that both read the same x row
+    span but different nibbles of the packed weight byte.
+  - Decode path (M=1) uses a memory-bound split-K GEMV kernel that loads
+    w_q, scales, and zeros on the fly and accumulates in fp32.
+  - Prefill path (M>=16) uses a tile-based GEMM kernel that dequantizes
+    BLOCK_K rows at a time and uses tensor cores.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+import triton
+import triton.language as tl
+
+OP_TYPE = "gemm_w4a16"
+SUPPORTED_PRECISIONS = ["int4_bf16"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+GROUP_SIZE = 128
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+# Autotune configs. We keep this small so compilation stays fast.
+_GEMM_CONFIGS = [
+    triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_SZ": 128},
+                  num_warps=nw, num_stages=ns)
+    for bm in (16, 32, 64, 128)
+    for bn in (64, 128, 256)
+    for bk in (128, 256)
+    for nw in (4, 8)
+    for ns in (2, 3)
+    if bm * bn <= 32768
+]
+
+
+@triton.autotune(configs=_GEMM_CONFIGS, key=["M", "N", "K"])
+@triton.jit
+def w4a16_gemm_kernel(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_mask = offs_m < M
+
+    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+    N_GROUPS: tl.constexpr = K // GROUP_SZ
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    # Loop over K in chunks of BLOCK_K.  We do two tl.dot calls per chunk: one
+    # for the even-K rows of the packed weights, one for the odd-K rows.
+    for k_blk in tl.range(0, K, BLOCK_K):
+        # Load packed weights: (BLOCK_K // 2, BLOCK_N) uint8
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N,
+            other=0,
+        )
+
+        # Unpack
+        w_lo = (wq & 0xF).to(tl.bfloat16)  # (BLOCK_K_HALF, BLOCK_N)  k even
+        w_hi = (wq >> 4).to(tl.bfloat16)  # (BLOCK_K_HALF, BLOCK_N)  k odd
+
+        # Load scales/zeros for the groups covered by this K block.
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)  # (n_groups_per_blk, BLOCK_N)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        # Broadcast to (BLOCK_K_HALF, BLOCK_N) by repeating each group GROUP_HALF times.
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        # Dequant: (w_int - zero) * scale
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        # Load x for the even/odd K rows in this block.
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_even[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+        x_odd = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_odd[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+
+        # Two dots per block.
+        acc += tl.dot(x_even, w_lo)
+        acc += tl.dot(x_odd, w_hi)
+
+    # Store
+    out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16),
+             mask=offs_m_mask[:, None] & (offs_n[None, :] < N))
+
+
+_GEMV_CONFIGS = [
+    triton.Config({"BLOCK_N": bn, "BLOCK_K": bk, "GROUP_SZ": 128},
+                  num_warps=nw, num_stages=ns)
+    for bn in (64, 128, 256)
+    for bk in (128, 256, 512)
+    for nw in (4, 8)
+    for ns in (2, 3, 4)
+]
+
+
+@triton.autotune(configs=_GEMV_CONFIGS, key=["M", "N", "K"])
+@triton.jit
+def w4a16_gemv_kernel(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    """M=1 specialized: each program computes a (1, BLOCK_N) output tile."""
+    pid = tl.program_id(0)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    # Load x once (full K) -- for M=1, this is small (K * 2 bytes).
+    offs_k = tl.arange(0, BLOCK_K)
+
+    for k_b

... (4572 more chars)

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/debug_test.py
@@ -0,0 +1,29 @@
+"""Quick debug to verify kernel correctness."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import reference
+import solution
+
+# Setup small test
+M, N, K = 1, 256, 128
+reference.M = M
+reference.N = N
+reference.K = K
+
+ref_model = reference.Model(M, N, K).cuda().eval()
+sol_model = solution.Model(M, N, K).cuda().eval()
+sol_model.load_state_dict(ref_model.state_dict())
+
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+
+with torch.no_grad():
+    ref_out = ref_model(x)
+    sol_out = sol_model(x)
+
+print("Ref:", ref_out[:8])
+print("Sol:", sol_out[:8])
+print("Diff:", (ref_out - sol_out).abs().max().item())
+print("Ref abs max:", ref_out.abs().max().item())
+print("Match:", torch.allclose(ref_out, sol_out, atol=1e-2, rtol=1e-2))

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/debug_test2.py
@@ -0,0 +1,35 @@
+"""Quick debug to verify kernel correctness on actual sizes."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import reference
+import solution
+
+# Test the actual problem size
+M, N, K = 1, 12288, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+
+ref_model = reference.Model(M, N, K).cuda().eval()
+sol_model = solution.Model(M, N, K).cuda().eval()
+sol_model.load_state_dict(ref_model.state_dict())
+
+# Test multiple seeds
+for seed in [42, 123, 456]:
+    torch.manual_seed(seed)
+    x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+    with torch.no_grad():
+        ref_out = ref_model(x)
+        sol_out = sol_model(x)
+    diff = (ref_out - sol_out).abs().max().item()
+    rel_diff = ((ref_out - sol_out).abs() / ref_out.abs().clamp_min(1e-6)).max().item()
+    print(f"seed={seed}: max_abs_diff={diff:.6f}, max_rel_diff={rel_diff:.4f}, ref_max={ref_out.abs().max().item():.4f}")
+    n_bad = ((ref_out - sol_out).abs() > 0.1).sum().item()
+    print(f"  bad={n_bad}/{ref_out.numel()}")
+    if n_bad > 0:
+        bad_idx = ((ref_out - sol_out).abs() > 0.1).nonzero()[:5]
+        for idx in bad_idx:
+            i, j = idx[0].item(), idx[1].item()
+            print(f"  [{i},{j}]: ref={ref_out[i,j].item():.4f}, sol={sol_out[i,j].item():.4f}")

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/debug_test4.py
@@ -0,0 +1,34 @@
+"""Test M=16 case."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import reference
+import solution
+
+# Test the actual problem size
+M, N, K = 16, 14336, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+
+ref_model = reference.Model(M, N, K).cuda().eval()
+sol_model = solution.Model(M, N, K).cuda().eval()
+sol_model.load_state_dict(ref_model.state_dict())
+
+torch.manual_seed(42)
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+x_scaled = (x.float() * 64.0).to(torch.bfloat16).cuda()
+
+with torch.no_grad():
+    ref_out = ref_model(x_scaled)
+    sol_out = sol_model(x_scaled)
+
+diff = (ref_out - sol_out).abs()
+print(f"max_abs_diff: {diff.max().item():.6f}")
+print(f"ref_max: {ref_out.abs().max().item():.4f}")
+
+allowed = 1.0 + 0.05 * ref_out.abs()
+bad = diff > allowed
+n_bad = bad.sum().item()
+print(f"bad: {n_bad}/{ref_out.numel()}")

--- a//workspace/problems/07_w4a16_gemm/debug_simple.py
+++ b//workspace/problems/07_w4a16_gemm/debug_simple.py
@@ -101,4 +101,9 @@
 
 print("ref:", ref_out[0, :8])
 print("sol:", out[0, :8])
-print("match:", torch.allclose(ref_out, out, atol=1e-2))
+print("diff max:", (ref_out - out).abs().max().item())
+print("match 1e-2:", torch.allclose(ref_out, out, atol=1e-2))
+print("match 1e-1:", torch.allclose(ref_out, out, atol=1e-1))
+print("match 1.0:", torch.allclose(ref_out, out, atol=1.0))
+print("ref numel:", ref_out.numel(), "out numel:", out.numel())
+print("ref shape:", ref_out.shape, "out shape:", out.shape)

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/debug_check.py
@@ -0,0 +1,114 @@
+"""Test M=1 case with a specific known config."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import reference
+
+# Setup test
+M, N, K = 1, 12288, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+
+ref_model = reference.Model(M, N, K).cuda().eval()
+
+# Build a custom solution that just uses our simple kernel
+import triton
+import triton.language as tl
+
+@triton.jit
+def w4a16_gemv_kernel_simple(
+    X, WQ, S, Z, OUT,
+    N, K,
+    stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_on,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    """Simplified gemv for M=1."""
+    pid = tl.program_id(0)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, BLOCK_K // GROUP_SZ)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(X + x_offs_even * stride_xk)
+        x_odd = tl.load(X + x_offs_odd * stride_xk)
+
+        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
+        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
+
+    out_ptrs = OUT + offs_n * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
+
+
+# Reference
+torch.manual_seed(42)
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+x_scaled = (x.float() * 64.0).to(torch.bfloat16).cuda()
+ref_out = ref_model(x_scaled)
+ref_model_state = {k: v.clone() for k, v in ref_model.state_dict().items()}
+
+# Try different configs
+for BLOCK_N, BLOCK_K, num_warps, num_stages in [
+    (64, 128, 4, 2),
+    (128, 128, 4, 2),
+    (256, 128, 8, 2),
+    (64, 256, 4, 2),
+    (128, 256, 4, 2),
+    (128, 512, 4, 2),
+    (64, 512, 4, 2),
+    (256, 256, 8, 2),
+    (32, 128, 4, 2),
+]:
+    out = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+    grid = (triton.cdiv(N, BLOCK_N),)
+    w4a16_gemv_kernel_simple[grid](
+        x_scaled, ref_model.w_q, ref_model.scales, ref_model.zeros, out,
+        N, K,
+        x_scaled.stride(1),
+        ref_model.w_q.stride(0), ref_model.w_q.stride(1),
+        ref_model.scales.stride(0), ref_model.scales.stride(1),
+        ref_model.zeros.stride(0), ref_model.zeros.stride(1),
+        out.stride(1),
+        BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+        num_warps=num_warps, num_stages=num_stages,
+    )
+    diff = (ref_out - out).abs().max().item()
+    allowed = 1.0 + 0.05 * ref_out.abs()
+    bad = ((ref_out - out).abs() > allowed).sum().item()
+    print(f"BLOCK_N={BLOCK_N:3d} BLOCK_K={BLOCK_K:3d} warps={num_warps} stages={num_stages}: max_diff={diff:.4f}, bad={bad}")

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/debug_fp32.py
@@ -0,0 +1,185 @@
+"""Test with fp32 dequant to see if precision is the issue."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import reference
+
+# Setup test
+M, N, K = 1, 12288, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+
+ref_model = reference.Model(M, N, K).cuda().eval()
+
+# Build a custom solution that does dequant in fp32
+import triton
+import triton.language as tl
+
+@triton.jit
+def w4a16_gemv_kernel_fp32_dequant(
+    X, WQ, S, Z, OUT,
+    N, K,
+    stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_on,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, BLOCK_K // GROUP_SZ)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        # Dequant in fp32
+        w_lo_f32 = (w_lo.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        w_hi_f32 = (w_hi.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        # Cast back to bf16
+        w_lo = w_lo_f32.to(tl.bfloat16)
+        w_hi = w_hi_f32.to(tl.bfloat16)
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(X + x_offs_even * stride_xk)
+        x_odd = tl.load(X + x_offs_odd * stride_xk)
+
+        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
+        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
+
+    out_ptrs = OUT + offs_n * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
+
+
+# Reference
+torch.manual_seed(42)
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+x_scaled = (x.float() * 64.0).to(torch.bfloat16).cuda()
+ref_out = ref_model(x_scaled)
+
+# Test
+out = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+BLOCK_N = 64
+BLOCK_K = 128
+grid = (triton.cdiv(N, BLOCK_N),)
+w4a16_gemv_kernel_fp32_dequant[grid](
+    x_scaled, ref_model.w_q, ref_model.scales, ref_model.zeros, out,
+    N, K,
+    x_scaled.stride(1),
+    ref_model.w_q.stride(0), ref_model.w_q.stride(1),
+    ref_model.scales.stride(0), ref_model.scales.stride(1),
+    ref_model.zeros.stride(0), ref_model.zeros.stride(1),
+    out.stride(1),
+    BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+    num_warps=4, num_stages=2,
+)
+
+diff = (ref_out - out).abs().max().item()
+allowed = 1.0 + 0.05 * ref_out.abs()
+bad = ((ref_out - out).abs() > allowed).sum().item()
+print(f"fp32 dequant: max_diff={diff:.4f}, bad={bad}")
+
+# Compare to bf16 dequant
+@triton.jit
+def w4a16_gemv_kernel_bf16_dequant(
+    X, WQ, S, Z, OUT,
+    N, K,
+    stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_on,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, BLOCK_K // GROUP_SZ)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        # Dequant in bf16
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(X + x_offs_even * stride_xk)
+        x_odd = tl.load(X + x_offs_odd * stride_xk)
+
+        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
+        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
+
+    out_ptrs = OUT + offs_n * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
+
+
+out2 = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+w4a16_gemv_kernel_bf16_dequant[grid](
+    x_scaled, ref_model.w_q, ref_model.scales, ref_model.zeros, out2,
+    N

... (507 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -15,6 +15,10 @@
     w_q, scales, and zeros on the fly and accumulates in fp32.
   - Prefill path (M>=16) uses a tile-based GEMM kernel that dequantizes
     BLOCK_K rows at a time and uses tensor cores.
+
+Numerical accuracy: we dequant in fp32 instead of bf16, which avoids the
+accumulation of rounding errors that the strict large-activation tolerance
+(0.10 abs/rel) catches.
 """
 from __future__ import annotations
 
@@ -104,9 +108,9 @@
         z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
 
-        # Dequant: (w_int - zero) * scale
-        w_lo = (w_lo - z) * s
-        w_hi = (w_hi - z) * s
+        # Dequant in fp32 to avoid bf16 rounding error accumulation.
+        w_lo = (w_lo.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        w_hi = (w_hi.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
 
         # Load x for the even/odd K rows in this block.
         x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
@@ -120,9 +124,12 @@
             mask=offs_m_mask[:, None], other=0.0,
         )
 
-        # Two dots per block.
-        acc += tl.dot(x_even, w_lo)
-        acc += tl.dot(x_odd, w_hi)
+        x_even_f = x_even.to(tl.float32)
+        x_odd_f = x_odd.to(tl.float32)
+
+        # Two dots per block. Use input_precision="ieee" for highest accuracy.
+        acc += tl.dot(x_even_f, w_lo, input_precision="ieee")
+        acc += tl.dot(x_odd_f, w_hi, input_precision="ieee")
 
     # Store
     out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
@@ -163,9 +170,6 @@
 
     acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
 
-    # Load x once (full K) -- for M=1, this is small (K * 2 bytes).
-    offs_k = tl.arange(0, BLOCK_K)
-
     for k_blk in tl.range(0, K, BLOCK_K):
         # Packed weights: (BLOCK_K // 2, BLOCK_N)
         wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
@@ -189,8 +193,9 @@
         z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
 
-        w_lo = (w_lo - z) * s
-        w_hi = (w_hi - z) * s
+        # Dequant in fp32
+        w_lo_f = (w_lo.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        w_hi_f = (w_hi.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
 
         # x_even: (BLOCK_K_HALF,), x_odd: (BLOCK_K_HALF,)
         x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
@@ -198,9 +203,12 @@
         x_even = tl.load(X + x_offs_even * stride_xk)
         x_odd = tl.load(X + x_offs_odd * stride_xk)
 
+        x_even_f = x_even.to(tl.float32)
+        x_odd_f = x_odd.to(tl.float32)
+
         # acc[n] += sum_k x_even[k] * w_lo[k, n] + x_odd[k] * w_hi[k, n]
-        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
-        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
+        acc += tl.sum(x_even_f[:, None] * w_lo_f, axis=0)
+        acc += tl.sum(x_odd_f[:, None] * w_hi_f, axis=0)
 
     out_ptrs = OUT + offs_n * stride_on
     tl.store(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
@@ -209,13 +217,6 @@
 # ---------------------------------------------------------------------------
 # Module
 # ---------------------------------------------------------------------------
-
-
-def _quant_args(M: int, N: int, K: int, group_size: int = GROUP_SIZE):
-    assert K % group_size == 0
-    assert K % 2 == 0
-    n_groups = K // group_size
-    return M, N, K, n_groups
 
 
 class Model(nn.Module):

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -108,9 +108,9 @@
         z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
 
-        # Dequant in fp32 to avoid bf16 rounding error accumulation.
-        w_lo = (w_lo.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
-        w_hi = (w_hi.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        # Dequant in bf16 (matches the reference's bf16 dequant)
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
 
         # Load x for the even/odd K rows in this block.
         x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
@@ -124,12 +124,10 @@
             mask=offs_m_mask[:, None], other=0.0,
         )
 
-        x_even_f = x_even.to(tl.float32)
-        x_odd_f = x_odd.to(tl.float32)
-
-        # Two dots per block. Use input_precision="ieee" for highest accuracy.
-        acc += tl.dot(x_even_f, w_lo, input_precision="ieee")
-        acc += tl.dot(x_odd_f, w_hi, input_precision="ieee")
+        # Two dots per block. Default precision (bf16 input) matches the
+        # reference's cuBLAS bf16 GEMM.
+        acc += tl.dot(x_even, w_lo)
+        acc += tl.dot(x_odd, w_hi)
 
     # Store
     out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -191,9 +191,9 @@
         z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
 
-        # Dequant in fp32
-        w_lo_f = (w_lo.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
-        w_hi_f = (w_hi.to(tl.float32) - z.to(tl.float32)) * s.to(tl.float32)
+        # Dequant in bf16 (matches the reference's bf16 dequant)
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
 
         # x_even: (BLOCK_K_HALF,), x_odd: (BLOCK_K_HALF,)
         x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
@@ -201,12 +201,10 @@
         x_even = tl.load(X + x_offs_even * stride_xk)
         x_odd = tl.load(X + x_offs_odd * stride_xk)
 
-        x_even_f = x_even.to(tl.float32)
-        x_odd_f = x_odd.to(tl.float32)
-
         # acc[n] += sum_k x_even[k] * w_lo[k, n] + x_odd[k] * w_hi[k, n]
-        acc += tl.sum(x_even_f[:, None] * w_lo_f, axis=0)
-        acc += tl.sum(x_odd_f[:, None] * w_hi_f, axis=0)
+        # Use bf16 products to match the reference's cuBLAS bf16 GEMM behavior.
+        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
+        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
 
     out_ptrs = OUT + offs_n * stride_on
     tl.store(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)

--- a//workspace/problems/07_w4a16_gemm/debug_best_config.py
+++ b//workspace/problems/07_w4a16_gemm/debug_best_config.py
@@ -148,10 +148,6 @@
         (128, 128, 256, 8, 2),
     ]
     for BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages in configs:
-        if BLOCK_M > M and M > 0:
-            # Skip if BLOCK_M > M (and M > 0) - this wastes too much
-            # But for M=1, BLOCK_M=16 wastes 15 rows
-            continue
         try:
             t = time_kernel(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages,
                             x, ref_model.w_q, ref_model.scales, ref_model.zeros, out)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -12,8 +12,8 @@
     group_size=128 K rows become two 64-wide dots that both read the same x
     row span but different nibbles of the packed weight byte.  This matches
     the bit-level rounding of the reference's cuBLAS bf16 GEMM.
-  - Modulo-based dequant indexing (no broadcast/reshape of scales/zeros).
-  - Autotune over a focused set of configs.
+  - Per-M config selection: small M wants small BLOCK_N for memory-bound
+    parallelism; large M wants larger BLOCK_N for compute reuse.
 """
 from __future__ import annotations
 
@@ -33,19 +33,23 @@
 # Triton kernel
 # ---------------------------------------------------------------------------
 
-_GEMM_CONFIGS = [
-    triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk, "GROUP_SZ": 128},
-                  num_warps=nw, num_stages=ns)
-    for bm in (16, 32, 64, 128)
-    for bn in (64, 128, 256)
-    for bk in (128, 256)
-    for nw in (4, 8)
-    for ns in (2, 3)
-    if bm * bn <= 32768
+# Configs to autotune over.  We keep the list focused.
+_CONFIGS = [
+    # M=1 / decode
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    # M=16..32 / small prefill
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    # M=256 / large prefill (compute-bound)
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
 ]
 
 
-@triton.autotune(configs=_GEMM_CONFIGS, key=["M", "N", "K"])
+@triton.autotune(configs=_CONFIGS, key=["M", "N", "K"])
 @triton.jit
 def w4a16_gemm_kernel(
     X, WQ, S, Z, OUT,
@@ -74,36 +78,29 @@
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
     for k_blk in tl.range(0, K, BLOCK_K):
-        # Load packed weights: (BLOCK_K // 2, BLOCK_N) uint8
         wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
         wq = tl.load(
             WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
-            mask=offs_n[None, :] < N,
-            other=0,
+            mask=offs_n[None, :] < N, other=0,
         )
 
-        # Unpack
-        w_lo = (wq & 0xF).to(tl.bfloat16)  # (BLOCK_K_HALF, BLOCK_N)  k even
-        w_hi = (wq >> 4).to(tl.bfloat16)  # (BLOCK_K_HALF, BLOCK_N)  k odd
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
 
-        # Load scales/zeros for the groups covered by this K block.
         g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
         s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
-                    mask=offs_n[None, :] < N, other=0.0)  # (n_groups_per_blk, BLOCK_N)
+                    mask=offs_n[None, :] < N, other=0.0)
         z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
                     mask=offs_n[None, :] < N, other=0.0)
 
-        # Broadcast to (BLOCK_K_HALF, BLOCK_N) by repeating each group GROUP_HALF times.
         s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
         z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
         z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
 
-        # Dequant in bf16 (matches the reference's bf16 dequant)
         w_lo = (w_lo - z) * s
         w_hi = (w_hi - z) * s
 
-        # Load x for the even/odd K rows in this block.
         x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
         x_offs_odd = x_offs_even + 1
         x_even = tl.load(
@@ -115,8 +112,6 @@
             mask=offs_m_mask[:, None], other=0.0,
         )
 
-        # Two dots per block. Default precision (bf16 input) matches the
-        # reference's cuBLAS bf16 GEMM.
         acc += tl.dot(x_even, w_lo)
         acc += tl.dot(x_odd, w_hi)
 

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -12,8 +12,9 @@
     group_size=128 K rows become two 64-wide dots that both read the same x
     row span but different nibbles of the packed weight byte.  This matches
     the bit-level rounding of the reference's cuBLAS bf16 GEMM.
-  - Per-M config selection: small M wants small BLOCK_N for memory-bound
-    parallelism; large M wants larger BLOCK_N for compute reuse.
+  - Per-shape hardcoded block sizes (we found autotune's single-run timing
+    is too noisy to reliably pick the right config; we benchmarked each shape
+    and picked the best config manually).
 """
 from __future__ import annotations
 
@@ -33,23 +34,30 @@
 # Triton kernel
 # ---------------------------------------------------------------------------
 
-# Configs to autotune over.  We keep the list focused.
-_CONFIGS = [
-    # M=1 / decode
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
-    # M=16..32 / small prefill
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    # M=256 / large prefill (compute-bound)
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-]
+# Block sizes picked per shape from local benchmarks on this GPU.
+# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages)
+_PICKED_CONFIGS = {
+    (1, 12288, 4096):  (16, 64, 128, 4, 2),
+    (32, 12288, 4096): (16, 128, 128, 4, 2),
+    (256, 12288, 4096):(16, 128, 128, 4, 2),
+    (1, 4096, 4096):   (16, 64, 256, 4, 2),
+    (16, 14336, 4096): (16, 128, 256, 4, 3),
+}
 
 
-@triton.autotune(configs=_CONFIGS, key=["M", "N", "K"])
+def _pick_config(M: int, N: int, K: int) -> tuple[int, int, int, int, int]:
+    """Pick (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for the shape."""
+    key = (M, N, K)
+    if key in _PICKED_CONFIGS:
+        return _PICKED_CONFIGS[key]
+    # Fallback heuristic
+    if M == 1:
+        return (16, 64, 128, 4, 2)
+    if M <= 32:
+        return (16, 128, 128, 4, 2)
+    return (16, 128, 128, 4, 2)
+
+
 @triton.jit
 def w4a16_gemm_kernel(
     X, WQ, S, Z, OUT,
@@ -149,10 +157,8 @@
         x = x.contiguous()
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        grid = lambda meta: (
-            triton.cdiv(M, meta["BLOCK_M"]),
-            triton.cdiv(N, meta["BLOCK_N"]),
-        )
+        BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _pick_config(M, N, K)
+        grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
         w4a16_gemm_kernel[grid](
             x, self.w_q, self.scales, self.zeros, out,
             M, N, K,
@@ -161,6 +167,8 @@
             self.scales.stride(0), self.scales.stride(1),
             self.zeros.stride(0), self.zeros.stride(1),
             out.stride(0), out.stride(1),
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+            num_warps=num_warps, num_stages=num_stages,
         )
         return out
 

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -10,11 +10,11 @@
 Key optimizations:
   - Two-dot pattern: process even/odd K rows with two separate dots. Each
     group_size=128 K rows become two 64-wide dots that both read the same x
-    row span but different nibbles of the packed weight byte.  This matches
-    the bit-level rounding of the reference's cuBLAS bf16 GEMM.
-  - Per-shape hardcoded block sizes (we found autotune's single-run timing
-    is too noisy to reliably pick the right config; we benchmarked each shape
-    and picked the best config manually).
+    row span but different nibbles of the packed weight byte.
+  - Per-shape hardcoded block sizes, picked from local benchmarks.
+  - M=1 kernel uses a smaller BLOCK_M=1 path to avoid wasted M tiles while
+    still feeding the tensor core's M=16 minimum (Triton handles the padding).
+  - For very small N (shape 3), we split-K to raise SM occupancy.
 """
 from __future__ import annotations
 
@@ -31,32 +31,8 @@
 
 
 # ---------------------------------------------------------------------------
-# Triton kernel
-# ---------------------------------------------------------------------------
-
-# Block sizes picked per shape from local benchmarks on this GPU.
-# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages)
-_PICKED_CONFIGS = {
-    (1, 12288, 4096):  (16, 64, 128, 4, 2),
-    (32, 12288, 4096): (16, 128, 128, 4, 2),
-    (256, 12288, 4096):(16, 128, 128, 4, 2),
-    (1, 4096, 4096):   (16, 64, 256, 4, 2),
-    (16, 14336, 4096): (16, 128, 256, 4, 3),
-}
-
-
-def _pick_config(M: int, N: int, K: int) -> tuple[int, int, int, int, int]:
-    """Pick (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) for the shape."""
-    key = (M, N, K)
-    if key in _PICKED_CONFIGS:
-        return _PICKED_CONFIGS[key]
-    # Fallback heuristic
-    if M == 1:
-        return (16, 64, 128, 4, 2)
-    if M <= 32:
-        return (16, 128, 128, 4, 2)
-    return (16, 128, 128, 4, 2)
-
+# Triton kernel — main path
+# ---------------------------------------------------------------------------
 
 @triton.jit
 def w4a16_gemm_kernel(
@@ -129,8 +105,90 @@
 
 
 # ---------------------------------------------------------------------------
+# Split-K kernel for M=1, small N
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def w4a16_splitk_kernel(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SPLIT_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    """Split-K version for very small N (e.g., N=4096, M=1)."""
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)  # split-K index
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+
+    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+
+    # Each split processes K / SPLIT_K elements of K
+    K_per_split: tl.constexpr = K // SPLIT_K
+    k_start = pid_k * K_per_split
+    k_end = k_start + K_per_split
+
+    for k_blk in tl.range(k_start, k_end, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(X + x_offs_even * stride_xk)
+        x_odd = tl.load(X + x_offs_odd * stride_xk)
+
+        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
+        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
+
+    # Atomic add (since multiple K splits write to the same output)
+    out_ptrs = OUT + offs_n * stride_on
+    tl.atomic_add(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
+
+
+# ---------------------------------------------------------------------------
 # Module
 # ---------------------------------------------------------------------------
+
+
+# Picked configs from local benchmarks.
+_PICKED_CONFIGS = {
+    (1, 12288, 4096):  (16, 64, 128, 4, 2),
+    (32, 12288, 4096): (16, 128, 128, 4, 2),
+    (256, 12288, 4096):(16, 128, 128, 4, 2),
+    (1, 4096, 4096):   (16, 64, 128, 4, 2),  # use split-K path
+    (16, 14336, 4096): (16, 128, 256, 4, 3),
+}
 
 
 class Model(nn.Module):
@@ -155,21 +213,42 @@
         assert x.dtype == torch.bfloat16
 
         x = x.contiguous()
-        out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
-
-        BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _pick_config(M, N, K)
-        grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
-        w4a16_gemm_kernel[grid](
-            x, self.w_q, self.scales, self.zeros, out,
-            M, N, K,
-            x.stride(0), x.stride(1),
-            self.w_q.stride(0), self.w_q.stride(1),
-            self.scales.stride(0), self.scales.stride(1),
-            self.zeros.stride(0), self.zeros.s

... (1938 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -10,11 +10,9 @@
 Key optimizations:
   - Two-dot pattern: process even/odd K rows with two separate dots. Each
     group_size=128 K rows become two 64-wide dots that both read the same x
-    row span but different nibbles of the packed weight byte.
+    row span but different nibbles of the packed weight byte.  This matches
+    the bit-level rounding of the reference's cuBLAS bf16 GEMM.
   - Per-shape hardcoded block sizes, picked from local benchmarks.
-  - M=1 kernel uses a smaller BLOCK_M=1 path to avoid wasted M tiles while
-    still feeding the tensor core's M=16 minimum (Triton handles the padding).
-  - For very small N (shape 3), we split-K to raise SM occupancy.
 """
 from __future__ import annotations
 
@@ -31,7 +29,7 @@
 
 
 # ---------------------------------------------------------------------------
-# Triton kernel — main path
+# Triton kernel
 # ---------------------------------------------------------------------------
 
 @triton.jit
@@ -105,88 +103,16 @@
 
 
 # ---------------------------------------------------------------------------
-# Split-K kernel for M=1, small N
-# ---------------------------------------------------------------------------
-
-@triton.jit
-def w4a16_splitk_kernel(
-    X, WQ, S, Z, OUT,
-    M, N, K,
-    stride_xm, stride_xk,
-    stride_wk, stride_wn,
-    stride_sg, stride_sn,
-    stride_zg, stride_zn,
-    stride_om, stride_on,
-    BLOCK_N: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-    SPLIT_K: tl.constexpr,
-    GROUP_SZ: tl.constexpr,
-):
-    """Split-K version for very small N (e.g., N=4096, M=1)."""
-    pid_n = tl.program_id(0)
-    pid_k = tl.program_id(1)  # split-K index
-
-    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-
-    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
-    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
-    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
-
-    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
-
-    # Each split processes K / SPLIT_K elements of K
-    K_per_split: tl.constexpr = K // SPLIT_K
-    k_start = pid_k * K_per_split
-    k_end = k_start + K_per_split
-
-    for k_blk in tl.range(k_start, k_end, BLOCK_K):
-        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
-        wq = tl.load(
-            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
-            mask=offs_n[None, :] < N, other=0,
-        )
-
-        w_lo = (wq & 0xF).to(tl.bfloat16)
-        w_hi = (wq >> 4).to(tl.bfloat16)
-
-        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
-        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
-                    mask=offs_n[None, :] < N, other=0.0)
-        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
-                    mask=offs_n[None, :] < N, other=0.0)
-
-        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
-        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
-        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
-        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
-
-        w_lo = (w_lo - z) * s
-        w_hi = (w_hi - z) * s
-
-        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
-        x_offs_odd = x_offs_even + 1
-        x_even = tl.load(X + x_offs_even * stride_xk)
-        x_odd = tl.load(X + x_offs_odd * stride_xk)
-
-        acc += tl.sum(x_even[:, None] * w_lo, axis=0)
-        acc += tl.sum(x_odd[:, None] * w_hi, axis=0)
-
-    # Atomic add (since multiple K splits write to the same output)
-    out_ptrs = OUT + offs_n * stride_on
-    tl.atomic_add(out_ptrs, acc.to(tl.bfloat16), mask=offs_n < N)
-
-
-# ---------------------------------------------------------------------------
 # Module
 # ---------------------------------------------------------------------------
 
 
-# Picked configs from local benchmarks.
+# Picked configs from local benchmarks (per-shape).
 _PICKED_CONFIGS = {
     (1, 12288, 4096):  (16, 64, 128, 4, 2),
     (32, 12288, 4096): (16, 128, 128, 4, 2),
     (256, 12288, 4096):(16, 128, 128, 4, 2),
-    (1, 4096, 4096):   (16, 64, 128, 4, 2),  # use split-K path
+    (1, 4096, 4096):   (16, 64, 128, 4, 2),
     (16, 14336, 4096): (16, 128, 256, 4, 3),
 }
 
@@ -213,42 +139,23 @@
         assert x.dtype == torch.bfloat16
 
         x = x.contiguous()
+        out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # For small N with M=1, use split-K to raise occupancy
-        if M == 1 and N <= 8192:
-            out = torch.zeros((M, N), dtype=torch.bfloat16, device=x.device)
-            BLOCK_N = 64
-            BLOCK_K = 128
-            SPLIT_K = 4
-            grid = (triton.cdiv(N, BLOCK_N), SPLIT_K)
-            w4a16_splitk_kernel[grid](
-                x, self.w_q, self.scales, self.zeros, out,
-                M, N, K,
-                x.stride(0), x.stride(1),
-                self.w_q.stride(0), self.w_q.stride(1),
-                self.scales.stride(0), self.scales.stride(1),
-                self.zeros.stride(0), self.zeros.stride(1),
-                out.stride(0), out.stride(1),
-                BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, SPLIT_K=SPLIT_K, GROUP_SZ=128,
-                num_warps=4, num_stages=2,
-            )
-        else:
-            out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
-            BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _PICKED_CONFIGS.get(
-                (M, N, K), (16, 128, 128, 4, 2)
-            )
-            grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
-            w4a16_gemm_kernel[grid](
-                x, self.w_q, self.scales, self.zeros, out,
-                M, N, K,
-                x.stride(0), x.stride(1),
-                self.w_q.stride(0), self.w_q.stride(1),
-                self.scales.stride(0), self.scales.stride(1),
-                self.zeros.stride(0),

... (965 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -12,7 +12,7 @@
     group_size=128 K rows become two 64-wide dots that both read the same x
     row span but different nibbles of the packed weight byte.  This matches
     the bit-level rounding of the reference's cuBLAS bf16 GEMM.
-  - Per-shape hardcoded block sizes, picked from local benchmarks.
+  - Triton autotune over a focused set of configs to find the best per shape.
 """
 from __future__ import annotations
 
@@ -32,6 +32,29 @@
 # Triton kernel
 # ---------------------------------------------------------------------------
 
+_CONFIGS = [
+    # M=1 / decode variants — small BN, more programs
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 32, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 32, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    # M=16-32 / small prefill
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    # M=256 / large prefill
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=8, num_stages=2),
+]
+
+
+@triton.autotune(configs=_CONFIGS, key=["M", "N", "K"])
 @triton.jit
 def w4a16_gemm_kernel(
     X, WQ, S, Z, OUT,
@@ -107,16 +130,6 @@
 # ---------------------------------------------------------------------------
 
 
-# Picked configs from local benchmarks (per-shape).
-_PICKED_CONFIGS = {
-    (1, 12288, 4096):  (16, 64, 128, 4, 2),
-    (32, 12288, 4096): (16, 128, 128, 4, 2),
-    (256, 12288, 4096):(16, 128, 128, 4, 2),
-    (1, 4096, 4096):   (16, 64, 128, 4, 2),
-    (16, 14336, 4096): (16, 128, 256, 4, 3),
-}
-
-
 class Model(nn.Module):
     def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
         super().__init__()
@@ -141,10 +154,10 @@
         x = x.contiguous()
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _PICKED_CONFIGS.get(
-            (M, N, K), (16, 128, 128, 4, 2)
+        grid = lambda meta: (
+            triton.cdiv(M, meta["BLOCK_M"]),
+            triton.cdiv(N, meta["BLOCK_N"]),
         )
-        grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
         w4a16_gemm_kernel[grid](
             x, self.w_q, self.scales, self.zeros, out,
             M, N, K,
@@ -153,8 +166,6 @@
             self.scales.stride(0), self.scales.stride(1),
             self.zeros.stride(0), self.zeros.stride(1),
             out.stride(0), out.stride(1),
-            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
-            num_warps=num_warps, num_stages=num_stages,
         )
         return out
 

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/test_bm1.py
@@ -0,0 +1,113 @@
+"""Test BLOCK_M=1."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def kernel_test(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_mask = offs_m < M
+
+    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_even[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+        x_odd = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_odd[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+
+        acc += tl.dot(x_even, w_lo)
+        acc += tl.dot(x_odd, w_hi)
+
+    out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16),
+             mask=offs_m_mask[:, None] & (offs_n[None, :] < N))
+
+
+import reference
+M, N, K = 1, 4096, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+ref_model = reference.Model(M, N, K).cuda().eval()
+
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+out = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+
+# Try BLOCK_M=1
+for BLOCK_M, BLOCK_N, BLOCK_K, nw, ns in [
+    (1, 64, 128, 4, 2),
+    (1, 128, 128, 4, 2),
+    (1, 64, 256, 4, 2),
+    (1, 32, 128, 4, 2),
+]:
+    try:
+        grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+        kernel_test[grid](
+            x, ref_model.w_q, ref_model.scales, ref_model.zeros, out,
+            M, N, K,
+            x.stride(0), x.stride(1),
+            ref_model.w_q.stride(0), ref_model.w_q.stride(1),
+            ref_model.scales.stride(0), ref_model.scales.stride(1),
+            ref_model.zeros.stride(0), ref_model.zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+            num_warps=nw, num_stages=ns,
+        )
+        ref_out = ref_model(x)
+        diff = (ref_out - out).abs().max().item()
+        print(f'BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}, BLOCK_K={BLOCK_K}, W={nw}, S={ns}: max diff = {diff:.4f}')
+    except Exception as e:
+        print(f'BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}, BLOCK_K={BLOCK_K}, W={nw}, S={ns}: ERROR {e}')

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/test_bm1_time.py
@@ -0,0 +1,161 @@
+"""Time BLOCK_M=1 for M=1 cases."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import triton
+import triton.language as tl
+import statistics
+import reference
+
+@triton.jit
+def kernel_test(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_mask = offs_m < M
+
+    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_even[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+        x_odd = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_odd[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+
+        acc += tl.dot(x_even, w_lo)
+        acc += tl.dot(x_odd, w_hi)
+
+    out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16),
+             mask=offs_m_mask[:, None] & (offs_n[None, :] < N))
+
+
+def time_kernel(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages, x, w_q, scales, zeros, out, n_iters=30):
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+    # Warmup
+    for _ in range(10):
+        kernel_test[grid](
+            x, w_q, scales, zeros, out,
+            M, N, K,
+            x.stride(0), x.stride(1),
+            w_q.stride(0), w_q.stride(1),
+            scales.stride(0), scales.stride(1),
+            zeros.stride(0), zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+            num_warps=num_warps, num_stages=num_stages,
+        )
+    torch.cuda.synchronize()
+
+    # L2 flush
+    l2 = torch.empty(128 * 1024 * 1024 // 4, dtype=torch.float32, device='cuda')
+
+    times = []
+    for _ in range(n_iters):
+        l2.zero_()
+        torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        kernel_test[grid](
+            x, w_q, scales, zeros, out,
+            M, N, K,
+            x.stride(0), x.stride(1),
+            w_q.stride(0), w_q.stride(1),
+            scales.stride(0), scales.stride(1),
+            zeros.stride(0), zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+            num_warps=num_warps, num_stages=num_stages,
+        )
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    return statistics.median(times)
+
+
+# Test each shape
+shapes = [
+    (1, 12288, 4096),
+    (1, 4096, 4096),
+]
+
+for M, N, K in shapes:
+    reference.M = M
+    reference.N = N
+    reference.K = K
+    ref_model = reference.Model(M, N, K).cuda().eval()
+    x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+    out = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+
+    print(f"\n=== M={M}, N={N}, K={K} ===")
+    for BLOCK_M, BLOCK_N, BLOCK_K, nw, ns in [
+        (1, 32, 128, 4, 2),
+        (1, 32, 256, 4, 2),
+        (1, 64, 128, 4, 2),
+        (1, 64, 256, 4, 2),
+        (1, 128, 128, 4, 2),
+        (1, 128, 256, 4, 2),
+        (1, 64, 128, 4, 3),
+        (1, 64, 128, 4, 4),
+        (1, 64, 256, 4, 3),
+        (1, 32, 128, 4, 3),
+        (1, 32, 256, 4, 3),
+        (1, 64, 512, 4, 2),
+    ]:
+        try:
+            t = time_kernel(M, N, K, BLOCK_M, BLOCK_N, BLOCK_K, nw, ns,
+                            x, ref_model.w_q, ref_model.scales, ref_model.zeros, out)
+            bytes_moved = M*K*2 + (K//2)*N + (K//128)*N*2 + (K//128)*N*2 + M*N*2
+            gbps = bytes_moved / (t * 1e-3) / 1e9
+            print(f"  BM={BLOCK_M:2d} BN={BLOCK_N:3d} BK={BLOCK_K:3d} W={nw} S={ns}: t={t*1000:.2f}us, {gbps:.0f} GB/s")
+        except Exception as e:
+            print(f"  BM={BLOCK_M:2d} BN={BLOCK_N:3d} BK={BLOCK_K:3d} W={nw} S={ns}: ERROR {type(e).__name__}: {e}")

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -33,24 +33,32 @@
 # ---------------------------------------------------------------------------
 
 _CONFIGS = [
-    # M=1 / decode variants — small BN, more programs
+    # M=1 / decode variants — small BN, BLOCK_M=1
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 32, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 32, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 32, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 512, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 512, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    # M>=16 / prefill (BLOCK_M=16 is tensor-core minimum)
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 32, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 32, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
-    # M=16-32 / small prefill
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 512, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    # M=256 / large prefill
-    triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
-    triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 256, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=8, num_stages=2),
     triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=8, num_stages=2),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=8, num_stages=2),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=8, num_stages=2),
 ]
 
 

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -42,9 +42,11 @@
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 512, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 1024, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 256, "GROUP_SZ": 128}, num_warps=4, num_stages=3),
     triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 512, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 1, "BLOCK_N": 128, "BLOCK_K": 1024, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     # M>=16 / prefill (BLOCK_M=16 is tensor-core minimum)
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 32, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),
     triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_SZ": 128}, num_warps=4, num_stages=2),

--- /dev/null
+++ b//workspace/problems/07_w4a16_gemm/test_single_dot.py
@@ -0,0 +1,120 @@
+"""Test single-dot pattern with permute."""
+import torch
+import sys
+sys.path.insert(0, '/workspace/problems/07_w4a16_gemm')
+
+import triton
+import triton.language as tl
+import reference
+
+
+@triton.jit
+def kernel_single_dot(
+    X, WQ, S, Z, OUT,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wk, stride_wn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SZ: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_m_mask = offs_m < M
+
+    n_groups_per_blk: tl.constexpr = BLOCK_K // GROUP_SZ
+    GROUP_HALF: tl.constexpr = GROUP_SZ // 2
+    BLOCK_K_HALF: tl.constexpr = BLOCK_K // 2
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k_blk in tl.range(0, K, BLOCK_K):
+        wq_offs = (k_blk // 2) + tl.arange(0, BLOCK_K_HALF)
+        wq = tl.load(
+            WQ + wq_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn,
+            mask=offs_n[None, :] < N, other=0,
+        )
+
+        w_lo = (wq & 0xF).to(tl.bfloat16)
+        w_hi = (wq >> 4).to(tl.bfloat16)
+
+        g_idx = (k_blk // GROUP_SZ) + tl.arange(0, n_groups_per_blk)
+        s = tl.load(S + g_idx[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+                    mask=offs_n[None, :] < N, other=0.0)
+        z = tl.load(Z + g_idx[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+                    mask=offs_n[None, :] < N, other=0.0)
+
+        s = tl.broadcast_to(s[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        s = tl.reshape(s, (BLOCK_K_HALF, BLOCK_N))
+        z = tl.broadcast_to(z[:, None, :], (n_groups_per_blk, GROUP_HALF, BLOCK_N))
+        z = tl.reshape(z, (BLOCK_K_HALF, BLOCK_N))
+
+        w_lo = (w_lo - z) * s
+        w_hi = (w_hi - z) * s
+
+        # Interleave: w_full[2i, n] = w_lo[i, n], w_full[2i+1, n] = w_hi[i, n]
+        # Use join + permute + reshape
+        w_full = tl.join(w_lo, w_hi)  # (BLOCK_K_HALF, BLOCK_N, 2) with last dim [lo, hi]
+        # We want to permute so the [lo, hi] dim is the K dim
+        w_full = tl.permute(w_full, (0, 2, 1))  # (BLOCK_K_HALF, 2, BLOCK_N) with middle dim [lo, hi]
+        w_full = tl.reshape(w_full, (BLOCK_K, BLOCK_N))
+
+        # Same for x
+        x_offs_even = k_blk + 2 * tl.arange(0, BLOCK_K_HALF)
+        x_offs_odd = x_offs_even + 1
+        x_even = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_even[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+        x_odd = tl.load(
+            X + offs_m[:, None] * stride_xm + x_offs_odd[None, :] * stride_xk,
+            mask=offs_m_mask[:, None], other=0.0,
+        )
+        x_full = tl.join(x_even, x_odd)  # (BLOCK_M, BLOCK_K_HALF, 2)
+        x_full = tl.permute(x_full, (0, 2, 1))  # (BLOCK_M, 2, BLOCK_K_HALF)
+        x_full = tl.reshape(x_full, (BLOCK_M, BLOCK_K))
+
+        acc += tl.dot(x_full, w_full)
+
+    out_ptrs = OUT + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(out_ptrs, acc.to(tl.bfloat16),
+             mask=offs_m_mask[:, None] & (offs_n[None, :] < N))
+
+
+M, N, K = 1, 12288, 4096
+reference.M = M
+reference.N = N
+reference.K = K
+ref_model = reference.Model(M, N, K).cuda().eval()
+
+x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+out = torch.zeros(M, N, dtype=torch.bfloat16).cuda()
+
+BLOCK_M, BLOCK_N, BLOCK_K = 1, 64, 128
+grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+kernel_single_dot[grid](
+    x, ref_model.w_q, ref_model.scales, ref_model.zeros, out,
+    M, N, K,
+    x.stride(0), x.stride(1),
+    ref_model.w_q.stride(0), ref_model.w_q.stride(1),
+    ref_model.scales.stride(0), ref_model.scales.stride(1),
+    ref_model.zeros.stride(0), ref_model.zeros.stride(1),
+    out.stride(0), out.stride(1),
+    BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, GROUP_SZ=128,
+    num_warps=4, num_stages=2,
+)
+ref_out = ref_model(x)
+diff = (ref_out - out).abs().max().item()
+print(f'Single dot: max diff = {diff:.4f}')
+
+# Check the layout
+import triton.language as tl as tl2
+print('tl.join creates a tensor with shape (..., 2) at the end')
+print('After permute and reshape, should be interleaved')

--- a//workspace/problems/07_w4a16_gemm/test_single_dot.py
+++ b//workspace/problems/07_w4a16_gemm/test_single_dot.py
@@ -113,8 +113,3 @@
 ref_out = ref_model(x)
 diff = (ref_out - out).abs().max().item()
 print(f'Single dot: max diff = {diff:.4f}')
-
-# Check the layout
-import triton.language as tl as tl2
-print('tl.join creates a tensor with shape (..., 2) at the end')
-print('After permute and reshape, should be interleaved')