chore: import upstream snapshot with attribution

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

@@ -0,0 +1,127 @@
+# 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
+
+# NOTE: TeaCacheContext and TeaCacheMixin have been moved to
+# sglang.multimodal_gen.runtime.cache.teacache
+# For backwards compatibility, re-export from the new location
+from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext  # noqa: F401
+from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin
+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"
+                )
+
+    def post_load_weights(self) -> None:
+        """Run model-specific post-load weight fixups after all parameters are materialized."""
+        return None
+
+    @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(TeaCacheMixin, BaseDiT):
+    """
+    An intermediate base class that adds TeaCache optimization functionality to DiT models.
+
+    Inherits TeaCacheMixin for cache logic and BaseDiT for core DiT functionality.
+    """
+
+    # 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._init_teacache_state()
+
+    @classmethod
+    def get_nunchaku_quant_rules(cls) -> dict[str, dict[str, Any]]:
+        """
+        Get quantization rules for Nunchaku quantization.
+
+        Returns a dict mapping layer name patterns to quantization configs:
+        {
+            "skip": [list of patterns to skip quantization],
+            "svdq_w4a4": [list of patterns for SVDQ W4A4],
+            "awq_w4a16": [list of patterns for AWQ W4A16],
+        }
+        """
+        return {}