From f8372201658083898930f9d35cc6c65a9b7e32a4 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 05:10:45 -0500 Subject: [PATCH 01/16] feat(vllm-plugin): map vLLM KV registration onto ATOM's KVCacheTensor First piece of the vLLM-plugin offload path. vLLM hands a connector a flat {layer_name: tensor} and says nothing about what is inside; ATOM's byte codec wants the movable tensors named apart per layer, because it moves each as its own contiguous segment. This is that translation and nothing else -- no transfer policy, no LMCache. MiniMax-M3 is the model that forced it. It registers three physical layouts at once, and the sparse layers are not contiguous as a whole: stride(1) jumps the entire K region (measured: shape (59454, 2, 128, 1, 128), stride (16384, 974094336, 128, 128, 1)). DenseKVByteCodec rejects non-contiguous segments, so the tensor cannot be one segment -- but t[:, 0] and t[:, 1] each are, and that is exactly the k_cache/v_cache pair the codec already expects. Dense layers stay whole: their K and V interleave inside a block, so the block is one opaque run. The DSA index caches, which vLLM registers as separate layers, fold into their owning layer's index_cache field, which the codec already enumerates. Layer order is numeric rather than dict order: segment order has to match between save and restore, and registration order is not a contract. The codec needs no change -- the tests build M3's real strides, run the mapping through DenseKVByteCodec, and check bytes_per_block against the layout by hand. (cherry picked from commit b723d764cd6f9051b6b893e93b70f6d22c85fb24) --- atom/plugin/vllm/kv_transfer/__init__.py | 2 + .../vllm/kv_transfer/kv_cache_layout.py | 114 ++++++++++++++++++ tests/plugin/test_vllm_kv_cache_layout.py | 112 +++++++++++++++++ 3 files changed, 228 insertions(+) create mode 100644 atom/plugin/vllm/kv_transfer/__init__.py create mode 100644 atom/plugin/vllm/kv_transfer/kv_cache_layout.py create mode 100644 tests/plugin/test_vllm_kv_cache_layout.py diff --git a/atom/plugin/vllm/kv_transfer/__init__.py b/atom/plugin/vllm/kv_transfer/__init__.py new file mode 100644 index 0000000000..2a78457700 --- /dev/null +++ b/atom/plugin/vllm/kv_transfer/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: MIT +"""vLLM-plugin side of ATOM's KV transfer paths.""" diff --git a/atom/plugin/vllm/kv_transfer/kv_cache_layout.py b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py new file mode 100644 index 0000000000..1e62171e5d --- /dev/null +++ b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: MIT +"""Map vLLM's flat KV-cache registration onto ATOM's ``KVCacheTensor`` list. + +vLLM hands a connector ``{layer_name: tensor}`` and says nothing about what is +inside each tensor. ATOM's offload codec instead wants, per layer, the movable +tensors named apart (``k_cache`` / ``v_cache`` / scales / ``index_cache``), +because it moves each one as its own contiguous byte segment. This module is +that translation, and nothing else -- no transfer policy, no LMCache. + +Why the split matters (MiniMax-M3, the model that forced this): + + dense layers (nb, 1, bs, 2*hd) K and V interleaved per token + sparse layers (nb, 2, bs, nh, hd) K and V in two SEPARATE regions + index caches (nb, bs, hd) DSA indexer keys, one per sparse layer + +A sparse layer's tensor is NOT contiguous as a whole: its ``stride(1)`` jumps +across the entire K region (measured on M3-MXFP4: shape ``(59454, 2, 128, 1, +128)``, stride ``(16384, 974094336, 128, 128, 1)``). ``DenseKVByteCodec`` +rejects non-contiguous segments, so the whole tensor cannot be one segment -- +but ``t[:, 0]`` and ``t[:, 1]`` each ARE contiguous, and that is exactly the +``k_cache`` / ``v_cache`` pair the codec already expects. Dense layers stay +whole: their K and V interleave inside a block, so the block's bytes are one +opaque run and splitting them would be meaningless. + +The codec never interprets these bytes, so "which tensor is K" only has to be +consistent between save and restore -- not semantically right. +""" + +import torch + +from atom.config import KVCacheTensor + +INDEX_CACHE_SUFFIX = ".index_cache" + + +def _layer_sort_key(layer_name: str) -> tuple: + """Order layers by their numeric position, falling back to name order. + + Segment order must be identical on save and restore, and dict insertion + order is whatever vLLM's registration happened to produce. + """ + parts = [] + for token in layer_name.replace("/", ".").split("."): + parts.append((0, int(token), "") if token.isdigit() else (1, 0, token)) + return tuple(parts) + + +def split_kv_tensor(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return ``(k_cache, v_cache)`` for one layer's registered KV tensor. + + ``v_cache`` is None when K and V share one opaque per-block run (dense + layers), in which case the whole tensor travels as ``k_cache``. + """ + # (nb, 2, ...) -- K/V split across dim 1, each half contiguous on its own. + if tensor.ndim >= 4 and tensor.shape[1] == 2: + return tensor[:, 0], tensor[:, 1] + # Anything else (incl. dense (nb, 1, bs, 2*hd)) is one opaque run. + return tensor, None + + +def build_kv_cache_tensors( + kv_caches: dict[str, torch.Tensor], +) -> list[KVCacheTensor]: + """Translate vLLM's ``{layer_name: tensor}`` into ATOM ``KVCacheTensor``s. + + Layers named ```` are folded into the owning + layer's ``index_cache`` rather than becoming layers of their own -- vLLM + registers M3's DSA indexer keys as separate entries, but they are part of + the same layer's movable bytes. + + Raises: + ValueError: a produced segment is not contiguous (the codec would + reject it later, with far less context about which layer). + """ + index_caches: dict[str, torch.Tensor] = {} + main: dict[str, torch.Tensor] = {} + for name, tensor in kv_caches.items(): + if name.endswith(INDEX_CACHE_SUFFIX): + index_caches[name[: -len(INDEX_CACHE_SUFFIX)]] = tensor + else: + main[name] = tensor + + orphans = set(index_caches) - set(main) + if orphans: + raise ValueError( + f"index caches registered without their owning layer: {sorted(orphans)}" + ) + + out: list[KVCacheTensor] = [] + for layer_num, name in enumerate(sorted(main, key=_layer_sort_key)): + k_cache, v_cache = split_kv_tensor(main[name]) + index_cache = index_caches.get(name) + + for role, seg in ( + ("k_cache", k_cache), + ("v_cache", v_cache), + ("index_cache", index_cache), + ): + if seg is not None and not seg.is_contiguous(): + raise ValueError( + f"{name}: {role} is not contiguous " + f"(shape={tuple(seg.shape)}, stride={seg.stride()}); " + "the byte codec can only move contiguous segments" + ) + + out.append( + KVCacheTensor( + layer_num=layer_num, + k_cache=k_cache, + v_cache=v_cache if v_cache is not None else torch.tensor([]), + index_cache=index_cache, + ) + ) + return out diff --git a/tests/plugin/test_vllm_kv_cache_layout.py b/tests/plugin/test_vllm_kv_cache_layout.py new file mode 100644 index 0000000000..a69f3d3de1 --- /dev/null +++ b/tests/plugin/test_vllm_kv_cache_layout.py @@ -0,0 +1,112 @@ +"""vLLM registration -> ATOM KVCacheTensor mapping, on MiniMax-M3's real layouts. + +M3 is the model that forced this mapping: it registers three different physical +layouts at once, and one of them (the sparse layers) is not contiguous as a +whole, so it cannot be handed to ``DenseKVByteCodec`` unsplit. The strides here +are the ones measured on M3-MXFP4 (see FINDINGS in the M3 offload notes), scaled +down in block count -- the contiguity properties depend on the axis order, not +on how many blocks there are. +""" + +from __future__ import annotations + +import pytest +import torch + +from atom.plugin.vllm.kv_transfer.kv_cache_layout import ( + build_kv_cache_tensors, + split_kv_tensor, +) + +NB, BS, HD = 8, 128, 128 +DENSE_LAYERS, SPARSE_LAYERS = 3, 5 + + +def _sparse_kv(nb: int = NB) -> torch.Tensor: + """K and V in two separate regions, exactly as M3 allocates them. + + ``stride(1)`` jumps the whole K region, so the tensor is not contiguous -- + which is the entire point of the split this test covers. + """ + k_block = BS * HD + k_total = nb * k_block + buf = torch.zeros(2 * k_total, dtype=torch.uint8) + return buf.as_strided((nb, 2, BS, 1, HD), (k_block, k_total, HD, HD, 1)) + + +def _m3_registration() -> dict[str, torch.Tensor]: + kv: dict[str, torch.Tensor] = {} + for i in range(DENSE_LAYERS): # K/V interleaved per token + kv[f"model.layers.{i}.self_attn.attn"] = torch.zeros( + (NB, 1, BS, 2 * HD), dtype=torch.uint8 + ) + for i in range(DENSE_LAYERS, DENSE_LAYERS + SPARSE_LAYERS): + name = f"model.layers.{i}.self_attn.attn" + kv[name] = _sparse_kv() + kv[f"{name}.index_cache"] = torch.zeros((NB, BS, HD), dtype=torch.float8_e4m3fn) + return kv + + +def test_sparse_layer_is_only_movable_once_split(): + sparse = _sparse_kv() + assert not sparse.is_contiguous(), "fixture no longer reproduces M3's layout" + + k, v = split_kv_tensor(sparse) + assert k.is_contiguous() and v.is_contiguous() + assert k.numel() == v.numel() == NB * BS * HD + + +def test_dense_layer_travels_whole(): + dense = torch.zeros((NB, 1, BS, 2 * HD), dtype=torch.uint8) + k, v = split_kv_tensor(dense) + assert v is None, "dense K/V interleave inside a block; splitting is meaningless" + assert k is dense + + +def test_index_caches_fold_into_their_owning_layer(): + tensors = build_kv_cache_tensors(_m3_registration()) + + assert ( + len(tensors) == DENSE_LAYERS + SPARSE_LAYERS + ), "index caches must not become layers of their own" + with_index = [t for t in tensors if t.index_cache is not None] + assert len(with_index) == SPARSE_LAYERS + + +def test_layer_order_is_numeric_not_dict_order(): + kv = _m3_registration() + shuffled = dict(reversed(list(kv.items()))) + + got = build_kv_cache_tensors(shuffled) + + # Segment order must not depend on registration order: a save written in + # one order and restored in another would scatter bytes to the wrong layers. + assert [t.layer_num for t in got] == list(range(DENSE_LAYERS + SPARSE_LAYERS)) + assert got[0].v_cache.numel() == 0, "layer 0 is dense -> no separate V" + assert got[-1].index_cache is not None, "last layer is sparse -> has an index cache" + + +def test_orphan_index_cache_is_rejected(): + with pytest.raises(ValueError, match="without their owning layer"): + build_kv_cache_tensors( + {"model.layers.0.self_attn.attn.index_cache": torch.zeros((NB, BS, HD))} + ) + + +def test_codec_accepts_the_mapped_tensors(): + """The mapping's whole purpose: make M3 pass ATOM's byte codec.""" + codec_mod = pytest.importorskip( + "atom.kv_transfer.offload.dense.kv_byte_codec", + reason="offload codec pulls aiter", + ) + tensors = build_kv_cache_tensors(_m3_registration()) + + codec = codec_mod.DenseKVByteCodec( + {str(t.layer_num): t for t in tensors}, num_blocks=NB + ) + + dense_bytes = BS * 2 * HD # one opaque K/V run + sparse_bytes = 2 * (BS * HD) + BS * HD # K + V + index + assert codec.bytes_per_block == ( + DENSE_LAYERS * dense_bytes + SPARSE_LAYERS * sparse_bytes + ) From 85c700dc999dcce3327978d31f2e57af94b0a0f2 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 08:34:37 -0500 Subject: [PATCH 02/16] feat(vllm-plugin): drive ATOM's byte-level LMCache offload from vLLM Adds the adapter that lets vLLM's connector API run ATOM's existing offload codec, so the plugin path gets the tier the native path already has. Why not point vLLM at LMCache's own connector: its GPU connectors accept only the clean NHD/HND family and pick ONE format per model. MiniMax-M3 registers three layouts at once, and on ROCm the paths that could express that are unavailable anyway -- the per-layer-format connector (V3) is off by default and hangs on M3, and the multi-process path imports cupy, which LMCache's platform/rocm does not provide. ATOM sidesteps all of it by never asking LMCache to understand the layout: DenseKVByteCodec gathers whole paged blocks into a chunk-major uint8 blob and LMCache stores opaque bytes. Pieces: - connector.py wraps ATOM's DenseOffloadConnector/Scheduler behind KVConnectorBase_V1. Layer-granular hooks are inert on purpose: ATOM moves a whole request's blocks per transfer, and blocking a forward on a save would put the offload tier on the critical path it exists to keep clear. - seq_view.py presents a vLLM Request as the seq ATOM's scheduler expects. It must be one stable instance per request: ATOM stores the object and compares identity to detect a recycled request id, so a fresh wrapper per call would reset the load lifecycle every step. - offload_config.py derives ATOM's three config fields from VllmConfig instead of duplicating block size and role. - AtomOffloadMetadata wraps ATOM's metadata, which derives from ATOM's own base rather than vLLM's and so cannot be returned from build_connector_meta. Registered by module path so importing the plugin does not drag LMCache into every run; vLLM resolves it only when --kv-transfer-config names it. Import and factory registration verified in the container against vLLM 0.27.2.dev; 12 unit tests cover the mapping and the identity contract. (cherry picked from commit 1ee388ebb81cf39d1b926b793f2b48989d39d2e0) --- atom/plugin/vllm/kv_transfer/connector.py | 207 ++++++++++++++++++ .../plugin/vllm/kv_transfer/offload_config.py | 63 ++++++ atom/plugin/vllm/kv_transfer/seq_view.py | 106 +++++++++ atom/plugin/vllm/register.py | 22 ++ tests/plugin/test_vllm_offload_seq_view.py | 78 +++++++ 5 files changed, 476 insertions(+) create mode 100644 atom/plugin/vllm/kv_transfer/connector.py create mode 100644 atom/plugin/vllm/kv_transfer/offload_config.py create mode 100644 atom/plugin/vllm/kv_transfer/seq_view.py create mode 100644 tests/plugin/test_vllm_offload_seq_view.py diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py new file mode 100644 index 0000000000..186de7e4bf --- /dev/null +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: MIT +"""vLLM KV connector that drives ATOM's byte-level LMCache offload. + +Why this exists rather than pointing vLLM at LMCache's own connector: + +LMCache's GPU connectors accept only the clean NHD/HND family and pick ONE +format for the whole model. MiniMax-M3 registers three physical layouts at once +(dense K/V interleaved, sparse K/V in separate regions, plus a DSA index cache), +so no single format describes it -- and on ROCm the paths that could describe it +are unavailable anyway: the per-layer-format connector (V3) is off by default +and hangs on M3, and the multi-process path needs cupy, which LMCache's +``platform/rocm`` does not provide. + +ATOM already solved this for its native engine by not asking LMCache to +understand the layout at all: ``DenseKVByteCodec`` gathers whole paged blocks +into a chunk-major uint8 blob, and LMCache only ever stores opaque bytes. That +codec is reused verbatim here; this module is the adapter that lets vLLM drive +it, so the plugin path gets the same guarantee the native path already tests +(byte-identical round-trip). + +Layer-granular hooks are deliberately inert: ATOM moves a whole request's blocks +per transfer, not one layer at a time. +""" + +import logging +from typing import TYPE_CHECKING, Any + +import torch +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + KVConnectorRole, +) + +from atom.plugin.vllm.kv_transfer.kv_cache_layout import build_kv_cache_tensors +from atom.plugin.vllm.kv_transfer.offload_config import build_offload_config +from atom.plugin.vllm.kv_transfer.seq_view import SeqViewRegistry + +if TYPE_CHECKING: + from vllm.forward_context import ForwardContext + +logger = logging.getLogger("atom") + + +class AtomOffloadMetadata(KVConnectorMetadata): + """Carries ATOM's offload metadata through vLLM's connector plumbing. + + ATOM's ``LMCacheOffloadMetadata`` derives from ATOM's own ConnectorMetadata, + not vLLM's, so it cannot be returned directly from ``build_connector_meta``. + Wrapping keeps ATOM's descriptors intact instead of flattening them into a + vLLM-shaped copy that would then have to be kept in sync. + """ + + def __init__(self, inner) -> None: + super().__init__() + self.inner = inner + + +class AtomLMCacheOffloadConnector(KVConnectorBase_V1): + """Drives ``atom.kv_transfer.offload`` from vLLM's connector API.""" + + def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): + super().__init__(vllm_config, role) + self._config = build_offload_config(vllm_config) + self._worker = None + self._scheduler = None + + self._seqs = SeqViewRegistry() + + if role == KVConnectorRole.WORKER: + from atom.kv_transfer.offload.dense.connector import DenseOffloadConnector + + self._worker = DenseOffloadConnector(self._config) + else: + from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler + + self._scheduler = DenseOffloadScheduler(self._config) + + # ---- worker side -------------------------------------------------- + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Translate vLLM's flat registration and hand it to ATOM's codec.""" + tensors = build_kv_cache_tensors(kv_caches) + if not tensors: + raise ValueError("ATOM offload connector: vLLM registered no KV caches") + + # Every segment's per-block stride is derived from num_blocks, so it has + # to be the physical block count -- not a token count. Block-major KV + # carries it in dim 0; taking it from the first mapped k_cache keeps the + # value consistent with the very tensors the codec will slice. + num_blocks = int(tensors[0].k_cache.shape[0]) + + self._worker.register_kv_caches( + {str(t.layer_num): t for t in tensors}, + num_blocks=num_blocks, + ) + logger.info( + "ATOM LMCache offload: registered %d layers, num_blocks=%d", + len(tensors), + num_blocks, + ) + + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: + metadata = self._get_connector_metadata() + inner = getattr(metadata, "inner", None) + if inner is not None: + self._worker.start_load_kv(inner) + + def wait_for_layer_load(self, layer_name: str) -> None: + """Inert: transfers are per-request, not per-layer. + + A load is published to the scheduler through ``get_finished`` only once + every block of the request has landed, so there is no partially-loaded + layer for a forward to wait on. + """ + + def save_kv_layer(self, layer_name: str, kv_layer, attn_metadata, **kwargs) -> None: + """Inert: saves are issued per request from ``build_connector_meta``.""" + + def wait_for_save(self) -> None: + """Inert: saves are fire-and-forget on ATOM's save executor. + + Blocking the forward on them would put offload on the critical path, + which is the opposite of what the tier is for. Completion still reaches + the scheduler via ``get_finished``. + """ + + def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: + return self._worker.get_finished(finished_req_ids) + + def shutdown(self) -> None: + for side in (self._worker, self._scheduler): + close = getattr(side, "close", None) or getattr(side, "shutdown", None) + if close is not None: + close() + + # ---- scheduler side ----------------------------------------------- + + def get_num_new_matched_tokens( + self, request, num_computed_tokens: int + ) -> tuple[int, bool]: + """How many extra prompt tokens the offload tier can supply. + + ``num_computed_tokens`` is vLLM's HBM-prefix-cache frontier; ATOM reads + the same quantity off the seq to avoid re-loading what is already + resident, so it has to be pushed in before the lookup runs. + """ + seq = self._seqs.get_or_create(request) + seq.set_num_cached_tokens(num_computed_tokens) + return self._scheduler.get_num_new_matched_tokens(seq) + + def update_state_after_alloc(self, request, blocks, num_external_tokens: int): + seq = self._seqs.get_or_create(request) + seq.set_block_table(_block_ids(blocks)) + self._scheduler.update_state_after_alloc(seq) + + def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata: + """Snapshot this step's transfers. + + The frontier of every scheduled request is refreshed first: ATOM decides + which chunks are safe to save by comparing against it, and a stale value + would either skip chunks or offer up tokens that are not computed yet. + """ + for req_id, num_tokens in _scheduled_frontiers(scheduler_output): + seq = self._seqs.get(req_id) + if seq is not None: + seq.set_num_cached_tokens(num_tokens) + return AtomOffloadMetadata(self._scheduler.build_connector_meta()) + + def request_finished(self, request, block_ids) -> tuple[bool, dict | None]: + seq = self._seqs.get(request.request_id) + if seq is not None: + self._scheduler.request_finished(seq) + # Blocks may still be pinned by an in-flight save; ATOM says when. + if self._scheduler.should_defer_free(seq): + return True, None + self._seqs.drop(request.request_id) + return False, None + + +def _block_ids(blocks) -> list[int]: + """Flatten vLLM's allocated-block structure to plain ids. + + vLLM has spelled this several ways across versions (KVCacheBlocks with + ``get_block_ids()``, a per-group tuple of lists, or already-flat ids), and + the offload codec only ever needs the ids. + """ + getter = getattr(blocks, "get_block_ids", None) + if getter is not None: + blocks = getter() + if ( + isinstance(blocks, (list, tuple)) + and blocks + and isinstance(blocks[0], (list, tuple)) + ): + return [int(b) for group in blocks for b in group] + return [int(b) for b in (blocks or [])] + + +def _scheduled_frontiers(scheduler_output): + """Yield ``(request_id, num_computed_tokens)`` for this step's requests.""" + for req in getattr(scheduler_output, "scheduled_new_reqs", ()) or (): + yield req.req_id, getattr(req, "num_computed_tokens", 0) + cached = getattr(scheduler_output, "scheduled_cached_reqs", None) + req_ids = getattr(cached, "req_ids", None) or [] + computed = getattr(cached, "num_computed_tokens", None) or [] + yield from zip(req_ids, computed) diff --git a/atom/plugin/vllm/kv_transfer/offload_config.py b/atom/plugin/vllm/kv_transfer/offload_config.py new file mode 100644 index 0000000000..15f55c72ba --- /dev/null +++ b/atom/plugin/vllm/kv_transfer/offload_config.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: MIT +"""Build the small config object ATOM's offload connector expects from vLLM's. + +ATOM's offload code was written against ``atom.config.Config``, but it only ever +reads three things from it. Rather than construct a whole ATOM Config inside the +plugin -- which would mean keeping two sources of truth for block size and role +in sync -- this hands it a shim carrying exactly those fields, read off the +vLLM config that is already authoritative in plugin mode. + +Kept apart from ``connector.py`` so it is importable (and testable) without +vLLM installed. +""" + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class OffloadConfigShim: + """The ATOM-offload-shaped view of a vLLM config. + + Field names match what ``_init_worker_common`` / ``DenseOffloadConnector`` + read; nothing else in ATOM's Config is touched by the offload path. + """ + + kv_cache_block_size: int + kv_transfer_config: dict[str, Any] = field(default_factory=dict) + decode_context_parallel_size: int = 1 + + +def _extra_config(kv_transfer_config: Any) -> dict[str, Any]: + extra = getattr(kv_transfer_config, "kv_connector_extra_config", None) + return dict(extra) if isinstance(extra, dict) else {} + + +def build_offload_config(vllm_config: Any) -> OffloadConfigShim: + """Derive ATOM's offload config from ``VllmConfig``. + + The role is normalized to ATOM's vocabulary: vLLM connectors are configured + with ``kv_role`` values like ``kv_both``, while ATOM's native path spells + the same intent ``offload``. ``_init_worker_common`` accepts both, so the + value is passed through unchanged rather than translated -- translating + would silently change save/load enablement if the two vocabularies drift. + """ + cache_config = getattr(vllm_config, "cache_config", None) + block_size = getattr(cache_config, "block_size", None) + if not block_size: + raise ValueError( + "ATOM offload connector: vLLM reported no cache_config.block_size; " + "the byte codec addresses KV by block and cannot proceed without it" + ) + + kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None) + role = getattr(kv_transfer_config, "kv_role", None) or "kv_both" + + parallel_config = getattr(vllm_config, "parallel_config", None) + dcp = getattr(parallel_config, "decode_context_parallel_size", 1) or 1 + + return OffloadConfigShim( + kv_cache_block_size=int(block_size), + kv_transfer_config={"kv_role": role, **_extra_config(kv_transfer_config)}, + decode_context_parallel_size=int(dcp), + ) diff --git a/atom/plugin/vllm/kv_transfer/seq_view.py b/atom/plugin/vllm/kv_transfer/seq_view.py new file mode 100644 index 0000000000..ae156b0e9f --- /dev/null +++ b/atom/plugin/vllm/kv_transfer/seq_view.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: MIT +"""Present a vLLM ``Request`` as the ``seq`` ATOM's offload scheduler expects. + +ATOM's offload scheduler was written against ``atom.model_engine.sequence. +Sequence``. It reads nine attributes off it, four of which are offload's own +mutable bookkeeping, and -- importantly -- it *stores the object* and later +compares identity (``previous is not seq``, ``entry[0] is not seq``) to detect +that a request id was reused by a new request. So a view must be created ONCE +per request and reused; handing out a fresh wrapper per call would look like a +new request every step and keep resetting the load lifecycle. + +The read-only half is projected from the vLLM Request; the mutable half lives +here, exactly as it lives on ATOM's Sequence. +""" + +from typing import Any + + +class SeqView: + """One vLLM request, shaped like an ATOM ``Sequence`` for offload.""" + + __slots__ = ( + "_load_operation", + "_num_cached_tokens", + "_request", + "block_table", + "offload_handoff_boundary_tokens", + "offload_loaded_tokens", + "prefix_hashes_published", + ) + + def __init__(self, request: Any) -> None: + self._request = request + self._num_cached_tokens = 0 + self.block_table: list[int] = [] + # Offload's own state, mirroring the fields on ATOM's Sequence. + self.offload_loaded_tokens = 0 + self.offload_handoff_boundary_tokens = 0 + self.prefix_hashes_published = False + self._load_operation = None + + # -- projected from the vLLM request -------------------------------- + + @property + def id(self) -> str: + return self._request.request_id + + @property + def token_ids(self) -> list[int]: + # LMCache keys are derived from prompt tokens; ``all_token_ids`` grows + # with decode output, which must not change a prefix's key mid-request. + return self._request.prompt_token_ids + + @property + def num_prompt_tokens(self) -> int: + return len(self._request.prompt_token_ids) + + @property + def num_cached_tokens(self) -> int: + """The scheduler's computed frontier for this request. + + ATOM reads this to decide which chunks are safe to save (everything + below the frontier has been computed) and how much of a lookup hit is + already resident in HBM. vLLM reports the same quantity as + ``num_computed_tokens``, pushed in by the connector each step. + """ + return self._num_cached_tokens + + def set_num_cached_tokens(self, value: int) -> None: + self._num_cached_tokens = int(value) + + def set_block_table(self, block_ids: list[int]) -> None: + self.block_table = list(block_ids) + + def __repr__(self) -> str: + return ( + f"SeqView(id={self.id!r}, prompt={self.num_prompt_tokens}, " + f"cached={self.num_cached_tokens}, blocks={len(self.block_table)})" + ) + + +class SeqViewRegistry: + """Keeps one ``SeqView`` per request id, for as long as offload needs it.""" + + def __init__(self) -> None: + self._views: dict[str, SeqView] = {} + + def get_or_create(self, request: Any) -> SeqView: + rid = request.request_id + view = self._views.get(rid) + # A reused id with a different underlying request must produce a NEW + # view: that identity change is precisely what tells ATOM's scheduler + # to drop the previous request's pending load. + if view is None or view._request is not request: + view = SeqView(request) + self._views[rid] = view + return view + + def get(self, request_id: str) -> SeqView | None: + return self._views.get(request_id) + + def drop(self, request_id: str) -> None: + self._views.pop(request_id, None) + + def __len__(self) -> int: + return len(self._views) diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index 25c9f9a317..f0bcca5775 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -171,10 +171,32 @@ def register_platform() -> str | None: apply_vllm_v4_block_reuse_patch() + _register_kv_connectors() + # return the ATOM platform to vllm return "atom.plugin.vllm.platform.ATOMPlatform" +def _register_kv_connectors() -> None: + """Expose ATOM's byte-level LMCache offload to vLLM's connector factory. + + Registered by module path so importing the plugin does not drag in the + offload stack (and LMCache) for every run -- vLLM resolves it lazily, only + when a --kv-transfer-config actually names it. + """ + from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory + + name = "AtomLMCacheOffloadConnector" + if name in getattr(KVConnectorFactory, "_registry", {}): + return + KVConnectorFactory.register_connector( + name, + "atom.plugin.vllm.kv_transfer.connector", + name, + ) + logger.info("Registered ATOM KV connector: %s", name) + + def _patch_vllm_attention_process_weights_after_loading(attention) -> None: orig = attention.process_weights_after_loading diff --git a/tests/plugin/test_vllm_offload_seq_view.py b/tests/plugin/test_vllm_offload_seq_view.py new file mode 100644 index 0000000000..661433041b --- /dev/null +++ b/tests/plugin/test_vllm_offload_seq_view.py @@ -0,0 +1,78 @@ +"""SeqView identity contract. + +ATOM's offload scheduler stores the seq object and compares identity to detect +that a request id was reused. Get this wrong and either every step looks like a +new request (load lifecycle reset forever) or a genuinely new request inherits +the previous one's pending load. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +from atom.plugin.vllm.kv_transfer.seq_view import SeqViewRegistry + + +def _request(rid: str = "r1", prompt=(1, 2, 3)): + return SimpleNamespace(request_id=rid, prompt_token_ids=list(prompt)) + + +def test_same_request_yields_the_same_view(): + reg = SeqViewRegistry() + req = _request() + + assert reg.get_or_create(req) is reg.get_or_create(req) + + +def test_reused_id_with_a_new_request_yields_a_new_view(): + reg = SeqViewRegistry() + first = reg.get_or_create(_request("r1")) + + second = reg.get_or_create(_request("r1")) # same id, different request + + assert second is not first, ( + "a recycled request id must present as a new seq, or the new request " + "inherits the old one's pending load" + ) + + +def test_mutable_offload_state_survives_across_lookups(): + reg = SeqViewRegistry() + req = _request() + view = reg.get_or_create(req) + view.offload_loaded_tokens = 256 + view.set_block_table([7, 8, 9]) + + again = reg.get_or_create(req) + + assert again.offload_loaded_tokens == 256 + assert again.block_table == [7, 8, 9] + + +def test_prompt_tokens_are_the_key_source_not_decode_output(): + reg = SeqViewRegistry() + req = _request(prompt=(1, 2, 3)) + req.all_token_ids = [1, 2, 3, 99, 100] # decode has appended output + view = reg.get_or_create(req) + + # LMCache keys come from these; letting decode output in would change a + # prefix's key mid-request and orphan everything already stored. + assert view.token_ids == [1, 2, 3] + assert view.num_prompt_tokens == 3 + + +def test_frontier_is_pushed_in_from_vllm(): + reg = SeqViewRegistry() + view = reg.get_or_create(_request()) + assert view.num_cached_tokens == 0 + + view.set_num_cached_tokens(128) + + assert view.num_cached_tokens == 128 + + +def test_drop_forgets_the_request(): + reg = SeqViewRegistry() + reg.get_or_create(_request("r1")) + reg.drop("r1") + assert reg.get("r1") is None and len(reg) == 0 From aca5bfa1a45a52b6c9b580894b7b44e4c2c11ab1 Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 08:47:32 -0500 Subject: [PATCH 03/16] fix(vllm-plugin): project every config field the offload path reads The first live run died with 'OffloadConfigShim has no attribute hf_config'. The shim carried three fields because that is all dense/connector.py reads directly; atom/kv_transfer/offload/config.py reads more -- hf_config (layer count and the LMCache page namespace), kv_cache_dtype, PP geometry, and the model name. Grepped the whole offload package for config reads rather than patching one attribute at a time. KV dtype needs translating: vLLM spells it fp8/auto/ bfloat16, ATOM indexes aiter.dtypes with fp8/bf16/fp16, and 'auto' means the model dtype. Unknown spellings raise instead of defaulting -- the codec sizes every byte segment from this, so a wrong guess mis-sizes transfers silently rather than failing. (cherry picked from commit fd9ede7acdcf5108934d2b11ed206190d9cd8b40) --- atom/plugin/vllm/kv_transfer/connector.py | 5 +- .../plugin/vllm/kv_transfer/offload_config.py | 128 ++++++++++++------ atom/plugin/vllm/register.py | 13 +- tests/plugin/test_vllm_offload_config.py | 81 +++++++++++ 4 files changed, 183 insertions(+), 44 deletions(-) create mode 100644 tests/plugin/test_vllm_offload_config.py diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 186de7e4bf..8b66d081b2 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -60,7 +60,10 @@ class AtomLMCacheOffloadConnector(KVConnectorBase_V1): """Drives ``atom.kv_transfer.offload`` from vLLM's connector API.""" def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): - super().__init__(vllm_config, role) + # kv_cache_config is required of out-of-tree v1 connectors: the factory + # rejects the 2-argument signature outright, and the base class stores + # it for the group-aware paths. + super().__init__(vllm_config, role, kv_cache_config) self._config = build_offload_config(vllm_config) self._worker = None self._scheduler = None diff --git a/atom/plugin/vllm/kv_transfer/offload_config.py b/atom/plugin/vllm/kv_transfer/offload_config.py index 15f55c72ba..e17954ecaf 100644 --- a/atom/plugin/vllm/kv_transfer/offload_config.py +++ b/atom/plugin/vllm/kv_transfer/offload_config.py @@ -1,63 +1,109 @@ # SPDX-License-Identifier: MIT -"""Build the small config object ATOM's offload connector expects from vLLM's. +"""Present a ``VllmConfig`` as the config ATOM's offload path expects. -ATOM's offload code was written against ``atom.config.Config``, but it only ever -reads three things from it. Rather than construct a whole ATOM Config inside the -plugin -- which would mean keeping two sources of truth for block size and role -in sync -- this hands it a shim carrying exactly those fields, read off the -vLLM config that is already authoritative in plugin mode. +ATOM's offload code was written against ``atom.config.Config``. It reads a +handful of fields off it -- block size, KV dtype, the HF config, PP geometry, +the transfer role -- all of which vLLM already owns in plugin mode. Rather than +build a second ATOM Config (two sources of truth that drift), this projects the +vLLM one, forwarding by reference wherever possible so there is nothing to keep +in sync. -Kept apart from ``connector.py`` so it is importable (and testable) without +Kept apart from ``connector.py`` so it stays importable, and testable, without vLLM installed. """ -from dataclasses import dataclass, field from typing import Any +# vLLM spells its KV cache dtype in its own vocabulary; ATOM indexes +# ``aiter.dtypes.d_dtypes`` with its own. Only the values a KV cache can +# actually hold are mapped -- an unknown one must raise rather than silently +# pick a width, because the codec sizes every byte segment from it. +_VLLM_TO_ATOM_KV_DTYPE = { + "fp8": "fp8", + "fp8_e4m3": "fp8", + "fp8_e5m2": "fp8", + "fp8_inc": "fp8", + "bfloat16": "bf16", + "float16": "fp16", + "half": "fp16", + "float32": "fp32", + "float": "fp32", +} + + +def _atom_kv_dtype(cache_config: Any, model_config: Any) -> str: + """Translate vLLM's ``cache_dtype`` into an ``aiter.dtypes`` key.""" + raw = str(getattr(cache_config, "cache_dtype", "auto") or "auto") + if raw == "auto": + # "auto" means "same as the model", which vLLM resolves on model_config. + raw = str(getattr(model_config, "dtype", "bfloat16")).replace("torch.", "") + key = _VLLM_TO_ATOM_KV_DTYPE.get(raw) + if key is None: + raise ValueError( + f"ATOM offload connector: unsupported KV cache dtype {raw!r}. " + "The byte codec sizes each segment from it, so it cannot be guessed." + ) + return key + -@dataclass class OffloadConfigShim: """The ATOM-offload-shaped view of a vLLM config. - Field names match what ``_init_worker_common`` / ``DenseOffloadConnector`` - read; nothing else in ATOM's Config is touched by the offload path. + Attribute names match what ``atom.kv_transfer.offload`` reads; everything + else on ATOM's Config is untouched by that path. """ - kv_cache_block_size: int - kv_transfer_config: dict[str, Any] = field(default_factory=dict) - decode_context_parallel_size: int = 1 - + def __init__(self, vllm_config: Any) -> None: + self._vllm_config = vllm_config + cache_config = getattr(vllm_config, "cache_config", None) + model_config = getattr(vllm_config, "model_config", None) + parallel_config = getattr(vllm_config, "parallel_config", None) -def _extra_config(kv_transfer_config: Any) -> dict[str, Any]: - extra = getattr(kv_transfer_config, "kv_connector_extra_config", None) - return dict(extra) if isinstance(extra, dict) else {} + block_size = getattr(cache_config, "block_size", None) + if not block_size: + raise ValueError( + "ATOM offload connector: vLLM reported no cache_config.block_size; " + "the codec addresses KV by block and cannot proceed without it" + ) + self.kv_cache_block_size = int(block_size) + self.kv_cache_dtype = _atom_kv_dtype(cache_config, model_config) + kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None) + role = getattr(kv_transfer_config, "kv_role", None) or "kv_both" + extra = getattr(kv_transfer_config, "kv_connector_extra_config", None) + self.kv_transfer_config = { + "kv_role": role, + **(dict(extra) if isinstance(extra, dict) else {}), + } -def build_offload_config(vllm_config: Any) -> OffloadConfigShim: - """Derive ATOM's offload config from ``VllmConfig``. + self.hf_config = getattr(model_config, "hf_config", None) + if self.hf_config is None: + raise ValueError( + "ATOM offload connector: vLLM reported no model_config.hf_config; " + "the LMCache namespace and layer count are derived from it" + ) - The role is normalized to ATOM's vocabulary: vLLM connectors are configured - with ``kv_role`` values like ``kv_both``, while ATOM's native path spells - the same intent ``offload``. ``_init_worker_common`` accepts both, so the - value is passed through unchanged rather than translated -- translating - would silently change save/load enablement if the two vocabularies drift. - """ - cache_config = getattr(vllm_config, "cache_config", None) - block_size = getattr(cache_config, "block_size", None) - if not block_size: - raise ValueError( - "ATOM offload connector: vLLM reported no cache_config.block_size; " - "the byte codec addresses KV by block and cannot proceed without it" + self.parallel_config = parallel_config + self.pipeline_parallel_size = int( + getattr(parallel_config, "pipeline_parallel_size", 1) or 1 ) + self.decode_context_parallel_size = int( + getattr(parallel_config, "decode_context_parallel_size", 1) or 1 + ) + # Feeds the LMCache page namespace, so two different models never share + # a key space. vLLM's served name is the closest analogue of ATOM's + # model_tag. + self.model = str(getattr(model_config, "model", "") or "atom-model") + self.model_tag = self.model - kv_transfer_config = getattr(vllm_config, "kv_transfer_config", None) - role = getattr(kv_transfer_config, "kv_role", None) or "kv_both" + def __repr__(self) -> str: + return ( + f"OffloadConfigShim(block_size={self.kv_cache_block_size}, " + f"kv_dtype={self.kv_cache_dtype!r}, role=" + f"{self.kv_transfer_config.get('kv_role')!r}, pp=" + f"{self.pipeline_parallel_size}, dcp={self.decode_context_parallel_size})" + ) - parallel_config = getattr(vllm_config, "parallel_config", None) - dcp = getattr(parallel_config, "decode_context_parallel_size", 1) or 1 - return OffloadConfigShim( - kv_cache_block_size=int(block_size), - kv_transfer_config={"kv_role": role, **_extra_config(kv_transfer_config)}, - decode_context_parallel_size=int(dcp), - ) +def build_offload_config(vllm_config: Any) -> OffloadConfigShim: + return OffloadConfigShim(vllm_config) diff --git a/atom/plugin/vllm/register.py b/atom/plugin/vllm/register.py index f0bcca5775..154535b99f 100644 --- a/atom/plugin/vllm/register.py +++ b/atom/plugin/vllm/register.py @@ -180,9 +180,18 @@ def register_platform() -> str | None: def _register_kv_connectors() -> None: """Expose ATOM's byte-level LMCache offload to vLLM's connector factory. + Convenience only. vLLM validates ``kv_transfer_config`` while building + VllmConfig, which happens BEFORE platform plugins are invoked, so a run + that names the connector by bare name fails config validation before this + ever runs. The supported way to select it is vLLM's out-of-tree entry + point, which takes priority over the registry and needs no registration: + + --kv-transfer-config '{"kv_connector": "AtomLMCacheOffloadConnector", + "kv_connector_module_path": "atom.plugin.vllm.kv_transfer.connector", + "kv_role": "kv_both"}' + Registered by module path so importing the plugin does not drag in the - offload stack (and LMCache) for every run -- vLLM resolves it lazily, only - when a --kv-transfer-config actually names it. + offload stack (and LMCache) for every run. """ from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory diff --git a/tests/plugin/test_vllm_offload_config.py b/tests/plugin/test_vllm_offload_config.py new file mode 100644 index 0000000000..6683182fb6 --- /dev/null +++ b/tests/plugin/test_vllm_offload_config.py @@ -0,0 +1,81 @@ +"""VllmConfig -> ATOM offload config projection. + +The offload path reads these fields off ATOM's Config; in plugin mode vLLM owns +them. Getting the KV dtype wrong is not a crash but a sizing error -- the codec +derives every byte segment from it -- so the mapping fails closed on anything +it does not recognise. +""" + +from __future__ import annotations + +from types import SimpleNamespace as NS + +import pytest + +from atom.plugin.vllm.kv_transfer.offload_config import build_offload_config + + +def _vllm_config(cache_dtype="fp8", model_dtype="bfloat16", block_size=128, **kw): + return NS( + cache_config=NS(block_size=block_size, cache_dtype=cache_dtype), + model_config=NS( + hf_config=NS(num_hidden_layers=60, model_type="minimax_m3"), + dtype=model_dtype, + model="/models/MiniMax-M3-MXFP4", + ), + parallel_config=NS( + pipeline_parallel_size=kw.get("pp", 1), + decode_context_parallel_size=kw.get("dcp", 1), + ), + kv_transfer_config=NS( + kv_role=kw.get("role", "kv_both"), + kv_connector_extra_config=kw.get("extra", {}), + ), + ) + + +def test_projects_the_fields_offload_reads(): + cfg = build_offload_config(_vllm_config()) + + assert cfg.kv_cache_block_size == 128 + assert cfg.kv_cache_dtype == "fp8" + assert cfg.hf_config.num_hidden_layers == 60 + assert cfg.kv_transfer_config["kv_role"] == "kv_both" + + +def test_auto_kv_dtype_resolves_to_the_model_dtype(): + cfg = build_offload_config(_vllm_config(cache_dtype="auto", model_dtype="bfloat16")) + assert cfg.kv_cache_dtype == "bf16" + + cfg = build_offload_config( + _vllm_config(cache_dtype="auto", model_dtype="torch.float16") + ) + assert cfg.kv_cache_dtype == "fp16" + + +@pytest.mark.parametrize("raw", ["fp8_e4m3", "fp8_e5m2", "fp8_inc"]) +def test_fp8_spellings_all_map_to_fp8(raw): + assert build_offload_config(_vllm_config(cache_dtype=raw)).kv_cache_dtype == "fp8" + + +def test_unknown_kv_dtype_fails_closed(): + # Guessing a width here would mis-size every byte segment silently. + with pytest.raises(ValueError, match="unsupported KV cache dtype"): + build_offload_config(_vllm_config(cache_dtype="int4_magic")) + + +def test_missing_block_size_is_rejected(): + with pytest.raises(ValueError, match="block_size"): + build_offload_config(_vllm_config(block_size=None)) + + +def test_missing_hf_config_is_rejected(): + bad = _vllm_config() + bad.model_config.hf_config = None + with pytest.raises(ValueError, match="hf_config"): + build_offload_config(bad) + + +def test_extra_connector_config_is_carried_through(): + cfg = build_offload_config(_vllm_config(extra={"lmcache.foo": 7})) + assert cfg.kv_transfer_config["lmcache.foo"] == 7 From 3972adf32212b019ceb4fcabc1a7f68992865fdb Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 08:53:54 -0500 Subject: [PATCH 04/16] fix(vllm-plugin): resolve nested HF configs for the offload namespace Second live run: 'MiniMaxM3Config object has no attribute num_hidden_layers'. Multimodal configs keep the transformer's own fields on a nested text config, and M3's outer config genuinely has no layer count. vLLM exposes the resolved inner one as model_config.hf_text_config while hf_config stays outer. ATOM's offload code reads both spellings -- config.py notes the skew itself: its namespace guard reads hf_config.model_type while is_qwen_next reads hf_text_config.model_type. A view that resolves inner-first and falls back to outer satisfies both without either side changing which attribute it asks for. (cherry picked from commit 07d031d243328ba43da567f4fe3db6a9cbd76daa) --- .../plugin/vllm/kv_transfer/offload_config.py | 40 ++++++++++++++++++- tests/plugin/test_vllm_offload_config.py | 34 ++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/kv_transfer/offload_config.py b/atom/plugin/vllm/kv_transfer/offload_config.py index e17954ecaf..69b4ad1062 100644 --- a/atom/plugin/vllm/kv_transfer/offload_config.py +++ b/atom/plugin/vllm/kv_transfer/offload_config.py @@ -46,6 +46,40 @@ def _atom_kv_dtype(cache_config: Any, model_config: Any) -> str: return key +class _HFConfigView: + """Read a model's HF config without caring whether it is nested. + + Multimodal configs (MiniMax-M3, Qwen3.5-VL, ...) keep the transformer's own + fields on a nested text config: ``MiniMaxM3Config`` has no + ``num_hidden_layers`` at all, it lives on ``.text_config``. vLLM exposes the + resolved inner config as ``model_config.hf_text_config`` while + ``hf_config`` stays the outer one, and ATOM's offload code reads both + spellings (``config.py`` notes this skew itself: its namespace guard reads + ``hf_config.model_type`` while ``is_qwen_next`` reads + ``hf_text_config.model_type``). + + Resolving inner-first and falling back to the outer config satisfies both + readers without asking either side to change which attribute it wants. + """ + + __slots__ = ("_inner", "_outer") + + def __init__(self, inner: Any, outer: Any) -> None: + self._inner = inner + self._outer = outer + + def __getattr__(self, name: str) -> Any: + if self._inner is not None: + try: + return getattr(self._inner, name) + except AttributeError: + pass + return getattr(self._outer, name) + + def __repr__(self) -> str: + return f"_HFConfigView(inner={type(self._inner).__name__}, outer={type(self._outer).__name__})" + + class OffloadConfigShim: """The ATOM-offload-shaped view of a vLLM config. @@ -76,7 +110,11 @@ def __init__(self, vllm_config: Any) -> None: **(dict(extra) if isinstance(extra, dict) else {}), } - self.hf_config = getattr(model_config, "hf_config", None) + outer_hf = getattr(model_config, "hf_config", None) + inner_hf = getattr(model_config, "hf_text_config", None) + self.hf_config = ( + _HFConfigView(inner_hf, outer_hf) if outer_hf is not None else None + ) if self.hf_config is None: raise ValueError( "ATOM offload connector: vLLM reported no model_config.hf_config; " diff --git a/tests/plugin/test_vllm_offload_config.py b/tests/plugin/test_vllm_offload_config.py index 6683182fb6..2d795458b9 100644 --- a/tests/plugin/test_vllm_offload_config.py +++ b/tests/plugin/test_vllm_offload_config.py @@ -79,3 +79,37 @@ def test_missing_hf_config_is_rejected(): def test_extra_connector_config_is_carried_through(): cfg = build_offload_config(_vllm_config(extra={"lmcache.foo": 7})) assert cfg.kv_transfer_config["lmcache.foo"] == 7 + + +def _nested_vllm_config(): + """A multimodal config: the transformer's fields live on the text config. + + MiniMaxM3Config genuinely has no num_hidden_layers -- reading it off the + outer config raises, which is how the first live run died. + """ + text = NS(num_hidden_layers=60, model_type="minimax_m3_text") + outer = NS(model_type="minimax_m3", text_config=text) # no num_hidden_layers + cfg = _vllm_config() + cfg.model_config.hf_config = outer + cfg.model_config.hf_text_config = text + return cfg + + +def test_nested_text_config_supplies_the_layer_count(): + cfg = build_offload_config(_nested_vllm_config()) + + assert cfg.hf_config.num_hidden_layers == 60 + + +def test_outer_only_fields_still_resolve(): + cfg = build_offload_config(_nested_vllm_config()) + cfg._vllm_config.model_config.hf_config.architectures = ["MiniMaxM3ForCausalLM"] + + # Inner-first, outer as fallback: both readers in ATOM's offload path are + # satisfied without either changing which attribute it asks for. + assert cfg.hf_config.architectures == ["MiniMaxM3ForCausalLM"] + + +def test_flat_config_is_unaffected(): + cfg = build_offload_config(_vllm_config()) + assert cfg.hf_config.num_hidden_layers == 60 From 17c8d2fc1c1a0bb2c57dc075b7051d34cdd3fb6d Mon Sep 17 00:00:00 2001 From: perzhang Date: Thu, 3 Sep 2026 09:35:09 -0500 Subject: [PATCH 05/16] fix(vllm-plugin): map ATOM's four completion sets onto vLLM's two First live request died with 'get_finished() takes 1 positional argument but 2 were given': vLLM passes finished_req_ids, ATOM's takes none and returns its own KVConnectorOutput. The mapping is not mechanical. finished_saving must surface as vLLM's finished_sending -- request_finished defers freeing while a save is in flight, and vLLM releases those blocks only when the id shows up there. ATOM's worker deliberately reports an EMPTY finished_sending because its own scheduler reads that as a P/D producer handoff and would deallocate live offload blocks. Two different contracts for one name, so the translation lives in the adapter rather than changing either side. Failed loads are reported as finished as well: the request is parked on that load and the alternative to waking it is a hang; vLLM then recomputes what it had counted as externally supplied. Completion ids come through as either a bare request id or one tagged with a generation, so req_id is extracted. (cherry picked from commit 639ae4b716405b97615c462087e4baf7e6c63346) --- atom/plugin/vllm/kv_transfer/connector.py | 31 ++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 8b66d081b2..33cdea28f4 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -129,7 +129,31 @@ def wait_for_save(self) -> None: """ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: - return self._worker.get_finished(finished_req_ids) + """Translate ATOM's four completion sets into vLLM's two. + + ``finished_saving`` HAS to surface as vLLM's ``finished_sending``: + ``request_finished`` defers freeing while a save is in flight, and vLLM + only releases those blocks once the id appears here. ATOM's own worker + deliberately reports an empty ``finished_sending`` because ITS scheduler + reads that as a P/D producer handoff -- a different contract from + vLLM's, so the mapping is done here rather than by changing ATOM. + + A failed load is reported as finished too: the request is parked + waiting on it, and the alternative to waking it is a hang. vLLM then + recomputes the tokens it had counted as externally supplied. + """ + out = self._worker.get_finished() + + finished_recving = {_req_id_of(c) for c in out.finished_loading} + failed = {_req_id_of(c) for c in out.failed_loading} + if failed: + logger.warning( + "ATOM LMCache offload: load failed for %s; recomputing", sorted(failed) + ) + finished_recving |= failed + + finished_sending = {_req_id_of(c) for c in out.finished_saving} + return finished_sending, finished_recving def shutdown(self) -> None: for side in (self._worker, self._scheduler): @@ -181,6 +205,11 @@ def request_finished(self, request, block_ids) -> tuple[bool, dict | None]: return False, None +def _req_id_of(completion_id) -> str: + """Completion ids are a bare request id, or one tagged with a generation.""" + return str(getattr(completion_id, "req_id", completion_id)) + + def _block_ids(blocks) -> list[int]: """Flatten vLLM's allocated-block structure to plain ids. From c6d5a6143de4152dd45ee7afce98134cdf5d22ff Mon Sep 17 00:00:00 2001 From: perzhang Date: Sun, 6 Sep 2026 03:29:24 -0500 Subject: [PATCH 06/16] fix(vllm-plugin): only report a save as finished_sending once the request is finished Previous run died on vLLM's 'assert request.is_finished()' inside _free_blocks. Reporting every finished_saving as finished_sending was too eager: vLLM treats that set as 'this request is done AND its KV has shipped, free its blocks', and asserts both req_id in self.requests and request.is_finished() before doing so. ATOM's saves are fire-and-forget and routinely land mid-decode. The two events arrive in either order -- a save can complete long before the request stops generating, and a request can finish with a save still in flight -- so each side is held until its counterpart shows up, and the pair is released together. Also adds a layout/dtype census at registration. The codec moves opaque bytes, so a mislabelled tensor never fails there; it surfaces much later inside an attention kernel ('Both operands must be same dtype') with nothing pointing back at registration. Confirmed with it that vLLM hands us 3 dense uint8, 57 sparse uint8 and 57 index float8_e4m3fn -- the index caches do carry a real fp8 dtype, which this path depends on: `fix/m3-vllm-plugin-index-cache-fp8-dtype` is what makes them so. --- atom/plugin/vllm/kv_transfer/connector.py | 49 +++++++++++++++++++---- 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 33cdea28f4..4ec10c6363 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -69,6 +69,11 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): self._scheduler = None self._seqs = SeqViewRegistry() + # finished_sending is only legal for a request vLLM has already + # finished AND whose save has landed; the two events arrive in either + # order, so both sides are accumulated until they meet. + self._saved_awaiting_finish: set[str] = set() + self._finished_awaiting_save: set[str] = set() if role == KVConnectorRole.WORKER: from atom.kv_transfer.offload.dense.connector import DenseOffloadConnector @@ -102,6 +107,26 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: len(tensors), num_blocks, ) + # Layout/dtype census. The codec moves opaque bytes, so a wrong dtype + # never surfaces here -- it surfaces much later inside an attention + # kernel ("Both operands must be same dtype"), with nothing pointing + # back at registration. One line here makes that diagnosable. + census: dict[tuple, int] = {} + for name, tensor in sorted(kv_caches.items()): + key = ( + "index" if name.endswith(".index_cache") else "kv", + tuple(tensor.shape[1:]), + str(tensor.dtype), + ) + census[key] = census.get(key, 0) + 1 + for (kind, shape, dtype), count in sorted(census.items(), key=str): + logger.info( + "ATOM LMCache offload: %d x %s tail_shape=%s dtype=%s", + count, + kind, + shape, + dtype, + ) def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: metadata = self._get_connector_metadata() @@ -131,12 +156,17 @@ def wait_for_save(self) -> None: def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: """Translate ATOM's four completion sets into vLLM's two. - ``finished_saving`` HAS to surface as vLLM's ``finished_sending``: - ``request_finished`` defers freeing while a save is in flight, and vLLM - only releases those blocks once the id appears here. ATOM's own worker - deliberately reports an empty ``finished_sending`` because ITS scheduler - reads that as a P/D producer handoff -- a different contract from - vLLM's, so the mapping is done here rather than by changing ATOM. + ``finished_saving`` surfaces as vLLM's ``finished_sending``, but only + once the request has ALSO finished. vLLM's scheduler asserts both + ``req_id in self.requests`` and ``request.is_finished()`` before + freeing, while ATOM's saves are fire-and-forget and routinely land + while the request is still decoding -- reporting those crashed the + engine on ``assert request.is_finished()``. The two events arrive in + either order, so each side is held until its counterpart shows up. + + ATOM's own worker deliberately reports an empty ``finished_sending`` + because ITS scheduler reads that as a P/D producer handoff. Same name, + two contracts; the translation lives here rather than in either side. A failed load is reported as finished too: the request is parked waiting on it, and the alternative to waking it is a hang. vLLM then @@ -152,7 +182,12 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: ) finished_recving |= failed - finished_sending = {_req_id_of(c) for c in out.finished_saving} + self._saved_awaiting_finish |= {_req_id_of(c) for c in out.finished_saving} + self._finished_awaiting_save |= set(finished_req_ids or ()) + finished_sending = self._saved_awaiting_finish & self._finished_awaiting_save + self._saved_awaiting_finish -= finished_sending + self._finished_awaiting_save -= finished_sending + return finished_sending, finished_recving def shutdown(self) -> None: From 7be17667df6a239dd6f288b0a72cea7dcd565f7e Mon Sep 17 00:00:00 2001 From: perzhang Date: Sat, 5 Sep 2026 10:40:41 -0500 Subject: [PATCH 07/16] fix(vllm-plugin): carry M3's fp8 KV scales through the offload transfer Every warm hit came back as garbage. The metrics all said success -- external_prefix_cache_hits_total counted 10496 of 10539 prompt tokens restored, LMCache reported the bytes stored and retrieved -- and the model answered with '564, tuple["torch.Tensor | None", "torch.Tensor | None"]: + """The fp8 scales a KV transfer must carry alongside this layer's bytes. + + The sparse cache stores fp8 mantissas whose scale is per token AND per + head -- `(num_blocks, num_kv_heads, block_size)`, one fp32 per element + of the paged cache. Move the mantissas without them and a restored + block is dequantised against whatever the block's previous occupant + left behind: fluent-looking garbage, no error anywhere. So any tier + that moves this layer's KV has to move these too, and it can only know + that by asking -- vLLM's `kv_caches` registration carries the KV + tensors alone, and these live on the layer. + + Named after the native path's `get_kv_transfer_tensors`, which reports + the same regions off `runner.kv_scale`. + + `kv_cache` overrides the layer's own tensor, for callers that hold it + before the layer does -- a connector registering at engine start runs + before the first forward, and the scales are allocated lazily. Passing + the tensor vLLM registered is also what keeps the allocation stable: + `_ensure_fp8_scales` reallocates on a shape or device change, which + would strand a pointer a tier had already registered. + + Returns `(None, None)` when the cache is not fp8 -- nothing to carry. + """ + cache = self.kv_cache if kv_cache is None else kv_cache + if self.kv_cache_dtype != "fp8": + return None, None + if cache is None or cache.numel() == 0: + raise RuntimeError( + f"{self.layer_name}: cannot size the fp8 KV scales before the " + "KV cache is allocated" + ) + return self._ensure_fp8_scales(cache) + def _page16_shuffle_cache_for_sparse_kernel( self, ) -> tuple[torch.Tensor, torch.Tensor, object, object]: diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 4ec10c6363..6d815195a1 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -64,6 +64,7 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): # rejects the 2-argument signature outright, and the base class stores # it for the group-aware paths. super().__init__(vllm_config, role, kv_cache_config) + self._vllm_config = vllm_config self._config = build_offload_config(vllm_config) self._worker = None self._scheduler = None @@ -87,8 +88,14 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): # ---- worker side -------------------------------------------------- def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: - """Translate vLLM's flat registration and hand it to ATOM's codec.""" - tensors = build_kv_cache_tensors(kv_caches) + """Translate vLLM's flat registration and hand it to ATOM's codec. + + The layer modules go along with the tensors: M3 keeps its fp8 KV scales + (one fp32 per token per head) on the layer, not in this dict, and a + transfer that moves the mantissas without them silently dequantises a + restored block against the previous occupant's scale. + """ + tensors = build_kv_cache_tensors(kv_caches, self._attention_layers()) if not tensors: raise ValueError("ATOM offload connector: vLLM registered no KV caches") @@ -128,6 +135,18 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: dtype, ) + def _attention_layers(self) -> dict[str, Any]: + """The layer modules behind vLLM's registered KV cache names. + + vLLM keeps them in the static forward context, which is where its own + attention-metadata builders read layers from; there is no per-layer + handle in the connector API itself. + """ + context = getattr( + self._vllm_config.compilation_config, "static_forward_context", None + ) + return dict(context) if context else {} + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: metadata = self._get_connector_metadata() inner = getattr(metadata, "inner", None) diff --git a/atom/plugin/vllm/kv_transfer/kv_cache_layout.py b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py index 1e62171e5d..7a2825c34e 100644 --- a/atom/plugin/vllm/kv_transfer/kv_cache_layout.py +++ b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py @@ -24,8 +24,17 @@ The codec never interprets these bytes, so "which tensor is K" only has to be consistent between save and restore -- not semantically right. + +What vLLM does NOT hand over is quantisation state. M3's sparse layers keep an +fp32 scale per token per head next to the fp8 cache, on the layer object; the +registration dict has only the caches. Restoring mantissas against the previous +occupant's scales is silent corruption, so the layers are asked for their scales +here and a layer that clearly owns per-block scales without a way to report them +is a hard error rather than a quiet omission. """ +from typing import Any + import torch from atom.config import KVCacheTensor @@ -58,8 +67,41 @@ def split_kv_tensor(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | return tensor, None +def _transfer_scales( + name: str, layer: Any, tensor: torch.Tensor +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Ask one layer for the scales that must travel with its KV bytes. + + Duck-typed on purpose: only the layers that HAVE movable scales implement + the hook, and the connector must not need a table of which model does. + + The fallback is not "assume none". A layer holding a multi-element + `k_scale`/`v_scale` has per-block quantisation state, and moving its KV + without that state restores mantissas against a stale scale -- wrong output + with nothing logged. Without the hook there is no way to move it, so say so + here rather than at the far end of a wrong answer. Scalar scales (vLLM's + usual per-tensor `_k_scale`) are constant and correctly ignored. + """ + getter = getattr(layer, "get_kv_transfer_scales", None) + if getter is not None: + k_scale, v_scale = getter(tensor) + return k_scale, v_scale + + for attr in ("k_scale", "v_scale", "_k_scale", "_v_scale"): + scale = getattr(layer, attr, None) + if isinstance(scale, torch.Tensor) and scale.numel() > 1: + raise ValueError( + f"{name}: layer holds a per-block {attr} " + f"(shape={tuple(scale.shape)}) but has no " + "get_kv_transfer_scales(); moving its KV without the scales " + "would restore mantissas against the previous occupant's scale" + ) + return None, None + + def build_kv_cache_tensors( kv_caches: dict[str, torch.Tensor], + layers: dict[str, Any] | None = None, ) -> list[KVCacheTensor]: """Translate vLLM's ``{layer_name: tensor}`` into ATOM ``KVCacheTensor``s. @@ -68,9 +110,16 @@ def build_kv_cache_tensors( registers M3's DSA indexer keys as separate entries, but they are part of the same layer's movable bytes. + Args: + kv_caches: vLLM's registration dict. + layers: the layer modules by name, for the quantisation state that does + not travel in ``kv_caches``. Omit only when no registered layer has + per-block scales. + Raises: ValueError: a produced segment is not contiguous (the codec would - reject it later, with far less context about which layer). + reject it later, with far less context about which layer), or a + layer owns per-block scales it cannot report. """ index_caches: dict[str, torch.Tensor] = {} main: dict[str, torch.Tensor] = {} @@ -90,10 +139,15 @@ def build_kv_cache_tensors( for layer_num, name in enumerate(sorted(main, key=_layer_sort_key)): k_cache, v_cache = split_kv_tensor(main[name]) index_cache = index_caches.get(name) + k_scale, v_scale = _transfer_scales( + name, (layers or {}).get(name), main[name] + ) for role, seg in ( ("k_cache", k_cache), ("v_cache", v_cache), + ("k_scale", k_scale), + ("v_scale", v_scale), ("index_cache", index_cache), ): if seg is not None and not seg.is_contiguous(): @@ -103,11 +157,25 @@ def build_kv_cache_tensors( "the byte codec can only move contiguous segments" ) + # The codec derives each segment's per-block stride as + # numel // num_blocks, so a scale whose leading axis is not the block + # axis would be sliced at the wrong granularity -- and, being the same + # dtype and roughly the right size, would not fail any later check. + num_blocks = int(k_cache.shape[0]) + for role, seg in (("k_scale", k_scale), ("v_scale", v_scale)): + if seg is not None and int(seg.shape[0]) != num_blocks: + raise ValueError( + f"{name}: {role} is not block-major " + f"(shape={tuple(seg.shape)}, num_blocks={num_blocks})" + ) + out.append( KVCacheTensor( layer_num=layer_num, k_cache=k_cache, v_cache=v_cache if v_cache is not None else torch.tensor([]), + k_scale=k_scale, + v_scale=v_scale, index_cache=index_cache, ) ) diff --git a/tests/plugin/test_vllm_kv_cache_layout.py b/tests/plugin/test_vllm_kv_cache_layout.py index a69f3d3de1..55866bfdab 100644 --- a/tests/plugin/test_vllm_kv_cache_layout.py +++ b/tests/plugin/test_vllm_kv_cache_layout.py @@ -34,6 +34,31 @@ def _sparse_kv(nb: int = NB) -> torch.Tensor: return buf.as_strided((nb, 2, BS, 1, HD), (k_block, k_total, HD, HD, 1)) +class _SparseLayer: + """A stand-in for M3's sparse attention: fp8 scales live on the layer.""" + + def __init__(self, nb: int = NB, heads: int = 1) -> None: + self.kv_scale = torch.zeros((2, nb, heads, BS), dtype=torch.float32) + + def get_kv_transfer_scales(self, kv_cache=None): + return self.kv_scale[0], self.kv_scale[1] + + +class _LayerWithUnreportableScales: + """Per-block scales and no way to report them -- must not pass silently.""" + + def __init__(self) -> None: + self.k_scale = torch.zeros((NB, 1, BS), dtype=torch.float32) + + +def _m3_layers(kv: dict[str, torch.Tensor]) -> dict[str, object]: + return { + name: _SparseLayer() + for name, tensor in kv.items() + if tensor.ndim == 5 # sparse layers are the fp8-scaled ones + } + + def _m3_registration() -> dict[str, torch.Tensor]: kv: dict[str, torch.Tensor] = {} for i in range(DENSE_LAYERS): # K/V interleaved per token @@ -93,6 +118,63 @@ def test_orphan_index_cache_is_rejected(): ) +def test_fp8_scales_travel_with_the_layer_they_belong_to(): + """Mantissas without their scales restore as fluent garbage, silently. + + M3's sparse cache scales fp8 per token AND per head, and vLLM's + registration dict does not carry that table -- the layer does. This is the + regression that made every warm hit produce garbage while every metric said + the transfer had succeeded. + """ + kv = _m3_registration() + tensors = build_kv_cache_tensors(kv, _m3_layers(kv)) + + sparse = [t for t in tensors if t.index_cache is not None] + assert len(sparse) == SPARSE_LAYERS + for t in sparse: + assert t.k_scale is not None and t.v_scale is not None + assert t.k_scale.shape[0] == NB, "scales must be block-major" + assert t.k_scale.is_contiguous() and t.v_scale.is_contiguous() + + dense = [t for t in tensors if t.index_cache is None] + assert all(t.k_scale is None for t in dense), "dense scales are per-tensor" + + +def test_scales_reach_the_codec_as_extra_segments(): + codec_mod = pytest.importorskip( + "atom.kv_transfer.offload.dense.kv_byte_codec", + reason="offload codec pulls aiter", + ) + kv = _m3_registration() + with_scales = build_kv_cache_tensors(kv, _m3_layers(kv)) + without = build_kv_cache_tensors(kv) + + def _bytes(tensors): + return codec_mod.DenseKVByteCodec( + {str(t.layer_num): t for t in tensors}, num_blocks=NB + ).bytes_per_block + + # 2 scales x 1 head x BS tokens x fp32, per sparse layer. + assert _bytes(with_scales) - _bytes(without) == SPARSE_LAYERS * 2 * BS * 4 + + +def test_per_block_scales_without_a_reporting_hook_are_rejected(): + kv = {"model.layers.0.self_attn.attn": torch.zeros((NB, 1, BS, 2 * HD))} + with pytest.raises(ValueError, match="per-block k_scale"): + build_kv_cache_tensors(kv, {"model.layers.0.self_attn.attn": _LayerWithUnreportableScales()}) + + +def test_non_block_major_scales_are_rejected(): + class _HeadMajor: + def get_kv_transfer_scales(self, kv_cache=None): + scale = torch.zeros((1, NB, BS), dtype=torch.float32) + return scale, scale + + kv = {"model.layers.0.self_attn.attn": torch.zeros((NB, 1, BS, 2 * HD))} + with pytest.raises(ValueError, match="not block-major"): + build_kv_cache_tensors(kv, {"model.layers.0.self_attn.attn": _HeadMajor()}) + + def test_codec_accepts_the_mapped_tensors(): """The mapping's whole purpose: make M3 pass ATOM's byte codec.""" codec_mod = pytest.importorskip( From ef5d38431ec5c930a79ad016803cf50e7201b836 Mon Sep 17 00:00:00 2001 From: perzhang Date: Sat, 5 Sep 2026 10:42:37 -0500 Subject: [PATCH 08/16] fix(vllm-plugin): return the worker's completions to the scheduler half vLLM runs a connector as two objects in two processes. ATOM's completion objects only ever reach the worker half; the scheduler half hears about them through update_connector_output, which this adapter did not implement. So on the scheduler side nothing ever finished: _save_inflight and the load lifecycle grew for the life of the process, has_pending_work() never went quiet, and every request whose save was still in flight at finish kept its SeqView -- which holds the request's prompt token ids. At M3's context lengths that is most of a megabyte retained per request, for the life of the server. vLLM's KVConnectorOutput carries request ids as plain strings, and both save_finished and load_finished deliberately refuse a raw id once the lifecycle has an exact operation identity, so that a delayed report cannot complete a newer lifecycle. Rather than weaken that guard for everyone, DenseOffloadScheduler grows save_finished_by_request / load_finished_by_request: they resolve the identity ATOM already parked and delegate. The guard keeps its meaning and the plugin path gets a way to speak. (cherry picked from commit 0b32c55e35d295b13490048343a2397b882dc1a5) --- atom/kv_transfer/offload/dense/connector.py | 27 +++++++ atom/plugin/vllm/kv_transfer/connector.py | 24 ++++++ tests/plugin/test_vllm_offload_completions.py | 81 +++++++++++++++++++ 3 files changed, 132 insertions(+) create mode 100644 tests/plugin/test_vllm_offload_completions.py diff --git a/atom/kv_transfer/offload/dense/connector.py b/atom/kv_transfer/offload/dense/connector.py index e6d7d5e183..7c18edaae6 100644 --- a/atom/kv_transfer/offload/dense/connector.py +++ b/atom/kv_transfer/offload/dense/connector.py @@ -742,6 +742,33 @@ def save_finished(self, req_id) -> None: self._save_inflight.pop(sid, None) self._finish_save_statistics(req_id) + def save_finished_by_request(self, req_id) -> None: + """Complete a save when only the plain request id is available. + + `save_finished` refuses a raw id once the lifecycle has an exact + `SaveOperationId`, so a delayed report cannot complete a newer + lifecycle. A vLLM-plugin scheduler cannot satisfy that: vLLM's + `KVConnectorOutput` carries request ids as plain strings, so the exact + identity never survives the trip back from the worker. + + Resolving the parked identity here keeps the guard meaningful instead of + weakening `save_finished` -- and without it the entry never clears, so + `_save_inflight` grows for the life of the process and + `has_pending_work()` never goes quiet. + """ + sid = str(req_id) + active = self._save_inflight.get(sid) + self.save_finished(active if active is not None else sid) + + def load_finished_by_request(self, req_id) -> bool: + """`load_finished` for a caller that has only the plain request id. + + Same reason as `save_finished_by_request`. + """ + sid = str(req_id) + entry = self._active_load_operations.get(sid) + return self.load_finished(entry[1] if entry is not None else sid) + def abandon_save(self, req_id) -> None: """Force-drop a save the scheduler reclaimed after it stalled. diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 6d815195a1..0499744dcf 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -248,6 +248,30 @@ def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata: seq.set_num_cached_tokens(num_tokens) return AtomOffloadMetadata(self._scheduler.build_connector_meta()) + def update_connector_output(self, connector_output) -> None: + """Feed the worker's completions back into ATOM's scheduler state. + + vLLM splits a connector across two processes and only the worker half + sees ATOM's completion objects; this is the scheduler half's only news + of them. Without it nothing ever clears: `_save_inflight` and the load + lifecycle grow for the life of the process, `has_pending_work()` never + goes quiet, and -- the one that actually hurts -- the SeqView of every + deferred request is retained, each holding that request's prompt token + ids. At M3's context lengths that is the difference between a bounded + server and one that grows by most of a megabyte per request. + + The ids arrive as plain strings (that is all vLLM's KVConnectorOutput + carries), so the `*_by_request` resolvers recover the exact operation + identity ATOM parked. + """ + for req_id in connector_output.finished_recving or (): + self._scheduler.load_finished_by_request(req_id) + for req_id in connector_output.finished_sending or (): + self._scheduler.save_finished_by_request(req_id) + # vLLM frees the blocks on this same report, so the request is over + # on both sides; the view was kept only for the deferred save. + self._seqs.drop(req_id) + def request_finished(self, request, block_ids) -> tuple[bool, dict | None]: seq = self._seqs.get(request.request_id) if seq is not None: diff --git a/tests/plugin/test_vllm_offload_completions.py b/tests/plugin/test_vllm_offload_completions.py new file mode 100644 index 0000000000..f03cd2b984 --- /dev/null +++ b/tests/plugin/test_vllm_offload_completions.py @@ -0,0 +1,81 @@ +"""The scheduler half's only news of what the worker finished. + +vLLM runs a connector as two objects in two processes: the worker sees ATOM's +completion objects, the scheduler sees `update_connector_output` and a set of +plain request-id strings. Miss that hook and nothing on the scheduler side ever +clears -- including the SeqView of every deferred request, each pinning that +request's prompt token ids. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from atom.plugin.vllm.kv_transfer.seq_view import SeqViewRegistry + +connector_mod = pytest.importorskip( + "atom.plugin.vllm.kv_transfer.connector", + reason="the adapter imports vLLM's connector base", +) + + +class _FakeScheduler: + """Records the resolver calls; those are the contract under test.""" + + def __init__(self) -> None: + self.saves: list[str] = [] + self.loads: list[str] = [] + + def save_finished_by_request(self, req_id) -> None: + self.saves.append(str(req_id)) + + def load_finished_by_request(self, req_id) -> bool: + self.loads.append(str(req_id)) + return True + + +def _adapter() -> tuple[object, _FakeScheduler]: + # Built without __init__: constructing it for real needs a VllmConfig and + # would pull the whole offload stack in, which is not what this covers. + adapter = object.__new__(connector_mod.AtomLMCacheOffloadConnector) + scheduler = _FakeScheduler() + adapter._scheduler = scheduler + adapter._seqs = SeqViewRegistry() + return adapter, scheduler + + +def _output(sending=(), recving=()): + return SimpleNamespace(finished_sending=set(sending), finished_recving=set(recving)) + + +def test_completions_reach_atoms_scheduler(): + adapter, scheduler = _adapter() + + adapter.update_connector_output(_output(sending=["a"], recving=["b"])) + + assert scheduler.saves == ["a"] + assert scheduler.loads == ["b"] + + +def test_a_finished_save_releases_the_seq_view(): + """The leak that matters: a deferred view holds the prompt token ids.""" + adapter, _ = _adapter() + request = SimpleNamespace(request_id="a", prompt_token_ids=[1, 2, 3]) + adapter._seqs.get_or_create(request) + assert len(adapter._seqs) == 1 + + adapter.update_connector_output(_output(sending=["a"])) + + assert len(adapter._seqs) == 0 + + +def test_empty_output_is_harmless(): + adapter, scheduler = _adapter() + + adapter.update_connector_output( + SimpleNamespace(finished_sending=None, finished_recving=None) + ) + + assert scheduler.saves == [] and scheduler.loads == [] From a1f8787dd6aa2021851a555f464e84c7400f6b1c Mon Sep 17 00:00:00 2001 From: perzhang Date: Sat, 5 Sep 2026 11:02:04 -0500 Subject: [PATCH 09/16] fix(vllm-plugin): only promise a load ATOM has agreed to issue Second gsm8k pass hung the server outright: 200 requests parked in WAITING_FOR_REMOTE_KVS, EngineCore spinning in schedule() at 100% CPU with every GPU at 0%, no log line for four minutes, client stuck at 0/200. A lookup hit is not a decision to load. ATOM weighs that separately and drops a hit that HBM already covers, one whose boundary is not chunk aligned, or one below its transfer floor -- its native scheduler asks should_park_for_load_after_alloc before parking anything. The adapter instead promised async=True on the strength of the lookup alone, and vLLM has no second chance: the only thing that releases a parked request is the worker reporting it in finished_recving, and a load that is never issued is never reported. Both routine drops fire constantly in exactly the configurations we want: - the floor defaults to 8192 tokens (OFFLOAD_MIN_LOAD_TOKENS), so every chat-sized prompt is dropped -- gsm8k's 2560 was a guaranteed hang; - with HBM prefix caching on, vLLM's frontier is block aligned (128) while LMCache chunks are 256, so roughly every other hit is unaligned. So the adapter asks before it promises, and reports no external tokens when the answer is no. ATOM clears its own pending-load state on that path, so the request just prefills normally. Not caught earlier because the round-trip test uses 10K-token prompts, which clear the 8192 floor. The tests here use a scheduler that reports a hit and declines to load it -- the exact shape that deadlocked. (cherry picked from commit a636c26b014fcced3c8465f95f18ea36f528f91f) --- atom/plugin/vllm/kv_transfer/connector.py | 27 ++++++++- tests/plugin/test_vllm_offload_completions.py | 57 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 0499744dcf..c540a9d7e7 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -225,10 +225,35 @@ def get_num_new_matched_tokens( ``num_computed_tokens`` is vLLM's HBM-prefix-cache frontier; ATOM reads the same quantity off the seq to avoid re-loading what is already resident, so it has to be pushed in before the lookup runs. + + A lookup hit is not yet a decision to load. ATOM weighs that separately + -- a hit already covered by HBM, one whose boundary is not chunk + aligned, or one too small to be worth a transfer is dropped -- and its + native scheduler asks `should_park_for_load_after_alloc` before parking + anything. + + vLLM has no such second chance: returning True here parks the request in + WAITING_FOR_REMOTE_KVS, and the ONLY thing that releases it is the + worker reporting the id in `finished_recving`. A load that is never + issued is never reported, so the request waits forever -- the engine + spins in `schedule()` with every GPU idle and the server is dead. Two of + the three drop reasons are routine: the default floor is 8192 tokens, + which is every prompt in a chat-sized workload, and with HBM prefix + caching on, vLLM's block-aligned frontier (128) is regularly not + chunk-aligned (256). + + So the decision is taken here, before the promise. When ATOM declines it + has already cleared its pending-load state, and reporting no external + tokens leaves the request to prefill normally. """ seq = self._seqs.get_or_create(request) seq.set_num_cached_tokens(num_computed_tokens) - return self._scheduler.get_num_new_matched_tokens(seq) + need, _ = self._scheduler.get_num_new_matched_tokens(seq) + if need <= 0: + return 0, False + if not self._scheduler.should_park_for_load_after_alloc(seq): + return 0, False + return need, True def update_state_after_alloc(self, request, blocks, num_external_tokens: int): seq = self._seqs.get_or_create(request) diff --git a/tests/plugin/test_vllm_offload_completions.py b/tests/plugin/test_vllm_offload_completions.py index f03cd2b984..47a1bdf3f1 100644 --- a/tests/plugin/test_vllm_offload_completions.py +++ b/tests/plugin/test_vllm_offload_completions.py @@ -79,3 +79,60 @@ def test_empty_output_is_harmless(): ) assert scheduler.saves == [] and scheduler.loads == [] + + +class _ParkScheduler: + """A scheduler that reports a hit and then declines to load it.""" + + def __init__(self, hit: int, park: bool) -> None: + self._hit = hit + self._park = park + self.asked_park = 0 + + def get_num_new_matched_tokens(self, seq): + return self._hit, True + + def should_park_for_load_after_alloc(self, seq) -> bool: + self.asked_park += 1 + return self._park + + +def _lookup_adapter(scheduler): + adapter = object.__new__(connector_mod.AtomLMCacheOffloadConnector) + adapter._scheduler = scheduler + adapter._seqs = SeqViewRegistry() + return adapter + + +def _req(rid="r1", prompt_len=4096): + return SimpleNamespace(request_id=rid, prompt_token_ids=list(range(prompt_len))) + + +def test_a_hit_atom_will_not_load_is_not_promised(): + """The deadlock: vLLM parks on async=True and only the worker can release. + + ATOM drops a hit that is below its transfer floor or not chunk aligned. If + the promise has already been made, nothing ever reports the load, the + request sits in WAITING_FOR_REMOTE_KVS forever and the engine spins with + every GPU idle. + """ + scheduler = _ParkScheduler(hit=2560, park=False) + adapter = _lookup_adapter(scheduler) + + assert adapter.get_num_new_matched_tokens(_req(), 0) == (0, False) + assert scheduler.asked_park == 1 + + +def test_a_hit_atom_will_load_is_promised_async(): + scheduler = _ParkScheduler(hit=10496, park=True) + adapter = _lookup_adapter(scheduler) + + assert adapter.get_num_new_matched_tokens(_req(), 0) == (10496, True) + + +def test_no_hit_does_not_ask_about_parking(): + scheduler = _ParkScheduler(hit=0, park=True) + adapter = _lookup_adapter(scheduler) + + assert adapter.get_num_new_matched_tokens(_req(), 0) == (0, False) + assert scheduler.asked_park == 0 From d9eef3e1095fefdc3a327a427c4677b01177b304 Mon Sep 17 00:00:00 2001 From: perzhang Date: Sat, 5 Sep 2026 11:07:48 -0500 Subject: [PATCH 10/16] feat(vllm-plugin): name a promised load that never gets dispatched The park deadlock has no symptom of its own: EngineCore spins in schedule(), every GPU sits at 0%, the stats logger goes quiet because nothing is running, and the last log line is whatever happened before the stall. Tracing it back to a connector promise took a py-spy dump of the engine and a read through vLLM's waiting-queue code. The promise is now gated on ATOM's own park decision, so this should stay empty. If it does not, one ERROR line names the requests instead of leaving a silent hang. Nothing here can rescue them -- only the worker's finished_recving releases a parked request, and this is the scheduler half -- so it reports and moves on rather than pretending to recover. Reported once per request, not once per step: a hung engine steps very fast. (cherry picked from commit 2fef7b963ec666ee91f8f734dbf39edf41c7936a) --- atom/plugin/vllm/kv_transfer/connector.py | 46 ++++++++++++++++++- tests/plugin/test_vllm_offload_completions.py | 30 ++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index c540a9d7e7..4111026cd7 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -75,6 +75,10 @@ def __init__(self, vllm_config, role: KVConnectorRole, kv_cache_config=None): # order, so both sides are accumulated until they meet. self._saved_awaiting_finish: set[str] = set() self._finished_awaiting_save: set[str] = set() + # Requests parked on a promised load, and how many steps ago. A promise + # that never turns into a dispatched load is an unrecoverable hang, so + # it is at least named in the log. See `_check_promised_loads`. + self._promised_loads: dict[str, int] = {} if role == KVConnectorRole.WORKER: from atom.kv_transfer.offload.dense.connector import DenseOffloadConnector @@ -253,6 +257,7 @@ def get_num_new_matched_tokens( return 0, False if not self._scheduler.should_park_for_load_after_alloc(seq): return 0, False + self._promised_loads[request.request_id] = 0 return need, True def update_state_after_alloc(self, request, blocks, num_external_tokens: int): @@ -271,7 +276,44 @@ def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata: seq = self._seqs.get(req_id) if seq is not None: seq.set_num_cached_tokens(num_tokens) - return AtomOffloadMetadata(self._scheduler.build_connector_meta()) + inner = self._scheduler.build_connector_meta() + self._check_promised_loads(inner) + return AtomOffloadMetadata(inner) + + # Steps a promised load may go undispatched before it is called out. Loads + # are emitted on the step after the promise, so anything past a handful of + # steps is already wrong; the margin is only so a busy scheduler does not + # produce noise. + _PROMISE_GRACE_STEPS = 50 + + def _check_promised_loads(self, inner) -> None: + """Name any request parked on a load that was never dispatched. + + Nothing here can rescue it: vLLM releases a parked request only when the + worker reports the id in `finished_recving`, and the scheduler half + cannot inject that. The promise is gated on ATOM's own park decision, so + this should stay empty -- but when it does not, the symptom is an engine + spinning in `schedule()` with every GPU idle and no log line at all, + which costs hours to trace back. One line here names the request. + """ + if not self._promised_loads: + return + for meta in getattr(inner, "requests", ()) or (): + self._promised_loads.pop(str(getattr(meta, "req_id", meta)), None) + stuck = [] + for req_id in list(self._promised_loads): + self._promised_loads[req_id] += 1 + if self._promised_loads[req_id] > self._PROMISE_GRACE_STEPS: + stuck.append(req_id) + del self._promised_loads[req_id] + if stuck: + logger.error( + "ATOM LMCache offload: promised a load for %s but none was " + "dispatched within %d steps; those requests are parked in " + "WAITING_FOR_REMOTE_KVS and cannot be released", + sorted(stuck), + self._PROMISE_GRACE_STEPS, + ) def update_connector_output(self, connector_output) -> None: """Feed the worker's completions back into ATOM's scheduler state. @@ -290,6 +332,7 @@ def update_connector_output(self, connector_output) -> None: identity ATOM parked. """ for req_id in connector_output.finished_recving or (): + self._promised_loads.pop(str(req_id), None) self._scheduler.load_finished_by_request(req_id) for req_id in connector_output.finished_sending or (): self._scheduler.save_finished_by_request(req_id) @@ -298,6 +341,7 @@ def update_connector_output(self, connector_output) -> None: self._seqs.drop(req_id) def request_finished(self, request, block_ids) -> tuple[bool, dict | None]: + self._promised_loads.pop(request.request_id, None) seq = self._seqs.get(request.request_id) if seq is not None: self._scheduler.request_finished(seq) diff --git a/tests/plugin/test_vllm_offload_completions.py b/tests/plugin/test_vllm_offload_completions.py index 47a1bdf3f1..05ba2314fa 100644 --- a/tests/plugin/test_vllm_offload_completions.py +++ b/tests/plugin/test_vllm_offload_completions.py @@ -43,6 +43,7 @@ def _adapter() -> tuple[object, _FakeScheduler]: scheduler = _FakeScheduler() adapter._scheduler = scheduler adapter._seqs = SeqViewRegistry() + adapter._promised_loads = {} return adapter, scheduler @@ -101,6 +102,7 @@ def _lookup_adapter(scheduler): adapter = object.__new__(connector_mod.AtomLMCacheOffloadConnector) adapter._scheduler = scheduler adapter._seqs = SeqViewRegistry() + adapter._promised_loads = {} return adapter @@ -136,3 +138,31 @@ def test_no_hit_does_not_ask_about_parking(): assert adapter.get_num_new_matched_tokens(_req(), 0) == (0, False) assert scheduler.asked_park == 0 + + +def test_a_promise_that_never_dispatches_is_named(caplog): + """The hang leaves no trace of its own; this is the only breadcrumb.""" + scheduler = _ParkScheduler(hit=10496, park=True) + adapter = _lookup_adapter(scheduler) + adapter.get_num_new_matched_tokens(_req("stuck"), 0) + + empty = SimpleNamespace(requests=[]) + with caplog.at_level("ERROR", logger="atom"): + for _ in range(adapter._PROMISE_GRACE_STEPS + 1): + adapter._check_promised_loads(empty) + + assert "stuck" in caplog.text + # Reported once, not every step afterwards. + caplog.clear() + adapter._check_promised_loads(empty) + assert caplog.text == "" + + +def test_a_dispatched_load_is_not_reported(): + scheduler = _ParkScheduler(hit=10496, park=True) + adapter = _lookup_adapter(scheduler) + adapter.get_num_new_matched_tokens(_req("ok"), 0) + + adapter._check_promised_loads(SimpleNamespace(requests=[SimpleNamespace(req_id="ok")])) + + assert adapter._promised_loads == {} From 269d5f47827018d0545e5480c4889d9173bee2db Mon Sep 17 00:00:00 2001 From: perzhang Date: Sun, 6 Sep 2026 21:15:03 -0500 Subject: [PATCH 11/16] style: satisfy Black on the vLLM offload adapter The Black CI gate reformats three lines in the new files: a call that now fits on one line, and two that no longer do. Co-Authored-By: Claude Opus 5 (1M context) --- atom/plugin/vllm/kv_transfer/kv_cache_layout.py | 4 +--- tests/plugin/test_vllm_kv_cache_layout.py | 4 +++- tests/plugin/test_vllm_offload_completions.py | 4 +++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/atom/plugin/vllm/kv_transfer/kv_cache_layout.py b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py index 7a2825c34e..877ec09e6d 100644 --- a/atom/plugin/vllm/kv_transfer/kv_cache_layout.py +++ b/atom/plugin/vllm/kv_transfer/kv_cache_layout.py @@ -139,9 +139,7 @@ def build_kv_cache_tensors( for layer_num, name in enumerate(sorted(main, key=_layer_sort_key)): k_cache, v_cache = split_kv_tensor(main[name]) index_cache = index_caches.get(name) - k_scale, v_scale = _transfer_scales( - name, (layers or {}).get(name), main[name] - ) + k_scale, v_scale = _transfer_scales(name, (layers or {}).get(name), main[name]) for role, seg in ( ("k_cache", k_cache), diff --git a/tests/plugin/test_vllm_kv_cache_layout.py b/tests/plugin/test_vllm_kv_cache_layout.py index 55866bfdab..3610d95c3f 100644 --- a/tests/plugin/test_vllm_kv_cache_layout.py +++ b/tests/plugin/test_vllm_kv_cache_layout.py @@ -161,7 +161,9 @@ def _bytes(tensors): def test_per_block_scales_without_a_reporting_hook_are_rejected(): kv = {"model.layers.0.self_attn.attn": torch.zeros((NB, 1, BS, 2 * HD))} with pytest.raises(ValueError, match="per-block k_scale"): - build_kv_cache_tensors(kv, {"model.layers.0.self_attn.attn": _LayerWithUnreportableScales()}) + build_kv_cache_tensors( + kv, {"model.layers.0.self_attn.attn": _LayerWithUnreportableScales()} + ) def test_non_block_major_scales_are_rejected(): diff --git a/tests/plugin/test_vllm_offload_completions.py b/tests/plugin/test_vllm_offload_completions.py index 05ba2314fa..58798a923a 100644 --- a/tests/plugin/test_vllm_offload_completions.py +++ b/tests/plugin/test_vllm_offload_completions.py @@ -163,6 +163,8 @@ def test_a_dispatched_load_is_not_reported(): adapter = _lookup_adapter(scheduler) adapter.get_num_new_matched_tokens(_req("ok"), 0) - adapter._check_promised_loads(SimpleNamespace(requests=[SimpleNamespace(req_id="ok")])) + adapter._check_promised_loads( + SimpleNamespace(requests=[SimpleNamespace(req_id="ok")]) + ) assert adapter._promised_loads == {} From bbfd3846e7961fcdc6ef53a18abdf7009e5f3041 Mon Sep 17 00:00:00 2001 From: perzhang Date: Sun, 6 Sep 2026 22:13:03 -0500 Subject: [PATCH 12/16] fix(vllm-plugin): break the breakable cudagraph at M3's sparse attention vLLM auto-enables VLLM_USE_BREAKABLE_CUDAGRAPH for this architecture, so plugin mode compiles nothing and splits nothing: one stream capture drives the whole forward, and only ops carrying @eager_break_during_capture end a segment. M3's sparse attention did not carry it, so 57 of the model's 60 attention layers were captured wholesale -- measured at capture time as graphs=4, eager_breaks=3, the three breaks being the dense layers that go through vLLM's own unified_attention. With the decorator: graphs=61, eager_breaks=60. Everything the sparse path reads per step was therefore frozen at whatever the capture batch held: the prefill/decode token counts, block_table, seq_lens, the topk indices. The damage stays invisible until a cache hit -- a cold prompt prefills past the largest captured size and runs eagerly, so it answers correctly; reuse a prefix and the short remainder lands inside a captured size, replays another batch's metadata, and answers fluently from the wrong KV. Any M3 run with prefix caching on is exposed, and the lower the hit rate the less it looks like a bug. The decorator requires the op to write into a caller-owned buffer, so the output allocation moves up into forward() and the op takes it as a mutated argument; a tensor allocated inside would land at a new address on every replay while the captured segments that consume it still read the address recorded at capture. Measured on MiniMax-M3-MXFP4, TP=4, cudagraph_mode=FULL_AND_PIECEWISE, prefix caching on, no LMCache. A 12,628-token prompt with markers at 10%, 50% and 90% depth, cold then replayed against 12,544 locally cached tokens: before, the replay missed the 50% and 90% markers and diverged; after, all three come back. Full decode graphs are untouched -- they dispatch with cudagraph_runtime_mode == FULL, which the decorator passes through. Co-Authored-By: Claude Opus 5 (1M context) --- .../vllm/attention/minimax_m3_attnetion.py | 65 +++++++++++++++---- 1 file changed, 54 insertions(+), 11 deletions(-) diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index b28a5102b8..506d22dcd1 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/atom/plugin/vllm/attention/minimax_m3_attnetion.py @@ -31,6 +31,7 @@ _register_vllm_static_forward_context, ) from atom.utils import mark_spliting_op +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -41,28 +42,35 @@ def minimax_m3_sparse_attention_fake( qkv: torch.Tensor, positions: torch.Tensor, layer_name: str, - output_hidden_size: int, -) -> torch.Tensor: - del positions, layer_name - return qkv.new_empty((qkv.shape[0], output_hidden_size)) + output: torch.Tensor, +) -> None: + del qkv, positions, layer_name, output @mark_spliting_op( is_custom=True, gen_fake=minimax_m3_sparse_attention_fake, - mutates_args=[], + mutates_args=["output"], ) def minimax_m3_sparse_attention( qkv: torch.Tensor, positions: torch.Tensor, layer_name: str, - output_hidden_size: int, -) -> torch.Tensor: + output: torch.Tensor, +) -> None: + """Write this layer's sparse attention into ``output``. + + The caller owns ``output`` rather than this op returning a fresh tensor, + because :func:`eager_break_during_capture` -- which the layer's + ``_sparse_attn_run`` carries -- replays the Python kernel on every + breakable-cudagraph replay. A tensor allocated inside would land at a new + address each replay while the captured segments that consume it still read + the address recorded at capture time. + """ from vllm.forward_context import get_forward_context layer = get_forward_context().no_compile_layers[layer_name] - output = qkv.new_empty((qkv.shape[0], output_hidden_size)) - return layer._forward_with_output(qkv, positions, output) + layer._sparse_attn_run(qkv, positions, output) class MiniMaxM3SparseIndexerCache(nn.Module, AttentionLayerBase): @@ -745,6 +753,39 @@ def _forward_with_output( ) return output + @eager_break_during_capture + def _sparse_attn_run( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Run sparse attention outside the breakable cudagraph segments. + + M3 reaches vLLM with ``VLLM_USE_BREAKABLE_CUDAGRAPH`` on (vLLM + auto-enables it for this architecture), so there is no FX splitting: + one stream capture drives the whole forward and only the ops carrying + this decorator end a segment. Without it the 57 sparse layers are + captured wholesale, and everything this path reads per step -- the + prefill/decode token counts, ``block_table``, ``seq_lens``, the topk + indices -- is frozen at whatever the capture batch happened to hold. + + The damage is silent and needs a cache hit to show: a cold prompt + prefills more tokens than the largest captured size and runs eagerly, + so it is correct; reuse a prefix and the short remainder lands inside a + captured size, replays another batch's metadata, and answers fluently + from the wrong KV. Any M3 run with prefix caching on is exposed. + + Decode is untouched: full decode graphs dispatch with + ``cudagraph_runtime_mode == FULL``, which the decorator passes through, + and that path builds its metadata through the backend's cudagraph-safe + persistent buffers. + + Mirrors what vLLM's own MiniMax-M3 does for its sparse attention, and + ATOM's Kimi-K3 plugin for the KDA mixer. + """ + self._forward_with_output(qkv, positions, output) + def forward( self, query: torch.Tensor, @@ -760,12 +801,14 @@ def forward( raise ValueError("MiniMax-M3 sparse vLLM attention requires packed qkv.") if positions is None: raise ValueError("positions is required for MiniMax-M3 sparse attention.") - return torch.ops.aiter.minimax_m3_sparse_attention( + output = qkv.new_empty((qkv.shape[0], self.q_size)) + torch.ops.aiter.minimax_m3_sparse_attention( qkv, positions, self.layer_name, - self.q_size, + output, ) + return output class MiniMaxM3DenseAttentionForVllm(nn.Module, AttentionLayerBase): From 3428952813eccc3f0f58e20c3c303cd504349a8b Mon Sep 17 00:00:00 2001 From: perzhang Date: Sun, 6 Sep 2026 22:22:54 -0500 Subject: [PATCH 13/16] test(vllm-plugin): guard M3's sparse attention eager break The defect this guards against is silent: drop the decorator and the sparse layers go back inside the captured graph, cold prompts still answer correctly, and only a prefix-cache hit exposes it. Confirming it costs a two-run accuracy sweep on four GPUs. The scan reads the source instead of importing it. The module imports aiter at module scope, so an importing test would skip on every CI runner -- precisely where the guard needs to fire. Both regression shapes were checked by hand: removing the decorator fails the first test, emptying mutates_args fails the second. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_vllm_minimax_m3_cudagraph_breaks.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/plugin/test_vllm_minimax_m3_cudagraph_breaks.py diff --git a/tests/plugin/test_vllm_minimax_m3_cudagraph_breaks.py b/tests/plugin/test_vllm_minimax_m3_cudagraph_breaks.py new file mode 100644 index 0000000000..d9eb855e71 --- /dev/null +++ b/tests/plugin/test_vllm_minimax_m3_cudagraph_breaks.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""MiniMax-M3's sparse attention must break the breakable cudagraph. + +vLLM auto-enables ``VLLM_USE_BREAKABLE_CUDAGRAPH`` for this architecture, so +plugin mode compiles nothing and splits nothing: one stream capture drives the +whole forward, and only ops carrying ``@eager_break_during_capture`` end a +segment. Drop the decorator and 57 of M3's 60 attention layers get captured +wholesale, freezing this batch's token counts, ``block_table``, ``seq_lens`` +and topk indices into every replay. + +Nothing downstream notices. A cold prompt prefills past the largest captured +size and runs eagerly, so it answers correctly; reuse a prefix and the short +remainder lands inside a captured size and answers fluently from the wrong KV. +Catching that costs a two-run accuracy sweep on four GPUs, which is why this +guard is here instead. + +The scan reads the source rather than importing it: the module imports ``aiter`` +at module scope, so an importing test would skip on every CI runner -- exactly +where the guard is supposed to fire. +""" + +import ast +from pathlib import Path + +M3_ATTENTION = ( + Path(__file__).resolve().parents[2] + / "atom" + / "plugin" + / "vllm" + / "attention" + / "minimax_m3_attnetion.py" +) + +BREAK_DECORATOR = "eager_break_during_capture" +SPARSE_OP = "minimax_m3_sparse_attention" + + +def _module() -> ast.Module: + return ast.parse(M3_ATTENTION.read_text()) + + +def _decorator_names(node: ast.FunctionDef) -> set[str]: + names = set() + for dec in node.decorator_list: + target = getattr(dec, "func", dec) + name = getattr(target, "id", None) or getattr(target, "attr", None) + if name: + names.add(name) + return names + + +def _find_function(tree: ast.Module, name: str) -> ast.FunctionDef | None: + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + return None + + +def test_sparse_attention_run_is_an_eager_break(): + tree = _module() + run = _find_function(tree, "_sparse_attn_run") + assert run is not None, ( + "MiniMaxM3SparseAttentionForVllm._sparse_attn_run is gone -- whatever " + "replaced it still has to carry the eager break, or M3's 57 sparse " + "layers go back inside the captured graph" + ) + assert BREAK_DECORATOR in _decorator_names(run), ( + f"_sparse_attn_run lost @{BREAK_DECORATOR}; sparse attention would be " + "captured with this batch's metadata frozen into every replay" + ) + + +def test_the_break_writes_into_a_caller_owned_buffer(): + """The decorator replays the Python kernel, so its output must be passed in. + + A tensor allocated inside lands at a new address on every replay while the + captured segments that consume it still read the address recorded at + capture -- the failure ``eager_break_during_capture`` documents by name. + """ + tree = _module() + + op = _find_function(tree, SPARSE_OP) + assert op is not None, f"{SPARSE_OP} is gone" + assert "output" in {a.arg for a in op.args.args}, ( + f"{SPARSE_OP} no longer takes `output` from its caller; allocating it " + "inside breaks every replay after the first" + ) + + mark = next( + ( + dec + for dec in op.decorator_list + if getattr(getattr(dec, "func", dec), "id", None) == "mark_spliting_op" + ), + None, + ) + assert mark is not None, f"{SPARSE_OP} is no longer registered as a custom op" + mutates = next((kw.value for kw in mark.keywords if kw.arg == "mutates_args"), None) + assert isinstance(mutates, ast.List), "mutates_args must be a literal list" + assert "output" in { + el.value for el in mutates.elts if isinstance(el, ast.Constant) + }, ( + f"{SPARSE_OP} must declare `output` in mutates_args, or torch records " + "it as a pure return and the schema stops marking the write" + ) + + run = _find_function(tree, "_sparse_attn_run") + assert run is not None and "output" in { + a.arg for a in run.args.args + }, "the eager break must receive the buffer rather than allocate one" From dcb8d5dd626fc7d543c287a26fb16b727b8f9f94 Mon Sep 17 00:00:00 2001 From: perzhang Date: Sun, 6 Sep 2026 22:34:23 -0500 Subject: [PATCH 14/16] style: sort the M3 attention import block for Ruff The block was already unsorted; adding an import to it put the existing violation inside the PR's diff context, which is what reviewdog gates on. vllm belongs in the third-party group with aiter and torch. Co-Authored-By: Claude Opus 5 (1M context) --- atom/plugin/vllm/attention/minimax_m3_attnetion.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/atom/plugin/vllm/attention/minimax_m3_attnetion.py b/atom/plugin/vllm/attention/minimax_m3_attnetion.py index 506d22dcd1..52ae9489a0 100644 --- a/atom/plugin/vllm/attention/minimax_m3_attnetion.py +++ b/atom/plugin/vllm/attention/minimax_m3_attnetion.py @@ -16,6 +16,9 @@ import torch from aiter import dtypes from torch import nn +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from atom.config import get_current_atom_config from atom.model_ops.minimax_m3.sparse_attn import ( @@ -31,9 +34,6 @@ _register_vllm_static_forward_context, ) from atom.utils import mark_spliting_op -from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.forward_context import get_forward_context -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase _MINIMAX_M3_TOPK_CACHE_STATE: dict = {} @@ -83,8 +83,8 @@ def __init__( head_dim: int, kv_cache_dtype: str, ) -> None: - from vllm.v1.attention.backend import AttentionType from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype + from vllm.v1.attention.backend import AttentionType super().__init__() atom_config = get_current_atom_config() From 83beae54e18ff2659743c88d5306c7ca67a6fe80 Mon Sep 17 00:00:00 2001 From: perzhang Date: Wed, 9 Sep 2026 04:45:30 -0500 Subject: [PATCH 15/16] fix(vllm-plugin): cap the offload frontier at what the block table covers A long prompt fails every save with "LMCache token range exceeds ATOM block table: needed_blocks=N+1, available_blocks=N" -- 88 times in one ISL-90k run, zero in any gsm8k-sized one. ATOM sizes a save from `seq.num_cached_tokens` and hands `seq.block_table` to LMCache alongside it. On the native path those two live on one Sequence and advance together. Here they arrive through two independent vLLM callbacks -- num_computed_tokens rides on scheduler_output, the block table on update_state_after_alloc -- and under chunked prefill they separate by a block. A chat-sized prompt is allocated in one go, so they never separate and every test we had passed. Capping is the semantics, not a guard: KV for tokens past the block table is not in any block this connector knows about, so offering it up was never right. The remainder saves on a later step, once the allocation catches up. Measured on MiniMax-M3-MXFP4 TP4, ISL 90,166 x 32 requests: the error count goes 88 -> 0. Cache hit rate is unchanged (59.09% -> 56.81% across runs, inside the 2.64-point run-to-run spread of this configuration), so this is a correctness fix, not a hit-rate one. Co-Authored-By: Claude Opus 5 (1M context) --- atom/plugin/vllm/kv_transfer/connector.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/atom/plugin/vllm/kv_transfer/connector.py b/atom/plugin/vllm/kv_transfer/connector.py index 4111026cd7..ae281da229 100644 --- a/atom/plugin/vllm/kv_transfer/connector.py +++ b/atom/plugin/vllm/kv_transfer/connector.py @@ -271,11 +271,32 @@ def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata: The frontier of every scheduled request is refreshed first: ATOM decides which chunks are safe to save by comparing against it, and a stale value would either skip chunks or offer up tokens that are not computed yet. + + The frontier is also capped at what the block table actually covers. + On ATOM's native path these two quantities live on one Sequence and + advance together; here they arrive through two independent vLLM + callbacks -- ``num_computed_tokens`` rides on scheduler_output, the + block table on ``update_state_after_alloc`` -- and under chunked + prefill they separate by a block. ATOM then sizes a save from the + frontier and hands the (shorter) block table to LMCache alongside it, + which fails the transfer with "LMCache token range exceeds ATOM block + table: needed_blocks=N+1, available_blocks=N". + + Capping is the correct semantics, not just a guard: KV for tokens past + the block table is not in any block this connector knows about, so + offering it up was never right. The remainder is saved on a later step, + once the allocation catches up. + + Only long prompts reach this: a chat-sized prompt is allocated in one + go, so the two quantities never separate and every gsm8k-scale test + passes. It took an ISL-90k aiperf run to surface. """ + block_size = int(self._config.kv_cache_block_size) for req_id, num_tokens in _scheduled_frontiers(scheduler_output): seq = self._seqs.get(req_id) if seq is not None: - seq.set_num_cached_tokens(num_tokens) + covered = len(seq.block_table) * block_size + seq.set_num_cached_tokens(min(int(num_tokens), covered)) inner = self._scheduler.build_connector_meta() self._check_promised_loads(inner) return AtomOffloadMetadata(inner) From 4071a77ac0b5398cef9fa2fd4aad6431d2e399d6 Mon Sep 17 00:00:00 2001 From: perzhang Date: Wed, 9 Sep 2026 06:32:15 -0500 Subject: [PATCH 16/16] docs(recipes): MiniMax-M3 LMCache byte-codec offload on the vLLM plugin M3 registers 117 KV tensors in three physical layouts, so LMCache's own GPU connector aborts on its single global format probe. This recipe covers the byte-codec path (AtomLMCacheOffloadConnector), which stores opaque bytes and never runs that probe. Beyond the launch command, it documents what took the longest to work out: - A four-check verification list plus two counter identities. Each check alone passes for the wrong reason -- the connector can be constructed and queried yet return nothing, so "queried" and "hit" must both be asserted. - The sizing budget. The same code measures 76.65% on a 262,144-token pool and 56.81% on a 524,288-token one, because a large pool lets HBM absorb the reuse and leaves the offload tier with nothing to do. A bigger pool reading as a worse tier is the fingerprint of this mistake. - The reachable ceiling, so a number is not called low against the wrong bar. Synthetic prefix load tops out at 78.50% (per-request unique tail can never hit); agentic replay measured 97.12%, which is the figure comparable to published ~96% results. - Denominator skew on agentic runs: server counters include warmup and trajectory reconstruction, aiperf counts only the profiling phase (27.3M vs 11.0M tokens here). Percentages compare, absolute counts do not. - Seven operational gotchas, including that the shm_broadcast warning appears in healthy startups too, and that pkill leaks GPU memory where podman restart does not. Measured on MiniMax-M3-MXFP4 TP4 / MI355X. Paired A/B with LMCache as the only variable: TTFT 1,247 ms vs 1,766 ms, throughput 133.19 vs 125.18 tok/s. Co-Authored-By: Claude Opus 5 (1M context) --- .../MiniMax-M3-LMCache-Byte-Offload.md | 209 ++++++++++++++++++ 1 file changed, 209 insertions(+) create mode 100644 recipes/atom_vllm/MiniMax-M3-LMCache-Byte-Offload.md diff --git a/recipes/atom_vllm/MiniMax-M3-LMCache-Byte-Offload.md b/recipes/atom_vllm/MiniMax-M3-LMCache-Byte-Offload.md new file mode 100644 index 0000000000..7774ad6a8c --- /dev/null +++ b/recipes/atom_vllm/MiniMax-M3-LMCache-Byte-Offload.md @@ -0,0 +1,209 @@ +# MiniMax-M3 — LMCache KV offload on the vLLM plugin (byte codec) + +MiniMax-M3 cannot use LMCache's own GPU connector. This recipe uses +`AtomLMCacheOffloadConnector`, which drives ATOM's `DenseKVByteCodec` from vLLM's +KV-connector API and leaves LMCache as a pure byte store. + +For the generic plugin + `LMCacheConnectorV1` path (works on M2.5 and other dense +models), see [LMCache KV Cache Offload](LMCache-KV-Cache-Offload.md). That path +**does not work on M3** — see *Why a separate connector* below. + +## Why a separate connector + +M3 registers 117 KV tensors in three different physical layouts at once: + +| layers | shape | note | +|---|---|---| +| 3 dense | `(nb, 1, 128, 256)` | K and V interleaved per token | +| 57 sparse | `(nb, 2, 128, 1, 128)` | K and V in two **separate** regions | +| 57 index caches | `(nb, 128, 128)` fp8 | DSA indexer keys | + +LMCache's `normalize_kv_and_discover_format()` probes for one **global** format, +so it aborts with `currently unsupported kv_caches format with list depth 1 and +tensor dimension 4`. The per-layer-format connector (V3) is off by default and +hangs on M3; the multi-process path needs cupy, which LMCache's `platform/rocm` +does not provide. + +`AtomLMCacheOffloadConnector` sidesteps the whole question: ATOM gathers whole +paged blocks into a chunk-major uint8 blob and LMCache only ever stores opaque +bytes, so no format probe ever runs. + +A sparse layer is also not contiguous as a whole (`stride(1)` jumps the entire K +region), so it is split into `t[:, 0]` / `t[:, 1]` before handing it to the codec. +M3's fp8 KV scales live on the layer, not in vLLM's `kv_caches` dict, and are +fetched through `get_kv_transfer_scales()` — moving mantissas without them +dequantises a restored block against the previous occupant's scale, which is +silent corruption. + +## Launch + +```bash +export PYTHONHASHSEED=0 # mandatory, see Gotchas +export LMCACHE_LOCAL_CPU=True +export LMCACHE_MAX_LOCAL_CPU_SIZE=20 # GiB **per TP rank** +export LMCACHE_CHUNK_SIZE=128 # must equal --block-size +export OFFLOAD_MIN_LOAD_TOKENS=256 # default 8192 disables the tier for chat-sized prompts +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=1 + +vllm serve /path/to/MiniMax-M3-MXFP4 \ + --served-model-name amd/MiniMax-M3-MXFP4 \ + --tensor-parallel-size 4 \ + --gpu-memory-utilization 0.8 \ + --block-size 128 \ + --max-model-len 131072 \ + --max-num-seqs 64 \ + --max-num-batched-tokens 32768 \ + --kv-cache-dtype fp8 \ + --no-async-scheduling \ + --language-model-only \ + --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --enable-prefix-caching \ + --enable-prompt-tokens-details \ + --kv-transfer-config '{"kv_connector":"AtomLMCacheOffloadConnector","kv_connector_module_path":"atom.plugin.vllm.kv_transfer.connector","kv_role":"kv_both"}' +``` + +Select the connector through vLLM's **out-of-tree entry point** (the +`kv_connector_module_path` above). vLLM validates `kv_transfer_config` while +building VllmConfig, which happens *before* platform plugins load, so naming the +connector without the module path fails config validation. + +## Verify it is actually on + +Four independent checks — all of them, because each can pass for the wrong reason: + +```bash +# 1. vLLM's own factory, on every worker AND the EngineCore +grep "Creating v1 connector with name: AtomLMCacheOffloadConnector" server.log + +# 2. the codec saw the whole model +grep "ATOM LMCache offload: registered 60 layers" server.log + +# 3+4. the tier is queried AND returns data (must be > 0) +curl -s localhost:8000/metrics | grep -E 'external_prefix_cache_(hits|queries)' +``` + +Two identities must hold on any interval; they catch a miscounting tier that +still looks plausible: + +``` +prefix_cache_queries - prefix_cache_hits == external_prefix_cache_queries +prefix_cache_hits + external_prefix_cache_hits == prompt_tokens_cached +``` + +The first says the two tiers are strictly serial (HBM first, LMCache only gets +what HBM missed); the second says nothing is double-counted. + +## Measured (MiniMax-M3-MXFP4, TP4, MI355X) + +Synthetic prefix-reuse load: 80,896-token shared prefix + 9,270 unique tail, +32 requests, concurrency 2, pool of 4 prefixes, KV capped to 262,144 tokens. +Paired runs — same seeds, same server flags, LMCache the only difference. + +| | LMCache on | LMCache off | +|---|---:|---:| +| Total cache read | **76.65%** | 55.17% | +| external (LMCache) supplied | 629,504 tok | 0 | +| **TTFT** | **1,247 ms** | 1,766 ms | +| Output throughput | 133.19 tok/s | 125.18 tok/s | + +**TTFT 29% lower, throughput 6% higher.** Run-to-run noise on this +configuration is 0.86% (throughput) and 0.11% (TTFT), so the effect is far +outside it. + +Cross-checked against an independent source: aiperf's +`usage.prompt_tokens_details.cached_tokens` summed to 2,211,584 against the +server's `prompt_tokens_cached` delta of 2,211,584 — exact. + +### Sizing matters more than anything else here + +The same code on a larger KV pool measures **56.81%** instead of 76.65%, because +HBM alone then holds ~2 of the 8 prefixes and LMCache only patches the edges. +Work out the budget before benchmarking: + +``` +free for caching = KV pool - concurrency x ISL +``` + +At KV 524,288 / conc 4 / ISL 90,166 that leaves 163,624 tokens = 2.02 prefixes; +at KV 262,144 / conc 2 it leaves 81,812 = 1.01 prefixes, which pushes ~75% of all +reuse onto LMCache. **If HBM can hold the whole working set, this tier cannot +help and a benchmark will show nothing.** + +Also compute the ceiling before calling a number low: + +``` +ceiling = prefix/ISL - (pool x prefix)/(ISL x requests) +``` + +which is 78.50% for the config above — the measured 76.65% is **97.6% of what is +physically reachable**. The per-request unique tail can never hit, and each +prefix must be computed once. Published agentic numbers (~96%) come from +multi-turn replay where the unique tail is a small fraction of each turn; they +are not comparable to a synthetic prefix load. + +### Agentic replay, and what an oversized pool looks like + +Same server, `--public-dataset semianalysis_cc_traces_weka_062126`, 8 trajectory +sessions, 900 s, but with the **default KV pool (7,610,112 tokens)**: + +| | | +|---|---:| +| Total cache read (aiperf) | **97.12%** | +| HBM tier | 91.36% | +| external tier queried | 2,578,390 tok | +| **external tier hit** | **0** | + +Two things worth reading carefully. + +**97.12% is the number to compare against published agentic results (~96%)** — +same kind of workload, same kind of measurement. The 76.65% from the synthetic +load above is not a worse result, it is a different ceiling. + +**The external tier hit nothing, and that is correct here.** The pool holds +7.6M tokens, so HBM alone absorbs 91.36%; what it misses is content appearing +for the first time, which no cache can hold. The connector was fully live — it +was queried 2,578,390 times, exactly the HBM miss volume — it simply had no +reusable bytes to offer. **An oversized pool makes this tier look useless.** +Cap the pool below the working set before concluding anything about it. + +**Do not compare the two sides' absolute token counts on an agentic run.** Here +the server reported 27,274,624 cached tokens and aiperf 10,992,256 — a 2.48x +gap, where every other run in this recipe matched exactly. The denominators +differ: aiperf counts only the profiling phase, the server counters also include +warmup and trajectory reconstruction. Percentages remain comparable; absolute +counts do not. + +## Gotchas + +- **`PYTHONHASHSEED=0` is mandatory.** Without it each TP rank derives a + different cache key for the same prompt and the hit ratio collapses to 0. +- **`LMCACHE_CHUNK_SIZE` must equal `--block-size` (128).** ATOM refuses a load + whose HBM frontier is not chunk-aligned; with prefix caching on that frontier + is block-aligned, so at chunk 256 roughly every other hit is dropped. +- **`OFFLOAD_MIN_LOAD_TOKENS` defaults to 8192**, which is above every prompt in + a chat-sized workload — the tier would never serve anything. +- **`LMCACHE_MAX_LOCAL_CPU_SIZE` is per rank.** TP4 x 20 GiB locks 80 GiB of + pinned memory; if free memory is low the allocation reclaims page cache and + each worker can take minutes. Startup looks hung — `EngineCore` prints + `No available shared memory broadcast block found in 60 seconds` once a minute. + That line alone is **not** a failure: a healthy startup prints it 5 times too. + Check whether the workers are still emitting log lines instead. +- **Saves are fire-and-forget.** A request returning does not mean its KV has + landed. Benchmarks that measure immediately after warm-up systematically + under-report external hits; allow a settle period. +- **`--enable-prompt-tokens-details` is required for client-side verification.** + Without it aiperf silently reports an empty prompt-cache column rather than an + error. +- **Stop the server with `podman restart`, not `pkill`.** Killing TP workers + leaves zombies holding GPU memory (82 GiB/card observed); only restarting the + container releases it. + +## Related + +- [LMCache KV Cache Offload](LMCache-KV-Cache-Offload.md) — generic plugin path + (`LMCacheConnectorV1`), does not support M3's layouts +- [MiniMax-M3](MiniMax-M3.md) — base serving recipe +- `recipes/MiniMax-M3-Agentic-Offload.md` — ATOM **native** backend offload with + agentic replay, DP2 sizing ladder