content: 'Read reference.py, check.py, benchmark.py, shapes.py to understand the problem'
status: 'in_progress'
priority: 'high'
content: 'Implement initial CUDA/Triton kernel in solution.py'
status: 'pending'
priority: 'high'
content: 'Run check.py to verify correctness'
status: 'pending'
priority: 'high'
content: 'Profile and benchmark initial implementation'
status: 'pending'
priority: 'high'
content: 'Iterate: optimize, verify, benchmark until PASS on check.py and good perf'
status: 'pending'
priority: 'high'

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope</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/09_fmha_preattn_mrope/reference.py</path>
<type>file</type>
<content>
1: """Naive multi-axis RoPE pre-attention reference (correctness only).
2: 
3: Mirrors the Qwen2-VL apply_multimodal_rotary_pos_emb semantics: three position
4: axes (temporal, height, width), each contributing a slice of the head_dim via
5: mrope_section. Standard rotate-half on q and k once per-axis cos/sin have been
6: mixed.
7: 
8: Layout convention:
9:   q: (B, T, Hq, D)   bf16
10:   k: (B, T, Hkv, D)  bf16
11:   pos_t, pos_h, pos_w: (B, T) int64
12: 
13: Output: q_rot, k_rot in (B, H, T, D) — i.e. the layout an attention kernel
14: expects (head dim second). The transpose is part of the "pre-attention prep".
15: """
16: import torch
17: import torch.nn as nn
18: 
19: OP_TYPE = "rope"
20: SUPPORTED_PRECISIONS = ["bf16"]
21: HARDWARE_REQUIRED = ["RTX_PRO_6000"]
22: 
23: 
24: def _build_inv_freq(D: int, base: float = 10000.0) -> torch.Tensor:
25:     return 1.0 / (base ** (torch.arange(0, D, 2, dtype=torch.float32) / D))
26: 
27: 
28: def _rotate_half(x: torch.Tensor) -> torch.Tensor:
29:     half = x.shape[-1] // 2
30:     return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
31: 
32: 
33: class Model(nn.Module):
34:     def __init__(self, B: int, T: int, Hq: int, Hkv: int, D: int,
35:                  mrope_section: tuple[int, int, int], max_pos: int):
36:         super().__init__()
37:         assert sum(mrope_section) == D // 2, \
38:          

