diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 8e5bd0074..6642038b1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -28,6 +28,7 @@ from typing import Callable, List, Optional, Tuple, Union import torch import torch.distributed as dist +from torch import nn from sglang.srt.configs import ( FalconH1Config, @@ -247,6 +248,13 @@ if _is_npu: torch_npu.npu.set_compile_mode(jit_compile=False) +def resolve_language_model(model: nn.Module) -> nn.Module: + model_cls_name = model.__class__.__name__ + if model_cls_name == "Qwen3OmniMoeForConditionalGeneration": + return model.thinker.model + return model.model + + class RankZeroFilter(logging.Filter): """Filter that only allows INFO level logs from rank 0, but allows all other levels from any rank.""" @@ -2434,6 +2442,7 @@ class ModelRunner: # Collect attention layers from the model self.attention_layers = [] + self.model.model = resolve_language_model(self.model) for layer in self.model.model.layers: if hasattr(layer, "self_attn"): if hasattr(layer.self_attn, "attn"): diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 49655d220..ecfe1fe31 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -97,7 +97,6 @@ suites = { TestFile("test_original_logprobs.py", 41), TestFile("test_page_size.py", 60), TestFile("test_penalty.py", 82), - TestFile("test_piecewise_cuda_graph.py", 850), TestFile("test_priority_scheduling.py", 130), TestFile("test_pytorch_sampling_backend.py", 66), TestFile("test_radix_attention.py", 105), @@ -158,6 +157,7 @@ suites = { TestFile("test_local_attn.py", 411), TestFile("test_multi_instance_release_memory_occupation.py", 64), TestFile("test_pp_single_node.py", 481), + TestFile("test_piecewise_cuda_graph.py", 1200), ], "per-commit-8-gpu-h200": [ TestFile("test_deepseek_v3_basic.py", 275), diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index 5e2619e32..b8f3511b8 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -371,5 +371,47 @@ class TestPiecewiseCudaGraphQwen25VLEmbedding(CustomTestCase): ) +class TestPiecewiseCudaGraphQwen3OmniMOE(CustomTestCase): + """Test piecewise CUDA graph with Qwen3-Omni-30B-A3B-Instruct model""" + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + 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=[ + "--enable-piecewise-cuda-graph", + "--piecewise-cuda-graph-compiler", + "eager", + "--disable-radix-cache", + "--tp=4", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k_accuracy(self): + """Test GSM8K accuracy with 8-shot setting""" + num_examples = 2000 + + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=num_examples, + num_threads=min(num_examples, 1024), + ) + + metrics = run_eval(args) + print(f"GSM8K Accuracy: {metrics['score']:.3f}") + + self.assertGreaterEqual(metrics["score"], 0.70) + + if __name__ == "__main__": unittest.main()