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: