From 443d7bcd836287c7d4a689083c4f9d0b5d66cfa1 Mon Sep 17 00:00:00 2001 From: Michelle Wu Date: Wed, 3 Dec 2025 19:37:16 +0800 Subject: [PATCH] [Ascend] fix AscendAttnMaskBuilder bug to support float16 models (#14271) --- python/sglang/srt/layers/attention/ascend_backend.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/ascend_backend.py b/python/sglang/srt/layers/attention/ascend_backend.py index 458e98016..d2b18c638 100644 --- 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