From 74f25e8e06d1208c0d7e9a59148dc612ceab0387 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 7 Aug 2026 10:24:02 +0800 Subject: [PATCH 1/2] refactor(diffusion): make ring admission a backend capability Replace the two name whitelists (USPAttention init's enum check and server_args' string check) with AttentionBackend.supports_ring_rotation: the per-hop online-softmax merge needs the kernel's softmax LSE, so backends declare it instead of call sites naming them. The server-args list stays name-level (it gates before backend classes are importable on every platform) but now mirrors the capability, with a test pinning the two views together. zimage's silent full-sequence gather under ring gets a one-shot warning. Co-Authored-By: Claude Fable 5 --- .../attention/backends/attention_backend.py | 6 +++ .../layers/attention/backends/flash_attn.py | 4 ++ .../layers/attention/backends/sage_attn.py | 4 ++ .../runtime/layers/attention/layer.py | 13 +++---- .../runtime/models/dits/zimage.py | 8 ++++ .../runtime/server_args/server_args.py | 20 ++++++---- .../test/unit/test_ring_admission.py | 39 +++++++++++++++++++ 7 files changed, 79 insertions(+), 15 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_ring_admission.py diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 45736398f795..86e0ea6e7c43 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -45,6 +45,12 @@ def get_impl_cls() -> type["AttentionImpl"]: def supports_packed_varlen(cls) -> bool: return cls.get_impl_cls().forward_varlen is not AttentionImpl.forward_varlen + @classmethod + def supports_ring_rotation(cls) -> bool: + """Whether this backend can serve as the ring-attention kernel; the + per-hop online-softmax merge needs the kernel's softmax LSE.""" + return False + @classmethod def unsupported_requirements( cls, requirements: AttentionRequirements 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 21f71185272f..a1d0b2445515 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 @@ -330,6 +330,10 @@ def build( # type: ignore class FlashAttentionBackend(AttentionBackend): + + @classmethod + def supports_ring_rotation(cls) -> bool: + return True accept_output_buffer: bool = True @staticmethod 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 ad785dd4acff..0ae7d02b640f 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 @@ -33,6 +33,10 @@ def _trailing_padding_used_len( class SageAttentionBackend(AttentionBackend): + + @classmethod + def supports_ring_rotation(cls) -> bool: + return True accept_output_buffer: bool = True @staticmethod diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index d88df2eeeae1..6a2963effb52 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -689,15 +689,12 @@ def __init__( head_size, dtype, supported_attention_backends=supported_attention_backends ) if get_ring_parallel_world_size() > 1: - backend_enum = attn_backend.get_enum() - if backend_enum not in ( - AttentionBackendEnum.FA, - AttentionBackendEnum.SAGE_ATTN, - ): + if not attn_backend.supports_ring_rotation(): raise RuntimeError( - f"Ring Attention is only supported for FlashAttention or SageAttention backends, " - f"but got {backend_enum.name}. " - f"Please ensure your platform supports these backends." + f"Ring Attention requires a backend whose kernel exposes the " + f"softmax LSE for the per-hop merge; " + f"{attn_backend.get_enum().name} does not declare support " + f"(see AttentionBackend.supports_ring_rotation)." ) impl_cls: Type[AttentionImpl] = attn_backend.get_impl_cls() self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False)) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index f558a7c15cb9..43e6e00a85ef 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -1516,6 +1516,14 @@ def forward( use_full_unified_sequence = ( get_sp_world_size() > 1 and get_ring_parallel_world_size() > 1 ) + if use_full_unified_sequence: + # Ring support for this attention layout is not implemented; the + # full-sequence gather is correct but gives up ring's memory and + # overlap benefits. + logger.warning_once( + "zimage under ring_degree > 1 falls back to a full-sequence " + "K/V gather" + ) x_local_seq_len = x.shape[1] if use_full_unified_sequence: x = sequence_model_parallel_all_gather(x.contiguous(), dim=1) diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index b8f3f4dd1665..e8d6cb055b30 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -71,6 +71,9 @@ # H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident. LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130 LORA_MERGE_MODES = ("auto", "merge", "dynamic") +# Mirrors AttentionBackend.supports_ring_rotation; the name-level check +# runs before backend classes are importable on every platform. +RING_CAPABLE_ATTENTION_BACKENDS = ("fa", "sage_attn") def _normalize_ltx2_two_stage_device_mode(mode: str | None) -> str | None: @@ -756,18 +759,21 @@ def _adjust_attention_backend(self): self.component_attention_backends["text_encoder"] = "torch_sdpa" if self.ring_degree > 1: - if self.attention_backend is not None and self.attention_backend not in ( - "fa", - "sage_attn", + if ( + self.attention_backend is not None + and self.attention_backend not in RING_CAPABLE_ATTENTION_BACKENDS ): raise ValueError( - "Ring Attention is only supported for flash attention or sage attention backend for now" + "Ring Attention requires one of the ring-capable backends " + f"({', '.join(RING_CAPABLE_ATTENTION_BACKENDS)}), got " + f"{self.attention_backend!r}" ) if self.attention_backend is None: - self.attention_backend = "fa" + self.attention_backend = RING_CAPABLE_ATTENTION_BACKENDS[0] logger.info( - "Ring Attention is currently only supported for flash attention or sage attention; " - "attention_backend has been automatically set to flash attention" + "Ring Attention requires a ring-capable backend; " + "attention_backend has been automatically set to %s", + self.attention_backend, ) if self.attention_backend is None and self.backend != Backend.DIFFUSERS: diff --git a/python/sglang/multimodal_gen/test/unit/test_ring_admission.py b/python/sglang/multimodal_gen/test/unit/test_ring_admission.py new file mode 100644 index 000000000000..0978d465458b --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_ring_admission.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ring admission is a backend capability, not a name whitelist.""" + +import unittest + +from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( + AttentionBackend, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + FlashAttentionBackend, +) +from sglang.multimodal_gen.runtime.layers.attention.backends.sdpa import SDPABackend +from sglang.multimodal_gen.runtime.server_args.server_args import ( + RING_CAPABLE_ATTENTION_BACKENDS, +) + + +class TestRingAdmission(unittest.TestCase): + def test_default_is_not_ring_capable(self): + self.assertFalse(AttentionBackend.supports_ring_rotation()) + self.assertFalse(SDPABackend.supports_ring_rotation()) + + def test_lse_backends_declare_support(self): + self.assertTrue(FlashAttentionBackend.supports_ring_rotation()) + + def test_server_args_names_match_capabilities(self): + # the name-level list gates before backend classes are importable on + # every platform; keep it consistent with the classes it mirrors + self.assertIn( + FlashAttentionBackend.get_enum().name.lower(), + RING_CAPABLE_ATTENTION_BACKENDS, + ) + self.assertNotIn( + SDPABackend.get_enum().name.lower(), RING_CAPABLE_ATTENTION_BACKENDS + ) + + +if __name__ == "__main__": + unittest.main() From 26d1794d0bb7113aeae4dd440a9225864b76ec34 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 7 Aug 2026 12:44:00 +0800 Subject: [PATCH 2/2] style: lint (black, isort, hoist backend import to module top) Co-Authored-By: Claude Fable 5 --- .../layers/attention/backends/flash_attn.py | 14 +++++++------- .../runtime/layers/attention/backends/sage_attn.py | 1 + 2 files changed, 8 insertions(+), 7 deletions(-) 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 a1d0b2445515..d42e1936bac0 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 @@ -6,6 +6,12 @@ import torch from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func +from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) from sglang.multimodal_gen.runtime.layers.utils import register_custom_op from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, @@ -285,13 +291,6 @@ def flash_attn_varlen_func_op_lse( ) -from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( - AttentionBackend, - AttentionImpl, - AttentionMetadata, - AttentionMetadataBuilder, -) - fa_ver = 3 @@ -334,6 +333,7 @@ class FlashAttentionBackend(AttentionBackend): @classmethod def supports_ring_rotation(cls) -> bool: return True + accept_output_buffer: bool = True @staticmethod 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 0ae7d02b640f..c4ffbaea56b8 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 @@ -37,6 +37,7 @@ class SageAttentionBackend(AttentionBackend): @classmethod def supports_ring_rotation(cls) -> bool: return True + accept_output_buffer: bool = True @staticmethod