From 0be539024faf50db2d4c220b68a082dd8ee817bf Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Sat, 16 May 2026 23:50:01 +0800 Subject: [PATCH] [BugFix]: Fix DeepSeek V4 HiCache layer count logic (#25477) --- .../hybrid_cache/hybrid_pool_assembler.py | 9 +- .../test_unified_radix_cache_kl_hicache.py | 155 ++++++++++++++++++ ...unified_radix_cache_kl_hicache_nightly.py} | 141 ---------------- 3 files changed, 161 insertions(+), 144 deletions(-) create mode 100644 test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py rename test/registered/radix_cache/{test_unified_radix_hicache_kl.py => test_unified_radix_cache_kl_hicache_nightly.py} (51%) diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index 3de18edde..1d4f48b13 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -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) diff --git a/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py new file mode 100644 index 000000000..2f3fc32ba --- /dev/null +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache.py @@ -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() diff --git a/test/registered/radix_cache/test_unified_radix_hicache_kl.py b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py similarity index 51% rename from test/registered/radix_cache/test_unified_radix_hicache_kl.py rename to test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py index 841af8733..799e94a31 100644 --- a/test/registered/radix_cache/test_unified_radix_hicache_kl.py +++ b/test/registered/radix_cache/test_unified_radix_cache_kl_hicache_nightly.py @@ -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.