[Perf][Model] Add optimized LTX-2.5 DiffVAE operators (#7308)

Signed-off-by: mglyn <1203789601@qq.com>

vllm_omni/diffusion/models/ltx2/vae/decoder.py

@@ -0,0 +1,430 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Native conversion rules and decoder behavior copied and modified from
+# Diffusers commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0).
+
+"""LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder."""
+
+import math
+from collections.abc import Mapping
+from functools import cache
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders import (
+    LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel,
+)
+from diffusers.models.autoencoders import ltx2_diffusion_decoder as diffusers_decoder
+from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
+    LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d,
+)
+from diffusers.models.autoencoders.vae import DecoderOutput
+from safetensors import safe_open
+
+from ..ops import is_ltx2_fna_eligible, is_ltx2_fusion_eligible, resolve_ltx2_vae_operators
+
+LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5"
+LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
+LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors"
+
+_NATIVE_DECODER_PREFIXES = ("vae.decoder.", "decoder.")
+_NATIVE_STATISTICS_KEYS = {
+    "per_channel_statistics.mean-of-means": "latents_mean",
+    "per_channel_statistics.std-of-means": "latents_std",
+}
+_NATIVE_KEY_REPLACEMENTS = (
+    ("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1."),
+    ("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2."),
+    (".attn.proj.", ".attn.to_out.0."),
+    (".attn.q_norm.", ".attn.norm_q."),
+    (".attn.k_norm.", ".attn.norm_k."),
+)
+_GATE_FOLD_TARGETS = {
+    ".attn.to_out.0.weight": ".gate_msa",
+    ".attn.to_out.0.bias": ".gate_msa",
+    ".mlp.w_down.weight": ".gate_mlp",
+    ".context_proj.weight": ".gate_ctx",
+    ".context_proj.bias": ".gate_ctx",
+}
+_GATE_SUFFIXES = tuple(_GATE_FOLD_TARGETS.values())
+
+
+def _strip_native_decoder_prefix(key: str) -> str | None:
+    for prefix in _NATIVE_DECODER_PREFIXES:
+        if key.startswith(prefix):
+            return key.removeprefix(prefix)
+    return None
+
+
+def _native_statistics_target(key: str) -> str | None:
+    return _NATIVE_STATISTICS_KEYS.get(key.removeprefix("vae."))
+
+
+def _rename_native_decoder_key(key: str) -> str:
+    for source, target in _NATIVE_KEY_REPLACEMENTS:
+        key = key.replace(source, target)
+    return key
+
+
+def _fold_native_gate(key: str, value: torch.Tensor, gates: Mapping[str, torch.Tensor]) -> torch.Tensor:
+    target_suffix = next((suffix for suffix in _GATE_FOLD_TARGETS if key.endswith(suffix)), None)
+    if target_suffix is None:
+        return value
+    gate_key = key[: -len(target_suffix)] + _GATE_FOLD_TARGETS[target_suffix]
+    gate = gates.get(gate_key)
+    if gate is None:
+        return value
+    gate = gate.to(device=value.device, dtype=torch.float32)
+    value_float = value.to(dtype=torch.float32)
+    folded = gate.unsqueeze(1) * value_float if value.ndim == 2 else gate * value_float
+    return folded.to(dtype=value.dtype)
+
+
+def convert_ltx25_native_diffusion_decoder_state_dict(
+    native_state_dict: Mapping[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    """Convert the canonical Native LTX-2.5 DiffVAE state dict to Diffusers."""
+    decoder_state_dict: dict[str, torch.Tensor] = {}
+    converted: dict[str, torch.Tensor] = {}
+    for key, value in native_state_dict.items():
+        statistics_target = _native_statistics_target(key)
+        if statistics_target is not None:
+            converted[statistics_target] = value
+            continue
+        decoder_key = _strip_native_decoder_prefix(key)
+        if decoder_key is not None:
+            decoder_state_dict[decoder_key] = value
+
+    gates = {key: value for key, value in decoder_state_dict.items() if key.endswith(_GATE_SUFFIXES)}
+    for key, value in decoder_state_dict.items():
+        # ``type_emb`` is an unused shipping-checkpoint residue. Preview heads
+        # and legacy static gates are not parameters of the runtime decoder.
+        if key == "type_emb" or key.startswith("coarse_") or ".coarse_" in key or key in gates:
+            continue
+
+        converted_key = _rename_native_decoder_key(key)
+        value = _fold_native_gate(converted_key, value, gates)
+        if converted_key.endswith((".qkv.weight", ".qkv.bias")):
+            leaf = "weight" if converted_key.endswith(".weight") else "bias"
+            prefix = converted_key[: -len(f"qkv.{leaf}")]
+            if value.shape[0] % 3 != 0:
+                raise ValueError(
+                    f"Fused LTX-2.5 DiffVAE parameter {key!r} has leading dimension "
+                    f"{value.shape[0]}, which is not divisible by 3."
+                )
+            chunk = value.shape[0] // 3
+            converted[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone()
+            converted[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone()
+            converted[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone()
+            continue
+
+        converted[f"decoder.{converted_key}"] = value
+
+    return converted
+
+
+def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch.Tensor]:
+    """Read and convert only DiffVAE tensors from the canonical Native VAE file."""
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        native_state_dict = {
+            key: handle.get_tensor(key)
+            for key in handle.keys()
+            if _strip_native_decoder_prefix(key) is not None or _native_statistics_target(key) is not None
+        }
+    return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict)
+
+
+@cache
+def _load_natten_token_permutation():
+    from kernels import get_kernel
+
+    natten = get_kernel("shi-labs/natten", version=1)
+    return natten.token_permute.token_permute_operation, natten.token_permute.token_unpermute_operation
+
+
+class LTX2VideoVaeNeighborhoodNattenProcessor(diffusers_decoder.LTX2VideoVaeNeighborhoodNattenProcessor):
+    """NATTEN dispatch with Omni FNA for eligible stage-5 grids."""
+
+    def __call__(
+        self, attn: "LTX2VideoVaeNeighborhoodAttention", hidden_states: torch.Tensor, block_mask=None
+    ) -> torch.Tensor:
+        batch_size, num_frames, height, width, channels = hidden_states.shape
+        query, key, value = attn.project_qkv(hidden_states)
+        # NATTEN's CUTLASS kernels silently produce wrong output for non-contiguous inputs.
+        query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+        operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fna_eligible(hidden_states) else None
+        # `scale=1.0`: the query is already scaled in `project_qkv`, as in the reference.
+        if (
+            tuple(attn.kernel_size) == (11, 11, 11)
+            and hidden_states.dtype == torch.bfloat16
+            and channels == 256
+            and attn.head_dim == 64
+            and operators is not None
+            and operators.fna is not None
+            # The native metadata encodes KV tile indices in 15 bits.
+            and math.prod((n + tile - 1) // tile for n, tile in zip((num_frames, height, width), (4, 4, 8))) <= 1 << 15
+        ):
+            token_permute, token_unpermute = _load_natten_token_permutation()
+            query_permuted, token_shape, _ = token_permute(query, (4, 4, 4), flip_tiled_dims=True)
+            key_permuted, _, _ = token_permute(key, (4, 4, 8), flip_tiled_dims=True)
+            value_permuted, _, _ = token_permute(value, (4, 4, 8), flip_tiled_dims=True)
+            output = operators.fna(
+                query_permuted,
+                key_permuted,
+                value_permuted,
+                shape=(num_frames, height, width),
+                window=tuple(attn.kernel_size),
+            )
+            hidden_states = token_unpermute(
+                output,
+                token_layout_shape=token_shape,
+                tile_shape=(4, 4, 4),
+                flip_tiled_dims=True,
+            )
+        else:
+            hidden_states = self._na3d(
+                query,
+                key,
+                value,
+                kernel_size=attn.kernel_size,
+                scale=1.0,
+                backend=self.backend,
+            )
+        hidden_states = hidden_states.reshape(batch_size, num_frames, height, width, channels)
+        return attn.to_out[0](hidden_states)
+
+
+class LTX2VideoVaeNeighborhoodAttention(diffusers_decoder.LTX2VideoVaeNeighborhoodAttention):
+    _available_processors = [
+        *diffusers_decoder.LTX2VideoVaeNeighborhoodAttention._available_processors,
+        LTX2VideoVaeNeighborhoodNattenProcessor,
+    ]
+
+    def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size, num_frames, height, width, _ = hidden_states.shape
+        shape = (batch_size, num_frames, height, width, self.heads, self.head_dim)
+        query = self.to_q(hidden_states).view(shape)
+        key = self.to_k(hidden_states).view(shape)
+        value = self.to_v(hidden_states).view(shape)
+
+        operators = resolve_ltx2_vae_operators(query.device) if is_ltx2_fusion_eligible(query) else None
+        optimized = (
+            None
+            if operators is None
+            else operators.qk_norm_rope(
+                query,
+                key,
+                self.norm_q.weight,
+                self.norm_k.weight,
+                self.norm_q.eps,
+                self.scale,
+                self.rope.rope_dim_split,
+                self.rope.base,
+            )
+        )
+        if optimized is not None:
+            query, key = optimized
+            return query, key, value
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+        query = query * self.scale
+        return self.rope(query), self.rope(key), value
+
+
+class LTX2VideoVaeSwiGLU(diffusers_decoder.LTX2VideoVaeSwiGLU):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fusion_eligible(hidden_states) else None
+        optimized = (
+            None
+            if operators is None
+            else operators.swiglu(hidden_states, self.w_gate.weight, self.w_up.weight, self.w_down.weight)
+        )
+        return super().forward(hidden_states) if optimized is None else optimized
+
+
+class LTX2VideoVaeDiffusionNABlock(diffusers_decoder.LTX2VideoVaeDiffusionNABlock):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        latent_context: torch.Tensor,
+        modulation: tuple[torch.Tensor, ...],
+        block_mask=None,
+    ) -> torch.Tensor:
+        if hidden_states.dtype is not torch.bfloat16 or not is_ltx2_fusion_eligible(hidden_states):
+            return super().forward(hidden_states, latent_context, modulation, block_mask)
+        operators = resolve_ltx2_vae_operators(hidden_states.device)
+        assert operators is not None
+
+        scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
+            modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(self.num_mod_params)
+        ]
+        context_output = self.context_proj(latent_context)
+        attention_input = operators.residual_norm(
+            hidden_states,
+            context_output,
+            None,
+            self.norm1.weight,
+            scale_msa,
+            shift_msa,
+            self.norm1.eps,
+        )
+        if attention_input is None:
+            hidden_states = hidden_states + context_output
+            hidden_states = hidden_states + self.attn(
+                self.norm1(hidden_states) * (1 + scale_msa) + shift_msa, block_mask
+            )
+            hidden_states = hidden_states + self.mlp(self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp)
+            return hidden_states
+
+        attention_output = self.attn(attention_input, block_mask)
+        mlp_input = operators.residual_norm(
+            hidden_states,
+            context_output,
+            attention_output,
+            self.norm2.weight,
+            scale_mlp,
+            shift_mlp,
+            self.norm2.eps,
+        )
+        if mlp_input is None:
+            residual = hidden_states + context_output
+            residual = residual + attention_output
+            mlp_input = self.norm2(residual) * (1 + scale_mlp) + shift_mlp
+        mlp_output = self.mlp(mlp_input)
+        output = operators.residual_add(
+            hidden_states,
+            context_output,
+            attention_output,
+            mlp_output,
+        )
+        if output is not None:
+            return output
+        residual = hidden_states + context_output
+        residual = residual + attention_output
+        return residual + mlp_output
+
+
+class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d):
+    """Diffusers decoder core with the short-clip NATTEN context fix."""
+
+    def forward_stage_4(
+        self,
+        hidden_states: torch.Tensor,
+        drop_leading_frame: bool = True,
+        crop_trailing_ghost: bool = True,
+    ) -> torch.Tensor:
+        blocks = self.det_stages[-1]
+        block_mask = blocks[0].attn.build_block_mask(hidden_states)
+        for block in blocks:
+            hidden_states = block(hidden_states, block_mask)
+        hidden_states = self.upsamples[-1](hidden_states, drop_leading_frame=drop_leading_frame)
+
+        num_pad = self.trailing_pad_latent_frames
+        if crop_trailing_ghost and num_pad > 0:
+            content_frames = max(hidden_states.shape[1] - num_pad * self.temporal_compression_ratio, 1)
+            keep_frames = min(hidden_states.shape[1], max(content_frames, self.stage5_kernel[0]))
+            hidden_states = hidden_states[:, :keep_frames]
+        return hidden_states
+
+
+class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel):
+    """Diffusers 0.40 decoder plus canonical Native loading and short-clip support."""
+
+    def __init__(
+        self,
+        out_channels: int = 3,
+        latent_channels: int = 128,
+        patch_size: int = 4,
+        scaling_factor: float = 1.0,
+        decoder_head_dim: int = 64,
+        decoder_stage_channels: tuple[int, ...] = (2048, 1024, 512, 512, 256),
+        decoder_stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8),
+        decoder_stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)),
+        decoder_upsample_strides: tuple[tuple[int, int, int], ...] = ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)),
+        decoder_upsample_channel_reductions: tuple[int, ...] = (2, 2, 1, 2),
+        decoder_stage5_kernel: tuple[int, int, int] = (11, 11, 11),
+        decoder_t_emb_dim: int = 384,
+        decoder_timestep_scale_multiplier: float = 1000.0,
+        decoder_model_output_type: str = "x0",
+        decoder_num_inference_steps: int = 1,
+        spatial_compression_ratio: int = 32,
+        temporal_compression_ratio: int = 8,
+    ) -> None:
+        super().__init__(
+            out_channels=out_channels,
+            latent_channels=latent_channels,
+            patch_size=patch_size,
+            scaling_factor=scaling_factor,
+            decoder_head_dim=decoder_head_dim,
+            decoder_stage_channels=decoder_stage_channels,
+            decoder_stage_depths=decoder_stage_depths,
+            decoder_stage_kernels=decoder_stage_kernels,
+            decoder_upsample_strides=decoder_upsample_strides,
+            decoder_upsample_channel_reductions=decoder_upsample_channel_reductions,
+            decoder_stage5_kernel=decoder_stage5_kernel,
+            decoder_t_emb_dim=decoder_t_emb_dim,
+            decoder_timestep_scale_multiplier=decoder_timestep_scale_multiplier,
+            decoder_model_output_type=decoder_model_output_type,
+            decoder_num_inference_steps=decoder_num_inference_steps,
+            spatial_compression_ratio=spatial_compression_ratio,
+            temporal_compression_ratio=temporal_compression_ratio,
+        )
+        # Diffusers constructs its private decoder core internally. Switching
+        # to this behavior-only subclass preserves parameters and state-dict keys.
+        self.decoder.__class__ = LTX2VideoDiffusionDecoder3d
+        self.decoder.stage5_kernel = tuple(decoder_stage5_kernel)
+        overrides = {
+            diffusers_decoder.LTX2VideoVaeNeighborhoodAttention: LTX2VideoVaeNeighborhoodAttention,
+            diffusers_decoder.LTX2VideoVaeSwiGLU: LTX2VideoVaeSwiGLU,
+            diffusers_decoder.LTX2VideoVaeDiffusionNABlock: LTX2VideoVaeDiffusionNABlock,
+        }
+        # Preserve every module and parameter object; only execution methods change.
+        for module in self.decoder.modules():
+            if replacement := overrides.get(type(module)):
+                module.__class__ = replacement
+
+    @classmethod
+    def from_ltx25_native_checkpoint(
+        cls,
+        checkpoint_path: str,
+        config: Mapping[str, Any],
+        dtype: torch.dtype,
+    ) -> "LTX2VideoDiffusionDecoderModel":
+        """Construct this decoder class directly from the canonical Native checkpoint."""
+        with torch.device("meta"):
+            model = cls.from_config(dict(config))
+        state_dict = load_ltx25_native_diffusion_decoder_state_dict(checkpoint_path)
+        try:
+            model.load_state_dict(state_dict, strict=True, assign=True)
+        except RuntimeError as exc:
+            raise ValueError(f"Invalid LTX-2.5 Native DiffVAE checkpoint {checkpoint_path!r}.") from exc
+        model.to(device="cpu", dtype=dtype)
+        return model
+
+    def decode(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | None = None,
+        num_inference_steps: int | None = None,
+        return_dict: bool = True,
+    ) -> DecoderOutput | tuple[torch.Tensor]:
+        decoded = super().decode(
+            z,
+            generator=generator,
+            num_inference_steps=num_inference_steps,
+            return_dict=False,
+        )[0]
+
+        # Short clips retain replicated trailing context through stage 5. Do
+        # not expose those context pixels to the caller.
+        target_num_frames = (z.shape[2] - 1) * self.temporal_compression_ratio + 1
+        target_height = z.shape[3] * self.spatial_compression_ratio
+        target_width = z.shape[4] * self.spatial_compression_ratio
+        decoded = decoded[:, :, :target_num_frames, :target_height, :target_width]
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)