diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index a2076a778c..21f1d0b82c 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -17,6 +17,11 @@ 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 def _make_block_stored( @@ -82,13 +87,57 @@ 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), mtp_k=_mtp_k, + 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: @@ -125,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): @@ -132,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. @@ -159,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 @@ -182,6 +274,10 @@ def can_allocate(self, seq: Sequence) -> int: # gone (#1417), while out-of-window front blocks (SWA-freed) don't block # the hit. num_cached_blocks = self.swa.bounded_hit(seq, compressed_hit, block_hashes) + # Instrumentation: record the pre-gate compressed hit so CacheStats can + # separate reuse lost to the SWA tail gate (compressed_hit - + # num_cached_blocks) from reuse lost to compressed eviction. + seq.num_compressed_hit_blocks = compressed_hit # Free-pool demand: blocks we actually reuse minus those already used # (shared ref); blocks we drop from the hit become fresh → counted. num_new_blocks = seq.num_blocks @@ -195,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): @@ -236,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: @@ -244,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 @@ -318,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 @@ -333,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): @@ -353,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/model_runner.py b/atom/model_engine/model_runner.py index f083b8b62f..8a34686c5d 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -1447,16 +1447,39 @@ def get_num_blocks(self) -> dict[str, int]: b = self.attn_metadata_builder swa_block_bytes = b.swa_pool_block_bytes() if swa_block_bytes > 0: - num_swa_blocks = b.swa_pool_num_blocks( - config.max_num_seqs, config.max_model_len - ) - swa_reserved = num_swa_blocks * swa_block_bytes # block_bytes (from _compute_block_bytes) currently includes the SWA # term; strip it so the compressed pool is sized on compressed bytes. compressed_block_bytes = block_bytes - swa_block_bytes - num_kvcache_blocks = max( - 0, (available_for_pool - swa_reserved) // compressed_block_bytes - ) + if envs.ATOM_SWA_FULL_RETAIN: + # Full-retain: give the SWA tail pool a small fraction `f` of the + # budget; the rest stays with the compressed pool. One SWA block is + # ~7x the bytes of one compressed block, so a 1:1 mirror + # (num_swa == num_kvcache) starves the compressed prefix index + # (measured: 298k -> 36.8k blocks -> hit rate collapsed). A small + # f keeps compressed near full while retaining the hot-boundary + # tail working set (LRU-evicted, same eviction discipline as + # vLLM's FreeKVCacheBlockQueue). Memory-bounded regardless of + # max_model_len. Live SWA footprint stays ~window/seq (window-free + # is kept); the tail pool holds lazily-freed-but-cached tails. + f = min(0.9, max(1e-3, envs.ATOM_SWA_TAIL_BUDGET_FRAC)) + swa_budget = int(available_for_pool * f) + compressed_budget = available_for_pool - swa_budget + num_swa_blocks = swa_budget // swa_block_bytes + num_kvcache_blocks = compressed_budget // compressed_block_bytes + swa_reserved = num_swa_blocks * swa_block_bytes + logger.info( + f"paged-SWA full-retain: tail_budget_frac={f:.3f}, " + f"swa_budget={swa_budget / (1 << 30):.2f}GB, " + f"compressed_budget={compressed_budget / (1 << 30):.2f}GB" + ) + else: + num_swa_blocks = b.swa_pool_num_blocks( + config.max_num_seqs, config.max_model_len + ) + swa_reserved = num_swa_blocks * swa_block_bytes + num_kvcache_blocks = max( + 0, (available_for_pool - swa_reserved) // compressed_block_bytes + ) config.num_swa_blocks = int(num_swa_blocks) config.swa_window_size = int( getattr(hf_config, "sliding_window", 128) or 128 diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index c84d218f13..c1418afe7b 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -138,9 +138,11 @@ class CacheStats: "total_requests", "total_cached_tokens", "total_full_tokens", + "total_compressed_tokens", "_interval_requests", "_interval_cached_tokens", "_interval_full_tokens", + "_interval_compressed_tokens", ) def __init__(self, log_interval: int = 100): @@ -148,18 +150,29 @@ def __init__(self, log_interval: int = 100): self.total_requests: int = 0 self.total_cached_tokens: int = 0 self.total_full_tokens: int = 0 + # Pre-SWA-gate compressed-prefix hit tokens. total - cached separates + # reuse lost to the SWA tail gate from reuse lost to compressed eviction. + self.total_compressed_tokens: int = 0 self._interval_requests: int = 0 self._interval_cached_tokens: int = 0 self._interval_full_tokens: int = 0 + self._interval_compressed_tokens: int = 0 - def update(self, num_cached_tokens: int, num_full_tokens: int) -> None: + def update( + self, + num_cached_tokens: int, + num_full_tokens: int, + num_compressed_tokens: int = 0, + ) -> None: """Record cache stats for one prefill sequence.""" self.total_requests += 1 self.total_cached_tokens += num_cached_tokens self.total_full_tokens += num_full_tokens + self.total_compressed_tokens += num_compressed_tokens self._interval_requests += 1 self._interval_cached_tokens += num_cached_tokens self._interval_full_tokens += num_full_tokens + self._interval_compressed_tokens += num_compressed_tokens if self.total_requests % self._log_interval == 0: self._log() @@ -175,22 +188,40 @@ def _reset_interval(self) -> None: self._interval_requests = 0 self._interval_cached_tokens = 0 self._interval_full_tokens = 0 + self._interval_compressed_tokens = 0 + + @staticmethod + def _rate(num: int, den: int) -> float: + return num / den if den > 0 else 0.0 def _log(self) -> None: - iv_rate = ( - self._interval_cached_tokens / self._interval_full_tokens - if self._interval_full_tokens > 0 - else 0.0 + # compressed = pre-SWA-gate prefix hit; cached = post-gate (admitted). + # (compressed - cached) is reuse lost to a missing SWA tail; (full - + # compressed) is reuse lost to compressed eviction / no logical reuse. + iv_hit = self._rate(self._interval_cached_tokens, self._interval_full_tokens) + iv_comp = self._rate( + self._interval_compressed_tokens, self._interval_full_tokens + ) + iv_gate = self._rate( + self._interval_compressed_tokens - self._interval_cached_tokens, + self._interval_full_tokens, ) logger.info( f"[Cache Stats Interval] Reqs: {self._interval_requests}, " - f"Cached/Total tokens: {self._interval_cached_tokens}/{self._interval_full_tokens}, " - f"Hit rate: {iv_rate:.2%}" + f"Cached/Total: {self._interval_cached_tokens}/{self._interval_full_tokens}, " + f"Hit: {iv_hit:.2%}, Compressed-hit: {iv_comp:.2%}, " + f"Lost-to-SWA-gate: {iv_gate:.2%}" + ) + tot_comp = self._rate(self.total_compressed_tokens, self.total_full_tokens) + tot_gate = self._rate( + self.total_compressed_tokens - self.total_cached_tokens, + self.total_full_tokens, ) logger.info( f"[Cache Stats ] Reqs: {self.total_requests}, " - f"Cached/Total tokens: {self.total_cached_tokens}/{self.total_full_tokens}, " - f"Hit rate: {self.hit_rate:.2%}" + f"Cached/Total: {self.total_cached_tokens}/{self.total_full_tokens}, " + f"Hit: {self.hit_rate:.2%}, Compressed-hit: {tot_comp:.2%}, " + f"Lost-to-SWA-gate: {tot_gate:.2%}" ) @@ -323,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() ] @@ -1298,7 +1339,11 @@ def _schedule_prefill_seq( ) -> tuple[int, int]: num_seqs_prefill += 1 if self.cache_stats: - self.cache_stats.update(seq.num_cached_tokens, seq.num_tokens) + self.cache_stats.update( + seq.num_cached_tokens, + seq.num_tokens, + seq.num_compressed_hit_blocks * self.block_manager.block_size, + ) num_batched_tokens += chunk seq.status = SequenceStatus.RUNNING seq.type = SequenceType.PREFILL diff --git a/atom/model_engine/sequence.py b/atom/model_engine/sequence.py index cfa58e5cf0..317515bb8f 100644 --- a/atom/model_engine/sequence.py +++ b/atom/model_engine/sequence.py @@ -78,6 +78,11 @@ def __init__( self.num_prompt_tokens = len(token_ids) self.num_rejected = 0 self.num_cached_tokens = 0 + # Instrumentation: compressed-prefix hash hit (blocks) BEFORE the SWA + # bounded_hit gate, recorded by BlockManager.can_allocate. The gap + # against the admitted num_cached_blocks is the reuse lost to a missing + # SWA tail (vs lost to compressed eviction). See CacheStats. + self.num_compressed_hit_blocks = 0 # True iff this seq is mid-prefill (chunked prefill produced KV for # some prompt tokens but not all). Maintained by the scheduler: # set in postprocess when an advance leaves prompt tokens remaining, @@ -91,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 1c122eabf5..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 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: @@ -36,11 +37,41 @@ def __init__( block_size: int, max_num_batched_tokens: int, mtp_k: int, + 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 self.block_size: int = block_size self.max_num_batched_tokens: int = max_num_batched_tokens + # Full-retention mode (ATOM_SWA_FULL_RETAIN): write + materialize EVERY + # SWA block of a prefill chunk (not only the trailing window), so the + # content-addressed cache holds the full history for cross-request replay + # hits. The live sliding-window free stays on (bounds active refs); the + # larger pool keeps freed-but-cached blocks resident until reuse. + self.full_retain: bool = bool(full_retain) + # Sparse checkpoint-tail retention (ATOM_SWA_RETENTION_INTERVAL, tokens). + # 0 = dense (retain every written tail, relies on pool size). >0 = keep a + # SWA tail only once per `retention_interval`-token segment plus at each + # request's prompt boundary, and PIN those checkpoint blocks (an extra + # ref so free_out_of_window / eviction skip them) so live-window churn + # cannot overwrite them. Mirrors vLLM's SlidingWindowManager sparse + # reachable_block_mask, adapted to ATOM's separate (small, isolated) SWA + # pool where masking alone is insufficient — the pin is required. LRU- + # capped at `checkpoint_frac` of the pool so live churn keeps headroom. + self.retention_blocks: int = ( + retention_interval // block_size + if (retention_interval > 0 and block_size > 0) + else 0 + ) + self.sparse_retain: bool = self.full_retain and self.retention_blocks > 0 + self.checkpoint_capacity: int = ( + int(num_blocks * checkpoint_frac) if self.sparse_retain else 0 + ) + # block_id -> None, ordered by pin/access recency (front = LRU). + self.checkpoint_lru: OrderedDict[int, None] = OrderedDict() # Prefix-cache hit gate: a hit only needs the trailing window before the # boundary to be SWA-present (SWA is local). `tail_blocks` = contiguous # blocks covering win_with_spec = window + mtp_k (spec-decode tail tokens @@ -53,34 +84,84 @@ 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: + """Whether logical block `i` is a retained checkpoint tail: it sits in the + trailing `tail_blocks` of a `retention_blocks`-sized segment, OR in the + trailing `tail_blocks` before the prompt boundary (a proven reuse point). + Mirrors vLLM SlidingWindowManager.reachable_block_mask (segment tails + + reachable-boundary tails).""" + if not self.sparse_retain: + return True + need = self.tail_blocks + rb = self.retention_blocks + if i % rb >= rb - need: # last `need` blocks of this segment + return True + prompt_blocks = seq.num_prompt_tokens // self.block_size + if i >= prompt_blocks - need: # trailing tail before the prompt boundary + return True + return False + + def _pin_checkpoint(self, block_id: int) -> None: + """Pin a checkpoint SWA block: hold an extra ref so free_out_of_window + never returns it to the free list, keeping its content-addressed tail + resident for cross-request reuse. LRU-evict the oldest pin when over + capacity. Idempotent per block_id (a re-pin just refreshes recency).""" + if block_id in self.checkpoint_lru: + self.checkpoint_lru.move_to_end(block_id) + return + self.blocks[block_id].ref_count += 1 # the pin ref + self.checkpoint_lru[block_id] = None + while len(self.checkpoint_lru) > self.checkpoint_capacity: + old_id, _ = self.checkpoint_lru.popitem(last=False) # LRU + blk = self.blocks[old_id] + blk.ref_count -= 1 # drop the pin ref + if blk.ref_count == 0: # no live seq holds it → reclaim + self._dealloc(old_id) # --------------------------- admission / hit --------------------------- # def has_free(self, n: int) -> bool: @@ -88,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. @@ -100,6 +181,15 @@ def admission_blocks(self, seq: Sequence) -> int: `seq.num_blocks` since SWA is filled incrementally + window-freed.""" if not self.enabled: return 0 + if self.full_retain: + # Full-retain materializes the whole current chunk (ensure_for_tokens + # free_before=0), so the peak concurrent footprint per prefill step is + # the chunk's block span (bounded by max_num_batched_tokens), NOT the + # trailing window. Freed after each chunk, so this is the per-step peak + # not the whole prompt. Capped by the prompt's block count. + bs = self.block_size + chunk_peak = (self.max_num_batched_tokens + bs - 1) // bs + 1 + return min(chunk_peak, seq.num_blocks) cap = self.tail_blocks + 1 return min(cap, seq.num_blocks) @@ -145,13 +235,15 @@ 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) seq.swa_block_table.append(swa_id) def alloc_placeholder(self, seq: Sequence): @@ -194,7 +286,11 @@ def ensure_for_tokens( # matching free_out_of_window's sentinel. Cuts prefill SWA allocation # from O(chunk_len/bs) to O(window/bs) — pairs with the window-only # swa_write in deepseek_v4.py. free_before mirrors free_out_of_window. - free_before = max(0, (seq_len - self.window) // bs) + # Full-retain: materialize every block of this chunk (free_before=0) so + # the full-chunk swa_write below has a valid physical dst for every token, + # and every block gets published for cross-request reuse. Default: + # window-only (only the trailing-window blocks the SWA read touches). + free_before = 0 if self.full_retain else max(0, (seq_len - self.window) // bs) start_blk = max(start_blk, free_before) table = seq.swa_block_table for i in range(start_blk, end_blk + 1): @@ -289,6 +385,12 @@ def publish_hash(self, seq: Sequence, i: int, h: int, token_ids: list[int]): block = self.blocks[swa_id] block.update(h, token_ids) self.hash_to_block_id[h] = block.block_id + # Sparse retention: pin this tail iff it is a checkpoint (segment or + # prompt-boundary tail), so live-window free/churn cannot overwrite it + # before a branch reuses it. Non-checkpoints stay unpinned and are + # reclaimed normally. + if self.sparse_retain and self._is_checkpoint(seq, i): + self._pin_checkpoint(swa_id) def release(self, seq: Sequence): """Release all of seq's SWA blocks (skipping -1 window-freed slots) and 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 3d919bd534..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]. @@ -523,6 +581,12 @@ def swa_pool_num_blocks(self, max_num_seqs: int, max_model_len: int) -> int: ceil(max_model_len/bs) (e.g. 1024 → ~66 blocks at 131072/8192/128/128). """ bs = self.block_size + # NOTE: full-retain (ATOM_SWA_FULL_RETAIN) does NOT size the pool here. + # It sizes num_swa_blocks == num_kvcache_blocks from the shared memory + # budget in ModelRunner._compute_kv_budget (lockstep with the compressed + # pool), which is memory-bounded. Sizing on max_model_len here would + # explode to ~TB at DSV4's 1M max_position_embeddings. This method is only + # consulted for the default (window-only) pool below. # per_decode uses win_with_spec (= window + max_spec_steps), not window # alone: under MTP each decoding seq writes up to `max_spec_steps` draft # tokens into the SWA pool before the next window-free, so its peak @@ -652,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, ) @@ -678,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, ) @@ -870,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: @@ -894,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: @@ -1365,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 ---- @@ -1397,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( @@ -1635,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 ) @@ -2259,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. @@ -2278,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( @@ -2343,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 @@ -2523,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] @@ -2572,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 ----- @@ -2596,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, @@ -2691,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 ): @@ -3006,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 86675f3b61..9b53b26896 100644 --- a/atom/models/deepseek_v4.py +++ b/atom/models/deepseek_v4.py @@ -135,6 +135,7 @@ _V4_USE_REF_QUANT = os.environ.get("V4_USE_REF_QUANT", "0") == "1" # Fused-kernel switches. Default off; flip via env to A/B against the eager path. _V4_USE_TRITON_FUSION = os.environ.get("ATOM_V4_USE_TRITON_FUSION", "0") == "1" +_V4_SWA_FULL_RETAIN = envs.ATOM_SWA_FULL_RETAIN ENABLE_DS_QKNORM_QUANT_FUSION = envs.ATOM_ENABLE_DS_QKNORM_QUANT_FUSION SPARSE_INDEXER_LOGITS_BUDGET_MB = envs.ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB @@ -1862,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 @@ -1899,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: @@ -1914,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 @@ -2046,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 @@ -2146,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 @@ -2386,7 +2411,17 @@ def _attn_core( swa_block_tables_gpu, self.swa_kv, swa_block_size, - min(self.window_size, attn_md.max_seqlen_q), + # Full-retain (ATOM_SWA_FULL_RETAIN): persist the WHOLE chunk's + # SWA KV (every token) so the content-addressed cache holds the + # full history for cross-request prefix reuse. Default: window-only + # (trailing `window` tokens) — see the OPT note above. Pairs with + # SlidingWindowPool.ensure_for_tokens materializing all chunk + # blocks (free_before=0) so every dst block is valid here. + ( + attn_md.max_seqlen_q + if _V4_SWA_FULL_RETAIN + else min(self.window_size, attn_md.max_seqlen_q) + ), k_packed=qkn.k_packed, k_rope=qkn.k_rope, swa_region_rope=self.swa_kv_rope, @@ -2467,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 962b85bb70..a1a489170f 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -84,6 +84,51 @@ "ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION", "1") == "1" ), + # DeepSeek-V4 paged-SWA: retain the FULL sliding-window KV history in the + # content-addressed SWA cache instead of the default window-only prefill + # write. Default OFF preserves the window-only optimization (writes only each + # prefill chunk's trailing `window` tokens). When ON, swa_write persists the + # whole chunk and ensure_for_tokens materializes every block, so cross-request + # prefix hits can reuse the middle SWA blocks (agentic branch/replay reuse) — + # the live sliding-window free is UNCHANGED (out-of-window refs still released + # each chunk/decode; freed blocks stay hash+KV resident until overwritten). + # 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 + # prefix index; a small fraction keeps compressed near full while retaining + # the hot-boundary tail working set (LRU-evicted). Only consulted when + # ATOM_SWA_FULL_RETAIN=1. Default 0.2; tune 0.15-0.25 with cache-hit + # instrumentation. Clamped to (0, 0.9). + "ATOM_SWA_TAIL_BUDGET_FRAC": lambda: float( + os.getenv("ATOM_SWA_TAIL_BUDGET_FRAC", "0.2") + ), + # DeepSeek-V4 paged-SWA sparse checkpoint retention (only with + # ATOM_SWA_FULL_RETAIN=1). Tokens per retained SWA-tail checkpoint: 0 = dense + # (retain every written tail, relies on pool size — floods a small pool); + # >0 = keep a tail only once per this-many-tokens segment plus at each prompt + # boundary, and PIN those so live-window churn cannot overwrite them (mirrors + # vLLM VLLM_PREFIX_CACHE_RETENTION_INTERVAL / SlidingWindowManager sparse + # reachable_block_mask). Should be a multiple of the KV block size; 32768 + # matches the vLLM trace-replay tuning. Default 0 (dense). + "ATOM_SWA_RETENTION_INTERVAL": lambda: int( + os.getenv("ATOM_SWA_RETENTION_INTERVAL", "0") + ), + # Fraction of the SWA pool that pinned checkpoint tails may occupy (LRU-capped) + # when sparse retention is on; the rest stays free for live-window churn. + "ATOM_SWA_CHECKPOINT_FRAC": lambda: float( + os.getenv("ATOM_SWA_CHECKPOINT_FRAC", "0.5") + ), # DSA sparse-indexer prefill: KV-dimension chunk size (in tokens) for # `fp8_mqa_logits`. The dense logits buffer is [prefill_tokens, total_kv]; # total_kv = sum of all co-scheduled prefill contexts and is NOT bounded by 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