Respect TensorRT-LLM input limits during generation

lm_eval/models/trtllm_causallms.py

@@ -2,10 +2,11 @@
 
 import logging
 from importlib.util import find_spec
-from typing import List, Literal, Optional, Tuple, Union, Any
+from typing import Any, Literal
 
 from tqdm import tqdm
 
+from lm_eval.api.instance import Instance
 from lm_eval.api.model import TemplateLM
 from lm_eval.api.registry import register_model
 from lm_eval.models.utils import (
@@ -22,7 +23,7 @@
     get_rolling_token_windows,
     make_disjoint_window,
 )
-from lm_eval.api.instance import Instance
+
 
 try:
     from tensorrt_llm import LLM, SamplingParams
@@ -46,21 +47,21 @@ def __init__(
         model: str,
         dtype: Literal["float16", "bfloat16", "float32", "auto"] = "auto",
         trust_remote_code: bool = False,
-        tokenizer: Optional[str] = None,
+        tokenizer: str | None = None,
         tokenizer_mode: Literal["auto", "slow"] = "auto",
-        add_bos_token: Optional[bool] = None,
+        add_bos_token: bool | None = None,
         prefix_token_id: int | None = None,
         max_batch_size: int = 32,
         max_input_len: int = 2048,
         max_output_len: int = 512,
         max_beam_width: int = 1,
         tensor_parallel_size: int = 1,
         pipeline_parallel_size: int = 1,
-        batch_size: Union[str, int] = "auto",
+        batch_size: str | int = "auto",
         max_gen_toks: int = 256,
         seed: int = 1234,
         truncation_side: Literal["left", "right", "middle"] = "left",
-        chat_template_args: Optional[dict] = None,
+        chat_template_args: dict | None = None,
         # TRT-LLM: enable thinking tags in apply_chat_template.
         enable_thinking: bool = True,
         # End marker for thinking tags - splits to get response after this token (if provided).
@@ -106,7 +107,7 @@ def __init__(
         from transformers import AutoTokenizer
 
         self.tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer if tokenizer else model,
+            tokenizer or model,
             trust_remote_code=trust_remote_code,
             use_fast=(tokenizer_mode == "auto"),
         )
@@ -122,10 +123,12 @@ def __init__(
                 test_tok = self.tok_encode("test", add_special_tokens=True)
                 test_tok_no_special = self.tok_encode("test", add_special_tokens=False)
                 self.add_bos_token = test_tok != test_tok_no_special
-            except Exception:
+            except Exception:  # noqa: BLE001 - retain fallback for tokenizer-specific failures
                 self.add_bos_token = False
 
-        eval_logger.info(f"BOS token will{' not' if not self.add_bos_token else ''} be added")
+        eval_logger.info(
+            f"BOS token will{' not' if not self.add_bos_token else ''} be added"
+        )
 
         if prefix_token_id is not None:
             eval_logger.info(
@@ -238,9 +241,9 @@ def apply_chat_template(
 
     def _loglikelihood_tokens(
         self,
-        requests: List[Tuple[Tuple[str, str], List[int], List[int]]],
+        requests: list[tuple[tuple[str, str], list[int], list[int]]],
         disable_tqdm: bool = False,
-    ) -> List[Tuple[float, bool]]:
+    ) -> list[tuple[float, bool]]:
         res = []
 
         def _collate(x):
@@ -260,7 +263,7 @@ def _collate(x):
         for chunk in chunks:
             inputs = []
             ctxlens = []
-            for cache_key, context_enc, continuation_enc in chunk:
+            for _cache_key, context_enc, continuation_enc in chunk:
                 inp = (context_enc + continuation_enc)[-(self.max_length) :]
                 ctxlen = len(context_enc) - max(
                     0, len(context_enc) + len(continuation_enc) - (self.max_length)
@@ -275,10 +278,10 @@ def _collate(x):
                     max_tokens=1,
                     temperature=0,
                     prompt_logprobs=1,
-                )
+                ),
             )
             for output, ctxlen, (cache_key, _, _), inp in zip(
-                outputs, ctxlens, chunk, inputs
+                outputs, ctxlens, chunk, inputs, strict=False
             ):
                 answer = self._parse_logprobs(
                     tokens=inp,
@@ -297,7 +300,7 @@ def _collate(x):
         return re_ord.get_original(res)
 
     @staticmethod
-    def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]:
+    def _parse_logprobs(tokens: list, outputs, ctxlen: int) -> tuple[float, bool]:
         """Process logprobs and tokens.
 
         :param tokens: list
@@ -331,7 +334,9 @@ def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]:
 
         continuation_logprobs_dicts = shifted_prompt_logprobs[ctxlen:]
         continuation_logprobs = 0.0
-        for token, logprob_dict in zip(tokens[ctxlen:], continuation_logprobs_dicts):
+        for token, logprob_dict in zip(
+            tokens[ctxlen:], continuation_logprobs_dicts, strict=False
+        ):
             if token in logprob_dict:
                 continuation_logprobs += logprob_dict[token].logprob
             else:
@@ -344,7 +349,9 @@ def _parse_logprobs(tokens: List, outputs, ctxlen: int) -> Tuple[float, bool]:
         # Determine is_greedy: the actual token must have rank=1 at every position.
         # The Logprob object has a .rank attribute (1 = highest probability).
         is_greedy = True
-        for token, logprob_dict in zip(tokens[ctxlen:], continuation_logprobs_dicts):
+        for token, logprob_dict in zip(
+            tokens[ctxlen:], continuation_logprobs_dicts, strict=False
+        ):
             if token in logprob_dict and logprob_dict[token].rank is not None:
                 if logprob_dict[token].rank != 1:
                     is_greedy = False
@@ -412,12 +419,13 @@ def _collate_gen(_requests):
                     gen_kwargs, eos=eos, default_max_gen_toks=self.max_gen_toks
                 )
 
-                # set the max length in tokens of inputs ("context_enc")
-                # max len for inputs = max length, minus room to generate the max new tokens
+                # Respect both the engine's input cap and its total sequence cap.
+                # Unused output capacity cannot be borrowed by the prompt.
                 toks, max_gen_toks = maybe_truncate(
                     toks,
                     max_gen_toks=max_gen_toks,
-                    max_model_len=self.max_input_len + self.max_output_len,
+                    max_model_len=self.max_input_len
+                    + min(self.max_output_len, max_gen_toks),
                     side=self.truncation_side,
                     verbose=True,
                 )
@@ -460,8 +468,8 @@ def loglikelihood_rolling(
         self,
         requests,
         disable_tqdm: bool = False,
-        override_bs: Optional[int] = None,
-    ) -> List[float]:
+        override_bs: int | None = None,
+    ) -> list[float]:
         """Compute perplexity using rolling windows.
 
         This method efficiently batches all rolling windows from all requests together
@@ -494,7 +502,7 @@ def loglikelihood_rolling(
 
             # Create rolling windows using the utility function
             # This handles the windowing logic properly with context
-            rolling_token_windows: List[Tuple[List[int], List[int]]] = list(
+            rolling_token_windows: list[tuple[list[int], list[int]]] = list(
                 map(
                     make_disjoint_window,
                     get_rolling_token_windows(
@@ -523,7 +531,7 @@ def loglikelihood_rolling(
         for i in range(0, len(all_windows), batch_size):
             batch = all_windows[i : i + batch_size]
             # Extract just the windows for processing, keeping track of request indices
-            batch_indices, batch_windows = zip(*batch)
+            batch_indices, batch_windows = zip(*batch, strict=False)
 
             # Process the batch using _loglikelihood_tokens
             batch_nlls = self._loglikelihood_tokens(
@@ -532,7 +540,7 @@ def loglikelihood_rolling(
             )
 
             # Store results with their request indices
-            all_nlls.extend(zip(batch_indices, batch_nlls))
+            all_nlls.extend(zip(batch_indices, batch_nlls, strict=False))
 
         # Reconstruct per-request loglikelihoods
         loglikelihoods = []
@@ -587,7 +595,7 @@ def vocab_size(self) -> int:
 
     def __del__(self):
         """Cleanup TRT-LLM resources."""
-        if hasattr(self, 'llm'):
+        if hasattr(self, "llm"):
             del self.llm
 
     @staticmethod

tests/models/test_trtllm_input_limit.py

@@ -0,0 +1,88 @@
+"""CPU coverage of TensorRT-LLM's separate input and sequence limits."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import transformers
+
+from lm_eval.models import trtllm_causallms as trtllm
+
+
+class CharacterTokenizer:
+    """Use one token per character without downloads or special tokens."""
+
+    pad_token_id = 0
+    eos_token_id = 0
+    bos_token_id = 1
+
+    def decode(self, token):
+        return "<bos>" if token == 1 else "<eos>"
+
+    def __call__(self, strings, **kwargs):
+        return SimpleNamespace(input_ids=[[ord(char) for char in s] for s in strings])
+
+
+@pytest.fixture
+def model(monkeypatch):
+    """Run the real constructor with a local tokenizer and a recording engine."""
+    engine = SimpleNamespace(
+        generate=Mock(
+            side_effect=lambda inputs, **kwargs: [
+                SimpleNamespace(outputs=[SimpleNamespace(text="answer")])
+                for _ in inputs
+            ]
+        )
+    )
+    factory = Mock(return_value=engine)
+    monkeypatch.setattr(trtllm, "find_spec", lambda name: True)
+    monkeypatch.setattr(trtllm, "LLM", factory, raising=False)
+    monkeypatch.setattr(trtllm, "SamplingParams", SimpleNamespace, raising=False)
+    monkeypatch.setattr(trtllm, "KvCacheConfig", SimpleNamespace, raising=False)
+    monkeypatch.setattr(
+        transformers.AutoTokenizer,
+        "from_pretrained",
+        Mock(return_value=CharacterTokenizer()),
+    )
+    model = trtllm.TRTLLM(
+        model="unused",
+        max_input_len=8,
+        max_output_len=4,
+        max_gen_toks=2,
+        add_bos_token=False,
+        enable_thinking=False,
+    )
+    assert factory.call_args.kwargs["max_input_len"] == 8
+    assert factory.call_args.kwargs["max_seq_len"] == 12
+    return model
+
+
+@pytest.mark.parametrize(
+    ("side", "expected"),
+    [("left", "mnopqrst"), ("right", "abcdefgh"), ("middle", "abcdqrst")],
+)
+def test_long_prompt_respects_input_limit(model, side, expected):
+    """Unused output capacity must not increase the permitted input length."""
+    model.truncation_side = side
+    request = SimpleNamespace(args=("abcdefghijklmnopqrst", {"until": ["\n"]}))
+    assert model.generate_until([request], disable_tqdm=True) == ["answer"]
+    call = model.llm.generate.call_args
+    assert call.kwargs["inputs"] == [[ord(char) for char in expected]]
+    assert call.kwargs["sampling_params"][0].max_tokens == 2
+
+
+@pytest.mark.parametrize("prompt_length", [3, 8, 9, 20])
+@pytest.mark.parametrize("generation_length", [1, 4, 6])
+def test_input_and_sequence_limits(model, prompt_length, generation_length):
+    """Keep short prompts and generation budgets while satisfying both caps."""
+    request = SimpleNamespace(
+        args=("a" * prompt_length, {"max_gen_toks": generation_length, "until": ["\n"]})
+    )
+    assert model.generate_until([request], disable_tqdm=True) == ["answer"]
+    call = model.llm.generate.call_args
+    tokens = call.kwargs["inputs"][0]
+    params = call.kwargs["sampling_params"][0]
+    assert len(tokens) <= 8
+    assert len(tokens) + params.max_tokens <= 12
+    assert len(tokens) == min(prompt_length, 8, 12 - generation_length)
+    assert params.max_tokens == generation_length