Helios: Real Real-Time Long Video Generation Model (#19782)

python/sglang/multimodal_gen/configs/models/dits/__init__.py

@@ -1,5 +1,6 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
+from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
 from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
 from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
@@ -7,6 +8,7 @@ from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
 from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig
 
 __all__ = [
+    "HeliosConfig",
     "HunyuanVideoConfig",
     "WanVideoConfig",
     "Hunyuan3DDiTConfig",

python/sglang/multimodal_gen/configs/models/dits/helios.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+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 HeliosArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Patch embeddings
+            r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+            # Condition embedder: text
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1",
+            # Condition embedder: time
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            # Blocks: self-attention (keep attn1. prefix, drop .0. from to_out)
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.attn1.to_out.\2",
+            # Blocks: cross-attention output (drop .0. from to_out)
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            # Blocks: feed-forward
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            # Blocks: cross-attn residual norm
+            r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.\2",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    lora_param_names_mapping: dict = field(default_factory=lambda: {})
+
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 226
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 13824
+    num_layers: int = 40
+    cross_attn_norm: bool = True
+    qk_norm: str = "rms_norm_across_heads"
+    eps: float = 1e-6
+    added_kv_proj_dim: int | None = None
+    rope_max_seq_len: int = 1024
+    pos_embed_seq_len: int | None = None
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
+
+    # Helios-specific
+    rope_dim: tuple[int, int, int] = (44, 42, 42)
+    rope_theta: float = 10000.0
+    guidance_cross_attn: bool = True
+    zero_history_timestep: bool = True
+    has_multi_term_memory_patch: bool = True
+    is_amplify_history: bool = False
+    history_scale_mode: str = "per_head"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class HeliosConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig)
+
+    prefix: str = "Helios"

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -15,6 +15,11 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import (
 from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import (
     Flux2FinetunedPipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.helios import (
+    HeliosDistilledConfig,
+    HeliosMidConfig,
+    HeliosT2VConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import (
     FastHunyuanConfig,
     HunyuanConfig,
@@ -35,6 +40,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipeline
 
 __all__ = [
     "DiffusersGenericPipelineConfig",
+    "HeliosDistilledConfig",
+    "HeliosMidConfig",
+    "HeliosT2VConfig",
     "HunyuanConfig",
     "FastHunyuanConfig",
     "Hunyuan3D2PipelineConfig",

python/sglang/multimodal_gen/configs/pipeline_configs/helios.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
+from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, T5Config
+from sglang.multimodal_gen.configs.models.encoders.t5 import T5ArchConfig
+from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# Helios UMT5 max sequence length (used for both tokenizer and post-processing padding)
+HELIOS_MAX_SEQUENCE_LENGTH = 226
+
+
+def umt5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor:
+    """Post-process UMT5 text encoder outputs, padding to HELIOS_MAX_SEQUENCE_LENGTH tokens."""
+    max_seq_len = HELIOS_MAX_SEQUENCE_LENGTH
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
+    seq_lens = mask.gt(0).sum(dim=1).long()
+    assert torch.isnan(hidden_state).sum() == 0
+    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)]
+    prompt_embeds_tensor: torch.Tensor = torch.stack(
+        [
+            torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))])
+            for u in prompt_embeds
+        ],
+        dim=0,
+    )
+    return prompt_embeds_tensor
+
+
+@dataclass
+class HeliosT2VConfig(PipelineConfig):
+    """Configuration for the Helios T2V pipeline."""
+
+    task_type: ModelTaskType = ModelTaskType.T2V
+
+    # DiT
+    dit_config: DiTConfig = field(default_factory=HeliosConfig)
+
+    # VAE (same as Wan)
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Denoising stage
+    flow_shift: float | None = 1.0
+
+    # Text encoding stage (UMT5 is T5-compatible)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (
+            T5Config(arch_config=T5ArchConfig(text_len=HELIOS_MAX_SEQUENCE_LENGTH)),
+        )
+    )
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = (
+        field(default_factory=lambda: (umt5_postprocess_text,))
+    )
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp32"
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",))
+
+    # Helios-specific chunked denoising params
+    num_latent_frames_per_chunk: int = 9
+    history_sizes: list[int] = field(default_factory=lambda: [16, 2, 1])
+    is_cfg_zero_star: bool = False
+    zero_steps: int = 1
+    keep_first_frame: bool = True
+
+    # Stage 2 (Pyramid SR) & Stage 3 (DMD) params
+    is_enable_stage2: bool = False
+    pyramid_num_stages: int = 3
+    pyramid_num_inference_steps_list: list[int] = field(
+        default_factory=lambda: [10, 10, 10]
+    )
+    is_distilled: bool = False
+    is_amplify_first_chunk: bool = False
+    scheduler_type: str = "unipc"
+    gamma: float = 1 / 3
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class HeliosMidConfig(HeliosT2VConfig):
+    """Configuration for Helios-Mid (Stage 1 + Stage 2 pyramid SR)."""
+
+    is_enable_stage2: bool = True
+    is_cfg_zero_star: bool = True
+    pyramid_num_inference_steps_list: list[int] = field(
+        default_factory=lambda: [20, 20, 20]
+    )
+
+
+@dataclass
+class HeliosDistilledConfig(HeliosT2VConfig):
+    """Configuration for Helios-Distilled (Stage 1 + Stage 2 + Stage 3 DMD)."""
+
+    is_enable_stage2: bool = True
+    is_distilled: bool = True
+    is_amplify_first_chunk: bool = True
+    scheduler_type: str = "dmd"
+    pyramid_num_inference_steps_list: list[int] = field(
+        default_factory=lambda: [10, 10, 10]
+    )

python/sglang/multimodal_gen/configs/sample/helios.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+
+
+@dataclass
+class HeliosT2VSamplingParams(SamplingParams):
+    # Video parameters
+    height: int = 384
+    width: int = 640
+    num_frames: int = 99
+    fps: int = 24
+
+    # Denoising stage
+    guidance_scale: float = 5.0
+    negative_prompt: str = (
+        "Bright tones, overexposed, static, blurred details, subtitles, style, "
+        "works, paintings, images, static, overall gray, worst quality, low quality, "
+        "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, "
+        "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, "
+        "still picture, messy background, three legs, many people in the background, "
+        "walking backwards"
+    )
+    num_inference_steps: int = 50
+
+    # Helios T2V supported resolutions
+    supported_resolutions: list[tuple[int, int]] | None = field(
+        default_factory=lambda: [
+            (640, 384),  # ~5:3
+            (384, 640),  # ~3:5
+            (832, 480),  # ~16:9-ish
+            (480, 832),  # ~9:16-ish
+        ]
+    )
+
+
+@dataclass
+class HeliosMidSamplingParams(HeliosT2VSamplingParams):
+    """Sampling params for Helios-Mid (Stage 2 pyramid SR)."""
+
+    num_inference_steps: int = 20
+
+
+@dataclass
+class HeliosDistilledSamplingParams(HeliosT2VSamplingParams):
+    """Sampling params for Helios-Distilled (DMD, no CFG needed)."""
+
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 10

python/sglang/multimodal_gen/registry.py

@@ -30,6 +30,9 @@ if TYPE_CHECKING:
 from sglang.multimodal_gen.configs.pipeline_configs import (
     FastHunyuanConfig,
     FluxPipelineConfig,
+    HeliosDistilledConfig,
+    HeliosMidConfig,
+    HeliosT2VConfig,
     HunyuanConfig,
     WanI2V480PConfig,
     WanI2V720PConfig,
@@ -74,6 +77,11 @@ from sglang.multimodal_gen.configs.sample.flux import (
     FluxSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.glmimage import GlmImageSamplingParams
+from sglang.multimodal_gen.configs.sample.helios import (
+    HeliosDistilledSamplingParams,
+    HeliosMidSamplingParams,
+    HeliosT2VSamplingParams,
+)
 from sglang.multimodal_gen.configs.sample.hunyuan import (
     FastHunyuanSamplingParam,
     HunyuanSamplingParams,
@@ -726,6 +734,34 @@ def _register_configs():
         model_detectors=[lambda hf_id: "hunyuan3d" in hf_id.lower()],
     )
 
+    # Helios
+    register_configs(
+        sampling_param_cls=HeliosT2VSamplingParams,
+        pipeline_config_cls=HeliosT2VConfig,
+        hf_model_paths=[
+            "BestWishYsh/Helios-Base",
+        ],
+        model_detectors=[
+            lambda hf_id: "helios" in hf_id.lower()
+            and "mid" not in hf_id.lower()
+            and "distill" not in hf_id.lower()
+        ],
+    )
+    register_configs(
+        sampling_param_cls=HeliosMidSamplingParams,
+        pipeline_config_cls=HeliosMidConfig,
+        hf_model_paths=[
+            "BestWishYsh/Helios-Mid",
+        ],
+    )
+    register_configs(
+        sampling_param_cls=HeliosDistilledSamplingParams,
+        pipeline_config_cls=HeliosDistilledConfig,
+        hf_model_paths=[
+            "BestWishYsh/Helios-Distilled",
+        ],
+    )
+
 
 _register_configs()
 

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

@@ -0,0 +1,825 @@
+# 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_tp_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    FP32LayerNorm,
+    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.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+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 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+
+    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.float()) * (1 + scale) + shift
+        ).type_as(hidden_states)
+        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)
+
+        if self.is_amplify_history and original_context_length is not None:
+            history_seq_len = hidden_states.shape[1] - original_context_length
+            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)
+        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,
+            is_cross_attention=True,
+        )
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        q, _ = self.to_q(hidden_states)
+        k, _ = self.to_k(encoder_hidden_states)
+        v, _ = self.to_v(encoder_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))
+
+        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 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        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 = (
+            FP32LayerNorm(dim, eps, elementwise_affine=True)
+            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 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+
+        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,
+    ):
+        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.float()) * (1 + scale_msa) + shift_msa
+        ).type_as(hidden_states)
+        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)
+            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)
+            hidden_states = hidden_states + attn_output
+
+        # 3. Feed-forward
+        norm_hidden_states = (
+            self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa
+        ).type_as(hidden_states)
+        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, OffloadableDiTMixin):
+    """
+    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"]
+
+    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:
+        orig_dtype = hidden_states.dtype
+        if not isinstance(encoder_hidden_states, torch.Tensor):
+            encoder_hidden_states = encoder_hidden_states[0]
+
+        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]
+
+        # 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] - original_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)
+            )
+
+        temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
+            timestep, encoder_hidden_states
+        )
+        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 = original_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 in self.blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states,
+                timestep_proj,
+                rotary_emb,
+                original_context_length,
+            )
+
+        self.cnt += 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

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py

@@ -0,0 +1,734 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from Helios diffusers scheduler:
+# https://github.com/BestWishYsh/Helios
+"""
+Helios scheduler implementing flow-matching with UniPC/Euler solvers.
+
+For Phase 1 T2V (stages=1), this simplifies to standard flow-matching
+with dynamic shifting and UniPC multistep solver.
+"""
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass
+class HeliosSchedulerOutput:
+    prev_sample: torch.FloatTensor
+    model_outputs: torch.FloatTensor | None = None
+    last_sample: torch.FloatTensor | None = None
+    this_order: int | None = None
+
+
+class HeliosSchedulerConfig:
+    """Mimics diffusers config interface for scheduler parameters."""
+
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+
+class HeliosScheduler:
+    """
+    Helios multi-stage scheduler supporting Euler, UniPC, and DMD solvers.
+
+    For Phase 1 T2V with stages=1, this is a standard flow-matching scheduler
+    with optional time shifting and UniPC multistep updates.
+    """
+
+    order = 1
+
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        stages: int = 1,
+        stage_range: list | None = None,
+        gamma: float = 1 / 3,
+        thresholding: bool = False,
+        prediction_type: str = "flow_prediction",
+        solver_order: int = 2,
+        predict_x0: bool = True,
+        solver_type: str = "bh2",
+        lower_order_final: bool = True,
+        disable_corrector: list[int] | None = None,
+        use_flow_sigmas: bool = True,
+        scheduler_type: str = "unipc",
+        use_dynamic_shifting: bool = False,
+        time_shift_type: str = "linear",
+        **kwargs,
+    ):
+        if stage_range is None:
+            # Evenly divide [0, 1] into 3 stages for pyramid SR
+            stage_range = [0, 1 / 3, 2 / 3, 1]
+        if disable_corrector is None:
+            disable_corrector = []
+
+        self.config = HeliosSchedulerConfig(
+            num_train_timesteps=num_train_timesteps,
+            shift=shift,
+            stages=stages,
+            stage_range=stage_range,
+            gamma=gamma,
+            thresholding=thresholding,
+            prediction_type=prediction_type,
+            solver_order=solver_order,
+            predict_x0=predict_x0,
+            solver_type=solver_type,
+            lower_order_final=lower_order_final,
+            disable_corrector=disable_corrector,
+            use_flow_sigmas=use_flow_sigmas,
+            scheduler_type=scheduler_type,
+            use_dynamic_shifting=use_dynamic_shifting,
+            time_shift_type=time_shift_type,
+        )
+
+        self.timestep_ratios = {}
+        self.timesteps_per_stage = {}
+        self.sigmas_per_stage = {}
+        self.start_sigmas = {}
+        self.end_sigmas = {}
+        self.ori_start_sigmas = {}
+
+        self.init_sigmas_for_each_stage()
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+        self.gamma = gamma
+
+        if solver_type not in ["bh1", "bh2"]:
+            raise NotImplementedError(f"{solver_type} is not implemented")
+
+        self.predict_x0 = predict_x0
+        self.model_outputs = [None] * solver_order
+        self.timestep_list = [None] * solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = disable_corrector
+        self.solver_p = None
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+        self.num_inference_steps = None
+
+    def init_sigmas(self):
+        num_train_timesteps = self.config.num_train_timesteps
+        shift = self.config.shift
+
+        alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1)
+        sigmas = 1.0 - alphas
+        sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy()
+        sigmas = torch.from_numpy(sigmas)
+        timesteps = (sigmas * num_train_timesteps).clone()
+
+        self._step_index = None
+        self._begin_index = None
+        self.timesteps = timesteps
+        self.sigmas = sigmas.to("cpu")
+
+    def init_sigmas_for_each_stage(self):
+        self.init_sigmas()
+
+        stage_distance = []
+        stages = self.config.stages
+        training_steps = self.config.num_train_timesteps
+        stage_range = self.config.stage_range
+
+        for i_s in range(stages):
+            start_indice = int(stage_range[i_s] * training_steps)
+            start_indice = max(start_indice, 0)
+            end_indice = int(stage_range[i_s + 1] * training_steps)
+            end_indice = min(end_indice, training_steps)
+            start_sigma = self.sigmas[start_indice].item()
+            end_sigma = (
+                self.sigmas[end_indice].item() if end_indice < training_steps else 0.0
+            )
+            self.ori_start_sigmas[i_s] = start_sigma
+
+            if i_s != 0:
+                ori_sigma = 1 - start_sigma
+                gamma = self.config.gamma
+                corrected_sigma = (
+                    1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
+                ) * ori_sigma
+                start_sigma = 1 - corrected_sigma
+
+            stage_distance.append(start_sigma - end_sigma)
+            self.start_sigmas[i_s] = start_sigma
+            self.end_sigmas[i_s] = end_sigma
+
+        tot_distance = sum(stage_distance)
+        for i_s in range(stages):
+            if i_s == 0:
+                start_ratio = 0.0
+            else:
+                start_ratio = sum(stage_distance[:i_s]) / tot_distance
+            if i_s == stages - 1:
+                # Use value just below 1.0 to avoid out-of-bounds indexing
+                end_ratio = 1.0 - 1e-16
+            else:
+                end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance
+            self.timestep_ratios[i_s] = (start_ratio, end_ratio)
+
+        for i_s in range(stages):
+            timestep_ratio = self.timestep_ratios[i_s]
+            # Clamp to max valid timestep (num_train_timesteps - 1)
+            timestep_max = min(
+                self.timesteps[int(timestep_ratio[0] * training_steps)], 999
+            )
+            timestep_min = self.timesteps[
+                min(int(timestep_ratio[1] * training_steps), training_steps - 1)
+            ]
+            timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1)
+            self.timesteps_per_stage[i_s] = (
+                timesteps[:-1]
+                if isinstance(timesteps, torch.Tensor)
+                else torch.from_numpy(timesteps[:-1])
+            )
+            # Sigma range [0.999, 0]: start just below 1.0 to avoid singularity
+            stage_sigmas = np.linspace(0.999, 0, training_steps + 1)
+            self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1])
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def time_shift(self, mu, sigma, t):
+        if self.config.time_shift_type == "exponential":
+            return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+        elif self.config.time_shift_type == "linear":
+            return mu / (mu + (1 / t - 1) ** sigma)
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        stage_index: int | None = None,
+        device: str | torch.device = None,
+        sigmas=None,
+        mu=None,
+        is_amplify_first_chunk: bool = False,
+    ):
+        if self.config.scheduler_type == "dmd":
+            if is_amplify_first_chunk:
+                num_inference_steps = num_inference_steps * 2 + 1
+            else:
+                num_inference_steps = num_inference_steps + 1
+
+        self.num_inference_steps = num_inference_steps
+        self.init_sigmas()
+
+        if self.config.stages == 1:
+            if sigmas is None:
+                sigmas = np.linspace(
+                    1,
+                    1 / self.config.num_train_timesteps,
+                    num_inference_steps + 1,
+                )[:-1].astype(np.float32)
+                if self.config.shift != 1.0:
+                    assert not self.config.use_dynamic_shifting
+                    sigmas = self.time_shift(self.config.shift, 1.0, sigmas)
+            timesteps = (sigmas * self.config.num_train_timesteps).copy()
+            sigmas = torch.from_numpy(sigmas)
+        else:
+            stage_timesteps = self.timesteps_per_stage[stage_index]
+            timesteps = np.linspace(
+                stage_timesteps[0].item(),
+                stage_timesteps[-1].item(),
+                num_inference_steps,
+            )
+            stage_sigmas = self.sigmas_per_stage[stage_index]
+            ratios = np.linspace(
+                stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps
+            )
+            sigmas = torch.from_numpy(ratios)
+
+        self.timesteps = torch.from_numpy(timesteps).to(device=device)
+        self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device)
+
+        self._step_index = None
+        self.reset_scheduler_history()
+
+        if self.config.scheduler_type == "dmd":
+            self.timesteps = self.timesteps[:-1]
+            self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]])
+
+        if self.config.use_dynamic_shifting:
+            assert self.config.shift == 1.0
+            self.sigmas = self.time_shift(mu, 1.0, self.sigmas)
+            if self.config.stages == 1:
+                self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps
+            else:
+                self.timesteps = self.timesteps_per_stage[
+                    stage_index
+                ].min() + self.sigmas[:-1] * (
+                    self.timesteps_per_stage[stage_index].max()
+                    - self.timesteps_per_stage[stage_index].min()
+                )
+
+    # ---------------------------------- Euler ----------------------------------
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step_euler(
+        self,
+        model_output: torch.FloatTensor,
+        timestep=None,
+        sample: torch.FloatTensor = None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.step_index is None:
+            self._step_index = 0
+
+        sample = sample.to(torch.float32)
+        sigma = self.sigmas[self.step_index]
+        sigma_next = self.sigmas[self.step_index + 1]
+
+        prev_sample = sample + (sigma_next - sigma) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- UniPC ----------------------------------
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        if self.config.use_flow_sigmas:
+            alpha_t = 1 - sigma
+            sigma_t = torch.clamp(sigma, min=1e-8)
+        else:
+            alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+            sigma_t = sigma * alpha_t
+        return alpha_t, sigma_t
+
+    def convert_model_output(self, model_output, sample=None, sigma=None, **kwargs):
+        flag = False
+        if sigma is None:
+            flag = True
+            sigma = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+        if self.predict_x0:
+            if self.config.prediction_type == "flow_prediction":
+                if flag:
+                    sigma_t = self.sigmas[self.step_index]
+                else:
+                    sigma_t = sigma
+                x0_pred = sample - sigma_t * model_output
+            elif self.config.prediction_type == "epsilon":
+                x0_pred = (sample - sigma_t * model_output) / alpha_t
+            elif self.config.prediction_type == "sample":
+                x0_pred = model_output
+            elif self.config.prediction_type == "v_prediction":
+                x0_pred = alpha_t * sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type {self.config.prediction_type} not supported"
+                )
+            return x0_pred
+        else:
+            if self.config.prediction_type == "epsilon":
+                return model_output
+            elif self.config.prediction_type == "sample":
+                return (sample - alpha_t * model_output) / sigma_t
+            elif self.config.prediction_type == "v_prediction":
+                return alpha_t * model_output + sigma_t * sample
+            else:
+                raise ValueError(
+                    f"prediction_type {self.config.prediction_type} not supported"
+                )
+
+    def multistep_uni_p_bh_update(
+        self, model_output, sample=None, order=None, sigma=None, sigma_next=None
+    ):
+        model_output_list = self.model_outputs
+        m0 = model_output_list[-1]
+        x = sample
+
+        if sigma_next is None and sigma is None:
+            sigma_t, sigma_s0 = (
+                self.sigmas[self.step_index + 1],
+                self.sigmas[self.step_index],
+            )
+        else:
+            sigma_t, sigma_s0 = sigma_next, sigma
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+        device = sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - i
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_k = h_phi_1 / hh - 1
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+            if order == 2:
+                rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+            else:
+                rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype)
+        else:
+            D1s = None
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            pred_res = (
+                torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
+            )
+            x_t = x_t_ - alpha_t * B_h * pred_res
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            pred_res = (
+                torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0
+            )
+            x_t = x_t_ - sigma_t * B_h * pred_res
+
+        return x_t.to(x.dtype)
+
+    def multistep_uni_c_bh_update(
+        self,
+        this_model_output,
+        last_sample=None,
+        this_sample=None,
+        order=None,
+        sigma_before=None,
+        sigma=None,
+    ):
+        model_output_list = self.model_outputs
+        m0 = model_output_list[-1]
+        x = last_sample
+        model_t = this_model_output
+
+        if sigma_before is None and sigma is None:
+            sigma_t, sigma_s0 = (
+                self.sigmas[self.step_index],
+                self.sigmas[self.step_index - 1],
+            )
+        else:
+            sigma_t, sigma_s0 = sigma, sigma_before
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+        device = this_sample.device
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            si = self.step_index - (i + 1)
+            mi = model_output_list[-(i + 1)]
+            alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+            lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+            rk = (lambda_si - lambda_s0) / h
+            rks.append(rk)
+            D1s.append((mi - m0) / rk)
+
+        rks.append(1.0)
+        rks = torch.tensor(rks, device=device)
+
+        R = []
+        b = []
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_k = h_phi_1 / hh - 1
+        factorial_i = 1
+
+        if self.config.solver_type == "bh1":
+            B_h = hh
+        elif self.config.solver_type == "bh2":
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= i + 1
+            h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+        R = torch.stack(R)
+        b = torch.tensor(b, device=device)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1)
+        else:
+            D1s = None
+
+        if order == 1:
+            rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+        else:
+            rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+        if self.predict_x0:
+            x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+            corr_res = (
+                torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+                if D1s is not None
+                else 0
+            )
+            D1_t = model_t - m0
+            x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+            corr_res = (
+                torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+                if D1s is not None
+                else 0
+            )
+            D1_t = model_t - m0
+            x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+
+        return x_t.to(x.dtype)
+
+    def step_unipc(
+        self,
+        model_output,
+        timestep=None,
+        sample=None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', run 'set_timesteps' first"
+            )
+
+        if self.step_index is None:
+            self._step_index = 0
+
+        use_corrector = (
+            self.step_index > 0
+            and self.step_index - 1 not in self.disable_corrector
+            and self.last_sample is not None
+        )
+
+        model_output_convert = self.convert_model_output(model_output, sample=sample)
+
+        if use_corrector:
+            sample = self.multistep_uni_c_bh_update(
+                this_model_output=model_output_convert,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+            self.timestep_list[i] = self.timestep_list[i + 1]
+        self.model_outputs[-1] = model_output_convert
+        self.timestep_list[-1] = timestep
+
+        if self.config.lower_order_final:
+            this_order = min(
+                self.config.solver_order, len(self.timesteps) - self.step_index
+            )
+        else:
+            this_order = self.config.solver_order
+        self.this_order = min(this_order, self.lower_order_nums + 1)
+        assert self.this_order > 0
+
+        self.last_sample = sample
+        prev_sample = self.multistep_uni_p_bh_update(
+            model_output=model_output,
+            sample=sample,
+            order=self.this_order,
+        )
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- DMD ----------------------------------
+    def add_noise(self, original_samples, noise, timestep, sigmas, timesteps):
+        sigmas = sigmas.to(noise.device)
+        timesteps = timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps):
+        original_dtype = flow_pred.dtype
+        device = flow_pred.device
+        flow_pred, xt, sigmas, timesteps = (
+            x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps)
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1
+        )
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1)
+        x0_pred = xt - sigma_t * flow_pred
+        return x0_pred.to(original_dtype)
+
+    def step_dmd(
+        self,
+        model_output: torch.FloatTensor,
+        timestep=None,
+        sample: torch.FloatTensor = None,
+        return_dict: bool = True,
+        cur_sampling_step: int = 0,
+        dmd_noisy_tensor: torch.FloatTensor | None = None,
+        dmd_sigmas: torch.FloatTensor | None = None,
+        dmd_timesteps: torch.FloatTensor | None = None,
+        all_timesteps: torch.FloatTensor | None = None,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        pred_image_or_video = self.convert_flow_pred_to_x0(
+            flow_pred=model_output,
+            xt=sample,
+            timestep=torch.full(
+                (model_output.shape[0],),
+                timestep,
+                dtype=torch.long,
+                device=model_output.device,
+            ),
+            sigmas=dmd_sigmas,
+            timesteps=dmd_timesteps,
+        )
+        if cur_sampling_step < len(all_timesteps) - 1:
+            prev_sample = self.add_noise(
+                pred_image_or_video,
+                dmd_noisy_tensor,
+                torch.full(
+                    (model_output.shape[0],),
+                    all_timesteps[cur_sampling_step + 1],
+                    dtype=torch.long,
+                    device=model_output.device,
+                ),
+                sigmas=dmd_sigmas,
+                timesteps=dmd_timesteps,
+            )
+        else:
+            prev_sample = pred_image_or_video
+
+        if not return_dict:
+            return (prev_sample,)
+        return HeliosSchedulerOutput(prev_sample=prev_sample)
+
+    # ---------------------------------- Main step ----------------------------------
+    def step(
+        self,
+        model_output,
+        timestep=None,
+        sample=None,
+        return_dict: bool = True,
+        **kwargs,
+    ) -> HeliosSchedulerOutput | tuple:
+        if self.config.scheduler_type == "euler":
+            return self.step_euler(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+            )
+        elif self.config.scheduler_type == "unipc":
+            return self.step_unipc(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+            )
+        elif self.config.scheduler_type == "dmd":
+            return self.step_dmd(
+                model_output=model_output,
+                timestep=timestep,
+                sample=sample,
+                return_dict=return_dict,
+                **kwargs,
+            )
+        else:
+            raise NotImplementedError(
+                f"Scheduler type '{self.config.scheduler_type}' not implemented"
+            )
+
+    def reset_scheduler_history(self):
+        self.model_outputs = [None] * self.config.solver_order
+        self.timestep_list = [None] * self.config.solver_order
+        self.lower_order_nums = 0
+        self.disable_corrector = self.config.disable_corrector
+        self.solver_p = None
+        self.last_sample = None
+        self._step_index = None
+        self._begin_index = None
+
+    def set_shift(self, shift: float):
+        """Update the shift parameter (called by SchedulerLoader after loading)."""
+        self.config.shift = shift
+        self.shift = shift
+
+    def __len__(self):
+        return self.config.num_train_timesteps
+
+
+EntryClass = HeliosScheduler

python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Helios video diffusion pipeline implementation.
+
+This module contains an implementation of the Helios video diffusion pipeline
+using the modular pipeline architecture. Phase 1: T2V only.
+"""
+
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    InputValidationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import (
+    HeliosChunkedDenoisingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class HeliosPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Helios video diffusion pipeline with LoRA support.
+
+    Implements the Helios T2V pipeline with chunked denoising,
+    multi-term memory history, and CFG Zero Star guidance.
+    """
+
+    pipeline_name = "HeliosPipeline"
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def initialize_pipeline(self, server_args: ServerArgs):
+        # Use the scheduler loaded from model's scheduler_config.json as-is.
+        # It contains critical config: use_dynamic_shifting=true,
+        # time_shift_type="exponential", etc.
+        scheduler = self.modules.get("scheduler")
+        if scheduler is not None and server_args.pipeline_config.flow_shift is not None:
+            scheduler.set_shift(server_args.pipeline_config.flow_shift)
+
+        # Configure scheduler for Stage 2/3 if enabled
+        pipeline_config = server_args.pipeline_config
+        if scheduler is not None and pipeline_config.is_enable_stage2:
+            scheduler.config.stages = pipeline_config.pyramid_num_stages
+            scheduler.config.scheduler_type = pipeline_config.scheduler_type
+            scheduler.config.gamma = pipeline_config.gamma
+            scheduler.init_sigmas_for_each_stage()
+
+    def create_pipeline_stages(self, server_args: ServerArgs) -> None:
+        self.add_stage(InputValidationStage())
+        self.add_standard_text_encoding_stage()
+        self.add_standard_latent_preparation_stage()
+        # Skip standard timestep preparation — the Helios denoising stage
+        # handles scheduler.set_timesteps internally per-chunk with mu.
+        self.add_stage(
+            HeliosChunkedDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.modules["scheduler"],
+            ),
+            "helios_chunked_denoising_stage",
+        )
+        # Standard DecodingStage handles VAE decode of the denoised latents
+        self.add_standard_decoding_stage()
+
+
+class HeliosPyramidPipeline(HeliosPipeline):
+    """Helios pyramid SR pipeline (used by Helios-Mid and Helios-Distilled)."""
+
+    pipeline_name = "HeliosPyramidPipeline"
+
+
+EntryClass = [HeliosPipeline, HeliosPyramidPipeline]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py

@@ -0,0 +1,666 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Helios-specific chunked denoising stage.
+
+Implements Stage 1 chunked denoising with multi-term memory history
+and CFG Zero Star guidance. VAE decoding is handled by the standard
+DecodingStage downstream.
+"""
+
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    PipelineStage,
+    StageParallelismType,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+def optimized_scale(positive_flat, negative_flat):
+    """CFG Zero Star: compute optimal guidance scale."""
+    dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
+    squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+    return dot_product / squared_norm
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+def sample_block_noise(
+    batch_size, channel, num_frames, height, width, gamma, patch_size=(1, 2, 2)
+):
+    """Generate spatially-correlated block noise for pyramid SR."""
+    _, ph, pw = patch_size
+    block_size = ph * pw
+
+    cov = (
+        torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma
+    )
+    dist = torch.distributions.MultivariateNormal(
+        torch.zeros(block_size, device=cov.device), covariance_matrix=cov
+    )
+    block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
+
+    noise = dist.sample((block_number,))
+    noise = noise.view(
+        batch_size, channel, num_frames, height // ph, width // pw, ph, pw
+    )
+    noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(
+        batch_size, channel, num_frames, height, width
+    )
+    return noise
+
+
+class HeliosChunkedDenoisingStage(PipelineStage):
+    """
+    Helios chunked denoising stage implementing Stage 1 loop.
+
+    Iterates over video chunks, manages history buffers (short/mid/long),
+    runs transformer per chunk with CFG guidance, scheduler step,
+    and accumulates denoised latents. VAE decoding is left to DecodingStage.
+    """
+
+    def __init__(self, transformer, scheduler):
+        super().__init__()
+        self.transformer = transformer
+        self.scheduler = scheduler
+
+    @property
+    def parallelism_type(self):
+        return StageParallelismType.REPLICATED
+
+    def _denoise_one_chunk(
+        self,
+        latents,
+        prompt_embeds,
+        negative_prompt_embeds,
+        timesteps,
+        guidance_scale,
+        indices_hidden_states,
+        indices_latents_history_short,
+        indices_latents_history_mid,
+        indices_latents_history_long,
+        latents_history_short,
+        latents_history_mid,
+        latents_history_long,
+        target_dtype,
+        device,
+        is_cfg_zero_star=True,
+        use_zero_init=True,
+        zero_steps=1,
+        batch=None,
+        server_args=None,
+    ):
+        """Denoise a single chunk with full timestep loop."""
+        batch_size = latents.shape[0]
+        do_cfg = guidance_scale > 1.0
+
+        for i, t in enumerate(timesteps):
+            timestep = t.expand(batch_size)
+            latent_model_input = latents.to(target_dtype)
+
+            with set_forward_context(
+                current_timestep=t,
+                forward_batch=None,
+                attn_metadata=None,
+            ):
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    indices_hidden_states=indices_hidden_states,
+                    indices_latents_history_short=indices_latents_history_short,
+                    indices_latents_history_mid=indices_latents_history_mid,
+                    indices_latents_history_long=indices_latents_history_long,
+                    latents_history_short=(
+                        latents_history_short.to(target_dtype)
+                        if latents_history_short is not None
+                        else None
+                    ),
+                    latents_history_mid=(
+                        latents_history_mid.to(target_dtype)
+                        if latents_history_mid is not None
+                        else None
+                    ),
+                    latents_history_long=(
+                        latents_history_long.to(target_dtype)
+                        if latents_history_long is not None
+                        else None
+                    ),
+                )
+
+            if do_cfg:
+                with set_forward_context(
+                    current_timestep=t,
+                    forward_batch=None,
+                    attn_metadata=None,
+                ):
+                    noise_uncond = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        indices_hidden_states=indices_hidden_states,
+                        indices_latents_history_short=indices_latents_history_short,
+                        indices_latents_history_mid=indices_latents_history_mid,
+                        indices_latents_history_long=indices_latents_history_long,
+                        latents_history_short=(
+                            latents_history_short.to(target_dtype)
+                            if latents_history_short is not None
+                            else None
+                        ),
+                        latents_history_mid=(
+                            latents_history_mid.to(target_dtype)
+                            if latents_history_mid is not None
+                            else None
+                        ),
+                        latents_history_long=(
+                            latents_history_long.to(target_dtype)
+                            if latents_history_long is not None
+                            else None
+                        ),
+                    )
+
+                if is_cfg_zero_star:
+                    noise_pred_text = noise_pred
+                    positive_flat = noise_pred_text.reshape(batch_size, -1)
+                    negative_flat = noise_uncond.reshape(batch_size, -1)
+
+                    alpha = optimized_scale(positive_flat, negative_flat)
+                    alpha = alpha.view(
+                        batch_size, *([1] * (len(noise_pred_text.shape) - 1))
+                    )
+                    alpha = alpha.to(noise_pred_text.dtype)
+
+                    if (i <= zero_steps) and use_zero_init:
+                        noise_pred = noise_pred_text * 0.0
+                    else:
+                        noise_pred = noise_uncond * alpha + guidance_scale * (
+                            noise_pred_text - noise_uncond * alpha
+                        )
+                else:
+                    noise_pred = noise_uncond + guidance_scale * (
+                        noise_pred - noise_uncond
+                    )
+
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+        return latents
+
+    def _denoise_one_chunk_stage2(
+        self,
+        latents,
+        prompt_embeds,
+        negative_prompt_embeds,
+        guidance_scale,
+        indices_hidden_states,
+        indices_latents_history_short,
+        indices_latents_history_mid,
+        indices_latents_history_long,
+        latents_history_short,
+        latents_history_mid,
+        latents_history_long,
+        target_dtype,
+        device,
+        pyramid_num_stages,
+        pyramid_num_inference_steps_list,
+        is_distilled,
+        is_amplify_first_chunk,
+        gamma,
+        is_cfg_zero_star=True,
+        use_zero_init=True,
+        zero_steps=1,
+        batch=None,
+        server_args=None,
+    ):
+        """Denoise a single chunk using pyramid super-resolution (Stage 2)."""
+        batch_size, num_channel, num_frames, height, width = latents.shape
+        patch_size = self.transformer.patch_size
+
+        # Downsample to lowest pyramid level
+        latents = latents.permute(0, 2, 1, 3, 4).reshape(
+            batch_size * num_frames, num_channel, height, width
+        )
+        for _ in range(pyramid_num_stages - 1):
+            height //= 2
+            width //= 2
+            latents = F.interpolate(latents, size=(height, width), mode="bilinear") * 2
+        latents = latents.reshape(
+            batch_size, num_frames, num_channel, height, width
+        ).permute(0, 2, 1, 3, 4)
+
+        start_point_list = None
+        if is_distilled:
+            start_point_list = [latents]
+
+        do_cfg = guidance_scale > 1.0
+
+        for i_s in range(pyramid_num_stages):
+            # Compute mu for current resolution
+            image_seq_len = (
+                latents.shape[-1]
+                * latents.shape[-2]
+                * latents.shape[-3]
+                // (patch_size[0] * patch_size[1] * patch_size[2])
+            )
+            mu = calculate_shift(image_seq_len)
+
+            self.scheduler.set_timesteps(
+                pyramid_num_inference_steps_list[i_s],
+                i_s,
+                device=device,
+                mu=mu,
+                is_amplify_first_chunk=is_amplify_first_chunk,
+            )
+            timesteps = self.scheduler.timesteps
+
+            if i_s > 0:
+                # Upsample 2x nearest-neighbor
+                height *= 2
+                width *= 2
+                latents = latents.permute(0, 2, 1, 3, 4).reshape(
+                    batch_size * num_frames,
+                    num_channel,
+                    height // 2,
+                    width // 2,
+                )
+                latents = F.interpolate(latents, size=(height, width), mode="nearest")
+                latents = latents.reshape(
+                    batch_size, num_frames, num_channel, height, width
+                ).permute(0, 2, 1, 3, 4)
+
+                # Renoise with correlated block noise
+                ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s]
+                alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma)
+                beta = alpha * (1 - ori_sigma) / math.sqrt(gamma)
+
+                bs, ch, nf, h, w = latents.shape
+                noise = sample_block_noise(bs, ch, nf, h, w, gamma, patch_size)
+                noise = noise.to(device=device, dtype=target_dtype)
+                latents = alpha * latents + beta * noise
+
+                if is_distilled:
+                    start_point_list.append(latents)
+
+            # Denoising loop for this pyramid stage
+            for idx, t in enumerate(timesteps):
+                timestep = t.expand(batch_size)
+                latent_model_input = latents.to(target_dtype)
+
+                with set_forward_context(
+                    current_timestep=t,
+                    forward_batch=None,
+                    attn_metadata=None,
+                ):
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        indices_hidden_states=indices_hidden_states,
+                        indices_latents_history_short=indices_latents_history_short,
+                        indices_latents_history_mid=indices_latents_history_mid,
+                        indices_latents_history_long=indices_latents_history_long,
+                        latents_history_short=(
+                            latents_history_short.to(target_dtype)
+                            if latents_history_short is not None
+                            else None
+                        ),
+                        latents_history_mid=(
+                            latents_history_mid.to(target_dtype)
+                            if latents_history_mid is not None
+                            else None
+                        ),
+                        latents_history_long=(
+                            latents_history_long.to(target_dtype)
+                            if latents_history_long is not None
+                            else None
+                        ),
+                    )
+
+                if do_cfg:
+                    with set_forward_context(
+                        current_timestep=t,
+                        forward_batch=None,
+                        attn_metadata=None,
+                    ):
+                        noise_uncond = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            indices_hidden_states=indices_hidden_states,
+                            indices_latents_history_short=indices_latents_history_short,
+                            indices_latents_history_mid=indices_latents_history_mid,
+                            indices_latents_history_long=indices_latents_history_long,
+                            latents_history_short=(
+                                latents_history_short.to(target_dtype)
+                                if latents_history_short is not None
+                                else None
+                            ),
+                            latents_history_mid=(
+                                latents_history_mid.to(target_dtype)
+                                if latents_history_mid is not None
+                                else None
+                            ),
+                            latents_history_long=(
+                                latents_history_long.to(target_dtype)
+                                if latents_history_long is not None
+                                else None
+                            ),
+                        )
+
+                    if is_cfg_zero_star:
+                        noise_pred_text = noise_pred
+                        positive_flat = noise_pred_text.reshape(batch_size, -1)
+                        negative_flat = noise_uncond.reshape(batch_size, -1)
+
+                        alpha_cfg = optimized_scale(positive_flat, negative_flat)
+                        alpha_cfg = alpha_cfg.view(
+                            batch_size,
+                            *([1] * (len(noise_pred_text.shape) - 1)),
+                        )
+                        alpha_cfg = alpha_cfg.to(noise_pred_text.dtype)
+
+                        if (i_s == 0 and idx <= zero_steps) and use_zero_init:
+                            noise_pred = noise_pred_text * 0.0
+                        else:
+                            noise_pred = noise_uncond * alpha_cfg + guidance_scale * (
+                                noise_pred_text - noise_uncond * alpha_cfg
+                            )
+                    else:
+                        noise_pred = noise_uncond + guidance_scale * (
+                            noise_pred - noise_uncond
+                        )
+
+                latents = self.scheduler.step(
+                    noise_pred,
+                    t,
+                    latents,
+                    return_dict=False,
+                    cur_sampling_step=idx,
+                    dmd_noisy_tensor=(
+                        start_point_list[i_s] if start_point_list is not None else None
+                    ),
+                    dmd_sigmas=self.scheduler.sigmas,
+                    dmd_timesteps=self.scheduler.timesteps,
+                    all_timesteps=timesteps,
+                )[0]
+
+        return latents
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """Run the Helios chunked denoising loop."""
+        pipeline_config = server_args.pipeline_config
+        device = (
+            batch.latents.device
+            if hasattr(batch, "latents") and batch.latents is not None
+            else torch.device("cuda")
+        )
+        target_dtype = PRECISION_TO_TYPE.get(
+            server_args.pipeline_config.precision, torch.bfloat16
+        )
+
+        # Get config params
+        num_latent_frames_per_chunk = pipeline_config.num_latent_frames_per_chunk
+        history_sizes = sorted(list(pipeline_config.history_sizes), reverse=True)
+        is_cfg_zero_star = pipeline_config.is_cfg_zero_star
+        zero_steps = pipeline_config.zero_steps
+        keep_first_frame = pipeline_config.keep_first_frame
+        guidance_scale = batch.guidance_scale
+        num_inference_steps = batch.num_inference_steps
+
+        # Stage 2 params
+        is_enable_stage2 = pipeline_config.is_enable_stage2
+        pyramid_num_stages = pipeline_config.pyramid_num_stages
+        pyramid_num_inference_steps_list = (
+            pipeline_config.pyramid_num_inference_steps_list
+        )
+        is_distilled = pipeline_config.is_distilled
+        is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk
+        gamma = pipeline_config.gamma
+
+        # Move transformer to GPU if CPU-offloaded
+        if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
+            if next(self.transformer.parameters()).device.type == "cpu":
+                self.transformer.to(get_local_torch_device())
+
+        # Get encoder outputs (prompt_embeds is a list of tensors, one per encoder)
+        prompt_embeds = batch.prompt_embeds
+        if isinstance(prompt_embeds, list):
+            prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.to(target_dtype)
+        negative_prompt_embeds = batch.negative_prompt_embeds
+        if isinstance(negative_prompt_embeds, list):
+            negative_prompt_embeds = (
+                negative_prompt_embeds[0] if negative_prompt_embeds else None
+            )
+        if negative_prompt_embeds is not None:
+            negative_prompt_embeds = negative_prompt_embeds.to(target_dtype)
+
+        # Scale factors inherited from the Wan VAE used by Helios
+        # (AutoencoderKLWan: temporal_compression_ratio=4, spatial_compression_ratio=8)
+        vae_scale_factor_temporal = 4
+        vae_scale_factor_spatial = 8
+
+        # Compute chunking
+        height = batch.height
+        width = batch.width
+        num_frames = batch.num_frames
+        num_channels_latents = self.transformer.in_channels
+
+        window_num_frames = (
+            num_latent_frames_per_chunk - 1
+        ) * vae_scale_factor_temporal + 1
+        num_latent_chunk = max(
+            1, (num_frames + window_num_frames - 1) // window_num_frames
+        )
+        num_history_latent_frames = sum(history_sizes)
+        batch_size = 1  # Helios processes one video at a time
+
+        # Prepare history latents
+        if not keep_first_frame:
+            history_sizes[-1] = history_sizes[-1] + 1
+        history_latents = torch.zeros(
+            batch_size,
+            num_channels_latents,
+            num_history_latent_frames,
+            height // vae_scale_factor_spatial,
+            width // vae_scale_factor_spatial,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        # Build frame indices
+        if keep_first_frame:
+            indices = torch.arange(
+                0, sum([1, *history_sizes, num_latent_frames_per_chunk])
+            )
+            (
+                indices_prefix,
+                indices_latents_history_long,
+                indices_latents_history_mid,
+                indices_latents_history_1x,
+                indices_hidden_states,
+            ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0)
+            indices_latents_history_short = torch.cat(
+                [indices_prefix, indices_latents_history_1x], dim=0
+            )
+        else:
+            indices = torch.arange(
+                0, sum([*history_sizes, num_latent_frames_per_chunk])
+            )
+            (
+                indices_latents_history_long,
+                indices_latents_history_mid,
+                indices_latents_history_short,
+                indices_hidden_states,
+            ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0)
+
+        indices_hidden_states = indices_hidden_states.unsqueeze(0)
+        indices_latents_history_short = indices_latents_history_short.unsqueeze(0)
+        indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0)
+        indices_latents_history_long = indices_latents_history_long.unsqueeze(0)
+
+        # Set up scheduler
+        patch_size = self.transformer.patch_size
+        image_seq_len = (
+            num_latent_frames_per_chunk
+            * (height // vae_scale_factor_spatial)
+            * (width // vae_scale_factor_spatial)
+            // (patch_size[0] * patch_size[1] * patch_size[2])
+        )
+        # Sigma schedule from near-1.0 (pure noise) to 0.0 (clean); 0.999 avoids singularity
+        sigmas = np.linspace(0.999, 0.0, num_inference_steps + 1)[:-1]
+        mu = calculate_shift(image_seq_len)
+
+        # Chunk loop
+        image_latents = None
+        total_generated_latent_frames = 0
+
+        self.log_info(
+            f"Starting chunked denoising: {num_latent_chunk} chunks, "
+            f"{num_inference_steps} steps each"
+        )
+
+        for k in range(num_latent_chunk):
+            is_first_chunk = k == 0
+
+            # Extract history
+            if keep_first_frame:
+                (
+                    latents_history_long,
+                    latents_history_mid,
+                    latents_history_1x,
+                ) = history_latents[:, :, -num_history_latent_frames:].split(
+                    history_sizes, dim=2
+                )
+                if image_latents is None and is_first_chunk:
+                    latents_prefix = torch.zeros(
+                        (
+                            batch_size,
+                            num_channels_latents,
+                            1,
+                            latents_history_1x.shape[-2],
+                            latents_history_1x.shape[-1],
+                        ),
+                        device=device,
+                        dtype=latents_history_1x.dtype,
+                    )
+                else:
+                    latents_prefix = image_latents
+                latents_history_short = torch.cat(
+                    [latents_prefix, latents_history_1x], dim=2
+                )
+            else:
+                (
+                    latents_history_long,
+                    latents_history_mid,
+                    latents_history_short,
+                ) = history_latents[:, :, -num_history_latent_frames:].split(
+                    history_sizes, dim=2
+                )
+
+            # Generate noise latents for this chunk
+            latents = torch.randn(
+                batch_size,
+                num_channels_latents,
+                (window_num_frames - 1) // vae_scale_factor_temporal + 1,
+                height // vae_scale_factor_spatial,
+                width // vae_scale_factor_spatial,
+                device=device,
+                dtype=torch.float32,
+            )
+
+            if is_enable_stage2:
+                # Stage 2: Pyramid SR denoising (handles scheduler internally)
+                latents = self._denoise_one_chunk_stage2(
+                    latents=latents,
+                    prompt_embeds=prompt_embeds,
+                    negative_prompt_embeds=negative_prompt_embeds,
+                    guidance_scale=guidance_scale,
+                    indices_hidden_states=indices_hidden_states,
+                    indices_latents_history_short=indices_latents_history_short,
+                    indices_latents_history_mid=indices_latents_history_mid,
+                    indices_latents_history_long=indices_latents_history_long,
+                    latents_history_short=latents_history_short,
+                    latents_history_mid=latents_history_mid,
+                    latents_history_long=latents_history_long,
+                    target_dtype=target_dtype,
+                    device=device,
+                    pyramid_num_stages=pyramid_num_stages,
+                    pyramid_num_inference_steps_list=pyramid_num_inference_steps_list,
+                    is_distilled=is_distilled,
+                    is_amplify_first_chunk=(is_amplify_first_chunk and is_first_chunk),
+                    gamma=gamma,
+                    is_cfg_zero_star=is_cfg_zero_star,
+                    use_zero_init=True,
+                    zero_steps=zero_steps,
+                    batch=batch,
+                    server_args=server_args,
+                )
+            else:
+                # Stage 1: Standard flat denoising
+                self.scheduler.set_timesteps(
+                    num_inference_steps, device=device, sigmas=sigmas, mu=mu
+                )
+                timesteps = self.scheduler.timesteps
+
+                latents = self._denoise_one_chunk(
+                    latents=latents,
+                    prompt_embeds=prompt_embeds,
+                    negative_prompt_embeds=negative_prompt_embeds,
+                    timesteps=timesteps,
+                    guidance_scale=guidance_scale,
+                    indices_hidden_states=indices_hidden_states,
+                    indices_latents_history_short=indices_latents_history_short,
+                    indices_latents_history_mid=indices_latents_history_mid,
+                    indices_latents_history_long=indices_latents_history_long,
+                    latents_history_short=latents_history_short,
+                    latents_history_mid=latents_history_mid,
+                    latents_history_long=latents_history_long,
+                    target_dtype=target_dtype,
+                    device=device,
+                    is_cfg_zero_star=is_cfg_zero_star,
+                    use_zero_init=True,
+                    zero_steps=zero_steps,
+                    batch=batch,
+                    server_args=server_args,
+                )
+
+            # Extract first frame as image_latents for subsequent chunks
+            if keep_first_frame and is_first_chunk and image_latents is None:
+                image_latents = latents[:, :, 0:1, :, :]
+
+            # Update history
+            total_generated_latent_frames += latents.shape[2]
+            history_latents = torch.cat([history_latents, latents], dim=2)
+
+        # Move transformer back to CPU after denoising
+        if server_args.dit_cpu_offload and not server_args.use_fsdp_inference:
+            if next(self.transformer.parameters()).device.type != "cpu":
+                self.transformer.to("cpu")
+        torch.cuda.empty_cache()
+
+        # Store denoised latents for the standard DecodingStage to decode
+        batch.latents = history_latents[:, :, -total_generated_latent_frames:]
+
+        return batch