[TRTLLM-9688][feat] Refactor processed logprobs to skip logprob calculation when not needed

- Added LogProbsMode class to define modes for log probabilities: RAW and PROCESSED.
- Updated SamplingParams and LlmRequest to utilize LogProbsMode for logprobs_mode parameter.
- Enhanced validation to check logprobs_mode against LogProbsMode values.
- Modified TorchSampler and related classes to support new logprobs_mode functionality.
- Modified TorchSampler to only calculate logprobs when a request needs it
- Updated tests to cover new logprobs_mode behavior and ensure correct processing of log probabilities.

Signed-off-by: Stefan Niebler <[email protected]>

Author: stnie

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -8,6 +8,7 @@
 from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 from tensorrt_llm.bindings import executor as tllm_executor
 from tensorrt_llm.executor.result import TokenLogprobs
+from tensorrt_llm.sampling_params import LogprobMode
 
 SamplingConfig = tensorrt_llm.bindings.SamplingConfig
 '''
@@ -460,7 +461,7 @@ def __init__(
             is_first_draft: bool = False,
             use_chunked_generation_logits: bool = True,
             logits_chunk_size: int = 8,
-            logprobs_mode: str = "raw",
+            logprobs_mode: LogprobMode | None = None,
             **kwargs):
 
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
@@ -539,7 +540,7 @@ def __init__(
         # currently, keep py_stop_words_list as python list, rather than tensor.
         self.py_stop_words_list = stop_words_list
 
-        self.py_logprobs_mode = logprobs_mode
+        self.py_logprobs_mode = LogprobMode.RAW if logprobs_mode is None else logprobs_mode
 
         self.py_result = PyResult(
             prompt_len=self.py_prompt_len,
@@ -568,6 +569,15 @@ def set_exclude_last_generation_logits(
         self.py_result.set_exclude_last_generation_logits(
             exclude_last_generation_logits)
 
+    def validate_logprobs_mode(self):
+        if self.py_logprobs_mode not in [
+                LogprobMode.RAW, LogprobMode.PROCESSED
+        ]:
+            raise ValueError(
+                f"Invalid logprobs_mode: {self.py_logprobs_mode} "
+                f"Expected one of {LogprobMode.RAW.value}, {LogprobMode.PROCESSED.value}"
+            )
+
     @property
     def cached_tokens(self) -> int:
         return self._cached_tokens
@@ -801,7 +811,7 @@ def executor_request_to_llm_request(
         py_multimodal_data=getattr(executor_request, "py_multimodal_data",
                                    None),
         kv_cache_retention_config=executor_request.kv_cache_retention_config,
-        logprobs_mode=getattr(executor_request, "py_logprobs_mode", "raw"),
+        logprobs_mode=getattr(executor_request, "py_logprobs_mode", None),
     )
     if child_req_ids:
         for child_id in child_req_ids:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1727,6 +1727,8 @@ def _validate_request(self, request: LlmRequest):
                     f"Request beam width {sampling_config.beam_width} "
                     f"is not equal to max_beam_width {self.max_beam_width}. This is not supported!"
                 )
+        # Validate logprobs mode
+        request.validate_logprobs_mode()
 
         # Check token ID ranges
         if isinstance(self.model_engine.model, DecoderModelForCausalLM):