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

[Auto Sync] Update logits_processor.py (20260209) (#18503)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2026-02-09 14:33:02 -08:00
committed by GitHub
parent 01e3f4682e
commit ce95f203b0
1 changed files with 7 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -247,6 +247,7 @@ class LogitsProcessor(nn.Module):
):
super().__init__()
self.config = config
self.vocab_size = config.vocab_size
self.logit_scale = logit_scale
self.use_attn_tp_group = get_global_server_args().enable_dp_lm_head
self.use_fp32_lm_head = get_global_server_args().enable_fp32_lm_head
@@ -906,23 +907,23 @@ class LogitsProcessor(nn.Module):
return hidden_states, hidden_states
def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor:
if self.config.vocab_size % self.attn_tp_size == 0:
if self.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,
self.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
logits.shape[0], self.vocab_size
)
else:
global_logits = torch.empty(
(self.config.vocab_size, logits.shape[0]),
(self.vocab_size, logits.shape[0]),
device=logits.device,
dtype=logits.dtype,
)
@@ -955,10 +956,10 @@ class LogitsProcessor(nn.Module):
if logits_metadata.next_token_logits_buffer is not None:
logits_buffer = logits_metadata.next_token_logits_buffer
assert logits_buffer.dtype == torch.float
logits_buffer.copy_(logits[:, : self.config.vocab_size])
logits_buffer.copy_(logits[:, : self.vocab_size])
logits = logits_buffer
else:
logits = logits[:, : self.config.vocab_size].float()
logits = logits[:, : self.vocab_size].float()
return logits
def _get_dllm_logits(