diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index bc516c13d6..21f1d0b82c 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -17,6 +17,10 @@ from atom.model_engine.kv_block import Block from atom.model_engine.sequence import Sequence from atom.model_engine.swa_pool import SlidingWindowPool +from atom.model_engine.unified_chunk_pool import ( + UnifiedKvCoordinator, + unified_slot_counts, +) from atom.utils import envs @@ -83,8 +87,35 @@ def __init__(self, config: Config): # byte-identical. See atom/model_engine/swa_pool.py. _spec = getattr(config, "speculative_config", None) _mtp_k = int(getattr(_spec, "num_speculative_tokens", 0) or 0) if _spec else 0 + _num_swa_blocks = getattr(config, "num_swa_blocks", 0) + + # DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE / plan §11, option-2). + # When on for a V4 (SWA) model, SWA and the per-type compressors (CSA k=32, + # HCA k=1) draw physical space from ONE coordinator so compress can borrow + # SWA-freed slots. `block_table` stays LOGICAL (prefix cache unchanged); the + # coordinator maps each live logical compress block to a physical id in the + # CSA and HCA pools (base-0, row = phys * k). Flag-off → coordinator is None + # and every path below stays byte-identical to today. + self._unified: UnifiedKvCoordinator | None = None + if envs.ATOM_UNIFIED_KV_SHARE and _num_swa_blocks > 0: + n_csa_slots, n_hca_slots = self._unified_slot_counts( + config, _num_swa_blocks, num_blocks, block_size + ) + self._unified = UnifiedKvCoordinator( + num_swa_base=_num_swa_blocks, + n_csa_slots=n_csa_slots, + n_hca_slots=n_hca_slots, + block_size=block_size, + ) + # Per-logical-block physical binding (index = logical compress id). + # -1 = unbound (block not currently holding content). A binding is + # created on true allocation, retained across cache-free (lazy reuse), + # and re-issued on evict-reuse; see _allocate_block / _deallocate_block. + self._logical_csa: list[int] = [-1] * num_blocks + self._logical_hca: list[int] = [-1] * num_blocks + self.swa = SlidingWindowPool( - num_blocks=getattr(config, "num_swa_blocks", 0), + num_blocks=_num_swa_blocks, window=getattr(config, "swa_window_size", 0), block_size=block_size, max_num_batched_tokens=getattr(config, "max_num_batched_tokens", 0), @@ -92,7 +123,21 @@ def __init__(self, config: Config): full_retain=envs.ATOM_SWA_FULL_RETAIN, retention_interval=envs.ATOM_SWA_RETENTION_INTERVAL, checkpoint_frac=envs.ATOM_SWA_CHECKPOINT_FRAC, + chunk_pool=self._unified, ) + if self._unified is not None: + # Compress borrowing an SWA-cached slot must drop the SWA cache entry + # first (plan §11.1), else a later SWA hash-hit double-claims it. + self._unified.set_swa_evict_cb(self.swa.evict_slot) + + @staticmethod + def _unified_slot_counts(config, num_swa_blocks, num_logical_blocks, block_size): + """Per-type compress-pool slot counts. Shares the SINGLE formula with the + attention builder's unified_kv tensor sizing (unified_slot_counts) so the + coordinator and the physical tensors agree exactly. block_ratio is asserted + 1 for V4 (logical == physical compress blocks) in allocate_kv_cache; here + num_logical_blocks == num_physical compress blocks.""" + return unified_slot_counts(num_swa_blocks, num_logical_blocks, block_size) @property def swa_enabled(self) -> bool: @@ -129,6 +174,9 @@ def _allocate_block(self, block_id: int) -> Block: block.reset() self.free_block_ids_set.discard(block_id) self.used_block_ids.add(block_id) + # Unified pool: this logical id now holds fresh content → (re)bind its + # per-type physical space (returns any prior binding to the coordinator). + self._rebind_physical(block_id) return self.blocks[block_id] def _deallocate_block(self, block_id: int): @@ -136,6 +184,42 @@ def _deallocate_block(self, block_id: int): self.used_block_ids.remove(block_id) self.free_block_ids.append(block_id) self.free_block_ids_set.add(block_id) + # Unified pool: RETAIN the physical CSA/HCA binding across cache-free so a + # later cache-hit reuses the resident KV (matches flag-off's lazy logical + # retention). Physical is returned to the coordinator only on evict-reuse + # (_rebind_physical) — that is when SWA can reclaim the freed slot. + + # --------------- unified pool: per-type physical binding --------------- # + def _rebind_physical(self, block_id: int) -> None: + """Give logical compress block `block_id` a fresh CSA+HCA physical binding + for new content (true alloc / evict-reuse). Frees any prior binding first + so the coordinator can hand a borrowed SWA slot back to SWA. Flag-off: no-op + (self._unified is None).""" + if self._unified is None: + return + old_csa = self._logical_csa[block_id] + if old_csa != -1: + self._unified.free_csa(old_csa) + self._unified.free_hca(self._logical_hca[block_id]) + self._logical_csa[block_id] = self._unified.alloc_csa() + self._logical_hca[block_id] = self._unified.alloc_hca() + + def _append_phys(self, seq: Sequence, block_id: int) -> None: + """Append `block_id`'s physical CSA/HCA ids to the seq's per-type tables, + positionally aligned with seq.block_table. Flag-off: no-op.""" + if self._unified is None: + return + seq.csa_block_table.append(self._logical_csa[block_id]) + seq.hca_block_table.append(self._logical_hca[block_id]) + + def _unified_has_free(self, n_new_blocks: int) -> bool: + """Whether the coordinator can supply `n_new_blocks` new CSA and HCA + physical blocks. Flag-off / n<=0: True.""" + if self._unified is None or n_new_blocks <= 0: + return True + return self._unified.has_free_csa(n_new_blocks) and self._unified.has_free_hca( + n_new_blocks + ) def can_allocate(self, seq: Sequence) -> int: """Return number of cache-hit blocks (>=0) if seq fits, else -1. @@ -163,6 +247,10 @@ def can_allocate(self, seq: Sequence) -> int: # when SWA disabled. if not self.swa.has_free(self.swa.admission_blocks(seq)): return -1 + # Unified pool: each new logical compress block needs a CSA + HCA + # physical block (no-op / True flag-off). + if not self._unified_has_free(seq.num_blocks): + return -1 return 0 # Step 1: compressed prefix (CSA/HCA/indexer share the block hash and # read the WHOLE history, so this stays a full front-to-back chained @@ -203,6 +291,10 @@ def can_allocate(self, seq: Sequence) -> int: # True when SWA disabled. if not self.swa.has_free(min(num_new_blocks, self.swa.admission_blocks(seq))): return -1 + # Unified pool: physical CSA/HCA blocks for the uncached logical blocks + # (cache-hit blocks keep their retained binding). No-op / True flag-off. + if not self._unified_has_free(num_new_blocks): + return -1 return num_cached_blocks def allocate(self, seq: Sequence, num_cached_blocks: int = 0): @@ -244,6 +336,7 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0): self.free_block_ids_set.discard(block_id) self.used_block_ids.add(block_id) seq.block_table.append(block_id) + self._append_phys(seq, block_id) # unified pool (no-op flag-off) if i < swa_hit_start: self.swa.alloc_placeholder(seq) # out of window: never read → -1 else: @@ -252,6 +345,7 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0): block_id = self._pop_free_block() self._allocate_block(block_id) seq.block_table.append(block_id) + self._append_phys(seq, block_id) # unified pool (no-op flag-off) # Uncached blocks: -1 placeholder keeps swa_block_table the same # length as block_table; ensure_for_tokens fills the current chunk's # window slots before each forward, free_after_prefill_chunk releases @@ -326,6 +420,13 @@ def deallocate(self, seq: Sequence): ) # release SWA blocks + clear swa_block_table (no-op if disabled) seq.num_cached_tokens = 0 seq.block_table.clear() + if self._unified is not None: + # Physical bindings are retained on the logical blocks (lazy reuse); + # only the per-seq positional tables are cleared here, mirroring + # block_table. (deallocate above dropped ref_counts; blocks that hit 0 + # went to the free pool with their CSA/HCA binding intact.) + seq.csa_block_table.clear() + seq.hca_block_table.clear() if seq.has_per_req_cache and seq.per_req_cache_group >= 0: self.free_per_req_cache_groups.append(seq.per_req_cache_group) seq.per_req_cache_group = -1 @@ -341,6 +442,8 @@ def can_append(self, seq: Sequence, num_new_tokens: int = 1) -> bool: return False if not self.swa.has_free(new_blocks_needed): # True when SWA disabled return False + if not self._unified_has_free(new_blocks_needed): # unified pool (no-op off) + return False return True def may_append(self, seq: Sequence, num_new_tokens: int = 1): @@ -361,6 +464,7 @@ def may_append(self, seq: Sequence, num_new_tokens: int = 1): block_id = self._pop_free_block() self._allocate_block(block_id) block_table.append(block_id) + self._append_phys(seq, block_id) # unified pool (no-op flag-off) self.swa.append_new(seq) # lockstep SWA block (no-op if disabled) # Reclaim SWA blocks that just fell out of the window (no-op if disabled). self.swa.free_out_of_window(seq, len(seq)) diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index dc9e7c62bc..c1418afe7b 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -354,6 +354,16 @@ def __init__( self.swa_block_tables = [ seq.swa_block_table for seq in seqs.values() if seq.block_table ] + # DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE / plan §11): per-type + # PHYSICAL compress tables (base-0, row = phys*k). CSA and HCA are allocated + # independently so each can borrow SWA-freed slots. `block_table` above + # stays LOGICAL. Empty (flag-off) → downstream reverts to block_tables. + self.csa_block_tables = [ + seq.csa_block_table for seq in seqs.values() if seq.block_table + ] + self.hca_block_tables = [ + seq.hca_block_table for seq in seqs.values() if seq.block_table + ] self.last_block_num_tokens = [ _seq.last_block_num_tokens for _seq in seqs.values() ] diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index 56a0c3bdf6..317515bb8f 100644 --- a/atom/model_engine/sequence.py +++ b/atom/model_engine/sequence.py @@ -96,6 +96,14 @@ def __init__( # out-of-window SWA blocks can be freed while compressed blocks persist). # Empty / unused for non-SWA models. self.swa_block_table = [] + # DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE / plan §11): per-type + # PHYSICAL compress block tables. `block_table` stays the LOGICAL compress + # id (prefix-cache / hashing unchanged); these hold the physical block ids + # in the CSA (k=32) and HCA (k=1) pools, which are allocated independently + # so each can borrow SWA-freed slots (base-0 addressing, row = phys * k). + # Positionally aligned with `block_table`. Empty unless the flag is on. + self.csa_block_table = [] + self.hca_block_table = [] # Per-request cache slot index (filled by BlockManager.allocate()). # -1 = unallocated. The slot indexes into the per-req cache tensors # owned by ModelRunner (e.g. mamba_k_cache for GDN). diff --git a/atom/model_engine/swa_pool.py b/atom/model_engine/swa_pool.py index 0e8f0cd300..8e68382520 100644 --- a/atom/model_engine/swa_pool.py +++ b/atom/model_engine/swa_pool.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -from collections import OrderedDict, deque +from collections import OrderedDict from atom.model_engine.kv_block import Block from atom.model_engine.sequence import Sequence +from atom.model_engine.unified_chunk_pool import UnifiedTypePool class SlidingWindowPool: @@ -39,6 +40,7 @@ def __init__( full_retain: bool = False, retention_interval: int = 0, checkpoint_frac: float = 0.5, + chunk_pool: "UnifiedTypePool | None" = None, ): self.enabled: bool = num_blocks > 0 self.window: int = window @@ -82,34 +84,49 @@ def __init__( ) self.blocks: list[Block] = [Block(i) for i in range(num_blocks)] self.hash_to_block_id: dict[int, int] = dict() - self.free_block_ids: deque[int] = deque(range(num_blocks)) - self.free_block_ids_set: set[int] = set(range(num_blocks)) - self.used_block_ids: set[int] = set() + # Free-list moved to the shared UnifiedChunkPool (plan §5.2): the SWA pool + # draws physical block ids from it, so freed SWA blocks can (Phase 2) be + # reused by the compressor and vice versa. Phase 1: SWA-only side, spc=1 → + # behaviorally identical to the old deque/set. When BlockManager injects a + # shared pool it is used directly; otherwise own a private SWA-only one. + # SWA draws slots from a per-type UnifiedTypePool. Flag-off (static split): + # a private SWA-only pool (k=0), num_swa_blocks slots → phys = slot = w, + # row = w*block_size, byte-identical to the old free-list. Flag-on: a + # shared per-type pool is injected (SWA + that type's compress share it). + self.chunk_pool: UnifiedTypePool = chunk_pool or UnifiedTypePool( + num_blocks, k=0, block_size=block_size + ) # ----------------------------- primitives ------------------------------ # def _pop(self) -> int: - while self.free_block_ids: - block_id = self.free_block_ids.popleft() - if block_id in self.free_block_ids_set: - self.free_block_ids_set.discard(block_id) - return block_id - raise AssertionError("No free SWA blocks available") + """Get a free SWA block id from the shared chunk pool (marks it used).""" + return self.chunk_pool.alloc_swa() def _alloc(self, block_id: int) -> Block: + """Reset block metadata for a freshly-popped id (lazy hash eviction). + Free-list state is owned by chunk_pool; _pop already marked it used.""" block = self.blocks[block_id] assert block.ref_count == 0 if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id: del self.hash_to_block_id[block.hash] block.reset() - self.free_block_ids_set.discard(block_id) - self.used_block_ids.add(block_id) return block def _dealloc(self, block_id: int): assert self.blocks[block_id].ref_count == 0 - self.used_block_ids.remove(block_id) - self.free_block_ids.append(block_id) - self.free_block_ids_set.add(block_id) + self.chunk_pool.free_swa(block_id) + + def evict_slot(self, block_id: int) -> None: + """Drop the content-cache entry for SWA block `block_id` because a compress + pool is about to overwrite its physical slot (unified pool, plan §11.1). + Registered as the coordinator's SWA evict callback. Only called for + released (ref-0) slots; pinned checkpoints hold a ref → never released, so + never reach here. Idempotent / safe if the slot holds no live hash.""" + block = self.blocks[block_id] + if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id: + del self.hash_to_block_id[block.hash] + block.hash = -1 + self.checkpoint_lru.pop(block_id, None) # ------------------------ sparse checkpoint pins ----------------------- # def _is_checkpoint(self, seq: Sequence, i: int) -> bool: @@ -152,7 +169,7 @@ def has_free(self, n: int) -> bool: blocks admission).""" if not self.enabled: return True - return len(self.free_block_ids_set) >= n + return self.chunk_pool.has_free_swa(n) def admission_blocks(self, seq: Sequence) -> int: """Peak concurrent SWA blocks one request holds during (chunked) prefill. @@ -218,13 +235,12 @@ def claim_cached(self, seq: Sequence, h: int, token_ids: list[int]): return swa_id = self.hash_to_block_id[h] block = self.blocks[swa_id] - if swa_id in self.used_block_ids: + if not self.chunk_pool.is_free_swa(swa_id): block.ref_count += 1 else: assert block.ref_count == 0 block.ref_count = 1 - self.free_block_ids_set.discard(swa_id) - self.used_block_ids.add(swa_id) + self.chunk_pool.claim_swa(swa_id) # Cross-request reuse of a pinned checkpoint → refresh its LRU recency. if swa_id in self.checkpoint_lru: self.checkpoint_lru.move_to_end(swa_id) diff --git a/atom/model_engine/unified_chunk_pool.py b/atom/model_engine/unified_chunk_pool.py new file mode 100644 index 0000000000..a1528afa35 --- /dev/null +++ b/atom/model_engine/unified_chunk_pool.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Unified per-layer-type KV slot allocator for DeepSeek-V4 (plan §10). + +Goal: within ONE layer type (all CSA layers, or all HCA layers), let SWA and that +type's compressor draw physical space from a SINGLE free-list, so space freed by +one is reusable by the other — replacing the static ``swa_pages`` split. + +Model (plan §10.1): + * Allocation unit = one **slot** = ``block_size`` (128) rows = exactly one SWA + block. SWA needs 128 contiguous rows, so the slot is the coarsest common unit + that avoids external fragmentation. + * A slot is homogeneous: held either by SWA (1 SWA block) or by compress + (``blocks_per_slot = block_size // k`` compress blocks; CSA k=32 → 4, HCA + k=1 → 128). A whole slot returns to the shared free-list only when empty. + * Addressing is base-0 (plan §10.2), so physical block ids are non-negative + (flydsl-safe, no reverse / negative offset): + SWA block : phys = slot (row = phys * 128) + compress block : phys = slot * blocks_per_slot + j (row = phys * k) + where ``j`` in ``[0, blocks_per_slot)`` is the sub-index within the slot. + +Per layer type there is one instance (CSA pool, HCA pool). All layers of that +type share it: identical k and structure → identical slot decisions apply to +every layer's own tensor. CSA↔HCA cannot share (different tensors / different k) +— that is physical, not a forced split. + +Self-guarding: ``num_slots == 0`` → disabled (non-V4 / no compress). SWA-only +(dense) pools set ``k == 0`` (no compress side). +""" + +from __future__ import annotations + +from collections import deque +from typing import Callable, Optional + + +def unified_slot_counts( + num_swa_blocks: int, num_compress_blocks: int, block_size: int = 128 +) -> tuple[int, int]: + """Per-type 128-row slot counts for the CSA and HCA unified tensors (plan §11). + + Correctness-first sizing: each type reserves the SWA region [0, num_swa_blocks) + plus an own-compress region large enough to hold every compress block WITHOUT + borrowing, so flag-on physical capacity >= flag-off (enabling the flag never + OOMs on its own; borrowing SWA-freed slots is pure upside). This equals the + flag-off tensor size to within one slot: n_csa*128 == swa_pages + ceil rounding + of num_compress_blocks*32. The SINGLE source of truth shared by BlockManager + (coordinator sizing) and the attention builder (unified_kv tensor sizing). + """ + csa_bps = block_size // 32 # CSA k=32 -> 4 compress blocks per 128-row slot + hca_bps = block_size // 1 # HCA k=1 -> 128 per slot + csa_own = (num_compress_blocks + csa_bps - 1) // csa_bps + hca_own = (num_compress_blocks + hca_bps - 1) // hca_bps + return num_swa_blocks + csa_own, num_swa_blocks + hca_own + + +class UnifiedTypePool: + def __init__(self, num_slots: int, k: int, block_size: int = 128): + self.enabled: bool = num_slots > 0 + self.num_slots: int = int(num_slots) + self.block_size: int = int(block_size) + self.k: int = int(k) # compress rows per block (0 = SWA-only / dense) + # compress blocks packed into one slot (0 when k==0) + self.blocks_per_slot: int = (block_size // k) if k > 0 else 0 + + # Shared slot free-list (both SWA and compress draw from it). + self._free_slots: deque[int] = deque(range(num_slots)) + self._free_slots_set: set[int] = set(range(num_slots)) + + # Slots currently held by SWA (each = 1 SWA block). + self._swa_slots: set[int] = set() + + # Compress bookkeeping: a slot handed to compress is sub-divided into + # blocks_per_slot sub-indices. Track, per open compress slot, the set of + # FREE sub-indices; a slot with all sub-indices free returns to the pool. + # `_cmp_open` holds slots with >=1 free sub-index (alloc source). + self._cmp_free_sub: dict[int, set[int]] = {} # slot -> {free sub-idx} + self._cmp_open: deque[int] = deque() + + # ------------------------------- slots -------------------------------- # + def _pop_free_slot(self) -> int: + while self._free_slots: + s = self._free_slots.popleft() + if s in self._free_slots_set: + self._free_slots_set.discard(s) + return s + raise AssertionError("No free KV slots available") + + def _return_slot(self, s: int) -> None: + self._free_slots.append(s) + self._free_slots_set.add(s) + + # ---- external reservation (option-2: SWA holds the SAME slot across the + # CSA and HCA pools; the coordinator reserves it in both so neither pool's + # compress side hands it out while SWA holds it). ---------------------- # + def reserve_slot(self, s: int) -> bool: + """Mark slot `s` as taken by an external owner (SWA). Returns True if it + was free and is now reserved; False if it wasn't free (can't reserve).""" + if s not in self._free_slots_set: + return False + self._free_slots_set.discard(s) + return True + + def release_slot(self, s: int) -> None: + """Release an externally-reserved slot back to this pool's free-list.""" + if s not in self._free_slots_set: + self._return_slot(s) + + def slot_is_free(self, s: int) -> bool: + return s in self._free_slots_set + + # --------------------------- capacity checks -------------------------- # + def has_free_swa(self, n_blocks: int) -> bool: + """Whether `n_blocks` SWA blocks can be admitted. Disabled → True.""" + if not self.enabled: + return True + return len(self._free_slots_set) >= n_blocks # 1 slot per SWA block + + def has_free_compress(self, n_blocks: int) -> bool: + """Whether `n_blocks` compress blocks fit in free sub-slots + free slots. + Disabled / SWA-only → True.""" + if not self.enabled or self.blocks_per_slot == 0: + return True + open_capacity = sum(len(v) for v in self._cmp_free_sub.values()) + slot_capacity = len(self._free_slots_set) * self.blocks_per_slot + return open_capacity + slot_capacity >= n_blocks + + def num_free_slots(self) -> int: + return len(self._free_slots_set) + + # ------------------------------- SWA ---------------------------------- # + def alloc_swa(self) -> int: + """Take a slot for one SWA block; return its physical block id (= slot).""" + s = self._pop_free_slot() + self._swa_slots.add(s) + return s + + def free_swa(self, phys: int) -> None: + """Return an SWA block's slot to the shared free-list.""" + self._swa_slots.discard(phys) + self._return_slot(phys) + + def is_free_swa(self, phys: int) -> bool: + """True if slot `phys` is currently on the free-list (for cached-hit claim).""" + return phys in self._free_slots_set + + def claim_swa(self, phys: int) -> None: + """Claim a specific free slot for an SWA cached-hit reuse.""" + self._free_slots_set.discard(phys) + self._swa_slots.add(phys) + + # ----------------------------- compress ------------------------------- # + def alloc_compress(self) -> int: + """Allocate one compress block; return its physical block id + (= slot * blocks_per_slot + sub). Opens a new slot when needed.""" + assert self.blocks_per_slot > 0, "compress alloc on an SWA-only pool" + # Reuse an open compress slot with a free sub-index. + while self._cmp_open: + s = self._cmp_open[0] + free_sub = self._cmp_free_sub.get(s) + if free_sub: + j = free_sub.pop() + if not free_sub: # slot now full → drop from open list + self._cmp_open.popleft() + del self._cmp_free_sub[s] + return s * self.blocks_per_slot + j + self._cmp_open.popleft() # stale + # No open slot: take a fresh one, use sub 0, keep the rest open. + s = self._pop_free_slot() + rest = set(range(1, self.blocks_per_slot)) + if rest: + self._cmp_free_sub[s] = rest + self._cmp_open.append(s) + return s * self.blocks_per_slot # sub 0 + + def free_compress(self, phys: int) -> None: + """Free one compress block; return its slot to the pool when fully empty.""" + s, j = divmod(phys, self.blocks_per_slot) + free_sub = self._cmp_free_sub.get(s) + if free_sub is None: + # Slot was full (not in open list): reopen it with this sub free. + free_sub = set() + self._cmp_free_sub[s] = free_sub + self._cmp_open.append(s) + free_sub.add(j) + if len(free_sub) == self.blocks_per_slot: + # All sub-blocks free → return the whole slot to the shared pool. + del self._cmp_free_sub[s] + # (leave the stale entry in _cmp_open; _pop/_alloc skip via set checks) + self._return_slot(s) + + +class UnifiedKvCoordinator: + """Option-2 (plan §10, SWA-single-table variant): SWA keeps ONE slot space + shared across layer types; compress is per-type (CSA / HCA). Compress freely + reuses slots SWA has freed (the main requirement: "compress reuses SWA's freed + space"); SWA does not grow past its base region (bounded by the dense/HCA + tensors), so SWA-side logic and tables stay unchanged. + + Slot = 128 rows (= one SWA block). SWA slot `s` maps to row `s*128` in EVERY + layer's tensor. Compress block phys (base-0): `phys = slot * blocks_per_slot + + sub`, row `phys * k` — non-negative (flydsl-safe). + + Shared region [0, num_swa_base): + * A slot is SWA-held, or compress-held (independently per type in that + type's tensor), or free. + * SWA takes a slot only if free in BOTH compress pools (so it's unused in + every tensor). Compress takes a slot if not SWA-held and free in its own + pool. Once compress grabs an SWA-freed slot, SWA cannot reclaim it until + compress frees it (shared budget → admission handles pressure). + Compress-only region [num_swa_base, n_*_total): each type's own compress space. + """ + + def __init__( + self, + num_swa_base: int, # SWA slot count (= num_swa_blocks); SWA region size + n_csa_slots: int, # total 128-row slots in a CSA-layer tensor + n_hca_slots: int, # total 128-row slots in an HCA-layer tensor + block_size: int = 128, + ): + self.enabled = num_swa_base > 0 + self.num_swa_base = int(num_swa_base) + self.block_size = int(block_size) + self.csa = UnifiedTypePool(n_csa_slots, k=32, block_size=block_size) + self.hca = UnifiedTypePool(n_hca_slots, k=1, block_size=block_size) + # Reserve the SWA region [0, num_swa_base) in both compress pools so + # compress starts only in its own region; SWA-freed slots are released in. + for s in range(min(self.num_swa_base, n_csa_slots)): + self.csa.reserve_slot(s) + for s in range(min(self.num_swa_base, n_hca_slots)): + self.hca.reserve_slot(s) + from collections import deque as _dq + self._swa_free = _dq(range(self.num_swa_base)) + self._swa_free_set = set(range(self.num_swa_base)) + self._swa_held: set[int] = set() + # SWA slots currently RELEASED into the compress pools (SWA freed them, + # compress may have taken them). A slot NOT in `_released` is still + # reserved for SWA in both pools (init state) → SWA takes it directly. + self._released: set[int] = set() + # SWA slots that were freed while still holding *lazily-cached* SWA content + # (hash live in the SWA pool, ref 0). If compress borrows such a slot, that + # SWA content is overwritten → the SWA prefix-cache entry MUST be dropped + # first, or a later SWA hash-hit would double-claim a compress-held slot + # and corrupt KV (plan §11.1). `_swa_evict_cb(slot)` performs that drop. + # Cleared when SWA reclaims the slot (its own reuse) or compress evicts it. + self._swa_cached: set[int] = set() + self._swa_evict_cb: Optional[Callable[[int], None]] = None + + def set_swa_evict_cb(self, cb: Callable[[int], None]) -> None: + """Register the SWA pool's hook to drop the cache entry for an SWA slot + that compress is about to overwrite (plan §11.1). `cb(slot)` must be a + no-op if the slot holds no live SWA cache entry.""" + self._swa_evict_cb = cb + + def _evict_swa_if_cached(self, slot: int) -> None: + if slot in self._swa_cached: + self._swa_cached.discard(slot) + if self._swa_evict_cb is not None: + self._swa_evict_cb(slot) + + # ------------------------------- SWA ---------------------------------- # + def alloc_swa(self) -> int: + """Take an SWA slot free in every tensor. Skips slots compress grabbed.""" + while self._swa_free: + s = self._swa_free.popleft() + if s not in self._swa_free_set: + continue + self._swa_free_set.discard(s) + if s not in self._released: + # Never released → still reserved for SWA in both pools. + self._swa_held.add(s) + self._swa_cached.discard(s) + return s + # Released earlier → must re-reserve; skip if compress grabbed it. + csa_ok = s >= self.csa.num_slots or self.csa.reserve_slot(s) + if not csa_ok: + continue + hca_ok = s >= self.hca.num_slots or self.hca.reserve_slot(s) + if not hca_ok: + if s < self.csa.num_slots: + self.csa.release_slot(s) # roll back + continue + self._released.discard(s) + self._swa_cached.discard(s) + self._swa_held.add(s) + return s + raise AssertionError("No free SWA slots available") + + def free_swa(self, s: int) -> None: + """Free an SWA slot; release it to both compress pools for reuse. The slot + may still hold lazily-cached SWA content (ref 0, hash live) → tracked in + `_swa_cached` so a borrowing compress alloc evicts that entry first.""" + self._swa_held.discard(s) + self._swa_free.append(s) + self._swa_free_set.add(s) + self._released.add(s) + self._swa_cached.add(s) + if s < self.csa.num_slots: + self.csa.release_slot(s) + if s < self.hca.num_slots: + self.hca.release_slot(s) + + def _is_swa_reclaimable(self, s: int) -> bool: + """A slot is SWA-reclaimable iff it's in the SWA region, not SWA-held, and + free in BOTH compress pools (SWA needs the whole 128-row slot free in every + layer tensor). Computed from the compress pools = always accurate.""" + if s >= self.num_swa_base or s in self._swa_held: + return False + csa_free = s >= self.csa.num_slots or self.csa.slot_is_free(s) + hca_free = s >= self.hca.num_slots or self.hca.slot_is_free(s) + return csa_free and hca_free + + def is_free_swa(self, s: int) -> bool: + """True if slot `s` is currently reclaimable by SWA (on the SWA free-list + AND not grabbed by compress). `_swa_free_set` is kept accurate as compress + borrows/returns slots, so membership is the source of truth. Mirrors + UnifiedTypePool.is_free_swa so the coordinator can back the SWA pool's + chunk_pool in flag-on mode.""" + return s in self._swa_free_set + + def claim_swa(self, s: int) -> None: + """SWA cache-hit reclaim of a specific free slot (caller checked + is_free_swa). Re-reserves it in the compress pools if it had been + released, and clears its released/cached tags.""" + self._swa_free_set.discard(s) + if s in self._released: + if s < self.csa.num_slots: + self.csa.reserve_slot(s) + if s < self.hca.num_slots: + self.hca.reserve_slot(s) + self._released.discard(s) + self._swa_cached.discard(s) + self._swa_held.add(s) + + def has_free_swa(self, n: int) -> bool: + if not self.enabled: + return True + return len(self._swa_free_set) >= n + + # ----------------------------- compress ------------------------------- # + def _on_compress_alloc(self, slot: int) -> None: + """A compress pool just grabbed `slot`. If it's a released SWA slot, it is + no longer SWA-reclaimable (that pool's 128 rows are now compress) — drop it + from the SWA free-set and evict any lazily-cached SWA content there.""" + if slot < self.num_swa_base and slot in self._released: + self._swa_free_set.discard(slot) + self._evict_swa_if_cached(slot) + + def _on_compress_free(self, slot: int) -> None: + """A compress pool returned space in `slot`. If it's a released SWA slot + that is now fully free in both compress pools again, it becomes + SWA-reclaimable — re-add it to the SWA free-set (empty; its SWA cache was + evicted when compress borrowed it).""" + if ( + slot < self.num_swa_base + and slot in self._released + and slot not in self._swa_free_set + and self._is_swa_reclaimable(slot) + ): + self._swa_free_set.add(slot) + self._swa_free.append(slot) + + def alloc_csa(self) -> int: + phys = self.csa.alloc_compress() + self._on_compress_alloc(phys // self.csa.blocks_per_slot) + return phys + + def free_csa(self, phys: int) -> None: + slot = phys // self.csa.blocks_per_slot + self.csa.free_compress(phys) + self._on_compress_free(slot) + + def alloc_hca(self) -> int: + phys = self.hca.alloc_compress() + self._on_compress_alloc(phys // self.hca.blocks_per_slot) + return phys + + def free_hca(self, phys: int) -> None: + slot = phys // self.hca.blocks_per_slot + self.hca.free_compress(phys) + self._on_compress_free(slot) + + def has_free_csa(self, n: int) -> bool: + return self.csa.has_free_compress(n) + + def has_free_hca(self, n: int) -> bool: + return self.hca.has_free_compress(n) diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index f96437aa44..cbd7a7a9ee 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -54,6 +54,7 @@ pcp_round_robin_query_indices, ) from atom.model_engine.scheduler import ScheduledBatch +from atom.utils import envs from atom.model_ops.attentions.backends import ( AttentionBackend, AttentionMetadataBuilder, @@ -179,10 +180,27 @@ class AttentionMetaData_DSV4(AttentionMetaData): swa_pages: int = 0 """Boundary in `unified_kv`: index < swa_pages → SWA region; index >= swa_pages → compress region. paged-SWA: `num_swa_blocks * block_size`.""" + num_compress_blocks: int = 0 + """Physical compress-block count (`num_physical_kvcache_blocks`). Only used + when `ATOM_UNIFIED_KV_SHARE` is on: the reverse compress layout maps physical + block `c` to slot `num_compress_blocks-1-c` within the compress region + (mirrors the negative-stride write view), so the read-side index builders + (csa_translate_pack / paged_prefill_indices) remap `phys` accordingly.""" + unified_kv_share: bool = False + """Snapshot of `envs.ATOM_UNIFIED_KV_SHARE` at metadata build time — passed to + the index-builder kernels so the compress offset (forward vs reverse-anchored) + matches the physical write layout.""" swa_block_tables: Optional[torch.Tensor] = None """[bs, max_blocks] int32 GPU — paged-SWA logical→physical block table for the independent SWA pool (parallel to `block_tables`, which addresses the compressed pool). -1 entries are window-freed blocks (never indexed).""" + csa_block_tables: Optional[torch.Tensor] = None + """[bs, max_blocks] int32 GPU — unified-share (ATOM_UNIFIED_KV_SHARE) PHYSICAL + CSA compress table (base-0, row = phys*k1). Replaces `block_tables` for CSA + read/write kernels when the flag is on; None (falls back to block_tables) off.""" + hca_block_tables: Optional[torch.Tensor] = None + """[bs, max_blocks] int32 GPU — unified-share PHYSICAL HCA compress table + (base-0, row = phys). Replaces `block_tables` for HCA kernels when on.""" # ----- Native 2buff fp8 per-token paged-decode index tensors ----- # Feed the aiter asm decode kernel `mla_decode_fwd_v4_nm` (op5), which treats @@ -283,6 +301,36 @@ def get_builder_cls() -> Type["AttentionMetadataBuilder"]: return DeepseekV4AttentionMetadataBuilder +def _bind_compress_view( + unified_1d: torch.Tensor, + swa_pages: int, + num_blocks: int, + k: int, + hd: int, +) -> torch.Tensor: + """Build the ``[N, k, hd]`` view the Compressor writes via + ``kv_cache[block_id, slot_in_block, :]`` and the index builders read. + + - Forward (default, ``ATOM_UNIFIED_KV_SHARE`` off): block ``c`` occupies rows + ``[swa_pages + c*k, swa_pages + (c+1)*k)`` — today's contiguous tail view, + ``N = num_blocks`` compress blocks after the SWA region. + - Unified-share (base-0, plan §11.2): the WHOLE tensor is the compress arena, + block ``c`` at rows ``[c*k, (c+1)*k)`` (``base = 0``). SWA (rows ``s*128``) + and compress share it; the coordinator guarantees a physical block id and an + SWA slot never alias the same rows. ``N = T_L // k`` covers every slot + (including SWA-region slots compress may borrow). All three compress-write + paths take the block stride from ``kv_cache.stride(0)``, so this base-0 view + needs NO write-kernel change; the read builders drop ``swa_pages`` to 0. + Non-negative offsets throughout → flydsl-safe (plan §5.3 wall avoided). + """ + if envs.ATOM_UNIFIED_KV_SHARE: + # Whole tensor as compress blocks (base-0). T_L is a multiple of k by + # construction (T_L = n_slots * 128, k divides 128), so view() is exact. + assert unified_1d.shape[0] % k == 0, (unified_1d.shape[0], k) + return unified_1d.view(-1, k, hd) + return unified_1d[swa_pages:].view(num_blocks, k, hd) + + class DeepseekV4AttentionMetadataBuilder(CommonAttentionBuilder): """Per-request cache owner for V4's state-cache buffers. @@ -304,6 +352,16 @@ def __init__(self, model_runner): ratios = list(getattr(hf, "compress_ratios", ())) assert ratios, "deepseek_v4 hf_config must define compress_ratios" self.compress_ratios = ratios + # Unified KV share (plan §11) does not yet wire the TBO ubatch decode + # buffers ({p}csa/hca_block_tables). Fail loud rather than KeyError / + # silently fall back to logical addressing (which would corrupt KV). + if envs.ATOM_UNIFIED_KV_SHARE and getattr( + model_runner.config, "enable_tbo_decode", False + ): + raise NotImplementedError( + "ATOM_UNIFIED_KV_SHARE + enable_tbo_decode not wired yet " + "(per-ubatch csa/hca block tables). Disable one of them." + ) self.num_layers = len(ratios) # Per-buffer-type layer indexing. # Buffers are layer-major: shape [num_layers_of_type, num_slots, *state_shape]. @@ -658,20 +716,44 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: head_dim = self.head_dim dtype = self._swa_dtype + # Unified KV share (plan §11): SWA and compress share ONE base-0 tensor per + # layer type. Rows = n_type_slots * block_size (SWA region + own compress + + # room for borrowing), addressed base-0 (row = phys*k). Sized from the SAME + # formula the coordinator uses so the physical tensor and BlockManager's + # slot pools agree exactly. Flag-off keeps today's swa_pages+compress tail. + unified = envs.ATOM_UNIFIED_KV_SHARE and swa_pages > 0 + if unified: + from atom.model_engine.unified_chunk_pool import unified_slot_counts + assert self.block_ratio == 1, ( + "ATOM_UNIFIED_KV_SHARE assumes logical==physical compress blocks " + f"(block_ratio must be 1, got {self.block_ratio})" + ) + n_csa_slots, n_hca_slots = unified_slot_counts( + self.model_runner.num_swa_blocks, num_blocks, self.block_size + ) + csa_rows = n_csa_slots * self.block_size + hca_rows = n_hca_slots * self.block_size + + def _layer_rows(ratio: int) -> int: + if unified: + if ratio == 4: + return csa_rows + if ratio == 128: + return hca_rows + return swa_pages # Dense: SWA-only + if ratio == 4: + return swa_pages + num_blocks * self.k1_csa + if ratio == 128: + return swa_pages + num_blocks * self.k2_hca + return swa_pages # Dense + # Per-layer unified_kv: SWA prefix + (CSA/HCA) compress tail. unified_kv: list[torch.Tensor] = [] ratios = self.compress_ratios for layer_id in range(self.num_layers): - ratio = ratios[layer_id] - if ratio == 4: - compress_pages = num_blocks * self.k1_csa - elif ratio == 128: - compress_pages = num_blocks * self.k2_hca - else: - compress_pages = 0 # Dense unified_kv.append( torch.zeros( - (swa_pages + compress_pages, head_dim), + (_layer_rows(ratios[layer_id]), head_dim), dtype=dtype, device=device, ) @@ -684,16 +766,9 @@ def allocate_per_req_cache(self, num_slots: int) -> dict[str, object]: unified_kv_rope: list[Optional[torch.Tensor]] = [] if self._kv_fp8: for layer_id in range(self.num_layers): - ratio = ratios[layer_id] - if ratio == 4: - compress_pages = num_blocks * self.k1_csa - elif ratio == 128: - compress_pages = num_blocks * self.k2_hca - else: - compress_pages = 0 # Dense unified_kv_rope.append( torch.zeros( - (swa_pages + compress_pages, self.rope_head_dim), + (_layer_rows(ratios[layer_id]), self.rope_head_dim), dtype=self._rope_dtype, device=device, ) @@ -876,19 +951,22 @@ def build_kv_cache_tensor(self, layer_id: int, module): pos = self.layer_id_to_csa_pos[layer_id_from_prefix] module.kv_state = runner.v4_csa_main_kv_state[pos] module.score_state = runner.v4_csa_main_score_state[pos] - # CSA Main compressed pool now lives in the tail of the - # owning layer's `unified_kv`. Compressor.forward writes via - # `kv_cache[block_id, slot_in_block, :] = entry`, so we hand - # it the standard [num_blocks, k1, head_dim] view. + # CSA Main compressed pool lives in the tail of the owning + # layer's `unified_kv`. Compressor.forward writes via + # `kv_cache[block_id, slot_in_block, :] = entry`. Forward: block + # c at rows [swa_pages+c*k, ..). Unified-share (reverse): block c + # at rows [T_L-(c+1)*k, ..) via a negative-stride view — the write + # kernels read block_stride from kv_cache.stride(0), so this flips + # the write with no kernel change. See _bind_compress_view. num_blocks = runner.num_physical_kvcache_blocks unified = runner.v4_unified_kv[layer_id_from_prefix] - module.kv_cache = unified[swa_pages:].view( - num_blocks, self.k1_csa, self.head_dim + module.kv_cache = _bind_compress_view( + unified, swa_pages, num_blocks, self.k1_csa, self.head_dim ) if self._kv_fp8: rope = runner.v4_unified_kv_rope[layer_id_from_prefix] - module.kv_cache_rope = rope[swa_pages:].view( - num_blocks, self.k1_csa, self.rope_head_dim + module.kv_cache_rope = _bind_compress_view( + rope, swa_pages, num_blocks, self.k1_csa, self.rope_head_dim ) module.write_mode = "main_2buff_fp8" else: @@ -900,13 +978,13 @@ def build_kv_cache_tensor(self, layer_id: int, module): module.score_state = runner.v4_hca_main_score_state[pos] num_blocks = runner.num_physical_kvcache_blocks unified = runner.v4_unified_kv[layer_id_from_prefix] - module.kv_cache = unified[swa_pages:].view( - num_blocks, self.k2_hca, self.head_dim + module.kv_cache = _bind_compress_view( + unified, swa_pages, num_blocks, self.k2_hca, self.head_dim ) if self._kv_fp8: rope = runner.v4_unified_kv_rope[layer_id_from_prefix] - module.kv_cache_rope = rope[swa_pages:].view( - num_blocks, self.k2_hca, self.rope_head_dim + module.kv_cache_rope = _bind_compress_view( + rope, swa_pages, num_blocks, self.k2_hca, self.rope_head_dim ) module.write_mode = "main_2buff_fp8" else: @@ -1371,6 +1449,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): # (model-forward swa_write + decode index kernel), keyed into the # separate num_swa_blocks pool. swa_bt_gpu = var["swa_block_tables"].copy_to_gpu(scheduled_bs) + # Unified KV share: per-type PHYSICAL compress tables (None flag-off). + csa_bt_gpu, hca_bt_gpu = self._populate_compress_block_tables( + batch, scheduled_bs + ) state_slot_gpu = ss_buf.copy_to_gpu(scheduled_bs) # ---- CPU numpy work, overlapped with prep_stream H2D ---- @@ -1403,6 +1485,8 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int): attn_metadata.state_slot_mapping_cpu = state_slot_np attn_metadata.compress_plans = compress_plans attn_metadata.swa_block_tables = swa_bt_gpu + attn_metadata.csa_block_tables = csa_bt_gpu # None flag-off + attn_metadata.hca_block_tables = hca_bt_gpu padded_bs = int(bs) self._attach_v4_per_fwd_meta( @@ -1641,6 +1725,11 @@ def prepare_prefill(self, batch: ScheduledBatch): attn_metadata.swa_block_tables = self._populate_swa_block_tables( batch, scheduled_bs ) + if attn_metadata.csa_block_tables is None: # unified share (None flag-off) + ( + attn_metadata.csa_block_tables, + attn_metadata.hca_block_tables, + ) = self._populate_compress_block_tables(batch, scheduled_bs) state_slot_gpu, state_slot_np = self._populate_state_slot_mapping( batch, scheduled_bs, return_cpu=True ) @@ -2265,7 +2354,18 @@ def _attach_v4_paged_decode_meta( # `_attach_v4_per_fwd_meta`. # ----- HCA compress paged offsets (CPU numpy, vectorized) ----- - block_tables_np_full = var[f"{buf_prefix_ubatch}block_tables"].np[:scheduled_bs] + # Unified KV share: HCA reads its PHYSICAL table with base 0 (row=phys*1); + # flag-off keeps logical block_tables + swa_pages boundary. + if envs.ATOM_UNIFIED_KV_SHARE: + block_tables_np_full = var[ + f"{buf_prefix_ubatch}hca_block_tables" + ].np[:scheduled_bs] + hca_base = 0 + else: + block_tables_np_full = var[ + f"{buf_prefix_ubatch}block_tables" + ].np[:scheduled_bs] + hca_base = swa_pages hca_total_indices = int(hca_indptr_np[T]) hca_indices_np = np.full(hca_total_indices, -1, dtype=np.int32) # n_committed_hca_per_seq is int32; gather stays int32. @@ -2284,7 +2384,7 @@ def _attach_v4_paged_decode_meta( write_pos = hca_indptr_np[token_indices] + entry_offsets bid_expanded = batch_id_per_token_np[token_indices] hca_indices_np[write_pos] = ( - swa_pages + block_tables_np_full[bid_expanded, entry_offsets] + hca_base + block_tables_np_full[bid_expanded, entry_offsets] ).astype(np.int32) # Stage to GPU (HCA compress section at head; SWA prefix scattered below). hca_indices_gpu = self._stage( @@ -2349,6 +2449,8 @@ def _attach_v4_paged_decode_meta( attn_metadata.kv_indptr_csa = csa_indptr_gpu attn_metadata.kv_indptr_hca = hca_indptr_gpu attn_metadata.swa_pages = swa_pages + attn_metadata.num_compress_blocks = self.model_runner.num_physical_kvcache_blocks + attn_metadata.unified_kv_share = envs.ATOM_UNIFIED_KV_SHARE # Per-token paged-decode index tensors for the fp8 asm decode kernel. The # kernel sees N = q_packed.shape[0] = T_pad (padded decode grid). Both @@ -2529,6 +2631,13 @@ def _build_paged_prefill_meta( cu_q_per_seq_gpu = var["cu_seqlens_q"].gpu[:scheduled_bs] if block_tables_gpu is None: block_tables_gpu = var["block_tables"].gpu[:scheduled_bs] + # Unified KV share: the prefill index kernel uses `block_tables` ONLY for + # the HCA-compress section → hand it the PHYSICAL HCA table with base 0 + # (hca_swa_pages below). Flag-off keeps logical block_tables + swa_pages. + hca_swa_pages = swa_pages + if attn_metadata.hca_block_tables is not None: + block_tables_gpu = attn_metadata.hca_block_tables[:scheduled_bs] + hca_swa_pages = 0 # paged-SWA: SWA-prefix offsets index the separate SWA pool via # swa_block_tables; HCA still uses the compressed block_tables. swa_block_tables_gpu = attn_metadata.swa_block_tables[:scheduled_bs] @@ -2578,7 +2687,7 @@ def _build_paged_prefill_meta( T=T, win=win, block_size=self.block_size, - swa_pages=swa_pages, + swa_pages=hca_swa_pages, ) # ----- skip_prefix_len_csa: per-token SWA prefix length ----- @@ -2602,6 +2711,8 @@ def _build_paged_prefill_meta( attn_metadata.kv_indptr_prefix_hca = hca_indptr attn_metadata.skip_prefix_len_csa = skip_csa_gpu attn_metadata.swa_pages = swa_pages + attn_metadata.num_compress_blocks = self.model_runner.num_physical_kvcache_blocks + attn_metadata.unified_kv_share = envs.ATOM_UNIFIED_KV_SHARE def _build_compress_plans( self, @@ -2697,6 +2808,26 @@ def _populate_swa_block_tables(self, batch: ScheduledBatch, scheduled_bs: int): swa_np[i, : len(bt)] = [max(0, b) for b in bt] return var["swa_block_tables"].copy_to_gpu(scheduled_bs) + def _populate_compress_block_tables(self, batch: ScheduledBatch, scheduled_bs: int): + """Unified KV share (plan §11): fill the per-type PHYSICAL compress tables + (CSA base-0 row=phys*k1, HCA base-0 row=phys) from batch.csa/hca_block_tables + and return (csa_gpu, hca_gpu) sliced to scheduled_bs. No-op (returns + (None, None)) when the flag is off — callers then fall back to block_tables.""" + if not envs.ATOM_UNIFIED_KV_SHARE: + return None, None + var = self.model_runner.forward_vars + out = [] + for name in ("csa_block_tables", "hca_block_tables"): + buf = var[name] + np_buf = buf.np + tables = getattr(batch, name, None) or [] + for i in range(scheduled_bs): + np_buf[i] = 0 + if i < len(tables) and len(tables[i]): + np_buf[i, : len(tables[i])] = tables[i] + out.append(buf.copy_to_gpu(scheduled_bs)) + return out[0], out[1] + def _populate_state_slot_mapping( self, batch: ScheduledBatch, scheduled_bs: int, return_cpu: bool = False ): @@ -3012,6 +3143,13 @@ def _alloc_v4_metadata_buffers(self) -> None: # (never indexed; SWA attention only reads in-window positions). _bt_cols = self.model_runner.forward_vars["block_tables"].np.shape[1] bufs["swa_block_tables"] = CpuGpuBuffer(bs, _bt_cols, **i32) + # Unified KV share (plan §11): per-type PHYSICAL compress tables (base-0). + # Same shape as block_tables; filled from batch.csa/hca_block_tables. Only + # populated/consumed when ATOM_UNIFIED_KV_SHARE is on (compress kernels then + # read these + swa_pages=0 instead of the logical block_tables). + if envs.ATOM_UNIFIED_KV_SHARE: + bufs["csa_block_tables"] = CpuGpuBuffer(bs, _bt_cols, **i32) + bufs["hca_block_tables"] = CpuGpuBuffer(bs, _bt_cols, **i32) self.model_runner.forward_vars.update(bufs) diff --git a/atom/model_ops/v4_kernels/csa_translate_pack.py b/atom/model_ops/v4_kernels/csa_translate_pack.py index c555ac2c92..70784824b6 100644 --- a/atom/model_ops/v4_kernels/csa_translate_pack.py +++ b/atom/model_ops/v4_kernels/csa_translate_pack.py @@ -63,7 +63,8 @@ def _csa_translate_pack_kernel( batch_id_per_token_ptr, # [T] int32 — token → seq, sentinel -1 skip_prefix_len_per_token_ptr, # [T] int32 — per-token write offset; ignored when INLINE_SKIP_FROM_POS kv_indices_csa_ptr, # [total_indices] int32 — destination - swa_pages, # i32 — SWA region size, runtime int + swa_pages, # i32 — compress base row. Static split: num_swa_blocks*block_size. + # Unified-share: 0 (compress addressed absolutely from unified row 0; phys = row//k). mnbps, # i32 — max blocks per seq, runtime int index_topk: tl.constexpr, csa_block_capacity: tl.constexpr, @@ -133,6 +134,9 @@ def _csa_translate_pack_kernel( mask=in_range, other=0, ) + # Static split: paged = swa_pages + phys*cap + slot. Unified-share passes + # swa_pages=0 and stores phys = absolute_row//cap (BlockManager), so this + # yields the absolute unified row — no reverse/negative offset (flydsl-safe). paged = swa_pages + phys * csa_block_capacity + slot tl.store( kv_indices_csa_ptr + write_base + k_offs, diff --git a/atom/model_ops/v4_kernels/paged_prefill_indices.py b/atom/model_ops/v4_kernels/paged_prefill_indices.py index 42ae6b4fe2..f309e1de8c 100644 --- a/atom/model_ops/v4_kernels/paged_prefill_indices.py +++ b/atom/model_ops/v4_kernels/paged_prefill_indices.py @@ -77,7 +77,8 @@ def _v4_paged_prefill_indices_kernel( # Constants. win: tl.constexpr, block_size, # paged-SWA: tokens per block (= V4 block_size, 128) - swa_pages, # num_blocks * block_size — boundary into HCA compress section + swa_pages, # HCA compress base row. Static split: num_blocks*block_size. + # Unified-share: 0 (HCA compress addressed absolutely from unified row 0). HCA_RATIO: tl.constexpr, # HCA compress ratio (128) for per-token causal cap BLOCK_N: tl.constexpr, # next_pow2(win) — covers SWA prefix and extend segments ): @@ -167,6 +168,9 @@ def _v4_paged_prefill_indices_kernel( k = j + i hca_mask = k < n_hca bt = tl.load(block_tables_ptr + bt_row_base + k, mask=hca_mask, other=0) + # HCA k2=1 → block c at row swa_pages + c. Unified-share passes + # swa_pages=0 and stores bt = absolute row (BlockManager), so this is the + # absolute unified row — no reverse/negative offset (flydsl-safe). tl.store( prefix_hca_indices_ptr + hca_dst_base + k, swa_pages + bt, diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py index 8eebba5f11..9b53b26896 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -1863,14 +1863,33 @@ def process_weights_after_loading(self) -> None: self.wo_a.quant_type = QuantType.No self.wo_a.need_normalize_e4m3fn_to_e4m3fnuz = False + def _main_compress_bt(self, attn_md): + """Main-compressor block table into unified_kv. Unified KV share (flag-on, + signalled by csa_block_tables being set): the PHYSICAL per-type table (CSA + k=32 / HCA k=1, base-0). Flag-off: the logical block_tables (today).""" + if attn_md.csa_block_tables is not None: + return ( + attn_md.csa_block_tables + if self.compress_ratio == 4 + else attn_md.hca_block_tables + ) + return attn_md.block_tables + def maybe_compressors_async( - self, x, plan, state_slot_mapping, block_tables + self, x, plan, state_slot_mapping, block_tables, idx_block_tables=None ) -> bool: """Fire Compressor(s) on side streams, return immediately. Main Compressor → alt_stream (CSA + HCA). Indexer Compressor → indexer_stream (CSA only). - Waits resolve instantly: side streams ~25us, main Q/KV chain ~87us.""" + Waits resolve instantly: side streams ~25us, main Q/KV chain ~87us. + + `block_tables` scatters the MAIN compressor into `unified_kv` (unified KV + share: per-type PHYSICAL table, base-0). `idx_block_tables` (defaults to + `block_tables`) scatters the Indexer inner compressor into the SEPARATE + `v4_csa_idx_kv` pool, which is NOT shared → always the LOGICAL table.""" + if idx_block_tables is None: + idx_block_tables = block_tables fc = get_forward_context() current_stream = fc.main_stream from atom.utils.tbo.ubatching import tbo_active @@ -1900,7 +1919,7 @@ def maybe_compressors_async( x, plan=plan, state_slot_mapping=state_slot_mapping, - block_tables=block_tables, + block_tables=idx_block_tables, ) else: if has_compressor: @@ -1915,7 +1934,7 @@ def maybe_compressors_async( x, plan=plan, state_slot_mapping=state_slot_mapping, - block_tables=block_tables, + block_tables=idx_block_tables, ) return use_async_compress @@ -2047,7 +2066,8 @@ def forward_impl( x, plan_for_layer, attn_md.state_slot_mapping, - attn_md.block_tables, + self._main_compress_bt(attn_md), # per-type physical (unified share) + attn_md.block_tables, # indexer inner compressor → separate idx pool ) # Q/KV projections (indexer projection deferred to the core's inline @@ -2147,7 +2167,11 @@ def _attn_core( ) else: use_async_compress = self.maybe_compressors_async( - x, plan_for_layer, state_slot_mapping, block_tables_gpu + x, + plan_for_layer, + state_slot_mapping, + self._main_compress_bt(attn_md), # per-type physical (unified) + block_tables_gpu, # indexer inner compressor → separate idx pool ) is_decode = attn_md.state is AttnState.DECODE # Single kernel fuses per-head Q RMSNorm (weightless) + KV RMSNorm @@ -2478,15 +2502,23 @@ def _fill_csa_paged_compress( skip_buf = attn_md.skip_prefix_len_csa window_size = 0 + # Unified KV share: CSA reads its PHYSICAL table with base 0 (row=phys*k1); + # flag-off keeps logical block_tables + swa_pages boundary. + _csa_bt = ( + attn_md.csa_block_tables + if attn_md.csa_block_tables is not None + else attn_md.block_tables + ) + _csa_base = 0 if attn_md.csa_block_tables is not None else attn_md.swa_pages csa_translate_pack( topk_local_raw, - attn_md.block_tables, + _csa_bt, positions, kv_indptr, attn_md.batch_id_per_token, skip_buf, kv_indices, - swa_pages=attn_md.swa_pages, + swa_pages=_csa_base, csa_block_capacity=csa_block_capacity, window_size=window_size, prefix=f"{self.layer_name}.csa_translate_pack", diff --git a/atom/utils/envs.py b/atom/utils/envs.py index f37d2e542c..a1a489170f 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -95,6 +95,14 @@ # Pairs with a larger SWA pool (swa_pool_num_blocks) so freed-but-cached # blocks survive until replay. Costs ~compressed-pool-magnitude SWA memory. "ATOM_SWA_FULL_RETAIN": lambda: (os.getenv("ATOM_SWA_FULL_RETAIN", "0") == "1"), + # DeepSeek-V4 unified KV pool (plan_atom_unified_kv_pool.md): SWA and the + # compressor share ONE per-layer physical tensor — SWA grows from the front + # (row = w*block_size), compress from the BACK (row = T_L - (c+1)*k_L), + # meeting in a shared middle so freed SWA space can back compress and vice + # versa. Default 0 = today's static swa_pages-forward split (no behavior + # change). 1 = reverse-anchored compress layout (Phase 2) + shared budget + # (Phase 3). Gated so the working path stays byte-identical when off. + "ATOM_UNIFIED_KV_SHARE": lambda: (os.getenv("ATOM_UNIFIED_KV_SHARE", "0") == "1"), # DeepSeek-V4 paged-SWA full-retain: fraction of the KV budget given to the # SWA tail pool (the rest goes to the compressed pool). One SWA block is ~7x # the bytes of one compressed block, so a 1:1 mirror starves the compressed diff --git a/tests/test_unified_chunk_pool.py b/tests/test_unified_chunk_pool.py new file mode 100644 index 0000000000..6019a43722 --- /dev/null +++ b/tests/test_unified_chunk_pool.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""CPU unit tests for the DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE, +plan_atom_unified_kv_pool.md §10/§11): the per-layer-type slot allocator +(``UnifiedTypePool``) and the SWA↔compress coordinator (``UnifiedKvCoordinator``), +including cross-role cache eviction and admission accuracy. No GPU required.""" + +from atom.model_engine.unified_chunk_pool import ( + UnifiedKvCoordinator, + UnifiedTypePool, + unified_slot_counts, +) + + +# --------------------------- UnifiedTypePool ---------------------------- # +def test_swa_only(): + # k=0 (SWA-only / flag-off): phys = slot, byte-identical to old free-list. + p = UnifiedTypePool(num_slots=4, k=0, block_size=128) + assert p.enabled and p.has_free_swa(4) and not p.has_free_swa(5) + a = p.alloc_swa() + b = p.alloc_swa() + assert a == 0 and b == 1 # FIFO from 0 + assert not p.is_free_swa(a) and p.num_free_slots() == 2 + p.free_swa(a) + assert p.is_free_swa(a) and p.num_free_slots() == 3 + # cached-hit claim path + c = p.alloc_swa() + p.free_swa(c) + assert p.is_free_swa(c) + p.claim_swa(c) + assert not p.is_free_swa(c) + d = UnifiedTypePool(0, k=0) + assert not d.enabled and d.has_free_swa(9) + + +def test_compress_csa(): + # CSA: k=32 -> blocks_per_slot=4. 2 slots -> 8 compress blocks. + p = UnifiedTypePool(num_slots=2, k=32, block_size=128) + assert ( + p.blocks_per_slot == 4 + and p.has_free_compress(8) + and not p.has_free_compress(9) + ) + got = [p.alloc_compress() for _ in range(8)] + # slot0 -> phys 0..3, slot1 -> phys 4..7 (row = phys*32 covers [0,256)) + assert sorted(got) == [0, 1, 2, 3, 4, 5, 6, 7], got + assert p.num_free_slots() == 0 + # free one block -> slot not fully free -> no slot returned + p.free_compress(2) + assert p.num_free_slots() == 0 + # realloc reuses the freed sub-slot + assert p.alloc_compress() == 2 + # free a whole slot's worth (0,1,2,3) -> slot0 returns to pool + for phys in (0, 1, 3): + p.free_compress(phys) + p.free_compress(2) + assert p.num_free_slots() == 1, p.num_free_slots() + + +def test_mutual_reuse(): + # SWA and compress share one pool: a slot freed by SWA is reusable by compress. + p = UnifiedTypePool(num_slots=3, k=32, block_size=128) # 3 slots + s0 = p.alloc_swa() + p.alloc_swa() # SWA takes 2 slots + assert p.num_free_slots() == 1 + # compress grabs the last free slot -> 4 CSA blocks + [p.alloc_compress() for _ in range(4)] + assert ( + p.num_free_slots() == 0 + and not p.has_free_compress(1) + and not p.has_free_swa(1) + ) + # SWA frees a slot -> compress can now borrow it + p.free_swa(s0) + assert p.has_free_compress(1) and p.has_free_swa(1) + more = [p.alloc_compress() for _ in range(4)] # compress borrows SWA's slot + assert p.num_free_slots() == 0 + # compress frees its borrowed slot -> SWA can take it back + for phys in more: + p.free_compress(phys) + assert p.has_free_swa(1) + assert p.alloc_swa() is not None + + +# ------------------------- UnifiedKvCoordinator ------------------------- # +def test_compress_borrows_swa_freed(): + c = UnifiedKvCoordinator(num_swa_base=4, n_csa_slots=8, n_hca_slots=5) + swa = [c.alloc_swa() for _ in range(4)] + assert sorted(swa) == [0, 1, 2, 3] + assert c.has_free_csa(16) and not c.has_free_csa(17) # own region [4,8) + c.free_swa(1) # SWA frees slot 1 -> CSA can borrow it + assert c.has_free_csa(17) + borrowed = [c.alloc_csa() for _ in range(20)] # 16 own + 4 borrowed + assert set([4, 5, 6, 7]).issubset(set(borrowed)) # slot 1 -> phys 4..7 + + +def test_swa_cannot_reclaim_taken_slot(): + c = UnifiedKvCoordinator(num_swa_base=3, n_csa_slots=4, n_hca_slots=4) + c.alloc_swa() + c.alloc_swa() + c.alloc_swa() + c.free_swa(1) + got = [c.alloc_csa() for _ in range(8)] # CSA grabs borrowed slot 1 + assert set([4, 5, 6, 7]).issubset(set(got)) + try: + c.alloc_swa() + raise AssertionError("SWA should have no reclaimable slot") + except AssertionError as e: + assert "No free SWA" in str(e), e + + +def test_compress_frees_swa_reclaims(): + c = UnifiedKvCoordinator(num_swa_base=3, n_csa_slots=4, n_hca_slots=4) + a = c.alloc_swa() + b = c.alloc_swa() + c.free_swa(a) + assert c.has_free_csa(8) and not c.has_free_csa(9) + csa = [c.alloc_csa() for _ in range(8)] + assert len(set(csa)) == 8 + assert set(range(a * 4, a * 4 + 4)).issubset(set(csa)) + for phys in csa: + c.free_csa(phys) + got = c.alloc_swa() # SWA reclaims after compress frees + assert got in (a, [s for s in range(3) if s not in (a, b)][0]), got + + +def test_compress_borrow_evicts_swa_cache(): + """Compress borrowing an SWA-cached slot must fire the SWA evict callback + (plan §11.1) — else a later SWA hash-hit double-claims a compress-held slot.""" + c = UnifiedKvCoordinator(num_swa_base=3, n_csa_slots=4, n_hca_slots=4) + evicted = [] + c.set_swa_evict_cb(lambda slot: evicted.append(slot)) + c.alloc_swa() + s1 = c.alloc_swa() + c.alloc_swa() + c.free_swa(s1) # slot 1 freed but SWA-cached + assert c.is_free_swa(s1) + got = [c.alloc_csa() for _ in range(8)] # slot3 own + slot1 borrowed + assert set([4, 5, 6, 7]).issubset(set(got)) + assert evicted == [1], evicted # exactly slot 1 evicted, once + assert not c.is_free_swa(s1) + + +def test_swa_hit_reclaims_before_borrow(): + c = UnifiedKvCoordinator(num_swa_base=3, n_csa_slots=4, n_hca_slots=4) + evicted = [] + c.set_swa_evict_cb(lambda slot: evicted.append(slot)) + c.alloc_swa() + s1 = c.alloc_swa() + c.alloc_swa() + c.free_swa(s1) + c.claim_swa(s1) # SWA cache-hit reclaim before compress borrows + assert not c.is_free_swa(s1) + assert c.has_free_csa(4) and not c.has_free_csa(5) # slot 1 re-reserved + got = [c.alloc_csa() for _ in range(4)] + assert set(range(12, 16)) == set(got) # only own region + assert evicted == [] + + +def test_has_free_swa_accurate_after_borrow(): + """has_free_swa must not count SWA slots compress has borrowed (admission).""" + c = UnifiedKvCoordinator(num_swa_base=3, n_csa_slots=4, n_hca_slots=4) + a = c.alloc_swa() + c.alloc_swa() + c.alloc_swa() + assert not c.has_free_swa(1) + c.free_swa(a) + assert c.has_free_swa(1) + [c.alloc_csa() for _ in range(8)] # compress borrows a + assert not c.has_free_swa(1), "borrowed slot must not count as SWA-free" + for phys in range(a * 4, a * 4 + 4): + c.free_csa(phys) + assert c.has_free_swa(1), "reclaimable after compress frees it" + assert c.alloc_swa() == a + + +# ----------------------------- sizing ----------------------------------- # +def test_unified_slot_counts_matches_flagoff_size(): + # n_csa*128 == swa_pages + num_blocks*32 when num_blocks % 4 == 0 (flag-on + # tensor size == flag-off), and covers rounding otherwise. + n_csa, n_hca = unified_slot_counts(num_swa_blocks=16, num_compress_blocks=64) + assert n_csa == 16 + 64 // 4 # 32 slots -> 4096 rows == 16*128 + 64*32 + assert n_hca == 16 + 1 # 64 HCA blocks pack into 1 slot (128/slot) + n_csa2, _ = unified_slot_counts(16, 65) # non-multiple rounds up + assert n_csa2 == 16 + 17