From 7eca0e1a649e9fb85e88c101432d8bfdbb5b9bfe Mon Sep 17 00:00:00 2001 From: Chang Guo Date: Sat, 25 Jul 2026 22:22:30 -0700 Subject: [PATCH] [KV Offload] Deduplicate replicated MLA KV in the shared CPU region (#48906) Signed-off-by: Change72 Co-authored-by: OpenAI Codex --- .buildkite/test_areas/lm_eval.yaml | 4 +- tests/evals/gsm8k/test_gsm8k_offloading.py | 30 +- .../unit/offloading_connector/test_config.py | 587 +++++++++++++++ .../offloading_connector/test_scheduler.py | 47 ++ .../unit/offloading_connector/test_worker.py | 185 ++++- .../unit/offloading_connector/utils.py | 4 + tests/v1/kv_offload/cpu/test_gpu_worker.py | 12 +- tests/v1/kv_offload/test_factory.py | 707 ++++++------------ .../kv_connector/v1/offloading/config.py | 35 +- .../kv_connector/v1/offloading/worker.py | 11 + vllm/v1/kv_offload/base.py | 1 + vllm/v1/kv_offload/config.py | 6 + vllm/v1/kv_offload/cpu/spec.py | 10 +- vllm/v1/kv_offload/tiering/spec.py | 10 +- 14 files changed, 1143 insertions(+), 506 deletions(-) create mode 100644 tests/v1/kv_connector/unit/offloading_connector/test_config.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 7d0f987d10..f891d9fd24 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -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 diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index f652dcaf1b..7ed97f7efc 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -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}" diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py new file mode 100644 index 0000000000..fc426ff318 --- /dev/null +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -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()) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index c65d02dda4..d1463aebb4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -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. diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index b4557bd7b6..f96fee36a8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -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. diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index c1979a0b1f..9acf5f5d49 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -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 diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 2f1ce67e9c..4e5e94c2bc 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -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, ) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index d08dbee576..6d57f1ce70 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -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()) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index b807b480f0..2a5659c2a9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -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, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index a61ec91754..bbdcb8d0d6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -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. diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 60c0aa4374..51f5162a99 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -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( diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py index cd7b3ee207..46ca489cc2 100644 --- a/vllm/v1/kv_offload/config.py +++ b/vllm/v1/kv_offload/config.py @@ -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 diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index f6d1c29aff..d755bfecbc 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -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 diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 7aae3a2eb7..0d200a4239 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -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,