[diffusion] chore: derive h3 attention admission from backend capabilities (#33707)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py

@@ -3,7 +3,6 @@ from dataclasses import dataclass, field
 
 from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
 from sglang.multimodal_gen.configs.models.fsdp import is_block
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64
 MINIMAX_H3_ADALN_MODALITY_NUM = 3
@@ -15,15 +14,6 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
 
     lora_param_names_mapping: dict = field(default_factory=dict)
 
-    _supported_attention_backends: set[AttentionBackendEnum] = field(
-        default_factory=lambda: {
-            AttentionBackendEnum.FA,
-            AttentionBackendEnum.SAGE_ATTN,
-            AttentionBackendEnum.AITER,
-            AttentionBackendEnum.TORCH_SDPA,
-        }
-    )
-
     num_layers: int = 50
     token_refiner_num_layers: int = 2
     hidden_size: int = 5376

python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py

@@ -2,6 +2,8 @@
 import os
 from dataclasses import dataclass, field
 
+import torch
+
 from sglang.multimodal_gen.configs.models.dits.minimax_h3 import MiniMaxH3DiTConfig
 from sglang.multimodal_gen.configs.models.encoders.minimax_h3_qwen3vl import (
     MiniMaxH3Qwen3VLConfig,
@@ -19,7 +21,14 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
 from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
     ModelDeploymentConfig,
 )
-from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionRequirements,
+)
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 
 
 @dataclass
@@ -169,6 +178,23 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
     def validate_server_args(self, server_args) -> None:
         # Reject known-inexact VAE modes before any large component download.
         self.vae_config.resolved_parallel_decode_mode()
+        component_backends = server_args.component_attention_backends or {}
+        attention_backend = component_backends.get(
+            "transformer", self._server_arg_value(server_args.attention_backend)
+        )
+        if attention_backend is None:
+            return
+        selected_backend = (
+            attention_backend
+            if isinstance(attention_backend, AttentionBackendEnum)
+            else AttentionBackendEnum[str(attention_backend).strip().upper()]
+        )
+        get_attn_backend(
+            self.dit_config.arch_config.attention_head_dim,
+            torch.bfloat16,
+            selected_attention_backend=selected_backend,
+            attention_requirements=AttentionRequirements(packed_varlen=True),
+        )
 
     def select_vae_weight_files(
         self,

python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py

@@ -16,6 +16,13 @@ from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 
+@dataclass(frozen=True)
+class AttentionRequirements:
+    """Semantic attention operations required by a caller."""
+
+    packed_varlen: bool = False
+
+
 class AttentionBackend(ABC):
     """Abstract class for attention backends."""
 
@@ -34,6 +41,18 @@ class AttentionBackend(ABC):
     def get_impl_cls() -> type["AttentionImpl"]:
         raise NotImplementedError
 
+    @classmethod
+    def supports_packed_varlen(cls) -> bool:
+        return cls.get_impl_cls().forward_varlen is not AttentionImpl.forward_varlen
+
+    @classmethod
+    def unsupported_requirements(
+        cls, requirements: AttentionRequirements
+    ) -> tuple[str, ...]:
+        if requirements.packed_varlen and not cls.supports_packed_varlen():
+            return ("packed varlen attention",)
+        return ()
+
     @staticmethod
     @abstractmethod
     def get_metadata_cls() -> type["AttentionMetadata"]:

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

@@ -14,6 +14,7 @@ import torch
 
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionBackend,
+    AttentionRequirements,
 )
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.server_args import get_global_server_args
@@ -151,6 +152,7 @@ def get_attn_backend(
     dtype: torch.dtype,
     supported_attention_backends: set[AttentionBackendEnum] | None = None,
     selected_attention_backend: AttentionBackendEnum | None = None,
+    attention_requirements: AttentionRequirements | None = None,
 ) -> type[AttentionBackend]:
     if supported_attention_backends is None:
         be_tuple = tuple()
@@ -188,6 +190,14 @@ def get_attn_backend(
     )
 
     backend_name = attention_backend_cls.get_enum().name.lower()
+    unsupported_requirements = attention_backend_cls.unsupported_requirements(
+        attention_requirements or AttentionRequirements()
+    )
+    if unsupported_requirements:
+        raise ValueError(
+            f"Attention backend '{backend_name}' does not implement "
+            f"{', '.join(unsupported_requirements)}"
+        )
     reason = "component constraint" if backend_name == constraint_backend else None
     if not _record_component_attn_backend(backend_name, reason):
         logger.info_once(f"Using {backend_name} attention backend")
@@ -218,12 +228,10 @@ def _cached_get_attn_backend(
             supported_attention_backend.__str__()
             for supported_attention_backend in supported_attention_backends
         ]
