[Feature] DeepSeek-V4 unified KV pool (option-2): SWA↔compress share physical KV

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,13 +174,52 @@ 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):
         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)
+        # 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))

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

atom/model_engine/scheduler.py

@@ -138,28 +138,41 @@ 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):
         self._log_interval = log_interval
         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