[Feature] OFFLOAD: standalone LMCache CPU/NVMe KV-offload connector - #1318
Merged
Conversation
yhl-amd
force-pushed
the
feature/lmcache-offload-merge
branch
2 times, most recently
from
June 23, 2026 03:14
b9b7c24 to
6bb1fbc
Compare
| if not self.kv_transfer_enabled: | ||
| return | ||
| connector = getattr(self.scheduler, "kv_connector", None) | ||
| if connector is None or not getattr(connector, "is_offload", False): |
Collaborator
There was a problem hiding this comment.
We use mooncake connector and moriio connector in PD, can not use connector, "is_offload" here? We should use a new variable and plugin in old connector
Contributor
Author
There was a problem hiding this comment.
it will support in later version
zufayu
requested review from
JiaoliangYu and
amd-ruitang3
and removed request for
amd-ruitang3
June 26, 2026 06:12
yhl-amd
force-pushed
the
feature/lmcache-offload-merge
branch
from
June 29, 2026 02:45
8a37b72 to
bd62ca0
Compare
valarLip
previously approved these changes
Jun 29, 2026
Collaborator
|
need rebase to latest main |
Add a standalone KV-offload subsystem that offloads ATOM KV cache to LMCache-backed CPU/NVMe storage and reloads it on cache hits, avoiding prefill recompute for evicted prefixes. - New atom/kv_transfer/offload package: LMCacheOffloadConnector, the ATOM<->LMCache GPU connector, a byte codec for ATOM's packed KV layout, a Triton staging kernel, plus metadata and config. - Wire the connector into the disaggregation base/factory/types and aggregate per-worker finished/failed transfer states. - Unit tests for the connector and byte-codec round-trip. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eduler Drive offload load/save from the engine loop and scheduler: - Dispatch async KV load after connector metadata, poll worker transfer status, and advance idle KV transfer when no forward batch runs. - Defer block free until a background D2H save has read the KV, and wake parked prefills for local recompute on a load miss (failed_recving). - Handle chunked-prefill deferred output across the offload park/resume boundary so stale sampled tokens are dropped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Honor max_completion_tokens (and the max_tokens alias) in the OpenAI completion/chat protocol and server so offload benchmarks can bound generation length. Adds protocol and server-helper unit tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document the ATOM standalone lmcache_offload connector: design, module map, scheduler/worker architecture, byte codec and AITER layout bridge, MemoryObj/segment layout, completion protocol, reload decision and chunk-alignment handoff, correctness/fp8/failure handling, the LMCache reuse-vs-override boundary, configuration, benchmarks, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
The byte codec assumed a block-major KV layout (tensor.shape[0] == block count), which holds for MHA/GQA but not for MLA: MLA stores a single per-layer latent cache (kv_lora_rank + qk_rope_head_dim, e.g. 576) viewed token-major as (num_blocks * block_size, 1, latent), with no separate V or scale tensors. So shape[0] is the token count and the codec computed a per-token (not per-block) byte stride, corrupting the offloaded KV. Both layouts share an identical contiguous byte layout (block b always starts at b * bytes_per_block), so instead of branching we take the physical block count explicitly and derive each segment's per-block stride as segment_bytes / num_blocks. The Triton fused staging kernel is byte-addressed and needs no change. - ATOMKVByteCodec: accept explicit num_blocks; per-block bytes from it; require contiguous + numel divisible by num_blocks (replaces the "same shape[0]" check). Falls back to shape[0] when num_blocks is None, preserving non-MLA behaviour. - Thread num_physical_kvcache_blocks: model_runner.allocate_kv_cache -> forward_context.set_kv_cache_data -> connector.register_kv_caches -> codec. Optional num_blocks kwarg added to the connector base + mooncake + moriio impls (ignored there). - build_lmcache_metadata: emit an MLA-shaped kv_shape (latent dim) for bookkeeping; storage stays opaque BINARY so use_mla remains False. - Tests: MLA token-major block accounting + byte-identical round-trip. Validated on DeepSeek-V3-5layer (real MLA, TP=2) end-to-end: offload save + reload (cxs multi-round, round 2 hits cached:[~33k], no recompute). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Document MLA (DeepSeek R1/V3, Kimi) KV-offload support in the connector README: token-major latent cache layout, explicit num_blocks threading, and BINARY opaque storage. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…_UNALIGNED_HANDOFF Unaligned HBM-prefix loads now always take the handoff path (recompute the misaligned head up to the next chunk boundary, then load the aligned remainder from CPU) instead of being gated behind the OFFLOAD_UNALIGNED_HANDOFF env var. The env read and the gate check in _maybe_start_unaligned_handoff are removed; the min_load / boundary guards are unchanged. - connector.py: drop _allow_unaligned_handoff + OFFLOAD_UNALIGNED_HANDOFF read; handoff is now unconditional. - README.md / README.zh-CN.md: remove the env var from the tuning table and the example commands, add a "removed" note, and document the always-on behaviour. Also track the previously-untracked zh-CN README. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
… MockConfig The rebase onto main left two enable_chunked_prefill keys in MockConfig's defaults dict (our branch moved it up and set =True; main kept the old =False line while adding hf_config after it). Python raises SyntaxError on a repeated keyword in a dict() call, so conftest failed to import and the whole unit-test suite aborted (exit 4). Keep the single enable_chunked_prefill=True (our intended default; all tests that depend on the value override it explicitly) plus main's hf_config stub. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…s-on handoff Commit "make unaligned handoff always-on; drop OFFLOAD_UNALIGNED_HANDOFF" removed _allow_unaligned_handoff from the connector but left three dead references in the test file, including test_load_is_skipped_if_hbm_floor_is_ not_chunk_aligned, which encodes the old default-skip behaviour. With handoff now unconditional, an unaligned HBM floor (hbm=6, chunk=4, min_load=0) takes the handoff path (park to boundary 8, emit the load once cached reaches 8) instead of skipping the CPU load, so that test's `lookup.cleared == ["654"]` no longer holds and it failed in CI. The scenario is already covered by the always-on tests: - test_unaligned_hbm_handoff_prefills_boundary_then_emits_load (handoff path) - test_unaligned_handoff_skips_if_boundary_remainder_is_too_small (the real skip+clear case, gated by min_load rather than alignment) Delete the stale test and the three no-op _allow_unaligned_handoff assignments that the connector no longer reads. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yhl-amd
force-pushed
the
feature/lmcache-offload-merge
branch
from
June 29, 2026 09:28
a195549 to
ac394b6
Compare
valarLip
approved these changes
Jun 29, 2026
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
3 participants
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.
Summary
Adds a standalone KV-offload subsystem that offloads the ATOM KV cache to
LMCache-backed CPU/NVMe storage and reloads it on cache hits, so evicted
prefixes are restored instead of recomputed.
Squashed into 3 logical commits:
[Feature] OFFLOAD: add LMCache CPU/NVMe KV-offload subsystem— newatom/kv_transfer/offloadpackage (LMCacheOffloadConnector, ATOM↔LMCacheGPU connector, byte codec for ATOM's packed KV layout, Triton staging
kernel, metadata/config) + disaggregation base/factory/types/aggregator
wiring.
[Feature] OFFLOAD: integrate KV offload load/save into engine and scheduler— async load dispatch + worker transfer polling + idle KV-transfer advance;
deferred block free until background D2H save completes; wake parked
prefills for local recompute on load miss; chunked-prefill deferred-output
handling across park/resume.
[Frontend] Support max completion tokens in OpenAI API— honormax_completion_tokens/max_tokensso offload benchmarks can boundgeneration length.
Test plan
pytest tests/test_lmcache_offload_connector.py(connector + byte-codec round-trip)pytest tests/test_scheduler.pypytest tests/entrypoints/test_protocol.py tests/entrypoints/test_api_server_helpers.py🤖 Generated with Claude Code