Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions docs/features/nixl_connector_compatibility.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
39 changes: 39 additions & 0 deletions tests/v1/kv_connector/unit/test_nixl_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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",
[
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down
15 changes: 13 additions & 2 deletions vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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,
Comment thread
tzulingk marked this conversation as resolved.
}

compat_hash = hash_factors(factors)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
hooks.
"""

_TRANSFER_MODE: str = "push"

def __init__(
self,
vllm_config: VllmConfig,
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading