Skip to content

[Feature] KV offload: hybrid bundle backend + dense/hybrid split#1683

Open
yhl-amd wants to merge 4 commits into
ROCm:mainfrom
yhl-amd:feat/dsv4-lmcache-bundle
Open

[Feature] KV offload: hybrid bundle backend + dense/hybrid split#1683
yhl-amd wants to merge 4 commits into
ROCm:mainfrom
yhl-amd:feat/dsv4-lmcache-bundle

Conversation

@yhl-amd

@yhl-amd yhl-amd commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a second KV-offload backend and reorganizes atom/kv_transfer/offload/
into two clearly-separated backends over a shared worker-side base.

  • hybrid/ — profile-driven bundle offload: one opaque offload unit per
    128-aligned terminal checkpoint. Targets hybrid-KV models (DSV4 CSA/SWA, Qwen3-Next).
    Adding a new hybrid-KV target = a new HybridProfile, not a new connector.
    Components: connector, admission (state-pool + max-inflight-saves gate),
    store, kv_bundle + codec, policy, and profiles/ (dsv4, qwen3_next).
  • dense/ — the existing token-chunked MHA/MLA LMCache offload path, moved out
    of the top-level connector.py (DenseOffloadConnector + GPU connector + byte codec).
  • _offload_common.py — worker-side plumbing shared by both backends:
    separate save/load copy executors (load never queues behind fire-and-forget saves),
    lock-guarded completion tallies, the KVConnectorOutput reporting contract, and the
    opaque-uint8 LMCache engine build (ATOMRawBytesLMCacheMetadata + KV_2LTD hack).
  • triton_offload_gather.py — Triton gather/scatter for staging compute-slot state
    into pool slots on the RPC thread before the next forward overwrites it.
  • Top-level connector.py is now a thin dispatch shell (927 → mostly removed);
    scheduler.py gets the small hooks (should_defer_free / save-frontier) both
    backends rely on.

Test plan

  • tests/test_hybrid_offload_* — admission, connector (GPU), gather, kv_bundle,
    kv_bundle codec, policy
  • tests/test_offload_profile_sources.py, tests/test_offload_qwen3_next.py,
    tests/test_offload_shell.py
  • tests/test_lmcache_offload_connector.py (dense path) updated for new layout
  • TP4 real-run smoke on DSV4 (GPU0-3) — save/load round-trip via serve

Notes

Squashes two commits: the DSV4 terminal-checkpoint offload unit and the
subsequent dense/hybrid split. No changes to the dense wire format; hybrid is
additive and gated behind its own connector/profile selection.

yhl-amd and others added 2 commits July 22, 2026 21:32
Add a standalone offload backend (dsv4_offload) that offloads a DeepSeek-V4
terminal checkpoint at a 128-aligned boundary B as one opaque LMCache "offload
unit": compressed KV (CSA main + HCA main + CSA indexer, keyed by block_table)
+ raw SWA tail [B-W,B) (swa_block_table) + CSA overlap state (v4_state_pool
slot via gather_slot). HCA (ratio=128, overlap=False) is offloaded compressed-
only; its ring is neither saved nor restored, correct by construction at a 128
boundary. On a prefix hit for the largest stored B<prompt_len, the seq parks,
loads + scatters the unit, and resumes suffix prefill [B, prompt_len).

New (atom/kv_transfer/offload/dsv4/): unit (container/layout/header/CRC/
fingerprint), unit_codec, gpu_connector (region gather via shared two-stream
pipeline), sources (B->SourceRegion + admission gate), admission (state-pool
slot + SWA IO-pin bookkeeping), policy (hit selection + checkpoint keying),
store (LMCache StorageManager batched_put/get), manager, connector
(worker + scheduler).

Unify shared machinery so the two offload backends no longer duplicate it:
- _offload_common: OffloadWorkerMixin + build_offload_engine (both use).
- atom_lmcache_staging.run_staged_pipeline + _PipelineStage: shared two-stream
  event-synced staging pipeline; DSV4 gpu connector and the chunked connector
  both drive it.
- triton_offload_gather: merged region gather/scatter (general) + chunk-major
  pack/unpack (specialization); replaces triton_kv_staging + the DSV4 gather.
- scheduler: materialize_window at offload_loaded_tokens (checkpoint B) so the
  SWA trailing window lands at the resume boundary (no-op for non-SWA offload).

Validated: 76 CPU tests (MHA 34 + DSV4 42) + 13 GPU unit tests; real
DeepSeek-V4-Flash TP=4 serve logits parity (B-scratch == B-resume, byte
identical) with HBM prefix cache disabled so offload is the only reuse path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
… base

Restructure atom/kv_transfer/offload: the former dsv4/ package becomes the
profile-driven hybrid/ backend, and the token-chunked MHA/MLA path moves into
dense/. Common worker-side plumbing (save/load executors, completion tallies,
KVConnectorOutput reporting, opaque-uint8 LMCache engine build) is factored into
_offload_common.py so neither backend duplicates the executor/engine boilerplate.

- add dense/ (DenseOffloadConnector + GPU connector + kv_byte_codec)
- add hybrid/ (profile-driven bundle connector, admission, store, codecs)
- drop dsv4/ and its tests; add hybrid/ + profile/shell test suites
- update offload/__init__.py, connector.py, README for the new layout

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
@github-actions

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

  • ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

  • ✅ Run after the PR is approved and Pre Checkin passes
  • ✅ Run immediately when an approval review is submitted
  • ✅ Can be requested before approval with labels
Label Tests
ci:full Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom Run native ATOM model accuracy tests
ci:vllm Run ATOM vLLM OOT model accuracy tests
ci:sglang Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1683 --add-label <label>

yhl-amd and others added 2 commits July 23, 2026 04:48
The gather/scatter file is a superset of the pre-existing triton_kv_staging.py
(old chunk-major pack/unpack kept for dense + new region gather/scatter for
hybrid). At 41% textual overlap git fell below its rename threshold and showed
it as delete+add, which is hard to review. Rename back to triton_kv_staging.py
so the PR shows it as a single in-place modification; update the 3 importers +
README references. No code behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
…f clean

Reorganize triton_kv_staging.py so the pre-existing chunk-major kernels/helpers
stay byte-identical to main (same position, indentation, names, signatures) and
the new region gather/scatter primitive (GatherRegion, gather_regions,
scatter_regions) is appended below. Drop the unnecessary _HAVE_TRITON guard that
had re-indented the old kernels (triton is always present in ATOM images), and
revert the cosmetic _build_meta rename + fused_* signature reflow.

Net effect on the PR diff: the file now reads as a near-pure addition
(+158/-1) instead of a full-file rewrite; no behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
@zufayu
zufayu requested a review from valarLip July 24, 2026 01:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant