Correct the asm prefill speedup, which the old benchmark overstated by demandal25 · Pull Request #370 · AMD-Ecosystem/flashinfer

demandal25 · 2026-09-13T19:55:03Z

Summary

The asm prefill speedup published by #365 was measured with a benchmark that released the allocator cache between timing the two arms, so the second arm paid hipMalloc inside its own measurement. #365 fixed the harness but shipped the numbers taken before the fix. These are the numbers the corrected harness gives, re-measured at e802be34e.

What changed

docs/rocm/backends.md — gfx950 gate-region figures, the per-seqlen curves for both architectures, the A/A floor they are judged against, and the non-square correctness claim.
README.md — the headline 1.21× becomes 1.17×.
flashinfer/rocm/arch_caps.py — the policy comment behind the threshold, which carried every superseded claim and is the single source of truth for it.
benchmarks/rocm/bench_asm_vs_cktile.py — docstring only; it told readers to disregard the gfx950 s=256 column as noise, which the tightened floor no longer supports.

No code changes; the routing gate and its threshold are unaffected.

Architecture / design notes

The threshold does not move. 23 of 24 gate-region cells still win or hold and the geomean is comfortably above 1, so qo_len >= 2048 on gfx950 remains the right cut. What changes is that the docs can no longer claim a clean sweep: one admitted cell is a small loss, and saying so is the difference between a measurement and a sales figure.

Two claims in the same paragraph turned out to need re-deriving, not just renumbering — they were written against the old harness and nothing in the first pass re-checked them:

claim	corrected
gfx950 "per-`seqlen` geomean rises monotonically above 1024"	1.01, 1.09, 1.15, 1.14, 1.19, 1.19 — a step down at 3072, inside that length's A/A floor of 0.995-1.003
gfx942 "1.34× at `seqlen` 1024 falls to 0.90× at 1536"	those points are 0.955 and 0.971; the non-monotonic conclusion survives on different data (1.08 at 2048, 0.97 at 3072, 1.06 at 4096, 0.99 at 6144)

The non-square hedge also goes. #365 opened by saying correctness on chunked-prefill shapes was "covered either way" and closed by adding kv_len == 4 * qo_len coverage on the asm arm — the doc still carried the opening position.

Benchmark results

AITER asm vs its CK Tile arm, bf16 head_dim 128, 60 cells (batch 1 and 4 × 16/32/64 q-heads × seqlen 256-6144), causal. Gate region is qo_len >= 2048, 24 cells.

single_prefill_with_kv_cache serves one request and the dispatcher pins args.batch = 1, so only the batch-1 half of the sweep is reachable through it. The batch-1 column is the one the routing decision rests on; the mixed column is what #365 published.

gfx950 / MI350X

	geomean	worst cell	cells below 1.00
published by #365 (mixed, 24 cells)	1.210	1.00	0
corrected harness, mixed 24 cells	1.166	0.975	1
corrected, reachable batch-1, 12 cells	1.138	0.975	1

The regressing cell is b1 hq16 s2048 at 0.9752, outside the batch-1 A/A floor of 0.995-1.008.

gfx942 / MI300X — conclusion unchanged, and the reachable subset makes it stronger: the gate region is a net loss there.

	geomean	worst cell	cells below 1.00
published by #365 (mixed, 24 cells)	1.032	0.80	8
corrected harness, mixed 24 cells	1.024	0.807	9
corrected, reachable batch-1, 12 cells	0.976	0.807	8

A/A control, same box and session. This is the part that makes the 0.975 cell meaningful rather than noise:

	A/A range (whole map)	gate-region geomean
old harness, gfx942	0.903 - 1.099	0.995
corrected, gfx942	0.987 - 1.013	0.998
corrected, gfx950	0.989 - 1.012	1.000

Removing empty_cache() from the per-arm loop cut the A/A spread from roughly ±10% to ±1.3%, which is the direct evidence that the old placement was injecting the noise rather than merely being untidy.

Reproduce: python benchmarks/rocm/bench_asm_vs_cktile.py --aa then without --aa.

Test plan

gfx942 — A/A, sweep and the seven-file prefill suite at the pinned commit, suite rc=0
gfx950 — A/A and sweep at e802be34e; suite running, will update if it is not clean
Numbers in the diff cross-checked against the raw sweep logs rather than retyped
pre-commit run on both changed files

…erstated bench_asm_vs_cktile.py called torch.cuda.empty_cache() between timing the two arms, so the second one re-acquired memory through hipMalloc inside its own measurement. #365 moved that out of the per-arm loop; these are the numbers the corrected harness gives at e802be3. gfx950 at and above qo_len 2048: published geomean 1.210, worst 1.00, 0 of 24 cells below 1.00 corrected geomean 1.166, worst 0.975, 1 of 24 (b1 hq16 s2048) The A/A floor also tightened, from 0.903-1.099 to 0.987-1.013, which is what makes the 0.975 cell a real loss rather than noise -- so "all 24 cells are wins or level" has to go, not just the headline figure. gfx942 moves 1.03 -> 1.02 with 9 rather than 8 cells regressing; the conclusion there is unchanged and it still ships nothing. Also drops the "correctness is covered either way" hedge on non-square shapes: #365 ended up adding kv_len == 4 * qo_len coverage on the asm arm, so that is now a statement of fact about a test rather than an assurance. Co-Authored-By: Claude <noreply@anthropic.com>

