[Bugfix][NIXL] Include transfer mode (push/pull) in the compatibility hash (vllm-project#50620)

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 2 people

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.

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",
     [

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -51,6 +51,11 @@
 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",

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -90,6 +90,11 @@
 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 @@ def _register_packed_kv_cache(
             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 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             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:

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

@@ -42,8 +42,9 @@
 #   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)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -277,4 +277,5 @@ def request_finished(
             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,
         )

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 @@ def request_finished(
             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(

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

@@ -77,6 +77,9 @@
 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",