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

2026-01-02 14:33:12 +08:00
parent 749736ba95
commit bd48ad5e6b
1 changed files with 4 additions and 0 deletions
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
@@ -217,7 +217,11 @@ class FlashAttentionImpl(AttentionImpl):
                seqlen,
                softmax_scale=self.softmax_scale,
                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)

        output = flash_attn_func(