chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: hebiao064 <hebiaobuaa@gmail.com>
This commit is contained in:
Qiaolin Yu
2025-09-21 19:36:08 -07:00
committed by GitHub
parent 86527a4799
commit e2ac7888b8
12 changed files with 117 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/sampling/sampling_batch_info.py
View File

@@ -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