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

DP Enhancement (#8280)

Browse Source
This commit is contained in:
Cheng Wan
2025-07-24 21:36:21 -07:00
committed by GitHub
parent 28d4d47280
commit c0fb25e949
20 changed files with 665 additions and 1116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
python/sglang/srt/layers/logits_processor.py
View File

@@ -27,7 +27,9 @@ from sglang.srt.distributed import (
tensor_model_parallel_all_gather,
)
from sglang.srt.layers.dp_attention import (
DPPaddingMode,
attn_tp_all_gather,
attn_tp_all_gather_into_tensor,
dp_gather_replicate,
dp_scatter,
get_attention_dp_rank,
@@ -111,7 +113,8 @@ class LogitsMetadata:
# Number of tokens to sample per DP rank
global_num_tokens_for_logprob_cpu: Optional[torch.Tensor] = None
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
# The gather mode for DP attention
dp_padding_mode: Optional[DPPaddingMode] = None
# for padding
padded_static_len: int = -1
@@ -163,12 +166,12 @@ class LogitsMetadata:
forward_batch_gathered_buffer=forward_batch.gathered_buffer,
global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
dp_padding_mode=DPPaddingMode.SUM_LEN,
)
def compute_dp_attention_metadata(self, hidden_states: torch.Tensor):
if self.global_num_tokens_for_logprob_cpu is None:
# we are capturing cuda graph
return
def compute_dp_attention_metadata(self):
# TODO(ch-wan): gathered_buffer here is larger than the actual required size in draft extend,
# we may use a smaller buffer in draft extend.
cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0)
dp_rank = get_attention_dp_rank()
@@ -179,18 +182,9 @@ class LogitsMetadata:
else:
dp_local_start_pos = cumtokens[dp_rank - 1]
dp_local_num_tokens = self.global_num_tokens_for_logprob_gpu[dp_rank]
gathered_buffer = torch.zeros(
(
sum(self.global_num_tokens_for_logprob_cpu),
hidden_states.shape[1],
),
dtype=hidden_states.dtype,
device=hidden_states.device,
)
self.dp_local_start_pos = dp_local_start_pos
self.dp_local_num_tokens = dp_local_num_tokens
self.gathered_buffer = gathered_buffer
class LogitsProcessor(nn.Module):
@@ -434,7 +428,7 @@ class LogitsProcessor(nn.Module):
guarantee the given hidden_states follow this constraint.
"""
if self.do_tensor_parallel_all_gather_dp_attn:
logits_metadata.compute_dp_attention_metadata(hidden_states)
logits_metadata.compute_dp_attention_metadata()
hidden_states, local_hidden_states = (
torch.empty_like(logits_metadata.gathered_buffer),
hidden_states,
@@ -463,15 +457,31 @@ class LogitsProcessor(nn.Module):
if self.do_tensor_parallel_all_gather:
if self.use_attn_tp_group:
global_logits = torch.empty(
(self.config.vocab_size, logits.shape[0]),
device=logits.device,
dtype=logits.dtype,
)
global_logits = global_logits.T
attn_tp_all_gather(
list(global_logits.tensor_split(self.attn_tp_size, dim=-1)), logits
)
if self.config.vocab_size % self.attn_tp_size == 0:
global_logits = torch.empty(
(
self.attn_tp_size,
logits.shape[0],
self.config.vocab_size // self.attn_tp_size,
),
device=logits.device,
dtype=logits.dtype,
)
attn_tp_all_gather_into_tensor(global_logits, logits)
global_logits = global_logits.permute(1, 0, 2).reshape(
logits.shape[0], self.config.vocab_size
)
else:
global_logits = torch.empty(
(self.config.vocab_size, logits.shape[0]),
device=logits.device,
dtype=logits.dtype,
)
global_logits = global_logits.T
attn_tp_all_gather(
list(global_logits.tensor_split(self.attn_tp_size, dim=-1)),
logits,
)
logits = global_logits
else:
logits = tensor_model_parallel_all_gather(logits)