[Fix] eagle/eagle3 speculative decoding conflicts with xgrammar in NPU (#20989)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/srt/constrained/torch_ops/bitmask_ops.py

@@ -0,0 +1,33 @@
+# Copyright 2023-2024 SGLang Team
+# 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 torch
+
+
+def apply_token_bitmask_inplace_torch(
+    logits: torch.Tensor,
+    bitmask: torch.Tensor,
+) -> None:
+    """Backend-agnostic torch fallback for packed-bitmask application.
+
+    This path is currently used as a fallback on NPU in xgrammar backend.
+    """
+    vocab_size = logits.shape[-1]
+    bitmask_cpu = bitmask.detach().cpu()
+    token_ids = torch.arange(vocab_size, device="cpu", dtype=torch.int32)
+    word_idx = token_ids // 32
+    bit_idx = token_ids % 32
+    words = bitmask_cpu[:, word_idx].to(torch.int32)
+    allowed = ((words >> bit_idx) & 1).to(torch.bool)
+    allowed = allowed.to(logits.device, non_blocking=True)
+    logits.masked_fill_(~allowed, float("-inf"))

python/sglang/srt/constrained/xgrammar_backend.py

@@ -35,6 +35,9 @@ from sglang.srt.constrained.base_grammar_backend import (
     GrammarStats,
     InvalidGrammarObject,
 )
+from sglang.srt.constrained.torch_ops.bitmask_ops import (
+    apply_token_bitmask_inplace_torch,
+)
 from sglang.srt.constrained.utils import is_legacy_structural_tag
 from sglang.srt.utils import is_hip
 
@@ -105,11 +108,13 @@ class XGrammarGrammar(BaseGrammarObject):
         return vocab_mask.to(device, non_blocking=True)
 
     def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
-        if logits.device.type in {"cuda", "npu", "xpu", "musa"}:
+        if logits.device.type in {"cuda", "xpu", "musa"}:
             if _is_hip:
                 apply_token_bitmask_inplace_cuda(logits, vocab_mask)
             else:
                 apply_token_bitmask_inplace_triton(logits, vocab_mask)
+        elif logits.device.type == "npu":
+            apply_token_bitmask_inplace_torch(logits, vocab_mask)
         else:
             raise RuntimeError(f"Unsupported device: {logits.device.type}")
 

python/sglang/test/ascend/test_ascend_vocab_mask.py

@@ -0,0 +1,81 @@
+import math
+
+import pytest
+import torch
+
+from sglang.srt.constrained import xgrammar_backend as xb
+
+
+def _pack_mask(allowed_ids, vocab_size, batch_size=1):
+    nwords = math.ceil(vocab_size / 32)
+    m = torch.zeros((batch_size, nwords), dtype=torch.int32)
+    for b in range(batch_size):
+        for tid in allowed_ids[b]:
+            m[b, tid // 32] |= 1 << (tid % 32)
+    return m
+
+
+def _apply_ref_cpu(logits, vocab_mask):
+    vocab_size = logits.shape[-1]
+    token_ids = torch.arange(vocab_size, device="cpu", dtype=torch.int64)
+    word_idx = token_ids // 32
+    bit_idx = (token_ids % 32).to(torch.int32)
+    words = vocab_mask.cpu()[:, word_idx].to(torch.int32)
+    allowed = ((words >> bit_idx) & 1).bool().to(logits.device)
+    out = logits.clone()
+    out.masked_fill_(~allowed, float("-inf"))
+    return out
+
+
+@pytest.mark.skipif(
+    not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
+)
+def test_mask_blocks_disallowed_token_on_npu():
+    device = "npu:0"
+    vocab_size = 64
+
+    logits = torch.zeros((1, vocab_size), device=device, dtype=torch.float32)
+    logits[0, 16] = 22.125
+    logits[0, 5] = 10.0
+
+    allowed = [[5, 6, 7, 8]]
+    vocab_mask = _pack_mask(allowed, vocab_size).to(device=device, dtype=torch.int32)
+
+    g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
+    out = logits.clone()
+    g.apply_vocab_mask(out, vocab_mask)
+
+    assert not torch.isfinite(out[0, 16])
+    assert int(torch.argmax(out[0]).item()) != 16
+
+
+@pytest.mark.skipif(
+    not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU required"
+)
+def test_npu_path_matches_reference_random():
+    device = "npu:0"
+    B, V = 4, 257
+    torch.manual_seed(0)
+
+    logits = torch.randn(B, V, device=device, dtype=torch.float32)
+
+    allowed = []
+    for _ in range(B):
+        ids = torch.randperm(V)[: V // 4].tolist()
+        allowed.append(ids)
+    vocab_mask = _pack_mask(allowed, V, B).to(device=device, dtype=torch.int32)
+
+    g = xb.XGrammarGrammar.__new__(xb.XGrammarGrammar)
+    out_npu = logits.clone()
+    g.apply_vocab_mask(out_npu, vocab_mask)
+
+    out_ref = _apply_ref_cpu(logits, vocab_mask)
+
+    assert torch.equal(torch.isfinite(out_npu), torch.isfinite(out_ref))
+    diff = (
+        torch.nan_to_num(out_npu - out_ref, nan=0.0, posinf=0.0, neginf=0.0)
+        .abs()
+        .max()
+        .item()
+    )
+    assert diff < 1e-5