[Feature] DSV4: elastic unified-KV arena + CSA boundary state fused into SWA block - #1704
Open
yhl-amd wants to merge 4 commits into
Open
[Feature] DSV4: elastic unified-KV arena + CSA boundary state fused into SWA block#1704yhl-amd wants to merge 4 commits into
yhl-amd wants to merge 4 commits into
Conversation
Foundation for the DeepSeek-V4 unified-KV pool: replace the fixed SWA / compressed split with one arena of equal-size physical **byte** chunks that owners (SWA, compressed KV) borrow elastically. Sizing chunks in bytes lets heterogeneous owners (different dtypes/page sizes) share the same physical memory; an under-used owner no longer wastes its reserved capacity. - `ChunkArena`: free-list of byte chunks. `ChunkBackedFreeList`: an owner's page allocator over the arena (two-list backed/unbacked pool) — a chunk returns to the arena only when ALL its pages are free (true eviction; cross-pool lending is pool-driven). Single-owner per chunk, so owners never alias. - `UnifiedKvArena`: per-ratio-group (c4 / c128 / dense) arenas; per-group physical page tables for batch shipping; owner capacity/admission gates. Pure control-plane data structures (no model/engine deps); unit-tested in isolation. Wired into DSV4 and used by the CSA-into-SWA fusion in the next commit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…tate fused into the SWA block Wire the byte-chunk arena into DeepSeek-V4 and add the native CSA prefix-cache boundary snapshot, fused into the SWA block so it shares the block's alloc/free/retention-pin lifecycle (no separate pool). GPU-validated on DeepSeek-V4-Pro TP4 fp8. Arena integration (elastic SWA / compressed KV): - ModelRunner sizing (per-group num_chunks from the reconciled num_physical_kvcache_blocks), BlockManager + SWA-pool wiring + cross-pool pool-driven eviction, Scheduler per-group physical-table shipping. - deepseek_v4_attn: one flat per-layer `unified_kv` arena tensor; per-group physical index builds; kernel addressing via per-group physical pages. CSA-into-SWA fusion (native prefix cache, bit-exact restore): - Layout (fp8): c4 chunk = 112KB = [SWA nope KV 64KB(128 rows) | CSA state 40KB | pad] = 224 rows; c128/dense stay 64KB/128. CSA is NOT an arena owner — when a c4 chunk is lent to compress its CSA bytes are reused (full elasticity). - SWA physical ROW stride (swa_phys * chunk_rows) decoupled from the token windowing modulus (block_size) via a `row_stride`/`swa_row_stride` param in swa_write + the paged prefill/decode index builders (default == block_size → byte-identical for c128/dense/non-arena). Decode SWA write is fused in aiter (one swa_block_size), so c4 layers un-fuse to a standalone swa_write with the chunk_rows stride; reads are ATOM index builders (gather by kv_indices), so no aiter change is needed. - Capture/restore index the fused segment (fp32 views strided by chunk == c4 physical SWA page); scheduler ships the c4 physical SWA table as the capture destination (unbacked blocks -> -1 so capture skips them; a 0 fallback would overwrite chunk-0's live KV) and the terminal cached block's c4 physical SWA page as the restore source. bound_hit collapses to the SWA trailing-window gate, so CSA prefix-hit correctness rides SWA retention (ATOM_SWA_FULL_RETAIN / ATOM_SWA_RETENTION_INTERVAL). The associative CsaStatePool is superseded and not present. Validation (GSM8K forced-prefix, n=80): arena+CSA-off 0.9750; SWA-224-stride-only 0.9625; fused capture 0.9875; fused + SWA full-retain 0.9875 / 0.9500 with ~70/81 prefix hits and restore firing (within TP4-fp8 greedy nondeterminism). CPU: full suite green. Design notes in CSA_SWA_FUSION_PLAN.md / UNIFIED_KV_ARENA_PLAN.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… owner) After the CSA-into-SWA fusion (bee9e5f) removed the associative CsaStatePool, StatePool became a 169-line pass-through: 14 of its 15 methods were one-line delegates to its SlidingWindowPool component, and the only method with logic, bound_hit, just wrapped an assert + a PrefixStateHit whose swa_hit_start field was dead (block_manager recomputes it locally). This collapses that vestigial layer. - Delete state_pool.py (StatePool + the dead PrefixStateHit dataclass). - BlockManager constructs SlidingWindowPool directly as self.swa — the single prefix-cache sidecar owner (CSA boundary rides the SWA block, no page pool). - requires_csa_boundary_state becomes a @Property (kept as a seam for a future non-SWA sidecar); bound_hit's assert guard is inlined into can_allocate. - Rewrite the 14 BlockManager + 3 Scheduler delegate call sites to self.swa.* using SlidingWindowPool's own method names; update stale StatePool docstrings in swa_pool.py / kv_block.py. - Tests: reach through bm.swa.* (matching the existing arena-test pattern). Pure control-plane cleanup, no behavior change: full CPU suite 733 passed / 33 skipped, unchanged from before. Net atom/+tests: +76 / -250. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: yihonglie <hyi@amd.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
…ority
Converge two coupling smells in the unified-KV arena wiring (no behavior change):
1. Group spec was a stringly-typed dict hand-indexed in 4 modules (config /
unified_kv_arena / model_runner / deepseek_v4_attn). Replace with a frozen
ArenaGroupSpec dataclass (single schema) carrying name / num_chunks /
bytes_per_chunk / chunk_rows / owners / csa_state_off, plus helpers
(compress_page_bytes, has_compress, max_compressed_blocks). It is pickle-safe
so it rides block_info across the runner→scheduler boundary unchanged; a
`coerce()` classmethod still accepts plain dicts (tests / legacy config).
- producer: compute_arena_group_specs._grp now builds ArenaGroupSpec.
- consumers: UnifiedKvArena.__init__, model_runner cmp_caps, attn row-map
now use attribute access; model_runner's per-group compressed-cap
re-derivation collapses to `spec.max_compressed_blocks`.
2. ratio→group ("c4"/"c128"/"dense") was re-implemented in 3 places. Add the
module-level group_of_ratio() authority; UnifiedKvArena.compress_group_of_ratio
and the two attn call sites (_arena_rows_for, _swa_stride) now delegate to it.
Verified: ArenaGroupSpec pickle round-trips (frozen), arena accepts typed +
dict specs, attn imports (no import cycle from the new dep); tests
test_unified_kv_arena + test_chunk_arena + test_block_manager_arena: 27 passed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: yihonglie <hyi@amd.com>
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
Elastic unified-KV memory management for DeepSeek-V4, plus fusing the CSA
prefix boundary-state snapshot directly into the SWA block (no separate page
pool). Three milestone commits over base #1697:
bcd07ca3— unified-KV byte-chunk arena (control plane).ChunkArena(byte chunks) +
ChunkBackedFreeList(backed/unbacked two-list) +UnifiedKvArena(per-group c4/c128/dense), single-owner-per-chunk enforcedat the free-pool level. Pure data structure, builds standalone (17 arena
unit tests).
bee9e5f5— arena integration + CSA-into-SWA fusion. Wires the arenainto DSV4 sizing / block_manager / scheduler / attention / kernel addressing,
and fuses the CSA boundary snapshot into the c4 SWA chunk: a 112 KB c4 chunk
is
[SWA nope KV 64 KB | CSA state 40 KB | pad 8 KB]= 224 rows (fp8), whilec128/dense stay 128 rows. Physical SWA row stride is decoupled from the
token-windowing modulus (
row_stride/swa_row_stride, default == block_sizeso c128/dense/non-arena stay byte-identical). c4-layer decode SWA writes
un-fuse to a standalone strided
swa_write(aiter untouched). Capture/restoreride the c4 physical SWA page; unbacked SWA ids map to -1 so capture skips
window-freed placeholders. CSA retention rides the SWA pin
(
ATOM_SWA_FULL_RETAIN/ATOM_SWA_RETENTION_INTERVAL). The associativeCsaStatePoolis deleted.da67dbbd— inline StatePool into BlockManager. With the fusion inplace,
StatePoolwas a 169-line pass-through (14/15 methods one-linedelegates; the lone logic method wrapped a dead
PrefixStateHit.swa_hit_startfield). Inlined:
SlidingWindowPoolis nowBlockManager.swa, the singleprefix-cache sidecar owner. Pure control-plane cleanup, +76 / -250.
Test plan
pytest tests/ --ignore=tests/plugin; the plugin dir has a pre-existing sglang env import failure unrelated to this change).bcd07ca3) builds + passes its 17 unit tests standalone.ATOM_SWA_FULL_RETAIN.row_stridedefaults toblock_size; SWA pool disabled → all methods no-op).Notes
🤖 Generated with Claude Code