From 41cd5a718942f97bc45b0b5d7fca82992e8ae529 Mon Sep 17 00:00:00 2001 From: DarkraiHL <16268229+DarkraiHL@users.noreply.github.com> Date: Sat, 15 Aug 2026 01:55:41 +0800 Subject: [PATCH] [Fix] Fix Qwen3.5 MTP startup with HiCache (#34560) Co-authored-by: hjzhang <76768149+1e4ves@users.noreply.github.com> Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com> --- python/sglang/srt/configs/model_config.py | 1 + .../hybrid_cache/hybrid_pool_assembler.py | 8 ++++++- .../registered/hicache/test_qwen35_hicache.py | 8 +++++++ .../unit/configs/test_model_config.py | 18 ++++++++++++++ .../mem_cache/test_hybrid_pool_assembler.py | 24 +++++++++++++++++-- 5 files changed, 56 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 7e3f53c7b..d28987651 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -707,6 +707,7 @@ class ModelConfig: ) self.hf_config.architectures[0] = "Qwen3_5ForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 + self.hf_text_config.num_nextn_predict_layers = 1 if is_draft_model and self.hf_config.architectures[0] == "ExaoneMoEForCausalLM": self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP" 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 099bd10cb..84b1e84fc 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 @@ -930,9 +930,15 @@ def build_full_draft_pools( server_args: ServerArgs, ) -> tuple[list[SidecarPoolSpec], list[PoolEntry]]: """Build draft KV/DSA sidecars whose indices follow target full KV.""" - from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + from sglang.srt.mem_cache.memory_pool import ( + DSATokenToKVPool, + HybridLinearKVPool, + ) pool = draft_kv_pool + if isinstance(pool, HybridLinearKVPool): + # Hybrid draft runners keep their sole attention layer in this sub-pool. + pool = pool.full_kv_pool if pool.layer_num == 0: return [], [] diff --git a/test/registered/hicache/test_qwen35_hicache.py b/test/registered/hicache/test_qwen35_hicache.py index 8398faaba..7b271917d 100644 --- a/test/registered/hicache/test_qwen35_hicache.py +++ b/test/registered/hicache/test_qwen35_hicache.py @@ -57,6 +57,14 @@ class TestQwen35WithHiCache(CustomTestCase): "128", "--mamba-ssm-dtype", "bfloat16", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", "--max-running-requests", "128", "--reasoning-parser", diff --git a/test/registered/unit/configs/test_model_config.py b/test/registered/unit/configs/test_model_config.py index 1014a11fc..778928154 100644 --- a/test/registered/unit/configs/test_model_config.py +++ b/test/registered/unit/configs/test_model_config.py @@ -4,6 +4,7 @@ import unittest from types import SimpleNamespace from sglang.srt.configs.model_config import ( + ModelConfig, get_hybrid_layer_ids, is_embedding_gemma, ) @@ -52,5 +53,22 @@ class TestEmbeddingGemmaConfig(CustomTestCase): self.assertFalse(is_embedding_gemma(config)) +class TestDraftModelConfig(CustomTestCase): + def test_qwen35_mtp_depth_is_synced_to_text_config(self): + config = object.__new__(ModelConfig) + config.is_draft_model = True + config.speculative_algorithm = "EAGLE" + config.hf_config = SimpleNamespace( + architectures=["Qwen3_5MoeForConditionalGeneration"] + ) + config.hf_text_config = SimpleNamespace() + + config._config_draft_model() + + self.assertEqual(config.hf_config.architectures, ["Qwen3_5ForCausalLMMTP"]) + self.assertEqual(config.hf_config.num_nextn_predict_layers, 1) + self.assertEqual(config.hf_text_config.num_nextn_predict_layers, 1) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py index 95a14ab4a..c93f837db 100644 --- a/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py +++ b/test/registered/unit/mem_cache/test_hybrid_pool_assembler.py @@ -1,8 +1,13 @@ -"""Unit test for hybrid HiCache fixed-size budget splitting.""" +"""Unit tests for hybrid HiCache pool assembly.""" import unittest +from types import SimpleNamespace -from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import _split_hicache_size +from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + _split_hicache_size, + build_full_draft_pools, +) +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -35,5 +40,20 @@ class TestSplitHicacheSize(CustomTestCase): self.assertEqual(sum(shares), 100) # total budget preserved, not doubled +class TestDraftSidecarPoolDispatch(CustomTestCase): + def test_full_builder_unwraps_empty_hybrid_linear_pool(self): + draft_kv_pool = object.__new__(HybridLinearKVPool) + draft_kv_pool.full_kv_pool = SimpleNamespace(layer_num=0) + + specs, entries = build_full_draft_pools( + draft_kv_pool=draft_kv_pool, + tree_cache=None, + server_args=None, + ) + + self.assertEqual(specs, []) + self.assertEqual(entries, []) + + if __name__ == "__main__": unittest.main()