From a9f02b0fa48be1642ad27960d6ee244ed2e8c056 Mon Sep 17 00:00:00 2001 From: Javen <68119726+zhujianwei-ops@users.noreply.github.com> Date: Tue, 22 Sep 2026 12:34:15 +0800 Subject: [PATCH] [NPU] Fix xgrammar apply_vocab_mask device dispatch to use torch.ops.npu (#36120) Co-authored-by: Even Zhou Co-authored-by: sglang-npu-bot --- python/sglang/srt/constrained/xgrammar_backend.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index 86293e2d7..167bee6fa 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -261,11 +261,15 @@ class XGrammarGrammarBackend(BaseGrammarBackend): @staticmethod def apply_vocab_mask(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": + import sgl_kernel_npu # noqa: F401 + + torch.ops.npu.apply_token_bitmask(logits, vocab_mask) else: raise RuntimeError(f"Unsupported device: {logits.device.type}")