PR6: NVFP4-tier KV + entropy cold pool (10L default table, +33% context ceiling) - #6
Open
Astrangemaninhere wants to merge 48 commits into
Open
PR6: NVFP4-tier KV + entropy cold pool (10L default table, +33% context ceiling)#6Astrangemaninhere wants to merge 48 commits into
Astrangemaninhere wants to merge 48 commits into
Conversation
added 30 commits
August 30, 2026 17:41
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).
…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.
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.
PR6: NVFP4-tier KV + entropy cold pool (10L default table, +33% context ceiling)
Background
The integrated NVFP4 work on top of PR1-5: NVFP4 (E2M1 + E4M3 g16 scales, ISO3 V) KV tier, GQA attention ops, a data-driven default layer table, and an entropy-coded cold pool that finally shrinks the device KV footprint.
Changes
GQA attention (batch 1) —
src/ops/kernel/gqa_attention_*:gqa_attentionreplaces the d256 causal attention path for main text.Default layer table —
src/targets/qwen3_6_27b/impl/variant.cpp:{0,1,3,4,6,7,8,9,13,14}, 6 layers INT8{2,5,10,11,12,15}.analyze_kv.pymetrics were shown not to predict combination outcomes.NVFP4-tier cold pool — entropy rANS slots:
Cold-pool fixes found while validating:
cold_i8pack/restore hard-coded the 9232 B slot stride, corrupting heads past slot 0 once the buffer grew; both kernels now takeslot_bytes.can_cold_transferrequiredwriter_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.Upstream sync:
fix(runtime): bound shared capture target expansioncherry-picked.Measured results (RTX 5090D 32 GB + 8 GB host, YaRN + MTP, --cold-policy window)
Needle recall verified at every point; ceiling is the device memory wall (380k needs 11.79 GiB > 11.60 GiB available).
Known limitations
--cold-policy host) were implemented and functionally verified (needle recall with 48 pages compressed to pinned host memory) but reverted: the device page pool stays at the logical capacity, so host slots do not reduce device memory. Shrinking the pool requires prefill to allocate pages chunk-progressively (currently it plans the full request up front), which is a larger scheduling change.cold checkpoint page has no source bookkeeping); active-source reuse works (restore inventory fix), single-sequence operation unaffected.--cold-policy hostis incompatible with speculative backends (their KV must stay resident).Dependency
Depends on PR1-5. Merge last.