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

[PD] Fix PP dynamic chunking with DP attention (#17339)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
2026-01-20 11:27:13 +08:00
committed by GitHub
parent 8fb45523f3
commit 71cb9d0302
1 changed files with 25 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -16,13 +16,18 @@ from sglang.srt.disaggregation.base.conn import KVPoll
from sglang.srt.disaggregation.utils import DisaggregationMode, poll_and_all_reduce
from sglang.srt.distributed.parallel_state import P2PWork
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_attention_dp_rank,
get_attention_dp_size,
is_dp_attention_enabled,
)
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.managers.utils import (
GenerationBatchResult,
get_logprob_dict_from_result,
get_logprob_from_pp_outputs,
)
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj
@@ -536,6 +541,8 @@ class SchedulerPPMixin:
latencies: List[float] = []
if self.pp_group.is_first_rank:
model_runner = self.tp_worker.model_runner
model_config = model_runner.model_config
input_ids_list = []
for i in range(128):
chunk_size = int(
@@ -582,25 +589,29 @@ class SchedulerPPMixin:
)
current_seq_len = len(req.fill_ids)
if is_dp_attention_enabled():
# For profiling, we only have one request on PP0
# Set global_num_tokens to indicate this rank has tokens, others have 0
dp_size = get_attention_dp_size()
global_num_tokens = [0] * dp_size
dp_rank = get_attention_dp_rank()
global_num_tokens[dp_rank] = current_seq_len
batch.global_num_tokens = global_num_tokens
batch.global_num_tokens_for_logprob = global_num_tokens
proxy_tensors = {
"hidden_states": torch.zeros(
(
current_seq_len,
self.tp_worker.model_runner.model_config.hidden_size,
),
dtype=self.tp_worker.model_runner.model_config.dtype,
(current_seq_len, model_config.hidden_size),
dtype=model_config.dtype,
device="cuda",
),
"residual": torch.zeros(
(
current_seq_len,
self.tp_worker.model_runner.model_config.hidden_size,
),
dtype=self.tp_worker.model_runner.model_config.dtype,
(current_seq_len, model_config.hidden_size),
dtype=model_config.dtype,
device="cuda",
),
}
from sglang.srt.managers.scheduler_pp_mixin import PPProxyTensors
pp_proxy = PPProxyTensors(proxy_tensors)
@@ -612,12 +623,9 @@ class SchedulerPPMixin:
start = time.perf_counter()
batch.prepare_for_extend()
model_worker_batch = batch.get_model_worker_batch()
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
forward_batch = ForwardBatch.init_new(
model_worker_batch, self.tp_worker.model_runner
)
_ = self.tp_worker.model_runner.forward(
forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner)
_ = model_runner.forward(
forward_batch=forward_batch, pp_proxy_tensors=pp_proxy
)