From a71e1ea2f00e4f3a06c473594feec941208e5b50 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Tue, 1 Sep 2026 18:02:43 +0000 Subject: [PATCH] Forward kv_cache_free_gpu_memory_fraction to the TensorRT-LLM engines (NVBug 6701763) huggingface_example.sh parsed --kv_cache_free_gpu_memory_fraction and printed it, but no consumer had existed since the trtllm-build path was removed, so every engine the script deploys sized its KV cache from a default: 0.9 for lm-eval's trtllm backend and trtllm-serve, 0.7 hardcoded in modelopt.deploy.llm.LLM. Wire the value through all four paths (quant smoke test, lm_eval, mmlu, simple_eval/livecodebench) and default it to 0.7, matching what the engines effectively used before. lm-eval's backend drops unknown --model_args keys, so lm_eval_trtllm.py now injects the fraction into the KvCacheConfig it builds. Signed-off-by: Chenjie Luo Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 2 + examples/hf_ptq/run_tensorrt_llm.py | 10 ++ .../hf_ptq/scripts/huggingface_example.sh | 9 +- examples/hf_ptq/scripts/parser.sh | 2 +- examples/llm_eval/README.md | 18 ++- examples/llm_eval/lm_eval_trtllm.py | 61 +++++++- examples/llm_eval/mmlu.py | 11 ++ modelopt/deploy/llm/generate.py | 16 ++- .../examples/hf_ptq/test_run_tensorrt_llm.py | 28 ++++ .../examples/llm_eval/test_lm_eval_trtllm.py | 130 +++++++++++++++++- 11 files changed, 268 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..d08315a14fb 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -43,6 +43,7 @@ Changelog - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. +- Fix ``--kv_cache_free_gpu_memory_fraction`` in ``examples/hf_ptq/scripts/huggingface_example.sh``: it was parsed but never forwarded, so every TensorRT-LLM engine the script deploys sized its KV cache from TensorRT-LLM's default 90% of free GPU memory and evaluation could run out of memory. The value now reaches the ``quant`` smoke test, ``lm_eval``, ``mmlu`` and ``simple_eval``/``livecodebench`` paths and defaults to 0.7; ``modelopt.deploy.llm.LLM`` and ``lm_eval_trtllm.py``'s ``--model_args`` accept ``kv_cache_free_gpu_memory_fraction`` directly. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index fea69221825..f96485cb13d 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -176,6 +176,8 @@ scripts/huggingface_example.sh --model $HF_PATH --quant --tp [1|2|4|8] > *If a GPU OOM error occurs during model quantization despite sufficient memory, setting the --use_seq_device_map flag can help. This enforces sequential device mapping, distributing the model across GPUs and utilizing up to 80% of each GPU's memory.* +> *If a GPU OOM error occurs after quantization, while the checkpoint is being deployed for the `quant` smoke test or an eval task, lower `--kv_cache_free_gpu_memory_fraction` (default 0.7). It caps the share of the free GPU memory the TensorRT-LLM KV cache may take, leaving the rest for context logits and logprob buffers.* + > *You can add `--low_memory_mode` to the command to lower the memory requirements of the PTQ process. With this mode, the script will compress model weights to low precision before calibration. This mode is only supported for FP8 and NVFP4 with max calibration.* #### Recipe-based Quantization diff --git a/examples/hf_ptq/run_tensorrt_llm.py b/examples/hf_ptq/run_tensorrt_llm.py index 323c2a7b13f..428153261f5 100644 --- a/examples/hf_ptq/run_tensorrt_llm.py +++ b/examples/hf_ptq/run_tensorrt_llm.py @@ -43,6 +43,15 @@ def parse_arguments(): default=False, action="store_true", ) + parser.add_argument( + "--kv_cache_free_gpu_memory_fraction", + type=float, + default=0.7, + help=( + "Fraction of the GPU memory left after loading the model weights that the KV cache " + "may use. Lower it if the engine runs out of memory." + ), + ) return parser.parse_args() @@ -74,6 +83,7 @@ def run(args): max_batch_size=len(input_texts), enable_kv_cache_reuse=False, trust_remote_code=args.trust_remote_code, + kv_cache_free_gpu_memory_fraction=args.kv_cache_free_gpu_memory_fraction, ) torch.cuda.cudart().cudaProfilerStart() outputs = llm.generate_text(input_texts, args.max_output_len) diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index adc4cf22ca3..35e54920661 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -229,6 +229,8 @@ if [[ $TASKS =~ "quant" ]] || [[ ! -d "$SAVE_PATH" ]] || [[ ! $(ls -A $SAVE_PATH RUN_ARGS+=" --trust_remote_code " fi + RUN_ARGS+=" --kv_cache_free_gpu_memory_fraction=$KV_CACHE_FREE_GPU_MEMORY_FRACTION " + # Only run the deploy+generate smoke test when "quant" is explicitly requested. Eval tasks # (lm_eval/mmlu/simple_eval) deploy the checkpoint themselves, so it is redundant there. if [[ $TASKS =~ "quant" ]]; then @@ -296,7 +298,7 @@ if [[ $TASKS =~ "lm_eval" ]]; then # explicitly; the engine's max_seq_len is max_input_len + max_output_len. python lm_eval_trtllm.py \ --model trtllm \ - --model_args "model=$SAVE_PATH,tokenizer=$MODEL_ABS_PATH,tensor_parallel_size=$LM_EVAL_TP,max_batch_size=$BUILD_MAX_BATCH_SIZE,max_gen_toks=$BUILD_MAX_OUTPUT_LEN,max_input_len=$BUILD_MAX_INPUT_LEN,max_output_len=$BUILD_MAX_OUTPUT_LEN" \ + --model_args "model=$SAVE_PATH,tokenizer=$MODEL_ABS_PATH,tensor_parallel_size=$LM_EVAL_TP,max_batch_size=$BUILD_MAX_BATCH_SIZE,max_gen_toks=$BUILD_MAX_OUTPUT_LEN,max_input_len=$BUILD_MAX_INPUT_LEN,max_output_len=$BUILD_MAX_OUTPUT_LEN,kv_cache_free_gpu_memory_fraction=$KV_CACHE_FREE_GPU_MEMORY_FRACTION" \ --tasks $LM_EVAL_TASKS \ --batch_size $BUILD_MAX_BATCH_SIZE $lm_eval_flags | tee $LM_EVAL_RESULT @@ -336,6 +338,7 @@ if [[ $TASKS =~ "mmlu" ]]; then --model_name causal \ --model_path $MODEL_ABS_PATH \ --checkpoint_dir $SAVE_PATH \ + --kv_cache_free_gpu_memory_fraction $KV_CACHE_FREE_GPU_MEMORY_FRACTION \ --data_dir $MMLU_DATA_PATH $mmlu_flags | tee $MMLU_RESULT popd @@ -348,7 +351,9 @@ if [[ $TASKS =~ "livecodebench" || $TASKS =~ "simple_eval" ]]; then HASH=$(echo -n "$SAVE_PATH" | md5sum | awk '{print $1}') PORT=$((10000 + (0x${HASH:0:4} % 50001))) echo "Starting trtllm-serve on $PORT" - trtllm-serve $SAVE_PATH --host 0.0.0.0 --port $PORT >$SAVE_PATH/serve.txt 2>&1 & + trtllm-serve $SAVE_PATH --host 0.0.0.0 --port $PORT \ + --kv_cache_free_gpu_memory_fraction $KV_CACHE_FREE_GPU_MEMORY_FRACTION \ + >$SAVE_PATH/serve.txt 2>&1 & SERVE_PID=$! # Poll the log instead of `tail -f | while ... break`: under `set -o pipefail` (set above), diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index e3eb0b18b63..584e4f69e8a 100644 --- a/examples/hf_ptq/scripts/parser.sh +++ b/examples/hf_ptq/scripts/parser.sh @@ -33,7 +33,7 @@ parse_options() { TASKS="quant" TRUST_REMOTE_CODE=false - KV_CACHE_FREE_GPU_MEMORY_FRACTION=0.8 + KV_CACHE_FREE_GPU_MEMORY_FRACTION=0.7 VERBOSE=true USE_SEQ_DEVICE_MAP=false CAST_MXFP4_TO_NVFP4=false diff --git a/examples/llm_eval/README.md b/examples/llm_eval/README.md index f4c53f49d81..65795e8196b 100644 --- a/examples/llm_eval/README.md +++ b/examples/llm_eval/README.md @@ -114,7 +114,7 @@ checkpoint directly with the TensorRT-LLM LLM API. ```sh python lm_eval_trtllm.py --model trtllm \ - --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512 \ + --model_args model=,tokenizer=,tensor_parallel_size=,max_batch_size=,max_input_len=4096,max_output_len=512,kv_cache_free_gpu_memory_fraction=0.7 \ --tasks \ --batch_size ``` @@ -137,11 +137,17 @@ python lm_eval_trtllm.py --model trtllm \ > loglikelihood task (hellaswag, mmlu, arc, ...) fails with a `KeyError`; > `lm_eval_trtllm.py` overrides the alignment. It goes away once the fix lands upstream. -> **_NOTE:_** The backend forwards only a fixed set of arguments to TensorRT-LLM, so the -> tuning the old `lm_eval_tensorrt_llm.py` applied is not reachable: expert parallelism is -> left at the TensorRT-LLM default (MoE checkpoints can fail in DeepEP kernels on some -> GPUs, e.g. SM 12.0) and the KV cache uses 90% of free GPU memory rather than 70%. Lower -> `tensor_parallel_size` if you hit either. +> **_NOTE:_** `kv_cache_free_gpu_memory_fraction` is the share of the GPU memory left after +> loading the weights that the KV cache may take. TensorRT-LLM defaults it to 0.9, which can +> leave too little room for the `prompt_logprobs` buffers and OOM on a large-memory GPU. +> lm-eval's backend drops the key, so `lm_eval_trtllm.py` forwards it to the engine; +> `huggingface_example.sh` passes 0.7 and exposes `--kv_cache_free_gpu_memory_fraction`. + +> **_NOTE:_** Other than the KV cache fraction, the backend forwards only a fixed set of +> arguments to TensorRT-LLM, so the remaining tuning the old `lm_eval_tensorrt_llm.py` +> applied is not reachable: expert parallelism is left at the TensorRT-LLM default (MoE +> checkpoints can fail in DeepEP kernels on some GPUs, e.g. SM 12.0). Lower +> `tensor_parallel_size` if you hit that. `lm_eval_tensorrt_llm.py` (`--model trt-llm`) has been removed; use the command above. diff --git a/examples/llm_eval/lm_eval_trtllm.py b/examples/llm_eval/lm_eval_trtllm.py index b9a4b2a92da..afcfc350541 100644 --- a/examples/llm_eval/lm_eval_trtllm.py +++ b/examples/llm_eval/lm_eval_trtllm.py @@ -16,17 +16,22 @@ """Run lm-evaluation-harness against a TensorRT-LLM checkpoint. Entry point around lm-eval's built-in ``trtllm`` backend -(``lm_eval.models.trtllm_causallms``, new in 0.4.12). It exists only to correct that -backend's ``prompt_logprobs`` handling -- everything else is upstream. Drop this file and -call ``lm_eval`` directly once the fix lands upstream. +(``lm_eval.models.trtllm_causallms``, new in 0.4.12). It exists to correct that backend's +``prompt_logprobs`` handling and to forward ``kv_cache_free_gpu_memory_fraction`` to +TensorRT-LLM -- everything else is upstream. Drop this file and call ``lm_eval`` directly +once both land upstream. python lm_eval_trtllm.py --model trtllm \ --model_args model=,tokenizer=,\ -tensor_parallel_size=,max_batch_size=,max_input_len=4096 \ +tensor_parallel_size=,max_batch_size=,max_input_len=4096,\ +kv_cache_free_gpu_memory_fraction=0.7 \ --tasks --batch_size """ +import inspect import sys +from contextlib import contextmanager +from functools import partial from importlib.metadata import version from lm_eval.__main__ import cli_evaluate @@ -36,6 +41,7 @@ # 0.4.12 is the first release shipping lm_eval.models.trtllm_causallms. raise ImportError(f"lm_eval_trtllm.py requires lm-eval >= 0.4.12; found {version('lm_eval')}.") +from lm_eval.models import trtllm_causallms from lm_eval.models.trtllm_causallms import TRTLLM # TensorRT-LLM only started passing the prompt token ids into `compute_logprobs` in @@ -125,6 +131,53 @@ def _parse_logprobs(tokens: list[int], outputs, ctxlen: int) -> tuple[float, boo TRTLLM._parse_logprobs = staticmethod(_parse_logprobs) +# The backend builds `KvCacheConfig(enable_block_reuse=False)` and hands `LLM(...)` a fixed +# set of keys, dropping every other `--model_args` entry, so the KV cache always takes +# TensorRT-LLM's default 90% of free memory. On a large-memory GPU that leaves too little +# room for the `prompt_logprobs` buffers and the run dies with a CUDA OOM. +_UPSTREAM_TAKES_KV_CACHE_FRACTION = ( + "kv_cache_free_gpu_memory_fraction" in inspect.signature(TRTLLM.__init__).parameters +) + + +@contextmanager +def _kv_cache_fraction_applied(fraction: float): + """Make the backend's single ``KvCacheConfig(...)`` call carry ``fraction``.""" + original = getattr(trtllm_causallms, "KvCacheConfig", None) + if original is None: + # tensorrt_llm is not installed; the backend raises its own ModuleNotFoundError. + yield + return + + trtllm_causallms.KvCacheConfig = partial(original, free_gpu_memory_fraction=fraction) + try: + yield + finally: + trtllm_causallms.KvCacheConfig = original + + +_UPSTREAM_INIT = TRTLLM.__init__ + + +def _init(self, *args, kv_cache_free_gpu_memory_fraction=None, **kwargs) -> None: + """``TRTLLM.__init__`` plus one ``--model_args`` key, forwarded to the KV cache.""" + if kv_cache_free_gpu_memory_fraction is None: + _UPSTREAM_INIT(self, *args, **kwargs) + return + + fraction = float(kv_cache_free_gpu_memory_fraction) + if not 0.0 < fraction <= 1.0: + # Out of range, TensorRT-LLM either allocates nothing or fails deep in the engine. + raise ValueError(f"kv_cache_free_gpu_memory_fraction must be in (0, 1]; got {fraction}.") + + with _kv_cache_fraction_applied(fraction): + _UPSTREAM_INIT(self, *args, **kwargs) + + +if not _UPSTREAM_TAKES_KV_CACHE_FRACTION: + TRTLLM.__init__ = _init + + if __name__ == "__main__": # Warn up front so an unusable container is obvious before the model loads, but do not # abort: generative tasks are unaffected by the old prompt_logprobs layout. diff --git a/examples/llm_eval/mmlu.py b/examples/llm_eval/mmlu.py index b5e31b37bfc..d11acad40a3 100755 --- a/examples/llm_eval/mmlu.py +++ b/examples/llm_eval/mmlu.py @@ -42,6 +42,7 @@ import random import warnings from argparse import Namespace +from typing import Any import numpy as np import pandas as pd @@ -274,10 +275,19 @@ def main( # Enable automatic save/load of modelopt state huggingface checkpointing mto.enable_huggingface_checkpointing() model_path = kwargs["model_path"] + # Pop before the branch below: only the TensorRT-LLM path knows this key, and the + # HuggingFace models built by select_model() reject what they do not declare. + kv_cache_free_gpu_memory_fraction = kwargs.pop("kv_cache_free_gpu_memory_fraction", None) tokenizer = get_tokenizer(model_path, trust_remote_code=kwargs.get("trust_remote_code", False)) if kwargs.get("checkpoint_dir"): assert LLM is not None, "tensorrt_llm APIs could not be imported." medusa_choices = kwargs.get("medusa_choices") + # Only override the wrapper's own default when the caller asked for a value. + llm_kwargs: dict[str, Any] = ( + {} + if kv_cache_free_gpu_memory_fraction is None + else {"kv_cache_free_gpu_memory_fraction": float(kv_cache_free_gpu_memory_fraction)} + ) model = LLM( checkpoint_dir=kwargs["checkpoint_dir"], tokenizer=tokenizer, @@ -285,6 +295,7 @@ def main( max_seq_len=MAX_SEQ_LEN, max_batch_size=1, trust_remote_code=kwargs.get("trust_remote_code", False), + **llm_kwargs, ) else: model = select_model( diff --git a/modelopt/deploy/llm/generate.py b/modelopt/deploy/llm/generate.py index 39306504137..b1221ea8b3e 100644 --- a/modelopt/deploy/llm/generate.py +++ b/modelopt/deploy/llm/generate.py @@ -63,6 +63,7 @@ def __init__( max_seq_len: int = 0, max_batch_size: int = 0, enable_kv_cache_reuse: bool = True, + kv_cache_free_gpu_memory_fraction: float = 0.7, ): """Initializes the LLM runner class. @@ -78,6 +79,10 @@ def __init__( requesting context logits (e.g. lm-eval loglikelihood tasks): with prefix block reuse, shared-prefix requests only return logits for the recomputed suffix, which breaks per-token logprob computation. + kv_cache_free_gpu_memory_fraction: fraction of the GPU memory left after loading the + model weights that the KV cache may use. Lower it when the engine runs out of + memory; TensorRT-LLM's own default of 0.9 is often too aggressive here, since + context logits and logprob buffers are allocated on top of the KV cache. """ with open(Path(checkpoint_dir) / "config.json") as config_file: config = json.load(config_file) @@ -124,8 +129,15 @@ def _find_max_position_embeddings(cfg: dict) -> int | None: enable_attention_dp = True break - # Sometimes 90% of the GPU memory is not enough for the TRT LLM torch engine. - trt_kv_cache_config = TRT_KvCacheConfig(free_gpu_memory_fraction=0.7) + # Sometimes 90% of the GPU memory is not enough for the TRT LLM torch engine, hence + # the lower default; callers can tune it further. + assert 0.0 < kv_cache_free_gpu_memory_fraction <= 1.0, ( + "kv_cache_free_gpu_memory_fraction must be in (0, 1]; got " + f"{kv_cache_free_gpu_memory_fraction}." + ) + trt_kv_cache_config = TRT_KvCacheConfig( + free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction + ) trt_kv_cache_config.max_tokens = self._max_seq_len * ( max_batch_size if max_batch_size > 0 else 8 ) diff --git a/tests/examples/hf_ptq/test_run_tensorrt_llm.py b/tests/examples/hf_ptq/test_run_tensorrt_llm.py index cf0042d470d..fb138ec4c01 100644 --- a/tests/examples/hf_ptq/test_run_tensorrt_llm.py +++ b/tests/examples/hf_ptq/test_run_tensorrt_llm.py @@ -83,6 +83,7 @@ def test_run_forwards_trust_remote_code(run_tensorrt_llm, monkeypatch, trust_rem max_output_len=8, input_texts="hello|world", trust_remote_code=trust_remote_code, + kv_cache_free_gpu_memory_fraction=0.7, ) ) @@ -91,3 +92,30 @@ def test_run_forwards_trust_remote_code(run_tensorrt_llm, monkeypatch, trust_rem assert tokenizer_kwargs["trust_remote_code"] is trust_remote_code # Context logits are requested below, so KV cache reuse must stay off. assert _RecordingLLM.last.kwargs["enable_kv_cache_reuse"] is False + + +def test_run_forwards_kv_cache_free_gpu_memory_fraction(run_tensorrt_llm, monkeypatch): + """The KV cache share has to reach the engine, not just be parsed (nvbug 6701763).""" + monkeypatch.setattr(run_tensorrt_llm, "get_tokenizer", lambda **kwargs: object()) + + run_tensorrt_llm.run( + SimpleNamespace( + tokenizer="", + checkpoint_dir="/fake/checkpoint", + max_output_len=8, + input_texts="hello", + trust_remote_code=False, + kv_cache_free_gpu_memory_fraction=0.5, + ) + ) + + assert _RecordingLLM.last.kwargs["kv_cache_free_gpu_memory_fraction"] == 0.5 + + +def test_kv_cache_free_gpu_memory_fraction_defaults_below_the_trtllm_default( + run_tensorrt_llm, monkeypatch +): + """TensorRT-LLM's own 0.9 leaves too little room for context logits; keep 0.7.""" + monkeypatch.setattr(sys, "argv", ["run_tensorrt_llm.py", "--checkpoint_dir", "/fake"]) + + assert run_tensorrt_llm.parse_arguments().kv_cache_free_gpu_memory_fraction == 0.7 diff --git a/tests/examples/llm_eval/test_lm_eval_trtllm.py b/tests/examples/llm_eval/test_lm_eval_trtllm.py index 19f0978e16f..548024eedf3 100644 --- a/tests/examples/llm_eval/test_lm_eval_trtllm.py +++ b/tests/examples/llm_eval/test_lm_eval_trtllm.py @@ -20,12 +20,16 @@ object, which these tests stub out. """ +import importlib.machinery +import inspect import sys import types from dataclasses import dataclass from pathlib import Path +from types import SimpleNamespace import pytest +import transformers # Skip on the backend module, not just `lm_eval`: it currently guards its `tensorrt_llm` # import, but if that ever becomes eager this should skip rather than error. @@ -38,6 +42,8 @@ sys.path.insert(0, str(_LLM_EVAL_DIR)) import lm_eval_trtllm +from lm_eval.models.trtllm_causallms import TRTLLM +from lm_eval.utils import simple_parse_args_string # TensorRT-LLM aligns prompt_logprobs to the *next* token, so entry i holds the # distribution that predicted tokens[i + 1]. These fixtures mirror that layout. @@ -153,12 +159,12 @@ def test_trust_remote_code_not_injected_when_unset(monkeypatch): def test_trtllm_backend_accepts_trust_remote_code(): - """The key lm-eval injects has to be a parameter the backend actually takes.""" - import inspect + """The key lm-eval injects has to be a parameter the backend actually takes. - from lm_eval.models.trtllm_causallms import TRTLLM - - assert "trust_remote_code" in inspect.signature(TRTLLM.__init__).parameters + Inspected on the upstream ``__init__``: ``TRTLLM.__init__`` is the wrapper below, which + takes ``**kwargs`` and forwards them, so its own signature proves nothing. + """ + assert "trust_remote_code" in inspect.signature(lm_eval_trtllm._UPSTREAM_INIT).parameters def _fake_trtllm(monkeypatch, version): @@ -191,6 +197,120 @@ def test_version_is_checked_before_scoring(monkeypatch): ) +class _RecordingKvCacheConfig: + """Stand-in for ``tensorrt_llm.llmapi.KvCacheConfig``, capturing its kwargs.""" + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +@pytest.fixture +def stub_engine(monkeypatch): + """Make the upstream ``__init__`` runnable without tensorrt_llm or a real tokenizer. + + Everything between the tokenizer load and the ``LLM(...)`` call is upstream code, so + this exercises the real construction rather than a stand-in for it. + """ + # The backend gates on find_spec(), which reads __spec__ off an already-imported module. + trtllm = types.ModuleType("tensorrt_llm") + trtllm.__spec__ = importlib.machinery.ModuleSpec("tensorrt_llm", loader=None) + monkeypatch.setitem(sys.modules, "tensorrt_llm", trtllm) + + class _LLM: + last_kwargs: dict = {} + + def __init__(self, **kwargs): + _LLM.last_kwargs = kwargs + + monkeypatch.setattr(lm_eval_trtllm.trtllm_causallms, "LLM", _LLM, raising=False) + monkeypatch.setattr( + lm_eval_trtllm.trtllm_causallms, "KvCacheConfig", _RecordingKvCacheConfig, raising=False + ) + + class _Tokenizer: + pad_token_id = 0 + eos_token_id = 0 + eos_token = "" + + monkeypatch.setattr( + transformers, "AutoTokenizer", SimpleNamespace(from_pretrained=lambda *a, **k: _Tokenizer()) + ) + return _LLM + + +def _build(**model_args): + """Instantiate the backend the way lm-eval does from a parsed `--model_args` string.""" + # add_bos_token is pinned only to skip the tokenizer round-trip that auto-detects it. + return TRTLLM.create_from_arg_obj({"model": "/ckpt", "add_bos_token": False, **model_args}) + + +def test_kv_cache_fraction_reaches_the_engine(stub_engine): + """`--model_args kv_cache_free_gpu_memory_fraction=...` must reach `KvCacheConfig`. + + Upstream drops unknown `--model_args` keys silently, so without the patch the KV cache + takes TensorRT-LLM's default 90% of free memory and large-memory GPUs OOM (nvbug + 6701763). + """ + _build(kv_cache_free_gpu_memory_fraction=0.5) + + kv_cache_config = stub_engine.last_kwargs["kv_cache_config"] + assert kv_cache_config.kwargs == {"enable_block_reuse": False, "free_gpu_memory_fraction": 0.5} + + +def test_kv_cache_fraction_is_not_injected_when_unset(stub_engine): + """Unset, TensorRT-LLM's own default must still apply -- we add no second default.""" + _build() + + assert "free_gpu_memory_fraction" not in stub_engine.last_kwargs["kv_cache_config"].kwargs + + +def test_kv_cache_fraction_is_parsed_from_a_model_args_string(): + """lm-eval coerces `--model_args` values itself; the key has to arrive as a float.""" + parsed = simple_parse_args_string("model=/ckpt,kv_cache_free_gpu_memory_fraction=0.5") + + assert parsed["kv_cache_free_gpu_memory_fraction"] == 0.5 + + +def test_kv_cache_fraction_patch_is_reverted(monkeypatch): + """Nothing may leak past the constructor; the patch is process-wide while held.""" + monkeypatch.setattr( + lm_eval_trtllm.trtllm_causallms, "KvCacheConfig", _RecordingKvCacheConfig, raising=False + ) + + with pytest.raises(RuntimeError), lm_eval_trtllm._kv_cache_fraction_applied(0.5): + raise RuntimeError("model load failed") + + assert lm_eval_trtllm.trtllm_causallms.KvCacheConfig is _RecordingKvCacheConfig + + +def test_kv_cache_fraction_without_tensorrt_llm(monkeypatch): + """Without tensorrt_llm the backend raises its own error; don't pre-empt it here.""" + monkeypatch.delattr(lm_eval_trtllm.trtllm_causallms, "KvCacheConfig", raising=False) + + with lm_eval_trtllm._kv_cache_fraction_applied(0.5): + pass + + +@pytest.mark.parametrize("fraction", [0.0, -0.1, 1.5]) +def test_init_rejects_an_out_of_range_fraction(fraction): + """Out of range, TensorRT-LLM either allocates nothing or fails deep in the engine.""" + with pytest.raises(ValueError, match="must be in"): + lm_eval_trtllm._init(object(), model="/ckpt", kv_cache_free_gpu_memory_fraction=fraction) + + +def test_init_patch_is_installed(): + assert TRTLLM.__init__ is lm_eval_trtllm._init + + +def test_upstream_still_drops_the_kv_cache_fraction(): + """Tripwire: when upstream takes the argument itself, drop the `__init__` patch.""" + assert lm_eval_trtllm._UPSTREAM_TAKES_KV_CACHE_FRACTION is False + + source = inspect.getsource(lm_eval_trtllm._UPSTREAM_INIT) + # It accepts **kwargs, and then builds the LLM kwargs from a fixed set of names. + assert "free_gpu_memory_fraction" not in source + + def test_upstream_is_still_misaligned(): """Tripwire: when this fails, upstream fixed the bug and this file can be deleted.