diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 7ad839b88a..36b8876c69 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -82,6 +82,8 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - Attention backend - KV cache dtype (`cache_dtype`) - EAGLE/MTP-style speculative method and draft-model configuration +- NIXL transfer mode (push vs pull) — a push (WRITE) connector and a pull (READ) + connector use incompatible transfer protocols and must never be paired !!! warning Disable the hash check with `--kv-transfer-config '{"kv_connector_extra_config": {"enforce_handshake_compat": false}}'` at your own risk. diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 244594d58e..a37b12f62a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -433,6 +433,9 @@ def test_kv_transfer_handshake(dist_init): ) ) assert delay + # Pull connector advertises its transfer mode in kv_transfer_params so + # an external router can distinguish it from a push producer. + assert kv_connector_metadata["transfer_mode"] == "pull" # Decode connector will be able to create handshake with the prefill connector. decode_connector = NixlConnector( @@ -2906,6 +2909,42 @@ def test_speculative_attention_backend_not_in_compatibility_hash(): assert local_hash == remote_hash +@pytest.mark.skip_global_cleanup +def test_transfer_mode_changes_compatibility_hash(): + # push (WRITE) and pull (READ) connectors use incompatible transfer + # protocols, so their compatibility hashes must differ; identical modes + # must match. The default mode is pull. + config = create_vllm_config() + + pull_hash = compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="pull" + ) + push_hash = compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="push" + ) + + assert pull_hash != push_hash + assert pull_hash == compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="pull" + ) + assert compute_nixl_compatibility_hash(config, "FLASH_ATTN", False) == pull_hash + + +@pytest.mark.skip_global_cleanup +def test_scheduler_advertises_transfer_mode(): + # Each scheduler advertises its transfer mode in kv_transfer_params so an + # external router can route pull (READ) vs push (WRITE) producers. + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( + NixlPushConnectorScheduler, + ) + + assert NixlPullConnectorScheduler._TRANSFER_MODE == "pull" + assert NixlPushConnectorScheduler._TRANSFER_MODE == "push" + + @pytest.mark.parametrize( "mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat", [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index afa16279ff..bd93c23233 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -51,6 +51,11 @@ logger = init_logger(__name__) class NixlBaseConnectorScheduler: """Base implementation of Scheduler side methods shared by pull and push.""" + # Emitted in kv_transfer_params so an external router can distinguish a + # pull (READ) producer from a push (WRITE) one. Overridden by the push + # scheduler. + _TRANSFER_MODE: str = "pull" + def __init__( self, vllm_config: "VllmConfig", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 815eaba7de..7d1f9466a8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -90,6 +90,11 @@ logger = init_logger(__name__) class NixlBaseConnectorWorker: """Base implementation of Worker side methods shared by pull and push.""" + # Transfer mode included in the NIXL compatibility hash so that a push + # (WRITE) connector and a pull (READ) connector never handshake together. + # Overridden by NixlPushConnectorWorker. + _TRANSFER_MODE: str = "pull" + def _compute_desc_ids( self, block_ids: BlockIds, @@ -976,6 +981,7 @@ class NixlBaseConnectorWorker: self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks, + transfer_mode=self._TRANSFER_MODE, ) total_size = storage.nbytes() @@ -1066,7 +1072,10 @@ class NixlBaseConnectorWorker: is_mamba=self._has_mamba, ) self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks + self.vllm_config, + self.backend_name, + self.transfer_topo.cross_layers_blocks, + transfer_mode=self._TRANSFER_MODE, ) if self.use_host_buffer: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 4758bab5d7..ed2e030c73 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -42,8 +42,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake # clock-sync timestamp # 6: Validate EAGLE/MTP speculative configuration compatibility +# 7: Include NIXL transfer mode (push vs pull) in the compatibility hash # -NIXL_CONNECTOR_VERSION: int = 6 +NIXL_CONNECTOR_VERSION: int = 7 @dataclass @@ -123,7 +124,10 @@ def _get_speculative_compatibility_factors( def compute_nixl_compatibility_hash( - vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool + vllm_config: VllmConfig, + attn_backend_name: str, + cross_layers_blocks: bool, + transfer_mode: str = "pull", ) -> str: """ Compute compatibility hash for NIXL KV transfer. @@ -137,6 +141,11 @@ def compute_nixl_compatibility_hash( - KV cache format (dtype, sliding window) - Attention backend - EAGLE/MTP configuration that affects transferred state + - Transfer mode (push vs pull) + + The transfer mode is included because the push (WRITE) and pull (READ) + connectors use incompatible transfer protocols; a push connector and a + pull connector must never complete a handshake with each other. Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout are validated at runtime in _validate_remote_agent_handshake and are not @@ -171,6 +180,8 @@ def compute_nixl_compatibility_hash( "cross_layers_blocks": cross_layers_blocks, "is_hma_enabled": is_hma_enabled, "speculative_config": _get_speculative_compatibility_factors(vllm_config), + # push (WRITE) and pull (READ) connectors are protocol-incompatible + "transfer_mode": transfer_mode, } compat_hash = hash_factors(factors) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 9beceeabad..d0a5383c5c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -277,4 +277,5 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): tp_size=self.vllm_config.parallel_config.tensor_parallel_size, remote_num_tokens=remote_num_tokens, remote_blocks_expiry_time=blocks_expiry_time, + transfer_mode=self._TRANSFER_MODE, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index a02491e487..6efe491be8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -60,6 +60,8 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): hooks. """ + _TRANSFER_MODE: str = "push" + def __init__( self, vllm_config: VllmConfig, @@ -290,6 +292,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): tp_size=self.vllm_config.parallel_config.tensor_parallel_size, pp_size=self.vllm_config.parallel_config.pipeline_parallel_size, remote_num_tokens=remote_num_tokens, + transfer_mode=self._TRANSFER_MODE, ) def build_connector_meta( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 7fdde27dcc..0721bf9d1d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -77,6 +77,9 @@ _PUSH_WRITER_POLL_INTERVAL_MS = 1.0 class NixlPushConnectorWorker(NixlBaseConnectorWorker): """Push-specific (WRITE) worker logic. See module docstring.""" + # Distinguishes push from pull in the NIXL compatibility hash. + _TRANSFER_MODE: str = "push" + def __init__( self, vllm_config: "VllmConfig",