From 2ab7d77e58f89870fa6c6fc9004f6390e4cc2596 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 8 Sep 2026 02:45:13 -0700 Subject: [PATCH] [https://nvbugs/6625851][fix] Skip the bitmask apply for zero-valid-token grammar rows A grammar state with no valid token (an xgrammar dead-end, or a mask corrupted upstream) leaves an all-zero bitmask row. logits_bitmask then writes -inf to every column of that logits row, softmax of an all--inf row is NaN, and the sampler's async NaN assert fires a device-side assert that the executor escalates to a peer-kill of every rank. Gate each row's token_mask entry on the row having at least one bit set, so such a row keeps its finite logits instead of poisoning sampling. The check spans the full bitmask row rather than the local shard, because an all-masked shard is legitimate when logits are vocab-sharded, and covers only the bits below vocab_size_padded since the trailing bits of a partial last mask word are unspecified. The affected request is not silently let through: it samples one unconstrained token, which the matcher rejects on the next build, so it fails through the existing guided-decoding error path while the rest of the deployment keeps serving. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/pyexecutor/guided_decoder.py | 29 ++++++++++--- .../test_logits_bitmask_op.py | 41 +++++++++++++++++++ 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py index b8d06c1cb97a..20cfcad018ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/guided_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/guided_decoder.py @@ -329,11 +329,30 @@ def _apply_bitmask(self, d2t_end = d2t_start + vocab_size_padded // tp_size d2t = d2t[d2t_start:d2t_end] - torch.ops.trtllm.logits_bitmask( - logits[:num_bitmask_tokens], - self.bitmask[:num_bitmask_tokens, bitmask_start:bitmask_end], - token_mask=self.token_mask[:num_bitmask_tokens], - d2t=d2t) + # A grammar state with zero valid tokens would mask the whole logits row + # to -inf; softmax then yields NaN for that row, which trips the + # sampler's async NaN assert and hard-kills every rank. Skip the apply + # for such a row instead: it samples one unconstrained token, which the + # matcher rejects at the next build, failing that single request through + # the regular guided-decoding error path. + # The check must span the full bitmask row rather than this rank's + # slice, since an all-masked local shard is legitimate when the logits + # are sharded in the vocabulary dimension. Only the bits below + # vocab_size_padded count: the trailing bits of a partial last word are + # never read by the kernel, and the backends do set them (xgrammar + # writes a full word), so reading them would mask the row as valid. + bitmask = self.bitmask[:num_bitmask_tokens] + num_words, num_tail_bits = divmod(self.vocab_size_padded, 32) + has_valid_token = bitmask[:, :num_words].any(dim=1) + if num_tail_bits > 0: + has_valid_token |= (bitmask[:, num_words] + & ((1 << num_tail_bits) - 1)) != 0 + token_mask = self.token_mask[:num_bitmask_tokens] * has_valid_token + + torch.ops.trtllm.logits_bitmask(logits[:num_bitmask_tokens], + bitmask[:, bitmask_start:bitmask_end], + token_mask=token_mask, + d2t=d2t) @nvtx_range("GuidedDecoder.add_batch") def add_batch(self, diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_logits_bitmask_op.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_logits_bitmask_op.py index 027ac882c4fc..78937bfd3aad 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_logits_bitmask_op.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_logits_bitmask_op.py @@ -3,6 +3,7 @@ import xgrammar import tensorrt_llm # noqa +from tensorrt_llm._torch.pyexecutor.guided_decoder import GuidedDecoder @pytest.mark.parametrize("batch_size", [1, 64]) @@ -59,3 +60,43 @@ def test_logits_bitmask_with_d2t(batch_size: int, vocab_size: int, stride: int, # Call logits bitmask op and evaluate torch.ops.trtllm.logits_bitmask(logits, bitmask, token_mask=token_mask, d2t=d2t) torch.testing.assert_close(logits, logits_reference) + + +# _apply_bitmask requires the vocab size to be divisible by the shard count, +# so 128001 (which exercises a partial last mask word) is only valid at tp 1. +@pytest.mark.parametrize("vocab_size, tp_size", [(128000, 1), (128000, 4), (128001, 1)]) +def test_apply_bitmask_skips_zero_valid_token_row(vocab_size: int, tp_size: int): + """A grammar row with no valid token must not be masked to all -inf. + + Masking the whole row makes softmax produce NaN, which trips the sampler's + async NaN assert and takes down every rank (https://nvbugs/6625851). + """ + batch_size, empty_row = 8, 3 + rank, local_vocab_size = tp_size - 1, vocab_size // tp_size + # Bypass __init__ to avoid needing a tokenizer / compiled grammar: only the + # attributes read by _apply_bitmask are required. + guided_decoder = object.__new__(GuidedDecoder) + guided_decoder.vocab_size_padded = vocab_size + guided_decoder.rank = rank + + bool_mask = torch.randint(0, 2, size=(batch_size, vocab_size), dtype=torch.bool, device="cuda") + # Keep one valid token in every shard of every other row, so that a fully + # masked-out shard can only come from the empty row. The stride is the shard + # width, not tp_size: at tp_size 1 the latter marks every token valid, so + # nothing is left to reject and the reference masking becomes a no-op. + bool_mask[:, ::local_vocab_size] = True + bool_mask[empty_row] = False + guided_decoder.bitmask = xgrammar.testing.bool_mask_to_bitmask(bool_mask) + guided_decoder.token_mask = torch.ones(batch_size, dtype=torch.int32, device="cuda") + + logits = torch.randn(batch_size, local_vocab_size, dtype=torch.float32, device="cuda") + logits_reference = logits.clone() + shard = slice(rank * local_vocab_size, (rank + 1) * local_vocab_size) + reject = ~bool_mask[:, shard] + reject[empty_row] = False # the empty row must be left untouched + logits_reference.masked_fill_(reject, -float("inf")) + + guided_decoder._apply_bitmask(None, logits, num_bitmask_tokens=batch_size) + + torch.testing.assert_close(logits, logits_reference) + assert not torch.isnan(torch.softmax(logits, dim=-1)).any()