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

atom/model_engine/block_manager.py

@@ -17,6 +17,10 @@
 from atom.model_engine.kv_block import Block
 from atom.model_engine.sequence import Sequence
 from atom.model_engine.swa_pool import SlidingWindowPool
+from atom.model_engine.unified_chunk_pool import (
+    UnifiedKvCoordinator,
+    unified_slot_counts,
+)
 from atom.utils import envs
 
 
@@ -83,16 +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:
@@ -129,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.
@@ -163,6 +247,10 @@ def can_allocate(self, seq: Sequence) -> int:
             # when SWA disabled.
             if not self.swa.has_free(self.swa.admission_blocks(seq)):
                 return -1
+            # Unified pool: each new logical compress block needs a CSA + HCA
+            # physical block (no-op / True flag-off).
+            if not self._unified_has_free(seq.num_blocks):
+                return -1
             return 0
         # Step 1: compressed prefix (CSA/HCA/indexer share the block hash and
         # read the WHOLE history, so this stays a full front-to-back chained
@@ -203,6 +291,10 @@ def can_allocate(self, seq: Sequence) -> int:
         # True when SWA disabled.
         if not self.swa.has_free(min(num_new_blocks, self.swa.admission_blocks(seq))):
             return -1
+        # Unified pool: physical CSA/HCA blocks for the uncached logical blocks
+        # (cache-hit blocks keep their retained binding). No-op / True flag-off.
+        if not self._unified_has_free(num_new_blocks):
+            return -1
         return num_cached_blocks
 
     def allocate(self, seq: Sequence, num_cached_blocks: int = 0):
@@ -244,6 +336,7 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0):
                 self.free_block_ids_set.discard(block_id)
                 self.used_block_ids.add(block_id)
             seq.block_table.append(block_id)
+            self._append_phys(seq, block_id)  # unified pool (no-op flag-off)
             if i < swa_hit_start:
                 self.swa.alloc_placeholder(seq)  # out of window: never read → -1
             else:
@@ -252,6 +345,7 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0):
             block_id = self._pop_free_block()
             self._allocate_block(block_id)
             seq.block_table.append(block_id)
+            self._append_phys(seq, block_id)  # unified pool (no-op flag-off)
             # Uncached blocks: -1 placeholder keeps swa_block_table the same
             # length as block_table; ensure_for_tokens fills the current chunk's
             # window slots before each forward, free_after_prefill_chunk releases
@@ -326,6 +420,13 @@ def deallocate(self, seq: Sequence):
         )  # release SWA blocks + clear swa_block_table (no-op if disabled)
         seq.num_cached_tokens = 0
         seq.block_table.clear()
+        if self._unified is not None:
+            # Physical bindings are retained on the logical blocks (lazy reuse);
+            # only the per-seq positional tables are cleared here, mirroring
+            # block_table. (deallocate above dropped ref_counts; blocks that hit 0
+            # went to the free pool with their CSA/HCA binding intact.)
+            seq.csa_block_table.clear()
+            seq.hca_block_table.clear()
         if seq.has_per_req_cache and seq.per_req_cache_group >= 0:
             self.free_per_req_cache_groups.append(seq.per_req_cache_group)
             seq.per_req_cache_group = -1
@@ -341,6 +442,8 @@ def can_append(self, seq: Sequence, num_new_tokens: int = 1) -> bool:
             return False
         if not self.swa.has_free(new_blocks_needed):  # True when SWA disabled
             return False
+        if not self._unified_has_free(new_blocks_needed):  # unified pool (no-op off)
+            return False
         return True
 
     def may_append(self, seq: Sequence, num_new_tokens: int = 1):
@@ -361,6 +464,7 @@ def may_append(self, seq: Sequence, num_new_tokens: int = 1):
                 block_id = self._pop_free_block()
                 self._allocate_block(block_id)
                 block_table.append(block_id)
+                self._append_phys(seq, block_id)  # unified pool (no-op flag-off)
                 self.swa.append_new(seq)  # lockstep SWA block (no-op if disabled)
         # Reclaim SWA blocks that just fell out of the window (no-op if disabled).
         self.swa.free_out_of_window(seq, len(seq))

atom/model_engine/scheduler.py

