A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix CPU memory leak related to not cleaning up old remotes data (#44424)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 11, 2026 at 16:54 UTC
750aab5b8e7f81b8d6d8dac33237cfb45a1f1455
4 files changed
+192
-8
tests/v1/kv_connector/unit/test_nixl_connector.py
+118
@@ -20,6 +20,7 @@ import torch
20
from vllm import LLM
21
from vllm.config import KVTransferConfig, set_current_vllm_config
22
from vllm.distributed.kv_transfer.kv_connector.utils import (
23
+ EngineTransferInfo,
24
KVOutputAggregator,
25
TransferTopology,
26
get_current_attn_backend,
@@ -1844,6 +1845,11 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
1845
worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]}
1846
worker.dst_xfer_side_handles = {"engine1": {0: 789}}
1847
worker._remote_agents = {"engine1": {0: "agent1"}}
1848
+ # _cleanup_remote_engine (called by shutdown) also clears these:
1849
+ worker.kv_caches_base_addr["engine1"] = {0: [0xABC]}
1850
+ worker.dst_num_blocks["engine1"] = 50
1851
+ worker.tp_mappings["engine1"] = MagicMock()
1852
+ worker._engine_last_active["engine1"] = time.perf_counter()
1853
worker._registered_descs = ["desc1", "desc2"]
1854
1855
mock_listener.is_alive.return_value = False
@@ -1874,6 +1880,118 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):
1880
mock_dereg.assert_any_call("desc2")
1881
1882
1883
+# ── TTL-based remote engine eviction tests ──────────────────────────
1884
+
1885
+
1886
+def _setup_worker_with_remote_engine(
1887
+ engine_ttl: float = 10.0,
1888
+) -> tuple[Any, str]:
1889
+ """Create a worker with one remote engine registered."""
1890
+ vllm_config = create_vllm_config(
1891
+ kv_connector_extra_config={"engine_ttl": engine_ttl},
1892
+ )
1893
+ worker = NixlConnectorWorker(
1894
+ vllm_config,
1895
+ vllm_config.kv_transfer_config.engine_id,
1896
+ make_kv_cache_config(block_size=16),
1897
+ )
1898
+
1899
+ engine_id = "remote-engine-1"
1900
+ worker._remote_agents[engine_id] = {0: "agent_0", 1: "agent_1"}
1901
+ worker.dst_xfer_side_handles[engine_id] = {0: 100, 1: 200}
1902
+ worker.kv_caches_base_addr[engine_id] = {0: [0xABC]}
1903
+ worker.dst_num_blocks[engine_id] = 50
1904
+ worker.tp_mappings[engine_id] = MagicMock()
1905
+ worker._engine_last_active[engine_id] = time.perf_counter()
1906
+
1907
+ worker.transfer_topo = MagicMock()
1908
+
1909
+ return worker, engine_id
1910
+
1911
+
1912
+@patch(
1913
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1914
+ FakeNixlWrapper,
1915
+)
1916
+def test_engine_ttl_eviction(default_vllm_config, dist_init):
1917
+ """Stale engines are evicted when TTL expires."""
1918
+ worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0)
1919
+ nixl_wrapper = worker.nixl_wrapper
1920
+
1921
+ with (
1922
+ patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel,
1923
+ patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem,
1924
+ ):
1925
+ # Make the engine stale.
1926
+ worker._engine_last_active[engine_id] = time.perf_counter() - 20.0
1927
+
1928
+ worker._evict_stale_engines()
1929
+
1930
+ assert engine_id not in worker._remote_agents
1931
+ assert engine_id not in worker.dst_xfer_side_handles
1932
+ assert engine_id not in worker.kv_caches_base_addr
1933
+ assert engine_id not in worker.dst_num_blocks
1934
+ assert engine_id not in worker.tp_mappings
1935
+ assert engine_id not in worker._engine_last_active
1936
+ worker.transfer_topo.unregister_remote_engine.assert_called_with(engine_id)
1937
+
1938
+ assert mock_rel.call_count == 2
1939
+ mock_rel.assert_any_call(100)
1940
+ mock_rel.assert_any_call(200)
1941
+
1942
+ assert mock_rem.call_count == 2
1943
+ mock_rem.assert_any_call("agent_0")
1944
+ mock_rem.assert_any_call("agent_1")
1945
+
1946
+
1947
+@patch(
1948
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1949
+ FakeNixlWrapper,
1950
+)
1951
+def test_engine_ttl_disabled(default_vllm_config, dist_init):
1952
+ """Eviction is disabled when engine_ttl <= 0."""
1953
+ worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=0.0)
1954
+
1955
+ # Make the engine stale.
1956
+ worker._engine_last_active[engine_id] = time.perf_counter() - 9999.0
1957
+
1958
+ worker._evict_stale_engines()
1959
+
1960
+ # Nothing should be evicted.
1961
+ assert engine_id in worker._remote_agents
1962
+ assert engine_id in worker.dst_xfer_side_handles
1963
+
1964
+
1965
+def test_transfer_topology_unregister():
1966
+ """TransferTopology.unregister_remote_engine removes the engine."""
1967
+ topo = TransferTopology(
1968
+ tp_rank=0,
1969
+ tp_size=1,
1970
+ block_size=16,
1971
+ engine_id="local",
1972
+ is_mla=False,
1973
+ is_mamba=False,
1974
+ total_num_kv_heads=4,
1975
+ attn_backends=[FlashAttentionBackend],
1976
+ )
1977
+
1978
+ info = EngineTransferInfo(
1979
+ remote_tp_size=1,
1980
+ remote_block_size=16,
1981
+ remote_block_len=64,
1982
+ remote_physical_blocks_per_logical=1,
1983
+ )
1984
+ topo.register_remote_engine("remote-1", info)
1985
+ assert topo.get_engine_info("remote-1") is info
1986
+
1987
+ topo.unregister_remote_engine("remote-1")
1988
+ with pytest.raises(KeyError):
1989
+ topo.get_engine_info("remote-1")
1990
+
1991
+ # Idempotent: no error on double-unregister
1992
+ topo.unregister_remote_engine("remote-1")
1993
+
1994
+
1995
@patch(
1996
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",
1997
FakeNixlWrapper,
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+1
@@ -162,6 +162,7 @@ def test_read_blocks_for_req_expands_remote_ids(
162
163
worker = object.__new__(NixlConnectorWorker)
164
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
165
+ worker._engine_last_active = {}
166
167
has_mamba = any(t is MambaSpec for t in resolved_types)
168
has_swa = any(t is SlidingWindowSpec for t in resolved_types)
vllm/distributed/kv_transfer/kv_connector/utils.py
+5
@@ -481,6 +481,11 @@ class TransferTopology:
481
) -> EngineTransferInfo:
482
return self._engines[(remote_engine_id, remote_pp_rank)]
483
484
+ def unregister_remote_engine(self, remote_engine_id: EngineId) -> None:
485
+ # Remove all pp_rank entries for the remote engine.
486
+ for key in [k for k in self._engines if k[0] == remote_engine_id]:
487
+ del self._engines[key]
488
+
489
# ============================================================
490
# Layout properties
491
# ============================================================
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+68
-8
@@ -408,6 +408,12 @@ class NixlConnectorWorker:
408
# Protects _handshake_futures and _remote_agents.
409
self._handshake_lock = threading.RLock()
410
411
+ # TTL-based eviction of stale remote engine state.
412
+ self._engine_last_active: dict[EngineId, float] = {}
413
+ self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config(
414
+ "engine_ttl", 3600.0
415
+ )
416
+
417
self.block_size = vllm_config.cache_config.block_size
418
self.model_config = vllm_config.model_config
419
@@ -711,6 +717,7 @@ class NixlConnectorWorker:
717
returned future.
718
Failures to handshake are logged and the request is marked as failed.
719
"""
720
+ self._evict_stale_engines()
721
with self._handshake_lock:
722
if engine_id in self._remote_agents:
723
return None
@@ -731,6 +738,7 @@ class NixlConnectorWorker:
738
del self._handshake_futures[eid]
739
try:
740
self._remote_agents[eid] = f.result()
741
+ self._engine_last_active[eid] = time.perf_counter()
742
except Exception as e:
743
self._log_failure(
744
failure_type="handshake_setup_failed",
@@ -2028,6 +2036,9 @@ class NixlConnectorWorker:
2036
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2037
assert meta.remote is not None and self.transfer_topo is not None
2038
engine_id = meta.remote.engine_id
2039
+ # Update last activity from this remote. Mind that cleanup is done on main
2040
+ # thread (this one), so we don't race on this structure.
2041
+ self._engine_last_active[engine_id] = time.perf_counter()
2042
plan = self.tp_mappings[engine_id]
2043
remote_info = self.transfer_topo.get_engine_info(engine_id)
2044
tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
@@ -2473,6 +2484,61 @@ class NixlConnectorWorker:
2484
break
2485
return result
2486
2487
+ def _evict_stale_engines(self) -> None:
2488
+ """Scan for and evict remote engines that have exceeded their TTL.
2489
+
2490
+ Called from the main thread in when a new remote engine appears.
2491
+ We can only go OOM as we discover and register a new remote, therefore we make
2492
+ sure we clean up stale engine data structures before then. This invariant
2493
+ prevents us from using background threads, though memory usage is not guaranteed
2494
+ to be "optimal" until a new handshake is performed.
2495
+
2496
+ Engines with active transfers or pending handshakes cannot be stale:
2497
+ - Active transfers touch _engine_last_active in start_load_kv.
2498
+ - Pending handshakes don't have an _engine_last_active entry yet
2499
+ """
2500
+ # NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
2501
+ # of remote engines is registered all at once (adding a background cleanup
2502
+ # thread wouldnt help either).
2503
+ # If that scenario is plausible, we can follow up with an LRU eviction policy.
2504
+ if self._engine_ttl <= 0:
2505
+ return
2506
+
2507
+ now = time.perf_counter()
2508
+ for eid, last_active in list(self._engine_last_active.items()):
2509
+ if now - last_active > self._engine_ttl:
2510
+ self._cleanup_remote_engine(eid)
2511
+
2512
+ def _cleanup_remote_engine(
2513
+ self, engine_id: EngineId, *, log_eviction: bool = True
2514
+ ) -> None:
2515
+ """Remove all state for a single remote engine.
2516
+
2517
+ Releases NIXL resources (dlist handles, remote agents) and clears
2518
+ all per-engine data structures. Used by both TTL eviction and
2519
+ shutdown.
2520
+ """
2521
+ assert engine_id in self._remote_agents
2522
+
2523
+ for handle in self.dst_xfer_side_handles.pop(engine_id).values():
2524
+ self.nixl_wrapper.release_dlist_handle(handle)
2525
+ for agent_name in self._remote_agents.pop(engine_id).values():
2526
+ self.nixl_wrapper.remove_remote_agent(agent_name)
2527
+
2528
+ del self.kv_caches_base_addr[engine_id]
2529
+ del self.dst_num_blocks[engine_id]
2530
+ del self.tp_mappings[engine_id]
2531
+ if self.transfer_topo is not None:
2532
+ self.transfer_topo.unregister_remote_engine(engine_id)
2533
+
2534
+ last_active = self._engine_last_active.pop(engine_id)
2535
+ if log_eviction:
2536
+ logger.info(
2537
+ "Evicted stale remote engine %s (inactive for %.1fs).",
2538
+ engine_id,
2539
+ time.perf_counter() - last_active,
2540
+ )
2541
+
2542
def __del__(self):
2543
self.shutdown()
2544
@@ -2493,14 +2559,8 @@ class NixlConnectorWorker:
2559
for handle in handles:
2560
self.nixl_wrapper.release_dlist_handle(handle)
2561
self.src_xfer_handles_by_tp_ratio.clear()
2496
- for dst_xfer_side_handles in self.dst_xfer_side_handles.values():
2497
- for dst_xfer_side_handle in dst_xfer_side_handles.values():
2498
- self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle)
2499
- self.dst_xfer_side_handles.clear()
2500
- for remote_agents in self._remote_agents.values():
2501
- for agent_name in remote_agents.values():
2502
- self.nixl_wrapper.remove_remote_agent(agent_name)
2503
- self._remote_agents.clear()
2562
+ for engine_id in list(self._remote_agents):
2563
+ self._cleanup_remote_engine(engine_id, log_eviction=False)
2564
for desc in self._registered_descs:
2565
self.nixl_wrapper.deregister_memory(desc)
2566
self._registered_descs.clear()