From 4615e446a8c258111111149700c4e56568a967a7 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Wed, 11 Jun 2025 20:26:27 +0000 Subject: [PATCH] Fix logprobs issues. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 85cbb9da3698..6c5089a45622 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -721,8 +721,8 @@ def update_requests(self, state: SampleStateTRTLLM): if request.py_return_log_probs: # NOTE: Log probs with drafting has not been tested yet. begin_log_probs_offset = request.prompt_len if request.sampling_config.beam_width == 1 else 0 - current_token = seq_len - request.prompt_len - len( - num_new_tokens[beam]) + step + current_token = seq_len - request.prompt_len - num_new_tokens[ + beam] + step log_probs.append({ new_token.item(): @@ -732,7 +732,7 @@ def update_requests(self, state: SampleStateTRTLLM): rank=1) }) - if num_new_tokens[beam] > 0 and request.py_return_log_probs: + if request.py_return_log_probs: cum_log_probs.append( state.host.cum_log_probs[seq_slot * beam_width + beam].item())