[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:
co-authored by
Claude Opus 4.8
parent
c05d75aaa9
commit
653cc6faca
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user