minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] refactor: move dit execution capabilities to runtime models (#34249)

Browse Source
This commit is contained in:
Mick
2026-08-11 18:20:43 +08:00
committed by GitHub
parent 071f0f1e9d
commit dd8c5849af
48 changed files with 191 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/base.py
+10 -8
View File
@@ -9,11 +9,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
# 2. ArchConfig should be inherited & overridden by each model arch_config
# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
class ArchConfig:
"""Static model metadata loaded from a Diffusers/Transformers config.
This includes architecture fields and checkpoint compatibility mappings.
Runtime capabilities, backend selection, hardware policies, and deployment
settings belong to the runtime model or server configuration instead.
"""
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=list
) # mapping from huggingface weight names to custom names
@@ -42,8 +46,8 @@ class ArchConfig:
@dataclass
class ModelConfig:
# Every model config parameter can be categorized into either ArchConfig or everything else
# Diffuser/Transformer parameters
# Static component metadata; this is intentionally separate from runtime
# capabilities and deployment settings.
arch_config: ArchConfig = field(default_factory=ArchConfig)
# sglang-diffusion-specific parameters here
@@ -69,9 +73,7 @@ class ModelConfig:
# This should be used only when loading from transformers/diffusers
def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
"""
Update arch_config with source_model_dict
"""
"""Load static architecture metadata from a source model config."""
arch_config = self.arch_config
for key, value in source_model_dict.items():
python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py
-9
View File
@@ -6,17 +6,8 @@ from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
def _is_conditioner_block(name: str, module) -> bool:
"""Check if module is a ConditionalCrossAttentionBlock."""
return "ConditionalCrossAttentionBlock" in type(module).__name__
@dataclass
class MOVADualTowerArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [_is_conditioner_block]
)
# Model architecture parameters
visual_layers: int = 40
audio_layers: int = 30
python/sglang/multimodal_gen/configs/models/dits/base.py
+7 -23
View File
@@ -6,13 +6,16 @@ from typing import Any
from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@dataclass
class DiTArchConfig(ArchConfig):
_fsdp_shard_conditions: list = field(default_factory=list)
_compile_conditions: list = field(default_factory=list)
"""Static DiT architecture metadata and checkpoint mappings.
This object is populated from the source model configuration. It must not
contain runtime implementation capabilities or the backend selected for a
particular deployment.
"""
# convert weights name from HF-format to SGLang-dit-format
param_names_mapping: dict = field(default_factory=dict)
@@ -23,24 +26,6 @@ class DiTArchConfig(ArchConfig):
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=dict)
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
AttentionBackendEnum.SLIDING_TILE_ATTN,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.FA,
AttentionBackendEnum.AITER,
AttentionBackendEnum.AITER_SAGE,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
AttentionBackendEnum.VMOBA_ATTN,
AttentionBackendEnum.SAGE_ATTN_3,
AttentionBackendEnum.LASER_ATTN,
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
AttentionBackendEnum.RAIN_FUSION_ATTN,
}
)
hidden_size: int = 0
num_attention_heads: int = 0
num_channels_latents: int = 0
@@ -48,8 +33,7 @@ class DiTArchConfig(ArchConfig):
boundary_ratio: float | None = None
def __post_init__(self) -> None:
if not self._compile_conditions:
self._compile_conditions = self._fsdp_shard_conditions.copy()
pass
@dataclass
python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py
-7
View File
@@ -4,11 +4,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_module_list_entry_in
def is_layers(n: str, m) -> bool:
return is_module_list_entry_in(n, ("layers", "gen_layers"))
def _build_cosmos3_param_names_mapping() -> dict:
@@ -128,8 +123,6 @@ class Cosmos3VideoArchConfig(DiTArchConfig):
Qwen3-8B-Instruct-derived and overridden by the checkpoint at load time.
"""
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_layers])
# Transformer architecture
hidden_size: int = 4096
num_hidden_layers: int = 36
python/sglang/multimodal_gen/configs/models/dits/ernie_image.py
-3
View File
@@ -4,7 +4,6 @@ from dataclasses import dataclass, field
from typing import Tuple
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_layer
@dataclass
@@ -31,8 +30,6 @@ class ErnieImageArchConfig(DiTArchConfig):
}
)
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
def __post_init__(self):
super().__post_init__()
self.hidden_size = self.num_attention_heads * self.attention_head_dim
python/sglang/multimodal_gen/configs/models/dits/helios.py
-3
View File
@@ -2,13 +2,10 @@
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
@dataclass
class HeliosArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
param_names_mapping: dict = field(
default_factory=lambda: {
# Patch embeddings
python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py
-14
View File
@@ -6,24 +6,10 @@ from dataclasses import dataclass, field
import torch
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import (
is_double_block,
is_refiner_block,
is_single_block,
is_txt_in,
)
@dataclass
class HunyuanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_double_block, is_single_block, is_refiner_block]
)
_compile_conditions: list = field(
default_factory=lambda: [is_double_block, is_single_block, is_txt_in]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# 1. context_embedder.time_text_embed submodules (specific rules, applied first):
python/sglang/multimodal_gen/configs/models/dits/ideogram.py
-9
View File
@@ -3,8 +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_layer
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
@dataclass
@@ -39,13 +37,6 @@ class Ideogram4DiTArchConfig(DiTArchConfig):
r"^(layers\.\d+\.attention)\.to_out\.0\.(.*)$": r"\1.o.\2",
}
)
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
}
)
def __post_init__(self) -> None:
super().__post_init__()
python/sglang/multimodal_gen/configs/models/dits/joy_image.py
-5
View File
@@ -3,15 +3,10 @@
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_blocks_or_double_blocks
@dataclass
class JoyImageArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_blocks_or_double_blocks]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# Condition embedder mappings
python/sglang/multimodal_gen/configs/models/dits/lingbot_video_moe.py
-6
View File
@@ -4,14 +4,8 @@ from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
def is_blocks(name: str, module) -> bool:
return "blocks" in name and str.isdigit(name.split(".")[-1])
@dataclass
class LingBotVideoMoEArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
param_names_mapping: dict = field(default_factory=dict)
reverse_param_names_mapping: dict = field(default_factory=dict)
lora_param_names_mapping: dict = field(default_factory=dict)
python/sglang/multimodal_gen/configs/models/dits/lingbot_world.py
-6
View File
@@ -7,14 +7,8 @@ from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
def is_blocks(n: str, m) -> bool:
return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
class LingBotWorldArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
python/sglang/multimodal_gen/configs/models/dits/ltx_2.py
-5
View File
@@ -3,7 +3,6 @@ from dataclasses import dataclass, field
from enum import Enum
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_transformer_blocks
class LTXModelType(Enum):
@@ -52,10 +51,6 @@ class LTX2AttentionFunction(str, Enum):
class LTX2ArchConfig(DiTArchConfig):
"""Architecture configuration for LTX-2 Video Transformer."""
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_blocks_or_transformer_blocks]
)
param_names_mapping: dict = field(
default_factory=lambda: {
# Parameter name mappings from HuggingFace checkpoint keys to SGLang module names.
python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py
-3
View File
@@ -2,7 +2,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
MINIMAX_H3_PACKED_SEQUENCE_ALIGNMENT = 64
MINIMAX_H3_ADALN_MODALITY_NUM = 3
@@ -10,8 +9,6 @@ MINIMAX_H3_ADALN_MODALITY_NUM = 3
@dataclass
class MiniMaxH3DiTArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
lora_param_names_mapping: dict = field(default_factory=dict)
num_layers: int = 50
python/sglang/multimodal_gen/configs/models/dits/mova_audio.py
-3
View File
@@ -4,13 +4,10 @@
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
@dataclass
class MOVAAudioArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
python/sglang/multimodal_gen/configs/models/dits/mova_video.py
-3
View File
@@ -4,13 +4,10 @@
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
@dataclass
class MOVAVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
python/sglang/multimodal_gen/configs/models/dits/qwenimage.py
-3
View File
@@ -5,7 +5,6 @@ from dataclasses import dataclass, field
from typing import Tuple
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
@dataclass
@@ -23,8 +22,6 @@ class QwenImageArchConfig(DiTArchConfig):
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
zero_cond_t: bool = False
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_block])
stacked_params_mapping: list[tuple[str, str, str]] = field(default_factory=list)
param_names_mapping: dict = field(
python/sglang/multimodal_gen/configs/models/dits/sana_wm.py
-5
View File
@@ -3,15 +3,10 @@
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_blocks_or_transformer_blocks
@dataclass
class SanaWMArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_blocks_or_transformer_blocks]
)
# --- Core dims (upstream: depth=20, hidden=2240, heads=20, linear_head_dim=112) ---
patch_size: int = 1
in_channels: int = 128 # LTX-2 VAE latent channels
python/sglang/multimodal_gen/configs/models/dits/sana_wm_refiner.py
-5
View File
@@ -3,15 +3,10 @@
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_blocks_or_transformer_blocks
@dataclass
class SanaWMRefinerArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(
default_factory=lambda: [is_blocks_or_transformer_blocks]
)
# Core dims
in_channels: int = 128
out_channels: int = 128
python/sglang/multimodal_gen/configs/models/dits/wanvideo.py
-3
View File
@@ -4,13 +4,10 @@
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
@dataclass
class WanVideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_block])
param_names_mapping: dict = field(
default_factory=lambda: {
r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
python/sglang/multimodal_gen/configs/models/dits/zimage.py
-3
View File
@@ -5,7 +5,6 @@ from dataclasses import dataclass, field
from typing import Tuple
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_zimage_layer
@dataclass
@@ -27,8 +26,6 @@ class ZImageArchConfig(DiTArchConfig):
axes_dims: Tuple[int, int, int] = (32, 48, 48)
axes_lens: Tuple[int, int, int] = (1024, 512, 512)
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_zimage_layer])
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)
python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py
+6 -3
View File
@@ -397,6 +397,10 @@ class ConditionalCrossAttentionBlock(nn.Module):
return self.inner(x=x, y=y, x_freqs=x_freqs, y_freqs=y_freqs)
def _is_conditioner_block(_name: str, module: nn.Module) -> bool:
return isinstance(module, ConditionalCrossAttentionBlock)
class DualTowerConditionalBridge(
CachableDiT,
LayerwiseOffloadableModuleMixin,
@@ -411,9 +415,8 @@ class DualTowerConditionalBridge(
layerwise_offload_dit_group_enabled = False
_fsdp_shard_conditions = MOVADualTowerConfig()._fsdp_shard_conditions
_compile_conditions = MOVADualTowerConfig()._compile_conditions
_supported_attention_backends = MOVADualTowerConfig()._supported_attention_backends
_fsdp_shard_conditions = [_is_conditioner_block]
_compile_conditions = [_is_conditioner_block]
param_names_mapping = MOVADualTowerConfig().param_names_mapping
reverse_param_names_mapping = MOVADualTowerConfig().reverse_param_names_mapping
lora_param_names_mapping = MOVADualTowerConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/base.py
+20 -10
View File
@@ -22,6 +22,9 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
# TODO
class BaseDiT(nn.Module, ABC):
# These are runtime implementation capabilities, not checkpoint metadata.
# Concrete DiT implementations override them when their tensor layout or
# execution semantics support only a subset of the available backends.
_fsdp_shard_conditions: list = []
_compile_conditions: list = []
param_names_mapping: dict
@@ -29,10 +32,21 @@ class BaseDiT(nn.Module, ABC):
hidden_size: int
num_attention_heads: int
num_channels_latents: int
# always supports torch_sdpa
_supported_attention_backends: set[AttentionBackendEnum] = (
DiTConfig()._supported_attention_backends
)
_supported_attention_backends: set[AttentionBackendEnum] = {
AttentionBackendEnum.SLIDING_TILE_ATTN,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.FA,
AttentionBackendEnum.AITER,
AttentionBackendEnum.AITER_SAGE,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.VIDEO_SPARSE_ATTN,
AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN,
AttentionBackendEnum.VMOBA_ATTN,
AttentionBackendEnum.SAGE_ATTN_3,
AttentionBackendEnum.LASER_ATTN,
AttentionBackendEnum.BLOCK_SPARSE_ATTN,
AttentionBackendEnum.RAIN_FUSION_ATTN,
}
def __init_subclass__(cls) -> None:
required_class_attrs = [
@@ -49,8 +63,8 @@ class BaseDiT(nn.Module, ABC):
def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
super().__init__()
# runtime models expose checkpoint architecture through `config`; load
# settings such as the model prefix stay separate
# `config.arch_config` contains static model metadata. Runtime
# capabilities remain class attributes on the model implementation.
self.config: DiTArchConfig = config.arch_config
self.prefix = config.prefix
self.hf_config = hf_config
@@ -111,10 +125,6 @@ class CachableDiT(SpectrumMixin, TeaCacheMixin, BaseDiT):
hidden_size: int
num_attention_heads: int
num_channels_latents: int
# always supports torch_sdpa
_supported_attention_backends: set[AttentionBackendEnum] = (
DiTConfig()._supported_attention_backends
)
def __init__(self, config: DiTConfig, **kwargs) -> None:
super().__init__(config, **kwargs)
python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
+3 -3
View File
@@ -26,6 +26,7 @@ flex_attention = torch.compile(
import torch.distributed as dist
from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_sp_world_size,
@@ -439,9 +440,8 @@ class CausalWanTransformerBlock(nn.Module):
class CausalWanTransformer3DModel(BaseDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
_compile_conditions = WanVideoConfig()._compile_conditions
_supported_attention_backends = WanVideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = WanVideoConfig().param_names_mapping
reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py
+8 -4
View File
@@ -15,6 +15,7 @@ import torch.nn as nn
import torch.nn.functional as F
from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoConfig
from sglang.multimodal_gen.configs.models.fsdp import is_module_list_entry_in
from sglang.multimodal_gen.runtime.distributed import (
get_sp_group,
get_sp_world_size,
@@ -55,6 +56,10 @@ from sglang.srt.utils import add_prefix
logger = init_logger(__name__)
def is_cosmos_layer(name: str, _module: object) -> bool:
return is_module_list_entry_in(name, ("layers", "gen_layers"))
# -----------------------------------------------------------------------------
# mRoPE position ID computation (Qwen3VL-style)
# -----------------------------------------------------------------------------
@@ -916,9 +921,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
- Generation (GEN): cross-attention from visual to UND K/V
"""
_fsdp_shard_conditions = Cosmos3VideoConfig()._fsdp_shard_conditions
_compile_conditions = Cosmos3VideoConfig()._compile_conditions
_supported_attention_backends = Cosmos3VideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_cosmos_layer]
_compile_conditions = [is_cosmos_layer]
param_names_mapping = Cosmos3VideoConfig().arch_config.param_names_mapping
reverse_param_names_mapping = (
Cosmos3VideoConfig().arch_config.reverse_param_names_mapping
@@ -1050,7 +1054,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
layer_idx=i,
prefix=f"gen_layers.{i}",
quant_config=quant_config,
supported_attention_backends=arch._supported_attention_backends,
supported_attention_backends=self._supported_attention_backends,
)
for i in range(arch.num_hidden_layers)
]
python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py
+2 -1
View File
@@ -40,6 +40,7 @@ from sglang.kernels.ops.diffusion.triton.rope_rotate_half_bitexact import (
from sglang.multimodal_gen.configs.models.dits.ernie_image import (
ErnieImageDitConfig,
)
from sglang.multimodal_gen.configs.models.fsdp import is_layer
from sglang.multimodal_gen.runtime.distributed import (
get_tp_world_size,
)
@@ -515,7 +516,7 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin)
_no_split_modules = ["ErnieImageSharedAdaLNBlock"]
_skip_layerwise_casting_patterns = ["pos_embed", "norm"]
_fsdp_shard_conditions = ErnieImageDitConfig().arch_config._fsdp_shard_conditions
_fsdp_shard_conditions = [is_layer]
_compile_conditions = []
param_names_mapping = ErnieImageDitConfig().arch_config.param_names_mapping
reverse_param_names_mapping = {}
python/sglang/multimodal_gen/runtime/models/dits/glm_image.py
-1
View File
@@ -987,7 +987,6 @@ class GlmImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
)
# 3. Transformer blocks
self._supported_attention_backends = arch_config._supported_attention_backends
self.transformer_blocks = nn.ModuleList(
[
GlmImageTransformerBlock(
python/sglang/multimodal_gen/runtime/models/dits/helios.py
+3 -3
View File
@@ -17,6 +17,7 @@ import torch.nn as nn
import torch.nn.functional as F
from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_sp_world_size,
@@ -559,9 +560,8 @@ class HeliosTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
with zero_history_timestep and guidance_cross_attn.
"""
_fsdp_shard_conditions = HeliosConfig()._fsdp_shard_conditions
_compile_conditions = HeliosConfig()._compile_conditions
_supported_attention_backends = HeliosConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = HeliosConfig().param_names_mapping
reverse_param_names_mapping = HeliosConfig().reverse_param_names_mapping
lora_param_names_mapping = HeliosConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py
+1 -1
View File
@@ -501,7 +501,7 @@ class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin):
qkv_bias = arch.qkv_bias
time_factor = arch.time_factor
guidance_embed = arch.guidance_embed
supported_attention_backends = arch._supported_attention_backends
supported_attention_backends = self._supported_attention_backends
self.in_channels = in_channels
self.context_in_dim = context_in_dim
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
+8 -3
View File
@@ -9,6 +9,12 @@ import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.fsdp import (
is_double_block,
is_refiner_block,
is_single_block,
is_txt_in,
)
from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams
from sglang.multimodal_gen.runtime.distributed import divide, get_tp_world_size
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
@@ -508,9 +514,8 @@ class HunyuanVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi
# PY: we make the input args the same as HF config
# shard single stream, double stream blocks, and refiner_blocks
_fsdp_shard_conditions = HunyuanVideoConfig()._fsdp_shard_conditions
_compile_conditions = HunyuanVideoConfig()._compile_conditions
_supported_attention_backends = HunyuanVideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_double_block, is_single_block, is_refiner_block]
_compile_conditions = [is_double_block, is_single_block, is_txt_in]
param_names_mapping = HunyuanVideoConfig().param_names_mapping
reverse_param_names_mapping = HunyuanVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = HunyuanVideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/ideogram.py
+8 -6
View File
@@ -14,6 +14,7 @@ from sglang.kernels.ops.diffusion.fused_gate_rmsnorm import (
mark_fused_gate_rmsnorm_site,
)
from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_layer
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_tp_world_size,
@@ -46,6 +47,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
OUTPUT_IMAGE_INDICATOR = 2
LLM_TOKEN_INDICATOR = 3
@@ -497,11 +499,12 @@ class Ideogram4FinalLayer(nn.Module):
class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin):
_repeated_blocks = ["Ideogram4TransformerBlock"]
layer_names = ["layers"]
_fsdp_shard_conditions = Ideogram4DiTConfig().arch_config._fsdp_shard_conditions
_compile_conditions = Ideogram4DiTConfig().arch_config._compile_conditions
_supported_attention_backends = (
Ideogram4DiTConfig().arch_config._supported_attention_backends
)
_fsdp_shard_conditions = [is_layer]
_compile_conditions = [is_layer]
_supported_attention_backends = {
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
}
param_names_mapping = Ideogram4DiTConfig().arch_config.param_names_mapping
reverse_param_names_mapping = {}
@@ -515,7 +518,6 @@ class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin):
super().__init__(config, hf_config, **kwargs)
cfg = self.config
use_weight_only_fp8_linears = config.use_weight_only_fp8_linears
self._supported_attention_backends = cfg._supported_attention_backends
hidden_size = cfg.num_attention_heads * cfg.attention_head_dim
self.hidden_size = hidden_size
self.num_attention_heads = cfg.num_attention_heads
python/sglang/multimodal_gen/runtime/models/dits/joy_image.py
+3 -3
View File
@@ -9,6 +9,7 @@ import torch.nn as nn
from einops import rearrange
from sglang.multimodal_gen.configs.models.dits.joy_image import JoyImageDiTConfig
from sglang.multimodal_gen.configs.models.fsdp import is_blocks_or_double_blocks
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_sp_group,
@@ -348,9 +349,8 @@ class JoyTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
"""
_supports_gradient_checkpointing = True
_fsdp_shard_conditions = JoyImageDiTConfig()._fsdp_shard_conditions
_compile_conditions = JoyImageDiTConfig()._compile_conditions
_supported_attention_backends = JoyImageDiTConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_blocks_or_double_blocks]
_compile_conditions = [is_blocks_or_double_blocks]
param_names_mapping = JoyImageDiTConfig().param_names_mapping
reverse_param_names_mapping = JoyImageDiTConfig().reverse_param_names_mapping
lora_param_names_mapping = JoyImageDiTConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/lingbot_video_moe.py
+6 -5
View File
@@ -60,6 +60,10 @@ LINGBOT_VIDEO_FP32_MODULES = (
)
def is_lingbot_block(name: str, _module: object) -> bool:
return "blocks" in name and name.split(".")[-1].isdigit()
def should_keep_in_fp32(name: str) -> bool:
return any(
module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES
@@ -344,11 +348,8 @@ class LingBotVideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixi
_no_split_modules = ("LingBotVideoBlock",)
_keep_in_fp32_modules = tuple(LINGBOT_VIDEO_FP32_MODULES)
_fsdp_shard_conditions = LingBotVideoMoEConfig()._fsdp_shard_conditions
_compile_conditions = LingBotVideoMoEConfig()._compile_conditions
_supported_attention_backends = (
LingBotVideoMoEConfig()._supported_attention_backends
)
_fsdp_shard_conditions = [is_lingbot_block]
_compile_conditions = [is_lingbot_block]
param_names_mapping = LingBotVideoMoEConfig().param_names_mapping
reverse_param_names_mapping = LingBotVideoMoEConfig().reverse_param_names_mapping
lora_param_names_mapping = LingBotVideoMoEConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py
+10 -10
View File
@@ -97,6 +97,12 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.srt.utils import add_prefix
logger = init_logger(__name__)
def is_lingbot_block(name: str, _module: object) -> bool:
return "blocks" in name and name.split(".")[-1].isdigit()
_is_cuda = current_platform.is_cuda()
@@ -574,11 +580,8 @@ class LingBotWorldTransformerBlock(nn.Module):
class LingBotWorldTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions
_compile_conditions = LingBotWorldVideoConfig()._compile_conditions
_supported_attention_backends = (
LingBotWorldVideoConfig()._supported_attention_backends
)
_fsdp_shard_conditions = [is_lingbot_block]
_compile_conditions = [is_lingbot_block]
param_names_mapping = LingBotWorldVideoConfig().param_names_mapping
reverse_param_names_mapping = LingBotWorldVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = LingBotWorldVideoConfig().lora_param_names_mapping
@@ -1133,11 +1136,8 @@ class CausalLingBotWorldTransformerBlock(CausalWanTransformerBlock):
class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel):
_fsdp_shard_conditions = LingBotWorldVideoConfig()._fsdp_shard_conditions
_compile_conditions = LingBotWorldVideoConfig()._compile_conditions
_supported_attention_backends = (
LingBotWorldVideoConfig()._supported_attention_backends
)
_fsdp_shard_conditions = [is_lingbot_block]
_compile_conditions = [is_lingbot_block]
param_names_mapping = LingBotWorldVideoConfig().param_names_mapping
reverse_param_names_mapping = LingBotWorldVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = LingBotWorldVideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/longlive2.py
+3 -3
View File
@@ -6,6 +6,7 @@ import torch.nn as nn
from torch.nn.attention.flex_attention import BlockMask
from sglang.multimodal_gen.configs.models.dits.longlive2 import LongLive2VideoConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.layers.kvcache.causal_attention_cache import (
CausalSelfAttentionKVCache,
CrossAttentionKVCache,
@@ -151,9 +152,8 @@ class LongLive2CausalWanTransformerBlock(CausalWanTransformerBlock):
class LongLive2Transformer3DModel(CausalWanTransformer3DModel):
_fsdp_shard_conditions = LongLive2VideoConfig()._fsdp_shard_conditions
_compile_conditions = LongLive2VideoConfig()._compile_conditions
_supported_attention_backends = LongLive2VideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = LongLive2VideoConfig().param_names_mapping
reverse_param_names_mapping = LongLive2VideoConfig().reverse_param_names_mapping
lora_param_names_mapping = LongLive2VideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+5 -3
View File
@@ -28,6 +28,9 @@ from sglang.kernels.ops.diffusion.ltx2_rmsnorm_modulate import (
)
from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config
from sglang.multimodal_gen.configs.models.fsdp import (
is_blocks_or_transformer_blocks,
)
from sglang.multimodal_gen.runtime.distributed import (
get_sp_parallel_rank,
get_sp_world_size,
@@ -1504,9 +1507,8 @@ class LTX2TransformerBlock(nn.Module):
class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = LTX2ArchConfig()._fsdp_shard_conditions
_compile_conditions = LTX2ArchConfig()._compile_conditions
_supported_attention_backends = LTX2ArchConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_blocks_or_transformer_blocks]
_compile_conditions = [is_blocks_or_transformer_blocks]
param_names_mapping = LTX2ArchConfig().param_names_mapping
reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping
lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+3 -2
View File
@@ -32,6 +32,7 @@ from sglang.multimodal_gen.configs.models.dits.minimax_h3 import (
MiniMaxH3DiTArchConfig,
MiniMaxH3DiTConfig,
)
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.distributed import (
get_tp_world_size,
tensor_model_parallel_all_gather,
@@ -1024,11 +1025,11 @@ class MiniMaxH3FinalLayer(nn.Module):
class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = _ARCH_DEFAULTS._fsdp_shard_conditions
_fsdp_shard_conditions = [is_block]
# parameters mix fp32 (patch projections, timestep embedder, and output
# heads) with bf16 blocks; FSDP must gather in each parameter's own dtype
_fsdp_mixed_dtype_params = True
_compile_conditions = _ARCH_DEFAULTS._compile_conditions
_compile_conditions = [is_block]
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
python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py
+3 -3
View File
@@ -13,6 +13,7 @@ from einops import rearrange
from torch.distributed.tensor import DTensor
from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
@@ -87,9 +88,8 @@ class Conv1dLocalIsland(nn.Conv1d):
class WanAudioModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = MOVAAudioConfig()._fsdp_shard_conditions
_compile_conditions = MOVAAudioConfig()._compile_conditions
_supported_attention_backends = MOVAAudioConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = MOVAAudioConfig().param_names_mapping
reverse_param_names_mapping = MOVAAudioConfig().reverse_param_names_mapping
lora_param_names_mapping = MOVAAudioConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py
+3 -3
View File
@@ -14,6 +14,7 @@ from einops import rearrange
from torch.distributed.tensor import DTensor
from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
@@ -417,9 +418,8 @@ class Conv3dLocalIsland(nn.Conv3d):
class WanModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = MOVAVideoConfig()._fsdp_shard_conditions
_compile_conditions = MOVAVideoConfig()._compile_conditions
_supported_attention_backends = MOVAVideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = MOVAVideoConfig().param_names_mapping
reverse_param_names_mapping = MOVAVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = MOVAVideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
+2 -1
View File
@@ -21,6 +21,7 @@ from sglang.kernels.ops.diffusion.fused_linear_gelu import (
mark_fused_gelu_site,
)
from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
from sglang.multimodal_gen.runtime.distributed import (
get_local_torch_device,
get_tp_world_size,
@@ -1315,7 +1316,7 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_repeated_blocks = ["QwenImageTransformerBlock"]
param_names_mapping = QwenImageDitConfig().arch_config.param_names_mapping
_fsdp_shard_conditions = QwenImageDitConfig().arch_config._fsdp_shard_conditions
_fsdp_shard_conditions = [is_transformer_block]
@classmethod
def get_nunchaku_quant_rules(cls) -> dict[str, list[str]]:
python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py
+5 -3
View File
@@ -7,6 +7,9 @@ import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits.sana_wm import SanaWMConfig
from sglang.multimodal_gen.configs.models.fsdp import (
is_blocks_or_transformer_blocks,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin,
)
@@ -344,9 +347,8 @@ class SanaWMTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
Returns: ``(B, C, T, H, W)`` predicted velocity / noise.
"""
_fsdp_shard_conditions = SanaWMConfig()._fsdp_shard_conditions
_compile_conditions = SanaWMConfig()._compile_conditions
_supported_attention_backends = SanaWMConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_blocks_or_transformer_blocks]
_compile_conditions = [is_blocks_or_transformer_blocks]
param_names_mapping = SanaWMConfig().param_names_mapping
reverse_param_names_mapping = SanaWMConfig().reverse_param_names_mapping
lora_param_names_mapping: dict = {}
python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py
+5 -5
View File
@@ -11,6 +11,9 @@ from sglang.multimodal_gen.configs.models.dits.sana_wm_refiner import (
SanaWMRefinerArchConfig,
SanaWMRefinerConfig,
)
from sglang.multimodal_gen.configs.models.fsdp import (
is_blocks_or_transformer_blocks,
)
from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
@@ -210,11 +213,8 @@ class SanaWMLTX2VideoRefiner(CachableDiT, LayerwiseOffloadableModuleMixin):
`strict=False` state_dict load.
"""
_fsdp_shard_conditions = SanaWMRefinerArchConfig()._fsdp_shard_conditions
_compile_conditions = SanaWMRefinerArchConfig()._compile_conditions
_supported_attention_backends = (
SanaWMRefinerArchConfig()._supported_attention_backends
)
_fsdp_shard_conditions = [is_blocks_or_transformer_blocks]
_compile_conditions = [is_blocks_or_transformer_blocks]
param_names_mapping = SanaWMRefinerArchConfig().param_names_mapping
reverse_param_names_mapping: dict = {}
lora_param_names_mapping: dict = {}
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py
+3 -3
View File
@@ -10,6 +10,7 @@ import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
from sglang.multimodal_gen.configs.models.fsdp import is_block
from sglang.multimodal_gen.runtime.distributed import (
divide,
get_sp_group,
@@ -857,9 +858,8 @@ class WanTransformerBlock_VSA(nn.Module):
class WanTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
_compile_conditions = WanVideoConfig()._compile_conditions
_supported_attention_backends = WanVideoConfig()._supported_attention_backends
_fsdp_shard_conditions = [is_block]
_compile_conditions = [is_block]
param_names_mapping = WanVideoConfig().param_names_mapping
reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+2 -1
View File
@@ -5,6 +5,7 @@ import torch
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
from sglang.multimodal_gen.configs.models.fsdp import is_zimage_layer
from sglang.multimodal_gen.runtime.distributed import (
get_sp_world_size,
get_tp_world_size,
@@ -768,7 +769,7 @@ class RopeEmbedder:
class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
_supports_gradient_checkpointing = True
_no_split_modules = ["ZImageTransformerBlock"]
_fsdp_shard_conditions = ZImageDitConfig().arch_config._fsdp_shard_conditions
_fsdp_shard_conditions = [is_zimage_layer]
param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
reverse_param_names_mapping = (
ZImageDitConfig().arch_config.reverse_param_names_mapping
python/sglang/multimodal_gen/test/server/test_server_common.py
+2 -3
View File
@@ -36,6 +36,7 @@ from sglang.multimodal_gen.test.server.test_server_utils import (
ServerContext,
ServerManager,
get_generate_fn,
is_missing_diffusers_pipeline_error,
)
from sglang.multimodal_gen.test.server.testcase_configs import (
BASELINE_CONFIG,
@@ -213,9 +214,7 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
# pipeline class. This avoids hard failures when a model needs a
# newer diffusers release than what is currently installed in CI.
msg = str(exc)
if "not found in diffusers" in msg or (
"has no attribute" in msg and "diffusers" in msg.lower()
):
if is_missing_diffusers_pipeline_error(msg):
pytest.skip(
f"Skipping {case.id}: required diffusers pipeline class "
f"is not available in the installed version. "
python/sglang/multimodal_gen/test/server/test_server_utils.py
+11
View File
@@ -65,6 +65,17 @@ FIRST_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 80.0
DECODING_STAGE_MIN_ABS_TOLERANCE_MS = 450.0
VIDEO_DENOISE_STEP_MIN_ABS_TOLERANCE_MS = 160.0
def is_missing_diffusers_pipeline_error(message: str) -> bool:
"""Return whether a server startup error is caused by a missing diffusers pipeline."""
normalized_message = message.lower()
return (
"not found in diffusers" in normalized_message
or "module 'diffusers' has no attribute" in normalized_message
or 'module "diffusers" has no attribute' in normalized_message
)
# Tracks mesh output file paths from generate_mesh for later correctness validation.
# Keyed by case_id, cleaned up after use.
MESH_OUTPUT_PATHS: dict[str, str] = {}
python/sglang/multimodal_gen/test/unit/test_dit_config_boundary.py
+42 -1
View File
@@ -1,7 +1,18 @@
from dataclasses import fields
import torch
from sglang.multimodal_gen.configs.models.dits.base import DiTConfig
from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
MOVADualTowerArchConfig,
)
from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
from sglang.multimodal_gen.runtime.models.bridges.mova_dual_tower import (
DualTowerConditionalBridge,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.test.server.test_server_utils import (
is_missing_diffusers_pipeline_error,
)
class _TestDiT(CachableDiT):
@@ -23,3 +34,33 @@ def test_dit_runtime_keeps_architecture_and_component_config_separate():
assert model.prefix == "Wan"
assert model._supports_cfg_cache
assert model._spectrum_supports_cfg_cache
def test_dit_arch_config_excludes_runtime_capabilities():
field_names = {field.name for field in fields(DiTArchConfig)}
assert "_fsdp_shard_conditions" not in field_names
assert "_compile_conditions" not in field_names
assert "_supported_attention_backends" not in field_names
def test_mova_bridge_declares_runtime_capabilities_on_model_class():
assert "_fsdp_shard_conditions" not in {
field.name for field in fields(MOVADualTowerArchConfig)
}
assert DualTowerConditionalBridge._fsdp_shard_conditions
assert DualTowerConditionalBridge._compile_conditions
assert DualTowerConditionalBridge._supported_attention_backends
def test_server_startup_skip_only_matches_missing_diffusers_pipeline():
assert is_missing_diffusers_pipeline_error(
"AttributeError: module 'diffusers' has no attribute 'MOVA'"
)
assert is_missing_diffusers_pipeline_error(
"Pipeline class MOVA not found in diffusers"
)
assert not is_missing_diffusers_pipeline_error(
"AttributeError: MOVADualTowerConfig has no attribute '_compile_conditions'\n"
"Loading pipeline modules from diffusers config"
)
python/sglang/multimodal_gen/test/unit/test_ideogram4.py
+3 -3
View File
@@ -741,12 +741,12 @@ class TestIdeogram4(unittest.TestCase):
["transformer", "unconditional_transformer"],
)
def test_ideogram_attention_backend_is_passed_from_config(self):
def test_ideogram_attention_backend_is_declared_by_runtime_model(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
config = Ideogram4DiTConfig()
self.assertEqual(
config.arch_config._supported_attention_backends,
Ideogram4Transformer2DModel._supported_attention_backends,
{AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
)
prev_args = server_args_module._global_server_args
@@ -763,7 +763,7 @@ class TestIdeogram4(unittest.TestCase):
self.assertEqual(
model.supported_attention_backends,
config.arch_config._supported_attention_backends,
Ideogram4Transformer2DModel._supported_attention_backends,
)
self.assertEqual(
model.layers[0].attention.attn.backend,
python/sglang/multimodal_gen/test/unit/test_regional_torch_compile.py
+4 -2
View File
@@ -4,7 +4,9 @@ from unittest.mock import patch
import pytest
import torch.nn as nn
from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig
from sglang.multimodal_gen.runtime.models.dits.ltx_2 import (
LTX2VideoTransformer3DModel,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
DenoisingStage,
)
@@ -39,7 +41,7 @@ class _RegionalModel(_CompilableModule):
def test_ltx2_compile_conditions_match_only_direct_blocks():
conditions = LTX2ArchConfig()._compile_conditions
conditions = LTX2VideoTransformer3DModel._compile_conditions
assert conditions
assert any(condition("transformer_blocks.0", object()) for condition in conditions)