From d9fa84b25b79203b08be54ef6f5e4ccebef8cd5f Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Fri, 15 May 2026 11:20:52 +0800 Subject: [PATCH] [UnifiedTree]: Support HiCache For DeepSeek_V4 (#24691) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: ispobock Co-authored-by: 晟海 --- .../srt/mem_cache/deepseek_v4_memory_pool.py | 13 +- .../sglang/srt/mem_cache/hicache_storage.py | 26 +- python/sglang/srt/mem_cache/hiradix_cache.py | 1 + .../hybrid_cache/hybrid_cache_controller.py | 81 ++- .../hybrid_cache/hybrid_pool_assembler.py | 373 +++++++++++-- .../sglang/srt/mem_cache/memory_pool_host.py | 493 +++++++++++++++++- .../srt/mem_cache/unified_radix_cache.py | 90 +++- python/sglang/test/kl_multiturn_utils.py | 67 ++- python/sglang/test/kl_test_utils.py | 12 +- .../test_unified_radix_cache_kl.py | 57 +- .../test_unified_radix_hicache_kl.py | 162 ++++++ 11 files changed, 1221 insertions(+), 154 deletions(-) create mode 100644 test/registered/radix_cache/test_unified_radix_hicache_kl.py diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index bc8aa8db1..94b930b89 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -605,7 +605,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): else: raise ValueError(f"Unsupported compression ratio: {ratio}") + def wait_layer_transfer(self, layer_id: int) -> None: + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + def get_attention_compress_states(self, layer_id: int) -> CompressStatePool: + self.wait_layer_transfer(layer_id) compress_state_pool = self.compress_state_pools[layer_id] assert ( compress_state_pool is not None @@ -613,6 +618,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return compress_state_pool def get_indexer_compress_states(self, layer_id: int) -> CompressStatePool: + self.wait_layer_transfer(layer_id) indexer_compress_state_pool = self.indexer_compress_state_pools[layer_id] assert ( indexer_compress_state_pool is not None @@ -620,6 +626,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return indexer_compress_state_pool def get_swa_key_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) return self.swa_kv_pool.get_key_buffer(layer_id) def set_swa_key_buffer( @@ -635,7 +642,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): assert compress_kv_pool is not None return compress_kv_pool.page_size - def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor: + def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor | None: + self.wait_layer_transfer(layer_id) _, compress_layer_id, compress_kv_pool = self.layer_mapping[layer_id] assert compress_kv_pool is not None return compress_kv_pool.get_key_buffer(compress_layer_id) @@ -656,6 +664,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return self.c4_indexer_kv_pool.page_size def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }" return self.c4_indexer_kv_pool.get_index_k_with_scale_buffer(compress_layer_id) @@ -666,6 +675,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): seq_len: int, page_indices: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: + self.wait_layer_transfer(layer_id) compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }" return self.c4_indexer_kv_pool.get_index_k_scale_buffer( @@ -709,6 +719,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): ) def get_swa_key_buffer_radix(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) return self.swa_kv_pool.get_key_buffer(layer_id) def set_swa_key_buffer_radix_fused( diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index f301a2c94..63ed143d7 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -1,14 +1,18 @@ +from __future__ import annotations + import logging import os from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum -from typing import Any, List, Optional, Set +from typing import TYPE_CHECKING, Any, List, Optional, Set import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.memory_pool_host import HostKVCache + +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool_host import HostKVCache logger = logging.getLogger(__name__) @@ -52,6 +56,14 @@ class PoolName(str, Enum): MAMBA = "mamba" SWA = "swa" INDEXER = "indexer" + # TODO(hzh0425): Current DeepSeek V4 pool naming is verbose; will be normalized to + # 'COMPRESSED_KV / COMPRESSED_INDEXER / COMPRESSED_STATE' in the next PR. + DEEPSEEK_V4_C4 = "deepseek_v4_c4" + DEEPSEEK_V4_C4_INDEXER = "deepseek_v4_c4_indexer" + DEEPSEEK_V4_C128 = "deepseek_v4_c128" + DEEPSEEK_V4_C4_STATE = "deepseek_v4_c4_state" + DEEPSEEK_V4_C4_INDEXER_STATE = "deepseek_v4_c4_indexer_state" + DEEPSEEK_V4_C128_STATE = "deepseek_v4_c128_state" def __str__(self) -> str: return self.value @@ -83,6 +95,16 @@ class PoolTransfer: keys: Optional[List[str]] = None hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES nodes_to_load: Optional[List[Any]] = None + indices_from_pool: Optional[PoolName] = None + + +@dataclass(frozen=True) +class SidecarPoolSpec: + """Pool whose transfer indices are reused from one real source pool.""" + + pool_name: PoolName + indices_from_pool: PoolName + hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES @dataclass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 56e1d52ba..3416d2a86 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -647,6 +647,7 @@ class HiRadixCache(RadixCache): pool = PoolTransfer( name=PoolName.INDEXER, hit_policy=PoolHitPolicy.ALL_PAGES, + indices_from_pool=PoolName.KV, ) return {"extra_pools": [pool]} else: diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index b4fdde959..a339eef5d 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -23,6 +23,7 @@ from sglang.srt.managers.cache_controller import ( from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageExtraInfo, PoolHitPolicy, + PoolName, PoolTransfer, PoolTransferResult, ) @@ -50,12 +51,12 @@ class CacheOperation(BaseCacheOperation): @staticmethod def merge_pool_transfers( - ops: List["CacheOperation"], + ops: List[CacheOperation], ) -> Optional[list[PoolTransfer]]: - grouped: dict[str, list[PoolTransfer]] = {} + grouped: dict[tuple[PoolName, Optional[PoolName]], list[PoolTransfer]] = {} for op in ops: for t in op.pool_transfers or []: - grouped.setdefault(t.name, []).append(t) + grouped.setdefault((t.name, t.indices_from_pool), []).append(t) if not grouped: return None @@ -65,16 +66,18 @@ class CacheOperation(BaseCacheOperation): return [ PoolTransfer( - name=name, + name=ts[0].name, host_indices=cat_or_none(t.host_indices for t in ts), device_indices=cat_or_none(t.device_indices for t in ts), keys=[k for t in ts if t.keys for k in t.keys] or None, + hit_policy=ts[0].hit_policy, + indices_from_pool=ts[0].indices_from_pool, ) - for name, ts in grouped.items() + for ts in grouped.values() ] @staticmethod - def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation": + def merge_ops(ops: List[CacheOperation]) -> CacheOperation: if len(ops) == 1: return ops[0] host_indices = torch.cat([op.host_indices for op in ops]) @@ -478,6 +481,7 @@ class HybridCacheController(BaseHiCacheController): device_indices=transfer_device_indices, keys=transfer.keys, hit_policy=transfer.hit_policy, + indices_from_pool=transfer.indices_from_pool, ) ) return host_indices, device_indices, resolved_pool_transfers @@ -485,7 +489,7 @@ class HybridCacheController(BaseHiCacheController): def _page_transfer(self, operation): # Transfer extra pools if operation.pool_transfers and not operation.is_terminated(): - self._resolve_shared_pool_transfers(operation) + self._resolve_sidecar_derived_pool_transfers(operation) results = self.storage_backend.batch_get_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) @@ -495,19 +499,26 @@ class HybridCacheController(BaseHiCacheController): def _page_backup(self, operation): # Backup extra pools if operation.pool_transfers: - self._resolve_shared_pool_transfers(operation) + self._resolve_sidecar_derived_pool_transfers(operation) results = self.storage_backend.batch_set_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) # Backup kv pools super()._page_backup(operation) - def _resolve_shared_pool_transfers(self, operation): + def _resolve_sidecar_derived_pool_transfers(self, operation): for transfer in operation.pool_transfers: - entry = self.mem_pool_host.entry_map.get(transfer.name) - if entry.share_indices_with_anchor: + if transfer.indices_from_pool is None: + continue + if transfer.indices_from_pool != PoolName.KV: + # TODO(hzh): Support storage sidecar derived pools from other sources + raise AssertionError( + "Storage sidecar derived pool currently only supports KV-shared " + f"indices, got {transfer.name} from {transfer.indices_from_pool}." + ) + transfer.host_indices = operation.host_indices + if transfer.keys is None: transfer.keys = operation.hash_value - transfer.host_indices = operation.host_indices def _sync_trailing_keys( self, @@ -541,14 +552,23 @@ class HybridCacheController(BaseHiCacheController): return None # (pool, free_fn, indices) for atomic rollback on failure. newly_allocated: list[tuple[PoolTransfer, Callable, torch.Tensor]] = [] + derived_transfers: list[PoolTransfer] = [] + + def rollback_allocated() -> None: + for prev_pool, prev_free_fn, prev_indices in newly_allocated: + prev_free_fn(prev_indices) + if alloc_host: + prev_pool.host_indices = None + else: + prev_pool.device_indices = None + for pool in extra_pools: + if pool.indices_from_pool is not None: + derived_transfers.append(pool) + continue entry = self.mem_pool_host.entry_map.get(pool.name) if entry is None: continue - if entry.share_indices_with_anchor: - pool.device_indices = kv_device_indices - pool.host_indices = kv_host_indices - continue if alloc_host: if pool.host_indices is not None or pool.device_indices is None: continue @@ -572,16 +592,33 @@ class HybridCacheController(BaseHiCacheController): indices = alloc_fn(size) if indices is None: # Atomic rollback: free everything we successfully allocated. - for prev_pool, prev_free_fn, prev_indices in newly_allocated: - prev_free_fn(prev_indices) - if alloc_host: - prev_pool.host_indices = None - else: - prev_pool.device_indices = None + rollback_allocated() return None if alloc_host: pool.host_indices = indices else: pool.device_indices = indices newly_allocated.append((pool, free_fn, indices)) + + # Assign indices to deferred pools from their source. + for pool in derived_transfers: + if pool.indices_from_pool == PoolName.KV: + pool.host_indices = kv_host_indices + pool.device_indices = kv_device_indices + continue + + source = next( + ( + transfer + for transfer in extra_pools + if transfer.indices_from_pool is None + and transfer.name == pool.indices_from_pool + ), + None, + ) + if source is None: + rollback_allocated() + return None + pool.host_indices = source.host_indices + pool.device_indices = source.device_indices return extra_pools diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 0a8b10e42..3de18edde 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -3,12 +3,15 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING, Any, Callable, Optional -from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName +from sglang.srt.mem_cache.hicache_storage import PoolName, SidecarPoolSpec from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) from sglang.srt.mem_cache.memory_pool_host import ( + DeepSeekV4PagedHostPool, + DeepSeekV4StateHostPool, HostPoolGroup, + LogicalHostPool, MambaPoolHost, MHATokenToKVPoolHost, MLATokenToKVPoolHost, @@ -71,7 +74,6 @@ def build_pool_entry( layer_mapping: dict[int, int], transfer_layer_num: int, is_anchor: bool = False, - share_indices_with_anchor: bool = False, host_evict_fn: Optional[Callable[[int], Any]] = None, device_evict_fn: Optional[Callable[[int], Any]] = None, device_alloc_fn: Optional[Callable[[int], Any]] = None, @@ -83,7 +85,6 @@ def build_pool_entry( device_pool=device_pool, layer_mapper=_make_layer_mapper(layer_mapping, transfer_layer_num), is_primary_index_anchor=is_anchor, - share_indices_with_anchor=share_indices_with_anchor, host_evict_fn=host_evict_fn, device_evict_fn=device_evict_fn, device_alloc_fn=device_alloc_fn, @@ -100,8 +101,8 @@ def build_kv_only_stack( page_size: int, tp_group, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, storage_backend: Optional[str], use_mla: bool, override_kv_cache_dim: Optional[int] = None, @@ -164,8 +165,8 @@ def build_hybrid_swa_stack( page_size: int, tp_group, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, storage_backend: Optional[str], use_mla: bool, host_swa_evict_fn: Optional[Callable[[int], Any]] = None, @@ -237,6 +238,249 @@ def build_hybrid_swa_stack( return host_pool_group, cache_controller +def _deepseek_v4_num_host_pages( + *, + params: CacheInitParams, + server_args: ServerArgs, + kvcache: Any, + page_size: int, + swa_page_size: int, +) -> tuple[int, int]: + allocator = params.token_to_kv_pool_allocator + device_full_size = getattr(allocator, "size_full", kvcache.size) + device_full_pages = (device_full_size + page_size - 1) // page_size + + device_swa_pages = (kvcache.swa_size + swa_page_size - 1) // swa_page_size + + if server_args.hicache_size > 0: + raise ValueError( + "DeepSeek V4 HiCache currently does not support --hicache-size; " + "use --hicache-ratio instead." + ) + ratio = server_args.hicache_ratio + full_host_pages = max(int(device_full_pages * ratio), device_full_pages + 1) + swa_host_pages = max(int(device_swa_pages * ratio), device_swa_pages + 1) + return full_host_pages, swa_host_pages + + +def build_deepseek_v4_hicache_stack( + *, + params: CacheInitParams, + server_args: ServerArgs, + kvcache: Any, + page_size: int, + tp_group, + load_cache_event, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, + storage_backend: Optional[str], + host_swa_evict_fn: Optional[Callable[[int], Any]] = None, + device_swa_evict_fn: Optional[Callable[[int], Any]] = None, + prefetch_threshold: int = 256, + model_name: Optional[str] = None, + storage_backend_extra_config: Optional[dict] = None, + pp_rank: int = 0, + pp_size: int = 1, + enable_storage_metrics: bool = False, +) -> tuple[HostPoolGroup, HybridCacheController]: + transfer_layer_num = len(kvcache.compression_ratios) + full_layer_mapping = {layer_id: layer_id for layer_id in range(transfer_layer_num)} + swa_layer_mapping = { + layer_id: layer_id for layer_id in range(len(kvcache.swa_kv_pool.kv_buffer)) + } + + c4_layer_mapping = {} + c128_layer_mapping = {} + c4_state_global_layers = [] + c128_state_global_layers = [] + for layer_id, layer_item in enumerate(kvcache.layer_mapping): + if layer_item.compress_ratio == 4: + c4_layer_mapping[layer_id] = layer_item.compress_layer_id + c4_state_global_layers.append(layer_id) + elif layer_item.compress_ratio == 128: + c128_layer_mapping[layer_id] = layer_item.compress_layer_id + c128_state_global_layers.append(layer_id) + + c4_state_mapping = { + layer_id: local_id for local_id, layer_id in enumerate(c4_state_global_layers) + } + c128_state_mapping = { + layer_id: local_id for local_id, layer_id in enumerate(c128_state_global_layers) + } + num_host_pages, swa_num_host_pages = _deepseek_v4_num_host_pages( + params=params, + server_args=server_args, + kvcache=kvcache, + page_size=page_size, + swa_page_size=kvcache.swa_page_size, + ) + + logical_host_pool = LogicalHostPool(num_host_pages * page_size, page_size) + swa_host_pool = DeepSeekV4PagedHostPool( + pool_name=str(PoolName.SWA), + device_buffers=kvcache.swa_kv_pool.kv_buffer, + item_bytes=kvcache.swa_kv_pool.bytes_per_page_padded, + num_host_pages=swa_num_host_pages, + slot_page_size=kvcache.swa_page_size, + allocator_type=server_args.hicache_storage_backend, + ) + swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator + entries = [ + build_pool_entry( + name=PoolName.KV, + host_pool=logical_host_pool, + device_pool=kvcache, + layer_mapping=full_layer_mapping, + transfer_layer_num=transfer_layer_num, + is_anchor=True, + ), + build_pool_entry( + name=PoolName.SWA, + host_pool=swa_host_pool, + device_pool=kvcache.swa_kv_pool, + layer_mapping=swa_layer_mapping, + transfer_layer_num=transfer_layer_num, + host_evict_fn=host_swa_evict_fn, + device_evict_fn=device_swa_evict_fn, + device_alloc_fn=swa_attn_allocator.alloc, + device_free_fn=swa_attn_allocator.free, + ), + ] + + if c4_layer_mapping: + c4_host_pool = DeepSeekV4PagedHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C4), + device_buffers=kvcache.c4_kv_pool.kv_buffer, + item_bytes=kvcache.c4_kv_pool.bytes_per_page_padded, + num_host_pages=num_host_pages, + slot_page_size=page_size, + allocator_type=server_args.hicache_storage_backend, + ) + c4_indexer_host_pool = DeepSeekV4PagedHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER), + device_buffers=kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer, + item_bytes=( + kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer[0].shape[1] + * kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer[0].element_size() + ), + num_host_pages=num_host_pages, + slot_page_size=page_size, + allocator_type=server_args.hicache_storage_backend, + ) + c4_state_host_pool = DeepSeekV4StateHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C4_STATE), + state_pools=[ + kvcache.compress_state_pools[layer_id] + for layer_id in c4_state_global_layers + ], + num_host_pages=swa_num_host_pages, + swa_page_size=kvcache.swa_page_size, + allocator_type=server_args.hicache_storage_backend, + ) + c4_indexer_state_host_pool = DeepSeekV4StateHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER_STATE), + state_pools=[ + kvcache.indexer_compress_state_pools[layer_id] + for layer_id in c4_state_global_layers + ], + num_host_pages=swa_num_host_pages, + swa_page_size=kvcache.swa_page_size, + allocator_type=server_args.hicache_storage_backend, + ) + entries.extend( + [ + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C4, + host_pool=c4_host_pool, + device_pool=kvcache.c4_kv_pool, + layer_mapping=c4_layer_mapping, + transfer_layer_num=transfer_layer_num, + ), + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C4_INDEXER, + host_pool=c4_indexer_host_pool, + device_pool=kvcache.c4_indexer_kv_pool, + layer_mapping=c4_layer_mapping, + transfer_layer_num=transfer_layer_num, + ), + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C4_STATE, + host_pool=c4_state_host_pool, + device_pool=None, + layer_mapping=c4_state_mapping, + transfer_layer_num=transfer_layer_num, + ), + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, + host_pool=c4_indexer_state_host_pool, + device_pool=None, + layer_mapping=c4_state_mapping, + transfer_layer_num=transfer_layer_num, + ), + ] + ) + + if c128_layer_mapping: + c128_host_pool = DeepSeekV4PagedHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C128), + device_buffers=kvcache.c128_kv_pool.kv_buffer, + item_bytes=kvcache.c128_kv_pool.bytes_per_page_padded, + num_host_pages=num_host_pages, + slot_page_size=page_size, + allocator_type=server_args.hicache_storage_backend, + ) + c128_state_host_pool = DeepSeekV4StateHostPool( + pool_name=str(PoolName.DEEPSEEK_V4_C128_STATE), + state_pools=[ + kvcache.compress_state_pools[layer_id] + for layer_id in c128_state_global_layers + ], + num_host_pages=swa_num_host_pages, + swa_page_size=kvcache.swa_page_size, + allocator_type=server_args.hicache_storage_backend, + ) + entries.extend( + [ + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C128, + host_pool=c128_host_pool, + device_pool=kvcache.c128_kv_pool, + layer_mapping=c128_layer_mapping, + transfer_layer_num=transfer_layer_num, + ), + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C128_STATE, + host_pool=c128_state_host_pool, + device_pool=None, + layer_mapping=c128_state_mapping, + transfer_layer_num=transfer_layer_num, + ), + ] + ) + + host_pool_group = HostPoolGroup(entries) + cache_controller = HybridCacheController( + params.token_to_kv_pool_allocator, + host_pool_group, + page_size, + tp_group, + load_cache_event=load_cache_event, + attn_cp_group=attn_cp_group, + attn_tp_group=attn_tp_group, + write_policy=server_args.hicache_write_policy, + io_backend=server_args.hicache_io_backend, + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + pp_rank=pp_rank, + pp_size=pp_size, + transfer_layer_num=transfer_layer_num, + enable_storage_metrics=enable_storage_metrics, + ) + return host_pool_group, cache_controller + + def build_hybrid_mamba_stack( *, params: CacheInitParams, @@ -248,8 +492,8 @@ def build_hybrid_mamba_stack( page_size: int, tp_group, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, storage_backend: Optional[str], use_mla: bool, host_mamba_evict_fn: Optional[Callable[[int], Any]] = None, @@ -317,22 +561,22 @@ def build_hybrid_mamba_stack( return host_pool_group, cache_controller -def build_shared_anchor_stack( +def build_anchor_sidecar_stack( *, params: CacheInitParams, server_args: ServerArgs, kv_pool: Any, - shared_pool_name: PoolName, + sidecar_pool_name: PoolName, full_layer_mapping: dict[int, int], page_size: int, tp_group, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, storage_backend: Optional[str], use_mla: bool, override_kv_cache_dim: Optional[int] = None, - shared_host_pool_factory: Callable[[Any], Any], + sidecar_host_pool_factory: Callable[[Any], Any], prefetch_threshold: int = 256, model_name: Optional[str] = None, storage_backend_extra_config: Optional[dict] = None, @@ -348,7 +592,7 @@ def build_shared_anchor_stack( use_mla=use_mla, override_kv_cache_dim=override_kv_cache_dim, ) - shared_host_pool = shared_host_pool_factory(kv_host_pool) + sidecar_host_pool = sidecar_host_pool_factory(kv_host_pool) entries = [ build_pool_entry( name=PoolName.KV, @@ -359,12 +603,11 @@ def build_shared_anchor_stack( is_anchor=True, ), build_pool_entry( - name=shared_pool_name, - host_pool=shared_host_pool, + name=sidecar_pool_name, + host_pool=sidecar_host_pool, device_pool=kv_pool, layer_mapping=full_layer_mapping, transfer_layer_num=transfer_layer_num, - share_indices_with_anchor=True, ), ] host_pool_group = HostPoolGroup(entries) @@ -396,11 +639,12 @@ def attach_hybrid_pool_to_unified_cache( server_args: ServerArgs, *, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> None: """Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache.""" from sglang.srt.mem_cache.base_prefix_cache import EvictParams + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, MLATokenToKVPool, @@ -414,8 +658,15 @@ def attach_hybrid_pool_to_unified_cache( swa_stack = isinstance(kvcache, SWAKVPool) mamba_stack = isinstance(kvcache, HybridLinearKVPool) nsa_stack = isinstance(kvcache, NSATokenToKVPool) + deepseek_v4_stack = isinstance(kvcache, DeepSeekV4TokenToKVPool) - if mamba_stack: + if deepseek_v4_stack: + use_mla = False + assert set(cache.components.keys()) == { + ComponentType.FULL, + ComponentType.SWA, + }, "DeepSeekV4TokenToKVPool requires FULL + SWA in UnifiedRadixCache." + elif mamba_stack: full_kv_pool = kvcache.full_kv_pool use_mla = kvcache.use_mla assert set(cache.components.keys()) == { @@ -436,7 +687,51 @@ def attach_hybrid_pool_to_unified_cache( ComponentType.FULL }, "Non-hybrid KV pool currently only supports FULL-only UnifiedRadixCache." - if mamba_stack: + if deepseek_v4_stack: + host_pool_group, cache_controller = build_deepseek_v4_hicache_stack( + params=params, + server_args=server_args, + kvcache=kvcache, + page_size=cache.page_size, + tp_group=params.tp_cache_group, + load_cache_event=load_cache_event, + attn_cp_group=attn_cp_group, + attn_tp_group=attn_tp_group, + storage_backend=None, + host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA), + device_swa_evict_fn=lambda n: cache.evict( + EvictParams(swa_num_tokens=n) + ), + pp_rank=params.pp_rank, + pp_size=params.pp_size, + ) + cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) + cache.host_pool_group = host_pool_group + cache.cache_controller = cache_controller + cache.components[ComponentType.FULL]._full_kv_pool_host = ( + cache.full_kv_pool_host + ) + cache.swa_kv_pool_host = host_pool_group.get_pool(PoolName.SWA) + cache.components[ComponentType.SWA]._swa_kv_pool_host = ( + cache.swa_kv_pool_host + ) + for pool_name, indices_from_pool in ( + (PoolName.DEEPSEEK_V4_C4, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), + (PoolName.DEEPSEEK_V4_C128, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_STATE, PoolName.SWA), + (PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, PoolName.SWA), + (PoolName.DEEPSEEK_V4_C128_STATE, PoolName.SWA), + ): + if pool_name in host_pool_group.entry_map: + cache.register_sidecar_pool( + SidecarPoolSpec( + pool_name=pool_name, + indices_from_pool=indices_from_pool, + ) + ) + transfer_layer_num = len(kvcache.compression_ratios) + elif mamba_stack: full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping) mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map) host_pool_group, cache_controller = build_hybrid_mamba_stack( @@ -519,11 +814,11 @@ def attach_hybrid_pool_to_unified_cache( full_layer_mapping = { layer_id: layer_id for layer_id in range(full_kv_pool.layer_num) } - host_pool_group, cache_controller = build_shared_anchor_stack( + host_pool_group, cache_controller = build_anchor_sidecar_stack( params=params, server_args=server_args, kv_pool=full_kv_pool, - shared_pool_name=PoolName.INDEXER, + sidecar_pool_name=PoolName.INDEXER, full_layer_mapping=full_layer_mapping, page_size=cache.page_size, tp_group=params.tp_cache_group, @@ -533,7 +828,7 @@ def attach_hybrid_pool_to_unified_cache( storage_backend=None, use_mla=use_mla, override_kv_cache_dim=full_kv_pool.kv_cache_dim, - shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( + sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( full_kv_pool, kv_host_pool, server_args.hicache_mem_layout, @@ -545,11 +840,11 @@ def attach_hybrid_pool_to_unified_cache( cache.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV) cache.host_pool_group = host_pool_group cache.cache_controller = cache_controller - # Register the NSA indexer pool as sharing anchor-KV indices so - # HiCache backup/load emits its PoolTransfer together with KV. - cache.register_hicache_anchor_kv_shared_indices_pool( - PoolName.INDEXER, - hit_policy=PoolHitPolicy.ALL_PAGES, + cache.register_sidecar_pool( + SidecarPoolSpec( + pool_name=PoolName.INDEXER, + indices_from_pool=PoolName.KV, + ) ) cache.components[ComponentType.FULL]._full_kv_pool_host = ( cache.full_kv_pool_host @@ -586,7 +881,9 @@ def attach_hybrid_pool_to_unified_cache( cache.cache_controller.layer_done_counter ) - if mamba_stack: + if deepseek_v4_stack: + pools_desc = "KV + SWA + DeepSeekV4 sidecars" + elif mamba_stack: pools_desc = "KV + MAMBA" elif swa_stack: pools_desc = "KV + SWA" @@ -613,8 +910,8 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( prefetch_threshold: int, enable_storage_metrics: bool, load_cache_event, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> None: """Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache. @@ -623,11 +920,11 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( try: kv = radix_cache.kv_cache layer_mapping = {layer_id: layer_id for layer_id in range(kv.layer_num)} - host_pool_group, cache_controller = build_shared_anchor_stack( + host_pool_group, cache_controller = build_anchor_sidecar_stack( params=params, server_args=server_args, kv_pool=kv, - shared_pool_name=PoolName.INDEXER, + sidecar_pool_name=PoolName.INDEXER, full_layer_mapping=layer_mapping, page_size=radix_cache.page_size, tp_group=radix_cache.tp_group, @@ -638,7 +935,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache( use_mla=True, override_kv_cache_dim=kv.kv_cache_dim, prefetch_threshold=prefetch_threshold, - shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( + sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost( kv, kv_host_pool, server_args.hicache_mem_layout, @@ -672,8 +969,8 @@ def attach_hybrid_pool_to_mamba_cache( prefetch_threshold: int, load_cache_event, enable_storage_metrics: bool = False, - attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None, - attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None, + attn_cp_group: Optional[torch.distributed.ProcessGroup] = None, + attn_tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> None: """Attach HostPoolGroup (KV + Mamba) + HybridCacheController for HiMambaRadixCache. diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 37fc5b4cc..c8d3b9a7b 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -1656,6 +1656,496 @@ class MambaPoolHost(HostKVCache): return ptr_list, element_size_list +# ---- V4 Compressed KV Host Pools ---- + + +class LogicalHostPool: + """Pure-logical anchor pool for V4 HiCache. + + The pool manages page-aligned token slots but holds no KV tensor. V4 + compressed side pools use these logical FULL indices as stable page anchors. + """ + + def __init__(self, size: int, page_size: int): + if size % page_size != 0: + raise ValueError( + "LogicalHostPool size must be page-aligned, " + f"got size={size}, page_size={page_size}" + ) + self.size = size + self.page_size = page_size + self.device = "cpu" + self.layout = "layer_first" + self.dtype = torch.uint8 + self.layer_num = 0 + self.start_layer = 0 + self.end_layer = 0 + self.kv_buffer = None + self.size_per_token = 0 + self.allocator = None + self.lock = threading.RLock() + self.clear() + + @synchronized + def clear(self): + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self): + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + if need_size % self.page_size != 0: + raise ValueError( + "LogicalHostPool allocation must be page-aligned, " + f"got need_size={need_size}, page_size={self.page_size}" + ) + if need_size > self.available_size(): + return None + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + return select_index + + @synchronized + def free(self, indices: torch.Tensor) -> int: + if len(indices) % self.page_size != 0: + raise ValueError( + "LogicalHostPool free must be page-aligned, " + f"got len(indices)={len(indices)}, page_size={self.page_size}" + ) + self.free_slots = torch.cat( + [self.free_slots, indices.to(dtype=torch.int64, device="cpu").flatten()] + ) + return len(indices) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + pass + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + pass + + def get_data_page(self, index, flat=True): + return torch.empty(0, dtype=torch.uint8) + + def get_dummy_flat_data_page(self): + return torch.empty(0, dtype=torch.uint8) + + def set_from_flat_data_page(self, index, data_page): + pass + + def get_page_buffer_meta(self, indices): + return None + + def get_ksize_per_token(self): + return 0 + + +class DeepSeekV4PagedHostPool(HostKVCache): + """Host mirror for a DeepSeek V4 paged KV/indexer sub-pool.""" + + def __init__( + self, + pool_name: str, + device_buffers: list[torch.Tensor], + item_bytes: int, + num_host_pages: int, + slot_page_size: int, + device: str = "cpu", + pin_memory: bool = True, + allocator_type: str = "default", + ): + self.pool_name = pool_name + self.layer_num = len(device_buffers) + self.item_bytes = item_bytes + self.num_host_pages = num_host_pages + self.slot_page_size = slot_page_size + self.dtype = torch.uint8 + self.device = device + self.pin_memory = pin_memory + self.allocator = get_allocator_from_storage(allocator_type) + self.page_size = slot_page_size + self.size = num_host_pages * slot_page_size + self.layout = "layer_first" + self.size_per_token = item_bytes + self.start_layer = 0 + self.end_layer = self.layer_num + self.lock = threading.RLock() + + self.device_buffers = device_buffers + self.gpu_device = device_buffers[0].device if device_buffers else device + + requested_bytes = self.layer_num * num_host_pages * self.item_bytes + host_mem = psutil.virtual_memory() + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory for V4 paged pool {pool_name}. " + f"Requesting {requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free." + ) + + alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] + self.kv_buffer = [ + alloc_func( + (num_host_pages, self.item_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + for _ in range(self.layer_num) + ] + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + + logger.info( + "Allocating %.2f GB host memory for V4 paged pool '%s' " + "(layers=%d, pages=%d, item_bytes=%d).", + requested_bytes / 1e9, + self.pool_name, + self.layer_num, + num_host_pages, + self.item_bytes, + ) + self.clear() + + def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: + if indices.numel() % self.slot_page_size != 0: + raise ValueError( + f"{self.pool_name} transfer indices must be page-aligned, " + f"got numel={indices.numel()}, slot_page_size={self.slot_page_size}" + ) + return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size + + def _check_io_backend(self, io_backend: str) -> None: + if io_backend != "direct": + raise NotImplementedError( + f"{self.pool_name} supports only direct io_backend, got {io_backend}" + ) + + def get_size_per_token(self): + return self.item_bytes + + def get_ksize_per_token(self): + return self.item_bytes + + def init_kv_buffer(self): + return self.kv_buffer + + def get_hybrid_pool_buffer(self): + return self.kv_buffer + + def clear(self): + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self): + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + need_size = ( + (need_size + self.slot_page_size - 1) // self.slot_page_size + ) * self.slot_page_size + if need_size > self.available_size(): + return None + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + return select_index + + @synchronized + def free(self, indices: torch.Tensor) -> int: + self.free_slots = torch.cat( + [self.free_slots, indices.to(dtype=torch.int64, device="cpu").flatten()] + ) + return len(indices) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + if host_indices is None or device_indices is None: + return + self._check_io_backend(io_backend) + host_rows = self._to_page_indices(host_indices) + device_rows = self._to_page_indices(device_indices) + transfer_kv_direct( + src_layers=self.device_buffers, + dst_layers=self.data_refs, + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + if host_indices is None or device_indices is None: + return + self._check_io_backend(io_backend) + host_rows = self._to_page_indices(host_indices) + device_rows = self._to_page_indices(device_indices) + transfer_kv_direct( + src_layers=[self.kv_buffer[layer_id]], + dst_layers=[self.device_buffers[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + page_size=1, + ) + + def get_data_page(self, index, flat=True): + index = int(index) // self.slot_page_size + data_page = torch.stack( + [self.kv_buffer[i][index] for i in range(self.layer_num)] + ) + return data_page.flatten() if flat else data_page + + def get_dummy_flat_data_page(self): + return torch.zeros( + (self.layer_num, self.item_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + ).flatten() + + def set_from_flat_data_page(self, index, data_page): + index = int(index) // self.slot_page_size + data = data_page.view(self.dtype).reshape(self.layer_num, self.item_bytes) + for i in range(self.layer_num): + self.kv_buffer[i][index].copy_(data[i]) + + def get_page_buffer_meta(self, indices): + ptr_list = [] + rows = self._to_page_indices(indices).tolist() + for row in rows: + for layer_id in range(self.layer_num): + ptr = ( + self.kv_buffer[layer_id].data_ptr() + + int(row) * self.item_bytes * self.dtype.itemsize + ) + ptr_list.append(ptr) + element_size = self.item_bytes * self.dtype.itemsize + return ptr_list, [element_size] * len(ptr_list) + + +class DeepSeekV4StateHostPool(HostKVCache): + """Host pool for V4 CompressStatePool page rows.""" + + def __init__( + self, + pool_name: str, + state_pools: list, + num_host_pages: int, + swa_page_size: int, + device: str = "cpu", + pin_memory: bool = True, + allocator_type: str = "default", + ): + if any(pool is None for pool in state_pools): + raise ValueError(f"{pool_name} state_pools must not contain None") + + self.pool_name = pool_name + self.state_pools = state_pools + self.layer_num = len(state_pools) + self.num_host_pages = num_host_pages + self.swa_page_size = swa_page_size + self.dtype = torch.uint8 + self.device = device + self.pin_memory = pin_memory + self.allocator = get_allocator_from_storage(allocator_type) + self.page_size = swa_page_size + self.size = num_host_pages * swa_page_size + self.layout = "layer_first" + self.start_layer = 0 + self.end_layer = self.layer_num + self.lock = threading.RLock() + + self.ring_size = 0 + self.state_page_bytes = 0 + self.device_page_views = [] + self.gpu_device = device + self._init_device_page_views() + self.size_per_token = self.state_page_bytes + + requested_bytes = self.layer_num * num_host_pages * self.state_page_bytes + host_mem = psutil.virtual_memory() + available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory for V4 state pool {pool_name}. " + f"Requesting {requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free." + ) + + alloc_func = ALLOC_MEMORY_FUNCS[self.gpu_device] + self.kv_buffer = [ + alloc_func( + (num_host_pages, self.state_page_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + for _ in range(self.layer_num) + ] + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + logger.info( + "Allocating %.2f GB host memory for V4 state pool '%s' " + "(layers=%d, pages=%d, state_page_bytes=%d).", + requested_bytes / 1e9, + self.pool_name, + self.layer_num, + num_host_pages, + self.state_page_bytes, + ) + + def _init_device_page_views(self) -> None: + expected_ring_size = None + expected_state_page_bytes = None + for pool in self.state_pools: + state_tensor = pool.kv_score_buffer.kv_score + if not state_tensor.is_contiguous(): + raise ValueError(f"{self.pool_name} state tensor must be contiguous") + ring_size = pool.ring_size + slot_bytes = state_tensor[0].nbytes + state_page_bytes = ring_size * slot_bytes + if expected_ring_size is None: + expected_ring_size = ring_size + expected_state_page_bytes = state_page_bytes + self.gpu_device = state_tensor.device + elif ( + expected_ring_size != ring_size + or expected_state_page_bytes != state_page_bytes + ): + raise ValueError( + f"{self.pool_name} state pools must share ring size and slot bytes" + ) + + state_bytes = state_tensor.view(torch.uint8).reshape( + state_tensor.shape[0], -1 + ) + usable_slots = (state_tensor.shape[0] // ring_size) * ring_size + self.device_page_views.append( + state_bytes[:usable_slots].reshape(-1, state_page_bytes) + ) + + self.ring_size = expected_ring_size or 0 + self.state_page_bytes = expected_state_page_bytes or 0 + + def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: + if indices.numel() % self.swa_page_size != 0: + raise ValueError( + f"{self.pool_name} transfer indices must be SWA-page-aligned, " + f"got numel={indices.numel()}, swa_page_size={self.swa_page_size}" + ) + return indices.reshape(-1, self.swa_page_size)[:, 0] // self.swa_page_size + + def _check_io_backend(self, io_backend: str) -> None: + if io_backend != "direct": + raise NotImplementedError( + f"{self.pool_name} supports only direct io_backend, got {io_backend}" + ) + + def get_size_per_token(self): + return self.state_page_bytes + + def get_ksize_per_token(self): + return self.state_page_bytes + + def init_kv_buffer(self): + return self.kv_buffer + + def get_hybrid_pool_buffer(self): + return self.kv_buffer + + def clear(self): + pass + + def available_size(self): + raise NotImplementedError( + f"{self.pool_name} reuses SWA transfer indices and has no allocator" + ) + + @synchronized + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + raise NotImplementedError( + f"{self.pool_name} reuses SWA transfer indices and has no allocator" + ) + + @synchronized + def free(self, indices: torch.Tensor) -> int: + raise NotImplementedError( + f"{self.pool_name} reuses SWA transfer indices and has no free list" + ) + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + if host_indices is None or device_indices is None: + return + self._check_io_backend(io_backend) + host_rows = self._to_page_indices(host_indices) + device_rows = self._to_page_indices(device_indices) + transfer_kv_direct( + src_layers=self.device_page_views, + dst_layers=self.data_refs, + src_indices=device_rows, + dst_indices=host_rows, + page_size=1, + ) + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + if host_indices is None or device_indices is None: + return + self._check_io_backend(io_backend) + host_rows = self._to_page_indices(host_indices) + device_rows = self._to_page_indices(device_indices) + transfer_kv_direct( + src_layers=[self.kv_buffer[layer_id]], + dst_layers=[self.device_page_views[layer_id]], + src_indices=host_rows, + dst_indices=device_rows, + page_size=1, + ) + + def get_data_page(self, index, flat=True): + index = int(index) // self.swa_page_size + data_page = torch.stack( + [self.kv_buffer[i][index] for i in range(self.layer_num)] + ) + return data_page.flatten() if flat else data_page + + def get_dummy_flat_data_page(self): + return torch.zeros( + (self.layer_num, self.state_page_bytes), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + ).flatten() + + def set_from_flat_data_page(self, index, data_page): + index = int(index) // self.swa_page_size + data = data_page.view(self.dtype).reshape(self.layer_num, self.state_page_bytes) + for i in range(self.layer_num): + self.kv_buffer[i][index].copy_(data[i]) + + def get_page_buffer_meta(self, indices): + ptr_list = [] + rows = self._to_page_indices(indices).tolist() + for row in rows: + for layer_id in range(self.layer_num): + ptr = ( + self.kv_buffer[layer_id].data_ptr() + + int(row) * self.state_page_bytes * self.dtype.itemsize + ) + ptr_list.append(ptr) + element_size = self.state_page_bytes * self.dtype.itemsize + return ptr_list, [element_size] * len(ptr_list) + + @dataclass class PoolEntry: name: PoolName @@ -1663,9 +2153,6 @@ class PoolEntry: device_pool: Any layer_mapper: Callable[[int], Optional[int]] is_primary_index_anchor: bool = False - # When True, host_pool uses the same logical slot indices as the anchor pool - # (e.g. DSA indexer); HostPoolGroup.free mirrors frees to this pool. - share_indices_with_anchor: bool = False # Optional eviction callbacks for auto-alloc in HybridCacheController. # host_evict_fn(n): evict n slots from the host pool (used by write()). # device_evict_fn(n): evict n slots from the device pool (used by load()). diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 5b264307a..f275c9925 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -22,7 +22,11 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + SidecarPoolSpec, +) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache_components import ( _NUM_COMPONENT_TYPES, @@ -221,9 +225,7 @@ class UnifiedRadixCache(BasePrefixCache): self._components_tuple: tuple[TreeComponent, ...] = tuple( self.components.values() ) - self.hicache_anchor_kv_shared_indices_pools: list[ - tuple[PoolName, PoolHitPolicy] - ] = [] + self.sidecar_pool_specs: list[SidecarPoolSpec] = [] # Streaming session: embedded StreamingSession with self as inner. # Always on -- zero overhead when no streaming session is open (the @@ -309,7 +311,7 @@ class UnifiedRadixCache(BasePrefixCache): ) self.load_cache_event = threading.Event() - self.hicache_anchor_kv_shared_indices_pools.clear() + self.sidecar_pool_specs.clear() attach_hybrid_pool_to_unified_cache( self, params, @@ -333,12 +335,8 @@ class UnifiedRadixCache(BasePrefixCache): f"transfer_layer_num={self.cache_controller.layer_num}" ) - def register_hicache_anchor_kv_shared_indices_pool( - self, - pool_name: PoolName, - hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES, - ) -> None: - self.hicache_anchor_kv_shared_indices_pools.append((pool_name, hit_policy)) + def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None: + self.sidecar_pool_specs.append(spec) def match_prefix(self, params: MatchPrefixParams) -> MatchResult: result = self.session.try_match_prefix(params) @@ -1164,7 +1162,10 @@ class UnifiedRadixCache(BasePrefixCache): ): return 0 - # Build aux transfers, keyed per component + device_value = node.component_data[BASE_COMPONENT_TYPE].value + kv_xfer = PoolTransfer(name=PoolName.KV, device_indices=device_value) + + # Build aux transfers, keyed per component. comp_xfers: dict[ComponentType, list] = {} for comp in self._components_tuple: if comp.component_type == BASE_COMPONENT_TYPE: @@ -1172,13 +1173,11 @@ class UnifiedRadixCache(BasePrefixCache): t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) if t: comp_xfers[comp.component_type] = t - anchor_kv_shared_indices_xfers = [ - PoolTransfer(name=pool_name, hit_policy=hit_policy) - for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools - ] + sidecar_xfers = self._build_sidecar_transfers( + CacheTransferPhase.BACKUP_HOST, kv_xfer, comp_xfers + ) # Pre-evict host if insufficient - device_value = node.component_data[BASE_COMPONENT_TYPE].value kv_tokens = len(device_value) host_avail = self.cache_controller.mem_pool_host.available_size() if host_avail < kv_tokens: @@ -1188,7 +1187,7 @@ class UnifiedRadixCache(BasePrefixCache): return 0 aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] - aux_xfers.extend(anchor_kv_shared_indices_xfers) + aux_xfers.extend(sidecar_xfers) host_indices = self.cache_controller.write( device_value, node_id=node.id, extra_pools=aux_xfers or None ) @@ -1245,10 +1244,9 @@ class UnifiedRadixCache(BasePrefixCache): ) if t: comp_xfers[comp.component_type] = t - anchor_kv_shared_indices_xfers = [ - PoolTransfer(name=pool_name, hit_policy=hit_policy) - for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools - ] + sidecar_xfers = self._build_sidecar_transfers( + CacheTransferPhase.LOAD_BACK, kv_xfer, comp_xfers + ) # Skip if there is nothing to load, or if the Full-KV transfer is too # small / exceeds memory quota. Aux transfers should still run even @@ -1269,7 +1267,7 @@ class UnifiedRadixCache(BasePrefixCache): # Load H→D aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] - aux_xfers.extend(anchor_kv_shared_indices_xfers) + aux_xfers.extend(sidecar_xfers) device_indices = self.cache_controller.load( host_indices=kv_xfer.host_indices, node_id=best_match_node.id, @@ -1301,6 +1299,52 @@ class UnifiedRadixCache(BasePrefixCache): ) return device_indices + def _build_sidecar_transfers( + self, + phase: CacheTransferPhase, + kv_xfer: PoolTransfer, + comp_xfers: dict[ComponentType, list[PoolTransfer]], + ) -> list[PoolTransfer]: + transfers: list[PoolTransfer] = [] + for spec in self.sidecar_pool_specs: + if spec.indices_from_pool == PoolName.KV: + indices_source = kv_xfer + else: + source_component = { + PoolName.SWA: ComponentType.SWA, + PoolName.MAMBA: ComponentType.MAMBA, + }.get(spec.indices_from_pool) + if source_component is None: + raise AssertionError( + f"Unsupported sidecar indices source pool " + f"{spec.indices_from_pool}." + ) + matching_sources = comp_xfers.get(source_component, ()) + if not matching_sources: + continue + indices_source = matching_sources[0] + if indices_source.name != spec.indices_from_pool: + raise AssertionError( + f"Sidecar indices source pool {spec.indices_from_pool} " + f"resolved to {indices_source.name} during {phase}." + ) + + indices = ( + indices_source.device_indices + if phase == CacheTransferPhase.BACKUP_HOST + else indices_source.host_indices + ) + if indices is None or len(indices) == 0: + continue + transfers.append( + PoolTransfer( + name=spec.pool_name, + hit_policy=spec.hit_policy, + indices_from_pool=spec.indices_from_pool, + ) + ) + return transfers + def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None: """Increment hit count; trigger write_backup when threshold reached.""" if self.cache_controller is None: diff --git a/python/sglang/test/kl_multiturn_utils.py b/python/sglang/test/kl_multiturn_utils.py index 9d95dc5fd..bd21c5023 100644 --- a/python/sglang/test/kl_multiturn_utils.py +++ b/python/sglang/test/kl_multiturn_utils.py @@ -104,6 +104,7 @@ def _replay_and_compare_kl( output_logprobs: list[list[float]], label: str, batch_size: int = 1, + sampling_temperature: float = 1, ): """Flush cache, run replay prefill in batches, compare KL divergence.""" all_input_logprobs = [] @@ -114,6 +115,7 @@ def _replay_and_compare_kl( base_url, replay_input_ids[start:end], output_logprobs[start:end], + temperature=sampling_temperature, ) ) acc = {model_name: {"kl_div": kl_threshold}} @@ -142,7 +144,9 @@ def _interleave_order(n: int, branches_per_group: int) -> list[int] | None: return order -def _generate_maybe_interleaved(base_url, inputs, max_new_tokens, order=None): +def _generate_maybe_interleaved( + base_url, inputs, max_new_tokens, order=None, sampling_temperature: float = 1 +): """Generate with optional interleaved submission order. Submits inputs reordered by ``order``, then maps results back to the @@ -150,9 +154,21 @@ def _generate_maybe_interleaved(base_url, inputs, max_new_tokens, order=None): inputs[i]. """ if order is None: - return _generate(base_url, inputs, max_new_tokens, return_logprob=True) + return _generate( + base_url, + inputs, + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, + ) ordered = [inputs[i] for i in order] - results = _generate(base_url, ordered, max_new_tokens, return_logprob=True) + results = _generate( + base_url, + ordered, + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, + ) unordered = [None] * len(results) for idx, orig in enumerate(order): unordered[orig] = results[idx] @@ -178,6 +194,7 @@ def test_input_output_logprobs_match_helper( # --- Cache assertion (for turns > 0) --- assert_decode_cached_tokens: Callable | None = None, replay_batch_size: int = 1, + sampling_temperature: float = 1, ): """Verify decode logprobs match prefill replay. @@ -213,7 +230,11 @@ def test_input_output_logprobs_match_helper( ] results = _generate( - base_url, current_input, max_new_tokens, return_logprob=True + base_url, + current_input, + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, ) assert len(results) == n @@ -242,6 +263,7 @@ def test_input_output_logprobs_match_helper( output_lps, label=label, batch_size=replay_batch_size, + sampling_temperature=sampling_temperature, ) @@ -269,6 +291,7 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper( # --- Interleaving for branch stress --- branches_per_group: int = 0, replay_batch_size: int = 1, + sampling_temperature: float = 1, ): """Verify logprobs when prefill cache is hit. @@ -305,10 +328,21 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper( # Seed cache with prefixes _flush_cache(base_url) - _generate(base_url, prefix_input_ids, max_new_tokens=0) + _generate( + base_url, + prefix_input_ids, + max_new_tokens=0, + temperature=sampling_temperature, + ) # Turn 0: prefill cache hit (NOT interleaved, matching original behavior) - results = _generate(base_url, full_input_ids, max_new_tokens, return_logprob=True) + results = _generate( + base_url, + full_input_ids, + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, + ) assert len(results) == n for i, result in enumerate(results): @@ -331,7 +365,11 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper( current_input[i] + last_outputs[i] + suffixes[i] for i in range(n) ] results = _generate_maybe_interleaved( - base_url, current_input, max_new_tokens, order + base_url, + current_input, + max_new_tokens, + order, + sampling_temperature=sampling_temperature, ) assert len(results) == n @@ -359,6 +397,7 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper( output_lps, label=label, batch_size=replay_batch_size, + sampling_temperature=sampling_temperature, ) @@ -383,6 +422,7 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( # --- Interleaving --- branches_per_group: int = 0, replay_batch_size: int = 1, + sampling_temperature: float = 1, ): """Verify logprobs when decode cache is hit. @@ -414,7 +454,11 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( # Turn 1: populate cache, no assertion, no interleaving _flush_cache(base_url) results = _generate( - base_url, first_turn_input_ids, max_new_tokens, return_logprob=True + base_url, + first_turn_input_ids, + max_new_tokens, + return_logprob=True, + temperature=sampling_temperature, ) assert len(results) == n @@ -429,7 +473,11 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( current_input[i] + last_outputs[i] + suffixes[i] for i in range(n) ] results = _generate_maybe_interleaved( - base_url, current_input, max_new_tokens, order + base_url, + current_input, + max_new_tokens, + order, + sampling_temperature=sampling_temperature, ) assert len(results) == n @@ -457,4 +505,5 @@ def test_input_output_logprobs_match_decode_cache_hit_helper( output_lps, label=label, batch_size=replay_batch_size, + sampling_temperature=sampling_temperature, ) diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py index de77fb7af..ef13354d8 100644 --- a/python/sglang/test/kl_test_utils.py +++ b/python/sglang/test/kl_test_utils.py @@ -128,13 +128,18 @@ def _flush_cache(base_url, timeout_s=30): def _generate( - base_url, input_ids, max_new_tokens, return_logprob=False, logprob_start_len=-1 + base_url, + input_ids, + max_new_tokens, + return_logprob=False, + logprob_start_len=-1, + temperature=1, ): """Send generate request and return results.""" json_data = { "input_ids": input_ids, "sampling_params": { - "temperature": 1, + "temperature": temperature, "max_new_tokens": max_new_tokens, "ignore_eos": True, }, @@ -151,7 +156,7 @@ def _generate( return response.json() -def _get_input_logprobs(base_url, new_input_ids, output_logprobs): +def _get_input_logprobs(base_url, new_input_ids, output_logprobs, temperature=1): """Run prefill to get input logprobs matching output logprobs.""" _flush_cache(base_url) results = _generate( @@ -160,6 +165,7 @@ def _get_input_logprobs(base_url, new_input_ids, output_logprobs): max_new_tokens=0, return_logprob=True, logprob_start_len=0, + temperature=temperature, ) assert len(results) == len(new_input_ids) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl.py index d4b055748..e6d38c6a2 100644 --- a/test/registered/radix_cache/test_unified_radix_cache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl.py @@ -48,6 +48,7 @@ class UnifiedRadixTreeTestMixin: prefix_len: int = 512 prefill_cache_assert = None decode_cache_assert = None + sampling_temperature: float = 1 gsm8k_threshold: float = 0.93 mmlu_threshold: float = 0.8 @@ -106,6 +107,7 @@ class UnifiedRadixTreeTestMixin: turn_suffixes=[t2, t3], assert_decode_cached_tokens=self.decode_cache_assert, max_new_tokens=self.max_new_tokens, + sampling_temperature=self.sampling_temperature, ) def test_multiturn_prefill_cache_hit_branching(self): @@ -135,6 +137,7 @@ class UnifiedRadixTreeTestMixin: assert_decode_cached_tokens=self.decode_cache_assert, branches_per_group=branches, max_new_tokens=self.max_new_tokens, + sampling_temperature=self.sampling_temperature, ) def test_multiturn_decode_cache_hit_branching(self): @@ -159,6 +162,7 @@ class UnifiedRadixTreeTestMixin: assert_decode_cached_tokens=self.decode_cache_assert, branches_per_group=branches, max_new_tokens=self.max_new_tokens, + sampling_temperature=self.sampling_temperature, ) @@ -267,58 +271,5 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase): kill_process_tree(cls.process.pid) -class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): - """Mamba hybrid + HiCache + UnifiedRadixCache.""" - - kl_threshold = 0.003 - prefill_cache_assert = staticmethod( - make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) - ) - decode_cache_assert = staticmethod( - make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) - ) - - @classmethod - def setUpClass(cls): - cls.model = MAMBA_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=[ - "--tp-size", - "4", - "--chunked-prefill-size", - "2048", - "--mem-fraction-static", - "0.85", - "--mamba-scheduler-strategy", - "extra_buffer", - "--mamba-track-interval", - str(MAMBA_TRACK_INTERVAL), - "--enable-hierarchical-cache", - "--hicache-ratio", - "4", - "--hicache-write-policy", - "write_through", - "--hicache-io-backend", - "direct", - "--hicache-mem-layout", - "page_first_direct", - "--max-total-tokens", - "12000", - "--max-running-requests", - "4", - ], - env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, - ) - cls.input_ids = get_input_ids(cls.model, num_samples=18) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - if __name__ == "__main__": unittest.main() diff --git a/test/registered/radix_cache/test_unified_radix_hicache_kl.py b/test/registered/radix_cache/test_unified_radix_hicache_kl.py new file mode 100644 index 000000000..81253f6a7 --- /dev/null +++ b/test/registered/radix_cache/test_unified_radix_hicache_kl.py @@ -0,0 +1,162 @@ +"""UnifiedRadixTree + HiCache KL divergence tests. + +Tests Mamba hybrid, DeepSeek V4 Flash, and GLM-5 models with HiCache L2 +offloading under UnifiedRadixTree, verifying multi-turn cache correctness +via KL divergence. +""" + +import unittest + +from test_unified_radix_cache_kl import UnifiedRadixTreeTestMixin + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_multiturn_utils import ( + get_input_ids, + make_mamba_decode_assert, + make_mamba_prefill_assert, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" +MAMBA_CHUNK_SIZE = 64 +MAMBA_TRACK_INTERVAL = 128 + +DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" +DSV4_FLASH_LAUNCH_TIMEOUT = 3600 + +DSV32_MODEL = "deepseek-ai/DeepSeek-V3.2" +DSV32_LAUNCH_TIMEOUT = 3600 + +register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True) + + +class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """Mamba hybrid + HiCache + UnifiedRadixCache.""" + + kl_threshold = 0.003 + prefill_cache_assert = staticmethod( + make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE) + ) + decode_cache_assert = staticmethod( + make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL) + ) + + @classmethod + def setUpClass(cls): + cls.model = MAMBA_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--chunked-prefill-size", + "2048", + "--mem-fraction-static", + "0.85", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + str(MAMBA_TRACK_INTERVAL), + "--enable-hierarchical-cache", + "--hicache-ratio", + "4", + "--hicache-write-policy", + "write_through", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--max-total-tokens", + "12000", + "--max-running-requests", + "4", + ], + env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"}, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label): + expected = history_len + output_len + actual = result["meta_info"]["cached_tokens"] + lower = max(0, expected - 256) + assert actual >= lower, f"{label}: expected cached_tokens>={lower}, got {actual}" + + +class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCase): + """DeepSeek V4 Flash FP8 + HiCache + UnifiedRadixCache.""" + + kl_threshold = 0.0035 + sampling_temperature = 0 + decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens) + gsm8k_threshold = 0.90 + num_gsm8k_questions = 100 + + @unittest.skip("no stable.") + def test_multiturn_logprobs_match(self): + pass + + @classmethod + def setUpClass(cls): + cls.model = DSV4_FLASH_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DSV4_FLASH_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp-size", + "4", + "--attention-backend", + "compressed", + "--page-size", + "256", + "--chunked-prefill-size", + "8192", + "--mem-fraction-static", + "0.9", + "--disable-shared-experts-fusion", + "--enable-hierarchical-cache", + "--hicache-ratio", + "4", + "--hicache-write-policy", + "write_through", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--swa-full-tokens-ratio", + "0.25", + "--max-total-tokens", + "20000", + "--max-running-requests", + "4", + ], + env={ + "SGLANG_DSV4_FP4_EXPERTS": "0", + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + }, + ) + cls.input_ids = get_input_ids(cls.model, num_samples=18) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + +if __name__ == "__main__": + unittest.main()