[MiniMax-M3][LMCache] Drive ATOM's byte-level KV offload from the vLLM plugin - #2146
Open
XiaobingSuper wants to merge 14 commits into
Open
[MiniMax-M3][LMCache] Drive ATOM's byte-level KV offload from the vLLM plugin#2146XiaobingSuper wants to merge 14 commits into
XiaobingSuper wants to merge 14 commits into
Conversation
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)
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)
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)
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)
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)
…uest 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.
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,</$=$,</$=$,...'. Cold runs on the same prompts were perfect. M3's sparse cache stores fp8 mantissas whose scale is per token AND per head: kv_scale is (2, num_blocks, num_kv_heads, block_size) fp32, one entry per element of the paged cache. It lives on the attention layer, allocated lazily on first forward. vLLM's register_kv_caches hands a connector the KV tensors and nothing else, so the transfer moved mantissas only: a restored block was dequantised against whatever scales its previous occupant had left in place. Nothing can fail loudly there -- both sides are the right size and the right dtype -- which is why 200 gsm8k questions scored 0.925 without ever exercising a load. The layer owns the layout, so the layer reports it: get_kv_transfer_scales() on the sparse attention, named after the native path's get_kv_transfer_tensors, which reports the same regions off runner.kv_scale. The connector passes the layer modules (from vLLM's static forward context) alongside the tensors, and the codec already enumerates k_scale/v_scale as segments -- no codec change. It takes the registered tensor as an argument because registration runs before the first forward, and _ensure_fp8_scales reallocates on a shape or device change: sizing against the very tensor vLLM registered is also what keeps the allocation stable, so a pointer handed to the tier stays live. Two guards, because this failure mode is invisible at runtime: - a layer holding a multi-element k_scale/v_scale with no way to report it is a hard error, not a silent omission (scalar per-tensor scales are constant and still ignored); - a reported scale whose leading axis is not the block axis is rejected: the codec slices per-block as numel // num_blocks, and a head-major scale is the right dtype and nearly the right size, so nothing downstream would notice. (cherry picked from commit 13b338565619e9c1151c70f1798352afcaeaa207)
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)
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)
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)
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
There is at least one Black formatting issue in newly added code that is likely to fail the repo’s black --check CI gate.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds a vLLM KV-connector adapter that lets vLLM drive ATOM’s existing byte-level LMCache offload path (including MiniMax-M3’s mixed KV layouts and fp8 per-token/per-head scale handling), aligning plugin-mode behavior with the native ATOM engine’s offload tier.
Changes:
- Introduces a vLLM
KVConnectorBase_V1implementation that bridges vLLM KV-transfer hooks to ATOM’s dense LMCache offload scheduler/worker, including completion plumbing and “promised load” hang detection. - Adds vLLM→ATOM shims for request identity (
SeqViewRegistry), config projection (OffloadConfigShim), and KV registration layout mapping (including splitting sparse K/V, folding index caches, and transporting fp8 scales). - Adds targeted unit tests covering identity semantics, config mapping, completion propagation/leak prevention, and M3-realistic stride/layout constraints.
File summaries
| File | Description |
|---|---|
| tests/plugin/test_vllm_offload_seq_view.py | Validates SeqView identity + prompt-token keying + registry lifecycle. |
| tests/plugin/test_vllm_offload_config.py | Tests vLLM→ATOM offload-config projection and dtype/block-size validation. |
| tests/plugin/test_vllm_offload_completions.py | Ensures scheduler-side completion propagation, leak prevention, and promised-load watchdog behavior. |
| tests/plugin/test_vllm_kv_cache_layout.py | Tests KV layout mapping for M3 (dense/sparse/index/scales) and codec compatibility. |
| atom/plugin/vllm/register.py | Registers the connector name with vLLM’s connector factory (convenience path). |
| atom/plugin/vllm/kv_transfer/seq_view.py | Provides a stable per-request seq wrapper matching ATOM offload scheduler expectations. |
| atom/plugin/vllm/kv_transfer/offload_config.py | Projects VllmConfig into the minimal ATOM offload config surface. |
| atom/plugin/vllm/kv_transfer/kv_cache_layout.py | Translates vLLM’s flat KV registration into ATOM KVCacheTensor segments (incl. fp8 scales/index). |
| atom/plugin/vllm/kv_transfer/connector.py | Main adapter implementing vLLM v1 connector API over ATOM dense offload. |
| atom/plugin/vllm/kv_transfer/init.py | Defines the kv_transfer package for the vLLM plugin adapter. |
| atom/plugin/vllm/attention/minimax_m3_attnetion.py | Adds get_kv_transfer_scales hook so fp8 scales travel with KV bytes. |
| atom/kv_transfer/offload/dense/connector.py | Adds request-id-based completion resolvers for the vLLM plugin path. |
Review details
- Files reviewed: 12/12 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+79
to
+80
| def __repr__(self) -> str: | ||
| return f"_HFConfigView(inner={type(self._inner).__name__}, outer={type(self._outer).__name__})" |
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this does
Lets vLLM drive ATOM's existing byte-level LMCache offload, so the vllm-atom plugin path gets the KV offload tier the native engine already has. MiniMax-M3 is the model that forced it; nothing here is M3-specific except the one attention hook.
Validated end to end on 4×MI355X (TP=4,
amd/MiniMax-M3-MXFP4, fp8 KV), with cudagraph and HBM prefix caching both enabled:external_kv_transferNo LMCache upgrade needed — validated on the image's stock LMCache 0.4.5, where the replay pass scores 0.945 against 0.945 cold with 621,312 tokens served externally, and the round trip is bit-exact. (0.4.5 is what rejected
LMCacheConnectorV1on M3 withunsupported kv_caches format with list depth 1 and tensor dimension 4; that is the path this branch does not use.) It also runs on 0.5.5rc4.Plus a targeted correctness check (markers planted at 10/50/90% depth, HBM flushed with filler prompts in between): 11,008 tokens restored entirely from LMCache, all three markers reproduced.
Why not point vLLM at LMCache's own connector
LMCache's GPU connectors accept only the clean NHD/HND family and pick one format per model. M3 registers three physical layouts at once (dense K/V interleaved, sparse K/V in separate regions, plus a DSA index cache), and on ROCm the paths that could express that are unavailable anyway: the per-layer-format connector (V3) is off by default upstream and hangs on M3, and the multi-process path imports
cupy, which LMCache'splatform/rocmdoes not provide.ATOM already solved this for its native engine by never asking LMCache to understand the layout:
DenseKVByteCodecgathers whole paged blocks into a chunk-majoruint8blob and LMCache stores opaque bytes. That codec is reused verbatim here — no change to LMCache, no change to vLLM. This branch is the adapter that lets vLLM drive it.Shape of the change
kv_cache_layout.py— vLLM hands a connector a flat{layer_name: tensor}and says nothing about what is inside. M3's sparse layers are not contiguous as a whole (stride(1)jumps the entire K region), so the tensor cannot be one segment, butt[:, 0]andt[:, 1]each are — exactly thek_cache/v_cachepair the codec expects. Dense layers stay whole; the DSA index caches fold into their owning layer.connector.py—KVConnectorBase_V1wrapper. 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 vLLMRequestas the seq ATOM's scheduler expects, as one stable instance per request (ATOM compares object identity to detect a recycled request id).offload_config.py— derives ATOM's config fromVllmConfigrather than duplicating block size and role.dense/connector.py— two small additions on the ATOM side, see below.Four defects found while validating (each is a commit with its own reasoning)
register_kv_cacheshands over the caches only. A restored block was dequantised against whatever its previous occupant left behind; both sides are the right size and dtype, so nothing can fail loudly. The layer now reports its own scales (get_kv_transfer_scales, named after the native path'sget_kv_transfer_tensors), with two hard errors so this cannot go quiet again: a layer holding multi-element scales it cannot report, and a reported scale whose leading axis is not the block axis.should_park_for_load_after_allocbefore parking anything. vLLM has no second chance:async=Trueparks the request until the worker reportsfinished_recving, and a load that is never issued is never reported. Both routine drops fire constantly in the configurations we want (the floor defaults to 8192 tokens; with prefix caching on, vLLM's block-aligned frontier is regularly not chunk aligned). Symptom: EngineCore spinning inschedule(), every GPU idle, no log line at all.update_connector_outputwas unimplemented, so nothing on the scheduler side ever finished — and every request whose save was still in flight kept itsSeqView, which pins that request's prompt token ids. At M3's context lengths that is most of a megabyte retained per request.save_finished_by_request/load_finished_by_requestare added toDenseOffloadSchedulerfor (3): vLLM'sKVConnectorOutputcarries request ids as plain strings, whilesave_finished/load_finisheddeliberately refuse a raw id once the lifecycle has an exact operation identity. Rather than weaken that guard for everyone, the plugin path gets a way to resolve the identity ATOM already parked.Depends on
fix/m3-vllm-plugin-index-cache-fp8-dtype— M3 with--kv-cache-dtype fp8needs the index cache to carry a real fp8 dtype, or the index-topk Triton kernel dots bf16 against uint8 on the first request. Independent of this branch; this one only depends on the result.Known constraint, not introduced here
Use
cudagraph_mode=FULL, notFULL_AND_PIECEWISE. M3 miscomputes under piecewise-captured graphs, and reusing a cached prefix is what exposes it: a resumed request prefills only a handful of tokens, which lands inside the capture sizes. Measured without LMCache and without this branch (plainvllm serve, default KV size, prefix caching on, same gsm8k twice):cudagraph_modeFULL_AND_PIECEWISEFULLPIECEWISEis worse still — it fails cold.NONEandFULLare both correct. This affects any M3 workload with prefix caching on, and stays invisible while the hit rate is near zero; it is being reported separately.