@@ -354,6 +354,16 @@ def __init__(
         self.swa_block_tables = [
             seq.swa_block_table for seq in seqs.values() if seq.block_table
         ]
+        # DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE / plan §11): per-type
+        # PHYSICAL compress tables (base-0, row = phys*k). CSA and HCA are allocated
+        # independently so each can borrow SWA-freed slots. `block_table` above
+        # stays LOGICAL. Empty (flag-off) → downstream reverts to block_tables.
+        self.csa_block_tables = [
+            seq.csa_block_table for seq in seqs.values() if seq.block_table
+        ]
+        self.hca_block_tables = [
+            seq.hca_block_table for seq in seqs.values() if seq.block_table
+        ]
         self.last_block_num_tokens = [
             _seq.last_block_num_tokens for _seq in seqs.values()
         ]

atom/model_engine/sequence.py

@@ -96,6 +96,14 @@ def __init__(
         # out-of-window SWA blocks can be freed while compressed blocks persist).
         # Empty / unused for non-SWA models.
         self.swa_block_table = []
+        # DeepSeek-V4 unified KV pool (ATOM_UNIFIED_KV_SHARE / plan §11): per-type
+        # PHYSICAL compress block tables. `block_table` stays the LOGICAL compress
+        # id (prefix-cache / hashing unchanged); these hold the physical block ids
+        # in the CSA (k=32) and HCA (k=1) pools, which are allocated independently
+        # so each can borrow SWA-freed slots (base-0 addressing, row = phys * k).
+        # Positionally aligned with `block_table`. Empty unless the flag is on.
+        self.csa_block_table = []
+        self.hca_block_table = []
         # Per-request cache slot index (filled by BlockManager.allocate()).
         # -1 = unallocated. The slot indexes into the per-req cache tensors
         # owned by ModelRunner (e.g. mamba_k_cache for GDN).

atom/model_engine/swa_pool.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
-from collections import OrderedDict, deque
+from collections import OrderedDict
 
 from atom.model_engine.kv_block import Block
 from atom.model_engine.sequence import Sequence
+from atom.model_engine.unified_chunk_pool import UnifiedTypePool
 
 
 class SlidingWindowPool:
@@ -39,6 +40,7 @@ def __init__(
         full_retain: bool = False,
         retention_interval: int = 0,
         checkpoint_frac: float = 0.5,
+        chunk_pool: "UnifiedTypePool | None" = None,
     ):
         self.enabled: bool = num_blocks > 0
         self.window: int = window
@@ -82,34 +84,49 @@ def __init__(
         )
         self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
         self.hash_to_block_id: dict[int, int] = dict()
-        self.free_block_ids: deque[int] = deque(range(num_blocks))
-        self.free_block_ids_set: set[int] = set(range(num_blocks))
-        self.used_block_ids: set[int] = set()
+        # Free-list moved to the shared UnifiedChunkPool (plan §5.2): the SWA pool
+        # draws physical block ids from it, so freed SWA blocks can (Phase 2) be
+        # reused by the compressor and vice versa. Phase 1: SWA-only side, spc=1 →
+        # behaviorally identical to the old deque/set. When BlockManager injects a
+        # shared pool it is used directly; otherwise own a private SWA-only one.
+        # SWA draws slots from a per-type UnifiedTypePool. Flag-off (static split):
+        # a private SWA-only pool (k=0), num_swa_blocks slots → phys = slot = w,
+        # row = w*block_size, byte-identical to the old free-list. Flag-on: a
+        # shared per-type pool is injected (SWA + that type's compress share it).
+        self.chunk_pool: UnifiedTypePool = chunk_pool or UnifiedTypePool(
+            num_blocks, k=0, block_size=block_size
+        )
 
     # ----------------------------- primitives ------------------------------ #
     def _pop(self) -> int:
-        while self.free_block_ids:
-            block_id = self.free_block_ids.popleft()
-            if block_id in self.free_block_ids_set:
-                self.free_block_ids_set.discard(block_id)
-                return block_id
-        raise AssertionError("No free SWA blocks available")
+        """Get a free SWA block id from the shared chunk pool (marks it used)."""
+        return self.chunk_pool.alloc_swa()
 
     def _alloc(self, block_id: int) -> Block:
+        """Reset block metadata for a freshly-popped id (lazy hash eviction).
+        Free-list state is owned by chunk_pool; _pop already marked it used."""
         block = self.blocks[block_id]
         assert block.ref_count == 0
         if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
             del self.hash_to_block_id[block.hash]
         block.reset()
-        self.free_block_ids_set.discard(block_id)
-        self.used_block_ids.add(block_id)
         return block
 
     def _dealloc(self, block_id: int):
         assert self.blocks[block_id].ref_count == 0
-        self.used_block_ids.remove(block_id)
-        self.free_block_ids.append(block_id)
-        self.free_block_ids_set.add(block_id)
+        self.chunk_pool.free_swa(block_id)
+
+    def evict_slot(self, block_id: int) -> None:
+        """Drop the content-cache entry for SWA block `block_id` because a compress
+        pool is about to overwrite its physical slot (unified pool, plan §11.1).
+        Registered as the coordinator's SWA evict callback. Only called for
+        released (ref-0) slots; pinned checkpoints hold a ref → never released, so
+        never reach here. Idempotent / safe if the slot holds no live hash."""
+        block = self.blocks[block_id]
+        if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
+            del self.hash_to_block_id[block.hash]
+        block.hash = -1
+        self.checkpoint_lru.pop(block_id, None)
 
     # ------------------------ sparse checkpoint pins ----------------------- #
     def _is_checkpoint(self, seq: Sequence, i: int) -> bool:
@@ -152,7 +169,7 @@ def has_free(self, n: int) -> bool:
         blocks admission)."""
         if not self.enabled:
             return True
-        return len(self.free_block_ids_set) >= n
+        return self.chunk_pool.has_free_swa(n)
 
     def admission_blocks(self, seq: Sequence) -> int:
         """Peak concurrent SWA blocks one request holds during (chunked) prefill.
@@ -218,13 +235,12 @@ def claim_cached(self, seq: Sequence, h: int, token_ids: list[int]):
             return
         swa_id = self.hash_to_block_id[h]
         block = self.blocks[swa_id]
-        if swa_id in self.used_block_ids:
+        if not self.chunk_pool.is_free_swa(swa_id):
             block.ref_count += 1
         else:
             assert block.ref_count == 0
             block.ref_count = 1
-            self.free_block_ids_set.discard(swa_id)
-            self.used_block_ids.add(swa_id)
+            self.chunk_pool.claim_swa(swa_id)
         # Cross-request reuse of a pinned checkpoint → refresh its LRU recency.
         if swa_id in self.checkpoint_lru:
             self.checkpoint_lru.move_to_end(swa_id)