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