iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
b89e85e99f7ae832877effa4eab43adebcca4813
vLLM-4x-RTX-3090/patches/0023-sched-preserve-per-group-eviction-order-in-deferred-.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

331 lines
14 KiB
Diff
Raw Blame History

From 0271354dc4a882fdcc5ef7eca2f5c55fba812925 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:17:29 -0400
Subject: [PATCH 023/103] sched: preserve per-group eviction order in deferred
block free (port of upstream #48532)
The deferred free path flattened every KV cache group into one list and
reversed it once, so group order came out backwards relative to the immediate
free path. The reversal moves into pop_blocks_for_free.
---
.../test_partial_prefix_cache_hits.py | 5 +-
tests/v1/core/test_deferred_block_free.py | 196 +++++++++++++++++-
vllm/v1/core/kv_cache_coordinator.py | 13 +-
vllm/v1/core/kv_cache_manager.py | 6 +-
vllm/v1/core/sched/scheduler.py | 9 +-
5 files changed, 217 insertions(+), 12 deletions(-)
diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
index 393ec48..4798e0c 100644
--- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
@@ -1656,10 +1656,11 @@ def test_free_cow_retained_blocks_defers_until_copy_step_processed():
)
free = Scheduler._free_cow_retained_blocks
- # Copy step still in flight: deferred with its fence.
+ # Copy step still in flight: deferred with its fence, in the same order
+ # the immediate path would free them (the drain no longer reverses).
free(mock, list(blocks), fence_seq=3)
assert not freed
- assert mock.deferred_frees == deque([(3, blocks[::-1])])
+ assert mock.deferred_frees == deque([(3, blocks)])
# Copy step processed: freed immediately.
mock.processed_step_seq = 3
diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py
index 0c4738a..cbe58f8 100644
--- a/tests/v1/core/test_deferred_block_free.py
+++ b/tests/v1/core/test_deferred_block_free.py
@@ -20,7 +20,8 @@ import pytest
import torch
from vllm.config import VllmConfig
-from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import BlockHashWithGroupId, KVCacheBlockCopy
from vllm.v1.kv_cache_interface import CircularBufferSpec
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
@@ -698,3 +699,196 @@ def test_deferred_victim_stops_the_preemption_cascade():
out2 = scheduler.schedule()
assert first.request_id in out2.num_scheduled_tokens
assert first.status == RequestStatus.RUNNING
+
+def _make_pool_with_cached_blocks(num_blocks: int, block_size: int = 16):
+ """Build a BlockPool holding `num_blocks` allocated, hashed blocks.
+
+ Returns (pool, blocks_by_id). Each block gets a distinct hash so it is a
+ prefix-cache eviction candidate once freed (blocks without a hash are
+ always evicted first and would mask ordering).
+ """
+ pool = BlockPool(
+ num_gpu_blocks=num_blocks + 1, enable_caching=True, hash_block_size=block_size
+ )
+ blocks = pool.get_new_blocks(num_blocks)
+ for i, block in enumerate(blocks):
+ block.set_block_hash(
+ BlockHashWithGroupId(f"hash{i}".encode()), num_tokens=block_size
+ )
+ return pool, {b.block_id: b for b in blocks}
+
+
+def _eviction_order(pool: BlockPool) -> list[int]:
+ """Drain the pool's free queue, returning block ids in eviction order
+ (first returned == evicted first)."""
+ order = []
+ while pool.get_num_free_blocks() > 0:
+ order.append(pool.get_new_blocks(1)[0].block_id)
+ return order
+
+
+def test_deferred_free_matches_immediate_free_eviction_order():
+ """Regression for the deferred-free path collapsing per-group eviction
+ order for multi-group (hybrid) KV caches.
+
+ Group A owns blocks [1, 2, 3] (allocation order) and group B owns
+ [4, 5, 6]. The immediate-free path frees each group in its own
+ reversed ``free_blocks`` call; the fixed deferred-free path must
+ reproduce exactly that order, not a single flattened reversal.
+ """
+ group_a_alloc = [1, 2, 3]
+ group_b_alloc = [4, 5, 6]
+
+ # Immediate path: KVCacheCoordinator.free() -> one reversed free_blocks
+ # call per single-type manager (group).
+ pool_imm, by_id = _make_pool_with_cached_blocks(6)
+ for group in (group_a_alloc, group_b_alloc):
+ pool_imm.free_blocks([by_id[i] for i in reversed(group)])
+ immediate_order = _eviction_order(pool_imm)
+
+ # Fixed deferred path: groups are kept separate (each already reversed)
+ # and freed in their own call on drain -- mirroring the immediate path.
+ pool_fixed, by_id = _make_pool_with_cached_blocks(6)
+ ordered_groups = [list(reversed(group_a_alloc)), list(reversed(group_b_alloc))]
+ for group in ordered_groups:
+ pool_fixed.free_blocks([by_id[i] for i in group])
+ fixed_order = _eviction_order(pool_fixed)
+
+ assert fixed_order == immediate_order == [3, 2, 1, 6, 5, 4]
+
+ # The old (buggy) deferred path flattened both groups then reversed once,
+ # which produces a different order. Guard against regressing back to it.
+ pool_buggy, by_id = _make_pool_with_cached_blocks(6)
+ flattened = group_a_alloc + group_b_alloc
+ pool_buggy.free_blocks([by_id[i] for i in reversed(flattened)])
+ buggy_order = _eviction_order(pool_buggy)
+ assert buggy_order == [6, 5, 4, 3, 2, 1]
+ assert buggy_order != immediate_order
+
+
+def test_pop_blocks_for_free_returns_eviction_order():
+ """``pop_blocks_for_free`` must hand back blocks already in eviction
+ order, so the deferred-free path can pass them to ``free_blocks``
+ unchanged. Reversing the flat list instead would emit whole groups in
+ the wrong order relative to the immediate-free path."""
+ scheduler = _create_deferring_scheduler()
+ coordinator = scheduler.kv_cache_manager.coordinator
+
+ request = create_requests(
+ num_requests=1,
+ num_tokens=NUM_PROMPT_TOKENS,
+ max_tokens=5,
+ )[0]
+ scheduler.add_request(request)
+ scheduler.schedule()
+
+ allocated = [
+ list(manager.req_to_blocks[request.request_id])
+ for manager in coordinator.single_type_managers
+ ]
+ expected = [block for group in allocated for block in reversed(group)]
+
+ assert coordinator.pop_blocks_for_free(request.request_id) == expected
+
+
+def _make_two_group_kv_cache_manager(num_blocks: int, block_size: int = 16):
+ """Build a KVCacheManager with two single-type managers (a hybrid
+ full-attention + sliding-window config), so CoW copies can be queued on
+ more than one group."""
+ from math import lcm
+
+ from vllm.v1.core.kv_cache_manager import KVCacheManager
+ from vllm.v1.kv_cache_interface import (
+ FullAttentionSpec,
+ KVCacheConfig,
+ KVCacheGroupSpec,
+ SlidingWindowSpec,
+ )
+
+ config = KVCacheConfig(
+ num_blocks=num_blocks,
+ kv_cache_tensors=[],
+ kv_cache_groups=[
+ KVCacheGroupSpec(
+ ["layer1"],
+ FullAttentionSpec(
+ block_size=block_size,
+ num_kv_heads=1,
+ head_size=1,
+ dtype=torch.float32,
+ ),
+ ),
+ KVCacheGroupSpec(
+ ["layer2"],
+ SlidingWindowSpec(
+ block_size=block_size,
+ num_kv_heads=1,
+ head_size=1,
+ dtype=torch.float32,
+ sliding_window=2 * block_size,
+ ),
+ ),
+ ],
+ )
+ return KVCacheManager(
+ config,
+ max_model_len=1024,
+ scheduler_block_size=lcm(
+ *(g.kv_cache_spec.block_size for g in config.kv_cache_groups)
+ ),
+ hash_block_size=block_size,
+ enable_caching=True,
+ )
+
+
+def test_cow_retained_blocks_span_groups_and_free_order_is_stable():
+ """CoW retentions from multiple KV cache groups are concatenated into one
+ flat list by ``take_kv_cache_block_copies``. Unlike per-request eviction,
+ ``_free_cow_retained_blocks`` never reverses, so the deferred path frees
+ that list in exactly the same order as the immediate path -- the
+ group-flattening does not reintroduce the eviction-order divergence.
+ """
+ manager = _make_two_group_kv_cache_manager(num_blocks=20)
+ coordinator = manager.coordinator
+ assert len(coordinator.single_type_managers) == 2
+
+ # Queue a CoW (source, cow) pair on each of the two groups.
+ blocks = manager.block_pool.get_new_blocks(4)
+ coordinator.single_type_managers[0]._pending_cow_copies.append(
+ (blocks[0], blocks[1])
+ )
+ coordinator.single_type_managers[1]._pending_cow_copies.append(
+ (blocks[2], blocks[3])
+ )
+
+ _copies, retained = manager.take_kv_cache_block_copies()
+ # Retentions from both groups are flattened into one list, group 0 first.
+ assert [b.block_id for b in retained] == [b.block_id for b in blocks]
+
+ # Drive the scheduler free helper on the same retained ordering, immediate
+ # vs deferred, and confirm identical eviction order.
+ retained_ids = [b.block_id for b in retained]
+
+ imm_scheduler = _create_deferring_scheduler()
+ imm_scheduler.defer_block_free = False
+ imm_pool, imm_by_id = _make_pool_with_cached_blocks(len(retained_ids))
+ imm_scheduler.kv_cache_manager.block_pool = imm_pool
+ imm_scheduler._free_cow_retained_blocks(
+ [imm_by_id[i] for i in retained_ids], fence_seq=1
+ )
+ immediate_order = _eviction_order(imm_pool)
+
+ def_scheduler = _create_deferring_scheduler()
+ def_pool, def_by_id = _make_pool_with_cached_blocks(len(retained_ids))
+ def_scheduler.kv_cache_manager.block_pool = def_pool
+ # fence in the future -> deferred; then drain once the fence is processed.
+ def_scheduler.processed_step_seq = 0
+ def_scheduler._free_cow_retained_blocks(
+ [def_by_id[i] for i in retained_ids], fence_seq=1
+ )
+ assert def_scheduler.deferred_frees
+ def_scheduler.processed_step_seq = 1
+ def_scheduler._drain_deferred_frees()
+ deferred_order = _eviction_order(def_pool)
+
+ assert deferred_order == immediate_order
diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py
index afb03ac..e3f31a6 100644
--- a/vllm/v1/core/kv_cache_coordinator.py
+++ b/vllm/v1/core/kv_cache_coordinator.py
@@ -394,18 +394,23 @@ class KVCacheCoordinator(ABC):
Pop the request's bookkeeping from all single-type managers and
return its blocks without returning them to the block pool. The
caller must eventually pass the returned blocks to
- `block_pool.free_blocks`, freeing them in reverse order (so that
- tail blocks are evicted first).
+ `block_pool.free_blocks` as-is: they are already ordered so that
+ tail blocks are evicted first, within each KV cache group.
Args:
request_id: The request ID.
Returns:
- The request's blocks in allocation order.
+ The request's blocks in eviction order (each group's blocks
+ reversed, then concatenated group by group).
"""
blocks: list[KVCacheBlock] = []
for manager in self.single_type_managers:
- blocks.extend(manager.pop_blocks_for_free(request_id))
+ # Reverse within each group so the concatenation is already in
+ # eviction order, matching `free()`, which frees each group in its
+ # own reversed `free_blocks()` call. Reversing the flat list
+ # instead would emit whole groups in the wrong order.
+ blocks.extend(reversed(manager.pop_blocks_for_free(request_id)))
return blocks
def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 226da3a..675ca1e 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -632,13 +632,15 @@ class KVCacheManager:
def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
"""Pop the request's bookkeeping and return its blocks without
returning them to the block pool. The caller must eventually free
- them in reverse order (so that tail blocks are evicted first).
+ them as-is: they are already ordered so that tail blocks are
+ evicted first, within each KV cache group.
Args:
request: The request to pop the blocks for.
Returns:
- The request's blocks in allocation order.
+ The request's blocks in eviction order (each group's blocks
+ reversed, then concatenated group by group).
"""
return self.coordinator.pop_blocks_for_free(request.request_id)
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index b16ae9d..ae9eeb6 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2710,7 +2710,9 @@ class Scheduler(SchedulerInterface):
if not self.defer_block_free or fence_seq <= self.processed_step_seq:
self.kv_cache_manager.block_pool.free_blocks(blocks)
return
- self.deferred_frees.append((fence_seq, blocks[::-1]))
+ # The immediate path above frees these as-is, so enqueue them as-is:
+ # the drain no longer reverses.
+ self.deferred_frees.append((fence_seq, blocks))
def _drain_deferred_frees(self):
"""Return deferred blocks whose fence step has completed.
@@ -2724,8 +2726,9 @@ class Scheduler(SchedulerInterface):
if fence > self.processed_step_seq:
break
_, blocks = self.deferred_frees.popleft()
- # Free in reverse order so that the tail blocks are evicted first.
- self.kv_cache_manager.block_pool.free_blocks(reversed(blocks))
+ # Already in eviction order (reversed per KV cache group by
+ # `pop_blocks_for_free`), so free as-is.
+ self.kv_cache_manager.block_pool.free_blocks(blocks)
def get_num_unfinished_requests(self) -> int:
if self._pause_state == PauseState.PAUSED_ALL:
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink