fix(rocm): correct AITER decode backend gaps — sliding window, CUDA graph, return_lse

[demandal25]

Summary

The AITER PA v1 decode backend on amd-integration has three call patterns that produce wrong output, hard crashes, or unhelpful NotImplementedErrors. This PR fixes each one at the level it's actually broken at, rather than blanket-disabling AITER.

Case	Behavior on amd-integration	This PR
Sliding-window attention (window_left >= 0)	AITER selected. Wrapper passes sliding_window = window_left (off-by-one), and window_left = 0 collides with AITER's "disabled" sentinel — silently wrong output.	AITER runs with corrected convention mapping (window_left + 1).
use_cuda_graph=True with explicit backend="aiter"	AITER selected. Per-plan scalars (max_kv_len, max_blocks_per_seq) are baked into the captured graph; replay against a larger batch launches with an under-sized grid.	Clear ValueError at plan() time. (auto-select already routes to fa2.)
run(return_lse=True)	Raises NotImplementedError("AITER decode backend does not currently return LSE").	Transparent dispatch through a pre-built FA2 shadow plan; one-time warning at plan() time so the per-call backend switch is not silent.

Why each fix

1. Sliding window — wrapper convention bug, not a kernel gap

AITER PA v1's kernel does implement window masking (csrc/cpp_itfs/pa/pa_kernels.cuh:457):

if (local_token_idx + i < context_len - sliding_window)
    tmp = -FLT_MAX;

gated by the sliding_window_enabled template flag (set by the compile step at csrc/cpp_itfs/pa/pa_v1.py:144). The wrapper already plumbs sliding_window through _aiter_pa_v1_resolve and the run-time call.

The bug on trunk is a convention difference:

• FlashInfer: window_left = W → query at position kv_len-1 sees positions [kv_len-1-W, kv_len-1] = W+1 tokens.
• AITER: sliding_window = S (with S > 0 enabling the mask) → admits S tokens.
• S = 0 is AITER's compile-time "disabled" sentinel.

Trunk passes sliding_window = window_left — off by one, plus window_left = 0 (one visible token) collides with AITER's disabled sentinel and silently returns full attention. Fixed: sliding_window = window_left + 1 when window_left >= 0, else 0.

This keeps AITER on the hot path for sliding-window models (Gemma, Mistral, etc.) instead of giving up perf to FA2.

2. CUDA graph — wrapper-level limitation, hard-rejected in explicit path

AITER's launch grid is computed at plan() time from max_kv_len and max_blocks_per_seq of the current batch and passed by value to the kernel launch. Under CUDA-graph capture these scalars are baked into the captured graph and can't be widened on replay against a larger batch.

Supporting this properly would require capturing with worst-case dimensions — a new API parameter (e.g. max_seq_len_per_request) — which is out of scope for this PR. The auto-select fallback to FA2 stays in place; the explicit backend="aiter" path (which on trunk silently produces broken launches) now raises:

ValueError: AITER decode backend is incompatible with CUDA-graph capture:
the kernel's launch grid is sized from per-plan scalars (max_kv_len,
max_blocks_per_seq) that are baked into the captured graph at capture time.
Use backend='fa2' for CUDA-graph workflows, or backend='auto' which routes
around this automatically.

3. return_lse — replace NotImplementedError with transparent fallback

AITER PA v1 does not output LSE (only out; the kernel computes per-partition exp_sums/max_logits internally for split-K but does not expose normalized LSE). Trunk raises NotImplementedError at run() time, breaking any caller that needs LSE under an AITER plan.

This PR pre-builds an FA2 decode plan at AITER plan() time and dispatches through it whenever return_lse=True arrives at run(). This is the only correct option since return_lse is per-call, not per-plan. Two details worth flagging:

• The shadow plan uses the real window_left (and the corresponding template flag), so it produces correct LSE under sliding-window AITER plans (now supported per fix Update Changelog #1).
• A one-time-per-device warning is emitted at AITER plan() time. Without it, a user toggling return_lse=True on a hot path would silently move from AITER → FA2 with no signal.

Tests added

tests/rocm_tests/test_batch_decode_aiter_hip.py:

• test_batch_decode_aiter_sliding_window_vs_fa2 — AITER↔FA2 parity over window_left ∈ {0, 31, 127, 1023}, fp16/bf16, batch sizes, GQA ratios, including the saturation regime (window_left >= max_kv_len-1) that exercises the kernel's no-op masking branch.
• test_batch_decode_aiter_return_lse_via_fa2 — verifies (a) return_lse=False still runs through AITER, (b) return_lse=True falls back to the shadow FA2 plan and returns (output, lse) matching the pure-FA2 reference, with and without sliding window.
• test_batch_decode_auto_routes_cuda_graph_to_fa2 — backend="auto" + use_cuda_graph=True resolves to fa2.
• Extended test_batch_decode_aiter_rejects_invalid_config — explicit backend="aiter" + use_cuda_graph=True raises a ValueError mentioning "CUDA-graph".

Test plan

• pytest tests/rocm_tests/test_batch_decode_aiter_hip.py -v — 178 passed in 69 s. New parity + LSE-fallback + CUDA-graph rejection cases.
• pytest tests/rocm_tests/test_sliding_window_hip.py -m "not slow" — 1248 passed in 60 s. Exercises the AITER path on every sliding-window decode shape that was previously silently wrong on trunk.
• pytest tests/rocm_tests/test_batch_decode_kernels_hip.py -m "not slow" -n auto --reruns 2 — 1872 passed in 174 s. Covers the broader decode matrix including return_lse=True (which on trunk raised NotImplementedError under AITER) and CUDA-graph wrappers (which now route to FA2 cleanly).

API impact

• BatchDecodeWithPagedKVCacheWrapper docstring updated to document the AITER-specific constraints (CUDA-graph incompatible; sliding-window supported transparently; return_lse falls back to FA2).
• No public signature changes.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR updates the ROCm batch-decode wrapper to avoid selecting the AITER decode backend in cases where it produces incorrect results or fails, and to transparently fall back to the FA2 decode path for those cases.

Changes:

• In backend="auto" mode, force FA2 when sliding-window attention is requested (window_left != -1) or when CUDA graph mode is enabled.
• For backend="aiter" execution with return_lse=True, precompute a FA2 decode plan during plan() and use it in run() to return (out, lse) instead of raising NotImplementedError.

Comments suppressed due to low confidence (1)

flashinfer/decode_rocm.py:1182

• The inline comment # window_left == -1 is enforced above is not universally true: the only enforcement is in the backend == "auto" resolution, so callers using backend="aiter" can still reach this code with window_left != -1. Either enforce window_left == -1 in the AITER branch (or handle sliding-window here) or update the comment to reflect the actual constraint.

                head_dim,
                PosEncodingMode[pos_encoding_mode].value,
                False,  # window_left == -1 is enforced above
                logits_soft_cap > 0,

---

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

[Copilot]

Pull request overview

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