[https://nvbugs/6625851][fix] Gate each row's token_mask on the full bitmask row having at least one bit…

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,

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