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

Share CP HiCache host budget across target and draft KV

Browse Source 
CP HiCache previously let the target host pool and the draft/MTP host pool each consume the full --hicache-size budget. With EAGLE/MTP enabled this doubled per-rank host allocation and could kill scheduler ranks during startup before Python emitted a traceback. The cache now treats target KV and draft KV as one logical host-cache object: target and draft capacities are computed from one per-rank byte budget, draft may receive more token capacity when its per-token footprint is smaller, and draft attachment remains tied to target residency.

Constraint: --hicache-size is a per-rank host budget and must not be multiplied by attaching draft KV.

Rejected: Give draft another independent --hicache-size allocation | repeats the observed host OOM failure mode.

Rejected: Disable draft HiCache attachment under CP | avoids OOM but breaks target/draft cache-hit consistency for MTP.

Confidence: medium

Scope-risk: moderate

Directive: Keep target and draft KV as one logical HiCache object; do not let draft host allocation consume an independent full hicache-size budget.

Tested: python -m py_compile on modified scheduler/cache/test files

Tested: remote g0034 container PYTHONPATH=python python -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q (45 passed)

Not-tested: full multi-rank GLM5 server restart after clearing existing remote router/defunct process state

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-27 04:28:50 +08:00
parent d14c02b0dc
commit 8571fe0cd9
6 changed files with 1365 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/managers/scheduler.py
View File

