NVIDIA · cjluo-nv · Sep 1, 2026 · coderabbitai · Sep 1, 2026
@@ -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)
 ^^^^^^^^^^^^^^^^^

@@ -176,6 +176,8 @@ scripts/huggingface_example.sh --model $HF_PATH --quant <QFORMAT> --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

@@ -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)

@@ -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),

@@ -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

@@ -114,7 +114,7 @@ checkpoint directly with the TensorRT-LLM LLM API.
 
 ```sh
 python lm_eval_trtllm.py --model trtllm \
-    --model_args model=<Quantized checkpoint dir>,tokenizer=<HF model folder>,tensor_parallel_size=<tp>,max_batch_size=<max batch size>,max_input_len=4096,max_output_len=512 \
+    --model_args model=<Quantized checkpoint dir>,tokenizer=<HF model folder>,tensor_parallel_size=<tp>,max_batch_size=<max batch size>,max_input_len=4096,max_output_len=512,kv_cache_free_gpu_memory_fraction=0.7 \
     --tasks <comma separated tasks> \
     --batch_size <max 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.
 

@@ -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=<quantized checkpoint dir>,tokenizer=<HF model folder>,\
-tensor_parallel_size=<tp>,max_batch_size=<max batch size>,max_input_len=4096 \
+tensor_parallel_size=<tp>,max_batch_size=<max batch size>,max_input_len=4096,\
+kv_cache_free_gpu_memory_fraction=0.7 \
         --tasks <comma separated tasks> --batch_size <max 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.

@@ -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,17 +275,27 @@ 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,
             medusa_choices=medusa_choices,
             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(

@@ -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
         )

@@ -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