feat(xqa): ragged Q and per-row sliding-window masking for speculative decode - #4137
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 4afe30f1914e75282cb49f7f7b66becb86d88a89 and a51867d. 📒 Files selected for processing (9)
🚧 Files skipped from review as they are similar to previous changes (8)
📝 WalkthroughWalkthroughThe XQA speculative-decode path now supports causal or full draft masks and packed ragged query lengths. Python APIs, tracing schemas, FFI bindings, JIT specialization, CUDA masking, benchmarks, and attention tests are updated. ChangesSpeculative Decode and Ragged Q
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant xqa_batch_decode_with_kv_cache
participant xqa
participant xqa_wrapper
participant MHA_launcher
xqa_batch_decode_with_kv_cache->>xqa: pass q_cu_seq_lens
xqa->>xqa_wrapper: invoke custom op with q_cu_seq_lens
xqa_wrapper->>MHA_launcher: forward cumulative draft lengths
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
2421726 to
c699f88
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@flashinfer/xqa.py`:
- Around line 261-268: Update the q_cu_seq_lens documentation in the public XQA
API to state that ragged-Q requests are unsupported on the SM90 FP8 execution
path, while preserving the existing description for supported paths.
- Around line 287-308: Strengthen validation in the use_ragged_q branch of the
XQA entry point before the CUDA launch: require q_cu_seq_lens to be contiguous
and on the same CUDA device as q, then validate its cumulative offsets are
non-negative, non-decreasing, each request length is at most q_seq_len, and the
final offset equals q.shape[0]. Preserve the existing dtype, size, shape, and
q_seq_len checks, and reject invalid metadata before passing
q_cu_seq_lens.data_ptr() to native code.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: b2c66a2c-ab03-4139-97ec-52a683138ad0
📥 Commits
Reviewing files that changed from the base of the PR and between cea7f46 and c699f88ed2c5913be2391a170b8abe4b8e12121d.
📒 Files selected for processing (9)
benchmarks/routines/attention.pycsrc/flashinfer_xqa_binding.cucsrc/xqa/mha.cucsrc/xqa/xqa_wrapper.cuflashinfer/decode.pyflashinfer/jit/xqa.pyflashinfer/trace/templates/attention.pyflashinfer/xqa.pytests/attention/test_xqa_batch_decode.py
| use_ragged_q = q_cu_seq_lens is not None | ||
| if use_ragged_q: | ||
| assert q_seq_len > 1, "q_cu_seq_lens requires q_seq_len > 1 (the max draft len)" | ||
| assert q.dim() == 3, ( | ||
| "With q_cu_seq_lens, q must be packed as " | ||
| f"[total_q_tokens, num_q_heads, head_dim], got {q.dim()}D" | ||
| ) | ||
| assert q_cu_seq_lens.dtype in (torch.int32, torch.uint32), ( | ||
| "q_cu_seq_lens must be int32 or uint32" | ||
| ) | ||
| assert q_cu_seq_lens.is_cuda, ( | ||
| "q_cu_seq_lens must be a device tensor; the kernel dereferences its " | ||
| "pointer on the GPU" | ||
| ) | ||
| assert q_cu_seq_lens.numel() == seq_lens.shape[0] + 1, ( | ||
| f"q_cu_seq_lens must have batch_size + 1 = {seq_lens.shape[0] + 1} " | ||
| f"entries, got {q_cu_seq_lens.numel()}" | ||
| ) | ||
| assert output.shape == q.shape, "Output must match packed ragged q shape" | ||
| # Not checked (needs a device sync): entries non-decreasing, each | ||
| # length <= q_seq_len, last entry == q.shape[0]. Violations return | ||
| # garbage rows rather than raising. |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Validate ragged offsets before passing their raw pointer to CUDA.
This accepts noncontiguous offsets, offsets on a different CUDA device, negative/non-monotonic values, oversized per-request lengths, and a final offset that disagrees with q.shape[0]. The native code treats data_ptr() as a contiguous unsigned array; invalid metadata can address incorrect rows or memory. Validate device/layout and cumulative-offset invariants before launch.
Proposed validation
assert q_cu_seq_lens.is_cuda, (
"q_cu_seq_lens must be a device tensor; the kernel dereferences its "
"pointer on the GPU"
)
+ assert q_cu_seq_lens.device == q.device
+ assert q_cu_seq_lens.dim() == 1 and q_cu_seq_lens.is_contiguous()
assert q_cu_seq_lens.numel() == seq_lens.shape[0] + 1, (
f"q_cu_seq_lens must have batch_size + 1 = {seq_lens.shape[0] + 1} "
f"entries, got {q_cu_seq_lens.numel()}"
)
+ offsets = q_cu_seq_lens.cpu()
+ lengths = offsets[1:] - offsets[:-1]
+ assert offsets[0].item() == 0
+ assert offsets[-1].item() == q.shape[0]
+ assert torch.all(lengths >= 0) and torch.all(lengths <= q_seq_len)
assert output.shape == q.shape, "Output must match packed ragged q shape"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| use_ragged_q = q_cu_seq_lens is not None | |
| if use_ragged_q: | |
| assert q_seq_len > 1, "q_cu_seq_lens requires q_seq_len > 1 (the max draft len)" | |
| assert q.dim() == 3, ( | |
| "With q_cu_seq_lens, q must be packed as " | |
| f"[total_q_tokens, num_q_heads, head_dim], got {q.dim()}D" | |
| ) | |
| assert q_cu_seq_lens.dtype in (torch.int32, torch.uint32), ( | |
| "q_cu_seq_lens must be int32 or uint32" | |
| ) | |
| assert q_cu_seq_lens.is_cuda, ( | |
| "q_cu_seq_lens must be a device tensor; the kernel dereferences its " | |
| "pointer on the GPU" | |
| ) | |
| assert q_cu_seq_lens.numel() == seq_lens.shape[0] + 1, ( | |
| f"q_cu_seq_lens must have batch_size + 1 = {seq_lens.shape[0] + 1} " | |
| f"entries, got {q_cu_seq_lens.numel()}" | |
| ) | |
| assert output.shape == q.shape, "Output must match packed ragged q shape" | |
| # Not checked (needs a device sync): entries non-decreasing, each | |
| # length <= q_seq_len, last entry == q.shape[0]. Violations return | |
| # garbage rows rather than raising. | |
| use_ragged_q = q_cu_seq_lens is not None | |
| if use_ragged_q: | |
| assert q_seq_len > 1, "q_cu_seq_lens requires q_seq_len > 1 (the max draft len)" | |
| assert q.dim() == 3, ( | |
| "With q_cu_seq_lens, q must be packed as " | |
| f"[total_q_tokens, num_q_heads, head_dim], got {q.dim()}D" | |
| ) | |
| assert q_cu_seq_lens.dtype in (torch.int32, torch.uint32), ( | |
| "q_cu_seq_lens must be int32 or uint32" | |
| ) | |
| assert q_cu_seq_lens.is_cuda, ( | |
| "q_cu_seq_lens must be a device tensor; the kernel dereferences its " | |
| "pointer on the GPU" | |
| ) | |
| assert q_cu_seq_lens.device == q.device | |
| assert q_cu_seq_lens.dim() == 1 and q_cu_seq_lens.is_contiguous() | |
| assert q_cu_seq_lens.numel() == seq_lens.shape[0] + 1, ( | |
| f"q_cu_seq_lens must have batch_size + 1 = {seq_lens.shape[0] + 1} " | |
| f"entries, got {q_cu_seq_lens.numel()}" | |
| ) | |
| offsets = q_cu_seq_lens.cpu() | |
| lengths = offsets[1:] - offsets[:-1] | |
| assert offsets[0].item() == 0 | |
| assert offsets[-1].item() == q.shape[0] | |
| assert torch.all(lengths >= 0) and torch.all(lengths <= q_seq_len) | |
| assert output.shape == q.shape, "Output must match packed ragged q shape" | |
| # Not checked (needs a device sync): entries non-decreasing, each | |
| # length <= q_seq_len, last entry == q.shape[0]. Violations return | |
| # garbage rows rather than raising. |
🤖 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 `@flashinfer/xqa.py` around lines 287 - 308, Strengthen validation in the
use_ragged_q branch of the XQA entry point before the CUDA launch: require
q_cu_seq_lens to be contiguous and on the same CUDA device as q, then validate
its cumulative offsets are non-negative, non-decreasing, each request length is
at most q_seq_len, and the final offset equals q.shape[0]. Preserve the existing
dtype, size, shape, and q_seq_len checks, and reject invalid metadata before
passing q_cu_seq_lens.data_ptr() to native code.
There was a problem hiding this comment.
Partially agree. The cheap metadata checks (same device as q, 1-D, contiguous) are worth adding and we will. The value checks need a .cpu() copy, which forces the device sync this path exists to avoid and breaks CUDA graph capture, so we're leaving those out. The docstring and an inline comment already state that violated offset invariants produce garbage output rather than an error.
There was a problem hiding this comment.
@yichengj0, agreed. Requiring a CPU copy here would undermine the asynchronous, CUDA-graph-compatible path, so the offset-value invariants should remain a documented caller contract rather than host-side validation.
The same-device, 1-D, and contiguous checks cover the pointer/layout hazards in the original finding. I’ll treat the remaining value validation as intentionally omitted for graph-capture compatibility.
✏️ Learnings added
Learnt from: yichengj0
Repo: flashinfer-ai/flashinfer PR: 4137
File: flashinfer/xqa.py:287-308
Timestamp: 2026-07-24T20:10:21.224Z
Learning: In `flashinfer/xqa.py`, ragged speculative-decoding `q_cu_seq_lens` value invariants (zero initial offset, monotonic offsets, per-request length bounded by `q_seq_len`, and final offset matching packed `q.shape[0]`) are intentionally not host-validated. Copying offsets to CPU would synchronize CUDA work and break CUDA graph capture; invalid values are a documented caller contract and may produce garbage output. Cheap metadata validation should still require `q_cu_seq_lens` to be on the same device as `q`, one-dimensional, and contiguous before passing its raw pointer to the CUDA kernel.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
bkryu
left a comment
There was a problem hiding this comment.
Thanks @yichengj0, left one comment about benchmark flags
|
https://github.com/flashinfer-ai/flashinfer/actions/runs/30125334423/job/89587528368?pr=4137 |
|
/bot run tests/attention |
|
+1 to @qsang-nv's comment on the failure above – I believe it's just a unsupported case that needs to be skipped (fp8 kv + enable_sink=false) |
saltyminty
left a comment
There was a problem hiding this comment.
Approved pending CI and fixing of the previously mentioned test failure.
|
@qsang-nv @saltyminty Thanks for flagging. We reproduced this on an H100 and root-caused it. All 36 failures are on the SM90 fp8 path, which dispatches to the Hopper-specific
Instead of skipping, the pushed fix extends that fallback: SM90 fp8 spec-dec with a sliding window or ragged Q also runs on the generic kernel, verified on an H100. The |
…e decode Plumb q_cu_seq_lens (cumulative per-request draft lengths) through the XQA wrapper and Python API so batched speculative decode can verify a different number of draft tokens per request, and make sliding-window masking exact per draft row: whole leading KV tiles are skipped only up to the earliest row's window begin, and the per-row window edge is masked in-kernel. Spec-dec modules compile with IS_SPEC_DEC_TREE=0; the packed mask API indexes linear draft chains (causal or full), and both mask modes are covered by tests and the benchmark.
Existing tests capped head_grp_size at 8. Add ratio-16 rows to the main decode/spec-dec matrix, the truncating sliding-window test, and the ragged-Q test. AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Prw1D3DGGRFKjju1GfFckK
…/test coverage Address code-review findings on the ragged-Q + per-row sliding-window work: - skip ragged requests with 0 draft tokens in the kernel (row clamps would underflow into an out-of-bounds mask read) - validate q_cu_seq_lens device residency and length on the host - reject ragged Q on the SM90 fp8 MHA path instead of routing into its unvalidated qCuSeqLens handling - restore the spec-dec benchmark's causal default; full (non-causal) draft masks are now opt-in via --spec_dec_mask - share the JIT module and torch op between ragged and uniform builds when ragged changes no compile flags; derive the op name from the JitSpec URI instead of duplicating the format string - add q_cu_seq_lens to the xqa trace template - extend tests to long contexts (multi-block split-KV path) and zero-length drafts AI-assisted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Prw1D3DGGRFKjju1GfFckK
…p8 limit Address review feedback: validate device, rank, and contiguity of q_cu_seq_lens before handing its pointer to the kernel, and document that ragged Q is rejected on the SM90 fp8 MHA path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…sliding window or ragged Q mha_sm90.cu's spec-dec path predates both features: its per-row sliding-window masking (activated by the IS_SPEC_DEC_TREE=0 build) is wrong for full draft masks, and it has no ragged-Q support. Route those cases to the generic kernel, matching the existing sinks fallback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
4afe30f to
a51867d
Compare
|
/bot run tests/attention |
|
/bot run tests/gemm |
|
/bot run tests/attention |
|
Thanks for the quick fix on the H100 failures — routing SM90 fp8 back to the generic kernel rather than skipping the cases is the right call. Three things I'd like resolved before this lands, plus two non-blocking notes. Blocking1.
|
….6.16rc2 (#4197) Cherry-picks for the `release-v0.6.16` branch, plus the version bump to `0.6.16rc2`. ## Commits | Commit | Source | Status | | --- | --- | --- | | `fix(norm): convert float2 to e4m3 directly in packed cast` | #4167 (`e683e307` on `main`) | merged upstream | | `feat(xqa): ragged Q and per-row sliding-window masking for speculative decode` | #4137 (`60783fb9` on `main`) | merged upstream | | `test(jit): assert BMM export symlink under GEN_SRC_DIR` | #4187 (`417bbd29` on `main`) | merged upstream | | `feat(mla): support packed low-head and variable-Q decode` | #4178 (PR head `0774943c`) | **not yet merged** to `main` | | `bump version to 0.6.16rc2` | — | — | Applied in `main` merge order; #4187 sits directly on top of #4137 upstream, so that ordering is preserved. ## Notes - **#4189 was not cherry-picked** — it was merged directly into `release-v0.6.16` (`34368112`) and is already the base of this branch. - **#4178 is still open against `main`.** It was cherry-picked at the request of the release owner. Its two head commits (`b6cc59491`, `0774943c2`) are squashed into one commit here, mirroring how it will land upstream. If the PR changes before merge, this commit should be dropped and re-picked. - The three merged cherry-picks apply cleanly with patch-ids identical to their sources. The #4178 squash is byte-identical to the PR's combined diff. - The #4178 commit was committed with `--no-verify`: the `mypy` pre-commit hook reports two pre-existing `[no-redef]` errors in `flashinfer/mla/_core.py` (conditional `if/else` imports of `_check_can_implement` / `_get_split_kv_and_workspace_size`). These reproduce identically on PR #4178's own branch and on the `release-v0.6.16` base — they are not introduced by this cherry-pick, and surface only when those modules are checked together. Left unfixed here to keep the cherry-pick faithful. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Lee Yongjun <jqueen.astro@gmail.com> Co-authored-by: yichengj <yichengj@nvidia.com> Co-authored-by: bryu <bryu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Ka-Hyun Nam <knam@nvidia.com> Co-authored-by: mingyangw <mingyangw@nvidia.com>
|
@qsang-nv Thanks again for the thorough review. All five items are addressed in #4199, and the SM90 kernel work (SWAP_AB mask support, For the SWAP_AB gap we took the fallback option. The generic path is already validated on H100. Suppressing the specialization instead would move the Hopper kernel to its unswapped layout, which would need its own correctness and perf validation at these small shapes, and that effort is better spent restoring SWAP_AB properly in #4198. Your scale-tensor side note turned out to be the same one-line move as the |
## 📌 Description Follow-up to #4137, addressing [the review feedback](#4137 (comment)) 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 #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 - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] 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; #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 #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](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
📌 Description
XQA is the FlashInfer decode kernel used on SM120/121 for models with attention sinks (#4070). Serving those models with speculative decoding runs the draft-verification step through XQA as well, and two gaps blocked that:
Fixes:
q_cu_seq_lensthrough the wrapper and Python API so each request can verify a different number of draft tokens, with host-side input validation. The SM90 fp8 path rejects ragged Q rather than run an unvalidated path.--spec_dec_mask), long contexts (split-KV path), zero-length drafts, and GQA group ratio 16.Kernel changes and investigation by @bkryu.
📈 Performance
The table shows XQA's speedup over each baseline, measured as kernel time with CUPTI at batch size 1 with head_dim 128, 32 query heads over 2 KV heads, and a bf16 KV cache. Each cell gives the speedup at context lengths 1k, 4k, and 32k. The column m is the number of query tokens per request: m=1 is plain decode, and m=4 or 8 is draft-block verification in speculative decode. Each baseline is measured under full attention and under sliding-window attention with a 1024-token window (SWA 1024).
Baselines:
🔍 Related Issues
#4070 (attention sinks on SM120/121).
🚀 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).tests/attention/test_xqa_batch_decode.pypasses on SM120 (RTX 5080, RTX PRO 6000) and SM121 (GB10). Outputs cross-checked against an independent reference kernel up to 64k context.Reviewer Notes
IS_SPEC_DEC_TREEcompile flag flips from 1 to 0). This is deliberate: the per-row window needs each draft token's sequence position, and only a linear chain defines one. The assumption only matters in sliding-window builds, and tree-shaped drafts never worked correctly with sliding windows, so no working caller changes behavior.Summary by CodeRabbit
New Features
q_cu_seq_lens.--spec_dec_mask(causal/full), including trace support.Bug Fixes
Tests
causal/fullmask modes, sliding-window behavior, and KV cache variants.