From b47366fbf9984fa25be62b48eb65d30b92644e60 Mon Sep 17 00:00:00 2001 From: yuefeng Wu <33725817+ChefWu551@users.noreply.github.com> Date: Fri, 29 May 2026 12:38:06 +0800 Subject: [PATCH] [NPU]: Optimize xgrammar token bitmask on NPU with AscendC (#24133) --- .../srt/constrained/torch_ops/bitmask_ops.py | 33 ------------------- .../srt/constrained/xgrammar_backend.py | 8 ++--- 2 files changed, 4 insertions(+), 37 deletions(-) delete mode 100644 python/sglang/srt/constrained/torch_ops/bitmask_ops.py diff --git a/python/sglang/srt/constrained/torch_ops/bitmask_ops.py b/python/sglang/srt/constrained/torch_ops/bitmask_ops.py deleted file mode 100644 index 25f07e5ca..000000000 --- a/python/sglang/srt/constrained/torch_ops/bitmask_ops.py +++ /dev/null @@ -1,33 +0,0 @@ -# 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 46ce3c305..c02021188 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -35,13 +35,11 @@ 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 _is_hip = is_hip() + if _is_hip: from sgl_kernel import apply_token_bitmask_inplace_cuda else: @@ -118,7 +116,9 @@ class XGrammarGrammar(BaseGrammarObject): else: apply_token_bitmask_inplace_triton(logits, vocab_mask) elif logits.device.type == "npu": - apply_token_bitmask_inplace_torch(logits, vocab_mask) + import sgl_kernel_npu # noqa: F401 + + torch.ops.npu.apply_token_bitmask(logits, vocab_mask) else: raise RuntimeError(f"Unsupported device: {logits.device.type}")