[None][fix] Fix bug of undefined py_topk_logprobs_vals (NVIDIA#8789)

Signed-off-by: Dong Cao <[email protected]>
Signed-off-by: FredricZ-2007 <[email protected]>

Author: dcaox

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -441,10 +441,12 @@ def _handle_request_broadcasting(self,
                 new_requests, "py_multimodal_data")
             py_scheduling_params = self._collect_py_objects_from_requests(
                 new_requests, "py_scheduling_params")
+            py_num_logprobs = self._collect_py_objects_from_requests(
+                new_requests, "py_num_logprobs")
             py_request_objects = tuple(
                 filter(None, [
                     py_logits_post_processors, py_multimodal_data,
-                    py_scheduling_params
+                    py_scheduling_params, py_num_logprobs
                 ]))
         else:
             py_request_objects = None

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -733,8 +733,7 @@ def executor_request_to_llm_request(
         mrope_position_deltas=mrope_position_deltas,
         lookahead_config=None,
         return_log_probs=executor_request.output_config.return_log_probs,
-        num_logprobs=executor_request.py_num_logprobs if hasattr(
-            executor_request, "py_num_logprobs") else 0,
+        num_logprobs=getattr(executor_request, "py_num_logprobs", 0),
         return_context_logits=executor_request.output_config.
         return_context_logits,
         return_perf_metrics=executor_request.output_config.return_perf_metrics,

tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py

@@ -11,6 +11,7 @@
     check_llama_7b_multi_lora_from_request_test_harness,
     check_phi3_lora_fused_modules_output_tp2_identical_to_tp1)
 from .test_llm import (_test_llm_capture_request_error, llama_model_path,
+                       llm_return_logprobs_test_harness,
                        tinyllama_logits_processor_test_harness)
 from .test_llm_pytorch import llama_7b_lora_from_dir_test_harness
 
@@ -104,3 +105,23 @@ async def test_llm_rpc_streaming_tp2():
                                                sampling_params=SamplingParams(
                                                    max_tokens=10, end_id=-1)):
             print(f"get result: {output}")
+
+
+@skip_ray
+@pytest.mark.gpu2
+@pytest.mark.parametrize(
+    "prompt_logprobs, logprobs, return_context_logits, return_generation_logits",
+    [
+        (None, 1, False,
+         False),  # generation logprobs only (top-1, PyTorch limit)
+    ])
+def test_llm_return_logprobs_streaming_tp2(prompt_logprobs, logprobs,
+                                           return_context_logits,
+                                           return_generation_logits):
+    llm_return_logprobs_test_harness(prompt_logprobs,
+                                     logprobs,
+                                     return_context_logits,
+                                     return_generation_logits,
+                                     streaming=True,
+                                     backend="pytorch",
+                                     tp_size=2)