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

[UnifiedTree]: Support HiCache For DeepSeek_V4 (#24691)

Browse Source 
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
Zhangheng
2026-05-14 20:20:52 -07:00
committed by GitHub
co-authored by ispobock 晟海
parent 974948a701
commit d9fa84b25b
11 changed files with 1221 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/radix_cache/test_unified_radix_cache_kl.py
+4 -53
View File
@@ -48,6 +48,7 @@ class UnifiedRadixTreeTestMixin:
prefix_len: int = 512
prefill_cache_assert = None
decode_cache_assert = None
sampling_temperature: float = 1
gsm8k_threshold: float = 0.93
mmlu_threshold: float = 0.8
@@ -106,6 +107,7 @@ class UnifiedRadixTreeTestMixin:
turn_suffixes=[t2, t3],
assert_decode_cached_tokens=self.decode_cache_assert,
max_new_tokens=self.max_new_tokens,
sampling_temperature=self.sampling_temperature,
)
def test_multiturn_prefill_cache_hit_branching(self):
@@ -135,6 +137,7 @@ class UnifiedRadixTreeTestMixin:
assert_decode_cached_tokens=self.decode_cache_assert,
branches_per_group=branches,
max_new_tokens=self.max_new_tokens,
sampling_temperature=self.sampling_temperature,
)
def test_multiturn_decode_cache_hit_branching(self):
@@ -159,6 +162,7 @@ class UnifiedRadixTreeTestMixin:
assert_decode_cached_tokens=self.decode_cache_assert,
branches_per_group=branches,
max_new_tokens=self.max_new_tokens,
sampling_temperature=self.sampling_temperature,
)
@@ -267,58 +271,5 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
kill_process_tree(cls.process.pid)
class TestUnifiedMambaRadixCacheWithHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache + UnifiedRadixCache."""
kl_threshold = 0.003
prefill_cache_assert = staticmethod(
make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE)
)
decode_cache_assert = staticmethod(
make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL)
)
@classmethod
def setUpClass(cls):
cls.model = MAMBA_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"--mem-fraction-static",
"0.85",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
str(MAMBA_TRACK_INTERVAL),
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--max-total-tokens",
"12000",
"--max-running-requests",
"4",
],
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/registered/radix_cache/test_unified_radix_hicache_kl.py
+162
View File
@@ -0,0 +1,162 @@
"""UnifiedRadixTree + HiCache KL divergence tests.
Tests Mamba hybrid, DeepSeek V4 Flash, and GLM-5 models with HiCache L2
offloading under UnifiedRadixTree, verifying multi-turn cache correctness
via KL divergence.
"""
import unittest
from test_unified_radix_cache_kl import UnifiedRadixTreeTestMixin
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
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,
)
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
MAMBA_CHUNK_SIZE = 64
MAMBA_TRACK_INTERVAL = 128
DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_LAUNCH_TIMEOUT = 3600
DSV32_MODEL = "deepseek-ai/DeepSeek-V3.2"
DSV32_LAUNCH_TIMEOUT = 3600
register_cuda_ci(est_time=900, suite="nightly-8-gpu-h200", nightly=True)
class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache + UnifiedRadixCache."""
kl_threshold = 0.003
prefill_cache_assert = staticmethod(
make_mamba_prefill_assert(chunk_size=MAMBA_CHUNK_SIZE)
)
decode_cache_assert = staticmethod(
make_mamba_decode_assert(track_interval=MAMBA_TRACK_INTERVAL)
)
@classmethod
def setUpClass(cls):
cls.model = MAMBA_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--tp-size",
"4",
"--chunked-prefill-size",
"2048",
"--mem-fraction-static",
"0.85",
"--mamba-scheduler-strategy",
"extra_buffer",
"--mamba-track-interval",
str(MAMBA_TRACK_INTERVAL),
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--max-total-tokens",
"12000",
"--max-running-requests",
"4",
],
env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label):
expected = history_len + output_len
actual = result["meta_info"]["cached_tokens"]
lower = max(0, expected - 256)
assert actual >= lower, f"{label}: expected cached_tokens>={lower}, got {actual}"
class TestUnifiedDeepSeekV4FlashHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""DeepSeek V4 Flash FP8 + HiCache + UnifiedRadixCache."""
kl_threshold = 0.0035
sampling_temperature = 0
decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens)
gsm8k_threshold = 0.90
num_gsm8k_questions = 100
@unittest.skip("no stable.")
def test_multiturn_logprobs_match(self):
pass
@classmethod
def setUpClass(cls):
cls.model = DSV4_FLASH_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DSV4_FLASH_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp-size",
"4",
"--attention-backend",
"compressed",
"--page-size",
"256",
"--chunked-prefill-size",
"8192",
"--mem-fraction-static",
"0.9",
"--disable-shared-experts-fusion",
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--swa-full-tokens-ratio",
"0.25",
"--max-total-tokens",
"20000",
"--max-running-requests",
"4",
],
env={
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()