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

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
co-authored by github-actions[bot]
parent 01e3f4682e
commit ce95f203b0
+7 -6
View File
@@ -247,6 +247,7 @@ class LogitsProcessor(nn.Module):
): ):
super().__init__() super().__init__()
self.config = config self.config = config
self.vocab_size = config.vocab_size
self.logit_scale = logit_scale self.logit_scale = logit_scale
self.use_attn_tp_group = get_global_server_args().enable_dp_lm_head 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 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 return hidden_states, hidden_states
def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor: 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( global_logits = torch.empty(
( (
self.attn_tp_size, self.attn_tp_size,
logits.shape[0], logits.shape[0],
self.config.vocab_size // self.attn_tp_size, self.vocab_size // self.attn_tp_size,
), ),
device=logits.device, device=logits.device,
dtype=logits.dtype, dtype=logits.dtype,
) )
attn_tp_all_gather_into_tensor(global_logits, logits) attn_tp_all_gather_into_tensor(global_logits, logits)
global_logits = global_logits.permute(1, 0, 2).reshape( global_logits = global_logits.permute(1, 0, 2).reshape(
logits.shape[0], self.config.vocab_size logits.shape[0], self.vocab_size
) )
else: else:
global_logits = torch.empty( global_logits = torch.empty(
(self.config.vocab_size, logits.shape[0]), (self.vocab_size, logits.shape[0]),
device=logits.device, device=logits.device,
dtype=logits.dtype, dtype=logits.dtype,
) )
@@ -955,10 +956,10 @@ class LogitsProcessor(nn.Module):
if logits_metadata.next_token_logits_buffer is not None: if logits_metadata.next_token_logits_buffer is not None:
logits_buffer = logits_metadata.next_token_logits_buffer logits_buffer = logits_metadata.next_token_logits_buffer
assert logits_buffer.dtype == torch.float assert logits_buffer.dtype == torch.float
logits_buffer.copy_(logits[:, : self.config.vocab_size]) logits_buffer.copy_(logits[:, : self.vocab_size])
logits = logits_buffer logits = logits_buffer
else: else:
logits = logits[:, : self.config.vocab_size].float() logits = logits[:, : self.vocab_size].float()
return logits return logits
def _get_dllm_logits( def _get_dllm_logits(