benchmark.py
check.py
problem.yaml
PROMPT.txt
reference.py
shapes.py
sota.py

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	            f"mrope_section must sum to D/2, got {mrope_section} sum={sum(mrope_section)} D/2={D//2}"
39	        self.B, self.T = B, T
40	        self.Hq, self.Hkv, self.

... (3109 more chars)

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 None
42	
43	    # --- Per-shape correctness --------------------------------------------
44	    all_shapes = shapes.SHAPES
45	    for shape_idx, shape in enume

... (3242 more chars)

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", "compute")
40	    flops_formula = meta["flops_formula"]
41	    bytes_formula = meta["bytes_formula"]
42	    num_perf_trials = int(meta.get("num_perf_trials", 3

... (3665 more chars)

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: "Qwen2-VL apply_multimodal_rotary_pos_emb (transformers)"
36	  url: "https://github.com/huggingface/transformers"
37	  function: "transformers.models.qwen2_v

... (83 more chars)

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	]
25	

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], cos_cache[pos_w]], dim=0)
38	    sin = torch.stack([sin_cache[pos_t], sin_cache[pos_h], sin_cache[pos_w]], dim=0)
39	
40	    # transformers wants (B, H, T, 

... (290 more chars)

1	I need you to optimize the multi-axis RoPE pre-attention prep step 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 op is the pre-attention prep used in Qwen2-VL-style vision-language models: take q (B, T, Hq, D) and k (B, T, Hkv, D) bf16 plus three int64 position-id tensors pos_t, pos_h, pos_w of shape (B, T) — temporal, height, width axes — gather per-axis cos/sin from a precomputed (max_pos, D) bf16 cache, mix the per-axis values across the head_dim according to mrope_section (a tuple [s0, s1, s2] summing to D/2 — the first s0 pair-indices use the temporal axis, next s1 use height, last s2 use width, mirrored into the second half of D), then apply standard rotate-half rotation (out = x*cos + rotate_half(x)*sin) on q and k. Return q_rot and k_rot in (B, H, T, D) layout — the transpose from (B, T, H, D) to (B, H, T, D) is part of the prep. Correctness tolerance on the bf16 outputs is 1e-2 abs/rel. Both q_rot and k_rot must match the reference within that.
4	
5	The shapes you have to handle are B=1 T=4096 Hq=32 Hkv=8 D=128 mrope_section=(16,24,24) max_pos=32768 (Qwen2-VL base), B=1 T=8192 Hq=28 Hkv=4 D=128 mrope_section=(16

... (1449 more chars)

Thu May  7 00:42:02 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 595.58.03              Driver Version: 595.58.03      CUDA Version: 13.2     |
+-----------------------------------------+------------------------+----------------------+
| 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 |
| 30%   25C    P8              6W /  600W |     326MiB /  97887MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

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

... (464 more chars)

2.11.0+cu130
13.0
(12, 0)

Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

/home/infatoshi/cuda/KernelBench-Hard/src/eval/:
__init__.py
__pycache__
correctness.py
report.py
roofline.py
shapes.py
timing.py

/home/infatoshi/cuda/KernelBench-Hard/src/hardware/:
__init__.py
__pycache__
m4_max.py
rtx_pro_6000.py

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
40	      - a dict {"atol": ..., "rtol": ...} -> used directly, with values coerced
41	    """
42	    # Accept several spellings of the same dty

... (2345 more chars)

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 statistics
30	
31	i

... (1271 more chars)

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

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	

01_fp8_gemm
02_kda_cutlass
03_paged_attention
04_kahan_softmax
05_topk_bitonic
06_sonic_moe_swiglu
07_w4a16_gemm
08_metal_lightning_attn
09_fmha_preattn_mrope
10_patch_embed_conv3d_gemm

20260427_023632_opencode_deepseek_deepseek-v4-flash_05_topk_bitonic
20260427_032119_opencode_deepseek_deepseek-v4-flash_05_topk_bitonic
20260427_044055_opencode_deepseek_deepseek-v4-flash_01_fp8_gemm
20260427_044258_opencode_deepseek_deepseek-v4-flash_01_fp8_gemm
20260427_051426_opencode_deepseek_deepseek-v4-pro_01_fp8_gemm
20260427_055936_opencode_openrouter-pinned_minimax_minimax-m2.7_01_fp8_gemm
20260427_060216_opencode_openrouter-pinned_qwen_qwen3.6-27b_01_fp8_gemm
20260427_060329_opencode_openrouter-pinned_qwen_qwen3.6-35b-a3b_01_fp8_gemm
20260427_060330_opencode_deepseek_deepseek-v4-flash_02_kda_cutlass
20260427_065343_opencode_deepseek_deepseek-v4-pro_02_kda_cutlass
20260427_073847_opencode_openrouter-pinned_minimax_minimax-m2.7_02_kda_cutlass
20260427_082347_opencode_openrouter-pinned_qwen_qwen3.6-27b_02_kda_cutlass
20260427_082405_opencode_openrouter-pinned_qwen_qwen3.6-35b-a3b_02_kda_cutlass
20260427_082406_opencode_deepseek_deepseek-v4-flash_03_paged_attention
20260427_085225_opencode_deepseek_deepseek-v4-pro_03_paged_attention
20260427_092504_opencode_openrouter-pinned_minimax_minimax-m2.7_03_paged_attention
20260427_093353_opencode_openrouter-pinned_qwen_qwen3.6-27b_03_paged_attention
20260427_094003_opencode_openrouter-pinned_qwen_qwen3.6-35b-a3b_03_paged_attention
20260427_094004_opencode_deepseek_deepseek-v4-flash_04_kahan_softmax
20260427_095552_opencode_deepseek_deepseek-v4-pro_04_kahan_softmax

__pycache__
benchmark.py
check.py
problem.yaml
PROMPT.txt
reference.py
shapes.py
sota.py

20260427_082406_opencode_deepseek_deepseek-v4-flash_03_paged_attention
20260427_085225_opencode_deepseek_deepseek-v4-pro_03_paged_attention
20260427_092504_opencode_openrouter-pinned_minimax_minimax-m2.7_03_paged_attention
20260427_093353_opencode_openrouter-pinned_qwen_qwen3.6-27b_03_paged_attention
20260427_094003_opencode_openrouter-pinned_qwen_qwen3.6-35b-a3b_03_paged_attention

benchmark.log
check.log
result.json
scratch
solution.py
stderr.log
transcript.jsonl

problem: '03_paged_attention'
harness: 'opencode'
model: 'deepseek/deepseek-v4-pro'
reasoning_effort: ''
has_solution: True
correct: True
peak_fraction: 0.0268
elapsed_seconds: 1933
harness_exit_code: 0
session_complete: True
usage:
  input_tokens: 65358
  output_tokens: 13136
  cache_read_tokens: 1116544
  cache_creation_tokens: 0
  reasoning_tokens: 45848
  total_cost_usd: None

"""Paged attention decode kernel for RTX PRO 6000 (SM120 Blackwell).

Single-query decode with grouped-query attention (GQA).
One 32-thread warp-block per query head for maximum parallelism.
Each lane handles D/32 dimensions; all reduction via __shfl_xor_sync.
Online softmax processes KV cache in a single streaming pass.
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16

# ---------------------------------------------------------------------------
# CUDA kernel + wrapper
# ---------------------------------------------------------------------------

_cpp_src = r"""
torch::Tensor launch_paged_attention_decode(
    torch::Tensor query,
    torch::Tensor kv_cache,
    torch::Tensor block_table,
    torch::Tensor seq_lens,
    int B, int H, int Hkv, int D, int P, int max_blocks, float scale);
"""

_cuda_src = r"""
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>

#define WARP_SIZE 32
#define MAX_STRIDE 8

extern "C" __global__ void paged_attention_decode_kernel(
    const __nv_bfloat16* __restrict__ query,
    const __nv_bfloat16* __restrict__ kv_cache,
    const int* __restrict__ block_table,
    const int* __restrict__ seq_lens,
    __nv_bfloat16* __restrict__ output,
    int B

... (4868 more chars)

"-O3",
                "--use_fast_math",
                "-gencode", "arch=compute_120,code=sm_120",
            ],
            extra_cflags=["-O3"],
            verbose=False,
        )
    return _module


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

class Model(nn.Module):
    """Single-query paged attention decode -- custom CUDA kernel."""

    def __init__(
        self,
        batch: int,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        seq_len: int,
        page_size: int,
    ):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.batch = batch
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.seq_len = seq_len
        self.page_size = page_size
        self.scale = 1.0 / math.sqrt(head_dim)
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)

    def forward(
        self,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
        block_table: torch.Tensor,
        seq_lens: torch.Tensor,
    ) -> torch.Tensor:
        B, H, D = query.shape
        max_blocks = block_table.shape[1]

        mod = _get_module()
        return mod.launch_paged_attention_decode(
            query.contiguous(),
            kv_cache.contiguous(),
            block_table.contiguous()

... (994 more chars)

SMs=188, name=NVIDIA RTX PRO 6000 Blackwell Workstation Edition, max_threads_per_sm=1536, max_blocks_per_sm=N/A, l2_cache_size=134217728

/usr/local/cuda-13/bin/nvcc
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2026 NVIDIA Corporation
Built on Mon_Mar_02_09:52:23_PM_PST_2026
Cuda compilation tools, release 13.2, V13.2.51
Build cuda_13.2.r13.2/compiler.37434383_0

benchmark.py
check.py
problem.yaml
PROMPT.txt
reference.py
shapes.py
sota.py

Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

File created successfully at: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py (file state is current in your context — no need to Read it back)

Traceback (most recent call last):
  File "/home/infatoshi/cuda/KernelBench-Hard/.venv/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 2774, in _run_ninja_build
    subprocess.run(
  File "/home/infatoshi/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/lib/python3.11/subprocess.py", line 571, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/check.py", line 119, in <module>
    main()
  File "/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/check.py", line 57, in main
    sol_model = solution.Model(*init_args).to(device).eval()
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py", line 200, in __init__
    _get_module()
  File "/home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py", line 165, in _get_module
    _module = load_inline(
              ^^^^^^^^^^^^
  File "/home/infatoshi/cuda/KernelBench-Hard/.venv/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 2174, in load_inline
    return _jit_compile(
           ^^^^^^^^^^^^^
  File "/home/infatoshi/cuda/KernelBench-Hard/.venv/lib/python3.11/site-packages/torch/utils/cpp_extension.

... (3104 more chars)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS

shape=0 variant=eager tflops=0.199 gbps=285.987 ms=0.316
shape=0 variant=compiled tflops=0.547 gbps=784.189 ms=0.115
shape=0 variant=solution tflops=0.823 gbps=1181.126 ms=0.076
shape=0 solution_peak_fraction=0.6562
shape=1 variant=eager tflops=0.191 gbps=279.047 ms=0.527
shape=1 variant=compiled tflops=0.533 gbps=778.719 ms=0.189
shape=1 variant=solution tflops=0.902 gbps=1316.802 ms=0.112
shape=1 solution_peak_fraction=0.7316
shape=2 variant=eager tflops=0.096 gbps=150.147 ms=0.147
shape=2 variant=compiled tflops=0.202 gbps=316.266 ms=0.070
shape=2 variant=solution tflops=0.659 gbps=1030.104 ms=0.021
shape=2 solution_peak_fraction=0.5723
shape=3 variant=eager tflops=0.172 gbps=246.739 ms=1.464
shape=3 variant=compiled tflops=0.482 gbps=692.257 ms=0.522
shape=3 variant=solution tflops=0.958 gbps=1374.647 ms=0.263
shape=3 solution_peak_fraction=0.7637
peak_fraction: 0.6768
RESULT: OK

/usr/local/cuda-13/bin/ncu
NVIDIA (R) Nsight Compute Command Line Profiler
Copyright (c) 2018-2026 NVIDIA Corporation
Version 2026.1.1.0 (build 37634170) (public-release)

Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

File created successfully at: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/_profile_run.py (file state is current in your context — no need to Read it back)

==PROF== Connected to process 23806 (/home/infatoshi/.local/share/uv/python/cpython-3.11.15-linux-x86_64-gnu/bin/python3.11)
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel": 0%....50%....100% - 10 passes
shape=0 done
==PROF== Disconnected from process 23806
==PROF== Report: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/_profile_v1.ncu-rep

[23806] python3.11@127.0.0.1
  void mrope_fused_kernel<128, 16, 24, 24>(const __nv_bfloat16 *, const __nv_bfloat16 *, const long *, const long *, const long *, const __nv_bfloat16 *, const __nv_bfloat16 *, __nv_bfloat16 *, __nv_bfloat16 *, int, int, int, int) (4096, 1, 1)x(128, 1, 1), Device 0, CC 12.0, Invocations 8
    Section: Launch Statistics
    -------------------------------- --------------- ---------- ---------- ----------
    Metric Name                          Metric Unit    Minimum    Maximum    Average
    -------------------------------- --------------- ---------- ---------- ----------
    Block Size                                           128.00     128.00     128.00
    Cluster Size                                           0.00       0.00       0.00
    Grid Size                                          4,096.00   4,096.00   4,096.00
    Preferred Cluster Size                                 0.00       0.00       0.00
    Registers Per Thread             register/thread      40.00      40.00      40.00
    Shared Memory Configuration Size           Kbyte      32.77      32.77      32.77
    Driver Shared Memory Per Block       Kbyte/block       1.02       1.02       1.02
    Dynamic Shared Memory Per Block       byte/block       0.00       0.00       0.00
    Static Shared Memory Per Block        byte/block       0.00       0.00       0.00
    # SMs                                         SM     188.00     188.00     188.00
    Stack Size                  

... (6119 more chars)

==ERROR== Option '--section PipeUtilization' did not match any section.
==ERROR== Use --list-sections to see the list of available sections.

[23806] python3.11@127.0.0.1
  void mrope_fused_kernel<128, 16, 24, 24>(const __nv_bfloat16 *, const __nv_bfloat16 *, const long *, const long *, const long *, const __nv_bfloat16 *, const __nv_bfloat16 *, __nv_bfloat16 *, __nv_bfloat16 *, int, int, int, int) (4096, 1, 1)x(128, 1, 1), Device 0, CC 12.0, Invocations 8
  Note: The shown averages are calculated as the arithmetic mean of the metric values after the evaluation of the    
  metrics for each individual kernel launch.                                                                         
  If aggregating across varying launch configurations (like shared memory, cache config settings), the arithmetic    
  mean can be misleading and looking at the individual results is recommended instead.

Theoretical Active Warps per SM        warp           48
    Theoretical Occupancy                     %          100
    Achieved Occupancy                        %        87.23
    Achieved Active Warps Per SM           warp        41.87
    Theoretical Warps Per Scheduler        warp           12
    GPU Maximum Warps Per Scheduler        warp           12
    ------------------------------- ----------- ------------

--
          result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can   
          occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices   
          Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on     
          optimizing occupancy.                                                                                         
--
    Theoretical Active Warps per SM        warp           48
    Theoretical Occupancy                     %          100
    Achieved Occupancy                        %        85.74
    Achieved Active Warps Per SM           warp        41.16
    Theoretical Warps Per Scheduler        warp           12
    GPU Maximum Warps Per Scheduler        warp           12
    ------------------------------- ----------- ------------

--
          result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can   
          occur between warps with

... (2411 more chars)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS

shape=0 variant=eager tflops=0.178 gbps=255.074 ms=0.354
shape=0 variant=compiled tflops=0.565 gbps=810.667 ms=0.111
shape=0 variant=solution tflops=0.938 gbps=1346.275 ms=0.067
shape=0 solution_peak_fraction=0.7479
shape=1 variant=eager tflops=0.191 gbps=279.293 ms=0.526
shape=1 variant=compiled tflops=0.529 gbps=773.019 ms=0.190
shape=1 variant=solution tflops=0.935 gbps=1364.725 ms=0.108
shape=1 solution_peak_fraction=0.7582
shape=2 variant=eager tflops=0.095 gbps=147.819 ms=0.150
shape=2 variant=compiled tflops=0.228 gbps=356.013 ms=0.062
shape=2 variant=solution tflops=0.796 gbps=1244.284 ms=0.018
shape=2 solution_peak_fraction=0.6913
shape=3 variant=eager tflops=0.176 gbps=252.187 ms=1.432
shape=3 variant=compiled tflops=0.487 gbps=698.427 ms=0.517
shape=3 variant=solution tflops=0.949 gbps=1361.627 ms=0.265
shape=3 solution_peak_fraction=0.7565
peak_fraction: 0.7379
RESULT: OK

==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v2": 0%....50%....100% - 10 passes
shape=0 done
==PROF== Disconnected from process 28978
==PROF== Report: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/_profile_v2.ncu-rep

[28978] python3.11@127.0.0.1
  void mrope_fused_kernel_v2<128, 16, 24, 24>(const __nv_bfloat16 *, const __nv_bfloat16 *, const long *, const long *, const long *, const __nv_bfloat16 *, const __nv_bfloat16 *, __nv_bfloat16 *, __nv_bfloat16 *, int, int, int, int) (4096, 1, 1)x(32, 1, 1), Device 0, CC 12.0, Invocations 8
    Section: Launch Statistics
    -------------------------------- --------------- ---------- ---------- ----------
    Metric Name                          Metric Unit    Minimum    Maximum    Average
    -------------------------------- --------------- ---------- ---------- ----------
    Block Size                                            32.00      32.00      32.00
    Cluster Size                                           0.00       0.00       0.00
    Grid Size                                          4,096.00   4,096.00   4,096.00
    Preferred Cluster Size                                 0.00       0.00       0.00
    Registers Per Thread             register/thread      40.00      40.00      40.00
    Shared Memory Configuration Size           Kbyte      32.77      32.77      32.77
    Driver Shared Memory Per Block       Kbyte/block       1.02       1.02       1.02
    Dynamic Shared Memory Per Block       byte/block       0.00       0.00       0.00
    Static Shared Memory Per Block        byte/block       0.00       0.00       0.00
    # SMs                                         SM     188.00     188.00     188.00
    Stack Size                

... (6103 more chars)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.200 gbps=286.335 ms=0.315
shape=0 variant=compiled tflops=0.527 gbps=755.826 ms=0.119
shape=0 variant=solution tflops=0.836 gbps=1199.456 ms=0.075
shape=0 solution_peak_fraction=0.6664
shape=1 variant=eager tflops=0.191 gbps=279.276 ms=0.526
shape=1 variant=compiled tflops=0.528 gbps=771.720 ms=0.190
shape=1 variant=solution tflops=0.950 gbps=1386.766 ms=0.106
shape=1 solution_peak_fraction=0.7704
shape=2 variant=eager tflops=0.096 gbps=150.752 ms=0.147
shape=2 variant=compiled tflops=0.225 gbps=350.952 ms=0.063
shape=2 variant=solution tflops=0.901 gbps=1407.739 ms=0.016
shape=2 solution_peak_fraction=0.7821
shape=3 variant=eager tflops=0.173 gbps=248.011 ms=1.456
shape=3 variant=compiled tflops=0.474 gbps=680.157 ms=0.531
shape=3 variant=solution tflops=0.951 gbps=1364.590 ms=0.265
shape=3 solution_peak_fraction=0.7581
peak_fraction: 0.7428
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.195 gbps=279.374 ms=0.323
shape=0 variant=compiled tflops=0.567 gbps=813.472 ms=0.111
shape=0 variant=solution tflops=0.950 gbps=1363.519 ms=0.066
shape=0 solution_peak_fraction=0.7575
shape=1 variant=eager tflops=0.192 gbps=279.803 ms=0.525
shape=1 variant=compiled tflops=0.523 gbps=763.574 ms=0.193
shape=1 variant=solution tflops=0.920 gbps=1343.766 ms=0.109
shape=1 solution_peak_fraction=0.7465
shape=2 variant=eager tflops=0.096 gbps=150.098 ms=0.147
shape=2 variant=compiled tflops=0.228 gbps=356.289 ms=0.062
shape=2 variant=solution tflops=0.932 gbps=1456.691 ms=0.015
shape=2 solution_peak_fraction=0.8093
shape=3 variant=eager tflops=0.175 gbps=250.694 ms=1.440
shape=3 variant=compiled tflops=0.479 gbps=687.533 ms=0.525
shape=3 variant=solution tflops=0.940 gbps=1349.253 ms=0.268
shape=3 solution_peak_fraction=0.7496
peak_fraction: 0.7653
RESULT: OK

done

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.200 gbps=286.583 ms=0.315
shape=0 variant=compiled tflops=0.570 gbps=818.190 ms=0.110
shape=0 variant=solution tflops=0.916 gbps=1314.288 ms=0.069
shape=0 solution_peak_fraction=0.7302
shape=1 variant=eager tflops=0.180 gbps=262.833 ms=0.559
shape=1 variant=compiled tflops=0.530 gbps=774.257 ms=0.190
shape=1 variant=solution tflops=0.911 gbps=1330.725 ms=0.110
shape=1 solution_peak_fraction=0.7393
shape=2 variant=eager tflops=0.096 gbps=150.196 ms=0.147
shape=2 variant=compiled tflops=0.217 gbps=339.823 ms=0.065
shape=2 variant=solution tflops=0.924 gbps=1443.006 ms=0.015
shape=2 solution_peak_fraction=0.8017
shape=3 variant=eager tflops=0.174 gbps=249.811 ms=1.446
shape=3 variant=compiled tflops=0.483 gbps=692.533 ms=0.521
shape=3 variant=solution tflops=0.949 gbps=1361.052 ms=0.265
shape=3 solution_peak_fraction=0.7561
peak_fraction: 0.7563
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.196 gbps=280.653 ms=0.322
shape=0 variant=compiled tflops=0.566 gbps=812.418 ms=0.111
shape=0 variant=solution tflops=0.940 gbps=1348.528 ms=0.067
shape=0 solution_peak_fraction=0.7492
shape=1 variant=eager tflops=0.191 gbps=279.293 ms=0.526
shape=1 variant=compiled tflops=0.530 gbps=774.518 ms=0.190
shape=1 variant=solution tflops=0.920 gbps=1342.980 ms=0.109
shape=1 solution_peak_fraction=0.7461
shape=2 variant=eager tflops=0.095 gbps=149.191 ms=0.148
shape=2 variant=compiled tflops=0.226 gbps=352.743 ms=0.063
shape=2 variant=solution tflops=0.911 gbps=1423.687 ms=0.016
shape=2 solution_peak_fraction=0.7909
shape=3 variant=eager tflops=0.174 gbps=250.066 ms=1.444
shape=3 variant=compiled tflops=0.482 gbps=692.193 ms=0.522
shape=3 variant=solution tflops=0.940 gbps=1348.769 ms=0.268
shape=3 solution_peak_fraction=0.7493
peak_fraction: 0.7587
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

Todos have been modified successfully. Ensure that you continue to use the todo list to track your progress. Please proceed with the current tasks if applicable

PASS
shape=0 variant=eager tflops=0.181 gbps=259.032 ms=0.349
shape=0 variant=compiled tflops=0.521 gbps=747.020 ms=0.121
shape=0 variant=solution tflops=0.895 gbps=1283.786 ms=0.070
shape=0 solution_peak_fraction=0.7132
shape=1 variant=eager tflops=0.206 gbps=300.938 ms=0.488
shape=1 variant=compiled tflops=0.489 gbps=713.912 ms=0.206
shape=1 variant=solution tflops=0.892 gbps=1302.982 ms=0.113
shape=1 solution_peak_fraction=0.7239
shape=2 variant=eager tflops=0.075 gbps=117.711 ms=0.188
shape=2 variant=compiled tflops=0.469 gbps=733.369 ms=0.030
shape=2 variant=solution tflops=0.994 gbps=1553.258 ms=0.014
shape=2 solution_peak_fraction=0.8629
shape=3 variant=eager tflops=0.172 gbps=246.313 ms=1.466
shape=3 variant=compiled tflops=0.472 gbps=677.096 ms=0.533
shape=3 variant=solution tflops=0.984 gbps=1412.414 ms=0.256
shape=3 solution_peak_fraction=0.7847
peak_fraction: 0.7689
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.232 gbps=332.836 ms=0.271
shape=0 variant=compiled tflops=0.737 gbps=1058.185 ms=0.085
shape=0 variant=solution tflops=1.132 gbps=1624.601 ms=0.056
shape=0 solution_peak_fraction=0.9026
shape=1 variant=eager tflops=0.206 gbps=301.342 ms=0.488
shape=1 variant=compiled tflops=0.675 gbps=986.188 ms=0.149
shape=1 variant=solution tflops=1.042 gbps=1521.837 ms=0.097
shape=1 solution_peak_fraction=0.8455
shape=2 variant=eager tflops=0.138 gbps=215.126 ms=0.103
shape=2 variant=compiled tflops=0.172 gbps=269.002 ms=0.082
shape=2 variant=solution tflops=0.807 gbps=1261.314 ms=0.018
shape=2 solution_peak_fraction=0.7007
shape=3 variant=eager tflops=0.172 gbps=246.160 ms=1.467
shape=3 variant=compiled tflops=0.533 gbps=764.583 ms=0.472
shape=3 variant=solution tflops=1.050 gbps=1506.405 ms=0.240
shape=3 solution_peak_fraction=0.8369
peak_fraction: 0.8179
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.180 gbps=257.848 ms=0.350
shape=0 variant=compiled tflops=0.487 gbps=699.336 ms=0.129
shape=0 variant=solution tflops=0.872 gbps=1251.328 ms=0.072
shape=0 solution_peak_fraction=0.6952
shape=1 variant=eager tflops=0.179 gbps=260.863 ms=0.564
shape=1 variant=compiled tflops=0.495 gbps=723.241 ms=0.203
shape=1 variant=solution tflops=0.906 gbps=1323.441 ms=0.111
shape=1 solution_peak_fraction=0.7352
shape=2 variant=eager tflops=0.072 gbps=112.262 ms=0.197
shape=2 variant=compiled tflops=0.155 gbps=242.186 ms=0.091
shape=2 variant=solution tflops=0.567 gbps=886.154 ms=0.025
shape=2 solution_peak_fraction=0.4923
shape=3 variant=eager tflops=0.178 gbps=255.207 ms=1.415
shape=3 variant=compiled tflops=0.533 gbps=765.049 ms=0.472
shape=3 variant=solution tflops=1.060 gbps=1520.307 ms=0.238
shape=3 solution_peak_fraction=0.8446
peak_fraction: 0.6790
RESULT: OK

97:    #pragma unroll Q_UNROLL
123:    #pragma unroll K_UNROLL

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated. All occurrences were successfully replaced. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.176 gbps=252.879 ms=0.357
shape=0 variant=compiled tflops=0.740 gbps=1061.369 ms=0.085
shape=0 variant=solution tflops=1.128 gbps=1618.543 ms=0.056
shape=0 solution_peak_fraction=0.8992
shape=1 variant=eager tflops=0.206 gbps=301.214 ms=0.488
shape=1 variant=compiled tflops=0.467 gbps=682.109 ms=0.216
shape=1 variant=solution tflops=0.887 gbps=1295.084 ms=0.114
shape=1 solution_peak_fraction=0.7195
shape=2 variant=eager tflops=0.073 gbps=113.423 ms=0.195
shape=2 variant=compiled tflops=0.165 gbps=258.441 ms=0.086
shape=2 variant=solution tflops=0.630 gbps=983.915 ms=0.022
shape=2 solution_peak_fraction=0.5466
shape=3 variant=eager tflops=0.172 gbps=246.823 ms=1.463
shape=3 variant=compiled tflops=0.468 gbps=671.615 ms=0.538
shape=3 variant=solution tflops=0.982 gbps=1409.239 ms=0.256
shape=3 solution_peak_fraction=0.7829
peak_fraction: 0.7254
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.173 gbps=248.897 ms=0.363
shape=0 variant=compiled tflops=0.440 gbps=631.194 ms=0.143
shape=0 variant=solution tflops=1.100 gbps=1578.249 ms=0.057
shape=0 solution_peak_fraction=0.8768
shape=1 variant=eager tflops=0.206 gbps=300.967 ms=0.488
shape=1 variant=compiled tflops=0.657 gbps=959.612 ms=0.153
shape=1 variant=solution tflops=0.932 gbps=1361.086 ms=0.108
shape=1 solution_peak_fraction=0.7562
shape=2 variant=eager tflops=0.060 gbps=93.601 ms=0.236
shape=2 variant=compiled tflops=0.159 gbps=248.947 ms=0.089
shape=2 variant=solution tflops=0.628 gbps=981.818 ms=0.023
shape=2 solution_peak_fraction=0.5455
shape=3 variant=eager tflops=0.177 gbps=254.464 ms=1.419
shape=3 variant=compiled tflops=0.528 gbps=758.136 ms=0.476
shape=3 variant=solution tflops=0.991 gbps=1421.756 ms=0.254
shape=3 solution_peak_fraction=0.7899
peak_fraction: 0.7311
RESULT: OK

==PROF== Profiling "mrope_fused_kernel_v3": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v3": 0%....50%....100% - 10 passes
shape=2 done
==PROF== Disconnected from process 97349
==PROF== Report: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/_profile_v3_s2.ncu-rep

[97349] python3.11@127.0.0.1
  void mrope_fused_kernel_v3<64, 8, 12, 12, 16, 2, 4>(const __nv_bfloat16 *, const __nv_bfloat16 *, const long *, const long *, const long *, const __nv_bfloat16 *, const __nv_bfloat16 *, __nv_bfloat16 *, __nv_bfloat16 *, int, int) (512, 2, 1)x(16, 4, 1), Device 0, CC 12.0, Invocations 8
    Section: Launch Statistics
    -------------------------------- --------------- --------- --------- ---------
    Metric Name                          Metric Unit   Minimum   Maximum   Average
    -------------------------------- --------------- --------- --------- ---------
    Block Size                                           64.00     64.00     64.00
    Cluster Size                                          0.00      0.00      0.00
    Grid Size                                         1,024.00  1,024.00  1,024.00
    Preferred Cluster Size                                0.00      0.00      0.00
    Registers Per Thread             register/thread     40.00     40.00     40.00
    Shared Memory Configuration Size           Kbyte     32.77     32.77     32.77
    Driver Shared Memory Per Block       Kbyte/block      1.02      1.02      1.02
    Dynamic Shared Memory Per Block       byte/block      0.00      0.00      0.00
    Static Shared Memory Per Block        byte/block      0.00      0.00      0.00
    # SMs                                         SM    188.00    188.00    188.00
    Stack Size                                        1,024.00  1,024.00

... (5923 more chars)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape=0 variant=eager tflops=0.175 gbps=251.157 ms=0.359
shape=0 variant=compiled tflops=0.502 gbps=720.777 ms=0.125
shape=0 variant=solution tflops=0.888 gbps=1274.794 ms=0.071
shape=0 solution_peak_fraction=0.7082
shape=1 variant=eager tflops=0.176 gbps=256.687 ms=0.573
shape=1 variant=compiled tflops=0.520 gbps=759.032 ms=0.194
shape=1 variant=solution tflops=0.914 gbps=1334.979 ms=0.110
shape=1 solution_peak_fraction=0.7417
shape=2 variant=eager tflops=0.070 gbps=108.679 ms=0.204
shape=2 variant=compiled tflops=0.161 gbps=251.941 ms=0.088
shape=2 variant=solution tflops=0.827 gbps=1291.963 ms=0.017
shape=2 solution_peak_fraction=0.7178
shape=3 variant=eager tflops=0.169 gbps=242.369 ms=1.490
shape=3 variant=compiled tflops=0.457 gbps=656.322 ms=0.550
shape=3 variant=solution tflops=0.942 gbps=1352.082 ms=0.267
shape=3 solution_peak_fraction=0.7512
peak_fraction: 0.7295
RESULT: OK

shape=0 variant=eager tflops=0.171 gbps=246.074 ms=0.367
shape=0 variant=compiled tflops=0.691 gbps=991.432 ms=0.091
shape=0 variant=solution tflops=0.891 gbps=1279.129 ms=0.071
shape=0 solution_peak_fraction=0.7106
shape=1 variant=eager tflops=0.178 gbps=259.845 ms=0.566
shape=1 variant=compiled tflops=0.678 gbps=989.587 ms=0.149
shape=1 variant=solution tflops=1.086 gbps=1586.484 ms=0.093
shape=1 solution_peak_fraction=0.8814
shape=2 variant=eager tflops=0.065 gbps=102.317 ms=0.216
shape=2 variant=compiled tflops=0.168 gbps=262.215 ms=0.084
shape=2 variant=solution tflops=0.654 gbps=1022.485 ms=0.022
shape=2 solution_peak_fraction=0.5680
shape=3 variant=eager tflops=0.174 gbps=249.499 ms=1.447
shape=3 variant=compiled tflops=0.528 gbps=757.119 ms=0.477
shape=3 variant=solution tflops=1.012 gbps=1452.127 ms=0.249
shape=3 solution_peak_fraction=0.8067
peak_fraction: 0.7319
RESULT: OK

peak_fraction: 0.7484
peak_fraction: 0.8334
peak_fraction: 0.7423
peak_fraction: 0.7327

shape | sol_gbps_median | sol_pct_peak | sol_ms_median
  0  |   1679.2 GB/s |  93.3%       | 0.0538 ms
  1  |   1364.1 GB/s |  75.8%       | 0.1078 ms
  2  |   1009.8 GB/s |  56.1%       | 0.0219 ms
  3  |   1423.7 GB/s |  79.1%       | 0.2536 ms

==PROF== Profiling "mrope_fused_kernel_v3": 0%....50%....100% - 10 passes
==PROF== Profiling "mrope_fused_kernel_v3": 0%....50%....100% - 10 passes
shape=2 done
==PROF== Disconnected from process 121601
==PROF== Report: /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/_profile_v3_s2_b.ncu-rep

Memory Throughput                           Tbyte/s         1.16
    Mem Busy                                          %        24.93
    Max Bandwidth                                     %        68.19
    Block Size                                                    32
    Grid Size                                                  2,048
    Threads                                   thread          65,536
    Theoretical Occupancy                          %              50
    Achieved Occupancy                             %           19.74
    Section: Occupancy
    Overall GPU Occupancy                     %            0
    Cluster Occupancy                         %            0
    Theoretical Occupancy                     %           50
    Achieved Occupancy                        %        19.74
    % Occupancy Graphs
    % Occupancy Per Register Count      %/register     8,170 (50; 50; 50; 50; 50;
    Table Name : Impact of Varying Block Size
    % Occupancy Per Block Size           %   5,270 (50; 50; 50; 100; 100;
    Block Size                                                         32
    % Occupancy Per Shared Mem Size      %/byte   585 (50; 50; 50; 43; 33; 25;
    % Occupancy Per Block Barriers Size                221 (50; 50; 25; 16; 12; 8;
    Table Name : Impact of Varying Block Size
    Block Size                                                     32
    # Warps Per Block Size                 2,540 (24; 24; 24; 48; 48;
    Memory Throughput                 

... (1005 more chars)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape | sol_gbps_median | sol_pct_peak | sol_ms_median
  0  |   1285.8 GB/s |  71.4%       | 0.0702 ms
  1  |   1346.3 GB/s |  74.8%       | 0.1092 ms
  2  |    956.7 GB/s |  53.1%       | 0.0231 ms
  3  |   1377.2 GB/s |  76.5%       | 0.2622 ms

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape | sol_gbps_median | sol_pct_peak | sol_ms_median
  0  |   1295.3 GB/s |  72.0%       | 0.0697 ms
  1  |   1347.3 GB/s |  74.9%       | 0.1091 ms
  2  |    983.2 GB/s |  54.6%       | 0.0225 ms
  3  |   1376.2 GB/s |  76.5%       | 0.2624 ms
shape | sol_gbps_median | sol_pct_peak | sol_ms_median
  0  |   1286.4 GB/s |  71.5%       | 0.0702 ms
  1  |   1342.2 GB/s |  74.6%       | 0.1095 ms
  2  |    981.1 GB/s |  54.5%       | 0.0225 ms
  3  |   1379.8 GB/s |  76.7%       | 0.2617 ms

shape | gbps_min | gbps_med | gbps_max
  0  |  1271.1 |  1291.4 |  1300.1 (71.7%)
  1  |  1344.6 |  1348.9 |  1365.3 (74.9%)
  2  |   962.0 |   986.0 |  1002.5 (54.8%)
  3  |  1373.8 |  1377.7 |  1382.1 (76.5%)

shape=0 variant=eager tflops=0.174 gbps=250.066 ms=0.361
shape=0 variant=compiled tflops=0.532 gbps=762.671 ms=0.118
shape=0 variant=solution tflops=0.888 gbps=1273.931 ms=0.071
shape=0 solution_peak_fraction=0.7077
shape=1 variant=eager tflops=0.176 gbps=257.363 ms=0.571
shape=1 variant=compiled tflops=0.478 gbps=697.966 ms=0.211
shape=1 variant=solution tflops=0.914 gbps=1335.367 ms=0.110
shape=1 solution_peak_fraction=0.7419
shape=2 variant=eager tflops=0.071 gbps=110.459 ms=0.200
shape=2 variant=compiled tflops=0.164 gbps=256.427 ms=0.086
shape=2 variant=solution tflops=0.842 gbps=1315.319 ms=0.017
shape=2 solution_peak_fraction=0.7307
shape=3 variant=eager tflops=0.169 gbps=242.325 ms=1.490
shape=3 variant=compiled tflops=0.465 gbps=666.931 ms=0.541
shape=3 variant=solution tflops=0.928 gbps=1331.109 ms=0.271
shape=3 solution_peak_fraction=0.7395
peak_fraction: 0.7298
RESULT: OK

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape | gbps_min | gbps_med | gbps_max
  0  |  1200.2 |  1213.6 |  1222.6 (67.4%)
  1  |  1313.0 |  1325.5 |  1332.3 (73.6%)
  2  |   946.2 |   971.5 |   980.4 (54.0%)
  3  |  1412.0 |  1428.1 |  1430.3 (79.3%)

The file /home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py has been updated successfully. (file state is current in your context — no need to Read it back)

PASS
shape | gbps_min | gbps_med | gbps_max
  0  |  1273.6 |  1293.5 |  1300.4 (71.9%)
  1  |  1329.2 |  1345.1 |  1354.1 (74.7%)
  2  |   968.7 |   984.6 |   998.1 (54.7%)
  3  |  1367.2 |  1376.2 |  1382.1 (76.5%)