fix(xqa): read the runtime draft mask on the SWAP_AB spec-dec path - #4229
Smallfu666 wants to merge 1 commit into
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 13a161e48890634d74142b94da63d0f7aed8e05b and 66d0f09. 📒 Files selected for processing (4)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughSM90 speculative-decoding masking now supports synthesized causal masks when no device mask is provided and combines runtime mask rows with causal or sliding-window boundaries. FP8 dispatch fallback now targets ragged-Q inputs and provided sinks instead of ChangesSM90 speculative decoding
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant xqa
participant warpGrpApplyMask
participant SpecDec
participant AccumulatorTiles
xqa->>warpGrpApplyMask: select SM90 speculative-decoding path
warpGrpApplyMask->>SpecDec: load per-tile mask rows
warpGrpApplyMask->>warpGrpApplyMask: compute causal and sliding-window masks
warpGrpApplyMask->>AccumulatorTiles: initialize masked elements to safeInitRowMax
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@csrc/xqa/mha_sm90.cu`:
- Around line 1768-1806: The null-mask fallback in the shown masking path must
reject or safely handle cacheSeqLen values below SPEC_Q_SEQ_LEN at runtime, not
only via assert. Add an appropriate runtime guard before computing maskStartRow
so underflow cannot bypass masking in release builds, while preserving the
existing fallback behavior for valid lengths.
🪄 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: 1a6d1c03-3c28-4e2d-903c-74461ecfc60d
📥 Commits
Reviewing files that changed from the base of the PR and between 6812ddb and 13a161e48890634d74142b94da63d0f7aed8e05b.
📒 Files selected for processing (2)
csrc/xqa/mha_sm90.cuflashinfer/xqa.py
…infer-ai#4198) The #if SWAP_AB && SPEC_DEC variant of warpGrpApplyMask synthesized a causal mask from bit position ((1<<(maskCol+1))-1) and never read specDec.params.mask, and ignored its tok0WinBeg parameter, so full draft masks and per-row sliding windows were silently wrong whenever the SM90 fp8 kernel took the SWAP_AB specialization (q_seq_len * head_group_ratio <= 32). Rework it to mirror the non-SWAP_AB version with the axes transposed (acc rows = KV positions in tile, acc cols = q heads): each q-token column selects one TileMaskRow via specDec.needMask/loadTileMaskRow (bit index = KV row), the sliding window masks rows with globalRow < tok0WinBeg + idxQTokInCta, and the end mask handles the cacheSeqLen tail. 64-bit shifts are explicitly guarded and window comparisons stay signed (tok0WinBeg can be deeply negative). The legacy synthesized-causal path is kept as the params.mask == nullptr fallback, with its cacheSeqLen - SPEC_Q_SEQ_LEN subtraction clamped so a direct C++ caller with cacheSeqLen < SPEC_Q_SEQ_LEN cannot underflow it and silently skip masking. TileMaskRow loads are deduplicated per q token across the unrolled column loop. Python side: drop the SM90 fp8 fallbacks for sliding window and SWAP_AB eligibility in flashinfer/xqa.py (the kernel now handles both); ragged Q and sinks fallbacks stay. Sync the flashinfer.decode docstring and the deterministic-test path comments accordingly. Verified on H200 and H100 (JIT, FLASHINFER_CUDA_ARCH_LIST=9.0a): the 5 failing numeric cases (deterministic full-fp8 + 4 sliding-window full-mask SWAP_AB cases) now pass on both; full -k fp8 regression 1380 passed / 0 failed; spec-dec bf16+fp8 matrix all green. The 16 sinks+spec-dec cases fail only when the SM90 kernel is forced and are a pre-existing separate TODO. Signed-off-by: Han-Yin Chang <nick20350@gmail.com>
13a161e to
66d0f09
Compare
|
Closing this PR. This is not a claim that the underlying problem is fixed. We are not continuing this work, so the PR should not stay open as though we were. #4198 Closing changes nothing in-tree. The dispatch guard from #4199 is still live in The work stays on the PR for whoever picks it up — the rewritten SWAP_AB |
2 participants
📌 Description
Partially addresses #4198: fixes the SWAP_AB runtime-mask and sliding-window gaps on the SM90 fp8 XQA spec-decode path. Ragged Q and the PDL scale-load hazard remain open (follow-ups).
On the SWAP_AB layout (
q_seq_len * head_group_ratio <= 32),warpGrpApplyMasksynthesized a causal mask from the column index and never readspecDec.params.mask; thetok0WinBegargument was accepted but unused, so sliding-window begin-masking was also lost on this path.Changes:
csrc/xqa/mha_sm90.cu— rewrite theSWAP_AB && SPEC_DECvariant ofwarpGrpApplyMask:SpecDec::needMask/loadTileMaskRowhelpers as the non-SWAP_AB variant, with the lookup transposed for this layout (mask row selected by the q-token column, bit index = the thread's KV row);TileMaskRowloads are deduplicated across adjacent columns sharing a q token;globalRow < tok0WinBeg + maskCol), with signed comparisons (tok0WinBegcan be deeply negative) and explicit shift clamping instead of relying on PTXshl.b64clamp behavior;tok0WinBegwas dead);params.mask == nullptr.flashinfer/xqa.py— remove theuse_sliding_windowandswap_ab_eligibleconditions from the SM90 fp8 MHA fallback guard (the kernel now handles both); keep theuse_ragged_qandsinksfallbacks; update the note.flashinfer/decode.py— drop the now-stale docstring sentence saying sliding-window falls back to the generic kernel.tests/attention/test_xqa_batch_decode.py— comment-only: label the deterministic-test parametrizations ((2,4)→ SM90 fp8 SWAP_AB path,(2,16)→ non-SWAP_AB path).The kernel change also clamps
maskStartRowin the null-mask fallback (cacheSeqLen >= SPEC_Q_SEQ_LEN ? cacheSeqLen - SPEC_Q_SEQ_LEN : 0) so acacheSeqLen < SPEC_Q_SEQ_LENcall from direct C++ users cannot underflow and skip masking in release builds.🔍 Related Issues
Partially addresses #4198 (gaps (a) and (b); ragged Q and the
qScalePtr/kvScalePtrPDL load-ordering item are intentionally left for follow-ups). Builds on the deterministic tests introduced in #4199.🚀 Pull Request Checklist
✅ 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
Validation on a matching reported CUDA/PyTorch stack — H100 (SXM) / SM90, CUDA 13.0 (nvcc 13.0.88), torch 2.12.0+cu130, source build (JIT), on current main + this patch. Notes: the original report is from an H100 PCIe; the node's kernel driver here is 550.127.08 with the CUDA 13.0 forward-compatibility package, not a native 580-series driver.
spec_dec_sliding_window[300-full-False-…-63-…-2-32-2-4]spec_dec_sliding_window[300-full-False-…-63-…-4-64-4-2]spec_dec_sliding_window[300-full-False-…-127-…-2-32-2-4]spec_dec_sliding_window[300-full-False-…-127-…-4-64-4-2]mask_mode_deterministic[2-4-full-fp8]pytest tests/attention/test_xqa_batch_decode.py -k fp8The 16 remaining before-failures in the forced-on targeted subset are the known sinks + spec-dec accuracy TODO (unchanged by this PR; sinks still fall back in normal dispatch).
Primary development validation — H200 / SM90, CUDA 13.0, torch 2.13.0, source build (JIT): same 5 cases fail before / pass after; post-rebase targeted regression 68 passed / 0 failed; full
-k fp8(pre-rebase tree) 1378 passed / 0 failed.unittest, etc.).Reviewer Notes
warpGrpApplyMaskis untouched.qScalePtr/kvScalePtrPDL load-ordering hazard added to [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 is not addressed here; happy to take it as a follow-up.Summary by CodeRabbit
Bug Fixes
Documentation