diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 5964e1cbd..958cff1cd 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -183,11 +183,6 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): pre_alloc_size=pre_alloc_size, ) - if envs.SGLANG_ENABLE_SPEC_V2.get() and not enable_mamba_extra_buffer: - raise ValueError( - "Spec v2 requires mamba scheduler strategy `extra_buffer` for mamba models. " - "Please set `--mamba-scheduler-strategy extra_buffer`." - ) self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 11c45ff1c..6d7fed4b7 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -465,11 +465,6 @@ class HybridReqToTokenPool(ReqToTokenPool): device=device, enable_memory_saver=enable_memory_saver, ) - if envs.SGLANG_ENABLE_SPEC_V2.get() and not enable_mamba_extra_buffer: - raise ValueError( - "Spec v2 requires mamba scheduler strategy `extra_buffer` for mamba models. " - "Please set `--mamba-scheduler-strategy extra_buffer`." - ) self.mamba_ping_pong_track_buffer_size = 2 if enable_overlap_schedule else 1 self.enable_mamba_extra_buffer = enable_mamba_extra_buffer diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 84c64b08d..a713f399b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -684,7 +684,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): mm_input, self.seq_lens_cpu[batch_idx] ) mrope_positions_list[batch_idx] = mrope_positions - elif self.forward_mode.is_extend(): + elif self.forward_mode.is_extend(include_draft_extend_v2=True): extend_seq_len, extend_prefix_len = ( batch.extend_seq_lens[batch_idx], batch.extend_prefix_lens[batch_idx], diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 939f61fc7..97d02274b 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -111,7 +111,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): if ( forward_batch.forward_mode.is_extend() and forward_batch.contains_mm_inputs() - and not forward_batch.forward_mode.is_draft_extend() + and not forward_batch.forward_mode.is_draft_extend(include_v2=True) ): assert input_embeds is not None input_embeds = torch.cat( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 02ee4c6b7..e346ae10b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1896,10 +1896,11 @@ class ServerArgs: self.disable_radix_cache = True self.disable_overlap_schedule = False else: - logger.warning( - f"Disabling radix cache since speculative decoding for {model_arch} is not supported with radix cache yet." - ) - self.disable_radix_cache = True + if not self.disable_radix_cache: + raise ValueError( + f"Speculative decoding for {model_arch} is not compatible with radix cache when using --mamba-scheduler-strategy no_buffer." + "To use radix cache with speculative decoding, please use --mamba-scheduler-strategy extra_buffer and set SGLANG_ENABLE_SPEC_V2=1." + ) def _handle_sampling_backend(self): if self.sampling_backend is None: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 03afe03ea..b6bf6d5ef 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -465,6 +465,7 @@ class EagleDraftWorker(BaseDraftWorker): batch: ModelWorkerBatch, target_hidden_states: torch.Tensor, next_token_ids: torch.Tensor, + mm_input_embeds: Optional[torch.Tensor] = None, ): """ Run draft model extend to correctly fill the KV cache. @@ -498,6 +499,8 @@ class EagleDraftWorker(BaseDraftWorker): # Run forward forward_batch = ForwardBatch.init_new(batch, self.draft_runner) + if mm_input_embeds is not None: + forward_batch.mm_input_embeds = mm_input_embeds logits_output = self.draft_runner.forward(forward_batch).logits_output # Update spec_info for the next draft step @@ -668,6 +671,7 @@ class EAGLEWorkerV2(BaseSpecWorker): model_worker_batch, batch_output.logits_output.hidden_states, batch_output.next_token_ids, + batch_output.logits_output.mm_input_embeds, ) ) return batch_output diff --git a/test/registered/4-gpu-models/test_qwen35_models.py b/test/registered/4-gpu-models/test_qwen35_models.py new file mode 100644 index 000000000..a0fc77f3b --- /dev/null +++ b/test/registered/4-gpu-models/test_qwen35_models.py @@ -0,0 +1,240 @@ +import unittest +from types import SimpleNamespace + +import requests + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci + +# This eval harness applies the chat_template, which is critical for qwen3.5 +# to get good accuracy on gsm8k +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=1000, suite="stage-c-test-4-gpu-b200") + +QWEN35_FP4_MODEL = "nvidia/Qwen3.5-397B-A17B-NVFP4" + +ACC_THRESHOLDS = { + QWEN35_FP4_MODEL: {"gsm8k": 0.95}, +} + + +class TestQwen35FP4(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN35_FP4_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", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + "--mamba-ssm-dtype", + "bfloat16", + "--max-running-requests", + "128", + "--reasoning-parser", + "qwen3", + "--attention-backend", + "trtllm_mha", + "--quantization", + "modelopt_fp4", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) + + +class TestQwen35FP4MTP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN35_FP4_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", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + "--mamba-ssm-dtype", + "bfloat16", + "--max-running-requests", + "128", + "--reasoning-parser", + "qwen3", + "--attention-backend", + "trtllm_mha", + "--quantization", + "modelopt_fp4", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--mem-fraction-static", + "0.8", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) + + server_info = requests.get(self.base_url + "/get_server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + self.assertGreater(avg_spec_accept_length, 3.3) + + +class TestQwen35FP4MTPV2(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN35_FP4_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + envs.SGLANG_ENABLE_SPEC_V2.set(True) + 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", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", + "--mamba-ssm-dtype", + "bfloat16", + "--max-running-requests", + "128", + "--reasoning-parser", + "qwen3", + "--attention-backend", + "trtllm_mha", + "--quantization", + "modelopt_fp4", + "--speculative-algorithm", + "NEXTN", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + "--mem-fraction-static", + "0.8", + "--model-loader-extra-config", + '{"enable_multithread_load": true,"num_threads": 64}', + ], + ) + + @classmethod + def tearDownClass(cls): + envs.SGLANG_ENABLE_SPEC_V2.set(False) + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], ACC_THRESHOLDS[self.model]["gsm8k"]) + + server_info = requests.get(self.base_url + "/get_server_info") + avg_spec_accept_length = server_info.json()["internal_states"][0][ + "avg_spec_accept_length" + ] + print(f"{avg_spec_accept_length=}") + self.assertGreater(avg_spec_accept_length, 3.3) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py index 754ea03c4..6f0784a83 100644 --- a/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py +++ b/test/registered/4-gpu-models/test_qwen3_next_models_mtp.py @@ -72,6 +72,7 @@ class TestQwen3NextMTP(CustomTestCase): "2048", "--mamba-scheduler-strategy", "no_buffer", + "--disable-radix-cache", ], )