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

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

Browse Source
This commit is contained in:
Michelle Wu
2025-12-03 19:37:16 +08:00
committed by GitHub
parent 16d8de2284
commit 443d7bcd83
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/ascend_backend.py
+5 -3
View File
@@ -105,9 +105,11 @@ class AscendAttnMaskBuilder:
) )
else: else:
mask_value = torch.finfo(torch.float32).min if dtype == torch.float16 else 1 mask_value = torch.finfo(torch.float32).min if dtype == torch.float16 else 1
attn_mask = torch.zeros( attn_mask = (
size=(max_seq_len, max_seq_len), dtype=dtype torch.zeros(size=(max_seq_len, max_seq_len))
).masked_fill_(mask_flag, mask_value) .masked_fill_(mask_flag, mask_value)
.to(dtype)
)
return attn_mask return attn_mask
@staticmethod @staticmethod