minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Diffusion] Ring Attention support sage backend (#16496)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-01-06 14:45:53 +08:00
committed by GitHub
parent 9bf76c11a9
commit 5e5b1183ed
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
+3
View File
@@ -68,4 +68,7 @@ class SageAttentionImpl(AttentionImpl):
sm_scale=self.softmax_scale, sm_scale=self.softmax_scale,
return_lse=return_softmax_lse, return_lse=return_softmax_lse,
) )
if return_softmax_lse:
output, softmax_lse = output
return output, softmax_lse
return output return output
python/sglang/multimodal_gen/runtime/server_args.py
+7 -4
View File
@@ -888,14 +888,17 @@ class ServerArgs:
logger.debug(f"Ring degree not set, using default value {self.ring_degree}") logger.debug(f"Ring degree not set, using default value {self.ring_degree}")
if self.ring_degree > 1: 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( 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" self.attention_backend = "fa"
logger.info( 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: if self.sp_degree == -1: