Fix SWA cache loc slicing for all attention backends - #29460
Merged
Merged
Conversation
Slice swa_out_cache_loc to match cache_loc length before passing to set_kv_buffer. Without this, the SWA location tensor can be longer than the KV tensors when sliding-window attention is active, causing shape mismatches.
merrymercy
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
hebiao064 and
ispobock
as code owners
June 26, 2026 21:51
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Contributor
Author
|
/tag-and-rerun-ci |
github-actions
Bot
added
blackwell
Move the swa_out_cache_loc slicing from the TRT-LLM MHA backend into KVWriteLoc.__post_init__ so it applies to all attention backends.
merrymercy
requested review from
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
hzh0425,
xiezhq-hermann and
yizhang2077
as code owners
June 27, 2026 00:52
merrymercy
changed the title
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
swa_out_cache_locslicing intoKVWriteLoc.__post_init__so it applies to all attention backends, not just TRT-LLM MHA.swa_out_cache_locis computed once at metadata-init time from the full (possibly padded)out_cache_loc. Piecewise CUDA graphs later narrowout_cache_loctoreal_num_tokensper layer viaradix_attention.py, butswa_out_cache_locon the metadata is never narrowed. This causes a shape mismatch whenKVWriteLocbundles them together forset_kv_buffer.swa_locto matchloc's length atKVWriteLocconstruction time, fixing the issue for all 8 backends (flashinfer, triton, flashattention, trtllm_mha, aiter, xpu, torch_native, musa).Original commits
53f21e2feCI States
Latest PR Test (Base): ❌ Run #28273505004
Latest PR Test (Extra): ❌ Run #28273504905