Skip to content

Add KV cache offload to disk (per-node NVMe staging) — fixes #57 - #58

Closed
jdecker76 wants to merge 8 commits into
MiaAI-Lab:mainfrom
jdecker76:feat/kv-cache-offload
Closed

jdecker76 wants to merge 8 commits into
MiaAI-Lab:mainfrom
jdecker76:feat/kv-cache-offload

Conversation

@jdecker76

@jdecker76 jdecker76 commented Aug 30, 2026 •

Copy link
Copy Markdown

Experimental offload candidate for #57; not current-image runtime qualification.

Off by default. GLM53_OFFLOAD_MMAP_DIR backs the CPU offload staging region with a sparse file on each node's own NVMe, so a prefix that has been evicted from the GPU pool is restored from disk instead of recomputed. Capacity becomes the size of that file rather than RAM, which is what makes a shared box practical.

Historical author measurements (August 2026 image)

105k-token prompt, evicted by 6 × 168k-token floods, 100 GB per-node store, thinking off, temp 0:

Restore 8.54 s vs 130 s to recompute (15×)
Output byte-correct
Read from disk 1.68 GB (kv_offload CPU_to_GPU)
Host stable, 6–8 GB free throughout
Capacity 3,162 blocks ≈ 11.3M tokens in 160 GB/node

Historical decode observations (not a current-stack equivalence claim) — tests/bench_decode.py, median of 5 × 400, offload on: structured 66.4 tok/s (0.938 accept / 6.57 per step), prose 26.7, against the 65.1 / 27.1 in the README. Also re-run at 1M context on our launcher: boots clean, GPU KV cache size: 1,347,826 tokens.

What the patch does

overlay/patch_kv_offload_groups.py, 27 anchored edits over 5 vLLM files:

  1. Unblocks the boot. build_offloading_config asserts every group is hash-aligned; KpoolTailSpec's 4-token block never can be. Non-alignable groups are treated as per-request scratch instead of aborting.
  2. Stops drafter starvation. Sub-block groups (kpool tail block=4, DFlash2 drafter block=64) emit ~56× more offload keys than the 3584-token groups and consume the whole staging tier, so MLA/mamba chunks never reach disk — on one 256 GB run every file on disk was _g6. Both are excluded. Drafter KV rebuilds inside its own window after a restore, and the target verifies every token, so this costs acceptance at worst, never correctness.
  3. Fixes restore coherence. The packed whole-row copy path registers one whole-row ref per group; since groups allocate block ids independently over shared rows, a restore's row-write clobbers other groups' live pages. Forces the per-layer path (also ~10× less data moved).
  4. Bounds the staging page cache (msync + MADV_DONTNEED every GLM53_OFFLOAD_RELEASE_BYTES). Required on GB10: CPU and GPU share one pool, so an unbounded region starves the GPU and the worker rank dies mid-prefill with no traceback.

Why per-node instead of an fs secondary tier

Secondary tiers and their promotions run scheduler-side on the head only. Under 2-node TP that means rank 1's shard is never written (every file lands ..._r0) and its restores read stale bytes — the request is admitted as a hit and generates confidently wrong output while kv_offload counters look healthy. Ours degenerated into "lockhandlehandle...". Backing the CPU primary tier per node sidesteps this by construction, since store/load jobs are broadcast to every rank. Detail in #57.

Safety

  • Per-file fail-closed on fresh source: missing/nonunique anchors leave that file unchanged; syntax failure restores that file and exits nonzero. Earlier files are not rolled back if a later file fails. Existing marker presence short-circuits revalidation; this is not an all-file transaction or canonical validation of already-marked source.
  • Idempotent, marked with # [glm53-kv-offload-scratch].
  • Kill switch: GLM53_SKIP_KV_OFFLOAD_GROUPS_PATCH=1.
  • Inert unless enabled — without GLM53_OFFLOAD_MMAP_DIR and the --kv-transfer-config flag, behaviour is unchanged.
  • tests/test_kv_offload_groups.py covers idempotency, fail-closed drift, the kill switch, and (inside the image) that every anchor is still present. The original 30 passed suite claim is historical, not the current review result.

Notes for anyone testing

  • Keep PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True and add --enable-cumem-allocator. The launcher uses its existing caller override: unset selects that default; explicit empty disables it. Do not disable expandable segments merely to bypass the connector guard.
  • Do not set PYTHONHASHSEED (the tier docs suggest it for cross-process key stability) — it invalidates the Triton JIT cache keys and the mid-collective recompile kills the worker rank.
  • Sizing: a restore is admitted only when every group's chunks are resident together, so keep cpu_bytes_to_use ≥ row_bytes × (restore_tokens / 3584 + 8). Undersized, stores still succeed and metrics still climb while no restore ever serves.
  • On GB10, bound vm.dirty_bytes (we use 4 GB / 1 GB background). Unbounded dirty pages livelocked both our Sparks twice; drop_caches can't help because the pages are dirty.

Tested on 2× DGX Spark (GB10, 121 GiB unified), TP=2 over CX7, DFlash2 k=7, KV_CACHE_DTYPE=fp8, at 400k and 1M context.

Maintainer source review — publication under HOLD

Published source head 6662b7d7aa5837340be6f3299e5310abf616c015 preserves author history through normal main merges and a fast-forward publication; integrated base is a9bbbadc4c7253870d0d519c3d445e4ee7b97540, not current main. Maintainer plotarmordev owns current-main integration and any later qualification scheduling; SIX-GlmLandingE alone owns the merge lane. merge-hold remains.

Corrected both allocator recipes, retained main caller precedence and replay/retention wiring, and added the missing KVOFF artifact-table/preflight admission. Main's supplied receipts report launcher parity passing and four offload tests passing on copies of all five Python sources from the immutable wheel-layer pin; absent runtime source explicitly skips its source-dependent test. These establish patch/compile/idempotence only, not allocator behavior or eviction/restore qualification. The final commit changes README historical scoping only.

Unresolved source contract: GLM53_OFFLOAD_RELEASE_BYTES is forwarded without launcher pre-stop validation. Runtime uses int(value or 0): empty/zero disables release, noninteger raises ValueError during worker construction, and negative values are accepted and trigger release on each completion poll with a region. Both load and store transfer sizes contribute, so the old “bytes stored” wording is imprecise. No positive-only restriction was invented in this review. Maintainer owns deciding/documenting the supported domain and pre-stop admission before live qualification.

Unblock requires separately reviewed current-main integration and separately authorized current-image forced eviction/restore showing actual disk-to-GPU transfer, correct output, and host allocator behavior under a selected scheduling policy. Historical tables are retained as historical author observations; stores/counters alone are insufficient, and no author rebase or benchmark is requested.

Backs vLLM's CPU offload staging region with a sparse file on each node's own
NVMe, so a prefix evicted from the GPU pool is restored from disk instead of
recomputed. Capacity becomes the file size rather than RAM. Off by default;
GLM53_OFFLOAD_MMAP_DIR turns it on.

Measured here (105k-token prompt, evicted by 6 x 168k floods, 100 GB store):
restore 8.54 s vs 130 s to recompute (15x), output byte-correct, 1.68 GB read
back from disk, host stable. 160 GB/node holds ~11.3M tokens. Decode unchanged:
bench_decode median of 5x400 gives structured 66.4 / prose 26.7.

overlay/patch_kv_offload_groups.py fixes what blocks offload on this model
(see MiaAI-Lab#57):

- the hash-alignment assert aborts the boot, because KpoolTailSpec's 4-token
  block can never align to the 64-token hash granularity
- sub-block groups (kpool tail, DFlash2 drafter) emit ~56x more keys than the
  3584-token groups and starve the staging tier, so MLA/mamba chunks never
  reach disk; both are excluded. Drafter KV rebuilds in its own window after a
  restore and the target verifies every token, so this costs acceptance at
  worst, never correctness.
- the packed whole-row copy path clobbers other groups' live pages on restore
  for a multi-block-size hybrid; the per-layer path is coherent and moves ~10x
  less data
- the staging region's page cache is bounded with msync + MADV_DONTNEED, which
  GB10 needs: CPU and GPU share one pool, and an unbounded region starves the
  GPU until the worker rank dies mid-prefill with no traceback

Per-node rather than an fs secondary tier because those run scheduler-side on
the head only: under 2-node TP rank 1's shard is never written and its restores
return stale bytes as confidently wrong output, with healthy-looking counters.
Backing the CPU primary tier per node is correct by construction.

The patch is fail-closed (a drifted anchor patches nothing and exits non-zero),
idempotent, and skippable via GLM53_SKIP_KV_OFFLOAD_GROUPS_PATCH=1. Wired into
start.sh on both ranks alongside the existing overlay patches.
@jdecker76

Copy link
Copy Markdown
Author

⚠️ Important addition: offload forces expandable_segments:False, which leaks badly on GB10

Running this in production surfaced something that belongs in the PR, and anyone enabling offload on a DGX Spark needs it.

vLLM's OffloadingConnector rejects PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (config/vllm.py, the "KV connector ... is incompatible with" check). So enabling offload silently drops you onto expandable_segments:False — and this repo's start.sh defaults to True for good reason.

Without expandable segments, PyTorch's caching allocator retains reserved blocks after a large prefill and never returns them. On GB10's unified memory that is not a slow leak — it is fatal:

same 337k-token request expandable_segments:False expandable_segments:True
host available at start 10358 MB 10358 MB
at t+60s 9244 MB 9682 MB
at t+240s 1352 MB ~10098 MB
outcome never completed (600 s timeout), host wedged completed 397 s, memory recovered

Memory did not return after the request finished (num_requests_running=0, still pinned at ~5 GB below baseline), which is the caching-allocator-retention signature rather than page cache or a mapping leak. We wedged both Sparks six times before isolating it — SSH and SSM die and it needs a physical power cycle, so it is worth flagging loudly.

The fix (upstream's own exemption, not a hack)

The same guard exempts the CuMem allocator, because CuMemAllocator.use_memory_pool disables expandable segments around the KV pool so KV lands on stable physical pages — which is precisely what the guard protects (registered/pinned KV for RDMA connectors like NIXL/Mooncake). So the working combination is:

GLM53_ALLOC_CONF=expandable_segments:True   # or just leave start.sh's default
# and add to the serve flags:
--enable-cumem-allocator

Verified live: offload active (OffloadingConnector, primary tier 3162 blocks), expandable_segments:True, 12.6 GB stored to the disk tier, and host memory flat through a 337k-token prefill.

I'll fold this into the PR (recipe comment + .env.example note + README section) so the offload path ships with the allocator guidance attached rather than as a footgun. Flagging it here first in case you'd rather handle the interaction differently — e.g. having start.sh set --enable-cumem-allocator automatically whenever --kv-transfer-config is present, which would make the combination impossible to get wrong.

…e on GB10)

The OffloadingConnector rejects PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,
so enabling offload pushes you onto expandable_segments:False — and this repo's
start.sh defaults to True for good reason. Without expandable segments the
caching allocator keeps reserved blocks after a large prefill and never returns
them; on GB10's unified memory that starves the host rather than merely
fragmenting it.

Measured on 2x DGX Spark, same 337k-token request:
  expandable_segments:False -> host available 10358 MB -> 1352 MB, no recovery
                               after the request finished, never completed
                               (600 s timeout), both Sparks wedged
  expandable_segments:True  -> dipped ~560 MB, completed in 397 s, recovered

Six wedges before this was isolated, each needing a physical power cycle, so it
is documented loudly rather than as a footnote.

The fix is the guard's own exemption: keep expandable_segments:True and pass
--enable-cumem-allocator. CuMemAllocator disables expandable segments around the
KV pool, so KV lands on stable physical pages -- which is what the guard exists
to protect (registered KV for RDMA connectors) -- while the rest of the engine
keeps expandable segments. Verified live with offload active: OffloadingConnector
loaded, primary tier 3162 blocks, 12.6 GB stored to the disk tier, host memory
flat through a 337k-token prefill.
@jdecker76

Copy link
Copy Markdown
Author

Status update: offload verified working — and this PR depends on #80

Since opening this, we ran it in production and then to a proper restore proof. Two things worth adding here.

It works, with numbers

169,328-token prompt, evicted from a 1,096,618-token GPU cache by 9 floods (1.9M tokens), then re-asked:

cold             228.7 s
GPU prefix hit     3.5 s
restored from disk 10.4 s   <- 22x faster than recompute, output byte-correct
second ask         2.3 s
disk->GPU        0.00 GB -> 2.66 GB
external hits          0 -> 164,864

Host memory flat ~10.4 GB across the flood. Decode unaffected: structured 65.63 tok/s, prose 27.91 (vs 65.1 / 27.1 published).

⚠️ This PR needs #80 to be usable

I previously reported that stores worked but restores never fired. That was a misdiagnosis. The cause is GLM53_MIXED_PREFILL_CHUNK=skip: it holds every prefill while any lane decodes, so requests sit as deferred and it looks exactly like the connector's load path failing. With #80 (mixed-prefill gate v2) applied, restores fire immediately and correctly.

Recommend treating #80 as a prerequisite — or merging them together. Offload without it will look broken to anyone who tries it, in a way that points at the wrong component. I lost two days to that.

Also required: the allocator pairing

Already documented in the .env.example/README commit here, restating because it is the part that bites hardest: the connector rejects expandable_segments:True, and running False is fatal on GB10 (10358 MB → 1352 MB on one 337k prefill, no recovery, both Sparks wedged). The fix is the guard's own exemption — expandable_segments:True plus --enable-cumem-allocator.

If you'd rather this be impossible to get wrong, start.sh could add --enable-cumem-allocator automatically whenever --kv-transfer-config is present. Happy to add that here if you want it.

One correction to my earlier claim

kv_offload_cpu_cache_usage_perc reads 0.0 even while restores succeed, so it is cosmetic on this build rather than a health signal. Anything I said implying otherwise was wrong.

Credit: the gate interaction surfaced via Reederey87's v1.3.2 release.

@jdecker76

Copy link
Copy Markdown
Author

Status update

Offload has now run in production on our kit and been root-caused properly. Full detail in #57, but the parts that matter for this PR:

The docs in this PR are still correct and still necessary. The allocator pairing (expandable_segments:True + --enable-cumem-allocator) remains mandatory — running with expandable segments off is fatal on GB10.

One claim of mine that this PR's discussion inherited was wrong. I reported that the restore path stalls the engine. It does not. That was queue wait behind another session's long chunked prefill (root cause + isolation in #110), and the restore itself fires in seconds — measured 318,976 tokens restored from disk, and 10.4 s vs 228.7 s cold on an earlier test. kv_offload_load_time_total was 70.39 s across a whole 90-minute run.

Worth adding to the docs if you take this PR: once #110's scheduling fix and #86's right-sizing are both in, offload largely stops firing — our final two-session workload wrote 80.12 GB to the store and restored 0.00 GB, because nothing was ever evicted. Offload becomes insurance rather than a throughput feature, and anyone benchmarking it needs to force an eviction first or they will measure nothing.

Happy to rebase this on current main if you want it — it has moved a long way since I opened this (E2 kernel, right-sizing, spinwait, TP4).

@smartprogrammer93

Copy link
Copy Markdown

@plotarmordev plotarmordev added the triage/P2 Community triage 2026-09-09: useful scoped work, schedule after P0/P1 label Sep 9, 2026
@plotarmordev plotarmordev added P2 Important scoped work; schedule after P0 and P1 merge-hold Do not merge until the stated maintainer dependency or hold is cleared. needs-final-review Prior audit is recorded; current-head review and merge disposition still need finalization. reviewed Our source review or technical verdict is recorded; not merge approval or runtime qualification. awaiting-test-evidence Reviewed; required exact-head correctness or performance evidence is still outstanding. and removed triage/P2 Community triage 2026-09-09: useful scoped work, schedule after P0/P1 needs-final-review Prior audit is recorded; current-head review and merge disposition still need finalization. labels Sep 12, 2026
@plotarmordev

Copy link
Copy Markdown
Collaborator

Review follow-up — ea645b228ea7

Your restore measurements, allocator-pairing findings, correction of the earlier stall attribution, and offer to rebase are acknowledged. We have not requested another measurement; the awaiting-test-evidence label was premature and is being removed.

The next decision is ours: whether to schedule disk offload now, and which scheduling policy to qualify it with. Please do not spend time rebasing until we make that call. If scheduled, we would test forced eviction/restore and allocator behavior on the current image, using TheGrill for the measurement record; that is not an outstanding request to you.

One documentation correction is worth making now: the PR body's testing notes still recommend expandable_segments:False, contradicting your later safety findings. Please align the body with the allocator pairing described in your updates.

Unless explicitly noted above, this update is a source/discussion review, not a new live qualification run.

@plotarmordev plotarmordev removed the awaiting-test-evidence Reviewed; required exact-head correctness or performance evidence is still outstanding. label Sep 15, 2026
@plotarmordev

Copy link
Copy Markdown
Collaborator

Maintainer direction — pursue the experimental offload path

We have selected #58 as the offload path to pursue, rather than adopting #99's store-only stage without a usable hybrid restore path. This is a direction decision, not current-image qualification, merge approval, or a scheduled live run. Thank you @jdecker76 for the per-node primary-tier design, restore measurements, allocator findings, and corrected stall attribution; those receipts remain historical evidence at their recorded identities.

The next work is maintainer-owned: correct both the PR body's testing note and README quickstart, then integrate the mechanism into canonical current main before deciding any live qualification schedule. The copy-paste quickstart at this head still recommends expandable_segments:False, contradicting the later safety note. The accepted integration uses main's PYTORCH_CUDA_ALLOC_CONF caller override, including its explicit-empty semantics, rather than retaining a parallel GLM53_ALLOC_CONF knob; offload guidance pairs expandable_segments:True with --enable-cumem-allocator.

We also own selecting and documenting the mixed-prefill policy for that integration: main's skip policy can defer prefills behind decode, so queue delay must not be reported as a restore failure. Preserve current replay/retention and both-rank behavior. Only after the documentation and current-main integration are reviewed will maintainers decide a separate forced-eviction/restore and allocator qualification window; stores or healthy counters alone do not establish a successful restore.

Please do not rebase or produce new measurements now. This supersedes the earlier author-side documentation request: maintainer @plotarmordev owns the correction and integration, and the PR stays open and unqualified in the meantime.

The quickstart and .env.example told readers GLM53_ALLOC_CONF=expandable_segments:False,
contradicting this PR's own warning block (False is fatal on GB10) and duplicating
main's PYTORCH_CUDA_ALLOC_CONF caller override with worse empty-value semantics
(:- falls back on empty; main's - honors an explicit empty as disable).

- README quickstart and .env.example: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  plus --enable-cumem-allocator in the offload serve flags; state the unset-only
  default and the explicit-empty disable instead of a second knob.
- .env.example: quote the EXTRA_ARGS example. .env is sourced as shell code, so the
  unquoted JSON ran `kv_connector:OffloadingConnector` as a command and killed the
  launcher under set -euo pipefail.
- Register GLM53_OFFLOAD_MMAP_DIR / GLM53_OFFLOAD_RELEASE_BYTES in the README knob
  table and scope the section to the TP=2 launcher (start-tp3.sh / start-tp4.sh do
  not forward these knobs or mount the overlay).

Docs only; no runtime behavior change.

@plotarmordev plotarmordev left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Source/publication review at 6662b7d7aa5837340be6f3299e5310abf616c015: fast-forward publication and updated description verified; merge-hold remains. Experimental offload is not approved for merge: current-main integration, RELEASE_BYTES admission/domain decision, and separately authorized forced-eviction/restore plus allocator qualification remain maintainer plotarmordev-owned. Four passing staged-source tests are not runtime qualification. No author rebase, download or benchmark requested; no merge performed.

@plotarmordev

Copy link
Copy Markdown
Collaborator

Documentation correction published — HOLD remains

Exact head 0bef1aa251b15036a85cff249125b0769b532190 is a normal fast-forward from 6662b7d7aa5837340be6f3299e5310abf616c015; only .env.example and README.md changed. They now describe completed store and load bytes, unset/empty/zero disabling release, Python-integer parsing (negative values cause release on each completion poll with a region; nonintegers fail worker construction), and absent launcher pre-stop validation. The historical 2 GB observation is not current-image qualification. No positive-only semantics or production behavior were added.

Next-action owner: maintainer plotarmordev. Decide the supported release-threshold domain and pre-stop handling before commissioning an authorized current-image forced-eviction/restore and allocator window. Unblock requires reviewed current-main integration and accepted actual disk-to-GPU restore, output-correctness, and host-memory evidence under the selected recipe; stores/counters alone do not qualify it. Existing marker short-circuit and per-file rather than transactional patching limits remain explicit in the description.

Main's prior four pinned-source patch/compile/idempotence checks apply to unchanged code, not runtime qualification. No tests, live operations, approval, or merge were performed in this documentation review. Experimental merge-hold remains; no immediate author rebase or benchmark request.

@plotarmordev plotarmordev left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hold update

Public head checked: 0bef1aa251b15036a85cff249125b0769b532190.

Why it is held: We observed disk-backed storage activity, but did not demonstrate a successful restore from it.

What we tested: The CPU/numeric checks and bounded shared-memory and file-backed probes passed their answer checks. The file-backed mappings gained allocated data. However, the repeated request reported zero cached tokens and no additional measured block-device reads. Allocation or writes alone are not restore proof; zero device reads alone also does not distinguish all page-cache behavior.

Next action — me, @plotarmordev: Finish integration and prepare a test that identifies the actual restore path. The test must store known data, remove resident-cache explanations within the supported design, observe restore-specific counters/events, and check the returned data/answers. Record source and runtime identities and both participating ranks. Use disposable storage, not live cache deletion.

Author action now: No immediate rebase or rerun. We own this qualification step. Mark ready only after the selected final version demonstrates the claimed restore behavior; the current result is incomplete evidence, not a proven restore failure.

Please keep the hold until the conditions above are met. When a new version is ready for review, include its commit ID, the checks actually run, their results, and anything not tested. Keep failed or incomplete runs in the report. Share only sanitized results—not private prompts, credentials, or machine configuration.

@plotarmordev

Copy link
Copy Markdown
Collaborator

Maintainer priority triage — 2026-09-26

Reviewed head: 0bef1aa251b15036a85cff249125b0769b532190.

P2, retained alongside #232 pending a maintainer integration choice. The previous selection of this path is not erased merely because a newer PR exists; conversely, #232 now reports restart-restore receipts worth evaluating. Next owner: @plotarmordev decides the supported offload architecture and restore-validation contract before requesting either author to rebase or rerun. Unblocks on that decision and the chosen path proving actual restored cache hits, safe invalidation, and integration with current-main cache behavior. Existing missing-restore evidence here is not relabeled as passing.

This is prioritization only; no merge or implementation work is authorized by it.

@plotarmordev

Copy link
Copy Markdown
Collaborator

Closing in favor of #232, which we have now selected as the KV-offload direction. When we chose this PR on 2026-09-15, its unblock condition was an accepted disk-to-GPU restore on current main; that was never shown here. #232 reports a restore after restart, so we are consolidating on one offload path instead of maintaining two. Thank you @jdecker76: your per-node primary-tier design, allocator findings (expandable_segments / cumem) and corrected stall attribution remain the reference for that work.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

merge-hold Do not merge until the stated maintainer dependency or hold is cleared. P2 Important scoped work; schedule after P0 and P1 reviewed Our source review or technical verdict is recorded; not merge approval or runtime qualification.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants