diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 6af99220c..34cbc1220 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -434,22 +434,29 @@ class MMReceiverBase(ABC): ) else: raise e + + # Skip mm_pool if not adaptive dispatch to encoder + enable_adaptive_dispatch_to_encoder = ( + server_args.enable_adaptive_dispatch_to_encoder + ) self.mm_processor = get_mm_processor( hf_config, server_args, _processor, transport_mode, - skip_mm_pool=True, + skip_mm_pool=not enable_adaptive_dispatch_to_encoder, ) @abstractmethod def process_waiting_requests(self, recv_reqs): pass - async def recv_mm_data(self, img_data, mm_processor, prompt): + async def recv_mm_data( + self, img_data, mm_processor, prompt, need_wait_for_image=True + ): req_id = None try: - if len(self.encode_urls) == 0: + if len(self.encode_urls) == 0 or not need_wait_for_image: return None req_id = uuid.uuid4().hex embedding_port, recv_socket = get_zmq_socket_on_host(self.context, zmq.PULL) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 22160f6a0..882da78d2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -503,6 +503,7 @@ class Envs: # EPD SGLANG_ENCODER_RECV_TIMEOUT = EnvFloat(180.0) SGLANG_ENCODER_SEND_TIMEOUT = EnvFloat(180.0) + SGLANG_ENCODER_DISPATCH_MIN_ITEMS = EnvInt(2) # Elastic EP Backup Port SGLANG_BACKUP_PORT_BASE = EnvInt(10000) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 672fe2384..48cd1e7fc 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1045,6 +1045,110 @@ def embed_mm_inputs( return input_embeds, other_info +def _embed_mm_inputs_with_split( + mm_inputs_list: List[MultimodalInputs], + extend_prefix_lens: List[int], + extend_seq_lens: List[int], + input_ids: torch.Tensor, + forward_batch: ForwardBatch, + input_embedding: nn.Embedding, + multimodal_model: nn.Module = None, + data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None, + placeholder_tokens: dict[Modality, List[int]] = None, + use_deepstack: Dict[Modality, bool] = {}, +): + """Split batch into precomputed vs non-precomputed, embed each group, merge back.""" + precomputed_req_indices = [] + non_precomputed_req_indices = [] + for idx, mm_input in enumerate(mm_inputs_list): + items = [item for item in mm_input.mm_items if item is not None] + if items and all( + getattr(item, "precomputed_embeddings", None) is not None for item in items + ): + precomputed_req_indices.append(idx) + else: + non_precomputed_req_indices.append(idx) + + embed_kwargs = dict( + multimodal_model=multimodal_model, + input_embedding=input_embedding, + data_embedding_func_mapping=data_embedding_func_mapping, + placeholder_tokens=placeholder_tokens, + use_deepstack=use_deepstack, + ) + + if not precomputed_req_indices or not non_precomputed_req_indices: + return 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, + **embed_kwargs, + ) + + all_seq_lens = forward_batch.extend_seq_lens_cpu + mm_batch_indices = [ + i for i, mm in enumerate(forward_batch.mm_inputs) if mm is not None + ] + token_starts = [] + cumulative = 0 + for sl in all_seq_lens: + token_starts.append(cumulative) + cumulative += sl + + vocab_size = input_embedding.num_embeddings + input_embeds = input_embedding(input_ids.clamp(min=0, max=vocab_size - 1)) + other_info = {} + + input_deepstack_embeds = None + if use_deepstack and multimodal_model is not None: + num_deepstack_embeddings = len(multimodal_model.deepstack_visual_indexes) + input_deepstack_embeds = torch.zeros( + input_ids.shape[0], + input_embedding.embedding_dim * num_deepstack_embeddings, + device=input_ids.device, + dtype=input_embedding.weight.dtype, + ) + other_info["input_deepstack_embeds"] = input_deepstack_embeds + + for group_req_indices in [precomputed_req_indices, non_precomputed_req_indices]: + sub_mm_inputs = [mm_inputs_list[i] for i in group_req_indices] + sub_prefix_lens = [extend_prefix_lens[i] for i in group_req_indices] + sub_seq_lens = [extend_seq_lens[i] for i in group_req_indices] + group_batch_indices = [mm_batch_indices[i] for i in group_req_indices] + sub_slices = [ + input_ids[token_starts[bi] : token_starts[bi] + all_seq_lens[bi]] + for bi in group_batch_indices + ] + sub_input_ids = torch.cat(sub_slices) + + sub_embeds, sub_info = embed_mm_inputs( + mm_inputs_list=sub_mm_inputs, + extend_prefix_lens=sub_prefix_lens, + extend_seq_lens=sub_seq_lens, + input_ids=sub_input_ids, + **embed_kwargs, + ) + + offset = 0 + for bi in group_batch_indices: + req_len = all_seq_lens[bi] + start = token_starts[bi] + input_embeds[start : start + req_len] = sub_embeds[ + offset : offset + req_len + ] + if ( + input_deepstack_embeds is not None + and "input_deepstack_embeds" in sub_info + ): + input_deepstack_embeds[start : start + req_len] = sub_info[ + "input_deepstack_embeds" + ][offset : offset + req_len] + offset += req_len + + return input_embeds, other_info + + def general_mm_embed_routine( input_ids: torch.Tensor, forward_batch: ForwardBatch, @@ -1091,17 +1195,34 @@ 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 ] - input_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, - ) + server_args = get_global_server_args() + if server_args and server_args.enable_adaptive_dispatch_to_encoder: + # Split by precomputed vs non-precomputed so get_embedding_and_mask only sees uniform batches + input_embeds, other_info = _embed_mm_inputs_with_split( + mm_inputs_list=mm_inputs_list, + extend_prefix_lens=extend_prefix_lens, + extend_seq_lens=extend_seq_lens, + input_ids=input_ids, + forward_batch=forward_batch, + input_embedding=embed_tokens, + multimodal_model=multimodal_model, + data_embedding_func_mapping=data_embedding_funcs, + placeholder_tokens=placeholder_tokens, + use_deepstack=use_deepstack, + ) + else: + input_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, + input_embedding=embed_tokens, + multimodal_model=multimodal_model, + 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"] diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 05a829002..7dca38d6c 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -726,6 +726,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi img_data=obj.image_data, mm_processor=self.mm_processor, prompt=(input_text or input_ids), + need_wait_for_image=obj.need_wait_for_image, ) if mm_inputs is None: mm_inputs: Dict = await self.mm_data_processor.process( @@ -735,6 +736,20 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi request_obj=obj, max_req_input_len=self.max_req_input_len, ) + elif ( + self.server_args.language_only + and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" + and not obj.need_wait_for_image + ): + # In language_only mode with zmq_to_scheduler, if we didn't dispatch + # to encoder (e.g., only one image), process locally like non-language_only mode + mm_inputs: Dict = await self.mm_data_processor.process( + image_data=obj.image_data, + audio_data=obj.audio_data, + input_text_or_ids=(input_text or input_ids), + request_obj=obj, + max_req_input_len=self.max_req_input_len, + ) if mm_inputs and "input_ids" in mm_inputs: input_ids = mm_inputs["input_ids"] @@ -2294,16 +2309,60 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi ) time_stats.set_created_time(created_time) + def _should_dispatch_to_encoder( + self, obj: Union[GenerateReqInput, EmbeddingReqInput] + ) -> bool: + """Check if the request should be dispatched to encoder for processing. + + Returns True if the request should be dispatched to encoder (multiple multimodal items), + False if it should be processed locally (single multimodal item or no multimodal items). + + Args: + obj: The request input object + + Returns: + bool: True if should dispatch to encoder, False otherwise + """ + if obj.batch_size > 1: + logger.warning( + "Batch request (batch_size=%d) is not supported in EPD disaggregation mode; skipping encoder dispatch.", + obj.batch_size, + ) + return False + if not isinstance(obj, GenerateReqInput) or not obj.contains_mm_input(): + return False + + # Count image / video / audio items for dispatch threshold + def _count_mm_items(data): + return ( + len(data) if isinstance(data, list) else (1 if data is not None else 0) + ) + + total_mm_items = ( + _count_mm_items(getattr(obj, "image_data", None)) + + _count_mm_items(getattr(obj, "video_data", None)) + + _count_mm_items(getattr(obj, "audio_data", None)) + ) + return total_mm_items >= envs.SGLANG_ENCODER_DISPATCH_MIN_ITEMS.get() + def _handle_epd_disaggregation_encode_request( self, obj: Union[GenerateReqInput, EmbeddingReqInput] ): """Handle EPD-disaggregation mode encoding request.""" - if ( - isinstance(obj, GenerateReqInput) - and self.server_args.encoder_transfer_backend == "zmq_to_scheduler" - and obj.contains_mm_input() - ): - self.mm_receiver.send_encode_request(obj) + if isinstance(obj, GenerateReqInput) and obj.contains_mm_input(): + # dispatch to encoder by default + should_dispatch = True + if self.server_args.enable_adaptive_dispatch_to_encoder: + should_dispatch = self._should_dispatch_to_encoder(obj) + + # Set need_wait_for_image flag based on whether we dispatch to encoder + # This flag will be used in _tokenize_one_request to determine processing path + if should_dispatch: + obj.need_wait_for_image = True + if self.server_args.encoder_transfer_backend == "zmq_to_scheduler": + self.mm_receiver.send_encode_request(obj) + else: + obj.need_wait_for_image = False def convert_to_span_attrs( self, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9893ee243..83c69bbb2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -685,6 +685,7 @@ class ServerArgs: language_only: bool = False encoder_transfer_backend: str = ENCODER_TRANSFER_BACKEND_CHOICES[0] encoder_urls: List[str] = dataclasses.field(default_factory=list) + enable_adaptive_dispatch_to_encoder: bool = False # For model weight update and weight loading custom_weight_loader: Optional[List[str]] = None @@ -5279,6 +5280,12 @@ class ServerArgs: default=[], help="List of encoder server urls.", ) + parser.add_argument( + "--enable-adaptive-dispatch-to-encoder", + default=ServerArgs.enable_adaptive_dispatch_to_encoder, + action="store_true", + help="When enabled, adaptively dispatch: multi-image requests go to encoder in language_only epd mode, single-image requests are processed locally.", + ) # Custom weight loader parser.add_argument(