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

2026-05-16 23:50:01 +08:00
parent 2fc217df4d
commit 0be539024f
3 changed files with 161 additions and 144 deletions
@@ -283,7 +283,8 @@ def build_deepseek_v4_hicache_stack(
    pp_size: int = 1,
    enable_storage_metrics: bool = False,
 ) -> tuple[HostPoolGroup, HybridCacheController]:
-    transfer_layer_num = len(kvcache.compression_ratios)
+    # TODO(hzh0425): Support PP for deepseek v4 with hicache
+    transfer_layer_num = kvcache.end_layer - kvcache.start_layer
    full_layer_mapping = {layer_id: layer_id for layer_id in range(transfer_layer_num)}
    swa_layer_mapping = {
        layer_id: layer_id for layer_id in range(len(kvcache.swa_kv_pool.kv_buffer))
@@ -293,7 +294,9 @@ def build_deepseek_v4_hicache_stack(
    c128_layer_mapping = {}
    c4_state_global_layers = []
    c128_state_global_layers = []
-    for layer_id, layer_item in enumerate(kvcache.layer_mapping):
+    for layer_id, layer_item in enumerate(
+        kvcache.layer_mapping[kvcache.start_layer : kvcache.end_layer]
+    ):
        if layer_item.compress_ratio == 4:
            c4_layer_mapping[layer_id] = layer_item.compress_layer_id
            c4_state_global_layers.append(layer_id)
@@ -730,7 +733,7 @@ def attach_hybrid_pool_to_unified_cache(
                            indices_from_pool=indices_from_pool,
                        )
                    )
-            transfer_layer_num = len(kvcache.compression_ratios)
+            transfer_layer_num = kvcache.end_layer - kvcache.start_layer
        elif mamba_stack:
            full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping)
            mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
@@ -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.