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
This commit is contained in:
laoyao0822
2026-06-27 20:37:05 +08:00
parent 132561a151
commit 2de822661c
2 changed files with 51 additions and 2 deletions

View File

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

View File

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