--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -0,0 +1,230 @@
+"""High-performance W4A16 weight-only quantized GEMM for SM120 Blackwell.
+
+AWQ/GPTQ-style asymmetric int4 with per-group bf16 scales and zero-points.
+Fuses unpack, dequant, and matmul in a single Triton kernel.
+
+Data layout:
+  x:      (M, K)                  bf16
+  w_q:    (K // 2, N)             uint8  (low nibble=even-K, high=odd-K)
+  scales: (K // 128, N)           bf16
+  zeros:  (K // 128, N)           bf16
+  out:    (M, N)                  bf16
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+GROUP_SIZE = 128
+
+
+# ---------------------------------------------------------------------------
+# Triton kernel — fused unpack + dequant + GEMM
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def _w4a16_gemm_kernel(
+    x_ptr,
+    wq_ptr,
+    scales_ptr,
+    zeros_ptr,
+    out_ptr,
+    M,
+    N,
+    K,
+    stride_xm,
+    stride_xk,
+    stride_wqk,
+    stride_wqn,
+    stride_sg,
+    stride_sn,
+    stride_zg,
+    stride_zn,
+    stride_om,
+    stride_on,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+):
+    """Each program computes out[pid_m*BM:(pid_m+1)*BM, pid_n*BN:(pid_n+1)*BN].
+
+    K is accumulated in tiles of BLOCK_SIZE_K (must divide GROUP_SIZE).
+    """
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    # --- offsets into each dimension ---
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)  # [BM]
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)  # [BN]
+    offs_k_base = tl.arange(0, BLOCK_SIZE_K)                     # [BK]
+
+    # --- validity mask for output row/col ---
+    m_mask = offs_m < M   # [BM]
+    n_mask = offs_n < N   # [BN]
+    out_mask = m_mask[:, None] & n_mask[None, :]                 # [BM, BN]
+
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    for k_base in range(0, K, BLOCK_SIZE_K):
+        k_mask = (k_base + offs_k_base) < K  # [BK]
+
+        # ---- load x[offs_m, k_base:...]  shape [BM, BK] ----
+        x_block = tl.load(
+            x_ptr
+            + offs_m[:, None] * stride_xm
+            + (k_base + offs_k_base)[None, :] * stride_xk,
+            mask=m_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        # ---- load w_q[k_base//2 + offs_k_base//2, offs_n]  shape [BK, BN] ----
+        wq_row_idx = (k_base + offs_k_base) // 2  # [BK]
+        wq_ptrs = (
+            wq_ptr
+            + wq_row_idx[:, None] * stride_wqk
+            + offs_n[None, :] * stride_wqn
+        )
+        wq_col_mask = k_mask[:, None] & n_mask[None, :]
+        wq_block = tl.load(wq_ptrs, mask=wq_col_mask, other=0)  # uint8
+
+        # unpack int4 from low/high nibble based on actual k parity
+        actual_k_idx = k_base + offs_k_base  # [BK]
+        is_odd_k = (actual_k_idx % 2).to(tl.int32)  # 0 or 1, broadcast [BK,1]
+        w_int = tl.where(
+            is_odd_k[:, None] == 0,
+            wq_block & 0xF,
+            wq_block >> 4,
+        ).to(tl.int32)  # [BK, BN]
+
+        # ---- load scales & zeros for this block  shape [BK, BN] ----
+        g_idx = (k_base + offs_k_base) // GROUP_SIZE  # [BK]
+        s_block = tl.load(
+            scales_ptr
+            + g_idx[:, None] * stride_sg
+            + offs_n[None, :] * stride_sn,
+            mask=k_mask[:, None] & n_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        z_block = tl.load(
+            zeros_ptr
+            + g_idx[:, None] * stride_zg
+            + offs_n[None, :] * stride_zn,
+            mask=k_mask[:, None] & n_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+
+        w_int = w_int.to(tl.float32)
+        w_deq = (w_int - z_block) * s_block  # [BK, BN] dequantized weight
+
+        # ---- matmul tile (BM, BK) @ (BK, BN) -> (BM, BN) ----
+        acc = tl.dot(x_block.to(tl.float32), w_deq, acc)
+
+    # ---- write result (M, N) bf16 ----
+    out_ptrs = (
+        out_ptr
+        + offs_m[:, None] * stride_om
+        + offs_n[None, :] * stride_on
+    )
+    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=out_mask)
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+class Model(nn.Module):
+    """W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""
+
+    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
+        super().__init__()
+        assert K % group_size == 0, "K must be divisible by group_size"
+        assert K % 2 == 0, "K must be even (int4 packing)"
+        self.M, self.N, self.K = M, N, K
+        self.group_size = group_size
+        n_groups = K // group_size
+
+        torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
+        w_full = torch.randn(K, N, dtype=torch.float32) * 0.02
+
+        w_g = w_full.view(n_groups, group_size, N)
+        w_min = w_g.min(dim=1, keepdim=True).values
+        w_max = w_g.max(dim=1, keepdim=True).values
+        scales = (w_max - w_min).clamp_min(1e-8) / 15.0
+        zeros = (-w_min / scales).round().clamp(0, 15)
+
+        w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8)
+        w_q = w_q.view(K, N)
+
+        # Pack: even row → low nibble, odd row → high nibble
+        lo = w_q[0::2].to(torch.uint8) & 0xF
+        hi = w_q[1::2].to(torch.uint8) & 0xF
+        w_packed = (lo | (hi << 4)).contiguous()
+
+        scales_2d = scales.squeeze(1).to(torch.bfloat16)
+        zeros_2d = zeros.squeeze(1).to(torch.bfloat16)
+
+        self.register_buffer("w_q", w_packed)
+        self.register_buffer("scales", scales_2d)
+        self.r

... (1759 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -179,16 +179,16 @@
 
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Block config: BM=64 is optimal for M>=2; M=1 special case
+        # Block sizes tuned to fit SM120 shared memory (101 KB)
         if M == 1:
             BLOCK_M = 1
             warps = 4
         else:
             BLOCK_M = 64
-            warps = 4
+            warps = 8
 
         BLOCK_N = 128
-        BLOCK_K = 128
+        BLOCK_K = 64
 
         grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -179,16 +179,16 @@
 
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Block sizes tuned to fit SM120 shared memory (101 KB)
         if M == 1:
             BLOCK_M = 1
             warps = 4
         else:
-            BLOCK_M = 64
-            warps = 8
-
-        BLOCK_N = 128
+            BLOCK_M = 32
+            warps = 4
+
+        BLOCK_N = 64
         BLOCK_K = 64
+        num_stages = 2
 
         grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -209,6 +209,7 @@
             BLOCK_SIZE_K=BLOCK_K,
             GROUP_SIZE=GROUP_SIZE,
             num_warps=warps,
+            num_stages=num_stages,
         )
         return out
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -181,9 +181,9 @@
 
         if M == 1:
             BLOCK_M = 1
-            warps = 4
+            warps = 1
         else:
-            BLOCK_M = 32
+            BLOCK_M = 64
             warps = 4
 
         BLOCK_N = 64

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -21,10 +21,6 @@
 GROUP_SIZE = 128
 
 
-# ---------------------------------------------------------------------------
-# Triton kernel — fused unpack + dequant + GEMM
-# ---------------------------------------------------------------------------
-
 @triton.jit
 def _w4a16_gemm_kernel(
     x_ptr,
@@ -35,102 +31,72 @@
     M,
     N,
     K,
-    stride_xm,
-    stride_xk,
-    stride_wqk,
-    stride_wqn,
-    stride_sg,
-    stride_sn,
-    stride_zg,
-    stride_zn,
-    stride_om,
-    stride_on,
+    stride_xm, stride_xk,
+    stride_wqk, stride_wqn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE: tl.constexpr,
+    GROUP_SIZE_CONST: tl.constexpr,
 ):
-    """Each program computes out[pid_m*BM:(pid_m+1)*BM, pid_n*BN:(pid_n+1)*BN].
-
-    K is accumulated in tiles of BLOCK_SIZE_K (must divide GROUP_SIZE).
-    """
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
 
-    # --- offsets into each dimension ---
-    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)  # [BM]
-    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)  # [BN]
-    offs_k_base = tl.arange(0, BLOCK_SIZE_K)                     # [BK]
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
 
