Let the benchmark harness run norm, rope and sampling on ROCm - #364
Conversation
The shared runner already ships routine modules for norm, rope and sampling -- every op in them has a HIP kernel -- but on ROCm all 23 routines ended at "[ERROR] No backends to test" and wrote a zero-row CSV. Two separate causes: `rocm_supported_backends` returned a hardcoded ["fa2", "auto"], which is the attention answer. norm/rope/sampling declare choices=["cuda"] (norm also "cute-dsl"), so the filter stripped the user's backend and offered two names their dispatch arms raise on. The registry is now routine -> (capability op, backend names), still deriving the op from arch_caps so the support list cannot drift from the matrix. `dtype_str_to_torch_dtype` mapped fp8_e4m3/fp8_e5m2 to the OCP encodings. CDNA implements fnuz, so the three fp8 rope routines died inside the kernel with "Output dtype must be float8" -- a message that names neither the dtype nor the argument that chose it. norm's four rmsnorm-family routines also rewrote ["cuda"] to ["cute-dsl"] before dispatching. Left alone the routines run on ROCm but label the row cute-dsl, naming a kernel family that does not exist on this backend. The rewrite is now gated on !IS_HIP, and every dispatch arm accepts "cuda" as well -- which also fixes `--backends cuda cute-dsl` raising "Unsupported backend: cuda" on CUDA. Measured on gfx942 (MI300X), all 23 routines, --refcheck where the routine supports it: before: 23/23 "No backends to test" after: 23/23 ran, backend column reads "cuda" Excluded deliberately: apply_rope_with_cos_sin_cache builds cos_sin_cache in --input_dtype but flashinfer/rope.py:1222 requires float32 and --input_dtype offers none, so it fails on CUDA too; top_k, top_k_page_table_transform and top_k_ragged_transform call flashinfer.topk, which has no csrc/rocm kernel. test_benchmark_harness.py's _backend_choices() read attention.py's parser for every routine, and attention never offers "cuda" -- so the "offered backends must be ones the CLI accepts" invariant had to become per-module before it could cover the new routines. Co-Authored-By: Claude <noreply@anthropic.com>
Two supported ops with no benchmark. Both upstream scripts are unusable here rather than merely unported: bench_block_sparse_attention.py builds a backend="fa3" wrapper and sweeps ["fa2", "fa3"], and fa3 raises "FA3 backend not currently supported for ROCm" from gen_customize_batch_prefill_module; bench_mixed_attention.py's third arm is flashinfer.BatchAttention, which is gated CUDA-only. Both follow bench_norm.py's standalone A/B shape rather than the rocprofv3 one: the question is where one arm overtakes the other, not where a single kernel sits on the roofline. bench_block_sparse_attention.py sweeps block density, which is the only axis that moves the answer -- dense is flat in it, sparse is linear. The density-1.0 row is a control rather than a data point: sparse has no work to skip there, so a speedup would mean the dense baseline is mis-specified, and the script says so rather than leaving it to the reader. Measured on gfx942, fixed and variable block wrappers: density 0.10 -> 4.1-6.4x density 0.50 -> 1.4-1.7x density 0.25 -> 2.6-3.2x density 1.00 -> 0.96-0.98x (control) Accuracy against a masked fp32 reference: max abs err 1.5e-4 to 5.0e-4. bench_mixed_attention.py sweeps decode batch size at a fixed prefill. It reports what the split arm's decode resolved to, and takes --decode-backend, because POD runs the HIP kernel for both halves: left on auto the split arm takes AITER and the ratio measures backend choice as much as fusion. Both readings say the same thing on gfx942 -- POD costs more than not fusing: --decode-backend fa2 (fusion only): 0.55-0.68x --decode-backend auto (what you'd run): 0.44-0.64x Each stream needs its own KV tensor with 0-based page indices. Sharing one cache and offsetting the prefill indices reads past the end of the slice and faults the GPU with HSA_STATUS_ERROR_MEMORY_FAULT -- and without an explicit synchronize the call returns correct-looking shapes first, so the setup error reads as a working benchmark. Co-Authored-By: Claude <noreply@anthropic.com>
benchmarks/README.md said "Only the attention routines are available", which stopped being true one commit ago. The per-routine table now carries norm, rope and sampling, names why three routines in those groups are deliberately unregistered, and explains what the "cuda" backend label means here -- the op's own default, which for norm and rope may itself be AITER, so the native-vs-AITER question belongs to rocm/bench_norm.py rather than this runner. The testlist gains 28 lines. Choices worth knowing: hidden_size 111 is in the norm sweep because the native kernel's vec_size is gcd(16/sizeof(T), d), so an ill-aligned d goes fully scalar -- measured 0.139 TB/s against 2.05 TB/s at d=4096, which a powers-of-two sweep never sees. The sampling rows carry no --refcheck because those routines sample. Verified end to end: 28 lines, 28 PERF rows, 28 CSV rows, no [ERROR] lines. The benchmark-kernel skill claimed MLA, cascade, POD and MoE were unavailable and that AITER was prefill-only. Both were already wrong before this branch -- bench_mla.py and bench_fused_moe_aiter.py have existed for weeks, and auto tries AITER for decode and MLA too. It now points at the generated README matrix rather than restating a list that drifts. Co-Authored-By: Claude <noreply@anthropic.com>
Self-review findings on the three commits below.
The fp8 override was applied inside dtype_str_to_torch_dtype, which
every routine shares. attention.py's kv_dtype accept-lists name the OCP
spellings, so an fnuz tensor was rejected and the row left the CSV
without failing the run. Measured on gfx942:
before: [PERF] fa2 :: median time 0.127 ms (--kv_dtype fp8_e4m3)
after: [ERROR] Unsupported kv_dtype: fp8_e4m3
Widening those accept-lists would mean 16 edits in an upstream file and
a behaviour change on a path this branch does not test, so the override
moved to hip_quant_dtype() at rope.py's three --quant_dtype sites
instead. The comment claiming CDNA has only the fnuz encodings was also
wrong -- moe_fp8_dtype() is arch-dependent because it follows AITER's
MFMA instructions. What is true, and what this actually relies on, is
that is_float8_tensor (csrc/rocm/pytorch_extension_utils.h) accepts the
fnuz encodings only, on both architectures.
norm.py's dispatch arms were widened to `backend in ("cute-dsl",
"cuda")`, which on CUDA silently turns `--backends cuda cute-dsl` into
two rows measuring the identical call. Now `(IS_HIP and backend ==
"cuda")`, matching the `and not IS_HIP` guard on the rewrite above it.
bench_block_sparse_attention.py:
- The dense arm called single_prefill_with_kv_cache_return_lse while
the sparse arm's return_lse defaults to False, so dense paid for an
extra [seq_len, heads] fp32 write in every ratio -- including the
density-1.0 control whose job is to catch exactly that.
- GQA was skipped for the variable-block wrapper on a false premise.
block_mask_map is shaped (num_kv_heads, MB, NB), so a group of qo
heads shares one kv head's pattern; the sweep now covers 32/8, which
is the shape SGLang's hybrid attention actually uses. 36 rows, up
from 27.
- The fixed-block wrapper was left on backend="auto" while everything
it is divided by is pinned to fa2.
- --num-blocks that does not divide seq_len truncated the block size,
shrinking the sparse problem below the dense baseline and crashing
--accuracy on a shape mismatch. Rejected up front.
- --kinds had no choices, so a typo ran the fixed case under the
typo'd label.
Both scripts: the per-case `finally: torch.cuda.synchronize()` re-raises
a deferred fault from outside the try, discarding every row measured so
far -- on a cold cache that is an hour of JIT builds. Guarded, and the
fault recorded on the row.
bench_mixed_attention.py: --causal was action="store_true" with
default=True and could never do anything; --accuracy's docstring called
the split arm a reference without noting that only holds at
--decode-backend fa2.
benchmarks/README.md said auto "may itself be AITER" for norm and rope
and that the ops take no backend argument. Both wrong, and the first
contradicted the skill file edited in the same commit:
_auto_select_norm_backend and _auto_select_rope_backend both return
"native" unconditionally, and flashinfer.rmsnorm does take backend=.
It also said "three routines" over a list of four.
Re-measured after the fixes on gfx942: fp8 attention row back at 0.121
ms, all three fp8 rope routines running, harness tests green, both
scripts' accuracy checks unchanged.
Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The block-sparse benchmark misreports effective density and does not validate its variable-wrapper arm.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Enables norm, RoPE, and sampling benchmarks on ROCm and adds standalone block-sparse and POD attention comparisons.
Changes:
- Adds architecture-aware backend registration and ROCm FP8 dtype mapping.
- Adds block-sparse and mixed-attention benchmarks.
- Expands tests, benchmark coverage, and documentation.
File summaries
| File | Description |
|---|---|
tests/rocm/test_benchmark_harness.py |
Expands backend and FP8 mapping tests. |
benchmarks/routines/rope.py |
Applies ROCm-compatible FP8 quantization dtypes. |
benchmarks/routines/rocm/support.py |
Registers supported ROCm routines and backends. |
benchmarks/routines/rocm/__init__.py |
Exports the dtype helper. |
benchmarks/routines/norm.py |
Dispatches native norm kernels correctly on ROCm. |
benchmarks/rocm/testlist_rocm.txt |
Adds norm, RoPE, and sampling cases. |
benchmarks/rocm/bench_mixed_attention.py |
Adds POD-versus-split attention benchmarking. |
benchmarks/rocm/bench_block_sparse_attention.py |
Adds sparse-versus-dense attention benchmarking. |
benchmarks/README.md |
Documents expanded ROCm benchmark support. |
.claude/skills/benchmark-kernel/SKILL.md |
Updates ROCm benchmarking guidance. |
Review details
- Files reviewed: 10/10 changed files
- Comments generated: 6
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Copilot review round 1 on #364; all six findings accepted. _block_mask forces the diagonal live after a Bernoulli draw, so the realized density is density + (1-density)/nb -- 0.126 for the row labelled 0.10 at nb=32. Since the file exists to locate the density crossover, every row and both accuracy paths now report the density that actually ran alongside the one requested. --accuracy only ever checked the fixed-block wrapper, leaving the variable arm -- half the default sweep, and the arm whose GQA support this branch changed in 261738f -- with no reference at all. It now covers whichever arms --kinds selected, and the variable check broadcasts each kv head's pattern across its qo group, which is the part a wrong setup gets wrong. Measured on gfx942, max abs err against a masked fp32 reference: fixed 8/8 1.2e-4 .. 4.5e-4 variable 32/32 1.1e-4 .. 4.6e-4 variable 32/8 1.2e-4 .. 5.1e-4 That GQA row is the first direct evidence the 261738f change is correct rather than merely plausible. --num-blocks 0 reached the divisibility check and raised ZeroDivisionError; negatives reached tensor construction. Both are rejected up front now. Provenance recorded neither torch.version.hip nor the amd-aiter version, which README.md asks for with every result. bench_mixed_ attention.py needs the aiter version most of all: --decode-backend auto routes its split arm through AITER. The _HEAD_PAIRS comment still said the variable wrapper requires equal head counts. That stopped being true in 261738f, which removed the skip. Co-Authored-By: Claude <noreply@anthropic.com>
Self-review of the review-response commit; ten findings taken, five declined as structural (noted on the PR). rope.py wrote args.quant_dtype into the CSV -- the requested OCP name -- while hip_quant_dtype had substituted fnuz for the kernel. A CUDA and a ROCm CSV both claimed fp8_e4m3 for measurements in different formats. The column now records what ran: float8_e4m3fnuz. hip_quant_dtype's off-HIP arm was unreachable by any test, since tests/rocm is skipif(not IS_HIP) at module scope -- and rope.py calls it on CUDA too, where inverting the mapping would silently turn every fp8 rope row fnuz. Split into a pure to_fnuz() plus an alias bound once at import, so the mapping is testable anywhere and `hip_quant_dtype is to_fnuz` pins which arm HIP got. --accuracy checked the fixed wrapper at MHA only while the sweep times it at 32/8. Both wrappers are now checked at both head shapes; a wrong qo-to-kv mapping was previously visible in neither. gfx942, max abs err against a masked fp32 reference: fixed 32/32 1.2e-4 .. 4.8e-4 variable 32/32 1.3e-4 .. 5.1e-4 fixed 32/8 1.1e-4 .. 5.5e-4 variable 32/8 1.3e-4 .. 5.1e-4 test_registered_routine_is_a_real_routine iterated ROCM_NATIVE_ROUTINES -- a second hardcoded copy -- so a name present only in the registry was invisible to it, which is the typo its docstring claims to catch. It now reads _ROCM_ROUTINE_TO_CAP_OP and asserts the two agree. Also: the summary pooled variable and fixed speedups into one median despite the docstring saying those timings are not comparable (now per kind); empty_cache() ran while the closures still owned the case's tensors and workspaces, so it freed nothing and the peak held two cases; _provenance stamped card 0 while _build allocates on the current device; the testlist pointed readers at bench_norm.py --aa for the native-vs-AITER comparison, which is the mode that cannot answer it; and the README said four routines were unregistered when the norm group alone has eight more. Co-Authored-By: Claude <noreply@anthropic.com>
|
Deliberately deferred, from the self-review on
|
|
Closing the bot-review loop. Round 1: 6 findings, all accepted, fixed in That verification has been run on both architectures, gfx942 (MI300X) and gfx950 (MI350X): Numbers and shapes are in the PR description; deferrals are in the comment above. Later rounds re-raising any of this should be read against these two comments rather than answered with another push. |
There was a problem hiding this comment.
🟡 Changes recommended
The incorrect Gemma capability key prevents two advertised routines from running.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 10/10 changed files
- Comments generated: 3
- Review effort level: Balanced
Follow-up to #364, addressing the three review comments that were still open when it merged. ## Summary #364 claimed to make 23 benchmark routines runnable on ROCm. It made 21: both Gemma routines were mapped to a capability op that `arch_caps.py` no longer declares, so they filtered to no backends and produced no rows. This fixes that, adds the guard that would have caught it, and clears two smaller items from the same review. ## What changed - **`benchmarks/routines/rocm/support.py`** — `gemma_rmsnorm` and `gemma_fused_add_rmsnorm` now reference the capability op `gemma_rmsnorm`; they referenced `layernorm`, which was removed from `arch_caps.py` once it turned out to have no ROCm kernel behind it. - **`tests/rocm/test_arch_caps.py`** — `test_benchmark_registry_ops_are_declared`, asserting every op the registry references has a `hip` row. It lives here because `arch-caps-conformance.yml` runs this file on every PR with no GPU; the registry is parsed with `ast` because that lane has no torch. - **`benchmarks/rocm/bench_block_sparse_attention.py`** — record `--seed` with the other case parameters, and actually seed the global RNG: `--seed` reached only the block mask, while every q/k/v came from the unseeded generator, so recording it alone would have overstated what a row pins. - **`benchmarks/rocm/testlist_rocm.txt`** — the norm heading claimed "bf16 and fp16"; all 7 of its commands pass `bfloat16`. ## Architecture / design notes **The missing thing was a CI lane, not a test.** `test_filter_only_offers_backends_the_cli_accepts` already fails on this defect for both Gemma routines — verified by restoring the key. But `tests/rocm/test_benchmark_harness.py` is `skipif(not IS_HIP)` and no workflow references it, so the suite that catches this runs only on a maintainer's ROCm box. That is how the `layernorm` removal merged. The new guard therefore goes in the one lane that runs everywhere, and the on-hardware test keeps its stronger check (a row can exist and still be gated `UNSUPPORTED` on an arch, which a declaration scan cannot see). **Why `ast` and not the regex scan already in that file.** `support.py` passes `op` as a variable to `capability_available`, so the op strings live in a dict literal and never appear at a call site for `test_every_op_the_library_asks_for_is_declared` to match. **Heading corrected rather than adding fp16 cases.** `vec_size` is `gcd(16/sizeof(T), d)` and fp16 and bf16 are both 2-byte, so fp16 cases would exercise the same code path and differ only in numerical noise — coverage in name only, paid for on every testlist run. ## Benchmark results No performance change: this makes two routines run that previously exited before timing anything, and adds a column to a CSV. For the record, on gfx942 at `c333f2003` both now measure (batch 989, hidden 4096, bf16): | routine | median | achieved | | --- | ---: | ---: | | `gemma_rmsnorm` | 0.007 ms | 2.35 TB/s | | `gemma_fused_add_rmsnorm` | 0.019 ms | 1.69 TB/s | ## Test plan - [x] A/B on the new guard, run exactly as CI does (`pytest --noconftest tests/rocm/test_arch_caps.py`): with the `layernorm` key restored it fails naming `layernorm`; with the fix the file's 85 cases pass. - [x] A/B end to end through the real runner, gfx942: with the old key `[ERROR] No backends to test. Exiting.`; with the fix, a `[PERF]` row for both Gemma routines. - [x] `tests/rocm/test_benchmark_harness.py` — exit 0 on gfx942 after removing the duplicate guard. With the bug restored it fails on `test_filter_only_offers_backends_the_cli_accepts[gemma_rmsnorm]` and `[gemma_fused_add_rmsnorm]`, which is the evidence that the pre-existing test already covered this. - [x] `--csv` output carries a `seed` column, verified by parsing the written file. - [x] `pre-commit run` on the changelist — clean, nothing reformatted. - [x] `/code-review xhigh` (the level `review-level.sh` computes for this changelist). Measured on gfx942 only. The fix is a capability-table lookup and a CSV field, neither of which is architecture-dependent, and the guard test does not touch a GPU. Co-authored-by: Claude <noreply@anthropic.com>
Summary
The shared benchmark runner already ships routine modules for norm, rope and sampling, and every op in them has a HIP kernel — but on ROCm all 23 routines ended at
[ERROR] No backends to testand wrote a zero-row CSV. This makes them run, adds benchmarks for two supported ops that had none (block-sparse and POD attention), and corrects the docs that said none of this existed.What changed
Harness
benchmarks/routines/rocm/support.py—rocm_supported_backendsreturned a hardcoded["fa2", "auto"], which is the attention answer; norm/rope/sampling declarechoices=["cuda"]and their dispatch arms raise on anything else, so the filter stripped the user's backend and offered two names nothing could run. The registry is now routine → (capability op, backend names), still deriving the op fromarch_caps.pyso the support list cannot drift from the generated matrix. Also hostship_quant_dtype().benchmarks/routines/rope.py—hip_quant_dtype()at the three--quant_dtypesites.dtype_str_to_torch_dtypemapsfp8_e4m3/fp8_e5m2to the OCP encodings, andis_float8_tensor(csrc/rocm/pytorch_extension_utils.h) accepts only fnuz, so the three fp8 rope routines died inside the kernel withOutput dtype must be float8. Scoped to the quantize path on purpose — see below.benchmarks/routines/norm.py— the four rmsnorm-family routines rewrote["cuda"] → ["cute-dsl"]before dispatch. Left alone they run on ROCm but label the rowcute-dsl, naming a kernel family that does not exist on this backend. The rewrite is gated on!IS_HIP, and the dispatch arms accept"cuda"underIS_HIPonly — widening them unconditionally would let--backends cuda cute-dslemit two CUDA rows measuring the identical call.New benchmarks
benchmarks/rocm/bench_block_sparse_attention.py— both ROCm wrappers against dense prefill, sweeping block density. The upstream script cannot be reused: it builds abackend="fa3"wrapper and sweeps["fa2", "fa3"], and fa3 raisesFA3 backend not currently supported for ROCm.benchmarks/rocm/bench_mixed_attention.py— POD against running prefill and decode separately, sweeping the decode batch size. The upstream script's third arm isflashinfer.BatchAttention, gated CUDA-only here.Both follow
bench_norm.py's standalone A/B shape rather than the rocprofv3 one, because the question is where one arm overtakes the other, not where a single kernel sits on the roofline.Tests and docs
tests/rocm/test_benchmark_harness.py—_backend_choices()readattention.py's parser for every routine, and attention never offerscuda, so the "offered backends must be ones the CLI accepts" invariant had to become per-module before it could cover the new routines. Plus a registry-typo guard and a case pinning the fnuz mapping.benchmarks/README.md,benchmarks/rocm/testlist_rocm.txt(+28 lines),.claude/skills/benchmark-kernel/SKILL.md.Architecture / design notes
Why the
backendcolumn sayscuda. The rule applied here is that the column names the kernel family that ran, using a name the routine's own CLI accepts.cute-dslbreaks it — there is no CuTe DSL kernel on ROCm.cudadoes not: it is upstream's name for "the library's own kernel", the same legacy spelling torch uses for the HIP device. Renaming it tohipwould mean editingchoicesand every dispatch arm across three upstream files (~34 edits, each a recurring merge conflict) to change a label that is already accurate under that rule. Declined.Four routines are deliberately unregistered.
apply_rope_with_cos_sin_cachebuildscos_sin_cachein--input_dtypewhileflashinfer/rope.py:1222requires float32 and--input_dtypeoffers none — it fails on CUDA too.top_k,top_k_page_table_transformandtop_k_ragged_transformcallflashinfer.topk, which has nocsrc/rocmkernel.Why the fp8 override is scoped to
--quant_dtype. Putting it in the shareddtype_str_to_torch_dtypefixed the rope routines and broke every fp8 attention row:attention.py's accept-lists name the OCP spellings, so an fnuz tensor is rejected and the row leaves the CSV without failing the run (measured:[PERF] fa2 :: 0.127 ms→[ERROR] Unsupported kv_dtype: fp8_e4m3). Widening those lists means 16 edits in an upstream file plus a behaviour change on a path this branch does not test, so the override sits at rope's three--quant_dtypesites instead. Note this is not the arch-dependentmoe_fp8_dtype()story —is_float8_tensortakes fnuz on both architectures.Upstream-file cost. Two files: eight one-line edits in
norm.py, four inrope.py(one import, three call sites).flashinfer_benchmark_utils.pyandsampling.pyare untouched by this branch. Everything else is underbenchmarks/routines/rocm/orbenchmarks/rocm/.Benchmark results
Block-sparse against dense, fixed-block wrapper, fp16, head_dim 128 — speedup tracks 1/density as it should, and the density-1.0 control comes in below 1.0 on both arches, which is what says the dense baseline is set up right. Densities are as requested; each row also reports the realized density, which sits slightly higher because the mask forces one live block per row:
POD against running the two separately — POD costs more than not fusing, on both arches and at every mix measured:
--decode-backend fa2(isolates fusion)--decode-backend auto(what you would otherwise run)The two readings differ because POD runs the HIP kernel for both halves while
autosends the split arm's decode to AITER; the script reports which one ran, so a row cannot be misread as a fusion result when it is a backend result.Test plan
tests/rocm/test_benchmark_harness.py— green on gfx942 and gfx950, plustest_batch_pod.py(141 together) on gfx950.--refcheckwhere supported: 23/23 ran on gfx942, labelledcuda. Baseline for the same four spot-checks on unmodified6202e794d: 4/4No backends to test.[PERF]rows → 28 CSV rows, zero[ERROR]lines.bench_block_sparse_attention.py --accuracyagainst a masked fp32 reference — both wrappers at both head shapes (32/32 and 32/8), all four densities: max abs err 1.1e-4 to 5.5e-4.bench_mixed_attention.py --accuracyagainst the separate wrappers: max abs err ≤ 7.8e-3 bf16 (one ulp at this magnitude), both arches.pre-commit runon every changed file./code-review xhighon the changelist — 15 findings, all 15 addressed; re-measured after.c5cd9cc55./code-review xhighon that review-response commit — 10 taken in3d50bb6ea, 5 declined as structural (see the PR comment).--kv_dtype fp8_e4m3attention row runs on both arches after the fix (0.121 ms gfx942, 0.135 ms gfx950) against 0.127 ms on unmodified6202e794d.Not run: the full
pytestsuite. The changelist touchesbenchmarks/plus one test file, and reaches no kernel or library code.