[Fix] Fix XQA V tile reading from wrong page when nbVItersPerXIter > 1 (#3022)
<!-- .github/pull_request_template.md --> ## ๐ Description <!-- What does this PR do? Briefly describe the changes and why theyโre needed. --> ## Summary Fix incorrect XQA attention results on architectures with `cacheVTileSeqLen = 32` (SM120/SM121) when `head_dim = 256` and `page_size < 64`. ## Bug On SM120 with `head_dim=256` and `page_size` of 16 or 32, the XQA `mha.cu` kernel produces incorrect attention outputs. All `page_size=64` cases pass; all `page_size < 64` cases fail. The bug is independent of dtype, `kv_layout`, `batch_size`, `window_left`, and other parameters. SM90 and SM100 are unaffected. ## Root Cause In `mha.cu`, the V tile page advancement logic (`loadVTilePart` lambda) has two branches based on `xIterSeqStride` vs `tokensPerPage`. The `else` branch (when `xIterSeqStride > tokensPerPage`) assumes `nbVItersPerXIter == 1`, meaning each warp X tile (64 tokens) is covered by a single V tile load. This assumption is violated under a specific combination of compile-time constants: - `cacheVTileSeqLen = 32` (SM120/SM121) - `head_dim = 256` โ `gemm1WarpsPerGrp = 4`, `gemm1NbWarpGrps = 1` - โ `cacheVTileSeqStride = 32 ร 1 = 32 < warpTile.x = 64` - โ **`nbVItersPerXIter = 2`** With `nbVItersPerXIter = 2`, each warp X tile requires two V iterations (vIter=0 covering tokens [0, 32) and vIter=1 covering tokens [32, 64)). When `page_size < 64`, these two V iterations land on different pages. However, `loadPages()` was only called after the last V iteration (`vIter == nbVItersPerXIter - 1`), leaving the page index stale for vIter=1 โ it reads KV cache data from the wrong page. On SM90/SM100 (`cacheVTileSeqLen = 64`), `cacheVTileSeqStride = 64` and `nbVItersPerXIter = 1`, so the bug never triggers. On SM120 with `head_dim = 128`, `gemm1NbWarpGrps = 2` makes `cacheVTileSeqStride = 64`, also avoiding the issue. ## Fix Replace the single page advancement at the end of each X iteration with per-V-iteration page advancement: - **Intermediate V iteration** (`vIter < nbVItersPerXIter - 1`): advance `idxPageBeg` by `step_per_viter = cacheVTileSeqStride / tokensPerPage` and reload pages. - **Last V iteration, last beam** (`isLastVIter && isLastBeam`): advance with CTA-tile boundary wrapping (multi-block mode), same as original but using `step_per_viter`. - **Last V iteration, not last beam** (`isLastVIter && !isLastBeam`): reset `idxPageBeg` backward so the next beam restarts from `vIter=0`'s page position. When `nbVItersPerXIter == 1`, only the first branch fires and the behavior is identical to the original code โ no performance or correctness impact on existing working paths. ## Test Plan - [x] `test_trtllm_batch_decode` with `backend=xqa`, `head_dim=256`, all page sizes (16, 32, 64) on SM120 - [x] `test_trtllm_batch_decode` with `backend=xqa`, `head_dim=128`, all page sizes on SM120 (regression check) - [x] `test_trtllm_batch_decode` with `backend=xqa`, `head_dim=256`, all page sizes on SM90 (regression check) ## ๐ Related Issues <!-- Link any related issues here --> ## ๐ 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. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## ๐งช Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved attention page-loading to correctly handle multi-iteration stride scenarios, including proper advance, rewind and reload behavior across iteration and beam boundaries for more reliable caching. * **Tests** * Re-enabled tests for configurations with head_dim == 256 that were previously skipped to verify multi-iteration page-loading. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Q
qsang-nv committed
103fcf862b902e100775ac781a0de36d067de3a4
Parent: 611af1b
Committed by GitHub <noreply@github.com>
on 5/13/2026, 1:37:31 AM