diff --git a/tests/v1/streaming_input/test_scheduler_streaming.py b/tests/v1/streaming_input/test_scheduler_streaming.py index bfef3dd33c..1f97d8a1c0 100644 --- a/tests/v1/streaming_input/test_scheduler_streaming.py +++ b/tests/v1/streaming_input/test_scheduler_streaming.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock import numpy as np import torch +import vllm.v1.core.kv_cache_utils as kv_cache_utils from vllm.config import DeviceConfig, VllmConfig from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -14,6 +15,12 @@ from vllm.multimodal.inputs import ( PlaceholderRange, ) from vllm.sampling_params import SamplingParams +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import ( + get_request_block_hasher, + hash_block_tokens, + init_none_hash, +) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.engine import FinishReason from vllm.v1.kv_cache_interface import ( @@ -365,6 +372,48 @@ class TestStreamingScheduler(unittest.TestCase): num_new_tokens = session.num_tokens - session.num_computed_tokens assert num_new_tokens == 2 + def test_update_request_as_session_drops_stale_block_hashes(self): + """A discarded sampled token must not stay fingerprinted in + block_hashes (#49377).""" + init_none_hash(sha256) + scheduler = create_scheduler() + hash_block_size = scheduler.hash_block_size + + # hash_block_size - 1 prompt tokens, so the discarded sampled token + # completes hash block 0. + prompt = list(range(1, hash_block_size)) + discarded_token = 999 + session = Request( + request_id="session", + prompt_token_ids=list(prompt), + sampling_params=SamplingParams(max_tokens=16), + pooling_params=None, + block_hasher=get_request_block_hasher(hash_block_size, sha256), + resumable=True, + ) + session.append_output_token_ids(discarded_token) + assert len(session.block_hashes) == 1 + stale_hash = session.block_hashes[0] + session.num_computed_tokens = len(prompt) + + next_token = 42 + update = StreamingUpdate.from_request( + DummyRequest(request_id="session", prompt_token_ids=[next_token]) + ) + scheduler._update_request_as_session(session, update) + + assert list(session.all_token_ids) == prompt + [next_token] + assert len(session.block_hashes) == 1 + expected = hash_block_tokens( + sha256, kv_cache_utils.NONE_HASH, tuple(prompt + [next_token]) + ) + assert session.block_hashes[0] == expected, ( + "block_hashes[0] does not identify the current token sequence" + ) + assert session.block_hashes[0] != stale_hash, ( + "block_hashes[0] still fingerprints the discarded sampled token" + ) + def test_streaming_e2e_lifecycle(self): """Comprehensive integration test covering complete streaming request lifecycle including scheduler state management and aliasing bug prevention. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index e25b5adda4..0f2425312c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1628,6 +1628,7 @@ class Scheduler(SchedulerInterface): session.num_prompt_tokens : num_computed_tokens ] del session._all_token_ids[num_computed_tokens:] + del session.block_hashes[num_computed_tokens // self.hash_block_size :] session._output_token_ids.clear() assert session.prompt_token_ids is not None # Extend prompt with kept output tokens.