[None][fix] Clamp very small non-zero temperature in SamplingParams

tensorrt_llm/sampling_params.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
+import math
 import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field, fields
@@ -27,6 +28,11 @@
 
 MAX_TOP_LOGPROBS = 20
 
+# Smallest non-zero temperature forwarded to the sampling backend. Values below
+# this overflow logits/temperature to inf/nan in float16/bfloat16 (mirrors the
+# guard vLLM uses). temperature == 0 is greedy decoding and is left untouched.
+MIN_SAMPLING_TEMPERATURE = 1e-2
+
 
 def validate_thinking_token_budget(value: Optional[Union[int, float, bool]]) -> Optional[int]:
     """Validate ``thinking_token_budget``; return ``None`` if unset."""
@@ -215,6 +221,7 @@ class SamplingParams:
         top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None.
         seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None.
         temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None.
+            Positive values smaller than 1e-2 are clamped up to 1e-2 to avoid inf/nan when dividing logits by the temperature in fp16/bf16; temperature = 0 stays greedy and non-finite values are rejected.
             The value None is treated as "not specified" in the following.
             If neither temperature, top_p, nor top_k are specified, sampling is greedy.
             If top_p < 1 and/or top_k > 1 are specified, sampling will proceed accordingly and temperature will default to temperature = 1.
@@ -370,8 +377,22 @@ 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 self.temperature < 0:
-            raise ValueError(f"require temperature >= 0, got temperature={self.temperature}")
+        if self.temperature is not None:
+            if not math.isfinite(self.temperature):
+                raise ValueError(
+                    f"require temperature to be finite, got temperature={self.temperature}"
+                )
+            if self.temperature < 0:
+                raise ValueError(f"require temperature >= 0, got temperature={self.temperature}")
+            # Clamp very small positive temperatures up to a numerically safe
+            # floor: logits / temperature otherwise overflows to inf/nan in
+            # fp16/bf16. temperature == 0 (greedy) is intentionally left as-is.
+            if 0 < self.temperature < MIN_SAMPLING_TEMPERATURE:
+                logger.debug(
+                    f"Clamping temperature from {self.temperature} to "
+                    f"{MIN_SAMPLING_TEMPERATURE} to avoid numerical instability."
+                )
+                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})")

tests/unittest/llmapi/test_sampling_params_temperature.py

@@ -0,0 +1,33 @@
+import math
+
+import pytest
+
+from tensorrt_llm.sampling_params import MIN_SAMPLING_TEMPERATURE, SamplingParams
+
+
+@pytest.mark.parametrize(
+    ("temperature", "expected"),
+    [
+        (None, None),  # unset -> unchanged
+        (0.0, 0.0),  # greedy -> untouched
+        (1e-12, MIN_SAMPLING_TEMPERATURE),  # tiny -> clamped (regression #15715)
+        (MIN_SAMPLING_TEMPERATURE - 1e-3, MIN_SAMPLING_TEMPERATURE),  # just below floor -> clamped
+        (MIN_SAMPLING_TEMPERATURE, MIN_SAMPLING_TEMPERATURE),  # boundary -> kept
+        (1.0, 1.0),  # normal -> kept
+    ],
+)
+def test_temperature_clamped_below_floor(temperature, expected):
+    # #15715: 0 < temperature < MIN reaches the backend and overflows
+    # logits / temperature to inf/nan in fp16/bf16; clamp it to the floor.
+    assert SamplingParams(temperature=temperature, top_k=2).temperature == expected
+
+
+def test_negative_temperature_rejected():
+    with pytest.raises(ValueError):
+        SamplingParams(temperature=-0.5)
+
+
+@pytest.mark.parametrize("temperature", [math.inf, -math.inf, math.nan])
+def test_non_finite_temperature_rejected(temperature):
+    with pytest.raises(ValueError):
+        SamplingParams(temperature=temperature)