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
@@ -57,6 +57,7 @@ _EXACT_VAL_ROWS = [None, [-0.5, -2.5], [-0.25, -1.5], [-0.125, -4.0]]
class _TokenizerManagerStub:
"""Borrow the real logprob meta_info methods without a full manager."""
convert_logprob_style = TokenizerManager.convert_logprob_style
add_logprob_to_meta_info = TokenizerManager.add_logprob_to_meta_info
detokenize_logprob_tokens = TokenizerManager.detokenize_logprob_tokens
detokenize_top_logprobs_tokens = TokenizerManager.detokenize_top_logprobs_tokens
@@ -543,6 +544,32 @@ class TestMetaInfoFromSchedulerArrays(CustomTestCase):
)
class TestTokenizerManagerLogprobs(CustomTestCase):
def test_output_logprobs_without_input_logprobs(self):
state = _make_state(return_logprob=True, top_logprobs_num=0)
recv_obj = SimpleNamespace(
input_token_logprobs_val=None,
input_token_logprobs_idx=None,
output_token_logprobs_val=[[-0.25]],
output_token_logprobs_idx=[[42]],
)
meta_info = {}
_TokenizerManagerStub().convert_logprob_style(
meta_info,
state,
top_logprobs_num=0,
token_ids_logprob=None,
return_text_in_logprobs=False,
recv_obj=recv_obj,
recv_obj_index=0,
)
self.assertEqual(meta_info["input_token_logprobs"], [])
self.assertEqual(meta_info["output_token_logprobs"], [(-0.25, 42, None)])
self.assertEqual(meta_info["output_token_logprobs_length"], 1)
def _make_batch_token_id_output(**overrides) -> BatchTokenIDOutput:
"""A two-request BatchTokenIDOutput with the required fields stubbed."""
n = 2