From 65bc839a5f3f63962d63857646708509b22775f4 Mon Sep 17 00:00:00 2001 From: yuefeng Wu <33725817+ChefWu551@users.noreply.github.com> Date: Thu, 16 Apr 2026 14:34:23 +0800 Subject: [PATCH] [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> --- .../srt/constrained/torch_ops/bitmask_ops.py | 33 ++++++++ .../srt/constrained/xgrammar_backend.py | 7 +- .../test/ascend/test_ascend_vocab_mask.py | 81 +++++++++++++++++++ 3 files changed, 120 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/constrained/torch_ops/bitmask_ops.py create mode 100644 python/sglang/test/ascend/test_ascend_vocab_mask.py diff --git a/python/sglang/srt/constrained/torch_ops/bitmask_ops.py b/python/sglang/srt/constrained/torch_ops/bitmask_ops.py new file mode 100644 index 000000000..25f07e5ca --- /dev/null +++ b/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")) diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index c04ad155e..542e920ce 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/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}") diff --git a/python/sglang/test/ascend/test_ascend_vocab_mask.py b/python/sglang/test/ascend/test_ascend_vocab_mask.py new file mode 100644 index 000000000..8128148ff --- /dev/null +++ b/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