From f142608408ed2a9b28aa649ce4c234bb1a2f30e1 Mon Sep 17 00:00:00 2001
From: ori <39351881+froststeam@users.noreply.github.com>
Date: Thu, 26 Mar 2026 12:00:58 +0800
Subject: [PATCH] [MUSA] apply_vocab_mask support musa device (#21296)

---
 python/sglang/srt/constrained/xgrammar_backend.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py
index 0012229d2..e2711e46d 100644
--- a/python/sglang/srt/constrained/xgrammar_backend.py
+++ b/python/sglang/srt/constrained/xgrammar_backend.py
@@ -104,11 +104,7 @@ 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 == "cuda"
-            or logits.device.type == "npu"
-            or logits.device.type == "xpu"
-        ):
+        if logits.device.type in {"cuda", "npu", "xpu", "musa"}:
             if _is_hip:
                 apply_token_bitmask_inplace_cuda(logits, vocab_mask)
             else: