From 669be5448bbbbcbd724a6a6c45787dbcccde843a Mon Sep 17 00:00:00 2001
From: zijiexia <37504505+zijiexia@users.noreply.github.com>
Date: Mon, 22 Jun 2026 12:20:14 -0700
Subject: [PATCH] [cuda graph] Enable prefill piecewise CUDA graph for
 Cohere2Vision (text path) (#28686)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 python/sglang/srt/configs/model_config.py     | 24 +++++++++
 python/sglang/srt/server_args.py              |  6 ++-
 ...st_multimodal_piecewise_cuda_graph_gate.py | 49 +++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 test/srt/test_multimodal_piecewise_cuda_graph_gate.py

diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
index 964bdcd6f..8921f71af 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -381,6 +381,14 @@ class ModelConfig:
         self.is_piecewise_cuda_graph_disabled_model = (
             is_piecewise_cuda_graph_disabled_model(self.hf_config.architectures)
         )
+        # Multimodal archs whose language-model prefill is verified safe to capture
+        # under piecewise CUDA graph. ServerArgs otherwise disables prefill piecewise
+        # CG for every multimodal model; this opt-in re-enables it for listed archs
+        # (the vision encoder still runs eagerly via general_mm_embed_routine, only the
+        # LM forward is captured).
+        self.is_multimodal_piecewise_cuda_graph_supported = enable_multimodal and (
+            is_multimodal_piecewise_cuda_graph_supported(self.hf_config.architectures)
+        )
         self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype)
 
         # Derive context length and model shapes
@@ -1651,6 +1659,14 @@ piecewise_cuda_graph_disabled_model_archs = [
     "LLaDAModelLM",
 ]
 
+# Multimodal archs allowed to keep prefill piecewise CUDA graph enabled. The
+# generic "multimodal model" rule in ServerArgs disables prefill piecewise CG for
+# all multimodal models; archs here opt back in because their LM prefill captures
+# cleanly (vision encoder runs eagerly outside the graph via general_mm_embed_routine).
+multimodal_piecewise_cuda_graph_supported_model_archs = [
+    "Cohere2VisionForConditionalGeneration",
+]
+
 if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get():
     multimodal_model_archs.append(external_mm_model_arch)
 
@@ -1709,6 +1725,14 @@ def is_piecewise_cuda_graph_disabled_model(model_architectures: List[str]):
     )
 
 
+def is_multimodal_piecewise_cuda_graph_supported(model_architectures: List[str]):
+    """Whether a multimodal arch may keep prefill piecewise CUDA graph enabled."""
+    return any(
+        arch in multimodal_piecewise_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/server_args.py b/python/sglang/srt/server_args.py
index 584d68b56..92fbcb575 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -3008,7 +3008,11 @@ class ServerArgs:
             ),
             ("MoE A2A backend", lambda: self.moe_a2a_backend != "none"),
             ("LoRA", lambda: bool(self.lora_paths) or self.enable_lora),
-            ("multimodal model", lambda: self.get_model_config().is_multimodal),
+            (
+                "multimodal model",
+                lambda: self.get_model_config().is_multimodal
+                and not self.get_model_config().is_multimodal_piecewise_cuda_graph_supported,
+            ),
             (
                 "GGUF quantization",
                 lambda: self.load_format == "gguf"
diff --git a/test/srt/test_multimodal_piecewise_cuda_graph_gate.py b/test/srt/test_multimodal_piecewise_cuda_graph_gate.py
new file mode 100644
index 000000000..c18910d77
--- /dev/null
+++ b/test/srt/test_multimodal_piecewise_cuda_graph_gate.py
@@ -0,0 +1,49 @@
+"""Unit test for the multimodal prefill piecewise-CUDA-graph opt-in gate.
+
+ServerArgs disables prefill piecewise CUDA graph for every multimodal model. Some
+multimodal archs (whose vision encoder runs eagerly outside the graph and whose LM
+prefill captures cleanly) opt back in via
+``multimodal_piecewise_cuda_graph_supported_model_archs``. This test pins that gate.
+"""
+
+import unittest
+
+from sglang.srt.configs.model_config import (
+    is_multimodal_piecewise_cuda_graph_supported,
+    multimodal_piecewise_cuda_graph_supported_model_archs,
+)
+
+
+class TestMultimodalPiecewiseCudaGraphGate(unittest.TestCase):
+    def test_cohere2_vision_opted_in(self):
+        # Cohere2-Vision (command-a / aya-vision family) LM prefill captures cleanly
+        # under piecewise CG; it must be opted back in.
+        self.assertTrue(
+            is_multimodal_piecewise_cuda_graph_supported(
+                ["Cohere2VisionForConditionalGeneration"]
+            )
+        )
+
+    def test_unlisted_multimodal_arch_stays_disabled(self):
+        # An arch not on the allow-list keeps the default (disabled) behavior.
+        self.assertFalse(
+            is_multimodal_piecewise_cuda_graph_supported(
+                ["SomeOtherVisionForConditionalGeneration"]
+            )
+        )
+        self.assertFalse(is_multimodal_piecewise_cuda_graph_supported([]))
+
+    def test_allow_list_entries_are_recognized(self):
+        for arch in multimodal_piecewise_cuda_graph_supported_model_archs:
+            self.assertTrue(is_multimodal_piecewise_cuda_graph_supported([arch]))
+
+    def test_match_within_mixed_arch_list(self):
+        self.assertTrue(
+            is_multimodal_piecewise_cuda_graph_supported(
+                ["OtherArch", "Cohere2VisionForConditionalGeneration"]
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()