[Feature] DeepSeek-V4 unified KV pool (option-2): SWA↔compress share physical KV#1659
Open
yhl-amd wants to merge 9 commits into
Open
[Feature] DeepSeek-V4 unified KV pool (option-2): SWA↔compress share physical KV#1659yhl-amd wants to merge 9 commits into
yhl-amd wants to merge 9 commits into
Conversation
…euse DeepSeek-V4's paged-SWA prefill uses a window-only write optimization: each prefill chunk only persists its trailing `window` tokens to the content- addressed SWA cache (deepseek_v4.py swa_write) and only materializes the trailing-window blocks (SlidingWindowPool.ensure_for_tokens). This is correct for a request's own decode (SWA reads only the last `window` tokens), but the middle SWA blocks are never written or hash-published, so a cross-request prefix hit whose reuse boundary lands in the middle of a prior sequence gets truncated by bounded_hit(). On agentic branch/replay workloads (long shared prefixes, boundaries deep inside earlier requests) this collapses the actual server prefix-cache hit rate. Add an opt-in env flag `ATOM_SWA_FULL_RETAIN` (default off; current window-only behavior unchanged) that makes the SWA cache retain the full history: - deepseek_v4.py: prefill swa_write writes the whole chunk (max_seqlen_q) instead of min(window, max_seqlen_q). - swa_pool.py: ensure_for_tokens materializes every block of the chunk (free_before=0) so the full-chunk write has a valid physical dst and every block gets hash-published; admission_blocks reserves the chunk-peak footprint. - deepseek_v4_attn.py: swa_pool_num_blocks sizes the pool to the full per-seq history so lazily-freed-but-cached blocks survive until a replay reuses them. - block_manager.py: plumbs the flag into SlidingWindowPool. The live sliding-window free (free_out_of_window / free_after_prefill_chunk) is intentionally unchanged: out-of-window refs are still released each chunk/decode step; freeing stays lazy (hash+KV resident until the slot is overwritten), and the enlarged pool keeps freed blocks reusable. Costs ~compressed-pool-magnitude SWA memory, so it pairs with a smaller max_num_seqs / max_model_len when memory is tight. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Fix Black code-style CI: line-wrapping in the ATOM_SWA_FULL_RETAIN env entry and the swa_write write_per_batch conditional. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
The initial ATOM_SWA_FULL_RETAIN pool sizing reserved max_num_seqs * ceil(max_model_len/bs) full-resolution SWA blocks. At DeepSeek-V4 Pro's max_position_embeddings=1048576 this is ~5953 GB (96 seqs * 8192 blocks * 8.13 MB), far exceeding device memory -> allocation failure / num_kvcache_blocks collapse. That is the wrong model. vLLM's SlidingWindowSpec accounts for SWA by the LIVE per-request footprint (min(window-1 + max_num_batched_tokens, max_model_len)), never seqs*max_ctx, and gets prefix-reuse retention from the shared block pool (LRU), bounded by the total KV budget. ATOM keeps a separate SWA pool, so size it from the shared budget instead: under full-retain, allocate num_swa_blocks == num_kvcache_blocks from available_for_pool // (compressed_block_bytes + swa_block_bytes) — a 1:1 lockstep with the compressed pool (every compressed block has a paired SWA block). Live SWA footprint stays ~window/seq (window-free is still on); the remaining pool holds lazily-freed-but-cached tails, so retention capacity == the compressed prefix cache. Memory-bounded regardless of max_model_len. - model_runner.py: full-retain branch sizes both pools equally from the budget. - deepseek_v4_attn.py: drop the max_model_len-based full-retain branch in swa_pool_num_blocks (sizing now lives in the runner); remove now-unused import. Verified on DeepSeek-V4-Pro TP=8, bf16 KV, default max_model_len=1048576: num_swa_blocks=num_kvcache_blocks=14516, swa_reserved=109.86 GB (was ~5953 GB), boots with no OOM; middle-of-sequence prefix reuse still 98.4%. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
The 1:1 full-retain sizing (num_swa_blocks == num_kvcache_blocks) starves the
compressed prefix index because one SWA block is ~7x the bytes of one compressed
block: on DSV4-Pro it cut the compressed pool from ~119.5k to ~14.5k blocks, and
on the real AgentX workload that collapsed the actual hit rate (5.76%, below the
13.43% default) — the reuse bottleneck moved from "SWA tail lost" to "compressed
prefix evicted".
Give the SWA tail pool a small fraction `f` (env ATOM_SWA_TAIL_BUDGET_FRAC,
default 0.2, clamped (0,0.9)) of the KV budget; the rest stays with the
compressed pool:
swa_budget = f * available_for_pool
num_swa_blocks = swa_budget // swa_block_bytes
num_kvcache_blocks = (available_for_pool - swa_budget) // compressed_block_bytes
The SWA pool uses the same eviction discipline as vLLM's FreeKVCacheBlockQueue
(free-queue, evict oldest-freed from the front, reused-on-hit blocks re-queued at
the back), so the tail pool retains the hot-boundary working set via LRU while
compressed stays near full.
Verified DSV4-Pro TP=8, bf16 KV, default max_model_len=1048576, f=0.2:
num_swa_blocks=3315 (25.09 GB), num_kvcache_blocks=93461 (vs 14516 at 1:1, vs
119500 default) — compressed only -22% instead of -88%; boots with no OOM;
middle-of-sequence prefix reuse still 98.4%.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
Instrument the prefix-cache path to separate the two failure modes when reuse is low: reuse lost to the SWA tail gate vs reuse lost to compressed eviction. - Sequence: record num_compressed_hit_blocks (pre-gate compressed-prefix hash hit) set by BlockManager.can_allocate. - CacheStats: track total/interval compressed tokens; log per interval and cumulative "Hit", "Compressed-hit", and "Lost-to-SWA-gate" rates. A high Lost-to-SWA-gate means tails are evicted (SWA pool too small/churny); a low Compressed-hit means compressed prefixes are evicted (compressed pool too small). Guides ATOM_SWA_TAIL_BUDGET_FRAC tuning. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Dense full-retain floods ATOM's small isolated SWA pool: at Flash c=48 the SWA gate rejected ~91% of an otherwise-97% compressed prefix hit because freed tails were overwritten before a branch reused them (measured; f-tuning 0.2->0.5 barely moved it). vLLM avoids this by retaining SWA tails SPARSELY in a shared pool (VLLM_PREFIX_CACHE_RETENTION_INTERVAL / SlidingWindowManager reachable_block_mask) and hits ~97.5% on the identical BF16 workload. Port the sparse policy, adapted to ATOM's separate small SWA pool where masking alone is insufficient — the checkpoint block must be PINNED so live-window churn cannot overwrite it: - ATOM_SWA_RETENTION_INTERVAL (tokens, default 0 = dense). >0 designates a tail a checkpoint iff it is in the trailing tail_blocks of a retention_interval- sized segment OR of the prompt boundary (mirrors vLLM segment + reachable boundary tails). - publish_hash pins checkpoint blocks (extra ref so free_out_of_window / eviction skip them); non-checkpoints are freed/churned normally. LRU-capped at ATOM_SWA_CHECKPOINT_FRAC of the pool; claim_cached refreshes recency. - bounded_hit is unchanged and stays correct: it truncates a hit to the nearest present checkpoint, so a reuse boundary whose tail was evicted is safely shortened, never mis-served. Default off (interval 0) → dense behavior byte-identical. Unit-tested: checkpoints survive window-free + churn, non-checkpoints are reclaimed, LRU eviction bounded. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…dinator Add the physical allocator for the DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE, plan_atom_unified_kv_pool.md §10/§11), letting SWA and the per-type compressors (CSA k=32, HCA k=1) draw from one 128-row-slot free-list so compress can reuse space SWA freed, instead of the static swa_pages split. - UnifiedTypePool: one layer type's shared slot free-list (SWA blocks + that type's compress blocks, base-0 addressing, flydsl-safe non-negative offsets). - UnifiedKvCoordinator (option-2): single SWA slot space shared across types, per-type compress pools, with cross-role cache eviction (a compress alloc that borrows an SWA-cached slot drops the SWA hash first, or a later SWA hash-hit would double-claim a compress-held slot) and accurate has_free_swa accounting. - unified_slot_counts: single sizing source shared by BlockManager and the attention builder (flag-on tensor size == flag-off). - New env ATOM_UNIFIED_KV_SHARE (default off = byte-identical to today). - tests/test_unified_chunk_pool.py: 10 CPU tests (pool, coordinator, eviction, admission accuracy, sizing). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
…scheduler Drive the unified KV pool from the block manager when ATOM_UNIFIED_KV_SHARE is on (no-op / byte-identical when off). block_table stays LOGICAL (prefix cache and hashing unchanged); each live logical compress block is bound to a physical CSA + HCA block via the coordinator. - Sequence: add per-type physical compress tables csa_block_table / hca_block_table (parallel to block_table / swa_block_table). - BlockManager: create UnifiedKvCoordinator (flag-on, V4 only), inject it into the SWA pool as its chunk_pool, register the SWA evict callback; bind physical on true alloc / evict-reuse, retain across cache-free for lazy reuse, free on evict-reuse (when SWA can reclaim); admission checks per-type capacity. - SlidingWindowPool: draw its free-list from the injected coordinator; add evict_slot() (drop a content-cache entry when compress overwrites its slot). - scheduler: gather csa/hca block tables alongside swa_block_tables. Non-V4 / flag-off paths are untouched (free_per_req_cache_groups not touched; SWA pool disabled for non-SWA models). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
… kernels) Physical/kernel side of the unified KV pool (ATOM_UNIFIED_KV_SHARE). Flag-off keeps today's swa_pages-offset layout byte-identical. - deepseek_v4_attn.py: size unified_kv per layer type (T_L = n_type_slots*128, via unified_slot_counts; == flag-off size); _bind_compress_view binds the WHOLE tensor as base-0 compress blocks (row = phys*k), so the three compress-write paths need no kernel change; add csa/hca CpuGpuBuffers + _populate_compress_ block_tables + metadata fields; route decode/prefill index builders to the per-type table with base 0. Guard TBO+unified (not yet wired) with a clear NotImplementedError; assert block_ratio==1. - deepseek_v4.py: route the MAIN compressor write to the per-type physical table (_main_compress_bt); keep the indexer inner compressor on the logical table (it writes the separate v4_csa_idx_kv pool); csa_translate_pack reads the CSA table with base 0. - csa_translate_pack.py / paged_prefill_indices.py: swa_pages=0 in unified mode yields the absolute unified row (phys*k), no reverse/negative offset. Verified (bf16 TP8): flag-off/flag-on cached_tokens=6528 identical; eager + cudagraph both correct; full GSM8K 1319 3-shot = 0.9500 flexible vs CI baseline 0.9522 (within noise); agentic c48 replay 0 errors, cache-hit parity. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Closed
8 tasks
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
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.
Summary
Lets DeepSeek-V4 SWA and the per-type compressors (CSA k=32, HCA k=1) share one
physical KV pool so compress can reuse space SWA freed (and vice versa), instead
of the static
swa_pagessplit that OOMs one side while the other sits idle.This is Milestone 1 (option-2) of
plan_atom_unified_kv_pool.md §11: singleSWA slot space + per-type compress pools; full per-type-SWA symmetry (approach A)
is a later milestone.
Gated behind a new env
ATOM_UNIFIED_KV_SHARE(default off → byte-identicalto today). Non-V4 / non-SWA models are unaffected;
free_per_req_cache_groups(GDN/compressor state) is untouched.
Design
block_tablestays logical (prefix cache + hashing unchanged). Each livelogical compress block is bound to a physical CSA + HCA block via a coordinator.
row = phys*k), all offsets non-negative → flydsl-safe;avoids the negative-stride / cross-layer-
kwall documented in the plan §5.3.drops the SWA hash first, else a later SWA hash-hit would double-claim a
compress-held slot and corrupt KV. Pinned sparse-checkpoint blocks hold a ref →
never on the free-list → never borrowed → retention (
cached_tokens) preserved.Files (new commits only)
unified_chunk_pool.py(new):UnifiedTypePool,UnifiedKvCoordinator,unified_slot_counts.block_manager.py/sequence.py/swa_pool.py/scheduler.py: wiring + per-type tables.deepseek_v4_attn.py/deepseek_v4.py/v4_kernels/{csa_translate_pack,paged_prefill_indices}.py: base-0 addressing + per-type routing.tests/test_unified_chunk_pool.py(new): 10 CPU tests.envs.py:ATOM_UNIFIED_KV_SHARE.Test plan
pytest tests/test_unified_chunk_pool.py(10 passed) — pool, coordinator, cross-role eviction, admission accuracy, sizing.cached_tokens=6528identical; GSM8K parity.--enforce-eager):cached_tokens=6528; capture path works (persistent buffers).