From 0825d7f4c6b30d6079b1c8b4cc3aabccfe837def Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Sun, 30 Nov 2025 21:43:09 -0800 Subject: [PATCH] [piecewise] Refactor VLM to support input embed buffer and remove external embedder hack (#14155) --- .../device_communicators/custom_all_reduce.py | 4 +- python/sglang/srt/managers/mm_utils.py | 391 +++--------------- .../sglang/srt/model_executor/model_runner.py | 14 - .../piecewise_cuda_graph_runner.py | 111 ++--- python/sglang/srt/models/internvl.py | 22 +- python/sglang/srt/models/qwen2_5_vl.py | 21 +- python/sglang/srt/models/qwen2_vl.py | 21 +- 7 files changed, 145 insertions(+), 439 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index c72d8b9a0..b768d46bc 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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 diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index cd566b7c8..f55f2e1d9 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/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,48 +660,65 @@ def general_mm_embed_routine( Returns: Hidden states from language model forward pass """ - 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: + + 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 forward_batch.forward_mode.is_decode() - and not forward_batch.forward_mode.is_target_verify() - and forward_batch.contains_mm_inputs() + not hasattr(language_model, "pp_group") + or language_model.pp_group.is_first_rank ): - 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 - ] - inputs_embeds, other_info = general_embed_mm_inputs( - mm_inputs_list=mm_inputs_list, - extend_prefix_lens=extend_prefix_lens, - extend_seq_lens=extend_seq_lens, - input_ids=input_ids, - multimodal_model=multimodal_model, - input_embedding=embed_tokens, - data_embedding_func_mapping=data_embedding_funcs, - placeholder_tokens=placeholder_tokens, - use_deepstack=use_deepstack, - ) - # add for qwen3_vl deepstack - if use_deepstack: - 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 + 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 + ] + 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, + input_ids=input_ids, + multimodal_model=multimodal_model, + input_embedding=embed_tokens, + data_embedding_func_mapping=data_embedding_funcs, + placeholder_tokens=placeholder_tokens, + use_deepstack=use_deepstack, + ) + # add for qwen3_vl deepstack + if use_deepstack: + 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 = embed_tokens(input_ids) - else: - inputs_embeds = None + inputs_embeds = None hidden_states = language_model( input_ids=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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index b70fb1fc2..8978b821b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index e80390896..a5c1f4d4e 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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 + 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 = old_disabled + tp_group.ca_comm.disabled = original_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,11 +226,23 @@ class PiecewiseCudaGraphRunner: else None ) 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() + 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.attention_layers = self.model_runner.attention_layers if get_global_graph_memory_pool() is None: @@ -258,11 +289,7 @@ class PiecewiseCudaGraphRunner: forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, batch_size=1, - input_ids=( - torch.randint(0, 100, (num_tokens,), device=self.device) - if not self.use_input_embeds - else None - ), + input_ids=(torch.randint(0, 100, (num_tokens,), device=self.device)), 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_ids = self.input_ids[:num_tokens] + 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,13 +413,9 @@ class PiecewiseCudaGraphRunner: else None ) positions = self.positions[:num_tokens] - mrope_positions = self.mrope_positions[:, :num_tokens] - - # pipeline parallelism - if self.pp_size > 1: - pp_proxy_tensors = PPProxyTensors( - {k: v[:num_tokens] for k, v in self.pp_proxy_tensors.items()} - ) + mrope_positions = ( + self.mrope_positions[:, :num_tokens] if self.is_multimodal else None + ) 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) - + 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.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_ids = self.input_ids[:static_num_tokens] + input_embeds = ( + self.input_embeds[:static_num_tokens] if self.is_multimodal else None + ) mrope_positions = ( self.mrope_positions[:, :static_num_tokens] diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index 105f2d33a..b23b1b61b 100644 --- a/python/sglang/srt/models/internvl.py +++ b/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 - # 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( + hidden_states = general_mm_embed_routine( 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, ) diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index 86df6e69f..fdc0c7650 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/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 - # 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( + hidden_states = general_mm_embed_routine( 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, ) diff --git a/python/sglang/srt/models/qwen2_vl.py b/python/sglang/srt/models/qwen2_vl.py index 4518d0879..fa4055c36 100644 --- a/python/sglang/srt/models/qwen2_vl.py +++ b/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 - # 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( + hidden_states = general_mm_embed_routine( input_ids=input_ids, forward_batch=forward_batch, - input_embeds=input_embeds, + language_model=self.model, + multimodal_model=self, positions=positions, )