[UnifiedTree]: Support l3 storage for swa and deepseek v4 (#26881)

Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
Zhangheng
2026-06-04 10:17:34 +08:00
committed by GitHub
co-authored by 晟海
parent 1a57145975
commit 736263f3dc
9 changed files with 517 additions and 65 deletions
+12 -3
View File
@@ -625,7 +625,15 @@ class HiCacheController:
self.dp_rank = 0 self.dp_rank = 0
# Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool. # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool) # DeepSeekV4TokenToKVPool has compressed MLA-style rank-replicated cache
# data. storage only needs rank 0 to write it back.
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
is_mla_model = isinstance(self.mem_pool_device, MLATokenToKVPool)
is_compressed_mla_model = isinstance(
self.mem_pool_device, DeepSeekV4TokenToKVPool
)
is_rank_replicated = is_mla_model or is_compressed_mla_model
# Least Common Multiple among heterogeneous tp size # Least Common Multiple among heterogeneous tp size
tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None) tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None)
should_split_heads = False should_split_heads = False
@@ -635,7 +643,7 @@ class HiCacheController:
tp_lcm_size % self.tp_size == 0 tp_lcm_size % self.tp_size == 0
), "tp_lcm_size must be divisible by tp_size." ), "tp_lcm_size must be divisible by tp_size."
should_split_heads = ( should_split_heads = (
not is_mla_backend not is_rank_replicated
and self.mem_pool_host.layout == "page_head" and self.mem_pool_host.layout == "page_head"
and tp_lcm_size > self.tp_size and tp_lcm_size > self.tp_size
) )
@@ -649,7 +657,8 @@ class HiCacheController:
pp_size=self.pp_size, pp_size=self.pp_size,
attn_cp_rank=attn_cp_rank, attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size, attn_cp_size=attn_cp_size,
is_mla_model=is_mla_backend, # TODO(hzh): Rename is_mla_model to is_rank_replicated.
is_mla_model=is_rank_replicated,
enable_storage_metrics=self.enable_storage_metrics, enable_storage_metrics=self.enable_storage_metrics,
is_page_first_layout=self.mem_pool_host.layout == "page_first", is_page_first_layout=self.mem_pool_host.layout == "page_first",
model_name=model_name, model_name=model_name,
@@ -74,6 +74,7 @@ class InsertResult:
"""Result of an insert operation""" """Result of an insert operation"""
prefix_len: int prefix_len: int
total_len: int = 0
mamba_exist: bool = False mamba_exist: bool = False
inserted_host_node: Any = None inserted_host_node: Any = None
@@ -136,6 +136,7 @@ class PrefetchOperation(StorageOperation):
prefix_keys=prefix_keys, prefix_keys=prefix_keys,
pool_transfers=pool_transfers, pool_transfers=pool_transfers,
) )
self.pool_transfers_done = not bool(pool_transfers)
def increment(self, num_tokens: int): def increment(self, num_tokens: int):
with self._lock: with self._lock:
@@ -581,9 +582,6 @@ class HybridCacheController(BaseHiCacheController):
kv_hit_pages = hit_result.kv_hit_pages kv_hit_pages = hit_result.kv_hit_pages
operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages) operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages)
if kv_hit_pages > 0 and operation.pool_transfers:
self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages)
return ( return (
hash_value[:kv_hit_pages], hash_value[:kv_hit_pages],
kv_hit_pages * self.page_size, kv_hit_pages * self.page_size,
@@ -618,14 +616,21 @@ class HybridCacheController(BaseHiCacheController):
return host_indices, device_indices, resolved_pool_transfers return host_indices, device_indices, resolved_pool_transfers
def _page_transfer(self, operation): def _page_transfer(self, operation):
# Transfer extra pools # KV pools first — determines actual completed page count
if operation.pool_transfers and not operation.is_terminated(): super()._page_transfer(operation)
# Extra pools only after KV fully completes. If KV terminated early
# (IO failure, timeout, TP mismatch), skip extra IO entirely to avoid
# data misalignment.
kv_completed_pages = operation.completed_tokens // self.page_size
if operation.pool_transfers and kv_completed_pages == len(operation.hash_value):
self._sync_trailing_keys(
operation.pool_transfers, operation.hash_value, kv_completed_pages
)
self._resolve_sidecar_derived_pool_transfers(operation) self._resolve_sidecar_derived_pool_transfers(operation)
results = self.storage_backend.batch_get_v2(operation.pool_transfers) results = self.storage_backend.batch_get_v2(operation.pool_transfers)
operation.pool_storage_result.update_extra_pool_hit_pages(results) operation.pool_storage_result.update_extra_pool_hit_pages(results)
operation.pool_transfers_done = True
# Transfer kv pools
super()._page_transfer(operation)
def _page_backup(self, operation): def _page_backup(self, operation):
# Backup extra pools # Backup extra pools
@@ -642,11 +647,24 @@ class HybridCacheController(BaseHiCacheController):
if transfer.indices_from_pool is None: if transfer.indices_from_pool is None:
continue continue
if transfer.indices_from_pool != PoolName.KV: if transfer.indices_from_pool != PoolName.KV:
# TODO(hzh): Support storage sidecar derived pools from other sources source = next(
raise AssertionError( (
"Storage sidecar derived pool currently only supports KV-shared " t
f"indices, got {transfer.name} from {transfer.indices_from_pool}." for t in operation.pool_transfers
if t.indices_from_pool is None
and t.name == transfer.indices_from_pool
),
None,
) )
if source is None:
raise AssertionError(
"Storage sidecar derived pool source missing: "
f"{transfer.name} from {transfer.indices_from_pool}."
)
transfer.host_indices = source.host_indices
if transfer.keys is None:
transfer.keys = source.keys
else:
transfer.host_indices = operation.host_indices transfer.host_indices = operation.host_indices
if transfer.keys is None: if transfer.keys is None:
transfer.keys = operation.hash_value transfer.keys = operation.hash_value
@@ -4,7 +4,11 @@ import logging
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Optional from typing import TYPE_CHECKING, Any, Callable, Optional
from sglang.srt.mem_cache.hicache_storage import PoolName, SidecarPoolSpec from sglang.srt.mem_cache.hicache_storage import (
PoolHitPolicy,
PoolName,
SidecarPoolSpec,
)
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
HybridCacheController, HybridCacheController,
) )
@@ -737,7 +741,15 @@ class _DeepSeekV4Strategy(StackStrategy):
enable_storage_metrics=enable_storage_metrics, enable_storage_metrics=enable_storage_metrics,
) )
sidecars = [ sidecars = [
SidecarPoolSpec(pool_name=name, indices_from_pool=src) SidecarPoolSpec(
pool_name=name,
indices_from_pool=src,
hit_policy=(
PoolHitPolicy.TRAILING_PAGES
if src == PoolName.SWA
else PoolHitPolicy.ALL_PAGES
),
)
for name, src in ( for name, src in (
(PoolName.DEEPSEEK_V4_C4, PoolName.KV), (PoolName.DEEPSEEK_V4_C4, PoolName.KV),
(PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV),
@@ -2078,6 +2078,7 @@ class DeepSeekV4PagedHostPool(HostKVCache):
return return
host_rows = self._to_page_indices(host_indices) host_rows = self._to_page_indices(host_indices)
device_rows = self._to_page_indices(device_indices) device_rows = self._to_page_indices(device_indices)
if io_backend == "kernel" and self.layout == "layer_first": if io_backend == "kernel" and self.layout == "layer_first":
transfer_kv_per_layer_mla( transfer_kv_per_layer_mla(
src=self.data_refs[layer_id], src=self.data_refs[layer_id],
@@ -601,6 +601,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
def register_mem_pool_host(self, mem_pool_host: HostKVCache): def register_mem_pool_host(self, mem_pool_host: HostKVCache):
super().register_mem_pool_host(mem_pool_host) super().register_mem_pool_host(mem_pool_host)
if getattr(self.mem_pool_host, "kv_buffer", None) is None:
# Hybrid logical anchors only own allocation indices. Their physical
# tensors are registered through register_mem_host_pool_v2().
return
assert self.mem_pool_host.layout in [ assert self.mem_pool_host.layout in [
"page_first", "page_first",
"page_first_direct", "page_first_direct",
@@ -631,10 +635,16 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
# the corresponding host pool implementation at runtime. # the corresponding host pool implementation at runtime.
self.registered_pools[host_pool_name] = host_pool self.registered_pools[host_pool_name] = host_pool
# Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O. # Non-anchor pools are either sidecar-specific pools with their own
# The storage backend only depends on this accessor, not concrete fields. # accessor, or ordinary KV-like host pools used as SWA side pools.
buf_list = host_pool.get_hybrid_pool_buffer() get_buffers = getattr(
for buf in buf_list: host_pool,
"get_hybrid_pool_buffer",
lambda: [getattr(host_pool, "kv_buffer", None)],
)
for buf in get_buffers():
if buf is None:
continue
super().register_buffer(buf) super().register_buffer(buf)
def _tag_keys(self, keys: List[str]) -> List[str]: def _tag_keys(self, keys: List[str]) -> List[str]:
@@ -645,23 +655,21 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
def _get_hybrid_page_component_keys( def _get_hybrid_page_component_keys(
self, page_keys: List[str], transfer: PoolTransfer self, page_keys: List[str], transfer: PoolTransfer
) -> Tuple[List[str], int]: ) -> Tuple[List[str], int]:
# A logical "page" may map to multiple physical objects in storage. host_pool = getattr(self, "registered_pools", {}).get(transfer.name)
# - INDEXER: one key per page if host_pool is None:
# - MAMBA : one temporal key + N conv keys per page raise ValueError(f"Unregistered Mooncake hybrid pool: {transfer.name}")
# key_multiplier records how many component keys are generated per page.
name = transfer.name # Suffix order must match get_page_buffer_meta() for one page, because
# Mooncake zips object keys with registered buffer pointers.
pool_name = transfer.name
suffixes = [] suffixes = []
if name == PoolName.INDEXER: if pool_name == PoolName.MAMBA:
suffixes = [f"_{self.mla_suffix}_{PoolName.INDEXER}"] # Mamba stores one temporal object plus one object per conv state.
elif name == PoolName.MAMBA: conv_num = len(getattr(host_pool, "conv_buffer", None) or [])
pools = getattr(self, "registered_pools", {}) suffixes = [f"_{self.mha_suffix}_temporal"] + [
mamba_pool = pools.get(PoolName.MAMBA) f"_{self.mha_suffix}_conv_{i}" for i in range(conv_num)
conv_num = len(getattr(mamba_pool, "conv_buffer", None) or [])
base_suffix = f"_{self.mha_suffix}"
suffixes = [f"{base_suffix}_temporal"] + [
f"{base_suffix}_conv_{i}" for i in range(conv_num)
] ]
elif name == PoolName.DRAFT: elif pool_name == PoolName.DRAFT:
# Draft pool's MLA/MHA layout is independent from the target # Draft pool's MLA/MHA layout is independent from the target
# (e.g. EAGLE-MHA draft on top of an MLA target), so pick the # (e.g. EAGLE-MHA draft on top of an MLA target), so pick the
# suffix scheme from the draft pool's own class. The `_draft` # suffix scheme from the draft pool's own class. The `_draft`
@@ -675,6 +683,33 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
f"_{self.mha_suffix}_{PoolName.DRAFT}_k", f"_{self.mha_suffix}_{PoolName.DRAFT}_k",
f"_{self.mha_suffix}_{PoolName.DRAFT}_v", f"_{self.mha_suffix}_{PoolName.DRAFT}_v",
] ]
elif pool_name in (
PoolName.INDEXER,
PoolName.DEEPSEEK_V4_C4,
PoolName.DEEPSEEK_V4_C4_INDEXER,
PoolName.DEEPSEEK_V4_C128,
PoolName.DEEPSEEK_V4_C4_STATE,
PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
PoolName.DEEPSEEK_V4_C128_STATE,
):
# DSA indexer and DeepSeek V4 side pools are page-packed
# single-object pools.
suffixes = [f"_{self.mla_suffix}_{pool_name}"]
elif pool_name == PoolName.SWA:
if not self.is_mla_backend and hasattr(host_pool, "v_buffer"):
# Ordinary MHA SWA mirrors a K/V pool.
suffixes = [
f"_{self.mha_suffix}_{pool_name}_k",
f"_{self.mha_suffix}_{pool_name}_v",
]
elif self.is_mla_backend:
suffixes = [f"_{self.mla_suffix}_{pool_name}"]
if not suffixes:
raise ValueError(
f"Unsupported Mooncake hybrid pool name: {pool_name}, "
f"host_pool={type(host_pool)}"
)
key_multiplier = len(suffixes) key_multiplier = len(suffixes)
component_keys = [ component_keys = [
f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes
@@ -687,6 +722,11 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
pool_transfers: Optional[List[PoolTransfer]] = None, pool_transfers: Optional[List[PoolTransfer]] = None,
extra_info: Optional[HiCacheStorageExtraInfo] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> PoolTransferResult: ) -> PoolTransferResult:
if self.mem_pool_host.kv_buffer is None:
# Logical anchor: no physical KV object exists in Mooncake, so the
# usable prefix is determined entirely by required sidecar objects.
kv_pages = len(keys)
else:
kv_pages = self.batch_exists(keys, extra_info) kv_pages = self.batch_exists(keys, extra_info)
hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {} hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {}
@@ -863,6 +903,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
host_indices: torch.Tensor, host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]: ) -> List[bool]:
if self.mem_pool_host.kv_buffer is None:
# DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
return [True] * len(keys)
# Apply extra_backend_tag prefix if available # Apply extra_backend_tag prefix if available
keys = self._tag_keys(keys) keys = self._tag_keys(keys)
@@ -888,6 +932,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
host_indices: torch.Tensor, host_indices: torch.Tensor,
extra_info: Optional[HiCacheStorageExtraInfo] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]: ) -> List[bool]:
if self.mem_pool_host.kv_buffer is None:
# DeepSeek V4's KV anchor is logical only; v2 side pools carry data.
return [True] * len(keys)
# Apply extra_backend_tag prefix if available # Apply extra_backend_tag prefix if available
keys = self._tag_keys(keys) keys = self._tag_keys(keys)
@@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchResult, MatchResult,
) )
from sglang.srt.mem_cache.hicache_storage import ( from sglang.srt.mem_cache.hicache_storage import (
PoolHitPolicy,
PoolName, PoolName,
PoolTransfer, PoolTransfer,
PoolTransferResult, PoolTransferResult,
@@ -415,7 +416,9 @@ class SWAComponent(TreeComponent):
root = self.cache.root_node root = self.cache.root_node
sliding_window_size = self.sliding_window_size sliding_window_size = self.sliding_window_size
swa_lock_size = 0 swa_lock_size = 0
swa_uuid_for_lock = None swa_uuid = None
uuid_key = "host_uuid" if lock_host else "uuid"
lru = self.cache.host_lru_lists[ct] if lock_host else self.cache.lru_lists[ct]
# Tombstoned nodes (cd.value is None) have no SWA chunk to protect # Tombstoned nodes (cd.value is None) have no SWA chunk to protect
# skip them and keep walking up. This path is hit when HiCache # skip them and keep walking up. This path is hit when HiCache
@@ -423,23 +426,36 @@ class SWAComponent(TreeComponent):
cur = node cur = node
while cur != root and swa_lock_size < sliding_window_size: while cur != root and swa_lock_size < sliding_window_size:
comp = cur.component_data[ct] comp = cur.component_data[ct]
if comp.value is None: value = comp.host_value if lock_host else comp.value
if value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id) result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
cur = cur.parent cur = cur.parent
continue continue
if comp.lock_ref == 0:
ref = comp.host_lock_ref if lock_host else comp.lock_ref
if ref == 0:
if lock_host:
if lru.in_list(cur):
lru.remove_node(cur)
else:
key_len = len(cur.key) key_len = len(cur.key)
self.cache.component_evictable_size_[ct] -= key_len self.cache.component_evictable_size_[ct] -= key_len
self.cache.component_protected_size_[ct] += key_len self.cache.component_protected_size_[ct] += key_len
comp.lock_ref += 1 if lock_host:
swa_lock_size += len(cur.key) comp.host_lock_ref = ref + 1
else:
comp.lock_ref = ref + 1
swa_lock_size += len(value)
if swa_lock_size >= sliding_window_size: if swa_lock_size >= sliding_window_size:
if comp.metadata.get("uuid") is None: if comp.metadata.get(uuid_key) is None:
comp.metadata["uuid"] = next_component_uuid() comp.metadata[uuid_key] = next_component_uuid()
swa_uuid_for_lock = comp.metadata["uuid"] swa_uuid = comp.metadata[uuid_key]
cur = cur.parent cur = cur.parent
result.swa_uuid_for_lock = swa_uuid_for_lock if lock_host:
result.swa_uuid_for_host_lock = swa_uuid
else:
result.swa_uuid_for_lock = swa_uuid
return result return result
def release_component_lock( def release_component_lock(
@@ -450,9 +466,14 @@ class SWAComponent(TreeComponent):
) -> None: ) -> None:
ct = self.component_type ct = self.component_type
root = self.cache.root_node root = self.cache.root_node
swa_uuid_for_lock = params.swa_uuid_for_lock if params else None swa_uuid_for_lock = (
(params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock)
if params
else None
)
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
dec_swa = True dec_swa = True
uuid_key = "host_uuid" if lock_host else "uuid"
# A node in skip_lock_node_ids was a tombstone when this lock was acquired. # A node in skip_lock_node_ids was a tombstone when this lock was acquired.
cur = node cur = node
@@ -461,15 +482,25 @@ class SWAComponent(TreeComponent):
if cur.id in skip_lock_node_ids: if cur.id in skip_lock_node_ids:
cur = cur.parent cur = cur.parent
continue continue
if comp.lock_ref == 0: ref = comp.host_lock_ref if lock_host else comp.lock_ref
if ref == 0:
cur = cur.parent cur = cur.parent
continue continue
if comp.lock_ref == 1: if ref == 1:
key_len = len(cur.key) if lock_host:
if comp.value is None and comp.host_value is not None:
host_lru = self.cache.host_lru_lists[ct]
if not host_lru.in_list(cur):
host_lru.insert_mru(cur)
else:
key_len = len(comp.value)
self.cache.component_evictable_size_[ct] += key_len self.cache.component_evictable_size_[ct] += key_len
self.cache.component_protected_size_[ct] -= key_len self.cache.component_protected_size_[ct] -= key_len
comp.lock_ref -= 1 if lock_host:
if swa_uuid_for_lock and comp.metadata.get("uuid") == swa_uuid_for_lock: comp.host_lock_ref = ref - 1
else:
comp.lock_ref = ref - 1
if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock:
dec_swa = False dec_swa = False
cur = cur.parent cur = cur.parent
@@ -546,6 +577,46 @@ class SWAComponent(TreeComponent):
) )
] ]
if phase == CacheTransferPhase.BACKUP_STORAGE:
cd = node.component_data[ct]
if cd.host_value is None or not node.hash_value:
return None
num_pages = len(cd.host_value) // self.cache.page_size
if num_pages == 0:
return None
return [
PoolTransfer(
name=PoolName.SWA,
host_indices=cd.host_value[-num_pages * self.cache.page_size :],
keys=node.hash_value[-num_pages:],
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
]
if phase == CacheTransferPhase.PREFETCH:
num_pages = min(
prefetch_tokens // self.cache.page_size,
(self.sliding_window_size + self.cache.page_size - 1)
// self.cache.page_size,
)
if num_pages == 0:
return None
num_tokens = num_pages * self.cache.page_size
host_indices = self._swa_kv_pool_host.alloc(num_tokens)
if host_indices is None:
self.cache.evict_host(num_tokens, ComponentType.SWA)
host_indices = self._swa_kv_pool_host.alloc(num_tokens)
if host_indices is None:
return []
return [
PoolTransfer(
name=PoolName.SWA,
host_indices=host_indices,
keys=["__placeholder__"] * num_pages,
hit_policy=PoolHitPolicy.TRAILING_PAGES,
)
]
return None return None
def commit_hicache_transfer( def commit_hicache_transfer(
@@ -586,6 +657,98 @@ class SWAComponent(TreeComponent):
assert offset == len(xfer.host_indices) assert offset == len(xfer.host_indices)
return return
if phase == CacheTransferPhase.PREFETCH:
self._commit_prefetch(
node,
transfers,
insert_result=insert_result,
pool_storage_result=pool_storage_result,
)
return
def _release_swa_host(self, host_indices: torch.Tensor) -> None:
if host_indices is not None and host_indices.numel() > 0:
self.cache.cache_controller.append_host_mem_release(
extra_pools=[PoolTransfer(name=PoolName.SWA, host_indices=host_indices)]
)
def _attach_swa_host_value(
self, node: "UnifiedTreeNode", host_indices: torch.Tensor
) -> None:
"""Write host_indices into node's SWA host_value and refresh tree state."""
ct = self.component_type
cd = node.component_data[ct]
cd.host_value = host_indices.clone()
host_lru = self.cache.host_lru_lists[ct]
if cd.value is None and not host_lru.in_list(node):
host_lru.insert_mru(node)
self.cache._update_evictable_leaf_sets(node)
if node.parent:
self.cache._update_evictable_leaf_sets(node.parent)
def _commit_prefetch(
self,
anchor,
transfers: list[PoolTransfer],
*,
insert_result: Optional[InsertResult] = None,
pool_storage_result: Optional[PoolTransferResult] = None,
) -> None:
"""Distribute the prefetched SWA buffer onto the leaf→anchor path.
The buffer holds the trailing ``loaded_pages`` of the completed KV
prefix, mapped to token range ``[loaded_start, total_len)``. We walk
upward from ``inserted_host_node`` to ``anchor`` and, for each node
whose token range overlaps the buffer:
- SWA tombstone (host_value is None) → fill from buffer (split if
the node only partially overlaps at the buffer's left edge)
- already has SWA host_value → release the corresponding slice
Any leftover buffer beyond the walked range is also released.
"""
if not transfers:
return
ct = self.component_type
host_indices = transfers[0].host_indices
loaded_pages = (
pool_storage_result.extra_pool_hit_pages.get(PoolName.SWA, 0)
if pool_storage_result
else 0
)
target = insert_result.inserted_host_node if insert_result else None
if not loaded_pages or target is None:
self._release_swa_host(host_indices)
return
# Buffer covers token range [loaded_start, total_len).
loaded_start = insert_result.total_len - loaded_pages * self.cache.page_size
# Walk leaf → anchor; ``pos`` is the right edge of ``cur`` in tokens.
pos, cur = insert_result.total_len, target
while cur is not anchor and pos > loaded_start:
node_start = pos - len(cur.key)
# Intersection of cur's range and the buffer.
fill_start = max(node_start, loaded_start)
fill_len = pos - fill_start
buf_off = fill_start - loaded_start
slice_ = host_indices[buf_off : buf_off + fill_len]
cd = cur.component_data[ct]
if cd.host_value is None and fill_len > 0:
# Tombstone: split off the in-buffer tail if needed, then fill.
if fill_start > node_start:
self.cache._split_node(cur.key, cur, fill_start - node_start)
self._attach_swa_host_value(cur, slice_)
else:
# Already has SWA (or empty overlap): drop this slice.
self._release_swa_host(slice_)
pos = node_start
cur = cur.parent
# Buffer prefix that fell outside the anchor→leaf path.
if pos > loaded_start:
self._release_swa_host(host_indices[: pos - loaded_start])
def drive_host_eviction( def drive_host_eviction(
self, num_tokens: int, tracker: dict[ComponentType, int] self, num_tokens: int, tracker: dict[ComponentType, int]
) -> None: ) -> None:
@@ -1121,9 +1121,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
if len(key): if len(key):
child_key = key.child_key(self.page_size) child_key = key.child_key(self.page_size)
result = InsertResult( result = InsertResult(prefix_len=matched_length, total_len=total_len)
prefix_len=matched_length,
)
if len(key) == 0: if len(key) == 0:
if ( if (
node is not self.root_node node is not self.root_node
@@ -1662,6 +1660,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
transfers.append( transfers.append(
PoolTransfer( PoolTransfer(
name=spec.pool_name, name=spec.pool_name,
keys=indices_source.keys,
hit_policy=spec.hit_policy, hit_policy=spec.hit_policy,
indices_from_pool=spec.indices_from_pool, indices_from_pool=spec.indices_from_pool,
) )
@@ -1850,6 +1849,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
) )
else: else:
return True return True
if (
completed
and getattr(operation, "pool_transfers", None)
and not getattr(operation, "pool_transfers_done", True)
):
can_terminate = False
operation_terminated = operation.is_terminated() operation_terminated = operation.is_terminated()
states = torch.tensor( states = torch.tensor(
@@ -1,5 +1,9 @@
"""Unit tests for UnifiedRadixCache""" """Unit tests for UnifiedRadixCache"""
import json
import shutil
import tempfile
import time
import unittest import unittest
from array import array from array import array
from dataclasses import dataclass, replace from dataclasses import dataclass, replace
@@ -1840,6 +1844,146 @@ class UnifiedRadixCacheSuite:
# HiCache Unit Tests (real cache_controller D<->H backup/load) # HiCache Unit Tests (real cache_controller D<->H backup/load)
# ================================================================ # ================================================================
# ---------- L3 storage (file backend) helpers ----------
def _path_chain(self, tree, node):
"""Return root->node node chain (excluding root)."""
chain = []
cur = node
while cur is not tree.root_node:
chain.append(cur)
cur = cur.parent
chain.reverse()
return chain
def _write_path_to_l3(self, tree, node):
"""Offload every node on root->node path from host to L3 storage."""
for n in self._path_chain(tree, node):
tree.write_backup_storage(n)
def _flush_l3_backups(self, tree, timeout: float = 10.0):
"""Wait for backup threads to finish, then drain acks (release locks)."""
deadline = time.time() + timeout
while tree.ongoing_backup and time.time() < deadline:
tree.drain_storage_control_queues()
if tree.ongoing_backup:
time.sleep(0.01)
tree.drain_storage_control_queues()
self.assertFalse(tree.ongoing_backup, "L3 backups did not complete in time")
def _run_prefetch_to_completion(self, tree, req_id, timeout: float = 10.0):
deadline = time.time() + timeout
while time.time() < deadline:
if tree.check_prefetch_progress(req_id):
return
time.sleep(0.01)
self.fail(f"prefetch {req_id} did not complete in time")
def _all_page_hashes(self, tree, node):
hashes = []
for n in self._path_chain(tree, node):
hashes.extend(list(n.hash_value))
return hashes
def test_hicache_l3_write_storage(self):
"""D->H->L3 offload: every KV page lands in the file storage backend."""
if self._skip_unsupported_hicache_test():
return
if self.cfg.has_mamba:
self.skipTest("mamba L3 offload is out of scope for this unit fixture")
storage_dir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
self._init_hicache(
tree,
storage_backend="file",
storage_dir=storage_dir,
prefetch_threshold=1,
)
seq = self._make_seq(1, 4)
self._insert(tree, allocator, req_to_token_pool, seq)
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
leaf = m.last_device_node
# D->H first, then H->L3.
self._backup_node(tree, leaf)
self.assertTrue(leaf.hash_value)
self._write_path_to_l3(tree, leaf)
self._flush_l3_backups(tree)
# Every KV page hash on the path must now exist in storage.
backend = tree.cache_controller.storage_backend
page_hashes = self._all_page_hashes(tree, leaf)
self.assertEqual(len(page_hashes), len(seq) // self.cfg.page_size)
self.assertEqual(backend.batch_exists(page_hashes), len(page_hashes))
tree.sanity_check()
def test_hicache_l3_prefetch(self):
"""L3 round trip: write with one tree, prefetch into a fresh tree.
Uses two independent trees that share the same file storage dir so the
prefetch path genuinely reloads from L3 (no host/device residue).
"""
if self._skip_unsupported_hicache_test():
return
if self.cfg.has_mamba:
self.skipTest("mamba L3 prefetch is out of scope for this unit fixture")
storage_dir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
seq = self._make_seq(1, 4)
# --- Producer tree: fill KV, backup D->H, offload H->L3. ---
prod, prod_alloc, prod_rtp = build_fixture(self.cfg)
self._init_hicache(
prod,
storage_backend="file",
storage_dir=storage_dir,
prefetch_threshold=1,
)
self._insert(prod, prod_alloc, prod_rtp, seq)
mp = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
prod_leaf = mp.last_device_node
self._fill_full_kv(prod_alloc, mp.device_indices, marker=7)
expected_k, expected_v = self._snapshot_full_kv(prod_alloc, mp.device_indices)
self._backup_node(prod, prod_leaf)
self._write_path_to_l3(prod, prod_leaf)
self._flush_l3_backups(prod)
# --- Consumer tree: prefetch the same tokens straight from L3. ---
cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
self._init_hicache(
cons,
storage_backend="file",
storage_dir=storage_dir,
prefetch_threshold=1,
)
req_id = "l3-prefetch-req"
cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None)
self._run_prefetch_to_completion(cons, req_id)
cons.drain_storage_control_queues()
# The full prefix must now be a host hit (loaded from L3).
mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(mc.host_hit_length, len(seq))
host_node = mc.last_host_node
self.assertIsNot(host_node, cons.root_node)
self.assertTrue(host_node.evicted)
# Load the reloaded host prefix back to device and verify KV bytes.
self._load_back_node(cons, host_node)
loaded_indices = cons.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq)))
).device_indices
self.assertEqual(len(loaded_indices), len(seq))
loaded_k, loaded_v = self._snapshot_full_kv(cons_alloc, loaded_indices)
self.assertTrue(torch.equal(loaded_k, expected_k))
self.assertTrue(torch.equal(loaded_v, expected_v))
cons.sanity_check()
def _skip_unsupported_hicache_test(self): def _skip_unsupported_hicache_test(self):
if self.cfg.has_swa and self.cfg.has_mamba: if self.cfg.has_swa and self.cfg.has_mamba:
self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks") self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")
@@ -1869,7 +2013,16 @@ class UnifiedRadixCacheSuite:
self._simulate_backup(tree, node) self._simulate_backup(tree, node)
stack.extend(node.children.values()) stack.extend(node.children.values())
def _init_hicache(self, tree, *, write_policy: str = "write_through"): def _init_hicache(
self,
tree,
*,
write_policy: str = "write_through",
storage_backend: Optional[str] = None,
storage_dir: Optional[str] = None,
prefetch_threshold: Optional[int] = None,
prefetch_policy: str = "wait_complete",
):
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
orig_kv_host_pool = assembler.MHATokenToKVPoolHost orig_kv_host_pool = assembler.MHATokenToKVPoolHost
@@ -1899,11 +2052,42 @@ class UnifiedRadixCacheSuite:
patcher.start() patcher.start()
self.addCleanup(patcher.stop) self.addCleanup(patcher.stop)
storage_extra_config = None
if storage_backend == "file":
import sglang.srt.managers.cache_controller as cache_controller
# The file-backend storage config records TP rank/size. These unit
# fixtures run without initializing distributed parallel state, so
# provide the local single-rank values that the fixture represents.
tp_rank_patcher = mock.patch.object(
cache_controller, "get_tensor_model_parallel_rank", return_value=0
)
tp_size_patcher = mock.patch.object(
cache_controller, "get_tensor_model_parallel_world_size", return_value=1
)
tp_rank_patcher.start()
tp_size_patcher.start()
self.addCleanup(tp_rank_patcher.stop)
self.addCleanup(tp_size_patcher.stop)
assert storage_dir is not None, "file backend needs a storage_dir"
# HiCacheFile reads the directory from this env var.
cm = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.override(storage_dir)
cm.__enter__()
self.addCleanup(cm.__exit__, None, None, None)
extra = {}
if prefetch_threshold is not None:
extra["prefetch_threshold"] = prefetch_threshold
storage_extra_config = json.dumps(extra) if extra else None
server_args = ServerArgs( server_args = ServerArgs(
model_path="dummy", model_path="dummy",
page_size=self.cfg.page_size, page_size=self.cfg.page_size,
hicache_io_backend="direct", hicache_io_backend="direct",
hicache_write_policy=write_policy, hicache_write_policy=write_policy,
hicache_storage_backend=storage_backend,
hicache_storage_backend_extra_config=storage_extra_config,
hicache_storage_prefetch_policy=prefetch_policy,
) )
# See build_fixture for why _mamba_cache_chunk_size is preset. # See build_fixture for why _mamba_cache_chunk_size is preset.
server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size)
@@ -1911,6 +2095,17 @@ class UnifiedRadixCacheSuite:
tree.init_hicache(server_args, tree.cache_init_params) tree.init_hicache(server_args, tree.cache_init_params)
tree.write_through_threshold = 1 << 30 tree.write_through_threshold = 1 << 30
tree.load_back_threshold = 0 tree.load_back_threshold = 0
if storage_backend is not None:
# Unit fixtures size host/device pools equally, which makes the
# production prefetch capacity limit (host - device) zero. Keep the
# L3 tests focused on storage round trips by allowing one fixture
# worth of prefetch tokens.
tree.cache_controller.prefetch_capacity_limit = max(
tree.cache_controller.prefetch_capacity_limit,
tree.cache_controller.mem_pool_host.size,
)
# Background prefetch/backup threads are daemon; stop them per-test.
self.addCleanup(tree.cache_controller._stop_storage_threads)
def _build_hicache_fixture(self): def _build_hicache_fixture(self):
fixture = build_fixture(self.cfg) fixture = build_fixture(self.cfg)