... (3384 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/check.py</path>
<type>file</type>
<content>
1: """Correctness runner for multi-axis RoPE pre-attention prep.
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

... (3519 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/benchmark.py</path>
<type>file</type>
<content>
1: """Roofline benchmark for multi-axis RoPE pre-attention prep.
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

... (3966 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/shapes.py</path>
<type>file</type>
<content>
1: """Canonical shape sweep for multi-axis RoPE pre-attention.
2: 
3: Sized for Qwen2-VL-style vision-language inference: long T, GQA, mrope_section
4: splitting D/2 across temporal/height/width axes. Mix of base, long-context,
5: small-head-dim, and very-long-context.
6: """
7: 
8: SHAPES = [
9:     # Qwen2-VL base: T=4k, head_dim=128, GQA 32:8, mrope [16,24,24]
10:     {"B": 1, "T": 4096,  "Hq": 32, "Hkv": 8, "D": 128,
11:      "mrope_section": (16, 24, 24), "max_pos": 32768},
12: 
13:     # Qwen2-VL 7B-style long context, narrower kv heads
14:     {"B": 1, "T": 8192,  "Hq": 28, "Hkv": 4, "D": 128,
15:      "mrope_section": (16, 24, 24), "max_pos": 32768},
16: 
17:     # Smaller head_dim and B=2 (tests batch dim and reduced D)
18:     {"B": 2, "T": 2048,  "Hq": 16, "Hkv": 2, "D": 64,
19:      "mrope_section": (8, 12, 12), "max_pos": 16384},
20: 
21:     # Very long context (16k tokens)
22:     {"B": 1, "T": 16384, "Hq": 32, "Hkv": 8, "D": 128,
23:      "mrope_section": (16, 24, 24), "max_pos": 65536},
24: ]

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

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/problem.yaml</path>
<type>file</type>
<content>
1: name: 09_fmha_preattn_mrope
2: display_name: "Multi-axis RoPE Pre-attention Prep"
3: precision: bf16
4: regime: memory  # pointwise rotation + cache gather; bandwidth-bound
5: 
6: # FLOPs per call: rotate-half on q and k.
7: # Per element: out = a*cos + b*sin -> 3 flops/elem (1 mul + 1 mul + 1 add).
8: # We do this on q ((B*T*Hq*D) elems) and k ((B*T*Hkv*D) elems).
9: flops_formula: "3 * B * T * (Hq + Hkv) * D"
10: 
11: # Bytes moved (lower bound):
12: #   read q, k:                  2 * B*T*(Hq+Hkv)*D
13: #   read 3 position id arrays:  3 * B*T*8        (int64)
14: #   read cos/sin from cache:    6 * B*T*D*2      (3 axes, 2 caches, bf16)
15: #   write q_rot, k_rot:         2 * B*T*(Hq+Hkv)*D
16: bytes_formula: "4 * B * T * (Hq + Hkv) * D + 12 * B * T * D + 24 * B * T"
17: 
18: hardware: [RTX_PRO_6000]
19: peak_tflops_key: bf16    # informational; regime=memory uses bandwidth
20: peak_bandwidth_key: dram
21: 
22: tolerance:
23:   bfloat16: 0.02
24: 
25: # Forbidden ops -- agent must implement the gather + per-axis mix + rotation.
26: # Banning vendor mRoPE entry points and obvious shortcuts.
27: forbidden:
28:   - "transformers.models.qwen2_vl.modeling_qwen2_vl.apply_multimodal_rotary_pos_emb"
29:   - "apply_multimodal_rotary_pos_emb"
30:   - "flash_attn.layers.rotary"
31:   - "flashinfer.rope"
32:   - "flashinfer.apply_rope"
33: 
34: sota:
35:   name:

... (284 more chars)

<path>/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/sota.py</path>
<type>file</type>
<content>
1: """SOTA reference for multi-axis RoPE pre-attention.
2: 
3: Wraps transformers' apply_multimodal_rotary_pos_emb. The full pipeline (gather
4: per-axis cos/sin from cache, then call the vendor rotation) is timed, matching
5: what the reference and the agent solution do.
6: 
7: If transformers isn't installed or the impl signature has drifted, returns
8: is_available()=False and the benchmark omits the SOTA line.
9: """
10: from __future__ import annotations
11: 
12: import torch
13: 
14: 
15: def _try_transformers():
16:     try:
17:         from transformers.models.qwen2_vl.modeling_qwen2_vl import (
18:             apply_multimodal_rotary_pos_emb as _impl,
19:         )
20:         return _impl
21:     except Exception:
22:         return None
23: 
24: 
25: def sota_forward(
26:     q: torch.Tensor, k: torch.Tensor,
27:     pos_t: torch.Tensor, pos_h: torch.Tensor, pos_w: torch.Tensor,
28:     cos_cache: torch.Tensor, sin_cache: torch.Tensor,
29:     mrope_section: tuple[int, int, int],
30: ):
31:     """q,k: (B, T, H, D); pos_*: (B, T); cos_cache,sin_cache: (max_pos, D)."""
32:     impl = _try_transformers()
33:     if impl is None:
34:         raise RuntimeError("transformers not available for SOTA")
35: 
36:     # Gather per-axis cos/sin and stack into (3, B, T, D) as transformers expects.
37:     cos = torch.stack([cos_cache[pos_t], cos_cache[pos_h], c

... (495 more chars)

Found 1 matches
/home/infatoshi/cuda/KernelBench-Hard/src/eval/correctness.py:
  Line 61: def check_correctness(

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

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

3.6.0

2.11.0+cu130
(12, 0)

no solution.py

content: 'Read reference.py, check.py, benchmark.py, shapes.py to understand the problem'
status: 'completed'
priority: 'high'
content: 'Implement initial Triton kernel in solution.py'
status: 'in_progress'
priority: 'high'
content: 'Run check.py to verify correctness'
status: 'pending'
priority: 'high'
content: 'Profile and benchmark initial implementation'
status: 'pending'
priority: 'high'
content: 'Iterate: optimize, verify, benchmark until PASS on check.py and good perf'
status: 'pending'
priority: 'high'