diff --git a/python/sglang/multimodal_gen/runtime/cache/spectrum.py b/python/sglang/multimodal_gen/runtime/cache/spectrum.py index abc03f84f..1bcb9b9de 100644 --- a/python/sglang/multimodal_gen/runtime/cache/spectrum.py +++ b/python/sglang/multimodal_gen/runtime/cache/spectrum.py @@ -331,12 +331,12 @@ class SpectrumMixin: - ``spectrum_record_features()`` — after a real forward, store block outputs. - ``spectrum_predict_features()`` — on skipped steps, return forecasted outputs. - Models with separate CFG branches (Wan, Hunyuan, SD3) list their config prefix + Models with separate CFG branches (Wan, Hunyuan, SD3) list their model prefix in ``_CFG_SUPPORTED_PREFIXES`` so cond/uncond maintain independent counters and forecasters. All other ``CachableDiT`` subclasses share one counter. """ - # DiT config prefixes that run separate cond/uncond forwards (see TeaCache). + # DiT model prefixes that run separate cond/uncond forwards (see TeaCache). _CFG_SUPPORTED_PREFIXES: set[str] = {"wan", "hunyuan", "sd3"} def _init_spectrum_state(self) -> None: @@ -366,8 +366,9 @@ class SpectrumMixin: # Runtime branch tracking self.spectrum_is_cfg_negative = False self._spectrum_ctx: Optional[SpectrumContext] = None - prefix = getattr(self.config, "prefix", "").lower() - self._spectrum_supports_cfg_cache = prefix in self._CFG_SUPPORTED_PREFIXES + self._spectrum_supports_cfg_cache = ( + self.prefix.lower() in self._CFG_SUPPORTED_PREFIXES + ) def reset_spectrum_state(self, spectrum_params: SpectrumParams) -> None: self.spectrum_cnt = 0 diff --git a/python/sglang/multimodal_gen/runtime/cache/teacache.py b/python/sglang/multimodal_gen/runtime/cache/teacache.py index 8830f7ec2..6b71cd9f2 100644 --- a/python/sglang/multimodal_gen/runtime/cache/teacache.py +++ b/python/sglang/multimodal_gen/runtime/cache/teacache.py @@ -23,8 +23,6 @@ from typing import TYPE_CHECKING, Any import numpy as np import torch -from sglang.multimodal_gen.configs.models import DiTConfig - if TYPE_CHECKING: from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @@ -127,7 +125,7 @@ class TeaCacheMixin: # Models that support CFG cache separation (wan/hunyuan/zimage) # Models not in this set (flux/qwen) auto-disable TeaCache when CFG is enabled _CFG_SUPPORTED_PREFIXES: set[str] = {"wan", "hunyuan", "zimage"} - config: DiTConfig + prefix: str def _init_teacache_state(self) -> None: """Initialize TeaCache state. Call this in subclass __init__.""" @@ -135,9 +133,7 @@ class TeaCacheMixin: self.cnt = 0 self.enable_teacache = True # Flag indicating if this model supports CFG cache separation - self._supports_cfg_cache = ( - self.config.prefix.lower() in self._CFG_SUPPORTED_PREFIXES - ) + self._supports_cfg_cache = self.prefix.lower() in self._CFG_SUPPORTED_PREFIXES # Always initialize positive cache fields (used in all modes) self.previous_modulated_input: torch.Tensor | None = None diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 08e01ae11..d778e1553 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -452,18 +452,21 @@ class FP32LayerNorm(CustomOp, nn.LayerNorm): ) self._forward_method = self.dispatch_forward() - try: - import attentions # noqa: F401 - except ImportError: - from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + if _is_npu: + try: + import attentions # noqa: F401 + except ImportError: + from sglang.multimodal_gen.runtime.utils.logging_utils import ( + init_logger, + ) - logger = init_logger(__name__) # pylint: disable=invalid-name - logger.warning( - "The 'attentions' library is not installed. Falling back to native layernorm. " - "Installing this library may improve performance on NPU." - "See: sgl-project/sgl-kernel-npu" - ) - self._forward_method = self.forward_native + logger = init_logger(__name__) # pylint: disable=invalid-name + logger.warning( + "The 'attentions' library is not installed. Falling back to native layernorm. " + "Installing this library may improve performance on NPU." + "See: sgl-project/sgl-kernel-npu" + ) + self._forward_method = self.forward_native def _cached_fp32_param( self, attr: str, param: torch.Tensor | None, device: torch.device diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index 75094ac98..1cc34150c 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -7,7 +7,7 @@ from typing import Any import torch from torch import nn -from sglang.multimodal_gen.configs.models import DiTConfig +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig # NOTE: SpectrumMixin lives in runtime.cache.spectrum from sglang.multimodal_gen.runtime.cache.spectrum import SpectrumMixin @@ -49,7 +49,10 @@ class BaseDiT(nn.Module, ABC): def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None: super().__init__() - self.config = config + # runtime models expose checkpoint architecture through `config`; load + # settings such as the model prefix stay separate + self.config: DiTArchConfig = config.arch_config + self.prefix = config.prefix self.hf_config = hf_config if not self.supported_attention_backends: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 4e8b300c2..872d17fe0 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py @@ -520,7 +520,7 @@ class CausalWanTransformer3DModel(BaseDiT, LayerwiseOffloadableModuleMixin): # Causal-specific self.block_mask = None - self.num_frame_per_block = config.arch_config.num_frames_per_block + self.num_frame_per_block = self.config.num_frames_per_block # Block size is bounded only by the causal block-mask construction, which # supports any positive value. assert self.num_frame_per_block >= 1 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 10c688517..997972be1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -933,7 +933,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): ) -> None: super().__init__(config=config, hf_config=hf_config) - arch = config.arch_config + arch = self.config self.hidden_size = arch.hidden_size self.num_hidden_layers = arch.num_hidden_layers self.num_attention_heads = arch.num_attention_heads diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 9e6acfe37..dce4aa256 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py @@ -427,7 +427,7 @@ class ErnieImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin) ): super().__init__(config=config, hf_config=hf_config) - arch = config.arch_config + arch = self.config self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.out_channels diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 7786b2827..12d671532 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -1202,7 +1202,6 @@ class FluxTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): quant_config: Optional[QuantizationConfig] = None, ) -> None: super().__init__(config=config, hf_config=hf_config) - self.config = config.arch_config self.out_channels = ( getattr(self.config, "out_channels", None) or self.config.in_channels diff --git a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py index ad2e10805..b26065c62 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/glm_image.py @@ -945,8 +945,7 @@ class GlmImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ): super().__init__(config=config, hf_config=hf_config) - self.config_data = config # Store config - arch_config = config.arch_config + arch_config = self.config self.in_channels = arch_config.in_channels self.out_channels = arch_config.out_channels diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py index 3474e76e0..dd58583f3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/hunyuan3d.py @@ -487,7 +487,7 @@ class Hunyuan3D2DiT(CachableDiT, LayerwiseOffloadableModuleMixin): **kwargs, ): super().__init__(config=config, hf_config=hf_config or {}, **kwargs) - arch = config.arch_config + arch = self.config in_channels = arch.in_channels context_in_dim = arch.context_in_dim diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py index 4e48932a9..d42a3a466 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ideogram.py @@ -513,7 +513,7 @@ class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin): **kwargs, ) -> None: super().__init__(config, hf_config, **kwargs) - cfg = config.arch_config + 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 @@ -586,7 +586,7 @@ class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin): def post_load_weights(self) -> None: if not self.rotary_emb.inv_freq.is_meta: return - cfg = self.config.arch_config + cfg = self.config inv_freq = 1.0 / ( cfg.rope_theta ** ( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py index 3ab91d0bd..c3d986ee6 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/krea2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/krea2.py @@ -522,7 +522,7 @@ class Krea2Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): quant_config: Optional[Any] = None, ) -> None: super().__init__(config=config, hf_config=hf_config) - ac = config.arch_config + ac = self.config self.arch_config = ac self.hidden_size = ac.features diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 77063c6c3..2992da867 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -1518,7 +1518,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): return timestep.amax(dim=tuple(range(1, timestep.ndim))) def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor: - ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) + ltx_variant = str(getattr(self.config, "ltx_variant", "ltx_2")) if ltx_variant == "ltx_2_3" and bool( getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False) ): @@ -1583,7 +1583,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ) -> None: super().__init__(config=config, hf_config=hf_config) - arch = config.arch_config + arch = self.config self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.audio_hidden_size = arch.audio_hidden_size @@ -1844,7 +1844,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): return video_coords.to(device=hidden_device) def _get_av_ca_gate_timestep_factor(self) -> float: - ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) + ltx_variant = str(getattr(self.config, "ltx_variant", "ltx_2")) if ltx_variant == "ltx_2_3": return self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier return float(self.av_ca_timestep_scale_multiplier) @@ -1856,7 +1856,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): prompt_timestep: torch.Tensor | None, audio_prompt_timestep: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: - ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) + ltx_variant = str(getattr(self.config, "ltx_variant", "ltx_2")) if ltx_variant != "ltx_2_3": return timestep, audio_timestep 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 d9524a0c2..947e9090d 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -1109,7 +1109,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config=config, hf_config=hf_config) - arch = config.arch_config + arch = self.config self.arch = arch self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 70e0cd4b3..6a1d68a17 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -1371,23 +1371,24 @@ class QwenImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): quant_config: Optional[QuantizationConfig] = None, ): super().__init__(config=config, hf_config=hf_config) - patch_size = config.arch_config.patch_size - in_channels = config.arch_config.in_channels - out_channels = config.arch_config.out_channels - num_layers = config.arch_config.num_layers - attention_head_dim = config.arch_config.attention_head_dim - num_attention_heads = config.arch_config.num_attention_heads - joint_attention_dim = config.arch_config.joint_attention_dim - axes_dims_rope = config.arch_config.axes_dims_rope - self.zero_cond_t = getattr(config.arch_config, "zero_cond_t", False) + arch = self.config + patch_size = arch.patch_size + in_channels = arch.in_channels + out_channels = arch.out_channels + num_layers = arch.num_layers + attention_head_dim = arch.attention_head_dim + num_attention_heads = arch.num_attention_heads + joint_attention_dim = arch.joint_attention_dim + axes_dims_rope = arch.axes_dims_rope + self.zero_cond_t = getattr(arch, "zero_cond_t", False) self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.use_additional_t_cond: bool = getattr( - config.arch_config, "use_additional_t_cond", False + arch, "use_additional_t_cond", False ) # For qwen-image-layered now self.use_layer3d_rope: bool = getattr( - config.arch_config, "use_layer3d_rope", False + arch, "use_layer3d_rope", False ) # For qwen-image-layered now if not self.use_layer3d_rope: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana.py b/python/sglang/multimodal_gen/runtime/models/dits/sana.py index dcc1dd2b1..98c3a693f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana.py @@ -401,7 +401,7 @@ class SanaTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): def __init__(self, config: SanaConfig, hf_config=None, **kwargs): super().__init__(config, hf_config=hf_config or {}, **kwargs) - arch = config.arch_config + arch = self.config self.out_channels = arch.out_channels self.patch_size = arch.patch_size self.inner_dim = arch.num_attention_heads * arch.attention_head_dim diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py index 387324dec..c029d3c94 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm.py @@ -353,7 +353,7 @@ class SanaWMTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin): def __init__(self, config: SanaWMConfig, hf_config=None, **kwargs) -> None: super().__init__(config, hf_config=hf_config or {}, **kwargs) - arch = config.arch_config + arch = self.config self.patch_size = (arch.patch_size_t, arch.patch_size, arch.patch_size) self.inner_dim = arch.num_attention_heads * arch.attention_head_dim diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py index 8d9b353b9..9ee6f888a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/sana_wm_refiner_transformer.py @@ -226,7 +226,7 @@ class SanaWMLTX2VideoRefiner(CachableDiT, LayerwiseOffloadableModuleMixin): quant_config: QuantizationConfig | None = None, ) -> None: super().__init__(config, hf_config=hf_config) - arch = config.arch_config + arch = self.config self.in_channels = int(arch.in_channels) self.out_channels = int(arch.out_channels) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py index 82236c919..8bcb04cf2 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/stablediffusion3.py @@ -40,8 +40,7 @@ class SD3Transformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): quant_config=None, ): super().__init__(config=config, hf_config=hf_config) - self.config = config - arch_config = config.arch_config + arch_config = self.config sample_size = arch_config.sample_size patch_size = arch_config.patch_size in_channels = arch_config.in_channels diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 52541cb01..16ca5c720 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -808,8 +808,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin): ) -> None: super().__init__(config=config, hf_config=hf_config) - self.config_data = config # Store config - arch_config = config.arch_config + arch_config = self.config self.in_channels = arch_config.in_channels self.out_channels = arch_config.out_channels diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index fcff3cdea..44408e550 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -157,12 +157,10 @@ class CausalDMDDenoisingStage(DenoisingStage): self.causal_kv_cache_neg: list | None = None self.crossattn_cache_neg: list | None = None # Model-dependent constants (aligned with causal_inference.py assumptions) - self.num_transformer_blocks = self.transformer.config.arch_config.num_layers - self.num_frames_per_block = ( - self.transformer.config.arch_config.num_frames_per_block - ) + self.num_transformer_blocks = self.transformer.config.num_layers + self.num_frames_per_block = self.transformer.config.num_frames_per_block self.sliding_window_num_frames = ( - self.transformer.config.arch_config.sliding_window_num_frames + self.transformer.config.sliding_window_num_frames ) try: @@ -171,7 +169,7 @@ class CausalDMDDenoisingStage(DenoisingStage): ) # type: ignore except Exception: self.local_attn_size = -1 - self.sink_size = self.transformer.config.arch_config.sink_size + self.sink_size = self.transformer.config.sink_size self._causal_attn_metadata_builder_cls = None self._causal_attn_metadata_builder = None @@ -190,8 +188,8 @@ class CausalDMDDenoisingStage(DenoisingStage): def _prepare_frame_seq_length(self, h: int, w: int) -> int: patch_ratio = ( - self.transformer.config.arch_config.patch_size[-1] - * self.transformer.config.arch_config.patch_size[-2] + self.transformer.config.patch_size[-1] + * self.transformer.config.patch_size[-2] ) self.num_token_per_frame = (h * w) // patch_ratio return self.num_token_per_frame diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py index c4d32f6eb..343866a54 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py @@ -27,8 +27,7 @@ class RealtimeChunkLatentPreparationStage(LatentPreparationStage): server_args: ServerArgs, ) -> int: return int( - batch.realtime_chunk_size - or self.transformer.config.arch_config.num_frames_per_block + batch.realtime_chunk_size or self.transformer.config.num_frames_per_block ) def get_latent_preparation_spec( @@ -47,7 +46,7 @@ class RealtimeChunkLatentPreparationStage(LatentPreparationStage): return LatentPreparationSpec( shape=( condition_latent.shape[0], - self.transformer.config.arch_config.out_channels, + self.transformer.config.out_channels, num_frames, condition_latent.shape[3], condition_latent.shape[4], diff --git a/python/sglang/multimodal_gen/test/unit/test_dit_config_boundary.py b/python/sglang/multimodal_gen/test/unit/test_dit_config_boundary.py new file mode 100644 index 000000000..39f893d79 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_dit_config_boundary.py @@ -0,0 +1,25 @@ +import torch + +from sglang.multimodal_gen.configs.models.dits.base import DiTConfig +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT + + +class _TestDiT(CachableDiT): + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor, + **kwargs, + ) -> torch.Tensor: + return hidden_states + + +def test_dit_runtime_keeps_architecture_and_component_config_separate(): + component_config = DiTConfig(prefix="Wan") + model = _TestDiT(config=component_config, hf_config={}) + + assert model.config is component_config.arch_config + assert model.prefix == "Wan" + assert model._supports_cfg_cache + assert model._spectrum_supports_cfg_cache