--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -77,11 +77,11 @@
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
     # The kernel assumes BLOCK_K divides GROUP_SIZE so that every K tile sits
-    # inside a single quantisation group.
-    half_k = BLOCK_K // 2
-    k_even_base = tl.arange(0, half_k) * 2       # 0, 2, 4, ..., BLOCK_K-2
-    k_odd_base = k_even_base + 1                  # 1, 3, 5, ..., BLOCK_K-1
-    kh_base = tl.arange(0, half_k)                # packed-byte offsets
+    # inside a single quantisation group.  We use BLOCK_K//2 directly in
+    # arange() calls because the argument must be a compile-time constant.
+    k_even_base = tl.arange(0, BLOCK_K // 2) * 2
+    k_odd_base = k_even_base + 1
+    kh_base = tl.arange(0, BLOCK_K // 2)
 
     for k_start in range(0, K, BLOCK_K):
         # --- packed weights for this K tile --------------------------------

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -1,7 +1,7 @@
 """W4A16 weight-only quantized GEMM — fused unpack + dequant + matmul via Triton.
 
 AWQ/GPTQ-style asymmetric int4 scheme with explicit zero-points and per-group
-(group_size=128) bf16 scales. The Triton kernel loads packed uint8 weights,
+(group_size=128) bf16 scales.  The Triton kernel loads packed uint8 weights,
 unpacks nibbles, applies the per-group affine dequant, and accumulates with
 bf16 activations in a single pass — no materialised bf16 weight matrix.
 """
@@ -20,7 +20,7 @@
 GROUP_SIZE = 128
 
 # ---------------------------------------------------------------------------
-# Triton kernel
+# Triton kernel  (fused unpack + dequant + gemm)
 # ---------------------------------------------------------------------------
 
 @triton.jit
@@ -48,6 +48,7 @@
     BLOCK_K: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     GROUP_M: tl.constexpr,
+    NUM_WARPS: tl.constexpr,
 ):
     """Fused W4A16 GEMM.
 
@@ -70,15 +71,12 @@
     mask_m = offs_m < M
     mask_n = offs_n < N
 
-    # Pointers to the start of the M / N slices
     x_base = x_ptr + offs_m[:, None] * stride_xm
     out_base = out_ptr + offs_m[:, None] * stride_om
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    # The kernel assumes BLOCK_K divides GROUP_SIZE so that every K tile sits
-    # inside a single quantisation group.  We use BLOCK_K//2 directly in
-    # arange() calls because the argument must be a compile-time constant.
+    # BLOCK_K divides GROUP_SIZE so every K tile is inside one quant group.
     k_even_base = tl.arange(0, BLOCK_K // 2) * 2
     k_odd_base = k_even_base + 1
     kh_base = tl.arange(0, BLOCK_K // 2)
@@ -98,15 +96,16 @@
         gid = k_start // GROUP_SIZE
         s_ptrs = scales_ptr + gid * stride_s_g + offs_n * stride_s_n
         z_ptrs = zeros_ptr + gid * stride_z_g + offs_n * stride_z_n
-        s_g = tl.load(s_ptrs, mask=mask_n, other=0.0)   # (BLOCK_N,)
-        z_g = tl.load(z_ptrs, mask=mask_n, other=0.0)   # (BLOCK_N,)
-
-        # --- dequant:  w_bf = (w_int4 - zero) * scale  ---------------------
-        s_g_f32 = s_g.to(tl.float32)[None, :]
-        z_g_f32 = z_g.to(tl.float32)[None, :]
-
-        w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32) * s_g_f32).to(tl.bfloat16)
-        w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32) * s_g_f32).to(tl.bfloat16)
+        s_g = tl.load(s_ptrs, mask=mask_n, other=0.0)
+        z_g = tl.load(z_ptrs, mask=mask_n, other=0.0)
+
+        # --- dequant:  w_bf16 = (w_int4 - zero) * scale  ----------------
+        # Do arithmetic in f32 for precision, then truncate to bf16 for dot.
+        s_g_f32 = s_g.to(tl.float32)
+        z_g_f32 = z_g.to(tl.float32)
+
+        w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
+        w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
 
         # --- activations (even / odd K rows) -------------------------------
         k_even = k_start + k_even_base
@@ -131,96 +130,55 @@
 
 
 # ---------------------------------------------------------------------------
-# Autotune configuration
-# ---------------------------------------------------------------------------
-
-_BLOCK_K_OPTIONS = [32, 64, 128]   # must divide GROUP_SIZE (128)
-
-
-def _configs():
-    """Return a list of Triton autotune Config objects.
-
-    We explore a range of tile sizes appropriate for the small-M decode regime
-    (where memory bandwidth matters most) through to larger prefill batches.
-    """
-    configs = []
-    for BLOCK_M in [1, 4, 8, 16, 32, 64]:
-        for BLOCK_N in [32, 64, 128]:
-            for BLOCK_K in _BLOCK_K_OPTIONS:
-                for GROUP_M in [1, 4, 8]:
-                    # Skip absurd combinations
-                    if BLOCK_M * BLOCK_N > 4096:
-                        continue
-                    configs.append(
-                        triton.Config(
-                            {
-                                "BLOCK_M": BLOCK_M,
-                                "BLOCK_N": BLOCK_N,
-                                "BLOCK_K": BLOCK_K,
-                                "GROUP_M": GROUP_M,
-                            },
-                            num_stages=3,
-                            num_warps=4,
-                        )
-                    )
-    return configs
-
-
-@triton.autotune(
-    configs=_configs(),
-    key=["M", "N", "K"],
-)
-@triton.jit
-def _gemm_w4a16(
-    x_ptr,
-    w_q_ptr,
-    scales_ptr,
-    zeros_ptr,
-    out_ptr,
-    M,
-    N,
-    K,
-    stride_xm,
-    stride_xk,
-    stride_wq_kh,
-    stride_wq_n,
-    stride_s_g,
-    stride_s_n,
-    stride_z_g,
-    stride_z_n,
-    stride_om,
-    stride_on,
-    BLOCK_M: tl.constexpr,
-    BLOCK_N: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-    GROUP_SIZE: tl.constexpr,
-    GROUP_M: tl.constexpr,
-):
-    _gemm_w4a16_kernel(
-        x_ptr,
-        w_q_ptr,
-        scales_ptr,
-        zeros_ptr,
-        out_ptr,
-        M,
-        N,
-        K,
-        stride_xm,
-        stride_xk,
-        stride_wq_kh,
-        stride_wq_n,
-        stride_s_g,
-        stride_s_n,
-        stride_z_g,
-        stride_z_n,
-        stride_om,
-        stride_on,
-        BLOCK_M=BLOCK_M,
-        BLOCK_N=BLOCK_N,
-        BLOCK_K=BLOCK_K,
-        GROUP_SIZE=GROUP_SIZE,
-        GROUP_M=GROUP_M,
-    )
+# Heuristic config selection — avoids autotune overhead and double-kernel
+# issues.  Tuned empirically on RTX PRO 6000 Blackwell.
+# ---------------------------------------------------------------------------
+
+def _pick_config(M: int, N: int, K: int):
+    """Return (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, num_warps) for this shape."""
+
+    # All shapes have K=4096, GROUP_SIZE=128 → 32 groups along K.
+    # BLOCK_K must divide GROUP_SIZE.  128 is the natural choice.
+    BLOCK_K = 128
+
+    if M <= 1:
+        # Decode: t

... (2307 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -48,7 +48,6 @@
     BLOCK_K: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     GROUP_M: tl.constexpr,
-    NUM_WARPS: tl.constexpr,
 ):
     """Fused W4A16 GEMM.
 

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -267,7 +267,6 @@
             BLOCK_K=BLOCK_K,
             GROUP_SIZE=self.group_size,
             GROUP_M=GROUP_M,
-            NUM_WARPS=num_warps,
             num_warps=num_warps,
         )
         return out

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -22,6 +22,7 @@
 # ---------------------------------------------------------------------------
 # Triton kernel  (fused unpack + dequant + gemm)
 # ---------------------------------------------------------------------------
