Repository navigation
return_log_probs slow down generation #2367
Description
System Info
Hi team,
when I try to add return_log_probs true and output_log_probs = true to return logprob, the generation speed degraded significantly. Below is my testing.
time curl -X POST localhost:8080/v2/models/ensemble/generate_stream -d '{"text_input": "Explain in detail 5 important events of WW2.", "parameters": {"max_tokens": 2000,"bad_words":[""],"stop_words":[""],"stream": true,"temperature": 0.0}}' | grep 'data: ' | wc -l
Result: 0m6.186s
time curl -X POST localhost:8080/v2/models/ensemble/generate_stream -d '{"text_input": "Explain in detail 5 important events of WW2.", "parameters": {"max_tokens": 2000,"bad_words":[""],"stop_words":[""],"stream": true,"temperature": 0.0, "return_log_probs": true}}' | grep 'data: ' | wc -l
Result: 1m9.191s
I am running Meta-Llama-3.1-8B-Instruct on a 8xH100 with TP8.
Who can help?
No response
Information
- The official example scripts
- My own modified scripts
Tasks
- An officially supported task in the
examplesfolder (such as GLUE/SQuAD, ...) - My own task or dataset (give details below)
Reproduction
See above
Expected behavior
NA
actual behavior
NA
additional notes
NA
Activity
Superjomn commentedon Oct 24, 2024
Thanks for sharing the issue.
I wonder what API you are using. The triton-backend, the openai_server.py or anything else? @Desmond819
Desmond819 commentedon Oct 24, 2024
I am using triton-backend
Desmond819 commentedon Oct 29, 2024
@Superjomn, any update on this?
CortexEdgeUser commentedon Nov 27, 2024
Same problem, on versions 0.12.0; 0.14.0; 0.16.0. return_log_probs slows down text generation
marcomancini1994 commentedon Jul 24, 2025
@kaiyux any idea on when this would be picked up?
I think I've understood why this happens.
The problem is during the decoding sampling strategy ( https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/samplingTopKKernels.cu ) and what happens is the following:
When output log probs is enabled: Raw logits → probabilities → sampling → log conversion
When disabled: Raw logits → sampling (direct comparison possible)
Was wondering if something like this could be an optimisation to take into account:
Populate the output probabilities using a mix of the two approaches where you first select the topK logits value and then calculate the outputlogprobs based on those selections? Avoiding to calculate the exponential before and just computing what's needed for the topK selected?
PS: I am a newbie on CUDA code, might have missed something or gave a wrong optimisation for the problematic path, sorry in advance in that's the case
karljang commentedon Sep 29, 2025
@Desmond819 ,
Apologies for the delayed response. Given the several updates made to the sampling feature, could you please try the latest version to see if the issue still persists?
github-actions commentedon Oct 14, 2025
Issue has not received an update in over 14 days. Adding stale label.
github-actions commentedon Oct 28, 2025
This issue was closed because it has been 14 days without activity since it has been marked as stale.
[None][ci] Restore the two Ray disagg tp2 waivers removed by NVIDIA#1…