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

[Diffusion][NPU][Bugfix] Ascend_fa crashes when sequence parallelism is used. (#23572)

Browse Source 
Co-authored-by: Napkin-AI <arseniy.mironov.dev@gmail.com>
This commit is contained in:
Arseniy Mironov
2026-04-23 19:21:30 +03:00
committed by GitHub
co-authored by Napkin-AI
parent 54e21bb3a5
commit 76e4c5a1f8
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py
+3 -4
View File
@@ -67,8 +67,6 @@ class AscendFAImpl(AttentionImpl):
) -> None: ) -> None:
self.causal = causal self.causal = causal
self.softmax_scale = softmax_scale self.softmax_scale = softmax_scale
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads or num_heads
def forward( def forward(
self, self,
@@ -79,6 +77,7 @@ class AscendFAImpl(AttentionImpl):
return_softmax_lse: bool = False, return_softmax_lse: bool = False,
) -> torch.Tensor: ) -> torch.Tensor:
mask = None mask = None
num_heads, num_key_value_heads = query.shape[2], key.shape[2]
if self.causal: if self.causal:
seq_len = query.shape[1] seq_len = query.shape[1]
mask = torch.triu( mask = torch.triu(
@@ -92,8 +91,8 @@ class AscendFAImpl(AttentionImpl):
query, query,
key, key,
value, value,
num_heads=self.num_heads, num_heads=num_heads,
num_key_value_heads=self.num_kv_heads, num_key_value_heads=num_key_value_heads,
scale=self.softmax_scale, scale=self.softmax_scale,
input_layout="BNSD", input_layout="BNSD",
softmax_lse_flag=return_softmax_lse, softmax_lse_flag=return_softmax_lse,