From bc6f0b5ce791dd198df94f6817a4c3d0222ab975 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 23 Jan 2026 17:57:41 -0800 Subject: [PATCH] [Auto Sync] Update logits_processor.py, test_logprobs.py (20260124) (#17664) Co-authored-by: github-actions[bot] Co-authored-by: yehu-ux --- python/sglang/srt/layers/logits_processor.py | 9 +++++++-- test/manual/test_logprobs.py | 3 ++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index fa7431048..fd272be6c 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -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( diff --git a/test/manual/test_logprobs.py b/test/manual/test_logprobs.py index 8f817372c..5aa68c5dd 100644 --- a/test/manual/test_logprobs.py +++ b/test/manual/test_logprobs.py @@ -334,7 +334,8 @@ class TestLogprobsDense(unittest.TestCase): print(f"Testing with config: {config}") # Sample records for this config - test_records = random.sample(records, k=min(NUM_SAMPLES, len(records))) + num_samples = config.get("num_samples", NUM_SAMPLES) + test_records = random.sample(records, k=min(num_samples, len(records))) random.shuffle(test_records) # Calculate how many samples should return logprobs