Skip to content

Commit 92070b9

Browse files
lancellylfr-0531
authored andcommitted
[None][fix] kv_cache_manager_v2: count only newly-added scratch blocks in resize (NVIDIA#14531)
Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
1 parent 2bab85b commit 92070b9

1 file changed

Lines changed: 8 additions & 2 deletions

File tree

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
div_up,
7272
expect_type,
7373
filled_list,
74+
intersect,
7475
make_typed,
7576
map_optional,
7677
stream_wait_events,
@@ -706,8 +707,13 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo
706707
continue
707708
stale_beg, stale_end = stale_ranges[lc]
708709
if enable_scratch:
709-
num_scratch_blocks = len(scratch_ranges[lc])
710-
num_new_normal_blocks = (new_num_blocks - old_num_blocks) - num_scratch_blocks
710+
# Only newly added blocks consume slots below; scratch range may
711+
# extend before old_num_blocks when history_length < old_capacity.
712+
new_block_range = HalfOpenRange(old_num_blocks, new_num_blocks)
713+
num_new_blocks_using_scratch = len(
714+
intersect(scratch_ranges[lc], new_block_range)
715+
)
716+
num_new_normal_blocks = len(new_block_range) - num_new_blocks_using_scratch
711717
num_new_slots[lc] = num_new_normal_blocks * beam_width
712718
else:
713719
if old_num_blocks < stale_beg:

0 commit comments

Comments
 (0)