diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 221f1b4b6ccc..aa96eccc8508 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -1,4 +1,5 @@ import json +import math import os from abc import ABC, abstractmethod from dataclasses import dataclass, field, fields @@ -122,6 +123,12 @@ def __call__( pass # noqa +# Smallest non-zero temperature forwarded to the sampling backend. Positive +# values below this overflow logits / temperature to inf/nan in fp16/bf16, so +# they are clamped up to this floor. temperature == 0 (greedy) is left as-is. +MIN_SAMPLING_TEMPERATURE = 1e-2 + + @dataclass(slots=True, kw_only=True) class SamplingParams: """Sampling parameters for text generation. @@ -318,8 +325,14 @@ def _validate(self): raise ValueError(f"require 0 <= top_p <= 1, got top_p={self.top_p}") if self.top_k is not None and self.top_k < 0: raise ValueError(f"require top_k >= 0, got top_k={self.top_k}") + if self.temperature is not None and not math.isfinite(self.temperature): + raise ValueError(f"temperature must be finite, got temperature={self.temperature}") if self.temperature is not None and self.temperature < 0: raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + # Clamp very small non-zero temperatures up to a numerically safe floor; + # keep 0.0 as-is so greedy decoding is unaffected. + if self.temperature is not None and 0 < self.temperature < MIN_SAMPLING_TEMPERATURE: + self.temperature = MIN_SAMPLING_TEMPERATURE if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py new file mode 100644 index 000000000000..f4b454beaea7 --- /dev/null +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math + +import pytest + +from tensorrt_llm.sampling_params import MIN_SAMPLING_TEMPERATURE, SamplingParams + + +def test_temperature_none_unchanged() -> None: + assert SamplingParams().temperature is None + + +def test_temperature_zero_kept_for_greedy() -> None: + # 0.0 means greedy decoding and must not be clamped. + assert SamplingParams(temperature=0.0).temperature == 0.0 + + +@pytest.mark.parametrize("tiny", [1e-12, 1e-6, MIN_SAMPLING_TEMPERATURE / 2]) +def test_tiny_positive_temperature_clamped(tiny: float) -> None: + # Assert against the contractual floor (1e-2), not the constant. + assert SamplingParams(temperature=tiny).temperature == 1e-2 + + +@pytest.mark.parametrize("temp", [MIN_SAMPLING_TEMPERATURE, 0.5, 1.0, 2.0]) +def test_normal_temperature_unchanged(temp: float) -> None: + assert SamplingParams(temperature=temp).temperature == temp + + +def test_negative_temperature_rejected() -> None: + with pytest.raises(ValueError, match="temperature"): + SamplingParams(temperature=-1.0) + + +@pytest.mark.parametrize("bad", [math.nan, math.inf, -math.inf]) +def test_non_finite_temperature_rejected(bad: float) -> None: + with pytest.raises(ValueError, match="finite"): + SamplingParams(temperature=bad)