From 082aac8fce2757a1bbe1049c7a44168ac8fb2030 Mon Sep 17 00:00:00 2001 From: milesial Date: Wed, 19 Aug 2026 15:29:22 -0700 Subject: [PATCH] [Bugfix] Fix min-new-token EOS handling (#31378) Signed-off-by: Alexandre Milesi --- .../srt/sampling/penaltylib/min_new_tokens.py | 9 ++++++--- test/registered/unit/sampling/test_penaltylib.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py index c6174de23..c77cb7851 100644 --- a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py +++ b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py @@ -25,9 +25,11 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer): padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence( sequences=[ torch.tensor( - data=( - list( + data=[ + token_id + for token_id in ( (req.sampling_params.stop_token_ids or set()) + | (req.eos_token_ids or set()) | (req.tokenizer.additional_stop_token_ids or set()) | ( {req.tokenizer.eos_token_id} @@ -35,7 +37,8 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer): else set() ) ) - ), + if token_id is not None + ], dtype=torch.int64, device=self.orchestrator.device, ) diff --git a/test/registered/unit/sampling/test_penaltylib.py b/test/registered/unit/sampling/test_penaltylib.py index d29b04b0f..39576702e 100644 --- a/test/registered/unit/sampling/test_penaltylib.py +++ b/test/registered/unit/sampling/test_penaltylib.py @@ -36,6 +36,7 @@ def _make_req(freq=0.0, presence=0.0, min_tokens=0, stop_ids=None, eos_id=2): req.sampling_params.presence_penalty = presence req.sampling_params.min_new_tokens = min_tokens req.sampling_params.stop_token_ids = stop_ids + req.eos_token_ids = None req.tokenizer.additional_stop_token_ids = None req.tokenizer.eos_token_id = eos_id return req @@ -347,6 +348,20 @@ class TestBatchedMinNewTokensPenalizer(CustomTestCase): # Non-stop tokens should be fine self.assertEqual(logits[0, 0].item(), 0.0) + def test_blocks_model_eos_without_tokenizer_eos(self): + """Model-config EOS remains blocked when tokenizer EOS metadata is missing.""" + req = _make_req(min_tokens=3, eos_id=None) + req.eos_token_ids = {6} + batch = _make_batch([req]) + orch = BatchedPenalizerOrchestrator( + VOCAB_SIZE, batch, {BatchedMinNewTokensPenalizer} + ) + pen = orch.penalizers[BatchedMinNewTokensPenalizer] + + logits = torch.zeros(1, VOCAB_SIZE) + pen.apply(logits) + self.assertTrue(torch.isneginf(logits[0, 6])) + def test_filter_keeps_subset(self): """Test that filter keeps the second request (min_tokens=5) and drops the first.""" orch, pen = self._setup([(3, None, 2), (5, None, 2)])