[HiCache] Buffer mode support sidecar pool (#37424)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

python/sglang/srt/mem_cache/buffer_mode/pipeline.py

@@ -42,7 +42,12 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     InsertParams,
     MatchPrefixParams,
 )
-from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    PoolTransfer,
+    SidecarPoolSpec,
+)
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.unified_cache.cache_action import RebuildFullToSWAMapping
 from sglang.srt.mem_cache.unified_cache.components import (
@@ -56,6 +61,7 @@ from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
 )
 
 if TYPE_CHECKING:
+    from sglang.srt.mem_cache.pool_host import HostPoolGroup
     from sglang.srt.mem_cache.unified_cache.components import SWAComponent
     from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
 
@@ -161,17 +167,34 @@ def staged_splice_tokens(f: _StagedPrefetch, device_prefix_len: int) -> int:
 
 
 def validate_buffer_only_stack(
-    sidecar_pool_specs: list, swa_component: Optional[SWAComponent]
+    sidecar_pool_specs: list[SidecarPoolSpec],
+    host_pool_group: HostPoolGroup,
+    swa_component: Optional[SWAComponent],
 ) -> None:
     """Post-assembly buffer-mode fences.
 
-    Sidecar pools (DSv4 compressed regions) and unified_kv SWA (device-only
+    Sidecars reuse their source pool's transient slot ids, so every sidecar
-    ring, never offloaded) have no per-pool staging path yet.
+    host pool must expose the full source slot namespace.  unified_kv SWA
+    (device-only ring, never offloaded) still has no staging path.
     """
-    if sidecar_pool_specs:
+    entry_map = host_pool_group.entry_map
+    for spec in sidecar_pool_specs:
+        source = entry_map.get(spec.indices_from_pool)
+        sidecar = entry_map.get(spec.pool_name)
+        if source is None or sidecar is None:
             raise ValueError(
-            "--hicache-host-memory-mode buffer_only does not support "
+                "--hicache-host-memory-mode buffer_only sidecar pool mapping "
-            "sidecar storage pools (DeepSeek-V4 compressed regions)."
+                f"is incomplete: pool={spec.pool_name}, "
+                f"indices_from_pool={spec.indices_from_pool}."
+            )
+        source_size = source.host_pool.logical_size
+        sidecar_size = sidecar.host_pool.logical_size
+        if sidecar_size < source_size:
+            raise ValueError(
+                "--hicache-host-memory-mode buffer_only sidecar host pool is "
+                "smaller than its index source: "
+                f"pool={spec.pool_name}, host_slots={sidecar_size}, "
+                f"source={spec.indices_from_pool}, source_slots={source_size}."
             )
     swa = swa_component
     if swa is not None and swa._swa_kv_pool_host is None:
@@ -547,6 +570,11 @@ class BufferModePipeline:
         cc = cache.cache_controller
         snapshot = intent.snapshot
         aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
+        # Sidecars reuse the source pool's transient host/device indices.  This
+        # includes both KV-derived pools and SWA-derived DSV4 state pools.  They
+        # allocate no additional staging, but must ride the same D2H operation
+        # so their bytes are present when the storage write starts.
+        aux_xfers.extend(cache._build_backup_sidecar(device_value, comp_xfers))
         host_indices = cc.write(
             device_value,
             node_id=snapshot.node_id,
@@ -625,21 +653,47 @@ class BufferModePipeline:
         snapshot = intent.snapshot
         self._cache.dec_lock_ref(snapshot.node_id, entry.lock_params)
 
-        # Every aux pool writes a trailing snapshot keyed by the last KV page
+        # Every independently staged aux pool writes a trailing snapshot keyed
-        # hashes it covers: the SWA window spans page_size-sized pages, the
+        # by the last KV page hashes it covers.  A derived sidecar writes the
-        # Mamba state is a single slot (host pool page_size 1 -> one key).
+        # exact key span of its source pool while reusing that source's slots.
         storage_xfers: list[PoolTransfer] = []
+        storage_sources = {
+            PoolName.KV: PoolTransfer(
+                name=PoolName.KV,
+                host_indices=entry.host_indices,
+                keys=snapshot.hash_values,
+            )
+        }
         for staged in entry.aux_xfers:
+            if staged.indices_from_pool is not None:
+                continue
             keys = self._aux_window_keys(snapshot.hash_values, staged)
             if keys is None:
                 continue
-            storage_xfers.append(
+            transfer = PoolTransfer(
-                PoolTransfer(
                 name=staged.name,
                 host_indices=staged.host_indices,
                 keys=keys,
                 hit_policy=PoolHitPolicy.TRAILING_PAGES,
             )
+            storage_xfers.append(transfer)
+            storage_sources.setdefault(staged.name, transfer)
+        for staged in entry.aux_xfers:
+            if staged.indices_from_pool is None:
+                continue
+            source = storage_sources.get(staged.indices_from_pool)
+            if source is None:
+                raise AssertionError(
+                    "Buffer-mode storage sidecar source missing: "
+                    f"{staged.name} from {staged.indices_from_pool}."
+                )
+            storage_xfers.append(
+                PoolTransfer(
+                    name=staged.name,
+                    keys=source.keys,
+                    hit_policy=staged.hit_policy,
+                    indices_from_pool=staged.indices_from_pool,
+                )
             )
         operation_id = self._cache.cache_controller.write_storage(
             entry.host_indices,
@@ -832,6 +886,14 @@ class BufferModePipeline:
         prefix_ctx = self._prefetch_prefix_ctx.pop(req_id, None)
         prefix_tokens = prefix_ctx[0] if prefix_ctx is not None else None
         aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
+        # Component transfers are already present in comp_xfers.  Preserve the
+        # derived sidecars from the storage operation as well; cc.load resolves
+        # them against the freshly allocated source host/device indices.
+        aux_xfers.extend(
+            transfer
+            for transfer in operation.pool_transfers or ()
+            if transfer.indices_from_pool is not None
+        )
 
         if num_tokens == 0 or prefix_tokens is None:
             # Nothing usable fetched: recompute.

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -439,7 +439,9 @@ class UnifiedRadixCache(BasePrefixCache):
         if self.host_memory_mode == "buffer_only":
             swa = self.components.get(ComponentType.SWA)
             validate_buffer_only_stack(
-                sidecar_pool_specs=self.sidecar_pool_specs, swa_component=swa
+                sidecar_pool_specs=self.sidecar_pool_specs,
+                host_pool_group=self.host_pool_group,
+                swa_component=swa,
             )
             self.buffer_pipeline = BufferModePipeline(
                 cache=self,

test/registered/unit/mem_cache/test_buffer_mode_sidecar.py

@@ -0,0 +1,282 @@
+"""Unit coverage for sidecar pools in HiCache buffer-only mode."""
+
+import unittest
+from array import array
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.mem_cache.buffer_mode.pipeline import (
+    BufferModePipeline,
+    _UnifiedBackupIntent,
+    validate_buffer_only_stack,
+)
+from sglang.srt.mem_cache.hicache_storage import (
+    PoolHitPolicy,
+    PoolName,
+    PoolTransfer,
+    SidecarPoolSpec,
+)
+from sglang.srt.mem_cache.radix_cache import RadixKey
+from sglang.srt.mem_cache.unified_cache.components.tree_component import (
+    ComponentType,
+)
+from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import (
+    BufferBackupSnapshot,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestBufferModeSidecar(unittest.TestCase):
+    @staticmethod
+    def _swa_component():
+        return SimpleNamespace(
+            full_window_pages=2,
+            _swa_kv_pool_host=SimpleNamespace(page_size=2, size=8),
+        )
+
+    @staticmethod
+    def _dsv4_specs():
+        return [
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C4,
+                indices_from_pool=PoolName.KV,
+            ),
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C4_INDEXER,
+                indices_from_pool=PoolName.KV,
+            ),
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C128,
+                indices_from_pool=PoolName.KV,
+            ),
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C4_STATE,
+                indices_from_pool=PoolName.SWA,
+                hit_policy=PoolHitPolicy.TRAILING_PAGES,
+            ),
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
+                indices_from_pool=PoolName.SWA,
+                hit_policy=PoolHitPolicy.TRAILING_PAGES,
+            ),
+            SidecarPoolSpec(
+                pool_name=PoolName.DEEPSEEK_V4_C128_STATE,
+                indices_from_pool=PoolName.SWA,
+                hit_policy=PoolHitPolicy.TRAILING_PAGES,
+            ),
+        ]
+
+    @classmethod
+    def _pool_group(
+        cls,
+        kv_size: int,
+        swa_size: int,
+        *,
+        override_size: dict[PoolName, int] | None = None,
+    ):
+        sizes = {PoolName.KV: kv_size, PoolName.SWA: swa_size}
+        for spec in cls._dsv4_specs():
+            sizes[spec.pool_name] = sizes[spec.indices_from_pool]
+        sizes.update(override_size or {})
+        return SimpleNamespace(
+            entry_map={
+                name: SimpleNamespace(host_pool=SimpleNamespace(logical_size=size))
+                for name, size in sizes.items()
+            }
+        )
+
+    def test_stack_accepts_dsv4_full_and_swa_sidecars(self):
+        validate_buffer_only_stack(
+            sidecar_pool_specs=self._dsv4_specs(),
+            host_pool_group=self._pool_group(kv_size=16, swa_size=8),
+            swa_component=self._swa_component(),
+        )
+
+    def test_stack_rejects_sidecar_smaller_than_source(self):
+        with self.assertRaisesRegex(ValueError, "smaller than its index source"):
+            validate_buffer_only_stack(
+                sidecar_pool_specs=[
+                    SidecarPoolSpec(
+                        pool_name=PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
+                        indices_from_pool=PoolName.SWA,
+                        hit_policy=PoolHitPolicy.TRAILING_PAGES,
+                    )
+                ],
+                host_pool_group=self._pool_group(
+                    kv_size=16,
+                    swa_size=8,
+                    override_size={PoolName.DEEPSEEK_V4_C4_INDEXER_STATE: 4},
+                ),
+                swa_component=self._swa_component(),
+            )
+
+    def test_write_stages_and_persists_dsv4_full_and_swa_sidecars(self):
+        page_size = 2
+        device_indices = torch.arange(4, dtype=torch.int64)
+        host_indices = torch.arange(10, 14, dtype=torch.int64)
+        swa_device_indices = torch.arange(30, 32, dtype=torch.int64)
+        swa_host_indices = torch.arange(20, 22, dtype=torch.int64)
+        swa = PoolTransfer(
+            name=PoolName.SWA,
+            device_indices=swa_device_indices,
+            hit_policy=PoolHitPolicy.TRAILING_PAGES,
+        )
+        sidecars = [
+            PoolTransfer(
+                name=spec.pool_name,
+                hit_policy=spec.hit_policy,
+                indices_from_pool=spec.indices_from_pool,
+            )
+            for spec in self._dsv4_specs()
+        ]
+
+        controller = MagicMock()
+        controller.mem_pool_host.entry_map = {
+            PoolName.SWA: SimpleNamespace(
+                host_pool=SimpleNamespace(page_size=page_size)
+            )
+        }
+
+        def _write(device_value, *, node_id, extra_pools):
+            self.assertEqual(node_id, 7)
+            self.assertEqual(
+                [transfer.name for transfer in extra_pools],
+                [PoolName.SWA, *[transfer.name for transfer in sidecars]],
+            )
+            # HostPoolGroup.resolve_host_transfers gives derived pools their
+            # source pool's indices without allocating another staging span.
+            swa.host_indices = swa_host_indices
+            for sidecar in sidecars:
+                if sidecar.indices_from_pool == PoolName.KV:
+                    sidecar.host_indices = host_indices
+                    sidecar.device_indices = device_value
+                else:
+                    sidecar.host_indices = swa_host_indices
+                    sidecar.device_indices = swa_device_indices
+            return host_indices
+
+        controller.write.side_effect = _write
+        controller.write_storage.return_value = 99
+
+        cache = MagicMock()
+        cache.cache_controller = controller
+        cache.page_size = page_size
+        cache._build_backup_sidecar.return_value = sidecars
+
+        pipeline = BufferModePipeline.__new__(BufferModePipeline)
+        pipeline._cache = cache
+        pipeline.ongoing_write_through = {}
+        pipeline.ongoing_backup = {}
+        pipeline.inflight_backup_hashes = {}
+        pipeline.write_staged_tokens_ = 0
+        pipeline.write_backlog_tokens_ = len(device_indices)
+
+        hashes = ["page-0", "page-1"]
+        snapshot = BufferBackupSnapshot(
+            node_id=7,
+            parent_node_id=0,
+            parent_is_root=True,
+            parent_last_hash=None,
+            hash_values=hashes,
+            key=RadixKey(array("q", [1, 2, 3, 4])),
+            prefix_keys=None,
+        )
+        intent = _UnifiedBackupIntent(snapshot=snapshot)
+
+        self.assertTrue(
+            pipeline._launch_backup_intent(
+                intent,
+                device_indices,
+                comp_xfers={ComponentType.SWA: [swa]},
+            )
+        )
+        self.assertEqual(pipeline.ongoing_write_through[7].aux_xfers, [swa, *sidecars])
+
+        pipeline.finish_backup_ack(7)
+
+        storage_transfers = {
+            transfer.name: transfer
+            for transfer in controller.write_storage.call_args.kwargs["extra_pools"]
+        }
+        self.assertEqual(storage_transfers[PoolName.SWA].keys, [hashes[-1]])
+        self.assertTrue(
+            torch.equal(storage_transfers[PoolName.SWA].host_indices, swa_host_indices)
+        )
+        for spec in self._dsv4_specs():
+            storage_sidecar = storage_transfers[spec.pool_name]
+            expected_keys = (
+                hashes if spec.indices_from_pool == PoolName.KV else [hashes[-1]]
+            )
+            self.assertEqual(storage_sidecar.keys, expected_keys)
+            self.assertEqual(storage_sidecar.hit_policy, spec.hit_policy)
+            self.assertEqual(storage_sidecar.indices_from_pool, spec.indices_from_pool)
+            self.assertIsNone(storage_sidecar.host_indices)
+        self.assertIn(99, pipeline.ongoing_backup)
+
+    def test_completed_prefetch_keeps_dsv4_full_and_swa_sidecars_for_h2d(self):
+        swa = PoolTransfer(
+            name=PoolName.SWA,
+            host_indices=torch.arange(20, 22, dtype=torch.int64),
+            hit_policy=PoolHitPolicy.TRAILING_PAGES,
+        )
+        sidecars = [
+            PoolTransfer(
+                name=spec.pool_name,
+                hit_policy=spec.hit_policy,
+                indices_from_pool=spec.indices_from_pool,
+            )
+            for spec in self._dsv4_specs()
+        ]
+        operation = SimpleNamespace(
+            id=23,
+            pool_transfers=sidecars,
+            storage_start=0,
+        )
+        host_indices = torch.arange(4, dtype=torch.int64)
+        req_id = "sidecar-prefetch"
+
+        cache = MagicMock()
+        cache.page_size = 2
+        cache.cache_controller.prefetch_tokens_occupied = len(host_indices)
+        cache.ongoing_prefetch = {
+            req_id: (
+                0,
+                RadixKey(array("q", [1, 2, 3, 4])),
+                host_indices,
+                operation,
+                None,
+                {ComponentType.SWA: [swa]},
+            )
+        }
+        cache.prefetch_loaded_tokens_by_reqid = {}
+        cache.prefetch_loaded_storage_start_by_reqid = {}
+        cache.storage_existence_cache = MagicMock()
+
+        pipeline = BufferModePipeline.__new__(BufferModePipeline)
+        pipeline._cache = cache
+        pipeline._prefetch_prefix_ctx = {req_id: ([], None, None)}
+        pipeline.staged_prefetches = {}
+
+        self.assertTrue(
+            pipeline.stage_completed_prefetch(
+                req_id=req_id,
+                num_tokens=len(host_indices),
+                hash_value=["page-0", "page-1"],
+            )
+        )
+
+        staged = pipeline.staged_prefetches[req_id]
+        self.assertEqual(staged.aux_xfers, [swa, *sidecars])
+        self.assertEqual(staged.num_tokens, len(host_indices))
+        self.assertEqual(
+            cache.prefetch_loaded_tokens_by_reqid[req_id], len(host_indices)
+        )
+        self.assertEqual(cache.prefetch_loaded_storage_start_by_reqid[req_id], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()