@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7eca0e1a64

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

Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Chang Guo committed Jul 25, 2026 at 22:22 UTC 7eca0e1a649e9fb85e88c101432d8bfdbb5b9bfe
14 files changed +1135 -498
.buildkite/test_areas/lm_eval.yaml
+2 -2
@@ -337,7 +337,7 @@ steps:
337
338 - label: LM Eval KV-Offload (2xH100)
339 key: kv-offload-medium
340 - timeout_in_minutes: 30
340 + timeout_in_minutes: 45
341 device: h100
342 num_devices: 2
343 source_file_dependencies:
@@ -347,7 +347,7 @@ steps:
347 - vllm/v1/simple_kv_offload/
348 - tests/evals/gsm8k/test_gsm8k_offloading.py
349 commands:
350 - - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b"
350 + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite"
351
352 - label: LM Eval KV-Offload (4xH100)
353 key: kv-offload-large
tests/evals/gsm8k/test_gsm8k_offloading.py
+25 -5
@@ -16,11 +16,13 @@ KV drops accuracy below threshold. The reload itself is not asserted, so
16 a silently skipped reload (e.g. offloading disabled) would not be flagged.
17
18 Covers both KV offloading connectors (OffloadingConnector and
19 -SimpleCPUOffloadConnector) across four architecture families:
19 +SimpleCPUOffloadConnector) across five architecture families:
20 - Hybrid Mamba (NemotronH: attention + Mamba)
21 - Heterogeneous head dim (Gemma 4)
22 - Hybrid GDN (Qwen 3.5: attention + GatedDeltaNet)
23 - Compressed attention (DeepSeek-V4-Flash: CSA)
24 + - Pure MLA (DeepSeek-V2-Lite: TieringOffloadingSpec, TP=2, replicated
25 + single-slot host layout)
26
27 Usage:
28 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py
@@ -47,14 +49,16 @@ NUM_FEWSHOT = 5
49 _OFFLOAD_SYNC_TIMEOUT = 60
50
51
50 -def _kv_transfer_config(connector: str, cpu_gib: int = 4) -> str:
52 +def _kv_transfer_config(
53 + connector: str, cpu_gib: int = 4, spec_name: str = "CPUOffloadingSpec"
54 +) -> str:
55 if connector == "OffloadingConnector":
56 return json.dumps(
57 {
58 "kv_connector": "OffloadingConnector",
59 "kv_role": "kv_both",
60 "kv_connector_extra_config": {
57 - "spec_name": "CPUOffloadingSpec",
61 + "spec_name": spec_name,
62 "cpu_bytes_to_use": cpu_gib << 30,
63 "eviction_policy": "lru",
64 },
@@ -115,8 +119,10 @@ class OffloadingModelConfig:
119 accuracy_threshold: float
120 tolerance: float = 0.05
121 extra_server_args: list[str] = field(default_factory=list)
122 + env_dict: dict[str, str] = field(default_factory=dict)
123 cpu_offload_gib: int = 4
124 startup_timeout: int = 600
125 + spec_name: str = "CPUOffloadingSpec"
126
127
128 MODELS = [
@@ -165,6 +171,20 @@ MODELS = [
171 cpu_offload_gib=16,
172 startup_timeout=1200,
173 ),
174 + OffloadingModelConfig(
175 + id="offloading-deepseek-v2-lite-tiering-tp2",
176 + model="deepseek-ai/DeepSeek-V2-Lite",
177 + connector="OffloadingConnector",
178 + # Baseline 0.360/0.345 over two runs (measured on 2xA100).
179 + accuracy_threshold=0.35,
180 + extra_server_args=[
181 + "--tensor-parallel-size",
182 + "2",
183 + ],
184 + cpu_offload_gib=8,
185 + startup_timeout=1200,
186 + spec_name="TieringOffloadingSpec",
187 + ),
188 # ── SimpleCPUOffloadConnector ────────────────────────────────────
189 OffloadingModelConfig(
190 id="simple-nemotron-h-8b",
@@ -229,7 +249,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
249 "--enable-prefix-caching",
250 "--no-disable-hybrid-kv-cache-manager",
251 "--kv-transfer-config",
232 - _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib),
252 + _kv_transfer_config(cfg.connector, cfg.cpu_offload_gib, cfg.spec_name),
253 "--trust-remote-code",
254 "--disable-uvicorn-access-log",
255 *cfg.extra_server_args,
@@ -239,7 +259,7 @@ def test_gsm8k_offloading_correctness(cfg: OffloadingModelConfig):
259 cfg.model,
260 server_args,
261 # /reset_prefix_cache requires dev mode.
242 - env_dict={"VLLM_SERVER_DEV_MODE": "1"},
262 + env_dict={"VLLM_SERVER_DEV_MODE": "1", **cfg.env_dict},
263 max_wait_seconds=cfg.startup_timeout,
264 ) as server:
265 base_url = f"http://{server.host}:{server.port}"
tests/v1/kv_connector/unit/offloading_connector/test_config.py new
+587
@@ -0,0 +1,587 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for translating vLLM cache metadata to native offloading config."""
4 +
5 +from typing import Any, cast
6 +from unittest.mock import MagicMock, patch
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
12 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
13 + build_offloading_config,
14 +)
15 +from vllm.platforms import current_platform
16 +from vllm.v1.kv_cache_interface import (
17 + FullAttentionSpec,
18 + HiddenStateCacheSpec,
19 + KVCacheConfig,
20 + KVCacheGroupSpec,
21 + KVCacheTensor,
22 + MambaSpec,
23 + MLAAttentionSpec,
24 + SlidingWindowMLASpec,
25 + SlidingWindowSpec,
26 + UniformTypeKVCacheSpecs,
27 +)
28 +
29 +
30 +def _make_vllm_config(
31 + *,
32 + extra_config: dict[str, Any] | None = None,
33 + tensor_parallel_size: int = 1,
34 + pipeline_parallel_size: int = 1,
35 + prefill_context_parallel_size: int = 1,
36 + decode_context_parallel_size: int = 1,
37 +) -> VllmConfig:
38 + config = MagicMock()
39 + config.cache_config.block_size = 16
40 + config.cache_config.enable_prefix_caching = True
41 + config.cache_config.prefix_match_unit = None
42 + config.cache_config.cache_dtype = torch.float16
43 + config.model_config.model = "test-model"
44 + config.model_config.use_mla = False
45 + world_size = (
46 + tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
47 + )
48 + with patch.object(current_platform, "device_count", return_value=world_size):
49 + config.parallel_config = ParallelConfig(
50 + tensor_parallel_size=tensor_parallel_size,
51 + pipeline_parallel_size=pipeline_parallel_size,
52 + prefill_context_parallel_size=prefill_context_parallel_size,
53 + decode_context_parallel_size=decode_context_parallel_size,
54 + )
55 + config.kv_events_config = None
56 + config.use_v2_model_runner = False
57 + config.kv_transfer_config = KVTransferConfig(
58 + kv_connector="OffloadingConnector",
59 + kv_role="kv_both",
60 + kv_connector_extra_config=dict(extra_config or {}),
61 + )
62 + return cast(VllmConfig, config)
63 +
64 +
65 +def _make_kv_cache_config() -> KVCacheConfig:
66 + num_blocks = 16
67 + kv_tensor = KVCacheTensor(
68 + size=num_blocks * 8,
69 + shared_by=["layer"],
70 + block_stride=0,
71 + )
72 + return KVCacheConfig(
73 + num_blocks=num_blocks,
74 + kv_cache_tensors=[kv_tensor],
75 + kv_cache_groups=[
76 + KVCacheGroupSpec(
77 + ["layer"],
78 + FullAttentionSpec(
79 + block_size=16,
80 + num_kv_heads=1,
81 + head_size=1,
82 + dtype=torch.float32,
83 + ),
84 + )
85 + ],
86 + )
87 +
88 +
89 +def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
90 + num_blocks = 4
91 + if packed:
92 + kv_cache_tensors = [
93 + KVCacheTensor(
94 + size=64,
95 + shared_by=[layer_name],
96 + block_stride=16,
97 + )
98 + for layer_name in ("layer0", "layer1")
99 + ]
100 + else:
101 + kv_cache_tensors = [
102 + KVCacheTensor(size=40, shared_by=["layer0"]),
103 + KVCacheTensor(size=24, shared_by=["layer1"]),
104 + ]
105 +
106 + return KVCacheConfig(
107 + num_blocks=num_blocks,
108 + kv_cache_tensors=kv_cache_tensors,
109 + kv_cache_groups=[
110 + KVCacheGroupSpec(
111 + ["layer0", "layer1"],
112 + FullAttentionSpec(
113 + block_size=16,
114 + num_kv_heads=1,
115 + head_size=1,
116 + dtype=torch.float32,
117 + ),
118 + )
119 + ],
120 + )
121 +
122 +
123 +def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
124 + return FullAttentionSpec(
125 + block_size=block_size,
126 + num_kv_heads=4,
127 + head_size=128,
128 + dtype=torch.float32,
129 + )
130 +
131 +
132 +def _mla_spec(
133 + block_size: int = 16,
134 + head_size: int = 512,
135 + dtype: torch.dtype = torch.float32,
136 +) -> MLAAttentionSpec:
137 + return MLAAttentionSpec(
138 + block_size=block_size,
139 + num_kv_heads=1,
140 + head_size=head_size,
141 + dtype=dtype,
142 + )
143 +
144 +
145 +def _make_mla_kv_cache_config(
146 + layer_names: list[str] | None = None,
147 + head_size: int = 512,
148 + dtype: torch.dtype = torch.float32,
149 + num_blocks: int = 4,
150 +) -> KVCacheConfig:
151 + if layer_names is None:
152 + layer_names = ["layer0", "layer1"]
153 + spec = _mla_spec(head_size=head_size, dtype=dtype)
154 + kv_cache_tensors = [
155 + KVCacheTensor(
156 + size=spec.page_size_bytes * num_blocks,
157 + shared_by=[layer_name],
158 + )
159 + for layer_name in layer_names
160 + ]
161 + return KVCacheConfig(
162 + num_blocks=num_blocks,
163 + kv_cache_tensors=kv_cache_tensors,
164 + kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
165 + )
166 +
167 +
168 +def _make_hybrid_kv_cache_config() -> KVCacheConfig:
169 + return KVCacheConfig(
170 + num_blocks=4,
171 + kv_cache_tensors=[
172 + KVCacheTensor(size=40, shared_by=["full_layer"]),
173 + KVCacheTensor(size=24, shared_by=["mla_layer"]),
174 + ],
175 + kv_cache_groups=[
176 + KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)),
177 + KVCacheGroupSpec(["mla_layer"], _mla_spec()),
178 + ],
179 + )
180 +
181 +
182 +def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
183 + config = _make_vllm_config()
184 + kv_cache_config = KVCacheConfig(
185 + num_blocks=0,
186 + kv_cache_tensors=[],
187 + kv_cache_groups=kv_cache_groups,
188 + )
189 + return build_offloading_config(
190 + config, kv_cache_config
191 + ).parallel.is_parallelism_agnostic
192 +
193 +
194 +def _replicated_layout(
195 + kv_cache_config: KVCacheConfig,
196 + *,
197 + tensor_parallel_size: int = 4,
198 + pipeline_parallel_size: int = 1,
199 + prefill_context_parallel_size: int = 1,
200 + decode_context_parallel_size: int = 1,
201 + use_mla: bool = True,
202 + use_v2_model_runner: bool = False,
203 + distributed_executor_backend: Any = "mp",
204 + nnodes: int = 1,
205 + world_size: int | None = None,
206 +) -> bool:
207 + config = _make_vllm_config(
208 + tensor_parallel_size=tensor_parallel_size,
209 + pipeline_parallel_size=pipeline_parallel_size,
210 + prefill_context_parallel_size=prefill_context_parallel_size,
211 + decode_context_parallel_size=decode_context_parallel_size,
212 + )
213 + config.model_config.use_mla = use_mla
214 + config.use_v2_model_runner = use_v2_model_runner
215 + config.parallel_config.distributed_executor_backend = distributed_executor_backend
216 + config.parallel_config.nnodes = nnodes
217 + if world_size is not None:
218 + config.parallel_config.world_size = world_size
219 + return build_offloading_config(config, kv_cache_config).replicated_layout
220 +
221 +
222 +@pytest.mark.parametrize("packed", [False, True])
223 +def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
224 + config = _make_vllm_config(
225 + extra_config={"block_size": 32},
226 + tensor_parallel_size=3,
227 + pipeline_parallel_size=2,
228 + )
229 +
230 + offloading_config = build_offloading_config(
231 + config, _make_sizing_kv_cache_config(packed)
232 + )
233 +
234 + assert offloading_config.worker_kv_bytes_per_block == 16
235 + assert offloading_config.parallel.world_size == 6
236 + assert offloading_config.cache.blocks_per_chunk == 2
237 +
238 +
239 +def test_rejects_partially_packed_tensor_layout():
240 + kv_cache_config = _make_sizing_kv_cache_config(packed=False)
241 + kv_cache_config.kv_cache_tensors[0].block_stride = 16
242 +
243 + with pytest.raises(AssertionError):
244 + build_offloading_config(_make_vllm_config(), kv_cache_config)
245 +
246 +
247 +def test_zero_blocks_skips_tensor_layout_validation():
248 + kv_cache_config = _make_sizing_kv_cache_config(packed=False)
249 + kv_cache_config.num_blocks = 0
250 + kv_cache_config.kv_cache_tensors[0].block_stride = 16
251 +
252 + offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config)
253 +
254 + assert offloading_config.worker_kv_bytes_per_block == 0
255 +
256 +
257 +def test_prefill_context_parallelism_does_not_scale_group_blocks():
258 + config = _make_vllm_config(
259 + extra_config={"block_size": 64},
260 + prefill_context_parallel_size=2,
261 + )
262 +
263 + offloading_config = build_offloading_config(config, _make_kv_cache_config())
264 +
265 + assert tuple(group.tokens_per_block for group in offloading_config.groups) == (16,)
266 + assert offloading_config.cache.tokens_per_hash == 16
267 + assert offloading_config.cache.blocks_per_chunk == 4
268 +
269 +
270 +def test_preserves_data_parallel_index():
271 + config = _make_vllm_config()
272 + config.parallel_config.data_parallel_index = 2
273 +
274 + offloading_config = build_offloading_config(config, _make_kv_cache_config())
275 +
276 + assert offloading_config.parallel.data_parallel_index == 2
277 +
278 +
279 +def test_resolves_heterogeneous_hybrid_block_sizes():
280 + config = _make_vllm_config()
281 + config.cache_config.block_size = 4
282 +
283 + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config())
284 +
285 + assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
286 + 12,
287 + 16,
288 + )
289 + assert offloading_config.cache.tokens_per_hash == 4
290 + assert offloading_config.cache.blocks_per_chunk == 1
291 +
292 +
293 +@pytest.mark.parametrize("world_size", [2, 4, 8])
294 +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
295 +def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
296 + world_size: int,
297 + use_v2_model_runner: bool,
298 +):
299 + assert _replicated_layout(
300 + _make_mla_kv_cache_config(),
301 + tensor_parallel_size=world_size,
302 + use_v2_model_runner=use_v2_model_runner,
303 + )
304 +
305 +
306 +@pytest.mark.parametrize(
307 + ("kv_cache_config", "case"),
308 + [
309 + (
310 + KVCacheConfig(
311 + num_blocks=4,
312 + kv_cache_tensors=[
313 + KVCacheTensor(
314 + size=_mla_spec().page_size_bytes * 4,
315 + shared_by=["layer"],
316 + )
317 + ],
318 + kv_cache_groups=[
319 + KVCacheGroupSpec(
320 + ["layer"],
321 + SlidingWindowMLASpec(
322 + block_size=16,
323 + num_kv_heads=1,
324 + head_size=512,
325 + dtype=torch.float32,
326 + sliding_window=128,
327 + ),
328 + )
329 + ],
330 + ),
331 + "sliding-window-mla",
332 + ),
333 + (
334 + KVCacheConfig(
335 + num_blocks=4,
336 + kv_cache_tensors=[
337 + KVCacheTensor(
338 + size=_mla_spec().page_size_bytes * 4,
339 + shared_by=["layer"],
340 + )
341 + ],
342 + kv_cache_groups=[
343 + KVCacheGroupSpec(
344 + ["layer"],
345 + HiddenStateCacheSpec(
346 + block_size=16,
347 + num_kv_heads=1,
348 + head_size=512,
349 + dtype=torch.float32,
350 + ),
351 + )
352 + ],
353 + ),
354 + "hidden-state",
355 + ),
356 + (
357 + KVCacheConfig(
358 + num_blocks=4,
359 + kv_cache_tensors=[
360 + KVCacheTensor(
361 + size=_mla_spec().page_size_bytes * 4,
362 + shared_by=["layer0"],
363 + ),
364 + KVCacheTensor(
365 + size=_mla_spec(head_size=256).page_size_bytes * 4,
366 + shared_by=["layer1"],
367 + ),
368 + ],
369 + kv_cache_groups=[
370 + KVCacheGroupSpec(
371 + ["layer0", "layer1"],
372 + UniformTypeKVCacheSpecs(
373 + block_size=16,
374 + kv_cache_specs={
375 + "layer0": _mla_spec(),
376 + "layer1": _mla_spec(head_size=256),
377 + },
378 + ),
379 + )
380 + ],
381 + ),
382 + "uniform-wrapper",
383 + ),
384 + (
385 + KVCacheConfig(
386 + num_blocks=4,
387 + kv_cache_tensors=[
388 + KVCacheTensor(
389 + size=_mla_spec().page_size_bytes * 4,
390 + shared_by=["mla"],
391 + ),
392 + KVCacheTensor(
393 + size=_full_attention_spec().page_size_bytes * 4,
394 + shared_by=["full"],
395 + ),
396 + ],
397 + kv_cache_groups=[
398 + KVCacheGroupSpec(["mla"], _mla_spec()),
399 + KVCacheGroupSpec(["full"], _full_attention_spec()),
400 + ],
401 + ),
402 + "mla-full-hybrid",
403 + ),
404 + (
405 + KVCacheConfig(
406 + num_blocks=4,
407 + kv_cache_tensors=[
408 + KVCacheTensor(
409 + size=_mla_spec().page_size_bytes * 4,
410 + shared_by=["mla"],
411 + ),
412 + KVCacheTensor(size=64 * 4, shared_by=["mamba"]),
413 + ],
414 + kv_cache_groups=[
415 + KVCacheGroupSpec(["mla"], _mla_spec()),
416 + KVCacheGroupSpec(
417 + ["mamba"],
418 + MambaSpec(
419 + block_size=16,
420 + shapes=((16, 1),),
421 + dtypes=(torch.float32,),
422 + ),
423 + ),
424 + ],
425 + ),
426 + "mla-mamba-hybrid",
427 + ),
428 + (
429 + KVCacheConfig(
430 + num_blocks=4,
431 + kv_cache_tensors=[
432 + KVCacheTensor(
433 + size=_mla_spec().page_size_bytes * 4,
434 + shared_by=["layer0"],
435 + ),
436 + KVCacheTensor(
437 + size=_mla_spec().page_size_bytes * 4,
438 + shared_by=["layer1"],
439 + ),
440 + ],
441 + kv_cache_groups=[
442 + KVCacheGroupSpec(["layer0"], _mla_spec()),
443 + KVCacheGroupSpec(["layer1"], _mla_spec()),
444 + ],
445 + ),
446 + "multi-group-mla",
447 + ),
448 + ],
449 + ids=[
450 + "sliding-window-mla",
451 + "hidden-state",
452 + "uniform-wrapper",
453 + "mla-full-hybrid",
454 + "mla-mamba-hybrid",
455 + "multi-group-mla",
456 + ],
457 +)
458 +def test_replicated_layout_excludes_unproven_cache_shapes(
459 + kv_cache_config: KVCacheConfig,
460 + case: str,
461 +):
462 + assert not _replicated_layout(kv_cache_config), case
463 +
464 +
465 +def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
466 + num_blocks = 4
467 + main_spec = _mla_spec(head_size=512)
468 + indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
469 + main_layers = [f"main_{i}" for i in range(61)]
470 + indexer_layers = [f"indexer_{i}" for i in range(61)]
471 + kv_cache_config = KVCacheConfig(
472 + num_blocks=num_blocks,
473 + kv_cache_tensors=[
474 + KVCacheTensor(
475 + size=main_spec.page_size_bytes * len(main_layers) * num_blocks,
476 + shared_by=main_layers,
477 + ),
478 + KVCacheTensor(
479 + size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks,
480 + shared_by=indexer_layers,
481 + ),
482 + ],
483 + kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
484 + )
485 +
486 + assert not _replicated_layout(kv_cache_config)
487 +
488 +
489 +@pytest.mark.parametrize(
490 + ("kwargs", "case"),
491 + [
492 + ({"tensor_parallel_size": 1}, "tp1"),
493 + ({"use_mla": False}, "use-mla-false"),
494 + ({"pipeline_parallel_size": 2, "world_size": 4}, "pp2"),
495 + ({"prefill_context_parallel_size": 2, "world_size": 4}, "pcp2"),
496 + ({"decode_context_parallel_size": 2}, "dcp2"),
497 + ({"world_size": 8}, "world-ne-tp"),
498 + ({"distributed_executor_backend": "ray"}, "ray"),
499 + ({"distributed_executor_backend": "uni"}, "uni"),
500 + ({"distributed_executor_backend": type("DummyExecutor", (), {})}, "class"),
501 + ({"nnodes": 2}, "multi-node"),
502 + ],
503 + ids=[
504 + "tp1",
505 + "use-mla-false",
506 + "pp2",
507 + "pcp2",
508 + "dcp2",
509 + "world-ne-tp",
510 + "ray",
511 + "uni",
512 + "class",
513 + "multi-node",
514 + ],
515 +)
516 +def test_replicated_layout_parallel_gate(kwargs: dict[str, Any], case: str):
517 + assert not _replicated_layout(_make_mla_kv_cache_config(), **kwargs), case
518 +
519 +
520 +def test_parallelism_agnostic_for_single_full_attention_group():
521 + assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
522 +
523 +
524 +@pytest.mark.parametrize(
525 + "kv_cache_groups",
526 + [
527 + [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))],
528 + [
529 + KVCacheGroupSpec(
530 + ["l0"],
531 + SlidingWindowSpec(
532 + block_size=16,
533 + num_kv_heads=4,
534 + head_size=128,
535 + dtype=torch.float32,
536 + sliding_window=128,
537 + ),
538 + )
539 + ],
540 + [
541 + KVCacheGroupSpec(["l0"], _full_attention_spec()),
542 + KVCacheGroupSpec(["l1"], _full_attention_spec()),
543 + ],
544 + ],
545 +)
546 +def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
547 + assert not _parallelism_agnostic(kv_cache_groups)
548 +
549 +
550 +def test_parallelism_agnostic_disabled_on_v2_model_runner():
551 + config = _make_vllm_config()
552 + config.use_v2_model_runner = True
553 + kv_cache_config = KVCacheConfig(
554 + num_blocks=0,
555 + kv_cache_tensors=[],
556 + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
557 + )
558 +
559 + offloading_config = build_offloading_config(config, kv_cache_config)
560 +
561 + assert not offloading_config.parallel.is_parallelism_agnostic
562 +
563 +
564 +def test_accepts_blocks_per_chunk_for_heterogeneous_groups():
565 + config = _make_vllm_config(extra_config={"blocks_per_chunk": 2})
566 +
567 + offloading_config = build_offloading_config(config, _make_hybrid_kv_cache_config())
568 +
569 + assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
570 + 12,
571 + 16,
572 + )
573 + assert offloading_config.cache.blocks_per_chunk == 2
574 +
575 +
576 +def test_block_size_and_blocks_per_chunk_are_mutually_exclusive():
577 + config = _make_vllm_config(extra_config={"block_size": 64, "blocks_per_chunk": 2})
578 +
579 + with pytest.raises(ValueError, match="Specify only one"):
580 + build_offloading_config(config, _make_kv_cache_config())
581 +
582 +
583 +def test_blocks_per_chunk_must_be_positive():
584 + config = _make_vllm_config(extra_config={"blocks_per_chunk": 0})
585 +
586 + with pytest.raises(ValueError, match="greater than 0"):
587 + build_offloading_config(config, _make_kv_cache_config())
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+47
@@ -14,6 +14,7 @@ from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
14 from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
15 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
16 OffloadingConnectorMetadata,
17 + OffloadingWorkerMetadata,
18 )
19 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
20 OffloadingConnectorStats,
@@ -40,6 +41,7 @@ from vllm.v1.kv_offload.base import (
41 get_offload_block_hash,
42 make_offload_key,
43 )
44 +from vllm.v1.outputs import KVConnectorOutput
45 from vllm.v1.request import RequestStatus
46
47
@@ -1539,6 +1541,51 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool):
1541 assert runner.manager.complete_store.call_count == 0
1542
1543
1544 +@pytest.mark.parametrize("async_scheduling", [True, False])
1545 +def test_complete_store_waits_for_all_worker_acks(
1546 + request_runner, async_scheduling: bool
1547 +):
1548 + tokens_per_block = 4
1549 + blocks_per_chunk = 3
1550 + tokens_per_chunk = tokens_per_block * blocks_per_chunk
1551 + runner = request_runner(
1552 + blocks_per_chunk=blocks_per_chunk,
1553 + block_size=tokens_per_block,
1554 + num_gpu_blocks=100,
1555 + async_scheduling=async_scheduling,
1556 + worker_count=3,
1557 + )
1558 + runner.new_request(token_ids=[0] * tokens_per_chunk)
1559 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (
1560 + generate_store_output(keys)
1561 + )
1562 + runner.run(decoded_tokens=[0, 0], complete_transfers=False)
1563 + assert len(runner.connector_scheduler._jobs) == 1
1564 + job_id = next(iter(runner.connector_scheduler._jobs))
1565 + assert runner.connector_scheduler._jobs[job_id].pending_count == 3
1566 + runner.manager.complete_store.reset_mock()
1567 +
1568 + runner.connector_scheduler.update_connector_output(
1569 + KVConnectorOutput(
1570 + kv_connector_worker_meta=OffloadingWorkerMetadata(
1571 + completed_jobs={job_id: 1}
1572 + )
1573 + )
1574 + )
1575 + assert runner.manager.complete_store.call_count == 0
1576 + assert runner.connector_scheduler._jobs[job_id].pending_count == 2
1577 +
1578 + runner.connector_scheduler.update_connector_output(
1579 + KVConnectorOutput(
1580 + kv_connector_worker_meta=OffloadingWorkerMetadata(
1581 + completed_jobs={job_id: 2}
1582 + )
1583 + )
1584 + )
1585 + assert runner.manager.complete_store.call_count == 1
1586 + assert job_id not in runner.connector_scheduler._jobs
1587 +
1588 +
1589 @pytest.mark.parametrize("async_scheduling", [True, False])
1590 def test_max_offload_tokens_validation(request_runner, async_scheduling: bool):
1591 """Validates max_offload_tokens: type coercion, boundary values, and capping.
tests/v1/kv_connector/unit/offloading_connector/test_worker.py
+184 -1
@@ -6,6 +6,10 @@ from unittest.mock import MagicMock
6 import pytest
7 import torch
8
9 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
10 + OffloadingConnectorMetadata,
11 + TransferJob,
12 +)
13 from vllm.platforms import current_platform
14 from vllm.utils.torch_utils import get_dtype_size
15 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -22,7 +26,17 @@ from vllm.v1.kv_cache_interface import (
26 from vllm.v1.kv_offload.base import (
27 CanonicalKVCacheRef,
28 CanonicalKVCaches,
29 + GPULoadStoreSpec,
30 + LoadStoreSpec,
31 + OffloadingManager,
32 OffloadingSpec,
33 + OffloadingWorker,
34 +)
35 +from vllm.v1.kv_offload.config import (
36 + OffloadingCacheConfig,
37 + OffloadingConfig,
38 + OffloadingModelConfig,
39 + OffloadingParallelConfig,
40 )
41
42 NUM_BLOCKS = 10
@@ -87,7 +101,11 @@ def _allocate_and_reshape_kv_caches(
101 set_kv_cache_layout(None)
102
103
90 -def _make_worker(kv_cache_config: KVCacheConfig):
104 +def _make_worker(
105 + kv_cache_config: KVCacheConfig,
106 + replicated_layout: bool = False,
107 + rank: int = 0,
108 +):
109 """
110 Create an OffloadingConnectorWorker with mocked dependencies.
111 """
@@ -96,6 +114,9 @@ def _make_worker(kv_cache_config: KVCacheConfig):
114 )
115
116 spec = MagicMock(spec=OffloadingSpec)
117 + spec.replicated_layout = replicated_layout
118 + spec.config = MagicMock()
119 + spec.config.parallel.rank = rank
120 spec.get_worker.return_value = MagicMock()
121
122 worker = OffloadingConnectorWorker(
@@ -107,11 +128,173 @@ def _make_worker(kv_cache_config: KVCacheConfig):
128 return worker, spec
129
130
131 +def _store_metadata(job_id: int) -> OffloadingConnectorMetadata:
132 + return OffloadingConnectorMetadata(
133 + load_jobs={},
134 + store_jobs={
135 + job_id: TransferJob(
136 + req_id="req",
137 + src_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)),
138 + dst_spec=LoadStoreSpec(),
139 + )
140 + },
141 + )
142 +
143 +
144 +def _load_metadata(job_id: int) -> OffloadingConnectorMetadata:
145 + return OffloadingConnectorMetadata(
146 + load_jobs={
147 + job_id: TransferJob(
148 + req_id="req",
149 + src_spec=LoadStoreSpec(),
150 + dst_spec=GPULoadStoreSpec([0], group_sizes=(1,), block_indices=(0,)),
151 + )
152 + },
153 + store_jobs={},
154 + )
155 +
156 +
157 +def _empty_metadata() -> OffloadingConnectorMetadata:
158 + return OffloadingConnectorMetadata(load_jobs={}, store_jobs={})
159 +
160 +
161 +def _offloading_config(rank: int = 0) -> OffloadingConfig:
162 + return OffloadingConfig(
163 + groups=(),
164 + worker_kv_bytes_per_block=0,
165 + enable_kv_cache_events=False,
166 + extra_config={},
167 + engine_id="test-engine",
168 + model=OffloadingModelConfig(name="test-model", dtype="float16"),
169 + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1),
170 + parallel=OffloadingParallelConfig(
171 + rank=rank,
172 + world_size=2,
173 + tp_size=2,
174 + pp_size=1,
175 + pcp_size=1,
176 + dcp_size=1,
177 + data_parallel_index=0,
178 + is_parallelism_agnostic=False,
179 + ),
180 + )
181 +
182 +
183 +class BareExternalOffloadingSpec(OffloadingSpec):
184 + def get_manager(self) -> OffloadingManager:
185 + raise NotImplementedError
186 +
187 + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker:
188 + raise NotImplementedError
189 +
190 +
191 # ---------------------------------------------------------------------------
192 # Tests
193 # ---------------------------------------------------------------------------
194
195
196 +def test_prepare_store_kv_non_writer_marks_completed_without_submit():
197 + worker, _ = _make_worker(
198 + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
199 + replicated_layout=True,
200 + rank=1,
201 + )
202 +
203 + worker.prepare_store_kv(_store_metadata(7))
204 + worker.start_kv_transfers(_empty_metadata())
205 +
206 + assert worker._unsubmitted_store_jobs == []
207 + assert worker.worker is not None
208 + worker.worker.submit_store.assert_not_called()
209 + meta = worker.build_connector_worker_meta()
210 + assert meta is not None
211 + assert meta.completed_jobs == {7: 1}
212 +
213 +
214 +def test_prepare_store_kv_writer_submits_store():
215 + worker, _ = _make_worker(
216 + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
217 + replicated_layout=True,
218 + rank=0,
219 + )
220 +
221 + worker.prepare_store_kv(_store_metadata(8))
222 + assert worker.build_connector_worker_meta() is None
223 + worker.start_kv_transfers(_empty_metadata())
224 +
225 + assert worker.worker is not None
226 + worker.worker.submit_store.assert_called_once()
227 +
228 +
229 +def test_prepare_store_kv_non_replicated_rank_gt_zero_queues_store():
230 + worker, _ = _make_worker(
231 + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
232 + replicated_layout=False,
233 + rank=1,
234 + )
235 +
236 + worker.prepare_store_kv(_store_metadata(9))
237 + assert worker.build_connector_worker_meta() is None
238 + assert len(worker._unsubmitted_store_jobs) == 1
239 +
240 + worker.start_kv_transfers(_empty_metadata())
241 +
242 + assert worker.worker is not None
243 + worker.worker.submit_store.assert_called_once()
244 + assert worker._unsubmitted_store_jobs == []
245 +
246 +
247 +def test_handle_preemptions_non_writer_acks_flushed_store():
248 + worker, _ = _make_worker(
249 + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
250 + replicated_layout=True,
251 + rank=1,
252 + )
253 + metadata = _store_metadata(10)
254 + metadata.jobs_to_flush = {10}
255 +
256 + worker.handle_preemptions(metadata)
257 +
258 + assert worker.worker is not None
259 + worker.worker.submit_store.assert_not_called()
260 + worker.worker.wait.assert_called_once_with({10})
261 + assert metadata.store_jobs == {}
262 + meta = worker.build_connector_worker_meta()
263 + assert meta is not None
264 + assert meta.completed_jobs == {10: 1}
265 +
266 +
267 +def test_start_kv_transfers_non_writer_still_submits_load():
268 + worker, _ = _make_worker(
269 + KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]),
270 + replicated_layout=True,
271 + rank=1,
272 + )
273 +
274 + worker.start_kv_transfers(_load_metadata(10))
275 +
276 + assert worker.worker is not None
277 + worker.worker.submit_load.assert_called_once()
278 + assert worker.build_connector_worker_meta() is None
279 +
280 +
281 +def test_offloading_connector_worker_accepts_plugin_spec_default_layout():
282 + from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
283 + OffloadingConnectorWorker,
284 + )
285 +
286 + spec = BareExternalOffloadingSpec(_offloading_config(rank=1))
287 +
288 + OffloadingConnectorWorker(
289 + spec=spec,
290 + kv_cache_config=KVCacheConfig(
291 + num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
292 + ),
293 + )
294 +
295 + assert spec.replicated_layout is False
296 +
297 +
298 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
299 def test_register_kv_caches(backend):
300 """Test register_kv_caches with multiple groups covering all layer types.
tests/v1/kv_connector/unit/offloading_connector/utils.py
+4
@@ -179,6 +179,7 @@ class RequestRunner:
179 async_scheduling: bool = True,
180 kv_cache_groups: list[KVCacheGroupSpec] | None = None,
181 extra_config_overrides: dict[str, Any] | None = None,
182 + worker_count: int = 1,
183 ):
184 assert blocks_per_chunk == 1 or kv_cache_groups is None, (
185 "blocks_per_chunk > 1 requires all groups to have the same "
@@ -198,6 +199,7 @@ class RequestRunner:
199 disable_hybrid_kv_cache_manager=False,
200 )
201 vllm_config.scheduler_config.async_scheduling = async_scheduling
202 + vllm_config.parallel_config.world_size = worker_count
203
204 extra_config: dict[str, Any] = {
205 "spec_name": "MockOffloadingSpec",
@@ -668,6 +670,7 @@ def request_runner():
670 blocks_per_chunk=1,
671 kv_cache_groups=None,
672 extra_config_overrides=None,
673 + worker_count=1,
674 ):
675 runner = RequestRunner(
676 block_size=block_size,
@@ -676,6 +679,7 @@ def request_runner():
679 async_scheduling=async_scheduling,
680 kv_cache_groups=kv_cache_groups,
681 extra_config_overrides=extra_config_overrides,
682 + worker_count=worker_count,
683 )
684 runners.append(runner)
685 return runner
tests/v1/kv_offload/cpu/test_gpu_worker.py
+10 -2
@@ -41,7 +41,10 @@ NUM_MAPPINGS_PER_GROUP = [2]
41 @pytest.mark.parametrize("num_tensors", NUM_TENSORS)
42 @pytest.mark.parametrize("seed", SEEDS)
43 @pytest.mark.parametrize("device", DEVICES)
44 -@pytest.mark.parametrize("use_shared_memory", [False, True])
44 +@pytest.mark.parametrize(
45 + ("use_shared_memory", "replicated_layout"),
46 + [(False, False), (True, False), (True, True)],
47 +)
48 @torch.inference_mode()
49 def test_transfer(
50 default_vllm_config,
@@ -55,6 +58,7 @@ def test_transfer(
58 seed: int,
59 device: str,
60 use_shared_memory: bool,
61 + replicated_layout: bool,
62 ) -> None:
63 set_random_seed(seed)
64
@@ -95,11 +99,15 @@ def test_transfer(
99 gpu_page_size_bytes * num_tensors * blocks_per_chunk,
100 SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT,
101 )
102 + simulated_world_size = 2
103 + kv_bytes_per_block = (
104 + cpu_page_size if replicated_layout else cpu_page_size * simulated_world_size
105 + )
106 mmap_region = SharedOffloadRegion(
107 engine_id=str(uuid.uuid4()),
108 num_blocks=num_cpu_blocks,
109 rank=0,
102 - kv_bytes_per_block=cpu_page_size,
110 + kv_bytes_per_block=kv_bytes_per_block,
111 cpu_page_size=cpu_page_size,
112 )
113
tests/v1/kv_offload/test_factory.py
+215 -476
@@ -1,35 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -Unit tests for OffloadingSpecFactory.
3 +"""Unit tests for native offloading specs and their factory."""
4
6 -These tests verify:
7 -1. Pre-registration integrity — registered module paths can actually import
8 - and yield correct OffloadingSpec subclasses (CI sentinel against file moves).
9 -2. End-to-end factory → spec construction with real configs.
10 -3. Downstream collaboration — build_metric_definitions delegation.
11 -4. Error paths — unregistered specs, missing config, duplicate registration.
12 -"""
13 -
14 -from typing import cast
15 -from unittest.mock import MagicMock, patch
5 +from typing import Any
6 +from unittest.mock import MagicMock
7
8 import pytest
18 -import torch
9
20 -from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
21 -from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
22 - build_offloading_config,
23 -)
24 -from vllm.platforms import current_platform
25 -from vllm.v1.kv_cache_interface import (
26 - FullAttentionSpec,
27 - KVCacheConfig,
28 - KVCacheGroupSpec,
29 - KVCacheTensor,
30 - MLAAttentionSpec,
31 - SlidingWindowSpec,
32 -)
10 from vllm.v1.kv_offload.base import (
11 CanonicalKVCaches,
12 OffloadingHistogramMetadata,
@@ -37,220 +14,81 @@ from vllm.v1.kv_offload.base import (
14 OffloadingSpec,
15 OffloadingWorker,
16 )
17 +from vllm.v1.kv_offload.config import (
18 + OffloadingCacheConfig,
19 + OffloadingConfig,
20 + OffloadingGroupConfig,
21 + OffloadingModelConfig,
22 + OffloadingParallelConfig,
23 +)
24 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
25 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
26 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
27 from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
28
45 -# ---------------------------------------------------------------------------
46 -# Fixtures
47 -# ---------------------------------------------------------------------------
48 -
29
30 @pytest.fixture(autouse=True)
31 def restore_registry():
52 - """Save and restore OffloadingSpecFactory._registry between tests."""
32 original = dict(OffloadingSpecFactory._registry)
33 yield
34 OffloadingSpecFactory._registry = original
35
36
58 -def _get_extra_config(config: VllmConfig) -> dict:
59 - assert config.kv_transfer_config is not None
60 - return config.kv_transfer_config.kv_connector_extra_config
61 -
62 -
63 -def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec:
64 - return OffloadingSpecFactory.create_spec(
65 - build_offloading_config(config, kv_cache_config)
66 - )
67 -
68 -
69 -def _make_vllm_config(
37 +def _make_offloading_config(
38 + *,
39 spec_name: str | None = "CPUOffloadingSpec",
71 - cpu_bytes_to_use: int | None = None,
72 - store_threshold: int = 0,
73 - extra_config: dict | None = None,
74 -):
75 - """Build a real VllmConfig with kv_transfer_config set for offloading."""
76 - from vllm.config import (
77 - CacheConfig,
78 - DeviceConfig,
79 - ModelConfig,
80 - SchedulerConfig,
81 - VllmConfig,
82 - )
83 -
84 - model_config = ModelConfig(
85 - model="facebook/opt-125m",
86 - trust_remote_code=True,
87 - dtype="float16",
88 - seed=42,
89 - )
90 - scheduler_config = SchedulerConfig(
91 - max_num_seqs=16,
92 - max_num_batched_tokens=64,
93 - max_model_len=10000,
94 - enable_chunked_prefill=True,
95 - is_encoder_decoder=model_config.is_encoder_decoder,
96 - )
97 - cache_config = CacheConfig(
98 - block_size=16,
99 - gpu_memory_utilization=0.9,
100 - cache_dtype="auto",
101 - enable_prefix_caching=True,
102 - )
103 -
104 - cfg = extra_config or {}
105 - if cpu_bytes_to_use is not None:
106 - cfg["cpu_bytes_to_use"] = cpu_bytes_to_use
107 - cfg["spec_name"] = spec_name
108 - if store_threshold > 0:
109 - cfg["store_threshold"] = store_threshold
110 -
111 - kv_transfer_config = KVTransferConfig(
112 - kv_connector="OffloadingConnector",
113 - kv_role="kv_both",
114 - kv_connector_extra_config=cfg,
115 - )
116 - return VllmConfig(
117 - scheduler_config=scheduler_config,
118 - model_config=model_config,
119 - cache_config=cache_config,
120 - kv_transfer_config=kv_transfer_config,
121 - device_config=DeviceConfig("cpu"),
122 - )
123 -
124 -
125 -def _make_layout_vllm_config(
126 - spec_name: str = "CPUOffloadingSpec",
127 - cpu_bytes_to_use: int | None = None,
128 - extra_config: dict | None = None,
129 - tensor_parallel_size: int = 1,
130 - pipeline_parallel_size: int = 1,
131 - prefill_context_parallel_size: int = 1,
132 - decode_context_parallel_size: int = 1,
133 -) -> VllmConfig:
134 - config = MagicMock()
135 - config.cache_config.block_size = 16
136 - config.cache_config.enable_prefix_caching = True
137 - config.cache_config.prefix_match_unit = None
138 - config.cache_config.cache_dtype = torch.float16
139 - config.model_config.model = "test-model"
140 - world_size = (
141 - tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
142 - )
143 - with patch.object(current_platform, "device_count", return_value=world_size):
144 - config.parallel_config = ParallelConfig(
145 - tensor_parallel_size=tensor_parallel_size,
146 - pipeline_parallel_size=pipeline_parallel_size,
147 - prefill_context_parallel_size=prefill_context_parallel_size,
148 - decode_context_parallel_size=decode_context_parallel_size,
149 - )
150 - config.kv_events_config = None
151 - config.use_v2_model_runner = False
152 -
153 - connector_extra_config = dict(extra_config or {})
154 - connector_extra_config["spec_name"] = spec_name
40 + cpu_bytes_to_use: int | None = 65536,
41 + worker_kv_bytes_per_block: int = 8,
42 + groups: tuple[OffloadingGroupConfig, ...] | None = None,
43 + tokens_per_hash: int = 16,
44 + blocks_per_chunk: int = 1,
45 + rank: int = 0,
46 + world_size: int = 1,
47 + tp_size: int | None = None,
48 + pp_size: int = 1,
49 + pcp_size: int = 1,
50 + dcp_size: int = 1,
51 + data_parallel_index: int = 0,
52 + is_parallelism_agnostic: bool = False,
53 + replicated_layout: bool = False,
54 + extra_config: dict[str, Any] | None = None,
55 +) -> OffloadingConfig:
56 + normalized_extra_config = dict(extra_config or {})
57 + if spec_name is not None:
58 + normalized_extra_config["spec_name"] = spec_name
59 if cpu_bytes_to_use is not None:
156 - connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
157 - config.kv_transfer_config = KVTransferConfig(
158 - kv_connector="OffloadingConnector",
159 - kv_role="kv_both",
160 - kv_connector_extra_config=connector_extra_config,
161 - )
162 - return cast(VllmConfig, config)
163 -
164 -
165 -def _make_kv_cache_config():
166 - """Build a minimal KVCacheConfig with one KV cache tensor."""
167 - num_blocks = 16
168 - num_kv_heads = 1
169 - head_size = 1
170 - dtype = torch.float32
171 - page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
172 - kv_tensor = KVCacheTensor(
173 - size=num_blocks * page_size, shared_by=["layer"], block_stride=0
60 + normalized_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use
61 +
62 + if groups is None:
63 + groups = (OffloadingGroupConfig(16, ("layer",)),)
64 +
65 + return OffloadingConfig(
66 + groups=groups,
67 + worker_kv_bytes_per_block=worker_kv_bytes_per_block,
68 + enable_kv_cache_events=False,
69 + extra_config=normalized_extra_config,
70 + engine_id="test-engine",
71 + model=OffloadingModelConfig(name="test-model", dtype="float16"),
72 + cache=OffloadingCacheConfig(
73 + tokens_per_hash=tokens_per_hash,
74 + blocks_per_chunk=blocks_per_chunk,
75 + ),
76 + parallel=OffloadingParallelConfig(
77 + rank=rank,
78 + world_size=world_size,
79 + tp_size=world_size if tp_size is None else tp_size,
80 + pp_size=pp_size,
81 + pcp_size=pcp_size,
82 + dcp_size=dcp_size,
83 + data_parallel_index=data_parallel_index,
84 + is_parallelism_agnostic=is_parallelism_agnostic,
85 + ),
86 + replicated_layout=replicated_layout,
87 )
175 - return KVCacheConfig(
176 - num_blocks=num_blocks,
177 - kv_cache_tensors=[kv_tensor],
178 - kv_cache_groups=[
179 - KVCacheGroupSpec(
180 - ["layer"],
181 - FullAttentionSpec(
182 - block_size=16,
183 - num_kv_heads=num_kv_heads,
184 - head_size=head_size,
185 - dtype=dtype,
186 - ),
187 - )
188 - ],
189 - )
190 -
88
192 -def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
193 - num_blocks = 4
194 - if packed:
195 - kv_cache_tensors = [
196 - KVCacheTensor(
197 - size=64,
198 - shared_by=[layer_name],
199 - block_stride=16,
200 - )
201 - for layer_name in ("layer0", "layer1")
202 - ]
203 - else:
204 - kv_cache_tensors = [
205 - KVCacheTensor(size=40, shared_by=["layer0"]),
206 - KVCacheTensor(size=24, shared_by=["layer1"]),
207 - ]
208 -
209 - return KVCacheConfig(
210 - num_blocks=num_blocks,
211 - kv_cache_tensors=kv_cache_tensors,
212 - kv_cache_groups=[
213 - KVCacheGroupSpec(
214 - ["layer0", "layer1"],
215 - FullAttentionSpec(
216 - block_size=16,
217 - num_kv_heads=1,
218 - head_size=1,
219 - dtype=torch.float32,
220 - ),
221 - )
222 - ],
223 - )
89
225 -
226 -def _make_hybrid_kv_cache_config() -> KVCacheConfig:
227 - return KVCacheConfig(
228 - num_blocks=4,
229 - kv_cache_tensors=[
230 - KVCacheTensor(size=40, shared_by=["full_layer"]),
231 - KVCacheTensor(size=24, shared_by=["mla_layer"]),
232 - ],
233 - kv_cache_groups=[
234 - KVCacheGroupSpec(
235 - ["full_layer"],
236 - FullAttentionSpec(
237 - block_size=12,
238 - num_kv_heads=1,
239 - head_size=1,
240 - dtype=torch.float32,
241 - ),
242 - ),
243 - KVCacheGroupSpec(
244 - ["mla_layer"],
245 - MLAAttentionSpec(
246 - block_size=16,
247 - num_kv_heads=1,
248 - head_size=576,
249 - dtype=torch.float32,
250 - ),
251 - ),
252 - ],
253 - )
90 +def _create_spec(**kwargs: Any) -> OffloadingSpec:
91 + return OffloadingSpecFactory.create_spec(_make_offloading_config(**kwargs))
92
93
94 class SingleArgExternalOffloadingSpec(OffloadingSpec):
@@ -261,104 +99,60 @@ class SingleArgExternalOffloadingSpec(OffloadingSpec):
99 raise NotImplementedError
100
101
264 -# ---------------------------------------------------------------------------
265 -# Pre-registration integrity (CI sentinel)
266 -# ---------------------------------------------------------------------------
267 -
268 -
102 def test_pre_registered_specs_can_be_imported():
270 - """If someone moves cpu/spec.py but forgets to update factory.py, CI fails."""
103 for name in OffloadingSpecFactory._registry:
104 cls = OffloadingSpecFactory._registry[name]()
105 assert issubclass(cls, OffloadingSpec)
106
107
108 def test_cpu_spec_registered():
277 - """CPUOffloadingSpec is registered and importable."""
109 cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]()
110 assert cls is CPUOffloadingSpec
111
112
113 def test_tiering_spec_registered():
283 - """TieringOffloadingSpec is registered and importable."""
114 cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]()
115 assert cls is TieringOffloadingSpec
116
117
288 -# ---------------------------------------------------------------------------
289 -# Normal path — get_spec_cls
290 -# ---------------------------------------------------------------------------
291 -
292 -
118 def test_get_spec_cls_returns_registered_class():
294 - """Registered spec_name returns correct class."""
295 - config = _make_vllm_config(spec_name="CPUOffloadingSpec")
296 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
119 + spec_cls = OffloadingSpecFactory.get_spec_cls(
120 + _make_offloading_config().extra_config
121 + )
122 assert spec_cls is CPUOffloadingSpec
123
124
300 -def test_get_spec_cls_default_to_cpu():
301 - """Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
302 - config = _make_vllm_config(spec_name=None)
303 - config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
304 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
125 +def test_get_spec_cls_defaults_to_cpu():
126 + spec_cls = OffloadingSpecFactory.get_spec_cls(
127 + _make_offloading_config(spec_name=None).extra_config
128 + )
129 assert spec_cls is CPUOffloadingSpec
130
131
308 -# ---------------------------------------------------------------------------
309 -# End-to-end — create_spec
310 -# ---------------------------------------------------------------------------
311 -
312 -
313 -def test_create_cpu_offloading_spec_end_to_end():
314 - """Full factory → spec construction with real VllmConfig/KVCacheConfig.
315 -
316 - Verifies:
317 - - cpu_bytes_to_use validation and num_blocks calculation
318 - - block_size % tokens_per_hash assertion
319 - - spec instance is CPUOffloadingSpec
320 - """
321 - config = _make_vllm_config(cpu_bytes_to_use=65536)
322 - kv_cache_config = _make_kv_cache_config()
323 - spec = _create_spec(config, kv_cache_config)
132 +def test_create_cpu_offloading_spec():
133 + spec = _create_spec()
134 assert isinstance(spec, CPUOffloadingSpec)
135 assert spec.num_blocks > 0
136
137
328 -@pytest.mark.parametrize("packed", [False, True])
329 -def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
330 - cpu_bytes_to_use = 1920
331 - config = _make_layout_vllm_config(
332 - cpu_bytes_to_use=cpu_bytes_to_use,
333 - extra_config={"block_size": 32},
334 - tensor_parallel_size=3,
335 - pipeline_parallel_size=2,
138 +def test_cpu_spec_sizes_normalized_worker_layout():
139 + spec = _create_spec(
140 + cpu_bytes_to_use=1920,
141 + worker_kv_bytes_per_block=16,
142 + blocks_per_chunk=2,
143 + world_size=6,
144 + tp_size=3,
145 + pp_size=2,
146 )
147
338 - spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
339 -
148 assert isinstance(spec, CPUOffloadingSpec)
149 assert spec.cpu_page_size_per_worker == 32
150 assert spec.kv_bytes_per_chunk == 192
343 - assert spec.num_blocks == cpu_bytes_to_use // 192
344 -
345 -
346 -def test_cpu_spec_rejects_partially_packed_tensor_layout():
347 - config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
348 - kv_cache_config = _make_sizing_kv_cache_config(packed=False)
349 - kv_cache_config.kv_cache_tensors[0].block_stride = 16
350 -
351 - with pytest.raises(AssertionError):
352 - _create_spec(config, kv_cache_config)
353 -
151 + assert spec.num_blocks == 10
152
355 -def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
356 - config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
357 - kv_cache_config = _make_sizing_kv_cache_config(packed=False)
358 - kv_cache_config.num_blocks = 0
359 - kv_cache_config.kv_cache_tensors[0].block_stride = 16
153
361 - spec = _create_spec(config, kv_cache_config)
154 +def test_cpu_spec_zero_worker_bytes_produces_empty_cache():
155 + spec = _create_spec(worker_kv_bytes_per_block=0, world_size=4)
156
157 assert isinstance(spec, CPUOffloadingSpec)
158 assert spec.cpu_page_size_per_worker == 0
@@ -368,225 +162,209 @@ def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
162
163 def test_tiering_spec_aligns_row_size():
164 alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
371 - cpu_bytes_to_use = alignment * 3
372 - config = _make_layout_vllm_config(
165 + spec = _create_spec(
166 spec_name="TieringOffloadingSpec",
374 - cpu_bytes_to_use=cpu_bytes_to_use,
375 - extra_config={"block_size": 32},
376 - tensor_parallel_size=3,
377 - pipeline_parallel_size=2,
167 + cpu_bytes_to_use=alignment * 3,
168 + worker_kv_bytes_per_block=16,
169 + blocks_per_chunk=2,
170 + world_size=6,
171 + tp_size=3,
172 + pp_size=2,
173 )
174
380 - spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
381 -
175 assert isinstance(spec, TieringOffloadingSpec)
176 assert spec.cpu_page_size_per_worker == 32
177 assert spec.kv_bytes_per_chunk == alignment
385 - assert spec.num_blocks == cpu_bytes_to_use // alignment
178 + assert spec.num_blocks == 3
179
180
388 -def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel():
389 - config = _make_layout_vllm_config(
390 - cpu_bytes_to_use=65536,
391 - extra_config={"block_size": 64},
392 - prefill_context_parallel_size=2,
181 +@pytest.mark.parametrize("world_size", [2, 4, 8])
182 +def test_tiering_spec_replicated_sizing_removes_world_factor(world_size: int):
183 + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
184 + spec = _create_spec(
185 + spec_name="TieringOffloadingSpec",
186 + cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
187 + worker_kv_bytes_per_block=worker_kv_bytes_per_block,
188 + world_size=world_size,
189 + replicated_layout=True,
190 )
191
395 - spec = _create_spec(config, _make_kv_cache_config())
192 + assert isinstance(spec, TieringOffloadingSpec)
193 + assert spec.replicated_layout is True
194 + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
195 + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block
196 + assert spec.num_blocks == 8
197
397 - assert spec.tokens_per_block == (16,)
398 - assert spec.tokens_per_hash == 16
399 - assert spec.blocks_per_chunk == 4
198
199 +def test_tiering_spec_create_worker_uses_single_slot_for_replicated_layout(monkeypatch):
200 + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module
201
402 -def test_offloading_config_preserves_data_parallel_index():
403 - config = _make_layout_vllm_config()
404 - config.parallel_config.data_parallel_index = 2
202 + worker_kv_bytes_per_block = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
203 + spec = _create_spec(
204 + spec_name="TieringOffloadingSpec",
205 + cpu_bytes_to_use=worker_kv_bytes_per_block * 8,
206 + worker_kv_bytes_per_block=worker_kv_bytes_per_block,
207 + world_size=4,
208 + replicated_layout=True,
209 + )
210 + assert isinstance(spec, TieringOffloadingSpec)
211
406 - offloading_config = build_offloading_config(config, _make_kv_cache_config())
212 + region = MagicMock()
213 + region_calls: list[dict[str, Any]] = []
214 + worker_calls: list[dict[str, Any]] = []
215
408 - assert offloading_config.parallel.data_parallel_index == 2
216 + def fake_region_ctor(**kwargs):
217 + region_calls.append(kwargs)
218 + return region
219
220 + def fake_worker_ctor(**kwargs):
221 + worker_calls.append(kwargs)
222 + return MagicMock()
223
411 -def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes():
412 - config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
413 - config.cache_config.block_size = 4
224 + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor)
225 + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", fake_worker_ctor)
226 + monkeypatch.setattr(
227 + tiering_spec_module.torch.accelerator, "current_device_index", lambda: 5
228 + )
229
415 - spec = _create_spec(config, _make_hybrid_kv_cache_config())
230 + kv_caches = MagicMock()
231 + spec.create_worker(kv_caches)
232 +
233 + assert region_calls[0]["rank"] == 0
234 + assert region_calls[0]["kv_bytes_per_block"] == worker_kv_bytes_per_block
235 + assert worker_calls[0]["kv_caches"] is kv_caches
236 + assert worker_calls[0]["mmap_region"] is region
237
417 - assert spec.tokens_per_block == (12, 16)
418 - assert spec.tokens_per_hash == 4
419 - assert spec.blocks_per_chunk == 1
238
239 +def test_tiering_spec_create_worker_folds_device_index_for_sharded_layout(monkeypatch):
240 + import vllm.v1.kv_offload.tiering.spec as tiering_spec_module
241
422 -def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec:
423 - return FullAttentionSpec(
424 - block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32
242 + spec = _create_spec(
243 + spec_name="TieringOffloadingSpec",
244 + worker_kv_bytes_per_block=4096,
245 + world_size=4,
246 )
247 + assert isinstance(spec, TieringOffloadingSpec)
248 +
249 + region_calls: list[dict[str, Any]] = []
250
251 + def fake_region_ctor(**kwargs):
252 + region_calls.append(kwargs)
253 + return MagicMock()
254
428 -def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
429 - config = _make_layout_vllm_config()
430 - kv_cache_config = KVCacheConfig(
431 - num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
255 + monkeypatch.setattr(tiering_spec_module, "SharedOffloadRegion", fake_region_ctor)
256 + monkeypatch.setattr(tiering_spec_module, "CPUOffloadingWorker", MagicMock())
257 + monkeypatch.setattr(
258 + tiering_spec_module.torch.accelerator,
259 + "current_device_index",
260 + lambda: 5,
261 )
433 - offloading_config = build_offloading_config(config, kv_cache_config)
434 - return offloading_config.parallel.is_parallelism_agnostic
435 -
436 -
437 -def test_parallelism_agnostic_for_single_full_attention_group():
438 - assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
439 -
440 -
441 -@pytest.mark.parametrize(
442 - "kv_cache_groups",
443 - [
444 - # MLA latent KV is replicated per rank, never head-sharded.
445 - [
446 - KVCacheGroupSpec(
447 - ["l0"],
448 - MLAAttentionSpec(
449 - block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32
450 - ),
451 - )
452 - ],
453 - # Sliding window is not full attention.
454 - [
455 - KVCacheGroupSpec(
456 - ["l0"],
457 - SlidingWindowSpec(
458 - block_size=16,
459 - num_kv_heads=4,
460 - head_size=128,
461 - dtype=torch.float32,
462 - sliding_window=128,
463 - ),
464 - )
465 - ],
466 - # Hybrid model: more than one KV cache group.
467 - [
468 - KVCacheGroupSpec(["l0"], _full_attention_spec()),
469 - KVCacheGroupSpec(["l1"], _full_attention_spec()),
470 - ],
471 - ],
472 -)
473 -def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
474 - assert not _parallelism_agnostic(kv_cache_groups)
262 +
263 + spec.create_worker(MagicMock())
264 +
265 + assert region_calls[0]["rank"] == 1
266
267
477 -def test_parallelism_agnostic_disabled_on_v2_model_runner():
478 - config = _make_layout_vllm_config()
479 - config.use_v2_model_runner = True
480 - kv_cache_config = KVCacheConfig(
481 - num_blocks=0,
482 - kv_cache_tensors=[],
483 - kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
268 +@pytest.mark.parametrize("world_size", [2, 4, 8])
269 +def test_cpu_spec_replicated_config_preserves_per_rank_sizing(world_size: int):
270 + worker_kv_bytes_per_block = 4096
271 + spec = _create_spec(
272 + cpu_bytes_to_use=worker_kv_bytes_per_block * world_size * 2,
273 + worker_kv_bytes_per_block=worker_kv_bytes_per_block,
274 + world_size=world_size,
275 + replicated_layout=True,
276 )
485 - offloading_config = build_offloading_config(config, kv_cache_config)
486 - assert not offloading_config.parallel.is_parallelism_agnostic
277
278 + assert isinstance(spec, CPUOffloadingSpec)
279 + assert spec.replicated_layout is False
280 + assert spec.cpu_page_size_per_worker == worker_kv_bytes_per_block
281 + assert spec.kv_bytes_per_chunk == worker_kv_bytes_per_block * world_size
282 + assert spec.num_blocks == 2
283
489 -def test_create_dynamic_spec_receives_translated_config():
490 - config = _make_layout_vllm_config(
491 - spec_name="SingleArgExternalOffloadingSpec",
492 - extra_config={
493 - "spec_module_path": "tests.v1.kv_offload.test_factory",
494 - },
284 +
285 +def test_offloading_spec_has_replicated_layout_default():
286 + spec = SingleArgExternalOffloadingSpec(_make_offloading_config())
287 + assert spec.replicated_layout is False
288 +
289 +
290 +def test_offloading_spec_uses_normalized_chunk_geometry():
291 + groups = (
292 + OffloadingGroupConfig(12, ("full_layer",)),
293 + OffloadingGroupConfig(16, ("mla_layer",)),
294 + )
295 + spec = _create_spec(
296 + groups=groups,
297 + tokens_per_hash=4,
298 + blocks_per_chunk=2,
299 )
496 - kv_cache_config = _make_kv_cache_config()
497 - offloading_config = build_offloading_config(config, kv_cache_config)
300
499 - spec = OffloadingSpecFactory.create_spec(offloading_config)
301 + assert spec.tokens_per_block == (12, 16)
302 + assert spec.tokens_per_hash == 4
303 + assert spec.blocks_per_chunk == 2
304 +
305
501 - assert isinstance(spec, SingleArgExternalOffloadingSpec)
502 - assert spec.config is offloading_config
306 +def test_create_dynamic_spec_receives_config():
307 + config = _make_offloading_config(
308 + spec_name="SingleArgExternalOffloadingSpec",
309 + extra_config={"spec_module_path": "tests.v1.kv_offload.test_factory"},
310 + )
311
312 + spec = OffloadingSpecFactory.create_spec(config)
313
505 -# ---------------------------------------------------------------------------
506 -# Dynamic import via spec_module_path
507 -# ---------------------------------------------------------------------------
314 + assert isinstance(spec, SingleArgExternalOffloadingSpec)
315 + assert spec.config is config
316
317
318 def test_dynamic_load_via_spec_module_path():
511 - """External spec loaded via spec_module_path.
512 -
513 - This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec)
514 - integrate with vLLM without being pre-registered in the factory.
515 - The fallback path: registry miss → spec_module_path → importlib.import_module.
516 - """
517 - config = _make_vllm_config(spec_name="CPUOffloadingSpec")
518 - # Delete from registry to force the dynamic import path
319 del OffloadingSpecFactory._registry["CPUOffloadingSpec"]
520 - # spec_name not in registry → falls through to spec_module_path
521 - config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
522 - "vllm.v1.kv_offload.cpu.spec"
320 + config = _make_offloading_config(
321 + extra_config={"spec_module_path": "vllm.v1.kv_offload.cpu.spec"}
322 )
524 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
525 - assert spec_cls is CPUOffloadingSpec
323
324 + spec_cls = OffloadingSpecFactory.get_spec_cls(config.extra_config)
325
528 -# ---------------------------------------------------------------------------
529 -# Error paths
530 -# ---------------------------------------------------------------------------
326 + assert spec_cls is CPUOffloadingSpec
327
328
329 def test_unregistered_spec_without_module_path_raises():
534 - """spec_name not in registry + no spec_module_path → ValueError."""
535 - config = _make_vllm_config(spec_name="NonexistentSpec")
330 + config = _make_offloading_config(spec_name="NonexistentSpec")
331 with pytest.raises(ValueError, match="Unsupported spec type"):
537 - OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
332 + OffloadingSpecFactory.get_spec_cls(config.extra_config)
333
539 - # create_spec should also fail (calls get_spec_cls internally)
540 - kv_cache_config = _make_kv_cache_config()
334 with pytest.raises(ValueError, match="Unsupported spec type"):
542 - _create_spec(config, kv_cache_config)
335 + OffloadingSpecFactory.create_spec(config)
336
337
338 def test_cpu_spec_missing_cpu_bytes_to_use_raises():
546 - """CPUOffloadingSpec requires cpu_bytes_to_use → Exception."""
547 - config = _make_vllm_config(cpu_bytes_to_use=None)
548 - config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
549 - kv_cache_config = _make_kv_cache_config()
339 with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
551 - _create_spec(config, kv_cache_config)
340 + _create_spec(cpu_bytes_to_use=None)
341
342
343 def test_duplicate_registration_raises():
555 - """register_spec with existing name → ValueError."""
344 with pytest.raises(ValueError, match="is already registered"):
345 OffloadingSpecFactory.register_spec(
346 "CPUOffloadingSpec", "some.module", "SomeClass"
347 )
348
349
562 -# ---------------------------------------------------------------------------
563 -# Downstream collaboration — build_metric_definitions
564 -# ---------------------------------------------------------------------------
565 -
566 -
350 def test_build_metric_definitions_below_threshold():
568 - """store_threshold < 2 keeps stores_skipped disabled."""
351 from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
352
571 - config = _make_vllm_config(store_threshold=1)
572 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
573 - metrics = spec_cls.build_metric_definitions(
574 - config.kv_transfer_config.kv_connector_extra_config
575 - )
353 + extra_config = {"store_threshold": 1}
354 + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
355 + metrics = spec_cls.build_metric_definitions(extra_config)
356 +
357 assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
358 assert CPUOffloadingMetrics.CPU_ALLOCATION_SIZE in metrics
359
360
361 def test_build_metric_definitions_allocation_size_histogram():
581 - """CPU allocation size is always reported as a histogram."""
362 from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
363
584 - config = _make_vllm_config(store_threshold=0)
585 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
586 - metrics = spec_cls.build_metric_definitions(
587 - config.kv_transfer_config.kv_connector_extra_config
588 - )
364 + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
365 + metrics = spec_cls.build_metric_definitions({})
366 metadata = metrics[CPUOffloadingMetrics.CPU_ALLOCATION_SIZE]
367 +
368 assert isinstance(metadata, OffloadingHistogramMetadata)
369 assert metadata.buckets == (
370 1,
@@ -603,49 +381,10 @@ def test_build_metric_definitions_allocation_size_histogram():
381
382
383 def test_build_metric_definitions_returns_counter_at_threshold():
606 - """store_threshold >= 2 → returns stores_skipped counter definition."""
384 from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
385
609 - config = _make_vllm_config(store_threshold=2)
610 - spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config))
611 - metrics = spec_cls.build_metric_definitions(
612 - config.kv_transfer_config.kv_connector_extra_config
613 - )
614 - assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
615 -
616 -
617 -def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups():
618 - config = _make_layout_vllm_config(
619 - cpu_bytes_to_use=65536,
620 - extra_config={"blocks_per_chunk": 2},
621 - )
622 -
623 - spec = _create_spec(config, _make_hybrid_kv_cache_config())
624 -
625 - assert spec.tokens_per_block == (12, 16)
626 - assert spec.blocks_per_chunk == 2
627 -
628 -
629 -def test_block_size_and_blocks_per_chunk_are_mutually_exclusive():
630 - config = _make_layout_vllm_config(
631 - cpu_bytes_to_use=65536,
632 - extra_config={
633 - "block_size": 64,
634 - "blocks_per_chunk": 2,
635 - },
636 - )
637 -
638 - with pytest.raises(ValueError, match="Specify only one"):
639 - _create_spec(config, _make_kv_cache_config())
640 -
386 + extra_config = {"store_threshold": 2}
387 + spec_cls = OffloadingSpecFactory.get_spec_cls({"spec_name": "CPUOffloadingSpec"})
388 + metrics = spec_cls.build_metric_definitions(extra_config)
389
642 -def test_blocks_per_chunk_must_be_positive():
643 - config = _make_layout_vllm_config(
644 - cpu_bytes_to_use=65536,
645 - extra_config={
646 - "blocks_per_chunk": 0,
647 - },
648 - )
649 -
650 - with pytest.raises(ValueError, match="greater than 0"):
651 - _create_spec(config, _make_kv_cache_config())
390 + assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+28 -7
@@ -102,19 +102,39 @@ def build_offloading_config(
102 )
103 worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
104
105 - # Only a single non-MLA full-attention group is parallelism-invariant:
106 - # MLA latent KV is replicated per rank (never head-sharded), and the V2
107 - # model runner's KV layout is not known to be parallelism-invariant.
108 - single_group = (
105 + single_group_spec = (
106 kv_cache_config.kv_cache_groups[0].kv_cache_spec
107 if len(kv_cache_config.kv_cache_groups) == 1
108 else None
109 )
110 + replicated_layout = (
111 + vllm_config.model_config.use_mla
112 + # Exact type: fail closed on wrappers and sliding-window variants.
113 + and type(single_group_spec) is MLAAttentionSpec
114 + # Page accounting: one MLA page per layer, no packed/mixed rows.
115 + and worker_kv_bytes_per_block > 0
116 + and worker_kv_bytes_per_block
117 + == single_group_spec.page_size_bytes
118 + * len(kv_cache_config.kv_cache_groups[0].layer_names)
119 + # Safe MVP boundary: TP-only, no other parallel axes.
120 + and parallel_config.tensor_parallel_size > 1
121 + and parallel_config.pipeline_parallel_size == 1
122 + and parallel_config.prefill_context_parallel_size == 1
123 + and parallel_config.decode_context_parallel_size == 1
124 + and parallel_config.world_size == parallel_config.tensor_parallel_size
125 + # Shared /dev/shm mmap layout is single-node mp only.
126 + and parallel_config.distributed_executor_backend == "mp"
127 + and parallel_config.nnodes_within_dp == 1
128 + )
129 +
130 + # Only a single non-MLA full-attention group is parallelism-invariant:
131 + # MLA latent KV is replicated per rank (never head-sharded), and the V2
132 + # model runner's KV layout is not known to be parallelism-invariant.
133 is_parallelism_agnostic = (
134 not vllm_config.use_v2_model_runner
115 - and single_group is not None
116 - and isinstance(single_group, FullAttentionSpec)
117 - and not isinstance(single_group, MLAAttentionSpec)
135 + and single_group_spec is not None
136 + and isinstance(single_group_spec, FullAttentionSpec)
137 + and not isinstance(single_group_spec, MLAAttentionSpec)
138 )
139
140 kv_events_config = vllm_config.kv_events_config
@@ -144,4 +164,5 @@ def build_offloading_config(
164 data_parallel_index=parallel_config.data_parallel_index,
165 is_parallelism_agnostic=is_parallelism_agnostic,
166 ),
167 + replicated_layout=replicated_layout,
168 )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py
+11
@@ -45,6 +45,10 @@ class OffloadingConnectorWorker:
45 self.spec = spec
46 self.kv_cache_config = kv_cache_config
47 self.worker: OffloadingWorker | None = None
48 + # Non-writers still ack: pending_count waits for world_size per job.
49 + self._is_store_writer = (
50 + not self.spec.replicated_layout or self.spec.config.parallel.rank == 0
51 + )
52
53 # job_id -> req_id for in-flight loads.
54 self._load_jobs: dict[int, ReqId] = {}
@@ -274,6 +278,9 @@ class OffloadingConnectorWorker:
278 for job_id in kv_connector_metadata.jobs_to_flush:
279 entry = kv_connector_metadata.store_jobs.pop(job_id, None)
280 if entry is not None:
281 + if not self._is_store_writer:
282 + self._connector_worker_meta.mark_completed(job_id)
283 + continue
284 assert isinstance(entry.src_spec, GPULoadStoreSpec)
285 self._unsubmitted_store_jobs.append(
286 (job_id, entry.src_spec, entry.dst_spec)
@@ -304,6 +311,10 @@ class OffloadingConnectorWorker:
311
312 def prepare_store_kv(self, metadata: OffloadingConnectorMetadata):
313 for job_id, entry in metadata.store_jobs.items():
314 + if not self._is_store_writer:
315 + # Gate before queueing: no _unsubmitted_store_jobs entry.
316 + self._connector_worker_meta.mark_completed(job_id)
317 + continue
318 # NOTE(orozery): defer the store to the beginning of the next
319 # engine step, so that offloading starts AFTER transfers related
320 # to token sampling, thereby avoiding delays to token generation.
vllm/v1/kv_offload/base.py
+1
@@ -496,6 +496,7 @@ class OffloadingSpec(ABC):
496 )
497 self.config = config
498 self.extra_config = config.extra_config
499 + self.replicated_layout: bool = False
500 self.kv_events_config = OffloadingKVEventsConfig(
501 enable_kv_cache_events=config.enable_kv_cache_events,
502 self_describing_kv_events=bool(
vllm/v1/kv_offload/config.py
+6
@@ -68,3 +68,9 @@ class OffloadingConfig:
68 model: OffloadingModelConfig
69 cache: OffloadingCacheConfig
70 parallel: OffloadingParallelConfig
71 + # True when the offloaded bytes of every worker are expected to be
72 + # byte-identical per block (pure-MLA model, single-node TP-only
73 + # parallelism), enabling a single-copy host layout in backends that
74 + # support it. Aggregate layout decision; per-layer replication metadata
75 + # is planned for CanonicalKVCacheRef (#48408).
76 + replicated_layout: bool = False
vllm/v1/kv_offload/cpu/spec.py
+8 -2
@@ -24,6 +24,7 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
24
25 class CPUOffloadingSpec(OffloadingSpec):
26 BLOCK_SIZE_ALIGNMENT = 1
27 + SUPPORTS_REPLICATED_LAYOUT = False
28
29 @classmethod
30 def build_metric_definitions(
@@ -85,12 +86,16 @@ class CPUOffloadingSpec(OffloadingSpec):
86 self.num_blocks = 0
87 self.kv_bytes_per_chunk = 0
88 self.cpu_page_size_per_worker = 0
89 + self.replicated_layout = (
90 + config.replicated_layout and self.SUPPORTS_REPLICATED_LAYOUT
91 + )
92 if config.worker_kv_bytes_per_block > 0 and world_size > 0:
89 - kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size
93 + num_copies = 1 if self.replicated_layout else world_size
94 + kv_bytes_per_block = config.worker_kv_bytes_per_block * num_copies
95 kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk
96
97 # calculate cpu_page_size_per_worker
93 - self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size
98 + self.cpu_page_size_per_worker = kv_bytes_per_chunk // num_copies
99
100 # calculate num_blocks
101 aligned_kv_bytes_per_chunk = round_up(
@@ -102,6 +107,7 @@ class CPUOffloadingSpec(OffloadingSpec):
107 # kv_bytes_per_chunk. Note that this might contain
108 # some padding. i.e. each offloaded block is of the form,
109 # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** |
110 + # or |--- B0 (single copy) ---| *** maybe-pad *** |
111 self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk
112
113 # scheduler-side
vllm/v1/kv_offload/tiering/spec.py
+7 -3
@@ -71,6 +71,7 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
71 """
72
73 BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
74 + SUPPORTS_REPLICATED_LAYOUT = True
75
76 @classmethod
77 @override
@@ -229,10 +230,13 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
230
231 @override
232 def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker:
232 - # Fold the global physical device index into the replica-local
233 - # [0, world_size) slot range.
233 world_size = self.config.parallel.world_size
235 - rank = torch.accelerator.current_device_index() % world_size
234 + if self.replicated_layout:
235 + rank = 0
236 + else:
237 + # Fold the global physical device index into the replica-local
238 + # [0, world_size) slot range.
239 + rank = torch.accelerator.current_device_index() % world_size
240 worker_mmap = SharedOffloadRegion(
241 engine_id=self._engine_id,
242 num_blocks=self.num_blocks,