WIP: initial multimodal-gen support (#12484)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: JiLi <leege233@gmail.com>
Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Co-authored-by: laixin <xielx@shanghaitech.edu.cn>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: jzhang38 <a1286225768@gmail.com>
Co-authored-by: BrianChen1129 <yongqichcd@gmail.com>
Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: rlsu9 <r3su@ucsd.edu>
Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com>
Co-authored-by: foreverpiano <pianoqwz@qq.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: PorridgeSwim <yz3883@columbia.edu>
Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>

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

@@ -0,0 +1,134 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models import DiTConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+# TODO
+class BaseDiT(nn.Module, ABC):
+    _fsdp_shard_conditions: list = []
+    _compile_conditions: list = []
+    param_names_mapping: dict
+    reverse_param_names_mapping: dict
+    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_subclass__(cls) -> None:
+        required_class_attrs = [
+            "_fsdp_shard_conditions",
+            "param_names_mapping",
+            "_compile_conditions",
+        ]
+        super().__init_subclass__()
+        for attr in required_class_attrs:
+            if not hasattr(cls, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' class variable"
+                )
+
+    def __init__(self, config: DiTConfig, hf_config: dict[str, Any], **kwargs) -> None:
+        super().__init__()
+        self.config = config
+        self.hf_config = hf_config
+        if not self.supported_attention_backends:
+            raise ValueError(
+                f"Subclass {self.__class__.__name__} must define _supported_attention_backends"
+            )
+
+    @abstractmethod
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        pass
+
+    def __post_init__(self) -> None:
+        required_attrs = ["hidden_size", "num_attention_heads", "num_channels_latents"]
+        for attr in required_attrs:
+            if not hasattr(self, attr):
+                raise AttributeError(
+                    f"Subclasses of BaseDiT must define '{attr}' instance variable"
+                )
+
+    @property
+    def supported_attention_backends(self) -> set[AttentionBackendEnum]:
+        return self._supported_attention_backends
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the model."""
+        return next(self.parameters()).device
+
+
+class CachableDiT(BaseDiT):
+    """
+    An intermediate base class that adds TeaCache optimization functionality to DiT models.
+    TeaCache accelerates inference by selectively skipping redundant computation when consecutive
+    diffusion steps are similar enough.
+    """
+
+    # These are required class attributes that should be overridden by concrete implementations
+    _fsdp_shard_conditions = []
+    param_names_mapping = {}
+    reverse_param_names_mapping = {}
+    lora_param_names_mapping: dict = {}
+    # Ensure these instance attributes are properly defined in subclasses
+    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)
+
+        self.cnt = 0
+        self.teacache_thresh = 0
+        self.coefficients: list[float] = []
+
+        # NOTE(will): Only wan2.1 needs these, so we are hardcoding it here
+        if self.config.prefix == "wan":
+            self.use_ret_steps = self.config.cache_config.use_ret_steps
+            self.is_even = False
+            self.previous_residual_even: torch.Tensor | None = None
+            self.previous_residual_odd: torch.Tensor | None = None
+            self.accumulated_rel_l1_distance_even = 0
+            self.accumulated_rel_l1_distance_odd = 0
+            self.should_calc_even = True
+            self.should_calc_odd = True
+        else:
+            self.accumulated_rel_l1_distance = 0
+            self.previous_modulated_input = None
+            self.previous_resiual = None
+        self.previous_e0_even: torch.Tensor | None = None
+        self.previous_e0_odd: torch.Tensor | None = None
+
+    def maybe_cache_states(
+        self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor
+    ) -> None:
+        pass
+
+    def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool:
+        return False
+
+    def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError("maybe_retrieve_cached_states is not implemented")