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

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Tzu-Ling Kan
2026-08-14 12:11:22 +08:00
committed by GitHub
co-authored by Claude Opus 4.8
parent c05d75aaa9
commit 653cc6faca
8 changed files with 76 additions and 3 deletions
@@ -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.
@@ -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",
[
@@ -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",
@@ -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:
@@ -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)
@@ -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,
)
@@ -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(
@@ -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",