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

[MUSA] apply_vocab_mask support musa device (#21296)

Browse Source
This commit is contained in:
ori
2026-03-25 21:00:58 -07:00
committed by GitHub
parent f420b9b4a5
commit f142608408
1 changed files with 1 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/xgrammar_backend.py
+1 -5
View File
@@ -104,11 +104,7 @@ class XGrammarGrammar(BaseGrammarObject):
return vocab_mask.to(device, non_blocking=True) return vocab_mask.to(device, non_blocking=True)
def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
if ( if logits.device.type in {"cuda", "npu", "xpu", "musa"}:
logits.device.type == "cuda"
or logits.device.type == "npu"
or logits.device.type == "xpu"
):
if _is_hip: if _is_hip:
apply_token_bitmask_inplace_cuda(logits, vocab_mask) apply_token_bitmask_inplace_cuda(logits, vocab_mask)
else: else: