SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

[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