Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#1,577 1 commento 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
55/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python
Ambito
ai-infra-agents, performance

Direzione di ricerca

Inizia da modelopt/deploy/llm/generate.py intorno alla riga 122 e da examples/llm_eval/lm_eval_tensorrt_llm.py, in particolare parse_logprobs. Riproduci il problema con il comando MMLU di lm-eval fornito e analizza richieste ripetute con il caching dei prefissi abilitato. Il lavoro è completato quando sono ripristinati i logits per l'intera lunghezza e una valutazione corretta della loglikelihood, mentre il riutilizzo della cache dei prefissi rimane abilitato, con risultati vicini alla baseline HF.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

bug

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

Tasks Version Filter n-shot Metric Value Stderr
philosophy 1 none 5 acc ↑ 0.1865 ± 0.0221 
 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

Tasks Version Filter n-shot Metric Value Stderr
philosophy 1 none 5 acc ↑ 0.3344 ± 0.0268
  • 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

Tasks Version Filter n-shot Metric Value Stderr
philosophy 1 none 5 acc ↑ 0.3151 ± 0.0264
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.
Lingua principale
Python
Stelle
5.2k
Fork
736
Merge medio
4g 2h
PR unite (30g)
160

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/Model-Optimizer

Tutte le issue di NVIDIA/Model-Optimizer

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.