content: 'Read reference.py, check.py, benchmark.py, shapes.py, problem.yaml to understand the problem'
status: 'in_progress'
priority: 'high'
content: 'Implement fused W4A16 GEMM kernel in solution.py'
status: 'pending'
priority: 'high'
content: 'Verify correctness with check.py'
status: 'pending'
priority: 'high'
content: 'Profile and benchmark, iterate for performance'
status: 'pending'
priority: 'high'

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm</path>
<type>directory</type>
<entries>
benchmark.py
check.py
problem.yaml
PROMPT.txt
reference.py
shapes.py
sota.py

(7 entries)
</entries>

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/reference.py</path>
<type>file</type>
<content>
1: """Naive W4A16 weight-only quantized GEMM reference (correctness only).
2: 
3: AWQ/GPTQ-style scheme:
4:   x:      (M, K)               bf16
5:   w_q:    (K // 2, N)          uint8   -- two int4 weights packed per byte (low nibble = even-K, high = odd-K)
6:   scales: (K // group, N)      bf16
7:   zeros:  (K // group, N)      bf16    -- asymmetric (stored already as float zero-point)
8:   out:    (M, N)                bf16
9: 
10: Dequant (per group along K):
11:   w_bf[k, n] = (w_q[k, n] - zeros[k // group, n]) * scales[k // group, n]
12: where w_q[k, n] is the unpacked 4-bit value (0..15).
13: 
14: This reference unpacks to a full bf16 matrix and then runs torch.matmul. Slow and
15: memory-heavy on the dequant; the agent's solution must fuse unpack+GEMM.
16: """
17: from __future__ import annotations
18: 
19: import torch
20: import torch.nn as nn
21: 
22: OP_TYPE = "gemm_w4a16"
23: SUPPORTED_PRECISIONS = ["int4_bf16"]
24: HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
25: 
26: GROUP_SIZE = 128
27: 
28: 
29: def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
30:     """Pack (K, N) uint8 in [0,15] into (K//2, N) uint8.
31: 
32:     Even rows go in the low nibble, odd rows in the high nibble.
33:     """
34:     K, N = w_q.shape
35:     assert K % 2 == 0
36:     lo = w_q[0::2].to(torch.uint8) & 0xF
37:     hi = w_q[1::2].to(torch.uint8) & 0xF
38:     ret