-    # --- validity mask for output row/col ---
-    m_mask = offs_m < M   # [BM]
-    n_mask = offs_n < N   # [BN]
-    out_mask = m_mask[:, None] & n_mask[None, :]                 # [BM, BN]
+    m_mask = offs_m < M
+    n_mask = offs_n < N
+    out_mask = m_mask[:, None] & n_mask[None, :]
 
     acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
 
     for k_base in range(0, K, BLOCK_SIZE_K):
-        k_mask = (k_base + offs_k_base) < K  # [BK]
+        k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
+        k_mask = k_offs < K
 
-        # ---- load x[offs_m, k_base:...]  shape [BM, BK] ----
-        x_block = tl.load(
-            x_ptr
-            + offs_m[:, None] * stride_xm
-            + (k_base + offs_k_base)[None, :] * stride_xk,
+        # x tile: (BM, BK)
+        x_tile = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
-        # ---- load w_q[k_base//2 + offs_k_base//2, offs_n]  shape [BK, BN] ----
-        wq_row_idx = (k_base + offs_k_base) // 2  # [BK]
-        wq_ptrs = (
-            wq_ptr
-            + wq_row_idx[:, None] * stride_wqk
-            + offs_n[None, :] * stride_wqn
+        # w_q tile: (BK, BN) — load packed bytes
+        wq_row = k_offs // 2
+        wq_tile = tl.load(
+            wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
+            mask=k_mask[:, None] & n_mask[None, :],
+            other=0,
         )
-        wq_col_mask = k_mask[:, None] & n_mask[None, :]
-        wq_block = tl.load(wq_ptrs, mask=wq_col_mask, other=0)  # uint8
 
-        # unpack int4 from low/high nibble based on actual k parity
-        actual_k_idx = k_base + offs_k_base  # [BK]
-        is_odd_k = (actual_k_idx % 2).to(tl.int32)  # 0 or 1, broadcast [BK,1]
-        w_int = tl.where(
-            is_odd_k[:, None] == 0,
-            wq_block & 0xF,
-            wq_block >> 4,
-        ).to(tl.int32)  # [BK, BN]
+        # unpack nibbles: even-k → low, odd-k → high
+        is_odd = (k_offs % 2)[:, None]
+        w_int = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # ---- load scales & zeros for this block  shape [BK, BN] ----
-        g_idx = (k_base + offs_k_base) // GROUP_SIZE  # [BK]
-        s_block = tl.load(
-            scales_ptr
-            + g_idx[:, None] * stride_sg
-            + offs_n[None, :] * stride_sn,
+        # scales / zeros: (BK, BN) via group index
+        g = k_offs // GROUP_SIZE_CONST
+        s_tile = tl.load(
+            scales_ptr + g[:, None] * stride_sg + offs_n[None, :] * stride_sn,
             mask=k_mask[:, None] & n_mask[None, :],
             other=0.0,
         ).to(tl.float32)
-        z_block = tl.load(
-            zeros_ptr
-            + g_idx[:, None] * stride_zg
-            + offs_n[None, :] * stride_zn,
+        z_tile = tl.load(
+            zeros_ptr + g[:, None] * stride_zg + offs_n[None, :] * stride_zn,
             mask=k_mask[:, None] & n_mask[None, :],
             other=0.0,
         ).to(tl.float32)
 
-        w_int = w_int.to(tl.float32)
-        w_deq = (w_int - z_block) * s_block  # [BK, BN] dequantized weight
+        w_deq = (w_int.to(tl.float32) - z_tile) * s_tile
+        acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
 
-        # ---- matmul tile (BM, BK) @ (BK, BN) -> (BM, BN) ----
-        acc = tl.dot(x_block.to(tl.float32), w_deq, acc)
-
-    # ---- write result (M, N) bf16 ----
-    out_ptrs = (
-        out_ptr
-        + offs_m[:, None] * stride_om
-        + offs_n[None, :] * stride_on
+    tl.store(
+        out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on,
+        acc.to(tl.bfloat16),
+        mask=out_mask,
     )
-    tl.store(out_ptrs, acc.to(tl.bfloat16), mask=out_mask)
 
 
 # ---------------------------------------------------------------------------
@@ -138,59 +104,48 @@
 # ---------------------------------------------------------------------------
 
 class Model(nn.Module):
-    """W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""
-
     def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
         super().__init__()
-        assert K % group_size == 0, "K must be divisible by group_size"
-        assert K % 2 == 0, "K must be even (int4 packing)"
+ 

... (3225 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -16,7 +16,6 @@
 import torch.nn as nn
 import triton
 import triton.language as tl
-
 
 GROUP_SIZE = 128
 
@@ -57,14 +56,14 @@
         k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
         k_mask = k_offs < K
 
-        # x tile: (BM, BK)
+        # x tile: (BM, BK) bf16
         x_tile = tl.load(
             x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
-        # w_q tile: (BK, BN) — load packed bytes
+        # w_q tile: (BK, BN) uint8 — packed
         wq_row = k_offs // 2
         wq_tile = tl.load(
             wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
@@ -72,19 +71,29 @@
             other=0,
         )
 
-        # unpack nibbles: even-k → low, odd-k → high
+        # unpack nibbles
         is_odd = (k_offs % 2)[:, None]
         w_int = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # scales / zeros: (BK, BN) via group index
+        # scales / zeros broadcast per group — load (1, BN) and broadcast
+        g_first = k_base // GROUP_SIZE_CONST
+        g_last = (k_base + BLOCK_SIZE_K - 1) // GROUP_SIZE_CONST  # inclusive check
+
+        # If this k range spans at most one group we can just load once
+        # For BK <= GROUP_SIZE this is always true when k_base is GROUP_SIZE-aligned,
+        # but our k_base increments by BLOCK_SIZE_K which may not be GROUP_SIZE-aligned.
+        # Safe approach: load per-k, but to save shared mem we'll use register broadcast
+        # since BN is small enough to fit in registers.
         g = k_offs // GROUP_SIZE_CONST
+        s_base = scales_ptr + g[:, None] * stride_sg + offs_n[None, :] * stride_sn
+        z_base = zeros_ptr + g[:, None] * stride_zg + offs_n[None, :] * stride_zn
         s_tile = tl.load(
-            scales_ptr + g[:, None] * stride_sg + offs_n[None, :] * stride_sn,
+            s_base,
             mask=k_mask[:, None] & n_mask[None, :],
             other=0.0,
         ).to(tl.float32)
         z_tile = tl.load(
-            zeros_ptr + g[:, None] * stride_zg + offs_n[None, :] * stride_zn,
+            z_base,
             mask=k_mask[:, None] & n_mask[None, :],
             other=0.0,
         ).to(tl.float32)
@@ -135,15 +144,16 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Small shared-memory configs per shape regime
-        if M == 1:
-            BM, BN, BK, warps, stages = 1, 128, 128, 4, 2
-        elif M <= 16:
-            BM, BN, BK, warps, stages = 16, 128, 128, 4, 2
-        elif M <= 64:
-            BM, BN, BK, warps, stages = 64, 64, 64, 4, 2
-        else:
-            BM, BN, BK, warps, stages = 128, 64, 64, 4, 2
+        # Shared memory budget ~101KB on SM120.
+        # Per-stage memory: x(BM,BK)*2 + wq(BK,BN)*1 + s(BK,BN)*2 + z(BK,BN)*2
+        # With stages=1 (no pipelining) and num_warps=4:
+        # Need to fit within 101KB total.
+
+        BM = 1  # M is small in all test shapes
+        BN = 64
+        BK = 64
+        warps = 4
+        stages = 1
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -40,6 +40,11 @@
     BLOCK_SIZE_K: tl.constexpr,
     GROUP_SIZE_CONST: tl.constexpr,
 ):
+    """Fused int4 dequant + GEMM.
+
+    BK must equal GROUP_SIZE_CONST so each K-tile belongs to exactly one group.
+    This scales/zeros are loaded once per K-tile instead of per-k-row.
+    """
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
@@ -56,14 +61,14 @@
         k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
         k_mask = k_offs < K
 
-        # x tile: (BM, BK) bf16
+        # x tile (BM, BK) bf16
         x_tile = tl.load(
             x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
-        # w_q tile: (BK, BN) uint8 — packed
+        # w_q tile (BK, BN) uint8 packed
         wq_row = k_offs // 2
         wq_tile = tl.load(
             wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
@@ -71,34 +76,25 @@
             other=0,
         )
 
-        # unpack nibbles
+        # unpack: even-K rows from low nibble, odd-K from high nibble
         is_odd = (k_offs % 2)[:, None]
-        w_int = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
+        w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # scales / zeros broadcast per group — load (1, BN) and broadcast
-        g_first = k_base // GROUP_SIZE_CONST
-        g_last = (k_base + BLOCK_SIZE_K - 1) // GROUP_SIZE_CONST  # inclusive check
-
-        # If this k range spans at most one group we can just load once
-        # For BK <= GROUP_SIZE this is always true when k_base is GROUP_SIZE-aligned,
-        # but our k_base increments by BLOCK_SIZE_K which may not be GROUP_SIZE-aligned.
-        # Safe approach: load per-k, but to save shared mem we'll use register broadcast
-        # since BN is small enough to fit in registers.
-        g = k_offs // GROUP_SIZE_CONST
-        s_base = scales_ptr + g[:, None] * stride_sg + offs_n[None, :] * stride_sn
-        z_base = zeros_ptr + g[:, None] * stride_zg + offs_n[None, :] * stride_zn
-        s_tile = tl.load(
-            s_base,
-            mask=k_mask[:, None] & n_mask[None, :],
+        # scales/zeros: ONE row per group (BK == GROUP_SIZE)
+        g = k_base // GROUP_SIZE_CONST
+        s_row = tl.load(
+            scales_ptr + g * stride_sg + offs_n * stride_sn,
+            mask=n_mask,
             other=0.0,
         ).to(tl.float32)
-        z_tile = tl.load(
-            z_base,
-            mask=k_mask[:, None] & n_mask[None, :],
+        z_row = tl.load(
+            zeros_ptr + g * stride_zg + offs_n * stride_zn,
+            mask=n_mask,
             other=0.0,
         ).to(tl.float32)
 
-        w_deq = (w_int.to(tl.float32) - z_tile) * s_tile
+        # dequant with broadcast: (1, BN) x (BK, BN) → (BK, BN)
+        w_deq = (w_int_4.to(tl.float32) - z_row[None, :]) * s_row[None, :]
         acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
 
     tl.store(
@@ -144,16 +140,26 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Shared memory budget ~101KB on SM120.
-        # Per-stage memory: x(BM,BK)*2 + wq(BK,BN)*1 + s(BK,BN)*2 + z(BK,BN)*2
-        # With stages=1 (no pipelining) and num_warps=4:
-        # Need to fit within 101KB total.
+        # BK=GROUP_SIZE=128 so each K-tile maps to exactly one group.
+        # scales/zeros loaded once per K-tile instead of per-row.
+        # Shared mem per stage: x(BM,128)*2 + wq(128,BN)*1 + s(1,BN)*2 + z(1,BN)*2
+        #   = BM*256 + BN*132
+        # 2 stages: 2*(BM*256 + BN*132) must fit in 101KB.
 
-        BM = 1  # M is small in all test shapes
-        BN = 64
-        BK = 64
-        warps = 4
-        stages = 1
+        # Shape-aware block configs
+        if M == 1:
+            # Decode: maximize N to amortize grid overhead
+            BM, BN, BK, warps, stages = 1, 256, 128, 4, 2
+        elif M <= 4:
+            BM, BN, BK, warps, stages = 2, 256, 128, 4, 2
+        elif M <= 16:
+            BM, BN, BK, warps, stages = 16, 128, 128, 4, 2
+        elif M <= 32:
+            BM, BN, BK, warps, stages = 32, 128, 128, 4, 2
+        elif M <= 128:
+            BM, BN, BK, warps, stages = 64, 128, 128, 4, 2
+        else:
+            BM, BN, BK, warps, stages = 128, 128, 128, 8, 2
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -146,20 +146,20 @@
         #   = BM*256 + BN*132
         # 2 stages: 2*(BM*256 + BN*132) must fit in 101KB.
 
-        # Shape-aware block configs
+        # Shape-aware block configs — constrained by 101KB shared mem.
+        # BK=GROUP_SIZE=128 so scale/zeros fit in 1 row per tile.
         if M == 1:
-            # Decode: maximize N to amortize grid overhead
-            BM, BN, BK, warps, stages = 1, 256, 128, 4, 2
+            BM, BN, BK, warps, stages = 1, 512, 128, 4, 1
         elif M <= 4:
-            BM, BN, BK, warps, stages = 2, 256, 128, 4, 2
+            BM, BN, BK, warps, stages = 4, 256, 128, 4, 1
         elif M <= 16:
-            BM, BN, BK, warps, stages = 16, 128, 128, 4, 2
+            BM, BN, BK, warps, stages = 16, 128, 128, 4, 1
         elif M <= 32:
-            BM, BN, BK, warps, stages = 32, 128, 128, 4, 2
+            BM, BN, BK, warps, stages = 32, 128, 128, 4, 1
         elif M <= 128:
-            BM, BN, BK, warps, stages = 64, 128, 128, 4, 2
+            BM, BN, BK, warps, stages = 64, 128, 128, 4, 1
         else:
-            BM, BN, BK, warps, stages = 128, 128, 128, 8, 2
+            BM, BN, BK, warps, stages = 128, 64, 128, 8, 1
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
 

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/test_configs.py
@@ -0,0 +1,155 @@
+"""Test different Triton configs for W4A16 GEMM."""
+import torch
+import triton
+import triton.language as tl
+
+GROUP_SIZE = 128
+
+
+@triton.jit
+def _w4a16_gemm_kernel_v2(
+    x_ptr, wq_ptr, scales_ptr, zeros_ptr, out_ptr,
+    M, N, K,
+    stride_xm, stride_xk,
+    stride_wqk, stride_wqn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    stride_om, stride_on,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_CONST: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+
+    m_mask = offs_m < M
+    n_mask = offs_n < N
+    out_mask = m_mask[:, None] & n_mask[None, :]
+
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    for k_base in range(0, K, BLOCK_SIZE_K):
+        k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
+        k_mask = k_offs < K
+
+        x_tile = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
+            mask=m_mask[:, None] & k_mask[None, :],
+            other=0.0,
+        )
+
+        wq_row = k_offs // 2
+        wq_tile = tl.load(
+            wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
+            mask=k_mask[:, None] & n_mask[None, :],
+            other=0,
+        )
+
+        is_odd = (k_offs % 2)[:, None]
+        w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
+
+        g = k_base // GROUP_SIZE_CONST
+        s_row = tl.load(
+            scales_ptr + g * stride_sg + offs_n * stride_sn,
+            mask=n_mask,
+            other=0.0,
+        ).to(tl.float32)
+        z_row = tl.load(
+            zeros_ptr + g * stride_zg + offs_n * stride_zn,
+            mask=n_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        w_deq = (w_int_4.to(tl.float32) - z_row[None, :]) * s_row[None, :]
+        acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
+
+    tl.store(
+        out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on,
+        acc.to(tl.bfloat16),
+        mask=out_mask,
+    )
+
+
+def test_config(M, N, K, BM, BN, BK, warps, stages, label=""):
+    """Test a specific config. Returns True on success."""
+    class Model(torch.nn.Module):
+        def __init__(self, M, N, K):
+            super().__init__()
+            self.M, self.N, self.K = M, N, K
+            n_groups = K // GROUP_SIZE
+            torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
+            w_full = torch.randn(K, N, dtype=torch.float32) * 0.02
+            w_g = w_full.view(n_groups, GROUP_SIZE, N)
+            w_min = w_g.min(dim=1, keepdim=True).values
+            w_max = w_g.max(dim=1, keepdim=True).values
+            scales = (w_max - w_min).clamp_min(1e-8) / 15.0
+            zeros = (-w_min / scales).round().clamp(0, 15)
+            w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
+            lo = w_q[0::2].to(torch.uint8) & 0xF
+            hi = w_q[1::2].to(torch.uint8) & 0xF
+            w_packed = (lo | (hi << 4)).contiguous()
+            self.register_buffer("w_q", w_packed)
+            self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16))
+            self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16))
+
+    model = Model(M, N, K).cuda()
+    x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
+    grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
+
+    try:
+        _w4a16_gemm_kernel_v2[grid](
+            x.contiguous(), model.w_q, model.scales, model.zeros, out,
+            M, N, K,
+            x.stride(0), x.stride(1),
+            model.w_q.stride(0), model.w_q.stride(1),
+            model.scales.stride(0), model.scales.stride(1),
+            model.zeros.stride(0), model.zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, BLOCK_SIZE_K=BK,
+            GROUP_SIZE_CONST=GROUP_SIZE,
+            num_warps=warps, num_stages=stages,
+        )
+        print(f"{label} M={M} BM={BM} BN={BN} BK={BK} w={warps} s={stages}: OK")
+        return True
+    except Exception as e:
+        print(f"{label} M={M} BM={BM} BN={BN} BK={BK} w={warps} s={stages}: FAIL - {e}")
+        return False
+
+
+if __name__ == "__main__":
+    shapes = [
+        (1, 12288, 4096),
+        (32, 12288, 4096),
+        (256, 12288, 4096),
+        (1, 4096, 4096),
+        (16, 14336, 4096),
+    ]
+
+    for M, N, K in shapes:
+        print(f"\n=== Testing shape ({M}, {N}, {K}) ===")
+        if M == 1:
+            # Try many configs for decode
+            for BN in [64, 128, 256]:
+                for stages in [1]:
+                    for warps in [1, 2, 4]:
+                        test_config(M, N, K, 1, BN, 128, warps, stages, "M1")
+        elif M <= 4:
+            test_config(M, N, K, M, 256, 128, 4, 1, f"M{M}")
+        elif M <= 16:
+            test_config(M, N, K, M, 128, 128, 4, 1, f"M{M}")
+            test_config(M, N, K, 16, 128, 128, 4, 1, f"M{M}a")
+        elif M <= 32:
+            test_config(M, N, K, M, 128, 128, 4, 1, f"M{M}")
+            test_config(M, N, K, 32, 64, 128, 4, 1, f"M{M}a")
+        elif M <= 128:
+            test_config(M, N, K, 32, 128, 128, 4, 1, f"M{M}a")
+            test_config(M, N, K, 64, 128, 128, 4, 1, f"M{M}b")
+            test_config(M, N, K, 128, 64, 128, 4, 1, f"M{M}c")
+        else:
+            test_config(M, N, K, 64, 128, 128, 4, 1, f"M{M}a")
+            test_config(M, N, K, 128, 64, 128, 4, 1, f"M{M}b")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -149,7 +149,7 @@
         # Shape-aware block configs — constrained by 101KB shared mem.
         # BK=GROUP_SIZE=128 so scale/zeros fit in 1 row per tile.
         if M == 1:
-            BM, BN, BK, warps, stages = 1, 512, 128, 4, 1
+            BM, BN, BK, warps, stages = 1, 128, 128, 4, 1
         elif M <= 4:
             BM, BN, BK, warps, stages = 4, 256, 128, 4, 1
         elif M <= 16:

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -40,11 +40,6 @@
     BLOCK_SIZE_K: tl.constexpr,
     GROUP_SIZE_CONST: tl.constexpr,
 ):
-    """Fused int4 dequant + GEMM.
-
-    BK must equal GROUP_SIZE_CONST so each K-tile belongs to exactly one group.
-    This scales/zeros are loaded once per K-tile instead of per-k-row.
-    """
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
 
@@ -76,11 +71,11 @@
             other=0,
         )
 
-        # unpack: even-K rows from low nibble, odd-K from high nibble
+        # unpack nibbles
         is_odd = (k_offs % 2)[:, None]
         w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # scales/zeros: ONE row per group (BK == GROUP_SIZE)
+        # scales/zeros: ONE row per group
         g = k_base // GROUP_SIZE_CONST
         s_row = tl.load(
             scales_ptr + g * stride_sg + offs_n * stride_sn,
@@ -93,7 +88,6 @@
             other=0.0,
         ).to(tl.float32)
 
-        # dequant with broadcast: (1, BN) x (BK, BN) → (BK, BN)
         w_deq = (w_int_4.to(tl.float32) - z_row[None, :]) * s_row[None, :]
         acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
 
@@ -102,6 +96,79 @@
         acc.to(tl.bfloat16),
         mask=out_mask,
     )
+
+
+# ---------------------------------------------------------------------------
+# Column-based GEMV kernel for M=1 (decode path)
+# Each CTA computes one output column. Warps share the reduction over K.
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def _w4a16_gemv_kernel(
+    x_ptr,
+    wq_ptr,
+    scales_ptr,
+    zeros_ptr,
+    out_ptr,
+    K,
+    N,
+    stride_xk,
+    stride_wqk, stride_wqn,
+    stride_sg, stride_sn,
+    stride_zg, stride_zn,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_WARPS: tl.constexpr,
+    GROUP_SIZE_CONST: tl.constexpr,
+):
+    col_idx = tl.program_id(0)
+    warpid = tl.program_id(1)
+    num_warps_in_grid = tl.num_programs(1)
+
+    # Each CTA has num_warps_in_grid blocks, each block processes a chunk of K
+    chunk_size = BLOCK_SIZE_K * NUM_WARPS
+
+    lane_id = tl.arange(0, BLOCK_SIZE_K)
+    base_k = warpid * BLOCK_SIZE_K
+
+    partial = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32)
+
+    for k_base in range(0, K, chunk_size):
+        k_offs = k_base + base_k + lane_id
+        k_mask = k_offs < K
+
+        # x vector
+        x_val = tl.load(
+            x_ptr + k_offs * stride_xk,
+            mask=k_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        # packed weight byte
+        wq_row = k_offs // 2
+        wq_byte = tl.load(
+            wq_ptr + wq_row * stride_wqk + col_idx * stride_wqn,
+            mask=k_mask,
+            other=0,
+        ).to(tl.int32)
+
+        is_odd = k_offs % 2
+        w_int = tl.where(is_odd == 0, wq_byte & 0xF, wq_byte >> 4)
+
+        # group index — load scale and zero once per chunk (since BK == GROUP_SIZE)
+        g = (k_base + base_k) // GROUP_SIZE_CONST
+        s = tl.load(scales_ptr + g * stride_sg + col_idx * stride_sn).to(tl.float32)
+        z = tl.load(zeros_ptr + g * stride_zg + col_idx * stride_zn).to(tl.float32)
+
+        w_deq = (w_int.to(tl.float32) - z) * s
+        partial = tl.where(k_mask, partial + x_val * w_deq, partial)
+
+    # Warp-level reduction using tl.reduce
+    if BLOCK_SIZE_K > 1:
+        acc = tl.sum(partial)
+    else:
+        acc = partial
+
+    tl.store(out_ptr + col_idx, acc.to(tl.bfloat16))
 
 
 # ---------------------------------------------------------------------------
@@ -140,48 +207,63 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # BK=GROUP_SIZE=128 so each K-tile maps to exactly one group.
-        # scales/zeros loaded once per K-tile instead of per-row.
-        # Shared mem per stage: x(BM,128)*2 + wq(128,BN)*1 + s(1,BN)*2 + z(1,BN)*2
-        #   = BM*256 + BN*132
-        # 2 stages: 2*(BM*256 + BN*132) must fit in 101KB.
-
-        # Shape-aware block configs — constrained by 101KB shared mem.
-        # BK=GROUP_SIZE=128 so scale/zeros fit in 1 row per tile.
         if M == 1:
-            BM, BN, BK, warps, stages = 1, 128, 128, 4, 1
-        elif M <= 4:
-            BM, BN, BK, warps, stages = 4, 256, 128, 4, 1
-        elif M <= 16:
-            BM, BN, BK, warps, stages = 16, 128, 128, 4, 1
-        elif M <= 32:
-            BM, BN, BK, warps, stages = 32, 128, 128, 4, 1
-        elif M <= 128:
-            BM, BN, BK, warps, stages = 64, 128, 128, 4, 1
+            # Column-based GEMV: one CTA per column, shared-mem free
+            # Each CTA processes K threads across multiple warps
+            BK = 128
+            warps_per_col = 8  # total threads per column = 128 * 8 = 1024
+            num_ctas = N * warps_per_col
+            grid = (N, warps_per_col)
+
+            _w4a16_gemv_kernel[grid](
+                x.contiguous(),
+                self.w_q,
+                self.scales,
+                self.zeros,
+                out,
+                K, N,
+                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),
+                BLOCK_SIZE_K=BK,
+                NUM_WARPS=warps_per_col,
+                GROUP_SIZE_CONST=GROUP_SIZE,
+                num_warps=1,
+                num_stages=1,
+            )
         else:
-            BM, BN, BK, warps, stages = 128, 64, 128, 8, 1
-
-        grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
-
-        _w4a16_gemm_kernel[grid](
-            x.contiguous(),
-            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),
-   

... (1598 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -1,14 +1,6 @@
-"""High-performance W4A16 weight-only quantized GEMM for SM120 Blackwell.
+"""W4A16 weight-only quantized GEMM for SM120 Blackwell.
 
-AWQ/GPTQ-style asymmetric int4 with per-group bf16 scales and zero-points.
 Fuses unpack, dequant, and matmul in a single Triton kernel.
-
-Data layout:
-  x:      (M, K)                  bf16
-  w_q:    (K // 2, N)             uint8  (low nibble=even-K, high=odd-K)
-  scales: (K // 128, N)           bf16
-  zeros:  (K // 128, N)           bf16
-  out:    (M, N)                  bf16
 """
 from __future__ import annotations
 
@@ -56,14 +48,12 @@
         k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
         k_mask = k_offs < K
 
-        # x tile (BM, BK) bf16
         x_tile = tl.load(
             x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
-        # w_q tile (BK, BN) uint8 packed
         wq_row = k_offs // 2
         wq_tile = tl.load(
             wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
@@ -71,11 +61,9 @@
             other=0,
         )
 
-        # unpack nibbles
         is_odd = (k_offs % 2)[:, None]
         w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # scales/zeros: ONE row per group
         g = k_base // GROUP_SIZE_CONST
         s_row = tl.load(
             scales_ptr + g * stride_sg + offs_n * stride_sn,
@@ -97,83 +85,6 @@
         mask=out_mask,
     )
 
-
-# ---------------------------------------------------------------------------
-# Column-based GEMV kernel for M=1 (decode path)
-# Each CTA computes one output column. Warps share the reduction over K.
-# ---------------------------------------------------------------------------
-
-@triton.jit
-def _w4a16_gemv_kernel(
-    x_ptr,
-    wq_ptr,
-    scales_ptr,
-    zeros_ptr,
-    out_ptr,
-    K,
-    N,
-    stride_xk,
-    stride_wqk, stride_wqn,
-    stride_sg, stride_sn,
-    stride_zg, stride_zn,
-    BLOCK_SIZE_K: tl.constexpr,
-    NUM_WARPS: tl.constexpr,
-    GROUP_SIZE_CONST: tl.constexpr,
-):
-    col_idx = tl.program_id(0)
-    warpid = tl.program_id(1)
-    num_warps_in_grid = tl.num_programs(1)
-
-    # Each CTA has num_warps_in_grid blocks, each block processes a chunk of K
-    chunk_size = BLOCK_SIZE_K * NUM_WARPS
-
-    lane_id = tl.arange(0, BLOCK_SIZE_K)
-    base_k = warpid * BLOCK_SIZE_K
-
-    partial = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32)
-
-    for k_base in range(0, K, chunk_size):
-        k_offs = k_base + base_k + lane_id
-        k_mask = k_offs < K
-
-        # x vector
-        x_val = tl.load(
-            x_ptr + k_offs * stride_xk,
-            mask=k_mask,
-            other=0.0,
-        ).to(tl.float32)
-
-        # packed weight byte
-        wq_row = k_offs // 2
-        wq_byte = tl.load(
-            wq_ptr + wq_row * stride_wqk + col_idx * stride_wqn,
-            mask=k_mask,
-            other=0,
-        ).to(tl.int32)
-
-        is_odd = k_offs % 2
-        w_int = tl.where(is_odd == 0, wq_byte & 0xF, wq_byte >> 4)
-
-        # group index — load scale and zero once per chunk (since BK == GROUP_SIZE)
-        g = (k_base + base_k) // GROUP_SIZE_CONST
-        s = tl.load(scales_ptr + g * stride_sg + col_idx * stride_sn).to(tl.float32)
-        z = tl.load(zeros_ptr + g * stride_zg + col_idx * stride_zn).to(tl.float32)
-
-        w_deq = (w_int.to(tl.float32) - z) * s
-        partial = tl.where(k_mask, partial + x_val * w_deq, partial)
-
-    # Warp-level reduction using tl.reduce
-    if BLOCK_SIZE_K > 1:
-        acc = tl.sum(partial)
-    else:
-        acc = partial
-
-    tl.store(out_ptr + col_idx, acc.to(tl.bfloat16))
-
-
-# ---------------------------------------------------------------------------
-# Model
-# ---------------------------------------------------------------------------
 
 class Model(nn.Module):
     def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
@@ -207,63 +118,38 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        if M == 1:
-            # Column-based GEMV: one CTA per column, shared-mem free
-            # Each CTA processes K threads across multiple warps
+        if M <= 16:
+            BM = min(M, 16)
+            BN = 256
             BK = 128
-            warps_per_col = 8  # total threads per column = 128 * 8 = 1024
-            num_ctas = N * warps_per_col
-            grid = (N, warps_per_col)
+            warps = 4
+        else:
+            BM = min(M, 128)
+            BN = 128
+            BK = 128
+            warps = 4
 
-            _w4a16_gemv_kernel[grid](
-                x.contiguous(),
-                self.w_q,
-                self.scales,
-                self.zeros,
-                out,
-                K, N,
-                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),
-                BLOCK_SIZE_K=BK,
-                NUM_WARPS=warps_per_col,
-                GROUP_SIZE_CONST=GROUP_SIZE,
-                num_warps=1,
-                num_stages=1,
-            )
-        else:
-            # Block configs — BK=GROUP_SIZE=128 so scale/zeros fit in 1 row per tile.
-            if M <= 4:
-                BM, BN, BK, warps, stages = 4, 256, 128, 4, 1
-            elif M <= 16:
-                BM, BN, BK, warps, stages = 16, 128, 128, 4, 1
-            elif M <= 32:
-                BM, BN, BK, warps, stages = 32, 128, 128, 4, 1
-            else:
-                BM, BN, BK, warps, stages = 64, 128, 128, 4, 1
+        grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
 
-            grid = (triton.c

... (1397 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -118,16 +118,14 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        if M <= 16:
-            BM = min(M, 16)
-            BN = 256
-            BK = 128
-            warps = 4
+        # BN=128 is safe for 101KB shared mem with stages=1
+        if M <= 4:
+            BM, BN = min(M, 4), 128
+        elif M <= 16:
+            BM, BN = min(M, 16), 128
         else:
-            BM = min(M, 128)
-            BN = 128
-            BK = 128
-            warps = 4
+            BM, BN = min(M, 64), 128
+        BK = 128
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
 
@@ -147,7 +145,7 @@
             BLOCK_SIZE_N=BN,
             BLOCK_SIZE_K=BK,
             GROUP_SIZE_CONST=GROUP_SIZE,
-            num_warps=warps,
+            num_warps=4,
             num_stages=1,
         )
         return out

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/profile_configs.py
@@ -0,0 +1,72 @@
+"""Profile Triton configs for W4A16 GEMM."""
+import torch
+import time
+import solution
+
+
+def benchmark_fn(model, x, n_iter=50):
+    for _ in range(5):
+        model(x)
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(n_iter):
+        model(x)
+    torch.cuda.synchronize()
+    return (time.perf_counter() - t0) / n_iter * 1000
+
+
+SHAPES = [
+    (1, 12288, 4096),
+    (32, 12288, 4096),
+    (256, 12288, 4096),
+    (1, 4096, 4096),
+    (16, 14336, 4096),
+]
+
+configs = [
+    (1, 64, 128),
+    (1, 128, 128),
+    (4, 64, 128),
+    (4, 128, 128),
+    (4, 256, 128),
+    (16, 64, 128),
+    (16, 128, 128),
+    (16, 256, 128),
+    (32, 64, 128),
+    (32, 128, 128),
+    (64, 64, 128),
+    (64, 128, 128),
+    (128, 64, 128),
+    (128, 128, 128),
+]
+
+for M, N, K in SHAPES:
+    print(f"\n=== Shape ({M}, {N}, {K}) ===")
+    model = solution.Model(M, N, K).cuda()
+    x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+
+    for BM, BN, BK in configs:
+        if BM > M and M > 0:
+            continue
+        try:
+            out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
+            grid = (solution.triton.cdiv(M, BM), solution.triton.cdiv(N, BN))
+            solution._w4a16_gemm_kernel[grid](
+                x.contiguous(), model.w_q, model.scales, model.zeros, out,
+                M, N, K,
+                x.stride(0), x.stride(1), model.w_q.stride(0), model.w_q.stride(1),
+                model.scales.stride(0), model.scales.stride(1), model.zeros.stride(0), model.zeros.stride(1),
+                out.stride(0), out.stride(1),
+                BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, BLOCK_SIZE_K=BK,
+                GROUP_SIZE_CONST=128, num_warps=4, num_stages=1,
+            )
+            ms = benchmark_fn(
+                type('M', (), {'x': out, 'model': model, 'BM': BM, 'BN': BN, 'BK': BK,
+                               '_w4a16_gemm_kernel': solution._w4a16_gemm_kernel,
+                               'triton': solution.triton, 'M': M, 'N': N, 'K': K,
+                               'w_q': model.w_q, 'scales': model.scales, 'zeros': model.zeros})(),
+                x, n_iter=50,
+            )
+            print(f"  BM={BM:3d} BN={BN:3d} BK={BK}: {ms:.3f} ms")
+        except Exception as e:
+            print(f"  BM={BM:3d} BN={BN:3d} BK={BK}: FAIL ({type(e).__name__})")

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/profile_configs.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/profile_configs.py
@@ -1,16 +1,94 @@
 """Profile Triton configs for W4A16 GEMM."""
 import torch
 import time
-import solution
+import triton
+import triton.language as tl
+
+GROUP_SIZE = 128
 
 
-def benchmark_fn(model, x, n_iter=50):
-    for _ in range(5):
-        model(x)
+# Import the reference kernel from test_configs
+# (we'll redefine it here to avoid import issues)
+@triton.jit
+def _w4a16_gemm_kernel(
+    x_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
+    M, N, K,
+    sxm, sxk, swqk, swqn, ssg, ssn, szg, szn, som, son,
+    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GS: tl.constexpr,
+):
+    pm = tl.program_id(0)
+    pn = tl.program_id(1)
+    om = pm * BM + tl.arange(0, BM)
+    on = pn * BN + tl.arange(0, BN)
+    mm = om < M
+    mn = on < N
+    outmask = mm[:, None] & mn[None, :]
+    acc = tl.zeros((BM, BN), dtype=tl.float32)
+    for kb in range(0, K, BK):
+        ko = kb + tl.arange(0, BK)
+        km = ko < K
+        xt = tl.load(x_ptr + om[:, None]*sxm + ko[None,:]*sxk, mask=mm[:,None]&km[None,:], other=0.0)
+        wqr = ko // 2
+        wqt = tl.load(wq_ptr + wqr[:,None]*swqk + on[None,:]*swqn, mask=km[:,None]&mn[None,:], other=0)
+        iodd = (ko % 2)[:, None]
+        wi4 = tl.where(iodd == 0, wqt & 0xF, wqt >> 4).to(tl.int32)
+        g = kb // GS
+        sr = tl.load(s_ptr + g*ssg + on*ssn, mask=mn, other=0.0).to(tl.float32)
+        zr = tl.load(z_ptr + g*szg + on*szn, mask=mn, other=0.0).to(tl.float32)
+        wd = (wi4.to(tl.float32) - zr[None,:]) * sr[None,:]
+        acc = tl.dot(xt.to(tl.float32), wd, acc)
+    tl.store(out_ptr + om[:,None]*som + on[None,:]*son, acc.to(tl.bfloat16), mask=outmask)
+
+
+def make_model(M, N, K):
+    class Model(torch.nn.Module):
+        def __init__(self, M, N, K):
+            super().__init__()
+            self.M, self.N, self.K = M, N, K
+            n_groups = K // GROUP_SIZE
+            torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
+            w_full = torch.randn(K, N, dtype=torch.float32) * 0.02
+            w_g = w_full.view(n_groups, GROUP_SIZE, N)
+            w_min = w_g.min(dim=1, keepdim=True).values
+            w_max = w_g.max(dim=1, keepdim=True).values
+            scales = (w_max - w_min).clamp_min(1e-8) / 15.0
+            zeros = (-w_min / scales).round().clamp(0, 15)
+            w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
+            lo = w_q[0::2].to(torch.uint8) & 0xF
+            hi = w_q[1::2].to(torch.uint8) & 0xF
+            w_packed = (lo | (hi << 4)).contiguous()
+            self.register_buffer("w_q", w_packed)
+            self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16))
+            self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16))
+    return Model(M, N, K).cuda()
+
+
+def benchmark(model, x, BM, BN, BK, n_iter=50):
+    M, N, K = x.shape[0], model.N, x.shape[1]
+    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
+    grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
+
+    def fwd():
+        _w4a16_gemm_kernel[grid](
+            x_c, model.w_q, model.scales, model.zeros, out,
+            M, N, K,
+            x.stride(0), x.stride(1), model.w_q.stride(0), model.w_q.stride(1),
+            model.scales.stride(0), model.scales.stride(1), model.zeros.stride(0), model.zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BM=BM, BN=BN, BK=BK, GS=GROUP_SIZE,
+            num_warps=4, num_stages=1,
+        )
+        return out
+
+    global x_c
+    x_c = x.contiguous()
+
+    # warmup
+    fwd()
     torch.cuda.synchronize()
     t0 = time.perf_counter()
     for _ in range(n_iter):
-        model(x)
+        fwd()
     torch.cuda.synchronize()
     return (time.perf_counter() - t0) / n_iter * 1000
 
@@ -24,49 +102,24 @@
 ]
 
 configs = [
-    (1, 64, 128),
-    (1, 128, 128),
-    (4, 64, 128),
-    (4, 128, 128),
-    (4, 256, 128),
-    (16, 64, 128),
-    (16, 128, 128),
-    (16, 256, 128),
-    (32, 64, 128),
-    (32, 128, 128),
-    (64, 64, 128),
-    (64, 128, 128),
-    (128, 64, 128),
-    (128, 128, 128),
+    (1, 64, 128), (1, 128, 128),
+    (4, 64, 128), (4, 128, 128), (4, 256, 128),
+    (16, 64, 128), (16, 128, 128), (16, 256, 128),
+    (32, 64, 128), (32, 128, 128),
+    (64, 64, 128), (64, 128, 128),
+    (128, 64, 128), (128, 128, 128),
 ]
 
 for M, N, K in SHAPES:
     print(f"\n=== Shape ({M}, {N}, {K}) ===")
-    model = solution.Model(M, N, K).cuda()
+    model = make_model(M, N, K)
     x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
 
     for BM, BN, BK in configs:
-        if BM > M and M > 0:
+        if BM > M:
             continue
         try:
-            out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
-            grid = (solution.triton.cdiv(M, BM), solution.triton.cdiv(N, BN))
-            solution._w4a16_gemm_kernel[grid](
-                x.contiguous(), model.w_q, model.scales, model.zeros, out,
-                M, N, K,
-                x.stride(0), x.stride(1), model.w_q.stride(0), model.w_q.stride(1),
-                model.scales.stride(0), model.scales.stride(1), model.zeros.stride(0), model.zeros.stride(1),
-                out.stride(0), out.stride(1),
-                BLOCK_SIZE_M=BM, BLOCK_SIZE_N=BN, BLOCK_SIZE_K=BK,
-                GROUP_SIZE_CONST=128, num_warps=4, num_stages=1,
-            )
-            ms = benchmark_fn(
-                type('M', (), {'x': out, 'model': model, 'BM': BM, 'BN': BN, 'BK': BK,
-                               '_w4a16_gemm_kernel': solution._w4a16_gemm_kernel,
-                               'triton': solution.triton, 'M': M, 'N': N, 'K': K,
-                               'w_q': model.w_q, 'scales': model.scales, 'zeros': model.zeros})(),
-                x, n_iter=50,
-            )
+   

... (239 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -118,13 +118,17 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # BN=128 is safe for 101KB shared mem with stages=1
-        if M <= 4:
+        # Profiled optimal configs for each shape class
+        if M == 1:
+            BM, BN = 1, 64
+        elif M <= 4:
             BM, BN = min(M, 4), 128
         elif M <= 16:
             BM, BN = min(M, 16), 128
+        elif M <= 64:
+            BM, BN = min(M, 64), 64
         else:
-            BM, BN = min(M, 64), 128
+            BM, BN = min(M, 128), 64
         BK = 128
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))

--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/profile_configs2.py
@@ -0,0 +1,113 @@
+"""Profile Triton configs with warps and stages variations."""
+import torch
+import time
+import triton
+import triton.language as tl
+
+GROUP_SIZE = 128
+
+
+@triton.jit
+def _w4a16_gemm_kernel(
+    x_ptr, wq_ptr, s_ptr, z_ptr, out_ptr,
+    M, N, K,
+    sxm, sxk, swqk, swqn, ssg, ssn, szg, szn, som, son,
+    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GS: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    offs_m = pid_m * BM + tl.arange(0, BM)
+    offs_n = pid_n * BN + tl.arange(0, BN)
+    m_mask = offs_m < M
+    n_mask = offs_n < N
+    out_mask = m_mask[:, None] & n_mask[None, :]
+    acc = tl.zeros((BM, BN), dtype=tl.float32)
+    for kb in range(0, K, BK):
+        ko = kb + tl.arange(0, BK)
+        km = ko < K
+        xt = tl.load(x_ptr + offs_m[:, None]*sxm + ko[None,:]*sxk, mask=m_mask[:, None]&km[None, :], other=0.0)
+        wqr = ko // 2
+        wqt = tl.load(wq_ptr + wqr[:, None]*swqk + offs_n[None, :]*swqn, mask=km[:, None]&n_mask[None, :], other=0)
+        iodd = (ko % 2)[:, None]
+        wi4 = tl.where(iodd == 0, wqt & 0xF, wqt >> 4).to(tl.int32)
+        g = kb // GS
+        sr = tl.load(s_ptr + g*ssg + offs_n*ssn, mask=n_mask, other=0.0).to(tl.float32)
+        zr = tl.load(z_ptr + g*szg + offs_n*szn, mask=n_mask, other=0.0).to(tl.float32)
+        wd = (wi4.to(tl.float32) - zr[None,:]) * sr[None,:]
+        acc = tl.dot(xt.to(tl.float32), wd, acc)
+    tl.store(out_ptr + offs_m[:, None]*som + offs_n[None, :]*son, acc.to(tl.bfloat16), mask=out_mask)
+
+
+def make_model(M, N, K):
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.N = N
+            n_groups = K // GROUP_SIZE
+            torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
+            w_full = torch.randn(K, N, dtype=torch.float32) * 0.02
+            w_g = w_full.view(n_groups, GROUP_SIZE, N)
+            w_min = w_g.min(dim=1, keepdim=True).values
+            w_max = w_g.max(dim=1, keepdim=True).values
+            scales = (w_max - w_min).clamp_min(1e-8) / 15.0
+            zeros = (-w_min / scales).round().clamp(0, 15)
+            w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
+            lo = w_q[0::2].to(torch.uint8) & 0xF
+            hi = w_q[1::2].to(torch.uint8) & 0xF
+            w_packed = (lo | (hi << 4)).contiguous()
+            self.register_buffer("w_q", w_packed)
+            self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16))
+            self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16))
+    return Model().cuda()
+
+
+def benchmark(model, x, BM, BN, BK, warps, stages, n_iter=50):
+    M, N, K = x.shape[0], model.N, x.shape[1]
+    out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
+    grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
+    x_c = x.contiguous()
+
+    def fwd():
+        _w4a16_gemm_kernel[grid](
+            x_c, model.w_q, model.scales, model.zeros, out,
+            M, N, K, x.stride(0), x.stride(1), model.w_q.stride(0), model.w_q.stride(1),
+            model.scales.stride(0), model.scales.stride(1), model.zeros.stride(0), model.zeros.stride(1),
+            out.stride(0), out.stride(1),
+            BM=BM, BN=BN, BK=BK, GS=GROUP_SIZE,
+            num_warps=warps, num_stages=stages,
+        )
+        return out
+
+    fwd()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(n_iter):
+        fwd()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - t0) / n_iter * 1000
+
+
+SHAPES = [
+    (256, 12288, 4096),
+    (32, 12288, 4096),
+]
+
+# Test varying warps and stages for shapes with larger M
+for M, N, K in SHAPES:
+    print(f"\n=== Shape ({M}, {N}, {K}) ===")
+    model = make_model(M, N, K)
+    x = torch.randn(M, K, dtype=torch.bfloat16).cuda()
+
+    for BM, BN, BK in [
+        (64, 64, 128), (64, 128, 128), (128, 64, 128), (128, 128, 128),
+        (32, 64, 128), (32, 128, 128), (M, 64, 128),
+    ]:
+        if BM > M:
+            continue
+        for warps in [4, 8]:
+            for stages in [1, 2, 3, 4]:
+                try:
+                    ms = benchmark(model, x, BM, BN, BK, warps, stages, n_iter=50)
+                    print(f"  BM={BM:3d} BN={BN:3d} BK={BK} w={warps} s={stages}: {ms:.3f} ms")
+                except Exception:
+                    pass  # silently skip OOM configs

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -118,13 +118,13 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Profiled optimal configs for each shape class
+        # Profiled optimal configs
         if M == 1:
             BM, BN = 1, 64
