chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

3
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
View File

@@ -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

11
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -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: