1	"""Naive top-k reference: torch.topk over the last dim.
2	
3	This is the correctness oracle. The agent's solution must produce the same
4	top-k values (and equivalent indices modulo ties) within the tolerance
5	declared in problem.yaml. Note that solution.py is FORBIDDEN from calling
6	torch.topk / torch.sort / torch.kthvalue (see problem.yaml).
7	"""
8	import torch
9	import torch.nn as nn
10	
11	OP_TYPE = "topk"
12	SUPPORTED_PRECISIONS = ["fp32"]
13	HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
14	
15	
16	class Model(nn.Module):
17	    """Top-k over the last dim of a 2D tensor.
18	
19	    Input:
20	        x: (batch, n) fp32
21	    Output:
22	        values:  (batch, k) fp32, sorted descending
23	        indices: (batch, k) int64, into the last dim of x
24	    """
25	
26	    def __init__(self, batch: int, n: int, k: int):
27	        super().__init__()
28	        self.batch, self.n, self.k = batch, n, k
29	        # No learned parameters, but declare a dummy buffer so state_dict
30	        # is non-empty and load_state_dict(strict=True) is meaningful.
31	        self.register_buffer("_dummy", torch.zeros(1))
32	
33	    def forward(self, x: torch.Tensor):
34	        values, indices = torch.topk(x, k=self.k, dim=-1, largest=True, sorted=True)
35	        return values, indices
36	
37	
38	# Module-level shims rebuilt by check.py / benchmark.py per shape.
39	batch = 64
40	n = 8192
41	k = 8
42	
43	
44	def get_inputs():
45	    # fp32 input drawn from a roughly Gaussian dist

... (211 more chars)

total 44
drwxrwxr-x 2 1002 1002 4096 Jun 10 23:47 .
drwxrwxr-x 3 1002 1002 4096 Jun 10 23:47 ..
-rw-rw-r-- 1 1002 1002 2084 Apr 27 09:21 PROMPT.txt
-rw-rw-r-- 1 1002 1002 5078 Jun  1 21:05 benchmark.py
-rw-rw-r-- 1 1002 1002 6288 Jun  2 01:09 check.py
-rw-rw-r-- 1 1002 1002 2199 Apr 26 04:20 problem.yaml
-rw-rw-r-- 1 1002 1002 1561 Apr 26 04:19 reference.py
-rw-rw-r-- 1 1002 1002  933 Apr 26 04:19 shapes.py
-rw-rw-r-- 1 1002 1002  945 Apr 26 04:20 sota.py

1	"""Correctness runner for TopK.
2	
3	Runs solution.Model vs reference.Model across all shapes in shapes.py, 3 seeds
4	each. Top-k correctness has two parts:
5	
6	  1. VALUES: sol_values must match ref_values within fp32 tol. Both are
7	     returned sorted descending, so positional comparison is well-defined.
8	  2. INDICES: lenient — we do NOT require sol_indices == ref_indices because
9	     ties in x can yield multiple valid index sets. Instead we gather x at
10	     sol_indices and check those values match ref_values within tol. This
11	     catches "wrong indices" without false-failing on legitimate tie-breaks.
12	
13	Also rejects forbidden ops by grep.
14	"""
15	import re
16	import sys
17	from pathlib import Path
18	
19	import torch
20	import yaml
21	
22	REPO_ROOT = Path(__file__).resolve().parents[2]
23	sys.path.insert(0, str(REPO_ROOT))
24	
25	from src.eval.correctness import check_correctness  # noqa: E402
26	from src.eval.numeric_stress import (  # noqa: E402
27	    numeric_stress_cases,
28	    numeric_stress_context,
29	    tolerance_for_case,
30	)
31	
32	
33	def main():
34	    try:
35	        import reference
36	        import shapes
37	        import solution
38	    except Exception as e:
39	        print(f"FAIL: import error: {e}")
40	        sys.exit(1)
41	
42	    problem_yaml = Path("problem.yaml")
43	    meta = yaml.safe_load(problem_yaml.read_text()) if problem_yaml.exists() else {}
44	
45	    # --- Forbidden-op check --------------------------------------

... (5310 more chars)

