From af40d043cdee085c572f41fcd4e2c435098e1491 Mon Sep 17 00:00:00 2001 From: Benedikt Falk Date: Tue, 25 Aug 2026 05:51:06 -0700 Subject: [PATCH 1/3] [Core] KV cache: block-range eviction for live requests Adds a modality-agnostic primitive to free a block-aligned token range of a request that stays alive, for long-running streaming-input sessions (#28973) that must bound per-session KV: - SingleTypeKVCacheManager.evict_blocks_for_request(request_id, block_ids): frees the given blocks from a live request's block table, removes them from the prefix-cache index so their content can never be served to another request, and keeps num_cached_block consistent. - KVCacheCoordinator.evict_blocks_for_request: fan-out per group. - KVCacheManager.evict_token_range_for_request(request_id, start, end) -> (aligned_start, aligned_end, num_freed): rounds the range INWARD to whole blocks and returns the range actually freed, so the caller compacts its own per-token arrays by exactly the same span and the block table and token bookkeeping cannot go out of sync. Guards the hybrid (multi-group) case with an assert rather than fanning one token range to groups of different block sizes. No behavior change for requests that do not call it. Consumer follows in the bounded-session retention PR (RFC #51948); the primitive is also usable for StreamingLLM-style sink+window text sessions or pruning stale tool-call context from a live session. Co-Authored-By: Claude Fable 5 Signed-off-by: Benedikt Falk --- tests/v1/core/test_kv_evict_for_request.py | 235 +++++++++++++++++++ vllm/v1/core/kv_cache_coordinator.py | 10 + vllm/v1/core/kv_cache_manager.py | 49 ++++ vllm/v1/core/single_type_kv_cache_manager.py | 55 +++++ 4 files changed, 349 insertions(+) create mode 100644 tests/v1/core/test_kv_evict_for_request.py diff --git a/tests/v1/core/test_kv_evict_for_request.py b/tests/v1/core/test_kv_evict_for_request.py new file mode 100644 index 000000000000..b9375b8c4f8d --- /dev/null +++ b/tests/v1/core/test_kv_evict_for_request.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for `KVCacheCoordinator.evict_blocks_for_request`. + +The primitive frees specific physical blocks from a still-alive request, +compacting the request's block table on the scheduler side. It is the +block-id-level free that `KVCacheManager.evict_token_range_for_request` +builds on. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch + +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig: + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _make_request( + request_id: str, + prompt_token_ids: list[int], + block_size: int, + hash_fn: Callable = sha256, +) -> Request: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + return Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=sp, + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, hash_fn), + ) + + +def _allocated_block_ids(manager: KVCacheManager, request_id: str) -> list[int]: + """First-group block IDs allocated for a request, in logical order, + excluding null_block sentinels.""" + blocks = manager.coordinator.single_type_managers[0].req_to_blocks.get( + request_id, [] + ) + return [b.block_id for b in blocks if not b.is_null] + + +def test_evict_blocks_compacts_block_list_no_caching(): + """Caller frees a subset; surviving blocks shift down. Freed blocks + return to the free pool (their ref_cnt drops to 0).""" + block_size = 16 + manager = KVCacheManager( + _make_kv_cache_config(block_size, num_blocks=8), + max_model_len=1024, + scheduler_block_size=block_size, + enable_caching=False, + hash_block_size=block_size, + ) + + # 4 full blocks worth of tokens (64 tokens). + req = _make_request("r0", [1] * 64, block_size) + computed, num_computed, _ = manager.get_computed_blocks(req) + assert num_computed == 0 + blocks_owned = manager.allocate_slots(req, 64, 0, computed) + assert blocks_owned is not None + + allocated = _allocated_block_ids(manager, "r0") + assert len(allocated) == 4, f"expected 4 allocated, got {allocated}" + + # Free the middle two. + to_evict = [allocated[1], allocated[2]] + free_q = manager.block_pool.free_block_queue + free_before = free_q.num_free_blocks + + n_freed = manager.coordinator.evict_blocks_for_request("r0", to_evict) + assert n_freed == 2 + + surviving = _allocated_block_ids(manager, "r0") + assert surviving == [allocated[0], allocated[3]], ( + f"block list should be compacted to [{allocated[0]}, {allocated[3]}], " + f"got {surviving}" + ) + + # Freed blocks must be back in the free pool. With caching disabled + # each block's ref_cnt drops to 0 immediately. + assert free_q.num_free_blocks == free_before + 2 + + # Request must remain alive. + manager.free(req) + + +def test_evict_blocks_idempotent_on_unknown_ids(): + """Block IDs not in the request's table are silently skipped.""" + block_size = 16 + manager = KVCacheManager( + _make_kv_cache_config(block_size, num_blocks=8), + max_model_len=1024, + scheduler_block_size=block_size, + enable_caching=False, + hash_block_size=block_size, + ) + req = _make_request("r1", [1] * 32, block_size) + computed, _, _ = manager.get_computed_blocks(req) + manager.allocate_slots(req, 32, 0, computed) + allocated = _allocated_block_ids(manager, "r1") + + # Block ID 999 is not in this request's table. + n_freed = manager.coordinator.evict_blocks_for_request("r1", [999, 999, 999]) + assert n_freed == 0 + assert _allocated_block_ids(manager, "r1") == allocated + manager.free(req) + + +def test_evict_blocks_empty_inputs_are_noop(): + block_size = 16 + manager = KVCacheManager( + _make_kv_cache_config(block_size, num_blocks=8), + max_model_len=1024, + scheduler_block_size=block_size, + enable_caching=False, + hash_block_size=block_size, + ) + req = _make_request("r2", [1] * 32, block_size) + manager.allocate_slots(req, 32, 0, manager.get_computed_blocks(req)[0]) + allocated = _allocated_block_ids(manager, "r2") + + assert manager.coordinator.evict_blocks_for_request("r2", []) == 0 + missing = manager.coordinator.evict_blocks_for_request("nonexistent-req", [1, 2, 3]) + assert missing == 0 + assert _allocated_block_ids(manager, "r2") == allocated + manager.free(req) + + +def test_evict_blocks_with_caching_removes_from_prefix_cache(): + """When caching is on, evicted blocks must be removed from the + prefix-cache hash map so future requests don't hit stale K/V.""" + block_size = 16 + manager = KVCacheManager( + _make_kv_cache_config(block_size, num_blocks=8), + max_model_len=1024, + scheduler_block_size=block_size, + enable_caching=True, + hash_block_size=block_size, + ) + + # Allocate 3 full blocks of identical tokens — guaranteed hashable. + req = _make_request("r3", [1] * 48, block_size) + computed, _, _ = manager.get_computed_blocks(req) + manager.allocate_slots(req, 48, 0, computed) + allocated = _allocated_block_ids(manager, "r3") + assert len(allocated) == 3 + + # Free the middle block. After eviction it must be neither in + # `cached_block_hash_to_block` nor have a block_hash. + middle_id = allocated[1] + pre_hash_count = len(manager.block_pool.cached_block_hash_to_block) + + n_freed = manager.coordinator.evict_blocks_for_request("r3", [middle_id]) + assert n_freed == 1 + assert manager.block_pool.blocks[middle_id].block_hash is None + assert len(manager.block_pool.cached_block_hash_to_block) < pre_hash_count, ( + "evicted block should have left the prefix-cache hash map" + ) + + # Surviving list is compacted. + assert _allocated_block_ids(manager, "r3") == [allocated[0], allocated[2]] + manager.free(req) + + +def test_evict_blocks_then_continue_allocating(): + """After eviction the request must still be allocatable (still alive).""" + block_size = 16 + manager = KVCacheManager( + _make_kv_cache_config(block_size, num_blocks=8), + max_model_len=1024, + scheduler_block_size=block_size, + enable_caching=False, + hash_block_size=block_size, + ) + req = _make_request("r4", [1] * 48, block_size) + manager.allocate_slots(req, 48, 0, manager.get_computed_blocks(req)[0]) + allocated_before = _allocated_block_ids(manager, "r4") + assert len(allocated_before) == 3 + + # Drop the oldest. + manager.coordinator.evict_blocks_for_request("r4", [allocated_before[0]]) + after_evict = _allocated_block_ids(manager, "r4") + assert len(after_evict) == 2 + + # Append more tokens; allocator should hand us a fresh block on top + # of the surviving two. + req.num_computed_tokens = 32 # 2 surviving blocks worth. + req._all_token_ids.extend([2] * 16) + req.prompt_token_ids = list(req._all_token_ids) + manager.allocate_slots(req, 16, 0) + after_grow = _allocated_block_ids(manager, "r4") + assert len(after_grow) == 3, f"expected 3 blocks after re-growth, got {after_grow}" + manager.free(req) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 767ea0033483..026e16b0f42f 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -351,6 +351,16 @@ def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: blocks.extend(manager.pop_blocks_for_free(request_id)) return blocks + def evict_blocks_for_request(self, request_id: str, block_ids: list[int]) -> int: + """Free a subset of the request's blocks; request stays alive. + + Used by streaming-video intra-session eviction. + """ + total_freed = 0 + for manager in self.single_type_managers: + total_freed += manager.evict_blocks_for_request(request_id, block_ids) + return total_freed + def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]: """ Get the number of common prefix blocks for all requests with allocated diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index d1af91b65993..f3c94e9eb14d 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -627,6 +627,55 @@ def evict_blocks(self, block_ids: set[int]) -> None: """ self.block_pool.evict_blocks(block_ids) + def evict_token_range_for_request( + self, request_id: str, token_start: int, token_end: int + ) -> tuple[int, int, int]: + """Convenience wrapper: evict the block-aligned span overlapping + the token range `[token_start, token_end)` on the first kv-cache + group, then free those blocks via the coordinator. + + Uses **inward** block alignment: only whole blocks fully within + `[token_start, token_end)` are freed. Block-aligned edges are + important so the caller can slice its per-token arrays + (`_all_token_ids`, and per-request position arrays if the model + tracks them explicitly) by the SAME range that was freed + from the block table — otherwise the two go out of sync and the + request's prompt grows past the worker's + `max_model_len`-sized buffers. + + Returns `(aligned_start, aligned_end, num_blocks_freed)` where + `aligned_start` and `aligned_end` describe the token range that + was actually evicted (block-aligned, possibly narrower than + `[token_start, token_end)`). Callers MUST use this returned + range — not the original — when compacting their own per-token + arrays in lockstep with the KV cache. + """ + assert len(self.coordinator.single_type_managers) == 1, ( + "evict_token_range_for_request derives block ids from kv-cache " + "group 0 only; fanning them to multiple groups is unsound. " + "Callers must ensure a single-group (non-hybrid) config." + ) + if token_end <= token_start: + return token_start, token_start, 0 + first = self.coordinator.single_type_managers[0] + block_size = first.block_size + block_idx_start = (token_start + block_size - 1) // block_size + block_idx_end = token_end // block_size + if block_idx_end <= block_idx_start: + return token_start, token_start, 0 + req_blocks = first.req_to_blocks.get(request_id) + if not req_blocks: + return token_start, token_start, 0 + aligned_start = block_idx_start * block_size + aligned_end = block_idx_end * block_size + block_ids = [ + req_blocks[i].block_id + for i in range(block_idx_start, min(block_idx_end, len(req_blocks))) + if not req_blocks[i].is_null + ] + num_freed = self.coordinator.evict_blocks_for_request(request_id, block_ids) + return aligned_start, aligned_end, num_freed + def reset_prefix_cache(self) -> bool: """Reset prefix cache. This function may be used in RLHF flows to invalidate prefix caching after the weights are updated, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index ceb56a59875f..2f0a45705d9c 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -523,6 +523,61 @@ def free(self, request_id: str) -> None: # Free blocks in reverse order so that the tail blocks are freed first. self.block_pool.free_blocks(reversed(self.pop_blocks_for_free(request_id))) + def evict_blocks_for_request(self, request_id: str, block_ids: list[int]) -> int: + """Free a subset of the request's blocks while the request stays alive. + + Used for intra-session eviction in long-running streaming + sessions (StreamingLLM-style sink+window for text, sliding-window + video, or any policy that retires a token range of a live + request). The primitive is modality-agnostic: it operates on + blocks and token ranges only. + + `num_cached_block` is decremented by the count of evicted blocks + that had a hash (i.e., were prefix-cacheable). + + Args: + request_id: The request to evict from. + block_ids: Physical block IDs to free. Need not be in any + particular order; the function looks them up in + `req_to_blocks`. Caller is responsible for selecting a + contiguous range that corresponds to a session segment. + + Returns: + Number of blocks whose memory was actually reclaimed, i.e. + blocks that reached `ref_cnt == 0` and were returned to the + free queue. Blocks still shared with another request + (`ref_cnt > 0` after the dereference) are dropped from this + request's table but NOT counted, since their memory is not + reclaimed. + """ + req_blocks = self.req_to_blocks.get(request_id) + if not req_blocks or not block_ids: + return 0 + block_id_set = set(block_ids) + surviving: list[KVCacheBlock] = [] + blocks_to_free: list[KVCacheBlock] = [] + for block in req_blocks: + if block.block_id in block_id_set and not block.is_null: + blocks_to_free.append(block) + else: + surviving.append(block) + if not blocks_to_free: + return 0 + self.req_to_blocks[request_id] = surviving + # Free in reverse order so the tail blocks (most recently allocated) + # are evicted first — matches the LRU convention in `free`. + self.block_pool.free_blocks(reversed(blocks_to_free)) + if self.enable_caching: + # Count cached blocks BEFORE evict_blocks clears their hashes. + cached_freed = sum(1 for b in blocks_to_free if b.block_hash is not None) + self.block_pool.evict_blocks({b.block_id for b in blocks_to_free}) + if cached_freed and request_id in self.num_cached_block: + self.num_cached_block[request_id] = max( + 0, self.num_cached_block[request_id] - cached_freed + ) + + return sum(1 for b in blocks_to_free if b.ref_cnt == 0) + @abstractmethod def get_num_common_prefix_blocks(self, running_request_id: str) -> int: """ From a489f9e981ba17ae4e72bf65d6e781a84733d4af Mon Sep 17 00:00:00 2001 From: Benedikt Falk Date: Tue, 25 Aug 2026 05:44:38 -0700 Subject: [PATCH 2/3] [Bugfix] Harden streaming-input session lifecycle Four independent fixes for resumable (streaming-input, #28973) sessions, found running multi-hour video-streaming sessions: 1. RequestOutputCollector merges DELTA outputs when the producer gets ahead of the consumer, and RequestOutput.add() overwrites finish_reason. Per-chunk boundaries are conveyed as transient finish_reason markers on non-finished outputs, so a merge across a boundary silently erases the marker and the consumer loses a chunk boundary. Replace the single-slot collector with a deque that never merges into an entry that is finished or carries a finish_reason. 2. process_outputs() treats ANY finish_reason on a streaming-input request as a chunk boundary. An engine-initiated ABORT or ERROR is terminal for the session; masking it wedges the consumer and desyncs the frontend and scheduler chunk queues. Gate boundary handling on the finish reason and route an abort back to the engine core. 3. Abort races: cancel AND await the input-stream task before aborting (close() only cancels, leaving a window where an in-flight chunk ADD is sent after the abort), and mark first chunks on the wire (EngineCoreRequest.first_chunk) so the scheduler rejects non-first chunks for unknown ids instead of resurrecting the finished session as an unabortable request that pins its KV blocks. 4. Zero-chunk streams put STREAM_FINISHED instead of running the placeholder final request as a real generation; sticky DP routing keeps all chunks of a session on the engine that holds it; sessions salt their prefix-cache hash chain (chunk folds mutate block content under unchanged hashes). Co-Authored-By: Claude Fable 5 Signed-off-by: Benedikt Falk --- tests/v1/engine/test_output_processor.py | 131 +++++++++ .../test_scheduler_streaming.py | 5 + .../streaming_input/test_session_hardening.py | 263 ++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 10 + vllm/v1/engine/__init__.py | 3 + vllm/v1/engine/async_llm.py | 40 ++- vllm/v1/engine/core_client.py | 3 + vllm/v1/engine/output_processor.py | 58 ++-- vllm/v1/request.py | 3 + 9 files changed, 492 insertions(+), 24 deletions(-) create mode 100644 tests/v1/streaming_input/test_session_hardening.py diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 51dbb9c98952..91c4e2691620 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -22,6 +22,7 @@ from vllm.v1.engine import ( EngineCoreEvent, EngineCoreEventType, + EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, FinishReason, @@ -1383,6 +1384,136 @@ async def test_cumulative_output_collector_n(): assert third[0].text == "c" +def _make_boundary_output( + token_id: int, finish_reason: str | None, finished: bool = False +) -> RequestOutput: + return RequestOutput( + request_id="my-request-id", + prompt=None, + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="a", + token_ids=[token_id], + cumulative_logprob=None, + logprobs=None, + finish_reason=finish_reason, + ) + ], + finished=finished, + ) + + +@pytest.mark.asyncio +async def test_collector_does_not_merge_across_finish_reason_boundary(): + """A pending chunk-boundary marker (finish_reason on a non-finished + output) must not be erased by merging the next chunk's output into it.""" + collector = RequestOutputCollector( + RequestOutputKind.DELTA, request_id="my-request-id-int" + ) + + collector.put(_make_boundary_output(0, "stop")) + collector.put(_make_boundary_output(1, None)) + collector.put(_make_boundary_output(2, None)) + + boundary = await collector.get() + assert boundary.outputs[0].finish_reason == "stop" + assert boundary.outputs[0].token_ids == [0] + + following = await collector.get() + assert following.outputs[0].finish_reason is None + assert following.outputs[0].token_ids == [1, 2] + + assert collector.output is None + assert not collector.ready.is_set() + + +def _make_streaming_input_request( + request_id: str, prompt_token_ids: list[int] +) -> EngineCoreRequest: + """One chunk of a streaming-input (resumable) request.""" + return EngineCoreRequest( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + mm_features=None, + sampling_params=SamplingParams( + detokenize=False, output_kind=RequestOutputKind.DELTA + ), + pooling_params=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + resumable=True, + external_req_id=f"{request_id}-ext", + ) + + +def _add_streaming_request_with_queued_chunk( + processor: OutputProcessor, request_id: str +) -> RequestOutputCollector: + queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id=request_id) + processor.add_request( + _make_streaming_input_request(request_id, [1, 2, 3]), None, queue=queue + ) + processor.add_request(_make_streaming_input_request(request_id, [7, 8]), None) + return queue + + +def test_streaming_input_stop_is_chunk_boundary(): + """STOP is a per-chunk boundary: finished is forced False, the queued + chunk is applied, and the request state stays alive.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[4], + finish_reason=FinishReason.STOP, + ) + ] + ) + + assert processed.reqs_to_abort == [] + assert processor.get_num_unfinished_requests() == 1 + output = queue.get_nowait() + assert output is not None + assert output.finished is False + assert output.outputs[0].finish_reason == "stop" + req_state = processor.request_states["stream-req"] + assert req_state.prompt_token_ids is not None + assert req_state.prompt_token_ids[-2:] == [7, 8] + + +@pytest.mark.parametrize("finish_reason", [FinishReason.ABORT, FinishReason.ERROR]) +def test_streaming_input_abort_error_is_terminal(finish_reason: FinishReason): + """ABORT/ERROR is terminal, not a chunk boundary: finished stays True, + the state is freed, and an abort is routed back to the engine core.""" + processor = OutputProcessor(None, log_stats=False) + queue = _add_streaming_request_with_queued_chunk(processor, "stream-req") + + processed = processor.process_outputs( + [ + EngineCoreOutput( + request_id="stream-req", + new_token_ids=[], + finish_reason=finish_reason, + ) + ] + ) + + assert processor.get_num_unfinished_requests() == 0 + assert processed.reqs_to_abort == ["stream-req"] + output = queue.get_nowait() + assert output is not None + assert output.finished is True + assert output.outputs[0].finish_reason == str(finish_reason) + + @pytest.mark.parametrize("runner", ["generate", "pooling"]) @pytest.mark.parametrize("abort_by", ["internal", "external"]) def test_abort_requests(runner: str, abort_by: str, dummy_test_vectors): diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index 822b4e49da34..9f348948d891 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -35,6 +35,7 @@ def __init__( prompt_token_ids=None, mm_features: list[MultiModalFeatureSpec] | None = None, max_tokens: int | None = 16, + first_chunk=True, ): super().__init__( request_id=request_id, @@ -45,6 +46,10 @@ def __init__( pooling_params=None, mm_features=mm_features, resumable=resumable, + # The frontend stamps first_chunk on the first chunk of every + # session (async_llm.handle_inputs); these tests construct + # session-opening requests directly, so stamp it here. + first_chunk=first_chunk, ) diff --git a/tests/v1/streaming_input/test_session_hardening.py b/tests/v1/streaming_input/test_session_hardening.py new file mode 100644 index 000000000000..c58a27fb9147 --- /dev/null +++ b/tests/v1/streaming_input/test_session_hardening.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Regression tests for streaming-input session lifecycle hardening: +post-abort chunk races, unbounded resume, zero-chunk streams, and DP +routing.""" + +from __future__ import annotations + +import asyncio +from unittest import mock +from unittest.mock import AsyncMock, MagicMock + +import pytest +import torch + +from vllm.config import DeviceConfig, VllmConfig +from vllm.engine.protocol import StreamingInput +from vllm.outputs import STREAM_FINISHED +from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.engine.core_client import DPLBAsyncMPClient +from vllm.v1.engine.output_processor import RequestOutputCollector +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm.v1.structured_output import StructuredOutputManager + +pytestmark = pytest.mark.cpu_test + +STOP_TOKEN = 128001 + + +def _create_scheduler() -> Scheduler: + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + vllm_config.model_config = MagicMock() + vllm_config.model_config.skip_tokenizer_init = True + vllm_config.model_config.is_multimodal_model = False + vllm_config.model_config.is_encoder_decoder = False + vllm_config.model_config.max_model_len = 1024 + vllm_config.model_config.enable_return_routed_experts = False + vllm_config.cache_config = MagicMock() + vllm_config.cache_config.num_gpu_blocks = 1000 + vllm_config.cache_config.enable_prefix_caching = False + kv_cache_config = KVCacheConfig( + num_blocks=1000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ) + ], + ) + return Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + log_stats=True, + structured_output_manager=StructuredOutputManager(vllm_config), + block_size=16, + hash_block_size=16, + ) + + +def _make_chunk( + request_id: str, + tokens: list[int], + max_tokens: int = 16, + first_chunk: bool = True, +) -> Request: + return Request( + request_id=request_id, + prompt_token_ids=tokens, + sampling_params=SamplingParams( + stop_token_ids=[STOP_TOKEN], max_tokens=max_tokens + ), + pooling_params=None, + resumable=True, + first_chunk=first_chunk, + ) + + +def _mro(req_id: str, token: int) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[req_id], + req_id_to_index={req_id: 0}, + sampled_token_ids=[[token]], + logprobs=None, + prompt_logprobs_dict={req_id: None}, + pooler_output=[], + ) + + +def test_non_first_chunk_for_unknown_session_is_dropped(): + """Only a session's first chunk may create a session: a non-first chunk + for an id the scheduler does not know belongs to a finished (or foreign) + session, and admitting it would resurrect an unabortable request.""" + scheduler = _create_scheduler() + + # Cold scheduler: a mid-session chunk with no known session is dropped. + scheduler.add_request(_make_chunk("sess-x", [7, 8], first_chunk=False)) + assert "sess-x" not in scheduler.requests + assert len(scheduler.waiting) == 0 + assert scheduler.num_waiting_for_streaming_input == 0 + + # The race this guards: session finishes, its late chunk must not revive. + scheduler.add_request(_make_chunk("sess-t", [1, 2, 3])) + scheduler.finish_requests("sess-t", RequestStatus.FINISHED_ABORTED) + scheduler.add_request(_make_chunk("sess-t", [7, 8], first_chunk=False)) + assert "sess-t" not in scheduler.requests + + # First chunks still open sessions. + scheduler.add_request(_make_chunk("sess-b", [4, 5, 6])) + assert "sess-b" in scheduler.requests + + +@pytest.mark.asyncio +async def test_stop_input_stream_cancels_and_awaits(): + """No chunk send may complete after _stop_input_stream returns.""" + queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id="r1") + sent = [] + + async def slow_send(): + await asyncio.sleep(30) + sent.append("chunk") + + task = asyncio.create_task(slow_send()) + queue._input_stream_task = task + await asyncio.sleep(0) + + await AsyncLLM._stop_input_stream(queue) + + assert task.cancelled() + assert queue._input_stream_task is None + assert sent == [] + + +@pytest.mark.asyncio +async def test_handle_inputs_stamps_first_chunk(): + """The frontend marks exactly the first chunk of a session.""" + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm._add_request = AsyncMock() + llm.model_config = MagicMock() + + reqs = [MagicMock(prompt_embeds=None) for _ in range(3)] + reqs[0].request_id = "r1-int" + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.side_effect = reqs + llm.input_processor.assign_request_id = MagicMock() + + async def two_chunks(): + yield StreamingInput(prompt="frame-1") + yield StreamingInput(prompt="frame-2") + + with mock.patch( + "vllm.v1.engine.async_llm.extract_prompt_components", + return_value=("", None, None), + ): + await llm._add_streaming_input_request( + "r1", two_chunks(), SamplingParams(max_tokens=4) + ) + for _ in range(6): + await asyncio.sleep(0) + + assert reqs[1].first_chunk is True + assert reqs[2].first_chunk is False + + +@pytest.mark.asyncio +async def test_input_stream_task_cleared_only_after_final_send(): + """The task pointer must stay set until the final request is sent, so a + concurrent abort can still cancel-and-await the task (no send may ever + happen after an abort).""" + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm.model_config = MagicMock() + + reqs = [MagicMock(prompt_embeds=None) for _ in range(2)] + reqs[0].request_id = "r1-int" + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.side_effect = reqs + llm.input_processor.assign_request_id = MagicMock() + + holder = {} + pointer_set_during_send = [] + + async def capture_add_request(*args, **kwargs): + pointer_set_during_send.append(holder["queue"]._input_stream_task is not None) + + llm._add_request = AsyncMock(side_effect=capture_add_request) + + async def one_chunk(): + yield StreamingInput(prompt="frame-1") + + with mock.patch( + "vllm.v1.engine.async_llm.extract_prompt_components", + return_value=("", None, None), + ): + holder["queue"] = await llm._add_streaming_input_request( + "r1", one_chunk(), SamplingParams(max_tokens=4) + ) + for _ in range(6): + await asyncio.sleep(0) + + # Two sends (the chunk, then the final request); the pointer must have + # been set for BOTH, and cleared only once everything was sent. + assert pointer_set_during_send == [True, True] + assert holder["queue"]._input_stream_task is None + + +@pytest.mark.asyncio +async def test_zero_chunk_stream_emits_finished_sentinel(): + """An input stream that closes without chunks must unblock the consumer + instead of submitting the placeholder final request as a generation.""" + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm._add_request = AsyncMock() + llm.model_config = MagicMock() + + final_req = MagicMock() + final_req.request_id = "r1-int" + final_req.prompt_embeds = None + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.return_value = final_req + llm.input_processor.assign_request_id = MagicMock() + + async def empty_stream(): + return + yield # pragma: no cover + + queue = await llm._add_streaming_input_request( + "r1", empty_stream(), SamplingParams(max_tokens=4) + ) + for _ in range(4): + await asyncio.sleep(0) + + assert queue.get_nowait() is STREAM_FINISHED + llm._add_request.assert_not_awaited() + + +def test_dp_sticky_routing_for_in_flight_request(): + """Chunks of an in-flight session route to the engine holding it.""" + client = DPLBAsyncMPClient.__new__(DPLBAsyncMPClient) + engine = object() + client.reqs_in_flight = {"sess-1": engine} + + request = MagicMock() + request.request_id = "sess-1" + assert client.get_core_engine_for_request(request) is engine diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c373d5cf5c87..352fc97eb0ef 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2361,6 +2361,16 @@ def add_request(self, request: Request) -> None: # Streaming-input session finished. self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) else: + if request.resumable and not request.first_chunk: + # A non-first chunk for an unknown id: its session is gone. + # Admitting it would resurrect the session as an unabortable + # request that pins its KV blocks until restart. + logger.warning( + "Dropping streaming request %s: its session already " + "finished (stale post-abort chunk).", + request.request_id, + ) + return if request.resumable: request.streaming_queue = deque() self._enqueue_waiting_request(request) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 5ae9ee0cac83..bad4b71f955c 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -139,6 +139,9 @@ class EngineCoreRequest( trace_headers: Mapping[str, str] | None = None resumable: bool = False + # True only for the first chunk of a streaming-input (resumable) request; + # the scheduler rejects non-first chunks for unknown ids as stale. + first_chunk: bool = False # The user-provided request ID. This field is set internally, # copied from the provided request_id that's originally assigned diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index dd5cb2c26287..876956a0bd59 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -484,6 +484,8 @@ async def _add_streaming_input_request( async def handle_inputs(): cancelled = False + errored = False + any_added = False try: async for input_chunk in input_stream: sp = input_chunk.sampling_params @@ -500,6 +502,10 @@ async def handle_inputs(): **inputs, # type: ignore[arg-type] ) req.external_req_id = request_id + req.first_chunk = not any_added + # Chunk folds mutate block content under unchanged + # prefix-cache hashes; session blocks are never shareable. + req.cache_salt = internal_req_id if req.prompt_embeds is not None: raise VLLMValidationError( "prompt_embeds not supported for streaming inputs" @@ -508,18 +514,30 @@ async def handle_inputs(): self.model_config, input_chunk.prompt ) await self._add_request(req, prompt_text, None, 0, queue) + any_added = True except (asyncio.CancelledError, GeneratorExit): cancelled = True except Exception as error: # Wrap in InputStreamError so generate() can propagate it # without wrapping in EngineGenerateError. queue.put(InputStreamError(error)) + errored = True finally: - queue._input_stream_task = None - if not cancelled: - # Send empty final request to indicate that inputs have - # finished. Don't send if cancelled (session was aborted). - await self._add_request(final_req, None, None, 0, queue) + try: + if not cancelled: + if any_added: + # Send empty final request to indicate that inputs + # have finished. Don't send if cancelled. + await self._add_request(final_req, None, None, 0, queue) + elif not errored: + # Zero-chunk stream: nothing was submitted + # engine-side; just unblock the generate() consumer. + queue.put(STREAM_FINISHED) + finally: + # Cleared only after the final send: until then, + # _stop_input_stream must be able to cancel-and-await + # this task so nothing is sent after an abort. + queue._input_stream_task = None # Ensure output handler is running. self._run_output_handler() @@ -543,6 +561,15 @@ def _validate_streaming_input_sampling_params( "or with stop strings." ) + @staticmethod + async def _stop_input_stream(q: RequestOutputCollector) -> None: + """Cancel and await the input-stream task, so no chunk ADD can be + sent after the request is aborted.""" + if (task := q._input_stream_task) is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + q._input_stream_task = None + # TODO: we should support multiple prompts in one call, as you # can do with LLM.generate. So that for multi-prompt completion # requests we don't need to send multiple messages to core proc, @@ -619,6 +646,7 @@ async def generate( # we abort the request if we end up here. except (asyncio.CancelledError, GeneratorExit): if q is not None: + await self._stop_input_stream(q) await self.abort(q.request_id, internal=True) if self.log_requests: logger.info("Request %s aborted.", request_id) @@ -639,6 +667,7 @@ async def generate( # Error from input stream generator - propagate directly. except InputStreamError as e: if q is not None: + await self._stop_input_stream(q) await self.abort(q.request_id, internal=True) if self.log_requests: logger.info("Request %s failed (input error): %s.", request_id, e) @@ -647,6 +676,7 @@ async def generate( # Unexpected error in the generate() task (possibly recoverable). except Exception as e: if q is not None: + await self._stop_input_stream(q) await self.abort(q.request_id, internal=True) if self.log_requests: try: diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 2bf7cf03c188..6ae6b30aec43 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -1470,6 +1470,9 @@ def __init__( ) // client_count def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity: + # Route all chunks of an in-flight (streaming) request to its engine. + if (engine := self.reqs_in_flight.get(request.request_id)) is not None: + return engine # Engines are in rank order. if (eng_index := request.data_parallel_rank) is None and ( eng_index := get_late_interaction_engine_index( diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 055e3af574c1..ca53347f23ee 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -57,42 +57,58 @@ class RequestOutputCollector: def __init__(self, output_kind: RequestOutputKind, request_id: str): self.aggregate = output_kind == RequestOutputKind.DELTA self.request_id = request_id - self.output: RequestOutput | PoolingRequestOutput | Exception | None = None + self.outputs: deque[RequestOutput | PoolingRequestOutput | Exception] = deque() self.ready = asyncio.Event() self._input_stream_task: asyncio.Task | None = None + @property + def output(self) -> RequestOutput | PoolingRequestOutput | Exception | None: + """Next output to be consumed (oldest pending), or None if drained.""" + return self.outputs[0] if self.outputs else None + def put(self, output: RequestOutput | PoolingRequestOutput | Exception) -> None: """Non-blocking put operation.""" - if self.output is None or isinstance(output, Exception): - self.output = output + if isinstance(output, Exception): + self.outputs.clear() + self.outputs.append(output) self.ready.set() - elif isinstance(self.output, RequestOutput) and isinstance( - output, RequestOutput + return + tail = self.outputs[-1] if self.outputs else None + if ( + isinstance(tail, RequestOutput) + and isinstance(output, RequestOutput) + and not tail.finished + and all(c.finish_reason is None for c in tail.outputs) ): # This ensures that request outputs with different request indexes # (if n > 1) do not override each other. - self.output.add(output, aggregate=self.aggregate) - elif isinstance(self.output, PoolingRequestOutput) and isinstance( + tail.add(output, aggregate=self.aggregate) + elif isinstance(tail, PoolingRequestOutput) and isinstance( output, PoolingRequestOutput ): - self.output = output + self.outputs[-1] = output + else: + self.outputs.append(output) + self.ready.set() async def get(self) -> RequestOutput | PoolingRequestOutput: """Get operation blocks on put event.""" - while (output := self.output) is None: + while not self.outputs: await self.ready.wait() - self.output = None - self.ready.clear() + output = self.outputs.popleft() + if not self.outputs: + self.ready.clear() if isinstance(output, Exception): raise output return output def get_nowait(self) -> RequestOutput | PoolingRequestOutput | None: """Non-blocking get operation.""" - output = self.output - if output is not None: - self.output = None + if not self.outputs: + return None + output = self.outputs.popleft() + if not self.outputs: self.ready.clear() if isinstance(output, Exception): raise output @@ -685,6 +701,10 @@ def process_outputs( # if required. req_state.logprobs_processor.update_from_output(engine_core_output) + streaming_chunk_boundary = req_state.streaming_input and ( + finish_reason not in (FinishReason.ABORT, FinishReason.ERROR) + ) + # 4) Create and handle RequestOutput objects. if request_output := req_state.make_request_output( new_token_ids, @@ -694,7 +714,7 @@ def process_outputs( kv_transfer_params, ec_transfer_params, ): - if req_state.streaming_input: + if streaming_chunk_boundary: request_output.finished = False if req_state.queue is not None: @@ -706,7 +726,7 @@ def process_outputs( # Free completed requests. if finish_reason is not None: - if req_state.streaming_input: + if streaming_chunk_boundary: if req_state.input_chunk_queue: update = req_state.input_chunk_queue.popleft() req_state.apply_streaming_update(update) @@ -714,9 +734,9 @@ def process_outputs( req_state.input_chunk_queue = None else: self._finish_request(req_state) - if not engine_core_output.finished: - # If req not finished in EngineCore, but Detokenizer - # detected stop string, abort needed in EngineCore. + if req_state.streaming_input or not engine_core_output.finished: + # EngineCore may still have the request live (stop + # string detected, streaming abort/error); abort it. reqs_to_abort.append(req_id) # Track per-request stats diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 8b453a09069e..4566165a8ee0 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -75,6 +75,7 @@ def __init__( block_hasher: Callable[["Request"], list["BlockHash"]] | None = None, resumable: bool = False, session_id: str | None = None, + first_chunk: bool = False, reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, abort_immediately: bool = False, @@ -227,6 +228,7 @@ def __init__( # Used for streaming self.resumable = resumable + self.first_chunk = first_chunk # None entry in the queue means finished. self.streaming_queue: deque[StreamingUpdate | None] | None = None @@ -257,6 +259,7 @@ def from_engine_core_request( block_hasher=block_hasher, resumable=request.resumable, session_id=request.session_id, + first_chunk=request.first_chunk, reasoning_ended=request.reasoning_ended, reasoning_parser_kwargs=request.reasoning_parser_kwargs, abort_immediately=request.abort_immediately, From 89bfabd975d9f006e4b2dfc2448846f346c2dbc4 Mon Sep 17 00:00:00 2001 From: Benedikt Falk Date: Tue, 25 Aug 2026 06:38:48 -0700 Subject: [PATCH 3/3] [Core][Draft] Bounded-memory streaming sessions: retention, in-place KV eviction, consolidation Engine core for RFC #51948: give a long-running streaming-input session (#28973) a constant memory footprint, for continuous video understanding and StreamingLLM-style sink+window text sessions. Mechanism (shared between text and multimodal sessions): - vllm/v1/streaming/: retention policy (per-session HistorySegment bookkeeping + sliding-window budgets), eviction (drop the oldest segments of a LIVE session in place via the KV range-free primitive, releasing encoder-cache entries and compacting all per-token arrays in lockstep by the block-aligned range actually freed), consolidation (self-preempting re-prefill of the surviving window before the model's trained position range is exceeded), chunk-relative position helpers. - Survivors keep their original positions; the position space develops gaps. Cached K stays rotationally consistent with no kernel changes. mRoPE models carry a per-token position tensor, extended per chunk from max_cached_position + 1 and sliced on eviction. 1D-RoPE (text) models need only a per-request scalar: RoPE position = index + cumulative evicted width, reset at re-prefill (the degenerate case of the same design, so LLM and VLM sessions share the mechanics end to end). - Scheduler: segment recording at chunk folds, over-budget eviction, re-prefill trigger + self-preemption (phantom-sample discard on idle triggers), fold-time cumulative context-length guard, grammar soft-stop for sessions, admission checks (single KV-cache group, no async scheduling / PP / DP for retention sessions). - Worker: position slice/extend/reset mirroring the scheduler, max_cached_position watermark reporting, persistent encoder cache for re-prefill (immediate physical eviction of dead entries), orphan mm-placeholder neutralization, streaming-aware encoder cache floor. Sessions that do not set streaming_retention (in SamplingParams.extra_args) take none of these paths. Draft caveats: 1D-RoPE position shift is applied to the shared positions buffer post-slot-mapping; backends consuming CommonAttentionMetadata.positions for masking are audited for the default FA path only. GPU E2E evidence to follow. 86 unit tests (tests/v1/streaming/). Existing suite failure sets are byte-identical to base on the same host. Co-Authored-By: Claude Fable 5 Signed-off-by: Benedikt Falk --- tests/v1/streaming/__init__.py | 2 + tests/v1/streaming/test_eviction.py | 1006 +++++++++++++++++ tests/v1/streaming/test_eviction_edges.py | 598 ++++++++++ tests/v1/streaming/test_mrope_streaming.py | 234 ++++ .../streaming/test_prefix_cache_eviction.py | 281 +++++ tests/v1/streaming/test_reprefill.py | 860 ++++++++++++++ .../streaming/test_streaming_regressions.py | 412 +++++++ .../test_streaming_session_eviction.py | 263 +++++ tests/v1/streaming/test_text_sessions.py | 165 +++ tests/v1/streaming/test_worker_slice.py | 208 ++++ vllm/multimodal/encoder_budget.py | 1 + vllm/v1/core/encoder_cache_manager.py | 57 + vllm/v1/core/kv_cache_utils.py | 17 +- vllm/v1/core/sched/output.py | 22 + vllm/v1/core/sched/scheduler.py | 401 ++++++- vllm/v1/outputs.py | 4 + vllm/v1/request.py | 66 ++ vllm/v1/streaming/__init__.py | 0 vllm/v1/streaming/eviction.py | 445 ++++++++ vllm/v1/streaming/mrope.py | 79 ++ vllm/v1/streaming/reprefill.py | 44 + vllm/v1/streaming/retention.py | 150 +++ vllm/v1/worker/gpu_input_batch.py | 7 + vllm/v1/worker/gpu_model_runner.py | 275 ++++- 24 files changed, 5573 insertions(+), 24 deletions(-) create mode 100644 tests/v1/streaming/__init__.py create mode 100644 tests/v1/streaming/test_eviction.py create mode 100644 tests/v1/streaming/test_eviction_edges.py create mode 100644 tests/v1/streaming/test_mrope_streaming.py create mode 100644 tests/v1/streaming/test_prefix_cache_eviction.py create mode 100644 tests/v1/streaming/test_reprefill.py create mode 100644 tests/v1/streaming/test_streaming_regressions.py create mode 100644 tests/v1/streaming/test_streaming_session_eviction.py create mode 100644 tests/v1/streaming/test_text_sessions.py create mode 100644 tests/v1/streaming/test_worker_slice.py create mode 100644 vllm/v1/streaming/__init__.py create mode 100644 vllm/v1/streaming/eviction.py create mode 100644 vllm/v1/streaming/mrope.py create mode 100644 vllm/v1/streaming/reprefill.py create mode 100644 vllm/v1/streaming/retention.py diff --git a/tests/v1/streaming/__init__.py b/tests/v1/streaming/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/tests/v1/streaming/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/streaming/test_eviction.py b/tests/v1/streaming/test_eviction.py new file mode 100644 index 000000000000..4a818068b0fc --- /dev/null +++ b/tests/v1/streaming/test_eviction.py @@ -0,0 +1,1006 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for `evict_segment` and `maybe_evict_old_segments` +operating on a real `KVCacheManager` +plus a synthetic request with pre-populated `session_history`. + +These tests exercise the full scheduler-side eviction flow without +running the actual scheduler loop. They validate: + + - `_all_token_ids` and `_mrope_positions` are sliced in lockstep. + - Subsequent `HistorySegment.token_range` fields shift down. + - `mm_features` offsets shift down past the evicted range; the + evicted segment's mm_feature is dropped. + - `num_prompt_tokens` / `num_computed_tokens` decrement correctly. + - The KV cache primitive runs (block IDs return to the pool). + - Pinned segments are NEVER evicted, even when over budget. + - `max_cached_position` is preserved (gaps are intentional). +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch + +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request +from vllm.v1.streaming.eviction import ( + evict_segment, + maybe_evict_old_segments, +) +from vllm.v1.streaming.retention import ( + HistorySegment, + StreamingRetentionParams, +) + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_managers(block_size: int = 16, num_blocks: int = 32): + kv = KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=1024, + enable_caching=False, + scheduler_block_size=block_size, + hash_block_size=block_size, + ) + enc = EncoderCacheManager(cache_size=4096) + return kv, enc + + +def _make_request( + request_id: str, + prompt_token_ids: list[int], + block_size: int = 16, + hash_fn: Callable = sha256, +) -> Request: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + return Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=sp, + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, hash_fn), + ) + + +def test_evict_segment_compacts_arrays_and_shifts_subsequent(): + """Three segments. Evict the middle one. Verify both arrays compact, + later segment's token_range shifts down, mm_features get shifted.""" + kv, enc = _make_managers(block_size=16, num_blocks=32) + req = _make_request("r-evict-mid", [9] * 64) + + # Pretend mRoPE positions were assigned linearly across the prompt. + req._mrope_positions = [(i, i, i) for i in range(64)] + req.max_cached_position = 63 + req.num_prompt_tokens = 64 + req.num_computed_tokens = 64 + + # Three segments (16, 32, 16 tokens each). Middle is a video. + seg0 = HistorySegment( + segment_type="system_prompt", token_range=(0, 16), pinned=True + ) + seg1 = HistorySegment( + segment_type="video", + token_range=(16, 48), + mm_item_id="vid-1", + ) + seg2 = HistorySegment(segment_type="user_text", token_range=(48, 64)) + req.session_history = [seg0, seg1, seg2] + req.mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="vid-1", + mm_position=PlaceholderRange(offset=16, length=32), + ) + ] + + # Allocate real KV blocks for this prompt. + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, 64, 0, manager_blocks) + + evict_segment(req, seg1, kv, enc) + + # Lockstep arrays compacted. + assert len(req._all_token_ids) == 32, len(req._all_token_ids) + assert len(req._mrope_positions) == 32, len(req._mrope_positions) + # Surviving mRoPE values keep their ORIGINAL coordinates (gaps). + assert req._mrope_positions[0] == (0, 0, 0) + assert req._mrope_positions[15] == (15, 15, 15) + assert req._mrope_positions[16] == (48, 48, 48), ( + f"surviving segment's first position should still be (48,48,48); " + f"got {req._mrope_positions[16]}" + ) + # max_cached_position preserved. + assert req.max_cached_position == 63 + # Later segment's token_range shifted down by 32. + assert seg2.token_range == (16, 32), seg2.token_range + # Earlier (pinned) segment unchanged. + assert seg0.token_range == (0, 16), seg0.token_range + # Evicted segment removed from history. + assert seg1 not in req.session_history + # Evicted segment's mm_feature dropped. + assert all(f.identifier != "vid-1" for f in req.mm_features) + # Counters decremented. + assert req.num_prompt_tokens == 32 + assert req.num_computed_tokens == 32 + + +def test_maybe_evict_respects_pinned_segments(): + """Pinned segments stay even when over budget. Only unpinned + video segments count towards `max_video_segments`.""" + kv, enc = _make_managers() + req = _make_request("r-pin", [9] * 80) + req._mrope_positions = [(i, i, i) for i in range(80)] + req.max_cached_position = 79 + req.num_prompt_tokens = 80 + req.num_computed_tokens = 80 + # 1 pinned system prompt + 4 unpinned video segments. age_chunks > 0 + # on every segment: the eviction phases skip age-0 (just-appended) + # segments (finding #0), and this fixture models already-aged prior + # chunks. Ages descend oldest-highest to preserve oldest-first order. + req.session_history = [ + HistorySegment( + segment_type="system_prompt", + token_range=(0, 16), + pinned=True, + age_chunks=5, + ), + HistorySegment( + segment_type="video", + token_range=(16, 32), + mm_item_id="v0", + age_chunks=4, + ), + HistorySegment( + segment_type="video", + token_range=(32, 48), + mm_item_id="v1", + age_chunks=3, + ), + HistorySegment( + segment_type="video", + token_range=(48, 64), + mm_item_id="v2", + age_chunks=2, + ), + HistorySegment( + segment_type="video", + token_range=(64, 80), + mm_item_id="v3", + age_chunks=1, + ), + ] + req.mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier=f"v{i}", + mm_position=PlaceholderRange(offset=16 + i * 16, length=16), + ) + for i in range(4) + ] + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, 80, 0, manager_blocks) + + retention = StreamingRetentionParams(max_video_segments=2, max_session_tokens=4000) + n = maybe_evict_old_segments(req, retention, kv, enc) + assert n == 2, f"expected 2 evictions to drop 4 video segments to 2; got {n}" + + # Pinned segment present and untouched. + pinned = [s for s in req.session_history if s.pinned] + assert len(pinned) == 1 and pinned[0].token_range == (0, 16) + # Remaining video segments are v2 and v3 (oldest dropped). + surviving_videos = [s for s in req.session_history if s.segment_type == "video"] + assert [s.mm_item_id for s in surviving_videos] == ["v2", "v3"] + + +def test_maybe_evict_noop_under_budget(): + """If video-segment count <= budget, eviction is a no-op.""" + kv, enc = _make_managers() + req = _make_request("r-under", [9] * 32) + req._mrope_positions = [(i, i, i) for i in range(32)] + req.max_cached_position = 31 + req.num_prompt_tokens = 32 + req.num_computed_tokens = 32 + req.session_history = [ + HistorySegment(segment_type="video", token_range=(0, 16), mm_item_id="v0"), + HistorySegment(segment_type="video", token_range=(16, 32), mm_item_id="v1"), + ] + n = maybe_evict_old_segments( + req, + StreamingRetentionParams(max_video_segments=4, max_session_tokens=4000), + kv, + enc, + ) + assert n == 0 + assert len(req.session_history) == 2 + + +def test_evict_pinned_segment_raises(): + """Direct `evict_segment` on a pinned segment is a programming error.""" + kv, enc = _make_managers() + req = _make_request("r-pinned", [9] * 16) + req.num_prompt_tokens = 16 + req.num_computed_tokens = 16 + pinned_seg = HistorySegment( + segment_type="system_prompt", token_range=(0, 16), pinned=True + ) + req.session_history = [pinned_seg] + with pytest.raises(ValueError, match="refusing to evict pinned"): + evict_segment(req, pinned_seg, kv, enc) + + +def test_evict_absorbs_orphans_into_next_segment(): + """When block-aligned eviction leaves orphan tokens at the segment's + boundaries, the next segment's `token_range` start extends backward + to absorb them. Without this, orphans pile up over many evictions + and `num_prompt_tokens` drifts past `max_session_tokens` until the + worker buffer overflows. + + Set up three segments where the middle one is NOT block-aligned: + - seg0: pinned anchor (0, 12) + - seg1: video (12, 44) ← evicted; size 32 + - seg2: video (44, 76) + + With block_size=16, inward alignment of [12, 44) yields [16, 32), + freeing 16 tokens but leaving 4 front-orphans at indices [12, 16) + and 12 back-orphans (originally [32, 44), now at [16, 28) after del). + seg2's range was [44, 76]; after shift it would be [28, 60]. With + absorption, seg2 instead becomes (12, 60), so the orphans are now + part of seg2 and will be reclaimed when seg2 is eventually evicted. + """ + kv, enc = _make_managers(block_size=16, num_blocks=16) + req = _make_request("r-orphan-absorb", [9] * 76) + req._mrope_positions = [(i, i, i) for i in range(76)] + req.max_cached_position = 75 + req.num_prompt_tokens = 76 + req.num_computed_tokens = 76 + + seg0 = HistorySegment(segment_type="user_text", token_range=(0, 12), pinned=True) + seg1 = HistorySegment( + segment_type="video", token_range=(12, 44), mm_item_id="v-mid" + ) + seg2 = HistorySegment( + segment_type="video", token_range=(44, 76), mm_item_id="v-last" + ) + req.session_history = [seg0, seg1, seg2] + req.mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v-mid", + mm_position=PlaceholderRange(offset=12, length=32), + ), + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v-last", + mm_position=PlaceholderRange(offset=44, length=32), + ), + ] + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, 76, 0, manager_blocks) + + evict_segment(req, seg1, kv, enc) + + # Pinned anchor unchanged. + assert seg0.token_range == (0, 12) + # seg2 absorbed orphans by extending its start back to seg1's raw + # start (12). Without absorption it would be (28, 60). + assert seg2.token_range == (12, 60), seg2.token_range + # seg1 retired. + assert seg1 not in req.session_history + # Arrays compacted by block-aligned 16 tokens (not seg1's raw 32). + assert req.num_prompt_tokens == 60 + assert len(req._all_token_ids) == 60 + + +def test_evict_token_range_inward_alignment(): + """`evict_token_range_for_request` rounds inward: a 30-token range + starting at 0 with block_size=16 frees ONE block (covering tokens + 0..15), not two. Tokens 16..29 stay in the partial second block.""" + kv, _ = _make_managers(block_size=16, num_blocks=8) + req = _make_request("r-align", [9] * 48) + kv.allocate_slots(req, 48, 0, kv.get_computed_blocks(req)[0]) + free_q = kv.block_pool.free_block_queue + free_before = free_q.num_free_blocks + + aligned_start, aligned_end, n = kv.evict_token_range_for_request( + "r-align", token_start=0, token_end=30 + ) + # Inward: only block 0 (tokens 0..15) freed; block 1 (16..31) stays. + assert n == 1, n + assert (aligned_start, aligned_end) == (0, 16), (aligned_start, aligned_end) + assert free_q.num_free_blocks == free_before + 1 + + +# --------------------------------------------------------------------------- +# Tests for asymmetric retention: split kept_output into assistant_text +# segments + separate text-token budget. Added with the change that +# removes the absorb-into-prior-video behavior. +# --------------------------------------------------------------------------- + + +def _dummy_video_mm_features(n_videos: int) -> list: + """Build n MultiModalFeatureSpec instances tagged v0..v{n-1}, each + occupying a 32-token placeholder. Tests that exercise eviction past + video segments need these so `evict_segment` can find the matching + mm_feature to drop.""" + return [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier=f"v{i}", + mm_position=PlaceholderRange(offset=i * 32, length=32), + ) + for i in range(n_videos) + ] + + +def _record_three_chunks( + *, anchor_len: int, video_len: int, caption_len: int +) -> list[HistorySegment]: + """Simulate three calls to `_record_streaming_segment` and return the + resulting `session_history`. Bypasses the rest of the scheduler — only + exercises the bookkeeping logic. + + Layout per call: + call A (chunk 2 arriving): anchor synthesized, ch1_caption + + ch2_video appended. + call B (chunk 3 arriving): ch2_caption + ch3_video appended. + call C (chunk 4 arriving): ch3_caption + ch4_video appended. + + Wait — caller asks for THREE chunks recorded. That maps to two + `_record_streaming_segment` calls (chunks 2 and 3 arriving), not three. + We do two calls and end up with: anchor, ch1_caption, ch2_video, + ch2_caption, ch3_video. Five segments. Matches the design doc's + "after the third chunk's recording" wording — "the third chunk + arrived" = chunk 3's call to `_record_streaming_segment`, which is + the second call. + + NOTE: We don't have a real scheduler instance, so we call + `Scheduler._record_streaming_segment(None, ...)` as an unbound + method. The implementation deliberately doesn't touch `self`. + """ + # Local import to avoid hard scheduler dependency at module import + # (the existing test file otherwise stays scheduler-free). + from vllm.v1.core.sched.scheduler import Scheduler + from vllm.v1.request import StreamingUpdate + + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + sess = Request( + request_id="r-record", + prompt_token_ids=[1] * anchor_len, # chunk 1's prompt only + sampling_params=sp, + pooling_params=None, + ) + sess.mm_features = [] + + # Call A (chunk 2 arrives): the scheduler has already extended + # prompt_token_ids with kept_output (chunk 1's caption) followed by + # chunk 2's video content. Replicate that final state before calling + # `_record_streaming_segment`. Lockstep arrays end at + # anchor_len + caption_len + video_len. + sess.prompt_token_ids.extend([2] * caption_len) + sess.prompt_token_ids.extend([3] * video_len) + sess._all_token_ids = list(sess.prompt_token_ids) + sess.num_prompt_tokens = len(sess.prompt_token_ids) + update_a = StreamingUpdate( + mm_features=[ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v2", + mm_position=PlaceholderRange( + offset=anchor_len + caption_len, length=video_len + ), + ) + ], + prompt_token_ids=[3] * video_len, + max_tokens=16, + arrival_time=0.0, + sampling_params=sp, + ) + sess.mm_features.extend(update_a.mm_features) + Scheduler._record_streaming_segment( + None, + sess, + update_a, + prior_response_start=anchor_len, + segment_start_idx=anchor_len + caption_len, + ) + + # Call B (chunk 3 arrives): prior_response_start = end of chunk 2's + # video segment (= anchor + caption + video), kept_output is + # chunk 2's caption, then chunk 3's video tokens. + chunk2_end = anchor_len + caption_len + video_len + sess.prompt_token_ids.extend([4] * caption_len) + sess._all_token_ids.extend([4] * caption_len) + sess.prompt_token_ids.extend([5] * video_len) + sess._all_token_ids.extend([5] * video_len) + sess.num_prompt_tokens = len(sess.prompt_token_ids) + update_b = StreamingUpdate( + mm_features=[ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v3", + mm_position=PlaceholderRange( + offset=chunk2_end + caption_len, length=video_len + ), + ) + ], + prompt_token_ids=[5] * video_len, + max_tokens=16, + arrival_time=0.0, + sampling_params=sp, + ) + sess.mm_features.extend(update_b.mm_features) + Scheduler._record_streaming_segment( + None, + sess, + update_b, + prior_response_start=chunk2_end, + segment_start_idx=chunk2_end + caption_len, + ) + return sess.session_history + + +def test_record_streaming_segment_splits_kept_output_into_assistant_text(): + """After two `_record_streaming_segment` calls (chunks 2 and 3 + arriving on top of chunk 1's anchor), `session_history` should be: + + [pinned anchor, ch1_caption, ch2_video, ch2_caption, ch3_video] + + Validates the split behavior — captions are NOT absorbed into the + prior video segment. Anchor is pinned; all others unpinned. + """ + history = _record_three_chunks(anchor_len=20, video_len=32, caption_len=12) + + assert len(history) == 5, [(s.segment_type, s.token_range) for s in history] + types = [s.segment_type for s in history] + assert types == [ + "user_text", # anchor (chunk 1 prompt) + "assistant_text", # chunk 1 caption + "video", # chunk 2 video + "assistant_text", # chunk 2 caption + "video", # chunk 3 video + ], types + + # Only the anchor is pinned. + pinned_flags = [s.pinned for s in history] + assert pinned_flags == [True, False, False, False, False], pinned_flags + + # Token ranges are contiguous and non-overlapping. + ranges = [s.token_range for s in history] + assert ranges == [ + (0, 20), # anchor: chunk 1 prompt + (20, 32), # chunk 1 caption (12 tokens) + (32, 64), # chunk 2 video (32 tokens) + (64, 76), # chunk 2 caption (12 tokens) + (76, 108), # chunk 3 video (32 tokens) + ], ranges + + # Video segments carry mm_item_id; caption + anchor don't. + assert history[0].mm_item_id is None # anchor + assert history[1].mm_item_id is None # chunk 1 caption + assert history[2].mm_item_id == "v2" + assert history[3].mm_item_id is None # chunk 2 caption + assert history[4].mm_item_id == "v3" + + +def _make_request_with_history( + request_id: str, + block_size: int, + segments_spec: list[tuple[str, int, str | None, bool]], +) -> tuple[Request, KVCacheManager, EncoderCacheManager]: + """Build a request with a synthetic, contiguous session_history. + + `segments_spec` is a list of (segment_type, length, mm_item_id, pinned) + tuples. The function lays them out back-to-back starting at index 0, + populates `_all_token_ids` / `_mrope_positions` / `num_prompt_tokens` + accordingly, allocates real KV blocks for the full range, and returns + (req, kv_manager, enc_manager) ready for `maybe_evict_old_segments`. + + Video segments get a corresponding `MultiModalFeatureSpec` entry on + the request so `evict_segment` can drop them. + """ + total = sum(length for (_, length, _, _) in segments_spec) + # Allocate enough blocks: ceil(total / block_size) plus generous + # headroom. KVCacheManager has internal reservations (null block, + # max_model_len-sized worker buffers) that fail allocate_slots + # silently if num_blocks is too tight. Use 4× the segment-count + # estimate so any reasonable test fixture fits. + n_blocks = max(64, ((total + block_size - 1) // block_size) * 4) + kv, enc = _make_managers(block_size=block_size, num_blocks=n_blocks) + req = _make_request(request_id, [9] * total, block_size=block_size) + req._mrope_positions = [(i, i, i) for i in range(total)] + req.max_cached_position = total - 1 + req.num_prompt_tokens = total + req.num_computed_tokens = total + + history: list[HistorySegment] = [] + mm_features: list = [] + cursor = 0 + # Mirror production aging (finding #0): the three eviction phases only + # pick victims with `age_chunks > 0`, because the just-appended, + # not-yet-computed segment (age 0 per `_record_streaming_segment`) has + # no KV blocks yet and must never be evicted. This synthetic history + # models a session whose ENTIRE content is already-computed, aged prior + # chunks being trimmed — there is no fresh age-0 append in the fixture + # — so every segment is given a strictly-positive age. Ages descend + # oldest-highest so list order still coincides with chronology + # (oldest-first eviction). Tests that specifically need the age-0 guard + # build their fixture inline. + n_segments = len(segments_spec) + for idx, (seg_type, length, mm_id, pinned) in enumerate(segments_spec): + # Oldest segment (idx 0) gets the highest age; the newest still + # gets a positive age (>= 1) so it remains evictable. + age = n_segments - idx + history.append( + HistorySegment( + segment_type=seg_type, + token_range=(cursor, cursor + length), + mm_item_id=mm_id, + pinned=pinned, + age_chunks=age, + ) + ) + if seg_type == "video" and mm_id is not None: + mm_features.append( + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier=mm_id, + mm_position=PlaceholderRange(offset=cursor, length=length), + ) + ) + cursor += length + req.session_history = history + req.mm_features = mm_features + + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, total, 0, manager_blocks) + return req, kv, enc + + +def test_text_budget_evicts_oldest_assistant_text_first(): + """Phase 2 brings the total non-pinned text tokens at or below + `max_text_tokens`, evicting oldest assistant_text segments first. + Video segments stay untouched because phase 1 doesn't trigger. + + Layout (all sizes ≥ block_size=16 so all are evictable by phase 2): + anchor(32, pinned) | text0(32) | video0(32, v0) | text1(32) | + video1(32, v1) | text2(32) | video2(32, v2) | text3(32) | + video3(32, v3) + + Total = 288 tokens, of which 4 text × 32 = 128 are non-pinned text. + Set max_text_tokens=64, so phase 2 should evict 2 text segments + (oldest two), leaving 64 text tokens. max_video_segments=10 keeps + phase 1 quiet; max_session_tokens left at None so phase 3 is quiet. + """ + req, kv, enc = _make_request_with_history( + "r-text-budget", + block_size=16, + segments_spec=[ + ("user_text", 32, None, True), # pinned anchor + ("assistant_text", 32, None, False), # text0 + ("video", 32, "v0", False), + ("assistant_text", 32, None, False), # text1 + ("video", 32, "v1", False), + ("assistant_text", 32, None, False), # text2 + ("video", 32, "v2", False), + ("assistant_text", 32, None, False), # text3 + ("video", 32, "v3", False), + ], + ) + + retention = StreamingRetentionParams( + max_video_segments=10, max_text_tokens=64, max_session_tokens=4000 + ) + n = maybe_evict_old_segments(req, retention, kv, enc) + # Two text segments evicted to bring 128 → 64. + assert n == 2, n + + # Remaining text segments are text2 and text3 (youngest two). + text_segs = [ + s + for s in req.session_history + if s.segment_type in ("user_text", "assistant_text") and not s.pinned + ] + text_tokens = sum(s.token_range[1] - s.token_range[0] for s in text_segs) + assert text_tokens <= 64, text_tokens + assert len(text_segs) == 2, [s.token_range for s in text_segs] + + # All 4 video segments still present. + video_ids = [s.mm_item_id for s in req.session_history if s.segment_type == "video"] + assert video_ids == ["v0", "v1", "v2", "v3"], video_ids + + +def test_text_budget_respects_pinned_segments(): + """Phase 2 skips pinned text segments even if dropping them would + bring the budget under target. Falls through to the next oldest + unpinned text segment. + + Layout: anchor(32, pinned) | text0(32, pinned) | text1(32) | text2(32). + Non-pinned text tokens = 64. Set max_text_tokens=16: phase 2 needs + to free at least 48. text0 is pinned → skip. text1 is oldest + unpinned → evict (32 tokens freed). text2 remains; non-pinned text + tokens now = 32, still over 16. Try again: text2 is now oldest + unpinned → evict. Non-pinned text = 0, budget met. + """ + req, kv, enc = _make_request_with_history( + "r-text-pinned", + block_size=16, + segments_spec=[ + ("user_text", 32, None, True), # anchor (pinned) + ("assistant_text", 32, None, True), # pinned mid-session text + ("assistant_text", 32, None, False), # text1 + ("assistant_text", 32, None, False), # text2 + ], + ) + retention = StreamingRetentionParams( + max_video_segments=10, max_text_tokens=16, max_session_tokens=4000 + ) + n = maybe_evict_old_segments(req, retention, kv, enc) + assert n == 2, n + # Both pinned segments still present. + pinned = [s for s in req.session_history if s.pinned] + assert len(pinned) == 2 + # Both unpinned text segments gone. + non_pinned_text = [ + s + for s in req.session_history + if s.segment_type in ("user_text", "assistant_text") and not s.pinned + ] + assert non_pinned_text == [] + + +def test_independent_video_and_text_budgets(): + """Phase 1 and phase 2 each hit their own budget independently. + + Layout: anchor(32, pinned) | v0..v4 (each 32) | text0..text4 (each 32). + Set max_video_segments=2 and max_text_tokens=64. Phase 1 should evict + 3 videos (5→2). Phase 2 should evict 3 texts (160→64). Final: anchor + + 2 videos + 2 texts. + """ + spec: list[tuple[str, int, str | None, bool]] = [ + ("user_text", 32, None, True), + ] + for i in range(5): + spec.append(("video", 32, f"v{i}", False)) + for i in range(5): + spec.append(("assistant_text", 32, None, False)) + req, kv, enc = _make_request_with_history( + "r-indep", block_size=16, segments_spec=spec + ) + retention = StreamingRetentionParams( + max_video_segments=2, max_text_tokens=64, max_session_tokens=4000 + ) + n = maybe_evict_old_segments(req, retention, kv, enc) + assert n == 6, n # 3 video + 3 text + + surviving_videos = [s for s in req.session_history if s.segment_type == "video"] + assert [s.mm_item_id for s in surviving_videos] == ["v3", "v4"] + + non_pinned_text = [ + s + for s in req.session_history + if s.segment_type in ("user_text", "assistant_text") and not s.pinned + ] + text_tokens = sum(s.token_range[1] - s.token_range[0] for s in non_pinned_text) + assert text_tokens == 64 + assert len(non_pinned_text) == 2 + + +def test_video_eviction_preserves_associated_caption_segment(): + """The point of asymmetric retention: dropping a video segment + doesn't drop its caption. Layout simulates the post-split + streaming flow: anchor | v0_caption | v0 | v1_caption | v1 | + v2_caption | v2 | v3_caption | v3. Set max_video_segments=2 and + leave text budget unset. Phase 1 evicts v0 and v1; their caption + segments survive untouched (modulo orphan absorption shifting + token_ranges). Phase 2 + 3 don't fire. + """ + spec: list[tuple[str, int, str | None, bool]] = [ + ("user_text", 32, None, True), # anchor + ] + for i in range(4): + spec.append(("assistant_text", 32, None, False)) # captionN + spec.append(("video", 32, f"v{i}", False)) + req, kv, enc = _make_request_with_history( + "r-video-evict-keeps-text", block_size=16, segments_spec=spec + ) + + retention = StreamingRetentionParams(max_video_segments=2, max_session_tokens=4000) + n = maybe_evict_old_segments(req, retention, kv, enc) + assert n == 2, n + + # Only 2 videos remain (v2, v3). + videos = [s for s in req.session_history if s.segment_type == "video"] + assert [s.mm_item_id for s in videos] == ["v2", "v3"], [ + s.mm_item_id for s in videos + ] + # All 4 captions still present — their tokens stayed in + # _all_token_ids / _mrope_positions (eviction didn't touch them). + captions = [s for s in req.session_history if s.segment_type == "assistant_text"] + assert len(captions) == 4, [s.token_range for s in captions] + + # Encoder cache entries for the evicted videos are gone; for the + # surviving ones they're still referenced. + remaining_mm_ids = {f.identifier for f in req.mm_features} + assert remaining_mm_ids == {"v2", "v3"} + + +def test_sports_preset_steady_state_under_budget(): + """Document the running-cleanly-under-budget case. Sports preset: + max_video_segments=30, max_text_tokens=4000, max_session_tokens=7000. + Simulate 10 chunks: anchor(32) + 10 video(64 each) + 10 caption(20 + each). 10 videos < 30 budget. 10 captions × 20 = 200 tokens < 4000. + Total = 32 + 640 + 200 = 872 < 7000. No phase should fire. + """ + spec: list[tuple[str, int, str | None, bool]] = [ + ("user_text", 32, None, True), # anchor + ] + for i in range(10): + spec.append(("video", 64, f"v{i}", False)) + spec.append(("assistant_text", 20, None, False)) + req, kv, enc = _make_request_with_history( + "r-steady", block_size=16, segments_spec=spec + ) + + history_before = list(req.session_history) + tokens_before = req.num_prompt_tokens + + retention = StreamingRetentionParams( + max_video_segments=30, + max_text_tokens=4000, + max_session_tokens=7000, + ) + n = maybe_evict_old_segments(req, retention, kv, enc) + assert n == 0, n + assert req.session_history == history_before + assert req.num_prompt_tokens == tokens_before + + +# --------------------------------------------------------------------------- +# Unpadded strict eviction (the only mode) +# +# Segments are block-UNALIGNED (nothing is padded). Eviction gates on the +# EXACT "frees a whole block" predicate, falls through past un-freeable +# victims, and coalesces adjacent short captions — so no segment ever hits +# PATH C (silent orphan, no decrement) and short captions can't stack +# unbounded. These tests build unpadded synthetic sessions directly. +# --------------------------------------------------------------------------- + + +def _owned_tokens(req: Request) -> int: + """Sum of segment lengths. Equals num_prompt_tokens when there are no + untracked orphans; a gap means a last-segment eviction left bounded + forward-absorption residue (<= 2*(block_size-1)). A LARGE gap would + mean PATH C dropped a segment — the leak the design prevents.""" + return sum(e - s for s, e in (seg.token_range for seg in req.session_history)) + + +def _retention(**kw) -> StreamingRetentionParams: + """Build a StreamingRetentionParams for the strict-eviction tests. + + Several strict tests deliberately use tiny/illegal budgets (e.g. + max_session_tokens=1, max_text_tokens=0) to force the phase-3 / coalesce + / pack code paths. `StreamingRetentionParams.__post_init__` now rejects + those (max_session_tokens must be >= the minimum; max_text_tokens must be + > 0 when set). To keep that coverage we build a VALID config and then + degrade `max_session_tokens` / `max_text_tokens` via direct attribute + assignment after construction (the dataclass is mutable, so this bypasses + __post_init__). + """ + kw.setdefault("max_video_segments", 30) + requested_mst = kw.pop("max_session_tokens", 4000) + requested_mtt = kw.pop("max_text_tokens", None) + # Construct with a valid (large) session budget, then override. + retention = StreamingRetentionParams(max_session_tokens=4000, **kw) + retention.max_session_tokens = requested_mst + retention.max_text_tokens = requested_mtt + return retention + + +def test_strict_guard_leaves_straddling_block_sized_segment(): + """A text segment of length == block_size at a MISALIGNED start frees + ZERO whole blocks (inward rounding: `[4,12)` at bs=8 → ceil(4/8)=1, + floor(12/8)=1). The exact `_frees_block` guard must DECLINE to evict + it — leave it in place, no PATH C, no orphan. A length>=block_size + heuristic would admit it and PATH-C-drop it, leaking. This is the + central guard-threshold fix.""" + req, kv, enc = _make_request_with_history( + "r-strict-straddle", + block_size=8, + segments_spec=[ + ("user_text", 4, None, True), # anchor ends at 4 (unaligned) + ("assistant_text", 8, None, False), # [4,12) straddles blocks + ], + ) + n = maybe_evict_old_segments(req, _retention(max_session_tokens=1), kv, enc) + assert n == 0, n + assert len(req.session_history) == 2, "must NOT drop the segment" + assert req.num_prompt_tokens == _owned_tokens(req), "must not orphan" + + +def test_strict_coalesce_evicts_stacked_short_captions(): + """Eight 2-token captions stacked after a block-aligned sink. Strict + mode coalesces the contiguous run into block-freeing segments so phase + 2 can evict them. Without coalescing each is sub-block, skipped + forever, and the budget can never be met.""" + spec: list[tuple[str, int, str | None, bool]] = [ + ("user_text", 8, None, True) # block-aligned sink (bs=8) + ] + spec += [("assistant_text", 2, None, False) for _ in range(8)] # [8,24) + req, kv, enc = _make_request_with_history( + "r-strict-coalesce", block_size=8, segments_spec=spec + ) + before = req.num_prompt_tokens # 24 + n = maybe_evict_old_segments(req, _retention(max_text_tokens=0), kv, enc) + assert n >= 1, n + assert req.num_prompt_tokens < before + assert req.num_prompt_tokens - _owned_tokens(req) <= 2 * (8 - 1) + text_left = sum( + e - s + for s, e in ( + seg.token_range + for seg in req.session_history + if seg.segment_type in ("user_text", "assistant_text") and not seg.pinned + ) + ) + assert text_left <= 2 * (8 - 1), text_left + + +def test_strict_phase3_fallthrough_skips_front_subblock(): + """Phase 3 must SKIP a front sub-block segment that can't free a block + (and can't be coalesced because a video separates it from other text) + and fall through to evict the freeable video — not PATH-C-drop the + front segment or stall. Layout (bs=8): anchor[0,8) pinned | + small[8,10) | video[10,42) | tail[42,44).""" + req, kv, enc = _make_request_with_history( + "r-strict-fallthrough", + block_size=8, + segments_spec=[ + ("user_text", 8, None, True), + ("assistant_text", 2, None, False), # sub-block front caption + ("video", 32, "v0", False), # freeable, not coalescable w/ text + ("assistant_text", 2, None, False), # tail + ], + ) + before = req.num_prompt_tokens # 44 + n = maybe_evict_old_segments( + req, _retention(max_video_segments=30, max_session_tokens=20), kv, enc + ) + assert n >= 1 + assert req.num_prompt_tokens < before + assert req.num_prompt_tokens - _owned_tokens(req) <= 2 * (8 - 1) + # The freeable video was evicted; the front sub-block caption survived + # (skipped by the fall-through, NOT PATH-C-dropped). + assert "v0" not in {f.identifier for f in req.mm_features} + assert any( + (s.token_range[1] - s.token_range[0]) == 2 + and s.segment_type == "assistant_text" + for s in req.session_history + ), "front sub-block caption must survive" + + +def test_strict_phase1_skips_subblock_video(): + """Phase 1 (video budget) evicts the oldest video that can free a + block, skipping a sub-block video that would PATH-C. The sub-block + video survives (bounded); the loop terminates.""" + req, kv, enc = _make_request_with_history( + "r-strict-subvideo", + block_size=8, + segments_spec=[ + ("user_text", 8, None, True), + ("video", 3, "v-small", False), # sub-block video [8,11) + ("video", 32, "v-big", False), # [11,43) freeable + ], + ) + maybe_evict_old_segments(req, _retention(max_video_segments=1), kv, enc) + ids = {f.identifier for f in req.mm_features} + assert "v-big" not in ids, "freeable video should be evicted" + assert "v-small" in ids, "sub-block video should survive (not PATH-C-dropped)" + videos = [s for s in req.session_history if s.segment_type == "video"] + assert len(videos) == 1 + assert req.num_prompt_tokens - _owned_tokens(req) <= 2 * (8 - 1) + + +def test_strict_does_not_coalesce_across_text_subtypes(): + """user_text and assistant_text both count as text but carry different + retention intent; coalesce must NOT merge across the subtype boundary, + and neither tiny segment is PATH-C-dropped.""" + req, kv, enc = _make_request_with_history( + "r-strict-no-cross", + block_size=8, + segments_spec=[ + ("user_text", 8, None, True), + ("user_text", 2, None, False), # [8,10) + ("assistant_text", 2, None, False), # [10,12) + ], + ) + n = maybe_evict_old_segments(req, _retention(max_session_tokens=8), kv, enc) + assert n == 0, n + types = [s.segment_type for s in req.session_history] + assert "user_text" in types and "assistant_text" in types + assert req.num_prompt_tokens == _owned_tokens(req), "no leak / no drop" + + +def test_strict_evs_packs_alternating_tiny_segments(): + """EVS regime: frames pruned so small that captions AND videos are all + sub-block and ALTERNATE (cap, vid, cap, vid). No single segment frees + a block and text-coalesce can't merge across the videos — so Phase 3 + packs the oldest run ACROSS types into a block-freeing blob, evicting + whole captions+frames together and releasing the frames' encoder + entries. Without this the session would drift unbounded (the original + crash, now in the EVS regime).""" + spec: list[tuple[str, int, str | None, bool]] = [ + ("user_text", 8, None, True) # block-aligned sink (bs=8) + ] + for k in range(3): + spec.append(("assistant_text", 2, None, False)) # tiny caption + spec.append(("video", 3, f"v{k}", False)) # EVS-pruned tiny frame + spec.append(("assistant_text", 2, None, False)) + req, kv, enc = _make_request_with_history( + "r-strict-evs", block_size=8, segments_spec=spec + ) + before = req.num_prompt_tokens # 25; every non-sink segment sub-block + n = maybe_evict_old_segments(req, _retention(max_session_tokens=8), kv, enc) + + # Packing fired and made progress toward budget (no stall). + assert n >= 1, n + assert req.num_prompt_tokens < before + # No unbounded orphan leak — only bounded last-blob edge residue. + assert req.num_prompt_tokens - _owned_tokens(req) <= 2 * (8 - 1) + # Drained to ~budget (bounded overshoot only), proving it didn't stall + # on the all-sub-block alternating layout. + assert req.num_prompt_tokens <= 8 + 2 * (8 - 1) + # Cross-type packing released evicted frames' encoder entries / + # mm_features (range-based free over the blob): at least one frame + # gone, so the multi-frame blob's release path ran. + assert len(req.mm_features) < 3, [f.identifier for f in req.mm_features] diff --git a/tests/v1/streaming/test_eviction_edges.py b/tests/v1/streaming/test_eviction_edges.py new file mode 100644 index 000000000000..5f40ac859113 --- /dev/null +++ b/tests/v1/streaming/test_eviction_edges.py @@ -0,0 +1,598 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Eviction edge-case tests from the streaming deep review (C15, C16, C26, +C31). + +Scenarios that the main `test_eviction.py` suite never constructs: + + - C26 (4): `evict_segment`'s computed-frontier clamp on a PARTIALLY + computed session (every existing fixture is fully computed). + - C31 / C15 (2): `_pack_oldest_run_to_free_block` must stop at the + computed frontier instead of destructively merging the just-appended, + uncomputed segment into a blob whose eviction then no-ops. + - C15 (3): evicting the LAST segment when the only predecessor is pinned + leaves bounded, non-compounding untracked residue. + - C15 (1): an mm placeholder straddling the block-aligned freed range is + dropped whole (feature + encoder entry) while its edge tokens survive + via absorption. + - C15 (4): duplicate mm identifiers within one session collapse to a + single encoder-cache reference (set semantics) — the eviction of one + occurrence physically frees the embedding out from under the survivor. + The test pins today's semantics (incl. no double-free) and documents + the residual hazard. + - C16: `_neutralize_orphan_mm_placeholders` must zero orphan IMAGE + placeholders too (the streaming REST endpoint pushes image frames), not + only video ones, and must not touch live/mm-backed or non-streaming + rows. + +Fixtures mirror tests/v1/streaming/test_eviction.py (real Request + +KVCacheManager + EncoderCacheManager). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request +from vllm.v1.streaming.eviction import ( + _pack_oldest_run_to_free_block, + evict_segment, + maybe_evict_old_segments, +) +from vllm.v1.streaming.retention import ( + HistorySegment, + StreamingRetentionParams, +) + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_managers(block_size: int = 16, num_blocks: int = 64): + kv = KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=1024, + enable_caching=False, + scheduler_block_size=block_size, + hash_block_size=block_size, + ) + enc = EncoderCacheManager(cache_size=4096) + return kv, enc + + +def _make_request_with_history( + request_id: str, + block_size: int, + segments_spec: list[tuple[str, int, str | None, bool, int]], + num_computed_tokens: int | None = None, +): + """Like test_eviction.py's helper, but with an EXPLICIT per-segment + `age_chunks` (last tuple field) and an overridable computed frontier — + the two knobs the edge cases here need. Duplicate mm_item_ids are + allowed (C15 (4)).""" + total = sum(length for (_, length, _, _, _) in segments_spec) + kv, enc = _make_managers(block_size=block_size, num_blocks=64) + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + req = Request( + request_id=request_id, + prompt_token_ids=[9] * total, + sampling_params=sp, + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, sha256), + ) + req._mrope_positions = [(i, i, i) for i in range(total)] + req.max_cached_position = total - 1 + req.num_prompt_tokens = total + req.num_computed_tokens = ( + total if num_computed_tokens is None else num_computed_tokens + ) + + history: list[HistorySegment] = [] + mm_features: list[MultiModalFeatureSpec] = [] + cursor = 0 + for seg_type, length, mm_id, pinned, age in segments_spec: + history.append( + HistorySegment( + segment_type=seg_type, + token_range=(cursor, cursor + length), + mm_item_id=mm_id, + pinned=pinned, + age_chunks=age, + ) + ) + if seg_type == "video" and mm_id is not None: + mm_features.append( + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier=mm_id, + mm_position=PlaceholderRange(offset=cursor, length=length), + ) + ) + cursor += length + req.session_history = history + req.mm_features = mm_features + + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, total, 0, manager_blocks) + return req, kv, enc + + +def _degraded_retention(**kw) -> StreamingRetentionParams: + """Valid construction, then degrade below the constructor floors via + attribute assignment (same pattern as test_eviction.py) so tiny budgets + can force phase 3 / packing.""" + requested_mst = kw.pop("max_session_tokens", 4000) + retention = StreamingRetentionParams(max_session_tokens=4000, **kw) + retention.max_session_tokens = requested_mst + return retention + + +def _assert_session_invariants( + req: Request, block_size: int, check_features: bool = True +) -> int: + """Shared post-eviction invariant helper (C15 fix (a)). Returns the + untracked-residue gap so callers can pin its exact value.""" + # Lockstep arrays. + assert len(req._all_token_ids) == req.num_prompt_tokens + if req._mrope_positions: + assert len(req._mrope_positions) == len(req._all_token_ids) + # Segments sorted, disjoint, in-range. + segs = sorted(req.session_history, key=lambda s: s.token_range[0]) + prev_end = 0 + owned = 0 + for seg in segs: + start, end = seg.token_range + assert 0 <= start <= end <= req.num_prompt_tokens, seg.token_range + assert start >= prev_end, f"overlapping segments at {seg.token_range}" + prev_end = end + owned += end - start + gap = req.num_prompt_tokens - owned + assert 0 <= gap <= 2 * (block_size - 1), ( + f"untracked residue {gap} exceeds the bounded-orphan guarantee" + ) + if check_features: + ranges = [s.token_range for s in segs] + for feature in req.mm_features: + f_start = feature.mm_position.offset + f_end = f_start + feature.mm_position.length + assert any(s <= f_start and f_end <= e for s, e in ranges), ( + f"feature {feature.identifier} placeholder [{f_start}," + f"{f_end}) not fully inside any surviving segment" + ) + return gap + + +# --------------------------------------------------------------------------- +# C26 (4): computed-frontier clamp on a partially-computed session +# --------------------------------------------------------------------------- + + +def test_evict_segment_clamps_to_computed_frontier(): + """Guards C26 gap (4): with the computed frontier mid-victim, + `evict_segment` frees only whole blocks BELOW the frontier and keeps + arrays / segments / counters in lockstep.""" + req, kv, enc = _make_request_with_history( + "r-frontier-clamp", + block_size=16, + segments_spec=[ + ("user_text", 16, None, True, 3), + ("video", 32, "v0", False, 1), # victim [16, 48) + ("user_text", 16, None, False, 0), + ], + num_computed_tokens=40, # frontier mid-victim + ) + victim = req.session_history[1] + tail = req.session_history[2] + + evict_segment(req, victim, kv, enc) + + # Clamped to min(48, 40) = 40, then inward-aligned to [16, 32). + assert req.pending_evicted_token_ranges == [(16, 32)] + assert req.num_prompt_tokens == 48 + assert req.num_computed_tokens == 24 # 40 - 16 + assert victim not in req.session_history + # Tail absorbed the orphans back to the victim's raw start (16). + assert tail.token_range == (16, 48), tail.token_range + # Victim's feature dropped even though only part of its range was freed. + assert all(f.identifier != "v0" for f in req.mm_features) + _assert_session_invariants(req, block_size=16) + + +def test_evict_segment_noops_when_victim_above_frontier(): + """Guards C26 gap (4): a victim entirely (or all but a sub-block part) + beyond `num_computed_tokens` must be left completely untouched — no + pending ranges, no history mutation, no counter drift.""" + for frontier in (16, 20): # fully above / sub-block below + req, kv, enc = _make_request_with_history( + f"r-frontier-noop-{frontier}", + block_size=16, + segments_spec=[ + ("user_text", 16, None, True, 1), + ("video", 32, "v0", False, 0), # [16, 48), uncomputed + ], + num_computed_tokens=frontier, + ) + victim = req.session_history[1] + history_before = list(req.session_history) + tokens_before = list(req._all_token_ids) + + evict_segment(req, victim, kv, enc) + + assert req.pending_evicted_token_ranges == [] + assert req.session_history == history_before + assert req._all_token_ids == tokens_before + assert req.num_prompt_tokens == 48 + assert req.num_computed_tokens == frontier + assert [f.identifier for f in req.mm_features] == ["v0"] + + +# --------------------------------------------------------------------------- +# C31 / C15 (2): packing must stop at the computed frontier +# --------------------------------------------------------------------------- + + +def _pack_fixture(num_computed_tokens: int): + """anchor(8, pinned) | text[8,12) computed | video[12,17) age-0. + With bs=8, the text alone frees no block; only extending over the + age-0 video would — exactly the run the frontier guard must refuse + while the video is uncomputed.""" + return _make_request_with_history( + "r-pack-frontier", + block_size=8, + segments_spec=[ + ("user_text", 8, None, True, 3), + ("assistant_text", 4, None, False, 1), + ("video", 5, "v-new", False, 0), + ], + num_computed_tokens=num_computed_tokens, + ) + + +def test_pack_refuses_to_cross_computed_frontier(): + """Guards C31: `_pack_oldest_run_to_free_block` must NOT extend the run + over the just-appended uncomputed segment. Before the fix it merged + text+video into one blob (destroying types/mm_item_id in + session_history) and the subsequent eviction no-op'd on the clamp.""" + req, kv, enc = _pack_fixture(num_computed_tokens=12) + + assert _pack_oldest_run_to_free_block(req, 8) is None + # History untouched: still three distinct segments. + assert [s.segment_type for s in req.session_history] == [ + "user_text", + "assistant_text", + "video", + ] + assert req.session_history[1].token_range == (8, 12) + assert req.session_history[2].mm_item_id == "v-new" + + +def test_phase3_with_uncomputed_tail_stalls_cleanly_without_blob(): + """Guards C31 via the full phase-3 path: over budget but with nothing + evictable below the frontier, `maybe_evict_old_segments` must return 0 + and leave `session_history` unmangled (no merged blob, no orphaned + tokens) instead of packing across the frontier.""" + req, kv, enc = _pack_fixture(num_computed_tokens=12) + history_before = [ + (s.segment_type, s.token_range, s.mm_item_id) for s in req.session_history + ] + + n = maybe_evict_old_segments( + req, _degraded_retention(max_session_tokens=8), kv, enc + ) + + assert n == 0 + assert [ + (s.segment_type, s.token_range, s.mm_item_id) for s in req.session_history + ] == history_before + assert req.num_prompt_tokens == 17 + assert req.pending_evicted_token_ranges == [] + _assert_session_invariants(req, block_size=8) + + +def test_pack_proceeds_once_run_is_computed(): + """Positive control for the C31 guard: the identical layout with the + frontier past the video packs and evicts normally (the frontier — not + age — is the invariant; a computed age-0 segment is safe to pack).""" + req, kv, enc = _pack_fixture(num_computed_tokens=17) + + n = maybe_evict_old_segments( + req, _degraded_retention(max_session_tokens=8), kv, enc + ) + + assert n == 1, n + # The packed blob [8, 17) evicted its aligned span [8, 16). + assert req.pending_evicted_token_ranges == [(8, 16)] + assert req.num_prompt_tokens == 9 + assert all(f.identifier != "v-new" for f in req.mm_features) + _assert_session_invariants(req, block_size=8) + + +# --------------------------------------------------------------------------- +# C15 (3): last-segment eviction with only a pinned predecessor +# --------------------------------------------------------------------------- + + +def test_last_segment_eviction_with_pinned_predecessor_bounded_residue(): + """Guards C15 (3): with no following segment and only a PINNED + predecessor, backward absorption cannot run, so the sub-block tail + residue becomes untracked. It must stay bounded (< block_size) and must + NOT compound across later eviction passes.""" + req, kv, enc = _make_request_with_history( + "r-pinned-prev-residue", + block_size=8, + segments_spec=[ + ("user_text", 8, None, True, 2), + ("video", 34, "v-last", False, 1), # [8, 42), unaligned end + ], + ) + victim = req.session_history[1] + + evict_segment(req, victim, kv, enc) + + # Inward alignment freed [8, 40); tokens [40, 42) survive at [8, 10) + # owned by NO segment (anchor is pinned, nothing follows). + assert req.pending_evicted_token_ranges == [(8, 40)] + assert len(req.session_history) == 1 + assert req.session_history[0].pinned + assert req.num_prompt_tokens == 10 + gap = _assert_session_invariants(req, block_size=8) + assert gap == 2 + assert gap < 8, "residue must stay below one block" + + # Later eviction passes find no unpinned victim; the gap must not grow. + for _ in range(3): + n = maybe_evict_old_segments( + req, _degraded_retention(max_session_tokens=8), kv, enc + ) + assert n == 0 + assert _assert_session_invariants(req, block_size=8) == 2 + + +# --------------------------------------------------------------------------- +# C15 (1): placeholder straddling the aligned freed boundary +# --------------------------------------------------------------------------- + + +def test_straddling_placeholder_dropped_whole_and_orphans_absorbed(): + """Guards C15 (1): a feature whose placeholder extends past the inward + block-aligned freed range is dropped ENTIRELY (feature + physical + encoder eviction) while its out-of-range pad tokens survive via + absorption into the next segment. Surviving features must remain fully + inside surviving segments; the freed mm_hash must reach `freed` so the + worker drops the GPU tensor (the orphan pad tokens are then zeroed at + re-prefill by the C16 neutralization — tested below).""" + req, kv, enc = _make_request_with_history( + "r-straddle", + block_size=16, + segments_spec=[ + ("user_text", 12, None, True, 3), + ("video", 32, "v-straddle", False, 2), # [12, 44), unaligned + ("video", 32, "v-last", False, 1), # [44, 76) + ], + ) + enc.allocate(req, input_id=0) + enc.allocate(req, input_id=1) + cache_size = enc.cache_size + victim = req.session_history[1] + tail = req.session_history[2] + + evict_segment(req, victim, kv, enc) + + # Aligned free of [16, 32): the placeholder [12, 44) straddles BOTH + # edges, yet the feature is dropped whole. + assert req.pending_evicted_token_ranges == [(16, 32)] + assert [f.identifier for f in req.mm_features] == ["v-last"] + # Physically evicted, exactly once, and drained. + assert "v-straddle" not in enc.cached + assert "v-straddle" not in enc.freeable + assert enc.get_freed_mm_hashes() == ["v-straddle"] + assert enc.get_freed_mm_hashes() == [] + # Slot accounting: only v-last's 32 embeds remain charged, and the + # gating invariant num_freeable_slots == num_free_slots + sum(freeable) + # holds with freeable empty. + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == enc.num_free_slots + # The orphan pad tokens ([12,16) and the survivors of [32,44)) were + # absorbed into the next segment, which extends back to the raw start. + assert tail.token_range == (12, 60) + # v-last's placeholder shifted down and still sits inside `tail`. + assert req.mm_features[0].mm_position.offset == 28 + _assert_session_invariants(req, block_size=16) + + +# --------------------------------------------------------------------------- +# C15 (4): duplicate mm identifier within one session +# --------------------------------------------------------------------------- + + +def test_duplicate_identifier_evicts_physically_under_live_reference(): + """Guards C15 (4) by PINNING today's semantics: EncoderCacheManager + refcounts are a per-mm_hash SET of request ids, so two features of one + request sharing an identifier (repeated/static frames) collapse to one + reference. Evicting the first occurrence therefore physically frees the + embedding while the surviving feature still references it — the + residual hazard the review documented (a re-prefill would need the + freed embedding). The test also pins that the second eviction is + idempotent: no double slot credit, no crash. If a guard lands (skip + `evict_unreferenced` while another live feature shares the identifier), + update the physical-evict asserts here.""" + req, kv, enc = _make_request_with_history( + "r-dup-id", + block_size=16, + segments_spec=[ + ("user_text", 16, None, True, 3), + ("video", 32, "dup", False, 2), # [16, 48) + ("video", 32, "dup", False, 1), # [48, 80), same content hash + ], + ) + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + enc.allocate(req, input_id=1) + # Set semantics: two allocations, ONE reference; both charged. + assert enc.cached["dup"] == {req.request_id} + assert enc.num_free_slots == cache_size - 64 + + seg_a, seg_b = req.session_history[1], req.session_history[2] + evict_segment(req, seg_a, kv, enc) + + # Physically evicted despite the live duplicate feature (the hazard). + assert "dup" not in enc.cached + assert enc.get_freed_mm_hashes() == ["dup"] + surviving = [f for f in req.mm_features if f.identifier == "dup"] + assert len(surviving) == 1, ( + "the second occurrence must still be live on the request" + ) + # Only ONE allocation's worth of slots was credited back; the + # survivor's charge remains (bounded leak, documented). + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == enc.num_free_slots + + # Evicting the surviving occurrence must be idempotent on the encoder + # cache: no double credit, no KeyError, nothing new in `freed`. + evict_segment(req, seg_b, kv, enc) + assert enc.get_freed_mm_hashes() == [] + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == enc.num_free_slots + assert req.mm_features == [] + _assert_session_invariants(req, block_size=16) + + +# --------------------------------------------------------------------------- +# C16: orphan-placeholder neutralization covers IMAGE (and video) markers +# --------------------------------------------------------------------------- + +IMAGE_TOKEN_ID = 151655 # Qwen3-VL <|image_pad|> +VIDEO_TOKEN_ID = 151656 # Qwen3-VL <|video_pad|> + + +def _run_neutralize( + token_ids: list[int], + req_specs: list[tuple[str, int, object]], + is_mm_flags: list[bool], + placeholder_ids: tuple[int, ...] = (IMAGE_TOKEN_ID, VIDEO_TOKEN_ID), +) -> torch.Tensor: + """Drive `GPUModelRunner._neutralize_orphan_mm_placeholders` on a stub + runner (same unbound-method pattern as test_worker_slice.py). + `req_specs` is [(req_id, num_scheduled_tokens, sampling_params)].""" + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + num_tokens = len(token_ids) + assert sum(n for _, n, _ in req_specs) == num_tokens + stub = SimpleNamespace( + mm_placeholder_ids_cpu=( + torch.tensor(sorted(placeholder_ids), dtype=torch.int32) + if placeholder_ids + else None + ), + input_batch=SimpleNamespace(req_ids=[rid for rid, _, _ in req_specs]), + requests={rid: SimpleNamespace(sampling_params=sp) for rid, _, sp in req_specs}, + input_ids=SimpleNamespace(gpu=torch.tensor(token_ids, dtype=torch.int32)), + _req_uses_streaming_retention=(GPUModelRunner._req_uses_streaming_retention), + ) + scheduler_output = SimpleNamespace( + num_scheduled_tokens={rid: n for rid, n, _ in req_specs} + ) + inputs_embeds = torch.ones((num_tokens, 4)) + is_mm_embed = torch.tensor(is_mm_flags, dtype=torch.bool) + GPUModelRunner._neutralize_orphan_mm_placeholders( + stub, scheduler_output, inputs_embeds, is_mm_embed, num_tokens + ) + return inputs_embeds + + +def _streaming_sp(): + return SimpleNamespace( + extra_args={"streaming_retention": {"reprefill_threshold": 0.7}} + ) + + +def _plain_sp(): + return SimpleNamespace(extra_args=None) + + +def test_neutralize_zeros_orphan_image_and_video_placeholders(): + """Guards C16: orphan <|image_pad|> markers (the deployed streaming + workload pushes IMAGE frames) must be zeroed alongside <|video_pad|>; + live (mm-backed) placeholder rows, plain text rows, and non-streaming + requests' rows must be untouched.""" + embeds = _run_neutralize( + token_ids=[ + IMAGE_TOKEN_ID, # streaming, orphan image pad -> zeroed + 7, # streaming, text -> kept + IMAGE_TOKEN_ID, # streaming, LIVE (is_mm_embed)-> kept + VIDEO_TOKEN_ID, # streaming, orphan video pad -> zeroed + IMAGE_TOKEN_ID, # NON-streaming request -> kept + ], + req_specs=[ + ("sess-stream", 4, _streaming_sp()), + ("req-plain", 1, _plain_sp()), + ], + is_mm_flags=[False, False, True, False, False], + ) + zeroed = (embeds == 0).all(dim=1).tolist() + assert zeroed == [True, False, False, True, False], zeroed + + +def test_neutralize_noop_without_placeholder_ids(): + """C16 companion: a model defining neither image_token_id nor + video_token_id short-circuits (mm_placeholder_ids_cpu is None).""" + embeds = _run_neutralize( + token_ids=[IMAGE_TOKEN_ID, 7], + req_specs=[("sess-stream", 2, _streaming_sp())], + is_mm_flags=[False, False], + placeholder_ids=(), + ) + assert (embeds == 1).all() + + +def test_neutralize_noop_without_streaming_requests(): + """C16 companion: non-streaming batches are never touched (the + `eligible` mask gates the zeroing to retention sessions).""" + embeds = _run_neutralize( + token_ids=[IMAGE_TOKEN_ID, VIDEO_TOKEN_ID], + req_specs=[("req-plain", 2, _plain_sp())], + is_mm_flags=[False, False], + ) + assert (embeds == 1).all() diff --git a/tests/v1/streaming/test_mrope_streaming.py b/tests/v1/streaming/test_mrope_streaming.py new file mode 100644 index 000000000000..32e195142b5f --- /dev/null +++ b/tests/v1/streaming/test_mrope_streaming.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""CPU-only unit tests for the mRoPE streaming helpers. + +These tests verify the position-arithmetic invariants of +`compute_chunk_mrope_positions` and `make_chunk_relative_mm_features` +without needing a real Qwen3-VL model: we drive them with a tiny stub +model that implements the `SupportsMRoPE` protocol with a deterministic, +trivially-checkable algorithm. + +The parity properties exercised here: + +1. `compute_chunk_mrope_positions` shifts positions by exactly + `base_position` (its first column equals `base_position`). +2. Computing positions chunk-by-chunk produces a tensor identical to + computing positions once over the cumulative prompt, as long as + chunk-relative offsets are computed correctly and `base_position` is + advanced by `prev_max + 1`. This is the load-bearing invariant for + Piece 1's no-eviction case. +3. After eviction (Piece 3 territory but already simulatable here), + surviving positions retain their original values, and new chunks get + positions strictly greater than any surviving position. +""" + +from dataclasses import dataclass, field + +import pytest +import torch + +from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange +from vllm.v1.streaming.mrope import ( + compute_chunk_mrope_positions, + make_chunk_relative_mm_features, +) + +# Synthetic token IDs. Anything works; the stub model below doesn't actually +# care about token semantics, only about mm_position offsets. +TEXT_TOKEN = 100 +VIDEO_PAD_TOKEN = 200 + + +@dataclass +class StubMRoPEModel: + """Tiny model implementing `SupportsMRoPE` for tests. + + For each `MultiModalFeatureSpec`, we assign a 3x4 block of positions + starting at the next available row. Text tokens get linear `(p, p, p)` + positions, like Qwen3-VL's standard rope path. The algorithm is + intentionally simpler than Qwen3-VL's but exercises the same call + interface and the same prefix-deterministic property. + """ + + supports_mrope: bool = field(default=True, init=False) + tokens_per_video: int = 4 + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + # Walk tokens left-to-right. For each token, if it's inside any + # mm_feature's range, assign a video-grid position; else a text one. + # video positions: t advances per feature, h/w cycle 0..1 in a 2x2 grid. + positions: list[tuple[int, int, int]] = [] + feature_for_index: dict[int, int] = {} + for fi, feat in enumerate(mm_features): + start = feat.mm_position.offset + length = feat.mm_position.length + for i in range(start, start + length): + feature_for_index[i] = fi + + cursor = 0 + for idx, _tok in enumerate(input_tokens): + if idx in feature_for_index: + fi = feature_for_index[idx] + # Stable 2x2 grid layout: (t=fi-offset-into-stream, h, w) + within_feature = idx - mm_features[fi].mm_position.offset + h = within_feature // 2 + w = within_feature % 2 + t = cursor + positions.append((t, cursor + h, cursor + w)) + # Advance cursor by feature length at end of feature. + if within_feature == mm_features[fi].mm_position.length - 1: + cursor += 2 + else: + positions.append((cursor, cursor, cursor)) + cursor += 1 + tensor = torch.tensor(positions, dtype=torch.long).t().contiguous() + max_val = int(tensor.max().item()) if tensor.numel() > 0 else -1 + return tensor, max_val + 1 - len(input_tokens) + + +def _make_video_feature(offset: int, length: int) -> MultiModalFeatureSpec: + return MultiModalFeatureSpec( + data=None, + modality="video", + identifier=f"vid-{offset}", + mm_position=PlaceholderRange(offset=offset, length=length), + ) + + +def test_compute_chunk_shifts_by_base_position(): + model = StubMRoPEModel() + chunk_tokens = [TEXT_TOKEN, TEXT_TOKEN, TEXT_TOKEN] + base = 50 + + positions, new_max = compute_chunk_mrope_positions( + model, chunk_tokens, [], base_position=base + ) + + assert positions.shape == (3, 3) + # Text-only chunk: first column should be (base, base, base). + assert positions[:, 0].tolist() == [base, base, base] + assert positions[:, 1].tolist() == [base + 1, base + 1, base + 1] + assert positions[:, 2].tolist() == [base + 2, base + 2, base + 2] + assert new_max == base + 2 + + +def test_compute_chunk_empty_tokens_returns_sentinel(): + model = StubMRoPEModel() + positions, new_max = compute_chunk_mrope_positions(model, [], [], base_position=10) + assert positions.shape == (3, 0) + # `new_max` for an empty chunk should leave `max_cached_position` unchanged + # when the caller does `max_cached_position = new_max`. Returning + # `base_position - 1` achieves that. + assert new_max == 9 + + +def test_chunk_relative_mm_features_rewrites_offsets(): + cumulative = [ + _make_video_feature(offset=2, length=4), + _make_video_feature(offset=20, length=4), + _make_video_feature(offset=40, length=4), + ] + # Suppose prior chunk ended at token 24 (covers first two features). + chunk_relative = make_chunk_relative_mm_features(cumulative, prev_num_tokens=24) + assert len(chunk_relative) == 1 + assert chunk_relative[0].mm_position.offset == 40 - 24 + assert chunk_relative[0].mm_position.length == 4 + # Original features unchanged (we copy, not mutate). + assert cumulative[2].mm_position.offset == 40 + + +def test_chunk_relative_mm_features_includes_boundary_offset(): + """A feature whose offset == prev_num_tokens belongs to the new chunk.""" + cumulative = [_make_video_feature(offset=10, length=4)] + chunk_relative = make_chunk_relative_mm_features(cumulative, prev_num_tokens=10) + assert len(chunk_relative) == 1 + assert chunk_relative[0].mm_position.offset == 0 + + +def test_per_chunk_compute_matches_full_recompute_no_eviction(): + """Load-bearing parity check. + + Computing positions chunk-by-chunk (with chunk-relative mm_features and + `base_position = prev_max + 1`) must produce a tensor identical to a + single-shot full recompute over the cumulative prompt, assuming no + eviction. This is the equivalence the existing `_init_mrope_positions` + relies on; we're verifying our per-chunk path stays equivalent. + """ + model = StubMRoPEModel() + + # Chunk 1: [text, text, video(4), text] + chunk1_tokens = [TEXT_TOKEN, TEXT_TOKEN] + [VIDEO_PAD_TOKEN] * 4 + [TEXT_TOKEN] + chunk1_features = [_make_video_feature(offset=2, length=4)] + + # Chunk 2: [text, video(4), text, text] + chunk2_tokens = [TEXT_TOKEN] + [VIDEO_PAD_TOKEN] * 4 + [TEXT_TOKEN, TEXT_TOKEN] + # Chunk-relative offset for chunk 2's video: 1. + chunk2_features_rel = [_make_video_feature(offset=1, length=4)] + + # Per-chunk path. + pos1, max1 = compute_chunk_mrope_positions( + model, chunk1_tokens, chunk1_features, base_position=0 + ) + pos2, max2 = compute_chunk_mrope_positions( + model, chunk2_tokens, chunk2_features_rel, base_position=max1 + 1 + ) + per_chunk = torch.cat([pos1, pos2], dim=1) + + # Full-recompute path: build cumulative prompt and cumulative mm_features + # with absolute offsets. + cumulative_tokens = chunk1_tokens + chunk2_tokens + cumulative_features = [ + _make_video_feature(offset=2, length=4), + _make_video_feature(offset=len(chunk1_tokens) + 1, length=4), + ] + full, _ = model.get_mrope_input_positions(cumulative_tokens, cumulative_features) + + assert per_chunk.shape == full.shape + assert torch.equal(per_chunk, full), ( + f"Per-chunk and full-recompute paths disagreed.\n" + f"per-chunk:\n{per_chunk}\nfull:\n{full}" + ) + + +def test_eviction_then_new_chunk_assigns_strictly_greater_positions(): + """Piece 3 invariant exercised at the Piece 1 helper level. + + After eviction the surviving positions retain their values. The next + chunk's positions must all be strictly greater than any surviving + position. This is what `base_position = max_cached_position + 1` + guarantees. + """ + model = StubMRoPEModel() + + # Chunk 1: text + video + text. + chunk1_tokens = [TEXT_TOKEN, TEXT_TOKEN] + [VIDEO_PAD_TOKEN] * 4 + [TEXT_TOKEN] + chunk1_features = [_make_video_feature(offset=2, length=4)] + pos1, max1 = compute_chunk_mrope_positions( + model, chunk1_tokens, chunk1_features, base_position=0 + ) + + # Simulate eviction of the video segment (tokens [2:6]). + surviving = torch.cat([pos1[:, :2], pos1[:, 6:]], dim=1) + surviving_max_pre = int(surviving.max().item()) + + # Next chunk: just text. Crucially, `max_cached_position` carried over + # from chunk 1 (NOT shrunk to the surviving max), so new positions are + # strictly greater than max1, leaving a gap where the video used to be. + next_chunk_tokens = [TEXT_TOKEN, TEXT_TOKEN] + next_pos, next_max = compute_chunk_mrope_positions( + model, next_chunk_tokens, [], base_position=max1 + 1 + ) + + assert int(next_pos.min().item()) > surviving_max_pre, ( + "new-chunk positions must exceed every surviving position" + ) + assert next_max == max1 + len(next_chunk_tokens) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/v1/streaming/test_prefix_cache_eviction.py b/tests/v1/streaming/test_prefix_cache_eviction.py new file mode 100644 index 000000000000..31d89fe33b95 --- /dev/null +++ b/tests/v1/streaming/test_prefix_cache_eviction.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Eviction x prefix-caching cross-request tests from the streaming deep +review (C17, C18). Prefix caching is ON in production, but until now every +streaming-eviction test ran with `enable_caching=False`. + + - C18 (a1): splice-miss regression — after a mid-stream eviction, an + identical-stream request must MISS at the spliced range (hash strip + + first-miss break), never chain across it. + - C18 (b): num_cached_block accounting when an evicted block is SHARED + (ref_cnt > 1): memory not reclaimed, sharer untouched, hash stripped, + exact per-request counters. + - C18 (d): the re-prefill KV sequence (free + evict_blocks, as + `_reprefill_streaming_session` does) makes a same-chain admission miss + completely. + - C17: session block-hash chains are salted (`cache_salt` = the + session's internal request id, set by AsyncLLM.handle_inputs), so a + session's KV — whose content/positions diverge from the content-chain + hashes after eviction/re-prefill — is unreachable by any other request + while preempt-resume self-hits still work. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch + +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request + +pytestmark = pytest.mark.cpu_test + +BLOCK_SIZE = 16 + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_manager(num_blocks: int = 64) -> KVCacheManager: + return KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=1024, + enable_caching=True, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + ) + + +def _make_request( + request_id: str, + prompt_token_ids: list[int], + cache_salt: str | None = None, + hash_fn: Callable = sha256, +) -> Request: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + return Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=sp, + pooling_params=None, + cache_salt=cache_salt, + block_hasher=get_request_block_hasher(BLOCK_SIZE, hash_fn), + ) + + +def _allocate_and_cache( + manager: KVCacheManager, req: Request, num_tokens: int +) -> list[int]: + """Admit `req`, allocate its whole prompt, and return the first-group + block ids. With enable_caching=True, allocate_slots also publishes the + full blocks into the prefix cache.""" + computed, num_computed, _ = manager.get_computed_blocks(req) + manager.allocate_slots(req, num_tokens - num_computed, num_computed, computed) + return [ + b.block_id + for b in manager.coordinator.single_type_managers[0].req_to_blocks[ + req.request_id + ] + if not b.is_null + ] + + +# The identical "original stream" every replay request uses. Varied tokens +# per block so each block hash is distinct. +STREAM = [i % 100 for i in range(96)] # 6 blocks + + +# --------------------------------------------------------------------------- +# C18 (a1): splice-miss regression +# --------------------------------------------------------------------------- + + +def test_identical_stream_misses_at_evicted_splice(): + """Guards C18 (a1): after session A evicts a middle token range, a + fresh request replaying the identical original stream must have its + cache hit stop strictly BEFORE the spliced range — the evicted blocks' + hashes were stripped and find_longest_cache_hit breaks at the first + miss, so it can never chain into A's post-splice blocks.""" + manager = _make_manager() + req_a = _make_request("sess-a", STREAM) + _allocate_and_cache(manager, req_a, 96) + # All 6 full blocks are now published in the prefix cache. + assert manager.coordinator.single_type_managers[0].num_cached_block["sess-a"] == 6 + + # Session A evicts the middle two blocks' worth ([32, 64)). + aligned_start, aligned_end, n_freed = manager.evict_token_range_for_request( + "sess-a", token_start=32, token_end=64 + ) + assert (aligned_start, aligned_end, n_freed) == (32, 64, 2) + + # Identical-stream replay request: the hit must stop at the splice. + req_b = _make_request("replay-b", STREAM) + _, num_hit_tokens, _ = manager.get_computed_blocks(req_b) + assert num_hit_tokens == 32, ( + f"cache hit must stop at the evicted range (32), got {num_hit_tokens}" + ) + + +# --------------------------------------------------------------------------- +# C18 (b): shared-block eviction accounting +# --------------------------------------------------------------------------- + + +def test_evicting_shared_block_reclaims_nothing_and_keeps_sharer_intact(): + """Guards C18 (b): evicting a block still shared with another request + (ref_cnt > 1) must return 0 (no memory reclaimed), strip the hash from + the prefix cache, leave the sharer's block table untouched, and keep + each request's num_cached_block exactly right (no reliance on the + max(0, ...) clamp).""" + manager = _make_manager() + single = manager.coordinator.single_type_managers[0] + tokens = [i % 100 for i in range(48)] # 3 blocks + + req_a = _make_request("sess-a", tokens) + a_blocks = _allocate_and_cache(manager, req_a, 48) + assert len(a_blocks) == 3 + + # B replays the same stream and shares A's cached prefix. The hit is + # capped at num_tokens - 1, so B shares the first TWO blocks. + req_b = _make_request("sess-b", tokens) + computed, num_hit, _ = manager.get_computed_blocks(req_b) + assert num_hit == 32 + manager.allocate_slots(req_b, 48 - num_hit, num_hit, computed) + b_blocks = [b.block_id for b in single.req_to_blocks["sess-b"]] + shared_id = a_blocks[1] + assert shared_id in b_blocks + assert manager.block_pool.blocks[shared_id].ref_cnt == 2 + + cached_a_before = single.num_cached_block["sess-a"] + cached_b_before = single.num_cached_block["sess-b"] + free_before = manager.block_pool.free_block_queue.num_free_blocks + + n_freed = manager.coordinator.evict_blocks_for_request("sess-a", [shared_id]) + + # Memory NOT reclaimed: the sharer still holds a reference. + assert n_freed == 0 + assert manager.block_pool.free_block_queue.num_free_blocks == free_before + assert manager.block_pool.blocks[shared_id].ref_cnt == 1 + # Dropped from A's table only; B's table untouched. + assert shared_id not in [b.block_id for b in single.req_to_blocks["sess-a"]] + assert [b.block_id for b in single.req_to_blocks["sess-b"]] == b_blocks + # Hash stripped so no third request can hit the stale entry. + assert manager.block_pool.blocks[shared_id].block_hash is None + # Exact accounting: A lost exactly the one cached block; B unchanged. + assert single.num_cached_block["sess-a"] == cached_a_before - 1 + assert single.num_cached_block["sess-b"] == cached_b_before + + # A third identical-stream request must now stop before the stripped + # block (only block 0 remains hittable). + req_c = _make_request("replay-c", tokens) + _, num_hit_c, _ = manager.get_computed_blocks(req_c) + assert num_hit_c == 16 + + +# --------------------------------------------------------------------------- +# C18 (d): re-prefill's evict_blocks makes a same-chain admission miss +# --------------------------------------------------------------------------- + + +def test_reprefill_block_eviction_prevents_same_chain_hit(): + """Guards C18 (d): `_reprefill_streaming_session` frees the session's + blocks and then evicts their prefix-cache entries + (kv_cache_manager.evict_blocks) — a request with the identical hash + chain admitted afterwards must get ZERO computed blocks; the session's + old KV (about to be recomputed at fresh dense positions) must not be + hittable.""" + manager = _make_manager() + req_a = _make_request("sess-a", STREAM) + a_block_ids = set(_allocate_and_cache(manager, req_a, 96)) + + # Sanity: before the re-prefill sequence the chain IS hittable. + probe = _make_request("probe", STREAM) + _, num_hit_before, _ = manager.get_computed_blocks(probe) + assert num_hit_before > 0 + + # The exact KV sequence of _reprefill_streaming_session + # (scheduler.py: gather ids -> free -> evict_blocks). + gathered: set[int] = set() + for per_group_ids in manager.get_block_ids("sess-a"): + gathered.update(per_group_ids) + assert gathered == a_block_ids + manager.free(req_a) + manager.evict_blocks(gathered) + + req_b = _make_request("replay-b", STREAM) + _, num_hit_after, _ = manager.get_computed_blocks(req_b) + assert num_hit_after == 0, ( + "same-chain admission after re-prefill eviction must miss entirely" + ) + + +# --------------------------------------------------------------------------- +# C17: session hash chains are salted and therefore globally unique +# --------------------------------------------------------------------------- + + +def test_cache_salt_isolates_session_hash_chain(): + """Guards C17: with the fix, every streaming session gets + `cache_salt = ` before its first chunk, so + its block-hash chain can never collide with (a) an unsalted request + replaying the same stream or (b) another session (different salt) — + while the session itself (same salt, e.g. preempt-resume) still + self-hits. This is what keeps post-eviction / post-re-prefill KV from + poisoning other requests despite the spliced content chain.""" + manager = _make_manager() + + salted = _make_request("sess-a", STREAM, cache_salt="sess-a-internal-id") + unsalted = _make_request("plain-b", STREAM) + other_session = _make_request("sess-c", STREAM, cache_salt="sess-c-internal-id") + same_salt = _make_request("sess-a-resume", STREAM, cache_salt="sess-a-internal-id") + + # The salt feeds the FIRST block's extra keys and chains into every + # later hash: the salted chain shares NO hash with the unsalted or + # differently-salted chains. + assert len(salted.block_hashes) == 6 + assert set(salted.block_hashes).isdisjoint(unsalted.block_hashes) + assert set(salted.block_hashes).isdisjoint(other_session.block_hashes) + assert salted.block_hashes == same_salt.block_hashes + + _allocate_and_cache(manager, salted, 96) + + # No cross-request read of the session's blocks... + _, hit_unsalted, _ = manager.get_computed_blocks(unsalted) + assert hit_unsalted == 0 + _, hit_other, _ = manager.get_computed_blocks(other_session) + assert hit_other == 0 + # ...but the session still self-hits (preempt-resume path), capped at + # num_tokens - 1 -> 5 of 6 blocks. + _, hit_self, _ = manager.get_computed_blocks(same_salt) + assert hit_self == 80 diff --git a/tests/v1/streaming/test_reprefill.py b/tests/v1/streaming/test_reprefill.py new file mode 100644 index 000000000000..a6344fe047fa --- /dev/null +++ b/tests/v1/streaming/test_reprefill.py @@ -0,0 +1,860 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for streaming re-prefill at position-range threshold. + +Two test surfaces: + + - `should_trigger_reprefill` (pure check, no side effects). Tested + directly with a synthetic Request. + + - `Scheduler._reprefill_streaming_session` (does the clearing + + re-queue). Tested with a real `KVCacheManager` + `EncoderCacheManager` + plus a minimal scheduler-stub that exposes the attributes the + helper reads (`kv_cache_manager`, `encoder_cache_manager`, + `waiting`). + +The helper is called as an unbound method, mirroring how +`test_eviction.py` calls `Scheduler._record_streaming_segment(None, ...)`. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch + +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + init_none_hash, +) +from vllm.v1.core.sched.request_queue import FCFSRequestQueue +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request, RequestStatus +from vllm.v1.streaming.reprefill import should_trigger_reprefill +from vllm.v1.streaming.retention import ( + HistorySegment, + StreamingRetentionParams, +) + +pytestmark = pytest.mark.cpu_test + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_managers(block_size: int = 16, num_blocks: int = 64): + kv = KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=1024, + enable_caching=False, + scheduler_block_size=block_size, + hash_block_size=block_size, + ) + enc = EncoderCacheManager(cache_size=4096) + return kv, enc + + +def _make_request( + request_id: str, + prompt_token_ids: list[int], + block_size: int = 16, + hash_fn: Callable = sha256, +) -> Request: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + req = Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=sp, + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, hash_fn), + ) + return req + + +class _SchedulerStub: + """Minimum-viable surface for unbound calls to + `Scheduler._reprefill_streaming_session`. The helper reads + `kv_cache_manager`, `encoder_cache_manager`, `waiting`, and + `prev_step_scheduled_req_ids`.""" + + def __init__( + self, + kv_cache_manager: KVCacheManager, + encoder_cache_manager: EncoderCacheManager, + ): + self.kv_cache_manager = kv_cache_manager + self.encoder_cache_manager = encoder_cache_manager + self.waiting = FCFSRequestQueue() + # `_reprefill_streaming_session` clears the request's prior-step + # scheduling bookkeeping so the resumed PREEMPTED request + # doesn't trip the `assert not scheduled_in_prev_step` in + # `_make_cached_request_data`. The stub must expose this + # attribute or the helper crashes with AttributeError. + self.prev_step_scheduled_req_ids: set[str] = set() + # `_free_encoder_inputs` consults `self.is_encoder_decoder` to + # decide between the VLM and enc-dec release paths. Tests of the + # gate-disabled fallback path reach this code; set False so they + # exercise the VLM branch. + self.is_encoder_decoder = False + self.log_stats = False + # Drifted upstream attrs read by _free_request / encoder paths. + self.num_prefill_lookahead = 0 + self._inflight_prefills: set = set() + + +def _bind_reprefill_helper(stub: _SchedulerStub): + """Return a callable invoking `Scheduler._reprefill_streaming_session` + as a bound method on the stub. The helper doesn't reference any + attribute outside what the stub exposes. + + `_reprefill_streaming_session` now calls + `self._reset_streaming_position_state(request)` (scheduler.py refactor: + the position/mRoPE reset shared with preemption-resume was extracted + into its own method). Bind the REAL method onto the stub so the reset + genuinely clears `_mrope_positions`, `max_cached_position`, and + `pending_evicted_token_ranges` exactly as production does (rather than a + hand-written no-op that wouldn't satisfy the clears-position-state + assertions).""" + import types + + from vllm.v1.core.sched.scheduler import Scheduler + + stub._reset_streaming_position_state = types.MethodType( + Scheduler._reset_streaming_position_state, stub + ) + + def _call(request: Request, discard_next_sample: bool = False) -> None: + # discard_next_sample=False models the folded-chunk case (the + # prompt ends with an unanswered chunk, so the first sample is a + # genuine caption token); these tests exercise the reset + # mechanics, which are identical for both values. + Scheduler._reprefill_streaming_session( + stub, request, discard_next_sample=discard_next_sample + ) + + return _call + + +# --------------------------------------------------------------------------- +# should_trigger_reprefill (pure check) +# --------------------------------------------------------------------------- + + +def test_trigger_fires_above_threshold_not_below(): + """`should_trigger_reprefill` returns True only when + `max_cached_position` exceeds `reprefill_threshold * model_max_position`.""" + req = _make_request("r-trigger", [9] * 8) + retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + model_max_position = 1000 + + req.max_cached_position = 699 # just below threshold (700) + assert should_trigger_reprefill(req, retention, model_max_position) is False + + req.max_cached_position = 700 # exactly at threshold — strict > + assert should_trigger_reprefill(req, retention, model_max_position) is False + + req.max_cached_position = 701 # just above threshold + assert should_trigger_reprefill(req, retention, model_max_position) is True + + +def test_trigger_disabled_when_threshold_at_or_above_one(): + """A threshold of 1.0 (or greater) means "disable re-prefill".""" + req = _make_request("r-disabled", [9] * 8) + retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=1.0, + ) + req.max_cached_position = 10_000_000 # arbitrarily large + assert should_trigger_reprefill(req, retention, model_max_position=1_000) is False + + +# --------------------------------------------------------------------------- +# _reprefill_streaming_session (the side-effecting helper) +# --------------------------------------------------------------------------- + + +def _build_request_with_kv_and_history( + request_id: str, total_tokens: int = 80, block_size: int = 16 +): + """Build a request with KV blocks allocated, mRoPE positions + populated, and a non-trivial session_history + mm_features. Returns + `(request, scheduler_stub, _call_reprefill)`.""" + kv, enc = _make_managers(block_size=block_size, num_blocks=64) + req = _make_request(request_id, [9] * total_tokens, block_size=block_size) + req._mrope_positions = [(i, i, i) for i in range(total_tokens)] + req.max_cached_position = total_tokens - 1 + req.num_prompt_tokens = total_tokens + req.num_computed_tokens = total_tokens + + # Populate session_history with one pinned anchor, one video segment, + # and one assistant_text caption. Mirrors what the streaming branch + # would produce after several chunks. age_chunks > 0 on every segment: + # the eviction phases only pick victims with age_chunks > 0 (finding + # #0 — the just-appended, not-yet-computed segment is never evicted), + # and this fixture models already-aged prior chunks, so the mvs=0 + # forced-eviction tests can actually drop the video segment. Ages + # descend oldest-highest to preserve oldest-first eviction order. + req.session_history = [ + HistorySegment( + segment_type="user_text", + token_range=(0, 16), + pinned=True, + age_chunks=4, + ), + HistorySegment( + segment_type="assistant_text", + token_range=(16, 32), + age_chunks=3, + ), + HistorySegment( + segment_type="video", + token_range=(32, 64), + mm_item_id="v0", + age_chunks=2, + ), + HistorySegment( + segment_type="assistant_text", + token_range=(64, 80), + age_chunks=1, + ), + ] + req.mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v0", + mm_position=PlaceholderRange(offset=32, length=32), + ) + ] + manager_blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, total_tokens, 0, manager_blocks) + + # Simulate an in-flight pending_evicted_token_ranges entry that hasn't + # yet been drained to the worker (so we can assert it's cleared). + req.pending_evicted_token_ranges.append((40, 56)) + + # Make sure the request looks like a streaming session. + req.streaming_retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + req.reprefill_count = 0 + + stub = _SchedulerStub(kv, enc) + return req, stub, _bind_reprefill_helper(stub) + + +def test_reprefill_preserves_content_clears_position_state(): + """`_reprefill_streaming_session` must clear K cache, mRoPE state, + num_computed_tokens, and pending evictions — but leave the content + (token ids, session_history, mm_features) untouched.""" + req, stub, call_reprefill = _build_request_with_kv_and_history("r-content-survives") + + # Snapshot content-bearing state before re-prefill. + all_token_ids_before = list(req._all_token_ids) + session_history_before = list(req.session_history) + mm_features_before = list(req.mm_features) + retention_before = req.streaming_retention + + call_reprefill(req) + + # Cleared fields. + assert req._mrope_positions == [] + assert req.max_cached_position == -1 + assert req.num_computed_tokens == 0 + assert req.pending_evicted_token_ranges == [] + assert req.status == RequestStatus.WAITING + + # Surviving content. + assert req._all_token_ids == all_token_ids_before + assert req.session_history == session_history_before + assert req.mm_features == mm_features_before + assert req.streaming_retention is retention_before + + +def test_reprefill_prepends_request_to_waiting_queue(): + """After re-prefill, the request must be at the HEAD of the + waiting queue (so the next scheduler step picks it up first).""" + req, stub, call_reprefill = _build_request_with_kv_and_history("r-queued-head") + + # Pre-populate the waiting queue with another fake request so we + # can verify "prepend" puts ours at the head. + other = _make_request("r-other", [1] * 8) + stub.waiting.add_request(other) + assert list(stub.waiting) == [other] + + call_reprefill(req) + + queued = list(stub.waiting) + assert queued[0] is req, [r.request_id for r in queued] + assert queued[1] is other + + +def test_reprefill_increments_count(): + """`reprefill_count` increments by one each call.""" + req, stub, call_reprefill = _build_request_with_kv_and_history("r-counter") + assert req.reprefill_count == 0 + + call_reprefill(req) + assert req.reprefill_count == 1 + + # Set up state for a second re-prefill (positions repopulated, etc.). + req.max_cached_position = 100 + req.num_computed_tokens = 100 + call_reprefill(req) + assert req.reprefill_count == 2 + + +def test_reprefill_frees_kv_blocks_and_preserves_encoder_cache(): + """The helper must release the request's KV blocks but MUST NOT + touch the encoder cache. Under the persistent-encoder-cache scheme + (Request.uses_persistent_encoder_cache), entries created during + chunk-N's prefill must survive re-prefill so subsequent re-prefill + prefill chunks reuse them without re-running the vision encoder. + + The earlier force-evict behavior was defensive against an + mm_features reassignment bug (now fixed via in-place mutation in + `eviction.py`); this test verifies the defensive logic is gone and + entries persist as designed. + """ + req, stub, call_reprefill = _build_request_with_kv_and_history("r-cache-persists") + kv = stub.kv_cache_manager + enc = stub.encoder_cache_manager + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + # Pre-state: entry is in `cached` with this request as the only + # referencer. + assert enc.cached.get("v0") == {req.request_id} + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == cache_size - 32 + + free_q = kv.block_pool.free_block_queue + free_before = free_q.num_free_blocks + + call_reprefill(req) + + # KV blocks returned to the pool (strict: a zero-free regression must fail). + free_after = free_q.num_free_blocks + assert free_after > free_before, (free_before, free_after) + # Encoder cache entry survives: still in `cached`, still referenced + # by this request, slot accounting unchanged. Worker-side + # encoder_cache[mm_hash] is similarly untouched (verified + # indirectly by the absence of an entry in `freed`). + assert enc.cached.get("v0") == {req.request_id} + assert "v0" not in enc.freeable + assert "v0" not in enc.freed + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == cache_size - 32 + + +def test_reprefill_persists_freeable_encoder_entry(): + """If the encoder entry happens to already be on `freeable` at + re-prefill time (e.g., a prior code path explicitly called + `free_encoder_input`), re-prefill leaves it alone — the entry is + still recoverable via `check_and_update_cache` on the next + scheduling pass, which is exactly what re-prefill prefill needs. + """ + req, stub, call_reprefill = _build_request_with_kv_and_history( + "r-cache-freeable-survives" + ) + enc = stub.encoder_cache_manager + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + enc.free_encoder_input(req, input_id=0) + # Pre-state: cached key still present with empty refs; freeable + # holds the entry; num_freeable_slots restored. + assert "v0" in enc.freeable + assert enc.cached["v0"] == set() + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == cache_size + + call_reprefill(req) + + # Entry survives in both `cached` and `freeable`. Slot accounting + # unchanged from the post-free_encoder_input state. + assert enc.cached.get("v0") == set() + assert "v0" in enc.freeable + assert "v0" not in enc.freed + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == cache_size + + +def test_reprefill_works_with_status_running_or_waiting(): + """The helper must NOT assert status == RUNNING (the way + `_preempt_request` does). Verify it accepts WAITING and + WAITING_FOR_STREAMING_REQ as starting states — those are the + states the streaming path actually fires from.""" + for starting_status in ( + RequestStatus.WAITING, + RequestStatus.WAITING_FOR_STREAMING_REQ, + RequestStatus.RUNNING, + ): + req, stub, call_reprefill = _build_request_with_kv_and_history( + f"r-status-{starting_status.name.lower()}" + ) + req.status = starting_status + call_reprefill(req) + # Always ends in WAITING regardless of where it started. The + # request goes back through the scheduled_new_reqs path so the + # worker re-runs `_update_streaming_request` (refreshing + # `num_prompt_tokens`/`output_token_ids` for the discard mask). + assert req.status == RequestStatus.WAITING + + +# --------------------------------------------------------------------------- +# Persistent encoder cache (encoder-cache-lifetime extension) +# --------------------------------------------------------------------------- + + +def test_uses_persistent_encoder_cache_predicate(): + """Request.uses_persistent_encoder_cache() returns True for + streaming sessions with re-prefill enabled, False otherwise. This + predicate gates `_free_encoder_inputs` and is the single source of + truth for the encoder-cache-lifetime extension.""" + req = _make_request("r-predicate", [9] * 8) + # No streaming retention → False. + assert req.uses_persistent_encoder_cache() is False + + # Streaming retention with threshold < 1.0 → True. + req.streaming_retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + assert req.uses_persistent_encoder_cache() is True + + # Streaming retention with threshold == 1.0 (re-prefill disabled) + # → False (falls back to old per-chunk freeing). + req.streaming_retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=1.0, + ) + assert req.uses_persistent_encoder_cache() is False + + +def test_encoder_cache_persists_across_chunks_for_streaming(): + """Encoder cache entries for a streaming request with persistent + encoder cache enabled must NOT move to `freeable` after the chunk + that scattered them completes. Verifies that `_free_encoder_inputs` + is gated for these sessions and that `evict_segment` (called from + retention) is the sole release path. + + We model this by calling `_free_encoder_inputs` directly on a + request whose mm_feature has been "consumed" by the chunk's + forward pass (i.e., `num_computed_tokens` exceeds the placeholder's + end position). With the gate active the entry stays referenced; + with the gate inactive it would transition to `freeable`. + """ + from vllm.v1.core.sched.scheduler import Scheduler + + req, stub, _ = _build_request_with_kv_and_history( + "r-cache-persists-chunk", total_tokens=80 + ) + enc = stub.encoder_cache_manager + enc.allocate(req, input_id=0) + # mm_feature is at offset=32, length=32 → ends at 64. Set + # num_computed_tokens past that to simulate "chunk done with this + # mm_feature". + req.num_computed_tokens = 80 + + # Streaming + re-prefill enabled (from fixture). + assert req.uses_persistent_encoder_cache() is True + + Scheduler._free_encoder_inputs(stub, req) + + # Gate active: entry remains in cached with our ref, NOT in + # freeable, NOT in freed. + assert enc.cached.get("v0") == {req.request_id} + assert "v0" not in enc.freeable + assert "v0" not in enc.freed + + +def test_encoder_cache_freeing_falls_back_when_reprefill_disabled(): + """If `reprefill_threshold == 1.0` (re-prefill disabled), the gate + in `_free_encoder_inputs` falls through to the original behavior: + once a chunk has consumed its mm_feature, the encoder cache entry + is released to `freeable`.""" + from vllm.v1.core.sched.scheduler import Scheduler + + req, stub, _ = _build_request_with_kv_and_history( + "r-cache-fallback", total_tokens=80 + ) + enc = stub.encoder_cache_manager + enc.allocate(req, input_id=0) + req.num_computed_tokens = 80 + # Disable re-prefill — gate should fall through to old behavior. + assert req.streaming_retention is not None + req.streaming_retention = StreamingRetentionParams( + max_video_segments=req.streaming_retention.max_video_segments, + max_session_tokens=req.streaming_retention.max_session_tokens, + reprefill_threshold=1.0, + ) + assert req.uses_persistent_encoder_cache() is False + + Scheduler._free_encoder_inputs(stub, req) + + # Gate inactive: entry moved to freeable; ref dropped. + assert enc.cached.get("v0") == set() + assert "v0" in enc.freeable + + +def test_segment_eviction_releases_encoder_cache(): + """Under the persistent-encoder-cache scheme, the intra-session + release path for streaming entries is `evict_segment` in + `vllm/v1/streaming/eviction.py`. Since the multi-session OOM fix it + must PHYSICALLY evict a zero-ref entry, not just park it in + `freeable`: the entry leaves `cached` and `freeable`, its mm_hash + reaches the worker exactly once via `get_freed_mm_hashes` (so the + GPU tensor is dropped), and slot accounting is fully restored.""" + from vllm.v1.streaming.eviction import evict_segment + + req, stub, _ = _build_request_with_kv_and_history( + "r-segment-evict", total_tokens=80 + ) + enc = stub.encoder_cache_manager + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + assert enc.cached.get("v0") == {req.request_id} + assert enc.num_free_slots == cache_size - 32 + + # Find the video segment. + video_seg = next(s for s in req.session_history if s.mm_item_id == "v0") + + evict_segment(req, video_seg, stub.kv_cache_manager, enc) + + # Physically evicted: gone from both `cached` and `freeable`. + assert "v0" not in enc.cached + assert "v0" not in enc.freeable + # Slot accounting restored to the pre-allocate value; with + # `freeable` empty the invariant num_freeable_slots == + # num_free_slots + sum(freeable.values()) collapses to equality. + assert enc.num_free_slots == cache_size + assert enc.num_freeable_slots == cache_size + # The worker is told to drop the GPU tensor exactly once; the freed + # list has drain semantics (second call returns nothing). + assert enc.get_freed_mm_hashes() == ["v0"] + assert enc.get_freed_mm_hashes() == [] + # Idempotence: re-evicting the same hash is a pop-miss no-op and + # must not double-credit num_free_slots. + assert enc.evict_unreferenced("v0") is False + assert enc.num_free_slots == cache_size + assert enc.num_freeable_slots == cache_size + + +def test_free_and_evict_releases_persistent_entries(): + """`free_and_evict` (the session close/abort path dispatched by + `Scheduler._free_request` for persistent-encoder-cache requests) + must drop the request's ref AND physically evict the now + unreferenced entry in one call: it leaves `cached`/`freeable`, + lands in `freed`, and the slots are restored.""" + req, stub, _ = _build_request_with_kv_and_history( + "r-free-and-evict", total_tokens=80 + ) + enc = stub.encoder_cache_manager + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + assert req.uses_persistent_encoder_cache() is True + + enc.free_and_evict(req) + + assert "v0" not in enc.cached + assert "v0" not in enc.freeable + assert enc.get_freed_mm_hashes() == ["v0"] + assert enc.get_freed_mm_hashes() == [] + assert enc.num_free_slots == cache_size + assert enc.num_freeable_slots == cache_size + + +def test_evict_segment_multi_ref_entry_survives(): + """If another request still references the mm_hash, `evict_segment` + must only drop THIS session's ref: the entry stays in `cached` with + the survivor's request_id and is neither parked in `freeable` nor + physically freed (`evict_unreferenced` early-returns on the + `freeable.pop` miss).""" + from vllm.v1.streaming.eviction import evict_segment + + req, stub, _ = _build_request_with_kv_and_history( + "r-multi-ref-evict", total_tokens=80 + ) + enc = stub.encoder_cache_manager + cache_size = enc.cache_size + enc.allocate(req, input_id=0) + + # A second request referencing the same cached entry. + other = _make_request("r-multi-ref-survivor", [9] * 8) + other.mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="v0", + mm_position=PlaceholderRange(offset=0, length=32), + ) + ] + assert enc.check_and_update_cache(other, 0) is True + assert enc.cached["v0"] == {req.request_id, other.request_id} + + video_seg = next(s for s in req.session_history if s.mm_item_id == "v0") + evict_segment(req, video_seg, stub.kv_cache_manager, enc) + + # Survivor keeps the entry: still cached, not freeable, not freed. + assert enc.cached.get("v0") == {other.request_id} + assert not enc.freeable + assert enc.get_freed_mm_hashes() == [] + # Slots stay charged for the still-live entry. + assert enc.num_free_slots == cache_size - 32 + assert enc.num_freeable_slots == cache_size - 32 + + +def test_free_request_dispatches_free_and_evict(): + """`Scheduler._free_request` must route persistent-encoder-cache + sessions (close/abort) to `free_and_evict` and everything else to + plain `free`. Uses the unbound-method-on-stub pattern with a mocked + encoder cache manager so only the dispatch is under test.""" + import types + from unittest.mock import Mock + + from vllm.v1.core.sched.scheduler import Scheduler + + for persistent in (True, False): + req, stub, _ = _build_request_with_kv_and_history( + f"r-free-dispatch-{persistent}" + ) + if not persistent: + req.streaming_retention = None + assert req.uses_persistent_encoder_cache() is persistent + req.status = RequestStatus.FINISHED_ABORTED + + # Minimal extra surface `_free_request` reads beyond the stub. + stub.encoder_cache_manager = Mock() + stub._connector_finished = lambda request: (False, None) + stub._free_blocks = types.MethodType(Scheduler._free_blocks, stub) + stub.finished_req_ids = set() + stub.finished_req_ids_dict = None + stub.requests = {req.request_id: req} + # Drifted upstream surface reached by _free_request. + stub.ec_connector = None + stub.defer_block_free = False + stub._pause_state = None + stub._free_request_blocks = types.MethodType( + Scheduler._free_request_blocks, stub + ) + + Scheduler._free_request(stub, req) + + if persistent: + stub.encoder_cache_manager.free_and_evict.assert_called_once_with(req) + stub.encoder_cache_manager.free.assert_not_called() + else: + stub.encoder_cache_manager.free.assert_called_once_with(req) + stub.encoder_cache_manager.free_and_evict.assert_not_called() + # Blocks freed and the request dropped from the table either way. + assert req.request_id in stub.finished_req_ids + assert req.request_id not in stub.requests + + +def test_encoder_cache_size_auto_sized_from_mm_limits(): + """compute_mm_encoder_budget should raise encoder_cache_size when + the per-modality concurrent-items cap implies a larger working + set than the default (= max_num_batched_tokens). This is the + auto-sizing path that keeps persistent encoder cache feasible + without forcing users to manually bump max_num_batched_tokens.""" + from types import SimpleNamespace + + from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget + + scheduler_config = SimpleNamespace( + disable_chunked_mm_input=False, + max_num_batched_tokens=2048, + max_num_encoder_input_tokens=2048, + encoder_cache_size=2048, + ) + mm_max_toks_per_item = {"video": 413} + + # Without limits — falls back to default. + _, cache_size_default = compute_mm_encoder_budget( + scheduler_config, mm_max_toks_per_item + ) + assert cache_size_default == 2048 + + # With limit_mm_per_prompt {video: 60} (sports-config default of + # max_video_segments * 2), required floor = 60 * 413 * 1.3 = 32214. + # Must be >= that. + _, cache_size_streaming = compute_mm_encoder_budget( + scheduler_config, mm_max_toks_per_item, {"video": 60} + ) + assert cache_size_streaming >= 60 * 413, ( + f"cache_size={cache_size_streaming} should accommodate 60 frames " + f"× 413 tokens = {60 * 413}" + ) + # Has the 30% headroom. + assert cache_size_streaming >= (60 * 413 * 13) // 10 + + +# --------------------------------------------------------------------------- +# is_reprefill IPC signal: distinguishes re-prefill from eviction +# --------------------------------------------------------------------------- + + +def test_eviction_alone_does_not_set_reprefill_flag(): + """Intra-session eviction must NOT set `pending_reprefill`. Eviction + is a structured shrink (the worker slices mRoPE preserving surviving + positions' values); re-prefill is a full reset (positions recomputed + from 0 because the K cache was freed). Conflating them broke RoPE + consistency on surviving tokens once the sliding window started + dropping segments.""" + from vllm.v1.streaming.eviction import maybe_evict_old_segments + + req, stub, _ = _build_request_with_kv_and_history( + "r-evict-not-reprefill", total_tokens=80 + ) + # Force eviction by setting max_video_segments to 0 — the one + # video segment in session_history will be dropped. The constructor + # guard now rejects max_video_segments < 1, so build a valid config + # and then degrade it by direct attribute assignment (the dataclass + # is mutable) to keep exercising the mvs=0 eviction path. + req.streaming_retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + req.streaming_retention.max_video_segments = 0 + assert req.pending_reprefill is False + + maybe_evict_old_segments( + req, + req.streaming_retention, + stub.kv_cache_manager, + stub.encoder_cache_manager, + ) + + # Eviction populated `pending_evicted_token_ranges` but DID NOT + # touch `pending_reprefill`. The worker will see evicted ranges + # in the next NewRequestData and slice mRoPE accordingly; the + # reset path stays untouched. + assert req.pending_evicted_token_ranges, "expected eviction to populate ranges" + assert req.pending_reprefill is False + + +def test_new_request_data_drains_reprefill_flag(): + """`NewRequestData.from_request` drains `request.pending_reprefill` + onto `is_reprefill`, then clears the request-side flag so subsequent + re-schedules of the same chunk don't re-fire the worker reset.""" + from vllm.v1.core.sched.output import NewRequestData + + req = _make_request("r-drain-flag", [9] * 8) + req.pending_reprefill = True + + data = NewRequestData.from_request(req, block_ids=([],)) + assert data.is_reprefill is True + assert req.pending_reprefill is False + + # Second drain returns False — flag stays cleared until the next + # `_reprefill_streaming_session` sets it again. + data2 = NewRequestData.from_request(req, block_ids=([],)) + assert data2.is_reprefill is False + + +def test_eviction_then_drain_carries_ranges_not_reprefill(): + """End-to-end check: after eviction, `NewRequestData` carries the + evicted ranges but `is_reprefill` stays False. This is the contract + the worker relies on to slice (not reset) mRoPE positions.""" + from vllm.v1.core.sched.output import NewRequestData + from vllm.v1.streaming.eviction import maybe_evict_old_segments + + req, stub, _ = _build_request_with_kv_and_history( + "r-evict-then-drain", total_tokens=80 + ) + # `_build_request_with_kv_and_history` seeds `pending_evicted_token_ranges` + # with (40, 56) to test re-prefill's clearing behavior — drain that + # here so the test only sees ranges produced by THIS test's eviction. + req.pending_evicted_token_ranges.clear() + # Constructor guard rejects max_video_segments < 1; build a valid + # config and degrade it post-construction to keep exercising the + # mvs=0 forced-eviction path (the dataclass is mutable). + req.streaming_retention = StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + req.streaming_retention.max_video_segments = 0 + + maybe_evict_old_segments( + req, + req.streaming_retention, + stub.kv_cache_manager, + stub.encoder_cache_manager, + ) + assert req.pending_evicted_token_ranges + + data = NewRequestData.from_request(req, block_ids=([],)) + assert data.evicted_token_ranges, "expected ranges to propagate" + assert data.is_reprefill is False, ( + "eviction must not set is_reprefill — that would cause the " + "worker to discard surviving mRoPE positions and recompute " + "them from 0, breaking RoPE consistency with cached K's" + ) + # Drained on the request side too. + assert req.pending_evicted_token_ranges == [] + + +def test_reprefill_clears_pending_ranges_and_sets_flag(): + """Combined re-prefill scenario: when eviction populated pending + ranges and then re-prefill fires in the same handle-stopped pass, + the helper clears the ranges (re-prefill subsumes them: all KV is + freed) and sets the explicit re-prefill flag. The worker then sees + is_reprefill=True with empty evicted_token_ranges and resets mRoPE.""" + from vllm.v1.core.sched.output import NewRequestData + + req, _, call_reprefill = _build_request_with_kv_and_history( + "r-reprefill-clears-ranges" + ) + # Fixture pre-seeds a range; verify it's actually there before re-prefill. + assert req.pending_evicted_token_ranges == [(40, 56)] + + call_reprefill(req) + + assert req.pending_evicted_token_ranges == [] + assert req.pending_reprefill is True + + data = NewRequestData.from_request(req, block_ids=([],)) + assert data.is_reprefill is True + assert not data.evicted_token_ranges # None or empty list diff --git a/tests/v1/streaming/test_streaming_regressions.py b/tests/v1/streaming/test_streaming_regressions.py new file mode 100644 index 000000000000..8d29ddc2784a --- /dev/null +++ b/tests/v1/streaming/test_streaming_regressions.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Regression tests from the streaming deep review (findings C2, C20, C26). + +Scheduler-level tests drive a REAL `Scheduler` (CPU, mocked model_config — +same pattern as tests/v1/streaming_input/test_scheduler_streaming.py) through +`schedule()` / `update_from_output()` so the exact production interleavings +are exercised: + + - C2: close-during-re-prefill must NOT assert-kill the engine — the + finish sentinel popped on the phantom-discard path frees the + request with a proper FINISHED_ABORTED status. + - C26: the phantom-sample discard (`reprefill_discard_next_sample`) drops + exactly one token and emits no EngineCoreOutput; and the re-prefill + TRIGGER path inside `_handle_stopped_request` (idle and folded + variants, including the C19 discard-flag contract). + - C20: `StreamingRetentionParams` rejects reprefill_threshold <= 0 and + NaN; `Scheduler.add_request` rejects retention whose + max_session_tokens would immediately re-trigger re-prefill. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.config import DeviceConfig, VllmConfig +from vllm.sampling_params import SamplingParams +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.outputs import ModelRunnerOutput +from vllm.v1.request import Request, RequestStatus +from vllm.v1.streaming.retention import StreamingRetentionParams +from vllm.v1.structured_output import StructuredOutputManager + +pytestmark = pytest.mark.cpu_test + +STOP_TOKEN = 128001 + +# Trained position range for the fake model. should_trigger_reprefill fires +# when max_cached_position > reprefill_threshold * this (0.7 * 10000 = 7000). +MODEL_MAX_POSITION = 10_000 + + +def _make_retention(**kw) -> StreamingRetentionParams: + kw.setdefault("max_video_segments", 30) + kw.setdefault("max_session_tokens", 4000) + kw.setdefault("reprefill_threshold", 0.7) + return StreamingRetentionParams(**kw) + + +def _create_scheduler() -> Scheduler: + vllm_config = VllmConfig(device_config=DeviceConfig("cpu")) + vllm_config.model_config = MagicMock() + vllm_config.model_config.skip_tokenizer_init = True + vllm_config.model_config.is_multimodal_model = False + # A truthy MagicMock here would trip the encoder-decoder assert in + # Scheduler.__init__ (mm_budget is None for non-multimodal configs). + vllm_config.model_config.is_encoder_decoder = False + vllm_config.model_config.max_model_len = 1024 + vllm_config.model_config.enable_return_routed_experts = False + vllm_config.model_config.uses_mrope = True + # Retention sessions are rejected under async scheduling; this suite + # models the deployed config (async scheduling off). + vllm_config.scheduler_config.async_scheduling = False + # Real int (a bare MagicMock would coerce via __int__ to a garbage + # value); this is the operand of should_trigger_reprefill. + vllm_config.model_config.hf_text_config.max_position_embeddings = MODEL_MAX_POSITION + vllm_config.cache_config = MagicMock() + vllm_config.cache_config.num_gpu_blocks = 1000 + vllm_config.cache_config.enable_prefix_caching = False + kv_cache_config = KVCacheConfig( + num_blocks=1000, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ) + ], + ) + return Scheduler( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, + log_stats=True, + structured_output_manager=StructuredOutputManager(vllm_config), + block_size=16, + hash_block_size=16, + ) + + +def _make_chunk( + request_id: str, + prompt_token_ids: list[int], + *, + resumable: bool = True, + retention: StreamingRetentionParams | None = None, + first_chunk: bool = True, +) -> Request: + extra_args = {"streaming_retention": retention} if retention is not None else None + return Request( + request_id=request_id, + prompt_token_ids=prompt_token_ids, + sampling_params=SamplingParams( + stop_token_ids=[STOP_TOKEN], max_tokens=16, extra_args=extra_args + ), + pooling_params=None, + resumable=resumable, + # The frontend stamps first_chunk on session-opening requests + # (async_llm.handle_inputs); these tests construct them directly. + first_chunk=first_chunk, + ) + + +def _mro(req_id: str, token: int, max_cached_position: int | None = None): + return ModelRunnerOutput( + req_ids=[req_id], + req_id_to_index={req_id: 0}, + sampled_token_ids=[[token]], + logprobs=None, + prompt_logprobs_dict={req_id: None}, + pooler_output=[], + max_cached_positions=( + {req_id: max_cached_position} if max_cached_position is not None else {} + ), + ) + + +def _tokens_in_outputs(eco_dict, req_id: str) -> list[int]: + """All token ids emitted for `req_id` across the returned client dict.""" + tokens: list[int] = [] + for engine_core_outputs in eco_dict.values(): + for output in engine_core_outputs.outputs: + if output.request_id == req_id: + tokens.extend(output.new_token_ids) + return tokens + + +def _drive_session_to_reprefill(scheduler: Scheduler, session: Request): + """Prefill + one decode step, then a stop with a high reported + max_cached_position so `_handle_stopped_request` takes the idle + re-prefill trigger path. Returns the stop step's outputs.""" + rid = session.request_id + scheduler.add_request(session) + + # Prefill the prompt, sample one ordinary caption token. + out = scheduler.schedule() + assert out.num_scheduled_tokens[rid] == len(session.prompt_token_ids) + scheduler.update_from_output(out, _mro(rid, 10)) + assert session.status == RequestStatus.RUNNING + + # Decode step ends the chunk with a stop token; the worker reports a + # position watermark above the 0.7 * MODEL_MAX_POSITION threshold, so + # the stop handler triggers a re-prefill. + out = scheduler.schedule() + eco = scheduler.update_from_output( + out, _mro(rid, STOP_TOKEN, max_cached_position=7500) + ) + return eco + + +# --------------------------------------------------------------------------- +# C2: close-during-re-prefill must not assert-kill the engine +# --------------------------------------------------------------------------- + + +def test_close_during_reprefill_finishes_cleanly(): + """Guards C2: a session-finish sentinel (None) queued while a re-prefill + is in flight is popped on the phantom-discard path of + `update_from_output`; before the fix `_free_request` was reached with the + request still RUNNING and its `assert request.is_finished()` killed the + whole EngineCore. Now the request must finish as FINISHED_ABORTED, leave + the registry, and free its blocks — with no output emitted.""" + scheduler = _create_scheduler() + free_q = scheduler.kv_cache_manager.block_pool.free_block_queue + free_blocks_initial = free_q.num_free_blocks + + session = _make_chunk("sess-c2", [1, 2, 3], retention=_make_retention()) + _drive_session_to_reprefill(scheduler, session) + assert session.pending_reprefill is True + assert session.reprefill_discard_next_sample is True + assert session.status == RequestStatus.WAITING + + # Client closes the session while the re-prefill is pending: the + # non-resumable final request becomes a None sentinel appended to + # streaming_queue (status is WAITING, not WAITING_FOR_STREAMING_REQ). + close_req = _make_chunk("sess-c2", [0], resumable=False) + scheduler.add_request(close_req) + assert list(session.streaming_queue) == [None] + + # Re-prefill executes and samples its one forced token. The discard path + # pops the sentinel; the fixed code must set FINISHED_ABORTED before + # freeing instead of tripping the is_finished() assert (engine death). + out = scheduler.schedule() + assert out.num_scheduled_tokens["sess-c2"] == session.num_tokens + eco = scheduler.update_from_output(out, _mro("sess-c2", 42)) + + assert session.status == RequestStatus.FINISHED_ABORTED + assert "sess-c2" not in scheduler.requests + assert session not in scheduler.running + assert session not in list(scheduler.waiting) + # The phantom sample was discarded, not delivered. + assert 42 not in session._all_token_ids + assert _tokens_in_outputs(eco, "sess-c2") == [] + # All KV blocks returned to the pool. + assert free_q.num_free_blocks == free_blocks_initial + # A stale post-abort chunk (non-first, unknown id) cannot resurrect the + # session: only first chunks may create sessions, by construction. + stale = _make_chunk( + "sess-c2", [7, 8], retention=_make_retention(), first_chunk=False + ) + scheduler.add_request(stale) + assert "sess-c2" not in scheduler.requests + + +def test_stale_chunk_after_close_during_reprefill_is_dropped(): + """Companion to C2: a chunk ADD racing behind the abort must not + resurrect the freed session as a zombie.""" + scheduler = _create_scheduler() + session = _make_chunk("sess-c2b", [1, 2, 3], retention=_make_retention()) + _drive_session_to_reprefill(scheduler, session) + scheduler.add_request(_make_chunk("sess-c2b", [0], resumable=False)) + out = scheduler.schedule() + scheduler.update_from_output(out, _mro("sess-c2b", 42)) + assert "sess-c2b" not in scheduler.requests + + stale = _make_chunk( + "sess-c2b", [7, 8], retention=_make_retention(), first_chunk=False + ) + scheduler.add_request(stale) + assert "sess-c2b" not in scheduler.requests + assert len(scheduler.waiting) == 0 + + +# --------------------------------------------------------------------------- +# C26 (2): phantom-sample discard +# --------------------------------------------------------------------------- + + +def test_phantom_sample_discarded_without_output(): + """Guards C26 gap (2): after an idle-triggered re-prefill the one forced + sample is a phantom — `update_from_output` must drop exactly that token, + emit NO EngineCoreOutput, clear the flag, and park the session waiting + for the next frame (a regression here shifts every caption by one + frame; previously only a 3h soak surfaced it).""" + scheduler = _create_scheduler() + session = _make_chunk("sess-phantom", [1, 2, 3], retention=_make_retention()) + _drive_session_to_reprefill(scheduler, session) + tokens_before = list(session._all_token_ids) + + # Re-prefill runs; the forced sample (42) must be discarded. + out = scheduler.schedule() + eco = scheduler.update_from_output(out, _mro("sess-phantom", 42)) + + assert session.reprefill_discard_next_sample is False + assert session._all_token_ids == tokens_before, "phantom token leaked in" + assert _tokens_in_outputs(eco, "sess-phantom") == [] + # Session is alive, idle, waiting for the next frame. + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + assert scheduler.num_waiting_for_streaming_input == 1 + assert "sess-phantom" in scheduler.requests + + # And the session still accepts the next frame normally. + scheduler.add_request( + _make_chunk("sess-phantom", [7, 8], retention=_make_retention()) + ) + assert session.status == RequestStatus.WAITING + assert list(session._all_token_ids[-2:]) == [7, 8] + assert scheduler.num_waiting_for_streaming_input == 0 + + +# --------------------------------------------------------------------------- +# C26 (3): re-prefill trigger path inside _handle_stopped_request +# --------------------------------------------------------------------------- + + +def test_reprefill_trigger_idle_path_state(): + """Guards C26 gap (3): the trigger integration inside + `_handle_stopped_request` (idle variant): kept output folded into the + prompt, position state reset, request re-queued at the HEAD of waiting, + the waiting-for-input counter balanced, and the phantom-discard flag set + (idle => the forced sample is a phantom).""" + scheduler = _create_scheduler() + session = _make_chunk("sess-trigger", [1, 2, 3], retention=_make_retention()) + eco = _drive_session_to_reprefill(scheduler, session) + + # The chunk's stop was still delivered to the frontend. + stop_tokens = _tokens_in_outputs(eco, "sess-trigger") + assert stop_tokens and stop_tokens[-1] == STOP_TOKEN + + assert session.status == RequestStatus.WAITING + assert list(scheduler.waiting)[0] is session + assert scheduler.num_waiting_for_streaming_input == 0 + assert session.pending_reprefill is True + assert session.reprefill_discard_next_sample is True + assert session.reprefill_count == 1 + # Position state reset for the dense-from-0 re-prefill. + assert session.num_computed_tokens == 0 + assert session.max_cached_position == -1 + # Kept output folded into the prompt; the stop token was discarded. + assert session.prompt_token_ids == [1, 2, 3, 10] + assert STOP_TOKEN not in session.prompt_token_ids + # The folded caption was recorded as an evictable assistant_text segment. + assert any( + seg.segment_type == "assistant_text" and seg.token_range == (3, 4) + for seg in session.session_history + ) + + +def test_reprefill_trigger_folded_path_keeps_first_sample(): + """Guards C26 gap (3) folded variant + the C19 contract: when the stop + handler folds a QUEUED frame before triggering re-prefill, the re-prefilled + prompt ends with that unanswered frame, so `reprefill_discard_next_sample` + must be False and the first post-re-prefill sample must be delivered + (discarding it wedges the session / shifts captions).""" + scheduler = _create_scheduler() + session = _make_chunk("sess-folded", [1, 2, 3], retention=_make_retention()) + rid = session.request_id + scheduler.add_request(session) + + out = scheduler.schedule() + scheduler.update_from_output(out, _mro(rid, 10)) + assert session.status == RequestStatus.RUNNING + + # Pipeline the next frame while the session is still decoding: it lands + # in streaming_queue. + scheduler.add_request(_make_chunk(rid, [7, 8], retention=_make_retention())) + assert len(session.streaming_queue) == 1 + + # Stop with the position watermark above threshold: the handler folds + # the queued frame FIRST, then triggers re-prefill. + out = scheduler.schedule() + scheduler.update_from_output(out, _mro(rid, STOP_TOKEN, max_cached_position=7500)) + + assert session.pending_reprefill is True + assert session.reprefill_discard_next_sample is False, ( + "folded trigger must keep the first post-re-prefill sample: the " + "prompt now ends with the unanswered queued frame" + ) + assert session.status == RequestStatus.WAITING + # Kept output (10) folded, stop discarded, queued frame appended. + assert session.prompt_token_ids == [1, 2, 3, 10, 7, 8] + + # Re-prefill completes: the first sample is frame N+1's first REAL + # caption token and must be appended + emitted. + out = scheduler.schedule() + assert out.num_scheduled_tokens[rid] == session.num_tokens + eco = scheduler.update_from_output(out, _mro(rid, 43)) + assert session._all_token_ids[-1] == 43 + assert _tokens_in_outputs(eco, rid) == [43] + assert session.status == RequestStatus.RUNNING + + +# --------------------------------------------------------------------------- +# C20: retention-params validation +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bad_threshold", [0.0, -0.5, float("nan")]) +def test_retention_rejects_nonpositive_or_nan_threshold(bad_threshold): + """Guards C20: engine-side StreamingRetentionParams must reject + reprefill_threshold <= 0 (immediate re-prefill livelock) and NaN (which + silently disables re-prefill -> unbounded position growth).""" + with pytest.raises(ValueError, match="reprefill_threshold"): + StreamingRetentionParams( + max_video_segments=30, + max_session_tokens=4000, + reprefill_threshold=bad_threshold, + ) + + +def test_retention_accepts_valid_thresholds(): + """C20 companion: in-range values construct; >= 1.0 means 'disable + re-prefill' and stays a warning, not an error.""" + StreamingRetentionParams( + max_video_segments=30, max_session_tokens=4000, reprefill_threshold=0.7 + ) + StreamingRetentionParams( + max_video_segments=30, max_session_tokens=4000, reprefill_threshold=1.0 + ) + + +def test_add_request_rejects_retention_that_would_loop(): + """Guards C20 fix (2): `Scheduler.add_request` must reject a retention + whose max_session_tokens >= reprefill_threshold * model_max_position — + the surviving prompt would re-cross the trigger right after every + re-prefill (GPU-burn livelock). Checked engine-side with the scheduler's + own _model_max_position, the exact trigger operand.""" + scheduler = _create_scheduler() + # 4000 >= 0.2 * 10000 = 2000 -> reject at admission. + bad = _make_chunk( + "sess-loop", + [1, 2, 3], + retention=_make_retention(reprefill_threshold=0.2), + ) + with pytest.raises(ValueError, match="max_session_tokens"): + scheduler.add_request(bad) + assert "sess-loop" not in scheduler.requests + + # The same budget with the default threshold is fine (4000 < 7000). + ok = _make_chunk("sess-ok", [1, 2, 3], retention=_make_retention()) + scheduler.add_request(ok) + assert "sess-ok" in scheduler.requests diff --git a/tests/v1/streaming/test_streaming_session_eviction.py b/tests/v1/streaming/test_streaming_session_eviction.py new file mode 100644 index 000000000000..a86db2365f8d --- /dev/null +++ b/tests/v1/streaming/test_streaming_session_eviction.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end streaming-session eviction tests (no padding). + +Segments are left block-unaligned (no synthetic pad tokens). These tests +drive the real `Scheduler._update_request_as_session` (via a stub) across +many chunk arrivals and verify strict eviction keeps the session leak-free: +the only prompt-vs-owned gap is bounded forward-absorption residue +(<= 2*(block_size-1)), never unbounded orphan drift. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import pytest +import torch + +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalKwargsItem, + PlaceholderRange, +) +from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash +from vllm.v1.core.sched.request_queue import FCFSRequestQueue +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm.v1.streaming.eviction import maybe_evict_old_segments +from vllm.v1.streaming.retention import StreamingRetentionParams + +pytestmark = pytest.mark.cpu_test + +BLOCK = 16 + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def _make_managers(block_size: int = BLOCK, num_blocks: int = 64): + kv = KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=4096, + enable_caching=False, + scheduler_block_size=block_size, + hash_block_size=block_size, + ) + enc = EncoderCacheManager(cache_size=4096) + return kv, enc + + +def _make_streaming_session( + request_id: str, + initial_prompt_tokens: list[int], + retention: StreamingRetentionParams, + block_size: int = BLOCK, + hash_fn: Callable = sha256, +) -> Request: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + req = Request( + request_id=request_id, + prompt_token_ids=initial_prompt_tokens, + sampling_params=sp, + pooling_params=None, + block_hasher=get_request_block_hasher(block_size, hash_fn), + ) + req.streaming_retention = retention + req.status = RequestStatus.WAITING_FOR_STREAMING_REQ + return req + + +class _SchedulerStub: + """Minimum surface for unbound calls to + `Scheduler._update_request_as_session` / `_record_streaming_segment`.""" + + def __init__( + self, + kv_cache_manager: KVCacheManager, + encoder_cache_manager: EncoderCacheManager, + block_size: int = BLOCK, + ): + self.kv_cache_manager = kv_cache_manager + self.encoder_cache_manager = encoder_cache_manager + self.block_size = block_size + self.log_stats = False + # The C29 cumulative context-length guard reads this; 4096 matches + # the file's engine config and stays inert for these sessions. + self.max_model_len = 4096 + self.num_waiting_for_streaming_input = 1 + self.waiting = FCFSRequestQueue() + self.prev_step_scheduled_req_ids: set[str] = set() + self.is_encoder_decoder = False + self.mm_receiver_cache = None + # Drifted upstream attrs read by _free_encoder_inputs/_free_request + # paths reachable from the session-update call. + self.num_prefill_lookahead = 0 + self._inflight_prefills: set = set() + from vllm.v1.core.sched.scheduler import Scheduler + + self._record_streaming_segment = ( + lambda session, update, prior_response_start, segment_start_idx: ( + Scheduler._record_streaming_segment( + self, session, update, prior_response_start, segment_start_idx + ) + ) + ) + + +def _apply_session_update( + stub: _SchedulerStub, session: Request, update: StreamingUpdate +) -> None: + from vllm.v1.core.sched.scheduler import Scheduler + + Scheduler._update_request_as_session(stub, session, update) + + +def _make_update( + new_tokens: list[int], + mm_features: list[MultiModalFeatureSpec] | None = None, +) -> StreamingUpdate: + sp = SamplingParams(max_tokens=1) + sp.update_from_generation_config({}, eos_token_id=100) + return StreamingUpdate( + mm_features=mm_features or [], + prompt_token_ids=new_tokens, + max_tokens=16, + arrival_time=0.0, + sampling_params=sp, + ) + + +def _retention(**kw) -> StreamingRetentionParams: + """Build a StreamingRetentionParams for the session-eviction tests. + + `test_unpadded_session_eviction_no_unbounded_orphan_leak` forces + aggressive eviction with a tiny `max_session_tokens` (16) that + `StreamingRetentionParams.__post_init__` now rejects (it must be >= the + minimum; and max_text_tokens must not exceed max_session_tokens). To + preserve that tiny-budget coverage we build a VALID config and then + degrade the token budgets via direct attribute assignment after + construction (the dataclass is mutable, so this bypasses __post_init__). + """ + base = dict( + max_video_segments=30, + max_text_tokens=4000, + max_session_tokens=4000, + reprefill_threshold=0.7, + ) + # Pull out any caller-supplied token budgets (which may be illegally + # small) and apply them after the valid construction. + requested_mst = kw.pop("max_session_tokens", base["max_session_tokens"]) + requested_mtt = kw.pop("max_text_tokens", base["max_text_tokens"]) + base.update(kw) + base["max_session_tokens"] = 4000 + base["max_text_tokens"] = 4000 + retention = StreamingRetentionParams(**base) + retention.max_session_tokens = requested_mst + retention.max_text_tokens = requested_mtt + return retention + + +def _drive_chunks(stub, session, real_sizes, out_len: int = 2) -> None: + """Apply a sequence of chunks; between chunks append a short 'output' + so the next chunk turns it into an assistant_text segment.""" + for chunk_idx, n_real in enumerate(real_sizes): + update = _make_update( + new_tokens=[20 + chunk_idx] * n_real, + mm_features=[ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier=f"v-{chunk_idx}", + mm_position=PlaceholderRange(offset=0, length=n_real), + ) + ], + ) + _apply_session_update(stub, session, update) + session._all_token_ids.extend([99] * out_len) + session.num_computed_tokens = session.num_prompt_tokens + out_len + + +def _owned(session: Request) -> int: + return sum(e - s for s, e in (seg.token_range for seg in session.session_history)) + + +def test_unpadded_session_is_gapless_and_unaligned(): + """With no padding, the per-chunk segments tile [0, num_prompt_tokens) + with no gaps and no synthetic tokens — and at least one segment is + block-UNALIGNED, proving nothing was padded to a block boundary.""" + kv, enc = _make_managers() + stub = _SchedulerStub(kv, enc) + session = _make_streaming_session( + "r-unpadded", initial_prompt_tokens=[1] * 13, retention=_retention() + ) + session.num_computed_tokens = 13 + _drive_chunks(stub, session, [5, 11, 7, 13, 9]) + + ranges = [s.token_range for s in session.session_history] + assert ranges[0][0] == 0 + for (s0, e0), (s1, e1) in zip(ranges, ranges[1:]): + assert e0 == s1, f"gap/overlap between {(s0, e0)} and {(s1, e1)}" + assert ranges[-1][1] == session.num_prompt_tokens + assert _owned(session) == session.num_prompt_tokens, "no orphans during build" + + lengths = [e - s for s, e in ranges] + assert any(length % BLOCK != 0 for length in lengths), ( + f"expected unpadded (block-unaligned) segments, got lengths {lengths}" + ) + + +def test_unpadded_session_eviction_no_unbounded_orphan_leak(): + """Build an unpadded session over several chunks, then force aggressive + eviction. Strict eviction reduces the session toward budget WITHOUT + PATH C — the only prompt-vs-owned gap is bounded forward-absorption + residue (<= 2*(block_size-1)), never unbounded orphan drift.""" + kv, enc = _make_managers() + stub = _SchedulerStub(kv, enc) + session = _make_streaming_session( + "r-evict", initial_prompt_tokens=[1] * 13, retention=_retention() + ) + session.num_computed_tokens = 13 + _drive_chunks(stub, session, [5, 11, 7, 13, 9, 6, 8]) + + manager_blocks, _, _ = kv.get_computed_blocks(session) + kv.allocate_slots(session, session.num_prompt_tokens, 0, manager_blocks) + + before = session.num_prompt_tokens + tight = _retention(max_video_segments=2, max_session_tokens=16) + n = maybe_evict_old_segments(session, tight, kv, enc) + + assert n > 0, "expected eviction to fire" + assert session.num_prompt_tokens < before, "expected progress toward budget" + gap = session.num_prompt_tokens - _owned(session) + assert 0 <= gap <= 2 * (BLOCK - 1), ( + f"unbounded orphan leak: num_prompt_tokens=" + f"{session.num_prompt_tokens} owned={_owned(session)} gap={gap}" + ) + assert len(session.prompt_token_ids) == session.num_prompt_tokens diff --git a/tests/v1/streaming/test_text_sessions.py b/tests/v1/streaming/test_text_sessions.py new file mode 100644 index 000000000000..540bc99bd7e9 --- /dev/null +++ b/tests/v1/streaming/test_text_sessions.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""1D-RoPE (text) streaming sessions: position-offset bookkeeping. + +Text sessions share the retention/eviction mechanics with mRoPE sessions; +instead of a per-token position array, a text token's RoPE position is its +index plus a per-request scalar offset (the cumulative width of evicted +ranges). These tests pin that bookkeeping: eviction grows the offset by the +block-aligned width, re-prefill resets it, the wire carries it, and the +re-prefill trigger sees the derived watermark. +""" + +import torch + +from vllm.sampling_params import SamplingParams +from vllm.v1.core.encoder_cache_manager import EncoderCacheManager +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, +) +from vllm.v1.request import Request +from vllm.v1.streaming.eviction import evict_segment +from vllm.v1.streaming.reprefill import should_trigger_reprefill +from vllm.v1.streaming.retention import HistorySegment, StreamingRetentionParams + +BLOCK = 4 + + +def _make_managers(num_blocks: int = 64): + kv = KVCacheManager( + KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + FullAttentionSpec( + block_size=BLOCK, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ), + max_model_len=1024, + enable_caching=False, + scheduler_block_size=BLOCK, + hash_block_size=BLOCK, + ) + return kv, EncoderCacheManager(cache_size=4096) + + +def _make_text_session(request_id: str, num_tokens: int) -> Request: + retention = StreamingRetentionParams( + max_video_segments=2, max_session_tokens=2048, reprefill_threshold=0.7 + ) + req = Request( + request_id=request_id, + prompt_token_ids=list(range(num_tokens)), + sampling_params=SamplingParams( + max_tokens=16, extra_args={"streaming_retention": retention} + ), + pooling_params=None, + resumable=True, + first_chunk=True, + ) + req.num_computed_tokens = num_tokens + return req + + +def test_eviction_accumulates_position_offset(): + """Evicting a text segment grows position_offset by exactly the + block-aligned width that was freed.""" + kv, enc = _make_managers() + req = _make_text_session("txt-1", 3 * BLOCK) + blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, 3 * BLOCK, 0, blocks) + assert req.position_offset == 0 + + seg = HistorySegment( + segment_type="user_text", token_range=(0, 2 * BLOCK), age_chunks=3 + ) + req.session_history.append(seg) + evict_segment(req, seg, kv, enc) + + # [0, 8) with block 4 is already aligned: full width evicted. + assert req.position_offset == 2 * BLOCK + assert req.pending_evicted_token_ranges == [(0, 2 * BLOCK)] + + # A second eviction accumulates on top. + req.num_computed_tokens = len(req._all_token_ids) + seg2 = HistorySegment( + segment_type="user_text", token_range=(0, BLOCK), age_chunks=3 + ) + req.session_history.append(seg2) + evict_segment(req, seg2, kv, enc) + assert req.position_offset == 3 * BLOCK + + +def test_eviction_offset_uses_block_aligned_width(): + """A non-aligned segment only frees whole blocks; the offset must grow + by the freed (aligned) width, not the raw segment width — otherwise the + RoPE positions and the token indices desynchronize.""" + kv, enc = _make_managers() + req = _make_text_session("txt-2", 4 * BLOCK) + blocks, _, _ = kv.get_computed_blocks(req) + kv.allocate_slots(req, 4 * BLOCK, 0, blocks) + + # Raw range [1, 2*BLOCK+1) -> inward alignment frees [BLOCK, 2*BLOCK). + seg = HistorySegment( + segment_type="user_text", token_range=(1, 2 * BLOCK + 1), age_chunks=3 + ) + req.session_history.append(seg) + evict_segment(req, seg, kv, enc) + assert req.position_offset == BLOCK + + +def test_new_request_data_carries_position_offset(): + req = _make_text_session("txt-3", 8) + req.position_offset = 12 + data = NewRequestData.from_request(req, block_ids=([],)) + assert data.position_offset == 12 + # And a plain request defaults to 0. + other = _make_text_session("txt-4", 8) + data2 = NewRequestData.from_request(other, block_ids=([],)) + assert data2.position_offset == 0 + + +def test_reprefill_trigger_uses_derived_text_watermark(): + """For text sessions the highest RoPE position is exactly + num_tokens - 1 + position_offset; the trigger must fire on that + derived watermark even though max_cached_position stays -1 (the + worker never reports it for non-mRoPE models).""" + retention = StreamingRetentionParams( + max_video_segments=2, max_session_tokens=2048, reprefill_threshold=0.5 + ) + req = _make_text_session("txt-5", 100) + assert req.max_cached_position == -1 + + model_max_position = 1000 + # Below threshold: 100 - 1 + 300 = 399 <= 500. + assert not should_trigger_reprefill( + req, + retention, + model_max_position, + highest_position=req.num_tokens - 1 + 300, + ) + # Above threshold: 100 - 1 + 450 = 549 > 500. + assert should_trigger_reprefill( + req, + retention, + model_max_position, + highest_position=req.num_tokens - 1 + 450, + ) + # Disabled threshold never fires. + retention_off = StreamingRetentionParams( + max_video_segments=2, max_session_tokens=2048, reprefill_threshold=1.0 + ) + assert not should_trigger_reprefill( + req, retention_off, model_max_position, highest_position=10**9 + ) diff --git a/tests/v1/streaming/test_worker_slice.py b/tests/v1/streaming/test_worker_slice.py new file mode 100644 index 000000000000..6f65a9c0f52c --- /dev/null +++ b/tests/v1/streaming/test_worker_slice.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Worker-side mRoPE streaming helpers on GPUModelRunner. + +Exercises `_slice_mrope_positions_for_evicted_ranges` (drop evicted columns) +and `_extend_mrope_positions_for_streaming_chunk` (append the new chunk's +positions from max_cached_position + 1) in isolation — no real model or +scheduler. Synthetic `CachedRequestState`s are driven through both; the +extend tests use a tiny `SupportsMRoPE` stub via a fake runner that exposes +only `get_model()`. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.v1.worker.gpu_input_batch import CachedRequestState +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +pytestmark = pytest.mark.cpu_test + + +def _make_req_state(num_tokens: int) -> CachedRequestState: + positions = torch.tensor( + [[i for i in range(num_tokens)] for _ in range(3)], + dtype=torch.long, + ) + return CachedRequestState( + req_id="r-slice", + prompt_token_ids=list(range(num_tokens)), + mm_features=[], + sampling_params=None, + generator=None, + block_ids=([],), + num_computed_tokens=num_tokens, + output_token_ids=[], + mrope_positions=positions, + mrope_position_delta=0, + max_cached_position=num_tokens - 1, + ) + + +def _slice_positions(state, ranges): + GPUModelRunner._slice_mrope_positions_for_evicted_ranges(state, ranges) + + +def _row(state): + return state.mrope_positions[0].tolist() + + +def test_single_middle_range(): + state = _make_req_state(16) + _slice_positions(state, [(4, 8)]) + assert state.mrope_positions.shape == (3, 12) + assert _row(state) == [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15] + + +def test_prefix_range(): + state = _make_req_state(8) + _slice_positions(state, [(0, 3)]) + assert state.mrope_positions.shape == (3, 5) + assert _row(state) == [3, 4, 5, 6, 7] + + +def test_suffix_range(): + state = _make_req_state(8) + _slice_positions(state, [(5, 8)]) + assert state.mrope_positions.shape == (3, 5) + assert _row(state) == [0, 1, 2, 3, 4] + + +def test_sequential_ranges_are_applied_post_shrink(): + """Two evictions in one chunk: first slices [4,8), then the SECOND + range refers to the post-first-shrink tensor — applied AFTER the + first range. Validates the sequential-apply contract.""" + state = _make_req_state(16) + # First: drop [4, 8) → tensor becomes [0,1,2,3, 8,9,...,15] (12 cols) + # Second: drop [4, 6) of the SHRUNK tensor → drop [8,9) of original + # → final tensor: [0,1,2,3, 10,11,12,13,14,15] + _slice_positions(state, [(4, 8), (4, 6)]) + assert state.mrope_positions.shape == (3, 10) + assert _row(state) == [0, 1, 2, 3, 10, 11, 12, 13, 14, 15] + + +def test_range_extending_past_width_is_clipped(): + state = _make_req_state(8) + _slice_positions(state, [(5, 100)]) + assert state.mrope_positions.shape == (3, 5) + assert _row(state) == [0, 1, 2, 3, 4] + + +def test_empty_and_invalid_ranges_are_no_op(): + state = _make_req_state(4) + _slice_positions(state, []) + _slice_positions(state, [(2, 2)]) + _slice_positions(state, [(3, 1)]) # invalid order + assert _row(state) == [0, 1, 2, 3] + + +def test_no_op_when_mrope_positions_is_none(): + state = _make_req_state(4) + state.mrope_positions = None + _slice_positions(state, [(0, 2)]) + assert state.mrope_positions is None + + +# --------------------------------------------------------------------------- +# GPUModelRunner._extend_mrope_positions_for_streaming_chunk +# +# Appends positions for the new chunk only (tokens past the prior chunk's +# mrope_positions width), starting at max_cached_position + 1. Driven with a +# tiny SupportsMRoPE stub and a fake runner exposing only get_model(). +# --------------------------------------------------------------------------- + + +class _StubMRoPEModel: + """Minimal SupportsMRoPE model: every token gets linear (p, p, p).""" + + supports_mrope = True + + def get_mrope_input_positions(self, input_tokens, mm_features): + n = len(input_tokens) + positions = ( + torch.tensor([[i, i, i] for i in range(n)], dtype=torch.long) + .t() + .contiguous() + ) + max_val = int(positions.max().item()) if n else -1 + return positions, max_val + 1 - n + + +class _FakeRunner: + """Stands in for GPUModelRunner — provides the only hooks _extend touches: + get_model() (main path) and _init_mrope_positions() (the None fallback, + delegated to the real method, which itself only needs get_model()).""" + + def __init__(self, model): + self._model = model + + def get_model(self): + return self._model + + def _init_mrope_positions(self, req_state): + GPUModelRunner._init_mrope_positions(self, req_state) + + +def _extend(model, state): + GPUModelRunner._extend_mrope_positions_for_streaming_chunk( + _FakeRunner(model), state + ) + + +def _make_extend_state(prev_width: int, num_prompt_tokens: int) -> CachedRequestState: + """A session with `prev_width` cached positions and a prompt grown to + `num_prompt_tokens` (the new chunk is the tail past prev_width).""" + positions = torch.tensor( + [list(range(prev_width)) for _ in range(3)], dtype=torch.long + ) + state = CachedRequestState( + req_id="r-extend", + prompt_token_ids=list(range(num_prompt_tokens)), + mm_features=[], + sampling_params=None, + generator=None, + block_ids=([],), + num_computed_tokens=prev_width, + output_token_ids=[], + mrope_positions=positions, + mrope_position_delta=0, + max_cached_position=prev_width - 1, + ) + # Set by the worker (`_update_streaming_request`) before _extend runs. + state.num_prompt_tokens = num_prompt_tokens + return state + + +def test_extend_appends_new_chunk_above_cached_max(): + state = _make_extend_state(prev_width=4, num_prompt_tokens=7) + _extend(_StubMRoPEModel(), state) + # 3 new columns appended; they start strictly above the prior max (3). + assert state.mrope_positions.shape == (3, 7) + assert _row(state) == [0, 1, 2, 3, 4, 5, 6] + assert min(_row(state)[4:]) > 3 + assert state.max_cached_position == 6 + # Decode continuation delta = new_max + 1 - num_prompt_tokens. + assert state.mrope_position_delta == 0 + + +def test_extend_noop_when_no_new_tokens(): + state = _make_extend_state(prev_width=5, num_prompt_tokens=5) + _extend(_StubMRoPEModel(), state) + assert state.mrope_positions.shape == (3, 5) + assert _row(state) == [0, 1, 2, 3, 4] + assert state.max_cached_position == 4 + + +def test_extend_falls_back_to_full_init_when_positions_missing(): + """No cached positions (e.g. a re-prefill reset null'd them): _extend + must fall back to a full _init_mrope_positions over the whole prompt.""" + state = _make_extend_state(prev_width=4, num_prompt_tokens=6) + state.mrope_positions = None + state.max_cached_position = None + _extend(_StubMRoPEModel(), state) + assert state.mrope_positions.shape == (3, 6) + assert _row(state) == [0, 1, 2, 3, 4, 5] + assert state.max_cached_position == 5 diff --git a/vllm/multimodal/encoder_budget.py b/vllm/multimodal/encoder_budget.py index 3042091611f9..dac3f7753e08 100644 --- a/vllm/multimodal/encoder_budget.py +++ b/vllm/multimodal/encoder_budget.py @@ -124,6 +124,7 @@ def __init__( encoder_compute_budget, encoder_cache_size = compute_mm_encoder_budget( scheduler_config, active_mm_max_toks_per_item, + mm_limits, ) self.encoder_compute_budget = encoder_compute_budget diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index 8d2a81a11b13..287b0fabbea8 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -260,6 +260,38 @@ def free(self, request: Request) -> None: for input_id in list(self.get_cached_input_ids(request)): self.free_encoder_input(request, input_id) + def evict_unreferenced(self, mm_hash: str) -> bool: + """Physically evict an entry whose reference set is empty. + + Entries in `freeable` are by construction unreferenced. Popping one + here appends its mm_hash to `freed`, so the next SchedulerOutput + (via `get_freed_mm_hashes`) tells the model runner to drop the GPU + tensor. Without this, an unreferenced entry stays GPU-resident until + `can_allocate` runs out of `num_free_slots` — which, with the + streaming-inflated cache size, can exceed physical GPU headroom. + + Returns True if the entry existed and was evicted. + """ + num_embeds = self.freeable.pop(mm_hash, None) + if num_embeds is None: + return False + del self.cached[mm_hash] + self.freed.append(mm_hash) + self.num_free_slots += num_embeds + return True + + def free_and_evict(self, request: Request) -> None: + """`free()` plus immediate physical eviction of zero-ref entries. + + Used when a streaming session (persistent encoder cache) finishes or + is aborted: its per-frame entries are content-unique and skipped by + the per-step free path, so parking them in `freeable` retains + ~one vision embedding per pushed frame of GPU memory across sessions. + """ + for input_id in list(self.get_cached_input_ids(request)): + self.free_encoder_input(request, input_id) + self.evict_unreferenced(request.mm_features[input_id].identifier) + def get_freed_mm_hashes(self) -> list[str]: """Get and clear the list of recently freed encoder cache entries. @@ -282,6 +314,7 @@ def get_manager_metadata(self) -> EncoderCacheManagerMetadata | None: def compute_mm_encoder_budget( scheduler_config: "SchedulerConfig", mm_max_toks_per_item: Mapping[str, int], + mm_max_items_per_prompt: Mapping[str, int] | None = None, ) -> tuple[int, int]: """Compute the encoder cache budget based on the model and scheduler configurations for a multimodal model. @@ -326,6 +359,30 @@ def compute_mm_encoder_budget( scheduler_config.encoder_cache_size, max_tokens_per_mm_item ) + # Streaming-aware floor: size the cache to hold all concurrent items at + # once plus 30% headroom. Required by streaming sessions with persistent + # encoder cache, whose retained frame embeddings must all stay resident + # across re-prefill events. + if mm_max_items_per_prompt: + streaming_floor = sum( + mm_max_items_per_prompt.get(m, 0) * tok + for m, tok in mm_max_toks_per_item.items() + ) + streaming_floor = (streaming_floor * 13) // 10 # 30% headroom + if streaming_floor > encoder_cache_size: + logger.info( + "Raising encoder_cache_size from %d to %d to accommodate " + "the limit_mm_per_prompt working set " + "(mm_max_items_per_prompt=%s, mm_max_toks_per_item=%s). " + "Required by streaming sessions with persistent encoder " + "cache.", + encoder_cache_size, + streaming_floor, + dict(mm_max_items_per_prompt), + dict(mm_max_toks_per_item), + ) + encoder_cache_size = streaming_floor + return encoder_compute_budget, encoder_cache_size diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 337a55aa6cd0..db1bab30a4df 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """KV-Cache Utilities.""" +import bisect import copy import hashlib import math @@ -734,12 +735,16 @@ def request_block_hasher(request: Request) -> list[BlockHash]: return [] curr_mm_idx = 0 - if start_token_idx > 0: - # Set curr_mm_idx = -1 to indicate the last mm input. - # Note that since we reach to this branch only when the block is - # completed with generated tokens, we only need to consider the - # last mm input. - curr_mm_idx = -1 + if start_token_idx > 0 and request.mm_features: + # Streaming extends mm_features past the prompt, so the feature + # overlapping a freshly-completed block isn't necessarily the last + # one. Bisect (features are sorted by offset) to the first feature + # whose span still reaches past start_token_idx. + mm_end_positions = [ + f.mm_position.offset + f.mm_position.length for f in request.mm_features + ] + curr_mm_idx = bisect.bisect_right(mm_end_positions, start_token_idx) + curr_mm_idx = min(curr_mm_idx, len(request.mm_features)) prev_block_hash_value = ( request.block_hashes[-1] if request.block_hashes else None diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 977dd34e8cbf..fbbb0447b3f7 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -48,6 +48,16 @@ class NewRequestData: # Only used for v2 model runner. prefill_token_ids: list[int] | None = None + # Streaming: token ranges evicted this step; the worker replays them to + # slice its position arrays in lockstep with the scheduler. + evicted_token_ranges: list[tuple[int, int]] | None = None + # Streaming: True on a re-prefill reset (worker drops position state and + # recomputes positions from 0). + is_reprefill: bool = False + # Streaming, 1D-RoPE (text) sessions: cumulative evicted width. A text + # token's RoPE position is its index plus this offset. + position_offset: int = 0 + @classmethod def from_request( cls, @@ -57,6 +67,15 @@ def from_request( uses_mrope: bool = False, uses_xdrope: bool = False, ) -> "NewRequestData": + # Drain the streaming side-channels the scheduler set on the request. + evicted_token_ranges: list[tuple[int, int]] | None = None + pending = getattr(request, "pending_evicted_token_ranges", None) + if pending: + evicted_token_ranges = list(pending) + pending.clear() + is_reprefill = bool(getattr(request, "pending_reprefill", False)) + if is_reprefill: + request.pending_reprefill = False return cls( req_id=request.request_id, prompt_token_ids=request.prompt_token_ids, @@ -74,6 +93,9 @@ def from_request( prompt_embeds=request.prompt_embeds, prompt_is_token_ids=request.prompt_is_token_ids, prefill_token_ids=prefill_token_ids, + evicted_token_ranges=evicted_token_ranges, + is_reprefill=is_reprefill, + position_offset=getattr(request, "position_offset", 0), ) @property diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 352fc97eb0ef..3c12e38f68eb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -64,6 +64,9 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.streaming.eviction import maybe_evict_old_segments +from vllm.v1.streaming.reprefill import should_trigger_reprefill +from vllm.v1.streaming.retention import HistorySegment, SegmentType from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext @@ -123,6 +126,15 @@ def __init__( else self.scheduler_config.max_num_batched_tokens ) self.max_model_len = vllm_config.model_config.max_model_len + # Model's trained position range; streaming re-prefill compares + # request.max_cached_position against a fraction of this. + self._model_max_position = int( + getattr( + vllm_config.model_config.hf_text_config, + "max_position_embeddings", + self.max_model_len, + ) + ) self.enable_kv_cache_events = ( self.kv_events_config is not None and self.kv_events_config.enable_kv_cache_events @@ -1071,10 +1083,21 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if new_blocks is None: # The request cannot be scheduled. - # NOTE: we need to untouch the request from the encode cache - # manager - if request.has_encoder_inputs: + # Free the request's encoder-cache entries on a failed + # alloc. Skip persistent-encoder-cache streaming sessions: + # their retained frame embeddings must survive a transient + # allocate-fail to avoid a full vision re-encode at + # re-prefill. + if ( + request.has_encoder_inputs + and not request.uses_persistent_encoder_cache() + ): self.encoder_cache_manager.free(request) + + if request.resumable: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue break # KVTransfer: the connector uses this info to determine @@ -1362,6 +1385,64 @@ def _get_new_block_ids_to_zero(self) -> list[int] | None: return new_block_ids_to_zero or None + def _reset_streaming_position_state(self, request: Request) -> None: + """Reset the scheduler-side position-watermark state shared by + re-prefill and preemption-resume.""" + request._mrope_positions = [] + request.max_cached_position = -1 + request.position_offset = 0 + request.pending_evicted_token_ranges.clear() + + def _reprefill_streaming_session( + self, request: Request, discard_next_sample: bool + ) -> None: + """Self-preempt for re-prefill: free the KV cache + reset position + state, prepend to waiting, and let the next step re-prefill the + surviving `_all_token_ids` at fresh dense-from-0 positions. + + `discard_next_sample` must be True only when the re-prefilled prompt + ends with already-answered content (idle trigger), so the one forced + sample is a phantom. When a queued chunk was just folded, the prompt + ends with that unanswered chunk and the first sample is a genuine + reply token that must be kept. + """ + prior_max_position = request.max_cached_position # for the log line + prior_num_tokens = request.num_tokens + + # Gather block IDs BEFORE freeing, then evict their prefix-cache + # entries so the stale content cannot be served to anyone. + reprefill_block_ids: set[int] = set() + for per_group_ids in self.kv_cache_manager.get_block_ids(request.request_id): + reprefill_block_ids.update(per_group_ids) + + self.kv_cache_manager.free(request) + if reprefill_block_ids: + self.kv_cache_manager.evict_blocks(reprefill_block_ids) + + request.num_computed_tokens = 0 + self._reset_streaming_position_state(request) + if request.spec_token_ids: + request.spec_token_ids = [] + + request.status = RequestStatus.WAITING + + self.prev_step_scheduled_req_ids.discard(request.request_id) + request.reprefill_count += 1 + + request.reprefill_discard_next_sample = discard_next_sample + + request.pending_reprefill = True + logger.info( + "streaming re-prefill triggered: req=%s " + "prior_max_position=%d surviving_tokens=%d count=%d", + request.request_id, + prior_max_position, + prior_num_tokens, + request.reprefill_count, + ) + + self.waiting.prepend_request(request) + def _preempt_request( self, request: Request, timestamp: float, drop_stale_output: bool = False ) -> None: @@ -1379,7 +1460,8 @@ def _preempt_request( "Only running requests can be preempted" ) self._free_request_blocks(request) - self.encoder_cache_manager.free(request) + if not request.uses_persistent_encoder_cache(): + self.encoder_cache_manager.free(request) self._inflight_prefills.discard(request) request.status = RequestStatus.PREEMPTED request.num_computed_tokens = 0 @@ -1398,6 +1480,9 @@ def _preempt_request( request.num_stale_output_tokens = request.num_in_flight_tokens request.num_output_placeholders = 0 request.num_preemptions += 1 + + if request.streaming_retention is not None: + self._reset_streaming_position_state(request) if self.log_stats: request.record_event(EngineCoreEventType.PREEMPTED, timestamp) @@ -1470,12 +1555,18 @@ def _update_request_as_session( kept_output_tokens = session._all_token_ids[ session.num_prompt_tokens : num_computed_tokens ] + + prior_response_start = session.num_prompt_tokens + del session._all_token_ids[num_computed_tokens:] session._output_token_ids.clear() assert session.prompt_token_ids is not None + # Extend prompt with kept output tokens. session.prompt_token_ids.extend(kept_output_tokens) + segment_start_idx = prior_response_start + len(kept_output_tokens) + if update.mm_features: base = session.num_tokens for mm_feature in update.mm_features: @@ -1486,6 +1577,17 @@ def _update_request_as_session( session._all_token_ids.extend(update.prompt_token_ids or ()) session.prompt_token_ids.extend(update.prompt_token_ids or ()) + + # Drop hashes covering the discarded sampled token (if it completed a + # block) so the chain is rebuilt from the new content; a stale hash + # would publish the block under the wrong prefix-cache key. Must run + # for retention-less sessions too, and at the HASH block granularity + # (single_type_managers[0].block_size need not match it on hybrid + # models, which non-retention resumable sessions may run on). + if session.block_hashes: + hash_block_size = self.kv_cache_manager.block_pool.hash_block_size + del session.block_hashes[num_computed_tokens // hash_block_size :] + # Update block hashes for the new tokens. session.update_block_hashes() session.num_prompt_tokens = len(session.prompt_token_ids) @@ -1495,9 +1597,155 @@ def _update_request_as_session( self.num_waiting_for_streaming_input -= 1 session.status = RequestStatus.WAITING + if ( + session.structured_output_request is not None + and session.structured_output_request.grammar is not None + ): + session.structured_output_request.grammar.reset() + + if session.streaming_retention is not None: + self._record_streaming_segment( + session, update, prior_response_start, segment_start_idx + ) + maybe_evict_old_segments( + session, + session.streaming_retention, + self.kv_cache_manager, + self.encoder_cache_manager, + ) + + if session.num_tokens >= self.max_model_len: + # Cumulative context-length guard: even after eviction the + # session prompt no longer fits the context window. Worker + # per-request buffers are max_model_len-sized, so scheduling it + # would crash the engine. Finish the session cleanly instead. + logger.error( + "Streaming session %s reached max_model_len (%d tokens >= " + "%d) after chunk append; finishing session.", + session.request_id, + session.num_tokens, + self.max_model_len, + ) + self.finish_requests( + session.request_id, RequestStatus.FINISHED_LENGTH_CAPPED + ) + return + if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) + def _fold_kept_output(self, session: Request) -> None: + """Fold the current chunk's generated output into the prompt, + discarding the final sampled token.""" + num_computed_tokens = session.num_computed_tokens + kept_output_tokens = session._all_token_ids[ + session.num_prompt_tokens : num_computed_tokens + ] + del session._all_token_ids[num_computed_tokens:] + session._output_token_ids.clear() + assert session.prompt_token_ids is not None + session.prompt_token_ids.extend(kept_output_tokens) + + # See _update_request_as_session: trim stale hashes unconditionally, + # at the hash block granularity. + if session.block_hashes: + hash_block_size = self.kv_cache_manager.block_pool.hash_block_size + del session.block_hashes[num_computed_tokens // hash_block_size :] + session.update_block_hashes() + session.num_prompt_tokens = len(session.prompt_token_ids) + + def _record_folded_output_segment( + self, session: Request, prior_response_start: int + ) -> None: + """Record the just-folded last reply as an evictable assistant_text + segment.""" + end = session.num_prompt_tokens + if end <= prior_response_start: + return # the folded chunk generated no kept output. + session.session_history.append( + HistorySegment( + segment_type="assistant_text", + token_range=(prior_response_start, end), + age_chunks=0, + ) + ) + + def _record_streaming_segment( + self, + session: Request, + update: StreamingUpdate, + prior_response_start: int, + segment_start_idx: int, + ) -> None: + """Append HistorySegment(s) for the just-appended chunk and age + all prior segments by one chunk. + + This function creates up to three segments: + + 1. **Pinned anchor** — on the very first call only (when + `session_history` is empty). Covers `[0, prior_response_start)`, + i.e., chunk 1's prompt only. Chunk 1's reply is split out + as a separate unpinned `assistant_text` segment. The + anchor is pinned and represents the attention sink (system + prompt + first user turn + first chunk content). + + 2. **assistant_text for the prior chunk's response** — if + `segment_start_idx > prior_response_start`. Unpinned, evictable + by retention phase 2, so replies are retained independently of + the content they describe. + + 3. **New chunk's content segment** — typed `"video"` if the + chunk added `mm_features`, else `"user_text"`. Its + corresponding response segment will be created on the NEXT + chunk's call. + + Aging order: synthesize the anchor first (so it gets aged like + any existing segment), then age all existing segments, then + append the two new segments (which start at age 0). + """ + # 1. Synthesize the pinned anchor on the very first call. + if not session.session_history and prior_response_start > 0: + session.session_history.append( + HistorySegment( + segment_type="user_text", + token_range=(0, prior_response_start), + pinned=True, + ) + ) + + # Age existing segments by one chunk boundary. Anything we + # append below will start at age 0. + for seg in session.session_history: + seg.age_chunks += 1 + + # 2. assistant_text segment for the prior chunk's response. + if segment_start_idx > prior_response_start: + session.session_history.append( + HistorySegment( + segment_type="assistant_text", + token_range=(prior_response_start, segment_start_idx), + ) + ) + + # 3. New chunk's content segment. + new_num_tokens = session.num_tokens + if new_num_tokens <= segment_start_idx: + return # chunk added no new tokens. + segment_type: SegmentType + mm_item_id: str | None = None + if update.mm_features: + segment_type = "video" + mm_item_id = update.mm_features[0].identifier + else: + segment_type = "user_text" + session.session_history.append( + HistorySegment( + segment_type=segment_type, + token_range=(segment_start_idx, new_num_tokens), + mm_item_id=mm_item_id, + ) + ) + def _make_cached_request_data( self, running_reqs: list[Request], @@ -1771,6 +2019,11 @@ def update_from_output( pooler_outputs = model_runner_output.pooler_output num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + + for req_id, pos in model_runner_output.max_cached_positions.items(): + req = self.requests.get(req_id) + if req is not None and pos > req.max_cached_position: + req.max_cached_position = pos ec_connector_output = model_runner_output.ec_connector_output cudagraph_stats = model_runner_output.cudagraph_stats @@ -1912,6 +2165,26 @@ def update_from_output( status_before_stop = request.status num_output_tokens_before = len(request._output_token_ids) + if request.reprefill_discard_next_sample and new_token_ids: + # The one forced sample completing an idle-triggered + # re-prefill is a phantom reply: drop it instead of + # delivering it (it would shift every later reply off by + # one chunk), then park/fold the session as usual. + request.reprefill_discard_next_sample = False + finished = self._handle_stopped_request(request) + if finished: + # The finish sentinel was popped without check_stop ever + # setting a finished status (the request is still + # RUNNING); set one before _free_request, which asserts + # is_finished(). + request.status = RequestStatus.FINISHED_ABORTED + self._free_request(request) + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + continue + # Check for stop and update request status. if new_token_ids: new_token_ids, stopped = self._update_request_with_output( @@ -1951,15 +2224,29 @@ def update_from_output( if advance_token_ids and not grammar.accept_tokens( req_id, advance_token_ids ): - logger.error( - "Unexpected: grammar rejected tokens %s for request %s. " - "Terminating request.", - advance_token_ids, - req_id, - ) - request.status = RequestStatus.FINISHED_ERROR - request.resumable = False - stopped = True + if request.resumable: + # Streaming session soft stop: cut this chunk's + # response short but keep the session alive; its + # context is still valid. _update_request_as_session + # resets the grammar FSM when the next chunk arrives. + logger.warning( + "Grammar rejected tokens %s for streaming request " + "%s; cutting chunk short, session continues.", + advance_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_STOPPED + stopped = True + else: + logger.error( + "Unexpected: grammar rejected tokens %s for " + "request %s. Terminating request.", + advance_token_ids, + req_id, + ) + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True routed_experts = None if ( @@ -2214,16 +2501,60 @@ def _handle_stopped_request(self, request: Request) -> bool: if not request.resumable: return True + if ( + request.streaming_retention is None + and request.status == RequestStatus.FINISHED_LENGTH_CAPPED + ): + # Without retention nothing bounds the session prompt; keeping + # the session alive would let the next chunk push num_tokens past + # max_model_len and overflow max_model_len-sized worker buffers. + return True + + output_folded = False if request.streaming_queue: update = request.streaming_queue.popleft() if update is None: # Streaming request finished. return True self._update_request_as_session(request, update) + if request.is_finished(): + # The fold hit the cumulative context-length guard and the + # session was finished (and freed) inside + # _update_request_as_session. Return False so the caller + # does not free it a second time. + return False + output_folded = True else: request.status = RequestStatus.WAITING_FOR_STREAMING_REQ self.num_waiting_for_streaming_input += 1 + # Streaming re-prefill: if the session's positions have approached + # the model's trained range, self-preempt and re-queue for a fresh + # full prefill at positions starting from 0. + if request.streaming_retention is not None and should_trigger_reprefill( + request, + request.streaming_retention, + self._model_max_position, + # mRoPE models report the watermark from the worker; 1D-RoPE + # (text) sessions derive it exactly: index + cumulative offset. + highest_position=max( + request.max_cached_position, + request.num_tokens - 1 + request.position_offset, + ), + ): + if not output_folded: + prior_response_start = request.num_prompt_tokens + self._fold_kept_output(request) + self._record_folded_output_segment(request, prior_response_start) + + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + self._reprefill_streaming_session( + request, discard_next_sample=not output_folded + ) + # The helper already prepended to the waiting queue. + return False + self._enqueue_waiting_request(request) return False @@ -2247,6 +2578,8 @@ def _update_request_with_output( return new_token_ids, stopped def _free_encoder_inputs(self, request: Request) -> None: + if request.uses_persistent_encoder_cache(): + return cached_encoder_input_ids = self.encoder_cache_manager.get_cached_input_ids( request ) @@ -2361,6 +2694,43 @@ def add_request(self, request: Request) -> None: # Streaming-input session finished. self.finish_requests(request.request_id, RequestStatus.FINISHED_ABORTED) else: + if request.streaming_retention is not None: + if len(self.kv_cache_config.kv_cache_groups) != 1: + raise ValueError( + "Streaming retention (intra-session KV eviction) " + "requires exactly one KV cache group, but this model " + f"has {len(self.kv_cache_config.kv_cache_groups)}. " + "Hybrid / attention-free models are not supported." + ) + retention = request.streaming_retention + if ( + retention.max_session_tokens + >= retention.reprefill_threshold * self._model_max_position + ): + raise ValueError( + "Streaming retention max_session_tokens " + f"({retention.max_session_tokens}) must stay below " + "reprefill_threshold * model_max_position " + f"({retention.reprefill_threshold} * " + f"{self._model_max_position}); otherwise a re-prefill " + "would re-trigger immediately and loop." + ) + if self.scheduler_config.async_scheduling: + raise ValueError( + "Streaming retention is incompatible with async " + "scheduling. Launch with --no-async-scheduling." + ) + if self.use_pp: + raise ValueError( + "Streaming retention is incompatible with pipeline " + "parallelism (step overlap). Use pipeline_parallel_" + "size=1." + ) + if self.vllm_config.parallel_config.data_parallel_size > 1: + raise ValueError( + "Streaming retention currently requires " + "data_parallel_size == 1." + ) if request.resumable and not request.first_chunk: # A non-first chunk for an unknown id: its session is gone. # Admitting it would resurrect the session as an unabortable @@ -2464,7 +2834,10 @@ def _free_request( ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request) connector_delay_free_blocks |= ec_delay_free - self.encoder_cache_manager.free(request) + if request.uses_persistent_encoder_cache(): + self.encoder_cache_manager.free_and_evict(request) + else: + self.encoder_cache_manager.free(request) request_id = request.request_id self.finished_req_ids.add(request_id) if self.finished_req_ids_dict is not None: diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 80b909afbc8d..e2a3c21631ce 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -345,6 +345,10 @@ class ModelRunnerOutput: # information related to cudagraph execution cudagraph_stats: CUDAGraphStat | None = None + # Streaming: worker-assigned `max_cached_position` per request whose + # position state was (re)initialized or extended this step. + max_cached_positions: dict[str, int] = field(default_factory=dict) + # Per-step routed experts data captured by the worker. # ``routing_data`` shape: (num_scheduled_tokens, num_layers, # num_experts_per_tok); expert IDs as uint8/uint16. diff --git a/vllm/v1/request.py b/vllm/v1/request.py index 4566165a8ee0..891a12b9fb3e 100644 --- a/vllm/v1/request.py +++ b/vllm/v1/request.py @@ -21,6 +21,7 @@ FinishReason, ) from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics +from vllm.v1.streaming.retention import HistorySegment, StreamingRetentionParams from vllm.v1.structured_output.request import StructuredOutputRequest from vllm.v1.utils import ConstantList @@ -102,6 +103,10 @@ def __init__( # P/D: Connector-specific KV transfer parameters. self.kv_transfer_params: dict[str, Any] | None = None + # Per-session retention config (read from + # `sampling_params.extra_args["streaming_retention"]` below); + # preserved across chunks by `_update_request_as_session`. + self.streaming_retention: StreamingRetentionParams | None = None # E/P/D: Connector-specific encoder-cache transfer parameters. self.ec_transfer_params: dict[str, Any] | None = None @@ -119,6 +124,12 @@ def __init__( self.kv_transfer_params = sampling_params.extra_args.get( "kv_transfer_params" ) + _retention = sampling_params.extra_args.get("streaming_retention") + if isinstance(_retention, StreamingRetentionParams): + self.streaming_retention = _retention # in-process API + elif isinstance(_retention, dict): + # Cross-process IPC: arrives as a plain dict. + self.streaming_retention = StreamingRetentionParams(**_retention) self.ec_transfer_params = sampling_params.extra_args.get( "ec_transfer_params" ) @@ -157,6 +168,23 @@ def __init__( for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids) ] + # Per-token positions, in lockstep with `_all_token_ids`. Populated + # lazily by the model runner on first prefill; extended per streaming + # chunk. Stored as 3-tuples: mRoPE models use all three coordinates, + # 1D-RoPE (text) models use coordinate 0 (vLLM's mRoPE convention + # already sets t == h == w for text tokens, so 1D is the degenerate + # case of the same array). `max_cached_position` is the highest + # position number in the KV cache; eviction never decrements it, so + # positions may have gaps after eviction. + self._mrope_positions: list[tuple[int, int, int]] = [] + self.max_cached_position: int = -1 + # 1D-RoPE (text) sessions: cumulative width of evicted token ranges. + # A text token's RoPE position is its index plus this offset, so a + # scalar replaces the full per-token array (the gaps are always + # "everything after the holes shifts up by the total hole width"). + # Reset to 0 on re-prefill; incremented by eviction. + self.position_offset: int = 0 + # Used in async scheduling. self.num_output_placeholders = 0 # Tokens of output in flight when the request was preempted: delivered @@ -232,6 +260,30 @@ def __init__( # None entry in the queue means finished. self.streaming_queue: deque[StreamingUpdate | None] | None = None + # Per-session segment bookkeeping: filled by the scheduler as chunks + # arrive, consumed by `streaming.eviction` when over budget. + self.session_history: list[HistorySegment] = [] + # Token ranges evicted scheduler-side since the last worker update, + # each in the _all_token_ids index space at eviction time. Drained + # into `NewRequestData.evicted_token_ranges` and replayed on the + # worker to slice its position arrays in lockstep. + self.pending_evicted_token_ranges: list[tuple[int, int]] = [] + + # Re-prefill bookkeeping. `reprefill_count` is observability only. + # `pending_reprefill` is set by `_reprefill_streaming_session` and + # drained into `NewRequestData.is_reprefill`, telling the worker to + # reset positions to 0 (vs. eviction, which only slices them). + self.reprefill_count: int = 0 + self.pending_reprefill: bool = False + # A re-prefill must sample >=1 token to complete. When the trigger + # found the session idle between chunks, that token is a phantom + # reply and this tells `update_from_output` to discard it + # (delivering it would shift every later reply off by one chunk). + # NOT set when the trigger folded a queued chunk first: the + # re-prefilled prompt then ends with that unanswered chunk, so the + # first sample is a genuine reply token that must be kept. + self.reprefill_discard_next_sample: bool = False + # If True, request should be aborted immediately after being added to # the scheduler so the connector's request_finished hook runs. self.abort_immediately = abort_immediately @@ -326,6 +378,20 @@ def is_finished(self) -> bool: def get_finished_reason(self) -> FinishReason | None: return RequestStatus.get_finished_reason(self.status) + def uses_persistent_encoder_cache(self) -> bool: + """True if encoder cache entries should live until their segment is + evicted, rather than be freed after each chunk's forward pass. + + Re-prefill-enabled sessions need the surviving embeddings available at + the next re-prefill so the encoder need not re-run on every retained + frame. A `reprefill_threshold >= 1.0` (disabled) falls back to + per-chunk freeing. + """ + return ( + self.streaming_retention is not None + and self.streaming_retention.reprefill_threshold < 1.0 + ) + def get_num_encoder_embeds(self, input_id: int) -> int: assert input_id < len(self.mm_features) return self.mm_features[input_id].mm_position.get_num_embeds() diff --git a/vllm/v1/streaming/__init__.py b/vllm/v1/streaming/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/v1/streaming/eviction.py b/vllm/v1/streaming/eviction.py new file mode 100644 index 000000000000..b93e6eac7e7b --- /dev/null +++ b/vllm/v1/streaming/eviction.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Intra-session KV eviction for streaming sessions. + +A sliding-window policy drops the oldest video segments from a live +session once the segment count exceeds the retention budget, keeping +long-running captioning within its KV budget. + +`evict_segment` is the unit step: it frees the segment's KV blocks, +drops its encoder cache entry, and removes the segment's tokens and +mRoPE positions from the request's lockstep arrays. Surviving segments' +`token_range` fields shift down by the evicted length so they stay valid +against the compacted arrays. + +mRoPE position VALUES are NOT shifted: surviving tokens keep their +original coordinates, so the position space develops gaps. That is what +keeps cached K/V rotationally consistent — each surviving token's K was +rotated under its original position, which is still what RoPE pairs with +the new chunk's queries. +""" + +from __future__ import annotations + +import contextlib +from dataclasses import replace +from typing import TYPE_CHECKING + +from vllm.v1.streaming.retention import HistorySegment, StreamingRetentionParams + +if TYPE_CHECKING: + from vllm.v1.core.encoder_cache_manager import EncoderCacheManager + from vllm.v1.core.kv_cache_manager import KVCacheManager + from vllm.v1.request import Request + + +def evict_segment( + request: Request, + segment: HistorySegment, + kv_cache_manager: KVCacheManager, + encoder_cache_manager: EncoderCacheManager, +) -> None: + """Drop one history segment from a live session. + + `segment` may be a single segment OR a packed blob from + `_pack_oldest_run_to_free_block`; the logic keys off the token range, so + it handles either. All effects run in lockstep, so the caller must have + validated the victim already: + + 1. Free the segment's KV blocks (back to the free pool). + 2. Release the encoder entry and drop the `MultiModalFeatureSpec` for + every mm item in the range (one per video, several per blob, none + for text). + 3. Delete the segment's slice of `_all_token_ids` / `_mrope_positions`. + 4. Shift later features' `mm_position.offset` down. + 5. Shift later segments' `token_range` down. + 6. Decrement `num_prompt_tokens` and `num_computed_tokens`. + 7. Remove the segment from `session_history`. + """ + if segment.pinned: + raise ValueError( + f"refusing to evict pinned segment {segment.segment_type} " + f"at {segment.token_range}" + ) + raw_start, raw_end = segment.token_range + if raw_end <= raw_start: + # Zombie: neighbour-clipping shrank it to nothing; just retire it. + with contextlib.suppress(ValueError): + request.session_history.remove(segment) + return + + # Clamp to the computed frontier before freeing (see Invariants); must + # precede `evict_token_range_for_request`, which frees blocks eagerly. + raw_end = min(raw_end, request.num_computed_tokens) + if raw_end <= raw_start: + return # nothing below the frontier is evictable + + # 1. Free KV blocks; the manager returns the block-aligned range freed. + start, end, _ = kv_cache_manager.evict_token_range_for_request( + request.request_id, raw_start, raw_end + ) + evicted_len = end - start + if evicted_len <= 0: + return # freed no whole block; leave it (the caller pre-filters) + + # Worker replays these ranges to slice mrope_positions, in arrival order. + request.pending_evicted_token_ranges.append((start, end)) + # 1D-RoPE (text) sessions: survivors shift down by the evicted width but + # keep their original RoPE positions, so the index->position offset grows. + request.position_offset += evicted_len + + # 2. Free encoder entries + drop mm_features overlapping the freed range + # (one per video, several per blob). Reverse + in-place: the worker shares + # this list, and front-to-back deletion would shift unvisited indices. + for input_id in range(len(request.mm_features) - 1, -1, -1): + feature = request.mm_features[input_id] + f_offset = feature.mm_position.offset + f_end = f_offset + feature.mm_position.length + if not (f_offset < end and f_end > start): + continue + encoder_cache_manager.free_encoder_input(request, input_id) + encoder_cache_manager.evict_unreferenced(feature.identifier) + del request.mm_features[input_id] + # Engine receiver cache is left alone (streaming forces mm cache off). + + # 3. Compact the lockstep token arrays by the block-aligned range. + del request._all_token_ids[start:end] + if len(request._mrope_positions) >= end: + del request._mrope_positions[start:end] + + # 4. Shift later mm_features' offsets down by the evicted length. + for feature in request.mm_features: + if feature.mm_position.offset >= end: + feature.mm_position = replace( + feature.mm_position, + offset=feature.mm_position.offset - evicted_len, + ) + + # 5. Shift later segments down, trim neighbours the freed range cut into, + # and absorb the edge orphans into a neighbour (forward normally, backward + # if the victim was the last segment). + zombies_to_drop: list[HistorySegment] = [] + next_after_evicted: HistorySegment | None = None + prev_before_evicted: HistorySegment | None = None + for other in request.session_history: + if other is segment: + continue + o_start, o_end = other.token_range + if o_start >= end and ( + next_after_evicted is None or o_start < next_after_evicted.token_range[0] + ): + next_after_evicted = other + if ( + not other.pinned + and o_end <= start + and ( + prev_before_evicted is None + or o_end > prev_before_evicted.token_range[1] + ) + ): + prev_before_evicted = other + + for other in request.session_history: + if other is segment: + continue + o_start, o_end = other.token_range + if o_start >= end: + new_start = o_start - evicted_len + if other is next_after_evicted: + new_start = raw_start # absorb the orphans against the victim + other.token_range = (new_start, o_end - evicted_len) + elif o_end > start: + # Freed range cut into this neighbour; trim it (drop if empty). + new_start = min(o_start, start) + new_end_calc = o_end - evicted_len + if new_end_calc <= new_start: + zombies_to_drop.append(other) + else: + other.token_range = (new_start, new_end_calc) + + # Backward absorption: no following segment, so extend the preceding one + # over the residue tail (ends at start + (raw_end - end)). + if next_after_evicted is None and prev_before_evicted is not None: + residue_end = start + (raw_end - end) + p_start, p_end = prev_before_evicted.token_range + if residue_end > p_end: + prev_before_evicted.token_range = (p_start, residue_end) + + for z in zombies_to_drop: + with contextlib.suppress(ValueError): + request.session_history.remove(z) + + # 6. Decrement counters. + request.num_prompt_tokens = max(0, request.num_prompt_tokens - evicted_len) + request.num_computed_tokens = max(0, request.num_computed_tokens - evicted_len) + # `all_token_ids` is a view onto `_all_token_ids`, so it tracks already. + if request.prompt_token_ids is not None: + del request.prompt_token_ids[start : min(end, len(request.prompt_token_ids))] + + # Trim block_hashes to match the compacted block table. NOTE: this only + # splices the list; the surviving hash VALUES still chain over the evicted + # content. Streaming sessions therefore set cache_salt (async_llm.py) so + # these blocks are never shareable with other requests. + bs = kv_cache_manager.coordinator.single_type_managers[0].block_size + del request.block_hashes[start // bs : end // bs] + + # 7. Remove the segment from history. + request.session_history.remove(segment) + + +_TEXT_SEGMENT_TYPES = ("user_text", "assistant_text") + + +def _block_size(kv_cache_manager: KVCacheManager) -> int: + """Smallest block size across kv-cache groups. The eviction primitive + rounds INWARD to whole blocks, so a segment shorter than this can't + free any KV — phase 2 uses this to filter unevictable text segments.""" + managers = kv_cache_manager.coordinator.single_type_managers + return min(m.block_size for m in managers) + + +def _frees_block(token_range: tuple[int, int], block_size: int) -> bool: + """True if inward block-rounding of `[start, end)` frees at least one + WHOLE block.""" + start, end = token_range + return (end // block_size) > ((start + block_size - 1) // block_size) + + +def _coalesce_adjacent_text(request: Request, block_size: int) -> int: + """Merge runs of adjacent, contiguous, same-type, unpinned text + segments so a stack of short captions becomes one segment big enough + for inward-rounded eviction to free a block. Returns segments removed. + """ + + hist = request.session_history + n = len(hist) + new_hist: list[HistorySegment] = [] + merges = 0 + i = 0 + while i < n: + seg = hist[i] + if seg.pinned or seg.segment_type not in _TEXT_SEGMENT_TYPES: + new_hist.append(seg) + i += 1 + continue + # Collect a maximal contiguous, same-type, unpinned text run. + run = [seg] + j = i + 1 + while j < n: + nxt = hist[j] + if ( + nxt.pinned + or nxt.segment_type != seg.segment_type + or nxt.token_range[0] != run[-1].token_range[1] + ): + break + run.append(nxt) + j += 1 + # Group consecutive members until each group's span frees a block. + group_start = 0 + k = 0 + while k < len(run): + group = run[group_start : k + 1] + span = (group[0].token_range[0], group[-1].token_range[1]) + at_run_end = k == len(run) - 1 + if _frees_block(span, block_size) or at_run_end: + if len(group) == 1: + new_hist.append(group[0]) + else: + new_hist.append( + HistorySegment( + segment_type=seg.segment_type, + token_range=span, + mm_item_id=None, + pinned=False, + age_chunks=min(g.age_chunks for g in group), + ) + ) + merges += len(group) - 1 + group_start = k + 1 + k += 1 + i = j + if merges: + hist[:] = new_hist + return merges + + +def _pack_oldest_run_to_free_block( + request: Request, block_size: int +) -> HistorySegment | None: + """Last-resort packing: glue the OLDEST run of adjacent unpinned + segments — regardless of type — into one block-freeing blob, swap it + into `session_history`, and return it for the caller to evict. Returns + None if even the maximally-extended oldest run can't free a block. + """ + hist = request.session_history + n = len(hist) + # Oldest unpinned segment (the run's anchor). + i = 0 + while i < n and hist[i].pinned: + i += 1 + if i >= n: + return None + # Extend a contiguous, unpinned run from i until its span frees a block. + j = i + while True: + span = (hist[i].token_range[0], hist[j].token_range[1]) + if _frees_block(span, block_size): + break + nxt = j + 1 + if ( + nxt >= n + or hist[nxt].pinned + or hist[nxt].token_range[0] != hist[j].token_range[1] + # Never pack past the computed frontier: the just-appended + # chunk's segments have no KV yet, so `evict_segment` would + # clamp to `num_computed_tokens` and no-op AFTER the run was + # already destructively merged in `session_history`. + or hist[nxt].token_range[1] > request.num_computed_tokens + ): + # Maximally-extended oldest run still can't free a block. + return None + j = nxt + if j == i: + return None # a single freeable segment; the caller already handles it + run = hist[i : j + 1] + merged = HistorySegment( + segment_type=run[0].segment_type, + token_range=(run[0].token_range[0], run[-1].token_range[1]), + mm_item_id=None, + pinned=False, + age_chunks=min(s.age_chunks for s in run), + ) + hist[i : j + 1] = [merged] + return merged + + +def maybe_evict_old_segments( + request: Request, + retention: StreamingRetentionParams, + kv_cache_manager: KVCacheManager, + encoder_cache_manager: EncoderCacheManager, +) -> int: + """Run the sliding-window retention policy; return segments evicted. + + Called after a chunk is appended, so `session_history` already includes + the new chunk's segments. Three sequential phases: + + 1. Video budget. Drop oldest unpinned `video` segments until + `len(videos) <= max_video_segments`. Vision tokens cost the most, + so this fires first. + 2. Text budget (if `max_text_tokens` set). Drop oldest unpinned text + segments until non-pinned text tokens <= `max_text_tokens`. + 3. Safety net (if `max_session_tokens` set). Drop the oldest unpinned + segment while `num_prompt_tokens > max_session_tokens` — the hard + guarantee against drift past `max_model_len`. + """ + if retention.eviction_policy != "sliding_window": + raise NotImplementedError( + f"eviction policy {retention.eviction_policy!r} is not yet supported" + ) + + evicted = 0 + block_size = _block_size(kv_cache_manager) + + # Phase 1: video budget (oldest freeable video first). + while True: + video_segments = [ + s + for s in request.session_history + if s.segment_type == "video" and not s.pinned + ] + if len(video_segments) <= retention.max_video_segments: + break + # age_chunks > 0 skips the new chunk (its tokens have no KV yet). + victim = next( + ( + s + for s in video_segments + if s.age_chunks > 0 and _frees_block(s.token_range, block_size) + ), + None, + ) + if victim is None: + break + evict_segment( + request, + victim, + kv_cache_manager, + encoder_cache_manager, + ) + evicted += 1 + + # Coalesce captions made adjacent by phase-1 eviction (see helper). + _coalesce_adjacent_text(request, block_size) + + # Phase 2: text budget. + if retention.max_text_tokens is not None: + while True: + non_pinned_text = [ + s + for s in request.session_history + if s.segment_type in _TEXT_SEGMENT_TYPES and not s.pinned + ] + text_token_count = sum( + s.token_range[1] - s.token_range[0] for s in non_pinned_text + ) + if text_token_count <= retention.max_text_tokens: + break + # `age_chunks > 0` skips the new chunk's own segment — see Phase 1. + evictable = [ + s + for s in non_pinned_text + if s.age_chunks > 0 and _frees_block(s.token_range, block_size) + ] + if not evictable: + break + oldest = evictable[0] + evict_segment( + request, + oldest, + kv_cache_manager, + encoder_cache_manager, + ) + evicted += 1 + + # Phase 3: total-size backstop. + if retention.max_session_tokens is not None: + while request.num_prompt_tokens > retention.max_session_tokens: + unpinned = [s for s in request.session_history if not s.pinned] + if not unpinned: + break + # `age_chunks > 0` skips the new chunk's own segment — see Phase 1. + victim = next( + ( + s + for s in unpinned + if s.age_chunks > 0 and _frees_block(s.token_range, block_size) + ), + None, + ) + if victim is None and _coalesce_adjacent_text(request, block_size): + victim = next( + ( + s + for s in request.session_history + if not s.pinned + and s.age_chunks > 0 + and _frees_block(s.token_range, block_size) + ), + None, + ) + if victim is None: + victim = _pack_oldest_run_to_free_block(request, block_size) + if victim is None: + break + evict_segment( + request, + victim, + kv_cache_manager, + encoder_cache_manager, + ) + evicted += 1 + + return evicted diff --git a/vllm/v1/streaming/mrope.py b/vllm/v1/streaming/mrope.py new file mode 100644 index 000000000000..dffb3eee58c0 --- /dev/null +++ b/vllm/v1/streaming/mrope.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""mRoPE position helpers for streaming sessions. + +A streaming session is one long-lived request that grows by one chunk per +frame. For each new chunk we compute mRoPE positions for the new tokens only, +starting just past the highest position already in the KV cache, and append +them to the session's position tensor. Positions for cached tokens are never +recomputed, so RoPE on surviving tokens stays consistent with their cached K/V +even after eviction leaves gaps in the position space. +""" + +from dataclasses import replace +from typing import TYPE_CHECKING + +import torch + +from vllm.model_executor.models.interfaces import SupportsMRoPE + +if TYPE_CHECKING: + from vllm.multimodal.inputs import MultiModalFeatureSpec + + +def compute_chunk_mrope_positions( + model: SupportsMRoPE, + chunk_tokens: list[int], + chunk_mm_features: list["MultiModalFeatureSpec"], + base_position: int, +) -> tuple[torch.Tensor, int]: + """Compute mRoPE positions for one chunk's new tokens. + + `chunk_tokens` and `chunk_mm_features` must cover the new chunk only, with + feature offsets relative to the chunk start (offset 0 == first chunk token). + Positions are shifted to begin at `base_position`; pass + `request.max_cached_position + 1` so the chunk's first position is strictly + greater than any position already in the KV cache. + + Returns the (3, len(chunk_tokens)) position tensor and the largest position + assigned (use it to update `max_cached_position`). + """ + if not chunk_tokens: + return torch.empty((3, 0), dtype=torch.long), base_position - 1 + + chunk_positions, _ = model.get_mrope_input_positions( + chunk_tokens, chunk_mm_features + ) + # The model numbers positions from 0 within the slice it sees; shift them + # uniformly so the lowest equals `base_position`. + chunk_positions = chunk_positions + base_position + new_max_position = int(chunk_positions.max().item()) + return chunk_positions, new_max_position + + +def make_chunk_relative_mm_features( + cumulative_mm_features: list["MultiModalFeatureSpec"], + prev_num_tokens: int, +) -> list["MultiModalFeatureSpec"]: + """Extract the latest chunk's mm-features with chunk-relative offsets. + + `cumulative_mm_features` is `req_state.mm_features` after the scheduler has + shifted offsets to be session-absolute. Features at offset >= + `prev_num_tokens` belong to the new chunk; they are copied with offset + reduced by `prev_num_tokens` so the model's mRoPE function sees chunk-local + indices. + """ + chunk_features: list[MultiModalFeatureSpec] = [] + for feature in cumulative_mm_features: + if feature.mm_position.offset >= prev_num_tokens: + chunk_features.append( + replace( + feature, + mm_position=replace( + feature.mm_position, + offset=feature.mm_position.offset - prev_num_tokens, + ), + ) + ) + return chunk_features diff --git a/vllm/v1/streaming/reprefill.py b/vllm/v1/streaming/reprefill.py new file mode 100644 index 000000000000..19b211c7241e --- /dev/null +++ b/vllm/v1/streaming/reprefill.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Self-preempting re-prefill at the position-range threshold. + +A streaming session's mRoPE positions grow monotonically and eventually +exceed the model's trained range (`max_position_embeddings`), degrading +behaviour. To keep sessions effectively unbounded, the session self-preempts +once positions cross a configurable fraction of that range: the scheduler +frees the KV cache and re-queues the request for a fresh full prefill of its +surviving tokens, recomputing all mRoPE positions from 0. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm.v1.request import Request + from vllm.v1.streaming.retention import StreamingRetentionParams + + +def should_trigger_reprefill( + request: Request, + retention: StreamingRetentionParams, + model_max_position: int, + highest_position: int | None = None, +) -> bool: + """Return True once the session's highest cached position exceeds + ``reprefill_threshold * model_max_position``. + + ``highest_position`` overrides the request's worker-reported + ``max_cached_position`` watermark; 1D-RoPE (text) sessions pass the + exactly-derivable ``num_tokens - 1 + position_offset``. + + A threshold >= 1.0 disables re-prefill entirely (the caller warned at + admission; see ``StreamingRetentionParams.__post_init__``). + """ + if retention.reprefill_threshold >= 1.0: + return False + if highest_position is None: + highest_position = request.max_cached_position + threshold = retention.reprefill_threshold * model_max_position + return highest_position > threshold diff --git a/vllm/v1/streaming/retention.py b/vllm/v1/streaming/retention.py new file mode 100644 index 000000000000..b6a778e3e8e0 --- /dev/null +++ b/vllm/v1/streaming/retention.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration and bookkeeping for streaming-session retention. + +A streaming session runs indefinitely; without eviction its KV cache grows +unbounded. `StreamingRetentionParams` configures how much to keep, and +`HistorySegment` tracks per-segment token ranges so the eviction code (in +`eviction.py`) can drop the oldest segments first. This module just holds the +data structures the scheduler maintains as chunks arrive. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Minimum total-token budget: must hold the pinned attention sink plus +# headroom for the next chunk's prefill. +_MIN_SESSION_TOKENS = 1024 + + +@dataclass +class StreamingRetentionParams: + """How long to keep streaming-session context in the KV cache.""" + + max_video_segments: int = 30 + """Hard cap on retained video segments; oldest above the cap are evicted + at chunk boundaries.""" + + max_text_tokens: int | None = None + """Soft cap on non-pinned text tokens (`user_text` + `assistant_text`); + None disables the text budget.""" + + max_session_tokens: int | None = 7000 + """Hard cap on total prompt length post-eviction; the safety net that + keeps an unbounded session below `max_model_len`. Required (must not be + None; see `__post_init__`).""" + + eviction_policy: Literal["sliding_window"] = "sliding_window" + """Currently only sliding-window is implemented.""" + + reprefill_threshold: float = 0.7 + """Fraction of the model's trained position range above which the session + self-preempts and re-prefills its surviving tokens at fresh positions from + 0 (see `reprefill.py`). Keeps mRoPE positions bounded over very long + sessions.""" + + def __post_init__(self) -> None: + """Validate the budget on every construction (in-process and + IPC-deserialized) so a bad config fails loudly at admission.""" + # < 1 would evict every chunk's video before the encoder runs. + if self.max_video_segments < 1: + raise ValueError( + "max_video_segments must be >= 1 " + f"(got {self.max_video_segments}); a value < 1 evicts " + "every chunk's video before it is computed." + ) + + # Required: the load-bearing safety net against block-alignment drift. + if self.max_session_tokens is None: + raise ValueError( + "max_session_tokens must be set: it is the hard token " + "safety net that keeps an unbounded streaming session " + "below max_model_len." + ) + if self.max_session_tokens < _MIN_SESSION_TOKENS: + raise ValueError( + f"max_session_tokens ({self.max_session_tokens}) must be " + f">= {_MIN_SESSION_TOKENS} to hold the pinned sink plus a " + "chunk's prefill headroom." + ) + + # Optional, but if set must be positive and within the session budget. + if self.max_text_tokens is not None: + if self.max_text_tokens <= 0: + raise ValueError( + f"max_text_tokens must be > 0 when set (got " + f"{self.max_text_tokens})." + ) + if self.max_text_tokens > self.max_session_tokens: + raise ValueError( + f"max_text_tokens ({self.max_text_tokens}) must not " + f"exceed max_session_tokens ({self.max_session_tokens})." + ) + + # A non-positive threshold would re-trigger a full re-prefill at + # every chunk boundary (unbounded GPU-burn livelock). Written as + # `not (0 < x)` so NaN — which fails every comparison and would + # silently disable re-prefill, letting positions grow unbounded — + # is rejected too. + if not (self.reprefill_threshold > 0): + raise ValueError( + f"reprefill_threshold ({self.reprefill_threshold}) must be " + "> 0; a non-positive value re-triggers re-prefill at every " + "chunk boundary." + ) + + # Disabling re-prefill is supported, but nothing else bounds mRoPE + # position growth, so a long session eventually exceeds the trained + # range (rotary cache out-of-bounds). Warn once. + if self.reprefill_threshold >= 1.0: + logger.warning_once( + "StreamingRetentionParams.reprefill_threshold=%s disables " + "re-prefill; mRoPE positions will grow without bound and a " + "long-running session will eventually exceed the model's " + "trained position range (rotary-cache out-of-bounds). Use a " + "value < 1.0 (default 0.7) to keep positions bounded.", + self.reprefill_threshold, + ) + + +SegmentType = Literal[ + "system_prompt", + "video", + "user_text", + "assistant_text", +] + + +@dataclass +class HistorySegment: + """One contiguous range of tokens in the streaming session prompt. + + Eviction keys purely off `token_range` (see eviction.py). After an + eviction, surviving segments' `token_range` is shifted down to stay + consistent with the compacted token arrays, but `_mrope_positions` values + are not — so the position space has gaps, preserving cached-K/V rotation. + """ + + segment_type: SegmentType + token_range: tuple[int, int] + """`[start, end)` into `_all_token_ids`. Shifts down as other segments are + evicted, so always read it fresh from the segment — never cache it.""" + + mm_item_id: str | None = None + """`MultiModalFeatureSpec` identifier (video segments only), used to free + the encoder cache entry.""" + + pinned: bool = False + """If True, eviction skips this segment. Set on the chunk-1 anchor.""" + + age_chunks: int = 0 + """Chunk boundaries elapsed since this segment was added; 0 means it + arrived with the current chunk (its KV may not be computed yet, so + eviction skips it).""" diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 7585f4d89245..b763c9b25e5a 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -46,6 +46,13 @@ class CachedRequestState: mrope_positions: torch.Tensor | None = None mrope_position_delta: int | None = None + # Highest position number currently in this request's KV cache + # (streaming sessions only; -1/None otherwise). + max_cached_position: int | None = None + # 1D-RoPE (text) streaming sessions: index -> RoPE position offset + # (cumulative evicted width). 0 for everything else. + position_offset: int = 0 + xdrope_positions: torch.Tensor | None = None lora_request: LoRARequest | None = None diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 93e19fca2152..2cd1191305c5 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -211,6 +211,10 @@ from vllm.v1.spec_decode.step3p5 import Step3p5MTPProposer from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change +from vllm.v1.streaming.mrope import ( + compute_chunk_mrope_positions, + make_chunk_relative_mm_features, +) from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext from vllm.v1.worker import mamba_utils @@ -579,6 +583,24 @@ def __init__( self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs( model_config ) + # Multimodal placeholder ids (image + video) used to neutralize + # orphan markers left by streaming eviction; None if the model + # defines neither. Precomputed once (int32 matches input_ids). + _mm_placeholder_ids = sorted( + { + t + for t in ( + getattr(model_config.hf_config, "image_token_id", None), + getattr(model_config.hf_config, "video_token_id", None), + ) + if t is not None + } + ) + self.mm_placeholder_ids_cpu: torch.Tensor | None = ( + torch.tensor(_mm_placeholder_ids, dtype=torch.int32, device="cpu") + if _mm_placeholder_ids + else None + ) if self.model_config.is_encoder_decoder: # Maximum length of the encoder input, only for encoder-decoder @@ -1501,15 +1523,37 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None # The request is resumed from preemption. # Replace the existing block IDs with the new ones. req_state.block_ids = new_block_ids + # A re-prefill resumes here at num_computed_tokens==0 with its + # outputs folded into the prompt; reset the stale position / + # prompt state (else the discard mask uses wrong values). + # Gated on num_output_tokens==0 so a mid-decode resume + # (outputs kept) is left to the restore path below. + if num_computed_tokens == 0 and num_output_tokens == 0: + req_state.output_token_ids.clear() + req_state.num_prompt_tokens = ( + length_from_prompt_token_ids_or_embeds( + req_state.prompt_token_ids, + req_state.prompt_embeds, + ) + ) + req_state.position_offset = 0 + if self.uses_mrope: + req_state.mrope_positions = None + req_state.max_cached_position = None + self._init_mrope_positions(req_state) if req_index is None: # The request is not in the persistent batch. # The request was either preempted and resumed later, or was not # scheduled in the previous step and needs to be added again. - if self.use_async_scheduling and num_output_tokens > 0: - # We must recover the output token ids for resumed requests in the - # async scheduling case, so that correct input_ids are obtained. + if num_output_tokens > 0 and ( + self.use_async_scheduling or resumed_from_preemption + ): + # Rebuild output_token_ids from the scheduler's + # authoritative all_token_ids — the worker's copy is stale + # under async scheduling and untrusted for a + # preempt-resumed request. resumed_token_ids = req_data.all_token_ids[req_id] req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] @@ -1688,6 +1732,16 @@ def _update_streaming_request( self.input_batch.remove_request(req_id) req_state = self.requests[req_id] + # Re-prefill: the scheduler's explicit signal to recompute positions + # from 0. + if new_req_data.is_reprefill and self.uses_mrope: + req_state.mrope_positions = None + req_state.max_cached_position = None + + # 1D-RoPE (text) sessions: index -> RoPE position offset, refreshed + # every streaming resume (grows on eviction, resets on re-prefill). + req_state.position_offset = new_req_data.position_offset + req_state.prompt_token_ids = new_req_data.prompt_token_ids req_state.mm_features = new_req_data.mm_features req_state.prompt_embeds = new_req_data.prompt_embeds @@ -1705,10 +1759,95 @@ def _update_streaming_request( req_state.output_token_ids.clear() if self.uses_mrope: - self._init_mrope_positions(req_state) + if new_req_data.evicted_token_ranges: + self._slice_mrope_positions_for_evicted_ranges( + req_state, new_req_data.evicted_token_ranges + ) + self._extend_mrope_positions_for_streaming_chunk(req_state) return req_state + @staticmethod + def _slice_mrope_positions_for_evicted_ranges( + req_state: CachedRequestState, + evicted_ranges: list[tuple[int, int]], + ) -> None: + """Drop the columns for evicted token ranges from + `req_state.mrope_positions`, mirroring the scheduler-side compaction. + + Ranges apply in order, each in the tensor's index space at its own + eviction time (later ranges already reflect earlier shrinks); ranges + past the current width are clipped. Surviving positions keep their + values and `max_cached_position` is untouched, so the position space + has gaps but cached K/V rotation stays consistent. + """ + if req_state.mrope_positions is None: + return + positions = req_state.mrope_positions + for start, end in evicted_ranges: + width = positions.shape[1] + end_clipped = min(end, width) + if start >= end_clipped: + continue + if start == 0 and end_clipped >= width: + positions = positions[:, :0] + elif start == 0: + positions = positions[:, end_clipped:] + elif end_clipped >= width: + positions = positions[:, :start] + else: + positions = torch.cat( + [positions[:, :start], positions[:, end_clipped:]], + dim=1, + ) + req_state.mrope_positions = positions + + def _extend_mrope_positions_for_streaming_chunk( + self, req_state: CachedRequestState + ) -> None: + """Append positions for the new chunk only (tokens past the prior + chunk's end), starting at `max_cached_position + 1` so they sit + strictly above any cached position — preserving cached K/V's rotation. + """ + assert req_state.prompt_token_ids is not None + if req_state.mrope_positions is None or req_state.max_cached_position is None: + self._init_mrope_positions(req_state) + return + + prev_num_tokens = req_state.mrope_positions.shape[1] + new_num_tokens = req_state.num_prompt_tokens + if new_num_tokens == prev_num_tokens: + return + assert new_num_tokens > prev_num_tokens, ( + f"streaming chunk shrank prompt from {prev_num_tokens} to " + f"{new_num_tokens}; eviction must use the slice path, not this one." + ) + + chunk_tokens = req_state.prompt_token_ids[prev_num_tokens:new_num_tokens] + cumulative_features = [ + f for f in req_state.mm_features if f.modality != "prompt_embeds" + ] + chunk_features = make_chunk_relative_mm_features( + cumulative_features, prev_num_tokens + ) + + model = self.get_model() + mrope_model = cast(SupportsMRoPE, model) + chunk_positions, new_max_position = compute_chunk_mrope_positions( + mrope_model, + chunk_tokens, + chunk_features, + base_position=req_state.max_cached_position + 1, + ) + + req_state.mrope_positions = torch.cat( + [req_state.mrope_positions, chunk_positions.to(req_state.mrope_positions)], + dim=1, + ) + req_state.max_cached_position = new_max_position + # Decode-time positions: linear continuation from new max. + req_state.mrope_position_delta = new_max_position + 1 - new_num_tokens + def _init_mrope_positions(self, req_state: CachedRequestState): model = self.get_model() assert supports_mrope(model), "M-RoPE support is not implemented." @@ -1740,6 +1879,12 @@ def _init_mrope_positions(self, req_state: CachedRequestState): mrope_features, ) ) + # Track the highest position number for streaming continuation. + # Empty prompts (only possible with prompt_embeds) leave it at -1. + if req_state.mrope_positions.numel() > 0: + req_state.max_cached_position = int(req_state.mrope_positions.max().item()) + else: + req_state.max_cached_position = -1 def _init_xdrope_positions(self, req_state: CachedRequestState): model = self.get_model() @@ -2257,6 +2402,23 @@ def _prepare_inputs( self.positions[:total_num_scheduled_tokens], ) + # 1D-RoPE (text) streaming sessions: shift the RoPE positions of this + # step's scheduled tokens by the session's cumulative evicted width, + # so new tokens continue past the evicted range and cached K keeps a + # consistent rotation. MUST run after compute_slot_mapping (physical + # slots are index-based). mRoPE models are untouched (their RoPE reads + # mrope_positions, not this buffer). NOTE(draft): attention backends + # that read CommonAttentionMetadata.positions for masking see the + # shifted values; audited for the default FA path, others pending. + if not self.uses_mrope: + for idx, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests.get(req_id) + if req_state is None or not req_state.position_offset: + continue + start = int(cu_num_tokens[idx - 1]) if idx > 0 else 0 + end = int(cu_num_tokens[idx]) + self.positions[start:end] += req_state.position_offset + # Copy the tensors to the GPU. self._prepare_input_ids( scheduler_output, @@ -3411,6 +3573,87 @@ def _gather_mm_embeddings( return mm_embeds, is_mm_embed + @staticmethod + def _req_uses_streaming_retention(req_state: CachedRequestState) -> bool: + """Worker-side mirror of `Request.uses_persistent_encoder_cache()`, + re-derived from `sampling_params` (which `CachedRequestState` carries, + unlike `streaming_retention`). True iff a `streaming_retention` in + `extra_args` (dataclass or dict) has `reprefill_threshold < 1.0`. + """ + sampling_params = req_state.sampling_params + if sampling_params is None: + return False + extra_args = sampling_params.extra_args + if not extra_args: + return False + retention = extra_args.get("streaming_retention") + if retention is None: + return False + if isinstance(retention, dict): + threshold = retention.get("reprefill_threshold", 1.0) + else: + threshold = getattr(retention, "reprefill_threshold", 1.0) + try: + return float(threshold) < 1.0 + except (TypeError, ValueError): + return False + + def _neutralize_orphan_mm_placeholders( + self, + scheduler_output: "SchedulerOutput", + inputs_embeds: torch.Tensor, + is_mm_embed: torch.Tensor, + num_scheduled_tokens: int, + ) -> None: + """Zero the embedding rows of orphan image/video placeholder markers + (e.g. `<|image_pad|>` / `<|video_pad|>`) left by block-aligned + streaming eviction. + + Eviction frees an mm feature at its raw range but compacts tokens at + the narrower block-aligned range, so up to `block_size - 1` markers + survive with no `mm_feature`; at re-prefill they miss + `_gather_mm_embeddings` and would get their raw (OOD) text embedding. + Quality-only — touches no KV-feeding state. No-op unless the model has + `image_token_id`/`video_token_id` and a streaming request has such an + orphan. + """ + placeholder_ids = self.mm_placeholder_ids_cpu + if placeholder_ids is None: + return + + eligible = torch.zeros(num_scheduled_tokens, dtype=torch.bool, device="cpu") + any_streaming = False + req_start_idx = 0 + for req_id in self.input_batch.req_ids: + n = scheduler_output.num_scheduled_tokens[req_id] + if self._req_uses_streaming_retention(self.requests[req_id]): + eligible[req_start_idx : req_start_idx + n] = True + any_streaming = True + req_start_idx += n + if not any_streaming: + return + + # Orphan mask (host-side): placeholder id, no backing feature, + # streaming. + token_ids_cpu = self.input_ids.gpu[:num_scheduled_tokens].to( + "cpu", non_blocking=False + ) + orphan_mask = ( + torch.isin(token_ids_cpu, placeholder_ids) + & (~is_mm_embed.to("cpu")) + & eligible + ) + if not bool(orphan_mask.any()): + return + + orphan_idx = ( + orphan_mask.nonzero(as_tuple=False) + .squeeze(1) + .to(inputs_embeds.device, non_blocking=True) + ) + # Zero only the orphan rows, removing their OOD text-embedding signal. + inputs_embeds[orphan_idx] = 0 + def get_model(self) -> nn.Module: if not hasattr(self, "model"): raise ValueError("Cannot get model before model has been initialized") @@ -3672,6 +3915,16 @@ def _preprocess( is_multimodal=is_mm_embed, ) + # Neutralize orphan image/video placeholder markers left by + # streaming eviction (see method); no-op for non-mm / + # non-streaming. + self._neutralize_orphan_mm_placeholders( + scheduler_output, + inputs_embeds_scheduled, + is_mm_embed, + num_scheduled_tokens, + ) + # TODO(woosuk): Avoid the copy. Optimize. self.inputs_embeds.gpu[:num_scheduled_tokens].copy_( inputs_embeds_scheduled @@ -4874,6 +5127,19 @@ def propose_draft_token_ids(sampled_token_ids): self.kv_connector_output = None with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + # Streaming: report each req's `max_cached_position` back to the + # scheduler so the re-prefill trigger sees the true watermark + # (else it stays -1 and never fires). Skip unset (None / -1). + max_cached_positions: dict[str, int] = {} + if self.uses_mrope: + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + if req_state is None: + continue + pos = req_state.max_cached_position + if pos is not None and pos >= 0: + max_cached_positions[req_id] = pos + output = ModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -4887,6 +5153,7 @@ def propose_draft_token_ids(sampled_token_ids): num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, routed_experts=None, + max_cached_positions=max_cached_positions, ) if not self.use_async_scheduling: