sgl-project--sglang/python/sglang/multimodal_gen/runtime/models/dits/helios.py

# SPDX-License-Identifier: Apache-2.0
# Adapted from Helios diffusers transformer:
# https://github.com/BestWishYsh/Helios
"""
Helios Transformer 3D model for video generation.

Implements the HeliosTransformer3DModel with multi-term memory patches,
3D rotary position embeddings, and per-block scale-shift modulation.
"""

import math
from functools import lru_cache
from typing import Any

import torch
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.runtime.distributed import (
    divide,
    get_sp_world_size,
    get_tp_world_size,
)
from sglang.multimodal_gen.runtime.distributed.communication_op import (
    sequence_model_parallel_all_gather,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group
from sglang.multimodal_gen.runtime.layers.attention import USPAttention
from sglang.multimodal_gen.runtime.layers.layernorm import (
    LayerNorm,
    LayerNormScaleShift,
    RMSNorm,
    tensor_parallel_rms_norm,
)
from sglang.multimodal_gen.runtime.layers.linear import (
    ColumnParallelLinear,
    RowParallelLinear,
)
from sglang.multimodal_gen.runtime.layers.mlp import MLP
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
    QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.visual_embedding import (
    ModulateProjection,
    PatchEmbed,
    TimestepEmbedder,
)
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------


def pad_for_3d_conv(x, kernel_size):
    """Pad input to make it divisible by kernel_size using replicate mode."""
    b, c, t, h, w = x.shape
    pt, ph, pw = kernel_size
    pad_t = (pt - (t % pt)) % pt
    pad_h = (ph - (h % ph)) % ph
    pad_w = (pw - (w % pw)) % pw
    return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate")


def center_down_sample_3d(x, kernel_size):
    """Average pooling for 3D downsampling."""
    return F.avg_pool3d(x, kernel_size, stride=kernel_size)


def apply_rotary_emb_transposed(hidden_states, freqs_cis):
    """Apply rotary positional embeddings with transposed cos/sin format."""
    x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
    cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1)
    out = torch.empty_like(hidden_states)
    out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2]
    out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2]
    return out.type_as(hidden_states)


# ---------------------------------------------------------------------------
# Output norm
# ---------------------------------------------------------------------------


class HeliosOutputNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
        self.norm = LayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )

    def forward(self, hidden_states, temb, original_context_length):
        temb = temb[:, -original_context_length:, :]
        shift, scale = (
            self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)
        ).chunk(2, dim=2)
        shift = shift.squeeze(2).to(hidden_states.device)
        scale = scale.squeeze(2).to(hidden_states.device)
        hidden_states = hidden_states[:, -original_context_length:, :]
        hidden_states = self.norm(hidden_states, shift, scale)
        return hidden_states


# ---------------------------------------------------------------------------
# Rotary Positional Embedding (3D)
# ---------------------------------------------------------------------------


class HeliosRotaryPosEmbed(nn.Module):
    """3D rotary position embeddings for (time, height, width)."""

    def __init__(self, rope_dim, theta):
        super().__init__()
        self.DT, self.DY, self.DX = rope_dim
        self.theta = theta
        # Store as plain attributes (not buffers) to avoid meta-device issues
        # during FSDP loading. They'll be re-created on the correct device in forward.
        self._freqs_base_t = None
        self._freqs_base_y = None
        self._freqs_base_x = None

    def _get_freqs_base(self, dim):
        return 1.0 / (
            self.theta
            ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)
        )

    def _ensure_freqs_base(self, device):
        """Lazily create frequency bases on the correct device."""
        if self._freqs_base_t is None or self._freqs_base_t.device != device:
            self._freqs_base_t = self._get_freqs_base(self.DT).to(device)
            self._freqs_base_y = self._get_freqs_base(self.DY).to(device)
            self._freqs_base_x = self._get_freqs_base(self.DX).to(device)

    @torch.no_grad()
    def get_frequency_batched(self, freqs_base, pos):
        freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos)
        freqs = freqs.repeat_interleave(2, dim=0)
        return freqs.cos(), freqs.sin()

    @torch.no_grad()
    @lru_cache(maxsize=32)
    def _get_spatial_meshgrid(self, height, width, device_str):
        device = torch.device(device_str)
        grid_y_coords = torch.arange(height, device=device, dtype=torch.float32)
        grid_x_coords = torch.arange(width, device=device, dtype=torch.float32)
        grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij")
        return grid_y, grid_x

    @torch.no_grad()
    def forward(self, frame_indices, height, width, device):
        self._ensure_freqs_base(device)
        batch_size = frame_indices.shape[0]
        num_frames = frame_indices.shape[1]

        frame_indices = frame_indices.to(device=device, dtype=torch.float32)
        grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device))

        grid_t = frame_indices[:, :, None, None].expand(
            batch_size, num_frames, height, width
        )
        grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1)
        grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1)

        freqs_cos_t, freqs_sin_t = self.get_frequency_batched(
            self._freqs_base_t, grid_t
        )
        freqs_cos_y, freqs_sin_y = self.get_frequency_batched(
            self._freqs_base_y, grid_y_batch
        )
        freqs_cos_x, freqs_sin_x = self.get_frequency_batched(
            self._freqs_base_x, grid_x_batch
        )

        result = torch.cat(
            [
                freqs_cos_t,
                freqs_cos_y,
                freqs_cos_x,
                freqs_sin_t,
                freqs_sin_y,
                freqs_sin_x,
            ],
            dim=0,
        )
        return result.permute(1, 0, 2, 3, 4)


# ---------------------------------------------------------------------------
# Condition Embedder
# ---------------------------------------------------------------------------


class HeliosTimeTextEmbedding(nn.Module):
    """Condition embedder combining timestep and text embeddings."""

    def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim):
        super().__init__()
        self.time_embedder = TimestepEmbedder(
            dim, frequency_embedding_size=time_freq_dim, act_layer="silu"
        )
        self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu")
        self.text_embedder = MLP(
            text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh"
        )

    def forward(
        self, timestep, encoder_hidden_states, is_return_encoder_hidden_states=True
    ):
        temb = self.time_embedder(timestep)
        timestep_proj = self.time_modulation(temb)

        if encoder_hidden_states is not None and is_return_encoder_hidden_states:
            encoder_hidden_states = self.text_embedder(encoder_hidden_states)

        return temb, timestep_proj, encoder_hidden_states


# ---------------------------------------------------------------------------
# Self-Attention for Helios
# ---------------------------------------------------------------------------


class HeliosSelfAttention(nn.Module):
    """Self-attention with RMSNorm Q/K, optional history key amplification."""

    def __init__(
        self,
        dim: int,
        num_heads: int,
        eps: float = 1e-6,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        tp_size = get_tp_world_size()
        self.local_num_heads = divide(num_heads, tp_size)

        self.to_q = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_k = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_v = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_out = RowParallelLinear(
            dim, dim, bias=True, reduce_results=True, quant_config=quant_config
        )
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
        self.tp_rmsnorm = tp_size > 1

        self.attn = USPAttention(
            num_heads=self.local_num_heads,
            head_size=self.head_dim,
            causal=False,
            is_cross_attention=False,
        )

        self.is_amplify_history = is_amplify_history
        if is_amplify_history:
            if history_scale_mode == "scalar":
                self.history_key_scale = nn.Parameter(torch.ones(1))
            elif history_scale_mode == "per_head":
                self.history_key_scale = nn.Parameter(torch.ones(num_heads))
            else:
                raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}")
            self.history_scale_mode = history_scale_mode
            self.max_scale = 10.0

    def forward(self, hidden_states, rotary_emb=None, original_context_length=None):
        q, _ = self.to_q(hidden_states)
        k, _ = self.to_k(hidden_states)
        v, _ = self.to_v(hidden_states)

        if self.tp_rmsnorm:
            q = tensor_parallel_rms_norm(q, self.norm_q)
            k = tensor_parallel_rms_norm(k, self.norm_k)
        else:
            q = self.norm_q(q)
            k = self.norm_k(k)

        q = q.unflatten(2, (self.local_num_heads, self.head_dim))
        k = k.unflatten(2, (self.local_num_heads, self.head_dim))
        v = v.unflatten(2, (self.local_num_heads, self.head_dim))

        if rotary_emb is not None:
            q = apply_rotary_emb_transposed(q, rotary_emb)
            k = apply_rotary_emb_transposed(k, rotary_emb)

        history_seq_len = (
            hidden_states.shape[1] - original_context_length
            if original_context_length is not None
            else 0
        )

        if self.is_amplify_history and original_context_length is not None:
            if history_seq_len > 0:
                scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * (
                    self.max_scale - 1.0
                )
                if self.history_scale_mode == "per_head":
                    scale_key = scale_key.view(1, 1, -1, 1)
                k = torch.cat(
                    [k[:, :history_seq_len] * scale_key, k[:, history_seq_len:]],
                    dim=1,
                )

        x = self.attn(q, k, v, num_replicated_prefix=history_seq_len)
        x = x.flatten(2)
        x, _ = self.to_out(x)
        return x