... (3312 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/check.py</path>
<type>file</type>
<content>
1: """Correctness runner for FP8 GEMM.
2: 
3: Runs solution.Model vs reference.Model across all shapes in shapes.py, 3 seeds
4: each, with per-dtype atol/rtol. Also rejects forbidden ops by grep.
5: """
6: import re
7: import sys
8: from pathlib import Path
9: 
10: import torch
11: import yaml
12: 
13: # Make the repo's src/ importable
14: REPO_ROOT = Path(__file__).resolve().parents[2]
15: sys.path.insert(0, str(REPO_ROOT))
16: 
17: from src.eval.correctness import check_correctness  # noqa: E402
18: 
19: 
20: def main():
21:     try:
22:         import reference
23:         import shapes
24:         import solution
25:     except Exception as e:
26:         print(f"FAIL: import error: {e}")
27:         sys.exit(1)
28: 
29:     problem_yaml = Path("problem.yaml")
30:     meta = yaml.safe_load(problem_yaml.read_text()) if problem_yaml.exists() else {}
31: 
32:     # --- Forbidden-op check ------------------------------------------------
33:     sol_src = Path("solution.py").read_text() if Path("solution.py").exists() else ""
34:     for forbidden in meta.get("forbidden", []):
35:         pat = re.escape(forbidden)
36:         if re.search(pat, sol_src):
37:             print(f"FAIL: forbidden op used: {forbidden}")
38:             sys.exit(1)
39: 
40:     device = torch.device("cuda:0")
41:     tol_override = meta.get("tolerance") or None
42: 
43:     # --- Per-shape

... (2909 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/benchmark.py</path>
<type>file</type>
<content>
1: """Roofline benchmark for FP8 GEMM.
2: 
3: For each shape: times eager reference, compiled reference, SOTA (if available),
4: and the agent's solution. Reports achieved TFLOPS, GB/s, and peak_fraction.
5: 
6: Output lines the harness picks up:
7:   shape=<idx> variant=<name> tflops=<N> gbps=<N> ms=<N>
8:   peak_fraction: <N>  (geomean over shapes of solution's peak_fraction)
9: """
10: import sys
11: from math import exp, log
12: from pathlib import Path
13: 
14: import torch
15: import yaml
16: 
17: REPO_ROOT = Path(__file__).resolve().parents[2]
18: sys.path.insert(0, str(REPO_ROOT))
19: 
20: from src.eval.roofline import compute_gbps, compute_tflops, peak_fraction  # noqa: E402
21: from src.eval.timing import time_fn  # noqa: E402
22: from src.hardware import get as get_hw  # noqa: E402
23: 
24: 
25: def _eval_formula(expr: str, vars: dict) -> float:
26:     # Very small eval: only names from `vars` are valid.
27:     return float(eval(expr, {"__builtins__": {}}, vars))
28: 
29: 
30: def main():
31:     import reference
32:     import shapes
33:     import solution
34: 
35:     meta = yaml.safe_load(Path("problem.yaml").read_text())
36:     hw = get_hw(meta["hardware"][0])
37:     peak_tflops = hw.peak_tflops_dense.get(meta["peak_tflops_key"], 0.0)
38:     peak_gbps = hw.peak_bandwidth_gb_s
39:     regime = meta.get("regime", "compute")
40:     flops_formul

... (3415 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/problem.yaml</path>
<type>file</type>
<content>
1: name: 07_w4a16_gemm
2: display_name: "W4A16 Weight-only Quantized GEMM"
3: precision: int4_bf16
4: regime: memory  # decode-dominant; M=1 is bandwidth-bound on the int4 weight stream
5: 
6: # Dense-equivalent FLOPs (matmul work, ignoring dequant arithmetic).
7: flops_formula: "2 * M * N * K"
8: 
9: # Bytes moved per call (memory roofline):
10: #   x:      M*K*2          (bf16 activations, streamed in once)
11: #   w_q:    (K/2)*N        (packed int4, 0.5 B/elem)
12: #   scales: (K/128)*N*2    (bf16 scales)
13: #   zeros:  (K/128)*N*2    (bf16 zero-points)
14: #   out:    M*N*2          (bf16 store)
15: bytes_formula: "M*K*2 + (K/2)*N + (K/128)*N*2 + (K/128)*N*2 + M*N*2"
16: 
17: hardware: [RTX_PRO_6000]
18: peak_tflops_key: bf16
19: peak_bandwidth_key: dram
20: 
21: tolerance:
22:   bfloat16: 0.10  # group-quant adds noise on top of bf16 accumulator slop
23: 
24: # Forbidden ops -- agent must write the unpack + GEMM themselves, not call a
25: # vendor library that does both.
26: forbidden:
27:   - "bitsandbytes.functional.dequantize_4bit"
28:   - "bitsandbytes.functional.gemv_4bit"
29:   - "marlin_kernel.gemm"
30:   - "torch.nn.functional.linear"
31: 
32: sota:
33:   name: "bitsandbytes NF4 (gemv_4bit / dequantize_4bit + matmul)"
34:   url: "https://github.com/TimDettmers/bitsandbytes"
35:   function: "bitsandbytes.functional.gemv_4bit"
36:   notes: |
37:     

... (757 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/shapes.py</path>
<type>file</type>
<content>
1: """Shape sweep for W4A16 GEMM.
2: 
3: Llama-style up_proj / qkv_proj shapes. Decode (M=1) is the bandwidth-bound
4: case every inference engine optimizes -- it's the bar to beat.
5: """
6: 
7: SHAPES = [
8:     {"M": 1,   "N": 12288, "K": 4096},   # decode: memory-bound on int4 weight read
9:     {"M": 32,  "N": 12288, "K": 4096},   # small prefill: mixed regime
10:     {"M": 256, "N": 12288, "K": 4096},   # larger prefill: approaching compute
11:     {"M": 1,   "N": 4096,  "K": 4096},   # decode: square shape
12:     {"M": 16,  "N": 14336, "K": 4096},   # speculative-decode-ish
13: ]

(End of file - total 13 lines)
</content>

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/sota.py</path>
<type>file</type>
<content>
1: """SOTA reference for W4A16 GEMM.
2: 
3: Library survey on RTX PRO 6000 Blackwell (SM120, CC 12.0):
4: 
5:   - Marlin (IST-DASLab):         no SM120 kernels (Ampere/Hopper only). Skip.
6:   - GPTQ-Triton (fpgaminer):     unmaintained; pure Triton path works on SM120
7:                                  but is not faster than Marlin on its target HW
8:                                  and has no Blackwell tuning. Skip as primary.
9:   - AWQ (mit-han-lab/llm-awq):   CUDA kernels not built for SM120 in the wheel.
10:                                  Skip.
11:   - bitsandbytes >= 0.49.2:      CUDA kernels compile and run on SM120 (verified
12:                                  on this machine). Different quant scheme (NF4,
13:                                  symmetric, blocksize 64) than our reference's
14:                                  AWQ-style asymmetric INT4 with group_size 128,
15:                                  but it occupies the same memory regime and is
16:                                  the only tuned W4A16-class kernel that runs on
17:                                  SM120 today. Used here as an *informational*
18:                                  SOTA line, not as a numerical reference.
19: 
20: The benchmark calls `sota_forward(x, ref_model)` and times it; correctness is
21: NOT checked against this path (the quant scheme differs).
22: """
23: from 

... (2725 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/07_w4a16_gemm/PROMPT.txt</path>
<type>file</type>
<content>
1: I need you to write a weight-only int4 quantized GEMM (W4A16) for the RTX PRO 6000 (SM120 Blackwell, GDDR7, 1.8 TB/s). The reference implementation is in reference.py and your kernel goes in solution.py. You can make whatever mess you want in this directory — scratch files, intermediate kernels, build artifacts, profiling traces — but the final answer has to be in solution.py with the same Model, get_inputs, and get_init_inputs interface as reference.py.
2: 
3: The scheme is AWQ/GPTQ-style asymmetric int4 with explicit zero-points and per-group bf16 scales. Inputs are x of shape (M, K) bf16, w_q of shape (K // 2, N) uint8 (two int4 weights packed per byte, low nibble = even-K row, high nibble = odd-K row), scales of shape (K // 128, N) bf16, and zeros of shape (K // 128, N) bf16. Group size is 128 along K. Dequant per group is w_bf[k, n] = (unpack(w_q)[k, n] - zeros[k // 128, n]) * scales[k // 128, n], and the output is (M, N) bf16. Correctness tolerance is 0.10 abs/rel — group-quant adds noise on top of bf16 accumulator slop. The shapes you have to handle are M=1 N=12288 K=4096 (decode, memory-bound on the int4 weight read), M=32 N=12288 K=4096 (small prefill, mixed regime), M=256 N=12288 K=4096 (larger prefill, approaching compute-bound), M=1 N=4096 K=4096 (decode, square), and M=16 N=14336 K=4096 (speculative-decode-ish).
4: 
5: This needs to be a real custom

... (1082 more chars)

NVIDIA RTX PRO 6000 Blackwell Workstation Edition, 12.0, 97887 MiB

3.6.0

2.11.0+cu130
(12, 0)

__init__.py
__pycache__
correctness.py
report.py
roofline.py
shapes.py
timing.py

3.6.0
tl.dot available

<path>/home/infatoshi/cuda/KernelBench-Hard/src/hardware</path>
<type>directory</type>
<entries>
__init__.py
__pycache__/
m4_max.py
rtx_pro_6000.py

(4 entries)
</entries>

<path>/home/infatoshi/cuda/KernelBench-Hard/src/eval/correctness.py</path>
<type>file</type>
<content>
1: """Per-dtype correctness tolerance.
2: 
3: Stanford's KernelBench uses 1e-4 atol/rtol for fp32 and 1e-2 for fp16/bf16.
4: We extend to fp8 (0.1, permissive) and int (exact).
5: """
6: from __future__ import annotations
7: 
8: import torch
9: 
10: DEFAULT_TOLERANCE = {
11:     torch.float32: {"atol": 1e-4, "rtol": 1e-4},
12:     torch.float16: {"atol": 1e-2, "rtol": 1e-2},
13:     torch.bfloat16: {"atol": 1e-2, "rtol": 1e-2},
14:     torch.float8_e4m3fn: {"atol": 1e-1, "rtol": 1e-1},
15:     torch.float8_e5m2: {"atol": 1e-1, "rtol": 1e-1},
16:     torch.int8: {"atol": 0, "rtol": 0},
17:     torch.int32: {"atol": 0, "rtol": 0},
18:     torch.int64: {"atol": 0, "rtol": 0},
19: }
20: 
21: 
22: def _coerce_float(x):
23:     """Accept int/float/str scientific-notation as float; pass through dicts."""
24:     if isinstance(x, (int, float)):
25:         return float(x)
26:     if isinstance(x, str):
27:         try:
28:             return float(x)
29:         except ValueError:
30:             return x
31:     return x
32: 
33: 
34: def tolerance_for_dtype(dtype: torch.dtype, override: dict | None = None) -> dict:
35:     """Lookup atol/rtol for a given dtype, with optional per-problem override.
36: 
37:     The override dict is keyed by str(dtype). Values may be:
38:       - a single number (or a string parseable as float, e.g. "5e-2") -> used
39:         as both atol and rtol
4

... (2580 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/src/eval/roofline.py</path>
<type>file</type>
<content>
1: """Roofline math: achieved TFLOPS / GB/s, peak fraction.
2: 
3: FLOPS and bytes formulas come from the problem's `problem.yaml`. They are the
4: dense-equivalent algorithmic work; kernels that skip work (sparsity, early
5: exit) are still measured against the dense cost.
6: """
7: from __future__ import annotations
8: 
9: 
10: def compute_tflops(flops: float, time_ms: float) -> float:
11:     """flops per invocation and wall time in ms -> TFLOPS."""
12:     if time_ms <= 0:
13:         return 0.0
14:     return flops / (time_ms * 1e-3) / 1e12
15: 
16: 
17: def compute_gbps(bytes_moved: float, time_ms: float) -> float:
18:     """bytes per invocation and wall time in ms -> GB/s."""
19:     if time_ms <= 0:
20:         return 0.0
21:     return bytes_moved / (time_ms * 1e-3) / 1e9
22: 
23: 
24: def peak_fraction(
25:     achieved: float,
26:     peak: float,
27: ) -> float:
28:     """Return achieved / peak, clamped to [0, infinity)."""
29:     if peak <= 0:
30:         return 0.0
31:     return max(0.0, achieved / peak)

(End of file - total 31 lines)
</content>

<path>/home/infatoshi/cuda/KernelBench-Hard/src/hardware/rtx_pro_6000.py</path>
<type>file</type>
<content>
1: """RTX PRO 6000 Blackwell Workstation — SM120, consumer-lineage Blackwell.
2: 
3: Peak tensor-core throughputs are dense-matrix advertised peaks. Actual kernels
4: will see 60-85% of peak on well-tuned code.
5: """
6: from dataclasses import dataclass
7: 
8: 
9: @dataclass(frozen=True)
10: class HardwareTarget:
11:     name: str
12:     sm: str
13:     vram_gb: int
14:     peak_bandwidth_gb_s: float  # DRAM
15:     peak_tflops_dense: dict[str, float]  # dtype -> TFLOPS
16: 
17: 
18: RTX_PRO_6000 = HardwareTarget(
19:     name="RTX PRO 6000 Blackwell Workstation",
20:     sm="sm_120a",
21:     vram_gb=96,
22:     peak_bandwidth_gb_s=1800.0,
23:     peak_tflops_dense={
24:         "fp4": 800.0,
25:         "nvfp4": 800.0,
26:         "mxfp4": 800.0,
27:         "fp6": 800.0,
28:         "fp8": 400.0,
29:         "bf16": 200.0,
30:         "fp16": 200.0,
31:         "tf32": 100.0,
32:         "fp32": 12.0,  # non-tensor-core SIMT fp32
33:         "int8": 400.0,
34:         "int4": 800.0,
35:     },
36: )

(End of file - total 36 lines)
</content>