-        elif M <= 4:
-            BM, BN = min(M, 4), 128
+        elif M == 4:
+            BM, BN = 4, 64
         elif M <= 16:
-            BM, BN = min(M, 16), 128
+            BM, BN = min(M, 16), 64
         elif M <= 64:
             BM, BN = min(M, 64), 64
         else:

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -1,6 +1,7 @@
 """W4A16 weight-only quantized GEMM for SM120 Blackwell.
 
 Fuses unpack, dequant, and matmul in a single Triton kernel.
+Uses make_block_ptr API for efficient vectorized memory access.
 """
 from __future__ import annotations
 
@@ -48,12 +49,14 @@
         k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
         k_mask = k_offs < K
 
+        # Load x tile
         x_tile = tl.load(
             x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
+        # Load packed w_q tile
         wq_row = k_offs // 2
         wq_tile = tl.load(
             wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
@@ -61,9 +64,11 @@
             other=0,
         )
 
+        # Unpack nibbles
         is_odd = (k_offs % 2)[:, None]
         w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
+        # Load scales and zeros (one row per group)
         g = k_base // GROUP_SIZE_CONST
         s_row = tl.load(
             scales_ptr + g * stride_sg + offs_n * stride_sn,
@@ -76,6 +81,7 @@
             other=0.0,
         ).to(tl.float32)
 
+        # Dequantize
         w_deq = (w_int_4.to(tl.float32) - z_row[None, :]) * s_row[None, :]
         acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
 
@@ -84,6 +90,40 @@
         acc.to(tl.bfloat16),
         mask=out_mask,
     )
