Support CUDA-graph capture on AITER batch decode (opt-in)

[demandal25]

Summary

Enable CUDA-graph capture for AITER batch decode via an explicit backend="aiter". Previously backend="aiter" raised under use_cuda_graph=True, so graph-captured decode always fell back to the in-tree fa2 kernel — forfeiting AITER's decode throughput on the serving hot path (vLLM/SGLang capture decode in graphs). backend="auto" is unchanged (still resolves to fa2 under capture), so this change is purely additive and opt-in.

What changed

• flashinfer/decode_rocm.py — Remove the hard error that rejected backend="aiter" under CUDA-graph capture; add a one-time-per-device warning documenting AITER's capture-at-max-sequence-length contract. backend="auto" still routes to fa2 under capture.
• tests/rocm_tests/test_batch_decode_aiter_hip.py — Drop the obsolete "aiter + graph raises" assertion; add test_batch_decode_aiter_cuda_graph_replay[fp16/bf16] (capture at a max seq len, replay shorter, assert parity vs eager AITER); keep test_batch_decode_auto_routes_cuda_graph_to_fa2.
• benchmarks/rocm_benchmarks/bench_decode_graph_hip.py — New: graph capture+replay correctness check plus aiter-vs-fa2 under-graph latency.
• benchmarks/rocm_benchmarks/bench_batch_decode_hip.py — New: eager fa2-vs-aiter paged-decode sweep (the baseline that motivates the feature).
• README.md, CHANGELOG.md — Document the opt-in graph path and its capture-at-max contract.

Architecture / design notes

Why opt-in rather than routed by auto: AITER PA v1's launch grid and compiled .so variant (npar_loops) are fixed at the shapes seen when the graph is captured. This is correct only when you capture at your maximum sequence length — the kernel early-exits per sequence on context_lens, so replays shorter than captured are correct, but longer ones are not. fa2's graph path, in contrast, is capacity-based and capture-order-independent. Because AITER's contract is stricter, auto must not select it silently: doing so regressed test_cuda_graph_batch_decode_with_paged_kv_cache, which captures at 1 page/seq and then replays at the full kv_len.

backend	use_cuda_graph	resolves to
auto	False	AITER when compatible (fp16/bf16 + NHD + NONE), else fa2
auto	True	fa2 (capacity-based, capture-order-independent)
aiter	False	AITER
aiter	True	AITER — caller must capture at max seq len

Known limitation: return_lse=True on the AITER decode path is served through the lazily-built fa2 shadow plan and is not covered under capture here; use return_lse=False (the decode norm) or backend="fa2" if you need LSE under graph.

Benchmark results

Paged decode, GQA 32/8, HD128, bf16, page_size 16, MI300X (gfx942). Under-graph replay latency at batch=128, seq=4096:

backend	ms / replay
fa2	0.889
aiter	0.679

AITER is 1.31× faster under graph at this shape (eager sweep shows ~1.3–1.45× at batch ≥ 32; near parity / fa2-favored at batch ≤ 8, where routing should still prefer fa2). Correctness: graph replay matches eager AITER exactly (max|diff| = 0.0000) for seq ∈ {512, 1024, 2048, 4096} within captured capacity.

Test plan

• tests/rocm_tests/test_batch_decode_aiter_hip.py (full file, --reruns 2, MI300X) — pass
• tests/rocm_tests/test_batch_decode_kernels_hip.py -k cuda_graph (--reruns 2) — pass (confirms auto still routes to fa2 under capture; no regression)
• benchmarks/rocm_benchmarks/bench_decode_graph_hip.py — replay parity vs eager + 1.31× under-graph speedup
• pre-commit run --files on all changed files — pass

[Copilot]

Pull request overview

This PR enables opt-in CUDA-graph capture for ROCm AITER batch decode when callers explicitly set backend="aiter", removing the prior hard error under use_cuda_graph=True. It keeps backend="auto" behavior unchanged (still routes to fa2 under capture) and documents AITER’s stricter “capture at max sequence length, replay shorter only” contract via warnings and docs.

Changes:

• Allow backend="aiter" to run under CUDA-graph capture, emitting a one-time-per-device warning describing the capture-at-max contract.
• Update ROCm tests to validate AITER graph capture + replay correctness, while ensuring backend="auto" still routes to fa2 under capture.
• Add ROCm benchmark scripts for eager aiter-vs-fa2 decode and for under-graph replay correctness + latency.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
flashinfer/decode_rocm.py	Removes the hard failure for AITER+CUDA-graph and adds a one-time warning explaining the capture-at-max contract; preserves auto→fa2 under capture.
tests/rocm_tests/test_batch_decode_aiter_hip.py	Drops the obsolete “AITER+graph raises” expectation and adds coverage for AITER CUDA-graph capture/replay correctness and auto routing.
benchmarks/rocm_benchmarks/bench_decode_graph_hip.py	New acceptance/benchmark script for graph-captured AITER decode correctness and under-graph latency comparisons to fa2.
benchmarks/rocm_benchmarks/bench_batch_decode_hip.py	New eager decode benchmark sweeping batch/kv_len to quantify aiter vs fa2 baseline differences.
README.md	Documents that AITER decode under CUDA-graph is opt-in and requires capture at max seq len; clarifies auto routing behavior.
CHANGELOG.md	Notes the new opt-in AITER CUDA-graph decode support and new benchmarks/tests.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

flashinfer/decode_rocm.py:1138

• The wrapper docstring earlier in this file still states that use_cuda_graph=True is not supported with backend="aiter" (see BatchDecodeWithPagedKVCacheWrapper.__init__ docs around decode_rocm.py:853-856). This PR changes the behavior to allow graph capture with explicit backend="aiter", so the docstring is now misleading for API consumers.

        # Resolve auto → concrete backend. AITER decode requires use_tensor_cores=False
        # (the AITER PA v1 kernel handles its own dispatch internally). Under CUDA-graph
        # capture, `auto` stays on fa2: fa2's graph path is capacity-based and correct
        # regardless of capture-vs-replay sizes, whereas AITER's launch grid and .so
        # variant are fixed at the shapes seen when the graph is captured and require
        # capturing at the maximum sequence length. AITER decode under CUDA graph is
        # therefore opt-in via an explicit backend="aiter" (see the capture-at-max
        # warning below), not something `auto` selects silently.