Skip to content

feat(xqa): ragged Q and per-row sliding-window masking for speculative decode - #4137

Merged
bkryu merged 5 commits into
flashinfer-ai:mainfrom
yichengj0:xqa-specdec-ragged-swa
Jul 28, 2026
Merged

bkryu merged 5 commits into
flashinfer-ai:mainfrom
yichengj0:xqa-specdec-ragged-swa

Conversation

@yichengj0

@yichengj0 yichengj0 commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

📌 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:

  • Every request in a batch had to verify the same number of draft tokens. The kernel indexes queries and masks through cumulative lengths internally, but the argument was never exposed.
  • Sliding-window masking computed one window start from the last draft token's position and applied it to the whole draft block. Each draft token sits at its own position, so this masked out KV that the earlier draft tokens should still attend to.

Fixes:

  • Plumb q_cu_seq_lens through 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.
  • Compute the window per draft row: whole KV tiles are skipped conservatively, and the exact per-row edge is masked in the kernel.
  • Extend tests and the benchmark to cover both draft-block mask modes (causal and full, selected with --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:

  • vLLM's Triton unified-attention kernel, sinks enabled on both sides.
  • FlashInfer fa2 wrappers, sinks disabled on both sides (fa2 has no sink support).
GPU m vs Triton, full attn vs Triton, SWA 1024 vs fa2, full attn vs fa2, SWA 1024
GB10 1 1.7 / 0.9 / 1.1 1.5 / 1.2 / 2.3 0.8 / 0.9 / 1.0 0.9 / 0.7 / 0.8
GB10 4 1.2 / 2.2 / 5.0 1.2 / 1.2 / 1.2 2.7 / 1.4 / 1.2 2.3 / 1.9 / 1.9
GB10 8 1.1 / 1.8 / 4.0 1.2 / 1.1 / 0.9 1.8 / 1.4 / 1.0 2.5 / 2.1 / 1.5
RTX PRO 6000 1 3.0 / 2.4 / 3.2 3.3 / 3.0 / 3.6 1.0 / 0.8 / 0.8 1.0 / 1.0 / 0.6
RTX PRO 6000 4 2.1 / 6.0 / 18.0 1.9 / 1.9 / 1.2 3.7 / 2.3 / 1.5 2.7 / 3.0 / 2.1
RTX PRO 6000 8 2.1 / 5.5 / 14.2 1.8 / 2.0 / 1.0 3.3 / 2.4 / 1.1 2.8 / 2.8 / 1.6

🔍 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

  • 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

  • Tests have been added or updated as needed.
  • All tests are passing (unittest, etc.).

tests/attention/test_xqa_batch_decode.py passes on SM120 (RTX 5080, RTX PRO 6000) and SM121 (GB10). Outputs cross-checked against an independent reference kernel up to 64k context.

Reviewer Notes

  • Speculative-decode builds now assume the draft tokens form a linear chain rather than a tree (the IS_SPEC_DEC_TREE compile 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.
  • Not validated on SM90/SM100 hardware. The changed paths are spec-dec only, and the SM90 fp8 wrapper rejects ragged explicitly; CI runs the tests on those arches.

Summary by CodeRabbit

  • New Features

    • Added ragged-Q speculative decoding support via optional q_cu_seq_lens.
    • Introduced configurable speculative draft attention masking with --spec_dec_mask (causal/full), including trace support.
    • Extended JIT/custom-op pathways to support ragged-Q specialization.
  • Bug Fixes

    • Improved speculative decoding masking for sliding-window cases and fixed an empty-query edge case for ragged requests.
  • Tests

    • Expanded coverage for ragged-Q, causal/full mask modes, sliding-window behavior, and KV cache variants.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

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

@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d360cc65-f0c7-45fc-a885-a2a78fdb00f3

📥 Commits

Reviewing files that changed from the base of the PR and between 4afe30f1914e75282cb49f7f7b66becb86d88a89 and a51867d.

📒 Files selected for processing (9)
  • benchmarks/routines/attention.py
  • csrc/flashinfer_xqa_binding.cu
  • csrc/xqa/mha.cu
  • csrc/xqa/xqa_wrapper.cu
  • flashinfer/decode.py
  • flashinfer/jit/xqa.py
  • flashinfer/trace/templates/attention.py
  • flashinfer/xqa.py
  • tests/attention/test_xqa_batch_decode.py
🚧 Files skipped from review as they are similar to previous changes (8)
  • csrc/flashinfer_xqa_binding.cu
  • flashinfer/trace/templates/attention.py
  • csrc/xqa/mha.cu
  • tests/attention/test_xqa_batch_decode.py
  • csrc/xqa/xqa_wrapper.cu
  • flashinfer/xqa.py
  • benchmarks/routines/attention.py
  • flashinfer/decode.py

📝 Walkthrough

Walkthrough

The 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.

Changes

Speculative Decode and Ragged Q

Layer / File(s) Summary
Configurable speculative masks
benchmarks/routines/attention.py, tests/attention/test_xqa_batch_decode.py
Benchmarks and tests support causal and full speculative masks, propagate the selected causal setting, and add sliding-window, ragged-Q, and NVFP4 coverage.
Ragged-Q API plumbing
flashinfer/decode.py, flashinfer/xqa.py, flashinfer/trace/templates/attention.py, csrc/flashinfer_xqa_binding.cu, csrc/xqa/xqa_wrapper.cu
Cumulative draft-length offsets are accepted, validated, described in trace schemas, and forwarded through Python, FFI, and launcher layers.
Ragged-Q JIT specialization
flashinfer/jit/xqa.py, flashinfer/xqa.py
Ragged-Q builds suppress incompatible speculative query-length specialization and use distinct affected JIT cache keys.
Kernel masking and validation
csrc/xqa/mha.cu, tests/attention/test_xqa_batch_decode.py
CUDA masking handles per-row sliding windows and zero-length ragged requests; corresponding speculative-decoding tests validate the behavior.

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
Loading

Possibly related PRs

Suggested reviewers: yzh119, jimmyzho, vinnie6167, jiahanc, yongwww

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the main change: ragged Q support and per-row sliding-window masking for speculative decode.
Description check ✅ Passed The description matches the template with description, related issues, checklist, tests, and reviewer notes filled in.
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.
✨ 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

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

@yichengj0
yichengj0 force-pushed the xqa-specdec-ragged-swa branch from 2421726 to c699f88 Compare July 24, 2026 18:35
@yichengj0
yichengj0 marked this pull request as ready for review July 24, 2026 18:35
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

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

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.py
  • csrc/flashinfer_xqa_binding.cu
  • csrc/xqa/mha.cu
  • csrc/xqa/xqa_wrapper.cu
  • flashinfer/decode.py
  • flashinfer/jit/xqa.py
  • flashinfer/trace/templates/attention.py
  • flashinfer/xqa.py
  • tests/attention/test_xqa_batch_decode.py

Comment thread flashinfer/xqa.py Outdated
Comment thread flashinfer/xqa.py
Comment on lines +287 to +308
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.

@coderabbitai coderabbitai Bot Jul 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.

Suggested change
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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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 bkryu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @yichengj0, left one comment about benchmark flags

Comment thread benchmarks/routines/attention.py
@bkryu bkryu added the run-ci label Jul 24, 2026
@qsang-nv

Copy link
Copy Markdown
Collaborator

https://github.com/flashinfer-ai/flashinfer/actions/runs/30125334423/job/89587528368?pr=4137
tests/attention/test_xqa_batch_decode.py failed on H100, please take a look.

@bkryu

bkryu commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

/bot run tests/attention

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

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

@saltyminty

Copy link
Copy Markdown
Collaborator

+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 saltyminty left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approved pending CI and fixing of the previously mentioned test failure.

@yichengj0

Copy link
Copy Markdown
Contributor Author

@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 mha_sm90.cu:

  • The 32 ragged_q failures are this PR's own guard raising: mha_sm90.cu has no ragged-Q support.
  • The 4 spec_dec_sliding_window failures are real. The IS_SPEC_DEC_TREE=0 build enabled a dormant per-row sliding-window path in mha_sm90.cu, and it mis-masks full (non-causal) draft masks. The enable_sink=True variants pass because the existing sinks workaround already routes them to the generic mha.cu.

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 mha_sm90.cu gaps could be tracked in a follow-up issue so the Hopper fast path can be restored later.

bkryu and others added 5 commits July 27, 2026 11:19
…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>
@yichengj0
yichengj0 force-pushed the xqa-specdec-ragged-swa branch from 4afe30f to a51867d Compare July 27, 2026 18:20
@bkryu

bkryu commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

/bot run tests/attention

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1050 has been updated with latest changes, and the CI pipeline #59824049 is currently running. I'll report back once the pipeline job completes.

@bkryu

bkryu commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

/bot run tests/gemm

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

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

@bkryu

bkryu commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

/bot run tests/attention

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

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

@bkryu
bkryu enabled auto-merge (squash) July 28, 2026 02:03
@qsang-nv

Copy link
Copy Markdown
Collaborator

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.


Blocking

1. qCuSeqLens is dereferenced before griddepcontrol.wait (PDL race)

Worth flagging up front that this one is on the SM12x path, not SM90: mha.cu is the only kernel compiled for SM120/121 (mha_sm90.cu is added to the sources and USE_SM90_MHA=1 only when has_sm90), ENABLE_PDL is 1 there (__CUDA_ARCH__ >= 900, and SM120 is 1200), and device_support_pdl() returns True for major >= 9. So this affects the PR's primary target architecture unconditionally.

mha.cu#L1599-L1610 reads qCuSeqLens[idxReq] and qCuSeqLens[idxReq + 1] before the acquire; acqBulk() (griddepcontrol.wait) is at L1685-L1688. PDL is on by default — makeLaunchConfig sets cudaLaunchAttributeProgrammaticStreamSerialization and the Python enable_pdl defaults to device_support_pdl(). Per the PTX contract, a dependent grid may begin executing before the prerequisite grid's writes are visible; griddepcontrol.wait is what establishes that visibility. Reading producer-written memory before it is a contract violation.

These are the first request-indexing loads placed before the acquire, and they are the dangerous ones: they gate an early return and they produce reqSeqOffset, which drives the mask row offset (L1626), the Q row offset (L1708) and the output row offset (L2780, &output[reqSeqOffset * nbQHeads]). A stale or garbage value doesn't merely perturb numerics — it can write into another request's output rows, or out of bounds if the buffer still holds allocator garbage, non-deterministically. Because the loaded value feeds a control-flow decision (the zero-length early return), this also isn't something the compiler could be sinking below the acquire.

q_cu_seq_lens is exactly the kind of buffer a preceding kernel produces on device — it's the prefix sum of accepted draft lengths. The API accepts it as a device tensor and performs no host-side synchronization, so it can naturally be produced by a preceding device kernel. If that producer is PDL-aware (issues launch_dependents before its writes retire), XQA can observe the pre-completion state.

No existing test intentionally supplies q_cu_seq_lens from a prerequisite kernel that triggers launch_dependents before completion. test_xqa_batch_decode_ragged_q constructs it through ordinary torch.cumsum/torch.cat, so the test does not establish the early-launch overlap needed to expose this race. A serving integration with a PDL-aware producer can.

Suggested fix — move the loads and everything derived from them below the acquire; the CTA-local shared-memory/barrier initialization can stay where it is:

  __syncthreads();

#if ENABLE_PDL
  preExit();
  acqBulk();
#endif

#if SPEC_DEC
  bool const variableQSeqLen = qCuSeqLens != nullptr;
  uint32_t const actualQSeqLen =
      variableQSeqLen ? uint32_t(qCuSeqLens[idxReq + 1] - qCuSeqLens[idxReq]) : qSeqLen;
  uint32_t const reqSeqOffset = variableQSeqLen ? uint32_t(qCuSeqLens[idxReq]) : (qSeqLen * idxReq);
  if (variableQSeqLen && actualQSeqLen == 0) {
    return;
  }
  // ... nbQHeadTokens / totalNbHeadTokensInGrp / nbValidHeadTokens / `mask +=` follow here
#endif

Returning after __syncthreads() is safe: the condition is CTA-uniform because all threads in a CTA use the same blockIdx.z and indptr entries, the barriers are CTA-local, and the request's CTAs have not entered any cross-CTA semaphore protocol yet (that only happens in the multi-block reduction at the end).

Side note, not a request for this PR: qScalePtr[0] / kvScalePtr[0] at L1582-L1583 are also global loads ahead of the acquire. They predate this PR and are far lower risk (a stale scalar scale gives wrong numbers, not misaddressed writes), but they're worth a look when someone next audits this kernel for PDL.

2. The fallback is still under-scoped: non-causal masks on SM90 fp8 without a sliding window

The new guard at xqa.py#L386-L393 keys on use_ragged_q or use_sliding_window. But the breakage is tied to SWAP_AB, not to the window.

The SWAP_AB mask implementation at mha_sm90.cu#L1755-L1801 does not use specDec at all. It synthesizes the draft mask from the column index:

uint32_t const maskCol = col / headGrpSize;
MaskType const bit_mask = (1ULL << (maskCol + 1)) - 1;   // hard-coded causal

The only consumers of the runtime mask (specDec.needMask / specDec.loadTileMaskRow) are at L2122/L2140-L2141, i.e. exclusively in the non-SWAP_AB variant. tok0WinBeg is likewise accepted but never used in the SWAP_AB version. And SWAP_AB is selected purely by SPEC_Q_SEQ_LEN, which jit/xqa.py#L91 emits whenever q_seq_len * head_group_ratio <= 32. That matches the failure pattern exactly: the 4 failing cases were (4,2,32,2,4) and (4,4,64,4,2) (product 8), while (4,5,16,2,8) (40) and (4,4,32,2,16) (64) passed.

So a full mask with sliding_win_size == 0 and q_seq_len * head_group_ratio <= 32 still reaches this path and returns causal results.

The main test already covers that combination and passes, and that is itself the problem. Every link is checkable: cc == 9 + fp8 KV sets run_sm90_fp8_mha = True; sinks is None; not ragged; window_left = -1 means no fallback (and in the previous revision there was no fallback at all); USE_SM90_MHA == 1 on that runner is proven by the ragged-Q ICHECK firing there in the previous revision's run; q_seq_len * head_group_ratio <= 32 gives SWAP_AB. So test_xqa_batch_decode[full][bf16-fp8-bf16][window_left=-1][enable_sink=False] ran a causal-masked kernel against a non-causal reference and still passed. That is a proof that the currently exercised inputs and fp8 tolerance did not distinguish causal from full, not evidence that the path is correct.

I want to be clear about what I think belongs here versus in a follow-up, because the SWAP_AB behavior itself predates this PR. What changed is the contract: before this PR the mask was documented as causal-only, so hard-coding causal was consistent with it. This PR promotes non-causal draft masks to an advertised feature —

-        Causal attention mask for speculative decoding mode (when ``q_seq_len > 1``).
+        Draft-block attention mask for speculative decoding mode (when
-        causal attention mask for xqa speculative decoding.
+        draft-block attention mask for xqa speculative decoding.

plus the new --spec_dec_mask {causal,full} benchmark flag and the spec_dec_mask_mode test parametrization. So the gap isn't inherited — the PR ships a documented mode that silently returns something else on one architecture.

In this PR:

  1. Extend the dispatch guard so the advertised behavior is correct everywhere. "Fall back only when the mask isn't causal" isn't implementable — xqa() can't inspect a device-side mask without a sync. Either fall back whenever SM90 + fp8 + spec-dec + q_seq_len * head_group_ratio <= 32, or keep the Hopper kernel and suppress -DSPEC_Q_SEQ_LEN for masked spec-dec so mha_sm90.cu takes the mask-honoring non-SWAP_AB path — the same suppression mechanism this PR already added for ragged Q. Either way it's a couple of lines in the guard you just added, on the fallback path you already validated on H100, so it needs no SM90 hardware to land.
  2. Add a deterministic mask-mode case. This is not an SM90 concern: the current parametrization ran a causal-masked kernel against a non-causal reference and passed, so it has no power to catch a mask-ignoring regression on any architecture — including SM120/121, where full is therefore exercised but not validated. E.g. use zero/equal Q and K so the logits are uniform, then give the future draft tokens clearly distinct V values.

Follow-up: making the SWAP_AB warpGrpApplyMask read the runtime mask (and actually use tok0WinBeg) is real kernel work that needs SM90 hardware to validate. Happy to see that tracked in the mha_sm90.cu issue you mentioned rather than done here.

3. The trace template's q_cu_seq_lens input is inert

attention.py#L2768 declares q_cu_seq_lens, but _xqa_batch_decode_reference forwards **kwargs into _trtllm_paged_attention_reference, which only reads cum_seq_lens_q (L1565). The offsets are silently dropped and the reference falls back to q_start = b * (num_tokens // batch_size).

Renaming the kwarg wouldn't be enough. That shared reference also hardcodes causal=False and has no notion of the packed bit-mask, and the template declares neither mask nor q_len_per_req — so a trace that sets q_cu_seq_lens trips assert q_len_per_req > 1 before it gets anywhere. Either implement ragged query partitioning and packed-mask semantics in the trace reference (including q_len_per_req and mask), or drop q_cu_seq_lens from the trace schema in this PR and add it with the rest.


Non-blocking

4. JIT module keying for ragged Q

Two small things in the same area.

Normalization lives in the caller. xqa.py#L373 computes ragged_build = use_ragged_q and q_seq_len * head_group_ratio <= 32 at the call site, while get_xqa_module is @functools.cached on the raw argument. If a process ever requests both raw False and raw True for a config with q_seq_len * head_group_ratio > 32, the two cache entries produce the same op_name and the second registration fails. Today's only in-tree caller normalizes, so this is latent. Note that normalizing inside the cached function body doesn't help — the cache key is computed from the raw arguments before the body runs. It needs an uncached wrapper:

def get_xqa_module(..., q_seq_len, use_ragged_q=False):
    normalized = use_ragged_q and q_seq_len * head_group_ratio <= 32
    return _get_xqa_module_cached(..., q_seq_len, normalized)

@functools.cache
def _get_xqa_module_cached(..., q_seq_len, use_ragged_q):
    ...

On SM12x the ragged variant is a redundant build. SPEC_Q_SEQ_LEN appears zero times in mha.cu; outside the static_assert in defines.h it is used only by mha_sm90.cu to select SWAP_AB. So on an SM120/121-only build, suppressing it for ragged Q changes no generated code — it just adds the _ragged_q URI suffix and produces a second, binary-equivalent module: another full XQA JIT compile, another registered torch op, another resident module for the same shape. gen_xqa_module already computes has_sm90 for the sources list; gating both the flag suppression and the URI suffix on it would keep non-SM90 builds down to one module.

5. decode.py docs don't reflect the SM90 fp8 fallback

The q_cu_seq_lens docstring at decode.py#L3606 does not disclose that ragged Q uses the generic kernel on SM90 fp8. This matters more than in the previous revision: it used to raise there, and now it silently drops to the generic kernel, so users lose the Hopper fp8 path with no signal. The same applies to sliding-window dispatch: the fallback keys on sliding_win_size > 0 rather than on whether the window actually truncates, so a windowed model on SM90 fp8 gives up the fast path even when kv_len <= window. Worth stating in both places.

Also +1 on tracking the mha_sm90.cu gaps in a follow-up issue so the Hopper path can be restored later.

@bkryu
bkryu merged commit 60783fb into flashinfer-ai:main Jul 28, 2026
48 of 65 checks passed
aleozlx added a commit that referenced this pull request Jul 28, 2026
….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>
@yichengj0

Copy link
Copy Markdown
Contributor Author

@qsang-nv Thanks again for the thorough review. All five items are addressed in #4199, and the SM90 kernel work (SWAP_AB mask support, tok0WinBeg, ragged Q) is tracked in #4198.

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 qCuSeqLens fix, so #4199 fixes it too.

bkryu pushed a commit that referenced this pull request Jul 28, 2026
## 📌 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants