[BugFix]: Fix DeepSeek V4 HiCache layer count logic (#25477)

This commit is contained in:
Zhangheng
2026-05-16 23:50:01 +08:00
committed by GitHub
parent 2fc217df4d
commit 0be539024f
3 changed files with 161 additions and 144 deletions
@@ -0,0 +1,155 @@
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,
is_in_ci,
popen_launch_server,
)
MAMBA_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
MAMBA_CHUNK_SIZE = 64
MAMBA_TRACK_INTERVAL = 128
DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8"
DSV4_FLASH_LAUNCH_TIMEOUT = 3600
register_cuda_ci(est_time=768, stage="base-c", runner_config="8-gpu-h200")
class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""Mamba hybrid + HiCache + UnifiedRadixCache."""
kl_threshold = 0.005
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-mamba-cache-size",
"500",
"--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.005
sampling_temperature = 0
decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens)
gsm8k_threshold = 0.90
num_gsm8k_questions = 100
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
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",
"2",
],
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()
@@ -13,162 +13,21 @@ from types import SimpleNamespace
from urllib.parse import urlparse
import requests
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
GLM5_MODEL = "zai-org/GLM-5.1-FP8"
GLM5_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-mamba-cache-size",
"500",
"--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",
"2",
],
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)
class GSM8KTwoPassMixin:
"""Mixin: run GSM8K twice with flush in between, verify accuracy diff.