Skip to content

[Fix] Fix XQA V tile reading from wrong page when nbVItersPerXIter > 1 #3022

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
qsang-nv wants to merge 2 commits into
base: main
Choose a base branch
from
+13 −11
Open

[Fix] Fix XQA V tile reading from wrong page when nbVItersPerXIter > 1 #3022

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f9f4680
fix xqa bug
qsang-nv Apr 9, 2026
bf66ca9
Merge branch 'main' into fix_xqa_sm120_headdim_256
qsang-nv May 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions csrc/xqa/mha.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2249,13 +2249,19 @@ CUBIN_EXPORT __global__
loadPages(idxPageBeg);
}
} else {
assert(nbVItersPerXIter == 1);
if ((idxBeam == beamWidth - 1 || isConvergedTile(seqIter)) &&
vIter == nbVItersPerXIter - 1) {
auto const step = exactDiv(xIterSeqStride, tokensPerPage);
idxPageBeg += (idxPageBeg % nbPagesPerCtaTile + step >= nbPagesPerCtaTile
? nbPagesPerCtaTile * (nbSubSeqPerSeq - 1) + step
: step);
constexpr auto step_per_viter = exactDiv(cacheVTileSeqStride, tokensPerPage);
bool const isLastVIter = (vIter == nbVItersPerXIter - 1);
bool const isLastBeam = (idxBeam == beamWidth - 1 || isConvergedTile(seqIter));
if (isLastVIter && isLastBeam) {
idxPageBeg += (idxPageBeg % nbPagesPerCtaTile + step_per_viter >= nbPagesPerCtaTile
? nbPagesPerCtaTile * (nbSubSeqPerSeq - 1) + step_per_viter
: step_per_viter);
loadPages(idxPageBeg);
} else if (isLastVIter) {
idxPageBeg -= step_per_viter * (nbVItersPerXIter - 1);
loadPages(idxPageBeg);
} else {
idxPageBeg += step_per_viter;
loadPages(idxPageBeg);
}
Comment on lines +2252 to 2266
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Apr 9, 2026

Choose a reason for hiding this comment

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

medium

The logic for idxPageBeg update is complex and involves multiple branches. Consider extracting this into a helper function or a more readable structure to improve maintainability, as this is a critical part of the V-tile loading logic.

}
Expand Down
4 changes: 0 additions & 4 deletions tests/attention/test_trtllm_gen_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1504,10 +1504,6 @@ def test_trtllm_batch_decode(
if backend == "xqa" and non_contiguous_query:
pytest.skip("xqa backend does not support non-contiguous query")

# fixme(qsang-nv): failing tests for xqa + head dim 256.
if backend == "xqa" and head_dim == 256:
pytest.skip("xqa backend + head dim 256 cases have precision issues")

# General set of tests for trtllm-gen decode
_test_trtllm_batch_decode(
backend,
Expand Down
Loading