From bd48ad5e6bb9d3cc121f37ad28cc305cd7ff6521 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Fri, 2 Jan 2026 14:33:12 +0800
Subject: [PATCH] [Diffusion] Fix broken ring_attention when use upstream fa3
 (#16270)

---
 .../runtime/layers/attention/backends/flash_attn.py           | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
index 17a9cd576..4b033e527 100644
--- 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(