# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
"""
Denoising stage for diffusion pipelines.
"""

import inspect
import math
import time
import weakref
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass, field, fields
from enum import Enum
from functools import lru_cache
from typing import Any

import torch
import torch.nn as nn

from sglang.multimodal_gen import envs
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode
from sglang.multimodal_gen.configs.pipeline_configs.flux import (
    Flux2PipelineConfig,
    FluxPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
    prompt_padding as bcg_utils,
)
from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
    CacheDitConfig,
    enable_cache_on_dual_transformer,
    enable_cache_on_transformer,
    get_scm_mask,
    refresh_context_on_dual_transformer,
    refresh_context_on_transformer,
)
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.distributed import (
    get_local_torch_device,
    get_sp_group,
    get_sp_world_size,
    get_tp_group,
    get_world_group,
    get_world_size,
    model_parallel_is_initialized,
)
from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
    run_cfg_parallel,
    run_two_branch_cfg_parallel,
)
from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
    CFGPolicy,
    _unwrap,
    _wrap,
)
from sglang.multimodal_gen.runtime.distributed.communication_op import (
    sequence_model_parallel_all_gather,
)
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
    get_classifier_free_guidance_world_size,
    world_group_is_initialized,
)
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
    configure_sta,
    save_mask_search_results,
)
from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
    TransformerLoader,
)
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import (
    ComponentUse,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
    LayerwiseOffloadableModuleMixin,
    is_layerwise_offloaded_module,
)
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.pipelines_core.stages.model_specific_stages.wan_ti2v import (
    blend_wan_ti2v_latents,
    expand_wan_ti2v_timestep,
    prepare_wan_ti2v_latents,
    prepare_wan_ti2v_sp_inputs,
    should_apply_wan_ti2v,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
    StageValidators as V,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
    VerificationResult,
)
from sglang.multimodal_gen.runtime.platforms import (
    AttentionBackendEnum,
    current_platform,
)
from sglang.multimodal_gen.runtime.post_training.rollout_denoising_mixin import (
    RolloutDenoisingMixin,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
from sglang.multimodal_gen.runtime.utils.precision import (
    autocast_enabled as precision_autocast_enabled,
)
from sglang.multimodal_gen.runtime.utils.precision import (
    resolve_precision,
)
from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
from sglang.multimodal_gen.runtime.utils.torch_compile import (
    CompiledModuleRegistry,
    build_torch_compile_kwargs,
    maybe_enable_inductor_compute_comm_overlap,
    resolve_torch_compile_mode,
)
from sglang.multimodal_gen.utils import dict_to_3d_list

logger = init_logger(__name__)


def _ensure_tensor_model_output(model_output):
    sample = getattr(model_output, "sample", None)
    if isinstance(sample, torch.Tensor):
        return sample
    return model_output


@dataclass(slots=True)
class DenoisingContext:
    """Loop-scoped state shared across the denoising skeleton and its hooks."""

    scheduler: Any
    extra_step_kwargs: dict[str, Any]
    target_dtype: torch.dtype
    autocast_enabled: bool
    timesteps: torch.Tensor
    num_inference_steps: int
    num_warmup_steps: int
    image_kwargs: dict[str, Any]
    pos_cond_kwargs: dict[str, Any]
    neg_cond_kwargs: dict[str, Any]
    latents: torch.Tensor
    boundary_timestep: float | None
    z: torch.Tensor | None
    reserved_frames_mask: torch.Tensor | None
    seq_len: int | None
    guidance: torch.Tensor
    is_warmup: bool
    cfg_policy: CFGPolicy | None = None
    trajectory_timesteps: list[torch.Tensor] = field(default_factory=list)
    trajectory_latents: list[torch.Tensor] = field(default_factory=list)
    extra: dict[str, Any] = field(default_factory=dict)

    def __getitem__(self, key: str) -> Any:
        return getattr(self, key)

    def get(self, key: str, default: Any = None) -> Any:
        return getattr(self, key, default)

    def to_kwargs(self) -> dict[str, Any]:
        """Return a shallow field mapping for derived context construction."""
        return {item.name: getattr(self, item.name) for item in fields(self)}


@dataclass(slots=True)
class DenoisingStepState:
    """Per-step hot-path state computed once and reused within a denoising step."""

    step_index: int
    t_host: torch.Tensor
    t_device: torch.Tensor
    t_int: int
    current_model: Any
    current_guidance_scale: Any
    attn_metadata: Any | None


class DualTransformerExecutionMode(str, Enum):
    """How a denoising stage uses a second DiT.

    BOUNDARY_EXPERTS means one transformer is selected per timestep.
    PAIRED_PER_STEP means both transformers participate in each denoising step.
    """

    BOUNDARY_EXPERTS = "boundary_experts"
    PAIRED_PER_STEP = "paired_per_step"


class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
    """
    Stage for running the denoising loop in diffusion pipelines.

    This stage handles the iterative denoising process that transforms
    the initial noise into the final output.
    """

    @property
    def role_affinity(self):
        return RoleType.DENOISER

    def __init__(
        self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None
    ) -> None:
        super().__init__()
        self.transformer = transformer
        self.transformer_2 = transformer_2
        # cache-dit state (for delayed mounting and idempotent control)
        self._cache_dit_enabled = False
        self._cached_num_steps = None
        self._torch_compile_registry = CompiledModuleRegistry()
        # Breakable CUDA graph runners, one per transformer module (lazy).
        self._bcg_runners: dict[int, Any] = {}

        hidden_size = self.server_args.pipeline_config.dit_config.hidden_size
        num_attention_heads = (
            self.server_args.pipeline_config.dit_config.num_attention_heads
        )
        attn_head_size = hidden_size // num_attention_heads

        # torch compile
        # list of offloaded dit modules if torch compile is enabled. cleared after compile and warmup
        self._offloaded_dit_modules_for_compile: list[torch.nn.Module] = []
        # layerwise-offload and then compile to avoid OOM
        for transformer in filter(None, [self.transformer, self.transformer_2]):
            self._maybe_offload_during_compile(transformer)
            self._maybe_torch_compile(transformer)

        self.scheduler = scheduler
        self.vae = vae
        self.pipeline = weakref.ref(pipeline) if pipeline else None

        selected_attention_backend = self._infer_transformer_attention_backend()
        # precision-constraint: attention backend metadata allocation currently assumes fp16;
        # do not replace with user precision policy without auditing backend support.
        self.attn_backend = get_attn_backend(
            head_size=attn_head_size,
            dtype=torch.float16,
            selected_attention_backend=selected_attention_backend,
        )

        # cfg
        self.guidance = None

        # misc
        self.profiler = None
        self._is_warmed_up = False
        self._extra_func_kwarg_names_cache: dict[int, tuple[bool, frozenset[str]]] = {}

    def _infer_transformer_attention_backend(self) -> AttentionBackendEnum | None:
        backends = {
            backend
            for transformer in (self.transformer, self.transformer_2)
            if transformer is not None
            for module in transformer.modules()
            if isinstance(
                (backend := getattr(module, "backend", None)), AttentionBackendEnum
            )
        }
        if not backends:
            return None
        if len(backends) > 1:
            sparse_backends = {backend for backend in backends if backend.is_sparse}
            selected_backend = (
                sorted(sparse_backends, key=lambda backend: backend.name)[0]
                if sparse_backends
                else sorted(backends, key=lambda backend: backend.name)[0]
            )
            logger.warning(
                "Multiple transformer attention backends detected: %s. "
                "Using %s for denoising metadata.",
                sorted(backend.name.lower() for backend in backends),
                selected_backend.name.lower(),
            )
            return selected_backend
        return next(iter(backends))

    def component_uses(
        self, server_args: ServerArgs, stage_name: str | None = None
    ) -> list[ComponentUse]:
        stage_name = self._component_stage_name(stage_name)
        uses: list[ComponentUse] = []
        if self.vae is not None:
            vae_dtype = resolve_precision(
                server_args, "vae", precision_attr="vae_precision"
            )
            uses.append(
                ComponentUse(
                    stage_name=stage_name,
                    component_name="vae",
                    target_dtype=vae_dtype,
                )
            )
        for default_name, module in (
            ("transformer", self.transformer),
            ("transformer_2", self.transformer_2),
        ):
            if module is None:
                continue
            component_name = self._component_name_for_stage_module(module, default_name)
            uses.append(
                ComponentUse(
                    stage_name=stage_name,
                    component_name=component_name,
                    phase=component_name,
                    preferred_ready_after_request=component_name == "transformer",
                    memory_intensive=True,
                )
            )
        return uses

    def _component_name_for_stage_module(
        self, module: nn.Module | None, default_name: str
    ) -> str:
        pipeline = self.pipeline() if self.pipeline else None
        if pipeline is None or module is None:
            return default_name
        for name, candidate in pipeline.modules.items():
            if candidate is module:
                return name
        return default_name

    def _maybe_offload_during_compile(self, module: object) -> None:
        """
        Layerwise-offload the DiT so the compile warmup autotunes each layer
        with only itself resident; forward() restores it after the warmup.
        """
        args = self.server_args
        if (
            not args.enable_torch_compile
            or not args.offload_during_compile
            or not args.warmup
            # a subclass with its own forward would never run the restore
            or type(self).forward is not DenoisingStage.forward
            or args.use_fsdp_inference
            or envs.SGLANG_CACHE_DIT_ENABLED
            or not isinstance(module, LayerwiseOffloadableModuleMixin)
            or is_layerwise_offloaded_module(module)
        ):
            return
        module.configure_layerwise_offload(args)
        if is_layerwise_offloaded_module(module):
            self._offloaded_dit_modules_for_compile.append(module)

    def _move_resident_components_for_warmup(self) -> list[torch.nn.Module]:
        """Move resident non-DiT components off-device while the warmup
        denoising (the compile/autotune peak) runs; forward() moves them back."""
        pipeline = self.pipeline() if self.pipeline is not None else None
        if pipeline is None:
            return []
        dit_ids = {id(self.transformer), id(self.transformer_2)}
        moved = []
        for module in pipeline.modules.values():
            if (
                isinstance(module, torch.nn.Module)
                and id(module) not in dit_ids
                and not is_layerwise_offloaded_module(module)
            ):
                param = next(module.parameters(), None)
                if param is not None and param.device.type != "cpu":
                    module.to("cpu")
                    moved.append(module)
        return moved

    def _maybe_torch_compile(self, module: object) -> None:
        """
        Compile a module with torch.compile, and enable inductor overlap tweak if available.
        No-op if torch compile is disabled or the object is not a nn.Module.
        """
        if self.server_args.enable_breakable_cuda_graph:
            # BCG captures the eager kernel stream itself; compiling first
            # would capture inductor's own cudagraph trees / guards.
            return
        if not getattr(
            self.server_args, "enable_torch_compile", False
        ) or not isinstance(module, nn.Module):
            return
        if envs.SGLANG_CACHE_DIT_ENABLED and not self._cache_dit_enabled:
            logger.debug("Deferring torch.compile until cache-dit is enabled")
            return
        if self._torch_compile_registry.is_compiled(module):
            return

        if current_platform.is_npu():
            compile_kwargs = build_torch_compile_kwargs(mode=None)
            logger.info("Compiling transformer with torchair backend on NPU")
        else:
            maybe_enable_inductor_compute_comm_overlap()
            dit_config = getattr(self.server_args.pipeline_config, "dit_config", None)
            mode = resolve_torch_compile_mode(
                "SGLANG_TORCH_COMPILE_MODE",
                config=dit_config,
                default="max-autotune-no-cudagraphs",
            )
            compile_kwargs = build_torch_compile_kwargs(mode=mode)
            logger.info(f"Compiling transformer with mode: {mode}")

        # TODO(triple-mu): support customized fullgraph and dynamic in the future
        self._torch_compile_registry.compile_once(
            module,
            compile_kwargs=compile_kwargs,
        )

    def _maybe_enable_cache_dit_and_torch_compile(
        self, num_inference_steps: int | tuple[int, int], batch: Req
    ) -> None:
        """Apply request-dependent transformer acceleration in trace-safe order."""
        self._maybe_enable_cache_dit(num_inference_steps, batch)
        for transformer in filter(None, [self.transformer, self.transformer_2]):
            self._maybe_torch_compile(transformer)

    def _cache_dit_dual_model_name(self) -> str:
        return "wan2.2"

    def _cache_dit_secondary_uses_primary_config(self) -> bool:
        return False

    def _dual_transformer_execution_mode(self) -> DualTransformerExecutionMode | None:
        if self.transformer_2 is None:
            return None
        return DualTransformerExecutionMode.BOUNDARY_EXPERTS

    def _cache_dit_step_counts(
        self, num_inference_steps: int | tuple[int, int]
    ) -> tuple[int, int | None]:
        if isinstance(num_inference_steps, tuple):
            primary_steps, secondary_steps = num_inference_steps
            return int(primary_steps), int(secondary_steps)
        steps = int(num_inference_steps)
        mode = self._dual_transformer_execution_mode()
        if mode is None:
            return steps, None
        if mode == DualTransformerExecutionMode.PAIRED_PER_STEP:
            return steps, steps
        raise ValueError("Boundary-expert dual transformers require split step counts.")

    @staticmethod
    def _parse_cache_dit_scm_bins() -> tuple[list[int] | None, list[int] | None, str]:
        scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET
        compute_bins_str = envs.SGLANG_CACHE_DIT_SCM_COMPUTE_BINS
        cache_bins_str = envs.SGLANG_CACHE_DIT_SCM_CACHE_BINS
        compute_bins = None
        cache_bins = None
        if compute_bins_str and cache_bins_str:
            try:
                compute_bins = [int(x.strip()) for x in compute_bins_str.split(",")]
                cache_bins = [int(x.strip()) for x in cache_bins_str.split(",")]
            except ValueError as exc:
                logger.warning("Failed to parse SCM bins: %s. SCM disabled.", exc)
                scm_preset = "none"
        elif compute_bins_str or cache_bins_str:
            logger.warning(
                "SCM custom bins require both compute_bins and cache_bins. "
                "Only one was provided (compute=%s, cache=%s). Falling back to preset '%s'.",
                compute_bins_str,
                cache_bins_str,
                scm_preset,
            )
        return compute_bins, cache_bins, scm_preset

    def _cache_dit_scm_masks(
        self, primary_num_steps: int, secondary_num_steps: int | None = None
    ) -> tuple[str, str, list[int] | None, list[int] | None]:
        scm_compute_bins, scm_cache_bins, scm_preset = self._parse_cache_dit_scm_bins()
        scm_policy = envs.SGLANG_CACHE_DIT_SCM_POLICY
        steps_computation_mask = get_scm_mask(
            preset=scm_preset,
            num_inference_steps=primary_num_steps,
            compute_bins=scm_compute_bins,
            cache_bins=scm_cache_bins,
        )

        steps_computation_mask_2 = None
        if secondary_num_steps is not None:
            if (
                self._cache_dit_secondary_uses_primary_config()
                and secondary_num_steps == primary_num_steps
            ):
                steps_computation_mask_2 = steps_computation_mask
            else:
                steps_computation_mask_2 = get_scm_mask(
                    preset=scm_preset,
                    num_inference_steps=secondary_num_steps,
                    compute_bins=scm_compute_bins,
                    cache_bins=scm_cache_bins,
                )
        return scm_preset, scm_policy, steps_computation_mask, steps_computation_mask_2

    @staticmethod
    def _build_cache_dit_config(
        num_inference_steps: int,
        steps_computation_mask: list[int] | None,
        scm_policy: str,
        *,
        secondary: bool = False,
    ) -> CacheDitConfig:
        return CacheDitConfig(
            enabled=True,
            Fn_compute_blocks=(
                envs.SGLANG_CACHE_DIT_SECONDARY_FN
                if secondary
                else envs.SGLANG_CACHE_DIT_FN
            ),
            Bn_compute_blocks=(
                envs.SGLANG_CACHE_DIT_SECONDARY_BN
                if secondary
                else envs.SGLANG_CACHE_DIT_BN
            ),
            max_warmup_steps=(
                envs.SGLANG_CACHE_DIT_SECONDARY_WARMUP
                if secondary
                else envs.SGLANG_CACHE_DIT_WARMUP
            ),
            residual_diff_threshold=(
                envs.SGLANG_CACHE_DIT_SECONDARY_RDT
                if secondary
                else envs.SGLANG_CACHE_DIT_RDT
            ),
            max_continuous_cached_steps=(
                envs.SGLANG_CACHE_DIT_SECONDARY_MC
                if secondary
                else envs.SGLANG_CACHE_DIT_MC
            ),
            enable_taylorseer=(
                envs.SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER
                if secondary
                else envs.SGLANG_CACHE_DIT_TAYLORSEER
            ),
            taylorseer_order=(
                envs.SGLANG_CACHE_DIT_SECONDARY_TS_ORDER
                if secondary
                else envs.SGLANG_CACHE_DIT_TS_ORDER
            ),
            num_inference_steps=num_inference_steps,
            steps_computation_mask=steps_computation_mask,
            steps_computation_policy=scm_policy,
        )

    def _maybe_enable_cache_dit(
        self, num_inference_steps: int | tuple[int, int], batch: Req
    ) -> None:
        """Enable cache-dit on the transformers if configured (idempotent).

        This method should be called after the transformer is fully loaded
        and before torch.compile is applied.

        For dual-transformer models (e.g., Wan2.2), this enables cache-dit on both
        transformers with (potentially) different configurations.

        """
        if self.server_args.enable_breakable_cuda_graph:
            # Cache-DiT wraps transformer.forward with step-skipping control
            # flow that must not be baked into a captured CUDA graph.
            return
        # NOTE: When a new request arrives, we need to refresh the cache-dit context.
        if self._cache_dit_enabled:
            primary_num_steps, secondary_num_steps = self._cache_dit_step_counts(
                num_inference_steps
            )
            scm_preset, scm_policy, steps_computation_mask, steps_computation_mask_2 = (
                self._cache_dit_scm_masks(primary_num_steps, secondary_num_steps)
            )
            if self.transformer_2 is not None:
                assert secondary_num_steps is not None
                refresh_context_on_dual_transformer(
                    self.transformer,
                    self.transformer_2,
                    primary_num_steps,
                    secondary_num_steps,
                    steps_computation_mask=steps_computation_mask,
                    steps_computation_mask_2=steps_computation_mask_2,
                    steps_computation_policy=scm_policy,
                )
            else:
                scm_preset = None if scm_preset == "none" else scm_preset
                refresh_context_on_transformer(
                    self.transformer,
                    primary_num_steps,
                    scm_preset=scm_preset,
                )
            return

        # Keep cache-dit disabled for ordinary warmup, but allow torch.compile
        # warmup to mount cache-dit before Dynamo traces the transformer.
        if not envs.SGLANG_CACHE_DIT_ENABLED:
            return
        if batch.is_warmup and not getattr(
            self.server_args, "enable_torch_compile", False
        ):
            return

        primary_num_steps, secondary_num_steps = self._cache_dit_step_counts(
            num_inference_steps
        )
        world_size = get_world_size()
        parallelized = world_size > 1

        sp_group = None
        tp_group = None
        if parallelized:
            sp_group_candidate = get_sp_group()
            tp_group_candidate = get_tp_group()

            sp_world_size = sp_group_candidate.world_size if sp_group_candidate else 1
            tp_world_size = tp_group_candidate.world_size if tp_group_candidate else 1

            has_sp = sp_world_size > 1
            has_tp = tp_world_size > 1

            sp_group = sp_group_candidate.device_group if has_sp else None
            tp_group = tp_group_candidate.device_group if has_tp else None

            logger.info(
                "cache-dit enabled in distributed environment (world_size=%d, has_sp=%s, has_tp=%s)",
                world_size,
                has_sp,
                has_tp,
            )
        _, scm_policy, steps_computation_mask, steps_computation_mask_2 = (
            self._cache_dit_scm_masks(primary_num_steps, secondary_num_steps)
        )
        primary_config = self._build_cache_dit_config(
            primary_num_steps,
            steps_computation_mask=steps_computation_mask,
            scm_policy=scm_policy,
        )

        if self.transformer_2 is not None:
            # dual transformer
            # build config for secondary transformer (low-noise expert)
            # uses secondary parameters which inherit from primary if not explicitly set
            assert secondary_num_steps is not None
            secondary_config = (
                primary_config
                if self._cache_dit_secondary_uses_primary_config()
                else self._build_cache_dit_config(
                    secondary_num_steps,
                    steps_computation_mask=steps_computation_mask_2,
                    scm_policy=scm_policy,
                    secondary=True,
                )
            )

            # for dual transformers, must use BlockAdapter to enable cache on both simultaneously.
            # Don't call enable_cache separately on each transformer.
            self.transformer, self.transformer_2 = enable_cache_on_dual_transformer(
                self.transformer,
                self.transformer_2,
                primary_config,
                secondary_config,
                model_name=self._cache_dit_dual_model_name(),
                sp_group=sp_group,
                tp_group=tp_group,
            )
            logger.info(
                "cache-dit enabled on dual transformers (steps=%d, %d)",
                primary_num_steps,
                secondary_num_steps,
            )
        else:
            # single transformer
            self.transformer = enable_cache_on_transformer(
                self.transformer,
                primary_config,
                model_name="transformer",
                sp_group=sp_group,
                tp_group=tp_group,
                has_separate_cfg=batch.do_classifier_free_guidance,
            )
            logger.info(
                "cache-dit enabled on transformer (steps=%d, Fn=%d, Bn=%d, rdt=%.3f)",
                primary_num_steps,
                envs.SGLANG_CACHE_DIT_FN,
                envs.SGLANG_CACHE_DIT_BN,
                envs.SGLANG_CACHE_DIT_RDT,
            )

        self._cache_dit_enabled = True
        self._cached_num_steps = num_inference_steps

    @lru_cache(maxsize=8)
    def _build_guidance(self, batch_size, target_dtype, device, guidance_val):
        """Builds a guidance tensor. This method is cached."""
        if isinstance(
            self.server_args.pipeline_config, FluxPipelineConfig
        ) and not isinstance(self.server_args.pipeline_config, Flux2PipelineConfig):
            guidance_val = guidance_val * 1000.0
        return torch.full(
            (batch_size,),
            guidance_val,
            dtype=target_dtype,
            device=device,
        )

    def get_or_build_guidance(self, bsz: int, dtype, device):
        """
        Get the guidance tensor, using a cached version if available.

        This method retrieves a cached guidance tensor using `_build_guidance`.
        The caching is based on batch size, dtype, device, and the guidance value,
        preventing repeated tensor creation within the denoising loop.
        """
        if self.server_args.pipeline_config.should_use_guidance:
            # TODO: should the guidance_scale be picked-up from sampling_params?
            guidance_val = self.server_args.pipeline_config.embedded_cfg_scale
            return self._build_guidance(bsz, dtype, device, guidance_val)
        else:
            return None

    @property
    def parallelism_type(self) -> StageParallelismType:
        # return StageParallelismType.CFG_PARALLEL if get_global_server_args().enable_cfg_parallel else StageParallelismType.REPLICATED
        return StageParallelismType.REPLICATED

    def _handle_boundary_ratio(
        self,
        server_args,
        batch,
        scheduler,
    ):
        """
        (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
        """
        boundary_ratio = server_args.pipeline_config.dit_config.boundary_ratio
        if batch.boundary_ratio is not None:
            logger.info(
                "Overriding boundary ratio from %s to %s",
                boundary_ratio,
                batch.boundary_ratio,
            )
            boundary_ratio = batch.boundary_ratio

        if boundary_ratio is not None:
            num_train_timesteps = getattr(scheduler, "num_train_timesteps", None)
            if num_train_timesteps is None:
                num_train_timesteps = scheduler.config.num_train_timesteps
            boundary_timestep = boundary_ratio * num_train_timesteps
        else:
            boundary_timestep = None

        return boundary_timestep

    def _prepare_denoising_loop(self, batch: Req, server_args: ServerArgs):
        """
        Prepare all necessary invariant variables for the denoising loop.

        Returns:
            A context object containing the invariant state for the denoising loop.
        """
        batch = server_args.pipeline_config.expand_conditioning_to_sample_batch(batch)

        assert self.transformer is not None
        pipeline = self.pipeline() if self.pipeline else None
        scheduler = batch.scheduler
        assert scheduler is not None

        dual_transformer_mode = self._dual_transformer_execution_mode()
        uses_boundary_transformer_2 = (
            dual_transformer_mode == DualTransformerExecutionMode.BOUNDARY_EXPERTS
        )
        boundary_timestep = (
            self._handle_boundary_ratio(server_args, batch, scheduler)
            if uses_boundary_transformer_2
            else None
        )
        # Get timesteps and calculate warmup steps
        timesteps = batch.timesteps
        num_inference_steps = batch.num_inference_steps
        num_warmup_steps = len(timesteps) - num_inference_steps * scheduler.order

        if uses_boundary_transformer_2:
            assert boundary_timestep is not None, "boundary_timestep must be provided"
            num_high_noise_steps = (timesteps >= boundary_timestep).sum().item()
            num_low_noise_steps = num_inference_steps - num_high_noise_steps
            cache_dit_num_inference_steps = (num_high_noise_steps, num_low_noise_steps)
        else:
            cache_dit_num_inference_steps = num_inference_steps

        transformer_was_loaded = server_args.model_loaded["transformer"]
        if not transformer_was_loaded:
            # FIXME: reuse more code
            loader = TransformerLoader()
            self.transformer = loader.load(
                server_args.model_paths["transformer"], server_args, "transformer"
            )

        self._maybe_enable_cache_dit_and_torch_compile(
            cache_dit_num_inference_steps, batch
        )

        if not transformer_was_loaded:
            if pipeline:
                pipeline.add_module("transformer", self.transformer)
            server_args.model_loaded["transformer"] = True

        if batch.rollout:
            self._maybe_prepare_rollout(batch)

        # Prepare extra step kwargs for scheduler
        extra_step_kwargs = self.prepare_extra_func_kwargs(
            scheduler.step,
            {"generator": batch.generator, "eta": batch.eta, "batch": batch},
        )

        # Setup precision and autocast settings
        target_dtype = resolve_precision(
            server_args, "dit", precision_attr="dit_precision"
        )
        autocast_enabled = precision_autocast_enabled(
            target_dtype, server_args.disable_autocast
        )

        # Prepare image latents and embeddings for I2V generation
        image_embeds = batch.image_embeds
        if len(image_embeds) > 0:
            image_embeds = [
                image_embed.to(target_dtype) for image_embed in image_embeds
            ]

        # Prepare STA parameters
        if self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN:
            self.prepare_sta_param(batch, server_args)

        # Get latents and embeddings
        latents = batch.latents
        # Removed Tensor truthiness assert to avoid GPU sync
        if batch.do_classifier_free_guidance:
            assert batch.negative_prompt_embeds is not None
            # Removed Tensor truthiness assert to avoid GPU sync

        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)

        # TI2V specific preparations - before SP sharding
        if should_preprocess_for_wan_ti2v:
            vae_dtype = resolve_precision(
                server_args, "vae", precision_attr="vae_precision"
            )
            with self.use_declared_component(
                component_name="vae",
                module=self.vae,
                target_dtype=vae_dtype,
            ) as vae:
                assert vae is not None
                self.vae = vae
                seq_len, z, reserved_frames_masks = prepare_wan_ti2v_latents(
                    self.vae,
                    latents,
                    target_dtype,
                    vae_dtype,
                    batch,
                    server_args,
                )
        else:
            seq_len, z, reserved_frames_masks = (
                None,
                None,
                None,
            )

        # Handle sequence parallelism after TI2V processing
        self._preprocess_sp_latents(batch, server_args)
        latents = batch.latents

        # Shard z and reserved_frames_mask for TI2V if SP is enabled
        if should_preprocess_for_wan_ti2v:
            reserved_frames_mask_sp, z_sp = prepare_wan_ti2v_sp_inputs(
                z, reserved_frames_masks, batch
            )
        else:
            reserved_frames_mask_sp, z_sp = (
                (
                    reserved_frames_masks[0]
                    if reserved_frames_masks is not None
                    else None
                ),
                z,
            )

        guidance = self.get_or_build_guidance(
            # TODO: replace with raw_latent_shape?
            latents.shape[0],
            latents.dtype,
            latents.device,
        )

        image_kwargs = self.prepare_extra_func_kwargs(
            getattr(self.transformer, "forward", self.transformer),
            {
                # TODO: make sure on-device
                "encoder_hidden_states_image": image_embeds,
                "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24),
            },
        )

        pos_cond_kwargs = self.prepare_extra_func_kwargs(
            getattr(self.transformer, "forward", self.transformer),
            {
                "encoder_hidden_states_2": batch.clip_embedding_pos,
                "encoder_attention_mask": batch.prompt_attention_mask,
                "encoder_hidden_states_mask": (
                    batch.prompt_embeds_mask
                    if batch.prompt_embeds_mask is not None
                    else batch.prompt_attention_mask
                ),
            }
            | server_args.pipeline_config.prepare_pos_cond_kwargs(
                batch,
                self.device,
                self._get_transformer_attr("rotary_emb"),
                dtype=target_dtype,
            )
            | dict(
                encoder_hidden_states=server_args.pipeline_config.get_pos_prompt_embeds(
                    batch
                )
            ),
        )

        if batch.do_classifier_free_guidance:
            neg_cond_kwargs = self.prepare_extra_func_kwargs(
                getattr(self.transformer, "forward", self.transformer),
                {
                    "encoder_hidden_states_2": batch.clip_embedding_neg,
                    "encoder_attention_mask": batch.negative_attention_mask,
                    "encoder_hidden_states_mask": (
                        batch.negative_prompt_embeds_mask
                        if batch.negative_prompt_embeds_mask is not None
                        else batch.negative_attention_mask
                    ),
                }
                | server_args.pipeline_config.prepare_neg_cond_kwargs(
                    batch,
                    self.device,
                    self._get_transformer_attr("rotary_emb"),
                    dtype=target_dtype,
                )
                | dict(
                    encoder_hidden_states=server_args.pipeline_config.get_neg_prompt_embeds(
                        batch
                    )
                ),
            )
        else:
            neg_cond_kwargs = {}

        cfg_policy = server_args.pipeline_config.cfg_policy.build(
            batch, image_kwargs, pos_cond_kwargs, neg_cond_kwargs
        )

        return DenoisingContext(
            scheduler=scheduler,
            extra_step_kwargs=extra_step_kwargs,
            target_dtype=target_dtype,
            autocast_enabled=autocast_enabled,
            timesteps=timesteps,
            num_inference_steps=num_inference_steps,
            num_warmup_steps=num_warmup_steps,
            image_kwargs=image_kwargs,
            pos_cond_kwargs=pos_cond_kwargs,
            neg_cond_kwargs=neg_cond_kwargs,
            latents=latents,
            boundary_timestep=boundary_timestep,
            z=z_sp,
            reserved_frames_mask=reserved_frames_mask_sp,
            seq_len=seq_len,
            guidance=guidance,
            is_warmup=batch.is_warmup,
            cfg_policy=cfg_policy,
        )

    def _before_denoising_loop(
        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
    ) -> None:
        """Prepare scheduler state before entering the shared denoising loop."""
        self._reset_scheduler_loop_state(ctx.scheduler)
        ctx.scheduler.set_begin_index(0)
        self._init_cfg_gate_state(ctx, batch, server_args)

    def _reset_scheduler_loop_state(self, scheduler) -> None:
        if hasattr(scheduler, "_step_index"):
            scheduler._step_index = None
        if hasattr(scheduler, "_begin_index"):
            scheduler._begin_index = None
        if hasattr(scheduler, "lower_order_nums"):
            scheduler.lower_order_nums = 0
        if hasattr(scheduler, "last_sample"):
            scheduler.last_sample = None
        if hasattr(scheduler, "this_order"):
            scheduler.this_order = 0

        solver_order = getattr(getattr(scheduler, "config", None), "solver_order", 0)
        if solver_order:
            if hasattr(scheduler, "model_outputs"):
                scheduler.model_outputs = [None] * solver_order
            if hasattr(scheduler, "timestep_list"):
                scheduler.timestep_list = [None] * solver_order

    def _init_cfg_gate_state(
        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
    ) -> None:
        """Initialize optional CFG residual reuse for the current denoising loop."""
        fraction = envs.SGLANG_DIFFUSION_CFG_GATE_STEP
        if not 0.0 <= fraction <= 1.0:
            raise ValueError(
                "SGLANG_DIFFUSION_CFG_GATE_STEP must be between 0.0 and 1.0, "
                f"got {fraction}."
            )

        num_steps = len(ctx.timesteps)
        requested = fraction < 1.0 and batch.do_classifier_free_guidance
        active = requested and not server_args.enable_cfg_parallel
        gate_step = int(num_steps * fraction) if active else num_steps + 1
        ctx.extra["cfg_gate_state"] = {
            "fraction": fraction,
            "requested": requested,
            "active": active,
            "gate_step": gate_step,
            "delta": None,
            "model_id": None,
            "fresh_uncond": 0,
            "reused": 0,
            "invalidations": 0,
        }

        if not (active or requested):
            return

        if ctx.is_warmup or (
            world_group_is_initialized() and get_world_group().local_rank != 0
        ):
            return

        if active:
            logger.info(
                "CFG gating enabled: reuse unconditioned residual after step %d/%d "
                "(fraction=%.3f).",
                gate_step,
                num_steps,
                fraction,
            )
            if batch.guidance_rescale > 0:
                logger.warning(
                    "CFG gating is enabled with guidance_rescale=%s; benchmark image "
                    "quality before using this setting in production.",
                    batch.guidance_rescale,
                )
        elif requested:
            logger.info(
                "CFG gating requested but skipped because CFG parallel is enabled."
            )

    def _get_transformer_attr(self, name: str) -> Any:
        seen: set[int] = set()
        stack = [self.transformer]
        while stack:
            module = stack.pop()
            if module is None or id(module) in seen:
                continue
            seen.add(id(module))

            value = getattr(module, name, None)
            if value is not None:
                return value

            for wrapper_attr in ("_fsdp_wrapped_module", "module", "_orig_mod"):
                wrapped = getattr(module, wrapper_attr, None)
                if wrapped is not None:
                    stack.append(wrapped)
        return None

    def _prepare_step_state(
        self,
        ctx: DenoisingContext,
        batch: Req,
        server_args: ServerArgs,
        step_index: int,
        t_host: torch.Tensor,
        timesteps_cpu: torch.Tensor,
    ) -> DenoisingStepState:
        """Build the per-step state shared by the loop and model-specific hooks."""
        t_int = int(t_host.item())
        t_device = ctx.timesteps[step_index]
        current_model, current_guidance_scale = self._select_and_manage_model(
            t_int=t_int,
            boundary_timestep=ctx.boundary_timestep,
            server_args=server_args,
            batch=batch,
        )
        attn_metadata = self._prepare_step_attn_metadata(
            ctx=ctx,
            batch=batch,
            server_args=server_args,
            step_index=step_index,
            t_int=t_int,
            timesteps_cpu=timesteps_cpu,
        )
        return DenoisingStepState(
            step_index=step_index,
            t_host=t_host,
            t_device=t_device,
            t_int=t_int,
            current_model=current_model,
            current_guidance_scale=current_guidance_scale,
            attn_metadata=attn_metadata,
        )

    def _prepare_step_attn_metadata(
        self,
        ctx: DenoisingContext,
        batch: Req,
        server_args: ServerArgs,
        step_index: int,
        t_int: int,
        timesteps_cpu: torch.Tensor,
    ) -> Any | None:
        """Build attention metadata for the current denoising step."""
        # Keep attention metadata preparation overridable so model-specific stages
        # can preserve their original semantics without duplicating step state setup.
        return self._build_attn_metadata(
            step_index,
            batch,
            server_args,
            timestep_value=t_int,
            timesteps=timesteps_cpu,
        )

    def _get_prompt_embeds_validator(
        self, batch: Req
    ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
        """Return the prompt-embedding validator used by verify_input."""
        return V.list_not_empty

    def _get_negative_prompt_embeds_validator(
        self, batch: Req
    ) -> Callable[[Any], bool] | list[Callable[[Any], bool]]:
        """Return the negative-prompt validator used by verify_input."""
        return lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)

    def _run_denoising_step(
        self,
        ctx: DenoisingContext,
        step: DenoisingStepState,
        batch: Req,
        server_args: ServerArgs,
    ) -> None:
        """Run one scheduler-backed denoising step in the shared base path.

        Model-specific stages should override this instead of the whole loop
        whenever possible to achieve better performance. Overrides that bypass
        ``_predict_noise_with_cfg`` / ``ctx.scheduler.step`` will lose the
        inner ``predict_noise`` / ``scheduler_step`` NVTX markers emitted
        below; mirror them in the override if those markers are needed.
        """
        use_nvtx = self.current_use_nvtx
        # 1. Prepare latent inputs in the model's compute dtype.
        latent_model_input = ctx.latents.to(ctx.target_dtype)
        if batch.image_latent is not None:
            assert (
                not server_args.pipeline_config.task_type == ModelTaskType.TI2V
            ), "image latents should not be provided for TI2V task"
            latent_model_input = torch.cat(
                [latent_model_input, batch.image_latent], dim=1
            ).to(ctx.target_dtype)

        # 2. Expand the timestep to the shape expected by the current model.
        timestep = self.expand_timestep_before_forward(
            batch,
            server_args,
            step.t_device,
            ctx.target_dtype,
            ctx.seq_len,
            ctx.reserved_frames_mask,
        )

        # 3. Apply scheduler-side input scaling before the model forward.
        latent_model_input = ctx.scheduler.scale_model_input(
            latent_model_input, step.t_device
        )

        # 4. Run the model prediction path, including CFG when enabled.
        with maybe_nvtx_range("predict_noise", use_nvtx):
            noise_pred = self._predict_noise_with_cfg(
                current_model=step.current_model,
                latent_model_input=latent_model_input,
                timestep=timestep,
                batch=batch,
                timestep_index=step.step_index,
                attn_metadata=step.attn_metadata,
                target_dtype=ctx.target_dtype,
                current_guidance_scale=step.current_guidance_scale,
                cfg_policy=ctx.cfg_policy,
                cfg_gate_state=ctx.extra.get("cfg_gate_state"),
                server_args=server_args,
                guidance=ctx.guidance,
                latents=ctx.latents,
            )
        if server_args.comfyui_mode:
            batch.noise_pred = noise_pred

        # 5. Advance the scheduler state with the predicted noise.
        with maybe_nvtx_range("scheduler_step", use_nvtx):
            ctx.latents = ctx.scheduler.step(
                model_output=noise_pred,
                timestep=step.t_device,
                sample=ctx.latents,
                **ctx.extra_step_kwargs,
                return_dict=False,
            )[0]

        # 6. Re-apply any model-specific latent constraints after the update.
        ctx.latents = self.post_forward_for_ti2v_task(
            batch,
            server_args,
            ctx.reserved_frames_mask,
            ctx.latents,
            ctx.z,
        )

    def _record_trajectory(
        self,
        ctx: DenoisingContext,
        step: DenoisingStepState,
        batch: Req,
        server_args: ServerArgs,
    ) -> None:
        """Append the current step to the returned latent trajectory, if requested."""
        if not batch.return_trajectory_latents:
            return
        ctx.trajectory_timesteps.append(step.t_host)
        ctx.trajectory_latents.append(ctx.latents)

    def _finalize_denoising_loop(
        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
    ) -> None:
        """Finalize the shared loop by handing state to post-denoising processing."""
        self._log_cfg_gate_summary(ctx, batch)
        self._post_denoising_loop(
            batch=batch,
            latents=ctx.latents,
            trajectory_latents=ctx.trajectory_latents,
            trajectory_timesteps=ctx.trajectory_timesteps,
            server_args=server_args,
            is_warmup=ctx.is_warmup,
        )

    def _log_cfg_gate_summary(self, ctx: DenoisingContext, batch: Req) -> None:
        state = ctx.extra.get("cfg_gate_state")
        if (
            not state
            or not state["requested"]
            or ctx.is_warmup
            or (world_group_is_initialized() and get_world_group().local_rank != 0)
        ):
            return

        logger.info(
            "CFG gating summary: fraction=%.3f, gate_step=%d/%d, "
            "fresh_uncond=%d, reused=%d, invalidations=%d, guidance_rescale=%s.",
            state["fraction"],
            state["gate_step"],
            len(ctx.timesteps),
            state["fresh_uncond"],
            state["reused"],
            state["invalidations"],
            batch.guidance_rescale,
        )

    def _post_denoising_loop(
        self,
        batch: Req,
        latents: torch.Tensor,
        trajectory_latents: list,
        trajectory_timesteps: list,
        server_args: ServerArgs,
        is_warmup: bool = False,
        *args,
        **kwargs,
    ):
        # Gather results if using sequence parallelism
        if trajectory_latents:
            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0)
        else:
            trajectory_tensor = None
            trajectory_timesteps_tensor = None

        # Gather results if using sequence parallelism
        latents, trajectory_tensor = self._postprocess_sp_latents(
            batch, latents, trajectory_tensor
        )

        # Gather noise_pred if using sequence parallelism
        # noise_pred has the same shape as latents (sharded along sequence dimension)
        if (
            self._sp_world_size() > 1
            and getattr(batch, "did_sp_shard_latents", False)
            and server_args.comfyui_mode
            and hasattr(batch, "noise_pred")
            and batch.noise_pred is not None
        ):
            batch.noise_pred = server_args.pipeline_config.gather_noise_pred_for_sp(
                batch, batch.noise_pred
            )

        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
            batch.trajectory_latents = trajectory_tensor.cpu()

        # Update batch with final latents
        batch.latents = self.server_args.pipeline_config.post_denoising_loop(
            latents, batch
        )

        # Save STA mask search results if needed
        if (
            not is_warmup
            and self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN
            and server_args.attention_backend_config.STA_mode == "STA_SEARCHING"
        ):
            self.save_sta_search_results(batch)

        # deallocate transformer if on mps
        pipeline = self.pipeline() if self.pipeline else None
        if torch.backends.mps.is_available() and not is_warmup:
            logger.info(
                "Memory before deallocating transformer: %s",
                torch.mps.current_allocated_memory(),
            )
            if self._component_residency_manager is not None:
                self._component_residency_manager.remove_nvtx_hooks_for_module(
                    self.transformer
                )
            del self.transformer
            if pipeline is not None and "transformer" in pipeline.modules:
                del pipeline.modules["transformer"]
            server_args.model_loaded["transformer"] = False
            logger.info(
                "Memory after deallocating transformer: %s",
                torch.mps.current_allocated_memory(),
            )

    def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
        """Shard latents for Sequence Parallelism if applicable."""
        if self._sp_world_size() <= 1:
            return

        if batch.latents is not None:
            (
                batch.latents,
                did_shard,
            ) = server_args.pipeline_config.shard_latents_for_sp(batch, batch.latents)
            batch.did_sp_shard_latents = did_shard
        else:
            batch.did_sp_shard_latents = False

        # image_latent must be sharded consistently with latents when it is
        # concatenated along the sequence dimension in the denoising loop.
        if batch.image_latent is not None:
            sp_video_metadata = {
                name: getattr(batch, name)
                for name in (
                    "sp_video_latent_num_frames",
                    "sp_video_start_frame",
                    "sp_video_tokens_per_frame",
                    "sp_video_valid_token_count",
                )
                if hasattr(batch, name)
            }
            batch.image_latent, _ = server_args.pipeline_config.shard_latents_for_sp(
                batch, batch.image_latent
            )
            for name, value in sp_video_metadata.items():
                setattr(batch, name, value)

    def _postprocess_sp_latents(
        self,
        batch: Req,
        latents: torch.Tensor,
        trajectory_tensor: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Gather latents after Sequence Parallelism if they were sharded."""
        if self._sp_world_size() > 1 and getattr(batch, "did_sp_shard_latents", False):
            latents = self.server_args.pipeline_config.gather_latents_for_sp(
                latents, batch=batch
            )
            if trajectory_tensor is not None:
                # trajectory_tensor shapes:
                # - video: [b, num_steps, c, t_local, h, w] -> gather on dim=3
                # - image: [b, num_steps, s_local, d] -> gather on dim=2
                trajectory_tensor = trajectory_tensor.to(get_local_torch_device())
                if isinstance(self.server_args.pipeline_config, ZImagePipelineConfig):
                    trajectory_tensor = (
                        self.server_args.pipeline_config.gather_latents_for_sp(
                            trajectory_tensor, batch=batch
                        )
                    )
                else:
                    gather_dim = 3 if trajectory_tensor.dim() >= 5 else 2
                    trajectory_tensor = sequence_model_parallel_all_gather(
                        trajectory_tensor, dim=gather_dim
                    )
                    if gather_dim == 2 and hasattr(batch, "raw_latent_shape"):
                        orig_s = batch.raw_latent_shape[1]
                        if trajectory_tensor.shape[2] > orig_s:
                            trajectory_tensor = trajectory_tensor[:, :, :orig_s, :]
        return latents, trajectory_tensor

    def _sp_world_size(self) -> int:
        if not model_parallel_is_initialized():
            return 1
        return get_sp_world_size()

    def step_profile(self):
        profiler = SGLDiffusionProfiler.get_instance()
        if profiler:
            profiler.step_denoising_step()

    def _manage_dit_use_site(
        self,
        current_model: nn.Module,
        current_phase: str,
        batch: Req,
    ) -> None:
        """
        manage dit's residency by reporting the active sequential use

        only applicable for dual-dit architecture like Wan

        Args:
            current_model: the next active dit, transformer_1 or transformer_2
        """
        manager = self._component_residency_manager

        component_name = manager.component_name_for_module(current_model, current_phase)
        phase = str(batch.extra.get("ltx2_phase", current_phase))
        use = ComponentUse(
            stage_name=self._active_component_stage_name(),
            component_name=component_name,
            phase=phase,
            preferred_ready_after_request=component_name == "transformer",
            memory_intensive=True,
        )
        manager.begin_use(use, module=current_model)

    def _select_and_manage_model(
        self,
        t_int: int,
        boundary_timestep: float | None,
        server_args: ServerArgs,
        batch: Req,
    ):
        if boundary_timestep is None or t_int >= boundary_timestep:
            # High-noise stage
            current_model = self.transformer
            current_guidance_scale = batch.guidance_scale
            current_phase = "transformer"
        else:
            # Low-noise stage
            current_model = self.transformer_2
            current_guidance_scale = batch.guidance_scale_2
            current_phase = "transformer_2"

        self._manage_dit_use_site(current_model, current_phase, batch)

        assert current_model is not None, "The model for the current step is not set."
        return current_model, current_guidance_scale

    def expand_timestep_before_forward(
        self,
        batch: Req,
        server_args: ServerArgs,
        t_device,
        target_dtype,
        seq_len: int | None,
        reserved_frames_mask,
    ):
        bsz = batch.raw_latent_shape[0]
        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)

        # expand timestep
        if should_preprocess_for_wan_ti2v:
            assert seq_len is not None, "Wan TI2V requires a token sequence length."
            timestep = expand_wan_ti2v_timestep(
                batch,
                t_device,
                target_dtype,
                seq_len,
                reserved_frames_mask,
                server_args.pipeline_config.dit_config.arch_config.patch_size,
            )
        else:
            timestep = t_device.repeat(bsz)
        return timestep

    def post_forward_for_ti2v_task(
        self, batch: Req, server_args: ServerArgs, reserved_frames_mask, latents, z
    ):
        """Re-apply Wan TI2V first-frame conditioning after each denoising step."""
        should_preprocess_for_wan_ti2v = should_apply_wan_ti2v(batch, server_args)
        if should_preprocess_for_wan_ti2v:
            latents = blend_wan_ti2v_latents(latents, reserved_frames_mask, z)

        return latents

    @contextmanager
    def _offload_for_torch_compile_warmup(self, batch: Req):
        """wrap a denoise so the torch.compile warmup has spare VRAM.

        No-op unless the DiT was layerwise-offloaded for compile.
        Warmup: move resident non-DiT components off-device for the autotune window, restore
        after.
        """
        if not self._offloaded_dit_modules_for_compile:
            yield
            return
        # if the dits have been layerwise-offloaded in preparation stage
        if batch.is_warmup:
            # if warmup is enabled, for warmup request, offload components to ensure sufficient VRAM headroom for torch compile
            moved = self._move_resident_components_for_warmup()
            try:
                yield
            finally:
                device = get_local_torch_device()
                for module in moved:
                    module.to(device)
            return
        # prepare for first real request: restore the user-configured residency
        for module in self._offloaded_dit_modules_for_compile:
            module.disable_offload()
        # clear the list for avoid overhead during real request
        self._offloaded_dit_modules_for_compile.clear()
        yield

    def forward(
        self,
        batch: Req,
        server_args: ServerArgs,
    ) -> Req:
        with self._offload_for_torch_compile_warmup(batch):
            return self._denoise(batch, server_args)

    @torch.no_grad()
    def _denoise(
        self,
        batch: Req,
        server_args: ServerArgs,
    ) -> Req:
        """
        Run the denoising loop.
        """
        ctx = self._prepare_denoising_loop(batch, server_args)
        if batch.rollout:
            self._maybe_init_denoising_env_collection(
                batch=batch,
                pipeline_config=server_args.pipeline_config,
                image_kwargs=ctx.image_kwargs,
                pos_cond_kwargs=ctx.pos_cond_kwargs,
                neg_cond_kwargs=ctx.neg_cond_kwargs,
                guidance=ctx.guidance,
            )
        denoising_start_time = time.time()
        self._before_denoising_loop(ctx, batch, server_args)
        # to avoid device-sync caused by timestep comparison
        timesteps_cpu = ctx.timesteps.cpu()
        num_timesteps = timesteps_cpu.shape[0]
        # Re-resolve the explicit-range gate so the per-step markers
        # below honor this request's is_warmup state. Layer hooks are
        # registered by the residency manager at the use-site.
        use_nvtx = self._apply_nvtx_gate(ctx.is_warmup)

        with (
            torch.autocast(
                device_type=current_platform.device_type,
                dtype=ctx.target_dtype,
                enabled=ctx.autocast_enabled,
            ),
            maybe_nvtx_range("denoising_loop", use_nvtx),
        ):
            with self.progress_bar(
                total=ctx.num_inference_steps, batch=batch
            ) as progress_bar:
                for step_index, t_host in enumerate(timesteps_cpu):
                    # Use ``:.4g`` so flow-matching schedulers (e.g. FLUX) that
                    # use non-integer timesteps keep their precision in markers.
                    step_marker = f"denoising_step_{step_index}_t{t_host.item():.4g}"
                    with (
                        maybe_nvtx_range(step_marker, use_nvtx),
                        StageProfiler(
                            f"denoising_step_{step_index}",
                            logger=logger,
                            metrics=batch.metrics,
                            perf_dump_path_provided=batch.perf_dump_path is not None,
                            record_as_step=True,
                        ),
                    ):
                        step = self._prepare_step_state(
                            ctx,
                            batch,
                            server_args,
                            step_index,
                            t_host,
                            timesteps_cpu,
                        )
                        # Capture the raw (pre-scale, pre-I2V-concat) noisy latent
                        # x_{t_i} for rollout trajectory collection. Must run
                        # BEFORE _run_denoising_step so ctx.latents is still the
                        # pre-step value. Gated on batch.rollout to keep the
                        # non-rollout path strictly untouched.
                        if batch.rollout:
                            batch._rollout_loop_step_index = step_index
                            self._maybe_append_dit_trajectory_step(
                                batch=batch,
                                latents=ctx.latents,
                                timestep_value=step.t_host,
                                step_index=step_index,
                            )
                        self._run_denoising_step(ctx, step, batch, server_args)
                        self._record_trajectory(ctx, step, batch, server_args)

                        if step_index == num_timesteps - 1 or (
                            (step_index + 1) > ctx.num_warmup_steps
                            and (step_index + 1) % ctx.scheduler.order == 0
                            and progress_bar is not None
                        ):
                            progress_bar.update()

                        if not ctx.is_warmup:
                            self.step_profile()

        denoising_end_time = time.time()

        if num_timesteps > 0 and not ctx.is_warmup:
            self.log_info(
                "average time per step: %.4f seconds",
                (denoising_end_time - denoising_start_time) / len(ctx.timesteps),
            )

        self._finish_active_component_use()

        # Rollout postprocessing must run BEFORE _finalize_denoising_loop so
        # the final scheduler.step output (ctx.latents) is still SP-sharded and
        # can be gathered uniformly alongside the per-step dit_trajectory via
        # gather_stacked_latents_for_sp.
        if batch.rollout:
            self._postprocess_rollout_outputs(
                batch=batch,
                latents=ctx.latents,
                num_inference_steps=num_timesteps,
                final_timestep=timesteps_cpu.new_zeros(()),
                server_args=server_args,
            )
        self._finalize_denoising_loop(ctx, batch, server_args)
        return batch

    def _get_extra_func_kwarg_names(self, func) -> tuple[bool, frozenset[str]]:
        import functools

        # Handle cache-dit's partial wrapping logic.
        # Cache-dit wraps the forward method with functools.partial where args[0] is the instance.
        # We access `_original_forward` if available to inspect the underlying signature.
        # See: https://github.com/vipshop/cache-dit
        if isinstance(func, functools.partial) and func.args:
            func = getattr(func.args[0], "_original_forward", func)

        # Unwrap any decorators (e.g. functools.wraps)
        target_func = inspect.unwrap(func)
        cache_target = (
            target_func.__func__ if inspect.ismethod(target_func) else target_func
        )
        cache_key = id(cache_target)
        cached = self._extra_func_kwarg_names_cache.get(cache_key)
        if cached is not None:
            return cached

        params = inspect.signature(target_func).parameters
        result = (
            any(
                param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()
            ),
            frozenset(params),
        )
        self._extra_func_kwarg_names_cache[cache_key] = result
        return result

    # TODO: this will extends the preparation stage, should let subclass/passed-in variables decide which to prepare
    def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]:
        """
        Prepare extra kwargs for the scheduler step / denoise step.

        Args:
            func: The function to prepare kwargs for.
            kwargs: The kwargs to prepare.
        """
        accepts_var_kwargs, param_names = self._get_extra_func_kwarg_names(func)
        if accepts_var_kwargs:
            return kwargs
        return {k: v for k, v in kwargs.items() if k in param_names}

    def _predict_noise_with_cfg(
        self,
        current_model: nn.Module,
        latent_model_input: torch.Tensor,
        timestep,
        batch: Req,
        timestep_index: int,
        attn_metadata,
        target_dtype,
        current_guidance_scale,
        cfg_policy: CFGPolicy,
        cfg_gate_state: dict[str, Any] | None,
        server_args: ServerArgs,
        guidance: torch.Tensor,
        latents: torch.Tensor,
    ) -> "torch.Tensor | tuple[torch.Tensor, ...]":
        """Run all CFG branch forward passes and combine into the final noise estimate."""
        cfg_scale = server_args.pipeline_config.get_classifier_free_guidance_scale(
            batch, current_guidance_scale
        )

        def predict_fn(branch):
            branch.configure_batch(batch)
            with set_forward_context(
                current_timestep=timestep_index,
                attn_metadata=attn_metadata,
                forward_batch=batch,
            ):
                raw = self._predict_noise(
                    current_model=current_model,
                    latent_model_input=latent_model_input,
                    timestep=timestep,
                    target_dtype=target_dtype,
                    guidance=guidance,
                    **branch.kwargs,
                )
            pred_t = _wrap(raw)
            if len(pred_t) == 1:
                pred_t = (
                    server_args.pipeline_config.slice_noise_pred(pred_t[0], latents),
                )
            return _unwrap(pred_t)

        if (
            cfg_gate_state
            and cfg_gate_state["active"]
            and not server_args.enable_cfg_parallel
            and type(cfg_policy) is CFGPolicy
            and len(cfg_policy.branches) == 2
            and cfg_policy.branches[0].is_conditional
            and not cfg_policy.branches[1].is_conditional
        ):
            model_id = id(current_model)
            if cfg_gate_state["model_id"] not in (None, model_id):
                cfg_gate_state["delta"] = None
                cfg_gate_state["invalidations"] += 1

            pos_pred = predict_fn(cfg_policy.branches[0])
            pos_t = _wrap(pos_pred)
            delta_t = cfg_gate_state["delta"]
            can_reuse = (
                timestep_index >= cfg_gate_state["gate_step"]
                and delta_t is not None
                and len(pos_t) == len(delta_t)
            )

            if can_reuse:
                neg_pred = _unwrap(tuple(p - d for p, d in zip(pos_t, delta_t)))
                cfg_gate_state["reused"] += 1
            else:
                neg_pred = predict_fn(cfg_policy.branches[1])
                neg_t = _wrap(neg_pred)
                cfg_gate_state["delta"] = tuple(
                    p.detach() - n.detach() for p, n in zip(pos_t, neg_t)
                )
                cfg_gate_state["model_id"] = model_id
                cfg_gate_state["fresh_uncond"] += 1

            return cfg_policy.combine(
                [pos_pred, neg_pred],
                batch,
                cfg_scale,
                server_args.pipeline_config,
            )

        if server_args.enable_cfg_parallel:
            if (
                len(cfg_policy.branches) == 2
                and get_classifier_free_guidance_world_size() == 2
            ):
                return run_two_branch_cfg_parallel(
                    cfg_policy,
                    predict_fn,
                    cfg_scale,
                    batch,
                    server_args.pipeline_config,
                )
            # perform cfg branches in parallel, following the cfg policy
            predictions = run_cfg_parallel(cfg_policy, predict_fn)
        else:
            # perform cfg branches one-by-one locally
            predictions = [predict_fn(branch) for branch in cfg_policy.branches]

        return cfg_policy.combine(
            predictions,
            batch,
            cfg_scale,
            server_args.pipeline_config,
            cfg_parallel=server_args.enable_cfg_parallel,
        )

    def _build_attn_metadata(
        self,
        i: int,
        batch: Req,
        server_args: ServerArgs,
        *,
        timestep_value: int | None = None,
        timesteps: torch.Tensor | None = None,
    ) -> Any | None:
        """
        Build attention metadata for custom attention backends.

        Args:
            i: The current timestep index.
        """
        attn_metadata = None
        self.attn_metadata_builder = None
        try:
            self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls()
        except NotImplementedError:
            self.attn_metadata_builder_cls = None
        if self.attn_metadata_builder_cls:
            self.attn_metadata_builder = self.attn_metadata_builder_cls()
        if (
            self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN
            or self.attn_backend.get_enum() == AttentionBackendEnum.VIDEO_SPARSE_ATTN
        ):
            attention_backend_config = server_args.attention_backend_config or {}
            vsa_sparsity = attention_backend_config.get(
                "VSA_sparsity", attention_backend_config.get("sparsity", 0.0)
            )
            attn_metadata = self.attn_metadata_builder.build(
                current_timestep=i,
                raw_latent_shape=batch.raw_latent_shape[2:5],
                patch_size=server_args.pipeline_config.dit_config.patch_size,
                STA_param=batch.STA_param,
                VSA_sparsity=vsa_sparsity,
                device=get_local_torch_device(),
            )
        elif (
            self.attn_backend.get_enum() == AttentionBackendEnum.SPARSE_VIDEO_GEN_2_ATTN
        ):
            if timestep_value is None or timesteps is None:
                raise ValueError(
                    "timestep_value and timesteps must be provided for SVG2 attention metadata"
                )

            svg2_cfg = server_args.attention_backend_config or {}
            num_layers = server_args.pipeline_config.dit_config.num_layers
            if (
                server_args.pipeline_config.dit_config.prefix.lower() == "hunyuan"
                and hasattr(server_args.pipeline_config.dit_config, "num_single_layers")
            ):
                num_layers += server_args.pipeline_config.dit_config.num_single_layers
            first_layers_fp = svg2_cfg.get("svg2_first_layers_fp", 0.03)
            if first_layers_fp <= 1.0:
                first_layers_fp = math.floor(first_layers_fp * num_layers)
            first_layers_fp = max(0, min(int(first_layers_fp), num_layers))

            first_times_fp = svg2_cfg.get("svg2_first_times_fp", 0.2)
            if first_times_fp <= 1.0:
                num_fp_steps = math.floor(first_times_fp * len(timesteps))
                if num_fp_steps > 0:
                    first_times_fp = float(timesteps[num_fp_steps - 1].item() - 1)
                else:
                    first_times_fp = float(timesteps.max().item() + 1)

            current_timestep = int(timestep_value)

            cache = batch.extra.get("svg2_cache")
            if cache is None:
                from sglang.multimodal_gen.runtime.layers.attention.backends.sparse_video_gen_2_attn import (
                    Svg2Cache,
                )

                cache = Svg2Cache()
                batch.extra["svg2_cache"] = cache

            patch_size = server_args.pipeline_config.dit_config.patch_size
            if isinstance(patch_size, list):
                patch_size = tuple(patch_size)
            if isinstance(patch_size, int):
                patch_size_t = getattr(
                    server_args.pipeline_config.dit_config, "patch_size_t", None
                )
                if patch_size_t is not None:
                    patch_size = (patch_size_t, patch_size, patch_size)

            context_length = 0
            prompt_length = None
            if server_args.pipeline_config.dit_config.prefix.lower() == "hunyuan":
                prompt_embeds = server_args.pipeline_config.get_pos_prompt_embeds(batch)
                if isinstance(prompt_embeds, list):
                    text_embeds = prompt_embeds[0] if prompt_embeds else None
                else:
                    text_embeds = prompt_embeds
                if isinstance(text_embeds, torch.Tensor) and text_embeds.ndim >= 2:
                    context_length = int(text_embeds.shape[1])
                if context_length > 0 and batch.prompt_attention_mask:
                    mask = batch.prompt_attention_mask[0]
                    if isinstance(mask, torch.Tensor):
                        if mask.shape[-1] > context_length:
                            mask = mask[:, -context_length:]
                        prompt_length = int(mask[0].sum().item())
                if prompt_length is None:
                    prompt_length = context_length

            attn_metadata = self.attn_metadata_builder.build(
                current_timestep=current_timestep,
                raw_latent_shape=batch.raw_latent_shape,
                patch_size=patch_size,
                num_q_centroids=svg2_cfg.get("svg2_num_q_centroids", 300),
                num_k_centroids=svg2_cfg.get("svg2_num_k_centroids", 1000),
                top_p_kmeans=svg2_cfg.get("svg2_top_p_kmeans", 0.9),
                min_kc_ratio=svg2_cfg.get("svg2_min_kc_ratio", 0.1),
                kmeans_iter_init=svg2_cfg.get("svg2_kmeans_iter_init", 50),
                kmeans_iter_step=svg2_cfg.get("svg2_kmeans_iter_step", 2),
                zero_step_kmeans_init=svg2_cfg.get("svg2_zero_step_kmeans_init", False),
                first_layers_fp=first_layers_fp,
                first_times_fp=first_times_fp,
                context_length=context_length,
                prompt_length=prompt_length,
                cache=cache,
                calculate_density=False,  # only need density when doing head load balancing
            )
        elif self.attn_backend.get_enum() == AttentionBackendEnum.VMOBA_ATTN:
            moba_params = server_args.attention_backend_config.moba_config.copy()
            moba_params.update(
                {
                    "current_timestep": i,
                    "raw_latent_shape": batch.raw_latent_shape[2:5],
                    "patch_size": server_args.pipeline_config.dit_config.patch_size,
                    "device": get_local_torch_device(),
                }
            )
        elif self.attn_backend.get_enum() == AttentionBackendEnum.FA:
            attn_metadata = self.attn_metadata_builder.build(
                raw_latent_shape=batch.raw_latent_shape
            )
        elif self.attn_backend.get_enum() in [
            AttentionBackendEnum.BLOCK_SPARSE_ATTN,
            AttentionBackendEnum.RAIN_FUSION_ATTN,
        ]:
            sparse_config = server_args.attention_backend_config

            current_timestep = i
            skip_first_steps = sparse_config.get("skip_first_steps", 10)
            sparsity = sparse_config.get("sparsity", 0.2)

            raw_latent_shape = batch.raw_latent_shape
            patch_size = server_args.pipeline_config.dit_config.patch_size

            if isinstance(patch_size, int):
                patch_size_t = getattr(
                    server_args.pipeline_config.dit_config, "patch_size_t", None
                )
                if patch_size_t is not None:
                    patch_size = (patch_size_t, patch_size, patch_size)
                else:
                    patch_size = (patch_size, patch_size, patch_size)

            attn_metadata = self.attn_metadata_builder.build(
                current_timestep=current_timestep,
                skip_first_steps=skip_first_steps,
                sparsity=sparsity,
                raw_latent_shape=raw_latent_shape,
                patch_size=patch_size,
            )
        else:
            # attn_metadata can be None for SDPA attention backend
            return None

        return attn_metadata

    def _predict_noise(
        self,
        current_model,
        latent_model_input,
        timestep,
        target_dtype,
        guidance: torch.Tensor,
        **kwargs,
    ):
        guidance_kwargs = self.prepare_extra_func_kwargs(
            getattr(current_model, "forward", current_model),
            {"guidance": guidance},
        )
        call_kwargs = dict(
            hidden_states=latent_model_input,
            timestep=timestep,
            **guidance_kwargs,
            **kwargs,
        )
        runner = self._maybe_get_bcg_runner(current_model)
        if runner is not None:
            model_output = self._bcg_run(runner, call_kwargs, current_model)
        else:
            model_output = current_model(**call_kwargs)
        return _ensure_tensor_model_output(model_output)

    @staticmethod
    def _bcg_is_warmup() -> bool:
        """True when the current forward is a warmup request."""
        from sglang.multimodal_gen.runtime.managers.forward_context import (
            get_forward_context,
        )

        try:
            forward_batch = get_forward_context().forward_batch
        except Exception:
            return False
        return bool(getattr(forward_batch, "is_warmup", False))

    def _bcg_run(self, runner, call_kwargs: dict, current_model):
        """Run the DiT through the BCG runner.

        During warmup we proactively capture one graph per text bucket (in
        addition to the request's own bucket) so that serving never records a
        fresh graph for a different prompt length — every bucket is already
        captured. Serving just replays (or runs eager for an uncaptured
        signature, never capturing).
        """
        if self._bcg_is_warmup():
            for bucket in self._bcg_text_buckets():
                runner.capture(
                    **self._bcg_pad_prompt_kwargs(
                        call_kwargs, current_model=current_model, force_bucket=bucket
                    )
                )
        return runner(
            **self._bcg_pad_prompt_kwargs(call_kwargs, current_model=current_model)
        )

    @staticmethod
    def _bcg_text_buckets() -> tuple[int, ...]:
        """Prompt sequence-length buckets, from --bcg-text-buckets."""
        from sglang.multimodal_gen.runtime.server_args import (
            DEFAULT_BCG_TEXT_BUCKETS,
            get_global_server_args,
        )

        try:
            resolver = get_global_server_args().resolved_bcg_text_buckets
            return resolver()
        except Exception:
            return DEFAULT_BCG_TEXT_BUCKETS

    def _bcg_pad_prompt_kwargs(
        self, call_kwargs: dict, current_model=None, force_bucket: int | None = None
    ):
        """Bucket prompt-conditioning inputs so BCG signatures ignore prompt length.

        Generic padding lives in ``breakable_cuda_graph.prompt_padding``;
        model-specific padders register from ``breakable_cuda_graph.model_padders``.

        ``force_bucket`` pads to exactly that bucket (used by warmup to capture
        every bucket); a prompt already longer than ``force_bucket`` is left
        unchanged, exactly as the normal bucket selection would do.
        """
        buckets = (
            (force_bucket,) if force_bucket is not None else self._bcg_text_buckets()
        )
        padder = bcg_utils.select_prompt_padder(current_model, call_kwargs)
        if padder is not None:
            return padder(call_kwargs, current_model, buckets)
        return bcg_utils.pad_masked_prompt_kwargs(call_kwargs, buckets)

    def _maybe_get_bcg_runner(self, current_model):
        """Return (lazily creating) the breakable CUDA graph runner for
        ``current_model``, or ``None`` if BCG is disabled / inapplicable.
        """
        if not self.server_args.enable_breakable_cuda_graph:
            return None
        if not isinstance(current_model, nn.Module):
            return None
        key = id(current_model)
        runner = self._bcg_runners.get(key)
        if runner is None:
            from sglang.multimodal_gen.runtime.breakable_cuda_graph.runner import (
                DiffusionBreakableCudaGraphRunner,
            )

            # DenoisingStage can switch between transformer and transformer_2;
            # each module owns separate graph state and static input buffers.
            runner = DiffusionBreakableCudaGraphRunner(
                current_model, get_local_torch_device()
            )
            self._bcg_runners[key] = runner
        return runner

    def prepare_sta_param(self, batch: Req, server_args: ServerArgs):
        """
        Prepare Sliding Tile Attention (STA) parameters and settings.
        """
        # TODO(kevin): STA mask search, currently only support Wan2.1 with 69x768x1280
        try:
            STA_mode = STA_Mode[server_args.attention_backend_config.STA_mode]
        except Exception as e:
            logger.error(f"Passed STA_mode: {STA_mode} doesn't exist")
            raise e
        skip_time_steps = server_args.attention_backend_config.skip_time_steps
        if batch.timesteps is None:
            raise ValueError("Timesteps must be provided")
        timesteps_num = batch.timesteps.shape[0]

        logger.info("STA_mode: %s", STA_mode)
        if (batch.num_frames, batch.height, batch.width) != (
            69,
            768,
            1280,
        ) and STA_mode != "STA_inference":
            raise NotImplementedError(
                "STA mask search/tuning is not supported for this resolution"
            )

        if (
            STA_mode == STA_Mode.STA_SEARCHING
            or STA_mode == STA_Mode.STA_TUNING
            or STA_mode == STA_Mode.STA_TUNING_CFG
        ):
            size = (batch.width, batch.height)
            if size == (1280, 768):
                # TODO: make it configurable
                sparse_mask_candidates_searching = [
                    "3, 1, 10",
                    "1, 5, 7",
                    "3, 3, 3",
                    "1, 6, 5",
                    "1, 3, 10",
                    "3, 6, 1",
                ]
                sparse_mask_candidates_tuning = [
                    "3, 1, 10",
                    "1, 5, 7",
                    "3, 3, 3",
                    "1, 6, 5",
                    "1, 3, 10",
                    "3, 6, 1",
                ]
                full_mask = ["3,6,10"]
            else:
                raise NotImplementedError(
                    "STA mask search is not supported for this resolution"
                )
        layer_num = self.transformer.config.num_layers
        # specific for HunyuanVideo
        if hasattr(self.transformer.config, "num_single_layers"):
            layer_num += self.transformer.config.num_single_layers
        head_num = self.transformer.config.num_attention_heads

        if STA_mode == STA_Mode.STA_SEARCHING:
            STA_param = configure_sta(
                mode=STA_Mode.STA_SEARCHING,
                layer_num=layer_num,
                head_num=head_num,
                time_step_num=timesteps_num,
                mask_candidates=sparse_mask_candidates_searching + full_mask,
                # last is full mask; Can add more sparse masks while keep last one as full mask
            )
        elif STA_mode == STA_Mode.STA_TUNING:
            STA_param = configure_sta(
                mode=STA_Mode.STA_TUNING,
                layer_num=layer_num,
                head_num=head_num,
                time_step_num=timesteps_num,
                mask_search_files_path=f"output/mask_search_result_pos_{size[0]}x{size[1]}/",
                mask_candidates=sparse_mask_candidates_tuning,
                full_attention_mask=[int(x) for x in full_mask[0].split(",")],
                skip_time_steps=skip_time_steps,  # Use full attention for first 12 steps
                save_dir=f"output/mask_search_strategy_{size[0]}x{size[1]}/",  # Custom save directory
                timesteps=timesteps_num,
            )
        elif STA_mode == STA_Mode.STA_TUNING_CFG:
            STA_param = configure_sta(
                mode=STA_Mode.STA_TUNING_CFG,
                layer_num=layer_num,
                head_num=head_num,
                time_step_num=timesteps_num,
                mask_search_files_path_pos=f"output/mask_search_result_pos_{size[0]}x{size[1]}/",
                mask_search_files_path_neg=f"output/mask_search_result_neg_{size[0]}x{size[1]}/",
                mask_candidates=sparse_mask_candidates_tuning,
                full_attention_mask=[int(x) for x in full_mask[0].split(",")],
                skip_time_steps=skip_time_steps,
                save_dir=f"output/mask_search_strategy_{size[0]}x{size[1]}/",
                timesteps=timesteps_num,
            )
        elif STA_mode == STA_Mode.STA_INFERENCE:
            import sglang.multimodal_gen.envs as envs

            config_file = envs.SGLANG_DIFFUSION_ATTENTION_CONFIG
            if config_file is None:
                raise ValueError("SGLANG_DIFFUSION_ATTENTION_CONFIG is not set")
            STA_param = configure_sta(
                mode=STA_Mode.STA_INFERENCE,
                layer_num=layer_num,
                head_num=head_num,
                time_step_num=timesteps_num,
                load_path=config_file,
            )

        batch.STA_param = STA_param
        batch.mask_search_final_result_pos = [[] for _ in range(timesteps_num)]
        batch.mask_search_final_result_neg = [[] for _ in range(timesteps_num)]

    def save_sta_search_results(self, batch: Req):
        """
        Save the STA mask search results.

        Args:
            batch: The current batch information.
        """
        size = (batch.width, batch.height)
        if size == (1280, 768):
            # TODO: make it configurable
            sparse_mask_candidates_searching = [
                "3, 1, 10",
                "1, 5, 7",
                "3, 3, 3",
                "1, 6, 5",
                "1, 3, 10",
                "3, 6, 1",
            ]
        else:
            raise NotImplementedError(
                "STA mask search is not supported for this resolution"
            )

        if batch.mask_search_final_result_pos is not None and batch.prompt is not None:
            save_mask_search_results(
                [dict(layer_data) for layer_data in batch.mask_search_final_result_pos],
                prompt=str(batch.prompt),
                mask_strategies=sparse_mask_candidates_searching,
                output_dir=f"output/mask_search_result_pos_{size[0]}x{size[1]}/",
            )
        if batch.mask_search_final_result_neg is not None and batch.prompt is not None:
            save_mask_search_results(
                [dict(layer_data) for layer_data in batch.mask_search_final_result_neg],
                prompt=str(batch.prompt),
                mask_strategies=sparse_mask_candidates_searching,
                output_dir=f"output/mask_search_result_neg_{size[0]}x{size[1]}/",
            )

    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
        """Verify denoising stage inputs."""
        result = VerificationResult()
        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
        # disable temporarily for image-generation models
        # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        result.add_check(
            "prompt_embeds",
            batch.prompt_embeds,
            self._get_prompt_embeds_validator(batch),
        )
        result.add_check("image_embeds", batch.image_embeds, V.is_list)
        # result.add_check(
        #     "image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)
        # )
        result.add_check(
            "num_inference_steps", batch.num_inference_steps, V.positive_int
        )
        result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float)
        result.add_check("eta", batch.eta, V.non_negative_float)
        result.add_check("generator", batch.generator, V.generator_or_list_generators)
        result.add_check(
            "do_classifier_free_guidance",
            batch.do_classifier_free_guidance,
            V.bool_value,
        )
        result.add_check(
            "negative_prompt_embeds",
            batch.negative_prompt_embeds,
            self._get_negative_prompt_embeds_validator(batch),
        )
        return result

    def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
        """Verify denoising stage outputs."""
        result = VerificationResult()
        # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        return result
