[UnifiedTree]: Support HiCache For DeepSeek_V4 (#24691)
Co-authored-by: ispobock <ispobaoke@gmail.com> Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
@@ -605,7 +605,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
|||||||
else:
|
else:
|
||||||
raise ValueError(f"Unsupported compression ratio: {ratio}")
|
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:
|
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]
|
compress_state_pool = self.compress_state_pools[layer_id]
|
||||||
assert (
|
assert (
|
||||||
compress_state_pool is not None
|
compress_state_pool is not None
|
||||||
@@ -613,6 +618,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
|||||||
return compress_state_pool
|
return compress_state_pool
|
||||||
|
|
||||||
def get_indexer_compress_states(self, layer_id: int) -> CompressStatePool:
|
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]
|
indexer_compress_state_pool = self.indexer_compress_state_pools[layer_id]
|
||||||
assert (
|
assert (
|
||||||
indexer_compress_state_pool is not None
|
indexer_compress_state_pool is not None
|
||||||
@@ -620,6 +626,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
|||||||
return indexer_compress_state_pool
|
return indexer_compress_state_pool
|
||||||
|
|
||||||
def get_swa_key_buffer(self, layer_id: int) -> torch.Tensor:
|
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)
|
return self.swa_kv_pool.get_key_buffer(layer_id)
|
||||||
|
|
||||||
def set_swa_key_buffer(
|
def set_swa_key_buffer(
|
||||||
@@ -635,7 +642,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
|||||||
assert compress_kv_pool is not None
|
assert compress_kv_pool is not None
|
||||||
return compress_kv_pool.page_size
|
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]
|
_, compress_layer_id, compress_kv_pool = self.layer_mapping[layer_id]
|
||||||
assert compress_kv_pool is not None
|
assert compress_kv_pool is not None
|
||||||
return compress_kv_pool.get_key_buffer(compress_layer_id)
|
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
|
return self.c4_indexer_kv_pool.page_size
|
||||||
|
|
||||||
def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
|
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]
|
compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
|
||||||
assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }"
|
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)
|
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,
|
seq_len: int,
|
||||||
page_indices: torch.Tensor,
|
page_indices: torch.Tensor,
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
self.wait_layer_transfer(layer_id)
|
||||||
compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
|
compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
|
||||||
assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }"
|
assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }"
|
||||||
return self.c4_indexer_kv_pool.get_index_k_scale_buffer(
|
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:
|
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)
|
return self.swa_kv_pool.get_key_buffer(layer_id)
|
||||||
|
|
||||||
def set_swa_key_buffer_radix_fused(
|
def set_swa_key_buffer_radix_fused(
|
||||||
|
|||||||
@@ -1,14 +1,18 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from typing import Any, List, Optional, Set
|
from typing import TYPE_CHECKING, Any, List, Optional, Set
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.srt.environ import envs
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -52,6 +56,14 @@ class PoolName(str, Enum):
|
|||||||
MAMBA = "mamba"
|
MAMBA = "mamba"
|
||||||
SWA = "swa"
|
SWA = "swa"
|
||||||
INDEXER = "indexer"
|
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:
|
def __str__(self) -> str:
|
||||||
return self.value
|
return self.value
|
||||||
@@ -83,6 +95,16 @@ class PoolTransfer:
|
|||||||
keys: Optional[List[str]] = None
|
keys: Optional[List[str]] = None
|
||||||
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
|
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
|
||||||
nodes_to_load: Optional[List[Any]] = None
|
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
|
@dataclass
|
||||||
|
|||||||
@@ -647,6 +647,7 @@ class HiRadixCache(RadixCache):
|
|||||||
pool = PoolTransfer(
|
pool = PoolTransfer(
|
||||||
name=PoolName.INDEXER,
|
name=PoolName.INDEXER,
|
||||||
hit_policy=PoolHitPolicy.ALL_PAGES,
|
hit_policy=PoolHitPolicy.ALL_PAGES,
|
||||||
|
indices_from_pool=PoolName.KV,
|
||||||
)
|
)
|
||||||
return {"extra_pools": [pool]}
|
return {"extra_pools": [pool]}
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -23,6 +23,7 @@ from sglang.srt.managers.cache_controller import (
|
|||||||
from sglang.srt.mem_cache.hicache_storage import (
|
from sglang.srt.mem_cache.hicache_storage import (
|
||||||
HiCacheStorageExtraInfo,
|
HiCacheStorageExtraInfo,
|
||||||
PoolHitPolicy,
|
PoolHitPolicy,
|
||||||
|
PoolName,
|
||||||
PoolTransfer,
|
PoolTransfer,
|
||||||
PoolTransferResult,
|
PoolTransferResult,
|
||||||
)
|
)
|
||||||
@@ -50,12 +51,12 @@ class CacheOperation(BaseCacheOperation):
|
|||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def merge_pool_transfers(
|
def merge_pool_transfers(
|
||||||
ops: List["CacheOperation"],
|
ops: List[CacheOperation],
|
||||||
) -> Optional[list[PoolTransfer]]:
|
) -> Optional[list[PoolTransfer]]:
|
||||||
grouped: dict[str, list[PoolTransfer]] = {}
|
grouped: dict[tuple[PoolName, Optional[PoolName]], list[PoolTransfer]] = {}
|
||||||
for op in ops:
|
for op in ops:
|
||||||
for t in op.pool_transfers or []:
|
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:
|
if not grouped:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
@@ -65,16 +66,18 @@ class CacheOperation(BaseCacheOperation):
|
|||||||
|
|
||||||
return [
|
return [
|
||||||
PoolTransfer(
|
PoolTransfer(
|
||||||
name=name,
|
name=ts[0].name,
|
||||||
host_indices=cat_or_none(t.host_indices for t in ts),
|
host_indices=cat_or_none(t.host_indices for t in ts),
|
||||||
device_indices=cat_or_none(t.device_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,
|
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
|
@staticmethod
|
||||||
def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation":
|
def merge_ops(ops: List[CacheOperation]) -> CacheOperation:
|
||||||
if len(ops) == 1:
|
if len(ops) == 1:
|
||||||
return ops[0]
|
return ops[0]
|
||||||
host_indices = torch.cat([op.host_indices for op in ops])
|
host_indices = torch.cat([op.host_indices for op in ops])
|
||||||
@@ -478,6 +481,7 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
device_indices=transfer_device_indices,
|
device_indices=transfer_device_indices,
|
||||||
keys=transfer.keys,
|
keys=transfer.keys,
|
||||||
hit_policy=transfer.hit_policy,
|
hit_policy=transfer.hit_policy,
|
||||||
|
indices_from_pool=transfer.indices_from_pool,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
return host_indices, device_indices, resolved_pool_transfers
|
return host_indices, device_indices, resolved_pool_transfers
|
||||||
@@ -485,7 +489,7 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
def _page_transfer(self, operation):
|
def _page_transfer(self, operation):
|
||||||
# Transfer extra pools
|
# Transfer extra pools
|
||||||
if operation.pool_transfers and not operation.is_terminated():
|
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)
|
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)
|
||||||
|
|
||||||
@@ -495,19 +499,26 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
def _page_backup(self, operation):
|
def _page_backup(self, operation):
|
||||||
# Backup extra pools
|
# Backup extra pools
|
||||||
if operation.pool_transfers:
|
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)
|
results = self.storage_backend.batch_set_v2(operation.pool_transfers)
|
||||||
operation.pool_storage_result.update_extra_pool_hit_pages(results)
|
operation.pool_storage_result.update_extra_pool_hit_pages(results)
|
||||||
|
|
||||||
# Backup kv pools
|
# Backup kv pools
|
||||||
super()._page_backup(operation)
|
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:
|
for transfer in operation.pool_transfers:
|
||||||
entry = self.mem_pool_host.entry_map.get(transfer.name)
|
if transfer.indices_from_pool is None:
|
||||||
if entry.share_indices_with_anchor:
|
continue
|
||||||
transfer.keys = operation.hash_value
|
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
|
transfer.host_indices = operation.host_indices
|
||||||
|
if transfer.keys is None:
|
||||||
|
transfer.keys = operation.hash_value
|
||||||
|
|
||||||
def _sync_trailing_keys(
|
def _sync_trailing_keys(
|
||||||
self,
|
self,
|
||||||
@@ -541,14 +552,23 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
return None
|
return None
|
||||||
# (pool, free_fn, indices) for atomic rollback on failure.
|
# (pool, free_fn, indices) for atomic rollback on failure.
|
||||||
newly_allocated: list[tuple[PoolTransfer, Callable, torch.Tensor]] = []
|
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:
|
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)
|
entry = self.mem_pool_host.entry_map.get(pool.name)
|
||||||
if entry is None:
|
if entry is None:
|
||||||
continue
|
continue
|
||||||
if entry.share_indices_with_anchor:
|
|
||||||
pool.device_indices = kv_device_indices
|
|
||||||
pool.host_indices = kv_host_indices
|
|
||||||
continue
|
|
||||||
if alloc_host:
|
if alloc_host:
|
||||||
if pool.host_indices is not None or pool.device_indices is None:
|
if pool.host_indices is not None or pool.device_indices is None:
|
||||||
continue
|
continue
|
||||||
@@ -572,16 +592,33 @@ class HybridCacheController(BaseHiCacheController):
|
|||||||
indices = alloc_fn(size)
|
indices = alloc_fn(size)
|
||||||
if indices is None:
|
if indices is None:
|
||||||
# Atomic rollback: free everything we successfully allocated.
|
# Atomic rollback: free everything we successfully allocated.
|
||||||
for prev_pool, prev_free_fn, prev_indices in newly_allocated:
|
rollback_allocated()
|
||||||
prev_free_fn(prev_indices)
|
|
||||||
if alloc_host:
|
|
||||||
prev_pool.host_indices = None
|
|
||||||
else:
|
|
||||||
prev_pool.device_indices = None
|
|
||||||
return None
|
return None
|
||||||
if alloc_host:
|
if alloc_host:
|
||||||
pool.host_indices = indices
|
pool.host_indices = indices
|
||||||
else:
|
else:
|
||||||
pool.device_indices = indices
|
pool.device_indices = indices
|
||||||
newly_allocated.append((pool, free_fn, 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
|
return extra_pools
|
||||||
|
|||||||
@@ -3,12 +3,15 @@ from __future__ import annotations
|
|||||||
import logging
|
import logging
|
||||||
from typing import TYPE_CHECKING, Any, Callable, Optional
|
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 (
|
from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import (
|
||||||
HybridCacheController,
|
HybridCacheController,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.memory_pool_host import (
|
from sglang.srt.mem_cache.memory_pool_host import (
|
||||||
|
DeepSeekV4PagedHostPool,
|
||||||
|
DeepSeekV4StateHostPool,
|
||||||
HostPoolGroup,
|
HostPoolGroup,
|
||||||
|
LogicalHostPool,
|
||||||
MambaPoolHost,
|
MambaPoolHost,
|
||||||
MHATokenToKVPoolHost,
|
MHATokenToKVPoolHost,
|
||||||
MLATokenToKVPoolHost,
|
MLATokenToKVPoolHost,
|
||||||
@@ -71,7 +74,6 @@ def build_pool_entry(
|
|||||||
layer_mapping: dict[int, int],
|
layer_mapping: dict[int, int],
|
||||||
transfer_layer_num: int,
|
transfer_layer_num: int,
|
||||||
is_anchor: bool = False,
|
is_anchor: bool = False,
|
||||||
share_indices_with_anchor: bool = False,
|
|
||||||
host_evict_fn: Optional[Callable[[int], Any]] = None,
|
host_evict_fn: Optional[Callable[[int], Any]] = None,
|
||||||
device_evict_fn: Optional[Callable[[int], Any]] = None,
|
device_evict_fn: Optional[Callable[[int], Any]] = None,
|
||||||
device_alloc_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,
|
device_pool=device_pool,
|
||||||
layer_mapper=_make_layer_mapper(layer_mapping, transfer_layer_num),
|
layer_mapper=_make_layer_mapper(layer_mapping, transfer_layer_num),
|
||||||
is_primary_index_anchor=is_anchor,
|
is_primary_index_anchor=is_anchor,
|
||||||
share_indices_with_anchor=share_indices_with_anchor,
|
|
||||||
host_evict_fn=host_evict_fn,
|
host_evict_fn=host_evict_fn,
|
||||||
device_evict_fn=device_evict_fn,
|
device_evict_fn=device_evict_fn,
|
||||||
device_alloc_fn=device_alloc_fn,
|
device_alloc_fn=device_alloc_fn,
|
||||||
@@ -100,8 +101,8 @@ def build_kv_only_stack(
|
|||||||
page_size: int,
|
page_size: int,
|
||||||
tp_group,
|
tp_group,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
storage_backend: Optional[str],
|
storage_backend: Optional[str],
|
||||||
use_mla: bool,
|
use_mla: bool,
|
||||||
override_kv_cache_dim: Optional[int] = None,
|
override_kv_cache_dim: Optional[int] = None,
|
||||||
@@ -164,8 +165,8 @@ def build_hybrid_swa_stack(
|
|||||||
page_size: int,
|
page_size: int,
|
||||||
tp_group,
|
tp_group,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
storage_backend: Optional[str],
|
storage_backend: Optional[str],
|
||||||
use_mla: bool,
|
use_mla: bool,
|
||||||
host_swa_evict_fn: Optional[Callable[[int], Any]] = None,
|
host_swa_evict_fn: Optional[Callable[[int], Any]] = None,
|
||||||
@@ -237,6 +238,249 @@ def build_hybrid_swa_stack(
|
|||||||
return host_pool_group, cache_controller
|
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(
|
def build_hybrid_mamba_stack(
|
||||||
*,
|
*,
|
||||||
params: CacheInitParams,
|
params: CacheInitParams,
|
||||||
@@ -248,8 +492,8 @@ def build_hybrid_mamba_stack(
|
|||||||
page_size: int,
|
page_size: int,
|
||||||
tp_group,
|
tp_group,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
storage_backend: Optional[str],
|
storage_backend: Optional[str],
|
||||||
use_mla: bool,
|
use_mla: bool,
|
||||||
host_mamba_evict_fn: Optional[Callable[[int], Any]] = None,
|
host_mamba_evict_fn: Optional[Callable[[int], Any]] = None,
|
||||||
@@ -317,22 +561,22 @@ def build_hybrid_mamba_stack(
|
|||||||
return host_pool_group, cache_controller
|
return host_pool_group, cache_controller
|
||||||
|
|
||||||
|
|
||||||
def build_shared_anchor_stack(
|
def build_anchor_sidecar_stack(
|
||||||
*,
|
*,
|
||||||
params: CacheInitParams,
|
params: CacheInitParams,
|
||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
kv_pool: Any,
|
kv_pool: Any,
|
||||||
shared_pool_name: PoolName,
|
sidecar_pool_name: PoolName,
|
||||||
full_layer_mapping: dict[int, int],
|
full_layer_mapping: dict[int, int],
|
||||||
page_size: int,
|
page_size: int,
|
||||||
tp_group,
|
tp_group,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
storage_backend: Optional[str],
|
storage_backend: Optional[str],
|
||||||
use_mla: bool,
|
use_mla: bool,
|
||||||
override_kv_cache_dim: Optional[int] = None,
|
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,
|
prefetch_threshold: int = 256,
|
||||||
model_name: Optional[str] = None,
|
model_name: Optional[str] = None,
|
||||||
storage_backend_extra_config: Optional[dict] = None,
|
storage_backend_extra_config: Optional[dict] = None,
|
||||||
@@ -348,7 +592,7 @@ def build_shared_anchor_stack(
|
|||||||
use_mla=use_mla,
|
use_mla=use_mla,
|
||||||
override_kv_cache_dim=override_kv_cache_dim,
|
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 = [
|
entries = [
|
||||||
build_pool_entry(
|
build_pool_entry(
|
||||||
name=PoolName.KV,
|
name=PoolName.KV,
|
||||||
@@ -359,12 +603,11 @@ def build_shared_anchor_stack(
|
|||||||
is_anchor=True,
|
is_anchor=True,
|
||||||
),
|
),
|
||||||
build_pool_entry(
|
build_pool_entry(
|
||||||
name=shared_pool_name,
|
name=sidecar_pool_name,
|
||||||
host_pool=shared_host_pool,
|
host_pool=sidecar_host_pool,
|
||||||
device_pool=kv_pool,
|
device_pool=kv_pool,
|
||||||
layer_mapping=full_layer_mapping,
|
layer_mapping=full_layer_mapping,
|
||||||
transfer_layer_num=transfer_layer_num,
|
transfer_layer_num=transfer_layer_num,
|
||||||
share_indices_with_anchor=True,
|
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
host_pool_group = HostPoolGroup(entries)
|
host_pool_group = HostPoolGroup(entries)
|
||||||
@@ -396,11 +639,12 @@ def attach_hybrid_pool_to_unified_cache(
|
|||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
*,
|
*,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
|
"""Attach HostPoolGroup + HybridCacheController to UnifiedRadixCache."""
|
||||||
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
|
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 (
|
from sglang.srt.mem_cache.memory_pool import (
|
||||||
HybridLinearKVPool,
|
HybridLinearKVPool,
|
||||||
MLATokenToKVPool,
|
MLATokenToKVPool,
|
||||||
@@ -414,8 +658,15 @@ def attach_hybrid_pool_to_unified_cache(
|
|||||||
swa_stack = isinstance(kvcache, SWAKVPool)
|
swa_stack = isinstance(kvcache, SWAKVPool)
|
||||||
mamba_stack = isinstance(kvcache, HybridLinearKVPool)
|
mamba_stack = isinstance(kvcache, HybridLinearKVPool)
|
||||||
nsa_stack = isinstance(kvcache, NSATokenToKVPool)
|
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
|
full_kv_pool = kvcache.full_kv_pool
|
||||||
use_mla = kvcache.use_mla
|
use_mla = kvcache.use_mla
|
||||||
assert set(cache.components.keys()) == {
|
assert set(cache.components.keys()) == {
|
||||||
@@ -436,7 +687,51 @@ def attach_hybrid_pool_to_unified_cache(
|
|||||||
ComponentType.FULL
|
ComponentType.FULL
|
||||||
}, "Non-hybrid KV pool currently only supports FULL-only UnifiedRadixCache."
|
}, "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)
|
full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping)
|
||||||
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
|
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
|
||||||
host_pool_group, cache_controller = build_hybrid_mamba_stack(
|
host_pool_group, cache_controller = build_hybrid_mamba_stack(
|
||||||
@@ -519,11 +814,11 @@ def attach_hybrid_pool_to_unified_cache(
|
|||||||
full_layer_mapping = {
|
full_layer_mapping = {
|
||||||
layer_id: layer_id for layer_id in range(full_kv_pool.layer_num)
|
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,
|
params=params,
|
||||||
server_args=server_args,
|
server_args=server_args,
|
||||||
kv_pool=full_kv_pool,
|
kv_pool=full_kv_pool,
|
||||||
shared_pool_name=PoolName.INDEXER,
|
sidecar_pool_name=PoolName.INDEXER,
|
||||||
full_layer_mapping=full_layer_mapping,
|
full_layer_mapping=full_layer_mapping,
|
||||||
page_size=cache.page_size,
|
page_size=cache.page_size,
|
||||||
tp_group=params.tp_cache_group,
|
tp_group=params.tp_cache_group,
|
||||||
@@ -533,7 +828,7 @@ def attach_hybrid_pool_to_unified_cache(
|
|||||||
storage_backend=None,
|
storage_backend=None,
|
||||||
use_mla=use_mla,
|
use_mla=use_mla,
|
||||||
override_kv_cache_dim=full_kv_pool.kv_cache_dim,
|
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,
|
full_kv_pool,
|
||||||
kv_host_pool,
|
kv_host_pool,
|
||||||
server_args.hicache_mem_layout,
|
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.full_kv_pool_host = host_pool_group.get_pool(PoolName.KV)
|
||||||
cache.host_pool_group = host_pool_group
|
cache.host_pool_group = host_pool_group
|
||||||
cache.cache_controller = cache_controller
|
cache.cache_controller = cache_controller
|
||||||
# Register the NSA indexer pool as sharing anchor-KV indices so
|
cache.register_sidecar_pool(
|
||||||
# HiCache backup/load emits its PoolTransfer together with KV.
|
SidecarPoolSpec(
|
||||||
cache.register_hicache_anchor_kv_shared_indices_pool(
|
pool_name=PoolName.INDEXER,
|
||||||
PoolName.INDEXER,
|
indices_from_pool=PoolName.KV,
|
||||||
hit_policy=PoolHitPolicy.ALL_PAGES,
|
)
|
||||||
)
|
)
|
||||||
cache.components[ComponentType.FULL]._full_kv_pool_host = (
|
cache.components[ComponentType.FULL]._full_kv_pool_host = (
|
||||||
cache.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
|
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"
|
pools_desc = "KV + MAMBA"
|
||||||
elif swa_stack:
|
elif swa_stack:
|
||||||
pools_desc = "KV + SWA"
|
pools_desc = "KV + SWA"
|
||||||
@@ -613,8 +910,8 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
|
|||||||
prefetch_threshold: int,
|
prefetch_threshold: int,
|
||||||
enable_storage_metrics: bool,
|
enable_storage_metrics: bool,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache.
|
"""Attach HostPoolGroup (KV + indexer) + HybridCacheController for HiRadixCache.
|
||||||
|
|
||||||
@@ -623,11 +920,11 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
|
|||||||
try:
|
try:
|
||||||
kv = radix_cache.kv_cache
|
kv = radix_cache.kv_cache
|
||||||
layer_mapping = {layer_id: layer_id for layer_id in range(kv.layer_num)}
|
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,
|
params=params,
|
||||||
server_args=server_args,
|
server_args=server_args,
|
||||||
kv_pool=kv,
|
kv_pool=kv,
|
||||||
shared_pool_name=PoolName.INDEXER,
|
sidecar_pool_name=PoolName.INDEXER,
|
||||||
full_layer_mapping=layer_mapping,
|
full_layer_mapping=layer_mapping,
|
||||||
page_size=radix_cache.page_size,
|
page_size=radix_cache.page_size,
|
||||||
tp_group=radix_cache.tp_group,
|
tp_group=radix_cache.tp_group,
|
||||||
@@ -638,7 +935,7 @@ def attach_hybrid_nsa_pool_to_hiradix_cache(
|
|||||||
use_mla=True,
|
use_mla=True,
|
||||||
override_kv_cache_dim=kv.kv_cache_dim,
|
override_kv_cache_dim=kv.kv_cache_dim,
|
||||||
prefetch_threshold=prefetch_threshold,
|
prefetch_threshold=prefetch_threshold,
|
||||||
shared_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost(
|
sidecar_host_pool_factory=lambda kv_host_pool: NSAIndexerPoolHost(
|
||||||
kv,
|
kv,
|
||||||
kv_host_pool,
|
kv_host_pool,
|
||||||
server_args.hicache_mem_layout,
|
server_args.hicache_mem_layout,
|
||||||
@@ -672,8 +969,8 @@ def attach_hybrid_pool_to_mamba_cache(
|
|||||||
prefetch_threshold: int,
|
prefetch_threshold: int,
|
||||||
load_cache_event,
|
load_cache_event,
|
||||||
enable_storage_metrics: bool = False,
|
enable_storage_metrics: bool = False,
|
||||||
attn_cp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_cp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
attn_tp_group: Optional["torch.distributed.ProcessGroup"] = None,
|
attn_tp_group: Optional[torch.distributed.ProcessGroup] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Attach HostPoolGroup (KV + Mamba) + HybridCacheController for HiMambaRadixCache.
|
"""Attach HostPoolGroup (KV + Mamba) + HybridCacheController for HiMambaRadixCache.
|
||||||
|
|
||||||
|
|||||||
@@ -1656,6 +1656,496 @@ class MambaPoolHost(HostKVCache):
|
|||||||
return ptr_list, element_size_list
|
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
|
@dataclass
|
||||||
class PoolEntry:
|
class PoolEntry:
|
||||||
name: PoolName
|
name: PoolName
|
||||||
@@ -1663,9 +2153,6 @@ class PoolEntry:
|
|||||||
device_pool: Any
|
device_pool: Any
|
||||||
layer_mapper: Callable[[int], Optional[int]]
|
layer_mapper: Callable[[int], Optional[int]]
|
||||||
is_primary_index_anchor: bool = False
|
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.
|
# Optional eviction callbacks for auto-alloc in HybridCacheController.
|
||||||
# host_evict_fn(n): evict n slots from the host pool (used by write()).
|
# 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()).
|
# device_evict_fn(n): evict n slots from the device pool (used by load()).
|
||||||
|
|||||||
@@ -22,7 +22,11 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
|||||||
MatchPrefixParams,
|
MatchPrefixParams,
|
||||||
MatchResult,
|
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.radix_cache import RadixKey
|
||||||
from sglang.srt.mem_cache.unified_cache_components import (
|
from sglang.srt.mem_cache.unified_cache_components import (
|
||||||
_NUM_COMPONENT_TYPES,
|
_NUM_COMPONENT_TYPES,
|
||||||
@@ -221,9 +225,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
self._components_tuple: tuple[TreeComponent, ...] = tuple(
|
self._components_tuple: tuple[TreeComponent, ...] = tuple(
|
||||||
self.components.values()
|
self.components.values()
|
||||||
)
|
)
|
||||||
self.hicache_anchor_kv_shared_indices_pools: list[
|
self.sidecar_pool_specs: list[SidecarPoolSpec] = []
|
||||||
tuple[PoolName, PoolHitPolicy]
|
|
||||||
] = []
|
|
||||||
|
|
||||||
# Streaming session: embedded StreamingSession with self as inner.
|
# Streaming session: embedded StreamingSession with self as inner.
|
||||||
# Always on -- zero overhead when no streaming session is open (the
|
# 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.load_cache_event = threading.Event()
|
||||||
self.hicache_anchor_kv_shared_indices_pools.clear()
|
self.sidecar_pool_specs.clear()
|
||||||
attach_hybrid_pool_to_unified_cache(
|
attach_hybrid_pool_to_unified_cache(
|
||||||
self,
|
self,
|
||||||
params,
|
params,
|
||||||
@@ -333,12 +335,8 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
f"transfer_layer_num={self.cache_controller.layer_num}"
|
f"transfer_layer_num={self.cache_controller.layer_num}"
|
||||||
)
|
)
|
||||||
|
|
||||||
def register_hicache_anchor_kv_shared_indices_pool(
|
def register_sidecar_pool(self, spec: SidecarPoolSpec) -> None:
|
||||||
self,
|
self.sidecar_pool_specs.append(spec)
|
||||||
pool_name: PoolName,
|
|
||||||
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES,
|
|
||||||
) -> None:
|
|
||||||
self.hicache_anchor_kv_shared_indices_pools.append((pool_name, hit_policy))
|
|
||||||
|
|
||||||
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
|
||||||
result = self.session.try_match_prefix(params)
|
result = self.session.try_match_prefix(params)
|
||||||
@@ -1164,7 +1162,10 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
):
|
):
|
||||||
return 0
|
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] = {}
|
comp_xfers: dict[ComponentType, list] = {}
|
||||||
for comp in self._components_tuple:
|
for comp in self._components_tuple:
|
||||||
if comp.component_type == BASE_COMPONENT_TYPE:
|
if comp.component_type == BASE_COMPONENT_TYPE:
|
||||||
@@ -1172,13 +1173,11 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST)
|
t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST)
|
||||||
if t:
|
if t:
|
||||||
comp_xfers[comp.component_type] = t
|
comp_xfers[comp.component_type] = t
|
||||||
anchor_kv_shared_indices_xfers = [
|
sidecar_xfers = self._build_sidecar_transfers(
|
||||||
PoolTransfer(name=pool_name, hit_policy=hit_policy)
|
CacheTransferPhase.BACKUP_HOST, kv_xfer, comp_xfers
|
||||||
for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools
|
)
|
||||||
]
|
|
||||||
|
|
||||||
# Pre-evict host if insufficient
|
# Pre-evict host if insufficient
|
||||||
device_value = node.component_data[BASE_COMPONENT_TYPE].value
|
|
||||||
kv_tokens = len(device_value)
|
kv_tokens = len(device_value)
|
||||||
host_avail = self.cache_controller.mem_pool_host.available_size()
|
host_avail = self.cache_controller.mem_pool_host.available_size()
|
||||||
if host_avail < kv_tokens:
|
if host_avail < kv_tokens:
|
||||||
@@ -1188,7 +1187,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
return 0
|
return 0
|
||||||
|
|
||||||
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
|
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(
|
host_indices = self.cache_controller.write(
|
||||||
device_value, node_id=node.id, extra_pools=aux_xfers or None
|
device_value, node_id=node.id, extra_pools=aux_xfers or None
|
||||||
)
|
)
|
||||||
@@ -1245,10 +1244,9 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
)
|
)
|
||||||
if t:
|
if t:
|
||||||
comp_xfers[comp.component_type] = t
|
comp_xfers[comp.component_type] = t
|
||||||
anchor_kv_shared_indices_xfers = [
|
sidecar_xfers = self._build_sidecar_transfers(
|
||||||
PoolTransfer(name=pool_name, hit_policy=hit_policy)
|
CacheTransferPhase.LOAD_BACK, kv_xfer, comp_xfers
|
||||||
for pool_name, hit_policy in self.hicache_anchor_kv_shared_indices_pools
|
)
|
||||||
]
|
|
||||||
|
|
||||||
# Skip if there is nothing to load, or if the Full-KV transfer is too
|
# 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
|
# small / exceeds memory quota. Aux transfers should still run even
|
||||||
@@ -1269,7 +1267,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
# Load H→D
|
# Load H→D
|
||||||
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
|
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(
|
device_indices = self.cache_controller.load(
|
||||||
host_indices=kv_xfer.host_indices,
|
host_indices=kv_xfer.host_indices,
|
||||||
node_id=best_match_node.id,
|
node_id=best_match_node.id,
|
||||||
@@ -1301,6 +1299,52 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
)
|
)
|
||||||
return device_indices
|
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:
|
def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None:
|
||||||
"""Increment hit count; trigger write_backup when threshold reached."""
|
"""Increment hit count; trigger write_backup when threshold reached."""
|
||||||
if self.cache_controller is None:
|
if self.cache_controller is None:
|
||||||
|
|||||||
@@ -104,6 +104,7 @@ def _replay_and_compare_kl(
|
|||||||
output_logprobs: list[list[float]],
|
output_logprobs: list[list[float]],
|
||||||
label: str,
|
label: str,
|
||||||
batch_size: int = 1,
|
batch_size: int = 1,
|
||||||
|
sampling_temperature: float = 1,
|
||||||
):
|
):
|
||||||
"""Flush cache, run replay prefill in batches, compare KL divergence."""
|
"""Flush cache, run replay prefill in batches, compare KL divergence."""
|
||||||
all_input_logprobs = []
|
all_input_logprobs = []
|
||||||
@@ -114,6 +115,7 @@ def _replay_and_compare_kl(
|
|||||||
base_url,
|
base_url,
|
||||||
replay_input_ids[start:end],
|
replay_input_ids[start:end],
|
||||||
output_logprobs[start:end],
|
output_logprobs[start:end],
|
||||||
|
temperature=sampling_temperature,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
acc = {model_name: {"kl_div": kl_threshold}}
|
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
|
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.
|
"""Generate with optional interleaved submission order.
|
||||||
|
|
||||||
Submits inputs reordered by ``order``, then maps results back to the
|
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].
|
inputs[i].
|
||||||
"""
|
"""
|
||||||
if order is None:
|
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]
|
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)
|
unordered = [None] * len(results)
|
||||||
for idx, orig in enumerate(order):
|
for idx, orig in enumerate(order):
|
||||||
unordered[orig] = results[idx]
|
unordered[orig] = results[idx]
|
||||||
@@ -178,6 +194,7 @@ def test_input_output_logprobs_match_helper(
|
|||||||
# --- Cache assertion (for turns > 0) ---
|
# --- Cache assertion (for turns > 0) ---
|
||||||
assert_decode_cached_tokens: Callable | None = None,
|
assert_decode_cached_tokens: Callable | None = None,
|
||||||
replay_batch_size: int = 1,
|
replay_batch_size: int = 1,
|
||||||
|
sampling_temperature: float = 1,
|
||||||
):
|
):
|
||||||
"""Verify decode logprobs match prefill replay.
|
"""Verify decode logprobs match prefill replay.
|
||||||
|
|
||||||
@@ -213,7 +230,11 @@ def test_input_output_logprobs_match_helper(
|
|||||||
]
|
]
|
||||||
|
|
||||||
results = _generate(
|
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
|
assert len(results) == n
|
||||||
|
|
||||||
@@ -242,6 +263,7 @@ def test_input_output_logprobs_match_helper(
|
|||||||
output_lps,
|
output_lps,
|
||||||
label=label,
|
label=label,
|
||||||
batch_size=replay_batch_size,
|
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 ---
|
# --- Interleaving for branch stress ---
|
||||||
branches_per_group: int = 0,
|
branches_per_group: int = 0,
|
||||||
replay_batch_size: int = 1,
|
replay_batch_size: int = 1,
|
||||||
|
sampling_temperature: float = 1,
|
||||||
):
|
):
|
||||||
"""Verify logprobs when prefill cache is hit.
|
"""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
|
# Seed cache with prefixes
|
||||||
_flush_cache(base_url)
|
_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)
|
# 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
|
assert len(results) == n
|
||||||
|
|
||||||
for i, result in enumerate(results):
|
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)
|
current_input[i] + last_outputs[i] + suffixes[i] for i in range(n)
|
||||||
]
|
]
|
||||||
results = _generate_maybe_interleaved(
|
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
|
assert len(results) == n
|
||||||
|
|
||||||
@@ -359,6 +397,7 @@ def test_input_output_logprobs_match_prefill_cache_hit_helper(
|
|||||||
output_lps,
|
output_lps,
|
||||||
label=label,
|
label=label,
|
||||||
batch_size=replay_batch_size,
|
batch_size=replay_batch_size,
|
||||||
|
sampling_temperature=sampling_temperature,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -383,6 +422,7 @@ def test_input_output_logprobs_match_decode_cache_hit_helper(
|
|||||||
# --- Interleaving ---
|
# --- Interleaving ---
|
||||||
branches_per_group: int = 0,
|
branches_per_group: int = 0,
|
||||||
replay_batch_size: int = 1,
|
replay_batch_size: int = 1,
|
||||||
|
sampling_temperature: float = 1,
|
||||||
):
|
):
|
||||||
"""Verify logprobs when decode cache is hit.
|
"""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
|
# Turn 1: populate cache, no assertion, no interleaving
|
||||||
_flush_cache(base_url)
|
_flush_cache(base_url)
|
||||||
results = _generate(
|
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
|
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)
|
current_input[i] + last_outputs[i] + suffixes[i] for i in range(n)
|
||||||
]
|
]
|
||||||
results = _generate_maybe_interleaved(
|
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
|
assert len(results) == n
|
||||||
|
|
||||||
@@ -457,4 +505,5 @@ def test_input_output_logprobs_match_decode_cache_hit_helper(
|
|||||||
output_lps,
|
output_lps,
|
||||||
label=label,
|
label=label,
|
||||||
batch_size=replay_batch_size,
|
batch_size=replay_batch_size,
|
||||||
|
sampling_temperature=sampling_temperature,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -128,13 +128,18 @@ def _flush_cache(base_url, timeout_s=30):
|
|||||||
|
|
||||||
|
|
||||||
def _generate(
|
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."""
|
"""Send generate request and return results."""
|
||||||
json_data = {
|
json_data = {
|
||||||
"input_ids": input_ids,
|
"input_ids": input_ids,
|
||||||
"sampling_params": {
|
"sampling_params": {
|
||||||
"temperature": 1,
|
"temperature": temperature,
|
||||||
"max_new_tokens": max_new_tokens,
|
"max_new_tokens": max_new_tokens,
|
||||||
"ignore_eos": True,
|
"ignore_eos": True,
|
||||||
},
|
},
|
||||||
@@ -151,7 +156,7 @@ def _generate(
|
|||||||
return response.json()
|
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."""
|
"""Run prefill to get input logprobs matching output logprobs."""
|
||||||
_flush_cache(base_url)
|
_flush_cache(base_url)
|
||||||
results = _generate(
|
results = _generate(
|
||||||
@@ -160,6 +165,7 @@ def _get_input_logprobs(base_url, new_input_ids, output_logprobs):
|
|||||||
max_new_tokens=0,
|
max_new_tokens=0,
|
||||||
return_logprob=True,
|
return_logprob=True,
|
||||||
logprob_start_len=0,
|
logprob_start_len=0,
|
||||||
|
temperature=temperature,
|
||||||
)
|
)
|
||||||
assert len(results) == len(new_input_ids)
|
assert len(results) == len(new_input_ids)
|
||||||
|
|
||||||
|
|||||||
@@ -48,6 +48,7 @@ class UnifiedRadixTreeTestMixin:
|
|||||||
prefix_len: int = 512
|
prefix_len: int = 512
|
||||||
prefill_cache_assert = None
|
prefill_cache_assert = None
|
||||||
decode_cache_assert = None
|
decode_cache_assert = None
|
||||||
|
sampling_temperature: float = 1
|
||||||
|
|
||||||
gsm8k_threshold: float = 0.93
|
gsm8k_threshold: float = 0.93
|
||||||
mmlu_threshold: float = 0.8
|
mmlu_threshold: float = 0.8
|
||||||
@@ -106,6 +107,7 @@ class UnifiedRadixTreeTestMixin:
|
|||||||
turn_suffixes=[t2, t3],
|
turn_suffixes=[t2, t3],
|
||||||
assert_decode_cached_tokens=self.decode_cache_assert,
|
assert_decode_cached_tokens=self.decode_cache_assert,
|
||||||
max_new_tokens=self.max_new_tokens,
|
max_new_tokens=self.max_new_tokens,
|
||||||
|
sampling_temperature=self.sampling_temperature,
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_multiturn_prefill_cache_hit_branching(self):
|
def test_multiturn_prefill_cache_hit_branching(self):
|
||||||
@@ -135,6 +137,7 @@ class UnifiedRadixTreeTestMixin:
|
|||||||
assert_decode_cached_tokens=self.decode_cache_assert,
|
assert_decode_cached_tokens=self.decode_cache_assert,
|
||||||
branches_per_group=branches,
|
branches_per_group=branches,
|
||||||
max_new_tokens=self.max_new_tokens,
|
max_new_tokens=self.max_new_tokens,
|
||||||
|
sampling_temperature=self.sampling_temperature,
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_multiturn_decode_cache_hit_branching(self):
|
def test_multiturn_decode_cache_hit_branching(self):
|
||||||
@@ -159,6 +162,7 @@ class UnifiedRadixTreeTestMixin:
|
|||||||
assert_decode_cached_tokens=self.decode_cache_assert,
|
assert_decode_cached_tokens=self.decode_cache_assert,
|
||||||
branches_per_group=branches,
|
branches_per_group=branches,
|
||||||
max_new_tokens=self.max_new_tokens,
|
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)
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user