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

[PP] Refactor PP to async mode (#11852)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: bluecoffee8 <jasperli2002@gmail.com>
Co-authored-by: zhangxiaolei123456 <zhangxiaolei.666@bytedance.com>
Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com>
This commit is contained in:
Xuchun Shang
2025-12-12 12:54:16 +08:00
committed by GitHub
co-authored by Lianmin Zheng Shangming Cai bluecoffee8 zhangxiaolei123456 ybyang
parent 8f5adac8c6
commit c01b2ee094
7 changed files with 1051 additions and 319 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
-36
View File
@@ -55,7 +55,6 @@ from sglang.srt.mem_cache.memory_pool import (
SWAKVPool,
)
from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj
if TYPE_CHECKING:
from torch.distributed import ProcessGroup
@@ -252,8 +251,6 @@ class PrefillBootstrapQueue:
# if req not in reqs_info_to_check, skip
if req.rid not in rids_to_check:
continue
# Either waiting for input or failed
assert poll == KVPoll.WaitingForInput or poll == KVPoll.Failed
if poll == KVPoll.Bootstrapping:
continue
@@ -710,36 +707,3 @@ class SchedulerDisaggregationPrefillMixin:
)
return
req.disagg_kv_sender.send(page_indices, state_indices)
def send_pyobj_to_next_stage(self, data):
if self.attn_tp_rank == 0:
dp_offset = self.attn_dp_rank * self.attn_tp_size
point_to_point_pyobj(
data,
self.pp_rank * self.tp_size + dp_offset,
self.world_group.device_group,
self.pp_rank * self.tp_size + dp_offset,
((self.pp_rank + 1) % self.pp_size) * self.tp_size + dp_offset,
)
def recv_pyobj_from_prev_stage(self):
if self.attn_tp_rank == 0:
dp_offset = self.attn_dp_rank * self.attn_tp_size
data = point_to_point_pyobj(
[],
self.pp_rank * self.tp_size + dp_offset,
self.world_group.device_group,
((self.pp_rank - 1) % self.pp_size) * self.tp_size + dp_offset,
self.pp_rank * self.tp_size + dp_offset,
)
else:
data = None
if self.attn_tp_size != 1:
data = broadcast_pyobj(
data,
self.attn_tp_group.rank,
self.attn_tp_cpu_group,
src=self.attn_tp_group.ranks[0],
)
return data
python/sglang/srt/environ.py
+1
View File
@@ -194,6 +194,7 @@ class Envs:
SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False)
SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1)
SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
# Test: pd-disaggregation
SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")
python/sglang/srt/managers/scheduler.py
+42 -7
View File
@@ -155,7 +155,7 @@ from sglang.srt.managers.utils import GenerationBatchResult, validate_input_leng
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.radix_cache import RadixCache
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
@@ -472,6 +472,21 @@ class Scheduler(
self.chunked_prefill_size is not None and server_args.enable_mixed_chunk
)
self.enable_dynamic_chunking = (
server_args.enable_dynamic_chunking and self.pp_size > 1
)
# Init the dynamic chunking predictor for PP
if self.enable_dynamic_chunking:
try:
self.profile_and_init_predictor()
except Exception as e:
logger.warning(
f"[PP Dynamic Chunk] Failed to profile prefill latency: {e}. "
"Dynamic chunking will be disabled."
)
self.enable_dynamic_chunking = False
# Init the grammar backend for constrained generation
self.grammar_queue: List[Req] = []
if not server_args.skip_tokenizer_init:
@@ -934,8 +949,7 @@ class Scheduler(
def init_overlap(self):
self.future_map = None
if not self.enable_overlap:
if not self.enable_overlap and self.pp_size == 1:
return
self.forward_stream: CudaStream = torch.get_device_module(self.device).Stream()
@@ -947,6 +961,9 @@ class Scheduler(
self.device
).stream(self.copy_stream)
if not self.enable_overlap:
return
self.future_map = FutureMap(
self.max_running_requests,
self.chunked_prefill_size,
@@ -1108,7 +1125,7 @@ class Scheduler(
recv_reqs = point_to_point_pyobj(
[],
self.pp_rank * self.tp_size + dp_offset,
self.world_group.device_group,
self.world_group.cpu_group,
(self.pp_rank - 1) * self.tp_size + dp_offset,
self.pp_rank * self.tp_size + dp_offset,
)
@@ -1766,6 +1783,16 @@ class Scheduler(
# in the waiting queue.
return None
# Determine chunked_prefill_size for this batch
chunked_prefill_size = self.chunked_prefill_size
if self.chunked_req is not None:
self.chunked_req.init_next_round_input()
if self.enable_dynamic_chunking:
history_len = len(self.chunked_req.prefix_indices)
dynamic_size = self.predict_next_chunk_size(history_len)
if dynamic_size is not None:
chunked_prefill_size = dynamic_size
# Prefill policy
adder = PrefillAdder(
self.page_size,
@@ -1774,7 +1801,7 @@ class Scheduler(
self.running_batch,
self.new_token_ratio,
self.max_prefill_tokens,
self.chunked_prefill_size,
chunked_prefill_size,
running_bs if self.is_mixed_chunk else 0,
self.priority_scheduling_preemption_threshold,
)
@@ -1966,7 +1993,9 @@ class Scheduler(
pass
def run_batch(
self, batch: ScheduleBatch
self,
batch: ScheduleBatch,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> Union[GenerationBatchResult, EmbeddingBatchResult]:
"""Run a batch."""
self.forward_ct += 1
@@ -2014,6 +2043,7 @@ class Scheduler(
self.future_map.resolve_future(model_worker_batch)
batch_result = self.model_worker.forward_batch_generation(
model_worker_batch
# here pp is not compatible with overlap
)
# FIXME(lsyin): maybe move this to forward_batch_generation
batch_result.copy_done = torch.get_device_module(
@@ -2047,8 +2077,13 @@ class Scheduler(
batch_result = self.tp_worker.forward_batch_split_prefill(batch)
future_indices_or_next_token_ids = batch_result.next_token_ids
else:
kwargs = (
{"pp_proxy_tensors": pp_proxy_tensors}
if self.spec_algorithm.is_none()
else {}
)
batch_result = self.model_worker.forward_batch_generation(
batch_or_worker_batch
batch_or_worker_batch, **kwargs
)
future_indices_or_next_token_ids = batch_result.next_token_ids
self.update_cache_from_scheduler(batch, batch_result)
python/sglang/srt/managers/scheduler_pp_mixin.py
+955 -250
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/managers/tp_worker.py
+1 -8
View File
@@ -376,6 +376,7 @@ class TpModelWorker(BaseTpWorker):
self,
model_worker_batch: ModelWorkerBatch,
forward_batch: Optional[ForwardBatch] = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
is_verify: bool = False,
skip_attn_backend_init=False,
) -> GenerationBatchResult:
@@ -391,14 +392,6 @@ class TpModelWorker(BaseTpWorker):
# FIXME(lsyin): unify the interface of forward_batch
assert forward_batch is not None
pp_proxy_tensors = None
if not self.pp_group.is_first_rank:
pp_proxy_tensors = PPProxyTensors(
self.pp_group.recv_tensor_dict(
all_gather_group=self.get_attention_tp_group()
)
)
if self.pp_group.is_last_rank:
if self.is_dllm():
return self._forward_batch_generation_dllm(forward_batch)
python/sglang/srt/server_args.py
+14
View File
@@ -290,6 +290,7 @@ class ServerArgs:
max_queued_requests: Optional[int] = None
max_total_tokens: Optional[int] = None
chunked_prefill_size: Optional[int] = None
enable_dynamic_chunking: bool = False
max_prefill_tokens: int = 16384
schedule_policy: str = "fcfs"
enable_priority_scheduling: bool = False
@@ -308,6 +309,7 @@ class ServerArgs:
tp_size: int = 1
pp_size: int = 1
pp_max_micro_batch_size: Optional[int] = None
pp_async_batch_depth: int = 0
stream_interval: int = 1
stream_output: bool = False
random_seed: Optional[int] = None
@@ -2516,6 +2518,12 @@ class ServerArgs:
default=ServerArgs.chunked_prefill_size,
help="The maximum number of tokens in a chunk for the chunked prefill. Setting this to -1 means disabling chunked prefill.",
)
parser.add_argument(
"--enable-dynamic-chunking",
action="store_true",
default=ServerArgs.enable_dynamic_chunking,
help="Enable dynamic chunk size adjustment for pipeline parallelism. When enabled, chunk sizes are dynamically calculated based on fitted function to maintain consistent execution time across chunks.",
)
parser.add_argument(
"--max-prefill-tokens",
type=int,
@@ -2624,6 +2632,12 @@ class ServerArgs:
default=ServerArgs.pp_max_micro_batch_size,
help="The maximum micro batch size in pipeline parallelism.",
)
parser.add_argument(
"--pp-async-batch-depth",
type=int,
default=ServerArgs.pp_async_batch_depth,
help="The async batch depth of pipeline parallelism.",
)
parser.add_argument(
"--stream-interval",
type=int,
python/sglang/srt/utils/common.py
+38 -18
View File
@@ -1262,41 +1262,61 @@ def point_to_point_pyobj(
group: Optional[torch.distributed.ProcessGroup] = None,
src: int = 0,
dst: int = 1,
async_send: bool = False,
):
"""Send data from src to dst in group using DeviceToDevice communication."""
device = torch.get_device_module().current_device()
"""Send data from src to dst in group."""
from sglang.srt.distributed.parallel_state import P2PWork
if async_send:
send_func = dist.isend
else:
send_func = dist.send
if rank == src:
p2p_works = []
if len(data) == 0:
tensor_size = torch.tensor([0], dtype=torch.long, device=device)
dist.send(tensor_size, dst=dst, group=group)
tensor_size = torch.tensor(
[0],
dtype=torch.long,
)
work = send_func(tensor_size, dst, group=group)
if async_send:
p2p_works.append(P2PWork(work, tensor_size))
else:
serialized_data = pickle.dumps(data)
size = len(serialized_data)
tensor_data = torch.ByteTensor(
np.frombuffer(serialized_data, dtype=np.uint8)
).to(
device=device
) # Move to GPU
tensor_size = torch.tensor([size], dtype=torch.long, device=device)
)
tensor_size = torch.tensor([size], dtype=torch.long)
dist.send(tensor_size, dst=dst, group=group)
dist.send(tensor_data, dst=dst, group=group)
return data
work = send_func(tensor_size, dst, group=group)
if async_send:
p2p_works.append(P2PWork(work, tensor_size))
work = send_func(tensor_data, dst, group=group)
if async_send:
p2p_works.append(P2PWork(work, tensor_data))
return p2p_works
elif rank == dst:
tensor_size = torch.tensor([0], dtype=torch.long, device=device)
dist.recv(tensor_size, src=src, group=group)
tensor_size = torch.tensor(
[0],
dtype=torch.long,
)
work = dist.irecv(tensor_size, src=src, group=group)
work.wait()
size = tensor_size.item()
if size == 0:
return []
tensor_data = torch.empty(size, dtype=torch.uint8, device=device)
dist.recv(tensor_data, src=src, group=group)
tensor_data = torch.empty(
size,
dtype=torch.uint8,
)
work = dist.irecv(tensor_data, src=src, group=group)
work.wait()
serialized_data = bytes(
tensor_data.cpu().numpy()
) # Move back to host for deserialization
serialized_data = bytes(tensor_data.cpu().numpy())
data = pickle.loads(serialized_data)
return data