[Bugfix] Validate NIXL speculative config compatibility (#49230)

Signed-off-by: Tzu-Ling <tzulingk@nvidia.com> Co-authored-by: GPT-5.6 Sol <noreply@openai.com> Co-authored-by: Cursor Grok 4.5 <noreply@cursor.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-08-04 02:51:04 +08:00
parent 65cf1276a2
commit 952694e384
3 changed files with 165 additions and 2 deletions
@@ -57,7 +57,9 @@ th:not(:first-child) {
 | Multimodal | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ |
 | Encoder-Decoder | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |

-<sup>1</sup> P and D instances must use the same speculation configuration.
+<sup>1</sup> P and D instances must use compatible speculation configurations.
+See [Configuration Notes](#configuration-notes) below for what must match and
+what may differ.

 <sup>2</sup> 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

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