-        logger.debug(
-            "Selected attention backend: '%s' not in supported attention backends: %s",
-            selected_backend,
-            supported_attention_backends_str,
+        raise ValueError(
+            f"Attention backend '{selected_backend}' is not supported by this "
+            f"attention layer; supported backends: {supported_attention_backends_str}"
         )
-        selected_backend = None
 
     attention_cls = current_platform.get_attn_backend_cls_str(
         selected_backend, head_size, dtype

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -36,6 +36,9 @@ from sglang.multimodal_gen.runtime.distributed import (
     get_tp_world_size,
     tensor_model_parallel_all_gather,
 )
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionRequirements,
+)
 from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
@@ -485,7 +488,7 @@ def _minimax_h3_attention_core_impl(
             get_attn_backend(
                 attention.head_dim,
                 q.dtype,
-                supported_attention_backends=attention._supported_attention_backends,
+                attention_requirements=AttentionRequirements(packed_varlen=True),
             )
         )
     out = attention._attention_impl.forward_varlen(
@@ -527,7 +530,6 @@ class MiniMaxH3Attention(nn.Module):
         self.inner_dim = self.total_num_heads * self.head_dim
         self.local_inner_dim = self.num_heads * self.head_dim
         self.softmax_scale = self.head_dim**-0.5
-        self._supported_attention_backends = arch._supported_attention_backends
         self._attention_impl = None
         # The checkpoint stores one fused qkv tensor. Each logical Q/K/V
         # matrix must be sharded independently; a plain ColumnParallelLinear
@@ -1015,7 +1017,6 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
     # heads) with bf16 blocks; FSDP must gather in each parameter's own dtype
     _fsdp_mixed_dtype_params = True
     _compile_conditions = _ARCH_DEFAULTS._compile_conditions
-    _supported_attention_backends = _ARCH_DEFAULTS._supported_attention_backends
     param_names_mapping = _ARCH_DEFAULTS.param_names_mapping
     reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping
     lora_param_names_mapping = _ARCH_DEFAULTS.lora_param_names_mapping
@@ -1178,7 +1179,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
         backend = get_attn_backend(
             self.arch.attention_head_dim,
             _BF16_DTYPE,
-            supported_attention_backends=self._supported_attention_backends,
+            attention_requirements=AttentionRequirements(packed_varlen=True),
         )
         for module in self.modules():
             if isinstance(module, MiniMaxH3Attention):

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -897,7 +897,6 @@ TWO_GPU_CASES = [
             cfg_parallel=True,
             extras=[
                 "--pipeline-class-name LTX2TwoStagePipeline",
-                "--component-attention-backends transformer=fa",
             ],
         ),
         DiffusionSamplingParams(prompt=T2V_PROMPT, extras={"seed": 42}),

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

@@ -3,6 +3,9 @@ from unittest.mock import patch
 
 import torch
 
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    _cached_get_attn_backend,
+)
 from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatformBase
 from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
 
@@ -38,6 +41,7 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
         FakeCudaPlatform.is_sm120_device = False
         FakeCudaPlatform.is_blackwell_device = False
         FakeCudaPlatform.supports_flash_attention = True
+        _cached_get_attn_backend.cache_clear()
 
     def resolve(
         self,
@@ -108,6 +112,17 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "Invalid attention backend"):
             self.resolve(AttentionBackendEnum.AITER_SAGE)
 
+    def test_explicit_backend_rejected_by_a_model_fails_closed(self):
+        with self.assertRaisesRegex(
+            ValueError, "not supported by this attention layer"
+        ):
+            _cached_get_attn_backend(
+                128,
+                torch.float16,
+                (AttentionBackendEnum.FA,),
+                AttentionBackendEnum.SAGE_ATTN,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

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

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
     MiniMaxH3PipelineConfig,
@@ -15,6 +16,9 @@ from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingPar
 from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
     VideoGenerationsRequest,
 )
+from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
+    AttentionRequirements,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
     MiniMaxH3PartitionAdmissionStage,
     MiniMaxH3ReleaseMetadata,
@@ -28,7 +32,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.task_profiles import (
     partition_for_task,
 )
-from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.server_args.server_args import Backend
 
 TARGET = {
@@ -307,3 +314,30 @@ def test_quality_admission_fails_closed_outside_validated_request():
     server_args.attention_backend = None
     with pytest.raises(ValueError, match="quality must be one of"):
         stage.forward(batch, server_args)
+
+
+def test_validate_server_args_requires_packed_varlen_backend():
+    config = SimpleNamespace(
+        vae_config=SimpleNamespace(resolved_parallel_decode_mode=lambda: None),
+        dit_config=SimpleNamespace(arch_config=SimpleNamespace(attention_head_dim=128)),
+        _server_arg_value=MiniMaxH3PipelineConfig._server_arg_value,
+    )
+    server_args = SimpleNamespace(
+        component_attention_backends={}, attention_backend="sage_attn"
+    )
+    with patch(
+        "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3.get_attn_backend"
+    ) as get_attn_backend:
+        MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
+    get_attn_backend.assert_called_once_with(
+        128,
+        torch.bfloat16,
+        selected_attention_backend=AttentionBackendEnum.SAGE_ATTN,
+        attention_requirements=AttentionRequirements(packed_varlen=True),
+    )
+    with patch(
+        "sglang.multimodal_gen.configs.pipeline_configs.minimax_h3.get_attn_backend",
+        side_effect=ValueError("does not implement packed varlen attention"),
+    ):
+        with pytest.raises(ValueError, match="does not implement packed varlen"):
+            MiniMaxH3PipelineConfig.validate_server_args(config, server_args)