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

[HiCache][HybridModel]: Support mamba state offloading & HybridCacheController (#20457)

Browse Source 
Co-authored-by: pansicheng <sicheng.pan.chn@gmail.com>
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
hzh0425
2026-03-23 20:02:50 -07:00
committed by GitHub
co-authored by pansicheng 晟海 ispobock
parent 2b1d3c935e
commit 0986bed8e2
13 changed files with 2012 additions and 220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+3
View File
@@ -195,6 +195,9 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool):
effective_mamba_size = (
mamba_size if mamba_size is not None else size
) + pre_alloc_size
# TODO: Support PP
self.start_layer = 0
self.layer_transfer_counter = None
self._init_mamba_pool(
size=effective_mamba_size,
mamba_spec_state_size=size + pre_alloc_size,
python/sglang/srt/managers/schedule_batch.py
-3
View File
@@ -641,7 +641,6 @@ class Req(ReqDllmMixin):
self.extend_logprob_start_len = 0
self.last_node: Any = None
self.last_host_node: Any = None
self.last_host_backup_node: Any = None
self.host_hit_length = 0
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
@@ -915,14 +914,12 @@ class Req(ReqDllmMixin):
self.prefix_indices,
self.last_node,
self.last_host_node,
self.last_host_backup_node,
self.host_hit_length,
self.mamba_branching_seqlen,
) = (
match_result.device_indices,
match_result.last_device_node,
match_result.last_host_node,
match_result.last_host_backup_node,
match_result.host_hit_length,
match_result.mamba_branching_seqlen,
)
python/sglang/srt/managers/schedule_policy.py
+1
View File
@@ -777,6 +777,7 @@ class PrefillAdder:
InitLoadBackParams(
last_host_node=req.last_host_node,
host_hit_length=req.host_hit_length,
req=req,
)
)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])
python/sglang/srt/managers/scheduler.py
+1 -5
View File
@@ -1818,11 +1818,7 @@ class Scheduler(
def _prefetch_kvcache(self, req: Req):
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache, cow_mamba=False)
last_host_node = (
req.last_host_backup_node
if req.last_host_backup_node is not None
else req.last_host_node
)
last_host_node = req.last_host_node
if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
last_hash = last_host_node.get_last_hash_value()
matched_len = len(req.prefix_indices) + req.host_hit_length
python/sglang/srt/mem_cache/base_prefix_cache.py
+4 -3
View File
@@ -129,8 +129,10 @@ class MatchResult(NamedTuple):
last_host_node : The last TreeNode on the host that was matched.
Note that if HiCache is not enabled,
this **must** be the same as `last_device_node`.
last_host_backup_node: The deepest backuped node for prefetch from storage.
host_hit_length : Length of the KV cache hit on the host, if applicable.
host_hit_length : Length of the host cache hit. For pure-KV caches this is the
number of evicted KV tokens on CPU. For hybrid Mamba models this
is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
host hit still triggers load-back without adding a separate field.
0 if HiCache is not enabled.
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
page-aligned position that could've been cache hit if there
@@ -140,7 +142,6 @@ class MatchResult(NamedTuple):
device_indices: torch.Tensor
last_device_node: Any
last_host_node: Any
last_host_backup_node: Any = None
host_hit_length: int = 0
mamba_branching_seqlen: Optional[int] = None
cache_protected_len: Optional[int] = None
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+636 -196
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/mem_cache/hicache_storage.py
+256 -3
View File
@@ -3,7 +3,8 @@ import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, List, Optional
from enum import Enum
from typing import Any, List, Optional, Set
import torch
@@ -65,6 +66,61 @@ class HiCacheStorageExtraInfo:
extra_info: Optional[dict] = None
class PoolName(str, Enum):
"""Well-known pool names used as PoolTransfer/PoolEntry identifiers."""
KV = "kv"
MAMBA = "mamba"
class PoolHitPolicy(str, Enum):
"""Hit policy for batch_exists_v2 per-pool prefix matching.
ALL_PAGES : every page in [0, kv_hit) must exist (default).
TRAILING_PAGES : only the last N pages must exist (e.g. Mamba/SWA states).
"""
ALL_PAGES = "all_pages"
TRAILING_PAGES = "trailing_pages"
@dataclass
class PoolTransfer:
"""Unified per-pool transfer descriptor for batch v2 interface.
device<->host path : host_indices + device_indices
host<->storage path: host_indices + keys
"""
name: PoolName
host_indices: Optional[torch.Tensor] = None
device_indices: Optional[torch.Tensor] = None
keys: Optional[List[str]] = None
hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES
@dataclass
class PoolTransferResult:
"""Tracks how many pages were successfully processed per pool."""
kv_hit_pages: int
extra_pool_hit_pages: dict[str, int]
@classmethod
def empty(cls) -> "PoolTransferResult":
return cls(0, {})
def update_kv_hit_pages(self, kv_hit_pages: int) -> None:
"""Accumulate kv_hit_pages across batches (max = last successful batch)."""
self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages)
def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None:
"""Record actual load/write success counts per extra pool."""
self.extra_pool_hit_pages.update(
{name: sum(rs) for name, rs in results.items()}
)
class HiCacheStorage(ABC):
"""
HiCacheStorage is a class that provides a generic key-value interface for storing and retrieving KV cache.
@@ -72,10 +128,69 @@ class HiCacheStorage(ABC):
"""
# todo, the page size of storage backend does not have to be the same as the same as host memory pool
def register_mem_pool_host(self, mem_pool_host: HostKVCache):
self.mem_pool_host = mem_pool_host
def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name):
if not hasattr(self, "registered_pools"):
self.registered_pools = {}
self.registered_pools[host_pool_name] = host_pool
def batch_exists_v2(
self,
keys: List[str],
pool_transfers: Optional[List[PoolTransfer]] = None,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> PoolTransferResult:
"""Check which cache pages exist in storage, respecting per-pool hit policies.
Longest-prefix semantics
Extra-pool hit policies (``PoolTransfer.hit_policy``)
------------------------------------------------------
Each ``PoolTransfer`` in ``pool_transfers`` describes a secondary
cache pool (e.g. Mamba SSM states) that must be co-present with the
KV pages. The final ``final_pages`` is the minimum across all pools,
so a missing auxiliary page shrinks the usable prefix.
- ``"all_pages"`` (default): every page in [0, kv_hit) must exist
for this pool. Used for pools that are required for every token
in the prefix (e.g. DeepSeek DSA pool).
- ``"trailing_pages"``: only the *last* ``len(transfer.keys)`` pages
of the KV prefix need to exist. Used for pools whose data covers
only the tail of a prefix (e.g. Mamba/SWA Pool).
Returns
-------
PoolTransferResult
``kv_hit_pages`` = length of the usable KV prefix.
``extra_pool_hit_pages`` maps each pool name to the number of pages
that were found.
"""
raise NotImplementedError()
def batch_get_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional["HiCacheStorageExtraInfo"] = None,
) -> dict[str, List[bool]]:
"""Read data from storage into host memory for each PoolTransfer.
Returns a dict mapping pool name to a per-entry success list.
"""
raise NotImplementedError()
def batch_set_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional["HiCacheStorageExtraInfo"] = None,
) -> dict[str, List[bool]]:
"""Write data from host memory to storage for each PoolTransfer.
Returns a dict mapping pool name to a per-entry success list.
"""
raise NotImplementedError()
def batch_get_v1(
self,
keys: List[str],
@@ -203,7 +318,6 @@ class HiCacheFile(HiCacheStorage):
self.config_suffix = f"_{model_name}"
else:
self.config_suffix = f"_{model_name}_{tp_rank}_{tp_size}"
if not os.path.exists(self.file_path) and tp_rank == 0:
os.makedirs(self.file_path)
logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
@@ -211,6 +325,18 @@ class HiCacheFile(HiCacheStorage):
def _get_suffixed_key(self, key: str) -> str:
return key + self.config_suffix
def _get_component_key(self, key: str, component_name: Optional[str] = None) -> str:
if component_name is None or component_name in ("__default__", PoolName.KV):
return self._get_suffixed_key(key)
return self._get_suffixed_key(f"{key}.{component_name}")
def _get_component_path(
self, key: str, component_name: Optional[str] = None
) -> str:
return os.path.join(
self.file_path, f"{self._get_component_key(key, component_name)}.bin"
)
def get(
self,
key: str,
@@ -280,6 +406,133 @@ class HiCacheFile(HiCacheStorage):
tensor_path = os.path.join(self.file_path, f"{key}.bin")
return os.path.exists(tensor_path)
def _collect_existing_component_keys(
self,
keys: List[str],
pool_transfers: Optional[List[PoolTransfer]] = None,
) -> Set[str]:
target_files = {f"{self._get_component_key(key)}.bin" for key in keys}
for transfer in pool_transfers or []:
for key in keys:
target_files.add(f"{self._get_component_key(key, transfer.name)}.bin")
existing_files = set()
with os.scandir(self.file_path) as entries:
for entry in entries:
if entry.is_file() and entry.name in target_files:
existing_files.add(entry.name)
return existing_files
def batch_exists_v2(
self,
keys: List[str],
pool_transfers: Optional[List[PoolTransfer]] = None,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> PoolTransferResult:
existing_files = self._collect_existing_component_keys(keys, pool_transfers)
def has_component(page_idx: int, name: str) -> bool:
return (
f"{self._get_component_key(keys[page_idx], name)}.bin" in existing_files
)
# Longest contiguous KV prefix present in storage.
kv_pages = next(
(
i
for i in range(len(keys))
if f"{self._get_component_key(keys[i])}.bin" not in existing_files
),
len(keys),
)
hit_count: dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {}
final_pages = kv_pages
for transfer in pool_transfers or []:
if final_pages == 0:
break
name = transfer.name
if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
boundary = next(
(i for i in range(kv_pages) if not has_component(i, name)), kv_pages
)
else: # trailing_pages
trailing = max(1, len(transfer.keys) if transfer.keys else 1)
boundary = 0
for prefix_len in range(kv_pages, 0, -1):
if all(
has_component(i, name)
for i in range(max(0, prefix_len - trailing), prefix_len)
):
boundary = prefix_len
break
if boundary:
hit_count[name] = boundary
final_pages = min(final_pages, boundary)
return PoolTransferResult(final_pages, hit_count)
def _log_key(self, pool_name: str, key: str) -> str:
return key if pool_name == PoolName.KV else f"{key}.{pool_name}"
def _read_page(self, pool_name: str, key: str, host_pool, page_offset: int) -> bool:
"""Read one page from storage into host_pool at page_offset."""
storage_key = self._log_key(pool_name, key)
data_page = self.get(storage_key, host_pool.get_dummy_flat_data_page())
if data_page is None:
return False
host_pool.set_from_flat_data_page(page_offset, data_page)
return True
def _write_page(
self, pool_name: str, key: str, host_pool, page_offset: int
) -> bool:
"""Write one page from host_pool at page_offset to storage as raw bytes."""
storage_key = self._log_key(pool_name, key)
data_page = host_pool.get_data_page(page_offset, flat=True)
return self.set(storage_key, data_page)
def _batch_io_v2(self, transfers: List[PoolTransfer], op_fn):
results: dict[str, List[bool]] = {}
for transfer in transfers:
host_pool = self.registered_pools[transfer.name]
keys = transfer.keys or []
page_size = getattr(host_pool, "page_size", 1) or 1
expected = len(keys) * page_size
host_indices = transfer.host_indices
if host_indices is None or host_indices.numel() != expected:
logger.error(
"%s indices length mismatch for %s: expected %s, got %s",
op_fn.__name__,
transfer.name,
expected,
host_indices.numel() if host_indices is not None else 0,
)
results[transfer.name] = [False] * len(keys)
continue
results[transfer.name] = [
op_fn(transfer.name, key, host_pool, host_indices[i * page_size].item())
for i, key in enumerate(keys)
]
return results
def batch_get_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional["HiCacheStorageExtraInfo"] = None,
) -> dict[str, List[bool]]:
return self._batch_io_v2(transfers, self._read_page)
def batch_set_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional["HiCacheStorageExtraInfo"] = None,
) -> dict[str, List[bool]]:
return self._batch_io_v2(transfers, self._write_page)
def clear(self) -> bool:
try:
for filename in os.listdir(self.file_path):
python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py
+500
View File
@@ -0,0 +1,500 @@
from __future__ import annotations
import logging
import threading
import time
from typing import TYPE_CHECKING, Any, List, Optional
import torch
from sglang.srt.managers.cache_controller import CacheOperation as BaseCacheOperation
from sglang.srt.managers.cache_controller import (
HiCacheAck,
)
from sglang.srt.managers.cache_controller import (
HiCacheController as BaseHiCacheController,
)
from sglang.srt.managers.cache_controller import (
LayerDoneCounter,
)
from sglang.srt.managers.cache_controller import (
StorageOperation as BaseStorageOperation,
)
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageExtraInfo,
PoolHitPolicy,
PoolTransfer,
PoolTransferResult,
)
from sglang.srt.mem_cache.memory_pool_host import PoolEntry
from sglang.srt.utils import get_device_module
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
logger = logging.getLogger(__name__)
device_module = get_device_module()
class CacheOperation(BaseCacheOperation):
def __init__(
self,
host_indices: torch.Tensor,
device_indices: torch.Tensor,
node_id: int,
priority: Optional[int] = None,
pool_transfers: Optional[list[PoolTransfer]] = None,
):
super().__init__(host_indices, device_indices, node_id, priority)
self.pool_transfers = pool_transfers
@staticmethod
def merge_pool_transfers(
ops: List["CacheOperation"],
) -> Optional[list[PoolTransfer]]:
grouped: dict[str, list[PoolTransfer]] = {}
for op in ops:
for t in op.pool_transfers or []:
grouped.setdefault(t.name, []).append(t)
if not grouped:
return None
def cat_or_none(tensors):
parts = [x for x in tensors if x is not None]
return torch.cat(parts) if parts else None
return [
PoolTransfer(
name=name,
host_indices=cat_or_none(t.host_indices for t in ts),
device_indices=cat_or_none(t.device_indices for t in ts),
keys=[k for t in ts if t.keys for k in t.keys] or None,
)
for name, ts in grouped.items()
]
@staticmethod
def merge_ops(ops: List["CacheOperation"]) -> "CacheOperation":
if len(ops) == 1:
return ops[0]
host_indices = torch.cat([op.host_indices for op in ops])
device_indices = torch.cat([op.device_indices for op in ops])
node_ids = []
priority = min(op.priority for op in ops)
for op in ops:
node_ids.extend(op.node_ids)
merged = CacheOperation(
host_indices,
device_indices,
-1,
priority,
pool_transfers=CacheOperation.merge_pool_transfers(ops),
)
merged.node_ids = node_ids
return merged
class StorageOperation(BaseStorageOperation):
def __init__(
self,
host_indices: torch.Tensor,
token_ids: List[int],
last_hash: Optional[str] = None,
hash_value: Optional[List[str]] = None,
prefix_keys: Optional[List[str]] = None,
pool_transfers: Optional[list[PoolTransfer]] = None,
):
super().__init__(host_indices, token_ids, last_hash, hash_value, prefix_keys)
self.pool_transfers = pool_transfers
self.pool_storage_result = PoolTransferResult.empty()
class PrefetchOperation(StorageOperation):
def __init__(
self,
request_id: str,
host_indices: torch.Tensor,
token_ids: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
pool_transfers: Optional[list[PoolTransfer]] = None,
):
self.request_id = request_id
self._lock = threading.Lock()
self._terminated_flag = False
self.start_time = time.monotonic()
super().__init__(
host_indices,
token_ids,
last_hash,
prefix_keys=prefix_keys,
pool_transfers=pool_transfers,
)
def increment(self, num_tokens: int):
with self._lock:
if self._terminated_flag:
return False
self.completed_tokens += num_tokens
return True
def mark_terminate(self):
with self._lock:
self._terminated_flag = True
def is_terminated(self) -> bool:
return self._terminated_flag
class HybridCacheController(BaseHiCacheController):
def __init__(
self,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
mem_pool_host: Any,
page_size: int,
tp_group: torch.distributed.ProcessGroup,
load_cache_event: threading.Event,
write_policy: str = "write_through_selective",
io_backend: str = "",
storage_backend: Optional[str] = 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,
transfer_layer_num: Optional[int] = None,
):
startup_storage_backend = storage_backend
super().__init__(
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
mem_pool_host=mem_pool_host,
page_size=page_size,
tp_group=tp_group,
load_cache_event=load_cache_event,
write_policy=write_policy,
io_backend=io_backend,
storage_backend=None,
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
pp_rank=pp_rank,
pp_size=pp_size,
)
# Override layer_num: hybrid models transfer all layers (For example, Linear Model (KV + Mamba)),
# not just the full attention layers reported by full_kv_pool.
if transfer_layer_num is not None and transfer_layer_num != self.layer_num:
self.layer_num = transfer_layer_num
self.layer_done_counter = LayerDoneCounter(self.layer_num)
if startup_storage_backend is not None:
self.attach_storage_backend(
storage_backend=startup_storage_backend,
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
host_pools=getattr(mem_pool_host, "entries", None),
)
def attach_storage_backend(
self,
storage_backend: str,
prefetch_threshold: int = 256,
model_name: Optional[str] = None,
storage_backend_extra_config: Optional[dict] = None,
host_pools: Optional[list[PoolEntry]] = None,
):
super().attach_storage_backend(
storage_backend=storage_backend,
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
)
for entry in host_pools or []:
self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name)
def reset(self):
super().reset()
if self.enable_storage:
self.host_mem_release_queue.queue.clear()
self.prefetch_tokens_occupied = 0
def write(
self,
device_indices: torch.Tensor,
priority: Optional[int] = None,
node_id: int = -1,
extra_pools: Optional[list[PoolTransfer]] = None,
) -> Optional[torch.Tensor]:
host_indices = self.mem_pool_host.alloc(len(device_indices))
if host_indices is None:
return None
pool_transfers = self._resolve_pool_transfers_allocation(
extra_pools, alloc_host=True
)
if pool_transfers is None and extra_pools:
self.mem_pool_host.free(host_indices)
return None
self.write_queue.append(
CacheOperation(
host_indices,
device_indices,
node_id,
priority,
pool_transfers=pool_transfers or None,
)
)
self.start_writing()
return host_indices
def start_writing(self) -> None:
if not self.write_queue:
return
op = CacheOperation.merge_ops(self.write_queue)
host_indices, device_indices = self.move_indices(op)
self.write_queue.clear()
start_event = device_module.Event()
finish_event = device_module.Event()
start_event.record()
with device_module.stream(self.write_stream):
start_event.wait(self.write_stream)
self.mem_pool_host.backup_from_device_all_layer(
self.mem_pool_device,
host_indices,
device_indices,
self.io_backend,
pool_transfers=op.pool_transfers,
)
finish_event.record()
if host_indices.is_cuda:
host_indices.record_stream(self.write_stream)
if device_indices.is_cuda:
device_indices.record_stream(self.write_stream)
self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids))
def load(
self,
host_indices: torch.Tensor,
priority: Optional[int] = None,
node_id: int = -1,
extra_pools: Optional[list[PoolTransfer]] = None,
) -> Optional[torch.Tensor]:
need_load_kv = host_indices.numel() > 0
if need_load_kv:
device_indices = self.mem_pool_device_allocator.alloc(len(host_indices))
if device_indices is None:
return None
else:
device_indices = torch.empty((0,), dtype=torch.int64, device=self.device)
pool_transfers = self._resolve_pool_transfers_allocation(
extra_pools, alloc_host=False
)
if pool_transfers is None and extra_pools:
if need_load_kv:
self.mem_pool_device_allocator.free(device_indices)
return None
self.load_queue.append(
CacheOperation(
host_indices,
device_indices,
node_id,
priority,
pool_transfers=pool_transfers or None,
)
)
return device_indices
def start_loading(self) -> int:
if not self.load_queue:
return -1
producer_id = self.layer_done_counter.update_producer()
op = CacheOperation.merge_ops(self.load_queue)
host_indices, device_indices = self.move_indices(op)
self.load_queue.clear()
producer_event = self.layer_done_counter.events[producer_id]
producer_event.start_event.record()
with device_module.stream(self.load_stream):
producer_event.start_event.wait(self.load_stream)
for i in range(self.layer_num):
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
device_indices,
i,
self.io_backend,
pool_transfers=op.pool_transfers,
)
producer_event.complete(i)
if host_indices.is_cuda:
host_indices.record_stream(self.load_stream)
if device_indices.is_cuda:
device_indices.record_stream(self.load_stream)
self.ack_load_queue.append(
HiCacheAck(
producer_event.start_event,
producer_event.finish_event,
op.node_ids,
)
)
return producer_id
def prefetch(
self,
request_id: str,
host_indices: torch.Tensor,
new_input_tokens: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
extra_pools: Optional[list[PoolTransfer]] = None,
) -> PrefetchOperation:
operation = PrefetchOperation(
request_id,
host_indices,
new_input_tokens,
last_hash,
prefix_keys=prefix_keys,
pool_transfers=extra_pools,
)
self.prefetch_queue.put(operation)
return operation
def write_storage(
self,
host_indices: torch.Tensor,
token_ids: List[int],
hash_value: Optional[List[str]] = None,
prefix_keys: Optional[List[str]] = None,
extra_pools: Optional[list[PoolTransfer]] = None,
) -> int:
operation = StorageOperation(
host_indices,
token_ids,
hash_value=hash_value,
prefix_keys=prefix_keys,
pool_transfers=extra_pools,
)
self.backup_queue.put(operation)
return operation.id
def _storage_hit_query(self, operation) -> tuple[list[str], int]:
last_hash = operation.last_hash
hash_value = []
for start in range(0, len(operation.token_ids), self.page_size):
last_hash = self.get_hash_str(
operation.token_ids[start : start + self.page_size], last_hash
)
hash_value.append(last_hash)
extra_info = HiCacheStorageExtraInfo(
prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None
)
if operation.pool_transfers:
hit_result = self.storage_backend.batch_exists_v2(
hash_value, operation.pool_transfers, extra_info
)
else:
kv_hit_count = self.storage_backend.batch_exists(hash_value, extra_info)
hit_result = PoolTransferResult(
kv_hit_pages=kv_hit_count, extra_pool_hit_pages={}
)
kv_hit_pages = hit_result.kv_hit_pages
operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages)
if kv_hit_pages > 0 and operation.pool_transfers:
self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages)
return (
hash_value[:kv_hit_pages],
kv_hit_pages * self.page_size,
)
def _page_transfer(self, operation):
# Transfer extra pools
if operation.pool_transfers and not operation.is_terminated():
results = self.storage_backend.batch_get_v2(operation.pool_transfers)
operation.pool_storage_result.update_extra_pool_hit_pages(results)
# Transfer kv pools
super()._page_transfer(operation)
def _page_backup(self, operation):
# Backup extra pools
if operation.pool_transfers:
results = self.storage_backend.batch_set_v2(operation.pool_transfers)
operation.pool_storage_result.update_extra_pool_hit_pages(results)
# Backup kv pools
super()._page_backup(operation)
def _sync_trailing_keys(
self,
pool_transfers: list[PoolTransfer],
all_hashes: list[str],
kv_hit_pages: int,
) -> None:
"""Re-align trailing-page sidecar keys after KV hit truncation.
When the storage hit is shorter than the original target prefix, each
pool transfer's keys must be updated to the last N hashes of the actual
hit range instead of the last N hashes of the original target range.
For mamba (N=1) this is just the last hit page hash; for SWA (N>1) it
is a sliding window of the last N hit pages.
"""
for transfer in pool_transfers:
if transfer.hit_policy != PoolHitPolicy.TRAILING_PAGES:
continue
trailing_n = len(transfer.keys) if transfer.keys else 1
transfer.keys = all_hashes[max(0, kv_hit_pages - trailing_n) : kv_hit_pages]
def _resolve_pool_transfers_allocation(
self,
extra_pools: Optional[list[PoolTransfer]],
alloc_host: bool,
) -> Optional[list[PoolTransfer]]:
"""Auto-alloc host or device indices for PoolTransfers where they are None."""
if not extra_pools:
return None
newly_allocated: list[tuple[PoolTransfer, Any, torch.Tensor]] = []
for pool in extra_pools:
entry = self.mem_pool_host.entry_map.get(pool.name)
if entry is None:
continue
if alloc_host:
if pool.host_indices is not None or pool.device_indices is None:
continue
entry_pool, evict_fn, size = (
entry.host_pool,
entry.host_evict_fn,
len(pool.device_indices),
)
else:
if pool.device_indices is not None or pool.host_indices is None:
continue
entry_pool, evict_fn, size = (
entry.device_pool,
entry.device_evict_fn,
len(pool.host_indices),
)
indices = entry_pool.alloc(size)
if indices is None and evict_fn:
evict_fn(size)
indices = entry_pool.alloc(size)
if indices is None:
# Roll back all previous allocations using each pool's own entry_pool.
for prev_pool, prev_entry_pool, prev_indices in newly_allocated:
prev_entry_pool.free(prev_indices)
if alloc_host:
prev_pool.host_indices = None
else:
prev_pool.device_indices = None
return None
if alloc_host:
pool.host_indices = indices
else:
pool.device_indices = indices
newly_allocated.append((pool, entry_pool, indices))
return extra_pools
python/sglang/srt/mem_cache/mamba_radix_cache.py
+11
View File
@@ -74,6 +74,7 @@ class TreeNode:
self.key: RadixKey = None
self.value: Optional[torch.Tensor] = None
self.mamba_value: Optional[torch.Tensor] = None
self.mamba_host_value: Optional[torch.Tensor] = None
# invariant: for any node, if mamba_lock_ref is locked, full_lock_ref must be locked;
# if full_lock_ref is locked, mamba_lock_ref doesn't need to be locked. So,
# full_lock_ref is always >= mamba_lock_ref.
@@ -98,6 +99,8 @@ class TreeNode:
self.next = None
self.mamba_prev = None
self.mamba_next = None
self.host_mamba_prev = None
self.host_mamba_next = None
self.id = TreeNode.counter if id is None else id
TreeNode.counter += 1
@@ -106,10 +109,18 @@ class TreeNode:
def evicted(self):
return self.value is None
@property
def mamba_evicted(self):
return self.mamba_value is None
@property
def backuped(self):
return self.host_value is not None
@property
def mamba_backuped(self):
return self.mamba_host_value is not None
def protect_host(self):
"""Protect the host value from eviction."""
self.host_ref_counter += 1
python/sglang/srt/mem_cache/memory_pool.py
+27 -2
View File
@@ -475,6 +475,9 @@ class HybridReqToTokenPool(ReqToTokenPool):
self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1
self.enable_mamba_extra_buffer = enable_mamba_extra_buffer
self.enable_memory_saver = enable_memory_saver
# TODO: Support PP
self.start_layer = 0
self.layer_transfer_counter = None
self._init_mamba_pool(
size=mamba_size,
mamba_spec_state_size=mamba_spec_state_size,
@@ -516,6 +519,11 @@ class HybridReqToTokenPool(ReqToTokenPool):
)
)
def register_layer_transfer_counter(
self, layer_transfer_counter: "LayerDoneCounter"
):
self.layer_transfer_counter = layer_transfer_counter
# For chunk prefill req, we do not need to allocate mamba cache,
# We could use allocated mamba cache instead.
def alloc(self, reqs: List["Req"]) -> Optional[List[int]]:
@@ -570,6 +578,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
def mamba2_layer_cache(self, layer_id: int):
assert layer_id in self.mamba_map
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return self.mamba_pool.mamba2_layer_cache(self.mamba_map[layer_id])
def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState:
@@ -1238,8 +1248,8 @@ class HybridLinearKVPool(KVCache):
self.device = device
self.full_layer_nums = len(full_attention_layer_ids)
self.page_size = page_size
# TODO support pp?
self.start_layer = 0
self.start_layer = 0 # TODO: Support PP
self.layer_transfer_counter = None
self.head_num = head_num
self.head_dim = head_dim
self.mamba_pool = mamba_pool
@@ -1323,15 +1333,30 @@ class HybridLinearKVPool(KVCache):
)
return self.full_attention_layer_id_mapping[layer_id]
def register_layer_transfer_counter(
self, layer_transfer_counter: "LayerDoneCounter"
):
self.layer_transfer_counter = layer_transfer_counter
# The layer-wise wait logic is executed at the Hybrid LinearPool level;
# no additional wait is needed in the full_kv_pool
self.full_kv_pool.register_layer_transfer_counter(None)
def _wait_for_layer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
def get_key_buffer(self, layer_id: int):
self._wait_for_layer(layer_id)
layer_id = self._transfer_full_attention_id(layer_id)
return self.full_kv_pool.get_key_buffer(layer_id)
def get_value_buffer(self, layer_id: int):
self._wait_for_layer(layer_id)
layer_id = self._transfer_full_attention_id(layer_id)
return self.full_kv_pool.get_value_buffer(layer_id)
def get_kv_buffer(self, layer_id: int):
self._wait_for_layer(layer_id)
layer_id = self._transfer_full_attention_id(layer_id)
return self.full_kv_pool.get_kv_buffer(layer_id)
python/sglang/srt/mem_cache/memory_pool_host.py
+529 -1
View File
@@ -1,10 +1,17 @@
from __future__ import annotations
import abc
import logging
import threading
from collections import defaultdict
from dataclasses import dataclass
from functools import wraps
from typing import Optional
from typing import TYPE_CHECKING, Any, Callable, Optional
if TYPE_CHECKING:
from sglang.srt.mem_cache.hicache_storage import PoolName
import numpy as np
import psutil
import torch
@@ -19,6 +26,7 @@ from sglang.jit_kernel.hicache import (
)
from sglang.srt.mem_cache.memory_pool import (
KVCache,
MambaPool,
MHATokenToKVPool,
MLATokenToKVPool,
NSATokenToKVPool,
@@ -1074,6 +1082,526 @@ class MLATokenToKVPoolHost(HostKVCache):
return ptr_list, element_size_list
class MambaPoolHost(HostKVCache):
def __init__(
self,
device_pool: MambaPool,
host_to_device_ratio: float,
host_size: int,
pin_memory: bool = True,
device: str = "cpu",
allocator_type: str = "default",
layout: str = "layer_first",
):
self.device_pool = device_pool
self.page_size = 1
assert layout in [
"page_first",
"page_first_direct",
"layer_first",
], "Unsupported layout: {layout}"
self.layout = layout
self.pin_memory = pin_memory
self.device = device
self.allocator = get_allocator_from_storage(allocator_type)
self.num_mamba_layers = device_pool.num_mamba_layers
self.conv_state_shapes = [
conv_state.shape[2:] for conv_state in device_pool.mamba_cache.conv
]
self.temporal_state_shape = device_pool.mamba_cache.temporal.shape[2:]
self.conv_dtype = device_pool.mamba_cache.conv[0].dtype
self.temporal_dtype = device_pool.mamba_cache.temporal.dtype
self.dtype = self.conv_dtype
self.size_per_token = self.get_size_per_token()
if host_size > 0:
self.size = int(host_size * 1e9 // self.size_per_token)
else:
self.size = int(device_pool.size * host_to_device_ratio)
self.page_num = self.size // self.page_size + 1
self.size = self.page_num * self.page_size
assert (
self.size > device_pool.size
), "The host memory should be larger than the device memory with the current protocol"
host_mem = psutil.virtual_memory()
requested_bytes = self.size * self.size_per_token
ten_gb = 10 * (1024**3)
available_bytes = host_mem.available - ten_gb
if requested_bytes > available_bytes:
raise ValueError(
f"Not enough host memory available. Requesting "
f"{requested_bytes / 1e9:.2f} GB but only have "
f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
f"size of the hierarchical cache."
)
logger.info(
"Allocating %.2f GB host memory for hierarchical Mamba cache (layout=%s).",
requested_bytes / 1e9,
self.layout,
)
self.init_kv_buffer()
self.lock = threading.RLock()
self.clear()
def init_kv_buffer(self):
alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
if self.layout in ["page_first", "page_first_direct"]:
# page-first: (page_num, num_layers, 1, *shape) — per-page data is contiguous
temporal_dims = (
self.size,
self.num_mamba_layers,
1,
) + self.temporal_state_shape
self.temporal_buffer = alloc_func(
temporal_dims,
dtype=self.temporal_dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.conv_buffer = []
for conv_shape in self.conv_state_shapes:
conv_dims = (self.size, self.num_mamba_layers, 1) + conv_shape
self.conv_buffer.append(
alloc_func(
conv_dims,
dtype=self.conv_dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
)
else:
# layer-first: (num_layers, size, *shape)
temporal_dims = (
self.num_mamba_layers,
self.size,
) + self.temporal_state_shape
self.temporal_buffer = alloc_func(
temporal_dims,
dtype=self.temporal_dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
self.conv_buffer = []
for conv_shape in self.conv_state_shapes:
conv_dims = (self.num_mamba_layers, self.size) + conv_shape
self.conv_buffer.append(
alloc_func(
conv_dims,
dtype=self.conv_dtype,
device=self.device,
pin_memory=self.pin_memory,
allocator=self.allocator,
)
)
def _iter_page_tensors(self, index: int):
if self.layout in ["page_first", "page_first_direct"]:
yield self.temporal_buffer[index]
for conv_buf in self.conv_buffer:
yield conv_buf[index]
else:
yield self.temporal_buffer[:, index : index + self.page_size]
for conv_buf in self.conv_buffer:
yield conv_buf[:, index : index + self.page_size]
@staticmethod
def _flatten_tensor_bytes(tensor: torch.Tensor) -> torch.Tensor:
return tensor.contiguous().view(torch.uint8).reshape(-1)
@synchronized
def clear(self):
self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device
)
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]:
assert (
need_size % self.page_size == 0
), "The requested size should be a multiple of the 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])
return len(indices)
def get_size_per_token(self):
conv_total_size = 0
for conv_shape in self.conv_state_shapes:
conv_total_size += int(np.prod(conv_shape)) * self.conv_dtype.itemsize
temporal_size = (
int(np.prod(self.temporal_state_shape)) * self.temporal_dtype.itemsize
)
return (conv_total_size + temporal_size) * self.num_mamba_layers
def get_ksize_per_token(self):
return self.get_size_per_token()
@staticmethod
def _item_size_per_index(tensor: torch.Tensor) -> int:
if tensor.shape[0] == 0:
return 0
return int(tensor[0].numel() * tensor.element_size())
@staticmethod
def _copy_tensor(
src: torch.Tensor,
dst: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
io_backend: str,
) -> None:
if src_indices.numel() == 0:
return
if io_backend == "kernel":
# TODO: Rename the interface for clarity.
# Here, transfer_kv_per_layer_mla is reused to transfer the Mamba state.
# This has nothing to do with MLA; it's only reused because this interface happens to transfer a single Pool.
transfer_kv_per_layer_mla(
src=src,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=MambaPoolHost._item_size_per_index(src),
)
elif io_backend == "direct":
transfer_kv_direct(
src_layers=[src],
dst_layers=[dst],
src_indices=src_indices,
dst_indices=dst_indices,
page_size=1,
)
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
@staticmethod
def _copy_tensor_pf_lf(
src: torch.Tensor,
dst: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
layer_id: int,
num_layers: int,
io_backend: str,
) -> None:
if src_indices.numel() == 0:
return
if io_backend == "kernel":
item_size = MambaPoolHost._item_size_per_index(dst)
transfer_kv_per_layer_mla_pf_lf(
src=src,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
layer_id=layer_id,
item_size=item_size,
src_layout_dim=item_size * num_layers,
)
elif io_backend == "direct":
transfer_kv_per_layer_direct_pf_lf(
src_ptrs=[src],
dst_ptrs=[dst],
src_indices=src_indices,
dst_indices=dst_indices,
layer_id=layer_id,
page_size=1,
)
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
@staticmethod
def _copy_tensor_all_layers_lf_pf(
src_layers: torch.Tensor,
dst: torch.Tensor,
src_indices: torch.Tensor,
dst_indices: torch.Tensor,
num_layers: int,
device: str,
io_backend: str,
) -> None:
if src_indices.numel() == 0:
return
if io_backend == "kernel":
item_size = MambaPoolHost._item_size_per_index(src_layers[0])
src_ptrs = torch.tensor(
[src_layers[i].data_ptr() for i in range(num_layers)],
dtype=torch.uint64,
device=device,
)
transfer_kv_all_layer_mla_lf_pf(
src_layers=src_ptrs,
dst=dst,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=item_size,
dst_layout_dim=item_size * num_layers,
num_layers=num_layers,
)
elif io_backend == "direct":
src_ptrs = [src_layers[i] for i in range(num_layers)]
transfer_kv_all_layer_direct_lf_pf(
src_ptrs=src_ptrs,
dst_ptrs=[dst],
src_indices=src_indices,
dst_indices=dst_indices,
page_size=1,
)
else:
raise ValueError(f"Unsupported io_backend: {io_backend}")
def load_to_device_per_layer(
self,
device_pool,
host_indices,
device_indices,
layer_id,
io_backend="kernel",
):
if self.layout in ["page_first", "page_first_direct"]:
self._copy_tensor_pf_lf(
src=self.temporal_buffer,
dst=device_pool.mamba_cache.temporal[layer_id],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
num_layers=self.num_mamba_layers,
io_backend=io_backend,
)
for conv_idx in range(len(self.conv_state_shapes)):
self._copy_tensor_pf_lf(
src=self.conv_buffer[conv_idx],
dst=device_pool.mamba_cache.conv[conv_idx][layer_id],
src_indices=host_indices,
dst_indices=device_indices,
layer_id=layer_id,
num_layers=self.num_mamba_layers,
io_backend=io_backend,
)
else:
self._copy_tensor(
self.temporal_buffer[layer_id],
device_pool.mamba_cache.temporal[layer_id],
host_indices,
device_indices,
io_backend,
)
for conv_idx in range(len(self.conv_state_shapes)):
self._copy_tensor(
self.conv_buffer[conv_idx][layer_id],
device_pool.mamba_cache.conv[conv_idx][layer_id],
host_indices,
device_indices,
io_backend,
)
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend="kernel"
):
if self.layout in ["page_first", "page_first_direct"]:
self._copy_tensor_all_layers_lf_pf(
src_layers=device_pool.mamba_cache.temporal,
dst=self.temporal_buffer,
src_indices=device_indices,
dst_indices=host_indices,
num_layers=self.num_mamba_layers,
device=self.device_pool.device,
io_backend=io_backend,
)
for conv_idx in range(len(self.conv_state_shapes)):
self._copy_tensor_all_layers_lf_pf(
src_layers=device_pool.mamba_cache.conv[conv_idx],
dst=self.conv_buffer[conv_idx],
src_indices=device_indices,
dst_indices=host_indices,
num_layers=self.num_mamba_layers,
device=self.device_pool.device,
io_backend=io_backend,
)
else:
for layer_id in range(self.num_mamba_layers):
self._copy_tensor(
device_pool.mamba_cache.temporal[layer_id],
self.temporal_buffer[layer_id],
device_indices,
host_indices,
io_backend,
)
for conv_idx in range(len(self.conv_state_shapes)):
self._copy_tensor(
device_pool.mamba_cache.conv[conv_idx][layer_id],
self.conv_buffer[conv_idx][layer_id],
device_indices,
host_indices,
io_backend,
)
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
data_page = torch.cat(
[
self._flatten_tensor_bytes(tensor)
for tensor in self._iter_page_tensors(index)
]
)
return data_page.flatten() if flat else data_page
def get_dummy_flat_data_page(self) -> torch.Tensor:
return torch.zeros(
self.page_size * self.size_per_token,
dtype=torch.uint8,
device=self.device,
pin_memory=self.pin_memory,
)
def set_from_flat_data_page(
self,
index: int,
data_page: torch.Tensor,
) -> None:
flat_bytes = data_page.contiguous().view(torch.uint8).reshape(-1)
start = 0
for tensor in self._iter_page_tensors(index):
num_bytes = tensor.numel() * tensor.element_size()
tensor_bytes = flat_bytes[start : start + num_bytes]
start += num_bytes
restored = tensor_bytes.view(dtype=tensor.dtype).reshape(tensor.shape)
tensor.copy_(restored)
@dataclass
class PoolEntry:
name: PoolName
host_pool: Any
device_pool: Any
layer_mapper: Callable[[int], Optional[int]]
is_primary_index_anchor: bool = False
# Optional eviction callbacks for auto-alloc in HybridCacheController.
# host_evict_fn(n): evict n slots from the host pool (used by write()).
# device_evict_fn(n): evict n slots from the device pool (used by load()).
host_evict_fn: Optional[Callable] = None
device_evict_fn: Optional[Callable] = None
class HostPoolGroup:
def __init__(self, entries: list[PoolEntry]):
if not entries:
raise ValueError("HostPoolGroup requires at least one pool entry.")
self.entries = entries
self.entry_map = {entry.name: entry for entry in entries}
self.anchor_entry = next(
(entry for entry in entries if entry.is_primary_index_anchor),
entries[0],
)
self.layout = self.anchor_entry.host_pool.layout
self.page_size = self.anchor_entry.host_pool.page_size
self.device = self.anchor_entry.host_pool.device
self.size = self.anchor_entry.host_pool.size
def clear(self) -> None:
for entry in self.entries:
entry.host_pool.clear()
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
return self.anchor_entry.host_pool.alloc(need_size)
def free(self, indices: torch.Tensor) -> int:
return self.anchor_entry.host_pool.free(indices)
def get_data_page(self, index, flat: bool = True):
return self.anchor_entry.host_pool.get_data_page(index, flat)
def get_dummy_flat_data_page(self):
return self.anchor_entry.host_pool.get_dummy_flat_data_page()
def set_from_flat_data_page(self, index: int, data_page) -> None:
return self.anchor_entry.host_pool.set_from_flat_data_page(index, data_page)
def load_to_device_per_layer(
self,
device_pool,
host_indices,
device_indices,
layer_id,
io_backend,
pool_transfers: Optional[list] = None,
) -> None:
# 1. Anchor (KV) transfer
anchor = self.anchor_entry
local_layer_id = anchor.layer_mapper(layer_id)
if local_layer_id is not None and host_indices.numel() > 0:
anchor.host_pool.load_to_device_per_layer(
anchor.device_pool,
host_indices,
device_indices,
local_layer_id,
io_backend,
)
# 2. Extra pool transfers
for transfer in pool_transfers or []:
entry = self.entry_map.get(transfer.name)
if entry is None or transfer.host_indices is None:
continue
local_layer_id = entry.layer_mapper(layer_id)
if local_layer_id is None:
continue
entry.host_pool.load_to_device_per_layer(
entry.device_pool,
transfer.host_indices,
transfer.device_indices,
local_layer_id,
io_backend,
)
def backup_from_device_all_layer(
self,
device_pool,
host_indices,
device_indices,
io_backend,
pool_transfers: Optional[list] = None,
) -> None:
# 1. Anchor (KV) backup
self.anchor_entry.host_pool.backup_from_device_all_layer(
self.anchor_entry.device_pool,
host_indices,
device_indices,
io_backend,
)
# 2. Extra pool backup
for transfer in pool_transfers or []:
entry = self.entry_map.get(transfer.name)
if entry is None or transfer.host_indices is None:
continue
entry.host_pool.backup_from_device_all_layer(
entry.device_pool,
transfer.host_indices,
transfer.device_indices,
io_backend,
)
class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
device_pool: NSATokenToKVPool