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

Preserve draft KV across CP HiCache hits

Browse Source 
Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix.  CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool.

Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV.

Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank.

Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint.

Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance.

Confidence: medium

Scope-risk: moderate

Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length.

Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings

Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery.
This commit is contained in:
laoyao0822
2026-05-26 23:58:40 +08:00
committed by leavelet
parent d655fad040
commit 315eaaff56
7 changed files with 681 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+114 -13
View File
@@ -1,14 +1,80 @@
import sys
import types
import unittest
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import torch
# Stub out sgl_kernel before any sglang import so this CPU unit test does not
# require CUDA extension libraries to be installed.
for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
if _mod not in sys.modules:
sys.modules[_mod] = MagicMock()
# Prefer the real sgl_kernel package when the test image provides it so custom
# Torch operators are registered. Fall back to stubs on local CPU-only hosts.
try:
import sgl_kernel # noqa: F401
import sgl_kernel.kvcacheio # noqa: F401
except (ImportError, RuntimeError):
if "sgl_kernel" not in sys.modules:
sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
sys.modules["sgl_kernel"].__path__ = []
if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
def _sgl_kernel_getattr(name):
if name.startswith("__"):
raise AttributeError(name)
fn = lambda *args, **kwargs: None
setattr(sys.modules["sgl_kernel"], name, fn)
return fn
sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
if "sgl_kernel.quantization" not in sys.modules:
quantization_stub = types.ModuleType("sgl_kernel.quantization")
quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
def _quantization_getattr(name):
if name.startswith("__"):
raise AttributeError(name)
fn = lambda *args, **kwargs: None
setattr(quantization_stub, name, fn)
return fn
quantization_stub.__getattr__ = _quantization_getattr
for _name in (
"ggml_dequantize",
"ggml_moe_a8",
"ggml_moe_a8_vec",
"ggml_moe_get_block_size",
"ggml_mul_mat_a8",
"ggml_mul_mat_vec_a8",
):
setattr(quantization_stub, _name, lambda *args, **kwargs: None)
sys.modules["sgl_kernel.quantization"] = quantization_stub
for _name in (
"sgl_per_token_group_quant_8bit",
"sgl_per_token_group_quant_fp8",
"sgl_per_token_quant_fp8",
"fp8_blockwise_scaled_mm",
"fp8_scaled_mm",
"silu_and_mul",
):
if not hasattr(sys.modules["sgl_kernel"], _name):
setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None)
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
for _schema in (
"sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
"fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
):
try:
_sgl_kernel_lib.define(_schema)
except RuntimeError as exc:
if (
"already" not in str(exc).lower()
and "duplicate" not in str(exc).lower()
):
raise
if "sgl_kernel.kvcacheio" not in sys.modules:
sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
@@ -70,6 +136,22 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
self.assertEqual(child.owned_positions.tolist(), [0, 4])
self.assertEqual(child.host_indices.tolist(), [22, 23])
def test_split_keeps_draft_host_indices_aligned_with_owned_positions(self):
metadata = CpHiCacheNodeMetadata(
logical_len=10,
owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64),
)
parent, child = metadata.split(5)
self.assertEqual(parent.host_indices.tolist(), [20, 21])
self.assertEqual(parent.draft_host_indices.tolist(), [120, 121])
self.assertEqual(child.owned_positions.tolist(), [0, 4])
self.assertEqual(child.host_indices.tolist(), [22, 23])
self.assertEqual(child.draft_host_indices.tolist(), [122, 123])
def test_zero_owned_metadata_is_valid(self):
metadata = CpHiCacheNodeMetadata(
logical_len=64,
@@ -150,6 +232,15 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
host_indices=torch.tensor([9], dtype=torch.int64),
)
def test_draft_host_length_mismatch_raises(self):
with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"):
CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
host_indices=torch.tensor([9, 10], dtype=torch.int64),
draft_host_indices=torch.tensor([109], dtype=torch.int64),
)
def test_out_of_range_positions_raise(self):
with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
CpHiCacheNodeMetadata(
@@ -222,6 +313,11 @@ class FakeEvictDeviceController:
return len(device_indices)
class FakeTokenAllocator:
def available_size(self):
return 0
class TestHiRadixCacheCPBackup(CustomTestCase):
def test_node_backuped_uses_cp_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
@@ -341,6 +437,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
cache.evictable_leaves = set()
cache.evictable_host_leaves = set()
cache.eviction_strategy = FakeEvictionStrategy()
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
cache.evictable_size_ = 4
cache.protected_size_ = 0
cache._record_remove_event = lambda node: (_ for _ in ()).throw(
@@ -602,14 +699,18 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
cache.root_node.children[1] = node
cache.evictable_host_leaves.add(node)
def mark_not_done(done_tensor, op=None, group=None):
self.assertIs(group, cache.tp_group)
done_tensor.fill_(0)
all_done_states = iter([False, True])
cache._cp_all_ranks_true = lambda done: next(all_done_states)
cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids]
cache._cp_filter_all_ranks_safe_node_ids = (
lambda node_ids, is_safe, **_kwargs: [
node_id
for node_id in node_ids
if is_safe(cache._cp_node_by_id(node_id))
]
)
with patch("torch.distributed.all_reduce", side_effect=mark_not_done):
physical_freed = cache._evict_host_for_physical_slots(
0, synchronize_across_ranks=True
)
physical_freed = cache._cp_evict_host_for_physical_slots(0)
self.assertEqual(physical_freed, 0)
self.assertEqual(freed, [])