[VLM] Support Piecewise CUDA Graph for InternVL (#13640)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/managers/mm_utils.py

@@ -983,12 +983,47 @@ def should_use_external_mm_preprocess(multimodal_model: nn.Module) -> bool:
 
     cls_name = multimodal_model.__class__.__name__
 
-    qwen_vl_classes = {
+    external_mm_preprocess_classes = {
         "Qwen2VLForConditionalGeneration",
         "Qwen2_5_VLForConditionalGeneration",
+        "InternVLChatModel",
     }
 
-    return cls_name in qwen_vl_classes
+    return cls_name in external_mm_preprocess_classes
+
+
+def resolve_external_mm_data_embedding_funcs(
+    multimodal_model: nn.Module,
+) -> Optional[Dict[Modality, Callable[[List[MultimodalDataItem]], torch.Tensor]]]:
+    """
+    Resolve the data_embedding_funcs mapping for external_mm_preprocess_routine
+    based on the given multimodal model. If this function returns None, the
+    external_mm_preprocess_routine will use its internal default behavior
+    (for example, for Qwen2_5_VL).
+
+    Resolution order:
+        1. If the model exposes external_mm_data_embedding_funcs explicitly,
+           adopt it.
+        2. TODO: Handle special classes with customized mm_data_embedding_funcs
+           (e.g. Qwen3_VL).
+        3. If not mapping, return None.
+    """
+
+    cls_name = multimodal_model.__class__.__name__
+
+    # High priority: model provides an explicit mapping attribute.
+    #    Example in InternVLChatModel.__init__:
+    #      self.external_mm_data_embedding_funcs = {
+    #          Modality.IMAGE: self.get_image_feature,
+    #      }
+    if hasattr(multimodal_model, "external_mm_data_embedding_funcs"):
+        funcs = getattr(multimodal_model, "external_mm_data_embedding_funcs")
+        # Allow an empty dict to mean "no data_embedding_funcs are needed".
+        return funcs or None
+
+    # If no mapping is found, return None so that external_mm_preprocess_routine
+    # can fall back to its default logic.
+    return None
 
 
 def external_mm_preprocess_routine(

python/sglang/srt/model_executor/model_runner.py

@@ -94,6 +94,7 @@ from sglang.srt.lora.lora_manager import LoRAManager
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.mm_utils import (
     external_mm_preprocess_routine,
+    resolve_external_mm_data_embedding_funcs,
     should_use_external_mm_preprocess,
 )
 from sglang.srt.mem_cache.allocator import (
@@ -2145,9 +2146,11 @@ class ModelRunner:
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
 
         if self.is_multimodal and should_use_external_mm_preprocess(self.model):
+            data_embedding_funcs = resolve_external_mm_data_embedding_funcs(self.model)
             forward_batch = external_mm_preprocess_routine(
                 forward_batch=forward_batch,
                 multimodal_model=self.model,
+                data_embedding_funcs=data_embedding_funcs,
             )
 
         kwargs = {}

python/sglang/srt/models/internvl.py

@@ -14,10 +14,7 @@ from sglang.srt.layers.attention import vision_utils
 from sglang.srt.layers.attention.vision import SingletonCache, VisionAttention
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
-from sglang.srt.managers.mm_utils import (
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternTokenPairs
-    MultiModalityDataPaddingPatternTokenPairs,
-    general_mm_embed_routine,
-)
 from sglang.srt.managers.schedule_batch import (
     Modality,
     MultimodalDataItem,
@@ -471,6 +468,12 @@ class InternVLChatModel(nn.Module):
             nn.Linear(llm_hidden_size, llm_hidden_size),
         )
 
+        self.external_mm_data_embedding_funcs = {
+            Modality.IMAGE: self.get_image_feature,
+        }
+
+        self.model = self.language_model.model
+
     def pixel_shuffle(self, x, scale_factor=0.5):
         n, w, h, c = x.size()
         # N, W, H, C --> N, W, H * scale, C // scale
@@ -531,17 +534,25 @@ class InternVLChatModel(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
 
-        hs = general_mm_embed_routine(
+        input_embeds = forward_batch.input_embeds
+        # It may seem strange to assign input_embeds again even after passing it as an argument.
+        # This is for compatibility considerations.
+        # In the 'extend' scenario, this forward function is called from two places:
+        # 1. model_runner calls forward directly,
+        # 2. piece_wise_cuda_graph_runner calls forward and replay.
+
+        # Currently,
+        # In 'extend', input_embeds is passed in.
+        # In 'decode', input_ids is passed in.
+
+        hidden_states = self.language_model(
             input_ids=input_ids,
             forward_batch=forward_batch,
-            language_model=self.language_model,
+            input_embeds=input_embeds,
-            data_embedding_funcs={
-                Modality.IMAGE: self.get_image_feature,
-            },
             positions=positions,
         )
 
-        return hs
+        return hidden_states
 
     def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
         # Get all special token IDs

test/srt/run_suite.py

@@ -95,7 +95,7 @@ suites = {
         TestFile("test_original_logprobs.py", 41),
         TestFile("test_page_size.py", 60),
         TestFile("test_penalty.py", 82),
-        TestFile("test_piecewise_cuda_graph.py", 750),
+        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),

test/srt/test_piecewise_cuda_graph.py

@@ -289,5 +289,46 @@ class TestPiecewiseCudaGraphQwen25VL(CustomTestCase):
         self.assertGreaterEqual(metrics["score"], 0.70)
 
 
+class TestPiecewiseCudaGraphInternVL25(CustomTestCase):
+    """Test piecewise CUDA graph with InternVL2.5-8B-Instruct model"""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = "OpenGVLab/InternVL2_5-8B"
+        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",
+            ],
+        )
+
+    @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()