Self-review caught that the first pass corrected every figure taken from the old harness but left two same-vintage *claims* standing. Both turn out wrong against the corrected data: - "the per-seqlen geomean rises monotonically above 1024" on gfx950. It does not: 1.01, 1.09, 1.15, 1.14, 1.19, 1.19 -- a step down at 3072. The dip is 0.4%, inside that seqlen's A/A floor of 0.995-1.003, so the shipping conclusion holds, but "monotonically" does not. - "1.34x at seqlen 1024 falls to 0.90x at 1536" on gfx942. Corrected, those points are 0.955 and 0.971. The non-monotonic conclusion survives on entirely different numbers: 1.08 at 2048, 0.97 at 3072, 1.06 at 4096, 0.99 at 6144, every step outside the floor. Also corrects the benchmark docstring, which told readers the gfx950 s=256 column does not survive the A/A floor. Under the tightened floor every per-seqlen geomean in the sweep sits outside it, s=256 included. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The policy rationale in arch_caps.py still contains the superseded measurements and conclusions.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Corrects overstated AITER asm prefill performance documentation without changing routing behavior.

Changes:

Revises gfx942/gfx950 benchmark results and noise-floor analysis.
Updates the headline speedup from 1.21× to 1.17×.
Clarifies non-square correctness coverage.

File summaries

File	Description
`README.md`	Updates headline performance.
`docs/rocm/backends.md`	Documents corrected benchmark conclusions.
`benchmarks/rocm/bench_asm_vs_cktile.py`	Updates benchmark guidance.

Review details

Files reviewed: 3/3 changed files
Comments generated: 1
Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

That comment is the single source of truth for the threshold, and it carried every superseded claim: monotone above 1024, 1.21x geomean, worst cell 1.00, "nothing regresses", and the gfx942 1.34x/0.90x points. Found by review, not by my own sweep: the repo-wide grep I ran to check for surviving copies was piped through `head -20` and cut off before reaching flashinfer/. Re-run untruncated, this was the last one. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟡 Changes recommended

The advertised single-prefill result mixes unreachable batch-4 measurements with batch-1 data, and the policy source comment remains stale.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

docs/rocm/backends.md:440

The corrected evidence here leaves the policy's stated source of truth stale: flashinfer/rocm/arch_caps.py:316-322 still records a monotonic gfx950 curve, 1.21×/1.00 with no regressions, and the old gfx942 1.34×→0.90× points. single_prefill.cuh:52-54 explicitly directs maintainers to that comment for the measurement, so it should be updated alongside these documentation figures.

* **gfx942** — non-monotonic, and the oscillation does not damp with length:
  1.08 at `seqlen` 2048, 0.97 at 3072, 1.06 at 4096, 0.99 at 6144. Every one of
  those steps is outside the A/A floor, so no threshold holds and the arm stays
  unreachable. The same 2048 cut scores geomean 1.02 with 9 of 24 cells
  regressing.

Files reviewed: 4/4 changed files
Comments generated: 2
Review effort level: Balanced

Copilot

🟡 Changes recommended

The benchmark docstring incorrectly claims every per-sequence geomean falls outside the stated noise floor.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 4/4 changed files
Comments generated: 1
Review effort level: Balanced

…n emit single_prefill_with_kv_cache serves one request and the dispatcher pins args.batch = 1 (single_prefill.cuh:181), but the sweep varies batch over {1, 4}. Half of every gate-region aggregate quoted so far describes a shape this path cannot produce. Gate region (qo_len >= 2048), reachable batch-1 subset against the mixed one: gfx950 mixed 24 cells 1.166 batch-1 12 cells 1.138, 1 below 1.00 gfx942 mixed 24 cells 1.024 batch-1 12 cells 0.976, 8 below 1.00 The threshold survives -- 11 of 12 reachable cells win on gfx950 -- and gfx942 reads worse on the reachable subset than on the mixed one, so both decisions stand. Only the published figures were measuring the wrong population. The batch-4 rows stay in the sweep: they are unreachable from here but are what the batched prefill paths see, so they are labelled rather than deleted. Also drops a blanket claim from the benchmark docstring that every per-seqlen geomean sits outside the noise floor. That was true against per-shape A/A ranges but invites reading a 0.1% margin as signal; the instruction is now to compare each cell against its own control. Co-Authored-By: Claude <noreply@anthropic.com>

demandal25 · 2026-09-13T22:02:41Z

Suppressed comment (docs/rocm/backends.md:440) — fixed in d2bd5fbb5: flashinfer/rocm/arch_caps.py now carries the corrected figures, and e073e8467 scopes them to the reachable batch-1 subset.

Copilot

🟡 Changes recommended

The batch-4 results are incorrectly described as representative of batched-prefill execution paths.

Get a fresh assessment by requesting another Copilot review.

Review details

Files reviewed: 4/4 changed files
Comments generated: 2
Review effort level: Balanced

They do not. This sweep times batch-mode mha_fwd against fmha_v3_fwd, while batched prefill goes to group-mode mha_varlen_fwd with args.use_asm_v3 = false (batch_prefill.cuh:60,63) or to mha_batch_prefill (batch_prefill.cuh:162-174). Introduced one commit ago while scoping the headline to batch-1: the rows needed a reason to stay and I gave them one that sounded plausible and is not true. They characterise no shipping path and are kept only so the previously published mixed aggregate stays reproducible. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🟢 Approval recommended

The documentation-only corrections are consistent with the reachable benchmark subset and existing routing behavior.

Review details

Files reviewed: 4/4 changed files
Comments generated: 0 new
Review effort level: Balanced

demandal25 and others added 2 commits September 13, 2026 15:49

Copilot AI balanced review requested due to automatic review settings September 13, 2026 19:55

Copilot started reviewing on behalf of demandal25 September 13, 2026 19:55 View session