diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py index 925f08934..0bc7e7892 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py index 2fe52d8d2..afb29a65d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py +++ b/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, 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 e0184a0e4..45736398f 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 @@ -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"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index aeea8ed2d..bed284e56 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 0abc01624..5f5709dd1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/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): diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 01c52acc1..534194e94 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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}), diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py index 4dae9e1ba..6eea8552a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index 181eb496a..8bc4b08dd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/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)