From f8e7f4b0138181f640a6cc039eced8f84322da7b Mon Sep 17 00:00:00 2001 From: "hewenqing.hwq" Date: Tue, 4 Aug 2026 16:08:15 +0800 Subject: [PATCH] performance(ring para): change sdpa backend to flash attention --- .../layers/attention/backends/sdpa.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/diffsynth_engine/layers/attention/backends/sdpa.py b/diffsynth_engine/layers/attention/backends/sdpa.py index cf7dd3a..bfd5073 100644 --- a/diffsynth_engine/layers/attention/backends/sdpa.py +++ b/diffsynth_engine/layers/attention/backends/sdpa.py @@ -9,7 +9,7 @@ AttentionType, ) -_scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention +_scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention class SDPABackend(AttentionBackend): @@ -88,6 +88,12 @@ def forward_with_lse( attn_metadata: AttentionMetadata | None = None, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: + if attn_mask is not None: + raise NotImplementedError( + "SDPA forward_with_lse does not support attn_mask; " + "masked Ring attention is not supported" + ) + query = rearrange(query, "b s n d -> b n s d") key = rearrange(key, "b s n d -> b n s d") value = rearrange(value, "b s n d -> b n s d") @@ -96,18 +102,15 @@ def forward_with_lse( key = torch.repeat_interleave(key, self.num_kv_groups, dim=1) value = torch.repeat_interleave(value, self.num_kv_groups, dim=1) - seq_len = query.shape[2] - output, lse = _scaled_dot_product_efficient_attention( + output, lse = _scaled_dot_product_flash_attention( query, key, value, - attn_bias=attn_mask, - compute_log_sumexp=True, + dropout_p=0.0, is_causal=self.causal, + return_debug_mask=False, scale=self.softmax_scale, )[:2] output = rearrange(output, "b n s d -> b s n d") - # the returned lse is padded but not restored, so we need to slice it - lse = lse[:, :, :seq_len] return output, lse