diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 5541cd99bd..7ad839b88a 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -57,7 +57,9 @@ th:not(:first-child) { | Multimodal | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ | | Encoder-Decoder | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -1 P and D instances must use the same speculation configuration. +1 P and D instances must use compatible speculation configurations. +See [Configuration Notes](#configuration-notes) below for what must match and +what may differ. 2 Requires `FLASH_ATTN` or `FLASHINFER` backend **and** `HND` KV cache layout. Enable via `--kv-transfer-config '{"kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'`. @@ -79,6 +81,7 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - Model (architecture, dtype, number of KV heads, head size, number of hidden layers) - Attention backend - KV cache dtype (`cache_dtype`) +- EAGLE/MTP-style speculative method and draft-model configuration !!! warning Disable the hash check with `--kv-transfer-config '{"kv_connector_extra_config": {"enforce_handshake_compat": false}}'` at your own risk. @@ -88,6 +91,10 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - `tensor-parallel-size` (heterogeneous TP, subject to model restrictions above) - `block-size` (heterogeneous block size, subject to restrictions above) - Number of KV cache blocks (determined by available memory on each instance) +- `num_speculative_tokens` (prefill and decode may use different draft depths) +- Draft-model `attention_backend` (each instance auto-selects independently; the + resulting KV block layout is validated at handshake time rather than via the + compatibility hash) ### KV cache layout diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index ba9f73f240..244594d58e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -9,6 +9,7 @@ import textwrap import time import uuid from collections import defaultdict +from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, patch @@ -2796,6 +2797,115 @@ def test_failed_request_skips_kv_postprocessing( assert invalid_blocks == {1, 2, 3} +def _set_test_speculative_config( + vllm_config, + *, + method: str = "eagle3", + model: str = "test/eagle3-drafter", + revision: str | None = None, + code_revision: str | None = None, + num_speculative_tokens: int = 1, + parallel_drafting: bool = False, + kv_cache_dtype: str | None = None, + attention_backend: str | None = None, + auxiliary_layer_ids: tuple[int, ...] = (2, 16, 29), +) -> None: + draft_model_config = SimpleNamespace( + model=model, + revision=revision, + code_revision=code_revision, + hf_config=SimpleNamespace(eagle_aux_hidden_state_layer_ids=auxiliary_layer_ids), + ) + vllm_config.speculative_config = SimpleNamespace( + method=method, + draft_model_config=draft_model_config, + num_speculative_tokens=num_speculative_tokens, + parallel_drafting=parallel_drafting, + kv_cache_dtype=kv_cache_dtype, + attention_backend=attention_backend, + use_eagle=lambda: True, + ) + + +@pytest.mark.parametrize( + "remote_overrides,should_match", + [ + ({}, True), + ({"num_speculative_tokens": 2}, True), + ({"method": "mtp"}, False), + ({"model": "test/different-drafter"}, False), + ({"revision": "different-revision"}, False), + ({"parallel_drafting": True}, False), + ({"kv_cache_dtype": "fp8"}, False), + # attention_backend is intentionally not part of the compat hash + # (see _get_speculative_compatibility_factors); overriding it must + # not change the hash. + ({"attention_backend": "FLASHINFER"}, True), + ({"auxiliary_layer_ids": (2, 16, 30)}, False), + ], +) +@pytest.mark.skip_global_cleanup +def test_speculative_config_compatibility_hash( + remote_overrides: dict[str, Any], should_match: bool +): + local_config = create_vllm_config() + remote_config = create_vllm_config() + _set_test_speculative_config(local_config) + _set_test_speculative_config(remote_config, **remote_overrides) + + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + + assert (local_hash == remote_hash) is should_match + + +@pytest.mark.skip_global_cleanup +def test_missing_speculative_config_changes_compatibility_hash(): + regular_config = create_vllm_config() + speculative_config = create_vllm_config() + _set_test_speculative_config(speculative_config) + + regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN", False) + speculative_hash = compute_nixl_compatibility_hash( + speculative_config, "FLASH_ATTN", False + ) + + assert regular_hash != speculative_hash + + +@pytest.mark.skip_global_cleanup +def test_speculative_kv_cache_dtype_resolves_to_target(): + # The draft kv_cache_dtype override defaults to None ("inherit the target's + # --kv-cache-dtype"). An explicit setting on one side that matches the + # other side's inherited (resolved) dtype must not spuriously mismatch. + local_config = create_vllm_config(cache_dtype="fp8") + remote_config = create_vllm_config(cache_dtype="fp8") + _set_test_speculative_config(local_config, kv_cache_dtype="fp8") # explicit + _set_test_speculative_config(remote_config, kv_cache_dtype=None) # inherits + + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + + assert local_hash == remote_hash + + +@pytest.mark.skip_global_cleanup +def test_speculative_attention_backend_not_in_compatibility_hash(): + # The draft attention_backend is intentionally excluded from the hash: the + # connector only has the raw override (auto-select), and its transfer- + # relevant effect is validated per region at runtime. Differing overrides + # must not change the hash. + local_config = create_vllm_config() + remote_config = create_vllm_config() + _set_test_speculative_config(local_config, attention_backend=None) + _set_test_speculative_config(remote_config, attention_backend="FLASHINFER") + + local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False) + remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False) + + assert local_hash == remote_hash + + @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/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 27eab825a1..4758bab5d7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -41,8 +41,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # 4: Add KV block lease renewal through heartbeats # 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake # clock-sync timestamp +# 6: Validate EAGLE/MTP speculative configuration compatibility # -NIXL_CONNECTOR_VERSION: int = 5 +NIXL_CONNECTOR_VERSION: int = 6 @dataclass @@ -78,6 +79,49 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata): agent_metadata_bytes: bytes # NixlAgentMetadata encoded +def _get_speculative_compatibility_factors( + vllm_config: VllmConfig, +) -> dict[str, Any] | None: + """Return NIXL compatibility factors for hidden-state-based speculators.""" + speculative_config = vllm_config.speculative_config + if speculative_config is None or not speculative_config.use_eagle(): + return None + + draft_model_config = speculative_config.draft_model_config + assert draft_model_config is not None + auxiliary_layer_ids = getattr( + draft_model_config.hf_config, + "eagle_aux_hidden_state_layer_ids", + None, + ) + + # kv_cache_dtype is a user override that defaults to None, meaning "inherit + # the target's --kv-cache-dtype". Resolve it to the effective value so an + # explicit setting on one side and inheritance on the other (same effective + # dtype) don't spuriously mismatch. + kv_cache_dtype = ( + speculative_config.kv_cache_dtype or vllm_config.cache_config.cache_dtype + ) + + # Note: the draft attention_backend is intentionally not hashed. Its only + # transfer-relevant effect is the KV block layout/size, which is validated + # per region at runtime in _validate_remote_agent_handshake. The connector + # only sees the raw override here (usually None = auto-select), never the + # resolved backend, so hashing it would cause false mismatches without + # catching anything the runtime layout check misses. + return { + "method": speculative_config.method, + "model": draft_model_config.model, + "revision": draft_model_config.revision, + "code_revision": draft_model_config.code_revision, + "parallel_drafting": speculative_config.parallel_drafting, + "kv_cache_dtype": str(kv_cache_dtype), + "auxiliary_layer_ids": ( + tuple(auxiliary_layer_ids) if auxiliary_layer_ids is not None else None + ), + } + + def compute_nixl_compatibility_hash( vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool ) -> str: @@ -92,6 +136,7 @@ def compute_nixl_compatibility_hash( - Model architecture (name, dtype, KV heads, layers) - KV cache format (dtype, sliding window) - Attention backend + - EAGLE/MTP configuration that affects transferred state Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout are validated at runtime in _validate_remote_agent_handshake and are not @@ -125,6 +170,7 @@ def compute_nixl_compatibility_hash( "cache_dtype": str(cache_config.cache_dtype), "cross_layers_blocks": cross_layers_blocks, "is_hma_enabled": is_hma_enabled, + "speculative_config": _get_speculative_compatibility_factors(vllm_config), } compat_hash = hash_factors(factors)