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

[Feat] DCP + HiCache L2 Support (ported from kimi-k3) (#33112)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yuwei An
2026-08-02 23:35:27 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent f8e62a9224
commit 1a3bea77f2
9 changed files with 442 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/metrics_reporter.py
+3 -4
View File
@@ -1008,10 +1008,9 @@ class SchedulerMetricsReporter:
self.scheduler.tree_cache, "token_to_kv_pool_host", None
) or getattr(self.scheduler.tree_cache, "full_kv_pool_host", None)
assert host_pool is not None, "Host pool not found"
self.stats.hicache_host_used_tokens = (
host_pool.size - host_pool.available_size()
)
self.stats.hicache_host_total_tokens = host_pool.size
host_total = host_pool.logical_size
self.stats.hicache_host_used_tokens = host_total - host_pool.available_size()
self.stats.hicache_host_total_tokens = host_total
def _update_lora_metrics(self):
"""Update LoRA pool metrics for monitoring and autoscaling."""
python/sglang/srt/mem_cache/hiradix_cache.py
+5
View File
@@ -99,6 +99,9 @@ class HiRadixCache(RadixCache):
# Filled by attach_hybrid_minimax_sparse_pool_to_hiradix_cache.
self.token_to_kv_pool_host = None
elif isinstance(self.kv_cache, MLATokenToKVPool):
from sglang.srt.runtime_context import get_parallel
_parallel = get_parallel()
self.token_to_kv_pool_host = MLATokenToKVPoolHost(
self.kv_cache,
server_args.hicache_ratio,
@@ -106,6 +109,8 @@ class HiRadixCache(RadixCache):
self.page_size,
server_args.hicache_mem_layout,
allocator_type=allocator_type,
dcp_size=_parallel.attn_dcp_size,
dcp_rank=_parallel.attn_dcp_rank,
)
else:
raise ValueError("HiRadixCache only supports MHA, MLA, DSA, and MSA models")
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+9
View File
@@ -28,6 +28,7 @@ from sglang.srt.mem_cache.pool_host.mha import (
)
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.runtime_context import get_parallel
if TYPE_CHECKING:
import torch
@@ -72,6 +73,14 @@ def build_kv_host_pool(
kwargs = {}
if override_kv_cache_dim is not None:
kwargs["override_kv_cache_dim"] = override_kv_cache_dim
parallel = get_parallel()
if parallel.dcp_enabled:
assert use_mla, (
"HiCache + DCP is only wired for the MLA host pool; the MHA host "
"pool has no DCP index translation."
)
kwargs["dcp_size"] = parallel.attn_dcp_size
kwargs["dcp_rank"] = parallel.attn_dcp_rank
return kv_host_pool_cls(
kv_pool,
server_args.hicache_ratio,
python/sglang/srt/mem_cache/memory_pool_host.py
+3
View File
@@ -650,6 +650,8 @@ class LogicalHostPool:
f"got size={size}, page_size={page_size}"
)
self.size = size
# Stands in for a host pool (and group anchor); DCP never widens it.
self.logical_size = size
self.page_size = page_size
self.device = "cpu"
self.layout = layout
@@ -1528,6 +1530,7 @@ class HostPoolGroup:
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
self.logical_size = self.anchor_entry.host_pool.logical_size
child_write_back_jit = [
getattr(entry.host_pool, "can_use_write_back_jit", False)
for entry in entries
python/sglang/srt/mem_cache/pool_host/base.py
+43 -5
View File
@@ -79,6 +79,8 @@ def synchronized(func):
class HostKVCache(abc.ABC):
dcp_size = 1
dcp_rank = 0
def __init__(
self,
@@ -90,9 +92,19 @@ class HostKVCache(abc.ABC):
pin_memory: bool,
device: str,
allocator_type: str = "default",
dcp_size: int = 1,
dcp_rank: int = 0,
):
self.device_pool = device_pool
self.page_size = page_size
# page_size arrives widened (x dcp_size); size/page_size/page_num are physical.
self.dcp_size = dcp_size
self.dcp_rank = dcp_rank
assert page_size % dcp_size == 0, (
f"HiCache host pool page_size ({page_size}) must be a multiple of "
f"dcp_size ({dcp_size}); expected the widened page from the DCP "
"paged allocator."
)
self.page_size = page_size // dcp_size
self.layout = layout
self.pin_memory = pin_memory
self.device = device
@@ -265,16 +277,16 @@ class HostKVCache(abc.ABC):
def clear(self):
# Initialize memory states and tracking structures.
self.mem_state = torch.zeros(
(self.size,), dtype=torch.uint8, device=self.device
(self.logical_size,), dtype=torch.uint8, device=self.device
)
self.free_slots = torch.arange(self.size, dtype=torch.int64)
self.free_slots = torch.arange(self.logical_size, dtype=torch.int64)
# Keep freed chunks aside and consume them lazily from alloc() to avoid
# concatenating a large free-list on every host-pool free.
self.release_slots = []
self.num_release_slots = 0
# Per-slot flag used to detect double-free.
# slot_used[k] is true if slot k is allocated.
self.slot_used = torch.zeros(self.size, dtype=torch.bool)
self.slot_used = torch.zeros(self.logical_size, dtype=torch.bool)
def available_size(self):
return len(self.free_slots) + self.num_release_slots
@@ -291,10 +303,36 @@ class HostKVCache(abc.ABC):
self.release_slots = []
self.num_release_slots = 0
@property
def logical_size(self) -> int:
"""Slots the radix/controller layer sees: dcp_size of them share a row."""
return self.size * self.dcp_size
@property
def logical_page_size(self) -> int:
"""Page size in that same logical space (the widened DCP page)."""
return self.page_size * self.dcp_size
def dcp_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor:
"""Transfer kernels index per-rank rows; callers hold widened logical slots.
Keep this rank's slots (% dcp_size == dcp_rank), then collapse (// dcp_size).
"""
if self.dcp_size == 1:
return indices
owned = indices[indices % self.dcp_size == self.dcp_rank] // self.dcp_size
assert owned.numel() * self.dcp_size == indices.numel(), (
"HiCache DCP translation expects runs of whole widened pages "
f"(every residue class equally represented); got {indices.numel()} "
f"logical slots -> {owned.numel()} owned rows with dcp_size="
f"{self.dcp_size}."
)
return owned
@synchronized
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
assert (
need_size % self.page_size == 0
need_size % self.logical_page_size == 0
), "The requested size should be a multiple of the page size."
if need_size > self.available_size():
return None
python/sglang/srt/mem_cache/pool_host/mla.py
+14
View File
@@ -62,6 +62,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
device: str = "cpu",
allocator_type: str = "default",
override_kv_cache_dim: Optional[int] = None,
dcp_size: int = 1,
dcp_rank: int = 0,
):
self.override_kv_cache_dim = override_kv_cache_dim
super().__init__(
@@ -73,6 +75,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
pin_memory,
device,
allocator_type,
dcp_size=dcp_size,
dcp_rank=dcp_rank,
)
# The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
# helpers in hicache.cuh are guarded by USE_ROCM and the staged
@@ -226,6 +230,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
):
if not self._is_device_layer_owned(device_pool, layer_id):
return
host_indices = self.dcp_kernel_indices(host_indices)
device_indices = self.dcp_kernel_indices(device_indices)
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
@@ -311,6 +317,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
def _backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
# Indices arrive already translated by backup_from_device_all_layer.
host_layer = self._host_layer_index(layer_id)
if io_backend == "kernel":
if self.layout == "layer_first":
@@ -370,6 +377,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
host_indices = self.dcp_kernel_indices(host_indices)
device_indices = self.dcp_kernel_indices(device_indices)
if self._is_device_layer_sharded(device_pool):
for layer_id in self._owned_device_layer_ids(device_pool):
self._backup_from_device_per_layer(
@@ -459,6 +468,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
raise ValueError(f"Unsupported IO backend: {io_backend}")
def get_data_page(self, index, flat: bool = True) -> torch.Tensor:
assert self.dcp_size == 1, (
"HiCache L3 storage paths are not yet DCP-aware (per-rank shards "
"need dcp_rank-scoped keys); --hicache-storage-backend with "
"--dcp-size > 1 should have been rejected at server start."
)
if self.layout == "layer_first":
data_page = self.kv_buffer[:, index : index + self.page_size, :, :]
elif self.layout == "page_first":
python/sglang/srt/server_args.py
+43
View File
@@ -7120,6 +7120,49 @@ class ServerArgs:
# Step 2: Storage-layout normalization without changing io backend.
self._resolve_storage_layout_compatibility()
# Step 3: DCP compatibility for the L2 (device<->host) path.
self._resolve_hicache_dcp_compatibility()
def _resolve_hicache_dcp_compatibility(self):
if self.dcp_size <= 1 or not self.enable_hierarchical_cache:
return
if self.hicache_storage_backend is not None:
raise NotImplementedError(
"--hicache-storage-backend (L3) with --dcp-size > 1 is not "
"supported yet: under DCP each rank holds a distinct "
"interleaved MLA KV shard, so the rank-0-only replicated-MLA "
"backup and the storage keys must become dcp_rank-aware "
"first. Run HiCache+DCP with L1/L2 only."
)
if self.speculative_algorithm is not None:
raise NotImplementedError(
"HiCache with --dcp-size > 1 does not support speculative "
"decoding yet (the draft-model host pool has no DCP index "
"translation)."
)
if self.enable_lmcache:
raise NotImplementedError(
"--enable-lmcache with --dcp-size > 1 is not supported: "
"LMCache has no DCP-aware index translation."
)
if self.enable_hisparse:
raise NotImplementedError(
"--enable-hisparse with --dcp-size > 1 is not supported: the "
"HiSparse host pool is constructed without DCP translation."
)
if not self.use_mla_backend():
raise NotImplementedError(
"HiCache with --dcp-size > 1 is only supported for MLA models: "
"the index translation lives in MLATokenToKVPoolHost, and the "
"MHA host pool has none."
)
logger.info(
"HiCache + DCP enabled (L1/L2 only): host pool uses widened "
"logical slot accounting with per-rank physical translation at "
"the transfer boundary (dcp_size=%d).",
self.dcp_size,
)
def _resolve_layout_io_compatibility(self):
if (
self.hicache_mem_layout == "page_first_direct"