@@ -613,6 +613,17 @@ class Scheduler(
DraftWorkerClass = self.spec_algorithm.create_worker(self.server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)
def _get_draft_token_to_kv_pool_for_hicache(self):
if self.draft_worker is None or self.spec_algorithm.is_ngram():
return None
if self.spec_algorithm.supports_spec_v2() and self.enable_overlap:
if self.server_args.enable_multi_layer_eagle:
draft_runner = self.draft_worker.draft_worker.draft_runner_list[0]
else:
draft_runner = self.draft_worker.draft_worker.draft_runner
return draft_runner.token_to_kv_pool
return self.draft_worker.model_runner.token_to_kv_pool
def init_model_worker(self):
self.init_tp_model_worker()
self.maybe_init_draft_worker()
@@ -727,6 +738,7 @@ class Scheduler(
pp_size=self.pp_size,
chunked_prefill_size=server_args.chunked_prefill_size,
sliding_window_size=self.sliding_window_size,
draft_token_to_kv_pool=self._get_draft_token_to_kv_pool_for_hicache(),
)
if (

3
python/sglang/srt/mem_cache/cache_init_params.py
View File

@@ -1,7 +1,7 @@
from __future__ import annotations
import dataclasses
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Any, Optional
import torch
@@ -33,6 +33,7 @@ class CacheInitParams:
chunked_prefill_size: Optional[int] = None
sliding_window_size: Optional[int] = None
draft_token_to_kv_pool: Optional[Any] = None
# Time-to-live for cache entries in seconds. If None, TTL is disabled.
cache_ttl_seconds: Optional[float] = None

132
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -53,6 +53,70 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _estimate_hicache_size_per_token(kv_cache) -> int:
dtype_size = kv_cache.store_dtype.itemsize
if isinstance(kv_cache, MHATokenToKVPool):
return kv_cache.head_dim * kv_cache.head_num * kv_cache.layer_num * dtype_size * 2
if isinstance(kv_cache, NSATokenToKVPool):
indexer_size_per_token = (
kv_cache.index_head_dim
+ kv_cache.index_head_dim // kv_cache.quant_block_size * 4
)
return (
kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num
+ indexer_size_per_token
* kv_cache.layer_num
* NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize
)
if isinstance(kv_cache, MLATokenToKVPool):
kv_cache_dim = kv_cache.kv_lora_rank + kv_cache.qk_rope_head_dim
return kv_cache_dim * dtype_size * kv_cache.layer_num
raise ValueError(f"Unsupported HiCache KV pool type: {type(kv_cache).__name__}")
def _compute_shared_hicache_token_capacities(
*,
total_host_bytes: int,
target_size_per_token: int,
draft_size_per_token: int,
page_size: int,
) -> Tuple[int, int]:
if total_host_bytes <= 0:
raise ValueError(f"total_host_bytes must be positive, got {total_host_bytes}")
if target_size_per_token <= 0:
raise ValueError(
f"target_size_per_token must be positive, got {target_size_per_token}"
)
if draft_size_per_token <= 0:
raise ValueError(
f"draft_size_per_token must be positive, got {draft_size_per_token}"
)
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
target_pages = total_host_bytes // (
(target_size_per_token + draft_size_per_token) * page_size
)
if target_pages <= 0:
raise ValueError(
"HiCache shared target+draft budget cannot hold a single page: "
f"total_host_bytes={total_host_bytes}, "
f"target_size_per_token={target_size_per_token}, "
f"draft_size_per_token={draft_size_per_token}, page_size={page_size}"
)
target_tokens = int(target_pages * page_size)
remaining_bytes = total_host_bytes - target_tokens * target_size_per_token
draft_pages = remaining_bytes // (draft_size_per_token * page_size)
draft_tokens = int(draft_pages * page_size)
if draft_tokens < target_tokens:
raise ValueError(
"HiCache shared budget computation produced draft capacity below target "
f"capacity: target_tokens={target_tokens}, draft_tokens={draft_tokens}"
)
return target_tokens, draft_tokens
@dataclass
class CpHiCacheNodeMetadata:
logical_len: int
@@ -170,43 +234,56 @@ class CpHiCacheNodeMetadata:
class HiRadixCache(RadixCache):
def _create_token_to_kv_pool_host(self, kv_cache, server_args: ServerArgs):
def _create_token_to_kv_pool_host(
self,
kv_cache,
server_args: ServerArgs,
*,
host_size: Optional[int] = None,
host_token_capacity: Optional[int] = None,
):
from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
)
host_size = server_args.hicache_size if host_size is None else host_size
if isinstance(kv_cache, MHATokenToKVPool):
return MHATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
server_args.hicache_size,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
)
if isinstance(kv_cache, NSATokenToKVPool):
return NSATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
server_args.hicache_size,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
)
if isinstance(kv_cache, MLATokenToKVPool):
return MLATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
server_args.hicache_size,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
)
raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools")
def attach_draft_kv_pool(self, draft_token_to_kv_pool) -> None:
def attach_draft_kv_pool(
self, draft_token_to_kv_pool, host_token_capacity: Optional[int] = None
) -> None:
if not self._uses_cp_hicache or draft_token_to_kv_pool is None:
return
if (
@@ -215,7 +292,10 @@ class HiRadixCache(RadixCache):
):
return
draft_token_to_kv_pool_host = self._create_token_to_kv_pool_host(
draft_token_to_kv_pool, self._server_args
draft_token_to_kv_pool,
self._server_args,
host_size=0 if host_token_capacity is not None else None,
host_token_capacity=host_token_capacity,
)
self.cache_controller.attach_draft_pool(
draft_token_to_kv_pool, draft_token_to_kv_pool_host
@@ -252,8 +332,41 @@ class HiRadixCache(RadixCache):
"CP shared KV HiCache host integration requires NSATokenToKVPool."
)
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
draft_token_to_kv_pool = getattr(params, "draft_token_to_kv_pool", None)
target_host_token_capacity = None
draft_host_token_capacity = None
if (
self._uses_cp_hicache
and draft_token_to_kv_pool is not None
and server_args.hicache_size > 0
):
target_size_per_token = _estimate_hicache_size_per_token(self.kv_cache)
draft_size_per_token = _estimate_hicache_size_per_token(
draft_token_to_kv_pool
)
target_host_token_capacity, draft_host_token_capacity = (
_compute_shared_hicache_token_capacities(
total_host_bytes=int(server_args.hicache_size * 1e9),
target_size_per_token=target_size_per_token,
draft_size_per_token=draft_size_per_token,
page_size=self.page_size,
)
)
logger.info(
"[HiCache-draft] sharing --hicache-size across target+draft: "
"budget_gb=%d target_size_per_token=%d draft_size_per_token=%d "
"target_tokens=%d draft_tokens=%d",
server_args.hicache_size,
target_size_per_token,
draft_size_per_token,
target_host_token_capacity,
draft_host_token_capacity,
)
self.token_to_kv_pool_host = self._create_token_to_kv_pool_host(
self.kv_cache, server_args
self.kv_cache,
server_args,
host_size=0 if target_host_token_capacity is not None else None,
host_token_capacity=target_host_token_capacity,
)
self.tp_group = params.tp_cache_group
@@ -302,6 +415,11 @@ class HiRadixCache(RadixCache):
enable_storage_metrics=self.enable_storage_metrics,
cp_shared_kv_layout=cp_shared_kv_layout,
)
if draft_token_to_kv_pool is not None:
self.attach_draft_kv_pool(
draft_token_to_kv_pool,
host_token_capacity=draft_host_token_capacity,
)
self._apply_storage_runtime_config(
storage_backend=server_args.hicache_storage_backend,
prefetch_threshold=prefetch_threshold,

21
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -147,6 +147,7 @@ class HostKVCache(abc.ABC):
pin_memory: bool,
device: str,
allocator_type: str = "default",
host_token_capacity: Optional[int] = None,
):
self.device_pool = device_pool
self.page_size = page_size
@@ -157,12 +158,20 @@ class HostKVCache(abc.ABC):
self.dtype = device_pool.store_dtype
self.size_per_token = self.get_size_per_token()
if host_size > 0:
if host_token_capacity is not None:
self.size = int(host_token_capacity)
elif host_size > 0:
self.size = int(host_size * 1e9 // self.size_per_token)
else:
self.size = int(device_pool.size * host_to_device_ratio)
# Align up the host memory pool size to the page size
self.page_num = self.size // self.page_size + 1
# Align up the host memory pool size to the page size. Preserve the
# historical extra page for ratio/byte-budget paths, but treat an
# explicit token capacity as already being the logical budget chosen by
# the caller.
if host_token_capacity is not None:
self.page_num = (self.size + self.page_size - 1) // self.page_size
else:
self.page_num = self.size // self.page_size + 1
self.size = self.page_num * self.page_size
self.start_layer = device_pool.start_layer
self.end_layer = device_pool.end_layer
@@ -286,6 +295,7 @@ class MHATokenToKVPoolHost(HostKVCache):
pin_memory: bool = True,
device: str = "cpu",
allocator_type: str = "default",
host_token_capacity: Optional[int] = None,
):
super().__init__(
device_pool,
@@ -296,6 +306,7 @@ class MHATokenToKVPoolHost(HostKVCache):
pin_memory,
device,
allocator_type,
host_token_capacity=host_token_capacity,
)
self.element_dim = self.device_pool.head_num * self.device_pool.head_dim
self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel(
@@ -747,6 +758,7 @@ class MLATokenToKVPoolHost(HostKVCache):
device: str = "cpu",
allocator_type: str = "default",
override_kv_cache_dim: Optional[int] = None,
host_token_capacity: Optional[int] = None,
):
self.override_kv_cache_dim = override_kv_cache_dim
super().__init__(
@@ -758,6 +770,7 @@ class MLATokenToKVPoolHost(HostKVCache):
pin_memory,
device,
allocator_type,
host_token_capacity=host_token_capacity,
)
self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)]
self.data_ptrs = torch.tensor(
@@ -1088,6 +1101,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
pin_memory: bool = True,
device: str = "cpu",
allocator_type: str = "default",
host_token_capacity: Optional[int] = None,
):
# Initialize indexer metadata before HostKVCache.__init__ calls get_size_per_token.
self.index_head_dim = device_pool.index_head_dim
@@ -1107,6 +1121,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
device,
allocator_type,
override_kv_cache_dim=device_pool.kv_cache_dim,
host_token_capacity=host_token_capacity,
)
self.indexer_page_stride_size = (
self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize