Skip to content

[None][fix] kv_cache_manager_v2: count only newly-added scratch blocks in resize - #14531

Merged
lancelly merged 2 commits into
NVIDIA:feat/deepseek_v4from
lancelly:fix/kv-cache-v2-swa-scratch-resize-slot-count
May 26, 2026
Merged

[None][fix] kv_cache_manager_v2: count only newly-added scratch blocks in resize#14531
lancelly merged 2 commits into
NVIDIA:feat/deepseek_v4from
lancelly:fix/kv-cache-v2-swa-scratch-resize-slot-count

Conversation

@lancelly

@lancelly lancelly commented May 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fixes IndexError: pop from empty list in KVCache.resize() when SWA scratch reuse is enabled and history_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, the enable_scratch branch of resize() was computing new-slot demand from the full scratch_ranges[lc]:

num_scratch_blocks = len(scratch_ranges[lc])
num_new_normal_blocks = (new_num_blocks - old_num_blocks) - num_scratch_blocks
num_new_slots[lc] = num_new_normal_blocks * beam_width

But the block construction loop below only iterates [old_num_blocks, new_num_blocks) and only those blocks pop slots:

for ordinal in typed_range(old_num_blocks, new_num_blocks):
    ...
    if enable_scratch and ordinal in scratch_ranges[lc]:
        continue
    slot = slots[lc].pop()

compute_scratch_range() derives its lower bound from history_length, not from old_capacity. After #14403 (Fix MTP by scratch reuse rewind), history_length is allowed to be smaller than old_capacity by up to max_rewind_len, so scratch_ranges[lc] can start before old_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's slots[lc].pop() underflows.

Concrete example

tokens_per_block = 128
old_capacity     = 1024  -> old_num_blocks = 8
new_capacity     = 1152  -> new_num_blocks = 9
history_length   = 896   -> input_range starts at block 7
scratch_ranges[lc] = [7, 8)        # block 7 is OLD, not newly added
newly added range  = [8, 9)        # construction loop iterates only block 8

old code:    num_new_slots = (9-8) - 1 = 0
loop pops:   block 8 needs 1 slot   -> IndexError

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_slots is intentionally unchanged — it sizes the scratch slot pool to serve the full scratch range, not just the newly-added portion.

…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>
@lancelly

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50206 [ run ] triggered by Bot. Commit: c3ec577 Link to invocation

Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py Outdated
Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py Outdated
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50206 [ run ] completed with state SUCCESS. Commit: c3ec577
/LLM/main/L0_MergeRequest_PR pipeline #39744 completed with status: 'SUCCESS'

CI Report

Link to invocation

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>
… 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
lancelly force-pushed the fix/kv-cache-v2-swa-scratch-resize-slot-count branch from 8713bfe to 2adfa0a Compare May 25, 2026 13:19
@lancelly

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50217 [ run ] triggered by Bot. Commit: 2adfa0a Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #50217 [ run ] completed with state SUCCESS. Commit: 2adfa0a
/LLM/main/L0_MergeRequest_PR pipeline #39754 completed with status: 'SUCCESS'

CI Report

Link to invocation

@lancelly
lancelly merged commit e4a5978 into NVIDIA:feat/deepseek_v4 May 26, 2026
6 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants