Fix min_p sampling

flashinfer backed won't return success states for min_p.
sgl-project · Jan 29, 2025 · 5bbc421 · 5bbc421
1 parent 9ae1db0
commit 5bbc421
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py
@@ -51,7 +51,7 @@ def forward(
             if sampling_info.need_min_p_sampling:
                 probs = top_k_renorm_prob(probs, sampling_info.top_ks)
                 probs = top_p_renorm_prob(probs, sampling_info.top_ps)
-                batch_next_token_ids, success = min_p_sampling_from_probs(
+                batch_next_token_ids = min_p_sampling_from_probs(
                     probs, uniform_samples, sampling_info.min_ps
                 )
             else:
@@ -63,9 +63,9 @@ def forward(
                     filter_apply_order="joint",
                 )
 
-            if not torch.all(success):
-                logger.warning("Detected errors during sampling!")
-                batch_next_token_ids = torch.zeros_like(batch_next_token_ids)
+                if not torch.all(success):
+                    logger.warning("Detected errors during sampling!")
+                    batch_next_token_ids = torch.zeros_like(batch_next_token_ids)
         elif global_server_args_dict["sampling_backend"] == "pytorch":
             # Here we provide a slower fallback implementation.
             batch_next_token_ids = top_k_top_p_min_p_sampling_from_probs_torch(