minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
huangtingwei
2026-09-04 23:24:18 +08:00
committed by GitHub
co-authored by Zhiqiang Xie
parent 8b1d8c1703
commit 0d0e2f92be
3 changed files with 362 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/buffer_mode/pipeline.py
+74 -12
View File
@@ -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
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:
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 "
"sidecar storage pools (DeepSeek-V4 compressed regions)."
"--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:
@@ -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
# 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
storage_xfers.append(
PoolTransfer(
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,
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
+3 -1
View File
@@ -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
+282
View File
@@ -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()