[MUSA] apply_vocab_mask support musa device (#21296)

2026-03-25 21:00:58 -07:00
parent f420b9b4a5
commit f142608408
1 changed files with 1 additions and 5 deletions
@@ -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: