[VLM] Support Piecewise CUDA Graph for Qwen2.5-VL (#13055)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: Yuhao Yang <yhyang201@gmail.com>

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -318,6 +318,10 @@ class CudaGraphRunner:
         # Graph inputs
         with torch.device(self.device):
             self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
+            self.input_embeds = torch.zeros(
+                (self.max_num_token, self.model_runner.model_config.hidden_size),
+                dtype=self.model_runner.model_config.dtype,
+            )
             self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
             self.seq_lens = torch.full(
                 (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32

python/sglang/srt/model_executor/model_runner.py

@@ -92,6 +92,10 @@ from sglang.srt.layers.sampler import Sampler
 from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
 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,
+    should_use_external_mm_preprocess,
+)
 from sglang.srt.mem_cache.allocator import (
     BaseTokenToKVPoolAllocator,
     PagedTokenToKVPoolAllocator,
@@ -2139,6 +2143,13 @@ class ModelRunner:
         skip_attn_backend_init: bool = False,
         pp_proxy_tensors=None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
+
+        if self.is_multimodal and should_use_external_mm_preprocess(self.model):
+            forward_batch = external_mm_preprocess_routine(
+                forward_batch=forward_batch,
+                multimodal_model=self.model,
+            )
+
         kwargs = {}
         if self.support_pp:
             kwargs["pp_proxy_tensors"] = pp_proxy_tensors

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -166,9 +166,15 @@ class PiecewiseCudaGraphRunner:
 
         self.max_num_tokens = max(self.capture_num_tokens)
 
+        self.use_input_embeds = model_runner.is_multimodal
+
         # Graph inputs
         with torch.device(self.device):
             self.input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+            self.input_embeds = torch.zeros(
+                (self.max_num_tokens, self.model_runner.model_config.hidden_size),
+                dtype=self.model_runner.dtype,
+            )
             self.out_cache_loc = torch.zeros(
                 (self.max_num_tokens,), dtype=self._cache_loc_dtype()
             )
@@ -176,6 +182,9 @@ class PiecewiseCudaGraphRunner:
                 (self.max_num_tokens,), dtype=self._cache_loc_dtype()
             )
             self.positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+            self.mrope_positions = torch.zeros(
+                (3, self.max_num_tokens), dtype=torch.int64
+            )
             self.tbo_plugin = TboCudaGraphRunnerPlugin()
 
         self.attention_layers = self.model_runner.attention_layers
@@ -216,7 +225,21 @@ class PiecewiseCudaGraphRunner:
             forward_batch = ForwardBatch(
                 forward_mode=ForwardMode.EXTEND,
                 batch_size=1,
-                input_ids=torch.randint(0, 100, (num_tokens,), device=self.device),
+                input_ids=(
+                    torch.randint(0, 100, (num_tokens,), device=self.device)
+                    if not self.use_input_embeds
+                    else None
+                ),
+                input_embeds=(
+                    torch.randn(
+                        num_tokens,
+                        self.model_runner.model_config.hidden_size,
+                        dtype=self.model_runner.dtype,
+                        device=self.device,
+                    )
+                    if self.use_input_embeds
+                    else None
+                ),
                 req_pool_indices=torch.arange(1, device=self.device),
                 seq_lens=torch.tensor([num_tokens], device=self.device),
                 next_token_logits_buffer=None,
@@ -246,7 +269,7 @@ class PiecewiseCudaGraphRunner:
                 global_num_tokens_for_logprob_gpu=None,
                 dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
                 global_dp_buffer_len=None,
-                mrope_positions=None,
+                mrope_positions=self.mrope_positions[:, :num_tokens],
                 spec_algorithm=None,
                 spec_info=None,
                 capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -326,10 +349,17 @@ class PiecewiseCudaGraphRunner:
         bs = 1
 
         # Graph inputs
-        input_ids = self.input_ids[:num_tokens]
+        if self.use_input_embeds:
+            input_ids = None
+            input_embeds = self.input_embeds[:num_tokens]
+        else:
+            input_ids = self.input_ids[:num_tokens]
+            input_embeds = None
+
         out_cache_loc = self.out_cache_loc[:num_tokens]
         out_cache_loc_swa = self.out_cache_loc_swa[:num_tokens]
         positions = self.positions[:num_tokens]
+        mrope_positions = self.mrope_positions[:, :num_tokens]
 
         # pipeline parallelism
         if self.pp_size > 1:
@@ -351,6 +381,7 @@ class PiecewiseCudaGraphRunner:
                 forward_mode=ForwardMode.EXTEND,
                 batch_size=bs,
                 input_ids=input_ids,
+                input_embeds=input_embeds,
                 req_pool_indices=torch.arange(bs, device=self.device),
                 seq_lens=torch.tensor([num_tokens], device=self.device),
                 next_token_logits_buffer=None,
@@ -376,7 +407,7 @@ class PiecewiseCudaGraphRunner:
                 global_num_tokens_for_logprob_gpu=None,
                 dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
                 global_dp_buffer_len=None,
-                mrope_positions=None,
+                mrope_positions=mrope_positions,
                 spec_algorithm=None,
                 spec_info=None,
                 capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -428,7 +459,11 @@ class PiecewiseCudaGraphRunner:
         forward_batch: ForwardBatch,
         **kwargs,
     ):
-        num_tokens = len(forward_batch.input_ids)
+        if self.use_input_embeds:
+            num_tokens = forward_batch.input_embeds.shape[0]
+        else:
+            num_tokens = len(forward_batch.input_ids)
+
         index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
         static_num_tokens = self.capture_num_tokens[index]
         self.raw_num_tokens = num_tokens
@@ -437,7 +472,11 @@ class PiecewiseCudaGraphRunner:
             self.out_cache_loc_swa.zero_()
         bs = forward_batch.batch_size
 
-        self.input_ids[:num_tokens].copy_(forward_batch.input_ids)
+        if self.use_input_embeds:
+            self.input_embeds[:num_tokens].copy_(forward_batch.input_embeds)
+        else:
+            self.input_ids[:num_tokens].copy_(forward_batch.input_ids)
+
         self.positions[:num_tokens].copy_(forward_batch.positions)
         self.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
         if forward_batch.out_cache_loc_swa is not None:
@@ -452,13 +491,32 @@ class PiecewiseCudaGraphRunner:
             else None
         )
 
+        if forward_batch.mrope_positions is not None:
+            self.mrope_positions[:, :num_tokens].copy_(forward_batch.mrope_positions)
+
+        if self.use_input_embeds:
+            input_ids = None
+            input_embeds = self.input_embeds[:static_num_tokens]
+        else:
+            input_ids = self.input_ids[:static_num_tokens]
+            input_embeds = None
+
+        positions = self.positions[:static_num_tokens]
+        out_cache_loc = self.out_cache_loc[:static_num_tokens]
+
+        mrope_positions = (
+            self.mrope_positions[:, :static_num_tokens]
+            if forward_batch.mrope_positions is not None
+            else None
+        )
+
         next_token_logits_buffer = None
-        mrope_positions = None
 
         static_forward_batch = ForwardBatch(
             forward_mode=forward_batch.forward_mode,
             batch_size=bs,
             input_ids=input_ids,
+            input_embeds=input_embeds,
             req_pool_indices=forward_batch.req_pool_indices,
             seq_lens=forward_batch.seq_lens,
             next_token_logits_buffer=next_token_logits_buffer,