From 5e5b1183ed2efc1a401a9f731ef688c463bb22ff Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Tue, 6 Jan 2026 14:45:53 +0800
Subject: [PATCH] [Diffusion] Ring Attention support sage backend (#16496)

---
 .../runtime/layers/attention/backends/sage_attn.py    |  3 +++
 python/sglang/multimodal_gen/runtime/server_args.py   | 11 +++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
index db4643096..973608f64 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
@@ -68,4 +68,7 @@ class SageAttentionImpl(AttentionImpl):
             sm_scale=self.softmax_scale,
             return_lse=return_softmax_lse,
         )
+        if return_softmax_lse:
+            output, softmax_lse = output
+            return output, softmax_lse
         return output
diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py
index 2c5c602dc..9095995a9 100644
--- a/python/sglang/multimodal_gen/runtime/server_args.py
+++ b/python/sglang/multimodal_gen/runtime/server_args.py
@@ -888,14 +888,17 @@ class ServerArgs:
             logger.debug(f"Ring degree not set, using default value {self.ring_degree}")
 
         if self.ring_degree > 1:
-            if self.attention_backend is not None and self.attention_backend != "fa":
+            if self.attention_backend is not None and self.attention_backend not in (
+                "fa",
+                "sage_attn",
+            ):
                 raise ValueError(
-                    "Ring Attention is only supported for flash attention backend for now"
+                    "Ring Attention is only supported for flash attention or sage attention backend for now"
                 )
-            else:
+            if self.attention_backend is None:
                 self.attention_backend = "fa"
                 logger.info(
-                    "Ring Attention is currently only supported for flash attention, attention_backend has been automatically set to flash attention"
+                    "Ring Attention is currently only supported for flash attention or sage attention; attention_backend has been automatically set to flash attention"
                 )
 
         if self.sp_degree == -1: