[2/2] Support deterministic inference for temperature > 0 (#10678)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: hebiao064 <hebiaobuaa@gmail.com>

python/sglang/srt/sampling/sampling_batch_info.py

@@ -60,6 +60,9 @@ class SamplingBatchInfo:
         Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]
     ] = None
 
+    # Used for deterministic sampling
+    sampling_seed: Optional[torch.Tensor] = None
+
     # Device
     device: str = "cuda"
 
@@ -93,6 +96,15 @@ class SamplingBatchInfo:
         min_ps = torch.tensor(
             [r.sampling_params.min_p for r in reqs], dtype=torch.float, device=device
         )
+        sampling_seed = (
+            torch.tensor(
+                [r.sampling_params.sampling_seed for r in reqs],
+                dtype=torch.int32,
+                device=device,
+            )
+            if enable_deterministic
+            else None
+        )
 
         logit_bias = None
         if any(r.sampling_params.logit_bias is not None for r in reqs):
@@ -158,6 +170,7 @@ class SamplingBatchInfo:
             top_ps=top_ps,
             top_ks=top_ks,
             min_ps=min_ps,
+            sampling_seed=sampling_seed,
             is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
             need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
             need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
@@ -239,9 +252,11 @@ class SamplingBatchInfo:
             "top_ps",
             "top_ks",
             "min_ps",
+            "sampling_seed",
         ]:
             value = getattr(self, item, None)
-            setattr(self, item, value[keep_indices_device])
+            if value is not None:
+                setattr(self, item, value[keep_indices_device])
 
         if self.logit_bias is not None:
             self.logit_bias = self.logit_bias[keep_indices_device]
@@ -343,10 +358,12 @@ class SamplingBatchInfo:
             "top_ps",
             "top_ks",
             "min_ps",
+            "sampling_seed",
         ]:
             self_val = getattr(self, item, None)
             other_val = getattr(other, item, None)
-            setattr(self, item, torch.cat([self_val, other_val]))
+            if self_val is not None and other_val is not None:
+                setattr(self, item, torch.cat([self_val, other_val]))
 
         self.is_all_greedy &= other.is_all_greedy
         self.need_top_p_sampling |= other.need_top_p_sampling

python/sglang/srt/sampling/sampling_params.py

@@ -15,8 +15,11 @@
 
 from typing import Any, Dict, List, Optional, Union
 
+from sglang.srt.utils import get_bool_env_var
+
 _SAMPLING_EPS = 1e-6
 TOP_K_ALL = 1 << 30
+DEFAULT_SAMPLING_SEED = 42
 
 
 class SamplingParams:
@@ -53,6 +56,7 @@ class SamplingParams:
         custom_params: Optional[Dict[str, Any]] = None,
         stream_interval: Optional[int] = None,
         logit_bias: Optional[Dict[str, float]] = None,
+        sampling_seed: Optional[int] = None,
     ) -> None:
         self.max_new_tokens = max_new_tokens
         self.stop_strs = stop
@@ -80,6 +84,14 @@ class SamplingParams:
         self.custom_params = custom_params
         self.stream_interval = stream_interval
         self.logit_bias = logit_bias
+        # Used for deterministic sampling
+        if (
+            get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE")
+            and sampling_seed is None
+        ):
+            # If deterministic inference is enabled and sampling_seed is not set, use the default seed
+            sampling_seed = DEFAULT_SAMPLING_SEED
+        self.sampling_seed = sampling_seed
 
         # Process some special cases
         if 0 <= self.temperature < _SAMPLING_EPS: