[demandal25]

Summary

With FLASHINFER_USE_TORCH_CUSTOM_OPS=1, torch.compile gave wrong attention, overwrote V, or faulted the GPU when K and V were views: two halves of one [2, pages, ...] tensor, one interleaved tensor under dynamic=True, or a layer of a multi-layer pool in any mode. The read-only attention ops declared the paged KV cache mutated, which exposes a torch 2.12 bug in how declared-mutated views are rebuilt. They now declare only what they write, under new op names so torch's on-disk compile cache cannot replay the old graphs. AITER decode shim defects found on the way are fixed first: an untracked allocation that broke mode="reduce-overhead", and K/V/q layouts the kernel cannot address being read wrongly without an error.

What changed

• csrc/rocm/batch_decode_aiter.cu (first commit): the unit k/v scale passed to PA v1 lives in static device memory instead of a lazily allocated tensor, which landed untracked in torch's CUDA-graph pool. K and V must share page shape, strides and q's device and have a dense head_dim; a q with spread-out heads is copied to a packed one. Before, a V laid out differently from K or a head-sliced q gave wrong attention with no error.
• flashinfer/rocm/decode.py (first commit): run() allocates its default out contiguous rather than like q. Every decode kernel writes o densely, so for a head-major q AITER refused the call and fa2 wrote its rows in the wrong order without an error.
• flashinfer/rocm/decode.py: fa2 batch decode and AITER decode no longer declare paged_k_cache/paged_v_cache, and register as {uri}_run_kvro.
• flashinfer/rocm/prefill.py: the same for the fa2 paged_run op, now {uri}_paged_run_kvro. AITER paged prefill already declared only o and lse. The jit_args variants keep their declarations: they take *args, which schema inference rejects, so they never register.
• README.md: a known-issue note under torch.compile for append over cache views, and the module-buffer advice for mode="reduce-overhead".
• tests/rocm/test_batch_decode_aiter.py: test_batch_decode_aiter_kv_and_q_layouts (a V with HND strides raises; a head-sliced q, a head_dim-strided q and a shorter V pool match fa2); test_batch_decode_aiter_batched_q_layouts (batch 3; fa2, tensor-core and AITER decode each read a q sliced from fused QKV and a head-major q into a contiguous out matching fa2 on a packed q).
• tests/rocm/test_torch_compile.py: test_torch_compile_aiter_decode_reduce_overhead_rerecords (first AITER call inside a reduce-overhead warm-up, with a head-sliced q, re-records, replay against an fa2 reference taken before compiling); test_torch_compile_kv_cache_views (fa2, tensor-core and AITER decode, fa2 paged prefill; fullgraph; [2, pages] views, an interleaved tensor under dynamic=True, layer 1 of a pool; reference before compiling, every cache tensor bit-identical after, each case reported); test_torch_compile_append_and_attention_order (append of half the positions before and after each route in one fullgraph graph, over separate caches and over views of one tensor under dynamic=True). The new compile tests turn off Inductor's on-disk caches.

Architecture / design notes

Max |compiled − eager| before this change, identical on gfx942 and gfx950 for fa2 decode, AITER decode and fa2 paged prefill:

K/V passed in	dynamic=False	default	dynamic=True
One interleaved [pages, 2, ...] tensor	0	0	1.05–1.26
(kv[:, 0], kv[:, 1])	0	0	0
(kv[0], kv[1]) of a [2, pages, ...] tensor	6.34	6.34	0
Separately allocated	0	0	0
Layer 1 of a two-layer pool	GPU memory fault	GPU memory fault	GPU memory fault

After it, every cell is exact for every op on both arches. Two alternatives were measured and rejected: enable_auto_functionalized_v2 = False fixes only the dynamic=True column, and declaring K alone on AITER decode still faults on layer 1 of a pool.

Why the ops are renamed. Inductor's FX-graph and AOTAutograd caches do not key on an op's alias schema. With a TORCHINDUCTOR_CACHE_DIR filled by the released ops, the fixed code replayed the old graph and was wrong again (6.34, both arches); under the new names it misses the cache and is exact. The default cache directory persists across runs, so without the rename an upgrade would keep the bug for anyone who had compiled with the flag on.

mode="reduce-overhead". The fa2 ops write their workspaces, so CUDA graphs stay skipped for them as before. AITER decode is now captured, which the shim fix makes safe. With a 1 GiB cache held as a module buffer it is neutral (67.87 → 67.86 µs gfx942, 42.82 → 42.51 µs gfx950); a cache passed as an argument is copied on every replay, as torch does for any non-static input (77–80 → 603 µs gfx942, 40–43 → 506 µs gfx950). The README now tells reduce-overhead users to register the cache as a module buffer.

Append is unchanged. It genuinely writes the cache and keeps both declarations, so it keeps the bug, including the layer-1 fault; the README now documents it and says to allocate K and V separately.

Benchmark results

Per call, minimum of 3 interleaved rounds on one idle card per arch, base caa2d8e → both commits, µs:

Call	gfx942	gfx950
Flag off, AITER decode	27.98 → 27.70	20.46 → 20.28
Flag off, fa2 decode	38.43 → 38.14	32.46 → 32.68
Flag off, tensor-core decode	88.26 → 88.17	42.96 → 43.03
Flag off, fa2 paged prefill	339.60 → 338.88	188.53 → 188.38
Flag on, eager AITER decode	41.95 → 41.33	26.71 → 25.19
Flag on, eager fa2 decode	41.42 → 42.15	32.55 → 32.58
Flag on, eager tensor-core decode	88.57 → 88.24	42.91 → 43.06
Flag on, eager fa2 paged prefill	339.58 → 339.31	188.57 → 188.55
Flag on, compiled AITER decode	69.29 → 70.67	39.33 → 35.64
Flag on, compiled fa2 decode	71.78 → 68.85	36.62 → 35.14
Flag on, compiled tensor-core decode	89.59 → 88.45	45.70 → 43.48
Flag on, compiled fa2 paged prefill	340.07 → 340.08	188.47 → 188.66

No regression: every slowdown is under 1.4 µs, and none exceeds 0.2 µs on both arches at once. The compiled gfx950 gains come from a single round and are not claimed.

Test plan

• A/B on caa2d8e's ops and shim, the first commit alone, and both commits: kv_cache_views fails in all 16 route × case cells on the old code and passes with both commits; the reduce-overhead test fails on the old shim (untracked pool tensor) and passes from the first commit on; the two layouts tests fail on the old shim (did not raise; non-contiguous output refused) and pass from the first commit on; append_and_attention_order passes throughout. gfx942 and gfx950.
• Layout × compile-mode matrix above, one process per cell, before and after: the first four layouts on both arches, the layer-1 row on gfx942; the new test covers layer 1 on both.
• Upgrade: a shared TORCHINDUCTOR_CACHE_DIR filled by the released ops, then read by this change: exact under the new names, both arches.
• test_torch_compile, test_batch_decode_aiter, test_batch_decode_kernels, test_batch_decode_vllm, test_batch_decode_speculative, test_decode_guards, test_cascade and test_deprecated_forward_api, flag off and on: gfx942 2758 / 2758 passed; gfx950 2758 / 2758 passed. Earlier in the branch, 15 decode, prefill and compile files passed on both arches (9050 / 9048).
• /code-review xhigh on origin/amd-integration...HEAD, nine rounds; findings applied in the commits or declined.
• pre-commit run -a

[Copilot]

🔵 Needs a closer look

Compiler aliasing and HIP graph-capture lifetime changes warrant final human review despite the reported multi-architecture validation.

No open issues 🎉

What changed in this PR

Fixes compiled read-only paged attention by correcting cache mutation declarations and hardening AITER decode.

Changes:

• Removes false K/V mutation declarations and renames ops to invalidate cached graphs.
• Fixes AITER scale storage and layout handling, and makes default decode outputs contiguous.
• Adds regression tests and documents remaining append limitations.

File	Description
tests/​rocm/​test_torch_compile.py	Tests cache views, append ordering, and graph re-recording.
tests/​rocm/​test_batch_decode_aiter.py	Tests query and cache layouts across decode routes.
README.md	Documents append risks and module-buffer advice.
flashinfer/​rocm/​prefill.py	Corrects and renames the read-only paged op.
flashinfer/​rocm/​decode.py	Corrects mutation declarations, renames ops, and allocates contiguous outputs.
csrc/​rocm/​batch_decode_aiter.cu	Uses static device scale storage and validates supported layouts.

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

GPU aliasing and graph-capture lifetime changes warrant hardware-backed human review despite the extensive reported validation.

No open issues 🎉

_{Review effort: Balanced}

[github-actions]

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

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.