From 1ba21100aa1d2561d6e1008bd6e3935e4fbe3d71 Mon Sep 17 00:00:00 2001 From: kliuae Date: Thu, 3 Sep 2026 11:10:20 +0000 Subject: [PATCH 1/2] vllm enable glm5.2 lmcache Signed-off-by: kliuae --- atom/plugin/vllm/lmcache_connector_patch.py | 366 ++++++++++++ atom/plugin/vllm/register.py | 6 + tests/plugin/test_lmcache_connector_guard.py | 566 +++++++++++++++++++ tests/test_lmcache_vllm_dsa_layout.py | 236 ++++++++ 4 files changed, 1174 insertions(+) create mode 100644 atom/plugin/vllm/lmcache_connector_patch.py create mode 100644 tests/plugin/test_lmcache_connector_guard.py create mode 100644 tests/test_lmcache_vllm_dsa_layout.py diff --git a/atom/plugin/vllm/lmcache_connector_patch.py b/atom/plugin/vllm/lmcache_connector_patch.py new file mode 100644 index 0000000000..e15a8b03c0 --- /dev/null +++ b/atom/plugin/vllm/lmcache_connector_patch.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Make LMCache usable with ATOM's DSA KV layouts under the vLLM plugin. + +GLM-5.2 registers a narrow indexer key cache next to each wide MLA cache and +vLLM folds them into one ``UniformTypeKVCacheSpecs`` group -- 99 tensors in two +page shapes for GLM-5.2-FP8. Three upstream behaviours get in the way: + +1. LMCache's default GPU connector sizes its pointer array from + ``model_config.get_num_layers()`` (78) and dies on the first store with + ``could not broadcast input array from shape (99,) into shape (78,)``. Its + V3 connector partitions by page shape instead, but is off by default. + +2. V3 builds its layer-group map inside the first transfer, while + ``LMCacheEngine.store`` sizes the memory object before that from + ``metadata.get_shapes()`` -- so the first store asks for a group the memory + object does not have (``IndexError`` at ``gpu_connectors.py:612``). + +3. LMCache builds no engine for the scheduler role, but its aborted-request + cleanup asserts one (``vllm_v1_adapter.py:1779``) from inside the + scheduler's ``_free_request``, so one client disconnect kills the + EngineCore. Not DSA-specific. + +``KVConnectorFactory.create_connector`` is the hook for all three: one choke +point for both roles, before LMCache's config singleton exists, and it receives +the resolved ``KVCacheConfig``. +""" + +import logging +import os +from typing import NamedTuple + +logger = logging.getLogger("atom") + +_LMCACHE_V3_ENV = "LMCACHE_USE_GPU_CONNECTOR_V3" +# Reference lmcache.v1.config_base._to_bool +_LMCACHE_V3_IS_SET = frozenset({"true", "1"}) + + +def _lmcache_config(): + try: + from lmcache.v1.config import LMCacheEngineConfig + + return LMCacheEngineConfig.from_env() + except Exception: + logger.debug("ATOM vLLM plugin: could not pre-read LMCache config", exc_info=True) + return None + + +def _is_lmcache_connector(kv_transfer_config) -> bool: + if kv_transfer_config is None: + return False + + names = [str(getattr(kv_transfer_config, "kv_connector", "") or "")] + extra = getattr(kv_transfer_config, "kv_connector_extra_config", None) or {} + names += [ + str(sub.get("kv_connector", "")) + for sub in extra.get("connectors", None) or () + if isinstance(sub, dict) + ] + return any("lmcache" in name.lower() for name in names) + + +def _layer_specs(kv_cache_config): + # In DSA, UniformTypeKVCacheSpecs can hold layers with different hidden sizes, + # so we need to walk over all layers to collect all specs. + specs = [] + for group in getattr(kv_cache_config, "kv_cache_groups", None) or (): + spec = getattr(group, "kv_cache_spec", None) + if spec is None: + continue + nested = getattr(spec, "kv_cache_specs", None) + specs.extend(nested.values() if nested else [spec]) + return specs + + +class _PageGeometry(NamedTuple): + # Reference lmcache.v1.kv_layer_groups.LayerGroupIdentity + + kv_size: int + num_kv_heads: int + head_size: int + block_size: "int | None" + dtype: str + + +def _page_geometries(kv_cache_config) -> "set[_PageGeometry]": + try: + from vllm.v1.kv_cache_interface import MLAAttentionSpec + except ImportError: + MLAAttentionSpec = () + + geometries = set() + for spec in _layer_specs(kv_cache_config): + num_kv_heads = getattr(spec, "num_kv_heads", None) + head_size = getattr(spec, "head_size", None) + if num_kv_heads is None or head_size is None: + continue + mla = isinstance(spec, MLAAttentionSpec) or "mla" in type(spec).__name__.lower() + geometries.add( + _PageGeometry( + kv_size=1 if mla else 2, + num_kv_heads=num_kv_heads, + head_size=head_size, + block_size=getattr(spec, "storage_block_size", None) + or getattr(spec, "block_size", None), + dtype=str(getattr(spec, "dtype", None)), + ) + ) + return geometries + + +def _v3_cannot_express(geometries): + """Why V3 would mis-address this layout, or None if it can carry it. + + V3 partitions by page shape but not by ``GPUKVFormat`` (probed once from + ``kv_caches[0]``) nor by the scalar ``block_size`` it hands to the transfer + kernel, so caches disagreeing on either are read with the wrong layout. + """ + for field, label in (("kv_size", "kv_size"), ("block_size", "physical block size")): + values = {getattr(g, field) for g in geometries} - {None} + if len(values) > 1: + return f"its caches disagree on {label} ({sorted(values)})" + return None + + +def _reject_unsafe_chunking(arch, block_size) -> None: + # Sparse MLA can preshuffle the indexer cache, but LMCache has settings that + # may chunk the indexer cache independently, oblivious to the preshuffle. + # So we need to reject chunk settings that would slice a preshuffled page + config = _lmcache_config() + # block_size=1 is not preshuffled + if config is None or not block_size or block_size <= 1: + return + + chunk_size = getattr(config, "chunk_size", 0) + if chunk_size and chunk_size % block_size: + problem = ( + f"LMCACHE_CHUNK_SIZE={chunk_size} is not a multiple of the " + f"{block_size}-token KV page" + ) + elif getattr(config, "save_unfull_chunk", False): + problem = "LMCACHE_SAVE_UNFULL_CHUNK stores a partial trailing page" + else: + return + + msg = ( + f"{problem}, which corrupts {arch}'s preshuffled DSA indexer cache while " + "leaving the MLA cache intact -- the model keeps running and silently " + "attends to the wrong tokens." + ) + logger.error(msg) + raise ValueError(msg) + + +def _model_declares_indexer_cache(vllm_config) -> bool: + model_config = getattr(vllm_config, "model_config", None) + hf_config = getattr(model_config, "hf_text_config", None) or getattr( + model_config, "hf_config", None + ) + return getattr(hf_config, "index_head_dim", None) is not None + + +def enforce_lmcache_gpu_connector(vllm_config, kv_cache_config=None) -> bool: + # Select V3 when the KV layout requires it + if not _is_lmcache_connector(getattr(vllm_config, "kv_transfer_config", None)): + return False + + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is not None and not getattr( + cache_config, "enable_prefix_caching", False + ): + logger.warning( + "An LMCache KV connector is configured but prefix caching is off. " + "The CPU/NVMe tier can still be filled, but vLLM will not reuse any " + "HBM prefix. Pass --enable-prefix-caching." + ) + + model_config = getattr(vllm_config, "model_config", None) + arch = (getattr(model_config, "architectures", None) or [""])[0] + + geometries = _page_geometries(kv_cache_config) + block_size = None + if geometries: + unsupported = _v3_cannot_express(geometries) + if unsupported is not None: + msg = ( + f"{arch} cannot be served with an LMCache KV connector: " + f"{unsupported}. Drop --kv-transfer-config to serve without KV " + "offload." + ) + logger.error(msg) + raise ValueError(msg) + needs_v3 = len(geometries) > 1 + reason = f"{len(geometries)} distinct KV page geometries" + block_size = next(iter(geometries)).block_size + else: + needs_v3 = _model_declares_indexer_cache(vllm_config) + reason = "a DSA indexer KV cache" + + if not needs_v3: + return False + + _reject_unsafe_chunking(arch, block_size) + + configured = os.environ.get(_LMCACHE_V3_ENV) + if configured is None: + os.environ[_LMCACHE_V3_ENV] = "True" + logger.info( + "ATOM plugin: %s has %s, which LMCache's default GPU connector " + "cannot express. Setting %s=True.", + arch, + reason, + _LMCACHE_V3_ENV, + ) + elif configured.strip().lower() not in _LMCACHE_V3_IS_SET: + # An explicit choice is never silently overridden, only refused. + msg = ( + f"{_LMCACHE_V3_ENV}={configured!r} is incompatible with {arch}, which " + f"has {reason}. LMCache reads only 'true' and '1' as enabled, so this " + "selects its default GPU connector, which fails on the first store. " + f"Set {_LMCACHE_V3_ENV}=True, or drop --kv-transfer-config." + ) + logger.error(msg) + raise ValueError(msg) + return True + + +def _prebuild_gpu_connector_layer_groups(connector, kv_caches) -> None: + # Run V3's group discovery at this point, rather than inside the first transfer + impl = getattr(connector, "_lmcache_engine", None) + gpu_connector = getattr( + getattr(impl, "lmcache_engine", None), "gpu_connector", None + ) + # Only V3 exposes _initialize_kv_cache_pointers, so use this for V3 lookup + initialize_pointers = getattr(gpu_connector, "_initialize_kv_cache_pointers", None) + if initialize_pointers is None or getattr(gpu_connector, "init", False): + return + + gpu_connector.initialize_kvcaches_ptr(kvcaches=list(kv_caches.values())) + initialize_pointers() + + manager = getattr( + getattr(gpu_connector, "metadata", None), "kv_layer_groups_manager", None + ) + groups = getattr(manager, "kv_layer_groups", None) or () + logger.info( + "ATOM plugin: primed LMCache KV layer groups from %d registered caches -> %s", + len(kv_caches), + [(g.num_layers, g.shape_desc.hs) for g in groups], + ) + + +def _mark_lmcache(connector, flag) -> bool: + if "lmcache" not in type(connector).__name__.lower(): + return False + if getattr(connector, flag, False): + return False + setattr(connector, flag, True) + return True + + +def _wrap_register_kv_caches(connector, required: bool = False) -> None: + # Prebuild the KV tensors once vLLM register_kv_caches + register_kv_caches = getattr(connector, "register_kv_caches", None) + if register_kv_caches is None or not _mark_lmcache( + connector, "_atom_lmcache_group_prebuilt" + ): + return + + def register_and_prime(kv_caches): + register_kv_caches(kv_caches) + try: + _prebuild_gpu_connector_layer_groups(connector, kv_caches) + except Exception: + if required: + raise + logger.warning( + "ATOM plugin: could not pre-build LMCache's KV layer groups; " + "falling back to its lazy discovery.", + exc_info=True, + ) + + connector.register_kv_caches = register_and_prime + + +class _NullLMCacheEngine: + """Stands in for the engine a scheduler-role connector does not have. + + ``storage_manager`` is what LMCache's aborted-request cleanup checks before + cancelling; ``is_healthy`` only because the stub is briefly visible to + ``LMCacheManager.is_healthy()``, which backs a Prometheus gauge. + """ + + storage_manager = None + + @staticmethod + def is_healthy() -> bool: + return True + + +def _wrap_request_finished(connector) -> None: + # Let an aborted request reach LMCache's cleanup without an engine + impl = getattr(connector, "_lmcache_engine", None) + request_finished = getattr(connector, "request_finished", None) + if impl is None or request_finished is None: + return + if not _mark_lmcache(connector, "_atom_lmcache_abort_guard"): + return + + def request_finished_without_engine(request, block_ids): + manager = getattr(impl, "_manager", None) + aborted = getattr(getattr(request, "status", None), "name", "") == ( + "FINISHED_ABORTED" + ) + if ( + not aborted + or manager is None + or getattr(manager, "_lmcache_engine", "missing") is not None + ): + return request_finished(request, block_ids) + + manager._lmcache_engine = _NullLMCacheEngine() + try: + return request_finished(request, block_ids) + finally: + manager._lmcache_engine = None + + connector.request_finished = request_finished_without_engine + + +def _wrap_lmcache_connectors(connector, required: bool) -> None: + + # vLLM's MultiConnector builds its children directly instead of coming + # back through the factory, so recursing here is the only way a nested + # LMCache connector is reached. + _wrap_register_kv_caches(connector, required) + _wrap_request_finished(connector) + for child in getattr(connector, "_connectors", None) or (): + _wrap_lmcache_connectors(child, required) + + +def apply_vllm_lmcache_connector_patch() -> None: + import functools + + from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory + + original = KVConnectorFactory.create_connector + if getattr(original, "_atom_lmcache_connector_patched", False): + return + + @functools.wraps(original.__func__) + def create_connector(cls, config, role, kv_cache_config=None, *args, **kwargs): + required = enforce_lmcache_gpu_connector(config, kv_cache_config) + extra = () if kv_cache_config is None else (kv_cache_config,) + connector = original.__func__(cls, config, role, *extra, *args, **kwargs) + _wrap_lmcache_connectors(connector, required) + return connector + + create_connector._atom_lmcache_connector_patched = True + KVConnectorFactory.create_connector = classmethod(create_connector) + logger.info( + "ATOM plugin: patched vLLM KVConnectorFactory to select and prime " + "LMCache's multi-geometry GPU connector for DSA KV layouts" + ) diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 25c9f9a317..3d816c072e 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -266,6 +266,12 @@ def register_model() -> None: apply_vllm_v4_block_reuse_patch() + from atom.plugin.vllm.lmcache_connector_patch import ( + apply_vllm_lmcache_connector_patch, + ) + + apply_vllm_lmcache_connector_patch() + from atom.plugin.vllm.gdn_backend import register_gdn_attention_backend register_gdn_attention_backend() diff --git a/tests/plugin/test_lmcache_connector_guard.py b/tests/plugin/test_lmcache_connector_guard.py new file mode 100644 index 0000000000..4a679cacfb --- /dev/null +++ b/tests/plugin/test_lmcache_connector_guard.py @@ -0,0 +1,566 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Pins the connector-selection logic in lmcache_connector_patch.""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest + +from atom.plugin.vllm.lmcache_connector_patch import ( + _LMCACHE_V3_ENV, + enforce_lmcache_gpu_connector, +) + +DSA_ARCH = "GlmMoeDsaForCausalLM" + + +@pytest.fixture(autouse=True) +def _clean_v3_env(): + """Keep the selection variable out of the rest of the session. + + ``monkeypatch.delenv(..., raising=False)`` records no undo when the variable + is absent, so the value the code under test then *sets* survives teardown + and leaks into every later test. Delete on both sides instead. + """ + os.environ.pop(_LMCACHE_V3_ENV, None) + yield + os.environ.pop(_LMCACHE_V3_ENV, None) + + +class MlaSpec(SimpleNamespace): + """A single-vector cache: kv_size 1, the shape GLM-5.2 uses for both families.""" + + +class KVSpec(SimpleNamespace): + """A separate-K-and-V cache: kv_size 2.""" + + +def _spec( + num_kv_heads=1, + head_size=576, + block_size=64, + dtype="torch.uint8", + cls=MlaSpec, + **extra, +): + return cls( + num_kv_heads=num_kv_heads, + head_size=head_size, + block_size=block_size, + dtype=dtype, + **extra, + ) + + +def _kv_cache_config(*specs): + """One vLLM KV cache group whose members may differ in hidden size. + + This is the shape vLLM produces for GLM-5.2: MLA and indexer layers land in + one ``UniformTypeKVCacheSpecs`` group, so they share block ids. + """ + group = SimpleNamespace( + kv_cache_spec=SimpleNamespace( + kv_cache_specs={f"layer.{i}": spec for i, spec in enumerate(specs)} + ) + ) + return SimpleNamespace(kv_cache_groups=[group]) + + +def _vllm_config( + *, + kv_connector: str | None = "LMCacheConnectorV1", + index_head_dim: int | None = 128, + enable_prefix_caching: bool = True, + nested_connectors: list[dict] | None = None, + architectures: list[str] | None = None, +): + hf_config = SimpleNamespace() + if index_head_dim is not None: + hf_config.index_head_dim = index_head_dim + + kv_transfer_config = None + if kv_connector is not None: + extra = {} + if nested_connectors is not None: + extra["connectors"] = nested_connectors + kv_transfer_config = SimpleNamespace( + kv_connector=kv_connector, + kv_connector_extra_config=extra, + ) + + return SimpleNamespace( + kv_transfer_config=kv_transfer_config, + model_config=SimpleNamespace( + hf_text_config=hf_config, + hf_config=hf_config, + architectures=architectures or [DSA_ARCH], + ), + cache_config=SimpleNamespace(enable_prefix_caching=enable_prefix_caching), + ) + + +def test_dsa_model_with_lmcache_selects_v3(): + enforce_lmcache_gpu_connector(_vllm_config()) + + assert os.environ[_LMCACHE_V3_ENV] == "True" + + +def test_lmcache_nested_in_multi_connector_is_detected(): + enforce_lmcache_gpu_connector( + _vllm_config( + kv_connector="multi", + nested_connectors=[ + {"kv_connector": "mooncake", "kv_role": "kv_producer"}, + {"kv_connector": "LMCacheConnectorV1", "kv_role": "kv_both"}, + ], + ) + ) + + assert os.environ[_LMCACHE_V3_ENV] == "True" + + +@pytest.mark.parametrize("value", ["0", "false", "False", "no", "off", ""]) +def test_explicitly_disabled_v3_fails_fast(monkeypatch, value): + monkeypatch.setenv(_LMCACHE_V3_ENV, value) + + with pytest.raises(ValueError, match=_LMCACHE_V3_ENV): + enforce_lmcache_gpu_connector(_vllm_config()) + + +def test_explicitly_enabled_v3_is_left_alone(monkeypatch): + monkeypatch.setenv(_LMCACHE_V3_ENV, "true") + + enforce_lmcache_gpu_connector(_vllm_config()) + + assert os.environ[_LMCACHE_V3_ENV] == "true" + + +def test_dense_model_is_untouched(): + enforce_lmcache_gpu_connector( + _vllm_config(index_head_dim=None, architectures=["DeepseekV3ForCausalLM"]) + ) + + assert _LMCACHE_V3_ENV not in os.environ + + +def test_dsa_model_without_lmcache_is_untouched(): + enforce_lmcache_gpu_connector(_vllm_config(kv_connector=None)) + enforce_lmcache_gpu_connector(_vllm_config(kv_connector="MooncakeConnector")) + + assert _LMCACHE_V3_ENV not in os.environ + + +def test_prefix_caching_off_warns_but_does_not_block(caplog): + with caplog.at_level("WARNING", logger="atom"): + enforce_lmcache_gpu_connector(_vllm_config(enable_prefix_caching=False)) + + assert any("prefix caching is off" in record.message for record in caplog.records) + + assert os.environ[_LMCACHE_V3_ENV] == "True" + + +def test_real_kv_cache_config_drives_the_decision(): + enforce_lmcache_gpu_connector( + _vllm_config(index_head_dim=None, architectures=["SomeOtherModel"]), + _kv_cache_config(_spec(head_size=576), _spec(head_size=132)), + ) + + assert os.environ[_LMCACHE_V3_ENV] == "True" + + +def test_single_geometry_config_beats_the_architecture_heuristic(): + enforce_lmcache_gpu_connector( + _vllm_config(), + _kv_cache_config(_spec(head_size=576), _spec(head_size=576)), + ) + + assert _LMCACHE_V3_ENV not in os.environ + + +def test_factory_patch_runs_the_guard_and_is_idempotent(monkeypatch): + factory = pytest.importorskip( + "vllm.distributed.kv_transfer.kv_connector.factory" + ).KVConnectorFactory + from atom.plugin.vllm.lmcache_connector_patch import ( + apply_vllm_lmcache_connector_patch, + ) + + seen = {} + monkeypatch.setattr( + factory, + "create_connector", + classmethod( + lambda cls, config, role, kv_cache_config=None: seen.setdefault( + "role", role + ) + ), + ) + + apply_vllm_lmcache_connector_patch() + # Bound classmethods are rebuilt on every attribute access, so compare the + # underlying function rather than the binding. + first = factory.create_connector.__func__ + apply_vllm_lmcache_connector_patch() + assert factory.create_connector.__func__ is first, "patch must not stack" + + factory.create_connector( + _vllm_config(), + "worker", + _kv_cache_config(_spec(head_size=576), _spec(head_size=132)), + ) + + assert os.environ[_LMCACHE_V3_ENV] == "True" + assert seen["role"] == "worker", "the original factory must still be called" + + +def test_register_kv_caches_primes_the_layer_groups(): + # V3 builds its layer-group map inside the first transfer, but the memory + # object is allocated before that and sized from the map -- so the map has + # to exist by the time registration returns. + from atom.plugin.vllm.lmcache_connector_patch import _wrap_register_kv_caches + + calls = [] + + class _FakeV3GPUConnector: + init = False + + def initialize_kvcaches_ptr(self, **kwargs): + calls.append(("initialize_kvcaches_ptr", len(kwargs["kvcaches"]))) + + def _initialize_kv_cache_pointers(self): + calls.append(("_initialize_kv_cache_pointers", None)) + self.init = True + + class _FakeLMCacheConnectorV1: + def __init__(self): + self._lmcache_engine = SimpleNamespace( + lmcache_engine=SimpleNamespace(gpu_connector=_FakeV3GPUConnector()) + ) + + def register_kv_caches(self, kv_caches): + calls.append(("register_kv_caches", len(kv_caches))) + + connector = _FakeLMCacheConnectorV1() + _wrap_register_kv_caches(connector) + _wrap_register_kv_caches(connector) # idempotent + + connector.register_kv_caches({"mla.0": object(), "indexer.0": object()}) + + assert calls == [ + ("register_kv_caches", 2), + ("initialize_kvcaches_ptr", 2), + ("_initialize_kv_cache_pointers", None), + ] + + +def test_priming_failure_does_not_take_the_server_down(caplog): + # LMCache internals are not a stable API; if they move, fall back to its own + # lazy discovery rather than killing startup. + from atom.plugin.vllm.lmcache_connector_patch import _wrap_register_kv_caches + + class _ExplodingGPUConnector: + init = False + + def initialize_kvcaches_ptr(self, **kwargs): + raise AttributeError("LMCache moved this") + + def _initialize_kv_cache_pointers(self): # pragma: no cover - unreachable + raise AssertionError("should not be reached") + + class _FakeLMCacheConnectorV1: + def __init__(self): + self._lmcache_engine = SimpleNamespace( + lmcache_engine=SimpleNamespace(gpu_connector=_ExplodingGPUConnector()) + ) + self.registered = None + + def register_kv_caches(self, kv_caches): + self.registered = kv_caches + + connector = _FakeLMCacheConnectorV1() + _wrap_register_kv_caches(connector) + + with caplog.at_level("WARNING", logger="atom"): + connector.register_kv_caches({"mla.0": object()}) + + assert connector.registered is not None, "registration itself must still happen" + assert any("lazy discovery" in record.message for record in caplog.records) + + +def test_non_lmcache_connector_is_not_wrapped(): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_register_kv_caches + + class MooncakeConnector: + def register_kv_caches(self, kv_caches): + pass + + connector = MooncakeConnector() + original = connector.register_kv_caches + _wrap_register_kv_caches(connector) + + assert connector.register_kv_caches == original + + +class _AbortedRequest: + status = SimpleNamespace(name="FINISHED_ABORTED") + request_id = "req-1" + + +class _LiveRequest: + status = SimpleNamespace(name="FINISHED_STOPPED") + request_id = "req-2" + + +def _scheduler_role_connector(seen): + """An LMCache connector as the scheduler role builds it: no engine.""" + from vllm.v1.request import RequestStatus # noqa: F401 (import sanity only) + + class _FakeLMCacheConnectorV1: + def __init__(self): + self._lmcache_engine = SimpleNamespace( + _manager=SimpleNamespace(_lmcache_engine=None) + ) + + def request_finished(self, request, block_ids): + # Mirrors LMCache: the aborted branch dereferences the engine. + if request.status.name == "FINISHED_ABORTED": + engine = self._lmcache_engine._manager._lmcache_engine + assert engine is not None + seen.append(("cancel_checked", engine.storage_manager)) + return False, None + + return _FakeLMCacheConnectorV1() + + +def test_aborted_request_does_not_kill_a_scheduler_role_connector(): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_request_finished + + seen = [] + connector = _scheduler_role_connector(seen) + _wrap_request_finished(connector) + _wrap_request_finished(connector) # idempotent + + assert connector.request_finished(_AbortedRequest(), [1, 2]) == (False, None) + assert seen == [("cancel_checked", None)], "the cancel must be a no-op" + assert ( + connector._lmcache_engine._manager._lmcache_engine is None + ), "the null engine must not outlive the call" + + +def test_non_aborted_request_takes_the_untouched_path(): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_request_finished + + seen = [] + connector = _scheduler_role_connector(seen) + _wrap_request_finished(connector) + + assert connector.request_finished(_LiveRequest(), [1]) == (False, None) + assert seen == [] + + +# --- truthiness must match LMCache's, not merely be reasonable --- + + +@pytest.mark.parametrize("value", ["yes", "y", "on", "enabled", "TRUE!", "2"]) +def test_values_lmcache_reads_as_off_are_treated_as_off(monkeypatch, value): + monkeypatch.setenv(_LMCACHE_V3_ENV, value) + + with pytest.raises(ValueError, match=_LMCACHE_V3_ENV): + enforce_lmcache_gpu_connector(_vllm_config()) + + +@pytest.mark.parametrize("value", ["true", "True", " TRUE ", "1"]) +def test_values_lmcache_reads_as_on_are_left_alone(monkeypatch, value): + monkeypatch.setenv(_LMCACHE_V3_ENV, value) + + assert enforce_lmcache_gpu_connector(_vllm_config()) is True + assert os.environ[_LMCACHE_V3_ENV] == value + + +# --- layouts V3 cannot address are refused, not enabled into --- + + +def test_mixed_kv_size_is_refused_rather_than_given_v3(): + with pytest.raises(ValueError, match="kv_size"): + enforce_lmcache_gpu_connector( + _vllm_config(architectures=["MiniMaxM3SparseForCausalLM"]), + _kv_cache_config( + _spec(head_size=132, cls=MlaSpec), + _spec(head_size=128, num_kv_heads=8, cls=KVSpec), + ), + ) + + assert _LMCACHE_V3_ENV not in os.environ + + +def test_mixed_physical_block_size_is_refused(): + with pytest.raises(ValueError, match="physical block size"): + enforce_lmcache_gpu_connector( + _vllm_config(architectures=["DeepseekV4ForCausalLM"]), + _kv_cache_config( + _spec(head_size=576, storage_block_size=64), + _spec(head_size=132, storage_block_size=16), + ), + ) + + assert _LMCACHE_V3_ENV not in os.environ + + +def test_physical_block_size_beats_the_logical_one(): + # Same logical block_size, different compress ratios: keying on + # spec.block_size alone would miss this. + with pytest.raises(ValueError, match="physical block size"): + enforce_lmcache_gpu_connector( + _vllm_config(), + _kv_cache_config( + _spec(head_size=576, block_size=64, storage_block_size=64), + _spec(head_size=132, block_size=64, storage_block_size=32), + ), + ) + + +def test_glm52_shape_is_still_accepted(): + assert ( + enforce_lmcache_gpu_connector( + _vllm_config(), + _kv_cache_config(_spec(head_size=576), _spec(head_size=132)), + ) + is True + ) + assert os.environ[_LMCACHE_V3_ENV] == "True" + + +# --- chunk settings that would slice a preshuffled page --- + + +def _dsa_config(): + return _vllm_config(), _kv_cache_config(_spec(head_size=576), _spec(head_size=132)) + + +def test_chunk_size_not_a_multiple_of_the_page_is_refused(monkeypatch): + monkeypatch.setenv("LMCACHE_CHUNK_SIZE", "100") + cfg, kv = _dsa_config() + + with pytest.raises(ValueError, match="LMCACHE_CHUNK_SIZE"): + enforce_lmcache_gpu_connector(cfg, kv) + + +def test_saving_unfull_chunks_is_refused(monkeypatch): + monkeypatch.setenv("LMCACHE_SAVE_UNFULL_CHUNK", "True") + cfg, kv = _dsa_config() + + with pytest.raises(ValueError, match="LMCACHE_SAVE_UNFULL_CHUNK"): + enforce_lmcache_gpu_connector(cfg, kv) + + +def test_default_chunking_is_accepted(monkeypatch): + monkeypatch.delenv("LMCACHE_CHUNK_SIZE", raising=False) + monkeypatch.delenv("LMCACHE_SAVE_UNFULL_CHUNK", raising=False) + cfg, kv = _dsa_config() + + assert enforce_lmcache_gpu_connector(cfg, kv) is True + + +def test_chunking_is_not_policed_for_single_geometry_models(monkeypatch): + # No indexer cache means no preshuffled page to slice. + monkeypatch.setenv("LMCACHE_CHUNK_SIZE", "100") + + assert ( + enforce_lmcache_gpu_connector( + _vllm_config(index_head_dim=None, architectures=["DeepseekV3ForCausalLM"]), + _kv_cache_config(_spec(head_size=576), _spec(head_size=576)), + ) + is False + ) + + +# --- a priming failure is only survivable when priming was optional --- + + +class _ExplodingConnectorLMCache: + def __init__(self): + self._lmcache_engine = SimpleNamespace( + lmcache_engine=SimpleNamespace( + gpu_connector=SimpleNamespace( + init=False, + initialize_kvcaches_ptr=_boom, + _initialize_kv_cache_pointers=lambda: None, + ) + ) + ) + self.registered = None + + def register_kv_caches(self, kv_caches): + self.registered = kv_caches + + +def _boom(**_kwargs): + raise AttributeError("LMCache moved this") + + +def test_priming_failure_is_fatal_when_v3_was_required(): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_register_kv_caches + + connector = _ExplodingConnectorLMCache() + _wrap_register_kv_caches(connector, required=True) + + with pytest.raises(AttributeError): + connector.register_kv_caches({"mla.0": object()}) + + assert connector.registered is not None, "registration still has to happen first" + + +def test_priming_failure_is_survivable_when_v3_was_optional(caplog): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_register_kv_caches + + connector = _ExplodingConnectorLMCache() + _wrap_register_kv_caches(connector, required=False) + + with caplog.at_level("WARNING", logger="atom"): + connector.register_kv_caches({"mla.0": object()}) + + assert any("lazy discovery" in r.message for r in caplog.records) + + +# --- the null engine has to be inert, not merely have storage_manager --- + + +def test_null_engine_answers_the_health_probe(): + from atom.plugin.vllm.lmcache_connector_patch import _NullLMCacheEngine + + engine = _NullLMCacheEngine() + assert engine.storage_manager is None + assert engine.is_healthy() is True + + +# --- MultiConnector children are reached --- + + +def test_lmcache_nested_in_multi_connector_is_wrapped(): + from atom.plugin.vllm.lmcache_connector_patch import _wrap_lmcache_connectors + + class LMCacheConnectorV1: + def __init__(self): + self._lmcache_engine = SimpleNamespace( + _manager=SimpleNamespace(_lmcache_engine=None) + ) + + def register_kv_caches(self, kv_caches): + pass + + def request_finished(self, request, block_ids): + return False, None + + class MultiConnector: + def __init__(self, children): + self._connectors = children + + child = LMCacheConnectorV1() + _wrap_lmcache_connectors(MultiConnector([child]), required=True) + + assert getattr(child, "_atom_lmcache_group_priming", False) + assert getattr(child, "_atom_lmcache_abort_guard", False) diff --git a/tests/test_lmcache_vllm_dsa_layout.py b/tests/test_lmcache_vllm_dsa_layout.py new file mode 100644 index 0000000000..97ea66ae5c --- /dev/null +++ b/tests/test_lmcache_vllm_dsa_layout.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""LMCache must move a DSA KV layout without disturbing it. + +The indexer pages are not token-major: aiter's +``indexer_k_quant_and_cache(preshuffle=True)`` scatters the fp8 keys into MFMA +16x16 tiles and parks the fp32 scales at the tail of the page, while LMCache's +transfer kernel addresses memory as ``slot * hidden_dim``. They still round-trip +because the scatter is a bijection within one page and LMCache moves whole +pages -- a property of its ``chunk_size`` / ``save_unfull_chunk`` defaults that +``enforce_lmcache_gpu_connector`` enforces rather than assumes. +""" + +from __future__ import annotations + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("moves real paged KV cache; needs a GPU", allow_module_level=True) + +lmcache_gpu_connectors = pytest.importorskip( + "lmcache.v1.gpu_connector.gpu_connectors", + reason="needs LMCache built for ROCm (BUILD_WITH_HIP=1)", +) +aiter = pytest.importorskip("aiter", reason="needs aiter for the indexer kernels") + +from lmcache.v1.memory_management import MemoryFormat +from lmcache.v1.metadata import LMCacheMetadata + +VLLMPagedMemGPUConnectorV2 = lmcache_gpu_connectors.VLLMPagedMemGPUConnectorV2 +VLLMPagedMemGPUConnectorV3 = lmcache_gpu_connectors.VLLMPagedMemGPUConnectorV3 + +BLOCK_SIZE = 64 # AiterSparseMlaBackendForVllm.get_preferred_block_size() +NUM_BLOCKS = 32 +CHUNK_SIZE = 256 # LMCACHE_CHUNK_SIZE, four whole pages +NUM_MLA_LAYERS = 3 +NUM_INDEXER_LAYERS = 2 +MLA_HEAD_SIZE = 576 # kv_lora_rank 512 + qk_rope_head_dim 64 +INDEXER_HEAD_DIM = 128 +INDEXER_HEAD_SIZE = INDEXER_HEAD_DIM + 4 # fp8 keys + one fp32 scale per token + +SRC_BLOCKS = [7, 2, 11, 5] +DST_BLOCKS = [1, 13, 4, 9] + + +class _MemoryObjStub: + """The slice of LMCache's MemoryObj that a GPU connector touches.""" + + def __init__(self, shapes, dtypes): + # The transfer kernel copies straight into host memory, so it has to + # be pinned -- an unpinned buffer fails with "Host tensor not + # registered/pinned". LMCache's real allocator pins its CPU pool. + self._tensors = [ + torch.zeros(shape, dtype=dtype, device="cpu").pin_memory() + for shape, dtype in zip(shapes, dtypes) + ] + self.metadata = type("_Meta", (), {"fmt": MemoryFormat.KV_MLA_FMT})() + + @property + def raw_tensor(self): + return self._tensors[0] + + @property + def tensor(self): + return self._tensors[0] + + def get_tensor(self, index): + return self._tensors[index] + + +def _slot_mapping(block_ids, device): + return torch.tensor( + [b * BLOCK_SIZE + offset for b in block_ids for offset in range(BLOCK_SIZE)], + dtype=torch.long, + device=device, + ) + + +def _metadata(): + """What lmcache.integration.vllm.utils.create_lmcache_metadata() builds. + + Note ``kv_shape`` counts only the model's transformer layers -- vLLM's model + config has no idea the indexer caches exist. That mismatch is the bug. + """ + return LMCacheMetadata( + model_name="dsa-layout-test", + world_size=1, + local_world_size=1, + worker_id=0, + local_worker_id=0, + kv_dtype=torch.uint8, + kv_shape=(NUM_MLA_LAYERS, 1, CHUNK_SIZE, 1, MLA_HEAD_SIZE), + use_mla=True, + chunk_size=CHUNK_SIZE, + ) + + +@pytest.fixture +def dsa_kv_caches(): + """MLA + preshuffled indexer caches, as the ATOM plugin allocates them.""" + device = torch.device("cuda:0") + torch.cuda.set_device(device) + torch.manual_seed(0) + + mla = [ + torch.randint( + 0, + 255, + (NUM_BLOCKS, BLOCK_SIZE, MLA_HEAD_SIZE), + dtype=torch.uint8, + device=device, + ) + for _ in range(NUM_MLA_LAYERS) + ] + indexer = [ + torch.zeros( + (NUM_BLOCKS, BLOCK_SIZE, INDEXER_HEAD_SIZE), + dtype=torch.uint8, + device=device, + ) + for _ in range(NUM_INDEXER_LAYERS) + ] + + slots = _slot_mapping(SRC_BLOCKS, device) + for cache in indexer: + keys = torch.randn( + len(slots), INDEXER_HEAD_DIM, dtype=torch.bfloat16, device=device + ) + aiter.indexer_k_quant_and_cache( + keys, + cache, + slots, + quant_block_size=INDEXER_HEAD_DIM, + scale_fmt="ue8m0", + preshuffle=True, + ) + torch.cuda.synchronize() + return device, mla, indexer + + +def _gather_indexer(cache, block_ids, device): + """Read the indexer cache back the way the sparse attention layer does.""" + from aiter import cp_gather_indexer_k_quant_cache, dtypes + + keys = torch.empty([CHUNK_SIZE, INDEXER_HEAD_DIM], device=device, dtype=dtypes.fp8) + scales = torch.empty([CHUNK_SIZE, 1], device=device, dtype=torch.float32) + cp_gather_indexer_k_quant_cache( + cache, + keys, + scales.view(dtypes.fp8), + torch.tensor([block_ids], dtype=torch.int32, device=device), + torch.tensor([0, CHUNK_SIZE], dtype=torch.int32, device=device), + preshuffle=True, + ) + return keys.view(torch.uint8).clone(), scales.clone() + + +def test_default_connector_cannot_express_a_dsa_layout(dsa_kv_caches): + device, mla, indexer = dsa_kv_caches + connector = VLLMPagedMemGPUConnectorV2.from_metadata( + _metadata(), use_gpu=False, device=device, layout_hints={"kv_layout": "NHD"} + ) + kv_caches = mla + indexer + memory_obj = _MemoryObjStub( + [torch.Size([1, NUM_MLA_LAYERS, CHUNK_SIZE, MLA_HEAD_SIZE])], [torch.uint8] + ) + + with pytest.raises(ValueError, match="broadcast"): + connector.from_gpu( + memory_obj, + 0, + CHUNK_SIZE, + kvcaches=kv_caches, + slot_mapping=_slot_mapping(SRC_BLOCKS, device), + ) + + +def test_v3_connector_round_trips_both_page_geometries(dsa_kv_caches): + device, mla, indexer = dsa_kv_caches + kv_caches = mla + indexer + + expected_mla = [cache[SRC_BLOCKS].clone() for cache in mla] + expected_indexer = [cache[SRC_BLOCKS].clone() for cache in indexer] + expected_gather = [_gather_indexer(cache, SRC_BLOCKS, device) for cache in indexer] + + metadata = _metadata() + connector = VLLMPagedMemGPUConnectorV3.from_metadata( + metadata, use_gpu=False, device=device, layout_hints={"kv_layout": "NHD"} + ) + connector.initialize_kvcaches_ptr(kvcaches=kv_caches) + connector._initialize_kv_cache_pointers() + + groups = metadata.kv_layer_groups_manager.kv_layer_groups + assert [(g.num_layers, g.shape_desc.hs) for g in groups] == [ + (NUM_MLA_LAYERS, MLA_HEAD_SIZE), + (NUM_INDEXER_LAYERS, INDEXER_HEAD_SIZE), + ] + assert {g.shape_desc.bs for g in groups} == {BLOCK_SIZE}, ( + "V3 forwards one scalar block_size to the transfer kernel; groups that " + "disagree would silently corrupt transfers" + ) + + memory_obj = _MemoryObjStub(metadata.get_shapes(CHUNK_SIZE), metadata.get_dtypes()) + connector.from_gpu( + memory_obj, + 0, + CHUNK_SIZE, + kvcaches=kv_caches, + slot_mapping=_slot_mapping(SRC_BLOCKS, device), + ) + torch.cuda.synchronize() + + for cache in kv_caches: + for block_id in DST_BLOCKS: + cache[block_id].zero_() + torch.cuda.synchronize() + + connector.to_gpu( + memory_obj, + 0, + CHUNK_SIZE, + kvcaches=kv_caches, + slot_mapping=_slot_mapping(DST_BLOCKS, device), + ) + torch.cuda.synchronize() + + for i, cache in enumerate(mla): + assert torch.equal(cache[DST_BLOCKS], expected_mla[i]), f"MLA layer {i}" + + for i, cache in enumerate(indexer): + assert torch.equal(cache[DST_BLOCKS], expected_indexer[i]), f"indexer {i} bytes" + keys, scales = _gather_indexer(cache, DST_BLOCKS, device) + assert torch.equal(keys, expected_gather[i][0]), f"indexer {i} keys" + assert torch.equal(scales, expected_gather[i][1]), f"indexer {i} scales" From 1ff6387968565613356b8620ee22d845207873b1 Mon Sep 17 00:00:00 2001 From: kliuae Date: Fri, 4 Sep 2026 09:57:01 +0000 Subject: [PATCH 2/2] fix deadlock and eviction Signed-off-by: kliuae --- atom/plugin/vllm/lmcache_connector_patch.py | 182 ++++++++++++- tests/plugin/test_lmcache_connector_guard.py | 264 ++++++++++++++++++- 2 files changed, 440 insertions(+), 6 deletions(-) diff --git a/atom/plugin/vllm/lmcache_connector_patch.py b/atom/plugin/vllm/lmcache_connector_patch.py index e15a8b03c0..8be5bdeecc 100644 --- a/atom/plugin/vllm/lmcache_connector_patch.py +++ b/atom/plugin/vllm/lmcache_connector_patch.py @@ -44,7 +44,9 @@ def _lmcache_config(): return LMCacheEngineConfig.from_env() except Exception: - logger.debug("ATOM vLLM plugin: could not pre-read LMCache config", exc_info=True) + logger.debug( + "ATOM vLLM plugin: could not pre-read LMCache config", exc_info=True + ) return None @@ -208,7 +210,7 @@ def enforce_lmcache_gpu_connector(vllm_config, kv_cache_config=None) -> bool: if configured is None: os.environ[_LMCACHE_V3_ENV] = "True" logger.info( - "ATOM plugin: %s has %s, which LMCache's default GPU connector " + "ATOM vLLM: %s has %s, which LMCache's default GPU connector " "cannot express. Setting %s=True.", arch, reason, @@ -227,6 +229,110 @@ def enforce_lmcache_gpu_connector(vllm_config, kv_cache_config=None) -> bool: return True +def _bound_broadcast_staging(engine, window: int = 32) -> None: + # In LMCache's save_only_first_rank node, the first rank copies its loaded + # cache to all other ranks through raw_tensors.to(device). In the upstream + # LMCache implementation, the receiving ranks gather all chunks in one go. + # If all data combined exceeds the GPU memory, the receiving ranks abort, + # and the first rank gets deadlocked waiting on the collectives. + # To avoid this, instead of accumulating all chunks, batch the staging + # into smaller number of chunks. + # lmcache.v1.cache_engine.LMCacheEngine._broadcast_or_receive_memory_objs + # but with batched chunks + if not getattr(engine, "save_only_first_rank", False): + return + original = getattr(engine, "_broadcast_or_receive_memory_objs", None) + if original is None or getattr(original, "_atom_bounded_staging", False): + return + + import torch + from lmcache import torch_device_type + from lmcache.v1.memory_management import MemoryObjMetadata, TensorMemoryObj + + original_retrieve = engine.retrieve + to_gpu_kwargs = None + + # batched_to_gpu needs retrieve's paged-KV kwargs, which the broadcast + # helper is not given + def retrieve(tokens, mask=None, **kwargs): + nonlocal to_gpu_kwargs + to_gpu_kwargs = kwargs + try: + return original_retrieve(tokens, mask, **kwargs) + finally: + to_gpu_kwargs = None + + def receive(reordered_chunks, ret_mask): + first_rank = engine.metadata.first_rank + local_rank = engine.metadata.worker_id % torch.cuda.device_count() + device = f"{torch_device_type}:{local_rank}" + chunk_count = engine.broadcast_object_fn(None, first_rank) + if chunk_count is None: + logger.warning("ATOM vLLM: rank %d received None chunk_count", local_rank) + return + + pending: list = [] + + def upload(): + if not pending: + return + engine.gpu_connector.batched_to_gpu( + [c[1] for c in pending], + [c[2] for c in pending], + [c[3] for c in pending], + **to_gpu_kwargs, + ) + if not engine.async_loading: + for _key, memory_obj, _start, _end in pending: + memory_obj.ref_count_down() + pending.clear() + + for _ in range(chunk_count): + combined_metadata = engine.broadcast_object_fn(None, first_rank) + if combined_metadata is None: + logger.warning( + "ATOM vLLM: rank %d received None chunk metadata", local_rank + ) + break + start, end, metadata_dict = combined_metadata + ret_mask[start:end] = True + metadata = MemoryObjMetadata.from_dict(metadata_dict) + raw_tensor = torch.empty( + torch.Size([metadata.get_size()]), dtype=torch.uint8, device=device + ) + engine.broadcast_fn(raw_tensor, first_rank) + pending.append( + ( + None, + TensorMemoryObj( + raw_data=raw_tensor, metadata=metadata, parent_allocator=None + ), + start, + end, + ) + ) + if len(pending) >= window: + upload() + upload() + + def broadcast_or_receive(reordered_chunks, ret_mask): + # Outside our retrieve wrapper the kwargs are unavailable, so windowing + # is impossible and upstream has to handle the whole step + if engine.metadata.is_first_rank() or to_gpu_kwargs is None: + original(reordered_chunks, ret_mask) + else: + receive(reordered_chunks, ret_mask) + + broadcast_or_receive._atom_bounded_staging = True + engine._broadcast_or_receive_memory_objs = broadcast_or_receive + engine.retrieve = retrieve + logger.info( + "ATOM vLLM: bounded LMCache's first-rank restore staging to %d chunks " + "per receiving rank", + window, + ) + + def _prebuild_gpu_connector_layer_groups(connector, kv_caches) -> None: # Run V3's group discovery at this point, rather than inside the first transfer impl = getattr(connector, "_lmcache_engine", None) @@ -240,13 +346,14 @@ def _prebuild_gpu_connector_layer_groups(connector, kv_caches) -> None: gpu_connector.initialize_kvcaches_ptr(kvcaches=list(kv_caches.values())) initialize_pointers() + _bound_broadcast_staging(getattr(impl, "lmcache_engine", None)) manager = getattr( getattr(gpu_connector, "metadata", None), "kv_layer_groups_manager", None ) groups = getattr(manager, "kv_layer_groups", None) or () logger.info( - "ATOM plugin: primed LMCache KV layer groups from %d registered caches -> %s", + "ATOM vLLM: primed LMCache KV layer groups from %d registered caches -> %s", len(kv_caches), [(g.num_layers, g.shape_desc.hs) for g in groups], ) @@ -277,7 +384,7 @@ def register_and_prime(kv_caches): if required: raise logger.warning( - "ATOM plugin: could not pre-build LMCache's KV layer groups; " + "ATOM vLLM: could not pre-build LMCache's KV layer groups; " "falling back to its lazy discovery.", exc_info=True, ) @@ -300,6 +407,70 @@ def is_healthy() -> bool: return True +def _release_unscheduled_lookup_pins(connector, grace_steps: int = 1) -> None: + # Mirroring atom.kv_transfer.offload.dense.connector, release lookup pins + # for requests that are never served so that they can be evicted. + start_load_kv = getattr(connector, "start_load_kv", None) + impl = getattr(connector, "_lmcache_engine", None) + if start_load_kv is None or impl is None: + return + if getattr(start_load_kv, "_atom_pin_release", False): + return + + idle_steps: dict = {} + + def release_stale_pins(): + engine = getattr(impl, "lmcache_engine", None) + pins = getattr(engine, "lookup_pins", None) + if not pins: + idle_steps.clear() + return + metadata = connector._get_connector_metadata() + # A request that is only being saved never reads the pinned chunks + loading = { + str(getattr(req, "req_id", "")) + for req in getattr(metadata, "requests", ()) + if getattr(req, "load_spec", None) is not None + } + for lookup_id in list(pins): + key = str(lookup_id) + if key in loading: + idle_steps.pop(key, None) + continue + idle = idle_steps.get(key, 0) + 1 + if idle > grace_steps: + engine.lookup_unpin(key) + idle_steps.pop(key, None) + logger.debug( + "ATOM vLLM: released the LMCache lookup pin for %s, which " + "was matched but never loaded", + key, + ) + else: + idle_steps[key] = idle + for key in list(idle_steps): + if key not in pins: + idle_steps.pop(key, None) + + def start_load_kv_releasing(forward_context, **kwargs): + try: + release_stale_pins() + except Exception: # optional third-party cleanup boundary + logger.debug( + "ATOM vLLM: releasing stale LMCache lookup pins failed", + exc_info=True, + ) + return start_load_kv(forward_context, **kwargs) + + start_load_kv_releasing._atom_pin_release = True + connector.start_load_kv = start_load_kv_releasing + logger.info( + "ATOM vLLM: releasing LMCache lookup pins that do not become loads " + "within %d step(s), so an unschedulable request cannot pin the CPU tier", + grace_steps + 1, + ) + + def _wrap_request_finished(connector) -> None: # Let an aborted request reach LMCache's cleanup without an engine impl = getattr(connector, "_lmcache_engine", None) @@ -337,6 +508,7 @@ def _wrap_lmcache_connectors(connector, required: bool) -> None: # LMCache connector is reached. _wrap_register_kv_caches(connector, required) _wrap_request_finished(connector) + _release_unscheduled_lookup_pins(connector) for child in getattr(connector, "_connectors", None) or (): _wrap_lmcache_connectors(child, required) @@ -361,6 +533,6 @@ def create_connector(cls, config, role, kv_cache_config=None, *args, **kwargs): create_connector._atom_lmcache_connector_patched = True KVConnectorFactory.create_connector = classmethod(create_connector) logger.info( - "ATOM plugin: patched vLLM KVConnectorFactory to select and prime " + "ATOM vLLM: patched vLLM KVConnectorFactory to select and prime " "LMCache's multi-geometry GPU connector for DSA KV layouts" ) diff --git a/tests/plugin/test_lmcache_connector_guard.py b/tests/plugin/test_lmcache_connector_guard.py index 4a679cacfb..a2f7560fdc 100644 --- a/tests/plugin/test_lmcache_connector_guard.py +++ b/tests/plugin/test_lmcache_connector_guard.py @@ -562,5 +562,267 @@ def __init__(self, children): child = LMCacheConnectorV1() _wrap_lmcache_connectors(MultiConnector([child]), required=True) - assert getattr(child, "_atom_lmcache_group_priming", False) + assert getattr(child, "_atom_lmcache_group_prebuilt", False) assert getattr(child, "_atom_lmcache_abort_guard", False) + + +def _fake_first_rank_engine(chunk_sizes, first_rank=True): + """An LMCache engine stub that records what the broadcast loop does.""" + import torch + + calls = SimpleNamespace(objects=[], tensors=[], delegated=0) + + def broadcast_object_fn(obj, src): + calls.objects.append((obj, src)) + + def broadcast_fn(tensor, src): + calls.tensors.append((tensor.data_ptr(), tensor.numel(), src)) + + def original(reordered_chunks, ret_mask): + calls.delegated += 1 + + original._atom_unblocked = False + + chunks = [] + for i, n in enumerate(chunk_sizes): + raw = torch.zeros(n, dtype=torch.uint8) + obj = SimpleNamespace( + raw_tensor=raw, metadata=SimpleNamespace(to_dict=lambda i=i: {"i": i}) + ) + chunks.append((None, obj, i * 8, i * 8 + 8)) + + engine = SimpleNamespace( + save_only_first_rank=True, + async_loading=False, + broadcast_fn=broadcast_fn, + broadcast_object_fn=broadcast_object_fn, + _broadcast_or_receive_memory_objs=original, + retrieve=lambda tokens, mask=None, **kw: None, + gpu_connector=SimpleNamespace(batched_to_gpu=lambda *a, **k: None), + metadata=SimpleNamespace( + is_first_rank=lambda: first_rank, first_rank=0, worker_id=0 + ), + ) + return engine, chunks, calls + + +def test_first_rank_delegates_to_upstream_send(): + """Rank 0 must stay on LMCache's own send path — we only window the peers.""" + from atom.plugin.vllm.lmcache_connector_patch import _bound_broadcast_staging + + engine, chunks, calls = _fake_first_rank_engine([64, 64], first_rank=True) + _bound_broadcast_staging(engine) + engine._broadcast_or_receive_memory_objs(chunks, None) + + assert calls.delegated == 1, "rank 0 should call through to upstream" + assert calls.objects == [] and calls.tensors == [] + + +def test_peer_without_retrieve_wrapper_falls_back_to_upstream(): + from atom.plugin.vllm.lmcache_connector_patch import _bound_broadcast_staging + + engine, chunks, calls = _fake_first_rank_engine([64], first_rank=False) + _bound_broadcast_staging(engine) + engine._broadcast_or_receive_memory_objs(chunks, None) + + assert calls.delegated == 1 + assert calls.objects == [] and calls.tensors == [] + + +def test_first_rank_copy_patch_is_skipped_and_idempotent(): + from atom.plugin.vllm.lmcache_connector_patch import _bound_broadcast_staging + + engine, _, _ = _fake_first_rank_engine([64]) + engine.save_only_first_rank = False + before = engine._broadcast_or_receive_memory_objs + _bound_broadcast_staging(engine) + assert engine._broadcast_or_receive_memory_objs is before + + engine, _, _ = _fake_first_rank_engine([64]) + _bound_broadcast_staging(engine) + patched = engine._broadcast_or_receive_memory_objs + _bound_broadcast_staging(engine) + assert engine._broadcast_or_receive_memory_objs is patched + + +def test_peer_receive_uploads_and_releases_every_window(): + """The peers must not hold the whole restore in device memory.""" + pytest.importorskip("torch") + import torch + + if not torch.cuda.is_available(): + pytest.skip("needs a GPU") + + import lmcache.v1.memory_management as mm + + from atom.plugin.vllm.lmcache_connector_patch import _bound_broadcast_staging + + n_chunks, window = 7, 3 + uploads: list[int] = [] + released: list[int] = [] + live: list[int] = [] + + class Obj: + def __init__(self, raw_data, metadata, parent_allocator): + self.raw_tensor = raw_data + live.append(1) + + def ref_count_down(self): + released.append(1) + live.pop() + + # rank 0 sends the count first, then one metadata tuple per chunk + feed = iter([n_chunks] + [(i * 8, i * 8 + 8, {}) for i in range(n_chunks)]) + ret_mask = torch.zeros(n_chunks * 8, dtype=torch.bool) + + def upload(objs, starts, ends, **kwargs): + assert kwargs == {"kvcaches": "sentinel"}, "retrieve kwargs must reach to_gpu" + uploads.append(len(objs)) + # nothing may be held beyond the window at any point + assert len(live) <= window + + engine = SimpleNamespace( + save_only_first_rank=True, + async_loading=False, + broadcast_fn=lambda t, src: None, + broadcast_object_fn=lambda obj, src: next(feed), + _broadcast_or_receive_memory_objs=lambda c, m: None, + gpu_connector=SimpleNamespace(batched_to_gpu=upload), + metadata=SimpleNamespace( + is_first_rank=lambda: False, first_rank=0, worker_id=1 + ), + ) + # the real retrieve is what drives the broadcast helper, so the stub must too + engine.retrieve = lambda tokens, mask=None, **kw: ( + engine._broadcast_or_receive_memory_objs([], ret_mask) + ) + + real_obj, real_meta = mm.TensorMemoryObj, mm.MemoryObjMetadata + mm.TensorMemoryObj = Obj + mm.MemoryObjMetadata = SimpleNamespace( + from_dict=staticmethod(lambda d: SimpleNamespace(get_size=lambda: 64)) + ) + try: + _bound_broadcast_staging(engine, window=window) + engine.retrieve([1], None, kvcaches="sentinel") + finally: + mm.TensorMemoryObj, mm.MemoryObjMetadata = real_obj, real_meta + + assert uploads == [3, 3, 1], uploads + assert sum(released) == n_chunks + assert not live, "every received chunk must be released" + assert bool(ret_mask.all()), "every chunk's span must be marked retrieved" + + +def test_staging_patch_is_skipped_and_idempotent(): + from atom.plugin.vllm.lmcache_connector_patch import _bound_broadcast_staging + + engine, _, _ = _fake_first_rank_engine([64]) + engine.save_only_first_rank = False + before = engine._broadcast_or_receive_memory_objs + _bound_broadcast_staging(engine) + assert engine._broadcast_or_receive_memory_objs is before + + engine, _, _ = _fake_first_rank_engine([64]) + _bound_broadcast_staging(engine) + patched = engine._broadcast_or_receive_memory_objs + _bound_broadcast_staging(engine) + assert engine._broadcast_or_receive_memory_objs is patched + + +def _pin_connector(pins, loading, grace=1): + """A worker-side connector stub holding `pins` with `loading` in the step.""" + from atom.plugin.vllm.lmcache_connector_patch import ( + _release_unscheduled_lookup_pins, + ) + + unpinned = [] + + engine = SimpleNamespace( + lookup_pins=dict(pins), + lookup_unpin=lambda rid: ( + unpinned.append(rid), + engine.lookup_pins.pop(rid, None), + ), + ) + meta = SimpleNamespace( + requests=[SimpleNamespace(req_id=r, load_spec=object()) for r in loading] + ) + started = [] + + conn = SimpleNamespace( + _lmcache_engine=SimpleNamespace(lmcache_engine=engine), + _get_connector_metadata=lambda: meta, + start_load_kv=lambda ctx, **kw: started.append(1), + ) + _release_unscheduled_lookup_pins(conn, grace_steps=grace) + return conn, engine, unpinned, started + + +def test_pins_that_never_became_loads_are_released(): + """Mirrors the native connector: a matched-but-unloaded lookup must unpin.""" + conn, engine, unpinned, started = _pin_connector( + pins={"a": {}, "b": {}, "c": {}}, loading=["b"] + ) + + conn.start_load_kv(None) # first sighting: inside the grace window + assert unpinned == [] + conn.start_load_kv(None) # still unscheduled -> release + assert sorted(unpinned) == ["a", "c"] + assert "b" in engine.lookup_pins, "a loading request must keep its pin" + assert len(started) == 2, "the real start_load_kv must still run" + + +def test_a_pin_that_starts_loading_is_kept(): + conn, _engine, unpinned, _ = _pin_connector(pins={"a": {}}, loading=[]) + conn.start_load_kv(None) + conn._get_connector_metadata = lambda: SimpleNamespace( + requests=[SimpleNamespace(req_id="a", load_spec=object())] + ) + conn.start_load_kv(None) + conn.start_load_kv(None) + assert unpinned == [], "the grace counter must reset once loading starts" + + +def test_pin_release_survives_a_broken_engine(): + """Cleanup is best-effort: it must never block the forward pass.""" + from atom.plugin.vllm.lmcache_connector_patch import ( + _release_unscheduled_lookup_pins, + ) + + started = [] + conn = SimpleNamespace( + _lmcache_engine=SimpleNamespace(lmcache_engine=None), + _get_connector_metadata=lambda: (_ for _ in ()).throw(RuntimeError("boom")), + start_load_kv=lambda ctx, **kw: started.append(1), + ) + _release_unscheduled_lookup_pins(conn) + conn.start_load_kv(None) + assert started == [1] + + +def test_a_save_only_request_does_not_keep_its_pin(): + """Only a load reads the pinned chunks; a save in flight must not hold them.""" + from atom.plugin.vllm.lmcache_connector_patch import ( + _release_unscheduled_lookup_pins, + ) + + unpinned = [] + engine = SimpleNamespace( + lookup_pins={"a": {}}, + lookup_unpin=lambda rid: ( + unpinned.append(rid), + engine.lookup_pins.pop(rid, None), + ), + ) + meta = SimpleNamespace(requests=[SimpleNamespace(req_id="a", load_spec=None)]) + conn = SimpleNamespace( + _lmcache_engine=SimpleNamespace(lmcache_engine=engine), + _get_connector_metadata=lambda: meta, + start_load_kv=lambda ctx, **kw: None, + ) + _release_unscheduled_lookup_pins(conn, grace_steps=1) + + conn.start_load_kv(None) + conn.start_load_kv(None) + assert unpinned == ["a"]