From 76e4c5a1f8565232584d46671aa07d3fe2c20422 Mon Sep 17 00:00:00 2001 From: Arseniy Mironov <98156294+Napkin-AI@users.noreply.github.com> Date: Thu, 23 Apr 2026 19:21:30 +0300 Subject: [PATCH] [Diffusion][NPU][Bugfix] Ascend_fa crashes when sequence parallelism is used. (#23572) Co-authored-by: Napkin-AI --- .../runtime/layers/attention/backends/ascend_fa.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py index 0fc3db29b..2c40bbda3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py @@ -67,8 +67,6 @@ class AscendFAImpl(AttentionImpl): ) -> None: self.causal = causal self.softmax_scale = softmax_scale - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads or num_heads def forward( self, @@ -79,6 +77,7 @@ class AscendFAImpl(AttentionImpl): return_softmax_lse: bool = False, ) -> torch.Tensor: mask = None + num_heads, num_key_value_heads = query.shape[2], key.shape[2] if self.causal: seq_len = query.shape[1] mask = torch.triu( @@ -92,8 +91,8 @@ class AscendFAImpl(AttentionImpl): query, key, value, - num_heads=self.num_heads, - num_key_value_heads=self.num_kv_heads, + num_heads=num_heads, + num_key_value_heads=num_key_value_heads, scale=self.softmax_scale, input_layout="BNSD", softmax_lse_flag=return_softmax_lse,