diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index f597c8fa27..acf55972df 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,6 +17,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( SchedulerOffloadConfig, ) from vllm.platforms import current_platform +from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config from vllm.v1.kv_cache_interface import ( FullAttentionSpec, HiddenStateCacheSpec, @@ -32,6 +33,8 @@ from vllm.v1.kv_cache_interface import ( UniformTypeKVCacheSpecs, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.kv_offload.file_mapper import FileMapper +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec def _make_vllm_config( @@ -876,6 +879,43 @@ def test_canonical_layout_certifies_v2_model_runner(): assert _parallelism_agnostic(groups, canonical=True, v2=True) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +@pytest.mark.parametrize("scheduler", [False, True]) +def test_canonical_mla_dsa_rows_are_tp_independent(tp_size, scheduler): + specs = { + f"l{i}": _mla_spec(block_size=64, head_size=size, dtype=torch.uint8) + for i, size in enumerate([656] * 4 + [132] * 3) + } + uniform = UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=4 * uniform.page_size_bytes, + layers=list(specs), + layer_stride=0, + block_stride=uniform.page_size_bytes, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(list(specs), uniform)], + ) + if scheduler: + kv_cache_config = generate_scheduler_kv_cache_config([kv_cache_config]) + config = _make_vllm_config( + tensor_parallel_size=tp_size, + extra_config={"canonical_layout": True, "cpu_bytes_to_use": 196608 * 8}, + ) + config.cache_config.kv_cache_layout = "LBHNC" + config.parallel_config.distributed_executor_backend = "mp" + offloading_config = build_offloading_config(config, kv_cache_config) + spec = TieringOffloadingSpec(offloading_config) + assert spec.kv_bytes_per_chunk == 196608 + assert spec.num_chunks == 8 + mapper = FileMapper.from_offloading_spec("/cache", spec, parallel_agnostic=True) + assert mapper.get_run_config()["tp_size"] == 1 + assert mapper.get_run_config()["replicated_layout"] is True + + def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() config.use_v2_model_runner = True 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 7a987fef3d..64f481d02e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -97,6 +97,7 @@ def _make_worker( kv_cache_config: KVCacheConfig, replicated_layout: bool = False, rank: int = 0, + canonical_layout: bool = False, ): """Create an OffloadingConnectorWorker with mocked dependencies.""" from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import ( @@ -106,6 +107,7 @@ def _make_worker( spec = MagicMock(spec=OffloadingSpec) spec.replicated_layout = replicated_layout spec.config = MagicMock() + spec.config.canonical_layout = canonical_layout spec.config.parallel.rank = rank spec.get_worker.return_value = MagicMock() @@ -204,11 +206,13 @@ def test_prepare_store_kv_non_writer_marks_completed_without_submit(): assert meta.completed_jobs == {7: 1} -def test_prepare_store_kv_writer_submits_store(): +@pytest.mark.parametrize("rank,canonical_layout", [(0, False), (0, True), (1, True)]) +def test_prepare_store_kv_writer_submits_store(rank, canonical_layout): worker, _ = _make_worker( KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]), replicated_layout=True, - rank=0, + rank=rank, + canonical_layout=canonical_layout, ) worker.prepare_store_kv(_store_metadata(8)) 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 c4e6d0f2a2..ca503d4e67 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -233,6 +233,20 @@ def build_offloading_config( and parallel_config.world_size == tp_size ) + if canonical_layout and is_parallelism_agnostic: + replicated_layout = ( + all( + type(spec) is MLAAttentionSpec + for _, group in selected_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ) + and parallel_config.nnodes_within_dp == 1 + and ( + parallel_config.world_size == 1 + or parallel_config.distributed_executor_backend == "mp" + ) + ) + kv_events_config = vllm_config.kv_events_config cache_dtype = ( vllm_config.model_config.dtype 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 94e8a8109d..5b14e50dda 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -53,7 +53,9 @@ class OffloadingConnectorWorker: 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 + not self.spec.replicated_layout + or self.spec.config.canonical_layout + or self.spec.config.parallel.rank == 0 ) # job_id -> req_id for in-flight loads.