[Ascend] fix AscendAttnMaskBuilder bug to support float16 models (#14271)

2025-12-03 19:37:16 +08:00
parent 16d8de2284
commit 443d7bcd83
1 changed files with 5 additions and 3 deletions
--- a/python/sglang/srt/layers/attention/ascend_backend.py
+++ b/python/sglang/srt/layers/attention/ascend_backend.py
@@ -105,9 +105,11 @@ class AscendAttnMaskBuilder:
            )
        else:
            mask_value = torch.finfo(torch.float32).min if dtype == torch.float16 else 1
-        attn_mask = torch.zeros(
-            size=(max_seq_len, max_seq_len), dtype=dtype
-        ).masked_fill_(mask_flag, mask_value)
+        attn_mask = (
+            torch.zeros(size=(max_seq_len, max_seq_len))
+            .masked_fill_(mask_flag, mask_value)
+            .to(dtype)
+        )
        return attn_mask

    @staticmethod