1	"""Roofline benchmark for TopK.
2	
3	For each shape: times eager reference (torch.topk), compiled reference, SOTA
4	(also torch.topk — see sota.py), and the agent's solution. Reports achieved
5	TFLOPS, GB/s, and peak_fraction (vs DRAM bandwidth, since this is memory-bound).
6	
7	Output lines the harness picks up:
8	  shape=<idx> variant=<name> tflops=<N> gbps=<N> ms=<N>
9	  peak_fraction: <N>  (geomean over shapes of solution's peak_fraction)
10	"""
11	import sys
12	from math import exp, log
13	from pathlib import Path
14	
15	import torch
16	import yaml
17	
18	REPO_ROOT = Path(__file__).resolve().parents[2]
19	sys.path.insert(0, str(REPO_ROOT))
20	
21	from src.eval.roofline import compute_gbps, compute_tflops, peak_fraction  # noqa: E402
22	from src.eval.timing import benchmark_baselines_enabled, time_variant  # noqa: E402
23	from src.hardware import get as get_hw  # noqa: E402
24	
25	
26	def _eval_formula(expr: str, vars: dict) -> float:
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", "memory")
40	    flops_formula = meta["flops_formula"]
41	    bytes_formula = meta["bytes_formula"]
42	    num_perf_trials = int(meta.

... (4052 more chars)

1	"""Canonical shape sweep for TopK.
2	
3	Mix of:
4	  - decoder vocab top-k (single sequence, very large n, moderate k) — pure
5	    bandwidth test; the input read dominates everything.
6	  - prefill / batched attention top-k (many rows, moderate n, small k) — tests
7	    per-row parallelism and shared-memory bitonic networks.
8	  - non-power-of-2 n stress case — bitonic sort networks naturally want
9	    powers of two; this forces the agent to handle padding or partial sorts.
10	  - small-k limit — k=1 (argmax) is a degenerate but useful sanity case.
11	"""
12	
13	SHAPES = [
14	    {"batch": 1,   "n": 131072, "k": 64},   # decoder vocab top-k (Llama vocab ~128k)
15	    {"batch": 64,  "n": 8192,   "k": 8},    # prefill / attention top-k
16	    {"batch": 32,  "n": 16384,  "k": 32},   # mid-size batched
17	    {"batch": 16,  "n": 12000,  "k": 16},   # non-power-of-2 n stress
18	    {"batch": 128, "n": 4096,   "k": 1},    # batched argmax (k=1 corner case)
19	]
20	

1	name: 05_topk_bitonic
2	display_name: "TopK via Bitonic Sort"
3	precision: fp32
4	regime: memory
5	
6	# Top-k is dominated by the input read (small output, no reduction over k).
7	# Comparator-network FLOPs are not the bottleneck on real hardware, so we
8	# track them but score on bandwidth.
9	flops_formula: "batch * n * 4"            # ~O(n log^2 n) compares total, but the
10	                                          # bitonic network is so cheap relative
11	                                          # to memory that a coarse 4*n estimate
12	                                          # is fine for telemetry only.
13	bytes_formula: "batch * n * 4 + batch * k * (4 + 8)"  # fp32 input read + (fp32 value + int64 idx) output
14	
15	hardware: [RTX_PRO_6000]
16	peak_tflops_key: fp32
17	peak_bandwidth_key: dram
18	
19	# Top-k correctness:
20	#   - VALUES must match within fp32 atol/rtol (the kth largest value is
21	#     well-defined modulo float-equal ties, so we use a loose-ish tol).
22	#   - INDICES are checked leniently: for each row, the multiset of returned
23	#     indices must select values that match ref values within tol. Direct
24	#     index equality is NOT required (ties in x can yield different valid
25	#     index sets).
26	tolerance:
27	  float32: 1.0e-4
28	
29	# Forbidden ops — using any of these in solution.py fails correctness post-hoc.
30	# This problem is about IMPLEMENTING the selection, not dispatching to PyTorch's
31	# tuned top-k. torch.sort is also banned b

... (859 more chars)

1	"""SOTA reference for TopK: torch.topk itself.
2	
3	torch.topk dispatches to a CUB-backed kernel that uses radix-select for
4	moderate k and a tuned bitonic sort for small n. It is the bar the agent's
5	hand-rolled bitonic kernel must beat. There is no obvious vendor library that
6	does better on the (batch, n, k) shape mix we evaluate — Faiss BlockSelect is
7	specialized for k>=32 with much larger n, and CUB's DeviceSegmentedRadixSort
8	sorts the full row (overkill for top-k).
9	
10	This file is INTENTIONALLY allowed to call torch.topk because it is the SOTA
11	oracle, not the agent's submission. The agent's solution.py is forbidden from
12	using torch.topk (see problem.yaml.forbidden).
13	"""
14	from __future__ import annotations
15	
16	import torch
17	
18	
19	def sota_forward(x: torch.Tensor, k: int):
20	    """Best-available top-k reference. x: (batch, n) fp32."""
21	    return torch.topk(x, k=k, dim=-1, largest=True, sorted=True)
22	
23	
24	def is_available() -> bool:
25	    return True
26	

problems
pyproject.toml
src
uv.lock
__init__.py
__pycache__
correctness.py
numeric_stress.py
report.py
roofline.py
shapes.py
timing.py

1	"""Centralized GPU timing for benchmark.py files.
2	
3	All problems' benchmark.py should call `time_fn` from here so we have one
4	implementation of warmup, L2 flush, and CUDA event capture to maintain.
5	
6	Methodology:
7	  - 10 warmup calls absorb Triton autotune (typical ~7 configs) and
8	    torch.compile reduce-overhead CUDA-graph capture.
9	  - Between each timed call, `_l2_flush()` writes 128 MB to evict L2
10	    (Blackwell L2 is 96 MB; we want any prior L2 contents flushed so
11	    we measure HBM-load bandwidth, not L2-cached bandwidth).
12	  - GPU timing via cuda.Event with synchronize() AFTER record() but
13	    BEFORE elapsed_time() — the canonical NVIDIA pattern.
14	  - Reported value is the median of `iters` trials, robust to outliers.
15	
16	Notes / known biases not addressed here:
17	  - torch.compile(mode="reduce-overhead") gets CUDA graphs which eliminate
18	    launch overhead. Custom Triton/CUDA kernels do NOT get this treatment.
19	    On small shapes where launch overhead matters, this gives compile an
20	    artificial advantage. We accept this as the cost of using the agreed
21	    "compiled" baseline policy.
22	  - cuBLAS / cuDNN allocate workspaces on first call. 10 warmup absorbs.
23	  - Median over a fairly small number of trials (default 30) is fine for
24	    headline numbers but won't catch bimodal distributions. Use --extra
25	    instrumentation if you ever care.
26	"""
27	from __future__ import annotations
28	
29	import os
30	import re
31	

... (3182 more chars)

1	"""Numeric distribution stress cases for correctness checks.
2	
3	These cases are not hidden shapes. They rerun the same shape/seed validation
4	under a few scale regimes that catch kernels that only work for the nominal
5	N(0, 1)-ish inputs.
6	"""
7	from __future__ import annotations
8	
9	import os
10	from collections.abc import Iterator, Mapping, Sequence
11	from contextlib import contextmanager
12	from dataclasses import dataclass, field
13	
14	import torch
15	
16	ToleranceOverride = Mapping[str, float | Mapping[str, float]]
17	
18	
19	@dataclass(frozen=True)
20	class NumericStressCase:
21	    name: str
22	    input_scales: Mapping[int, float] = field(default_factory=dict)
23	    state_scales: Mapping[str, float] = field(default_factory=dict)
24	    tolerance: ToleranceOverride | None = None
25	
26	
27	NOMINAL_CASE = NumericStressCase("nominal")
28	
29	_SMALL_BF16 = {"bfloat16": {"atol": 1e-4, "rtol": 5e-2}}
30	_MED_BF16 = {"bfloat16": {"atol": 5e-4, "rtol": 5e-2}}
31	_LARGE_BF16 = {"bfloat16": {"atol": 2e-1, "rtol": 5e-2}}
32	_TINY_FP32 = {"float32": {"atol": 1e-7, "rtol": 1e-4}}
33	_LARGE_FP32 = {"float32": {"atol": 1e-1, "rtol": 1e-4}}
34	
35	_CASES: dict[str, tuple[NumericStressCase, ...]] = {
36	    "01_fp8_gemm": (
37	        NumericStressCase("small_input", input_scales={0: 1e-3}, tolerance=_SMALL_BF16),
38	        NumericStressCase("large_input", input_scales={0: 64.0}, tolerance=_LARGE_BF16),
39	        NumericStressCase(
40	            "small_weight",
41	      

... (5156 more chars)

Wed Jun 10 23:58:51 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 610.43.02              KMD Version: 610.43.02     CUDA UMD Version: 13.3     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA RTX PRO 6000 Blac...    On  |   00000000:01:00.0 Off |                  Off |
| 32%   45C    P0             90W /  600W |      15MiB /  97887MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|======================================================================================

... (1043 more chars)

"""Roofline math: achieved TFLOPS / GB/s, peak fraction.

FLOPS and bytes formulas come from the problem's `problem.yaml`. They are the
dense-equivalent algorithmic work; kernels that skip work (sparsity, early
exit) are still measured against the dense cost.
"""
from __future__ import annotations


def compute_tflops(flops: float, time_ms: float) -> float:
    """flops per invocation and wall time in ms -> TFLOPS."""
    if time_ms <= 0:
        return 0.0
    return flops / (time_ms * 1e-3) / 1e12


def compute_gbps(bytes_moved: float, time_ms: float) -> float:
    """bytes per invocation and wall time in ms -> GB/s."""
    if time_ms <= 0:
        return 0.0
    return bytes_moved / (time_ms * 1e-3) / 1e9


def peak_fraction(
    achieved: float,
    peak: float,
) -> float:
    """Return achieved / peak, clamped to [0, infinity)."""
    if peak <= 0:
        return 0.0
    return max(0.0, achieved / peak)
"""Per-dtype correctness tolerance.

Stanford's KernelBench uses 1e-4 atol/rtol for fp32 and 1e-2 for fp16/bf16.
We extend to fp8 (0.1, permissive) and int (exact).
"""
from __future__ import annotations

import torch

DEFAULT_TOLERANCE = {
    torch.float32: {"atol": 1e-4, "rtol": 1e-4},
    torch.float16: {"atol": 1e-2, "rtol": 1e-2},
    torch.bfloat16: {"atol": 1e-2, "rtol": 1e-2},
    torch.float8_e4m3fn: {"atol": 1e-1, "rtol": 1e-1},
    torch.float8_e5m2: {"atol": 1e-1, "rtol": 1e-1},
    torch.int8: {"atol": 0, "rtol": 0},
    torch.int32: {"atol": 0, "rtol": 0},

... (3095 more chars)

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	)
37	

empty kernel: 4.67 us
read 0.5MB blocks=148:    5.60 us     93.6 GB/s
read 0.5MB blocks=188:    5.73 us     91.5 GB/s
read 0.5MB blocks=376:    5.57 us     94.2 GB/s
read 0.5MB blocks=752:    5.39 us     97.2 GB/s
read 0.5MB blocks=1504:    5.62 us     93.4 GB/s
read 2MB blocks=148:    5.71 us    367.1 GB/s
read 2MB blocks=188:    5.60 us    374.5 GB/s
read 2MB blocks=376:    5.42 us    386.6 GB/s
read 2MB blocks=752:    5.60 us    374.5 GB/s
read 2MB blocks=1504:    5.76 us    364.1 GB/s
read 8MB blocks=148:    7.86 us   1067.8 GB/s
read 8MB blocks=188:    8.37 us   1002.5 GB/s
read 8MB blocks=376:    7.92 us   1059.2 GB/s
read 8MB blocks=752:    7.58 us   1106.1 GB/s
read 8MB blocks=1504:    7.82 us   1072.2 GB/s
read 64MB blocks=148:   47.90 us   1400.9 GB/s
read 64MB blocks=188:   47.49 us   1413.2 GB/s
read 64MB blocks=376:   46.94 us   1429.6 GB/s
read 64MB blocks=752:   46.74 us   1435.9 GB/s
read 64MB blocks=1504:   46.90 us   1431.0 GB/s
event pair only: 2.29 us

Updated task #1 status

Updated task #2 status

File created successfully at: /workspace/problems/05_topk_bitonic/solution.py (file state is current in your context — no need to Read it back)

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

File created successfully at: /workspace/problems/05_topk_bitonic/scratch/test_sol.py (file state is current in your context — no need to Read it back)

done b=1 n=131072 k=64
done b=64 n=8192 k=8
done b=32 n=16384 k=32
done b=16 n=12000 k=16
done b=128 n=4096 k=1
done b=3 n=700 k=5
done b=2 n=64 k=64
done b=1 n=65 k=64
done b=5 n=4097 k=33
done b=7 n=12345 k=7
ALL OK

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /workspace/problems/05_topk_bitonic/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)