fix(xqa): fix PDL load ordering and SM90 fp8 draft-mask dispatch

[yichengj0]

📌 Description

Follow-up to #4137, addressing the review feedback that arrived after auto-merge. Thanks @qsang-nv for the detailed analysis.

Issues:

• Under programmatic dependent launch (PDL), the generic XQA kernel (csrc/xqa/mha.cu) read q_cu_seq_lens and the scale tensors before the acquire that makes a producer kernel's writes visible. q_cu_seq_lens drives the output row offset, so a stale read could write into the wrong request's rows.
• On SM90 with fp8 KV cache, the small-batch layout of the Hopper XQA kernel (csrc/xqa/mha_sm90.cu, used when q_seq_len * head_group_ratio <= 32) hardcodes a causal draft mask, so a full draft mask silently returned causal results. Existing refchecks use random data and cannot tell the two modes apart within fp8 tolerance.
• The XQA trace template declared a q_cu_seq_lens input that the trace reference silently dropped.

Fixes:

• Move the q_cu_seq_lens and scale-tensor loads below the PDL acquire in mha.cu.
• Extend the SM90 fp8 fallback so small-batch speculative decode also runs on the generic kernel, matching the sliding-window and ragged-Q cases. Restoring the Hopper fast path is tracked in [Bug] mha_sm90.cu: spec-dec draft mask ignored in the SWAP_AB layout, wrong sliding-window masking for full masks, no ragged Q, scale loads before the PDL acquire #4198.
• Add a deterministic mask test: zero Q and K make each output row an exact mean of the visible V values, so any deviation from the requested mask fails loudly on every architecture. It runs two shapes, one that falls back to the generic kernel and one that stays on the Hopper kernel on SM90 fp8.
• Remove q_cu_seq_lens from the trace template until the trace reference supports ragged Q.
• Normalize the ragged-Q module key inside the module getter, and build a separate ragged variant only when an SM90 target is compiled; on other targets it is identical to the uniform module.
• Document the SM90 fp8 fallback in the xqa(), xqa_batch_decode_with_kv_cache, and trtllm_batch_decode_with_kv_cache docstrings.

🔍 Related Issues

#4198 (restore the Hopper fp8 fast path for speculative decode). Review thread: #4137.

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Added test_xqa_batch_decode_mask_mode_deterministic (causal/full mask, bf16/fp8 KV, two head-group shapes) with exact expected outputs.
• On SM120 (RTX 5080): the new test, the ragged-Q and sliding-window suites (256 cases), and the trace suite (970 cases) pass.
• SM90/SM100 are covered by CI. On SM90 fp8 the new test's head_grp_size=16 shape exercises the Hopper kernel and the head_grp_size=4 shape exercises the widened fallback.

Reviewer Notes

• Behavior change on SM90 with fp8 KV cache: small-batch speculative decode now runs on the generic kernel, including causal masks. The mask lives on the device, so dispatch cannot check its content without a sync. The Hopper kernel previously returned causal results regardless of the requested mask; [Bug] mha_sm90.cu: spec-dec draft mask ignored in the SWAP_AB layout, wrong sliding-window masking for full masks, no ragged Q, scale loads before the PDL acquire #4198 restores that fast path.
• The SPEC_Q_SEQ_LEN build specialization is no longer reachable at run time but is kept: the fast-path restoration in [Bug] mha_sm90.cu: spec-dec draft mask ignored in the SWAP_AB layout, wrong sliding-window masking for full masks, no ragged Q, scale loads before the PDL acquire #4198 re-enables it.
• Ragged workloads cannot be represented in trace dumps until the trace reference supports ragged Q; removing the inert input beats shipping a wrong reference.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Improved speculative decoding handling for variable-length queries, including safer execution ordering for mask/sequence offsets.
  • Refined SM90 FP8 KV-cache kernel selection, extending the conditions that fall back to the generic kernel (ragged queries, attention sinks, sliding-window, and small head-group bound).
  • Improved compilation/caching behavior to better match supported speculative-decoding configurations.

• Documentation

  • Updated the XQA documentation to clarify when the generic kernel is used for SM90 FP8 speculative decoding.

• Tests

  • Added a deterministic test covering causal/full speculative-decoding masks for both BF16 and FP8 KV caches.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

📝 Walkthrough

Walkthrough

XQA speculative-decoding execution now computes launch state after acqBulk(), centralizes JIT specialization decisions, expands SM90 FP8 fallback conditions, removes q_cu_seq_lens from the trace schema, updates documentation, and adds deterministic mask-mode tests.

Changes

XQA speculative decoding

