[piecewise] Refactor VLM to support input embed buffer and remove external embedder hack (#14155)

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -76,7 +76,8 @@ class CustomAllreduce:
         are in the same node.
         """
         self._IS_CAPTURING = False
-        self.disabled = True
+        self.disabled = True  # This can be modified in-place by context manager in piecewise cuda graph runner
+        self.original_disabled = True  # To store the original state
 
         if not custom_ar:
             # disable because of missing custom allreduce library
@@ -206,6 +207,7 @@ class CustomAllreduce:
             self.register_buffer(self.buffer)
 
         self.disabled = False
+        self.original_disabled = False  # Ensure original_disabled == disabled
         self.tms_cudagraph = envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
 
     @staticmethod

python/sglang/srt/managers/mm_utils.py

@@ -11,6 +11,7 @@ import numpy as np
 import torch
 from torch import nn
 
+from sglang.srt.distributed.parallel_state import get_tp_group
 from sglang.srt.layers.multimodal import gpu_tensor_hash
 from sglang.srt.managers.schedule_batch import (
     CudaIpcTensorTransportProxy,
@@ -20,6 +21,7 @@ from sglang.srt.managers.schedule_batch import (
 )
 from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.model_executor.piecewise_cuda_graph_runner import use_original_ca_comm
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once
 from sglang.utils import logger
@@ -493,7 +495,7 @@ def get_embedding_and_mask(
     return embedding, special_multimodal_mask
 
 
-def general_embed_mm_inputs(
+def embed_mm_inputs(
     mm_inputs_list: List[MultimodalInputs],
     extend_prefix_lens: List[int],
     extend_seq_lens: List[int],
@@ -658,16 +660,27 @@ def general_mm_embed_routine(
     Returns:
         Hidden states from language model forward pass
     """
+
+    tp_group = get_tp_group()
+
+    with use_original_ca_comm(tp_group):
+        # We disable custom allreduce in piecewise cuda graph.
+        # However, because we only capture the language model part, the multimodal can still use custom allreduce.
         assert hasattr(language_model, "get_input_embeddings")
         embed_tokens = language_model.get_input_embeddings()
-    if not hasattr(language_model, "pp_group") or language_model.pp_group.is_first_rank:
+        if (
+            not hasattr(language_model, "pp_group")
+            or language_model.pp_group.is_first_rank
+        ):
             if (
                 not forward_batch.forward_mode.is_decode()
                 and not forward_batch.forward_mode.is_target_verify()
                 and forward_batch.contains_mm_inputs()
             ):
                 mm_inputs_list = [
-                mm_input for mm_input in forward_batch.mm_inputs if mm_input is not None
+                    mm_input
+                    for mm_input in forward_batch.mm_inputs
+                    if mm_input is not None
                 ]
                 extend_prefix_lens = [
                     prefix_len
@@ -679,7 +692,7 @@ def general_mm_embed_routine(
                     for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu)
                     if forward_batch.mm_inputs[i] is not None
                 ]
-            inputs_embeds, other_info = general_embed_mm_inputs(
+                inputs_embeds, other_info = embed_mm_inputs(
                     mm_inputs_list=mm_inputs_list,
                     extend_prefix_lens=extend_prefix_lens,
                     extend_seq_lens=extend_seq_lens,
@@ -692,12 +705,18 @@ def general_mm_embed_routine(
                 )
                 # add for qwen3_vl deepstack
                 if use_deepstack:
-                kwargs["input_deepstack_embeds"] = other_info["input_deepstack_embeds"]
+                    kwargs["input_deepstack_embeds"] = other_info[
+                        "input_deepstack_embeds"
+                    ]
                 # once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models
                 # just being defensive here
                 forward_batch.mm_inputs = None
             else:
                 inputs_embeds = embed_tokens(input_ids)
+            # Copy to pre-allocated buffer if available (for CUDA graph address stability)
+            if forward_batch.input_embeds is not None:
+                forward_batch.input_embeds.copy_(inputs_embeds)
+                inputs_embeds = forward_batch.input_embeds
         else:
             inputs_embeds = None
 
@@ -816,295 +835,3 @@ def hash_feature(f):
         reconstruct_t = f.reconstruct_on_target_device(torch.cuda.current_device())
         return tensor_hash([reconstruct_t])
     return data_hash(f)
-
-
-def resolve_language_model(multimodal_model: nn.Module) -> Optional[nn.Module]:
-    # Qwen2-VL / Qwen3-VL Style
-    if hasattr(multimodal_model, "model"):
-        lm = getattr(multimodal_model, "model")
-        if hasattr(lm, "get_input_embeddings"):
-            return lm
-
-    # Llava / OneVision Style
-    if hasattr(multimodal_model, "language_model"):
-        lm = getattr(multimodal_model, "language_model")
-        if hasattr(lm, "get_input_embeddings"):
-            return lm
-
-    if hasattr(multimodal_model, "get_input_embeddings"):
-        return multimodal_model
-
-    return None
-
-
-def external_embed_mm_inputs(
-    forward_batch: ForwardBatch,
-    mm_inputs_list: List[MultimodalInputs],
-    extend_prefix_lens: List[int],
-    extend_seq_lens: List[int],
-    input_ids: torch.Tensor,
-    input_embedding: nn.Embedding,
-    multimodal_model: nn.Module = None,
-    data_embedding_func_mapping: Dict[
-        Modality, Callable[[List[MultimodalDataItem]], torch.Tensor]
-    ] = None,
-) -> Optional[torch.Tensor]:
-    """
-    Embed multimodal inputs and integrate them with text token embeddings.
-
-    Args:
-        mm_inputs_list: List of multimodal inputs to process
-        extend_prefix_lens: Prefix lengths for each request
-        extend_seq_lens: Sequence lengths for each request
-        input_ids: Input token IDs tensor
-        input_embedding: Embedding layer for text tokens
-
-    Returns:
-        Combined embedding tensor with multimodal content integrated
-    """
-    if mm_inputs_list is None:
-        return None
-
-    # 1. Calculate the multimodal data which exists in input_ids, with the help of pad_values
-    # we assume that multimodal data are represented with its pad_values in input_ids
-    item_flatten_list = []
-    for mm_inputs in mm_inputs_list:
-        item_flatten_list += [item for item in mm_inputs.mm_items if item is not None]
-
-    modalities, embeddings, masks = [], [], []
-
-    # 2. Get multimodal embedding separately
-    # Try get mm embedding if any
-    for modality in Modality.all():
-        items = [
-            item for item in item_flatten_list if item.is_modality(modality=modality)
-        ]
-        embedder = (
-            None
-            if data_embedding_func_mapping is None
-            else data_embedding_func_mapping.get(modality, None)
-        )
-        if embedder is None:
-            # "image", "video", etc
-            modality_id = modality.name.lower()
-            embedder = getattr(multimodal_model, f"get_{modality_id}_feature", None)
-        if len(items) != 0:
-            assert embedder is not None, f"no embedding method found for {modality}"
-            placeholder_tensor = torch.as_tensor(
-                [item.pad_value for item in items],
-                device=input_ids.device,
-            )
-            # calculate per request items length offset
-            items_size = torch.zeros(len(mm_inputs_list) + 1, dtype=int)
-            items_offsets = []
-            for i, mm_inputs in enumerate(mm_inputs_list):
-                mm_items = [
-                    item
-                    for item in mm_inputs.mm_items
-                    if item.is_modality(modality=modality)
-                ]
-                items_size[i + 1] = len(mm_items)
-                items_offsets.append(
-                    flatten_nested_list([item.offsets for item in mm_items])
-                )
-            items_size = torch.cumsum(items_size, dim=0).tolist()
-
-            embedding, mask = get_embedding_and_mask(
-                data_embedding_func=embedder,
-                embedding_items=items,
-                placeholder_tensor=placeholder_tensor,
-                input_ids=input_ids,
-                items_size=items_size,
-                prefix_length=extend_prefix_lens,
-                extend_length=extend_seq_lens,
-                items_offset_list=items_offsets,
-            )
-
-            modalities += [modality]
-            embeddings += [embedding]
-            masks += [mask]
-
-    # 3. Get input embeddings
-    vocab_size = input_embedding.num_embeddings
-    # Important: clamp after getting original multimodal regions
-    # Clamp input ids. This is because the input_ids for the multimodal tokens are
-    # filled with the hash values of the multimodal for the prefix matching in the radix attention.
-    # There values are useless because their embeddings will be replaced by vision embeddings anyway.
-    input_ids.clamp_(min=0, max=vocab_size - 1)
-    inputs_embeds = input_embedding(input_ids)
-
-    indices = []
-    for mask in masks:
-        if mask is not None:
-            indices.append(torch.where(mask.squeeze(dim=-1))[0])
-        else:
-            indices.append(None)
-
-    # only for qwen3vl right now,  replace the original use_deepstack with this method.
-    if hasattr(multimodal_model, "post_process"):
-        embeddings, forward_batch = multimodal_model.post_process(
-            inputs_embeds, modalities, embeddings, indices, forward_batch
-        )
-
-    # 4. scatter embeddings into input embedding
-    for i, modality, embedding, index in zip(
-        range(len(embeddings)), modalities, embeddings, indices
-    ):
-        if embedding is None or index is None:
-            continue
-        # in-place update
-        inputs_embeds[index] = embedding.to(inputs_embeds.device, inputs_embeds.dtype)
-
-    return inputs_embeds, forward_batch
-
-
-def should_use_external_mm_preprocess(multimodal_model: nn.Module) -> bool:
-    """Decide whether we should use our generic "external_mm_preprocess_routine".
-
-    We are adapting VLM for piecewise CUDA graph. Since the encoder's forward
-    pass cannot be executed within the model's forward pass, we need to
-    precompute image embeddings using the encoder within the model runner.
-    For models that have already been adjusted, there is a member called
-    should_use_external_mm_preprocess, which is set to True. In practice,
-    the external_mm_preprocess_routine function will be called in the
-    model_runner.forward_extend to handle multimodal inputs.
-
-    For models that have not yet been adapted, the general_mm_embed_routine
-    will still be called in the model class's forward function for processing.
-
-    Current strategy:
-        - Llava family (models with vision_tower + multi_modal_projector):
-        Their forward already calls general_mm_embed_routine and includes
-        built-in multimodal processing. If we run it again in ModelRunner,
-        it will conflict with the internal logic, so we skip it here.
-        - Others (such as Qwen2-VL / Qwen2.5-VL): use the multimodal
-        preprocessing.
-    """
-
-    cls_name = multimodal_model.__class__.__name__
-
-    external_mm_preprocess_classes = {
-        "Qwen2VLForConditionalGeneration",
-        "Qwen2_5_VLForConditionalGeneration",
-        "InternVLChatModel",
-    }
-
-    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(
-    forward_batch: ForwardBatch,
-    multimodal_model: Optional[nn.Module] = None,
-    data_embedding_funcs: Dict[
-        Modality, Callable[[List[MultimodalDataItem]], torch.Tensor]
-    ] = None,
-) -> torch.Tensor:
-    """
-    Process multimodal inputs and forward through language model.
-    Args:
-        input_ids: Input token IDs tensor
-        forward_batch: Batch information for model forward pass
-        data_embedding_funcs: A dictionary mapping from modality type to the corresponding embedding function.
-        **kwargs: Additional arguments passed to language model
-    Returns:
-        Hidden states from language model forward pass
-    """
-
-    language_model = resolve_language_model(multimodal_model)
-    if language_model is None:
-        raise ValueError(
-            f"Cannot resolve language model from {type(multimodal_model).__name__}. "
-            f"Please ensure the model has 'model' or 'language_model' attribute."
-        )
-
-    assert hasattr(language_model, "get_input_embeddings")
-    embed_tokens = language_model.get_input_embeddings()
-    if not hasattr(language_model, "pp_group") or language_model.pp_group.is_first_rank:
-
-        input_ids = forward_batch.input_ids
-        if (
-            not forward_batch.forward_mode.is_decode()
-            and not forward_batch.forward_mode.is_target_verify()
-            and forward_batch.contains_mm_inputs()
-        ):
-            mm_inputs_list = [
-                mm_input for mm_input in forward_batch.mm_inputs if mm_input is not None
-            ]
-            extend_prefix_lens = [
-                prefix_len
-                for i, prefix_len in enumerate(forward_batch.extend_prefix_lens_cpu)
-                if forward_batch.mm_inputs[i] is not None
-            ]
-            extend_seq_lens = [
-                seq_len
-                for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu)
-                if forward_batch.mm_inputs[i] is not None
-            ]
-            input_embeds, forward_batch = external_embed_mm_inputs(
-                forward_batch=forward_batch,
-                mm_inputs_list=mm_inputs_list,
-                extend_prefix_lens=extend_prefix_lens,
-                extend_seq_lens=extend_seq_lens,
-                input_ids=forward_batch.input_ids,
-                multimodal_model=multimodal_model,
-                input_embedding=embed_tokens,
-                data_embedding_func_mapping=data_embedding_funcs,
-            )
-            # once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models
-            # just being defensive here
-            forward_batch.mm_inputs = None
-        else:
-            # NOTE: This may reduce the performance for only-text inputs.
-            # Using a fixed-address buffer might be better, though it could be a bit dirty.
-            input_embeds = embed_tokens(input_ids)
-            # only for qwen3vl
-            if getattr(multimodal_model, "use_deepstack", False):
-                forward_batch.input_deepstack_embeds = torch.zeros(
-                    (
-                        len(input_ids),
-                        multimodal_model.config.hidden_size
-                        * len(multimodal_model.deepstack_visual_indexes),
-                    ),
-                    device=input_embeds.device,
-                    dtype=input_embeds.dtype,
-                )
-
-        forward_batch.input_embeds = input_embeds
-    else:
-        forward_batch.input_embeds = None
-
-    return forward_batch

python/sglang/srt/model_executor/model_runner.py

@@ -97,11 +97,6 @@ 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,
-    resolve_external_mm_data_embedding_funcs,
-    should_use_external_mm_preprocess,
-)
 from sglang.srt.mem_cache.allocator import (
     BaseTokenToKVPoolAllocator,
     PagedTokenToKVPoolAllocator,
@@ -2548,15 +2543,6 @@ class ModelRunner:
         skip_attn_backend_init: bool = False,
         pp_proxy_tensors=None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
-
-        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 = {}
         if self.support_pp:
             kwargs["pp_proxy_tensors"] = pp_proxy_tensors

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -70,15 +70,38 @@ def disable_ca_comm(tp_group):
 
     TODO(yuwei): Fix this
     """
-    old_disabled = None
+    if tp_group.ca_comm is None:
+        yield
+        return
+
+    original_disabled = tp_group.ca_comm.disabled
+    tp_group.ca_comm.original_disabled = original_disabled
     try:
-        if tp_group.ca_comm is not None:
-            old_disabled = tp_group.ca_comm.disabled
         tp_group.ca_comm.disabled = True
         yield
     finally:
-        if tp_group.ca_comm is not None and old_disabled is not None:
+        tp_group.ca_comm.disabled = original_disabled
-            tp_group.ca_comm.disabled = old_disabled
+
+
+@contextmanager
+def use_original_ca_comm(tp_group):
+    """
+    For the module not in piecewise cuda graph capture, use the original custom allreduce communication.
+    This is a no-op if not using piecewise cuda graph because .disabled == .original_disabled
+
+    TODO(Byron): remove this once custom allreduce is enabled in piecewise cuda graph
+    """
+    if tp_group.ca_comm is None:
+        yield
+        return
+
+    current_disabled = tp_group.ca_comm.disabled
+    original_disabled = tp_group.ca_comm.original_disabled
+    try:
+        tp_group.ca_comm.disabled = original_disabled
+        yield
+    finally:
+        tp_group.ca_comm.disabled = current_disabled
 
 
 @contextmanager
@@ -189,15 +212,11 @@ class PiecewiseCudaGraphRunner:
 
         self.max_num_tokens = max(self.capture_num_tokens)
 
-        self.use_input_embeds = model_runner.is_multimodal
+        self.is_multimodal = 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()
             )
@@ -207,10 +226,22 @@ class PiecewiseCudaGraphRunner:
                 else None
             )
             self.positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
+
+            self.tbo_plugin = TboCudaGraphRunnerPlugin()
+
+            if (
+                self.is_multimodal
+            ):  # Only create input_embeds and mrope_positions for multimodal model to save memory
+                # 1. In multimodal, we only compile and capture the language model part.
+                # 2. The embedder is outside of the graph, but cuda graph requires the input embeds to have a fixed memory address.
+                # 3. Input embeds is a pre-allocated buffer. In model.forward, we copy the embed output to this buffer.
+                self.input_embeds = torch.zeros(
+                    (self.max_num_tokens, self.model_runner.model_config.hidden_size),
+                    dtype=self.model_runner.dtype,
+                )
                 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
 
@@ -258,11 +289,7 @@ class PiecewiseCudaGraphRunner:
             forward_batch = ForwardBatch(
                 forward_mode=ForwardMode.EXTEND,
                 batch_size=1,
-                input_ids=(
+                input_ids=(torch.randint(0, 100, (num_tokens,), device=self.device)),
-                    torch.randint(0, 100, (num_tokens,), device=self.device)
-                    if not self.use_input_embeds
-                    else None
-                ),
                 input_embeds=(
                     torch.randn(
                         num_tokens,
@@ -270,7 +297,7 @@ class PiecewiseCudaGraphRunner:
                         dtype=self.model_runner.dtype,
                         device=self.device,
                     )
-                    if self.use_input_embeds
+                    if self.is_multimodal
                     else None
                 ),
                 req_pool_indices=torch.arange(1, device=self.device),
@@ -298,7 +325,9 @@ 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=self.mrope_positions[:, :num_tokens],
+                mrope_positions=(
+                    self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
+                ),
                 spec_algorithm=None,
                 spec_info=None,
                 capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -374,12 +403,8 @@ class PiecewiseCudaGraphRunner:
         bs = 1
 
         # Graph inputs
-        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
+        input_embeds = self.input_embeds[:num_tokens] if self.is_multimodal else None
 
         out_cache_loc = self.out_cache_loc[:num_tokens]
         out_cache_loc_swa = (
@@ -388,12 +413,8 @@ class PiecewiseCudaGraphRunner:
             else None
         )
         positions = self.positions[:num_tokens]
-        mrope_positions = self.mrope_positions[:, :num_tokens]
+        mrope_positions = (
-
+            self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
-        # pipeline parallelism
-        if self.pp_size > 1:
-            pp_proxy_tensors = PPProxyTensors(
-                {k: v[:num_tokens] for k, v in self.pp_proxy_tensors.items()}
         )
 
         global_dp_buffer_len = None
@@ -488,11 +509,7 @@ class PiecewiseCudaGraphRunner:
         forward_batch: ForwardBatch,
         **kwargs,
     ):
-        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
@@ -502,11 +519,7 @@ class PiecewiseCudaGraphRunner:
                 self.out_cache_loc_swa.zero_()
         bs = forward_batch.batch_size
 
-        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 self.out_cache_loc_swa is not None:
@@ -528,12 +541,10 @@ class PiecewiseCudaGraphRunner:
         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
+        input_embeds = (
+            self.input_embeds[:static_num_tokens] if self.is_multimodal else None
+        )
 
         mrope_positions = (
             self.mrope_positions[:, :static_num_tokens]

python/sglang/srt/models/internvl.py

@@ -19,7 +19,10 @@ from sglang.srt.layers.attention.vision import SingletonCache, VisionAttention
 from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 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
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternTokenPairs,
+    general_mm_embed_routine,
+)
 from sglang.srt.managers.schedule_batch import (
     Modality,
     MultimodalDataItem,
@@ -592,21 +595,12 @@ class InternVLChatModel(nn.Module):
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
 
-        input_embeds = forward_batch.input_embeds
+        hidden_states = general_mm_embed_routine(
-        # 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,
-            input_embeds=input_embeds,
+            language_model=self.language_model,
+            multimodal_model=self,
+            data_embedding_funcs=self.external_mm_data_embedding_funcs,
             positions=positions,
         )
 

python/sglang/srt/models/qwen2_5_vl.py

@@ -57,7 +57,10 @@ from sglang.srt.layers.pooler import Pooler, PoolingType
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
-from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternMultimodalTokens,
+    general_mm_embed_routine,
+)
 from sglang.srt.managers.schedule_batch import (
     Modality,
     MultimodalDataItem,
@@ -624,21 +627,11 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
                     f"(3, seq_len) positions, but got {positions.size()}"
                 )
 
-        input_embeds = forward_batch.input_embeds
+        hidden_states = general_mm_embed_routine(
-        # 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.model(
             input_ids=input_ids,
             forward_batch=forward_batch,
-            input_embeds=input_embeds,
+            language_model=self.model,
+            multimodal_model=self,
             positions=positions,
             pp_proxy_tensors=pp_proxy_tensors,
         )

python/sglang/srt/models/qwen2_vl.py

@@ -39,7 +39,10 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.pooler import Pooler, PoolingType
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
-from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternMultimodalTokens,
+    general_mm_embed_routine,
+)
 from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
@@ -532,21 +535,11 @@ class Qwen2VLForConditionalGeneration(nn.Module):
                     f"(3, seq_len) positions, but got {positions.size()}"
                 )
 
-        input_embeds = forward_batch.input_embeds
+        hidden_states = general_mm_embed_routine(
-        # 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.model(
             input_ids=input_ids,
             forward_batch=forward_batch,
-            input_embeds=input_embeds,
+            language_model=self.model,
+            multimodal_model=self,
             positions=positions,
         )