+
+
+def _compute_block_config(M, N, K, shared_mem_limit=101376):
+    """Compute optimal block config given shared memory constraint.
+    
+    Shared memory per stage (bytes):
+      x tile: BM * BK * 2 (bf16)
+      wq tile: BK * BN * 1 (uint8)
+      scales/zeros: BN * 2 * 2 (bf16, loaded once per K-tile)
+    
+    With stages=1, total = BM*BK*2 + BK*BN + 4*BN
+    
+    Strategy: maximize computation density (BM*BN) within memory budget.
+    For memory-bound (small M), maximize BN. For compute-bound (large M), maximize BM.
+    """
+    BK = GROUP_SIZE  # Always use group size for efficient scale loading
+    
+    if M == 1:
+        # Pure memory-bound: one row, maximize N coverage
+        return 1, 256
+    elif M <= 4:
+        return min(M, 4), 256
+    elif M <= 16:
+        return min(M, 16), 256
+    elif M <= 32:
+        # Mixed regime
+        return min(M, 32), 256
+    elif M <= 64:
+        # Approaching compute-bound
+        return min(M, 64), 256
+    elif M <= 128:
+        return min(M, 128), 256
+    else:
+        return 128, 256
 
 
 class Model(nn.Module):
@@ -118,20 +158,32 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Profiled optimal configs
-        if M == 1:
-            BM, BN = 1, 64
-        elif M == 4:
-            BM, BN = 4, 64
-        elif M <= 16:
-            BM, BN = min(M, 16), 64
-        elif M <= 64:
-            BM, BN = min(M, 64), 64
-        else:
-            BM, BN = min(M, 128), 64
-        BK = 128
+        # Use optimized block config
+        BM, BN = _compute_block_config(M, N, K)
+        BK = GROUP_SIZE
+
+        # Handle cases where large BN causes shared memory overflow
+        # Check if config would fit in shared memory
+        shared_per_stage = BM * BK * 2 + BK * BN + 4 * BN
+        
+        # If too large, reduce BN
+        if shared_per_stage > 80000:  # Leave some headroom
+            BN = 128
+            shared_per_stage = BM * BK * 2 + BK * BN + 4 * BN
+            # Still too large, reduce BM
+            if shared_per_stage > 80000:
+                BM = min(BM, 16)
+                BN = 64
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
+        
+        # Use 4 warps, 2 stages for best performance
+        num_warps = 4 if BM * BN <= 4096 else 8
+        num_stages = 2 if shared_per_stage * 2 < 90000 else 1
+        
+        # Override for small M where stages=2 doesn't help
+        if M <= 4:
+            num_stages = 1
 
         _w4a16_gemm_kernel[grid](
             x.contiguous(),
@@ -149,8 +201,8 @@
             BLOCK_SIZE_N=BN,
             BLOCK_SIZE_K=BK,
             GROUP_SIZE_CONST=GROUP_SIZE,
-            num_warps=4,
-            num_stages=1,
+            num_warps=num_warps,
+            num_stages=num_stages,
         )
         return out
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -1,7 +1,6 @@
 """W4A16 weight-only quantized GEMM for SM120 Blackwell.
 
 Fuses unpack, dequant, and matmul in a single Triton kernel.
-Uses make_block_ptr API for efficient vectorized memory access.
 """
 from __future__ import annotations
 
