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

Revert "[FEAT] optimize tensor zmq transfer for multimodal inputs" (#16386)

Browse Source
This commit is contained in:
Yuhao Yang
2026-01-05 14:05:26 +08:00
committed by GitHub
parent e53160bb31
commit 2138ff48c6
5 changed files with 23 additions and 271 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
python/sglang/srt/managers/data_parallel_controller.py
View File

@@ -16,7 +16,6 @@
import faulthandler
import logging
import multiprocessing as mp
import pickle
import signal
import threading
import time
@@ -26,7 +25,6 @@ from typing import Callable, List, Optional
import psutil
import setproctitle
import torch
import zmq
from sglang.srt.environ import envs
@@ -210,15 +208,13 @@ class DataParallelController:
)
def send_to_all_workers(self, obj):
msg = [b"NORM", pickle.dumps(obj)]
for worker in self.workers:
worker.send_multipart(msg, copy=False)
worker.send_pyobj(obj)
def send_control_message(self, obj):
msg = [b"NORM", pickle.dumps(obj)]
# Send control messages to first worker of tp group
for worker in self.workers[:: self.control_message_step]:
worker.send_multipart(msg, copy=False)
worker.send_pyobj(obj)
def handle_load_update_req(self, obj):
self.dp_budget.update_budget(obj)
@@ -503,9 +499,8 @@ class DataParallelController:
def maybe_external_dp_rank_routing(self, req: Req):
if req.data_parallel_rank is not None:
msg = [b"NORM", pickle.dumps(req)]
logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
self.workers[req.data_parallel_rank].send_multipart(msg, copy=False)
self.workers[req.data_parallel_rank].send_pyobj(req)
return True
return False
@@ -513,8 +508,7 @@ class DataParallelController:
if self.maybe_external_dp_rank_routing(req):
return
msg = [b"NORM", pickle.dumps(req)]
self.workers[self.round_robin_counter].send_multipart(msg, copy=False)
self.workers[self.round_robin_counter].send_pyobj(req)
self.round_robin_counter = (self.round_robin_counter + 1) % len(self.workers)
def follow_bootstrap_room_scheduler(self, req: Req):
@@ -536,8 +530,7 @@ class DataParallelController:
"prefill or decode instances; send to the router instead."
)
target_rank = req.bootstrap_room % len(self.workers)
msg = [b"NORM", pickle.dumps(req)]
self.workers[target_rank].send_multipart(msg, copy=False)
self.workers[target_rank].send_pyobj(req)
def shortest_queue_scheduler(self, req):
if self.maybe_external_dp_rank_routing(req):
@@ -549,8 +542,7 @@ class DataParallelController:
else:
self.follow_bootstrap_room_scheduler(req)
else:
msg = [b"NORM", pickle.dumps(req)]
self.workers[target_worker].send_multipart(msg, copy=False)
self.workers[target_worker].send_pyobj(req)
def minimum_tokens_scheduler(self, req):
if self.maybe_external_dp_rank_routing(req):
@@ -565,87 +557,12 @@ class DataParallelController:
else:
self.follow_bootstrap_room_scheduler(req)
def _parse_multipart_message(self, parts):
# Check message type
msg_type = bytes(parts[0])
if msg_type == b"NORM":
# Normal message
recv_req = pickle.loads(parts[1])
elif msg_type == b"FEAT":
# Message with optimized feature tensors
recv_req = pickle.loads(parts[1])
feature_infos = pickle.loads(parts[2])
# Reconstruct tensors
for i, feature_info in enumerate(feature_infos):
buffer_idx = 3 + i
buffer = (
parts[buffer_idx].buffer
if hasattr(parts[buffer_idx], "buffer")
else parts[buffer_idx]
)
dtype = feature_info["dtype"]
shape = feature_info["shape"]
tensor = torch.frombuffer(buffer, dtype=dtype).reshape(shape)
idx = feature_info["idx"]
if hasattr(recv_req, "mm_inputs") and recv_req.mm_inputs:
mm_items = recv_req.mm_inputs.get("mm_items", [])
if idx < len(mm_items):
mm_items[idx].feature = tensor
else:
logger.warning(f"Unknown message type: {msg_type}")
return recv_req
def _parse_multipart_message(self, parts):
# Check message type
msg_type = bytes(parts[0])
if msg_type == b"NORM":
# Normal message
recv_req = pickle.loads(parts[1])
elif msg_type == b"FEAT":
# Message with optimized feature tensors
recv_req = pickle.loads(parts[1])
feature_infos = pickle.loads(parts[2])
# Reconstruct tensors
for i, feature_info in enumerate(feature_infos):
buffer_idx = 3 + i
buffer = (
parts[buffer_idx].buffer
if hasattr(parts[buffer_idx], "buffer")
else parts[buffer_idx]
)
dtype = feature_info["dtype"]
shape = feature_info["shape"]
tensor = torch.frombuffer(buffer, dtype=dtype).reshape(shape)
idx = feature_info["idx"]
if hasattr(recv_req, "mm_inputs") and recv_req.mm_inputs:
mm_items = recv_req.mm_inputs.get("mm_items", [])
if idx < len(mm_items):
mm_items[idx].feature = tensor
else:
logger.warning(f"Unknown message type: {msg_type}")
return recv_req
def event_loop(self):
while True:
while True:
self.soft_watchdog.feed()
try:
parts = self.recv_from_tokenizer.recv_multipart(
flags=zmq.NOBLOCK, copy=False
)
if not parts:
break
recv_req = self._parse_multipart_message(parts)
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
except zmq.ZMQError:
break
self._request_dispatcher(recv_req)