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"
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py
+112
View File
@@ -0,0 +1,112 @@
"""HiCache L2 under decode context parallelism (DCP) + UnifiedRadixCache.
Under DCP the radix layer allocates widened logical indices while each rank's
buffers hold only its 1/dcp_size shard, so a missing translation makes cache
hits return another rank's KV. The KL cases catch that as a large divergence.
Blackwell-only: the MLA DCP decode path needs ``tokenspeed_mla`` (SM100/12x).
"""
import subprocess
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
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,
)
register_cuda_ci(est_time=1500, stage="extra-b", runner_config="4-gpu-b200")
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
DCP_SIZE = 4
PAGE_SIZE = 64
WIDENED_PAGE = PAGE_SIZE * DCP_SIZE
MAX_MAMBA_CACHE_SIZE = 256
# Bound the host pools directly. Sizing them off the device pool would need
# --max-total-tokens, which triggers out-of-range KV writes under DCP.
HICACHE_SIZE_GB = 10
class TestUnifiedKimiLinearDcpHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Kimi Linear + DCP4 + HiCache L2 + UnifiedRadixCache."""
kl_threshold = 0.01
gsm8k_threshold = 0.85
mmlu_threshold = 0.4
prefill_cache_assert = staticmethod(
make_mamba_prefill_assert(chunk_size=WIDENED_PAGE)
)
decode_cache_assert = staticmethod(
make_mamba_decode_assert(track_interval=WIDENED_PAGE)
)
@classmethod
def setUpClass(cls):
cls.model = KIMI_LINEAR_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
other_args=[
"--trust-remote-code",
"--tp-size",
"4",
"--dcp-size",
str(DCP_SIZE),
"--page-size",
str(PAGE_SIZE),
"--attention-backend",
"tokenspeed_mla",
"--kv-cache-dtype",
"fp8_e4m3",
"--dcp-comm-backend",
"a2a",
"--dcp-replicate-q-proj",
"--dtype",
"bfloat16",
"--random-seed",
"0",
"--cuda-graph-max-bs-decode",
"64",
"--cuda-graph-backend-prefill",
"disabled",
"--mem-fraction-static",
"0.80",
"--enable-hierarchical-cache",
"--hicache-size",
str(HICACHE_SIZE_GB),
"--hicache-write-policy",
"write_through",
"--max-running-requests",
"64",
"--max-mamba-cache-size",
str(MAX_MAMBA_CACHE_SIZE),
"--enable-metrics",
],
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18, trust_remote_code=True)
@classmethod
def tearDownClass(cls):
cls.process.terminate()
try:
cls.process.wait(timeout=60)
except subprocess.TimeoutExpired:
pass
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py
+210
View File
@@ -0,0 +1,210 @@
"""HiCache under decode context parallelism (DCP): host-pool index math.
Under DCP the radix/controller layer works in a widened logical index space
(page_size * dcp_size wide pages, dcp_size * physical capacity), while each
rank's device and host buffers only materialize the owned 1/dcp_size token
shard (owner rule: index % dcp_size == dcp_rank, physical row = index //
dcp_size — the same rule the device-side KV write and page-table kernels
use). These tests cover the translation helper, the logical/physical host
pool sizing, and that the transfer entry points hand *physical* rows to the
kernels.
"""
import unittest
from types import SimpleNamespace
from unittest import mock
import torch
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
DCP_SIZE = 8
PHYSICAL_PAGE = 64
WIDENED_PAGE = PHYSICAL_PAGE * DCP_SIZE
def _fake_mla_device_pool(size: int = 1024) -> SimpleNamespace:
return SimpleNamespace(
size=size,
store_dtype=torch.float16,
kv_lora_rank=8,
qk_rope_head_dim=4,
layer_num=2,
start_layer=0,
end_layer=1,
device="cpu",
layers_to_capture=None,
layer_shard_enabled=False,
)
def _make_host_pool(dcp_rank: int, device_size: int = 1024) -> MLATokenToKVPoolHost:
return MLATokenToKVPoolHost(
_fake_mla_device_pool(device_size),
host_to_device_ratio=2.0,
host_size=0,
page_size=WIDENED_PAGE,
layout="layer_first",
pin_memory=False,
device="cpu",
dcp_size=DCP_SIZE,
dcp_rank=dcp_rank,
)
class TestDcpKernelIndices(CustomTestCase):
def _bare_pool(self, dcp_size: int, dcp_rank: int) -> MLATokenToKVPoolHost:
pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
pool.dcp_size = dcp_size
pool.dcp_rank = dcp_rank
return pool
def test_identity_without_dcp(self):
pool = self._bare_pool(1, 0)
indices = torch.arange(37)
self.assertIs(pool.dcp_kernel_indices(indices), indices)
def test_aligned_page_translates_to_full_physical_page(self):
# One widened page starting at logical 512 covers physical rows
# 64..127 on every rank.
indices = torch.arange(WIDENED_PAGE, 2 * WIDENED_PAGE)
for rank in range(DCP_SIZE):
pool = self._bare_pool(DCP_SIZE, rank)
out = pool.dcp_kernel_indices(indices)
torch.testing.assert_close(
out, torch.arange(PHYSICAL_PAGE, 2 * PHYSICAL_PAGE)
)
def test_matches_owner_rule_on_merged_unordered_pages(self):
# Concatenation of non-adjacent widened pages in arbitrary order, as
# produced by merged CacheOperations after allocator churn.
pages = [3, 0, 5]
indices = torch.cat(
[torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in pages]
)
for rank in range(DCP_SIZE):
pool = self._bare_pool(DCP_SIZE, rank)
out = pool.dcp_kernel_indices(indices)
expected = (
indices[indices % DCP_SIZE == rank] // DCP_SIZE
) # owner rule, same as filter_dcp_local_kv_indices
torch.testing.assert_close(out, expected)
self.assertEqual(out.numel() * DCP_SIZE, indices.numel())
def test_ragged_run_is_rejected(self):
pool = self._bare_pool(DCP_SIZE, 0)
with self.assertRaises(AssertionError):
pool.dcp_kernel_indices(torch.arange(WIDENED_PAGE + 1))
def test_positional_residue_pairing_survives_host_sort(self):
# move_indices (direct/layer_first) sorts host indices and permutes
# device indices to match. Independent residue filtering of both
# tensors must keep the same token positions on every rank.
g = torch.Generator().manual_seed(0)
host_pages = [7, 2]
device_pages = [1, 4]
host = torch.cat(
[torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE) for p in host_pages]
)
device = torch.cat(
[
torch.arange(p * WIDENED_PAGE, (p + 1) * WIDENED_PAGE)
for p in device_pages
]
)
# token identity: position i pairs host[i] <-> device[i]
perm = torch.randperm(host.numel(), generator=g)
# sort host as move_indices does, permuting device alongside
host_sorted, order = host[perm].sort()
device_matched = device[perm][order]
for rank in range(DCP_SIZE):
pool = self._bare_pool(DCP_SIZE, rank)
host_mask = host_sorted % DCP_SIZE == rank
device_mask = device_matched % DCP_SIZE == rank
# same positions selected on both sides -> pairing preserved
torch.testing.assert_close(host_mask, device_mask)
self.assertEqual(
pool.dcp_kernel_indices(host_sorted).numel(),
host.numel() // DCP_SIZE,
)
class TestHostPoolSizingUnderDcp(CustomTestCase):
def test_logical_and_physical_sizing(self):
pool = _make_host_pool(dcp_rank=3)
# kernel-facing page is physical
self.assertEqual(pool.page_size, PHYSICAL_PAGE)
self.assertEqual(pool.logical_page_size, WIDENED_PAGE)
# physical rows = ratio * device physical size, page aligned
self.assertEqual(pool.size, pool.page_num * PHYSICAL_PAGE)
self.assertEqual(pool.logical_size, pool.size * DCP_SIZE)
# buffers materialize physical rows only
self.assertEqual(pool.kv_buffer.shape[1], pool.size)
# allocator surface is logical
self.assertEqual(pool.free_slots.numel(), pool.logical_size)
self.assertEqual(pool.mem_state.numel(), pool.logical_size)
def test_alloc_is_widened_page_granular(self):
pool = _make_host_pool(dcp_rank=0)
out = pool.alloc(WIDENED_PAGE)
self.assertEqual(out.numel(), WIDENED_PAGE)
with self.assertRaises(AssertionError):
pool.alloc(PHYSICAL_PAGE) # not a multiple of the widened page
def test_non_dcp_pool_unchanged(self):
pool = MLATokenToKVPoolHost(
_fake_mla_device_pool(),
host_to_device_ratio=2.0,
host_size=0,
page_size=PHYSICAL_PAGE,
layout="layer_first",
pin_memory=False,
device="cpu",
)
self.assertEqual(pool.page_size, PHYSICAL_PAGE)
self.assertEqual(pool.logical_size, pool.size)
self.assertEqual(pool.logical_page_size, PHYSICAL_PAGE)
class TestTransferEntryPointsTranslate(CustomTestCase):
def _run_backup(self, pool, host_indices, device_indices):
device_pool = SimpleNamespace(
data_ptrs=torch.zeros(2, dtype=torch.uint64),
kv_buffer=[torch.zeros(1)] * 2,
)
# create=True: mla.py imports the kernel only under `if _is_cuda or
# _is_hip`, so the name is absent on the CPU runner this test targets.
with mock.patch(
"sglang.srt.mem_cache.pool_host.mla.transfer_kv_all_layer_mla",
create=True,
) as kernel:
pool.can_use_jit = False
pool.can_use_write_back_jit = False
with mock.patch.object(
MLATokenToKVPoolHost, "_is_device_layer_sharded", return_value=False
):
pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend="kernel"
)
return kernel.call_args.kwargs
def test_backup_receives_physical_rows(self):
pool = _make_host_pool(dcp_rank=5)
logical = torch.arange(2 * WIDENED_PAGE)
kwargs = self._run_backup(pool, logical, logical.clone())
expected = torch.arange(2 * PHYSICAL_PAGE)
torch.testing.assert_close(kwargs["src_indices"], expected)
torch.testing.assert_close(kwargs["dst_indices"], expected)
def test_l3_data_page_is_guarded(self):
pool = _make_host_pool(dcp_rank=0)
with self.assertRaises(AssertionError):
pool.get_data_page(0)
if __name__ == "__main__":
unittest.main()