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