From 2de822661ce5e882d8c7a84da17e71487ceaf655 Mon Sep 17 00:00:00 2001
From: laoyao0822 <laoyao0822@qq.com>
Date: Sat, 27 Jun 2026 20:37:05 +0800
Subject: [PATCH] Avoid decode-step syncs in min-new-token penalties

The old min_new_tokens penalizer updated logits through boolean-mask indexing. That indexing is data-dependent and can force synchronization on the decode hot path.

Use an elementwise torch.where followed by inplace add so the operation stays tensorized and avoids the mask-index update path.

Constraint: Keep the numeric behavior for active rows equivalent without multiplying by zero, which would turn -inf penalties into NaN.

Rejected: Expanding the mask and assigning through logits[mask] | this is the synchronization pattern being removed.

Confidence: high

Scope-risk: narrow

Directive: Do not reintroduce boolean-mask writes in decode-step penalty paths without profiling the synchronization behavior.

Tested: RED/GREEN local pytest test/registered/unit/sampling/test_min_new_tokens_penalizer.py

Tested: RED/GREEN remote pytest in cjy-glm5-new for min_new_tokens penalizer together with related regression tests

Tested: git diff --check; py_compile for min_new_tokens.py

Not-tested: Full decode throughput benchmark
---
 .../srt/sampling/penaltylib/min_new_tokens.py |  4 +-
 .../sampling/test_min_new_tokens_penalizer.py | 49 +++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 test/registered/unit/sampling/test_min_new_tokens_penalizer.py

diff --git a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py
index 08f76e1f1..ea27b05cb 100644
--- a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py
+++ b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py
@@ -67,8 +67,8 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
         self.len_output_tokens += 1
 
     def _apply(self, logits: torch.Tensor):
-        mask = (self.len_output_tokens < self.min_new_tokens).expand_as(logits)
-        logits[mask] += self.stop_token_penalties[mask]
+        mask = self.len_output_tokens < self.min_new_tokens
+        logits.add_(torch.where(mask, self.stop_token_penalties, 0.0))
 
     def _filter(self, keep_indices: torch.Tensor):
         self.min_new_tokens = self.min_new_tokens[keep_indices]
diff --git a/test/registered/unit/sampling/test_min_new_tokens_penalizer.py b/test/registered/unit/sampling/test_min_new_tokens_penalizer.py
new file mode 100644
index 000000000..087dbb083
--- /dev/null
+++ b/test/registered/unit/sampling/test_min_new_tokens_penalizer.py
@@ -0,0 +1,49 @@
+import unittest
+from unittest.mock import patch
+
+from sglang.srt.sampling.penaltylib.min_new_tokens import BatchedMinNewTokensPenalizer
+
+
+class _Mask:
+    def expand_as(self, logits):  # old implementation calls this before boolean indexing
+        raise AssertionError("min_new_tokens penalty must not use boolean-mask indexing")
+
+
+class _Comparable:
+    def __lt__(self, other):
+        return _Mask()
+
+
+class _FakeLogits:
+    def __init__(self):
+        self.added = None
+
+    def __getitem__(self, key):
+        raise AssertionError("min_new_tokens penalty must not index logits with a mask")
+
+    def add_(self, value):
+        self.added = value
+        return self
+
+
+class TestMinNewTokensPenalizer(unittest.TestCase):
+    def test_apply_uses_elementwise_where_instead_of_boolean_indexing(self):
+        penalizer = BatchedMinNewTokensPenalizer.__new__(BatchedMinNewTokensPenalizer)
+        penalizer.len_output_tokens = _Comparable()
+        penalizer.min_new_tokens = object()
+        penalizer.stop_token_penalties = object()
+        logits = _FakeLogits()
+        sentinel = object()
+
+        with patch(
+            "sglang.srt.sampling.penaltylib.min_new_tokens.torch.where",
+            return_value=sentinel,
+        ) as where:
+            penalizer._apply(logits)
+
+        where.assert_called_once()
+        self.assertIs(logits.added, sentinel)
+
+
+if __name__ == "__main__":
+    unittest.main()