From 808685080b4eec905c17ea873605e0e0ec97712b Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 8 Sep 2026 05:03:54 -0700 Subject: [PATCH 1/4] [https://nvbugs/6625851][fix] Fail guided-decoding requests reaching a dead-end grammar state A grammar state with no valid next token fills an all-zero bitmask row. The apply kernel then masks the whole logits row to -inf, softmax returns NaN for it, and the sampler's device-side NaN assert fires. Because that assert is a global device assert, a single request's dead-end grammar state hard-kills every rank of the deployment (MPI_Abort). Detect the condition on the host, right after fill_next_token_bitmask, and fail that one request through the existing guided-decoding error path. This matches how vLLM and SGLang handle a stuck grammar matcher: both terminate only the affected request (RequestStatus.FINISHED_ERROR and FINISH_ABORT respectively) and never take down the engine. Detecting it at fill time rather than at apply time keeps the check off the GPU critical path and, unlike skipping the mask for that row, avoids emitting a token that violates the grammar. Only the bits below vocab_size_padded are counted: the trailing bits of a partial last word are never read by the apply kernel and are not guaranteed to be cleared by the backends. For draft requests the dead end terminates drafting instead of failing the request, consistent with the existing unacceptable-draft-token path. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/guided_decoder.py | 43 ++++++++++++ .../misc/test_guided_decoder_bitmask.py | 68 +++++++++++++++++++ 2 files changed, 111 insertions(+) create mode 100644 tests/unittest/_torch/misc/test_guided_decoder_bitmask.py diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index b8d06c1cb97a..420df343544f 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -18,6 +18,25 @@ from .scheduler import ScheduledRequests +def row_has_valid_token(row: torch.Tensor, vocab_size_padded: int) -> bool: + """Whether a filled next-token bitmask row allows at least one token. + + A grammar state with no valid continuation produces an all-zero row, which + would mask the whole logits row to -inf and make softmax return NaN for it. + + Only the bits below `vocab_size_padded` are counted: the trailing bits of a + partial last word are never read by the apply kernel and are not guaranteed + to be cleared by the backends, so counting them could report a dead-end row + as valid. + """ + num_words, num_tail_bits = divmod(vocab_size_padded, 32) + if torch.any(row[:num_words]): + return True + if num_tail_bits == 0: + return False + return bool(row[num_words].item() & ((1 << num_tail_bits) - 1)) + + @dataclass(slots=True) class GuidedRequest: """A snapshot of an LlmRequest that contains relevant fields for guided decoding. @@ -206,6 +225,11 @@ def __init__(self, def bitmask_size(self) -> int: return math.ceil(self.vocab_size_padded / 32) + def _has_valid_token(self, index: int) -> bool: + """Whether the bitmask row just filled at `index` allows any token.""" + return row_has_valid_token(self.bitmask_host[index], + self.vocab_size_padded) + def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]: """Build the bitmask for requests with guided decoding enabled. @@ -257,6 +281,20 @@ def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]: self.num_advanced_tokens[slot] += 1 if not matcher.is_terminated(): matcher.fill_next_token_bitmask(self.bitmask_host, offset) + if not self._has_valid_token(offset): + if req.is_draft: + self.is_draft_terminated[slot] = True + logger.debug( + f"Draft request {req.request_id} at slot {slot} reached a grammar state with no valid token." + ) + continue + # The request is about to be terminated, so exclude it + # from the rollback pass: the matcher advance made + # above is never verified against accepted tokens. + self.num_advanced_tokens[slot] = 0 + raise ValueError( + f"Request {req.request_id} at slot {slot} reached a grammar state with no valid token." + ) self.token_mask_host[offset] = 1 self.num_guided_tokens[slot] += 1 # Process draft tokens. Bound by the layout's draft length: @@ -273,6 +311,11 @@ def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]: break matcher.fill_next_token_bitmask(self.bitmask_host, offset + i) + if not self._has_valid_token(offset + i): + # Stop guiding here rather than failing the + # request: the remaining draft positions are left + # unconstrained and get verified as usual. + break self.token_mask_host[offset + i] = 1 self.num_guided_tokens[slot] += 1 diff --git a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py new file mode 100644 index 000000000000..928d21b37473 --- /dev/null +++ b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.guided_decoder import row_has_valid_token + + +def _empty_row(vocab_size_padded: int) -> torch.Tensor: + return torch.zeros(math.ceil(vocab_size_padded / 32), dtype=torch.int32) + + +# 128000 is word-aligned; 128001 and 128031 leave a partial last word. +@pytest.mark.parametrize("vocab_size_padded", [128000, 128001, 128031]) +def test_dead_end_row_has_no_valid_token(vocab_size_padded: int): + """A grammar state with no valid continuation must be reported as dead.""" + assert not row_has_valid_token(_empty_row(vocab_size_padded), vocab_size_padded) + + +@pytest.mark.parametrize("vocab_size_padded", [128000, 128001, 128031]) +@pytest.mark.parametrize("token_id_from_end", [0, 1, 32, 12345]) +def test_single_valid_token_is_detected(vocab_size_padded: int, token_id_from_end: int): + """A single set bit anywhere below vocab_size_padded keeps the row alive.""" + token_id = vocab_size_padded - 1 - token_id_from_end + row = _empty_row(vocab_size_padded) + row[token_id // 32] |= 1 << (token_id % 32) + assert row_has_valid_token(row, vocab_size_padded) + + +@pytest.mark.parametrize("vocab_size_padded", [128001, 128031]) +def test_trailing_padding_bits_do_not_count(vocab_size_padded: int): + """Padding bits above vocab_size_padded must not mark a dead row valid. + + The apply kernel never reads them and the backends do not guarantee they + are cleared, so counting them would let a fully masked row reach the + sampler and produce a NaN logits row. + """ + row = _empty_row(vocab_size_padded) + num_words, num_tail_bits = divmod(vocab_size_padded, 32) + # Set every bit of the last word that lies at or above vocab_size_padded. + row[num_words] = torch.tensor(-1 << num_tail_bits, dtype=torch.int32) + assert not row_has_valid_token(row, vocab_size_padded) + + # The highest in-range bit of that same partial word must still count. + row[num_words] |= 1 << (num_tail_bits - 1) + assert row_has_valid_token(row, vocab_size_padded) + + +def test_sign_bit_counts_as_valid_token(): + """Bit 31 of a word makes the int32 negative; it is still a valid token.""" + vocab_size_padded = 128000 + row = _empty_row(vocab_size_padded) + row[0] = torch.tensor(-2147483648, dtype=torch.int32) # only bit 31 set + assert row_has_valid_token(row, vocab_size_padded) From 07972f7456e7c8d7216c3ab6d1a7ff3fc90fe1e3 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 8 Sep 2026 05:23:00 -0700 Subject: [PATCH 2/4] [https://nvbugs/6625851][fix] Roll back the draft matcher advance on a dead-end row Addresses review feedback on the dead-end handling: - A draft request hitting a dead end skipped the num_advanced_draft_tokens accumulation, so _rollback_draft_tokens did not undo the accept_token(new_token) that had already advanced the matcher, leaving the target model one token ahead. Unlike the unacceptable-token path, the matcher does advance here, so record it. - Add _build coverage for the three changed paths: a regular request failing on an immediate dead end (and clearing its rollback accounting), a draft request terminating drafting and rolling the advance back, and a dead end at a later draft position leaving the remaining positions unguided without failing the request. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/guided_decoder.py | 5 + .../misc/test_guided_decoder_bitmask.py | 150 +++++++++++++++++- 2 files changed, 154 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index 420df343544f..683514bcd083 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -287,6 +287,11 @@ def _build(self, requests: GuidedRequests) -> List[Tuple[int, str]]: logger.debug( f"Draft request {req.request_id} at slot {slot} reached a grammar state with no valid token." ) + # Unlike the unacceptable-token path above, the + # matcher did advance past new_token here, so the + # drafting loop must still roll that advance back. + self.num_advanced_draft_tokens[ + slot] += self.num_advanced_tokens[slot] continue # The request is about to be terminated, so exclude it # from the rollback pass: the matcher advance made diff --git a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py index 928d21b37473..9d1172ef01f8 100644 --- a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py +++ b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py @@ -13,11 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +from typing import List import pytest import torch -from tensorrt_llm._torch.pyexecutor.guided_decoder import row_has_valid_token +from tensorrt_llm._torch.pyexecutor import guided_decoder as guided_decoder_module +from tensorrt_llm._torch.pyexecutor.guided_decoder import ( + GuidedDecoder, + GuidedRequest, + GuidedRequests, + row_has_valid_token, +) +from tensorrt_llm.llmapi.llm_args import GuidedDecodingConfig def _empty_row(vocab_size_padded: int) -> torch.Tensor: @@ -66,3 +74,143 @@ def test_sign_bit_counts_as_valid_token(): row = _empty_row(vocab_size_padded) row[0] = torch.tensor(-2147483648, dtype=torch.int32) # only bit 31 set assert row_has_valid_token(row, vocab_size_padded) + + +# --- _build dead-end handling ------------------------------------------------- + +_VOCAB_SIZE = 128000 +_SLOT = 0 +# Any non-None value works: the tests pre-seed the matcher, so the grammar +# matcher factory is never asked to compile these params. +_GUIDED_PARAMS = object() + + +class _ScriptedMatcher: + """Grammar matcher whose per-position bitmask rows are scripted by the test. + + `rows[i]` says whether the i-th `fill_next_token_bitmask` call should + produce a row with a valid token; a False entry is a dead-end state. + """ + + def __init__(self, rows: List[bool]): + self._rows = rows + self.accepted: List[int] = [] + self.num_rolled_back = 0 + self._num_fills = 0 + + def accept_token(self, token_id: int) -> bool: + self.accepted.append(token_id) + return True + + def rollback(self, num_tokens: int) -> None: + self.num_rolled_back += num_tokens + del self.accepted[len(self.accepted) - num_tokens :] + + def fill_next_token_bitmask(self, bitmask: torch.Tensor, index: int) -> None: + has_valid_token = self._rows[self._num_fills] + self._num_fills += 1 + bitmask[index].zero_() + if has_valid_token: + bitmask[index][0] = 1 + + def is_terminated(self) -> bool: + return False + + +def _make_decoder(monkeypatch, max_num_draft_tokens: int) -> GuidedDecoder: + # The factory needs a real tokenizer and a compiled grammar; stub it out so + # __init__ runs unchanged while the tests drive a scripted matcher instead. + monkeypatch.setattr( + guided_decoder_module, "XGrammarMatcherFactory", lambda *args, **kwargs: None + ) + return GuidedDecoder( + GuidedDecodingConfig(), + max_num_sequences=4, + vocab_size_padded=_VOCAB_SIZE, + max_num_draft_tokens=max_num_draft_tokens, + ) + + +def _generation_request(*, is_draft: bool, draft_tokens: List[int]) -> GuidedRequest: + return GuidedRequest( + guided_decoding_params=_GUIDED_PARAMS, + request_id=7, + seq_slot=_SLOT, + is_generation_in_progress_state=True, + new_token=11, + is_draft=is_draft, + draft_tokens=draft_tokens, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_build_fails_request_on_dead_end_row(monkeypatch): + """A dead end fails that one request instead of reaching the sampler.""" + decoder = _make_decoder(monkeypatch, max_num_draft_tokens=0) + decoder.grammar_matchers[_SLOT] = _ScriptedMatcher([False]) + requests = GuidedRequests( + [_generation_request(is_draft=False, draft_tokens=[])], + num_contexts=0, + num_generations=1, + max_num_draft_tokens=0, + ) + + failed_requests = decoder._build(requests) + + assert [req_id for req_id, _ in failed_requests] == [7] + # The row must stay unguided so the apply kernel skips it entirely. + assert decoder.token_mask_host[0].item() == 0 + # Cleared so the failed request is skipped by _rollback_rejected_tokens, + # which would otherwise raise on a negative rollback count. + assert decoder.num_advanced_tokens[_SLOT] == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_build_draft_dead_end_is_rolled_back(monkeypatch): + """A draft dead end terminates drafting and stays rollback-accurate.""" + decoder = _make_decoder(monkeypatch, max_num_draft_tokens=4) + matcher = _ScriptedMatcher([False]) + decoder.grammar_matchers[_SLOT] = matcher + requests = GuidedRequests( + [_generation_request(is_draft=True, draft_tokens=[])], + num_contexts=0, + num_generations=1, + max_num_draft_tokens=4, + ) + + failed_requests = decoder._build(requests) + + assert failed_requests == [] + assert decoder.is_draft_terminated[_SLOT] + assert decoder.token_mask_host[0].item() == 0 + # The matcher accepted new_token before hitting the dead end, so the + # drafting loop must roll that advance back; otherwise the target model + # resumes from a matcher that is one token ahead. + assert decoder.num_advanced_draft_tokens[_SLOT] == 1 + decoder._rollback_draft_tokens(requests) + assert matcher.num_rolled_back == 1 + assert matcher.accepted == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_build_draft_position_dead_end_stops_guiding(monkeypatch): + """A dead end at a draft position leaves later positions unguided.""" + decoder = _make_decoder(monkeypatch, max_num_draft_tokens=2) + decoder.grammar_matchers[_SLOT] = _ScriptedMatcher([True, False]) + requests = GuidedRequests( + [_generation_request(is_draft=False, draft_tokens=[21, 22])], + num_contexts=0, + num_generations=1, + max_num_draft_tokens=2, + ) + + failed_requests = decoder._build(requests) + + # The request itself is still viable: only drafting stops early. + assert failed_requests == [] + assert decoder.token_mask_host[0].item() == 1 + assert decoder.token_mask_host[1].item() == 0 + assert decoder.token_mask_host[2].item() == 0 + assert decoder.num_guided_tokens[_SLOT] == 1 + # new_token plus the one accepted draft token. + assert decoder.num_advanced_tokens[_SLOT] == 2 From f4ed47ede5beefed6777ded850556288b3e8b351 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 8 Sep 2026 05:31:38 -0700 Subject: [PATCH 3/4] [https://nvbugs/6625851][test] Make the scripted matcher tolerant past its script Without this, removing the dead-end check makes the draft-position test fail with an IndexError from the test helper rather than on the assertion it is meant to prove. Rows past the end of the script now produce a valid bitmask row. Signed-off-by: ZhaoyangWang --- tests/unittest/_torch/misc/test_guided_decoder_bitmask.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py index 9d1172ef01f8..9cc3e6a1c090 100644 --- a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py +++ b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py @@ -89,7 +89,10 @@ class _ScriptedMatcher: """Grammar matcher whose per-position bitmask rows are scripted by the test. `rows[i]` says whether the i-th `fill_next_token_bitmask` call should - produce a row with a valid token; a False entry is a dead-end state. + produce a row with a valid token; a False entry is a dead-end state. Calls + past the end of the script produce a valid row, so that a regression which + fails to stop at a dead end is caught by an assertion rather than by this + helper running out of scripted rows. """ def __init__(self, rows: List[bool]): @@ -107,7 +110,7 @@ def rollback(self, num_tokens: int) -> None: del self.accepted[len(self.accepted) - num_tokens :] def fill_next_token_bitmask(self, bitmask: torch.Tensor, index: int) -> None: - has_valid_token = self._rows[self._num_fills] + has_valid_token = self._rows[self._num_fills] if self._num_fills < len(self._rows) else True self._num_fills += 1 bitmask[index].zero_() if has_valid_token: From d8ba420e08d41a2469fa61b5cdfe37b72c44cdeb Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 8 Sep 2026 05:37:36 -0700 Subject: [PATCH 4/4] [https://nvbugs/6625851][test] Cover per-request isolation in a mixed batch The dead-end tests used single-request batches, so they could not tell a per-request failure apart from one that also disrupts the rest of the batch - which is the property this fix exists for. Add a two-request batch and assert the healthy request keeps its guided row. Signed-off-by: ZhaoyangWang --- .../misc/test_guided_decoder_bitmask.py | 42 +++++++++++++++++-- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py index 9cc3e6a1c090..9abaeb0c9da8 100644 --- a/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py +++ b/tests/unittest/_torch/misc/test_guided_decoder_bitmask.py @@ -134,11 +134,17 @@ def _make_decoder(monkeypatch, max_num_draft_tokens: int) -> GuidedDecoder: ) -def _generation_request(*, is_draft: bool, draft_tokens: List[int]) -> GuidedRequest: +def _generation_request( + *, + is_draft: bool, + draft_tokens: List[int], + request_id: int = 7, + seq_slot: int = _SLOT, +) -> GuidedRequest: return GuidedRequest( guided_decoding_params=_GUIDED_PARAMS, - request_id=7, - seq_slot=_SLOT, + request_id=request_id, + seq_slot=seq_slot, is_generation_in_progress_state=True, new_token=11, is_draft=is_draft, @@ -217,3 +223,33 @@ def test_build_draft_position_dead_end_stops_guiding(monkeypatch): assert decoder.num_guided_tokens[_SLOT] == 1 # new_token plus the one accepted draft token. assert decoder.num_advanced_tokens[_SLOT] == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_build_dead_end_isolates_the_failing_request(monkeypatch): + """Only the dead-end request fails; the rest of the batch keeps decoding. + + This is the property the fix exists for: before it, a single dead-end row + took down the whole deployment via the sampler's global NaN assert. + """ + decoder = _make_decoder(monkeypatch, max_num_draft_tokens=0) + decoder.grammar_matchers[0] = _ScriptedMatcher([False]) + decoder.grammar_matchers[1] = _ScriptedMatcher([True]) + requests = GuidedRequests( + [ + _generation_request(is_draft=False, draft_tokens=[], request_id=7, seq_slot=0), + _generation_request(is_draft=False, draft_tokens=[], request_id=8, seq_slot=1), + ], + num_contexts=0, + num_generations=2, + max_num_draft_tokens=0, + ) + + failed_requests = decoder._build(requests) + + assert [req_id for req_id, _ in failed_requests] == [7] + # The healthy request keeps its guided row and stays constrained. + assert decoder.token_mask_host[0].item() == 0 + assert decoder.token_mask_host[1].item() == 1 + assert decoder.num_guided_tokens[1] == 1 + assert decoder.num_advanced_tokens[1] == 1