From 475962a139d16f46e62a130f55138dfeb9104006 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Fri, 21 Nov 2025 13:21:51 +0800 Subject: [PATCH] [VLM] Support Piecewise CUDA Graph for InternVL (#13640) Co-authored-by: luoyuan.luo --- python/sglang/srt/managers/mm_utils.py | 39 +++++++++++++++++- .../sglang/srt/model_executor/model_runner.py | 3 ++ python/sglang/srt/models/internvl.py | 31 +++++++++----- test/srt/run_suite.py | 2 +- test/srt/test_piecewise_cuda_graph.py | 41 +++++++++++++++++++ 5 files changed, 103 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 99a909312..cd566b7c8 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/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( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 76ad65724..c8dc396ec 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 = {} diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index 389c2a884..458cd95f2 100644 --- a/python/sglang/srt/models/internvl.py +++ b/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 diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 43728034e..511142b0c 100644 --- a/test/srt/run_suite.py +++ b/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), diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index 50b2cb34d..ea2db8fe1 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/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()