feat(runtime): NVMe cold tier (ColdPolicy::Disk) under the cold-pool planner - #149
Closed
Astrangemaninhere wants to merge 45 commits into
Closed
feat(runtime): NVMe cold tier (ColdPolicy::Disk) under the cold-pool planner#149Astrangemaninhere wants to merge 45 commits into
Astrangemaninhere wants to merge 45 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d85c8461be
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Astrangemaninhere
force-pushed
the
disk-cold-tier
branch
3 times, most recently
from
September 1, 2026 01:57
08b84f9 to
1f40594
Compare
Author
|
All six review findings are addressed in the updated commit (1f40594). Functional verification on the 5090/WSL2 box with
Per finding:
Thanks for the review. |
added 25 commits
September 1, 2026 10:10
Fixed raw slots (9232 B: header + E2M1 nibbles + E4M3 g16 scales) hold requantized cold pages for both the INT8 and NVFP4 tiers. Requantizing INT8 planes to g64 E2M1 measures NMSE 0.012-0.014 (inside the accepted NVFP4-layer envelope) at 1.85-1.99x per head-page, ~1.66x aggregate cold KV on the 27B production table. The pack/restore kernels, the per-layer dtype dispatch, the decode and prefill cold staging (inline nibble->int8 adapter preserving the int8 QK tensor cores), and the length-based slot sizing are all included; --cold-policy window|host plus --cold-keep-tokens/--cold-host-bytes control activation. Three latent v1 cold-addressing bugs (compress_page slot scaling, decode and prefill flat slot indices) are fixed on the way.
Fixed raw slots (9232 B: header + E2M1 nibbles + E4M3 g16 scales) hold requantized cold pages for both the INT8 and NVFP4 tiers. Requantizing INT8 planes to g64 E2M1 measures NMSE 0.012-0.014 (inside the accepted NVFP4-layer envelope) at 1.85-1.99x per head-page, ~1.66x aggregate cold KV on the 27B production table. The pack/restore kernels, the per-layer dtype dispatch, the decode and prefill cold staging (inline nibble->int8 adapter preserving the int8 QK tensor cores), and the length-based slot sizing are all included; --cold-policy window|host plus --cold-keep-tokens/--cold-host-bytes control activation. Three latent v1 cold-addressing bugs (compress_page slot scaling, decode and prefill flat slot indices) are fixed on the way.
…only) The paged-KV cold mechanism (sentinel pages, slot pool, compress) does not exist in upstream master yet; the cold-compress pass and its member state are removed until that mechanism lands in a follow-up PR. This PR keeps the entropy codec ops, the ColdPolicy option surface, the per-layer KV plumbing, and the op tests.
Cold slots are allocated as per-layer regions (9232 B raw slots + I32 validity) by the decoder state; the pool exposes allocate/ release with a used bitmap. A decode-boundary pass packs the retired tail (valid - cold_keep_tokens) of a sequence's text KV into raw entropy slots, shrinks the address space entitlement, returns the physical pages to the pool, and publishes block-table sentinels (entry <= -2, slot base = -2 - entry) via publish_indices. Attention producers decode sentinel pages inline from the slots in the cold staging branches (INT8 adapter preserves int8 QK cores; NVFP4 tier keeps its native E2M1/ISO3 nibble semantics).
added 20 commits
September 1, 2026 10:10
…ensitive layers to I8)
…ctx8k, parity 8L)
Measured 13.3k zh perplexity at ctx 4096 and 8192. 8/8 split
(NVFP4 {0,1,3,4,6,7,8,9}, I8 {2,5,10-15}) scores 1.335 vs all-I8
1.432 at ctx 4096; demoting L13/14 to NVFP4 is free (10L 1.3344,
11L 1.3385, 12L 1.345, 14L 1.375 - all below all-I8). Page pool
shrinks another 12.5%. L2/L5 are the true outliers (demote alone
+0.20/+0.09); the old 10L 0-4,6-9,15 combo degraded to 1.551 by
demoting both. analyze_kv.py NMSE/rank metrics do NOT predict these
combo outcomes - neighborhood marginals do.
The cold pool only carried INT8 planes; NVFP4 layers (10 of 16 in the default table) skipped cold transfer entirely, so the page pool could not shrink on the layers that need it most. This wires the previously orphaned page-slot rANS codec (entropy_nvfp4_slot, ported in batch 1) into the egress/restore/decode paths: - egress: NVFP4 layers requantize K (Nvfp4G16) and V (Iso3VG16) to g64 scales and rANS-encode into entropy slots; INT8 layers keep the raw 9232 B nibble slots. Per-layer dtype dispatch replaces the all-I8 gate. - restore: rANS-decode + scale-tail scatter back into native planes. - slot buffer: one 9536 B size serves both codecs (rANS max = 320 B header + 32x256 B streams + 1024 B scale tail). Fixed pre-existing cold-pool bugs found while validating: - cold_i8 pack/restore hard-coded the 9232 B slot stride, so heads past slot 0 landed at the wrong offset once the buffer grew to 9536 B; both kernels now take slot_bytes. - can_cold_transfer required writer_references == 0, which the paged store never clears for live pages, so cold transfer never ran; the gate now only requires a device replica and no pins/fork ties (cold pages sit before the decode frontier and are never written again). Verified: 13.1k-token zh needle (紫电青霜) recalled identically with the cold pool compressing 17/48 prefix pages, INT8-only and mixed 10L tables; cold off matches cold on exactly.
The --cold-policy host tier was parsed but never wired: cold slots always lived on device, so the device KV footprint never shrank and the 8 GB host budget (host-kv) sat unused as a backup-only copy. Host mode allocates the cold-slot payload buffers with cudaHostAllocMapped (one pinned allocation per layer) instead of device tensors. Pack kernels write them through UVA; decode kernels read them back over PCIe. Only the validity plane stays device-resident. The cold-pool policy gates (egress, warm-restore, decode trigger, scratch allocation) now accept Host alongside Window. Functionally verified: 13.1k-token zh needle recalled identically with 48 prefix pages compressed into host slots; device runtime for the cold payload drops to zero (only hot pages + validity planes occupy device memory), which is what lets the context ceiling exceed the device KV budget.
…t memory" This reverts commit d8854a4.
device_kv_prefix_pages, shared_device_kv_prefix_pages and
missing_shared_device_kv_prefix_pages treated cold-compressed pages as
non-resident, so a prefix reuse across sequences (request B reusing
request A's checkpoint) computed an inconsistent restore inventory and
aborted with 'Text KV restore inventory is inconsistent'. Cold pages
restore in place from their raw slots (missing_kv_restore already skips
them), so the resident counts must include them.
Known limitation (documented, not fixed): reusing a *catalogued* (shared)
checkpoint that contains cold pages still fails ('cold checkpoint page has
no source bookkeeping') because the catalog entry carries no cold-slot
record; single-sequence cold-pool operation is unaffected. Concurrent
cold-pool reuse of an active source works.
Two DFlash2 defects found while validating the DFlash2/MTP switching path:
1. CUDA-graph preparation segfaulted with --spec dflash2. The graph
representative (prepare_representative) initialized the DFlash host
ingress/egress unconditionally under io.dflash_decode, but DFlash2 keeps
its own dflash2_host_ingress/egress buffers and dflash_host_ingress is
null when only DFlash2 is active, so *dflash_host_ingress = {} crashed.
The block now selects ingress/egress by backend, and DFlash2's
pending_features are zeroed like DFlash's. Verified: dflash2 CLI and
serve generation run (16 tok, 116 tok/s serve decode).
2. --spec auto failed to load ('loaded weights do not match the frozen
startup features') because only plan_load resolved auto internally;
make_sequence_planner and construct_loaded_model still saw
SpeculativeBackend::Auto, so the planner built an Auto-features plan and
the startup-features consistency check rejected the loaded weights.
resolved_auto_speculative is now a public Package static and the
registry resolves auto once up front, passing the concrete options
through plan_load, the planner, the loaded model and the instance. The
35b target gets the same hook (auto -> MTP, no DFlash2 weights).
Verified: dflash2 artifact with --spec auto serves 360k context (MTP),
explicit --spec dflash2 serves short context at 116 tok/s.
--spec auto previously chose DFlash2 only when max_context fit the draft capacity, and otherwise fell back to MTP. For a DFlash2 artifact that fallback is wrong: the artifact has no MTP draft head (the two are mutually exclusive), so auto at a long context selected MTP and failed to load with a confusing weights mismatch. auto now keys purely off the artifact weights profile: a DFlash2 artifact always picks DFlash2 (a memory shortfall surfaces as a clear reservation error), any other artifact defaults to MTP. Verified: base artifact + auto serves 360k (MTP); DFlash2 artifact + auto serves 16k with speculative=dflash2.
…planner Cold pages spill to per-layer disk files: the device slot pool becomes a working set, egress mirrors compressed slots (file offset = slot * stride), restore reads back before decode. The warm path prefetches pending slots on the transfer stream with double-buffered staging so H2D overlaps decode. - ColdPolicy::Disk + --cold-policy disk / --cold-disk-path / --cold-disk-bytes - per-layer spill files opened once, closed on teardown - prefetch_cold_pages: async H2D prefetch in warm_cold_prefix - egress/restore keep the slot codecs (rANS / raw) byte-identical Responds to upstream Neroued#143 (NVMe cold tier under the Device/Host planner).
Astrangemaninhere
force-pushed
the
disk-cold-tier
branch
from
September 1, 2026 02:10
1f40594 to
4e9f5d5
Compare
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
2 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.
NVMe cold tier under the cold-pool planner (ColdPolicy::Disk)
Adds a disk placement to the existing cold-pool degradation chain:
Device (hot) -> pinned Host (warm) -> NVMe file (cold) -> evict.
What it does
ColdPolicy::Disk+--cold-policy diskwith--cold-disk-pathand--cold-disk-bytes(default 32 GiB). Empty path uses the system temp dir(WSL2: keep it on ext4, not
/mnt/c).(compressed bytes, rANS or raw codec unchanged), so the file offset is
slot * stride. The device slot pool becomes a working set; capacity isbounded by the disk budget instead of
max_cold_pages.file (D2H via the device stream, then write). Restore reads the file back
before the decode kernels run.
warm_cold_prefixprefetches every pending slot on thetransfer stream with double-buffered pinned staging, so the H2D legs of the
warm path overlap with the previous decode step; the decode stream waits
once before restoring the batch.
restore_cold_pagegains adisk_prefetchedflag; the checkpoint/rewritepath falls back to the synchronous read.
Why
Responds to the design in upstream #143 (NVMe cold tier under the
Device/Host planner). Same shape as vLLM/LMCache, TRT-LLM host offload and
NVIDIA ICMSP: more parked conversations survive (reuse probability),
not higher live decode concurrency. Restore of a few GiB from NVMe is far
cheaper than re-prefilling an 80-120k agent chat (upstream #73 measured
host restore at 422 ms vs 53 s prefill).
Notes
device pool (no approximate / block-hash cache).
materialize-to-device, then run.
Verified: builds clean, baseline ppl unchanged (1.3353 on the zh corpus),
disk spill/restore exercised with
--cold-policy disk --cold-disk-path /tmp.Verification status
--cold-policy disk --cold-disk-path /tmp):egress compressed 17 prefix pages and mirrored them (1,296,896 B/layer,
exact
pages * slot_stride); restore + prefix reuse verified (4.3k-tokenprompt 2.5 s first request, 0.9 s on reuse).
(engine + ops). Limitation: the rebased branch (on latest master) has
not been built standalone yet — the same 8-file diff was built and linked
in the pre-rebase tree; a clean-master build is pending.
are pre-existing integration-tree fixes included in the diff.
maintainer is still pending (comment posted on Future: NVMe cold tier under Device/Host planner for parked continuations #143).