[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>
Tzu-Ling Kan committed
Aug 3, 2026 at 13:51 UTC
952694e3843e478dd99cffd132c756d582fe8a94
3 files changed
+165
-2
docs/features/nixl_connector_compatibility.md
+8
-1
@@ -57,7 +57,9 @@ th:not(:first-child) {
57
| Multimodal | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ | ❔ |
58
| Encoder-Decoder | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
59
60
-<sup>1</sup> P and D instances must use the same speculation configuration.
60
+<sup>1</sup> P and D instances must use compatible speculation configurations.
61
+See [Configuration Notes](#configuration-notes) below for what must match and
62
+what may differ.
63
64
<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"}}'`.
65
@@ -79,6 +81,7 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
81
- Model (architecture, dtype, number of KV heads, head size, number of hidden layers)
82
- Attention backend
83
- KV cache dtype (`cache_dtype`)
84
+- EAGLE/MTP-style speculative method and draft-model configuration
85
86
!!! warning
87
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
91
- `tensor-parallel-size` (heterogeneous TP, subject to model restrictions above)
92
- `block-size` (heterogeneous block size, subject to restrictions above)
93
- Number of KV cache blocks (determined by available memory on each instance)
94
+- `num_speculative_tokens` (prefill and decode may use different draft depths)
95
+- Draft-model `attention_backend` (each instance auto-selects independently; the
96
+ resulting KV block layout is validated at handshake time rather than via the
97
+ compatibility hash)
98
99
### KV cache layout
100
tests/v1/kv_connector/unit/test_nixl_connector.py
+110
@@ -9,6 +9,7 @@ import textwrap
9
import time
10
import uuid
11
from collections import defaultdict
12
+from types import SimpleNamespace
13
from typing import Any, cast
14
from unittest.mock import MagicMock, patch
15
@@ -2796,6 +2797,115 @@ def test_failed_request_skips_kv_postprocessing(
2797
assert invalid_blocks == {1, 2, 3}
2798
2799
2800
+def _set_test_speculative_config(
2801
+ vllm_config,
2802
+ *,
2803
+ method: str = "eagle3",
2804
+ model: str = "test/eagle3-drafter",
2805
+ revision: str | None = None,
2806
+ code_revision: str | None = None,
2807
+ num_speculative_tokens: int = 1,
2808
+ parallel_drafting: bool = False,
2809
+ kv_cache_dtype: str | None = None,
2810
+ attention_backend: str | None = None,
2811
+ auxiliary_layer_ids: tuple[int, ...] = (2, 16, 29),
2812
+) -> None:
2813
+ draft_model_config = SimpleNamespace(
2814
+ model=model,
2815
+ revision=revision,
2816
+ code_revision=code_revision,
2817
+ hf_config=SimpleNamespace(eagle_aux_hidden_state_layer_ids=auxiliary_layer_ids),
2818
+ )
2819
+ vllm_config.speculative_config = SimpleNamespace(
2820
+ method=method,
2821
+ draft_model_config=draft_model_config,
2822
+ num_speculative_tokens=num_speculative_tokens,
2823
+ parallel_drafting=parallel_drafting,
2824
+ kv_cache_dtype=kv_cache_dtype,
2825
+ attention_backend=attention_backend,
2826
+ use_eagle=lambda: True,
2827
+ )
2828
+
2829
+
2830
+@pytest.mark.parametrize(
2831
+ "remote_overrides,should_match",
2832
+ [
2833
+ ({}, True),
2834
+ ({"num_speculative_tokens": 2}, True),
2835
+ ({"method": "mtp"}, False),
2836
+ ({"model": "test/different-drafter"}, False),
2837
+ ({"revision": "different-revision"}, False),
2838
+ ({"parallel_drafting": True}, False),
2839
+ ({"kv_cache_dtype": "fp8"}, False),
2840
+ # attention_backend is intentionally not part of the compat hash
2841
+ # (see _get_speculative_compatibility_factors); overriding it must
2842
+ # not change the hash.
2843
+ ({"attention_backend": "FLASHINFER"}, True),
2844
+ ({"auxiliary_layer_ids": (2, 16, 30)}, False),
2845
+ ],
2846
+)
2847
+@pytest.mark.skip_global_cleanup
2848
+def test_speculative_config_compatibility_hash(
2849
+ remote_overrides: dict[str, Any], should_match: bool
2850
+):
2851
+ local_config = create_vllm_config()
2852
+ remote_config = create_vllm_config()
2853
+ _set_test_speculative_config(local_config)
2854
+ _set_test_speculative_config(remote_config, **remote_overrides)
2855
+
2856
+ local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2857
+ remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2858
+
2859
+ assert (local_hash == remote_hash) is should_match
2860
+
2861
+
2862
+@pytest.mark.skip_global_cleanup
2863
+def test_missing_speculative_config_changes_compatibility_hash():
2864
+ regular_config = create_vllm_config()
2865
+ speculative_config = create_vllm_config()
2866
+ _set_test_speculative_config(speculative_config)
2867
+
2868
+ regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN", False)
2869
+ speculative_hash = compute_nixl_compatibility_hash(
2870
+ speculative_config, "FLASH_ATTN", False
2871
+ )
2872
+
2873
+ assert regular_hash != speculative_hash
2874
+
2875
+
2876
+@pytest.mark.skip_global_cleanup
2877
+def test_speculative_kv_cache_dtype_resolves_to_target():
2878
+ # The draft kv_cache_dtype override defaults to None ("inherit the target's
2879
+ # --kv-cache-dtype"). An explicit setting on one side that matches the
2880
+ # other side's inherited (resolved) dtype must not spuriously mismatch.
2881
+ local_config = create_vllm_config(cache_dtype="fp8")
2882
+ remote_config = create_vllm_config(cache_dtype="fp8")
2883
+ _set_test_speculative_config(local_config, kv_cache_dtype="fp8") # explicit
2884
+ _set_test_speculative_config(remote_config, kv_cache_dtype=None) # inherits
2885
+
2886
+ local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2887
+ remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2888
+
2889
+ assert local_hash == remote_hash
2890
+
2891
+
2892
+@pytest.mark.skip_global_cleanup
2893
+def test_speculative_attention_backend_not_in_compatibility_hash():
2894
+ # The draft attention_backend is intentionally excluded from the hash: the
2895
+ # connector only has the raw override (auto-select), and its transfer-
2896
+ # relevant effect is validated per region at runtime. Differing overrides
2897
+ # must not change the hash.
2898
+ local_config = create_vllm_config()
2899
+ remote_config = create_vllm_config()
2900
+ _set_test_speculative_config(local_config, attention_backend=None)
2901
+ _set_test_speculative_config(remote_config, attention_backend="FLASHINFER")
2902
+
2903
+ local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
2904
+ remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
2905
+
2906
+ assert local_hash == remote_hash
2907
+
2908
+
2909
@pytest.mark.parametrize(
2910
"mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat",
2911
[
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+47
-1
@@ -41,8 +41,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
41
# 4: Add KV block lease renewal through heartbeats
42
# 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake
43
# clock-sync timestamp
44
+# 6: Validate EAGLE/MTP speculative configuration compatibility
45
#
45
-NIXL_CONNECTOR_VERSION: int = 5
46
+NIXL_CONNECTOR_VERSION: int = 6
47
48
49
@dataclass
@@ -78,6 +79,49 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata):
79
agent_metadata_bytes: bytes # NixlAgentMetadata encoded
80
81
82
+def _get_speculative_compatibility_factors(
83
+ vllm_config: VllmConfig,
84
+) -> dict[str, Any] | None:
85
+ """Return NIXL compatibility factors for hidden-state-based speculators."""
86
+ speculative_config = vllm_config.speculative_config
87
+ if speculative_config is None or not speculative_config.use_eagle():
88
+ return None
89
+
90
+ draft_model_config = speculative_config.draft_model_config
91
+ assert draft_model_config is not None
92
+ auxiliary_layer_ids = getattr(
93
+ draft_model_config.hf_config,
94
+ "eagle_aux_hidden_state_layer_ids",
95
+ None,
96
+ )
97
+
98
+ # kv_cache_dtype is a user override that defaults to None, meaning "inherit
99
+ # the target's --kv-cache-dtype". Resolve it to the effective value so an
100
+ # explicit setting on one side and inheritance on the other (same effective
101
+ # dtype) don't spuriously mismatch.
102
+ kv_cache_dtype = (
103
+ speculative_config.kv_cache_dtype or vllm_config.cache_config.cache_dtype
104
+ )
105
+
106
+ # Note: the draft attention_backend is intentionally not hashed. Its only
107
+ # transfer-relevant effect is the KV block layout/size, which is validated
108
+ # per region at runtime in _validate_remote_agent_handshake. The connector
109
+ # only sees the raw override here (usually None = auto-select), never the
110
+ # resolved backend, so hashing it would cause false mismatches without
111
+ # catching anything the runtime layout check misses.
112
+ return {
113
+ "method": speculative_config.method,
114
+ "model": draft_model_config.model,
115
+ "revision": draft_model_config.revision,
116
+ "code_revision": draft_model_config.code_revision,
117
+ "parallel_drafting": speculative_config.parallel_drafting,
118
+ "kv_cache_dtype": str(kv_cache_dtype),
119
+ "auxiliary_layer_ids": (
120
+ tuple(auxiliary_layer_ids) if auxiliary_layer_ids is not None else None
121
+ ),
122
+ }
123
+
124
+
125
def compute_nixl_compatibility_hash(
126
vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
127
) -> str:
@@ -92,6 +136,7 @@ def compute_nixl_compatibility_hash(
136
- Model architecture (name, dtype, KV heads, layers)
137
- KV cache format (dtype, sliding window)
138
- Attention backend
139
+ - EAGLE/MTP configuration that affects transferred state
140
141
Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout
142
are validated at runtime in _validate_remote_agent_handshake and are not
@@ -125,6 +170,7 @@ def compute_nixl_compatibility_hash(
170
"cache_dtype": str(cache_config.cache_dtype),
171
"cross_layers_blocks": cross_layers_blocks,
172
"is_hma_enabled": is_hma_enabled,
173
+ "speculative_config": _get_speculative_compatibility_factors(vllm_config),
174
}
175
176
compat_hash = hash_factors(factors)