[https://nvbugs/5513423][chore] Enhance min length penalty application in TorchSampler to work with speculative decoding

- Refactored `_apply_min_length_penalty` method to support step-wise application of minimum length penalties based on the number of steps per request.
- Updated the test for `min_tokens` to include a parameterized approach for speculative decoding configurations.

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

Author: stnie

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -655,17 +655,33 @@ def _apply_embedding_bias(
         return logits
 
     @staticmethod
+    @torch.inference_mode()
     def _apply_min_length_penalty(logits: torch.Tensor,
-                                  requests: list[LlmRequest]):
+                                  requests: list[LlmRequest],
+                                  num_steps: list[int]) -> torch.Tensor:
+        """Inplace apply min_length_penalty to logits.
 
-        if not any(
-                r.py_min_length and r.max_beam_num_tokens < r.py_min_length[0]
-                for r in requests):
-            return logits
-        logits = logits.clone()
-        for index, r in enumerate(requests):
-            if r.py_min_length and r.max_beam_num_tokens < r.py_min_length[0]:
-                logits[index, [r.py_end_id]] = float('-inf')
+            Args:
+                logits: The logits to apply min length penalty to
+                requests: The requests to apply min length penalty to
+                num_steps: The number of steps per request
+
+            Returns:
+                The logits with min length penalty applied
+        """
+        if any(r.py_min_length and r.max_beam_num_tokens < r.py_min_length[0]
+               for r in requests):
+            current_offset = 0
+            for index, r in enumerate(requests):
+                if r.py_min_length:
+                    for step in range(num_steps[index]):
+                        if r.max_beam_num_tokens + step < r.py_min_length[0]:
+                            logits[current_offset + step,
+                                   r.py_end_id] = float('-inf')
+                        else:
+                            #early exit
+                            break
+                current_offset += num_steps[index]
         return logits
 
     def _process_requests(self,
@@ -696,8 +712,9 @@ def _process_requests(self,
             raw_logits = model_outputs["logits"]
 
         requests = scheduled_requests.all_requests()
-        raw_logits = self._apply_min_length_penalty(raw_logits, requests)
         num_steps = [1 + get_draft_token_length(req) for req in requests]
+        raw_logits = self._apply_min_length_penalty(raw_logits, requests,
+                                                    num_steps)
         sum_steps = sum(num_steps)
         no_draft_tokens = len(requests) == sum_steps
         fast_path = not self.enable_mixed_sampler and no_draft_tokens and log_probs_host is None

tests/unittest/llmapi/test_llm_pytorch.py

@@ -6,7 +6,7 @@
 from tensorrt_llm import LLM
 from tensorrt_llm.executor import GenerationExecutorWorker
 from tensorrt_llm.llmapi import KvCacheConfig
-from tensorrt_llm.llmapi.llm_args import PeftCacheConfig
+from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig, PeftCacheConfig
 from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer
 from tensorrt_llm.metrics import MetricNames
 from tensorrt_llm.sampling_params import SamplingParams
@@ -861,12 +861,28 @@ def test_llm_with_proxy_error():
 
 
 @pytest.mark.part0
-def test_min_tokens():
+@pytest.mark.parametrize("use_speculative", [True, False])
+def test_min_tokens(use_speculative: bool):
     """Check min_tokens is respected."""
-    llm = LLM(model=llama_model_path,
-              kv_cache_config=global_kvcache_config,
-              enable_mixed_sampler=True,
-              max_seq_len=2048)
+    llm_common_config = dict(
+        model=llama_model_path,
+        max_batch_size=2,
+        kv_cache_config=global_kvcache_config,
+        max_num_tokens=2048,
+        enable_mixed_sampler=True,
+    )
+
+    if use_speculative:
+        spec_config = NGramDecodingConfig(
+            max_draft_len=4,
+            max_matching_ngram_size=2,
+            is_keep_all=True,
+            is_use_oldest=True,
+            is_public_pool=True,
+        )
+        llm = LLM(**llm_common_config, speculative_config=spec_config)
+    else:
+        llm = LLM(**llm_common_config)
 
     output_len = 2000
     sampling_params = SamplingParams(max_tokens=output_len,