iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core] Fix ValueError on KV load failure with a hybrid KV cache (#50388)

Browse Source 
Signed-off-by: Kebe <mail@kebe7jun.com>
This commit is contained in:
Kebe
2026-09-13 21:03:05 +08:00
committed by GitHub
parent fa1b3b1922
commit c351fd3c64
3 changed files with 380 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
+340 -2
View File
@@ -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
tests/v1/kv_connector/unit/utils.py
+2
View File
@@ -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,
)
vllm/v1/core/sched/scheduler.py
+38 -2
View File
@@ -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