fix: preserve output logprobs without input logprobs (#34627)

Signed-off-by: jain-ria <riajain@NVIDIA.com>
This commit is contained in:
jain-ria
2026-08-18 12:07:13 -05:00
committed by GitHub
parent 3f26febaff
commit 83d7d45330
2 changed files with 39 additions and 10 deletions
+12 -10
View File
@@ -2655,11 +2655,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
recv_obj: BatchStrOutput,
recv_obj_index: int,
):
if recv_obj.input_token_logprobs_val is None:
return
if (
len(recv_obj.input_token_logprobs_val) > 0
recv_obj.input_token_logprobs_val is not None
and len(recv_obj.input_token_logprobs_val) > 0
and recv_obj.input_token_logprobs_val[recv_obj_index] is not None
):
state.input_token_logprobs_val.extend(
@@ -2668,12 +2666,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
state.input_token_logprobs_idx.extend(
recv_obj.input_token_logprobs_idx[recv_obj_index]
)
state.output_token_logprobs_val.extend(
recv_obj.output_token_logprobs_val[recv_obj_index]
)
state.output_token_logprobs_idx.extend(
recv_obj.output_token_logprobs_idx[recv_obj_index]
)
if (
recv_obj.output_token_logprobs_val is not None
and recv_obj.output_token_logprobs_val[recv_obj_index] is not None
):
state.output_token_logprobs_val.extend(
recv_obj.output_token_logprobs_val[recv_obj_index]
)
state.output_token_logprobs_idx.extend(
recv_obj.output_token_logprobs_idx[recv_obj_index]
)
if top_logprobs_num > 0:
if len(recv_obj.input_top_logprobs_val) > 0: