Allow prefill breakable CUDA graph for Qwen3.5 via multimodal opt-in allowlist (#30620)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user