+
 
 @triton.jit
 def _gemm_w4a16_kernel(
@@ -75,16 +76,17 @@
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    # BLOCK_K divides GROUP_SIZE so every K tile is inside one quant group.
-    k_even_base = tl.arange(0, BLOCK_K // 2) * 2
-    k_odd_base = k_even_base + 1
-    kh_base = tl.arange(0, BLOCK_K // 2)
+    # Pre-compute offsets for the K dimension
+    # BLOCK_K divides GROUP_SIZE, so every K tile is inside one quant group.
+    k_offs = tl.arange(0, BLOCK_K)                      # 0, 1, ..., BLOCK_K-1
+    kh_offs = tl.arange(0, BLOCK_K // 2)                 # packed-byte indices
 
     for k_start in range(0, K, BLOCK_K):
         # --- packed weights for this K tile --------------------------------
-        k_half_offs = (k_start // 2) + kh_base
-        wq_ptrs = w_q_ptr + k_half_offs[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
-        mask_wq = (k_half_offs[:, None] < K // 2) & mask_n[None, :]
+        k_half_base = k_start // 2
+        wq_offs_k = k_half_base + kh_offs
+        wq_ptrs = w_q_ptr + wq_offs_k[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
+        mask_wq = (wq_offs_k[:, None] < K // 2) & mask_n[None, :]
         w_packed = tl.load(wq_ptrs, mask=mask_wq, other=0)
 
         # Unpack nibbles: low = even K rows, high = odd K rows
@@ -98,26 +100,45 @@
         s_g = tl.load(s_ptrs, mask=mask_n, other=0.0)
         z_g = tl.load(z_ptrs, mask=mask_n, other=0.0)
 
-        # --- dequant:  w_bf16 = (w_int4 - zero) * scale  ----------------
-        # Do arithmetic in f32 for precision, then truncate to bf16 for dot.
+        # --- dequant:  w_bf16 = (w_int4 - zero) * scale  -------------------
         s_g_f32 = s_g.to(tl.float32)
         z_g_f32 = z_g.to(tl.float32)
 
         w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
         w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
 
-        # --- activations (even / odd K rows) -------------------------------
-        k_even = k_start + k_even_base
-        k_odd = k_start + k_odd_base
-
-        x_even_ptrs = x_base + k_even[None, :] * stride_xk
-        x_odd_ptrs = x_base + k_odd[None, :] * stride_xk
-
-        mask_ke = mask_m[:, None] & (k_even[None, :] < K)
-        mask_ko = mask_m[:, None] & (k_odd[None, :] < K)
-
-        x_even = tl.load(x_even_ptrs, mask=mask_ke, other=0.0)
-        x_odd = tl.load(x_odd_ptrs, mask=mask_ko, other=0.0)
+        # --- activations: load contiguously then split even/odd ------------
+        k_abs = k_start + k_offs
+        x_ptrs = x_base + k_abs[None, :] * stride_xk
+        mask_x = mask_m[:, None] & (k_abs[None, :] < K)
+        x_full = tl.load(x_ptrs, mask=mask_x, other=0.0)   # (BLOCK_M, BLOCK_K)
+
+        # Split into even and odd columns
+        # x_full[:, 0::2] → even K rows, x_full[:, 1::2] → odd K rows
+        x_even = tl.reshape(
+            tl.permute(
+                tl.reshape(x_full, (BLOCK_M, BLOCK_K // 2, 2)),
+                (0, 2, 1),
+            ),
+            (BLOCK_M, BLOCK_K // 2),
+        )  # this is actually [:, 1, :] from the permute since (0,2,1) moves dim 2 to dim 1
+        # Wait — let me think about this more carefully.
+        # x_full: (BLOCK_M, BLOCK_K) where BLOCK_K = 2 * (BLOCK_K//2)
+        # After first reshape: (BLOCK_M, BLOCK_K//2, 2)
+        #   result[m, k2, 0] = x_full[m, 2*k2]     (even)
+        #   result[m, k2, 1] = x_full[m, 2*k2 + 1] (odd)
+        # After permute (0, 2, 1): (BLOCK_M, 2, BLOCK_K//2)
+        #   result[m, 0, k2] = x_full[m, 2*k2]     (even)
+        #   result[m, 1, k2] = x_full[m, 2*k2 + 1] (odd)
+        # After final reshape: (BLOCK_M, BLOCK_K//2)
+        #   This takes the first "slice", which is result[m, 0, :] = even
+
+        # Hmm, actually I need to be more careful with the reshape semantics.
+        # Let me just use the simple stride-slicing approach which should work.
+
+        # Actually, in Triton we CAN use Python-style slicing on tensors:
+        x_even = x_full[:, 0::2]    # every other column starting at 0
+        x_odd = x_full[:, 1::2]     # every other column starting at 1
 
         # --- accumulate ----------------------------------------------------
         acc += tl.dot(x_even, w_lo_bf)
@@ -129,55 +150,33 @@
 
 
 # ---------------------------------------------------------------------------
-# Heuristic config selection — avoids autotune overhead and double-kernel
-# issues.  Tuned empirically on RTX PRO 6000 Blackwell.
+# Heuristic config selection — tuned empirically on RTX PRO 6000 Blackwell.
 # ---------------------------------------------------------------------------
 
 def _pick_config(M: int, N: int, K: int):
-    """Return (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M, num_warps) for this shape."""
-
-    # All shapes have K=4096, GROUP_SIZE=128 → 32 groups along K.
-    # BLOCK_K must divide GROUP_SIZE.  128 is the natural choice.
-    BLOCK_K = 128
-
-    if M <= 1:
-        # Decode: tiny M, memory-bound on weight read.
-        # BLOCK_M=4 is faster than BLOCK_M=1 because it improves tensor-core
-        # utilisation and amortises control overhead with minimal waste
-        # (three masked-out rows per tile).
-        BLOCK_M = 4 if N >= 4096 else 1
-        BLOCK_N = 64
-        num_warps = 4
+    """Return (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M) for this shape."""
+
+    BLOCK_K = 128  # must divide GROUP_SIZE=128; one group per K-tile is natural
+
+    if M == 1 and N >= 8192:
+        # Decode, wide: memory-bound. BLOCK_M=4 amortises launch overhead
+        # despite wasting 3 masked rows.
+        return 4, 64, BL

... (1860 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -76,20 +76,20 @@
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    # Pre-compute offsets for the K dimension
-    # BLOCK_K divides GROUP_SIZE, so every K tile is inside one quant group.
-    k_offs = tl.arange(0, BLOCK_K)                      # 0, 1, ..., BLOCK_K-1
-    kh_offs = tl.arange(0, BLOCK_K // 2)                 # packed-byte indices
+    # Offsets for even/odd K rows and packed-byte indices.
+    # BLOCK_K divides GROUP_SIZE → every K tile inside one quant group.
+    k_even_base = tl.arange(0, BLOCK_K // 2) * 2
+    k_odd_base = k_even_base + 1
+    kh_base = tl.arange(0, BLOCK_K // 2)
 
     for k_start in range(0, K, BLOCK_K):
         # --- packed weights for this K tile --------------------------------
-        k_half_base = k_start // 2
-        wq_offs_k = k_half_base + kh_offs
-        wq_ptrs = w_q_ptr + wq_offs_k[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
-        mask_wq = (wq_offs_k[:, None] < K // 2) & mask_n[None, :]
+        k_half_offs = (k_start // 2) + kh_base
+        wq_ptrs = w_q_ptr + k_half_offs[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
+        mask_wq = (k_half_offs[:, None] < K // 2) & mask_n[None, :]
         w_packed = tl.load(wq_ptrs, mask=mask_wq, other=0)
 
-        # Unpack nibbles: low = even K rows, high = odd K rows
+        # Unpack nibbles: low nibble = even K rows, high nibble = odd K rows
         w_lo = w_packed & 0x0F
         w_hi = (w_packed >> 4) & 0x0F
 
@@ -107,38 +107,18 @@
         w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
         w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
 
-        # --- activations: load contiguously then split even/odd ------------
-        k_abs = k_start + k_offs
-        x_ptrs = x_base + k_abs[None, :] * stride_xk
-        mask_x = mask_m[:, None] & (k_abs[None, :] < K)
-        x_full = tl.load(x_ptrs, mask=mask_x, other=0.0)   # (BLOCK_M, BLOCK_K)
-
-        # Split into even and odd columns
-        # x_full[:, 0::2] → even K rows, x_full[:, 1::2] → odd K rows
-        x_even = tl.reshape(
-            tl.permute(
-                tl.reshape(x_full, (BLOCK_M, BLOCK_K // 2, 2)),
-                (0, 2, 1),
-            ),
-            (BLOCK_M, BLOCK_K // 2),
-        )  # this is actually [:, 1, :] from the permute since (0,2,1) moves dim 2 to dim 1
-        # Wait — let me think about this more carefully.
-        # x_full: (BLOCK_M, BLOCK_K) where BLOCK_K = 2 * (BLOCK_K//2)
-        # After first reshape: (BLOCK_M, BLOCK_K//2, 2)
-        #   result[m, k2, 0] = x_full[m, 2*k2]     (even)
-        #   result[m, k2, 1] = x_full[m, 2*k2 + 1] (odd)
-        # After permute (0, 2, 1): (BLOCK_M, 2, BLOCK_K//2)
-        #   result[m, 0, k2] = x_full[m, 2*k2]     (even)
-        #   result[m, 1, k2] = x_full[m, 2*k2 + 1] (odd)
-        # After final reshape: (BLOCK_M, BLOCK_K//2)
-        #   This takes the first "slice", which is result[m, 0, :] = even
-
-        # Hmm, actually I need to be more careful with the reshape semantics.
-        # Let me just use the simple stride-slicing approach which should work.
-
-        # Actually, in Triton we CAN use Python-style slicing on tensors:
-        x_even = x_full[:, 0::2]    # every other column starting at 0
-        x_odd = x_full[:, 1::2]     # every other column starting at 1
+        # --- activations (even / odd K rows) -------------------------------
+        k_even = k_start + k_even_base
+        k_odd = k_start + k_odd_base
+
+        x_even_ptrs = x_base + k_even[None, :] * stride_xk
+        x_odd_ptrs = x_base + k_odd[None, :] * stride_xk
+
+        mask_ke = mask_m[:, None] & (k_even[None, :] < K)
+        mask_ko = mask_m[:, None] & (k_odd[None, :] < K)
+
+        x_even = tl.load(x_even_ptrs, mask=mask_ke, other=0.0)
+        x_odd = tl.load(x_odd_ptrs, mask=mask_ko, other=0.0)
 
         # --- accumulate ----------------------------------------------------
         acc += tl.dot(x_even, w_lo_bf)
@@ -155,28 +135,28 @@
 
 def _pick_config(M: int, N: int, K: int):
     """Return (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M) for this shape."""
-
-    BLOCK_K = 128  # must divide GROUP_SIZE=128; one group per K-tile is natural
-
-    if M == 1 and N >= 8192:
-        # Decode, wide: memory-bound. BLOCK_M=4 amortises launch overhead
-        # despite wasting 3 masked rows.
-        return 4, 64, BLOCK_K, 8
+    BLOCK_K = 128  # one group per K-tile
+
+    if M == 1:
+        if N >= 8192:
+            return (4, 64, BLOCK_K, 8)
+        else:
+            return (1, 64, BLOCK_K, 8)
     elif M <= 4:
-        return 4, 64, BLOCK_K, 8
+        return (4, 64, BLOCK_K, 8)
     elif M <= 8:
-        return 8, 64, BLOCK_K, 8
+        return (8, 64, BLOCK_K, 8)
     elif M <= 16:
         if N >= 12288:
-            return 8, 64, BLOCK_K, 8
+            return (8, 64, BLOCK_K, 8)
         else:
-            return 8, 128, BLOCK_K, 8
+            return (8, 128, BLOCK_K, 8)
     elif M <= 32:
-        return 8, 128, BLOCK_K, 8
+        return (8, 128, BLOCK_K, 8)
     elif M <= 64:
-        return 16, 128, BLOCK_K, 8
+        return (16, 128, BLOCK_K, 8)
     else:
-        return 16, 128, BLOCK_K, 8
+        return (16, 128, BLOCK_K, 8)
 
 
 # ---------------------------------------------------------------------------

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -2,8 +2,8 @@
 
 AWQ/GPTQ-style asymmetric int4 scheme with explicit zero-points and per-group
 (group_size=128) bf16 scales.  The Triton kernel loads packed uint8 weights,
-unpacks nibbles, applies the per-group affine dequant, and accumulates with
-bf16 activations in a single pass — no materialised bf16 weight matrix.
+unpacks nibbles into an interleaved bf16 tile, and accumulates with bf16
+activations in a single fused pass — no materialised bf16 weight matrix.
 """
 
 from __future__ import annotations
@@ -52,9 +52,9 @@
 ):
     """Fused W4A16 GEMM.
 
-    Each program computes a (BLOCK_M, BLOCK_N) tile of the output.  The
-    reduction loops over K in steps of BLOCK_K, which must divide GROUP_SIZE
-    so that every K-block sits entirely inside a single quantisation group.
+    Each program computes a (BLOCK_M, BLOCK_N) tile of the output.
+    BLOCK_K must divide GROUP_SIZE so that every K-block sits entirely
+    inside a single quantisation group.
     """
     pid = tl.program_id(0)
     num_pid_m = tl.cdiv(M, BLOCK_M)
@@ -76,29 +76,27 @@
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    # Offsets for even/odd K rows and packed-byte indices.
-    # BLOCK_K divides GROUP_SIZE → every K tile inside one quant group.
-    k_even_base = tl.arange(0, BLOCK_K // 2) * 2
-    k_odd_base = k_even_base + 1
-    kh_base = tl.arange(0, BLOCK_K // 2)
+    # Pre-compute K-dimension offsets and packed-byte indices
+    k_offs = tl.arange(0, BLOCK_K)                         # 0 .. BLOCK_K-1
+    kh_offs = tl.arange(0, BLOCK_K // 2)                    # packed-byte indices
 
     for k_start in range(0, K, BLOCK_K):
         # --- packed weights for this K tile --------------------------------
-        k_half_offs = (k_start // 2) + kh_base
-        wq_ptrs = w_q_ptr + k_half_offs[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
-        mask_wq = (k_half_offs[:, None] < K // 2) & mask_n[None, :]
+        k_half_base = k_start // 2
+        wq_ptrs = w_q_ptr + (k_half_base + kh_offs)[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
+        mask_wq = ((k_half_base + kh_offs)[:, None] < K // 2) & mask_n[None, :]
         w_packed = tl.load(wq_ptrs, mask=mask_wq, other=0)
 
-        # Unpack nibbles: low nibble = even K rows, high nibble = odd K rows
-        w_lo = w_packed & 0x0F
+        # Unpack nibbles: low = even K rows, high = odd K rows
+        w_lo = w_packed & 0x0F    # (BLOCK_K//2, BLOCK_N)  uint8
         w_hi = (w_packed >> 4) & 0x0F
 
         # --- per-group scales & zeros --------------------------------------
         gid = k_start // GROUP_SIZE
-        s_ptrs = scales_ptr + gid * stride_s_g + offs_n * stride_s_n
-        z_ptrs = zeros_ptr + gid * stride_z_g + offs_n * stride_z_n
-        s_g = tl.load(s_ptrs, mask=mask_n, other=0.0)
-        z_g = tl.load(z_ptrs, mask=mask_n, other=0.0)
+        s_g = tl.load(scales_ptr + gid * stride_s_g + offs_n * stride_s_n,
+                      mask=mask_n, other=0.0)
+        z_g = tl.load(zeros_ptr + gid * stride_z_g + offs_n * stride_z_n,
+                      mask=mask_n, other=0.0)
 
         # --- dequant:  w_bf16 = (w_int4 - zero) * scale  -------------------
         s_g_f32 = s_g.to(tl.float32)
@@ -107,22 +105,24 @@
         w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
         w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
 
-        # --- activations (even / odd K rows) -------------------------------
-        k_even = k_start + k_even_base
-        k_odd = k_start + k_odd_base
-
-        x_even_ptrs = x_base + k_even[None, :] * stride_xk
-        x_odd_ptrs = x_base + k_odd[None, :] * stride_xk
-
-        mask_ke = mask_m[:, None] & (k_even[None, :] < K)
-        mask_ko = mask_m[:, None] & (k_odd[None, :] < K)
-
-        x_even = tl.load(x_even_ptrs, mask=mask_ke, other=0.0)
-        x_odd = tl.load(x_odd_ptrs, mask=mask_ko, other=0.0)
-
-        # --- accumulate ----------------------------------------------------
-        acc += tl.dot(x_even, w_lo_bf)
-        acc += tl.dot(x_odd, w_hi_bf)
+        # --- interleave even/odd weights into (BLOCK_K, BLOCK_N) -----------
+        # tl.join adds a trailing dim → (BLOCK_K//2, BLOCK_N, 2)
+        # tl.permute → (BLOCK_K//2, 2, BLOCK_N)
+        # tl.reshape → (BLOCK_K, BLOCK_N)  with even/odd interleaved along K
+        w_interleaved = tl.permute(
+            tl.join(w_lo_bf, w_hi_bf),
+            (0, 2, 1),
+        )
+        w_deq = tl.reshape(w_interleaved, (BLOCK_K, BLOCK_N))
+
+        # --- activations: single contiguous load ---------------------------
+        k_abs = k_start + k_offs
+        x_ptrs = x_base + k_abs[None, :] * stride_xk
+        mask_x = mask_m[:, None] & (k_abs[None, :] < K)
+        x_full = tl.load(x_ptrs, mask=mask_x, other=0.0)   # (BLOCK_M, BLOCK_K)
+
+        # --- accumulate  (single dot instead of two) -----------------------
+        acc += tl.dot(x_full, w_deq)
 
     # --- store ------------------------------------------------------------
     out_ptrs = out_base + offs_n[None, :] * stride_on
@@ -138,11 +138,6 @@
     BLOCK_K = 128  # one group per K-tile
 
     if M == 1:
-        if N >= 8192:
-            return (4, 64, BLOCK_K, 8)
-        else:
-            return (1, 64, BLOCK_K, 8)
-    elif M <= 4:
         return (4, 64, BLOCK_K, 8)
     elif M <= 8:
         return (8, 64, BLOCK_K, 8)
@@ -164,11 +159,7 @@
 # ---------------------------------------------------------------------------
 
 class Model(nn.Module):
-    """W4A16 GEMM:  y = x @ dequant(w_q, scales, zeros).
-
-    Weight quantisation is synthetic (matches the reference), producing a
-    correct (w_q, scales, zeros) triple that round-trips cleanly.
-    """
+    """W4A16 GEMM:  y = x @ dequant(w_q, scales, zeros)."""
 
     def __init__(self, M: in

... (78 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -2,8 +2,8 @@
 
 AWQ/GPTQ-style asymmetric int4 scheme with explicit zero-points and per-group
 (group_size=128) bf16 scales.  The Triton kernel loads packed uint8 weights,
-unpacks nibbles into an interleaved bf16 tile, and accumulates with bf16
-activations in a single fused pass — no materialised bf16 weight matrix.
+unpacks nibbles, applies the per-group affine dequant, and accumulates with
+bf16 activations in a single fused pass — no materialised bf16 weight matrix.
 """
 
 from __future__ import annotations
@@ -20,12 +20,37 @@
 GROUP_SIZE = 128
 
 # ---------------------------------------------------------------------------
-# Triton kernel  (fused unpack + dequant + gemm)
-# ---------------------------------------------------------------------------
-
-
+# Autotuning configs — slim set covering all shape regimes
+# ---------------------------------------------------------------------------
+
+@triton.autotune(
+    configs=[
+        # decode-style: tiny M, memory-bound
+        triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 4, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        # small prefill / spec-decode
+        triton.Config({"BLOCK_M": 8, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        # larger prefill
+        triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
+        # wide-N variants with larger blocks
+        triton.Config({"BLOCK_M": 4, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
+        triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
+        # extra: num_stages=2 for less shared memory pressure
+        triton.Config({"BLOCK_M": 4, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_M": 8, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
+        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
+    ],
+    key=["M", "N", "K"],
+    prune_configs_by={"early_config_prune": None},
+)
 @triton.jit
-def _gemm_w4a16_kernel(
+def _gemm_w4a16(
     x_ptr,
     w_q_ptr,
     scales_ptr,
@@ -50,12 +75,7 @@
     GROUP_SIZE: tl.constexpr,
     GROUP_M: tl.constexpr,
 ):
-    """Fused W4A16 GEMM.
-
-    Each program computes a (BLOCK_M, BLOCK_N) tile of the output.
-    BLOCK_K must divide GROUP_SIZE so that every K-block sits entirely
-    inside a single quantisation group.
-    """
+    """Fused W4A16 GEMM — unpack + dequant + matmul in one pass."""
     pid = tl.program_id(0)
     num_pid_m = tl.cdiv(M, BLOCK_M)
     num_pid_n = tl.cdiv(N, BLOCK_N)
@@ -76,19 +96,18 @@
 
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
 
-    # Pre-compute K-dimension offsets and packed-byte indices
-    k_offs = tl.arange(0, BLOCK_K)                         # 0 .. BLOCK_K-1
-    kh_offs = tl.arange(0, BLOCK_K // 2)                    # packed-byte indices
+    k_even_base = tl.arange(0, BLOCK_K // 2) * 2
+    k_odd_base = k_even_base + 1
+    kh_base = tl.arange(0, BLOCK_K // 2)
 
     for k_start in range(0, K, BLOCK_K):
         # --- packed weights for this K tile --------------------------------
-        k_half_base = k_start // 2
-        wq_ptrs = w_q_ptr + (k_half_base + kh_offs)[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
-        mask_wq = ((k_half_base + kh_offs)[:, None] < K // 2) & mask_n[None, :]
+        k_half_offs = (k_start // 2) + kh_base
+        wq_ptrs = w_q_ptr + k_half_offs[:, None] * stride_wq_kh + offs_n[None, :] * stride_wq_n
+        mask_wq = (k_half_offs[:, None] < K // 2) & mask_n[None, :]
         w_packed = tl.load(wq_ptrs, mask=mask_wq, other=0)
 
-        # Unpack nibbles: low = even K rows, high = odd K rows
-        w_lo = w_packed & 0x0F    # (BLOCK_K//2, BLOCK_N)  uint8
+        w_lo = w_packed & 0x0F
         w_hi = (w_packed >> 4) & 0x0F
 
         # --- per-group scales & zeros --------------------------------------
@@ -98,60 +117,33 @@
         z_g = tl.load(zeros_ptr + gid * stride_z_g + offs_n * stride_z_n,
                       mask=mask_n, other=0.0)
 
-        # --- dequant:  w_bf16 = (w_int4 - zero) * scale  -------------------
+        # --- dequant -------------------------------------------------------
         s_g_f32 = s_g.to(tl.float32)
         z_g_f32 = z_g.to(tl.float32)
 
         w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
         w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
 
-        # --- interleave even/odd weights into (BLOCK_K, BLOCK_N) -----------
-        # tl.join adds a trailing dim → (BLOCK_K//2, BLOCK_N, 2)
-        # tl.permute → (BLOCK_K//2, 2, BLOCK_N)
-        # tl.reshape → (BLOCK_K, BLOCK_N)  with even/odd interleaved along K
-        w_interleaved = tl.permute(
-            tl.join(w_lo_bf, w_hi_bf),
-            (0, 2, 1),
-        )
-        w_deq = tl.reshape(w_interleaved, (BLOCK_K, BLOCK_N))
-
-        #

... (3045 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -2,8 +2,9 @@
 
 AWQ/GPTQ-style asymmetric int4 scheme with explicit zero-points and per-group
 (group_size=128) bf16 scales.  The Triton kernel loads packed uint8 weights,
-unpacks nibbles, applies the per-group affine dequant, and accumulates with
-bf16 activations in a single fused pass — no materialised bf16 weight matrix.
+unpacks nibbles, applies the per-group affine dequant (using bf16 arithmetic
+to reduce register pressure), and accumulates with bf16 activations in a
+single fused pass.
 """
 
 from __future__ import annotations
@@ -20,37 +21,12 @@
 GROUP_SIZE = 128
 
 # ---------------------------------------------------------------------------
-# Autotuning configs — slim set covering all shape regimes
-# ---------------------------------------------------------------------------
-
-@triton.autotune(
-    configs=[
-        # decode-style: tiny M, memory-bound
-        triton.Config({"BLOCK_M": 1, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 4, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        # small prefill / spec-decode
-        triton.Config({"BLOCK_M": 8, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        # larger prefill
-        triton.Config({"BLOCK_M": 16, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 32, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=3),
-        # wide-N variants with larger blocks
-        triton.Config({"BLOCK_M": 4, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 16, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        triton.Config({"BLOCK_M": 32, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=8, num_stages=2),
-        # extra: num_stages=2 for less shared memory pressure
-        triton.Config({"BLOCK_M": 4, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_M": 8, "BLOCK_N": 64, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
-        triton.Config({"BLOCK_M": 8, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, num_warps=4, num_stages=2),
-    ],
-    key=["M", "N", "K"],
-    prune_configs_by={"early_config_prune": None},
-)
+# Triton kernel  (fused unpack + dequant + gemm)
+# ---------------------------------------------------------------------------
+
+
 @triton.jit
-def _gemm_w4a16(
+def _gemm_w4a16_kernel(
     x_ptr,
     w_q_ptr,
     scales_ptr,
@@ -117,12 +93,9 @@
         z_g = tl.load(zeros_ptr + gid * stride_z_g + offs_n * stride_z_n,
                       mask=mask_n, other=0.0)
 
-        # --- dequant -------------------------------------------------------
-        s_g_f32 = s_g.to(tl.float32)
-        z_g_f32 = z_g.to(tl.float32)
-
-        w_lo_bf = ((w_lo.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
-        w_hi_bf = ((w_hi.to(tl.float32) - z_g_f32[None, :]) * s_g_f32[None, :]).to(tl.bfloat16)
+        # --- dequant in bf16 (avoids f32 intermediate to save registers) ---
+        w_lo_bf = (w_lo.to(tl.float32) - z_g[None, :]).to(tl.bfloat16) * s_g[None, :]
+        w_hi_bf = (w_hi.to(tl.float32) - z_g[None, :]).to(tl.bfloat16) * s_g[None, :]
 
         # --- activations (even / odd K rows) -------------------------------
         k_even = k_start + k_even_base
@@ -144,6 +117,32 @@
     # --- store ------------------------------------------------------------
     out_ptrs = out_base + offs_n[None, :] * stride_on
     tl.store(out_ptrs, acc.to(tl.bfloat16), mask=mask_m[:, None] & mask_n[None, :])
+
+
+# ---------------------------------------------------------------------------
+# Heuristic config — empirically best configs per shape on RTX PRO 6000
+# ---------------------------------------------------------------------------
+
+def _pick_config(M: int, N: int, K: int):
+    """Return (BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M) for this shape."""
+    BK = 128  # = GROUP_SIZE — one group per K-tile
+
+    if M <= 1:
+        if N >= 8192:
+            return (1, 64, BK, 8)
+        return (4, 64, BK, 8)
+    elif M <= 8:
+        return (8, 64, BK, 8)
+    elif M <= 16:
+        if N >= 12288:
+            return (8, 64, BK, 8)
+        return (8, 128, BK, 8)
+    elif M <= 32:
+        return (8, 128, BK, 8)
+    elif M <= 64:
+        return (16, 128, BK, 8)
+    else:
+        return (16, 128, BK, 8)
 
 
 # ---------------------------------------------------------------------------
@@ -197,12 +196,11 @@
         scales = self.scales.contiguous()
         zeros = self.zeros.contiguous()
 
-        grid = lambda meta: (
-            triton.cdiv(M_out, meta["BLOCK_M"])
-            * triton.cdiv(self.N, meta["BLOCK_N"]),
-        )
-
-        _gemm_w4a16[grid](
+        BLOCK_M, BLOCK_N, BLOCK_K, GROUP_M = _pick_config(M_out, self.N, self.K)
+
+        grid = (triton.cdiv(M_out, BLOCK_M) * triton.cdiv(self.N, BLOCK_N),)
+
+        _gemm_w4a16_kernel[grid](
             x,
             w_q,
             scales,
@@ -221,7 +219,11 @@
             zeros.stride(1),
             out.stride(0),
             out.stride(1),
+            BLOCK_M=BLOCK_M,
+            BLOCK_N=BLOCK_N,
+            BLOCK_K=BLOCK_K,
             GROUP_SIZE=self.group_size,
+            GROUP_M=GROUP_M,
         

... (24 more chars)

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -128,21 +128,27 @@
     BK = 128  # = GROUP_SIZE — one group per K-tile
 
     if M <= 1:
+        # Decode: very small M, maximise N tile for throughput.
         if N >= 8192:
-            return (1, 64, BK, 8)
-        return (4, 64, BK, 8)
+            return (1, 64, BK, 4)
+        return (1, 64, BK, 4)
+    elif M <= 4:
+        return (4, 64, BK, 4)
     elif M <= 8:
-        return (8, 64, BK, 8)
+        return (8, 64, BK, 4)
     elif M <= 16:
         if N >= 12288:
-            return (8, 64, BK, 8)
-        return (8, 128, BK, 8)
+            return (8, 64, BK, 4)
+        return (8, 128, BK, 4)
     elif M <= 32:
-        return (8, 128, BK, 8)
+        return (16, 128, BK, 4)
     elif M <= 64:
-        return (16, 128, BK, 8)
+        return (32, 128, BK, 4)
+    elif M <= 128:
+        return (32, 128, BK, 4)
     else:
-        return (16, 128, BK, 8)
+        # M >= 256 — larger prefill
+        return (32, 128, BK, 4)
 
 
 # ---------------------------------------------------------------------------

--- a//workspace/problems/07_w4a16_gemm/solution.py
+++ b//workspace/problems/07_w4a16_gemm/solution.py
@@ -128,27 +128,21 @@
     BK = 128  # = GROUP_SIZE — one group per K-tile
 
     if M <= 1:
-        # Decode: very small M, maximise N tile for throughput.
-        if N >= 8192:
-            return (1, 64, BK, 4)
-        return (1, 64, BK, 4)
-    elif M <= 4:
-        return (4, 64, BK, 4)
+        # Decode: maximise memory throughput.  BLOCK_M={1,4} both measured
+        # at ~525 GB/s; using 4 amortises launch overhead slightly better.
+        return (4, 64, BK, 8)
     elif M <= 8:
-        return (8, 64, BK, 4)
+        return (8, 64, BK, 8)
     elif M <= 16:
         if N >= 12288:
-            return (8, 64, BK, 4)
-        return (8, 128, BK, 4)
+            return (8, 64, BK, 8)
+        return (8, 128, BK, 8)
     elif M <= 32:
-        return (16, 128, BK, 4)
+        return (8, 128, BK, 8)
     elif M <= 64:
-        return (32, 128, BK, 4)
-    elif M <= 128:
-        return (32, 128, BK, 4)
+        return (16, 128, BK, 8)
     else:
-        # M >= 256 — larger prefill
-        return (32, 128, BK, 4)
+        return (16, 128, BK, 8)
 
 
 # ---------------------------------------------------------------------------