# ---------------------------------------------------------------------------
# Cross-Attention for Helios
# ---------------------------------------------------------------------------


class HeliosCrossAttention(nn.Module):
    """Cross-attention with RMSNorm Q/K normalization."""

    def __init__(
        self,
        dim: int,
        num_heads: int,
        eps: float = 1e-6,
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        tp_size = get_tp_world_size()
        self.local_num_heads = divide(num_heads, tp_size)

        self.to_q = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_k = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_v = ColumnParallelLinear(
            dim, dim, bias=True, gather_output=False, quant_config=quant_config
        )
        self.to_out = RowParallelLinear(
            dim, dim, bias=True, reduce_results=True, quant_config=quant_config
        )
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
        self.tp_rmsnorm = tp_size > 1

        self.attn = USPAttention(
            num_heads=self.local_num_heads,
            head_size=self.head_dim,
            causal=False,
            skip_sequence_parallel=True,
        )

    def project_kv(self, encoder_hidden_states):
        """Project encoder states to this block's cross-attn (key, value)."""
        k, _ = self.to_k(encoder_hidden_states)
        v, _ = self.to_v(encoder_hidden_states)
        if self.tp_rmsnorm:
            k = tensor_parallel_rms_norm(k, self.norm_k)
        else:
            k = self.norm_k(k)
        k = k.unflatten(2, (self.local_num_heads, self.head_dim))
        v = v.unflatten(2, (self.local_num_heads, self.head_dim))
        return k, v

    def forward(
        self, hidden_states, encoder_hidden_states=None, encoder_key_value=None
    ):
        q, _ = self.to_q(hidden_states)
        if self.tp_rmsnorm:
            q = tensor_parallel_rms_norm(q, self.norm_q)
        else:
            q = self.norm_q(q)
        q = q.unflatten(2, (self.local_num_heads, self.head_dim))

        if encoder_key_value is None:
            if encoder_hidden_states is None:
                raise ValueError(
                    "encoder_hidden_states is required when encoder_key_value"
                    " is not provided."
                )
            encoder_key_value = self.project_kv(encoder_hidden_states)
        k, v = encoder_key_value

        x = self.attn(q, k, v)
        x = x.flatten(2)
        x, _ = self.to_out(x)
        return x


# ---------------------------------------------------------------------------
# Transformer Block
# ---------------------------------------------------------------------------


class HeliosTransformerBlock(nn.Module):
    """
    Single transformer block with self-attention, cross-attention, FFN,
    and scale-shift modulation from timestep embeddings.
    """

    def __init__(
        self,
        dim: int,
        ffn_dim: int,
        num_heads: int,
        cross_attn_norm: bool = True,
        eps: float = 1e-6,
        guidance_cross_attn: bool = True,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()

        # 1. Self-attention
        self.norm1 = LayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )
        self.attn1 = HeliosSelfAttention(
            dim=dim,
            num_heads=num_heads,
            eps=eps,
            is_amplify_history=is_amplify_history,
            history_scale_mode=history_scale_mode,
            quant_config=quant_config,
        )

        # 2. Cross-attention
        self.attn2 = HeliosCrossAttention(
            dim=dim,
            num_heads=num_heads,
            eps=eps,
            quant_config=quant_config,
        )
        self.self_attn_residual_norm = (
            LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=torch.float32)
            if cross_attn_norm
            else nn.Identity()
        )

        # 3. Feed-forward
        self.ffn = MLP(
            dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config
        )
        self.norm3 = LayerNormScaleShift(
            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
        )

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)

        # 4. Guidance cross-attention flag
        self.guidance_cross_attn = guidance_cross_attn

    def forward(
        self,
        hidden_states,
        encoder_hidden_states,
        temb,
        rotary_emb,
        original_context_length=None,
        cross_attn_key_value=None,
    ):
        if temb.ndim == 4:
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
                self.scale_shift_table.unsqueeze(0) + temb.float()
            ).chunk(6, dim=2)
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
            gate_msa = gate_msa.squeeze(2)
            c_shift_msa = c_shift_msa.squeeze(2)
            c_scale_msa = c_scale_msa.squeeze(2)
            c_gate_msa = c_gate_msa.squeeze(2)
        else:
            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
                self.scale_shift_table + temb.float()
            ).chunk(6, dim=1)

        # 1. Self-attention
        norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa)
        attn_output = self.attn1(
            norm_hidden_states, rotary_emb, original_context_length
        )
        hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(
            hidden_states
        )

        # 2. Cross-attention
        if self.guidance_cross_attn:
            history_seq_len = hidden_states.shape[1] - original_context_length
            history_hidden_states, current_hidden_states = torch.split(
                hidden_states, [history_seq_len, original_context_length], dim=1
            )
            norm_hidden_states = self.self_attn_residual_norm(
                current_hidden_states.float()
            ).type_as(current_hidden_states)
            attn_output = self.attn2(
                norm_hidden_states,
                encoder_hidden_states,
                encoder_key_value=cross_attn_key_value,
            )
            current_hidden_states = current_hidden_states + attn_output
            hidden_states = torch.cat(
                [history_hidden_states, current_hidden_states], dim=1
            )
        else:
            norm_hidden_states = self.self_attn_residual_norm(
                hidden_states.float()
            ).type_as(hidden_states)
            attn_output = self.attn2(
                norm_hidden_states,
                encoder_hidden_states,
                encoder_key_value=cross_attn_key_value,
            )
            hidden_states = hidden_states + attn_output

        # 3. Feed-forward
        norm_hidden_states = self.norm3(hidden_states, c_shift_msa, c_scale_msa)
        ff_output = self.ffn(norm_hidden_states)
        hidden_states = (
            hidden_states.float() + ff_output.float() * c_gate_msa
        ).type_as(hidden_states)

        return hidden_states


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------


class HeliosTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
    """
    Helios Transformer 3D model for video generation.

    Implements multi-scale history patches, 3D RoPE, and chunked denoising
    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
    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

    def __init__(
        self,
        config: HeliosConfig,
        hf_config: dict[str, Any],
        quant_config: QuantizationConfig | None = None,
    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.in_channels = config.in_channels
        self.out_channels = config.out_channels
        self.num_channels_latents = config.num_channels_latents
        self.patch_size = config.patch_size
        self.text_len = config.text_len
        self.inner_dim = inner_dim

        # Helios-specific config
        self.zero_history_timestep = config.zero_history_timestep
        self.has_multi_term_memory_patch = config.has_multi_term_memory_patch
        self.guidance_cross_attn = config.guidance_cross_attn

        # 1. Patch & position embedding
        self.patch_embedding = PatchEmbed(
            in_chans=config.in_channels,
            embed_dim=inner_dim,
            patch_size=config.patch_size,
            flatten=False,
        )

        # 2. Rotary position embeddings
        self.rope = HeliosRotaryPosEmbed(
            rope_dim=config.rope_dim, theta=config.rope_theta
        )

        # 3. Multi-term memory patches
        if self.has_multi_term_memory_patch:
            self.patch_short = nn.Conv3d(
                config.in_channels,
                inner_dim,
                kernel_size=config.patch_size,
                stride=config.patch_size,
            )
            self.patch_mid = nn.Conv3d(
                config.in_channels,
                inner_dim,
                kernel_size=tuple(2 * p for p in config.patch_size),
                stride=tuple(2 * p for p in config.patch_size),
            )
            self.patch_long = nn.Conv3d(
                config.in_channels,
                inner_dim,
                kernel_size=tuple(4 * p for p in config.patch_size),
                stride=tuple(4 * p for p in config.patch_size),
            )

        # 4. Condition embeddings
        self.condition_embedder = HeliosTimeTextEmbedding(
            dim=inner_dim,
            time_freq_dim=config.freq_dim,
            time_proj_dim=inner_dim * 6,
            text_embed_dim=config.text_dim,
        )

        # 5. Transformer blocks
        self.blocks = nn.ModuleList(
            [
                HeliosTransformerBlock(
                    dim=inner_dim,
                    ffn_dim=config.ffn_dim,
                    num_heads=config.num_attention_heads,
                    cross_attn_norm=config.cross_attn_norm,
                    eps=config.eps,
                    guidance_cross_attn=config.guidance_cross_attn,
                    is_amplify_history=config.is_amplify_history,
                    history_scale_mode=config.history_scale_mode,
                    quant_config=quant_config,
                )
                for _ in range(config.num_layers)
            ]
        )

        # 6. Output norm & projection
        self.norm_out = HeliosOutputNorm(inner_dim, config.eps)
        self.proj_out = ColumnParallelLinear(
            inner_dim,
            config.out_channels * math.prod(config.patch_size),
            bias=True,
            gather_output=True,
            quant_config=quant_config,
        )

        self.cnt = 0
        self.__post_init__()
        self.layer_names = ["blocks"]
        self.sp_size = get_sp_world_size()

    # Cross-attention K/V cache.
    #
    # Text conditioning is constant across the denoise loop, so the text
    # projection and every block's cross-attn K/V are computed once per request
    # (keyed by encoder-tensor identity) and reused across steps.

    @staticmethod
    def _request_cache(forward_batch, name):
        """Per-request cache dict on ``forward_batch.extra``.

        Returns None (-> caller recomputes, caching disabled) when there is no
        forward batch or gradients are enabled."""
        if forward_batch is None or torch.is_grad_enabled():
            return None
        extra = getattr(forward_batch, "extra", None)
        return None if extra is None else extra.setdefault(name, {})

    @staticmethod
    def _tensor_key(t):
        """Identity key for ``t``; equal only for the same underlying tensor."""
        return (
            t.data_ptr(),
            tuple(t.shape),
            tuple(t.stride()),
            t.dtype,
            t.device.type,
            t.device.index,
        )

    def _get_cross_attn_key_values(self, encoder_hidden_states, forward_batch):
        """Per-block cross-attn (key, value) for ``encoder_hidden_states``.

        Cached per request, keyed on the encoder tensor's identity
        (``_tensor_key``). The same object — ``batch.prompt_embeds`` — is passed
        every denoise step, so the key is stable and steps after the first hit
        the cache.
        """
        cache = self._request_cache(forward_batch, "helios_cross_attn_kv")
        key = self._tensor_key(encoder_hidden_states) if cache is not None else None
        kvs = cache.get(key) if key is not None else None
        if kvs is None:
            projected = self.condition_embedder.text_embedder(encoder_hidden_states)
            kvs = [block.attn2.project_kv(projected) for block in self.blocks]
            if key is not None:
                cache[key] = kvs
        return kvs

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
        timestep: torch.LongTensor,
        # Stage 1 history inputs
        indices_hidden_states=None,
        indices_latents_history_short=None,
        indices_latents_history_mid=None,
        indices_latents_history_long=None,
        latents_history_short=None,
        latents_history_mid=None,
        latents_history_long=None,
        **kwargs,
    ) -> torch.Tensor:
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]

        # Check if sequence parallelism is enabled
        forward_batch = get_forward_context().forward_batch
        if forward_batch is not None:
            sequence_shard_enabled = (
                forward_batch.enable_sequence_shard and self.sp_size > 1
            )
        else:
            sequence_shard_enabled = False

        batch_size = hidden_states.shape[0]
        p_t, p_h, p_w = self.patch_size

        # 1. Patch embed the noisy latents
        hidden_states = self.patch_embedding(hidden_states)
        (
            _,
            _,
            post_patch_num_frames,
            post_patch_height,
            post_patch_width,
        ) = hidden_states.shape

        if indices_hidden_states is None:
            indices_hidden_states = (
                torch.arange(0, post_patch_num_frames)
                .unsqueeze(0)
                .expand(batch_size, -1)
            )

        hidden_states = hidden_states.flatten(2).transpose(1, 2)

        # 2. Compute rotary embeddings
        rotary_emb = self.rope(
            frame_indices=indices_hidden_states,
            height=post_patch_height,
            width=post_patch_width,
            device=hidden_states.device,
        )
        rotary_emb = rotary_emb.flatten(2).transpose(1, 2)
        original_context_length = hidden_states.shape[1]

        # Sequence parallelism: shard current tokens and RoPE across SP ranks
        seq_shard_pad = 0
        if sequence_shard_enabled:
            sp_rank = get_sp_group().rank_in_group
            seq_len = hidden_states.shape[1]
            if seq_len % self.sp_size != 0:
                seq_shard_pad = self.sp_size - (seq_len % self.sp_size)
                hs_pad = torch.zeros(
                    batch_size,
                    seq_shard_pad,
                    hidden_states.shape[2],
                    dtype=hidden_states.dtype,
                    device=hidden_states.device,
                )
                re_pad = torch.zeros(
                    batch_size,
                    seq_shard_pad,
                    rotary_emb.shape[2],
                    dtype=rotary_emb.dtype,
                    device=rotary_emb.device,
                )
                hidden_states = torch.cat([hidden_states, hs_pad], dim=1)
                rotary_emb = torch.cat([rotary_emb, re_pad], dim=1)
            local_seq_len = hidden_states.shape[1] // self.sp_size
            hidden_states = hidden_states.view(
                batch_size, self.sp_size, local_seq_len, -1
            )[:, sp_rank, :, :].contiguous()
            rotary_emb = rotary_emb.view(batch_size, self.sp_size, local_seq_len, -1)[
                :, sp_rank, :, :
            ].contiguous()
            effective_context_length = local_seq_len
        else:
            effective_context_length = original_context_length

        # 3. Process short history
        if (
            latents_history_short is not None
            and indices_latents_history_short is not None
        ):
            latents_history_short = latents_history_short.to(hidden_states)
            latents_history_short = self.patch_short(latents_history_short)
            _, _, _, H1, W1 = latents_history_short.shape
            latents_history_short = latents_history_short.flatten(2).transpose(1, 2)

            rotary_emb_history_short = self.rope(
                frame_indices=indices_latents_history_short,
                height=H1,
                width=W1,
                device=latents_history_short.device,
            )
            rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(
                1, 2
            )
            hidden_states = torch.cat([latents_history_short, hidden_states], dim=1)
            rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1)

        # 4. Process mid history
        if latents_history_mid is not None and indices_latents_history_mid is not None:
            latents_history_mid = latents_history_mid.to(hidden_states)
            latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4))
            latents_history_mid = self.patch_mid(latents_history_mid)
            latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2)

            rotary_emb_history_mid = self.rope(
                frame_indices=indices_latents_history_mid,
                height=H1,
                width=W1,
                device=latents_history_mid.device,
            )
            rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2))
            rotary_emb_history_mid = center_down_sample_3d(
                rotary_emb_history_mid, (2, 2, 2)
            )
            rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2)

            hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1)
            rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1)

        # 5. Process long history
        if (
            latents_history_long is not None
            and indices_latents_history_long is not None
        ):
            latents_history_long = latents_history_long.to(hidden_states)
            latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8))
            latents_history_long = self.patch_long(latents_history_long)
            latents_history_long = latents_history_long.flatten(2).transpose(1, 2)

            rotary_emb_history_long = self.rope(
                frame_indices=indices_latents_history_long,
                height=H1,
                width=W1,
                device=latents_history_long.device,
            )
            rotary_emb_history_long = pad_for_3d_conv(
                rotary_emb_history_long, (4, 4, 4)
            )
            rotary_emb_history_long = center_down_sample_3d(
                rotary_emb_history_long, (4, 4, 4)
            )
            rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2)

            hidden_states = torch.cat([latents_history_long, hidden_states], dim=1)
            rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1)

        history_context_length = hidden_states.shape[1] - effective_context_length

        # 6. Compute condition embeddings
        if indices_hidden_states is not None and self.zero_history_timestep:
            timestep_t0 = torch.zeros(
                (1,), dtype=timestep.dtype, device=timestep.device
            )
            temb_t0, timestep_proj_t0, _ = self.condition_embedder(
                timestep_t0,
                encoder_hidden_states,
                is_return_encoder_hidden_states=False,
            )
            temb_t0 = temb_t0.unsqueeze(1).expand(
                batch_size, history_context_length, -1
            )
            timestep_proj_t0 = (
                timestep_proj_t0.unflatten(-1, (6, -1))
                .view(1, 6, 1, -1)
                .expand(batch_size, -1, history_context_length, -1)
            )

        # Take only the time embeddings (temb, timestep_proj); skip the text
        # projection (is_return_encoder_hidden_states=False) since it is computed
        # once per request and cached by _get_cross_attn_key_values below.
        temb, timestep_proj, _ = self.condition_embedder(
            timestep, encoder_hidden_states, is_return_encoder_hidden_states=False
        )
        cross_attn_key_values = self._get_cross_attn_key_values(
            encoder_hidden_states, forward_batch
        )
        timestep_proj = timestep_proj.unflatten(-1, (6, -1))

        if indices_hidden_states is not None and not self.zero_history_timestep:
            main_repeat_size = hidden_states.shape[1]
        else:
            main_repeat_size = effective_context_length
        temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1)
        timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(
            batch_size, 6, main_repeat_size, -1
        )

        if indices_hidden_states is not None and self.zero_history_timestep:
            temb = torch.cat([temb_t0, temb], dim=1)
            timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2)

        if timestep_proj.ndim == 4:
            timestep_proj = timestep_proj.permute(0, 2, 1, 3)

        # 7. Transformer blocks
        hidden_states = hidden_states.contiguous()
        encoder_hidden_states = encoder_hidden_states.contiguous()
        rotary_emb = rotary_emb.contiguous()

        for block, key_value in zip(self.blocks, cross_attn_key_values):
            hidden_states = block(
                hidden_states,
                encoder_hidden_states,
                timestep_proj,
                rotary_emb,
                effective_context_length,
                cross_attn_key_value=key_value,
            )

        self.cnt += 1

        # SP: all-gather current tokens before output
        if sequence_shard_enabled:
            current_tokens = hidden_states[:, -local_seq_len:, :].contiguous()
            current_tokens = sequence_model_parallel_all_gather(current_tokens, dim=1)
            if seq_shard_pad > 0:
                current_tokens = current_tokens[:, :original_context_length, :]
            hidden_states = current_tokens
            # Re-create temb for norm_out (all current tokens share same timestep)
            temb = temb[:, :1, :].expand(batch_size, original_context_length, -1)

        # 8. Output norm & projection
        hidden_states = self.norm_out(hidden_states, temb, original_context_length)
        hidden_states, _ = self.proj_out(hidden_states)

        # 9. Unpatchify
        hidden_states = hidden_states.reshape(
            batch_size,
            post_patch_num_frames,
            post_patch_height,
            post_patch_width,
            p_t,
            p_h,
            p_w,
            -1,
        )
        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

        return output


EntryClass = HeliosTransformer3DModel