[Auto Sync] Update logits_processor.py, test_logprobs.py (20260124) (#17664)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: yehu-ux <yehu@x.ai>
This commit is contained in:
Lianmin Zheng
2026-01-23 17:57:41 -08:00
committed by GitHub
parent e1833c4f5a
commit bc6f0b5ce7
2 changed files with 9 additions and 3 deletions

View File

@@ -766,14 +766,19 @@ class LogitsProcessor(nn.Module):
# Compute the logprobs of the chunk
chunk_input_logprobs = chunk_logits[chunk_indices]
# Only index per-token arrays when the corresponding feature is active.
# Otherwise these tensors can be per-sequence (or scalars), which can
# cause out-of-bounds indexing on GPU.
chunk_temperature = (
logits_metadata.temperature[global_indices]
if logits_metadata.temperature is not None
if logits_metadata.temp_scaled_logprobs
and logits_metadata.temperature is not None
else None
)
chunk_top_p = (
logits_metadata.top_p[global_indices]
if logits_metadata.top_p is not None
if logits_metadata.top_p_normalized_logprobs
and logits_metadata.top_p is not None
else None
)
chunk_input_logprobs = compute_temp_top_p_normalized_logprobs(