minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU] Fix xgrammar apply_vocab_mask device dispatch to use torch.ops.npu (#36120)

Browse Source 
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
This commit is contained in:
Javen
2026-09-21 21:34:15 -07:00
committed by GitHub
co-authored by Even Zhou sglang-npu-bot
parent 877a293d6d
commit a9f02b0fa4
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/xgrammar_backend.py
+5 -1
View File
@@ -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}")