Layer / File(s)	Summary
Post-acquisition kernel state and validation csrc/xqa/mha.cu, tests/attention/test_xqa_batch_decode.py	Scale and variable query-length values are computed after bulk acquisition, and deterministic causal/full-mask tests cover BF16 and FP8 KV caches.
JIT specialization and module caching flashinfer/jit/xqa.py, flashinfer/xqa.py	SM90 detection, SWAP_AB eligibility, ragged-Q build changes, specialization flags, and cached module inputs are centralized.
Runtime kernel routing and trace contract flashinfer/xqa.py, flashinfer/decode.py, flashinfer/trace/templates/attention.py	Additional speculative-decoding cases use the generic kernel, documentation reflects routing, and q_cu_seq_lens is removed from trace inputs.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related issues

• [Bug] mha_sm90.cu: spec-dec draft mask ignored in the SWAP_AB layout, wrong sliding-window masking for full masks, no ragged Q, scale loads before the PDL acquire #4198 — Addresses the PDL ordering hazard by moving scale and speculative-decoding length/mask handling after acqBulk().

Possibly related PRs

• flashinfer-ai/flashinfer#4137 — Both modify SPEC_DEC and ragged-Q handling in csrc/xqa/mha.cu.
• flashinfer-ai/flashinfer#4161 — Both modify XQA multi-block synchronization and kernel ordering.
• flashinfer-ai/flashinfer#2157 — Both modify speculative sequence-length specialization and SM90 MHA selection.

Suggested reviewers: yzh119, jiahanc, aleozlx, dhiraj113, sricketts

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main changes: PDL load ordering and SM90 fp8 draft-mask dispatch.
Description check	✅ Passed	The description matches the template with description, related issues, checklist, tests, and reviewer notes filled out.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/attention/test_xqa_batch_decode.py (1)

1126-1141: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy lift

Cover the PDL ordering regression.

This call uses uniform Q and scalar scales, so it never exercises the moved qCuSeqLens, qScalePtr, or kvScalePtr reads. Add an SM90a-gated ragged-Q test with tensor scales and enable_pdl=True; otherwise the pre-acqBulk() regression can return unnoticed. Confidence: high.

As per coding guidelines, “Skip architecture-specific tests using the appropriate flashinfer.utils checks, such as is_sm90a_supported() or get_compute_capability().”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/attention/test_xqa_batch_decode.py` around lines 1126 - 1141, Add an
SM90a-gated regression test alongside the existing XQA batch decode coverage
that uses ragged per-request Q lengths, tensor-based q/kv scales, and
enable_pdl=True. Ensure the test exercises qCuSeqLens, qScalePtr, and kvScalePtr
reads before the PDL acqBulk path, and skip it using the appropriate
flashinfer.utils architecture check such as is_sm90a_supported().

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/attention/test_xqa_batch_decode.py`:
- Around line 1126-1141: Add an SM90a-gated regression test alongside the
existing XQA batch decode coverage that uses ragged per-request Q lengths,
tensor-based q/kv scales, and enable_pdl=True. Ensure the test exercises
qCuSeqLens, qScalePtr, and kvScalePtr reads before the PDL acqBulk path, and
skip it using the appropriate flashinfer.utils architecture check such as
is_sm90a_supported().

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: c192fce8-0845-4a84-84a2-6e40eee66729

📥 Commits

Reviewing files that changed from the base of the PR and between 417bbd2 and dda315f5446acde347625a919add6cf31550d2b2.

📒 Files selected for processing (6)

• csrc/xqa/mha.cu
• flashinfer/decode.py
• flashinfer/jit/xqa.py
• flashinfer/trace/templates/attention.py
• flashinfer/xqa.py
• tests/attention/test_xqa_batch_decode.py

💤 Files with no reviewable changes (1)

• flashinfer/trace/templates/attention.py

[qsang-nv]

Verified the three items from the #4137 thread on dda315f5:

• PDL ordering — the qCuSeqLens and scale-tensor loads, the zero-length early return, and every derived index (nbValidHeadTokens, mask +=, …) all moved below acqBulk() in mha.cu#L1653-L1691. I checked the window that was vacated: L1582-L1651 now contains only asserts, blockIdx/threadIdx arithmetic and shared-memory barrier init — no global load is left ahead of the acquire — and the first uses of qScaleValue/kvCacheScaleValue are at L1806 and L2572.
• SM90 fp8 dispatch — the swap_ab_eligible(...) term closes the non-causal-mask hole, and the deterministic test's construction checks out (zero Q/K gives a uniform softmax; V values 1/2/4/8/16 are exact in e4m3; causal-vs-full expectations differ by 1.0 and 3.5 against atol=5e-2).
• Module keying — the uncached wrapper plus cached normalized helper is the right shape, and gating the ragged variant on an SM90 target is a real win: SPEC_Q_SEQ_LEN appears only in defines.h's static_assert and mha_sm90.cu, and mha.cu doesn't include specDec.h, so re-emitting it on SM120/121 changes nothing and the ragged and uniform builds correctly collapse to one module there.