@@ -49,14 +48,12 @@
         k_offs = k_base + tl.arange(0, BLOCK_SIZE_K)
         k_mask = k_offs < K
 
-        # Load x tile
         x_tile = tl.load(
             x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk,
             mask=m_mask[:, None] & k_mask[None, :],
             other=0.0,
         )
 
-        # Load packed w_q tile
         wq_row = k_offs // 2
         wq_tile = tl.load(
             wq_ptr + wq_row[:, None] * stride_wqk + offs_n[None, :] * stride_wqn,
@@ -64,11 +61,9 @@
             other=0,
         )
 
-        # Unpack nibbles
         is_odd = (k_offs % 2)[:, None]
         w_int_4 = tl.where(is_odd == 0, wq_tile & 0xF, wq_tile >> 4).to(tl.int32)
 
-        # Load scales and zeros (one row per group)
         g = k_base // GROUP_SIZE_CONST
         s_row = tl.load(
             scales_ptr + g * stride_sg + offs_n * stride_sn,
@@ -81,7 +76,6 @@
             other=0.0,
         ).to(tl.float32)
 
-        # Dequantize
         w_deq = (w_int_4.to(tl.float32) - z_row[None, :]) * s_row[None, :]
         acc = tl.dot(x_tile.to(tl.float32), w_deq, acc)
 
