[Bugfix] Fix sliding-window KV allocation scaling

tests/v1/core/test_kv_cache_utils.py

@@ -1547,7 +1547,6 @@ def test_get_max_concurrency_for_kv_cache_config():
         enable_chunked_prefill=True,
         max_model_len=model_config.max_model_len,
         is_encoder_decoder=model_config.is_encoder_decoder,
-        # Pin to sync: SWA per-request bounds grow with overlapping batches.
         async_scheduling=False,
     )
 
@@ -1585,7 +1584,7 @@ def test_get_max_concurrency_for_kv_cache_config():
     assert max_concurrency_full_attention == 1.5
 
     kv_cache_config_sliding_window = KVCacheConfig(
-        num_blocks=129 * 3,
+        num_blocks=65 * 3,
         kv_cache_tensors=[],
         kv_cache_groups=[
             KVCacheGroupSpec([f"layer_{i}" for i in range(32)], sliding_window_spec),
@@ -1597,7 +1596,7 @@ def test_get_max_concurrency_for_kv_cache_config():
     assert max_concurrency_sliding_window == 3
 
     kv_cache_config_hybrid_model = KVCacheConfig(
-        num_blocks=(1024 + 129) * 3,
+        num_blocks=(1024 + 65) * 3,
         kv_cache_tensors=[],
         kv_cache_groups=[
             KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
@@ -1617,11 +1616,11 @@ def test_get_max_concurrency_for_kv_cache_config():
     assert max_concurrency == max_concurrency_hybrid_model
 
     # Unequal group sizes in the standard layout: each group's pages cost
-    # whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
+    # whole pool blocks, so a request needs 1024 + 65 = 1089 blocks — the
     # same as the equal-hybrid case above, regardless of the second group
     # holding only 2 layers.
     kv_cache_config_unequal_groups = KVCacheConfig(
-        num_blocks=1153 * 3,
+        num_blocks=1089 * 3,
         kv_cache_tensors=[],
         kv_cache_groups=[
             KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
@@ -1637,16 +1636,16 @@ def test_get_max_concurrency_for_kv_cache_config():
 
     # UniformTypeKVCacheSpecs group (worker config shape): the aggregated
     # spec's memory/page ratio equals a single layer's page count, so the
-    # group needs 1024 blocks and the request 1153 in total. The previous
+    # group needs 1024 blocks and the request 1089 in total. The previous
     # formula normalized both groups' memory by the first group's page size,
-    # reporting 3459/1057 = 3.27 here instead of 3 — and a different value
-    # again for the scheduler-config shape below.
+    # reporting a value above 3 here — and a different value again for the
+    # scheduler-config shape below.
     uniform_full_spec = UniformTypeKVCacheSpecs(
         block_size=full_attention_spec.block_size,
         kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
     )
     kv_cache_config_uniform_group = KVCacheConfig(
-        num_blocks=1153 * 3,
+        num_blocks=1089 * 3,
         kv_cache_tensors=[],
         kv_cache_groups=[
             KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -629,7 +629,7 @@ def test_get_num_blocks_to_allocate():
         num_kv_heads=1,
         head_size=1,
         dtype=torch.float32,
-        sliding_window=4,  # Placeholder value, not related to test result
+        sliding_window=100,  # Large enough not to cap this generic allocation test
     )
 
     block_pool = BlockPool(
@@ -732,14 +732,13 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
     )
 
 
-def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
+def test_sliding_window_caps_predictor_and_allocator_at_window_size():
     """In forward steps, `get_num_blocks_to_allocate` must return exactly what
     `allocate_new_blocks` will pull; otherwise `block_pool.get_new_blocks`
     raises `ValueError: Cannot get N free blocks from the pool`.
     """
     block_size = 2
     sliding_window = 8  # 4-block live window
-    cap = sliding_window // block_size
 
     spec = SlidingWindowSpec(
         block_size=block_size,
@@ -757,11 +756,12 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
         enable_caching=False,
         kv_cache_group_id=0,
         scheduler_block_size=spec.block_size,
-        max_admission_blocks_per_request=cap,
+        max_admission_blocks_per_request=10**9,
     )
 
     request_id = "req"
     total_computed = 0
+    allocated_per_step = []
     # Walk through request forward steps. Check num_blocks returned by
     # `get_num_blocks_to_allocate` matches what `allocate_new_blocks` pulls
     for num_tokens in (4, 8, 12, 16):
@@ -780,4 +780,15 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
             f"num_tokens={num_tokens}: predictor returned {predicted} "
             f"but allocator pulled {len(new_blocks)}"
         )
+        allocated_per_step.append(len(new_blocks))
         total_computed = num_tokens
+
+    assert allocated_per_step == [2, 2, 1, 0]
+    assert len(manager.req_to_blocks[request_id]) == 8
+    assert (
+        sum(
+            block is block_pool.null_block
+            for block in manager.req_to_blocks[request_id]
+        )
+        == 3
+    )

tests/v1/core/test_swa_inflight_window_free.py

@@ -1,12 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Out-of-window block frees vs in-flight GPU steps.
+"""Out-of-window block allocation vs in-flight GPU steps.
 
-With async scheduling / PP, `num_computed_tokens` optimistically includes
-tokens of unprocessed steps whose attention windows still read the blocks
-just below the optimistic boundary (and rejected spec tokens can roll it
-back), so `allocate_slots` frees on the processed-token basis:
-`num_computed_tokens - num_in_flight_tokens`.
+Sliding-window allocation preserves absolute block-table positions with null
+blocks while retaining only the window tail. Chunked-local attention still
+frees on the processed-token basis because it retains the current full chunk.
 """
 
 import torch
@@ -95,9 +93,8 @@ def test_num_in_flight_tokens_accounting():
     assert request.num_in_flight_tokens == 0
 
 
-def test_swa_free_waits_for_in_flight_step():
-    """Async: out-of-window blocks stay allocated until the step that still
-    reads them has been processed."""
+def test_swa_caps_blocks_for_in_flight_step():
+    """Async: logical slots outside the retained window use null blocks."""
     scheduler = _create_swa_scheduler(async_scheduling=True)
     request = create_requests(
         num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
@@ -109,30 +106,25 @@ def test_swa_free_waits_for_in_flight_step():
     out0 = scheduler.schedule()  # prefill, in flight from here on
     free_after_prefill = block_pool.get_num_free_blocks()
 
-    # Decode scheduled while the prefill still reads the out-of-window blocks:
-    # they must not be freed yet.
+    # The prefill allocation retains only the window tail while preserving
+    # absolute block-table positions with null blocks.
     out1 = scheduler.schedule()
-    assert _num_null_blocks(scheduler, req_id) == 0
+    assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
     assert block_pool.get_num_free_blocks() == free_after_prefill
 
-    # Prefill output processed; the next allocate frees the out-of-window
-    # blocks.
+    # Processing the prefill does not free the already-null slots again.
     scheduler.update_from_output(out0, _make_model_runner_output(out0))
     scheduler.schedule()
     assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
-    assert (
-        block_pool.get_num_free_blocks()
-        == free_after_prefill + NUM_OUT_OF_WINDOW_BLOCKS
-    )
+    assert block_pool.get_num_free_blocks() == free_after_prefill
     # Not double-freed on the following steps.
     scheduler.update_from_output(out1, _make_model_runner_output(out1))
     scheduler.schedule()
     assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
 
 
-def test_swa_free_immediate_when_sync():
-    """Sync: no in-flight step at schedule time, frees happen at the first
-    decode allocation as before."""
+def test_swa_caps_blocks_immediately_when_sync():
+    """Sync: the initial allocation retains only the window tail."""
     scheduler = _create_swa_scheduler(async_scheduling=False)
     request = create_requests(
         num_requests=1, num_tokens=NUM_PROMPT_TOKENS, block_size=BLOCK_SIZE
@@ -141,14 +133,15 @@ def test_swa_free_immediate_when_sync():
     req_id = request.request_id
 
     out0 = scheduler.schedule()
+    assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
     scheduler.update_from_output(out0, _make_model_runner_output(out0))
     assert request.num_in_flight_tokens == 0
 
     scheduler.schedule()
     assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
 
 
-def test_swa_admission_cap_accounts_for_overlapping_batches():
+def test_swa_admission_cap_is_independent_of_in_flight_tokens():
     spec = SlidingWindowSpec(
         block_size=16,
         num_kv_heads=1,
@@ -159,13 +152,12 @@ def test_swa_admission_cap_accounts_for_overlapping_batches():
     base = spec.max_admission_blocks_per_request(
         max_in_flight_tokens=1024, max_model_len=16384
     )
-    # (1024 - 1 + 1024) tokens -> 128 blocks, +1 for window misalignment.
-    assert base == 129
+    # 1024 - 1 tokens -> 64 blocks, +1 for window misalignment.
+    assert base == 65
     overlapped = spec.max_admission_blocks_per_request(
         max_in_flight_tokens=2 * 1024, max_model_len=16384
     )
-    # One extra in-flight chunk is held back: (1024 - 1 + 2 * 1024) tokens.
-    assert overlapped == 193
+    assert overlapped == base
 
 
 def test_chunked_local_free_waits_for_in_flight_step():
@@ -248,12 +240,11 @@ def test_connector_finish_frees_on_settled_basis():
     out0 = scheduler.schedule()  # prefill, in flight from here on
     scheduler.schedule()  # decode over-scheduled: num_computed_tokens optimistic
 
-    # Finishing now (connector store) must NOT prune out-of-window blocks the
-    # in-flight prefill still reads.
+    # The allocation already represents out-of-window logical slots as nulls.
     scheduler._connector_finished(request)
-    assert _num_null_blocks(scheduler, req_id) == 0
+    assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS
 
-    # Once the in-flight step settles, the same prune releases them.
+    # Settling the in-flight step must not change the capped allocation.
     scheduler.update_from_output(out0, _make_model_runner_output(out0))
     scheduler._connector_finished(request)
     assert _num_null_blocks(scheduler, req_id) == NUM_OUT_OF_WINDOW_BLOCKS

vllm/v1/core/single_type_kv_cache_manager.py

@@ -151,6 +151,12 @@ def _has_partial_local_hit(
             and num_local_computed_tokens % self.block_size != 0
         )
 
+    def _get_num_required_blocks(self, num_tokens: int) -> int:
+        return cdiv(num_tokens, self.block_size)
+
+    def _get_num_allocated_blocks(self, request_id: str) -> int:
+        return len(self.req_to_blocks.get(request_id, ()))
+
     def get_num_blocks_to_allocate(
         self,
         request_id: str,
@@ -185,7 +191,7 @@ def get_num_blocks_to_allocate(
             The number of blocks to allocate.
         """
 
-        num_required_blocks = cdiv(num_tokens, self.block_size)
+        num_required_blocks = self._get_num_required_blocks(num_tokens)
         if apply_admission_cap and self._max_admission_blocks_per_request is not None:
             # Recycling-aware specs (SWA, chunked-local) cap the per-request
             # reservation here so admission matches the startup pool sizer
@@ -199,7 +205,8 @@ def get_num_blocks_to_allocate(
             num_required_blocks = min(
                 num_required_blocks, self._max_admission_blocks_per_request
             )
-        num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
+        num_req_block_slots = len(self.req_to_blocks.get(request_id, ()))
+        num_req_blocks = self._get_num_allocated_blocks(request_id)
 
         if request_id in self.num_cached_block:
             # Fast-path: a running request won't have any new prefix-cache hits.
@@ -210,7 +217,7 @@ def get_num_blocks_to_allocate(
             return max(num_required_blocks - num_req_blocks, 0)
 
         num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
-        num_local_computed_blocks = len(new_computed_blocks) + num_req_blocks
+        num_local_computed_blocks = len(new_computed_blocks) + num_req_block_slots
         # Number of whole blocks that are skipped by the attention window.
         # If nothing is skipped, this is 0.
         num_skipped_blocks = num_skipped_tokens // self.block_size
@@ -223,9 +230,11 @@ def get_num_blocks_to_allocate(
         )
 
         # Among the `new_computed_blocks`, the first `num_skipped_blocks` worth
-        # of blocks are skipped; `num_req_blocks` of those may already be in
+        # of blocks are skipped; `num_req_block_slots` of those may already be in
         # `req_to_blocks`, so only skip the remainder from `new_computed_blocks`.
-        num_skipped_new_computed_blocks = max(0, num_skipped_blocks - num_req_blocks)
+        num_skipped_new_computed_blocks = max(
+            0, num_skipped_blocks - num_req_block_slots
+        )
 
         # If a computed block is an eviction candidate (in the free queue and
         # ref_cnt == 0), it will be removed from the free queue when touched by
@@ -368,9 +377,14 @@ def allocate_new_blocks(
             cow_blocks.append(cow_block)
 
         req_blocks = self.req_to_blocks[request_id]
-        num_required_blocks = cdiv(num_tokens, self.block_size)
-        num_new_blocks = num_required_blocks - len(req_blocks)
-        if num_new_blocks <= 0:
+        num_required_block_slots = cdiv(num_tokens, self.block_size)
+        num_required_blocks = self._get_num_required_blocks(num_tokens)
+        num_new_blocks = max(
+            num_required_blocks - self._get_num_allocated_blocks(request_id), 0
+        )
+        num_null_blocks = num_required_block_slots - len(req_blocks) - num_new_blocks
+        req_blocks.extend([self._null_block] * num_null_blocks)
+        if num_new_blocks == 0:
             return cow_blocks
         else:
             new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
@@ -916,6 +930,19 @@ def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:
         # cached blocks.
         self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens
 
+    def _get_num_required_blocks(self, num_tokens: int) -> int:
+        num_window_blocks = self.kv_cache_spec.max_admission_blocks_per_request(
+            max_in_flight_tokens=0,
+            max_model_len=num_tokens,
+        )
+        return min(super()._get_num_required_blocks(num_tokens), num_window_blocks)
+
+    def _get_num_allocated_blocks(self, request_id: str) -> int:
+        return sum(
+            block is not self._null_block
+            for block in self.req_to_blocks.get(request_id, ())
+        )
+
     @classmethod
     def _contiguous_blocks_for_hit(
         cls, window_size: int, block_size: int, use_eagle: bool

vllm/v1/kv_cache_interface.py

@@ -718,22 +718,16 @@ def max_admission_blocks_per_request(
     ) -> int:
         """Per-request admission cap, in blocks.
 
-        Single source of truth for both startup pool sizing
-        (`max_memory_usage_bytes`) and the runtime admission gate. Per-request
-        real-held blocks plateau at this bound because
-        `SlidingWindowManager.remove_skipped_blocks` runs from `allocate_slots`
-        before each chunk's `get_num_blocks_to_allocate`.
-
-        `max_in_flight_tokens` is the max tokens scheduled but not yet settled
-        (one batch per concurrent step); see `VllmConfig.max_in_flight_tokens`.
+        Single source of truth for startup pool sizing, the runtime admission
+        gate, and execution allocation. Logical positions outside this bound
+        are represented by null blocks in the request block table.
         """
-        # During chunked prefill, we hold KV for the last `sliding_window-1`
-        # computed tokens plus the in-flight tokens (frees happen on the
-        # processed-token basis); never more than `max_model_len`. An additional
-        # `extra_retained_tokens` trailing tokens are kept alive below the
-        # window for multi-module spec decoding, and must be accounted here too.
+        del max_in_flight_tokens
+        # Keep the last `sliding_window-1` tokens plus the trailing extension
+        # used by multi-module speculative decoding; never more than
+        # `max_model_len`.
         num_tokens = min(
-            self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens,
+            self.sliding_window - 1 + self.extra_retained_tokens,
             max_model_len,
         )
         # +1 because the sliding window may not start from the beginning of