diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index f0c7d64b4..4d2dd894d 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -452,6 +452,9 @@ class ModelConfig: self.is_multimodal_piecewise_cuda_graph_supported = enable_multimodal and ( is_multimodal_piecewise_cuda_graph_supported(self.hf_config.architectures) ) + self.is_multimodal_breakable_cuda_graph_supported = enable_multimodal and ( + is_multimodal_breakable_cuda_graph_supported(self.hf_config.architectures) + ) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) # Derive context length and model shapes @@ -1758,6 +1761,13 @@ multimodal_piecewise_cuda_graph_supported_model_archs = [ "MiniMaxM3SparseForConditionalGeneration", ] +# Multimodal archs whose LM prefill is validated under breakable CUDA graph; +# embed-carrying batches are rejected at replay (can_run_graph) and run eager. +multimodal_breakable_cuda_graph_supported_model_archs = [ + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", +] + if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): multimodal_model_archs.append(external_mm_model_arch) @@ -1824,6 +1834,14 @@ def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]) ) +def is_multimodal_breakable_cuda_graph_supported(model_architectures: List[str]): + """Whether a multimodal arch may keep prefill breakable CUDA graph enabled.""" + return any( + arch in multimodal_breakable_cuda_graph_supported_model_archs + for arch in model_architectures + ) + + # SequenceClassification models that use CrossEncodingPooler _cross_encoding_pooler_archs = [ "BertForSequenceClassification", diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 8fe294763..bdd430d60 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -612,6 +612,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return False if forward_batch.replace_embeds is not None: return False + # The captured graph embeds from input_ids only; multimodal batches + # merge mm embeddings in the outer wrapper, which capture bypasses. + if forward_batch.mm_inputs is not None and any( + x is not None for x in forward_batch.mm_inputs + ): + return False # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None. if forward_batch.forward_mode.is_target_verify(): return False diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a9b6340c7..48ec4791d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3645,8 +3645,12 @@ class ServerArgs: ), # DP-attn × BCG capture/replay not yet validated. ("DP attention", lambda: self._resolved().enable_dp_attention), - # Multimodal prefill replay faults under BCG. - ("multimodal model", lambda: self.get_model_config().is_multimodal), + # Multimodal prefill replay faults under BCG; allowlisted archs opt back in. + ( + "multimodal model", + lambda: self.get_model_config().is_multimodal + and not self.get_model_config().is_multimodal_breakable_cuda_graph_supported, + ), ] for name, predicate in rules: if predicate(): diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 5a7cabfc0..f7029e221 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1197,6 +1197,64 @@ class TestCudaGraphDisaggregationRoles(CustomTestCase): self.assertIn((Phase.DECODE, "backend"), args._cuda_graph_config_locked) +class TestBreakableCudaGraphMultimodalAllowlist(CustomTestCase): + """The BCG "multimodal model" rule exempts archs on the BCG multimodal + opt-in allowlist (multimodal_breakable_cuda_graph_supported_model_archs).""" + + def _handled_args(self, *, architectures, is_multimodal, allowlisted): + args = ServerArgs(model_path="dummy") + args.model_config = SimpleNamespace( + hf_config=SimpleNamespace(architectures=architectures), + is_piecewise_cuda_graph_disabled_model=False, + is_multimodal=is_multimodal, + is_multimodal_piecewise_cuda_graph_supported=False, + is_multimodal_breakable_cuda_graph_supported=allowlisted, + ) + with ( + patch("sglang.srt.utils.is_cuda", return_value=True), + patch.object(ServerArgs, "use_mla_backend", return_value=False), + ): + args._handle_cuda_graph_config() + return args + + def test_multimodal_arch_disables_prefill_breakable(self): + args = self._handled_args( + architectures=["Qwen3VLForConditionalGeneration"], + is_multimodal=True, + allowlisted=False, + ) + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED) + + def test_allowlisted_multimodal_arch_keeps_prefill_breakable(self): + args = self._handled_args( + architectures=["Qwen3_5MoeForConditionalGeneration"], + is_multimodal=True, + allowlisted=True, + ) + self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.BREAKABLE) + + def test_allowlist_membership(self): + from sglang.srt.configs.model_config import ( + is_multimodal_breakable_cuda_graph_supported, + ) + + self.assertTrue( + is_multimodal_breakable_cuda_graph_supported( + ["Qwen3_5MoeForConditionalGeneration"] + ) + ) + self.assertTrue( + is_multimodal_breakable_cuda_graph_supported( + ["Qwen3_5ForConditionalGeneration"] + ) + ) + self.assertFalse( + is_multimodal_breakable_cuda_graph_supported( + ["Qwen3VLForConditionalGeneration"] + ) + ) + + class TestCutedslMoeMaxNumTokens(CustomTestCase): """The shared CuteDSL MoE per-forward token bound. Fields are set directly to exercise the math independently of __post_init__ resolution.