Skip to content

[Bug]: Clamp very small non-zero temperature values to avoid numerical instability #15715

Description

@chfeng-cs

System Info

CPU architecture: x86_64

CPU: Intel Xeon Gold 6462C, 8 vCPUs (4 cores × 2 threads), max 3900 MHz

GPU name: NVIDIA L20
GPU memory size: 48G

TensorRT-LLM branch or tag: main
TensorRT-LLM commit: 28acbad
PyTorch: 2.10.0+cu130
CUDA: 13.0.88
cuDNN: 9.15.0
NVIDIA driver version: 550.163.01

OS: Ubuntu 24.04.3 LTS

Who can help?

No response

Information

  • The official example scripts
  • My own modified scripts

Tasks

  • An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
  • My own task or dataset (give details below)

Reproduction

from tensorrt_llm.sampling_params import SamplingParams

# Before the fix: very small non-zero temperature bypasses validation
# and enters the sampling backend, causing potential inf/nan in logits/temperature
sp = SamplingParams(temperature=1e-12, top_k=50)
print(sp.temperature)  # prints 1e-12 — dangerously small

# Simulating what happens inside the backend:
import torch
logits = torch.tensor([1.0, 2.0, 3.0])
temperature = 1e-12
scaled = logits / temperature
print(scaled)  # tensor([1.0e+12, 2.0e+12, 3.0e+12]) → leads to inf/nan in softmax
print(torch.softmax(scaled, dim=-1))  # tensor([0., 0., 1.]) or nan depending on dtype

Expected behavior

SamplingParams should reject or safely normalize temperature values that are
positive but below a numerical stability threshold.

Specifically:

  • temperature=0.0 → kept as 0.0 (greedy decoding)
  • temperature=1e-12 → clamped to MIN_SAMPLING_TEMPERATURE (1e-2)
  • temperature=1.0 → kept as 1.0 (normal sampling)

actual behavior

amplingParams only validates temperature >= 0, so values like 1e-12
pass validation unchanged and are forwarded to the sampling backend.

Inside the backend, logits / temperature with temperature=1e-12 produces
extremely large logit values (e.g., 1e+12), which overflow to inf or nan
in float16/bfloat16 precision, leading to undefined sampling behavior.

Example:

import torch
logits = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)
scaled = logits / 1e-12
print(scaled)
# tensor([inf, inf, inf], dtype=torch.float16)

print(torch.softmax(scaled, dim=-1))
# tensor([nan, nan, nan], dtype=torch.float16)

additional notes

  1. This fix mirrors the same numerical stability guard used by vLLM (temperature = max(temperature, 1e-2) for non-zero values).

  2. The issue is most severe in float16/bfloat16: typical logit values (~10) divided by 1e-6 already exceed float16's max range (~65504), producing inf/nan in softmax.

  3. temperature=0.0 (greedy decoding) is intentionally excluded from the clamp and its behavior is unchanged.

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Decoding/Sampling<NV>Token sampling algorithms in TRTLLM for text gen (top-k, top-p, beam).bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions