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

[KV Offload] Deduplicate replicated MLA KV in the shared CPU region (#48906)

Browse Source 
Signed-off-by: Change72 <changg@nvidia.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Chang Guo
2026-07-26 08:22:30 +03:00
committed by GitHub
co-authored by OpenAI Codex
parent 7a29a3c54c
commit 7eca0e1a64
14 changed files with 1143 additions and 506 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/lm_eval.yaml
+2 -2
View File
@@ -337,7 +337,7 @@ steps:
- label: LM Eval KV-Offload (2xH100)
key: kv-offload-medium
timeout_in_minutes: 30
timeout_in_minutes: 45
device: h100
num_devices: 2
source_file_dependencies:
@@ -347,7 +347,7 @@ steps:
- vllm/v1/simple_kv_offload/
- tests/evals/gsm8k/test_gsm8k_offloading.py
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
- pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
- label: LM Eval KV-Offload (4xH100)
key: kv-offload-large
tests/evals/gsm8k/test_gsm8k_offloading.py
+25 -5
View File
@@ -16,11 +16,13 @@ KV drops accuracy below threshold. The reload itself is not asserted, so
a silently skipped reload (e.g. offloading disabled) would not be flagged.
Covers both KV offloading connectors (OffloadingConnector and
SimpleCPUOffloadConnector) across four architecture families:
SimpleCPUOffloadConnector) across five architecture families:
- Hybrid Mamba (NemotronH: attention + Mamba)
- Heterogeneous head dim (Gemma 4)
- Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
- Compressed attention (DeepSeek-V4-Flash: CSA)
- Pure MLA (DeepSeek-V2-Lite: TieringOffloadingSpec, TP=2, replicated
single-slot host layout)
Usage:
pytest -s -v evals/gsm8k/test_gsm8k_offloading.py
@@ -47,14 +49,16 @@ NUM_FEWSHOT = 5
_OFFLOAD_SYNC_TIMEOUT = 60
def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str:
def _kv_transfer_config(
connector: str, cpu_gib: int = 4, spec_name: str = "CPUOffloadingSpec"
) -> str:
if connector == "OffloadingConnector":
return json.dumps(
{
"kv_connector": "OffloadingConnector",
"kv_role": "kv_both",
"kv_connector_extra_config": {
"spec_name": "CPUOffloadingSpec",
"spec_name": spec_name,
"cpu_bytes_to_use": cpu_gib << 30,
"eviction_policy": "lru",
},
@@ -115,8 +119,10 @@ class OffloadingModelConfig:
accuracy_threshold: float
tolerance: float = 0.05
extra_server_args: list[str] = field(default_factory=list)
env_dict: dict[str, str] = field(default_factory=dict)
cpu_offload_gib: int = 4
startup_timeout: int = 600
spec_name: str = "CPUOffloadingSpec"
MODELS = [
@@ -165,6 +171,20 @@ MODELS = [
cpu_offload_gib=16,
startup_timeout=1200,
),
OffloadingModelConfig(
id="offloading-deepseek-v2-lite-tiering-tp2",
model="deepseek-ai/DeepSeek-V2-Lite",
connector="OffloadingConnector",
# Baseline 0.360/0.345 over two runs (measured on 2xA100).
accuracy_threshold=0.35,
extra_server_args=[
"--tensor-parallel-size",
"2",
],
cpu_offload_gib=8,
startup_timeout=1200,
spec_name="TieringOffloadingSpec",
),
# ── SimpleCPUOffloadConnector ────────────────────────────────────
OffloadingModelConfig(
id="simple-nemotron-h-8b",
@@ -229,7 +249,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
"--enable-prefix-caching",
"--no-disable-hybrid-kv-cache-manager",
"--kv-transfer-config",
_kv_transfer_config(cfg.connector, cfg.cpu_offload_gib),
_kv_transfer_config(cfg.connector, cfg.cpu_offload_gib, cfg.spec_name),
"--trust-remote-code",
"--disable-uvicorn-access-log",
*cfg.extra_server_args,
@@ -239,7 +259,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
cfg.model,
server_args,
# /reset_prefix_cache requires dev mode.
env_dict={"VLLM_SERVER_DEV_MODE": "1"},
env_dict={"VLLM_SERVER_DEV_MODE": "1", **cfg.env_dict},
max_wait_seconds=cfg.startup_timeout,
) as server:
base_url = f"http://{server.host}:{server.port}"
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+587
View File
@@ -0,0 +1,587 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for translating vLLM cache metadata to native offloading config."""
from typing import Any, cast
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
from vllm.platforms import current_platform
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
HiddenStateCacheSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MambaSpec,
MLAAttentionSpec,
SlidingWindowMLASpec,
SlidingWindowSpec,
UniformTypeKVCacheSpecs,
)
def _make_vllm_config(
*,
extra_config: dict[str, Any] | None = None,
tensor_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
prefill_context_parallel_size: int = 1,
decode_context_parallel_size: int = 1,
) -> VllmConfig:
config = MagicMock()
config.cache_config.block_size = 16
config.cache_config.enable_prefix_caching = True
config.cache_config.prefix_match_unit = None
config.cache_config.cache_dtype = torch.float16
config.model_config.model = "test-model"
config.model_config.use_mla = False
world_size = (
tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
)
with patch.object(current_platform, "device_count", return_value=world_size):
config.parallel_config = ParallelConfig(
tensor_parallel_size=tensor_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
prefill_context_parallel_size=prefill_context_parallel_size,
decode_context_parallel_size=decode_context_parallel_size,
)
config.kv_events_config = None
config.use_v2_model_runner = False
config.kv_transfer_config = KVTransferConfig(
kv_connector="OffloadingConnector",
kv_role="kv_both",
kv_connector_extra_config=dict(extra_config or {}),
)
return cast(VllmConfig, config)
def _make_kv_cache_config() -> KVCacheConfig:
num_blocks = 16
kv_tensor = KVCacheTensor(
size=num_blocks * 8,
shared_by=["layer"],
block_stride=0,
)
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[kv_tensor],
kv_cache_groups=[
KVCacheGroupSpec(
["layer"],
FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
)
],
)
def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
num_blocks = 4
if packed:
kv_cache_tensors = [
KVCacheTensor(
size=64,
shared_by=[layer_name],
block_stride=16,
)
for layer_name in ("layer0", "layer1")
]
else:
kv_cache_tensors = [
KVCacheTensor(size=40, shared_by=["layer0"]),
KVCacheTensor(size=24, shared_by=["layer1"]),
]
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=[
KVCacheGroupSpec(
["layer0", "layer1"],
FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
)
],
)
def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
return FullAttentionSpec(
block_size=block_size,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
)
def _mla_spec(
block_size: int = 16,
head_size: int = 512,
dtype: torch.dtype = torch.float32,
) -> MLAAttentionSpec:
return MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=head_size,
dtype=dtype,
)
def _make_mla_kv_cache_config(
layer_names: list[str] | None = None,
head_size: int = 512,
dtype: torch.dtype = torch.float32,
num_blocks: int = 4,
) -> KVCacheConfig:
if layer_names is None:
layer_names = ["layer0", "layer1"]
spec = _mla_spec(head_size=head_size, dtype=dtype)
kv_cache_tensors = [
KVCacheTensor(
size=spec.page_size_bytes * num_blocks,
shared_by=[layer_name],
)
for layer_name in layer_names
]
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
)
def _make_hybrid_kv_cache_config() -> KVCacheConfig:
return KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(size=40, shared_by=["full_layer"]),
KVCacheTensor(size=24, shared_by=["mla_layer"]),
],
kv_cache_groups=[
KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)),
KVCacheGroupSpec(["mla_layer"], _mla_spec()),
],
)
def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
config = _make_vllm_config()
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=kv_cache_groups,
)
return build_offloading_config(
config, kv_cache_config
).parallel.is_parallelism_agnostic
def _replicated_layout(
kv_cache_config: KVCacheConfig,
*,
tensor_parallel_size: int = 4,
pipeline_parallel_size: int = 1,
prefill_context_parallel_size: int = 1,
decode_context_parallel_size: int = 1,
use_mla: bool = True,
use_v2_model_runner: bool = False,
distributed_executor_backend: Any = "mp",
nnodes: int = 1,
world_size: int | None = None,
) -> bool:
config = _make_vllm_config(
tensor_parallel_size=tensor_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
prefill_context_parallel_size=prefill_context_parallel_size,
decode_context_parallel_size=decode_context_parallel_size,
)
config.model_config.use_mla = use_mla
config.use_v2_model_runner = use_v2_model_runner
config.parallel_config.distributed_executor_backend = distributed_executor_backend
config.parallel_config.nnodes = nnodes
if world_size is not None:
config.parallel_config.world_size = world_size
return build_offloading_config(config, kv_cache_config).replicated_layout
@pytest.mark.parametrize("packed", [False, True])
def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
config = _make_vllm_config(
extra_config={"block_size": 32},
tensor_parallel_size=3,
pipeline_parallel_size=2,
)
offloading_config = build_offloading_config(
config, _make_sizing_kv_cache_config(packed)
)
assert offloading_config.worker_kv_bytes_per_block == 16
assert offloading_config.parallel.world_size == 6
assert offloading_config.cache.blocks_per_chunk == 2
def test_rejects_partially_packed_tensor_layout():
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.kv_cache_tensors[0].block_stride = 16
with pytest.raises(AssertionError):
build_offloading_config(_make_vllm_config(), kv_cache_config)
def test_zero_blocks_skips_tensor_layout_validation():
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.num_blocks = 0
kv_cache_config.kv_cache_tensors[0].block_stride = 16
offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config)
assert offloading_config.worker_kv_bytes_per_block == 0
def test_prefill_context_parallelism_does_not_scale_group_blocks():
config = _make_vllm_config(
extra_config={"block_size": 64},
prefill_context_parallel_size=2,
)
offloading_config = build_offloading_config(config, _make_kv_cache_config())
assert tuple(group.tokens_per_block for group in offloading_config.groups) == (16,)
assert offloading_config.cache.tokens_per_hash == 16
assert offloading_config.cache.blocks_per_chunk == 4
def test_preserves_data_parallel_index():
config = _make_vllm_config()
config.parallel_config.data_parallel_index = 2
offloading_config = build_offloading_config(config, _make_kv_cache_config())
assert offloading_config.parallel.data_parallel_index == 2
def test_resolves_heterogeneous_hybrid_block_sizes():
config = _make_vllm_config()
config.cache_config.block_size = 4
offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config())
assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
12,
16,
)
assert offloading_config.cache.tokens_per_hash == 4
assert offloading_config.cache.blocks_per_chunk == 1
@pytest.mark.parametrize("world_size", [2, 4, 8])
@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
world_size: int,
use_v2_model_runner: bool,
):
assert _replicated_layout(
_make_mla_kv_cache_config(),
tensor_parallel_size=world_size,
use_v2_model_runner=use_v2_model_runner,
)
@pytest.mark.parametrize(
("kv_cache_config", "case"),
[
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["layer"],
)
],
kv_cache_groups=[
KVCacheGroupSpec(
["layer"],
SlidingWindowMLASpec(
block_size=16,
num_kv_heads=1,
head_size=512,
dtype=torch.float32,
sliding_window=128,
),
)
],
),
"sliding-window-mla",
),
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["layer"],
)
],
kv_cache_groups=[
KVCacheGroupSpec(
["layer"],
HiddenStateCacheSpec(
block_size=16,
num_kv_heads=1,
head_size=512,
dtype=torch.float32,
),
)
],
),
"hidden-state",
),
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["layer0"],
),
KVCacheTensor(
size=_mla_spec(head_size=256).page_size_bytes * 4,
shared_by=["layer1"],
),
],
kv_cache_groups=[
KVCacheGroupSpec(
["layer0", "layer1"],
UniformTypeKVCacheSpecs(
block_size=16,
kv_cache_specs={
"layer0": _mla_spec(),
"layer1": _mla_spec(head_size=256),
},
),
)
],
),
"uniform-wrapper",
),
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["mla"],
),
KVCacheTensor(
size=_full_attention_spec().page_size_bytes * 4,
shared_by=["full"],
),
],
kv_cache_groups=[
KVCacheGroupSpec(["mla"], _mla_spec()),
KVCacheGroupSpec(["full"], _full_attention_spec()),
],
),
"mla-full-hybrid",
),
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["mla"],
),
KVCacheTensor(size=64 * 4, shared_by=["mamba"]),
],
kv_cache_groups=[
KVCacheGroupSpec(["mla"], _mla_spec()),
KVCacheGroupSpec(
["mamba"],
MambaSpec(
block_size=16,
shapes=((16, 1),),
dtypes=(torch.float32,),
),
),
],
),
"mla-mamba-hybrid",
),
(
KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["layer0"],
),
KVCacheTensor(
size=_mla_spec().page_size_bytes * 4,
shared_by=["layer1"],
),
],
kv_cache_groups=[
KVCacheGroupSpec(["layer0"], _mla_spec()),
KVCacheGroupSpec(["layer1"], _mla_spec()),
],
),
"multi-group-mla",
),
],
ids=[
"sliding-window-mla",
"hidden-state",
"uniform-wrapper",
"mla-full-hybrid",
"mla-mamba-hybrid",
"multi-group-mla",
],
)
def test_replicated_layout_excludes_unproven_cache_shapes(
kv_cache_config: KVCacheConfig,
case: str,
):
assert not _replicated_layout(kv_cache_config), case
def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
num_blocks = 4
main_spec = _mla_spec(head_size=512)
indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
main_layers = [f"main_{i}" for i in range(61)]
indexer_layers = [f"indexer_{i}" for i in range(61)]
kv_cache_config = KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[
KVCacheTensor(
size=main_spec.page_size_bytes * len(main_layers) * num_blocks,
shared_by=main_layers,
),
KVCacheTensor(
size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks,
shared_by=indexer_layers,
),
],
kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
)
assert not _replicated_layout(kv_cache_config)
@pytest.mark.parametrize(
("kwargs", "case"),
[
({"tensor_parallel_size": 1}, "tp1"),
({"use_mla": False}, "use-mla-false"),
({"pipeline_parallel_size": 2, "world_size": 4}, "pp2"),
({"prefill_context_parallel_size": 2, "world_size": 4}, "pcp2"),
({"decode_context_parallel_size": 2}, "dcp2"),
({"world_size": 8}, "world-ne-tp"),
({"distributed_executor_backend": "ray"}, "ray"),
({"distributed_executor_backend": "uni"}, "uni"),
({"distributed_executor_backend": type("DummyExecutor", (), {})}, "class"),
({"nnodes": 2}, "multi-node"),
],
ids=[
"tp1",
"use-mla-false",
"pp2",
"pcp2",
"dcp2",
"world-ne-tp",
"ray",
"uni",
"class",
"multi-node",
],
)
def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str):
assert not _replicated_layout(_make_mla_kv_cache_config(), **kwargs), case
def test_parallelism_agnostic_for_single_full_attention_group():
assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
@pytest.mark.parametrize(
"kv_cache_groups",
[
[KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))],
[
KVCacheGroupSpec(
["l0"],
SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
),
)
],
[
KVCacheGroupSpec(["l0"], _full_attention_spec()),
KVCacheGroupSpec(["l1"], _full_attention_spec()),
],
],
)
def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
assert not _parallelism_agnostic(kv_cache_groups)
def test_parallelism_agnostic_disabled_on_v2_model_runner():
config = _make_vllm_config()
config.use_v2_model_runner = True
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
)
offloading_config = build_offloading_config(config, kv_cache_config)
assert not offloading_config.parallel.is_parallelism_agnostic
def test_accepts_blocks_per_chunk_for_heterogeneous_groups():
config = _make_vllm_config(extra_config={"blocks_per_chunk": 2})
offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config())
assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
12,
16,
)
assert offloading_config.cache.blocks_per_chunk == 2
def test_block_size_and_blocks_per_chunk_are_mutually_exclusive():
config = _make_vllm_config(extra_config={"block_size": 64, "blocks_per_chunk": 2})
with pytest.raises(ValueError, match="Specify only one"):
build_offloading_config(config, _make_kv_cache_config())
def test_blocks_per_chunk_must_be_positive():
config = _make_vllm_config(extra_config={"blocks_per_chunk": 0})
with pytest.raises(ValueError, match="greater than 0"):
build_offloading_config(config, _make_kv_cache_config())
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+47
View File
@@ -14,6 +14,7 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
OffloadingWorkerMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
OffloadingConnectorStats,
@@ -40,6 +41,7 @@ from vllm.v1.kv_offload.base import (
get_offload_block_hash,
make_offload_key,
)
from vllm.v1.outputs import KVConnectorOutput
from vllm.v1.request import RequestStatus
@@ -1539,6 +1541,51 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
assert runner.manager.complete_store.call_count == 0
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_complete_store_waits_for_all_worker_acks(
request_runner, async_scheduling: bool
):
tokens_per_block = 4
blocks_per_chunk = 3
tokens_per_chunk = tokens_per_block * blocks_per_chunk
runner = request_runner(
blocks_per_chunk=blocks_per_chunk,
block_size=tokens_per_block,
num_gpu_blocks=100,
async_scheduling=async_scheduling,
worker_count=3,
)
runner.new_request(token_ids=[0] * tokens_per_chunk)
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
generate_store_output(keys)
)
runner.run(decoded_tokens=[0, 0], complete_transfers=False)
assert len(runner.connector_scheduler._jobs) == 1
job_id = next(iter(runner.connector_scheduler._jobs))
assert runner.connector_scheduler._jobs[job_id].pending_count == 3
runner.manager.complete_store.reset_mock()
runner.connector_scheduler.update_connector_output(
KVConnectorOutput(
kv_connector_worker_meta=OffloadingWorkerMetadata(
completed_jobs={job_id: 1}
)
)
)
assert runner.manager.complete_store.call_count == 0
assert runner.connector_scheduler._jobs[job_id].pending_count == 2
runner.connector_scheduler.update_connector_output(
KVConnectorOutput(
kv_connector_worker_meta=OffloadingWorkerMetadata(
completed_jobs={job_id: 2}
)
)
)
assert runner.manager.complete_store.call_count == 1
assert job_id not in runner.connector_scheduler._jobs
@pytest.mark.parametrize("async_scheduling", [True, False])
def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
"""Validates max_offload_tokens: type coercion, boundary values, and capping.
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+184 -1
View File
@@ -6,6 +6,10 @@ from unittest.mock import MagicMock
import pytest
import torch
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
OffloadingConnectorMetadata,
TransferJob,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import get_dtype_size
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -22,7 +26,17 @@ from vllm.v1.kv_cache_interface import (
from vllm.v1.kv_offload.base import (
CanonicalKVCacheRef,
CanonicalKVCaches,
GPULoadStoreSpec,
LoadStoreSpec,
OffloadingManager,
OffloadingSpec,
OffloadingWorker,
)
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
NUM_BLOCKS = 10
@@ -87,7 +101,11 @@ def _allocate_and_reshape_kv_caches(
set_kv_cache_layout(None)
def _make_worker(kv_cache_config: KVCacheConfig):
def _make_worker(
kv_cache_config: KVCacheConfig,
replicated_layout: bool = False,
rank: int = 0,
):
"""
Create an OffloadingConnectorWorker with mocked dependencies.
"""
@@ -96,6 +114,9 @@ def _make_worker(kv_cache_config: KVCacheConfig):
)
spec = MagicMock(spec=OffloadingSpec)
spec.replicated_layout = replicated_layout
spec.config = MagicMock()
spec.config.parallel.rank = rank
spec.get_worker.return_value = MagicMock()
worker = OffloadingConnectorWorker(
@@ -107,11 +128,173 @@ def _make_worker(kv_cache_config: KVCacheConfig):
return worker, spec
def _store_metadata(job_id: int) -> OffloadingConnectorMetadata:
return OffloadingConnectorMetadata(
load_jobs={},
store_jobs={
job_id: TransferJob(
req_id="req",
src_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)),
dst_spec=LoadStoreSpec(),
)
},
)
def _load_metadata(job_id: int) -> OffloadingConnectorMetadata:
return OffloadingConnectorMetadata(
load_jobs={
job_id: TransferJob(
req_id="req",
src_spec=LoadStoreSpec(),
dst_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)),
)
},
store_jobs={},
)
def _empty_metadata() -> OffloadingConnectorMetadata:
return OffloadingConnectorMetadata(load_jobs={}, store_jobs={})
def _offloading_config(rank: int = 0) -> OffloadingConfig:
return OffloadingConfig(
groups=(),
worker_kv_bytes_per_block=0,
enable_kv_cache_events=False,
extra_config={},
engine_id="test-engine",
model=OffloadingModelConfig(name="test-model", dtype="float16"),
cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
parallel=OffloadingParallelConfig(
rank=rank,
world_size=2,
tp_size=2,
pp_size=1,
pcp_size=1,
dcp_size=1,
data_parallel_index=0,
is_parallelism_agnostic=False,
),
)
class BareExternalOffloadingSpec(OffloadingSpec):
def get_manager(self) -> OffloadingManager:
raise NotImplementedError
def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
raise NotImplementedError
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def test_prepare_store_kv_non_writer_marks_completed_without_submit():
worker, _ = _make_worker(
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
replicated_layout=True,
rank=1,
)
worker.prepare_store_kv(_store_metadata(7))
worker.start_kv_transfers(_empty_metadata())
assert worker._unsubmitted_store_jobs == []
assert worker.worker is not None
worker.worker.submit_store.assert_not_called()
meta = worker.build_connector_worker_meta()
assert meta is not None
assert meta.completed_jobs == {7: 1}
def test_prepare_store_kv_writer_submits_store():
worker, _ = _make_worker(
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
replicated_layout=True,
rank=0,
)
worker.prepare_store_kv(_store_metadata(8))
assert worker.build_connector_worker_meta() is None
worker.start_kv_transfers(_empty_metadata())
assert worker.worker is not None
worker.worker.submit_store.assert_called_once()
def test_prepare_store_kv_non_replicated_rank_gt_zero_queues_store():
worker, _ = _make_worker(
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
replicated_layout=False,
rank=1,
)
worker.prepare_store_kv(_store_metadata(9))
assert worker.build_connector_worker_meta() is None
assert len(worker._unsubmitted_store_jobs) == 1
worker.start_kv_transfers(_empty_metadata())
assert worker.worker is not None
worker.worker.submit_store.assert_called_once()
assert worker._unsubmitted_store_jobs == []
def test_handle_preemptions_non_writer_acks_flushed_store():
worker, _ = _make_worker(
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
replicated_layout=True,
rank=1,
)
metadata = _store_metadata(10)
metadata.jobs_to_flush = {10}
worker.handle_preemptions(metadata)
assert worker.worker is not None
worker.worker.submit_store.assert_not_called()
worker.worker.wait.assert_called_once_with({10})
assert metadata.store_jobs == {}
meta = worker.build_connector_worker_meta()
assert meta is not None
assert meta.completed_jobs == {10: 1}
def test_start_kv_transfers_non_writer_still_submits_load():
worker, _ = _make_worker(
KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
replicated_layout=True,
rank=1,
)
worker.start_kv_transfers(_load_metadata(10))
assert worker.worker is not None
worker.worker.submit_load.assert_called_once()
assert worker.build_connector_worker_meta() is None
def test_offloading_connector_worker_accepts_plugin_spec_default_layout():
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
OffloadingConnectorWorker,
)
spec = BareExternalOffloadingSpec(_offloading_config(rank=1))
OffloadingConnectorWorker(
spec=spec,
kv_cache_config=KVCacheConfig(
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
),
)
assert spec.replicated_layout is False
@pytest.mark.parametrize("backend", ATTN_BACKENDS)
def test_register_kv_caches(backend):
"""Test register_kv_caches with multiple groups covering all layer types.
tests/v1/kv_connector/unit/offloading_connector/utils.py
+4
View File
@@ -179,6 +179,7 @@ class RequestRunner:
async_scheduling: bool = True,
kv_cache_groups: list[KVCacheGroupSpec] | None = None,
extra_config_overrides: dict[str, Any] | None = None,
worker_count: int = 1,
):
assert blocks_per_chunk == 1 or kv_cache_groups is None, (
"blocks_per_chunk > 1 requires all groups to have the same "
@@ -198,6 +199,7 @@ class RequestRunner:
disable_hybrid_kv_cache_manager=False,
)
vllm_config.scheduler_config.async_scheduling = async_scheduling
vllm_config.parallel_config.world_size = worker_count
extra_config: dict[str, Any] = {
"spec_name": "MockOffloadingSpec",
@@ -668,6 +670,7 @@ def request_runner():
blocks_per_chunk=1,
kv_cache_groups=None,
extra_config_overrides=None,
worker_count=1,
):
runner = RequestRunner(
block_size=block_size,
@@ -676,6 +679,7 @@ def request_runner():
async_scheduling=async_scheduling,
kv_cache_groups=kv_cache_groups,
extra_config_overrides=extra_config_overrides,
worker_count=worker_count,
)
runners.append(runner)
return runner
tests/v1/kv_offload/cpu/test_gpu_worker.py
+10 -2
View File
@@ -41,7 +41,10 @@ NUM_MAPPINGS_PER_GROUP = [2]
@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", DEVICES)
@pytest.mark.parametrize("use_shared_memory", [False, True])
@pytest.mark.parametrize(
("use_shared_memory", "replicated_layout"),
[(False, False), (True, False), (True, True)],
)
@torch.inference_mode()
def test_transfer(
default_vllm_config,
@@ -55,6 +58,7 @@ def test_transfer(
seed: int,
device: str,
use_shared_memory: bool,
replicated_layout: bool,
) -> None:
set_random_seed(seed)
@@ -95,11 +99,15 @@ def test_transfer(
gpu_page_size_bytes * num_tensors * blocks_per_chunk,
SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
)
simulated_world_size = 2
kv_bytes_per_block = (
cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size
)
mmap_region = SharedOffloadRegion(
engine_id=str(uuid.uuid4()),
num_blocks=num_cpu_blocks,
rank=0,
kv_bytes_per_block=cpu_page_size,
kv_bytes_per_block=kv_bytes_per_block,
cpu_page_size=cpu_page_size,
)
tests/v1/kv_offload/test_factory.py
+223 -484
View File
@@ -1,35 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for OffloadingSpecFactory.
"""Unit tests for native offloading specs and their factory."""
These tests verify:
1. Pre-registration integrity — registered module paths can actually import
and yield correct OffloadingSpec subclasses (CI sentinel against file moves).
2. End-to-end factory → spec construction with real configs.
3. Downstream collaboration — build_metric_definitions delegation.
4. Error paths — unregistered specs, missing config, duplicate registration.
"""
from typing import cast
from unittest.mock import MagicMock, patch
from typing import Any
from unittest.mock import MagicMock
import pytest
import torch
from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
build_offloading_config,
)
from vllm.platforms import current_platform
from vllm.v1.kv_cache_interface import (
FullAttentionSpec,
KVCacheConfig,
KVCacheGroupSpec,
KVCacheTensor,
MLAAttentionSpec,
SlidingWindowSpec,
)
from vllm.v1.kv_offload.base import (
CanonicalKVCaches,
OffloadingHistogramMetadata,
@@ -37,220 +14,81 @@ from vllm.v1.kv_offload.base import (
OffloadingSpec,
OffloadingWorker,
)
from vllm.v1.kv_offload.config import (
OffloadingCacheConfig,
OffloadingConfig,
OffloadingGroupConfig,
OffloadingModelConfig,
OffloadingParallelConfig,
)
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture(autouse=True)
def restore_registry():
"""Save and restore OffloadingSpecFactory._registry between tests."""
original = dict(OffloadingSpecFactory._registry)
yield
OffloadingSpecFactory._registry = original
def _get_extra_config(config: VllmConfig) -> dict:
assert config.kv_transfer_config is not None
return config.kv_transfer_config.kv_connector_extra_config
def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec:
return OffloadingSpecFactory.create_spec(
build_offloading_config(config, kv_cache_config)
)
def _make_vllm_config(
def _make_offloading_config(
*,
spec_name: str | None = "CPUOffloadingSpec",
cpu_bytes_to_use: int | None = None,
store_threshold: int = 0,
extra_config: dict | None = None,
):
"""Build a real VllmConfig with kv_transfer_config set for offloading."""
from vllm.config import (
CacheConfig,
DeviceConfig,
ModelConfig,
SchedulerConfig,
VllmConfig,
)
model_config = ModelConfig(
model="facebook/opt-125m",
trust_remote_code=True,
dtype="float16",
seed=42,
)
scheduler_config = SchedulerConfig(
max_num_seqs=16,
max_num_batched_tokens=64,
max_model_len=10000,
enable_chunked_prefill=True,
is_encoder_decoder=model_config.is_encoder_decoder,
)
cache_config = CacheConfig(
block_size=16,
gpu_memory_utilization=0.9,
cache_dtype="auto",
enable_prefix_caching=True,
)
cfg = extra_config or {}
cpu_bytes_to_use: int | None = 65536,
worker_kv_bytes_per_block: int = 8,
groups: tuple[OffloadingGroupConfig, ...] | None = None,
tokens_per_hash: int = 16,
blocks_per_chunk: int = 1,
rank: int = 0,
world_size: int = 1,
tp_size: int | None = None,
pp_size: int = 1,
pcp_size: int = 1,
dcp_size: int = 1,
data_parallel_index: int = 0,
is_parallelism_agnostic: bool = False,
replicated_layout: bool = False,
extra_config: dict[str, Any] | None = None,
) -> OffloadingConfig:
normalized_extra_config = dict(extra_config or {})
if spec_name is not None:
normalized_extra_config["spec_name"] = spec_name
if cpu_bytes_to_use is not None:
cfg["cpu_bytes_to_use"] = cpu_bytes_to_use
cfg["spec_name"] = spec_name
if store_threshold > 0:
cfg["store_threshold"] = store_threshold
normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
kv_transfer_config = KVTransferConfig(
kv_connector="OffloadingConnector",
kv_role="kv_both",
kv_connector_extra_config=cfg,
)
return VllmConfig(
scheduler_config=scheduler_config,
model_config=model_config,
cache_config=cache_config,
kv_transfer_config=kv_transfer_config,
device_config=DeviceConfig("cpu"),
if groups is None:
groups = (OffloadingGroupConfig(16, ("layer",)),)
return OffloadingConfig(
groups=groups,
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
enable_kv_cache_events=False,
extra_config=normalized_extra_config,
engine_id="test-engine",
model=OffloadingModelConfig(name="test-model", dtype="float16"),
cache=OffloadingCacheConfig(
tokens_per_hash=tokens_per_hash,
blocks_per_chunk=blocks_per_chunk,
),
parallel=OffloadingParallelConfig(
rank=rank,
world_size=world_size,
tp_size=world_size if tp_size is None else tp_size,
pp_size=pp_size,
pcp_size=pcp_size,
dcp_size=dcp_size,
data_parallel_index=data_parallel_index,
is_parallelism_agnostic=is_parallelism_agnostic,
),
replicated_layout=replicated_layout,
)
def _make_layout_vllm_config(
spec_name: str = "CPUOffloadingSpec",
cpu_bytes_to_use: int | None = None,
extra_config: dict | None = None,
tensor_parallel_size: int = 1,
pipeline_parallel_size: int = 1,
prefill_context_parallel_size: int = 1,
decode_context_parallel_size: int = 1,
) -> VllmConfig:
config = MagicMock()
config.cache_config.block_size = 16
config.cache_config.enable_prefix_caching = True
config.cache_config.prefix_match_unit = None
config.cache_config.cache_dtype = torch.float16
config.model_config.model = "test-model"
world_size = (
tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
)
with patch.object(current_platform, "device_count", return_value=world_size):
config.parallel_config = ParallelConfig(
tensor_parallel_size=tensor_parallel_size,
pipeline_parallel_size=pipeline_parallel_size,
prefill_context_parallel_size=prefill_context_parallel_size,
decode_context_parallel_size=decode_context_parallel_size,
)
config.kv_events_config = None
config.use_v2_model_runner = False
connector_extra_config = dict(extra_config or {})
connector_extra_config["spec_name"] = spec_name
if cpu_bytes_to_use is not None:
connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
config.kv_transfer_config = KVTransferConfig(
kv_connector="OffloadingConnector",
kv_role="kv_both",
kv_connector_extra_config=connector_extra_config,
)
return cast(VllmConfig, config)
def _make_kv_cache_config():
"""Build a minimal KVCacheConfig with one KV cache tensor."""
num_blocks = 16
num_kv_heads = 1
head_size = 1
dtype = torch.float32
page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
kv_tensor = KVCacheTensor(
size=num_blocks * page_size, shared_by=["layer"], block_stride=0
)
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=[kv_tensor],
kv_cache_groups=[
KVCacheGroupSpec(
["layer"],
FullAttentionSpec(
block_size=16,
num_kv_heads=num_kv_heads,
head_size=head_size,
dtype=dtype,
),
)
],
)
def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
num_blocks = 4
if packed:
kv_cache_tensors = [
KVCacheTensor(
size=64,
shared_by=[layer_name],
block_stride=16,
)
for layer_name in ("layer0", "layer1")
]
else:
kv_cache_tensors = [
KVCacheTensor(size=40, shared_by=["layer0"]),
KVCacheTensor(size=24, shared_by=["layer1"]),
]
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=[
KVCacheGroupSpec(
["layer0", "layer1"],
FullAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
)
],
)
def _make_hybrid_kv_cache_config() -> KVCacheConfig:
return KVCacheConfig(
num_blocks=4,
kv_cache_tensors=[
KVCacheTensor(size=40, shared_by=["full_layer"]),
KVCacheTensor(size=24, shared_by=["mla_layer"]),
],
kv_cache_groups=[
KVCacheGroupSpec(
["full_layer"],
FullAttentionSpec(
block_size=12,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
),
),
KVCacheGroupSpec(
["mla_layer"],
MLAAttentionSpec(
block_size=16,
num_kv_heads=1,
head_size=576,
dtype=torch.float32,
),
),
],
)
def _create_spec(**kwargs: Any) -> OffloadingSpec:
return OffloadingSpecFactory.create_spec(_make_offloading_config(**kwargs))
class SingleArgExternalOffloadingSpec(OffloadingSpec):
@@ -261,104 +99,60 @@ class SingleArgExternalOffloadingSpec(OffloadingSpec):
raise NotImplementedError
# ---------------------------------------------------------------------------
# Pre-registration integrity (CI sentinel)
# ---------------------------------------------------------------------------
def test_pre_registered_specs_can_be_imported():
"""If someone moves cpu/spec.py but forgets to update factory.py, CI fails."""
for name in OffloadingSpecFactory._registry:
cls = OffloadingSpecFactory._registry[name]()
assert issubclass(cls, OffloadingSpec)
def test_cpu_spec_registered():
"""CPUOffloadingSpec is registered and importable."""
cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]()
assert cls is CPUOffloadingSpec
def test_tiering_spec_registered():
"""TieringOffloadingSpec is registered and importable."""
cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]()
assert cls is TieringOffloadingSpec
# ---------------------------------------------------------------------------
# Normal path — get_spec_cls
# ---------------------------------------------------------------------------
def test_get_spec_cls_returns_registered_class():
"""Registered spec_name returns correct class."""
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
spec_cls = OffloadingSpecFactory.get_spec_cls(
_make_offloading_config().extra_config
)
assert spec_cls is CPUOffloadingSpec
def test_get_spec_cls_default_to_cpu():
"""Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
config = _make_vllm_config(spec_name=None)
config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
def test_get_spec_cls_defaults_to_cpu():
spec_cls = OffloadingSpecFactory.get_spec_cls(
_make_offloading_config(spec_name=None).extra_config
)
assert spec_cls is CPUOffloadingSpec
# ---------------------------------------------------------------------------
# End-to-end — create_spec
# ---------------------------------------------------------------------------
def test_create_cpu_offloading_spec_end_to_end():
"""Full factory → spec construction with real VllmConfig/KVCacheConfig.
Verifies:
- cpu_bytes_to_use validation and num_blocks calculation
- block_size % tokens_per_hash assertion
- spec instance is CPUOffloadingSpec
"""
config = _make_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_kv_cache_config()
spec = _create_spec(config, kv_cache_config)
def test_create_cpu_offloading_spec():
spec = _create_spec()
assert isinstance(spec, CPUOffloadingSpec)
assert spec.num_blocks > 0
@pytest.mark.parametrize("packed", [False, True])
def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
cpu_bytes_to_use = 1920
config = _make_layout_vllm_config(
cpu_bytes_to_use=cpu_bytes_to_use,
extra_config={"block_size": 32},
tensor_parallel_size=3,
pipeline_parallel_size=2,
def test_cpu_spec_sizes_normalized_worker_layout():
spec = _create_spec(
cpu_bytes_to_use=1920,
worker_kv_bytes_per_block=16,
blocks_per_chunk=2,
world_size=6,
tp_size=3,
pp_size=2,
)
spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
assert isinstance(spec, CPUOffloadingSpec)
assert spec.cpu_page_size_per_worker == 32
assert spec.kv_bytes_per_chunk == 192
assert spec.num_blocks == cpu_bytes_to_use // 192
assert spec.num_blocks == 10
def test_cpu_spec_rejects_partially_packed_tensor_layout():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.kv_cache_tensors[0].block_stride = 16
with pytest.raises(AssertionError):
_create_spec(config, kv_cache_config)
def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
kv_cache_config = _make_sizing_kv_cache_config(packed=False)
kv_cache_config.num_blocks = 0
kv_cache_config.kv_cache_tensors[0].block_stride = 16
spec = _create_spec(config, kv_cache_config)
def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
assert isinstance(spec, CPUOffloadingSpec)
assert spec.cpu_page_size_per_worker == 0
@@ -368,225 +162,209 @@ def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
def test_tiering_spec_aligns_row_size():
alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
cpu_bytes_to_use = alignment * 3
config = _make_layout_vllm_config(
spec = _create_spec(
spec_name="TieringOffloadingSpec",
cpu_bytes_to_use=cpu_bytes_to_use,
extra_config={"block_size": 32},
tensor_parallel_size=3,
pipeline_parallel_size=2,
cpu_bytes_to_use=alignment * 3,
worker_kv_bytes_per_block=16,
blocks_per_chunk=2,
world_size=6,
tp_size=3,
pp_size=2,
)
spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
assert isinstance(spec, TieringOffloadingSpec)
assert spec.cpu_page_size_per_worker == 32
assert spec.kv_bytes_per_chunk == alignment
assert spec.num_blocks == cpu_bytes_to_use // alignment
assert spec.num_blocks == 3
def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={"block_size": 64},
prefill_context_parallel_size=2,
@pytest.mark.parametrize("world_size", [2, 4, 8])
def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int):
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
spec = _create_spec(
spec_name="TieringOffloadingSpec",
cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
world_size=world_size,
replicated_layout=True,
)
spec = _create_spec(config, _make_kv_cache_config())
assert spec.tokens_per_block == (16,)
assert spec.tokens_per_hash == 16
assert spec.blocks_per_chunk == 4
assert isinstance(spec, TieringOffloadingSpec)
assert spec.replicated_layout is True
assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block
assert spec.num_blocks == 8
def test_offloading_config_preserves_data_parallel_index():
config = _make_layout_vllm_config()
config.parallel_config.data_parallel_index = 2
def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch):
import vllm.v1.kv_offload.tiering.spec as tiering_spec_module
offloading_config = build_offloading_config(config, _make_kv_cache_config())
worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
spec = _create_spec(
spec_name="TieringOffloadingSpec",
cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
world_size=4,
replicated_layout=True,
)
assert isinstance(spec, TieringOffloadingSpec)
assert offloading_config.parallel.data_parallel_index == 2
region = MagicMock()
region_calls: list[dict[str, Any]] = []
worker_calls: list[dict[str, Any]] = []
def fake_region_ctor(**kwargs):
region_calls.append(kwargs)
return region
def fake_worker_ctor(**kwargs):
worker_calls.append(kwargs)
return MagicMock()
monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor)
monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
monkeypatch.setattr(
tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5
)
kv_caches = MagicMock()
spec.create_worker(kv_caches)
assert region_calls[0]["rank"] == 0
assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block
assert worker_calls[0]["kv_caches"] is kv_caches
assert worker_calls[0]["mmap_region"] is region
def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes():
config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
config.cache_config.block_size = 4
def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch):
import vllm.v1.kv_offload.tiering.spec as tiering_spec_module
spec = _create_spec(config, _make_hybrid_kv_cache_config())
spec = _create_spec(
spec_name="TieringOffloadingSpec",
worker_kv_bytes_per_block=4096,
world_size=4,
)
assert isinstance(spec, TieringOffloadingSpec)
region_calls: list[dict[str, Any]] = []
def fake_region_ctor(**kwargs):
region_calls.append(kwargs)
return MagicMock()
monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor)
monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock())
monkeypatch.setattr(
tiering_spec_module.torch.accelerator,
"current_device_index",
lambda: 5,
)
spec.create_worker(MagicMock())
assert region_calls[0]["rank"] == 1
@pytest.mark.parametrize("world_size", [2, 4, 8])
def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
worker_kv_bytes_per_block = 4096
spec = _create_spec(
cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
worker_kv_bytes_per_block=worker_kv_bytes_per_block,
world_size=world_size,
replicated_layout=True,
)
assert isinstance(spec, CPUOffloadingSpec)
assert spec.replicated_layout is False
assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size
assert spec.num_blocks == 2
def test_offloading_spec_has_replicated_layout_default():
spec = SingleArgExternalOffloadingSpec(_make_offloading_config())
assert spec.replicated_layout is False
def test_offloading_spec_uses_normalized_chunk_geometry():
groups = (
OffloadingGroupConfig(12, ("full_layer",)),
OffloadingGroupConfig(16, ("mla_layer",)),
)
spec = _create_spec(
groups=groups,
tokens_per_hash=4,
blocks_per_chunk=2,
)
assert spec.tokens_per_block == (12, 16)
assert spec.tokens_per_hash == 4
assert spec.blocks_per_chunk == 1
assert spec.blocks_per_chunk == 2
def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
return FullAttentionSpec(
block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32
)
def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
config = _make_layout_vllm_config()
kv_cache_config = KVCacheConfig(
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
)
offloading_config = build_offloading_config(config, kv_cache_config)
return offloading_config.parallel.is_parallelism_agnostic
def test_parallelism_agnostic_for_single_full_attention_group():
assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
@pytest.mark.parametrize(
"kv_cache_groups",
[
# MLA latent KV is replicated per rank, never head-sharded.
[
KVCacheGroupSpec(
["l0"],
MLAAttentionSpec(
block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
),
)
],
# Sliding window is not full attention.
[
KVCacheGroupSpec(
["l0"],
SlidingWindowSpec(
block_size=16,
num_kv_heads=4,
head_size=128,
dtype=torch.float32,
sliding_window=128,
),
)
],
# Hybrid model: more than one KV cache group.
[
KVCacheGroupSpec(["l0"], _full_attention_spec()),
KVCacheGroupSpec(["l1"], _full_attention_spec()),
],
],
)
def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
assert not _parallelism_agnostic(kv_cache_groups)
def test_parallelism_agnostic_disabled_on_v2_model_runner():
config = _make_layout_vllm_config()
config.use_v2_model_runner = True
kv_cache_config = KVCacheConfig(
num_blocks=0,
kv_cache_tensors=[],
kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
)
offloading_config = build_offloading_config(config, kv_cache_config)
assert not offloading_config.parallel.is_parallelism_agnostic
def test_create_dynamic_spec_receives_translated_config():
config = _make_layout_vllm_config(
def test_create_dynamic_spec_receives_config():
config = _make_offloading_config(
spec_name="SingleArgExternalOffloadingSpec",
extra_config={
"spec_module_path": "tests.v1.kv_offload.test_factory",
},
extra_config={"spec_module_path": "tests.v1.kv_offload.test_factory"},
)
kv_cache_config = _make_kv_cache_config()
offloading_config = build_offloading_config(config, kv_cache_config)
spec = OffloadingSpecFactory.create_spec(offloading_config)
spec = OffloadingSpecFactory.create_spec(config)
assert isinstance(spec, SingleArgExternalOffloadingSpec)
assert spec.config is offloading_config
# ---------------------------------------------------------------------------
# Dynamic import via spec_module_path
# ---------------------------------------------------------------------------
assert spec.config is config
def test_dynamic_load_via_spec_module_path():
"""External spec loaded via spec_module_path.
This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec)
integrate with vLLM without being pre-registered in the factory.
The fallback path: registry miss → spec_module_path → importlib.import_module.
"""
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
# Delete from registry to force the dynamic import path
del OffloadingSpecFactory._registry["CPUOffloadingSpec"]
# spec_name not in registry → falls through to spec_module_path
config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
"vllm.v1.kv_offload.cpu.spec"
config = _make_offloading_config(
extra_config={"spec_module_path": "vllm.v1.kv_offload.cpu.spec"}
)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
spec_cls = OffloadingSpecFactory.get_spec_cls(config.extra_config)
assert spec_cls is CPUOffloadingSpec
# ---------------------------------------------------------------------------
# Error paths
# ---------------------------------------------------------------------------
def test_unregistered_spec_without_module_path_raises():
"""spec_name not in registry + no spec_module_path → ValueError."""
config = _make_vllm_config(spec_name="NonexistentSpec")
config = _make_offloading_config(spec_name="NonexistentSpec")
with pytest.raises(ValueError, match="Unsupported spec type"):
OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
OffloadingSpecFactory.get_spec_cls(config.extra_config)
# create_spec should also fail (calls get_spec_cls internally)
kv_cache_config = _make_kv_cache_config()
with pytest.raises(ValueError, match="Unsupported spec type"):
_create_spec(config, kv_cache_config)
OffloadingSpecFactory.create_spec(config)
def test_cpu_spec_missing_cpu_bytes_to_use_raises():
"""CPUOffloadingSpec requires cpu_bytes_to_use → Exception."""
config = _make_vllm_config(cpu_bytes_to_use=None)
config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
kv_cache_config = _make_kv_cache_config()
with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
_create_spec(config, kv_cache_config)
_create_spec(cpu_bytes_to_use=None)
def test_duplicate_registration_raises():
"""register_spec with existing name → ValueError."""
with pytest.raises(ValueError, match="is already registered"):
OffloadingSpecFactory.register_spec(
"CPUOffloadingSpec", "some.module", "SomeClass"
)
# ---------------------------------------------------------------------------
# Downstream collaboration — build_metric_definitions
# ---------------------------------------------------------------------------
def test_build_metric_definitions_below_threshold():
"""store_threshold < 2 keeps stores_skipped disabled."""
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=1)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
extra_config = {"store_threshold": 1}
spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
metrics = spec_cls.build_metric_definitions(extra_config)
assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics
def test_build_metric_definitions_allocation_size_histogram():
"""CPU allocation size is always reported as a histogram."""
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=0)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
metrics = spec_cls.build_metric_definitions({})
metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE]
assert isinstance(metadata, OffloadingHistogramMetadata)
assert metadata.buckets == (
1,
@@ -603,49 +381,10 @@ def test_build_metric_definitions_allocation_size_histogram():
def test_build_metric_definitions_returns_counter_at_threshold():
"""store_threshold >= 2 → returns stores_skipped counter definition."""
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
config = _make_vllm_config(store_threshold=2)
spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
metrics = spec_cls.build_metric_definitions(
config.kv_transfer_config.kv_connector_extra_config
)
extra_config = {"store_threshold": 2}
spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
metrics = spec_cls.build_metric_definitions(extra_config)
assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={"blocks_per_chunk": 2},
)
spec = _create_spec(config, _make_hybrid_kv_cache_config())
assert spec.tokens_per_block == (12, 16)
assert spec.blocks_per_chunk == 2
def test_block_size_and_blocks_per_chunk_are_mutually_exclusive():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={
"block_size": 64,
"blocks_per_chunk": 2,
},
)
with pytest.raises(ValueError, match="Specify only one"):
_create_spec(config, _make_kv_cache_config())
def test_blocks_per_chunk_must_be_positive():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={
"blocks_per_chunk": 0,
},
)
with pytest.raises(ValueError, match="greater than 0"):
_create_spec(config, _make_kv_cache_config())
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+28 -7
View File
@@ -102,19 +102,39 @@ def build_offloading_config(
)
worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
# Only a single non-MLA full-attention group is parallelism-invariant:
# MLA latent KV is replicated per rank (never head-sharded), and the V2
# model runner's KV layout is not known to be parallelism-invariant.
single_group = (
single_group_spec = (
kv_cache_config.kv_cache_groups[0].kv_cache_spec
if len(kv_cache_config.kv_cache_groups) == 1
else None
)
replicated_layout = (
vllm_config.model_config.use_mla
# Exact type: fail closed on wrappers and sliding-window variants.
and type(single_group_spec) is MLAAttentionSpec
# Page accounting: one MLA page per layer, no packed/mixed rows.
and worker_kv_bytes_per_block > 0
and worker_kv_bytes_per_block
== single_group_spec.page_size_bytes
* len(kv_cache_config.kv_cache_groups[0].layer_names)
# Safe MVP boundary: TP-only, no other parallel axes.
and parallel_config.tensor_parallel_size > 1
and parallel_config.pipeline_parallel_size == 1
and parallel_config.prefill_context_parallel_size == 1
and parallel_config.decode_context_parallel_size == 1
and parallel_config.world_size == parallel_config.tensor_parallel_size
# Shared /dev/shm mmap layout is single-node mp only.
and parallel_config.distributed_executor_backend == "mp"
and parallel_config.nnodes_within_dp == 1
)
# Only a single non-MLA full-attention group is parallelism-invariant:
# MLA latent KV is replicated per rank (never head-sharded), and the V2
# model runner's KV layout is not known to be parallelism-invariant.
is_parallelism_agnostic = (
not vllm_config.use_v2_model_runner
and single_group is not None
and isinstance(single_group, FullAttentionSpec)
and not isinstance(single_group, MLAAttentionSpec)
and single_group_spec is not None
and isinstance(single_group_spec, FullAttentionSpec)
and not isinstance(single_group_spec, MLAAttentionSpec)
)
kv_events_config = vllm_config.kv_events_config
@@ -144,4 +164,5 @@ def build_offloading_config(
data_parallel_index=parallel_config.data_parallel_index,
is_parallelism_agnostic=is_parallelism_agnostic,
),
replicated_layout=replicated_layout,
)
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+11
View File
@@ -45,6 +45,10 @@ class OffloadingConnectorWorker:
self.spec = spec
self.kv_cache_config = kv_cache_config
self.worker: OffloadingWorker | None = None
# Non-writers still ack: pending_count waits for world_size per job.
self._is_store_writer = (
not self.spec.replicated_layout or self.spec.config.parallel.rank == 0
)
# job_id -> req_id for in-flight loads.
self._load_jobs: dict[int, ReqId] = {}
@@ -274,6 +278,9 @@ class OffloadingConnectorWorker:
for job_id in kv_connector_metadata.jobs_to_flush:
entry = kv_connector_metadata.store_jobs.pop(job_id, None)
if entry is not None:
if not self._is_store_writer:
self._connector_worker_meta.mark_completed(job_id)
continue
assert isinstance(entry.src_spec, GPULoadStoreSpec)
self._unsubmitted_store_jobs.append(
(job_id, entry.src_spec, entry.dst_spec)
@@ -304,6 +311,10 @@ class OffloadingConnectorWorker:
def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
for job_id, entry in metadata.store_jobs.items():
if not self._is_store_writer:
# Gate before queueing: no _unsubmitted_store_jobs entry.
self._connector_worker_meta.mark_completed(job_id)
continue
# NOTE(orozery): defer the store to the beginning of the next
# engine step, so that offloading starts AFTER transfers related
# to token sampling, thereby avoiding delays to token generation.
vllm/v1/kv_offload/base.py
+1
View File
@@ -496,6 +496,7 @@ class OffloadingSpec(ABC):
)
self.config = config
self.extra_config = config.extra_config
self.replicated_layout: bool = False
self.kv_events_config = OffloadingKVEventsConfig(
enable_kv_cache_events=config.enable_kv_cache_events,
self_describing_kv_events=bool(
vllm/v1/kv_offload/config.py
+6
View File
@@ -68,3 +68,9 @@ class OffloadingConfig:
model: OffloadingModelConfig
cache: OffloadingCacheConfig
parallel: OffloadingParallelConfig
# True when the offloaded bytes of every worker are expected to be
# byte-identical per block (pure-MLA model, single-node TP-only
# parallelism), enabling a single-copy host layout in backends that
# support it. Aggregate layout decision; per-layer replication metadata
# is planned for CanonicalKVCacheRef (#48408).
replicated_layout: bool = False
vllm/v1/kv_offload/cpu/spec.py
+8 -2
View File
@@ -24,6 +24,7 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
class CPUOffloadingSpec(OffloadingSpec):
BLOCK_SIZE_ALIGNMENT = 1
SUPPORTS_REPLICATED_LAYOUT = False
@classmethod
def build_metric_definitions(
@@ -85,12 +86,16 @@ class CPUOffloadingSpec(OffloadingSpec):
self.num_blocks = 0
self.kv_bytes_per_chunk = 0
self.cpu_page_size_per_worker = 0
self.replicated_layout = (
config.replicated_layout and self.SUPPORTS_REPLICATED_LAYOUT
)
if config.worker_kv_bytes_per_block > 0 and world_size > 0:
kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size
num_copies = 1 if self.replicated_layout else world_size
kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies
kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk
# calculate cpu_page_size_per_worker
self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size
self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies
# calculate num_blocks
aligned_kv_bytes_per_chunk = round_up(
@@ -102,6 +107,7 @@ class CPUOffloadingSpec(OffloadingSpec):
# kv_bytes_per_chunk. Note that this might contain
# some padding. i.e. each offloaded block is of the form,
# |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
# or |--- B0 (single copy) ---| *** maybe-pad *** |
self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
# scheduler-side
vllm/v1/kv_offload/tiering/spec.py
+7 -3
View File
@@ -71,6 +71,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
"""
BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
SUPPORTS_REPLICATED_LAYOUT = True
@classmethod
@override
@@ -229,10 +230,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
@override
def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
# Fold the global physical device index into the replica-local
# [0, world_size) slot range.
world_size = self.config.parallel.world_size
rank = torch.accelerator.current_device_index() % world_size
if self.replicated_layout:
rank = 0
else:
# Fold the global physical device index into the replica-local
# [0, world_size) slot range.
rank = torch.accelerator.current_device_index() % world_size
worker_mmap = SharedOffloadRegion(
engine_id=self._engine_id,
num_blocks=self.num_blocks,