diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 4f35527b0e..74bac55598 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -2,18 +2,28 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from unittest.mock import Mock +from typing import Literal +from unittest.mock import Mock, patch import pytest +import torch +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler -from vllm.v1.request import Request, RequestStatus +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) +from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( create_model_runner_output, create_request, create_scheduler, create_vllm_config, + make_kv_cache_config, ) @@ -34,6 +44,18 @@ def scheduler(): return create_scheduler(vllm_config) +@pytest.fixture +def hybrid_scheduler(): + """Scheduler for a hybrid model: full attention plus a Mamba group.""" + vllm_config = create_vllm_config(kv_load_failure_policy="recompute") + kv_cache_config = make_kv_cache_config( + block_size=vllm_config.cache_config.block_size, + mamba_enabled=True, + num_blocks=10000, + ) + return create_scheduler(vllm_config, kv_cache_config=kv_cache_config) + + @pytest.mark.parametrize( "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs", [ @@ -337,3 +359,319 @@ def test_async_progressive_load_failure( assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS assert scheduler.failed_recving_kv_req_ids == {request.request_id} assert scheduler.connector.get_num_new_matched_tokens.call_count == 1 + + +def _schedule_hybrid_async_load( + scheduler: Scheduler, num_prompt_blocks: int, num_external_computed_blocks: int +) -> Request: + """Put one request into an async external load on a hybrid model.""" + request = create_request(num_tokens=num_prompt_blocks * scheduler.block_size) + scheduler.add_request(request=request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_computed_blocks * scheduler.block_size}, + async_load=True, + ) + ) + scheduler.connector.take_events.return_value = () + return request + + +def _report_invalid_blocks( + scheduler: Scheduler, scheduler_output, invalid_block_ids: set[int] +) -> None: + scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids=invalid_block_ids, + use_eos=True, + ), + ) + + +@pytest.mark.parametrize("failed_group_idx", [0, 1]) +def test_hybrid_load_failure_recomputes_whole_request( + hybrid_scheduler: Scheduler, failed_group_idx: int +): + """A hybrid model must recover, not crash. + + Its groups sit on different block grids, so there is no single truncation + point; the request restarts from scratch instead. + """ + num_prompt_blocks, num_external_computed_blocks = 10, 9 + request = _schedule_hybrid_async_load( + hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks + ) + scheduler_output = hybrid_scheduler.schedule() + + block_ids_per_group = hybrid_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert len(block_ids_per_group) == 2, "expected a hybrid (multi-group) layout" + null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id + # The Mamba group backs only its aligned snapshot with a real block, so + # pick a block that actually belongs to this request. + failed_block = next( + block_id + for block_id in block_ids_per_group[failed_group_idx] + if block_id != null_block_id + ) + + _report_invalid_blocks(hybrid_scheduler, scheduler_output, {failed_block}) + + assert request.num_computed_tokens == 0 + assert hybrid_scheduler.failed_recving_kv_req_ids == {request.request_id} + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_hybrid_load_failure_ignores_null_block(hybrid_scheduler: Scheduler): + """The null block is shared by every request and holds no KV, so reporting + it must not restart anyone.""" + num_prompt_blocks, num_external_computed_blocks = 10, 9 + request = _schedule_hybrid_async_load( + hybrid_scheduler, num_prompt_blocks, num_external_computed_blocks + ) + scheduler_output = hybrid_scheduler.schedule() + + null_block_id = hybrid_scheduler.kv_cache_manager.block_pool.null_block.block_id + _, mamba_block_ids = hybrid_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + # The Mamba group parks its unaligned positions on the null block, so a + # request that never failed would otherwise be restarted by this report. + assert mamba_block_ids[0] == null_block_id + + num_computed_tokens = request.num_computed_tokens + _report_invalid_blocks(hybrid_scheduler, scheduler_output, {null_block_id}) + + assert request.num_computed_tokens == num_computed_tokens + assert not hybrid_scheduler.failed_recving_kv_req_ids + + +def _make_deepseek_v4_five_group_config(num_blocks: int) -> KVCacheConfig: + """Build the scheduler-side projection of V4-Flash's packed KV layout.""" + + def mla_spec() -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + tokens_per_state=4, + model_version="deepseek_v4", + ) + + def swa_spec( + *, block_size: int, sliding_window: int, head_size: int + ) -> SlidingWindowMLASpec: + return SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + head_size=head_size, + dtype=torch.float32, + sliding_window=sliding_window, + model_version="deepseek_v4", + ) + + # DeepSeek-V4 packs these into five physical groups. The two SWA groups + # have identical scheduling semantics but are split to align layer tuples. + groups = [ + KVCacheGroupSpec(["full_mla_and_indexer"], mla_spec()), + KVCacheGroupSpec( + ["c4_state"], + swa_spec(block_size=4, sliding_window=8, head_size=2048), + ), + KVCacheGroupSpec( + ["c128_state"], + swa_spec(block_size=8, sliding_window=128, head_size=2048), + ), + KVCacheGroupSpec( + ["swa.0"], + swa_spec(block_size=64, sliding_window=512, head_size=512), + ), + KVCacheGroupSpec( + ["swa.1"], + swa_spec(block_size=64, sliding_window=512, head_size=512), + ), + ] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=groups, + ) + + +def _create_deepseek_v4_scheduler( + kv_load_failure_policy: Literal["recompute", "fail"] = "recompute", +) -> Scheduler: + num_blocks = 10000 + vllm_config = create_vllm_config( + block_size=256, + max_num_batched_tokens=4096, + max_model_len=8192, + # max_model_len only reaches SchedulerConfig; opt-125m still caps the + # model at 2048, which is below the 9-block requests used here. + hf_overrides={"max_position_embeddings": 8192}, + kv_load_failure_policy=kv_load_failure_policy, + ) + return create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=_make_deepseek_v4_five_group_config(num_blocks), + hash_block_size=4, + ) + + +@pytest.fixture +def deepseek_v4_scheduler() -> Scheduler: + return _create_deepseek_v4_scheduler() + + +def _schedule_deepseek_v4_async_load( + scheduler: Scheduler, +) -> tuple[Request, SchedulerOutput]: + num_external_tokens = 8 * scheduler.block_size + request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_tokens}, async_load=True + ) + ) + scheduler.connector.take_events.return_value = () + return request, scheduler.schedule() + + +@pytest.mark.parametrize("failed_group_idx", range(5)) +def test_deepseek_v4_load_failure_recomputes_whole_request( + deepseek_v4_scheduler: Scheduler, failed_group_idx: int +): + """A failed load in any V4-Flash KV group must not kill EngineCore.""" + request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) + + block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert len(block_ids_per_group) == 5 + null_block_id = ( + deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id + ) + failed_block_id = next( + block_id + for block_id in block_ids_per_group[failed_group_idx] + if block_id != null_block_id + ) + + deepseek_v4_scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids={failed_block_id}, + use_eos=True, + ), + ) + + assert request.num_computed_tokens == 0 + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert deepseek_v4_scheduler.failed_recving_kv_req_ids == {request.request_id} + + +def test_deepseek_v4_load_failure_ignores_shared_null_block( + deepseek_v4_scheduler: Scheduler, +): + """The shared null block contains no request KV and cannot fail a load.""" + request, scheduler_output = _schedule_deepseek_v4_async_load(deepseek_v4_scheduler) + null_block_id = ( + deepseek_v4_scheduler.kv_cache_manager.block_pool.null_block.block_id + ) + block_ids_per_group = deepseek_v4_scheduler.kv_cache_manager.get_block_ids( + request.request_id + ) + assert any(null_block_id in group for group in block_ids_per_group[1:]) + num_computed_tokens = request.num_computed_tokens + + deepseek_v4_scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + reqs=[], + finished_recving=set(), + invalid_block_ids={null_block_id}, + use_eos=True, + ), + ) + + assert request.num_computed_tokens == num_computed_tokens + assert not deepseek_v4_scheduler.failed_recving_kv_req_ids + + +def test_deepseek_v4_sync_load_failure_fails_and_evicts_real_blocks(): + """Fail policy must terminate cleanly and never evict the null block.""" + scheduler = _create_deepseek_v4_scheduler(kv_load_failure_policy="fail") + num_external_tokens = 8 * scheduler.block_size + request = create_request(num_tokens=9 * scheduler.block_size, block_size=4) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + {request.request_id: num_external_tokens}, async_load=False + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.take_events.return_value = () + scheduler_output = scheduler.schedule() + + block_ids_per_group = scheduler.kv_cache_manager.get_block_ids(request.request_id) + assert len(block_ids_per_group) == 5 + null_block_id = scheduler.kv_cache_manager.block_pool.null_block.block_id + real_block_ids = { + block_id + for group_block_ids in block_ids_per_group + for block_id in group_block_ids + if block_id != null_block_id + } + failed_block_id = next( + block_id for block_id in block_ids_per_group[2] if block_id != null_block_id + ) + + evicted_block_ids: list[set[int]] = [] + original_evict_blocks = scheduler.kv_cache_manager.evict_blocks + + def evict_blocks_spy(block_ids: set[int]) -> None: + evicted_block_ids.append(set(block_ids)) + original_evict_blocks(block_ids) + + with ( + patch.object( + scheduler.kv_cache_manager, + "evict_blocks", + side_effect=evict_blocks_spy, + ), + patch.object(scheduler, "_connector_finished", return_value=(False, None)), + ): + outputs = scheduler.update_from_output( + scheduler_output, + create_model_runner_output( + [request], + invalid_block_ids={failed_block_id}, + use_eos=True, + ), + ) + + assert request.status == RequestStatus.FINISHED_ERROR + assert request.get_finished_reason() == FinishReason.ERROR + assert request.request_id not in scheduler.requests + assert evicted_block_ids == [real_block_ids] + assert null_block_id not in evicted_block_ids[0] + + engine_outputs = next(iter(outputs.values())) + assert len(engine_outputs.outputs) == 1 + assert engine_outputs.outputs[0].request_id == request.request_id + assert engine_outputs.outputs[0].finish_reason == FinishReason.ERROR diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 7b70065d7f..07a40dc4ec 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -155,6 +155,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -185,6 +186,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 8220373845..1b17eb3b15 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3058,6 +3058,11 @@ class Scheduler(SchedulerInterface): For observability, it also accumulates the total number of tokens that will need to be recomputed across all affected requests. + A hybrid model spreads one token span over KV cache groups that disagree + on block size, so a block index maps to a different token position in + each group and there is no single longest valid prefix. Those requests + recompute in full instead. + Args: requests: The set of requests to scan for invalid blocks. invalid_block_ids: IDs of invalid blocks. @@ -3081,18 +3086,49 @@ class Scheduler(SchedulerInterface): # these requests must be rescheduled, but only the first will recompute # it. This set tracks blocks already marked for recomputation. marked_invalid_block_ids: set[int] = set() + # Positions a group holds no KV for -- skipped sliding-window blocks, + # and every Mamba block but the aligned snapshot -- all point at this + # one block, shared by every request. It carries no request's data, so + # it can neither fail for a request nor be evicted on its behalf. + null_block_id = self.kv_cache_manager.block_pool.null_block.block_id for request in requests: is_affected = False marked_invalid_block = False req_id = request.request_id - # TODO (davidb): add support for hybrid memory allocator - (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) + # One block-id list per KV cache group. + req_block_ids_per_group = self.kv_cache_manager.get_block_ids(req_id) # We iterate only over blocks that may contain externally computed # tokens req_num_computed_tokens = ( request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) ) + if len(req_block_ids_per_group) > 1: + # No single truncation point exists across groups, so the whole + # request is recomputed. Blocks are still marked, so a request + # sharing one of them is rescheduled too. + request_invalid_block_ids = { + block_id + for req_block_ids in req_block_ids_per_group + for block_id in req_block_ids + if block_id != null_block_id and block_id in invalid_block_ids + } + if request_invalid_block_ids: + marked_invalid_block_ids |= request_invalid_block_ids + total_affected_tokens += req_num_computed_tokens + request.num_computed_tokens = 0 + if evict_blocks: + for req_block_ids in req_block_ids_per_group: + blocks_to_evict.update( + block_id + for block_id in req_block_ids + if block_id != null_block_id + ) + affected_req_ids.add(req_id) + continue + + (req_block_ids,) = req_block_ids_per_group + req_num_computed_blocks = ( req_num_computed_tokens + self.block_size - 1 ) // self.block_size