PR2: entropy-coded cold pool (window policy) - #2
Open
Astrangemaninhere wants to merge 11 commits into
Open
Conversation
added 11 commits
August 30, 2026 20:58
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).
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
1 participant
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.
PR2: entropy-coded cold pool (window policy)
Background
Long contexts keep every prefix page physically resident even when only the recent window is actively attended to. This PR adds a cold-pool window mechanism: retired prefix pages compress into fixed raw slots and return their physical pages to the pool.
Changes
--cold-keep-tokensrequantize + pack into slots, publish a sentinel block-table entry, and dematerialize the device replica back to the pool.Files
src/ops/launcher/cold_i8.cu/src/ops/kernel/cold_i8_kernels.cuh(slot codec)src/ops/kernel/entropy_cold_requant_kernels.cuh(INT8 -> E2M1 g64 requant)src/targets/qwen3_6/impl/runtime/program_impl.h(egress / restore / warm paths)src/targets/qwen3_6/impl/runtime/logical_kv_store.h(cold transfer bookkeeping)Known limitations
Dependency
Depends on PR1 (per-layer planes feed the requant). Merge after PR1.