feat(kvcache): host-resident prefix tier for hybrid GDN models (snapshots + KV pages) - #525
Open
alvarorsouza-arch wants to merge 3 commits into
Open
alvarorsouza-arch wants to merge 3 commits into
alvarorsouza-arch wants to merge 3 commits into
Conversation
- mirror_store_quant (new Triton kernel, graph-safe UVA): the write-through mirror now copies e4m3 codes + fp32 scales read from the pool AFTER quantization (the pool is the source of truth). The bf16-era mirror corrupted fp8 pages silently on rehydrate. - mirror_scales host bank + fused-copy descriptor extended with the scale banks (per-bank row sizes). - set_select_width takes max_bs: the whole batch selection must fit the physical pool by construction (fixes assert-crash on multi-row decode with small pools). - measured: 299,648 logical tokens with kv-cache-dtype fp8 + kv-host-pages, needle at 253k tokens answered correctly (~1026 t/s), quality suite identical to baseline.
Prefix reuse on hybrid GDN models needs KV pages AND a GDN state snapshot alive at the match point; the small global snapshot pool gets drained by interleaved auxiliary calls (agent title generators) and every match truncates to zero. Measured on real WorkBuddy agent turns: TTFT 20-150s recompute per message. This adds a host-resident tier: evicted GDN snapshots and KV prefix pages move to pinned RAM instead of dying. Matches rehydrate from host in ~4s (24k-token prefix) instead of a full recompute. - HybridRadixCache hooks: on_evict_node / on_match_dead_snapshot; evicted leaves become tombstones instead of cascade-freeing KV - HostPrefixTier: pinned GDN snapshot bank + cold-prefix index with its own KV slot space (logical ids get reused by the scheduler, so content is copied D2H on eviction, never referenced by id) - Admission rehydration; with KVHostOffloader active the copy is host-to-host into its mirror (no bulk H2D; lazy fetch covers residency) - Eviction copies run in 128-page slices (staging the whole prefix OOMed exactly under memory pressure, which is when eviction runs) Measured: churn test cached-token 0 -> 23,936 (16.5s -> 3.8s); total pool flush rehydrates at 24,000 cached tokens in 4.4s; quality suite identical (39/40 QA, 30/30 cloze); decode steady 33-37 t/s. All off by default (FT_GDN_HOST_TIER / FT_PREFIX_HOST envs).
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 19, 2026
…ybrid GDN models (snapshots + KV pages)
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 19, 2026
…cales they cannot mirror Upstream FlashML-org#499/FlashML-org#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4 pool (FlashML-org#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable. Assisted-by: Claude Fable 5.1
This was referenced Sep 19, 2026
This branch has not been deployed
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.
Note: stacked on #499 (KV host-offload). The GDN-snapshot tier works standalone; the KV-page tier composes with the offloader.
Problem
On hybrid GDN models (qwen4_exp), prefix reuse needs two things alive at once: the KV pages in the radix tree and a GDN state snapshot at that point. The snapshot pool is small and global LRU. Any auxiliary call (WorkBuddy's title generator is our real-world case) drains the main conversation's snapshots between turns, and the match truncates to zero. Measured before the patch: agent turns with 20 to 150 s TTFT. After: stable single-digit seconds.
What the patch does
HybridRadixCache:on_evict_node(fired before unlink/free) andon_match_dead_snapshot(fired during the match walk). With the tier on, an evicted leaf becomes a tombstone and its KV is no longer cascade-freed.HostPrefixTier(new,kvcache/host_prefix_tier.py): a pinned host bank of GDN snapshots plus an index of cold prefixes with its own KV slot space. Scheduler logical page ids get reused, so content is copied D2H on eviction and never referenced by logical id afterwards.Flags:
FT_GDN_HOST_TIER=1,FT_GDN_HOST_SLOTS=32,FT_PREFIX_HOST=1,FT_PREFIX_HOST_GB=16. All off by default.Measured numbers (RTX 4070 Ti SUPER 16 GB, 251 GB RAM, PCIe Gen4)
Churn test (24k-token conversation, 10 auxiliary calls in between, then back):
Total flush test (two ~130k-token prefills through a 200k pool; every page of the conversation leaves the GPU): the conversation comes back with cached-token 24,000 in 4.4 s.
Regressions: our quality suite (40 known-answer questions + 30 cloze) is identical to baseline (39/40 both). Decode steady at 33-37 t/s.
Known limitations
Files
kvcache/host_prefix_tier.py,tests/kvcache/test_host_prefix_tier.py,tests/kvcache/test_tree_hooks.py,tests/kvcache/test_tier3.pykvcache/hybrid_radix_cache.py,scheduler/cache.py,engine/config.py