diff --git a/tests/test_config.py b/tests/test_config.py index ea688f7d8c9c..5254327f71bd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -114,7 +114,34 @@ def test_per_request_spec_decode_metrics_requires_spec_decode(): ) -def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog): +@pytest.mark.parametrize( + "kv_transfer_config", + [ + KVTransferConfig( + kv_connector="NixlConnector", + kv_role="kv_both", + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + { + "kv_connector": "NixlConnector", + "kv_role": "kv_both", + }, + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + }, + ] + }, + ), + ], +) +def test_pd_dcp_interleave_size_is_adjusted_to_block_size( + caplog, disable_log_dedup, kv_transfer_config +): config = VllmConfig( cache_config=CacheConfig(block_size=16), device_config=DeviceConfig(device="cpu"), @@ -124,10 +151,7 @@ def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog): cp_kv_cache_interleave_size=3, distributed_executor_backend="mp", ), - kv_transfer_config=KVTransferConfig( - kv_connector="NixlConnector", - kv_role="kv_both", - ), + kv_transfer_config=kv_transfer_config, ) kv_cache_config = SimpleNamespace( @@ -140,6 +164,51 @@ def test_pd_dcp_interleave_size_is_adjusted_to_block_size(caplog): assert "automatically adjusted from 3 to block_size 16" in caplog.text +def test_kv_offloading_does_not_adjust_dcp_interleave_size(): + config = VllmConfig( + cache_config=CacheConfig(block_size=16), + device_config=DeviceConfig(device="cpu"), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + decode_context_parallel_size=2, + cp_kv_cache_interleave_size=1, + distributed_executor_backend="mp", + ), + kv_transfer_config=KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + ), + ) + + kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))] + ) + config.adjust_dcp_kv_cache_interleave_size(kv_cache_config) + + assert config.parallel_config.cp_kv_cache_interleave_size == 1 + + +def test_kv_offloading_does_not_skip_dcp_interleave_validation(): + config = SimpleNamespace( + cache_config=SimpleNamespace( + block_size=16, + mamba_cache_mode="none", + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=2, + cp_kv_cache_interleave_size=3, + ), + scheduler_config=SimpleNamespace(disable_chunked_mm_input=False), + kv_transfer_config=KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + ), + ) + + with pytest.raises(AssertionError, match="divisible by"): + VllmConfig.validate_block_size(config) + + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object config = VllmConfig() diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7fb53cdc24e6..86d4285b0399 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2742,7 +2742,7 @@ def _validate_v1_model_runner(self) -> None: def adjust_dcp_kv_cache_interleave_size( self, kv_cache_config: "KVCacheConfig" ) -> None: - """Normalize DCP interleave size against the resolved block_size for PD. + """Normalize DCP interleave size against block_size for NIXL P/D. Called by each worker (via ensure_kv_transfer_initialized), once it knows its own final block_size via kv_cache_config. @@ -2750,11 +2750,6 @@ def adjust_dcp_kv_cache_interleave_size( dcp_size = self.parallel_config.decode_context_parallel_size if dcp_size <= 1: return - # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid - # scaling by dcp_size (we need the local block_size here). - local_block_size = min( - g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups - ) if self.parallel_config.dcp_kv_cache_interleave_size > 1 and ( self.parallel_config.cp_kv_cache_interleave_size != self.parallel_config.dcp_kv_cache_interleave_size @@ -2768,11 +2763,17 @@ def adjust_dcp_kv_cache_interleave_size( "deprecated when PCP is fully supported." ) - if ( - self.kv_transfer_config is not None - and self.kv_transfer_config.kv_connector is not None - and self.parallel_config.cp_kv_cache_interleave_size != local_block_size + if self.kv_transfer_config is None or not self.kv_transfer_config.has_connector( + "NixlConnector" ): + return + + # Get the kernel block_size, but don't use resolve_kv_cache_block_size to avoid + # scaling by dcp_size (we need the local block_size here). + local_block_size = min( + g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups + ) + if self.parallel_config.cp_kv_cache_interleave_size != local_block_size: interleave = self.parallel_config.cp_kv_cache_interleave_size self.parallel_config.cp_kv_cache_interleave_size = local_block_size logger.info_once( @@ -2793,14 +2794,13 @@ def validate_block_size(self) -> None: """ block_size = self.cache_config.block_size - # Skip DCP interleave-size compatibility when a KV connector is configured: - # cp_kv_cache_interleave_size is pinned to block_size for PD by each worker - pd_active = ( + # Skip DCP interleave-size compatibility for NIXL P/D: the interleave + # size is pinned to block_size by each worker. + nixl_pd_active = ( self.kv_transfer_config is not None - and self.kv_transfer_config.kv_connector is not None - and self.kv_transfer_config.is_kv_transfer_instance + and self.kv_transfer_config.has_connector("NixlConnector") ) - if self.parallel_config.decode_context_parallel_size > 1 and not pd_active: + if self.parallel_config.decode_context_parallel_size > 1 and not nixl_pd_active: assert ( self.parallel_config.cp_kv_cache_interleave_size <= block_size and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0 diff --git a/vllm/distributed/kv_transfer/kv_transfer_state.py b/vllm/distributed/kv_transfer/kv_transfer_state.py index 33193380f210..211d33fff215 100644 --- a/vllm/distributed/kv_transfer/kv_transfer_state.py +++ b/vllm/distributed/kv_transfer/kv_transfer_state.py @@ -85,7 +85,7 @@ def ensure_kv_transfer_initialized( vllm_config.kv_transfer_config.is_kv_transfer_instance and _KV_CONNECTOR_AGENT is None ): - # PD only supports an interleave_size equal to block_size. + # NIXL P/D requires an interleave_size equal to block_size. vllm_config.adjust_dcp_kv_cache_interleave_size(kv_cache_config) _sync_engine_id_across_tp(vllm_config)