generate_context_logits returns truncated logits when KV prefix cache is enabled, causing silently wrong loglikelihood evaluation #1577

Before submitting an issue, please make sure it hasn't been already addressed by searching through the existing and past issues.

Describe the bug

• modelopt.deploy.llm.LLM.generate_context_logits() is built on top of TensorRT-LLM's return_context_logits=True. When KV cache prefix reuse is active (which is the default), TensorRT-LLM skips the forward pass for cached prefix tokens and returns only the logits for the un-cached tail of the prompt — without any error or warning. This silently breaks every downstream loglikelihood-style evaluation (lm-eval-harness MMLU etc.) that uses generate_context_logits, because parse_logprobs slices into a logits tensor that is much shorter than the prompt and produces meaningless logprobs.

Steps/Code to reproduce bug

python examples/llm_eval/lm_eval_tensorrt_llm.py \
      --model trt-llm \
      --model_args tokenizer=/path/to/Llama-3.2-1B,checkpoint_dir=/path/to/Llama-3.2-1B \
      --tasks mmlu_philosophy \
      --num_fewshot 5 \
      --batch_size 1

❯ trt-llm (tokenizer=/root/infra/hf/Llama-3.2-1B,checkpoint_dir=/root/infra/hf/Llama-3.2-1B), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1

 lm_eval --model hf \
      --model_args pretrained=/path/to/Llama-3.2-1B,dtype=float16 \
      --tasks mmlu_philosophy \
      --num_fewshot 5 \
      --batch_size 1

❯ hf (pretrained=/root/infra/hf/Llama-3.2-1B,dtype=float16), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1

• The TRT-LLM result is below the 25% random baseline because every prompt after the first one returns truncated logits → parse_logprobs ends up summing to ~0.0 for every candidate → lm-eval picks option A by stable-sort tie-breaking, and accuracy collapses to "% of questions where A is the gold answer".

Root cause

• Adding debug prints to parse_logprobs shows the smoking gun:

 # request 1 (cold)
 ctxlen=469, logits_len=470, tokens_len=470   ← OK

 # request 2+ (prefix cache hit)
 [TensorRT-LLM][WARNING] [kv cache manager] storeContextBlocks: Can not find sequence for request 2048
 ctxlen=469, logits_len=1, tokens_len=470     ← BROKEN: only 1 logit returned

• The first request does a full prefill and returns 470 logits as expected. Subsequent requests in the same MMLU subtask share the 5-shot demo prefix → KV cache prefix reuse hits → the forward pass is skipped for the cached portion → TensorRT-LLM returns only the logits for the un-cached tail (often just 1 token, the last differing token of the question).

• parse_logprobs in examples/llm_eval/lm_eval_tensorrt_llm.py then does:

 logits_single_batch[(ctxlen_single_batch - 1) : -1]   # slice [468:-1] on a length-1 tensor

which yields an empty tensor → logprob_sum=0.0 for every continuation candidate → lm-eval cannot distinguish A/B/C/D.

Confirmed fix

• Setting enable_block_reuse=False in the KV cache config restores correctness. Verified locally: mmlu_philosophy 5-shot recovers from 0.1865 to 0.3151, matching the HF golden baseline.

# In modelopt/deploy/llm/generate.py, line ~122:

  trt_kv_cache_config = TRT_KvCacheConfig(
      free_gpu_memory_fraction=0.7,
      enable_block_reuse=False,
  )

❯ trt-llm (tokenizer=/root/infra/hf/Llama-3.2-1B,checkpoint_dir=/root/infra/hf/Llama-3.2-1B), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: 1

Things tried that did NOT fix it

• Switching to TensorRT-LLM's prompt_logprobs API (SamplingParams(prompt_logprobs=k)): on the tested TRT-LLM version, prompt_logprobs shares the same forward-output path as1 return_context_logits and is also truncated when prefix cache hits. The hope that prompt_logprobs would semantically guarantee full-length output (as it does in vLLM) does not hold for TensorRT-LLM in this version.

Expected behavior

Who can help?

• ?

System information

• Container used (if applicable): nvcr.io/nvidia/tensorrt-llm/release:1.2.0
• OS (e.g., Ubuntu 22.04, CentOS 7, Windows 10): Ubuntu 22.04
• CPU architecture (x86_64, aarch64): x86_64
• GPU name (e.g. H100, A100, L40S): NVIDIA RTX A4000
• GPU memory size: 16GB
• Number of GPUs: 1
• Library versions (if applicable):
  • Python: ?
  • ModelOpt version or commit hash: tag0.37.0 + 824c9457 ([NVBUG: 5617733] Update LLM generate API for modelopt LLM eval)
  • CUDA: ?
  • PyTorch: ?
  • Transformers: ?
  • TensorRT-LLM: v1.2.0
  • ONNXRuntime: ?
  • TensorRT: ?
• Any other details that may help: I'm curious if there is a way that can both fix this bug and use prefix caching.