Three things below: one small ask for this PR, one for the #4198 follow-up, one doc nit.

---

The deterministic test never reaches mha_sm90.cu on any architecture

The new test is pinned to q_len=4, head_grp_size=4 (L1085), so q_seq_len * head_group_ratio == 16 and the new swap_ab_eligible term routes it to the generic kernel on SM90. SM100 and SM120/121 use mha.cu unconditionally. So across the whole matrix this test only ever validates mha.cu.

After the widened fallback, exactly one speculative-decode configuration still selects the Hopper kernel: fp8 KV, q_seq_len * head_group_ratio > 32, no sliding window, no ragged Q, no sinks — the non-SWAP_AB warpGrpApplyMask that does read the runtime mask. That path currently has only random-data refchecks, which is the exact class of coverage that failed to notice an entirely ignored mask mode last time.

One extra parametrize entry would cover it, e.g. a second shape with num_kv_heads=2, head_grp_size=16 (q_len=4 → 64 > 32). That gives the Hopper masking path a deterministic assertion today and a regression guard for when #4198 restores the fast path.

Follow-up: the Hopper kernel still reads both scale tensors ahead of every acquire

This PR moves qScalePtr / kvScalePtr behind acqBulk() in mha.cu, but the same two loads sit at the very top of the Hopper kernel, mha_sm90.cu#L643-L644 — before even the __CUDA_ARCH__ == 900 guard, and before every acquire site in the file.

Worth being precise about where those sites are, because the acqBulk() at L1241-L1243 is under #if ENABLE_PDL == 1 and is dead here: defines.h selects ENABLE_PDL 2 for __CUDA_ARCH__ == 900, which is the only configuration in which this kernel body exists. The live acquires are the four aggressive-mode sites at L1269, L1290, L1316 and L1375, all in the IO-warp branch and all after L643. The compute warps that consume qScaleValue at L781-L783 never execute an acquire at all.

Scope, so it can be prioritized properly: both ternaries short-circuit on a null pointer, so no load happens unless the caller passes tensor-backed scales. xqa_batch_decode_with_kv_cache passes Python floats, so the pointers are null there; the exposure is direct xqa() callers with device-tensor scales — a dynamic-quant serving stack, say. The route stays reachable after this PR for plain decode as well as for the q_seq_len * head_group_ratio > 32 spec-dec case above.

The cheap fix keeps aggressive PDL intact: sink the two loads to their first use rather than the top of the kernel, placed after the compute warps' first produced.wait() — note that L776-L779 are arrive() calls, not waits, so L781 itself isn't yet ordered. Routing tensor-backed scales to the generic kernel would also work but is a heavy hammer for SM90 fp8 decode. Predates both PRs and lower risk than the qCuSeqLens case, so this reads as #4198 material rather than something to redo here.

Nit: document the fallback under window_left too

The consolidated note in xqa() covers all the conditions, but the high-level window_left docs don't mention that SM90 fp8 speculative decode drops to the generic kernel for any configured window, including one that never truncates. The new note only sits under q_cu_seq_lens, so ragged-Q users see it and sliding-window users don't. While there: the xqa() note lists ragged Q, a positive sliding_win_size and the <= 32 shape condition, but not the pre-existing sinks is not None fallback — worth listing all four in one place.

[yichengj0]

@qsang-nv Thanks for the careful verification. All three items are handled:

• Added a head_grp_size=16 shape (q_len * ratio = 64) to the deterministic mask test, so the non-SWAP_AB Hopper masking path gets a deterministic assertion. It passes on SM120, and H100 CI covers the Hopper leg.
• Agreed the mha_sm90.cu scale loads are [Bug] mha_sm90.cu: spec-dec draft mask ignored in the SWAP_AB layout, wrong sliding-window masking for full masks, no ragged Q, scale loads before the PDL acquire #4198 material. Added them to that issue with your acquire-site analysis linked.
• Doc nits fixed: window_left now mentions the fallback, and the xqa() note lists all four conditions including sinks.

[bkryu]

/bot run tests/attention

[flashinfer-bot]

GitLab MR !1064 has been created, and the CI pipeline #59992940 is currently running. I'll report back once the pipeline job completes.

[saltyminty]

approved pending CI