[diffusion] cli: support component attention backend overrides (#24320)

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

@@ -6,8 +6,9 @@
 import os
 from collections.abc import Generator
 from contextlib import contextmanager
+from contextvars import ContextVar
 from functools import cache
-from typing import cast
+from typing import NamedTuple, cast
 
 import torch
 
@@ -63,6 +64,16 @@ def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
 forced_attn_backend: AttentionBackendEnum | None = None
 
 
+class ComponentAttnBackendContext(NamedTuple):
+    backend: AttentionBackendEnum | None
+    component_name: str | None
+
+
+component_attn_backend_context: ContextVar[ComponentAttnBackendContext | None] = (
+    ContextVar("component_attn_backend_context", default=None)
+)
+
+
 def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None:
     """
     Force all attention operations to use a specified backend.
@@ -86,10 +97,25 @@ def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
     return forced_attn_backend
 
 
+def get_component_attn_backend_context() -> ComponentAttnBackendContext | None:
+    return component_attn_backend_context.get()
+
+
+def get_component_forced_attn_backend() -> AttentionBackendEnum | None:
+    context = get_component_attn_backend_context()
+    return context.backend if context is not None else None
+
+
+def get_component_attn_backend_name() -> str | None:
+    context = get_component_attn_backend_context()
+    return context.component_name if context is not None else None
+
+
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
     supported_attention_backends: set[AttentionBackendEnum] | None = None,
+    selected_attention_backend: AttentionBackendEnum | None = None,
 ) -> type[AttentionBackend]:
     if supported_attention_backends is None:
         be_tuple = tuple()
@@ -98,7 +124,41 @@ def get_attn_backend(
         be_tuple = tuple(
             sorted(list(supported_attention_backends), key=lambda b: b.name)
         )
-    return _cached_get_attn_backend(head_size, dtype, be_tuple)
+
+    selected_backend = selected_attention_backend or get_global_forced_attn_backend()
+    if selected_backend is None:
+        selected_backend = get_component_forced_attn_backend()
+    if selected_backend is None:
+        server_args = get_global_server_args()
+        if server_args.attention_backend is not None:
+            try:
+                selected_backend = AttentionBackendEnum[
+                    server_args.attention_backend.upper()
+                ]
+            except KeyError:
+                raise ValueError(
+                    f"Invalid attention backend '{server_args.attention_backend}' specified via command line. "
+                    f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
+                )
+
+    component_name = get_component_attn_backend_name()
+    backend_not_specified = selected_backend is None
+    attention_backend_cls = _cached_get_attn_backend(
+        head_size,
+        dtype,
+        be_tuple,
+        selected_backend,
+    )
+    if component_name:
+        backend_name = attention_backend_cls.get_enum().name.lower()
+        if backend_not_specified:
+            logger.info_once(
+                f"Attention backend not specified for {component_name}, "
+                f"using {backend_name} backend for {component_name}"
+            )
+        else:
+            logger.info_once(f"Using {backend_name} backend for {component_name}")
+    return attention_backend_cls
 
 
 @cache
@@ -106,32 +166,11 @@ def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
     supported_attention_backends: tuple[AttentionBackendEnum],
+    selected_backend: AttentionBackendEnum | None,
 ) -> type[AttentionBackend]:
-    # Check whether a particular choice of backend was
-    # previously forced via global_force_attn_backend() or --attention-backend CLI arg.
     from sglang.multimodal_gen.runtime.platforms import current_platform
 
     supported_attention_backends = set(supported_attention_backends)
-    selected_backend = None
-    backend_by_global_setting: AttentionBackendEnum | None = (
-        get_global_forced_attn_backend()
-    )
-    if backend_by_global_setting is not None:
-        selected_backend = backend_by_global_setting
-    else:
-        # Check the server arguments for a backend override
-        server_args = get_global_server_args()
-        if server_args.attention_backend is not None:
-            try:
-                selected_backend = AttentionBackendEnum[
-                    server_args.attention_backend.upper()
-                ]
-
-            except KeyError:
-                raise ValueError(
-                    f"Invalid attention backend '{server_args.attention_backend}' specified via command line. "
-                    f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
-                )
 
     # get device-specific attn_backend
     if len(supported_attention_backends) == 0:
@@ -140,14 +179,16 @@ def _cached_get_attn_backend(
     elif selected_backend is None and len(supported_attention_backends) == 1:
         selected_backend = next(iter(supported_attention_backends))
     elif selected_backend is None:
-        logger.debug(f"Attention backend not specified")
+        logger.debug("Attention backend not specified")
     elif selected_backend not in supported_attention_backends:
         supported_attention_backends_str = [
             supported_attention_backend.__str__()
             for supported_attention_backend in supported_attention_backends
         ]
         logger.debug(
-            f"Selected attention backend: '{selected_backend}' not in supported attention backends: {supported_attention_backends_str}"
+            "Selected attention backend: '%s' not in supported attention backends: %s",
+            selected_backend,
+            supported_attention_backends_str,
         )
         selected_backend = None
 
@@ -161,6 +202,24 @@ def _cached_get_attn_backend(
     return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
 
 
+@contextmanager
+def component_attn_backend_context_manager(
+    attn_backend: AttentionBackendEnum | None,
+    component_name: str | None = None,
+) -> Generator[None, None, None]:
+    if attn_backend is None and component_name is None:
+        yield
+        return
+
+    token = component_attn_backend_context.set(
+        ComponentAttnBackendContext(attn_backend, component_name)
+    )
+    try:
+        yield
+    finally:
+        component_attn_backend_context.reset(token)
+
+
 @contextmanager
 def global_force_attn_backend_context_manager(
     attn_backend: AttentionBackendEnum,

python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py

@@ -16,6 +16,10 @@ from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 
 from sglang.multimodal_gen.configs.models import ModelConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    component_attn_backend_context_manager,
+    get_component_attn_backend_context,
+)
 from sglang.multimodal_gen.runtime.loader.utils import (
     _normalize_component_type,
     component_name_to_loader_cls,
@@ -114,10 +118,26 @@ class ComponentLoader(ABC):
             component_model_path,
             gpu_mem_before_loading,
         )
-        try:
-            component = self.load_customized(
-                component_model_path, server_args, component_name
+        attn_backend = None
+        component_attn_name = None
+        if get_component_attn_backend_context() is None:
+            attn_backend, matched_backend_key = (
+                server_args.resolve_component_attention_backend(component_name)
             )
+            component_attn_name = matched_backend_key or component_name
+            if attn_backend is not None:
+                logger.info(
+                    "Using %s backend for component: %s",
+                    attn_backend.name.lower(),
+                    matched_backend_key,
+                )
+        try:
+            with component_attn_backend_context_manager(
+                attn_backend, component_name=component_attn_name
+            ):
+                component = self.load_customized(
+                    component_model_path, server_args, component_name
+                )
             source = "sgl-diffusion"
         except Exception as e:
             if "Unsupported model architecture" in str(e):
@@ -130,9 +150,12 @@ class ComponentLoader(ABC):
                     f"Error while loading customized {component_name}, falling back to native version"
                 )
             # fallback to native version
-            component = self.load_native(
-                component_model_path, server_args, transformers_or_diffusers
-            )
+            with component_attn_backend_context_manager(
+                attn_backend, component_name=component_attn_name
+            ):
+                component = self.load_native(
+                    component_model_path, server_args, transformers_or_diffusers
+                )
             should_offload = self.should_offload(server_args)
             target_device = self.target_device(should_offload)
             component = component.to(device=target_device)

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -18,6 +18,9 @@ from sglang.multimodal_gen.runtime.disaggregation.roles import (
     RoleType,
     filter_modules_for_role,
 )
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    component_attn_backend_context_manager,
+)
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     PipelineComponentLoader,
 )
@@ -411,13 +414,27 @@ class ComposedPipelineBase(ABC):
             component_model_path = self._resolve_component_path(
                 server_args, module_name, load_module_name
             )
-            module, memory_usage = PipelineComponentLoader.load_component(
-                component_name=load_module_name,
-                component_model_path=component_model_path,
-                transformers_or_diffusers=transformers_or_diffusers,
-                server_args=server_args,
-                component_architecture=architecture,
+            attn_backend, matched_backend_key = (
+                server_args.resolve_component_attention_backend(
+                    module_name, load_module_name
+                )
             )
+            if attn_backend is not None:
+                logger.info(
+                    "Using %s backend for component: %s",
+                    attn_backend.name.lower(),
+                    matched_backend_key,
+                )
+            with component_attn_backend_context_manager(
+                attn_backend, component_name=matched_backend_key or module_name
+            ):
+                module, memory_usage = PipelineComponentLoader.load_component(
+                    component_name=load_module_name,
+                    component_model_path=component_model_path,
+                    transformers_or_diffusers=transformers_or_diffusers,
+                    server_args=server_args,
+                    component_architecture=architecture,
+                )
 
             self.memory_usages[load_module_name] = memory_usage
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -179,10 +179,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
 
-        # TODO(will): hack, should use the actual one in dit
+        selected_attention_backend = self._infer_transformer_attention_backend()
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
             dtype=torch.float16,
+            selected_attention_backend=selected_attention_backend,
         )
 
         # cfg
@@ -195,6 +196,26 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self._cached_num_steps = None
         self._is_warmed_up = False
 
+    def _infer_transformer_attention_backend(self) -> AttentionBackendEnum | None:
+        backends = {
+            backend
+            for transformer in (self.transformer, self.transformer_2)
+            if transformer is not None
+            for module in transformer.modules()
+            if isinstance(
+                (backend := getattr(module, "backend", None)), AttentionBackendEnum
+            )
+        }
+        if not backends:
+            return None
+        if len(backends) > 1:
+            logger.warning(
+                "Multiple transformer attention backends detected: %s. "
+                "Using one backend for denoising metadata.",
+                sorted(backend.name.lower() for backend in backends),
+            )
+        return sorted(backends, key=lambda backend: backend.name)[0]
+
     def component_uses(
         self, server_args: ServerArgs, stage_name: str | None = None
     ) -> list[ComponentUse]:

python/sglang/multimodal_gen/runtime/server_args.py

@@ -131,6 +131,9 @@ class ServerArgs(DisaggArgsMixin):
     # Attention
     attention_backend: str = None
     attention_backend_config: addict.Dict | None = None
+    component_attention_backends: dict[str, str] | str | None = field(
+        default_factory=dict
+    )
     cache_dit_config: str | dict[str, Any] | None = (
         None  # cache-dit config for diffusers
     )
@@ -470,6 +473,11 @@ class ServerArgs(DisaggArgsMixin):
     def _adjust_attention_backend(self):
         if self.attention_backend in ["fa3", "fa4"]:
             self.attention_backend = "fa"
+        self.component_attention_backends = (
+            self._normalize_component_attention_backends(
+                self.component_attention_backends
+            )
+        )
 
         # attention_backend_config
         if self.attention_backend_config is None:
@@ -512,6 +520,82 @@ class ServerArgs(DisaggArgsMixin):
                 return
             self._set_default_attention_backend()
 
+    @staticmethod
+    def _normalize_attention_backend_name(backend: str) -> str:
+        if not isinstance(backend, str):
+            raise ValueError("Attention backend name must be a string")
+        normalized = backend.strip().lower()
+        if normalized in ("fa3", "fa4"):
+            normalized = "fa"
+        try:
+            return AttentionBackendEnum[normalized.upper()].name.lower()
+        except KeyError:
+            raise ValueError(
+                f"Invalid attention backend '{backend}'. "
+                f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
+            ) from None
+
+    @staticmethod
+    def _parse_component_attention_backend_map(
+        value: dict[str, str] | str | None,
+    ) -> dict[str, str]:
+        if value is None or value == "":
+            return {}
+        if isinstance(value, dict):
+            return dict(value)
+        if not isinstance(value, str):
+            raise ValueError(
+                "component_attention_backends must be a dict or a comma-separated component=backend string"
+            )
+
+        try:
+            parsed = json.loads(value)
+            if not isinstance(parsed, dict):
+                raise ValueError
+            return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+        result: dict[str, str] = {}
+        for pair in value.split(","):
+            pair = pair.strip()
+            if not pair:
+                continue
+            if "=" not in pair:
+                raise ValueError(
+                    "component_attention_backends must use component=backend entries"
+                )
+            component, backend = pair.split("=", 1)
+            result[component.strip()] = backend.strip()
+        return result
+
+    @classmethod
+    def _normalize_component_attention_backends(
+        cls, value: dict[str, str] | str | None
+    ) -> dict[str, str]:
+        raw = cls._parse_component_attention_backend_map(value)
+        normalized: dict[str, str] = {}
+        for component, backend in raw.items():
+            if not isinstance(component, str):
+                raise ValueError("Component attention backend key must be a string")
+            component_name = component.strip().replace("-", "_")
+            if not component_name:
+                raise ValueError("Component attention backend key must not be empty")
+            normalized[component_name] = cls._normalize_attention_backend_name(backend)
+        return normalized
+
+    def resolve_component_attention_backend(
+        self, *component_names: str | None
+    ) -> tuple[AttentionBackendEnum | None, str | None]:
+        for component_name in component_names:
+            if component_name is None:
+                continue
+            key = component_name.replace("-", "_")
+            backend = self.component_attention_backends.get(key)
+            if backend is not None:
+                return AttentionBackendEnum[backend.upper()], key
+        return None, None
+
     def _adjust_warmup(self):
         if self.warmup_resolutions is not None:
             self.warmup = True
@@ -808,6 +892,16 @@ class ServerArgs(DisaggArgsMixin):
             default=None,
             help="Configuration for the attention backend. Can be a JSON string, a path to a JSON/YAML file, or key=value pairs.",
         )
+        parser.add_argument(
+            "--component-attention-backends",
+            type=str,
+            default=None,
+            help=(
+                "Per-component attention backend overrides for native pipelines. "
+                "Use component names from model_index.json, e.g. "
+                "'text_encoder=torch_sdpa,transformer=fa'."
+            ),
+        )
         parser.add_argument(
             "--cache-dit-config",
             type=str,
@@ -1267,6 +1361,43 @@ class ServerArgs(DisaggArgsMixin):
             component_paths[component] = path
         return component_paths, remaining
 
+    @staticmethod
+    def _extract_component_attention_backends(
+        unknown_args: list[str],
+    ) -> tuple[dict[str, str], list[str]]:
+        component_attention_backends: dict[str, str] = {}
+        remaining: list[str] = []
+        i = 0
+        while i < len(unknown_args):
+            arg = unknown_args[i]
+            key_part = arg.split("=", 1)[0] if "=" in arg else arg
+            component = None
+            if key_part.startswith("--component-attention-backends."):
+                component = key_part[len("--component-attention-backends.") :].replace(
+                    "-", "_"
+                )
+            elif key_part.startswith("--component_attention_backends."):
+                component = key_part[len("--component_attention_backends.") :].replace(
+                    "-", "_"
+                )
+
+            if component is not None:
+                if "=" in arg:
+                    component_attention_backends[component] = arg.split("=", 1)[1]
+                elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
+                    "-"
+                ):
+                    i += 1
+                    component_attention_backends[component] = unknown_args[i]
+                else:
+                    remaining.append(arg)
+                    i += 1
+                    continue
+            else:
+                remaining.append(arg)
+            i += 1
+        return component_attention_backends, remaining
+
     @classmethod
     def from_cli_args(
         cls, args: argparse.Namespace, unknown_args: list[str] | None = None
@@ -1276,6 +1407,9 @@ class ServerArgs(DisaggArgsMixin):
 
         # extract dynamic --<component>-path from unknown args
         dynamic_paths, remaining = cls._extract_component_paths(unknown_args)
+        dynamic_attention_backends, remaining = (
+            cls._extract_component_attention_backends(remaining)
+        )
         if remaining:
             raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}")
 
@@ -1291,6 +1425,12 @@ class ServerArgs(DisaggArgsMixin):
             existing = dict(provided_args.get("component_paths") or {})
             existing.update(dynamic_paths)
             provided_args["component_paths"] = existing
+        if dynamic_attention_backends:
+            existing = cls._parse_component_attention_backend_map(
+                provided_args.get("component_attention_backends")
+            )
+            existing.update(dynamic_attention_backends)
+            provided_args["component_attention_backends"] = existing
 
         return cls.from_dict(provided_args)
 

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -48,6 +48,71 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             args.component_paths["vae"], os.path.expanduser("~/fake/local/vae")
         )
 
+    def test_component_attention_backends_are_normalized(self):
+        args = self._from_dict_without_model_resolution(
+            {
+                "model_path": "/data/my-model",
+                "component_attention_backends": "text-encoder=torch_sdpa,transformer=fa3",
+            }
+        )
+
+        self.assertEqual(
+            args.component_attention_backends,
+            {"text_encoder": "torch_sdpa", "transformer": "fa"},
+        )
+
+    def test_component_attention_backend_lookup(self):
+        args = self._from_dict_without_model_resolution(
+            {
+                "model_path": "/data/my-model",
+                "component_attention_backends": {"text_encoder": "torch_sdpa"},
+            }
+        )
+
+        backend, matched_key = args.resolve_component_attention_backend(
+            "text_encoder", "transformer"
+        )
+
+        self.assertEqual(backend.name, "TORCH_SDPA")
+        self.assertEqual(matched_key, "text_encoder")
+
+    def test_invalid_component_attention_backend_raises(self):
+        with self.assertRaises(ValueError):
+            self._from_dict_without_model_resolution(
+                {
+                    "model_path": "/data/my-model",
+                    "component_attention_backends": {"text_encoder": "bad_backend"},
+                }
+            )
+        with self.assertRaises(ValueError):
+            self._from_dict_without_model_resolution(
+                {
+                    "model_path": "/data/my-model",
+                    "component_attention_backends": "text_encoder",
+                }
+            )
+
+    def test_dynamic_component_attention_backend_cli_args(self):
+        parser = FlexibleArgumentParser()
+        ServerArgs.add_cli_args(parser)
+        argv = [
+            "--model-path",
+            "/fake",
+            "--component-attention-backends.text-encoder",
+            "torch_sdpa",
+        ]
+
+        with patch.object(sys, "argv", ["sglang"] + argv):
+            args, unknown_args = parser.parse_known_args(argv)
+            with patch.object(
+                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+            ):
+                server_args = ServerArgs.from_cli_args(args, unknown_args)
+
+        self.assertEqual(
+            server_args.component_attention_backends, {"text_encoder": "torch_sdpa"}
+        )
+
 
 class TestOffloadDefaults(unittest.TestCase):
     def _from_dict_with_task_type(