A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 02:30 UTC
[Bugfix][KV Transfer] Reject NixlConnector + expandable_segments:True (#41237)
Roy Wang committed
May 9, 2026 at 07:47 UTC
60851b1d22749de50da36287301fe994118941c2
2 files changed
+98
tests/v1/kv_connector/unit/test_config.py
+55
@@ -65,6 +65,61 @@ def test_kv_connector(
65
assert "existing_key" not in kv_connector_extra_config
66
67
68
+def _build_config(
69
+ *, kv_connector: str | None, enable_sleep_mode: bool = False
70
+) -> VllmConfig:
71
+ """Build a VllmConfig that exercises _verify_kv_transfer_compat without
72
+ requiring a real model (avoids HF downloads in CI)."""
73
+ from types import SimpleNamespace
74
+
75
+ kv_transfer_config = (
76
+ KVTransferConfig(kv_connector=kv_connector, kv_role="kv_both")
77
+ if kv_connector is not None
78
+ else None
79
+ )
80
+ cfg = VllmConfig.__new__(VllmConfig)
81
+ cfg.kv_transfer_config = kv_transfer_config
82
+ cfg.model_config = SimpleNamespace(enable_sleep_mode=enable_sleep_mode)
83
+ cfg._verify_kv_transfer_compat()
84
+ return cfg
85
+
86
+
87
+@pytest.mark.parametrize(
88
+ "kv_connector", ["NixlConnector", "MooncakeConnectorV1", "SomeOOTConnector"]
89
+)
90
+def test_kv_connector_rejects_expandable_segments(monkeypatch, kv_connector):
91
+ """KV connectors that pin KV cache memory (e.g. via ibv_reg_mr) are
92
+ invalidated when expandable_segments lets the CUDA VMM allocator remap
93
+ the underlying physical pages. We can't enumerate every connector that
94
+ does this (especially OOT ones), so reject the combination whenever any
95
+ connector is configured."""
96
+ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
97
+ with pytest.raises(ValueError, match="expandable_segments"):
98
+ _build_config(kv_connector=kv_connector)
99
+
100
+
101
+def test_kv_connector_allows_expandable_segments_with_sleep_mode(monkeypatch):
102
+ """Sleep mode routes KV allocations through CuMemAllocator's pool, which
103
+ auto-disables expandable_segments (see #40812)."""
104
+ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
105
+ _build_config(kv_connector="NixlConnector", enable_sleep_mode=True)
106
+
107
+
108
+def test_kv_connector_allows_other_alloc_conf(monkeypatch):
109
+ """Other PYTORCH_CUDA_ALLOC_CONF values must not be rejected."""
110
+ monkeypatch.setenv(
111
+ "PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:512,expandable_segments:False"
112
+ )
113
+ _build_config(kv_connector="NixlConnector")
114
+
115
+
116
+def test_no_kv_connector_ignores_expandable_segments(monkeypatch):
117
+ """The expandable_segments check only applies when a KV connector is
118
+ configured."""
119
+ monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
120
+ _build_config(kv_connector=None)
121
+
122
+
123
def test_kv_offloading_size_only_uses_native_default():
124
"""Test that setting only kv_offloading_size enables native offloading."""
125
vllm_config = VllmConfig(
vllm/config/vllm.py
+43
@@ -713,6 +713,48 @@ class VllmConfig:
713
# This is the same for all backends
714
self.kv_transfer_config.kv_role = "kv_both"
715
716
+ def _verify_kv_transfer_compat(self) -> None:
717
+ """Reject configurations that silently corrupt KV transfers."""
718
+ if (
719
+ self.kv_transfer_config is None
720
+ or self.kv_transfer_config.kv_connector is None
721
+ ):
722
+ return
723
+
724
+ # PyTorch's expandable_segments allocator uses CUDA VMM, which can
725
+ # remap a virtual address range to different physical pages over the
726
+ # engine's lifetime. KV connectors that pin KV cache memory (e.g.
727
+ # NixlConnector via ibv_reg_mr, MooncakeConnector) end up with their
728
+ # registrations pointing at stale physical pages after any remap,
729
+ # producing RDMA failures like IBV_WC_REM_ACCESS_ERR /
730
+ # NIXL_ERR_REMOTE_DISCONNECT at the first inter-node KV transfer.
731
+ # We can't enumerate every in-tree and out-of-tree connector that
732
+ # pins memory, so we conservatively reject the combination whenever
733
+ # any KV connector is configured.
734
+ #
735
+ # Sleep mode is exempt: CuMemAllocator.use_memory_pool toggles
736
+ # expandable_segments off around its pool (see #40812), so the KV
737
+ # cache allocated within that context lands on stable physical pages
738
+ # even when the env var is set.
739
+ if "expandable_segments:True" not in os.environ.get(
740
+ "PYTORCH_CUDA_ALLOC_CONF", ""
741
+ ):
742
+ return
743
+ if self.model_config is not None and self.model_config.enable_sleep_mode:
744
+ return
745
+
746
+ raise ValueError(
747
+ f"KV connector {self.kv_transfer_config.kv_connector} is "
748
+ "incompatible with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True "
749
+ "unless enable_sleep_mode is also enabled. PyTorch's CUDA VMM "
750
+ "allocator can remap KV cache virtual addresses to different "
751
+ "physical pages, invalidating any pinned/registered KV memory "
752
+ "(e.g. IB memory regions registered by NIXL or Mooncake). Either "
753
+ "unset expandable_segments:True or enable sleep mode (which "
754
+ "routes KV allocations through CuMemAllocator's pool, where "
755
+ "expandable_segments is automatically disabled)."
756
+ )
757
+
758
def __post_init__(self):
759
"""Verify configs are valid & consistent with each other."""
760
@@ -1353,6 +1395,7 @@ class VllmConfig:
1395
1396
# Handle the KV connector configs
1397
self._post_init_kv_transfer_config()
1398
+ self._verify_kv_transfer_compat()
1399
1400
# Log the custom passes that are enabled
1401
self.compilation_config.pass_config.log_enabled_passes()