[Feature] DeepSeek-V4 unified KV pool (option-2): SWA↔compress share physical KV#5
Closed
yhl-amd wants to merge 3 commits into
Closed
[Feature] DeepSeek-V4 unified KV pool (option-2): SWA↔compress share physical KV#5yhl-amd wants to merge 3 commits into
yhl-amd wants to merge 3 commits into
Conversation
…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>
Owner
Author
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
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).