Refuse attention inputs the ROCm kernels would compute wrongly by demandal25 · Pull Request #399 · AMD-Ecosystem/flashinfer

demandal25 · 2026-10-27T02:46:52Z

Summary

Closes eight correctness and diagnostics gaps in the ROCm attention entry points. Five of them returned wrong answers or killed the process with no error: a q/k/v whose last dimension is strided was read as if packed; run() tensors that disagree with plan() ran silently or crashed; head counts the GQA group size cannot divide ran with a truncated group; AITER decode's return_lse fallback planned from a caller buffer that may have changed since plan(); and a graph-mode partial_state merge wrote past a batch smaller than the first one planned. The other three make API logging reach the ROCm wrappers, give the compile guard's message its real cause, and record that the documented view miscompile is a PyTorch defect.

What changed

Strided layouts

csrc/rocm/*.cu (single/batch prefill, batch/single decode, POD, batch POD, and the three AITER prefill shims) — CHECK_LAST_DIM_CONTIGUOUS on what each shim reads; the fa2 and POD shims also require a contiguous o/lse, and paged prefill's partial_state must be contiguous and shaped like o.
flashinfer/rocm/decode.py — single_decode_with_kv_cache copies its one-token q instead of refusing it.
tests/rocm/test_strided_last_dim.py — every route, input, output and POD variant.

Plan-vs-run shapes

flashinfer/rocm/prefill.py — _check_run_matches_plan (rank, heads, head dims, page size, minimum token counts), used by paged and ragged prefill and batch decode; a plan() that raises after its first write leaves run() refusing.
flashinfer/rocm/{decode,prefill}.py plan() — batch and argument validation now precede the first write, so a rejected batch leaves the previous plan usable; the CUDA-graph row capacity is stored as an int, not a view of the caller's CPU qo_indptr.
tests/rocm/test_run_refusals.py, tests/rocm/test_torch_compile.py — each mismatch on each route, rejected plans, and a fullgraph compile over re-plans.

GQA head divisibility

csrc/rocm/*.cu, pytorch_extension_utils.h — CHECK_GQA_HEAD_DIVISIBLE in the fa2 prefill and decode plans and every attention run shim; it also refuses zero kv heads.

AITER decode LSE fallback

flashinfer/rocm/decode.py — the lazily built FA2 shadow plan keeps a clone of the planned indptr, not the caller's CPU tensor.
tests/rocm/test_batch_decode_aiter.py — a caller that restages its CPU indptr before the first return_lse run.

Graph-mode `partial_state` merge

include/flashinfer/rocm/attention/{cascade,prefill}.cuh — MergeStateInPlace takes an optional device row count and the paged run passes the graph's total_num_rows, the bound the VariableLengthMergeStates launch beside it already uses.
csrc/rocm/batch_prefill.cu — partial_state on a split-KV plan requires maybe_lse, which the merge reads.

API logging

flashinfer/rocm/{prefill,decode,mla,fused_moe}.py — plain @flashinfer_api on the public entry points.
tests/rocm/test_api_logging_rocm.py, docs/rocm/environment.md, .claude/skills/debug-rocm-crash/SKILL.md.

Compile guard message, view miscompile

flashinfer/rocm/torch_compile.py — with the flag on, an op torch cannot register says why and to call it outside the compiled region, instead of telling the user to set the flag that is already set.
docs/rocm/torch-compile.md — one sentence: the miscompile reproduces without FlashInfer.

Architecture / design notes

Strided last dim raises rather than copies, except single_decode's q. A KV cache can be gigabytes, so a silent .contiguous() would trade a wrong answer for a hidden copy per call; a one-token q is a few KB.
Token counts are minimums, not equalities, because CUDA-graph callers pad q.
Not covered by the plan check: MLA (its run() derives shapes from q_nope and calls AITER's own op) and the POD, batch POD and block-sparse wrappers, whose Python is upstream's. Their C++ shims still get the layout and GQA checks.
Plain @flashinfer_api, not upstream's trace=. With a template the decorator wraps even at level 0; plain returns the function itself, so the default path costs nothing. fi_trace stays unavailable for ROCm attention.
No copy-back guard for the view miscompile. A trivial custom op reproduces it, and append writes scattered pages, so staging a cache view would copy the whole layer cache in and out per call.

Before this PR (gfx950):

Case	Result
q, k, v or cache with `stride(-1) == 2`	max \|diff\| 0.98–1.40 vs contiguous on fa2 single/paged/ragged prefill and batch decode, and AITER single/paged-q/ragged prefill
`run()` heads / head_dim / page_size / rank / tokens off the plan	ran without error, or died: SIGFPE, HSA hardware exception, abort
12 q heads over 8 kv heads	wrong output on fa2 single prefill and batch decode; CK assert and abort on AITER prefill; 4 over 8 is a SIGFPE
AITER decode `return_lse` after the caller restages its CPU `indptr`	7503 of 8192 output elements wrong, max \|diff\| 0.10
Graph-mode paged prefill re-planned from 16 to 8 rows, with `partial_state`	rows 8-15 of the buffer behind an 8-row `out` overwritten

View miscompile, trivial torch.library.custom_op, no FlashInfer code, max |compiled − eager|:

	`dynamic=None`	`False`	`True`
torch 2.12, gfx950	0	0	5.96–12.42
torch 2.13, gfx942	0	0	5.96–12.42
torch 2.13, CPU	0	0	5.26–12.77

Not changed: with the flag off, AITER prefill and aiter_fused_moe still fail under torch.compile with torch's fake-tensor error. A Python guard made a default compile fall back to eager silently, and torch refuses register_fake on these extension ops.

Benchmark results

Kernels are unchanged; the cost is host time per call. Median wall time per call over 5000 calls at small, host-bound shapes, minimum of three interleaved rounds per arm, base 5f2e2b2 vs head 868faa19 (same tree as the final head), bf16, one idle card per arch (GPU use 0% before and after), both arms on the same mount path. All values µs.

Shapes: 8 q and 8 kv heads, head_dim 128, page size 16. Decode is one page per request at batch 1 and 64; paged prefill is 16 q rows over 4 pages; ragged prefill is 16 q rows over 64 kv rows; single_decode attends over 64 kv rows; MoE is 4 tokens, 8 experts, top-2, hidden 512, intermediate 256.

Route	gfx942 base	gfx942 head	Δ	gfx950 base	gfx950 head	Δ
decode fa2, batch 1	14.44	15.14	+0.70	8.95	9.37	+0.42
decode fa2, batch 64	14.60	15.52	+0.92	9.03	9.48	+0.45
decode AITER, batch 1	22.35	23.12	+0.77	15.02	15.38	+0.36
decode AITER, batch 64	23.00	23.82	+0.82	15.27	15.60	+0.33
paged prefill fa2	19.92	20.39	+0.47	11.78	12.21	+0.43
paged prefill AITER	63.95	64.36	+0.40	34.37	34.43	+0.06
ragged prefill AITER	15.86	16.64	+0.78	9.53	10.01	+0.48
Mean, the seven checked routes			+0.70 (+3.7%)			+0.36 (+3.3%)
`single_decode` (control: no plan check)	11.61	11.76	+0.15	7.43	7.48	+0.05
`aiter_fused_moe` (control: decorator only)	30.67	30.68	+0.02	19.78	19.69	−0.09

The added work is the plan-vs-run compare: 0.67 µs in isolation on the gfx942 host and 0.35 µs on the gfx950 one, about half of it reading three .shapes. Decode under CUDA graphs, which is how vLLM and SGLang run it, replays none of it. The controls confirm @flashinfer_api costs nothing at the default level. The two commits after 868faa19 change only plan() (one small clone) and the graph-mode partial_state merge, which none of these routes reach.

Test plan

Full pytest -n auto --reruns 2 -m "not slow" at 868faa19: gfx942 33,977 passed / 4,026 skipped, gfx950 33,987 passed / 4,016 skipped, 0 failed
At the head (1c7bafd), touched files plus test_cascade.py, test_shared_prefix_kernels.py and test_batch_decode_aiter.py: gfx950 845 passed; gfx942 882 passed and 6 cascade/shared-prefix failures at --reruns 0, which fail as often on 868faa19 on the same card (5 vs 1 over three rounds of the six tests)
FLASHINFER_USE_TORCH_CUSTOM_OPS=1 on the compile, logging, refusal and stride test files at the head: gfx942 166 passed, gfx950 166 passed
A/B: every new refusal, rejected-plan, LSE-fallback and partial_state test fails on the tree before its fix (or kills its worker), except the head-sliced q on the tensor-core path, which already honoured strides, and the GQA cases fa2 batch prefill and AITER decode already refused
pre-commit run -a

The attention kernels honour the token and head strides of their inputs but read head_dim as packed, and the fa2 kernels write o and lse packed. Nothing checked either, so a strided tensor returned a wrong answer with no error. Measured on gfx950, max |strided - contiguous| for an input with stride(-1) == 2 (e.g. big[..., ::2]): single prefill fa2 / AITER, q or k 1.03 / 1.03, 1.23 / 1.18 ragged prefill fa2 / AITER, q or k 1.01 / 1.22, 1.12 / 1.31 paged prefill fa2, q or cache; AITER q 1.22, 1.40; 1.17 batch decode fa2 (+ tensor cores), q 1.09 (1.03) single_decode, tensor cores, q 0.98 Each shim now applies CHECK_LAST_DIM_CONTIGUOUS to what it reads -- on the AITER paged route only on the native branch, since the gather branch copies a cache that is not row-contiguous -- and the fa2 and POD shims require a contiguous o and lse, so a strided out= is refused rather than filled wrongly; paged prefill's partial_state, merged with o's packed strides, must be contiguous and shaped like o. POD shares the prefill and decode kernels, so it gets the same checks. AITER decode already copied q. The macro gains the missing space in its message. single_decode copies its one-token q instead of refusing it, so its fa2 path also accepts a non-contiguous q, which it used to reject with "q must be contiguous". Every new test fails on the unfixed tree except the head-sliced q on the tensor-core path, which honoured strides already. Co-Authored-By: Claude <noreply@anthropic.com>

The prefill and decode kernels are compiled and scheduled for the heads, head dims, page size and token counts given to plan(), but run() never compared its tensors against them. On the unfixed tree a run() tensor smaller than planned either ran silently (pytest: DID NOT RAISE) or killed the process with a floating-point exception, an HSA hardware exception or an abort -- five xdist workers died on gfx950 -- across fa2, fa2 tensor-core and AITER decode and fa2/AITER paged and ragged prefill. run() now raises ValueError naming the field, the planned value and the tensor's. Prefill also requires at least the planned q rows (and, ragged, k/v rows): out is sized from q, so a short q was written past its end. At least, not exactly, since a CUDA-graph caller may pad. The check runs after the paged cache is unpacked, so tuple, single-tensor and HND caches compare by layout; K is checked against head_dim_qk and V against head_dim_vo. Rank is checked too, since every shim reads heads at q.size(1): a q with an extra leading dim would match on its last two dims and still reach the SIGFPE. A plan() that raises after its first write to the wrapper leaves it half-planned, so run() now refuses until plan() succeeds. Batch decode and paged and ragged prefill now finish validating the batch -- graph capacity, CUDA-graph batch size, row capacity and indices length, kv_indptr length, and decode's explicit-AITER argument checks -- before that first write, so a rejected batch still leaves the previous plan usable; decode already promised this for replayed graphs. Later rejections (the C++ plan, custom-mask buffers) still leave run() refusing. Decode's run() checks before it infers q_len_per_req, so a mismatch names the field rather than a q_len_per_req it never had. The CUDA-graph row capacity is now stored as an int: it was a 0-d view of a CPU qo_indptr, which is the caller's own buffer, so a caller that reused it moved the capacity the captured graph was sized for. The refusal is a length test rather than an identity test because Dynamo cannot trace `is` between tuples, and a compiled run() over re-plans with new token counts compiles no more frames than the base (2 for 4 plans, gfx942). MLA is not covered: its run() derives shapes from q_nope and calls AITER's own op. Nor are POD, batch POD and block-sparse attention, whose Python wrappers are upstream's. The common path is one flat compare: 0.35 us in isolation on the gfx950 host and 0.67 us on the gfx942 one, about half of it reading the three .shape attributes. Per call, host time on the checked routes rises 0.3 to 0.5 us on gfx950 and 0.4 to 0.9 us on gfx942; single_decode and aiter_fused_moe, which skip the check, are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>

No function under flashinfer/rocm/ carried @flashinfer_api, so on ROCm the API log covered norm, RoPE, sampling and page but never prefill, decode, MLA or MoE -- the calls a crash report most needs. Decorated: the single prefill/decode functions and their jit-module variants, __init__/plan/run of the paged and ragged prefill, batch decode and MLA wrappers, and aiter_fused_moe. The CUDA-graph decode subclass inherits its parent's. The decorator is plain, unlike upstream's trace= on run and single_*: with a template it wraps even at level 0, while plain hands the function back unchanged, so the default path costs nothing (tested). The trade is that fi_trace and FLASHINFER_TRACE_DUMP stay unavailable for ROCm attention. Levels 1, 3, 5 and 10 are tested on the fa2 routes and level 1 on MLA and MoE. Above 0 the wrapper calls datetime.now, which Dynamo cannot trace, so a logged call breaks the graph and fullgraph=True fails, as it already did for the upstream-decorated ops; environment.md now says so. Co-Authored-By: Claude <noreply@anthropic.com>

With FLASHINFER_USE_TORCH_CUSTOM_OPS=1, an op whose signature torch's schema inference rejects (the fp8-quantizing norms, the sampling ops that take a generator) falls back to the compile guard, whose message told the user to set the flag that was already set. The guard now carries the inference error and says to call the op outside the compiled region. Not changed: with the flag off, AITER prefill and aiter_fused_moe still fail under torch.compile with torch's fake-tensor error rather than one naming the flag. Guarding them in Python made a default compile fall back to eager silently and cost every eager call a frame; a raising fake would be free, but torch refuses register_fake on these ops because the extension registers them as catch-all CompositeImplicitAutograd. Measured on gfx950 before reverting: base fails loudly in every case, the guard made all four AITER ops run eagerly under a default compile. Co-Authored-By: Claude <noreply@anthropic.com>

docs/rocm/torch-compile.md says an op writing into a view of a larger tensor can miscompile. The same cases reproduce with a trivial custom op, no FlashInfer code (torch.library.custom_op whose impl is out.copy_(x * 2 + 1) or an index_copy_ scatter), max |compiled - eager|: dynamic=None False True torch 2.12, gfx950 0 0 5.96 - 12.42 torch 2.13, gfx942 0 0 5.96 - 12.42 torch 2.13, CPU 0 0 5.26 - 12.77 over an out= slice, two views of one base, an unbind of a [pages, 2, ...] cache and kv[0]/kv[1] of a [2, pages, ...] one; a separately allocated output is exact everywhere. Two views of one buffer passed as separate graph inputs fault the GPU with the trivial op too (torch 2.12, gfx950), under auto_functionalized v1 and v2 alike. FlashInfer's own ops match the trivial op case for case, so there is nothing to fix on our side and no copy-back guard: append writes scattered pages, so staging a cache view would copy the whole layer cache in and out per call. Co-Authored-By: Claude <noreply@anthropic.com>

The attention shims compute num_qo_heads / num_kv_heads, and only single_decode, the fa2 prefill scheduler and AITER decode's plan() checked the remainder. Measured on gfx950, 12 q heads over 8 kv heads: single prefill, fa2 runs, wrong output batch decode, fa2 runs, wrong output single/paged/ragged prefill, AITER CK assert, abort and 4 over 8 on single prefill fa2 is a SIGFPE. CHECK_GQA_HEAD_DIVISIBLE now runs in the fa2 prefill and decode plans and in every prefill, decode and POD run shim, fa2 and AITER, so each route refuses with the same message. It also refuses zero kv heads, where the modulo itself would trap. Co-Authored-By: Claude <noreply@anthropic.com>

github-actions · 2026-10-27T02:54:51Z

Security scan found no critical or high findings on lines changed by this pull request

2 medium, low, or unrated finding(s) are not listed; the scan run has them.

View the security scan run for complete reports.

“New” means the finding is located in a changed file on an added or modified line; it is not a base-versus-head findings delta.

Copilot

🟡 Changes recommended

Partial-state validation still permits undersized buffers after a shrinking graph-mode replan, allowing out-of-bounds memory access.

1 open issue

Prevent split-KV buffer overrun after shrinking replan · New

What changed in this PR

This PR hardens ROCm attention entry points against unsupported layouts and mismatched plans, replacing silent incorrect results or crashes with explicit errors.

Changes:

Add tensor-layout, plan/run shape, and grouped-query head-divisibility checks.
Preserve usable plans after early validation failures and reject partially updated plans.
Enable API logging, clarify compile diagnostics, and add regression tests and documentation.

File	Description
`tests/rocm/test_torch_compile.py`	Test recompilation behavior and registration-error diagnostics.
`tests/rocm/test_strided_last_dim.py`	Cover strided inputs, outputs, and partial-state validation.
`tests/rocm/test_run_refusals.py`	Cover shape mismatches, invalid heads, and rejected plans.
`tests/rocm/test_api_logging_rocm.py`	Verify logging levels and compile interaction.
`flashinfer/rocm/torch_compile.py`	Report the actual custom-op registration failure.
`flashinfer/rocm/prefill.py`	Validate run shapes, track plan validity, and enable logging.
`flashinfer/rocm/mla.py`	Enable public API logging.
`flashinfer/rocm/fused_moe.py`	Enable MoE API logging.
`flashinfer/rocm/decode.py`	Validate plans, normalize single-token queries, and enable logging.
`docs/rocm/torch-compile.md`	Identify the view miscompile as a PyTorch defect.
`docs/rocm/environment.md`	Document logging’s effect on compiled graphs.
`csrc/rocm/single_prefill.cu`	Check tensor layouts and head divisibility.
`csrc/rocm/single_prefill_aiter.cu`	Check AITER input layouts and head divisibility.
`csrc/rocm/single_decode.cu`	Require contiguous output and log-sum-exp buffers.
`csrc/rocm/pytorch_extension_utils.h`	Harden shared head and stride checks.
`csrc/rocm/pod.cu`	Validate combined prefill/decode layouts and heads.
`csrc/rocm/batch_ragged_prefill_aiter.cu`	Validate ragged AITER layouts and heads.
`csrc/rocm/batch_prefill.cu`	Validate layouts, heads, and partial-state buffers.
`csrc/rocm/batch_prefill_paged_aiter.cu`	Validate query and native paged-cache layouts.
`csrc/rocm/batch_pod.cu`	Validate batched POD layouts and heads.
`csrc/rocm/batch_decode.cu`	Check plan/run head divisibility and tensor layouts.
`csrc/rocm/batch_decode_aiter.cu`	Check AITER decode head divisibility.
`.claude/skills/debug-rocm-crash/SKILL.md`	Update debugging guidance to include API logging.

_{Review effort: Balanced}

AITER paged attention has no log-sum-exp output, so a return_lse=True run on an AITER-planned decode wrapper goes through an FA2 shadow plan that is built lazily, at the first such run. plan() kept indptr_host for it, and indptr.to("cpu") returns the caller's own tensor when indptr is already on the CPU. A caller that reused that buffer for the next batch before the first return_lse run got an FA2 plan for the new batch run against the old one's page table. Measured on gfx950, four requests of 128/256/512/1024 tokens re-staged as 64/32/16/8 pages (same total): 7503 of 8192 output elements wrong, max |diff| 0.10. plan() now keeps a clone, a batch_size + 1 int copy at plan time only. Co-Authored-By: Claude <noreply@anthropic.com>

Paged prefill merges partial_state into the output with MergeStateInPlace over plan_info.total_num_rows, unmasked. In CUDA-graph mode that is the capacity fixed by the first plan(), and the graph path always splits KV, so after a re-plan for fewer rows the merge read and wrote past an output and partial_state sized for the new batch. Measured on gfx950, a 16-row capacity re-planned to 8 rows: rows 8-15 of the buffer behind an 8-row out were overwritten. The merge kernel now takes an optional device-side row count and returns early past it; the paged run passes params.total_num_rows, the same bound the VariableLengthMergeStates launch beside it already uses. Eager plans pass null and keep their exact capacity. That leaves exact-size buffers valid and skips the junk rows a padded caller would otherwise merge. The run also requires maybe_lse whenever partial_state meets a split-KV plan, since the merge reads it unconditionally. Reported by the PR review bot. Co-Authored-By: Claude <noreply@anthropic.com>

Copilot

🔵 Needs a closer look

The cross-backend GPU memory-safety and graph-replanning changes warrant final human review despite the reported regression results and resolved prior overflow.

No open issues 🎉

1 resolved since last review

Prevent split-KV buffer overrun after shrinking replan

_{Review effort: Balanced}

demandal25 and others added 6 commits October 23, 2026 09:10

Copilot AI balanced review requested due to automatic review settings October 27, 2026 02:46

Copilot started reviewing on behalf of demandal25 October 27, 2026 02:47

View session

Copilot AI reviewed Oct 27, 2026

View reviewed changes

Comment thread csrc/rocm/batch_prefill.cu Resolved

demandal25 and others added 2 commits October 26, 2026 23:12

Copilot AI balanced review requested due to automatic review settings October 27, 2026 04:06

Copilot started reviewing on behalf of demandal25 October 27, 2026 04:06

View session

Copilot AI reviewed Oct 27, 2026

View reviewed changes

demandal25 merged commit 8764a46 into amd-integration Oct 27, 2026
15 checks passed

demandal25 deleted the attention-input-guards branch October 27, 2026 05:15

Navigation Menu

Repository navigation

Refuse attention inputs the ROCm kernels would compute wrongly - #399

Refuse attention inputs the ROCm kernels would compute wrongly#399
demandal25 merged 8 commits into
amd-integrationfrom
attention-input-guards

demandal25 commented Oct 27, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

demandal25 commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What changed

Strided layouts

Plan-vs-run shapes

GQA head divisibility

AITER decode LSE fallback

Graph-mode partial_state merge

API logging

Compile guard message, view miscompile

Architecture / design notes

Benchmark results

Test plan

Uh oh!

github-actions Bot commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

demandal25 commented Oct 27, 2026 •

edited

Loading

Graph-mode `partial_state` merge

github-actions Bot commented Oct 27, 2026 •

edited

Loading