[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 (
-    MultiModalityDataPaddingPatternTokenPairs,
-    general_mm_embed_routine,
-)
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternTokenPairs
 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,
-            data_embedding_funcs={
-                Modality.IMAGE: self.get_image_feature,
-            },
+            input_embeds=input_embeds,
             positions=positions,
         )
 
-        return hs
+        return hidden_states
 
     def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
         # Get all special token IDs