Allow prefill breakable CUDA graph for Qwen3.5 via multimodal opt-in allowlist (#30620)

This commit is contained in:
YAMY
2026-07-13 18:39:44 -07:00
committed by GitHub
parent 7431f35fd8
commit 41ad0d9c26
4 changed files with 88 additions and 2 deletions
+18
View File
@@ -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",
@@ -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
+6 -2
View File
@@ -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():
@@ -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.