@@ -90,40 +84,6 @@
         acc.to(tl.bfloat16),
         mask=out_mask,
     )
-
-
-def _compute_block_config(M, N, K, shared_mem_limit=101376):
-    """Compute optimal block config given shared memory constraint.
-    
-    Shared memory per stage (bytes):
-      x tile: BM * BK * 2 (bf16)
-      wq tile: BK * BN * 1 (uint8)
-      scales/zeros: BN * 2 * 2 (bf16, loaded once per K-tile)
-    
-    With stages=1, total = BM*BK*2 + BK*BN + 4*BN
-    
-    Strategy: maximize computation density (BM*BN) within memory budget.
-    For memory-bound (small M), maximize BN. For compute-bound (large M), maximize BM.
-    """
-    BK = GROUP_SIZE  # Always use group size for efficient scale loading
-    
-    if M == 1:
-        # Pure memory-bound: one row, maximize N coverage
-        return 1, 256
-    elif M <= 4:
-        return min(M, 4), 256
-    elif M <= 16:
-        return min(M, 16), 256
-    elif M <= 32:
-        # Mixed regime
-        return min(M, 32), 256
-    elif M <= 64:
-        # Approaching compute-bound
-        return min(M, 64), 256
-    elif M <= 128:
-        return min(M, 128), 256
-    else:
-        return 128, 256
 
 
 class Model(nn.Module):
@@ -158,32 +118,22 @@
         N = self.N
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
-        # Use optimized block config
-        BM, BN = _compute_block_config(M, N, K)
-        BK = GROUP_SIZE
-
-        # Handle cases where large BN causes shared memory overflow
-        # Check if config would fit in shared memory
-        shared_per_stage = BM * BK * 2 + BK * BN + 4 * BN
-        
-        # If too large, reduce BN
-        if shared_per_stage > 80000:  # Leave some headroom
-            BN = 128
-            shared_per_stage = BM * BK * 2 + BK * BN + 4 * BN
-            # Still too large, reduce BM
-            if shared_per_stage > 80000:
-                BM = min(BM, 16)
-                BN = 64
+        # Profiled optimal configs — all with stages=1, warps=4, BK=128
+        if M == 1:
+            BM, BN = 1, 64
+        elif M <= 4:
+            BM, BN = min(M, 4), 128
+        elif M <= 16:
+            BM, BN = min(M, 16), 128
+        elif M <= 32:
+            BM, BN = 32, 128
+        elif M <= 64:
+            BM, BN = min(M, 64), 128
+        else:
+            BM, BN = 64, 128
+        BK = 128
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
-        
-        # Use 4 warps, 2 stages for best performance
-        num_warps = 4 if BM * BN <= 4096 else 8
-        num_stages = 2 if shared_per_stage * 2 < 90000 else 1
-        
-        # Override for small M where stages=2 doesn't help
-        if M <= 4:
-            num_stages = 1
 
         _w4a16_gemm_kernel[grid](
             x.contiguous(),
@@ -201,8 +151,8 @@
             BLOCK_SIZE_N=BN,
             BLOCK_SIZE_K=BK,
             GROUP_SIZE_CONST=GROUP_SIZE,
-            num_warps=num_warps,
-            num_stages=num_stages,
+            num_warps=4,
+            num_stages=1,
         )
         return out
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/solution.py
@@ -119,18 +119,19 @@
         out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device)
 
         # Profiled optimal configs — all with stages=1, warps=4, BK=128
+        # BN=64 beats BN=128 for most shapes (better register usage / occupancy)
         if M == 1:
             BM, BN = 1, 64
         elif M <= 4:
-            BM, BN = min(M, 4), 128
+            BM, BN = min(M, 4), 64
         elif M <= 16:
-            BM, BN = min(M, 16), 128
+            BM, BN = min(M, 16), 64
         elif M <= 32:
-            BM, BN = 32, 128
+            BM, BN = 32, 64
         elif M <= 64:
-            BM, BN = min(M, 64), 128
+            BM, BN = min(M, 64), 64
         else:
-            BM, BN = 64, 128
+            BM, BN = 64, 64
         BK = 128
 
         grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))