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:
@@ -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]
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user