chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[EPD][Feat]support adaptive forward (#18118)

Browse Source
This commit is contained in:
Zheng Wengang
2026-03-05 21:12:30 +08:00
committed by GitHub
parent 806d41ab65
commit 0de0d74195
5 changed files with 215 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/disaggregation/encode_receiver.py
View File

@@ -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)

1
python/sglang/srt/environ.py
View File

@@ -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)

143
python/sglang/srt/managers/mm_utils.py
View File

@@ -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"]

71
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -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,

7
python/sglang/srt/server_args.py
View File

@@ -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(