[Diffusion] Fix broken ring_attention when use upstream fa3 (#16270)

This commit is contained in:
Xiaoyu Zhang
2026-01-02 14:33:12 +08:00
committed by GitHub
parent 749736ba95
commit bd48ad5e6b
@@ -217,7 +217,11 @@ class FlashAttentionImpl(AttentionImpl):
seqlen, seqlen,
softmax_scale=self.softmax_scale, softmax_scale=self.softmax_scale,
causal=self.causal, causal=self.causal,
return_attn_probs=return_softmax_lse,
) )
if return_softmax_lse:
out, softmax_lse = out
return out.reshape(bsz, seqlen, nheads_q, -1), softmax_lse
return out.reshape(bsz, seqlen, nheads_q, d) return out.reshape(bsz, seqlen, nheads_q, d)
output = flash_attn_func( output = flash_attn_func(