diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index cd440c099..c5f5e0561 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/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,18 +167,35 @@ 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 - ring, never offloaded) have no per-pool staging path yet. + Sidecars reuse their source pool's transient slot ids, so every sidecar + 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: - raise ValueError( - "--hicache-host-memory-mode buffer_only does not support " - "sidecar storage pools (DeepSeek-V4 compressed regions)." - ) + 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 sidecar pool mapping " + 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: # Only reachable on SWA models with the unified_kv layout (SWA as @@ -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,20 +653,46 @@ 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 - # hashes it covers: the SWA window spans page_size-sized pages, the - # Mamba state is a single slot (host pool page_size 1 -> one key). + # Every independently staged aux pool writes a trailing snapshot keyed + # by the last KV page hashes it covers. A derived sidecar writes the + # 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 + transfer = 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, - host_indices=staged.host_indices, - keys=keys, - hit_policy=PoolHitPolicy.TRAILING_PAGES, + keys=source.keys, + hit_policy=staged.hit_policy, + indices_from_pool=staged.indices_from_pool, ) ) operation_id = self._cache.cache_controller.write_storage( @@ -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. diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index d100a416d..6c13b7caf 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py new file mode 100644 index 000000000..68d634617 --- /dev/null +++ b/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()