fix(llm_eval): repair test_qwen3_eval_fp8 end-to-end (#1650)

### What does this PR do?

Type of change: Bug fix

`tests/examples/llm_eval/test_llm_eval.py::test_qwen3_eval_fp8` was
silently passing while its evals crashed, then began failing as a
timeout. This repairs the whole pipeline:

- **lm_eval `IndexError` (root cause):** TRT-LLM KV-cache prefix reuse
returns truncated `context_logits` for shared-prefix requests (e.g.
hellaswag's one-context / many-endings), which breaks `parse_logprobs`.
Add an `enable_kv_cache_reuse` flag to `modelopt.deploy.llm.LLM`
(default `True`, unchanged) and disable it for the eval deployment so
full-length context logits are returned.
- **Silent CI green:** `python eval.py | tee result.txt` returns `tee`'s
exit code, so a crashing eval was masked. Add `set -o pipefail` to
`huggingface_example.sh` so failures fail the test.
- **Long-prompt overflows:** with the tiny test model's toy tokenizer,
gsm8k/MMLU prompts exceed `max_seq_len`. Bump test
`max_position_embeddings` to 8192, skip MMLU prompts that don't fit even
at zero-shot, and add an MMLU sample limit (`--mmlu_limit`).
- **human-eval build failures:** install with `--no-build-isolation`
(`pkg_resources` is absent in pip's isolated build env), patch its
malformed `console_scripts` entry point, and pin the clone.
- **Cleanups:** gate the post-quant `run_tensorrt_llm.py` smoke test
behind the `quant` task (eval tasks deploy on their own; ~45s saved for
eval-only runs); replace the SIGPIPE-prone serve-readiness `tail -f |
while` with a poll loop (required under `pipefail`).

### Usage

N/A — example/test fix.

### Testing

All four eval tasks verified end-to-end in the CI container (TRT-LLM
1.3.0rc17, RTX 6000 Ada): lm_eval (hellaswag + gsm8k), MMLU, and
simple_eval (humaneval) all complete with exit 0 and no
`IndexError`/overflow. Cold full run ≈ 340s on this GPU.

CI test on 2-gpu:
https://github.com/NVIDIA/Model-Optimizer/actions/runs/27154417497/job/80153551154

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅ (new `enable_kv_cache_reuse`
defaults to current behavior; new script flags are optional)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (no new
dependencies)
- Did you write any new necessary tests?: N/A (fixes and strengthens an
existing test)
- Did you update Changelog?: N/A (bug fix to examples/tests)
- Did you get Claude approval on this PR?: ❌ (pending)

### Additional Information

The full test runs ~340s on an RTX 6000 Ada; CI runners are historically
slower, while `@pytest.mark.timeout` is set to 600 — worth watching the
first CI run and bumping if it's close.

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
  * Added an option to limit MMLU evaluation length.

* **Bug Fixes**
* Disabled KV-cache prefix reuse for evaluations needing per-token
context logits to prevent truncated/incorrect logprobs.
* Skip examples whose prompts remain too long; warn and report accuracy
as NaN if all examples are skipped.

* **Chores / Scripts**
* Improved example scripts for reproducible installs, patched entry
point handling, pipeline failure detection, conditional test invocation,
polling-based log wait, and a new CLI flag for MMLU limits.

* **Tests**
* Increased timeout and prompt headroom; capped MMLU smoke tests for
speed.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: kevalmorabia97

examples/llm_eval/lm_eval_tensorrt_llm.py

@@ -64,6 +64,10 @@ def __init__(
             tokenizer=self.tokenizer,
             max_batch_size=int(batch_size),
             max_seq_len=max_length,
+            # Loglikelihood tasks request context logits. KV cache prefix reuse would return
+            # logits only for the recomputed suffix on shared-prefix requests (e.g. hellaswag),
+            # truncating context_logits and breaking parse_logprobs. Disable it.
+            enable_kv_cache_reuse=False,
         )
         self.max_length = max_length - 1
         logger.info("Loaded TRT-LLM")

examples/llm_eval/mmlu.py

@@ -183,7 +183,10 @@ def gen_prompt(train_df, subject, k=-1):
 def evaluate(args, subject, model: EvalModel | LLM, dev_df, test_df):
     cors = []
     all_probs = []
-    for i in range(test_df.shape[0]):
+    num_examples = test_df.shape[0]
+    if args.limit is not None:
+        num_examples = min(num_examples, args.limit)
+    for i in range(num_examples):
         # get prompt and make sure it fits
         k = args.ntrain
         prompt_end = format_example(test_df, i, include_answer=False)
@@ -201,6 +204,12 @@ def check_valid_length(model, prompt):
             train_prompt = gen_prompt(dev_df, subject, k)
             prompt = train_prompt + prompt_end
 
+        # Skip examples that do not fit even at zero-shot, otherwise the backend rejects
+        # prompts longer than max_seq_len and aborts the whole evaluation.
+        if not check_valid_length(model, prompt):
+            print(f"Skipping {subject} example {i}: prompt exceeds max_seq_len even at 0-shot.")
+            continue
+
         label = test_df.iloc[i, test_df.shape[1] - 1]
         if isinstance(model, EvalModel):
             pred = model.run(prompt)
@@ -212,7 +221,11 @@ def check_valid_length(model, prompt):
         cors.append(cor)
         all_probs.append(probs)
 
-    acc = np.mean(cors)
+    if not cors:
+        # Every example was skipped (all prompts exceeded max_seq_len). Surface it instead of
+        # silently producing a nan accuracy downstream.
+        print(f"WARNING: all {subject} examples were skipped; reporting accuracy as nan.")
+    acc = np.mean(cors) if cors else float("nan")
     cors = np.array(cors)
 
     all_probs = np.array(all_probs)
@@ -233,8 +246,12 @@ def main(
     auto_quantize_score_size: int = 128,
     auto_quantize_checkpoint: str | None = None,
     sparse_cfg: str | None = None,
+    limit: int | None = None,
     **kwargs,
 ):
+    if limit is not None and limit <= 0:
+        raise ValueError(f"limit must be a positive integer when provided, got {limit}.")
+
     random.seed(RAND_SEED)
     np.random.seed(RAND_SEED)
 

examples/llm_eval/run_simple_eval.sh

@@ -28,11 +28,23 @@ if [ ! -d "human-eval" ]; then
     git clone https://github.com/openai/human-eval.git
 fi
 
+# Pin to a known commit for reproducibility (and so the entry-point patch below matches), forcing
+# it every run so a reused checkout cannot drift to an arbitrary revision. -f discards the patch
+# applied to setup.py on a previous run before re-applying it below.
+git -C human-eval checkout -q -f 6d43fb980f9fee3c892a914eda09951f772ad10d
+
+# human-eval's console_scripts entry point lacks the ":callable" suffix, which newer pip/setuptools
+# reject ("A callable suffix is required"). The target module defines main(), so point at it.
+sed -i 's|human_eval\.evaluate_functional_correctness"|human_eval.evaluate_functional_correctness:main"|' human-eval/setup.py
+
 if [ ! -d "simple-evals" ]; then
     git clone https://github.com/openai/simple-evals.git
 fi
 
-pip install -e human-eval
+# --no-build-isolation: human-eval's legacy setup.py imports pkg_resources at build time,
+# which pip's isolated build env does not provide with newer setuptools. Build against the
+# base environment (which has setuptools/pkg_resources) instead.
+pip install -e human-eval --no-build-isolation
 pip install openai
 
 pushd simple-evals

examples/llm_ptq/run_tensorrt_llm.py

@@ -66,7 +66,14 @@ def run(args):
 
     print("TensorRT-LLM example outputs:")
 
-    llm = LLM(args.checkpoint_dir, tokenizer=tokenizer, max_batch_size=len(input_texts))
+    # generate_context_logits() below requires KV cache reuse disabled: with prefix block reuse,
+    # shared-prefix inputs return truncated (silently incorrect) context logits.
+    llm = LLM(
+        args.checkpoint_dir,
+        tokenizer=tokenizer,
+        max_batch_size=len(input_texts),
+        enable_kv_cache_reuse=False,
+    )
     torch.cuda.cudart().cudaProfilerStart()
     outputs = llm.generate_text(input_texts, args.max_output_len)
     torch.cuda.cudart().cudaProfilerStop()

examples/llm_ptq/scripts/huggingface_example.sh

@@ -29,6 +29,10 @@ for i in $(env | grep ^SLURM_ | cut -d"=" -f 1); do unset -v $i; done
 for i in $(env | grep ^PMI_ | cut -d"=" -f 1); do unset -v $i; done
 for i in $(env | grep ^PMIX_ | cut -d"=" -f 1); do unset -v $i; done
 
+# Fail on errors inside pipelines (e.g. `python eval.py | tee result.txt`), otherwise a crashing
+# eval is masked by tee's exit code and the script passes silently.
+set -o pipefail
+
 if [ -z "$MODEL_PATH" ]; then
     echo "Unsupported model argument: Expected a huggingface model path or model name" >&2
     exit 1
@@ -216,7 +220,11 @@ if [[ $TASKS =~ "quant" ]] || [[ ! -d "$SAVE_PATH" ]] || [[ ! $(ls -A $SAVE_PATH
         RUN_ARGS+=" --trust_remote_code "
     fi
 
-    python run_tensorrt_llm.py --checkpoint_dir=$SAVE_PATH $RUN_ARGS
+    # 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
+        python run_tensorrt_llm.py --checkpoint_dir=$SAVE_PATH $RUN_ARGS
+    fi
 fi
 
 if [[ -d "${MODEL_PATH}" ]]; then
@@ -285,11 +293,16 @@ if [[ $TASKS =~ "mmlu" ]]; then
         tar -xf /tmp/mmlu.tar -C data && mv data/data $MMLU_DATA_PATH
     fi
 
+    mmlu_flags=""
+    if [ -n "$MMLU_LIMIT" ]; then
+        mmlu_flags+=" --limit $MMLU_LIMIT "
+    fi
+
     python mmlu.py \
         --model_name causal \
         --model_path $MODEL_ABS_PATH \
         --checkpoint_dir $SAVE_PATH \
-        --data_dir $MMLU_DATA_PATH | tee $MMLU_RESULT
+        --data_dir $MMLU_DATA_PATH $mmlu_flags | tee $MMLU_RESULT
     popd
 
 fi
@@ -304,16 +317,16 @@ if [[ $TASKS =~ "livecodebench" || $TASKS =~ "simple_eval" ]]; then
     trtllm-serve $SAVE_PATH --host 0.0.0.0 --port $PORT >$SAVE_PATH/serve.txt 2>&1 &
     SERVE_PID=$!
 
-    tail -f $SAVE_PATH/serve.txt | while read line; do
-        if echo "$line" | grep -q "Application startup complete"; then
-            echo "Application startup complete."
-            break
-        fi
+    # Poll the log instead of `tail -f | while ... break`: under `set -o pipefail` (set above),
+    # breaking out of that pipeline leaves tail to die by SIGPIPE, which would abort the script.
+    while ! grep -q "Application startup complete" $SAVE_PATH/serve.txt 2>/dev/null; do
         if ! kill -0 $SERVE_PID 2>/dev/null; then
             echo "trtllm-serve has exited."
             exit 1
         fi
+        sleep 2
     done
+    echo "Application startup complete."
 
     pushd ../llm_eval/
 

examples/llm_ptq/scripts/parser.sh

@@ -28,6 +28,7 @@ parse_options() {
     LM_EVAL_TASKS="mmlu,gsm8k"
     LM_EVAL_LIMIT=
     SIMPLE_EVAL_TASKS="mmlu"
+    MMLU_LIMIT=
 
     TASKS="quant"
 
@@ -38,7 +39,7 @@ parse_options() {
     CAST_MXFP4_TO_NVFP4=false
 
   # Parse command-line options
-  ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4" -n "$0" -- "$@")
+  ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4" -n "$0" -- "$@")
 
   eval set -- "$ARGS"
   while true; do
@@ -61,6 +62,7 @@ parse_options() {
       --lm_eval_limit ) LM_EVAL_LIMIT="$2"; shift 2;;
       --simple_eval_tasks ) SIMPLE_EVAL_TASKS="$2"; shift 2;;
       --simple_eval_limit ) SIMPLE_EVAL_LIMIT="$2"; shift 2;;
+      --mmlu_limit ) MMLU_LIMIT="$2"; shift 2;;
       --trust_remote_code ) TRUST_REMOTE_CODE=true; shift;;
       --use_seq_device_map ) USE_SEQ_DEVICE_MAP=true; shift;;
       --gpu_max_mem_percentage ) GPU_MAX_MEM_PERCENTAGE="$2"; shift 2;;
@@ -161,6 +163,7 @@ parse_options() {
   echo "lm_eval_limit: $LM_EVAL_LIMIT"
   echo "simple_eval_tasks: $SIMPLE_EVAL_TASKS"
   echo "simple_eval_limit: $SIMPLE_EVAL_LIMIT"
+  echo "mmlu_limit: $MMLU_LIMIT"
   echo "num_sample: $NUM_SAMPLES"
   echo "use_seq_device_map: $USE_SEQ_DEVICE_MAP"
   echo "gpu_max_mem_percentage: $GPU_MAX_MEM_PERCENTAGE"

modelopt/deploy/llm/generate.py

@@ -62,6 +62,7 @@ def __init__(
         trust_remote_code: bool = False,
         max_seq_len: int = 0,
         max_batch_size: int = 0,
+        enable_kv_cache_reuse: bool = True,
     ):
         """Initializes the LLM runner class.
 
@@ -73,6 +74,10 @@ def __init__(
             trust_remote_code: whether to trust the remote code (for the torch backend).
             max_seq_len: Max sequence length for the LLM backend. If 0, it is not specified.
             max_batch_size: Max batch size for the LLM backend. If 0, it is not specified.
+            enable_kv_cache_reuse: whether to enable KV cache block reuse. Must be disabled when
+                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.
         """
         with open(Path(checkpoint_dir) / "config.json") as config_file:
             config = json.load(config_file)
@@ -124,6 +129,8 @@ def _find_max_position_embeddings(cfg: dict) -> int | None:
         trt_kv_cache_config.max_tokens = self._max_seq_len * (
             max_batch_size if max_batch_size > 0 else 8
         )
+        trt_kv_cache_config.enable_block_reuse = enable_kv_cache_reuse
+        self._enable_kv_cache_reuse = enable_kv_cache_reuse
 
         cuda_graph_config = None
         if max_batch_size > 0:
@@ -281,6 +288,11 @@ def generate_context_logits(
         assert self._support_context_logits_and_stop_words, (
             "Context logits are not supported with the current tensorrt_llm version."
         )
+        assert not self._enable_kv_cache_reuse, (
+            "Context logits require enable_kv_cache_reuse=False: with KV cache prefix reuse, "
+            "shared-prefix requests only return logits for the recomputed suffix, producing "
+            "truncated (and silently incorrect) context logits."
+        )
         assert temperature >= 0.0, "Temperature must be greater than 0.0."
 
         kwargs = _sanitize_temperature_and_top_p(temperature, top_p)

tests/examples/llm_eval/test_llm_eval.py

@@ -41,11 +41,13 @@ def test_lm_eval_hf(tmp_path):
 
 
 @minimum_sm(89)
-@pytest.mark.timeout(600)
+@pytest.mark.timeout(900)
 def test_qwen3_eval_fp8(tmp_path):
-    # Bump max_position_embeddings: TRT-LLM serve rejects prompts longer than
-    # max_seq_len, and the default (32) is shorter than even simple MMLU prompts.
-    model_dir = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, max_position_embeddings=2048)
+    # Bump max_position_embeddings: TRT-LLM serve rejects prompts longer than max_seq_len.
+    # The default (32) is shorter than even simple MMLU prompts, and 2048 is shorter than
+    # 5-shot gsm8k prompts (~3.9k tokens). The eval LLM caps max_seq_len at max_gen_toks + 4096,
+    # so 8192 leaves headroom for the longest prompts we evaluate.
+    model_dir = create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, max_position_embeddings=8192)
     try:
         run_llm_ptq_command(
             model=str(model_dir),
@@ -56,6 +58,9 @@ def test_qwen3_eval_fp8(tmp_path):
             simple_eval_tasks="humaneval",
             lm_eval_limit=16,
             simple_eval_limit=16,
+            # MMLU has no inherent sample cap and otherwise evaluates the full ~14k-question
+            # test set; limit it like lm_eval/simple_eval to keep this a fast smoke test.
+            mmlu_limit=16,
             output=128,  # Cap generation length: gsm8k/humaneval otherwise generate up to 1024 tokens/sample
             batch=8,
         )