[None][fix] kv_cache_manager_v2: count only newly-added scratch blocks in resize - #14531
Merged
lancelly merged 2 commits intoMay 26, 2026
Conversation
…s in resize KVCache.resize() under enable_swa_scratch_reuse computes new-slot demand using the full length of scratch_ranges[lc], but only blocks in [old_num_blocks, new_num_blocks) actually pop slots in the block construction loop. Since NVIDIA#14403 relaxed the assertion to allow history_length < old_capacity (rewind window for MTP / speculative decoding), scratch_ranges[lc] can begin before old_num_blocks. In that case the allocation phase under-requests slots and the construction phase trips `IndexError: pop from empty list` on slots[lc].pop(). Intersect scratch_ranges[lc] with [old_num_blocks, new_num_blocks) before counting so the request matches what the loop will actually consume. _take_excess_scratch_slots is intentionally unchanged: it sizes the scratch slot pool to serve the full scratch range, not just the newly added portion. Reproduces with DSv4 + SWA scratch reuse + MTP; workarounds were TRTLLM_DSV4_ENABLE_SWA_SCRATCH_REUSE=0 or a build pre-NVIDIA#14403. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50206 [ run ] triggered by Bot. Commit: |
lowsfer
approved these changes
May 25, 2026
Collaborator
|
PR_Github #50206 [ run ] completed with state |
lowsfer
approved these changes
May 25, 2026
lancelly
added a commit
to lancelly/TensorRT-LLM
that referenced
this pull request
May 25, 2026
… names in resize Address review nits on NVIDIA#14531: - Rename num_scratch_blocks -> num_new_blocks_using_scratch to make clear this counts scratch blocks within the newly-added range only, not the full scratch_ranges[lc]. - Use len(new_block_range) instead of recomputing (new_num_blocks - old_num_blocks), since new_block_range is already in scope. No behavior change. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lowsfer
approved these changes
May 25, 2026
… names in resize Address review nits on NVIDIA#14531: - Rename num_scratch_blocks -> num_new_blocks_using_scratch to make clear this counts scratch blocks within the newly-added range only, not the full scratch_ranges[lc]. - Use len(new_block_range) instead of recomputing (new_num_blocks - old_num_blocks), since new_block_range is already in scope. No behavior change. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lancelly
force-pushed
the
fix/kv-cache-v2-swa-scratch-resize-slot-count
branch
from
May 25, 2026 13:19
8713bfe to
2adfa0a
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50217 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50217 [ run ] completed with state |
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 10, 2026
…s in resize (NVIDIA#14531) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 11, 2026
…s in resize (NVIDIA#14531) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
lfr-0531
pushed a commit
to lfr-0531/TensorRT-LLM
that referenced
this pull request
Jun 12, 2026
…s in resize (NVIDIA#14531) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
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
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
Fixes
IndexError: pop from empty listinKVCache.resize()when SWA scratch reuse is enabled andhistory_length < old_capacity(rewind window allowed since #14403). Reproduces with DSv4 + MTP + block_reuse.Root cause
In
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py, theenable_scratchbranch ofresize()was computing new-slot demand from the fullscratch_ranges[lc]:But the block construction loop below only iterates
[old_num_blocks, new_num_blocks)and only those blocks pop slots:compute_scratch_range()derives its lower bound fromhistory_length, not fromold_capacity. After #14403 (Fix MTP by scratch reuse rewind),history_lengthis allowed to be smaller thanold_capacityby up tomax_rewind_len, soscratch_ranges[lc]can start beforeold_num_blocks. When that happens, scratch blocks that sit in the old range are still subtracted from the allocation count but are never re-encountered in the construction loop, so the allocation under-requests by exactly that overlap and the loop'sslots[lc].pop()underflows.Concrete example
Fix
Intersect
scratch_ranges[lc]with[old_num_blocks, new_num_blocks)before counting, so the request matches what the loop will actually consume._take_excess_scratch_slotsis intentionally unchanged — it sizes the scratch slot pool to serve the full scratch range, not just the newly-added portion.