chore: import upstream snapshot with attribution

utils/DATASET_PROMPT_LOGIC.md

@@ -0,0 +1,243 @@
+# Dataset Prompt 分配逻辑说明
+
+本文档说明 `MultiTextConcatDataset`（纯文本）和 `MultiVideoConcatDataset`（视频训练）在不同配置下产生 `prompts` 的行为。
+
+---
+
+## 架构总览
+
+| 使用场景 | Dataset 类 | 输出 |
+|---|---|---|
+| diffusion.py 训练 | `MultiVideoConcatDataset` | `frames` + `prompts` |
+| diffusion.py 推理 | `MultiTextConcatDataset` | `prompts` |
+| distillation.py 训练（backward_sim） | `MultiTextConcatDataset` | `prompts` |
+| distillation.py 可视化 | `MultiTextConcatDataset` | `prompts` |
+| inference.py | `MultiTextConcatDataset` | `prompts` |
+
+---
+
+## 一、MultiTextConcatDataset（纯文本）
+
+### 关键参数
+
+| 参数 | 含义 |
+|---|---|
+| `num_blocks` | 输出 prompt 列表的固定长度 |
+| `chunks_per_shot` | 每个 shot 重复的 block 数（0=使用 even_durations 均分） |
+| `scene_cut_prefix` | 切镜标记前缀，默认 `"The scene transitions. "` |
+
+简写约定：`0`, `1`, `2` 表示不同 caption；`p+X` 表示 `scene_cut_prefix + X`
+
+### 1. txt 模式（data_path 指向 .txt 文件）
+
+每行一个 caption，每个 sample 取 `idx` 行的 caption，重复 `num_blocks` 次。
+**不加** scene_cut_prefix（单 shot 语义）。
+
+```
+data_path="prompts.txt", line[3]="A dog runs", num_blocks=12
+→ [A, A, A, A, A, A, A, A, A, A, A, A]
+```
+
+无论 `chunks_per_shot` 设为多少，txt 模式始终单 shot 重复。
+
+### 2. 目录模式（data_path 指向目录）
+
+读取 `caption/<subfolder>/*.json`（不需要 `video/` 目录）。
+每个 JSON 的 `caption` 字段作为一个 shot 的文本。
+
+#### Shot duration 三级 fallback
+
+按优先级决定每个 shot 占多少个 block：
+
+1. **`shot_durations.txt`**（per-folder 文件）— 每行或逗号分隔的整数
+2. **`chunks_per_shot`**（全局 config）— 所有 shot 统一重复固定次数
+3. **`_even_durations`**（均分）— 将 `num_blocks` 均匀分给所有 caption
+
+#### 长度处理
+
+输出始终恰好 `num_blocks` 个 prompt：
+- **超过** → 截断尾部
+- **不足** → 用最后一个 caption 直接 padding（**不加** scene_cut_prefix）
+
+#### 示例
+
+**3 caption, num_blocks=12, chunks_per_shot=4**
+
+```
+captions: [0, 1, 2]
+shot_durations: [4, 4, 4]
+→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2]
+```
+
+**3 caption, num_blocks=12, chunks_per_shot=0（even_durations）**
+
+```
+captions: [0, 1, 2]
+even_durations: base=4, extra=0 → [4, 4, 4]
+→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2]
+```
+
+**2 caption, num_blocks=12, chunks_per_shot=4**
+
+```
+captions: [0, 1]
+shot_durations: [4, 4], sum=8 < 12 → 最后一个 shot 扩展到 8
+→ [0, 0, 0, 0, p+1, 1, 1, 1, 1, 1, 1, 1]
+                                  ↑ padding 不加 prefix
+```
+
+**5 caption, num_blocks=12, chunks_per_shot=4**
+
+```
+captions: [0, 1, 2, 3, 4]
+shot_durations: [4, 4, 4, 4, 4], 但 num_blocks=12 → clamped 到 [4, 4, 4]
+→ [0, 0, 0, 0, p+1, 1, 1, 1, p+2, 2, 2, 2]
+   caption 3 和 4 被截断
+```
+
+**2 caption, num_blocks=12, chunks_per_shot=0（even_durations）**
+
+```
+captions: [0, 1]
+even_durations: base=6, extra=0 → [6, 6]
+→ [0, 0, 0, 0, 0, 0, p+1, 1, 1, 1, 1, 1]
+```
+
+**1 caption, num_blocks=12**
+
+```
+captions: [0]
+→ [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+   单 shot，不加 prefix
+```
+
+**shot_durations.txt 覆盖**
+
+```
+captions: [0, 1, 2], shot_durations.txt 内容: "2, 6, 4", num_blocks=12
+→ [0, 0, p+1, 1, 1, 1, 1, 1, p+2, 2, 2, 2]
+```
+
+---
+
+## 二、MultiVideoConcatDataset（视频训练）
+
+### 关键参数
+
+| 参数 | 含义 |
+|---|---|
+| `total_segments` | 总 segment 数量（= 1 + num_subsequent_segments） |
+| `single_video_only` | config 中的 `uniform_prompt`，True 时只从一个视频采样 |
+| `max_chunks_per_shot` | 单镜头最大连续 chunk 数，超过则跳 1 秒做虚拟切镜（0=不限制） |
+| `scene_cut_prefix` | 切镜标记前缀 |
+| `allow_padding` | 视频不够时是否允许 padding（否则跳过该 folder） |
+
+Prompt 由实际视频采样决定，逐 segment 从视频文件加载对应的 per-video caption。
+
+### 1. 多视频自然拼接（`max_chunks_per_shot=0`，默认）
+
+按视频文件顺序采样，切换视频文件时加 `scene_cut_prefix`：
+
+```
+video A 够采 3 chunks, video B 够采 4 chunks, total_segments=7
+→ [A, A, A, p+B, B, B, B]
+```
+
+### 2. `single_video_only=True`
+
+强制只从一个视频文件采样。视频不够长则整个 folder 被跳过：
+
+```
+video A 够采 7 chunks, total_segments=7
+→ [A, A, A, A, A, A, A]
+
+video A 只够采 5 chunks, total_segments=7
+→ (失败，跳到下一个 folder)
+```
+
+### 3. `max_chunks_per_shot=3`（限制单镜头时长）
+
+从同一视频连续采超过 3 chunks 后，跳 1 秒做虚拟切镜，加 `scene_cut_prefix`：
+
+```
+video A 很长, video B, total_segments=7, max_chunks_per_shot=3
+→ [A, A, A, p+A, A, A, p+B]
+             ↑跳1秒虚拟切镜   ↑换视频
+```
+
+如果跳 1 秒后 A 不够了，直接跳到 B：
+
+```
+video A 够 4 chunks（跳1秒后不够）, video B, total_segments=7, max_chunks_per_shot=3
+→ [A, A, A, p+B, B, B, B]
+             ↑A跳1秒后不够，换B
+```
+
+### 4. `single_video_only=True` + `max_chunks_per_shot=3`
+
+单视频内也可以做虚拟切镜：
+
+```
+video A 很长, total_segments=7, single_video_only=True, max_chunks_per_shot=3
+→ [A, A, A, p+A, A, A, p+A]
+```
+
+### 5. 训练 padding（`allow_padding=True`）
+
+视频不够时用最后一个 caption 直接 padding（不加 prefix）：
+
+```
+video A 够采 3 chunks, video B 够采 2 chunks, total_segments=7, allow_padding=True
+→ [A, A, A, p+B, B, B, B]
+                      ↑后 2 个用最后一个 caption padding
+```
+
+---
+
+## 三、Multi-Shot Sink
+
+通过 config 的 `multi_shot_sink: true` 开启。
+开启后，当检测到某个 block 处于新场景的起始位置时，会在该 block 去噪完成并更新 cache 后，
+将 KV cache 的 attention sink 从旧场景的第一帧迁移到新场景的第一帧。
+同时会自动把全局 sink 长度设为 `sink_size`，不需要单独配置 global sink。
+
+**配置方式**（yaml config 中添加）：
+
+```yaml
+sink_size: 8
+multi_shot_sink: true    # 默认 false，不迁移 sink
+```
+
+此选项在以下所有场景均生效：
+
+| 场景 | Pipeline | 检测方式 |
+|---|---|---|
+| **Diffusion trainer evaluation** | `CausalDiffusionInferencePipeline` | 检查 prompt 是否以 `scene_cut_prefix` 开头 |
+| **Distillation trainer evaluation** | `CausalDiffusionInferencePipeline` | 同上 |
+| **离线推理 `inference.py`** | `CausalDiffusionInferencePipeline` | 同上 |
+| **Distillation backward simulation** | `SelfForcingTrainingPipeline` | trainer 预计算 `scene_cut_mask` 传入 `conditional_dict` |
+| **Streaming long tuning** | `SelfForcingTrainingPipeline` | 同上（mask 随 chunk 自动 slice） |
+
+### 实现机制
+
+**Inference pipeline**（`CausalDiffusionInferencePipeline`）：
+直接检查每个 chunk 的 raw prompt 是否以 `scene_cut_prefix` 开头。
+
+**Training pipeline**（`SelfForcingTrainingPipeline`）：
+由于 prompts 在进入 pipeline 前已编码为 embedding，无法从 embedding 反推文本。
+因此 trainer 在编码前从原始 prompts 计算布尔列表 `scene_cut_mask`，
+放入 `conditional_dict["scene_cut_mask"]` 一路透传到 pipeline。
+对于 streaming training，`scene_cut_mask` 在 `_slice_cond_dict_for_chunk` 中随 prompt 一起 slice。
+
+```
+block:  [0]  [1]  [2]  [p+3]  [4]  [5]  [p+6]  [7]
+mask:    F    F    F    T      F    F    T       F
+sink:    0    0    0   →3      3    3   →6       6
+                       ↑更新sink          ↑更新sink
+```
+
+当 `multi_shot_sink: false`（默认）时，sink 始终锚定在视频的第一帧，不做任何迁移。
+
+`scene_cut_prefix` 由 `DEFAULT_SCENE_CUT_PREFIX` 常量定义（`dataset.py`），
+inference pipeline 引用同一常量，确保训练和推理的切镜检测一致。
+可通过 config 中的 `scene_cut_prefix` 字段统一覆盖。

utils/__init__.py

@@ -0,0 +1,5 @@
+# Marker file: turn `utils/` from a namespace package into a regular package.
+# torch 2.12's torchrun + multiprocessing has trouble resolving namespace
+# packages from cwd in subprocesses; making this an explicit regular package
+# makes `from utils.position_embedding_utils import ...` (used inside
+# wan_5b/modules/model.py) reliable across torch versions.

utils/adaln_triton.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Triton fused adaLN-modulation kernel.
+
+iter-43: replaces the
+    (norm(x).unflatten(1, (F, frame_seqlen)) * (1 + e_scale) + e_shift).flatten(1, 2)
+chain (LayerNorm + 2 broadcast elementwise ops per call, x2 per transformer block
+for norm1/norm2) with a single Triton kernel. Each token does one fp32 pass:
+mean/var reduce, normalize, multiply by (1+scale), add shift, cast back.
+
+LayerNorm matches `WanLayerNorm` (nn.LayerNorm with elementwise_affine=False,
+eps=1e-6, output cast back to input dtype).
+"""
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+# iter-44: autotune over num_warps + num_stages. Fixed BLOCK_C
+# (next_power_of_2(C)) — varying it would change the reduce semantics. Provide
+# enough configs to cover the (small_B, large_L) regime of inference chunks.
+_ADALN_CONFIGS = [
+    triton.Config({}, num_warps=nw, num_stages=ns)
+    for nw in (4, 8, 16)
+    for ns in (1, 2, 3)
+]
+
+
+@triton.autotune(configs=_ADALN_CONFIGS, key=["C", "FRAME_SEQLEN"])
+@triton.jit
+def _adaln_modulate_kernel(
+    x_ptr,        # [B, L, C]
+    scale_ptr,    # [B, F, 1, C] (or any layout, indexed via strides)
+    shift_ptr,    # [B, F, 1, C]
+    out_ptr,      # [B, L, C]
+    C,
+    FRAME_SEQLEN,
+    x_stride_b, x_stride_l,
+    scale_stride_b, scale_stride_f,
+    shift_stride_b, shift_stride_f,
+    eps,
+    ADD_ONE_TO_SCALE: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_l = tl.program_id(1)
+    pid_f = pid_l // FRAME_SEQLEN
+
+    offs_c = tl.arange(0, BLOCK_C)
+    mask = offs_c < C
+
+    x_off = pid_b * x_stride_b + pid_l * x_stride_l + offs_c
+    x = tl.load(x_ptr + x_off, mask=mask, other=0.0).to(tl.float32)
+
+    inv_C = 1.0 / C
+    mean = tl.sum(x, axis=0) * inv_C
+    x_centered = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(x_centered * x_centered, axis=0) * inv_C
+    rstd = 1.0 / tl.sqrt(var + eps)
+    # Match WanLayerNorm: nn.LayerNorm casts back to input dtype via .type_as(x)
+    # before downstream ops. Round-trip through bf16 to keep numerics identical
+    # to the eager path so latent diff stays in run-to-run noise floor.
+    x_norm = (x_centered * rstd).to(tl.bfloat16).to(tl.float32)
+
+    s_off = pid_b * scale_stride_b + pid_f * scale_stride_f + offs_c
+    h_off = pid_b * shift_stride_b + pid_f * shift_stride_f + offs_c
+    scale = tl.load(scale_ptr + s_off, mask=mask, other=0.0).to(tl.float32)
+    shift = tl.load(shift_ptr + h_off, mask=mask, other=0.0).to(tl.float32)
+
+    if ADD_ONE_TO_SCALE:
+        # Match eager: `1 + e_scale` happens in bf16 (autocast off at this site)
+        scale = (scale + 1.0)
+        scale = scale.to(tl.bfloat16).to(tl.float32)
+
+    # Match eager: bf16 * bf16, bf16 + bf16
+    prod = (x_norm * scale).to(tl.bfloat16).to(tl.float32)
+    y = prod + shift
+
+    tl.store(out_ptr + x_off, y, mask=mask)
+
+
+def adaln_modulate_triton(
+    x: torch.Tensor,        # [B, L, C] contiguous
+    e_scale: torch.Tensor,  # [B, F, 1, C]
+    e_shift: torch.Tensor,  # [B, F, 1, C]
+    frame_seqlen: int,
+    eps: float = 1e-6,
+    add_one_to_scale: bool = True,
+) -> torch.Tensor:
+    """Fused (LayerNorm + (1+e_scale)*x + e_shift) over [B, F*frame_seqlen, C].
+
+    Replaces the WanAttentionBlock norm1/norm2 + modulate pattern. Output dtype
+    follows x.dtype; internal arithmetic is fp32. eps matches `WanLayerNorm`.
+    """
+    assert x.dim() == 3, f"x must be [B, L, C], got {x.shape}"
+    B, L, C = x.shape
+    assert L % frame_seqlen == 0, (L, frame_seqlen)
+    F = L // frame_seqlen
+    assert e_scale.dim() == 4 and e_scale.shape[0] == B and e_scale.shape[1] == F \
+        and e_scale.shape[2] == 1 and e_scale.shape[3] == C, \
+        f"e_scale {tuple(e_scale.shape)} vs expected {(B, F, 1, C)}"
+    assert e_shift.shape == e_scale.shape
+
+    out = torch.empty_like(x)
+    BLOCK_C = triton.next_power_of_2(C)
+
+    grid = (B, L)
+    # num_warps / num_stages picked by @triton.autotune (iter-44).
+    _adaln_modulate_kernel[grid](
+        x, e_scale, e_shift, out,
+        C, frame_seqlen,
+        x.stride(0), x.stride(1),
+        e_scale.stride(0), e_scale.stride(1),
+        e_shift.stride(0), e_shift.stride(1),
+        eps,
+        ADD_ONE_TO_SCALE=add_one_to_scale,
+        BLOCK_C=BLOCK_C,
+    )
+    return out

utils/config.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from omegaconf import OmegaConf
+
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，"
+    "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，"
+    "画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，"
+    "手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+)
+
+
+wan_default_config = {
+    "Wan2.2-TI2V-5B": {
+        "resolution": [1280, 704],
+        "temporal_compression_ratio": 4,
+        "spatial_compression_ratio": 16,
+        "num_heads": 24,
+        "head_dim": 128,
+        "num_transformer_blocks": 30,
+        "fps": 24,
+    }
+}
+
+
+SECTION_KEYS = (
+    "infra",
+    "algorithm",
+    "training",
+    "data",
+    "evaluation",
+    "inference",
+    "logging",
+    "checkpoints",
+)
+
+
+def _set_once(config, key, value, source):
+    if value is None:
+        return
+    if key in config and config[key] != value:
+        raise ValueError(
+            f"{key} is defined more than once with different values: "
+            f"{config[key]} vs {value} from {source}."
+        )
+    config[key] = value
+
+
+def section_get(config, section_key, key, default=None, aliases=()):
+    """Read a grouped config value, falling back to legacy flat names."""
+    section = config.get(section_key, None)
+    candidate_keys = (key, *aliases)
+    if section is not None:
+        for candidate in candidate_keys:
+            if candidate in section:
+                return section[candidate]
+    for candidate in candidate_keys:
+        if candidate in config:
+            return config[candidate]
+    return default
+
+
+def normalize_config(config):
+    """Expand grouped release configs into the flat runtime schema.
+
+    The training and inference code historically reads fields such as
+    ``config.batch_size`` and ``config.model_kwargs`` directly.  Release
+    configs can group those fields for readability, then call this function at
+    the entry point to preserve the existing runtime contract.
+    """
+    for section_key in SECTION_KEYS:
+        section = config.get(section_key, None)
+        if section is None:
+            continue
+        for key, value in section.items():
+            config[key] = value
+
+    evaluation = config.get("evaluation", None)
+    if evaluation is not None:
+        if "interval" in evaluation:
+            _set_once(config, "generate_interval", evaluation.interval, "evaluation.interval")
+            _set_once(config, "vis_interval", evaluation.interval, "evaluation.interval")
+        if "num_frames" in evaluation:
+            num_frames = evaluation.num_frames
+            if isinstance(num_frames, (list, tuple)):
+                vis_lengths = list(num_frames)
+                inference_num_frames = vis_lengths[0] if vis_lengths else 0
+            else:
+                inference_num_frames = int(num_frames)
+                vis_lengths = [inference_num_frames]
+            _set_once(config, "inference_num_frames", inference_num_frames, "evaluation.num_frames")
+            _set_once(config, "vis_video_lengths", vis_lengths, "evaluation.num_frames")
+        if "use_ema" in evaluation:
+            _set_once(config, "vis_ema", evaluation.use_ema, "evaluation.use_ema")
+
+    model_section = config.get("model", None)
+    base_model_kwargs = config.get("model_kwargs", None)
+    model_kwargs = OmegaConf.create({})
+    if base_model_kwargs is not None:
+        model_kwargs = OmegaConf.merge(model_kwargs, base_model_kwargs)
+
+    if model_section is not None:
+        section_kwargs = model_section.get("kwargs", None)
+        if section_kwargs is not None:
+            model_kwargs = OmegaConf.merge(model_kwargs, section_kwargs)
+
+        model_name = model_section.get("name", None)
+        if model_name is not None:
+            model_kwargs.model_name = model_name
+            config.model_name = model_name
+
+        _set_once(
+            config,
+            "num_frame_per_block",
+            model_section.get("num_frame_per_block", None),
+            "model.num_frame_per_block",
+        )
+
+    if "model_name" in config and "model_name" not in model_kwargs:
+        model_kwargs.model_name = config.model_name
+    if "timestep_shift" in config and "timestep_shift" not in model_kwargs:
+        model_kwargs.timestep_shift = config.timestep_shift
+    if "timestep_shift" in model_kwargs:
+        _set_once(config, "timestep_shift", model_kwargs.timestep_shift, "model_kwargs.timestep_shift")
+
+    model_num_frame_per_block = model_kwargs.get("num_frame_per_block", None)
+    if model_num_frame_per_block is not None:
+        _set_once(config, "num_frame_per_block", model_num_frame_per_block, "model_kwargs.num_frame_per_block")
+
+    if len(model_kwargs) > 0:
+        config.model_kwargs = model_kwargs
+
+    if "wandb_host" not in config:
+        config.wandb_host = "https://api.wandb.ai"
+
+    if "negative_prompt" not in config:
+        config.negative_prompt = DEFAULT_NEGATIVE_PROMPT
+
+    if config.get("trainer", None) == "score_distillation":
+        dmd_defaults = {
+            "i2v": False,
+            "teacher_forcing": False,
+            "backward_simulation": True,
+            "independent_first_frame": False,
+            "num_train_timestep": 1000,
+            "denoising_loss_type": "flow",
+            "real_guidance_scale": 3.0,
+            "fake_guidance_scale": 0.0,
+        }
+        for key, value in dmd_defaults.items():
+            if key not in config:
+                config[key] = value
+        if "causal" not in config:
+            config.causal = bool(config.get("all_causal", True))
+
+    # Causal DMD uses the same Wan backbone for generator/teacher/critic unless
+    # a role-specific override is explicitly provided.
+    if getattr(config, "all_causal", False) and "model_kwargs" in config:
+        for role_key in ("real_model_kwargs", "fake_model_kwargs"):
+            if config.get(role_key, None) is None:
+                config[role_key] = OmegaConf.create(
+                    OmegaConf.to_container(config.model_kwargs, resolve=True)
+                )
+
+    return config

utils/distributed.py

@@ -0,0 +1,149 @@
+from datetime import timedelta
+from functools import partial
+import os
+import torch
+import torch.distributed as dist
+from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, StateDictType
+from torch.distributed.fsdp.api import CPUOffload
+from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
+
+
+def fsdp_state_dict(model):
+    fsdp_fullstate_save_policy = FullStateDictConfig(
+        offload_to_cpu=True, rank0_only=True
+    )
+    with FSDP.state_dict_type(
+        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
+    ):
+        checkpoint = model.state_dict()
+
+    return checkpoint
+
+
+def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None, cpu_offload=False):
+    if mixed_precision:
+        mixed_precision_policy = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            buffer_dtype=torch.float32,
+            cast_forward_inputs=False
+        )
+    else:
+        mixed_precision_policy = None
+
+    if wrap_strategy == "transformer":
+        auto_wrap_policy = partial(
+            transformer_auto_wrap_policy,
+            transformer_layer_cls=transformer_module
+        )
+    elif wrap_strategy == "size":
+        auto_wrap_policy = partial(
+            size_based_auto_wrap_policy,
+            min_num_params=min_num_params
+        )
+    else:
+        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")
+
+    os.environ["NCCL_CROSS_NIC"] = "1"
+
+    sharding_strategy = {
+        "full": ShardingStrategy.FULL_SHARD,
+        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
+        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
+        "no_shard": ShardingStrategy.NO_SHARD,
+    }[sharding_strategy]
+
+    module = FSDP(
+        module,
+        auto_wrap_policy=auto_wrap_policy,
+        sharding_strategy=sharding_strategy,
+        mixed_precision=mixed_precision_policy,
+        device_id=torch.cuda.current_device(),
+        limit_all_gathers=True,
+        use_orig_params=True,
+        cpu_offload=CPUOffload(offload_params=cpu_offload),
+        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks
+    )
+    return module
+
+
+def barrier():
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def launch_distributed_job(backend: str = "nccl"):
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    host = os.environ["MASTER_ADDR"]
+    port = int(os.environ["MASTER_PORT"])
+
+    if ":" in host:  # IPv6
+        init_method = f"tcp://[{host}]:{port}"
+    else:  # IPv4
+        init_method = f"tcp://{host}:{port}"
+    # Use a long timeout so that slow collectives during checkpoint saving
+    # (e.g. FSDP.optim_state_dict all-gather + rank0-only disk write for a
+    # multi-GB full optimizer state) do not trip the NCCL watchdog on other
+    # ranks while they wait at the post-save barrier.
+    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,
+                            init_method=init_method, timeout=timedelta(minutes=60))
+    torch.cuda.set_device(local_rank)
+
+
+class EMA_FSDP:
+    def __init__(self, fsdp_module: torch.nn.Module, decay: float = 0.999):
+        self.decay = decay
+        self.shadow = {}
+        self._init_shadow(fsdp_module)
+
+    @staticmethod
+    def _clean_param_name(name: str) -> str:
+        """Remove FSDP wrapper prefixes from parameter names."""
+        return name.replace("_fsdp_wrapped_module.", "").replace("_checkpoint_wrapped_module.", "").replace("_orig_mod.", "")
+
+    @torch.no_grad()
+    def _init_shadow(self, fsdp_module):
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                # Clean the parameter name to remove FSDP prefixes
+                # This ensures shadow keys are compatible with unwrapped models for inference
+                cleaned_name = self._clean_param_name(n)
+                self.shadow[cleaned_name] = p.detach().clone().float().cpu()
+
+    @torch.no_grad()
+    def update(self, fsdp_module):
+        d = self.decay
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=False):
+            for n, p in fsdp_module.module.named_parameters():
+                cleaned_name = self._clean_param_name(n)
+                if cleaned_name in self.shadow:
+                    self.shadow[cleaned_name].mul_(d).add_(p.detach().float().cpu(), alpha=1. - d)
+
+    # Optional helpers ---------------------------------------------------
+    def state_dict(self):
+        # Return shadow dict directly - keys are already cleaned during init/update
+        # This makes the state_dict directly usable for inference with unwrapped models
+        return self.shadow            # picklable
+
+    def load_state_dict(self, sd):
+        # Handle both cases: with or without FSDP prefixes
+        # This ensures backward compatibility and flexibility
+        cleaned_sd = {}
+        for k, v in sd.items():
+            # Remove FSDP prefixes if present to match internal naming convention
+            cleaned_key = self._clean_param_name(k)
+            cleaned_sd[cleaned_key] = v.clone()
+        self.shadow = cleaned_sd
+
+    def copy_to(self, fsdp_module):
+        # load EMA weights into an (unwrapped) copy of the generator
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+        with FSDP.summon_full_params(fsdp_module, writeback=True):
+            for n, p in fsdp_module.module.named_parameters():
+                cleaned_name = self._clean_param_name(n)
+                if cleaned_name in self.shadow:
+                    p.data.copy_(self.shadow[cleaned_name].to(p.dtype, device=p.device))

utils/error_buffer.py

@@ -0,0 +1,321 @@
+# Adopted from https://github.com/vita-epfl/Stable-Video-Infinity
+# SPDX-License-Identifier: Apache-2.0
+
+import random
+import torch
+
+
+class ErrorBuffer:
+    """Bucketed ring buffer for storing prediction errors on CPU.
+
+    Two layouts are supported:
+
+      * **1D (timestep-only)** — when ``num_blocks <= 0``. Buckets are keyed by
+        the diffusion timestep. This is the original SVI behavior.
+
+      * **2D (position × timestep)** — when ``num_blocks > 0``. Each entry is
+        keyed by both the global block position along the sequence and the
+        timestep. Inject paths can then choose:
+          - ``sample(pos, t)``: match BOTH position and timestep
+            (E_vid / E_noise — noise-level dependent errors)
+          - ``sample_pos_any_t(pos)``: match position, sample uniformly across
+            timesteps (E_img — position-dependent context corruption that is
+            agnostic to the current denoising step)
+          - ``sample_global()``: legacy fallback, samples uniformly everywhere
+
+    The 2D layout encodes the teacher-forcing insight that ``noisy_suffix[i]``
+    looks at clean_prefix[0..i] during training but at model rollouts during
+    inference; storing prediction errors per-position therefore lets later
+    blocks self-feed larger errors without any manual position ramp.
+
+    **Sharded timestep buckets** (``shard_size > 1``):
+    Each rank only allocates the timestep buckets it owns
+    (``t_bucket % shard_size == shard_rank``), reducing per-rank CPU memory
+    by ~``shard_size`` times.  Typically ``shard_rank/shard_size`` are set to
+    ``sp_rank/sp_size`` so that sharding is per-SP-rank and saving follows
+    the same per-SP-rank pattern as the 2D position split.  On ``add()``,
+    non-owned buckets are silently skipped; on ``sample()``, non-owned buckets
+    are remapped to the nearest owned one.
+    """
+
+    def __init__(
+        self,
+        num_buckets=40,
+        max_size_per_bucket=50,
+        num_train_timesteps=1000,
+        modulate_factor=0.3,
+        replacement_strategy="random",
+        num_blocks=0,
+        global_block_offset=0,
+        shard_rank=0,
+        shard_size=1,
+    ):
+        self.num_buckets = num_buckets
+        self.max_size = max_size_per_bucket
+        self.num_train_timesteps = num_train_timesteps
+        self.modulate_factor = modulate_factor
+        self.replacement_strategy = replacement_strategy
+        self.bucket_width = num_train_timesteps / num_buckets
+        self.num_blocks = int(num_blocks) if num_blocks else 0
+        # ``global_block_offset`` is only used for stats / debug display so
+        # users can tell which absolute positions of the full sequence this
+        # buffer covers (the LAST SP rank carries the highest accumulated
+        # error positions). It does NOT participate in bucket keying.
+        self.global_block_offset = int(global_block_offset)
+
+        self.shard_rank = int(shard_rank)
+        self.shard_size = max(int(shard_size), 1)
+        self._owned_t_buckets = sorted(
+            t for t in range(num_buckets) if t % self.shard_size == self.shard_rank
+        )
+
+        if self.num_blocks > 0:
+            self.buckets = {
+                (p, t): []
+                for p in range(self.num_blocks)
+                for t in self._owned_t_buckets
+            }
+        else:
+            self.buckets = {t: [] for t in self._owned_t_buckets}
+        self.total_added = 0
+
+    # ------------------------------------------------------------------ keys
+    def _t_bucket(self, timestep_index):
+        b = int(timestep_index / self.bucket_width)
+        return max(0, min(b, self.num_buckets - 1))
+
+    def _is_owned_t(self, t_bucket):
+        return self.shard_size <= 1 or (t_bucket % self.shard_size == self.shard_rank)
+
+    def _nearest_owned_t(self, t_bucket):
+        """Remap ``t_bucket`` to the closest owned timestep bucket."""
+        if self.shard_size <= 1 or self._is_owned_t(t_bucket):
+            return t_bucket
+        fwd = (self.shard_rank - t_bucket % self.shard_size) % self.shard_size
+        bwd = self.shard_size - fwd
+        t_up, t_down = t_bucket + fwd, t_bucket - bwd
+        up_ok = 0 <= t_up < self.num_buckets
+        down_ok = 0 <= t_down < self.num_buckets
+        if up_ok and down_ok:
+            return t_up if fwd <= bwd else t_down
+        return t_up if up_ok else t_down
+
+    def _make_key(self, t_bucket, block_pos):
+        if self.num_blocks > 0:
+            assert block_pos is not None, "block_pos required when num_blocks>0"
+            p = max(0, min(int(block_pos), self.num_blocks - 1))
+            return (p, t_bucket)
+        return t_bucket
+
+    # ------------------------------------------------------------------ add
+    def add(self, error_block, timestep_index, block_pos=None):
+        """Store a single block error into the matching bucket.
+
+        Args:
+            error_block: (block_size, C, H, W) tensor
+            timestep_index: int, raw index in [0, num_train_timesteps)
+            block_pos: int, global block position; required iff num_blocks>0
+        """
+        t = self._t_bucket(timestep_index)
+        if not self._is_owned_t(t):
+            return
+        key = self._make_key(t, block_pos)
+        # Store in the source dtype on CPU to match SVI (which keeps bf16),
+        # cutting buffer memory in half vs. casting to fp32.
+        entry = error_block.detach().to("cpu", copy=True)
+
+        buf = self.buckets[key]
+        if len(buf) < self.max_size:
+            buf.append(entry)
+        else:
+            if self.replacement_strategy == "fifo":
+                buf.pop(0)
+                buf.append(entry)
+            elif self.replacement_strategy == "l2":
+                stacked = torch.stack(buf)
+                dists = (stacked - entry.unsqueeze(0)).flatten(1).norm(dim=1)
+                most_similar = torch.argmin(dists).item()
+                buf[most_similar] = entry
+            else:  # "random" (default)
+                idx = random.randint(0, self.max_size - 1)
+                buf[idx] = entry
+        self.total_added += 1
+
+    # ------------------------------------------------------------------ sample
+    def sample(self, timestep_index, device, dtype, block_pos=None):
+        """Sample one entry matching (block_pos, timestep_index) when 2D, or
+        just timestep_index when 1D.  Non-owned timestep buckets are
+        transparently remapped to the nearest owned one.  Returns None if the
+        (remapped) bucket is empty."""
+        t = self._nearest_owned_t(self._t_bucket(timestep_index))
+        key = self._make_key(t, block_pos)
+        buf = self.buckets[key]
+        if not buf:
+            return None
+        err = random.choice(buf)
+        return self._modulate(err).to(device=device, dtype=dtype)
+
+    def sample_pos_any_t(self, block_pos, device, dtype):
+        """For 2D buffers: sample at the given position, with random timestep.
+
+        This is the natural choice for context (E_img) injection — the clean
+        prefix is the result of a full ODE rollout so its accumulated error
+        could have originated at any timestep, but its magnitude scales with
+        position along the sequence.
+
+        Falls back to ``sample_global`` when the buffer is 1D.
+        Only owned timestep buckets are scanned.
+        """
+        if self.num_blocks <= 0:
+            return self.sample_global(device, dtype)
+        p = max(0, min(int(block_pos), self.num_blocks - 1))
+        all_entries = []
+        for t in self._owned_t_buckets:
+            all_entries.extend(self.buckets[(p, t)])
+        if not all_entries:
+            return None
+        err = random.choice(all_entries)
+        return self._modulate(err).to(device=device, dtype=dtype)
+
+    def sample_global(self, device, dtype):
+        """Sample one entry uniformly from ALL buckets (legacy SVI E_img)."""
+        all_entries = []
+        for buf in self.buckets.values():
+            all_entries.extend(buf)
+        if not all_entries:
+            return None
+        err = random.choice(all_entries)
+        return self._modulate(err).to(device=device, dtype=dtype)
+
+    # ------------------------------------------------------------------ misc
+    def _modulate(self, err):
+        if self.modulate_factor > 0:
+            lo = 1.0 - self.modulate_factor
+            hi = 1.0 + self.modulate_factor
+            err = err * random.uniform(lo, hi)
+        return err
+
+    def is_empty(self):
+        return self.total_added == 0
+
+    def has_pos(self, block_pos):
+        """Whether ANY owned timestep bucket at ``block_pos`` has samples (2D only)."""
+        if self.num_blocks <= 0:
+            return not self.is_empty()
+        p = max(0, min(int(block_pos), self.num_blocks - 1))
+        return any(len(self.buckets[(p, t)]) > 0 for t in self._owned_t_buckets)
+
+    def stats(self):
+        filled = sum(1 for b in self.buckets.values() if len(b) > 0)
+        total = sum(len(b) for b in self.buckets.values())
+        num_owned_t = len(self._owned_t_buckets)
+        denom = self.num_blocks * num_owned_t if self.num_blocks > 0 else num_owned_t
+        out = {
+            "total_added": self.total_added,
+            "filled_buckets": f"{filled}/{denom}",
+            "total_entries": total,
+        }
+        if self.shard_size > 1:
+            out["shard"] = f"shard_rank={self.shard_rank}/{self.shard_size} ({num_owned_t}/{self.num_buckets} t-buckets)"
+        if self.num_blocks > 0:
+            lo = self.global_block_offset
+            hi = self.global_block_offset + self.num_blocks
+            out["global_block_range"] = f"[{lo},{hi})"
+        return out
+
+    def state_dict(self):
+        # Keys are tuples (pos, t) when 2D — torch.save handles them fine
+        # via pickle. We serialize the bucket layout so loaders can validate.
+        return {
+            "buckets": {k: list(v) for k, v in self.buckets.items()},
+            "total_added": self.total_added,
+            "num_blocks": self.num_blocks,
+            "num_buckets": self.num_buckets,
+            "global_block_offset": self.global_block_offset,
+            "shard_rank": self.shard_rank,
+            "shard_size": self.shard_size,
+        }
+
+    def load_state_dict(self, state, strict_offset=True):
+        """Restore buckets from a serialized state.
+
+        Args:
+            state: dict produced by ``state_dict``.
+            strict_offset: when True (default) and the buffer is 2D,
+                refuse to load if the saved ``global_block_offset`` does
+                not match the current one. This prevents the silent
+                position-misalignment bug under SP, where a checkpoint
+                saved by SP rank 0 (covering global blocks ``[0, B)``)
+                would otherwise be loaded into SP rank 1 (which expects
+                ``[B, 2B)``) and corrupt position-bucketed sampling.
+                Pass ``strict_offset=False`` only for backward-compat
+                with checkpoints saved before this field existed.
+        """
+        if self.num_blocks > 0 and strict_offset:
+            saved_off = state.get("global_block_offset", None)
+            if saved_off is None:
+                raise RuntimeError(
+                    "Refusing to load: this is a 2D position-bucketed buffer "
+                    "but the checkpoint has no `global_block_offset` field. "
+                    "Pass strict_offset=False if you accept the misalignment risk."
+                )
+            if int(saved_off) != self.global_block_offset:
+                raise RuntimeError(
+                    f"Refusing to load: checkpoint covers global blocks "
+                    f"starting at {saved_off}, but this rank covers blocks "
+                    f"starting at {self.global_block_offset}. Make sure each "
+                    f"SP rank loads its own per-rank checkpoint file."
+                )
+        # Shard check: warn but don't crash if shard layout changed (e.g.
+        # resuming a non-sharded checkpoint into a sharded buffer is fine —
+        # we just load whichever buckets overlap).
+        saved_shard_size = int(state.get("shard_size", state.get("dp_size", 1)))
+        saved_shard_rank = int(state.get("shard_rank", state.get("dp_rank", 0)))
+        if saved_shard_size != self.shard_size or saved_shard_rank != self.shard_rank:
+            import logging
+            logging.warning(
+                f"[ErrorBuffer] Shard layout changed: checkpoint was "
+                f"shard_rank={saved_shard_rank}/{saved_shard_size}, current is "
+                f"shard_rank={self.shard_rank}/{self.shard_size}. "
+                f"Loading overlapping buckets only."
+            )
+
+        saved = state["buckets"]
+        # Lenient match: ignore keys that don't exist in the current layout.
+        for k in self.buckets:
+            if k in saved:
+                self.buckets[k] = saved[k]
+            elif isinstance(k, tuple):
+                # Try string-form key from older serializations
+                continue
+            elif str(k) in saved:
+                self.buckets[k] = saved[str(k)]
+        self.total_added = int(state.get("total_added", 0))
+
+
+def build_error_buffer(config, num_blocks=0, global_block_offset=0,
+                       shard_rank=0, shard_size=1):
+    """Build an ErrorBuffer from an OmegaConf/dict config node.
+
+    When ``num_blocks > 0`` the buffer becomes 2D (position × timestep),
+    enabling teacher-forcing-aware position-dependent error injection.
+    Pass ``global_block_offset`` so logs can identify which absolute slice
+    of the full sequence this rank's buffer covers (e.g. the last SP rank
+    is responsible for the most error-accumulated tail blocks).
+
+    ``shard_rank`` / ``shard_size`` shard timestep buckets: each rank only
+    allocates the buckets it owns, reducing per-rank CPU memory by
+    ~``shard_size`` times.  Typically set to ``(sp_rank, sp_size)``.
+    """
+    cfg = config if isinstance(config, dict) else dict(config)
+    return ErrorBuffer(
+        num_buckets=cfg.get("num_buckets", 40),
+        max_size_per_bucket=cfg.get("buffer_size_per_bucket", 50),
+        num_train_timesteps=cfg.get("num_train_timesteps", 1000),
+        modulate_factor=cfg.get("modulate_factor", 0.3),
+        replacement_strategy=cfg.get("replacement_strategy", "random"),
+        num_blocks=num_blocks,
+        global_block_offset=global_block_offset,
+        shard_rank=shard_rank,
+        shard_size=shard_size,
+    )

utils/fp8.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""TorchAO FP8 post-training quantization helpers."""
+
+import torch
+import torch.nn as nn
+
+
+# Small conditioning/output projections are both numerically sensitive and too
+# small to amortize dynamic activation quantization on H100.
+_BF16_MODULES = {
+    "text_embedding.0",
+    "text_embedding.2",
+    "time_embedding.0",
+    "time_embedding.2",
+    "time_projection.1",
+    "head.head",
+}
+
+
+def quantize_model_fp8(model: nn.Module, *, verbose: bool = False) -> int:
+    """Quantize compatible BF16 linear layers to row-wise dynamic FP8 in-place."""
+    if not torch.cuda.is_available():
+        raise RuntimeError("TorchAO FP8 inference requires a CUDA GPU.")
+
+    device = next(model.parameters()).device
+    if device.type != "cuda":
+        raise ValueError("Move the BF16 model to CUDA before applying FP8 quantization.")
+    if torch.cuda.get_device_capability(device) < (8, 9):
+        raise RuntimeError("TorchAO FP8 inference requires compute capability 8.9 or newer.")
+
+    try:
+        from torchao.quantization import (
+            Float8DynamicActivationFloat8WeightConfig,
+            PerRow,
+            quantize_,
+        )
+    except ImportError as exc:
+        raise ImportError(
+            "FP8 inference requires a TorchAO version compatible with the installed PyTorch."
+        ) from exc
+
+    quantized_names = []
+    skipped_names = []
+
+    def filter_fn(module: nn.Module, fqn: str) -> bool:
+        if not isinstance(module, nn.Linear):
+            return False
+        if fqn in _BF16_MODULES:
+            skipped_names.append(fqn)
+            return False
+        if module.weight.dtype != torch.bfloat16:
+            raise TypeError(f"FP8 layer {fqn!r} must be BF16, got {module.weight.dtype}.")
+        out_features, in_features = module.weight.shape
+        if in_features % 16 or out_features % 16:
+            skipped_names.append(fqn)
+            return False
+        quantized_names.append(fqn)
+        return True
+
+    quantize_(
+        model,
+        Float8DynamicActivationFloat8WeightConfig(granularity=PerRow()),
+        filter_fn=filter_fn,
+    )
+    if verbose:
+        print(
+            f"[FP8] TorchAO W8A8 quantized {len(quantized_names)} linear layers "
+            f"with row-wise scaling; kept {len(skipped_names)} layers in BF16"
+        )
+    return len(quantized_names)

utils/i2v_conditioning.py

@@ -0,0 +1,71 @@
+import torch
+
+
+def _get_i2v_context_frames(
+    image_or_video: torch.Tensor,
+    initial_latent: torch.Tensor | None,
+) -> int:
+    if initial_latent is None:
+        return 0
+    if image_or_video.ndim != initial_latent.ndim:
+        raise ValueError(
+            f"initial_latent rank {initial_latent.ndim} must match "
+            f"image/video rank {image_or_video.ndim}."
+        )
+    if image_or_video.shape[0] != initial_latent.shape[0]:
+        raise ValueError(
+            f"initial_latent batch {initial_latent.shape[0]} must match "
+            f"image/video batch {image_or_video.shape[0]}."
+        )
+    if image_or_video.shape[2:] != initial_latent.shape[2:]:
+        raise ValueError(
+            f"initial_latent shape after frames {tuple(initial_latent.shape[2:])} "
+            f"must match image/video {tuple(image_or_video.shape[2:])}."
+        )
+
+    context_frames = int(initial_latent.shape[1])
+    if context_frames <= 0:
+        return 0
+    if context_frames >= image_or_video.shape[1]:
+        raise ValueError(
+            f"initial_latent has {context_frames} frames but clip has "
+            f"{image_or_video.shape[1]} frames."
+        )
+    return context_frames
+
+
+def _overwrite_i2v_context(
+    image_or_video: torch.Tensor,
+    initial_latent: torch.Tensor | None,
+    context_frames: int,
+) -> torch.Tensor:
+    if context_frames <= 0:
+        return image_or_video
+    output = image_or_video.clone()
+    output[:, :context_frames] = initial_latent[:, :context_frames].to(
+        device=output.device,
+        dtype=output.dtype,
+    )
+    return output
+
+
+def _zero_i2v_context_timestep(
+    timestep: torch.Tensor,
+    context_frames: int,
+) -> torch.Tensor:
+    if context_frames <= 0:
+        return timestep
+    output = timestep.clone()
+    output[:, :context_frames] = 0
+    return output
+
+
+def _i2v_loss_mask_like(
+    image_or_video: torch.Tensor,
+    context_frames: int,
+) -> torch.Tensor | None:
+    if context_frames <= 0:
+        return None
+    mask = torch.ones_like(image_or_video, dtype=torch.bool)
+    mask[:, :context_frames] = False
+    return mask

utils/inference_utils.py

@@ -0,0 +1,312 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Small helpers for release inference examples."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Sequence
+
+import torch
+from einops import rearrange
+from torchvision.io import write_video
+
+from utils.nvfp4_checkpoint import (
+    clean_fsdp_state_dict_keys,
+    drop_fouroversix_master_weights,
+    is_nvfp4_state_dict,
+    is_te_nvfp4_checkpoint,
+    quantize_model_for_fouroversix_nvfp4,
+    quantize_model_for_transformer_engine_nvfp4,
+    unwrap_generator_state_dict,
+)
+
+
+def _torch_load(path: str):
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+def load_generator_checkpoint(generator, checkpoint_path: str, *, use_ema: bool = False, strict: bool | None = None):
+    """Load a LongLive generator checkpoint into ``generator``."""
+    checkpoint = _torch_load(checkpoint_path)
+    state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema)
+    if use_ema:
+        state_dict = clean_fsdp_state_dict_keys(state_dict)
+    if strict is None:
+        strict = not use_ema
+    return generator.load_state_dict(state_dict, strict=strict)
+
+
+def _load_lora_state_dict(lora_ckpt_path: str) -> Mapping[str, torch.Tensor]:
+    """Load a LoRA checkpoint, unwrapping ``generator_lora`` when present."""
+    checkpoint = _torch_load(lora_ckpt_path)
+    if isinstance(checkpoint, Mapping) and "generator_lora" in checkpoint:
+        return checkpoint["generator_lora"]
+    return checkpoint
+
+
+def apply_and_merge_lora(
+    pipeline,
+    config,
+    *,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype = torch.bfloat16,
+    verbose: bool = False,
+):
+    """Wrap ``pipeline.generator.model`` with a LoRA adapter, load weights, and merge.
+
+    The merged module ends up structurally identical to the original generator
+    (``nn.Linear`` layers carrying the base + LoRA delta), which is what NVFP4
+    quantization needs as its starting point.
+
+    Returns ``True`` when LoRA was applied and merged, ``False`` when the config
+    did not request a LoRA adapter.
+    """
+    adapter_cfg = getattr(config, "adapter", None)
+    lora_ckpt = getattr(config, "lora_ckpt", None)
+    if adapter_cfg is None or not lora_ckpt:
+        return False
+
+    import peft
+    from utils.lora_utils import configure_lora_for_model
+
+    if device is not None:
+        pipeline.generator.to(device=torch.device(device), dtype=dtype)
+    else:
+        pipeline.generator.to(dtype=dtype)
+
+    if verbose:
+        print(f"[LoRA] Wrapping generator with adapter config: {adapter_cfg}")
+    pipeline.generator.model = configure_lora_for_model(
+        pipeline.generator.model,
+        model_name="generator",
+        lora_config=adapter_cfg,
+        is_main_process=verbose,
+    )
+
+    if verbose:
+        print(f"[LoRA] Loading LoRA weights from: {lora_ckpt}")
+    lora_state = _load_lora_state_dict(lora_ckpt)
+    peft.set_peft_model_state_dict(pipeline.generator.model, lora_state)  # type: ignore[arg-type]
+
+    if verbose:
+        print("[LoRA] Merging LoRA delta into base weights (merge_and_unload)...")
+    pipeline.generator.model = pipeline.generator.model.merge_and_unload(safe_merge=True)
+    pipeline.generator.model.eval().requires_grad_(False)
+    pipeline.is_lora_enabled = False
+    pipeline.is_lora_merged = True
+    return True
+
+
+def place_vae_for_streaming(pipeline, config) -> torch.device | None:
+    """Move ``pipeline.vae`` to ``config.vae_device`` for streaming-pipeline decode.
+
+    Only acts when both ``streaming_vae`` and ``vae_device`` are set; otherwise
+    leaves the VAE on whatever device the rest of the pipeline already uses.
+    Mirrors the relocation done in ``inference.py`` so that quick-start scripts
+    can opt in to the streaming-pipeline VAE simply by enabling those config
+    fields.
+    """
+    if not bool(getattr(config, "streaming_vae", False)):
+        return None
+    vae_device_str = getattr(config, "vae_device", None)
+    if not vae_device_str:
+        return None
+
+    vae_device = torch.device(vae_device_str)
+    pipeline.vae.to(device="cpu")
+    pipeline.vae.to(device=vae_device)
+    if hasattr(pipeline.vae, "mean"):
+        pipeline.vae.mean = pipeline.vae.mean.to(device=vae_device)
+        pipeline.vae.std = pipeline.vae.std.to(device=vae_device)
+    return vae_device
+
+
+def setup_nvfp4_pipeline(
+    pipeline,
+    config,
+    device: torch.device | str,
+    *,
+    verbose: bool = False,
+):
+    """Configure ``pipeline`` for NVFP4 inference from a generator checkpoint.
+
+    Handles both supported NVFP4 backends:
+
+    * ``model_quant_use_transformer_engine=True`` -> a BF16 generator checkpoint
+      that gets wrapped with TransformerEngine NVFP4 modules and materialized
+      after moving to ``device``.
+    * ``model_quant_use_transformer_engine=False`` -> either a BF16 generator
+      checkpoint that gets quantized with FourOverSix at load time, or a
+      pre-materialized FourOverSix NVFP4 state dict loaded directly into the
+      already-quantized architecture.
+
+    Optional LoRA support (BF16 base only): when ``config.adapter`` and
+    ``config.lora_ckpt`` are both set, the LoRA adapter is loaded on the BF16
+    base generator, merged via ``merge_and_unload``, and the resulting weights
+    are then quantized — so the same yaml can swap between TE and FourOverSix
+    backends without pre-merging the LoRA checkpoint.
+
+    For materialized FourOverSix checkpoints LoRA cannot be applied (the master
+    weights have already been quantized away); ``lora_ckpt``/``adapter`` are
+    ignored in that case with a printed warning.
+    """
+    if not bool(getattr(config, "model_quant", False)):
+        raise ValueError("setup_nvfp4_pipeline requires model_quant=true in the config.")
+
+    generator_ckpt = getattr(config, "generator_ckpt", None)
+    if not generator_ckpt:
+        raise ValueError("checkpoints.generator_ckpt is required for NVFP4 inference.")
+
+    use_te = bool(getattr(config, "model_quant_use_transformer_engine", False))
+    device = torch.device(device)
+    use_ema = bool(getattr(config, "use_ema", False))
+
+    checkpoint = _torch_load(generator_ckpt)
+    state_dict = unwrap_generator_state_dict(checkpoint, use_ema=use_ema)
+    if use_ema:
+        state_dict = clean_fsdp_state_dict_keys(state_dict)
+
+    if is_te_nvfp4_checkpoint(checkpoint):
+        raise ValueError(
+            "Detected a TransformerEngine module state_dict export (no longer supported). "
+            "Re-export with `--backend transformer_engine` (merged BF16) or `--backend fouroversix`."
+        )
+
+    is_prequantized = is_nvfp4_state_dict(state_dict)
+    has_lora_request = bool(getattr(config, "adapter", None)) and bool(getattr(config, "lora_ckpt", None))
+
+    pipeline.is_lora_enabled = False
+    pipeline.is_lora_merged = False
+
+    if is_prequantized:
+        if has_lora_request and verbose:
+            print(
+                "[NVFP4] generator_ckpt is a materialized FourOverSix NVFP4 checkpoint; "
+                "ignoring lora_ckpt/adapter because the master weights are already quantized. "
+                "Use a BF16 base checkpoint if you need to load a LoRA on top."
+            )
+        if use_te:
+            raise ValueError(
+                "generator_ckpt is a materialized NVFP4 (FourOverSix) checkpoint; set "
+                "model_quant_use_transformer_engine: false."
+            )
+        pipeline.generator.model, _ = quantize_model_for_fouroversix_nvfp4(
+            pipeline.generator.model,
+            config=config,
+            keep_master_weights=False,
+            verbose=verbose,
+        )
+        drop_fouroversix_master_weights(pipeline.generator.model)
+        pipeline.generator.load_state_dict(state_dict, strict=True)
+
+        pipeline.text_encoder.to(dtype=torch.bfloat16)
+        pipeline.vae.to(dtype=torch.bfloat16)
+    else:
+        load_strict = not use_ema
+        pipeline.generator.load_state_dict(state_dict, strict=load_strict)
+
+        if has_lora_request:
+            # Apply + merge LoRA on the BF16 base before quantization. Move the
+            # generator to CUDA first so the TE wrapper (which requires CUDA
+            # modules) can later replace the merged Linear layers in-place.
+            apply_and_merge_lora(
+                pipeline,
+                config,
+                device=device,
+                dtype=torch.bfloat16,
+                verbose=verbose,
+            )
+
+        if use_te:
+            pipeline.generator.model, _ = quantize_model_for_transformer_engine_nvfp4(
+                pipeline.generator.model,
+                config=config,
+                keep_master_weights=False,
+                verbose=verbose,
+            )
+            te_fallback = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False))
+            if te_fallback:
+                from utils.quant import _materialize_mixed_quantized_weights_for_inference as materialize_fn
+            else:
+                from utils.quant import _materialize_transformer_engine_weights_for_inference as materialize_fn
+        else:
+            pipeline.generator.model, _ = quantize_model_for_fouroversix_nvfp4(
+                pipeline.generator.model,
+                config=config,
+                keep_master_weights=False,
+                verbose=verbose,
+            )
+            from utils.quant import _materialize_quantized_weights_for_inference as materialize_fn
+
+        pipeline.to(dtype=torch.bfloat16)
+        materialize_fn(pipeline.generator.model, target_device=device)
+
+    pipeline.generator.to(device=device)
+    pipeline.text_encoder.to(device=device)
+    pipeline.vae.to(device=device)
+    place_vae_for_streaming(pipeline, config)
+
+    return pipeline
+
+
+def prepare_single_prompt_inputs(
+    config,
+    prompt: str,
+    device: torch.device | str,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    batch_size: int = 1,
+    generator: torch.Generator | None = None,
+):
+    """Create the per-block prompt list and latent noise for one text prompt."""
+    num_frames = int(getattr(config, "num_output_frames", config.image_or_video_shape[1]))
+    frames_per_block = int(getattr(config, "num_frame_per_block", 1))
+    if num_frames % frames_per_block != 0:
+        raise ValueError(f"num_frames={num_frames} must be divisible by num_frame_per_block={frames_per_block}")
+
+    latent_shape = list(config.image_or_video_shape[2:])
+    if len(latent_shape) != 3:
+        raise ValueError(f"Expected latent shape [C, H, W], got {latent_shape}")
+
+    num_blocks = num_frames // frames_per_block
+    prompts = [[prompt] * num_blocks for _ in range(batch_size)]
+    noise = torch.randn(
+        [batch_size, num_frames, *latent_shape],
+        device=device,
+        dtype=dtype,
+        generator=generator,
+    )
+    return noise, prompts
+
+
+def video_to_uint8(video: torch.Tensor) -> torch.Tensor:
+    """Convert a generated video tensor from [T, C, H, W] or [1, T, C, H, W] to uint8 THWC."""
+    if video.ndim == 5:
+        if video.shape[0] != 1:
+            raise ValueError("video_to_uint8 expects a single sample when a batch dimension is present.")
+        video = video[0]
+    if video.ndim != 4:
+        raise ValueError(f"Expected video tensor with 4 dims, got shape={tuple(video.shape)}")
+    if video.shape[1] in (1, 3):
+        video = rearrange(video, "t c h w -> t h w c")
+    return (255.0 * video.cpu()).clamp(0, 255).to(torch.uint8)
+
+
+def save_video(video: torch.Tensor, output_path: str | os.PathLike, *, fps: int = 24) -> None:
+    """Save a generated LongLive video tensor as an mp4 file."""
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    write_video(str(output_path), video_to_uint8(video), fps=fps)

utils/kernel/README.md

@@ -0,0 +1,54 @@
+<!--
+Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+
+Licensed under the Apache License, Version 2.0 (the "License").
+You may not use this file except in compliance with the License.
+To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+
+No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# LongLive KV Dequant CUDA Extension
+
+Build from this directory:
+
+```bash
+cd utils/kernel
+OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 MAX_JOBS=4 \
+  python setup.py build_ext --inplace
+```
+
+Runtime import:
+
+```python
+from utils.kernel.kv_dequant import dequantize_kv_cache_fp4
+```
+
+`utils.quant.dequantize_kv_cache()` already calls this extension first and falls
+back to the original Triton path if the extension is not built.
+
+For direct calls, pass the same scale limits used by the QuantizedTensor's
+`scale_rule`:
+
+- `static_6`: `e2m1_max=6.0`, `e4m3_max=448.0`
+- `static_4`: `e2m1_max=4.0`, `e4m3_max=448.0`
+- `mse` / `l1_norm` / `abs_max` 4o6 modes: `e2m1_max=6.0`, `e4m3_max=256.0`
+
+The normal `utils.quant.dequantize_kv_cache()` path reads these values from
+`qt.scale_rule`, so manual selection is not needed there.
+
+You can also pass `scale_rule` directly:
+
+```python
+out = dequantize_kv_cache_fp4(
+    values,
+    scale_factors,
+    amax,
+    num_heads=num_heads,
+    block_token_size=block_token_size,
+    dtype=torch.bfloat16,
+    scale_rule="static_6",
+)
+```

utils/kernel/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Custom CUDA kernels used by LongLive."""

utils/kernel/kv_dequant.cpp

@@ -0,0 +1,20 @@
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+//
+// Licensed under the Apache License, Version 2.0 (the "License").
+// You may not use this file except in compliance with the License.
+// To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+//
+// No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+//
+// SPDX-License-Identifier: Apache-2.0
+#include <torch/extension.h>
+
+TORCH_LIBRARY(longlive_kernels, m)
+{
+    m.def("dequantize_kv_cache_fp4(Tensor[] values, Tensor[] scale_factors, Tensor[] amax, int num_heads, int block_token_size, int dtype_code, float e2m1_max, float e4m3_max) -> Tensor");
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.doc() = "LongLive custom CUDA kernels";
+}

utils/kernel/kv_dequant.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+try:
+    from . import longlive_kv_dequant_cuda  # noqa: F401
+except ImportError:
+    import longlive_kv_dequant_cuda  # noqa: F401
+
+
+def _dtype_to_code(dtype: torch.dtype) -> int:
+    if dtype == torch.bfloat16:
+        return 0
+    if dtype == torch.float16:
+        return 1
+    if dtype == torch.float32:
+        return 2
+    raise ValueError(f"Unsupported fused KV dequant dtype: {dtype}")
+
+
+def scale_rule_to_fp4_limits(scale_rule) -> tuple[float, float]:
+    """Return the dequant denominator limits used by FourOverSix ScaleRule."""
+    if hasattr(scale_rule, "max_allowed_e2m1_value") and hasattr(
+        scale_rule, "max_allowed_e4m3_value",
+    ):
+        return (
+            float(scale_rule.max_allowed_e2m1_value()),
+            float(scale_rule.max_allowed_e4m3_value()),
+        )
+
+    normalized = str(scale_rule).lower()
+    if "." in normalized:
+        normalized = normalized.rsplit(".", 1)[-1]
+    normalized = normalized.strip().strip("\"'")
+
+    if normalized == "static_4":
+        return 4.0, 448.0
+    if normalized == "static_6":
+        return 6.0, 448.0
+    if normalized in {"mse", "mae", "l1_norm", "abs_max"}:
+        return 6.0, 256.0
+
+    raise ValueError(f"Unsupported FP4 scale_rule: {scale_rule}")
+
+
+def dequantize_kv_cache_fp4(
+    values: list[torch.Tensor],
+    scale_factors: list[torch.Tensor],
+    amax: list[torch.Tensor],
+    *,
+    num_heads: int,
+    block_token_size: int,
+    dtype: torch.dtype,
+    e2m1_max: float | None = None,
+    e4m3_max: float | None = None,
+    scale_rule=None,
+) -> torch.Tensor:
+    """Dequantize multiple AR KV-cache chunks with one CUDA launch."""
+    if e2m1_max is None or e4m3_max is None:
+        if scale_rule is None:
+            raise ValueError(
+                "Either e2m1_max/e4m3_max or scale_rule must be provided.",
+            )
+        e2m1_max, e4m3_max = scale_rule_to_fp4_limits(scale_rule)
+
+    return torch.ops.longlive_kernels.dequantize_kv_cache_fp4.default(
+        values,
+        scale_factors,
+        amax,
+        num_heads,
+        block_token_size,
+        _dtype_to_code(dtype),
+        e2m1_max,
+        e4m3_max,
+    )

utils/kernel/kv_dequant_cuda.cu

@@ -0,0 +1,244 @@
+// Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+//
+// Licensed under the Apache License, Version 2.0 (the "License").
+// You may not use this file except in compliance with the License.
+// To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+//
+// No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+//
+// SPDX-License-Identifier: Apache-2.0
+#include <ATen/Dispatch.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_fp16.h>
+#include <cuda_fp4.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace {
+
+#define CHECK_CUDA_TENSOR(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
+
+__device__ __constant__ float kE2M1ToFloat[16] = {
+    0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
+    -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f,
+};
+
+// iter-37: hardware FP4→FP16x2 via CUDA 12.8+ built-in API
+// __nv_cvt_fp4x2_to_halfraw2 (wraps cvt.rn.f16x2.e2m1x2 PTX instruction).
+// Returns __half2_raw with 2 fp16 values from 1 packed byte.
+__device__ __forceinline__ __half2_raw e2m1x2_to_halfraw2(uint8_t byte) {
+    return __nv_cvt_fp4x2_to_halfraw2(
+        static_cast<__nv_fp4x2_storage_t>(byte), __NV_E2M1);
+}
+
+__device__ __forceinline__ int64_t blocked_scale_index(
+    const int row,
+    const int scale_col,
+    const int scale_cols)
+{
+    // Inverse of fouroversix.quantize.utils.to_blocked for a scale matrix
+    // shaped [rows_padded, scale_cols].
+    const int row_block = row / 128;
+    const int row_in_block = row - row_block * 128;
+    const int scale_col_block = scale_col / 4;
+    const int scale_col_in_block = scale_col - scale_col_block * 4;
+    const int scale_col_blocks = scale_cols / 4;
+
+    const int logical_block = row_block * scale_col_blocks + scale_col_block;
+    return (((int64_t)logical_block * 32 + (row_in_block & 31)) * 16
+            + (row_in_block >> 5) * 4 + scale_col_in_block);
+}
+
+template <typename scalar_t>
+__global__ void fp4_kv_dequant_kernel(
+    const uint64_t* __restrict__ value_ptrs,
+    const uint64_t* __restrict__ scale_ptrs,
+    const uint64_t* __restrict__ amax_ptrs,
+    scalar_t* __restrict__ output,
+    const int64_t total_packed_values,
+    const int block_token_size,
+    const int num_heads,
+    const int packed_cols,
+    const int scale_cols,
+    const float inv_global_scale_denom)
+{
+    const int64_t packed_idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+    if (packed_idx >= total_packed_values) {
+        return;
+    }
+
+    const int col_pair = packed_idx % packed_cols;
+    const int64_t global_row = packed_idx / packed_cols;
+    const int rows_per_cache_block = block_token_size * num_heads;
+    const int cache_block = global_row / rows_per_cache_block;
+    const int row_in_cache_block = global_row - (int64_t)cache_block * rows_per_cache_block;
+
+    const int token_in_block = row_in_cache_block / num_heads;
+    const int head = row_in_cache_block - token_in_block * num_heads;
+    const int out_token = cache_block * block_token_size + token_in_block;
+
+    const auto* values = reinterpret_cast<const uint8_t*>(value_ptrs[cache_block]);
+    const auto* scales = reinterpret_cast<const __nv_fp8_e4m3*>(scale_ptrs[cache_block]);
+    const auto* amax = reinterpret_cast<const float*>(amax_ptrs[cache_block]);
+
+    const uint8_t packed = values[(int64_t)row_in_cache_block * packed_cols + col_pair];
+    const int scale_col = (col_pair * 2) / 16;
+    const int64_t scale_idx = blocked_scale_index(row_in_cache_block, scale_col, scale_cols);
+    const float scale = static_cast<float>(scales[scale_idx]);
+    const float global_scale = amax[0] * inv_global_scale_denom;
+
+    // iter-37: hardware FP4→FP16x2 via CUDA 12.8 built-in (wraps cvt.rn.f16x2.e2m1x2).
+    const __half2_raw f16x2 = e2m1x2_to_halfraw2(packed);
+    // __half2_raw layout: .x = low nibble's fp16 (unsigned short), .y = high nibble's.
+    const float low = __half2float(__ushort_as_half(f16x2.x)) * scale * global_scale;
+    const float high = __half2float(__ushort_as_half(f16x2.y)) * scale * global_scale;
+
+    const int out_col = col_pair * 2;
+    const int64_t out_base = (((int64_t)out_token * num_heads + head) * (packed_cols * 2)) + out_col;
+    output[out_base] = static_cast<scalar_t>(low);
+    output[out_base + 1] = static_cast<scalar_t>(high);
+}
+
+at::ScalarType dtype_code_to_scalar_type(const int64_t dtype_code)
+{
+    switch (dtype_code) {
+    case 0:
+        return at::ScalarType::BFloat16;
+    case 1:
+        return at::ScalarType::Half;
+    case 2:
+        return at::ScalarType::Float;
+    default:
+        TORCH_CHECK(false, "Unsupported KV dequant dtype code: ", dtype_code);
+    }
+    return at::ScalarType::Float;
+}
+
+at::Tensor make_device_pointer_tensor(at::TensorList tensors)
+{
+    auto options = at::TensorOptions()
+                       .dtype(at::ScalarType::Long)
+                       .device(tensors.front().device());
+    at::Tensor ptrs = at::empty({static_cast<int64_t>(tensors.size())}, options);
+
+    std::vector<int64_t> host_ptrs(tensors.size());
+    for (size_t i = 0; i < tensors.size(); ++i) {
+        host_ptrs[i] = reinterpret_cast<int64_t>(tensors[i].data_ptr());
+    }
+
+    // The pointer table is tiny; use a synchronous copy so the temporary host
+    // vector cannot outlive an async H2D transfer.
+    C10_CUDA_CHECK(cudaMemcpy(
+        ptrs.data_ptr<int64_t>(),
+        host_ptrs.data(),
+        host_ptrs.size() * sizeof(int64_t),
+        cudaMemcpyHostToDevice));
+    return ptrs;
+}
+
+}  // namespace
+
+at::Tensor dequantize_kv_cache_fp4_cuda(
+    at::TensorList values,
+    at::TensorList scale_factors,
+    at::TensorList amax,
+    int64_t num_heads,
+    int64_t block_token_size,
+    int64_t dtype_code,
+    double e2m1_max,
+    double e4m3_max)
+{
+    TORCH_CHECK(!values.empty(), "values must contain at least one cache block");
+    TORCH_CHECK(values.size() == scale_factors.size(),
+                "values and scale_factors must have the same length");
+    TORCH_CHECK(values.size() == amax.size(),
+                "values and amax must have the same length");
+    TORCH_CHECK(num_heads > 0, "num_heads must be positive");
+    TORCH_CHECK(block_token_size > 0, "block_token_size must be positive");
+    TORCH_CHECK(e2m1_max > 0.0 && e4m3_max > 0.0,
+                "e2m1_max and e4m3_max must be positive");
+
+    const auto device = values.front().device();
+    c10::cuda::CUDAGuard device_guard(device);
+    const int64_t max_blocks = static_cast<int64_t>(values.size());
+    const int64_t packed_cols = values.front().size(1);
+    const int64_t head_dim = packed_cols * 2;
+    const int64_t rows_padded = values.front().size(0);
+    const int64_t logical_rows = block_token_size * num_heads;
+    const int64_t scale_cols = head_dim / 16;
+
+    TORCH_CHECK(head_dim == 128, "KV dequant currently expects head_dim=128, got ", head_dim);
+    TORCH_CHECK(scale_cols % 4 == 0, "scale column count must be a multiple of 4");
+    TORCH_CHECK(rows_padded >= logical_rows,
+                "values rows are smaller than logical KV block rows");
+    TORCH_CHECK(rows_padded % 128 == 0, "values rows must be padded to a multiple of 128");
+
+    for (int64_t i = 0; i < max_blocks; ++i) {
+        CHECK_CUDA_TENSOR(values[i]);
+        CHECK_CUDA_TENSOR(scale_factors[i]);
+        CHECK_CUDA_TENSOR(amax[i]);
+        CHECK_CONTIGUOUS(values[i]);
+        CHECK_CONTIGUOUS(scale_factors[i]);
+        CHECK_CONTIGUOUS(amax[i]);
+        TORCH_CHECK(values[i].device() == device, "all values tensors must be on the same device");
+        TORCH_CHECK(scale_factors[i].device() == device,
+                    "all scale_factors tensors must be on the same device");
+        TORCH_CHECK(amax[i].device() == device, "all amax tensors must be on the same device");
+        TORCH_CHECK(values[i].scalar_type() == at::ScalarType::Byte,
+                    "values tensors must be uint8");
+        TORCH_CHECK(amax[i].scalar_type() == at::ScalarType::Float,
+                    "amax tensors must be float32");
+        TORCH_CHECK(values[i].dim() == 2, "values tensors must be 2D");
+        TORCH_CHECK(values[i].size(0) == rows_padded && values[i].size(1) == packed_cols,
+                    "all values tensors must have the same shape");
+    }
+
+    const auto out_dtype = dtype_code_to_scalar_type(dtype_code);
+    at::Tensor output = at::empty(
+        {1, max_blocks * block_token_size, num_heads, head_dim},
+        values.front().options().dtype(out_dtype));
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+    at::Tensor value_ptrs = make_device_pointer_tensor(values);
+    at::Tensor scale_ptrs = make_device_pointer_tensor(scale_factors);
+    at::Tensor amax_ptrs = make_device_pointer_tensor(amax);
+
+    const int64_t total_packed_values = max_blocks * logical_rows * packed_cols;
+    const int threads = 256;
+    const dim3 blocks((total_packed_values + threads - 1) / threads);
+    const float inv_global_scale_denom = static_cast<float>(1.0 / (e2m1_max * e4m3_max));
+
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half,
+        at::ScalarType::BFloat16,
+        output.scalar_type(),
+        "fp4_kv_dequant_kernel",
+        [&] {
+            fp4_kv_dequant_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
+                reinterpret_cast<const uint64_t*>(value_ptrs.data_ptr<int64_t>()),
+                reinterpret_cast<const uint64_t*>(scale_ptrs.data_ptr<int64_t>()),
+                reinterpret_cast<const uint64_t*>(amax_ptrs.data_ptr<int64_t>()),
+                output.data_ptr<scalar_t>(),
+                total_packed_values,
+                static_cast<int>(block_token_size),
+                static_cast<int>(num_heads),
+                static_cast<int>(packed_cols),
+                static_cast<int>(scale_cols),
+                inv_global_scale_denom);
+        });
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return output;
+}
+
+TORCH_LIBRARY_IMPL(longlive_kernels, CUDA, m)
+{
+    m.impl("dequantize_kv_cache_fp4", &dequantize_kv_cache_fp4_cuda);
+}

utils/kernel/setup.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+from pathlib import Path
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+THIS_DIR = Path(__file__).resolve().parent
+
+setup(
+    name="longlive_kv_dequant_cuda",
+    ext_modules=[
+        CUDAExtension(
+            name="longlive_kv_dequant_cuda",
+            sources=[
+                str(THIS_DIR / "kv_dequant.cpp"),
+                str(THIS_DIR / "kv_dequant_cuda.cu"),
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-std=c++17"],
+                "nvcc": [
+                    "-O3",
+                    "-std=c++17",
+                    "--expt-relaxed-constexpr",
+                    # iter-37: need sm_100a (Blackwell arch-specific) for
+                    # cvt.rn.f16x2.e2m1x2 instruction. Plain sm_100 lacks it.
+                    "-gencode=arch=compute_100a,code=sm_100a",
+                ],
+            },
+        ),
+    ],
+    cmdclass={"build_ext": BuildExtension},
+)

utils/lightvae_5b_wrapper.py

@@ -0,0 +1,423 @@
+import logging
+import os
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from wan_5b.modules.vae2_2 import (
+    CausalConv3d,
+    Decoder3d,
+    Encoder3d,
+    count_conv3d,
+    patchify,
+    unpatchify,
+)
+
+
+def _extract_checkpoint_state_dict(raw):
+    state = raw
+    if isinstance(state, dict) and "state_dict" in state:
+        state = state["state_dict"]
+    if isinstance(state, dict) and "gen_model" in state:
+        state = state["gen_model"]
+    if isinstance(state, dict) and "generator" in state:
+        state = state["generator"]
+    if not isinstance(state, dict):
+        raise ValueError("Unsupported checkpoint format: expected a dict-like state_dict.")
+    return state
+
+
+def _map_lightvae_key_to_wanvae(key):
+    def _map_resnet_tail(tail):
+        if tail.startswith("norm1."):
+            return "residual.0." + tail[len("norm1."):]
+        if tail.startswith("conv1."):
+            return "residual.2." + tail[len("conv1."):]
+        if tail.startswith("norm2."):
+            return "residual.3." + tail[len("norm2."):]
+        if tail.startswith("conv2."):
+            return "residual.6." + tail[len("conv2."):]
+        if tail.startswith("conv_shortcut."):
+            return "shortcut." + tail[len("conv_shortcut."):]
+        return tail
+
+    if key.startswith("dynamic_feature_projection_heads."):
+        return None
+
+    if key.startswith("quant_conv."):
+        return key.replace("quant_conv.", "conv1.", 1)
+    if key.startswith("post_quant_conv."):
+        return key.replace("post_quant_conv.", "conv2.", 1)
+
+    if key.startswith("encoder.conv_in."):
+        return key.replace("encoder.conv_in.", "encoder.conv1.", 1)
+    if key.startswith("encoder.mid_block.resnets.0."):
+        tail = key[len("encoder.mid_block.resnets.0."):]
+        return "encoder.middle.0." + _map_resnet_tail(tail)
+    if key.startswith("encoder.mid_block.attentions.0."):
+        return key.replace("encoder.mid_block.attentions.0.", "encoder.middle.1.", 1)
+    if key.startswith("encoder.mid_block.resnets.1."):
+        tail = key[len("encoder.mid_block.resnets.1."):]
+        return "encoder.middle.2." + _map_resnet_tail(tail)
+    if key.startswith("encoder.norm_out."):
+        return key.replace("encoder.norm_out.", "encoder.head.0.", 1)
+    if key.startswith("encoder.conv_out."):
+        return key.replace("encoder.conv_out.", "encoder.head.2.", 1)
+
+    if key.startswith("encoder.down_blocks."):
+        parts = key.split(".")
+        if len(parts) >= 6 and parts[3] == "resnets":
+            tail = ".".join(parts[5:])
+            return f"encoder.downsamples.{parts[2]}.downsamples.{parts[4]}." + _map_resnet_tail(tail)
+        if len(parts) >= 7 and parts[3] == "downsampler" and parts[4] == "resample":
+            return f"encoder.downsamples.{parts[2]}.downsamples.2.resample.{parts[5]}." + ".".join(parts[6:])
+        if len(parts) >= 6 and parts[3] == "downsampler" and parts[4] == "time_conv":
+            return f"encoder.downsamples.{parts[2]}.downsamples.2.time_conv." + ".".join(parts[5:])
+
+    if key.startswith("decoder.conv_in."):
+        return key.replace("decoder.conv_in.", "decoder.conv1.", 1)
+    if key.startswith("decoder.mid_block.resnets.0."):
+        tail = key[len("decoder.mid_block.resnets.0."):]
+        return "decoder.middle.0." + _map_resnet_tail(tail)
+    if key.startswith("decoder.mid_block.attentions.0."):
+        return key.replace("decoder.mid_block.attentions.0.", "decoder.middle.1.", 1)
+    if key.startswith("decoder.mid_block.resnets.1."):
+        tail = key[len("decoder.mid_block.resnets.1."):]
+        return "decoder.middle.2." + _map_resnet_tail(tail)
+    if key.startswith("decoder.norm_out."):
+        return key.replace("decoder.norm_out.", "decoder.head.0.", 1)
+    if key.startswith("decoder.conv_out."):
+        return key.replace("decoder.conv_out.", "decoder.head.2.", 1)
+
+    if key.startswith("decoder.up_blocks."):
+        parts = key.split(".")
+        if len(parts) >= 6 and parts[3] == "resnets":
+            tail = ".".join(parts[5:])
+            return f"decoder.upsamples.{parts[2]}.upsamples.{parts[4]}." + _map_resnet_tail(tail)
+        if len(parts) >= 7 and parts[3] == "upsampler" and parts[4] == "resample":
+            return f"decoder.upsamples.{parts[2]}.upsamples.3.resample.{parts[5]}." + ".".join(parts[6:])
+        if len(parts) >= 6 and parts[3] == "upsampler" and parts[4] == "time_conv":
+            return f"decoder.upsamples.{parts[2]}.upsamples.3.time_conv." + ".".join(parts[5:])
+
+    return key
+
+
+def _normalize_vae_state_dict(raw_state):
+    state = _extract_checkpoint_state_dict(raw_state)
+    normalized = {}
+    for key, value in state.items():
+        mapped_key = _map_lightvae_key_to_wanvae(key)
+        if mapped_key is None:
+            continue
+        normalized[mapped_key] = value
+    return normalized
+
+
+def infer_lightvae_pruning_rate_from_ckpt(vae_path, full_decoder_conv1_out=1024):
+    if vae_path is None or not os.path.exists(vae_path):
+        return None
+    try:
+        raw_state = torch.load(vae_path, map_location="cpu")
+        state = _extract_checkpoint_state_dict(raw_state)
+    except Exception as exc:
+        logging.warning("Failed to load checkpoint for pruning-rate inference: %s", exc)
+        return None
+
+    weight = None
+    if isinstance(state, dict):
+        if "decoder.conv_in.weight" in state:
+            weight = state["decoder.conv_in.weight"]
+        elif "decoder.conv1.weight" in state:
+            weight = state["decoder.conv1.weight"]
+
+    if weight is None:
+        try:
+            normalized_state = _normalize_vae_state_dict(state)
+            weight = normalized_state.get("decoder.conv1.weight", None)
+        except Exception:
+            weight = None
+
+    if weight is None or not hasattr(weight, "shape") or len(weight.shape) < 1:
+        return None
+
+    student_out = int(weight.shape[0])
+    if full_decoder_conv1_out <= 0:
+        return None
+
+    pruning_rate = 1.0 - (float(student_out) / float(full_decoder_conv1_out))
+    pruning_rate = max(0.0, min(0.99, pruning_rate))
+    return round(pruning_rate, 6)
+
+
+def convert_to_channels_last_3d(module):
+    for child in module.children():
+        if isinstance(child, nn.Conv3d):
+            child.weight.data = child.weight.data.to(memory_format=torch.channels_last_3d)
+        else:
+            convert_to_channels_last_3d(child)
+
+
+class PrunableWanVAE(nn.Module):
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=48,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0,
+        pruning_rate=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        dim = max(1, int(round(dim * (1.0 - pruning_rate))))
+        dec_dim = max(1, int(round(dec_dim * (1.0 - pruning_rate))))
+
+        self.encoder = Encoder3d(
+            dim,
+            z_dim * 2,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_downsample,
+            dropout,
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim,
+            z_dim,
+            dim_mult,
+            num_res_blocks,
+            attn_scales,
+            self.temperal_upsample,
+            dropout,
+        )
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        x = patchify(x, patch_size=2)
+        total_steps = 1 + (x.shape[2] - 1) // 4
+        for step in range(total_steps):
+            self._enc_conv_idx = [0]
+            if step == 0:
+                out = self.encoder(
+                    x[:, :, :1, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+            else:
+                out_chunk = self.encoder(
+                    x[:, :, 1 + 4 * (step - 1):1 + 4 * step, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_chunk], 2)
+        mu, _ = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1
+            )
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1
+            )
+        else:
+            z = z / scale[1] + scale[0]
+        total_steps = z.shape[2]
+        x = self.conv2(z)
+        for step in range(total_steps):
+            self._conv_idx = [0]
+            if step == 0:
+                out = self.decoder(
+                    x[:, :, step:step + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            else:
+                out_chunk = self.decoder(
+                    x[:, :, step:step + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_chunk], 2)
+        out = unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def cached_decode(self, z, scale):
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1
+            )
+        else:
+            z = z / scale[1] + scale[0]
+        total_steps = z.shape[2]
+        x = self.conv2(z)
+        is_first = self._feat_map[0] is None
+        for step in range(total_steps):
+            self._conv_idx = [0]
+            if step == 0 and is_first:
+                out = self.decoder(
+                    x[:, :, step:step + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                    first_chunk=True,
+                )
+            elif step == 0:
+                out = self.decoder(
+                    x[:, :, step:step + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+            else:
+                out_chunk = self.decoder(
+                    x[:, :, step:step + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
+                out = torch.cat([out, out_chunk], 2)
+        return unpatchify(out, patch_size=2)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _load_lightvae_model(pretrained_path=None, z_dim=48, dim=160, device="cpu", **kwargs):
+    cfg = dict(
+        dim=dim,
+        z_dim=z_dim,
+        dim_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        attn_scales=[],
+        temperal_downsample=[False, True, True],
+        dropout=0.0,
+    )
+    cfg.update(**kwargs)
+
+    with torch.device("meta"):
+        model = PrunableWanVAE(**cfg)
+
+    if pretrained_path is None or not os.path.exists(pretrained_path):
+        raise FileNotFoundError(f"VAE checkpoint not found at {pretrained_path}")
+
+    logging.info("loading %s", pretrained_path)
+    raw_state = torch.load(pretrained_path, map_location="cpu")
+    state_dict = _normalize_vae_state_dict(raw_state)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False, assign=True)
+    logging.info(
+        "LightVAE checkpoint loaded with strict=False (missing=%d, unexpected=%d)",
+        len(missing),
+        len(unexpected),
+    )
+
+    convert_to_channels_last_3d(model)
+    return model
+
+
+class LightVAE5BWrapper(nn.Module):
+    def __init__(
+        self,
+        vae_path: str,
+        pruning_rate: Optional[float] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if pruning_rate is None:
+            pruning_rate = infer_lightvae_pruning_rate_from_ckpt(vae_path)
+            if pruning_rate is None:
+                pruning_rate = 0.75
+                logging.warning(
+                    "Unable to infer LightVAE pruning rate from checkpoint; fallback to 0.75."
+                )
+
+        mean = [
+            -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557,
+            -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825,
+            -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502,
+            -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230,
+            -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748,
+            0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667,
+        ]
+        std = [
+            0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,
+            0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,
+            0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,
+            0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,
+            0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,
+            0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744,
+        ]
+        self.mean = torch.tensor(mean, dtype=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+        self.vae_path = os.path.abspath(vae_path)
+        self.pruning_rate = pruning_rate
+        self.device = torch.device(device)
+        self.dtype = dtype
+
+        self.model = _load_lightvae_model(
+            pretrained_path=self.vae_path,
+            pruning_rate=self.pruning_rate,
+        ).eval().requires_grad_(False)
+        self.to(device=self.device, dtype=self.dtype)
+
+    def to(self, device=None, dtype=None):
+        device = self.device if device is None else torch.device(device)
+        dtype = self.dtype if dtype is None else dtype
+        self.model.to(device=device, dtype=dtype)
+        self.mean = self.mean.to(device=device, dtype=dtype)
+        self.std = self.std.to(device=device, dtype=dtype)
+        self.device = device
+        self.dtype = dtype
+        return self
+
+    def eval(self):
+        super().eval()
+        self.model.eval()
+        return self
+
+    @torch.no_grad()
+    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
+
+        scale = [self.mean, 1.0 / self.std]
+        decode_fn = self.model.cached_decode if use_cache else self.model.decode
+
+        output = []
+        for item in zs:
+            output.append(
+                decode_fn(item.unsqueeze(0).to(device=self.device, dtype=self.dtype), scale)
+                .float()
+                .clamp_(-1, 1)
+                .squeeze(0)
+            )
+        output = torch.stack(output, dim=0)
+        return output.permute(0, 2, 1, 3, 4)

utils/lora_utils.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+import peft
+from peft import get_peft_model_state_dict
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
+    StateDictType, FullStateDictConfig
+)
+
+
+def configure_lora_for_model(transformer, model_name, lora_config, is_main_process=True, all_causal=False):
+    """Configure LoRA for a WanDiffusionWrapper model
+    
+    Args:
+        transformer: The transformer model to apply LoRA to
+        model_name: 'generator' or 'fake_score'
+        lora_config: LoRA configuration
+        is_main_process: Whether this is the main process (for logging)
+        all_causal: Whether all models use causal attention blocks
+    
+    Returns:
+        lora_model: The LoRA-wrapped model
+    """
+    target_linear_modules = set()
+    
+    if model_name == 'generator':
+        adapter_target_modules = ['CausalWanAttentionBlock']
+    elif model_name == 'fake_score':
+        adapter_target_modules = ['CausalWanAttentionBlock'] if all_causal else ['WanAttentionBlock']
+    else:
+        raise ValueError(f"Invalid model name: {model_name}")
+    
+    for name, module in transformer.named_modules():
+        if module.__class__.__name__ in adapter_target_modules:
+            for full_submodule_name, submodule in module.named_modules(prefix=name):
+                if isinstance(submodule, torch.nn.Linear):
+                    target_linear_modules.add(full_submodule_name)
+    
+    target_linear_modules = list(target_linear_modules)
+    
+    if is_main_process:
+        print(f"LoRA target modules for {model_name}: {len(target_linear_modules)} Linear layers")
+        if getattr(lora_config, 'verbose', False):
+            for module_name in sorted(target_linear_modules):
+                print(f"  - {module_name}")
+    
+    # Create LoRA config
+    adapter_type = lora_config.get('type', 'lora')
+    if adapter_type == 'lora':
+        peft_config = peft.LoraConfig(
+            r=lora_config.get('rank', 16),
+            lora_alpha=lora_config.get('alpha', None) or lora_config.get('rank', 16),
+            lora_dropout=lora_config.get('dropout', 0.0),
+            target_modules=target_linear_modules,
+        )
+    else:
+        raise NotImplementedError(f'Adapter type {adapter_type} is not implemented')
+    
+    # Apply LoRA to the transformer
+    lora_model = peft.get_peft_model(transformer, peft_config)
+
+    if is_main_process:
+        print('peft_config', peft_config)
+        lora_model.print_trainable_parameters()
+    
+    return lora_model
+
+
+def gather_lora_state_dict(lora_model):
+    with FSDP.state_dict_type(
+        lora_model,                  
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(rank0_only=True, offload_to_cpu=True)
+    ):
+        full = lora_model.state_dict()
+    return get_peft_model_state_dict(lora_model, state_dict=full)
+
+
+def load_lora_checkpoint(lora_model, lora_state_dict, model_name, is_main_process=True):
+    """Load LoRA weights from state dict
+    
+    Args:
+        lora_model: The LoRA-wrapped model
+        lora_state_dict: LoRA state dict to load
+        model_name: 'generator' or 'critic'
+        is_main_process: Whether this is the main process (for logging)
+    """
+    if is_main_process:
+        print(f"Loading LoRA {model_name} weights: {len(lora_state_dict)} keys in checkpoint")
+    
+    peft.set_peft_model_state_dict(lora_model, lora_state_dict)
+    
+    if is_main_process:
+        print(f"LoRA {model_name} weights loaded successfully") 

utils/loss.py

@@ -0,0 +1,98 @@
+from abc import ABC, abstractmethod
+import torch
+
+
+class DenoisingLoss(ABC):
+    @abstractmethod
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        gradient_mask: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Base class for denoising loss.
+        Input:
+            - x: the clean data with shape [B, F, C, H, W]
+            - x_pred: the predicted clean data with shape [B, F, C, H, W]
+            - noise: the noise with shape [B, F, C, H, W]
+            - noise_pred: the predicted noise with shape [B, F, C, H, W]
+            - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]
+            - timestep: the current timestep with shape [B, F]
+        """
+        pass
+
+
+class X0PredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        gradient_mask: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        err = (x - x_pred) ** 2
+        if gradient_mask is not None:
+            return err[gradient_mask].mean()
+        return err.mean()
+
+
+class VPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        gradient_mask: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))
+        err = weights * (x - x_pred) ** 2
+        if gradient_mask is not None:
+            return err[gradient_mask].mean()
+        return err.mean()
+
+
+class NoisePredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        gradient_mask: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        err = (noise - noise_pred) ** 2
+        if gradient_mask is not None:
+            return err[gradient_mask].mean()
+        return err.mean()
+
+
+class FlowPredLoss(DenoisingLoss):
+    def __call__(
+        self, x: torch.Tensor, x_pred: torch.Tensor,
+        noise: torch.Tensor, noise_pred: torch.Tensor,
+        alphas_cumprod: torch.Tensor,
+        timestep: torch.Tensor,
+        gradient_mask: torch.Tensor = None,
+        **kwargs
+    ) -> torch.Tensor:
+        err = (kwargs["flow_pred"] - (noise - x)) ** 2
+        if gradient_mask is not None:
+            return err[gradient_mask].mean()
+        return err.mean()
+
+
+NAME_TO_CLASS = {
+    "x0": X0PredLoss,
+    "v": VPredLoss,
+    "noise": NoisePredLoss,
+    "flow": FlowPredLoss
+}
+
+
+def get_denoising_loss(loss_type: str) -> DenoisingLoss:
+    return NAME_TO_CLASS[loss_type]

utils/memory.py

@@ -0,0 +1,146 @@
+# Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils
+# Apache-2.0 License
+# By lllyasviel
+
+import torch
+
+
+cpu = torch.device('cpu')
+gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+gpu_complete_modules = []
+
+
+class DynamicSwapInstaller:
+    @staticmethod
+    def _install_module(module: torch.nn.Module, **kwargs):
+        original_class = module.__class__
+        module.__dict__['forge_backup_original_class'] = original_class
+
+        def hacked_get_attr(self, name: str):
+            if '_parameters' in self.__dict__:
+                _parameters = self.__dict__['_parameters']
+                if name in _parameters:
+                    p = _parameters[name]
+                    if p is None:
+                        return None
+                    if p.__class__ == torch.nn.Parameter:
+                        return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad)
+                    else:
+                        return p.to(**kwargs)
+            if '_buffers' in self.__dict__:
+                _buffers = self.__dict__['_buffers']
+                if name in _buffers:
+                    return _buffers[name].to(**kwargs)
+            return super(original_class, self).__getattr__(name)
+
+        module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
+            '__getattr__': hacked_get_attr,
+        })
+
+        return
+
+    @staticmethod
+    def _uninstall_module(module: torch.nn.Module):
+        if 'forge_backup_original_class' in module.__dict__:
+            module.__class__ = module.__dict__.pop('forge_backup_original_class')
+        return
+
+    @staticmethod
+    def install_model(model: torch.nn.Module, **kwargs):
+        for m in model.modules():
+            DynamicSwapInstaller._install_module(m, **kwargs)
+        return
+
+    @staticmethod
+    def uninstall_model(model: torch.nn.Module):
+        for m in model.modules():
+            DynamicSwapInstaller._uninstall_module(m)
+        return
+
+
+def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.device):
+    if hasattr(model, 'scale_shift_table'):
+        model.scale_shift_table.data = model.scale_shift_table.data.to(target_device)
+        return
+
+    for k, p in model.named_modules():
+        if hasattr(p, 'weight'):
+            p.to(target_device)
+            return
+
+
+def get_cuda_free_memory_gb(device=None):
+    if device is None:
+        device = gpu
+
+    memory_stats = torch.cuda.memory_stats(device)
+    bytes_active = memory_stats['active_bytes.all.current']
+    bytes_reserved = memory_stats['reserved_bytes.all.current']
+    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+    bytes_inactive_reserved = bytes_reserved - bytes_active
+    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+    return bytes_total_available / (1024 ** 3)
+
+
+
+def log_gpu_memory(stage: str, device=None, rank=0):
+    """Log GPU memory usage at a given training stage."""
+    free_gb = get_cuda_free_memory_gb(device)
+    total_gb = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3)
+    used_gb = total_gb - free_gb
+    print(f"[rank {rank}] [GPU Memory][{stage}] Used: {used_gb:.2f} GB | Free: {free_gb:.2f} GB | Total: {total_gb:.2f} GB")
+
+
+
+
+def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=target_device)
+
+    model.to(device=target_device)
+    torch.cuda.empty_cache()
+    return
+
+
+def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
+    print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
+
+    for m in model.modules():
+        if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
+            torch.cuda.empty_cache()
+            return
+
+        if hasattr(m, 'weight'):
+            m.to(device=cpu)
+
+    model.to(device=cpu)
+    torch.cuda.empty_cache()
+    return
+
+
+def unload_complete_models(*args):
+    for m in gpu_complete_modules + list(args):
+        m.to(device=cpu)
+        print(f'Unloaded {m.__class__.__name__} as complete.')
+
+    gpu_complete_modules.clear()
+    torch.cuda.empty_cache()
+    return
+
+
+def load_model_as_complete(model, target_device, unload=True):
+    if unload:
+        unload_complete_models()
+
+    model.to(device=target_device)
+    print(f'Loaded {model.__class__.__name__} to {target_device} as complete.')
+
+    gpu_complete_modules.append(model)
+    return

utils/misc.py

@@ -0,0 +1,39 @@
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed: int, deterministic: bool = False):
+    """
+    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.
+
+    Args:
+        seed (`int`):
+            The seed to set.
+        deterministic (`bool`, *optional*, defaults to `False`):
+            Whether to use deterministic algorithms where available. Can slow down training.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    if deterministic:
+        torch.use_deterministic_algorithms(True)
+
+
+def merge_dict_list(dict_list):
+    if len(dict_list) == 1:
+        return dict_list[0]
+
+    merged_dict = {}
+    for k, v in dict_list[0].items():
+        if isinstance(v, torch.Tensor):
+            if v.ndim == 0:
+                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)
+            else:
+                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)
+        else:
+            # for non-tensor values, we just copy the value from the first item
+            merged_dict[k] = v
+    return merged_dict

utils/nvfp4_checkpoint.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+from torch import nn
+
+
+NVFP4_CHECKPOINT_FORMAT = "longlive_generator_nvfp4"
+TE_NVFP4_CHECKPOINT_FORMAT = "longlive_generator_te_nvfp4"
+NVFP4_CHECKPOINT_VERSION = 1
+
+
+def is_nvfp4_state_dict(state_dict: object) -> bool:
+    """Return True when a state dict contains materialized FourOverSix NVFP4 buffers."""
+    if not isinstance(state_dict, Mapping):
+        return False
+    return any(str(key).endswith("quantized_weight_values") for key in state_dict)
+
+
+def is_te_nvfp4_checkpoint(checkpoint: object) -> bool:
+    """Return True for checkpoints saved with TransformerEngine module state."""
+    return (
+        isinstance(checkpoint, Mapping)
+        and checkpoint.get("checkpoint_format") == TE_NVFP4_CHECKPOINT_FORMAT
+    )
+
+
+def unwrap_generator_state_dict(checkpoint: object, use_ema: bool = False) -> object:
+    """Extract the generator state dict from common LongLive checkpoint layouts."""
+    if not isinstance(checkpoint, Mapping):
+        return checkpoint
+    if "generator" in checkpoint or "generator_ema" in checkpoint:
+        ema_key = "generator_ema" if use_ema and "generator_ema" in checkpoint else "generator"
+        return checkpoint[ema_key]
+    if "model" in checkpoint:
+        return checkpoint["model"]
+    return checkpoint
+
+
+def clean_fsdp_state_dict_keys(state_dict: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Remove FSDP wrapper prefixes used by some EMA checkpoints."""
+    return {str(key).replace("_fsdp_wrapped_module.", ""): value for key, value in state_dict.items()}
+
+
+def build_model_quantization_config(config, keep_master_weights: bool = False):
+    from utils.quant import ModelQuantizationConfig
+
+    quant_cfg = ModelQuantizationConfig(
+        scale_rule=getattr(config, "model_quant_scale_rule", "static_6"),
+        quantize_backend=getattr(config, "model_quant_backend", None),
+        activation_scale_rule=getattr(
+            config,
+            "model_quant_activation_scale_rule",
+            getattr(config, "model_quant_scale_rule", "static_6"),
+        ),
+        weight_scale_rule=getattr(config, "model_quant_weight_scale_rule", None),
+        gradient_scale_rule=getattr(config, "model_quant_gradient_scale_rule", None),
+    )
+    quant_cfg.keep_master_weights = keep_master_weights
+    return quant_cfg
+
+
+def _maybe_to_dict(value):
+    if value is None:
+        return None
+    try:
+        from omegaconf import OmegaConf
+
+        if OmegaConf.is_config(value):
+            value = OmegaConf.to_container(value, resolve=True)
+    except ImportError:
+        pass
+    return dict(value)
+
+
+def quantize_model_for_fouroversix_nvfp4(model: nn.Module, config, *, keep_master_weights: bool = False, verbose: bool = True):
+    """Replace eligible modules with FourOverSix NVFP4 modules using the runtime config."""
+    from utils.quant import quantize_model_with_filter
+
+    return quantize_model_with_filter(
+        model,
+        quant_config=build_model_quantization_config(config, keep_master_weights=keep_master_weights),
+        filtered_modules=getattr(config, "model_quant_filtered_modules", None),
+        use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True),
+        cast_model_to_bf16=True,
+        materialize_for_inference=False,
+        use_transformer_engine=False,
+        verbose=verbose,
+    )
+
+
+def quantize_model_for_transformer_engine_nvfp4(
+    model: nn.Module,
+    config,
+    *,
+    keep_master_weights: bool = False,
+    verbose: bool = True,
+):
+    """Replace eligible modules with TransformerEngine NVFP4 wrappers."""
+    from utils.quant import quantize_model_with_filter
+
+    use_transformer_engine = True
+    te_inference_only = bool(getattr(config, "model_quant_te_inference_only", use_transformer_engine))
+    te_low_precision_weights = bool(getattr(config, "model_quant_te_low_precision_weights", te_inference_only))
+    te_fallback_to_fouroversix = bool(getattr(config, "model_quant_te_fallback_to_fouroversix", False))
+
+    return quantize_model_with_filter(
+        model,
+        quant_config=build_model_quantization_config(config, keep_master_weights=keep_master_weights),
+        filtered_modules=getattr(config, "model_quant_filtered_modules", None),
+        use_default_filtered_modules=getattr(config, "model_quant_use_default_filtered_modules", True),
+        cast_model_to_bf16=True,
+        materialize_for_inference=False,
+        use_transformer_engine=True,
+        te_inference_only=te_inference_only,
+        te_low_precision_weights=te_low_precision_weights,
+        te_recipe_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_recipe_kwargs", None)),
+        te_module_kwargs=_maybe_to_dict(getattr(config, "model_quant_te_module_kwargs", None)),
+        te_fallback_to_fouroversix=te_fallback_to_fouroversix,
+        verbose=verbose,
+    )
+
+
+def drop_fouroversix_master_weights(model: nn.Module) -> list[str]:
+    """Drop high-precision master weights after loading materialized NVFP4 buffers."""
+    materialized_modules = []
+    for module_name, module in model.named_modules():
+        if not hasattr(module, "parameters_to_quantize"):
+            continue
+
+        parameters_to_quantize = getattr(module, "parameters_to_quantize", ())
+        if callable(parameters_to_quantize):
+            parameters_to_quantize = parameters_to_quantize()
+        if not parameters_to_quantize:
+            continue
+
+        dropped_any = False
+        for parameter_name in parameters_to_quantize:
+            if isinstance(getattr(module, parameter_name, None), nn.Parameter):
+                module.register_parameter(parameter_name, None)
+                dropped_any = True
+            elif hasattr(module, parameter_name):
+                setattr(module, parameter_name, None)
+                dropped_any = True
+
+        if not dropped_any:
+            continue
+        for cache_name in ("_quantized_weight", "_quantized_weight_transposed", "_quantized_weights"):
+            if hasattr(module, cache_name):
+                delattr(module, cache_name)
+        if hasattr(module, "config") and hasattr(module.config, "keep_master_weights"):
+            module.config.keep_master_weights = False
+        materialized_modules.append(module_name)
+    return materialized_modules
+
+
+def cpu_state_dict(module: nn.Module) -> dict[str, torch.Tensor]:
+    """Return a detached CPU state dict suitable for torch.save."""
+    return {key: value.detach().cpu() for key, value in module.state_dict().items()}

utils/nvfp4_kernel.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""NVFP4 Fake Quantization Triton Implementation.
+
+This module provides high-performance GPU implementations of NVFP4 fake quantization
+operations using Triton kernels.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+__all__ = ["fp4_dequantize", "static_blockwise_fp4_fake_quant"]
+
+
+_TORCH_TO_TL_DTYPE = {
+    torch.float32: tl.float32,
+    torch.float: tl.float32,
+    torch.float16: tl.float16,
+    torch.half: tl.float16,
+    torch.bfloat16: tl.bfloat16,
+}
+
+
+def _torch_dtype_to_tl(dtype: torch.dtype):
+    if dtype not in _TORCH_TO_TL_DTYPE:
+        raise ValueError(f"Unsupported dtype for fp4 fake quantization: {dtype}")
+    return _TORCH_TO_TL_DTYPE[dtype]
+
+
+@triton.jit
+def fp4_dequantize_kernel(
+    packed_ptr,
+    scale_ptr,
+    global_scale_ptr,
+    output_ptr,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+    TILE_SIZE: tl.constexpr,
+):
+    """Dequantizes FP4 packed data using per-block scaling factors.
+
+    Args:
+        packed_ptr (tl.pointer): Pointer to packed uint8 tensor (M x N//2)
+        scale_ptr (tl.pointer): Pointer to per-block scale tensor (M x N//BLOCK_SIZE)
+        output_ptr (tl.pointer): Pointer to output tensor (M x N)
+        global_scale_ptr (tl.pointer): Pointer to global scale tensor
+        N (int): Number of columns in unpacked tensor
+        BLOCK_SIZE (tl.constexpr): Size of each FP4 quantization block
+        TILE_SIZE (tl.constexpr): Size of the processing tile (in packed elements)
+    """
+    # Get program ID for processing packed elements
+    pid = tl.program_id(0)
+
+    # Calculate packed element offsets (each packed element contains 2 FP4 values)
+    packed_start = pid * TILE_SIZE
+    packed_offs = packed_start + tl.arange(0, TILE_SIZE)
+
+    # Calculate 2D coordinates for packed data
+    packed_row_idx = packed_offs // (N // 2)
+    packed_col_idx = packed_offs % (N // 2)
+
+    # Create mask for packed data bounds checking
+    packed_mask = packed_col_idx < (N // 2)
+
+    # Load global scale
+    global_scale = tl.load(global_scale_ptr)
+
+    # Load packed data
+    packed_data = tl.load(packed_ptr + packed_offs, mask=packed_mask, other=0)
+
+    # Unpack packed FP4 values (uint8) to float16x2
+    x_f16x2_packed = tl.inline_asm_elementwise(
+        asm="""
+        {
+            .reg .b8 byte0, byte1, byte2, byte3;
+            mov.b32 {byte0, byte1, byte2, byte3}, $4;
+            cvt.rn.f16x2.e2m1x2 $0, byte0;
+            cvt.rn.f16x2.e2m1x2 $1, byte1;
+            cvt.rn.f16x2.e2m1x2 $2, byte2;
+            cvt.rn.f16x2.e2m1x2 $3, byte3;
+        }
+        """,
+        constraints="=r,=r,=r,=r,r",
+        args=[packed_data],
+        dtype=tl.uint32,
+        is_pure=True,
+        pack=4,
+    )
+    val_low = (
+        (x_f16x2_packed & 0xFFFF).cast(tl.uint16).cast(tl.float16, bitcast=True).cast(tl.float32)
+    )
+    val_high = (
+        (x_f16x2_packed >> 16).cast(tl.uint16).cast(tl.float16, bitcast=True).cast(tl.float32)
+    )
+
+    # Calculate output positions for both values
+    out_col_low = packed_col_idx * 2
+    out_col_high = packed_col_idx * 2 + 1
+    out_offs_low = packed_row_idx * N + out_col_low
+    out_offs_high = packed_row_idx * N + out_col_high
+
+    # Calculate block indices for scaling
+    block_col_low = out_col_low // BLOCK_SIZE
+    block_col_high = out_col_high // BLOCK_SIZE
+    scale_offs_low = packed_row_idx * (N // BLOCK_SIZE) + block_col_low
+    scale_offs_high = packed_row_idx * (N // BLOCK_SIZE) + block_col_high
+
+    # Load scaling factors
+    scale_low = tl.load(scale_ptr + scale_offs_low, mask=packed_mask & (out_col_low < N), other=1.0)
+    scale_high = tl.load(
+        scale_ptr + scale_offs_high, mask=packed_mask & (out_col_high < N), other=1.0
+    )
+
+    # Apply scaling
+    result_low = val_low * scale_low.to(tl.float32) * global_scale
+    result_high = val_high * scale_high.to(tl.float32) * global_scale
+
+    # Store results
+    out_mask_low = packed_mask & (out_col_low < N)
+    out_mask_high = packed_mask & (out_col_high < N)
+
+    tl.store(output_ptr + out_offs_low, result_low, mask=out_mask_low)
+    tl.store(output_ptr + out_offs_high, result_high, mask=out_mask_high)
+
+
+def fp4_dequantize(
+    packed_tensor: torch.Tensor,
+    scale_tensor: torch.Tensor,
+    global_scale: torch.Tensor,
+    block_size: int = 16,
+    tile_size: int = 128,
+    dtype: torch.dtype = torch.get_default_dtype(),
+) -> torch.Tensor:
+    """Dequantizes FP4 packed tensor using per-block scaling factors.
+
+    Args:
+        packed_tensor (torch.Tensor): Packed uint8 tensor of shape (M, N//2)
+        scale_tensor (torch.Tensor): Per-block scale tensor of shape (M, N//block_size)
+        global_scale (torch.Tensor): Global scaling factor tensor
+        block_size (int): Size of FP4 quantization blocks
+        tile_size (int): Size of processing tiles
+
+    Returns:
+        torch.Tensor: Dequantized tensor of shape (M, N)
+    """
+    packed_N = packed_tensor.shape[-1]
+    N = packed_N * 2
+    # Create output tensor with proper shape handling
+    output_shape = list(packed_tensor.shape)
+    output_shape[-1] = N
+    output = torch.empty(output_shape, dtype=dtype, device=packed_tensor.device)
+
+    # Calculate total number of elements and grid size
+    grid = lambda meta: (triton.cdiv(packed_tensor.numel(), meta["TILE_SIZE"]),)
+
+    fp4_dequantize_kernel[grid](
+        packed_tensor,
+        scale_tensor,
+        global_scale,
+        output,
+        N,
+        BLOCK_SIZE=block_size,
+        TILE_SIZE=tile_size,
+    )
+
+    return output
+
+
+@triton.jit
+def static_blockwise_fp4_fake_quant_kernel(
+    x_ptr,  # [NUM_FP4_BLOCKS * BLOCK_SIZE]
+    y_ptr,  # [NUM_FP4_BLOCKS * BLOCK_SIZE]
+    scale_ptr,  # [NUM_FP4_BLOCKS]
+    NUM_FP4_BLOCKS,
+    BLOCK_SIZE: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    if pid >= NUM_FP4_BLOCKS:
+        return
+
+    block_offset = pid * BLOCK_SIZE
+    idx = block_offset + tl.arange(0, BLOCK_SIZE)
+
+    scale = tl.load(scale_ptr + pid).to(tl.float32)
+
+    x = tl.load(x_ptr + idx).to(tl.float32)
+
+    x_abs = tl.abs(x)
+    # If scale is 0, inf, or nan, use 1.0 (matching CUDA kernel behavior)
+    # Note: (x != x) checks if x is NaN per IEEE 754
+    scale_safe = tl.where(
+        (scale == 0) | (scale != scale) | (tl.abs(scale) == float("inf")),  # noqa: PLR0124
+        1.0,
+        scale,
+    )
+    abs_scaled = x_abs / scale_safe
+
+    # FP4 values: 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0
+    q_val = tl.where(
+        abs_scaled <= 0.25,
+        0.0,
+        tl.where(
+            abs_scaled < 0.75,
+            0.5,
+            tl.where(
+                abs_scaled <= 1.25,
+                1.0,
+                tl.where(
+                    abs_scaled < 1.75,
+                    1.5,
+                    tl.where(
+                        abs_scaled <= 2.5,
+                        2.0,
+                        tl.where(
+                            abs_scaled < 3.5,
+                            3.0,
+                            tl.where(abs_scaled <= 5.0, 4.0, 6.0),
+                        ),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+    x_rescaled = q_val * scale_safe
+    x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled)
+
+    tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE))
+
+
+def static_blockwise_fp4_fake_quant(
+    x: torch.Tensor,
+    amax: torch.Tensor,
+    global_amax: torch.Tensor | None = None,
+    quantize_block_scales: bool = True,
+    out_dtype: torch.dtype | None = None,
+):
+    """Static blockwise FP4 fake quantization using Triton kernel.
+
+    Args:
+        x: [NUM_FP4_BLOCKS, BLOCK_SIZE] on CUDA.
+        amax: [NUM_FP4_BLOCKS] or [NUM_FP4_BLOCKS, 1] per-block amax values.
+        global_amax: FP32 scalar global amax. If provided, used to compute scale_fp8_quant_amax.
+        quantize_block_scales: If True, quantize block scales to FP8.
+        out_dtype: Output dtype. Defaults to x.dtype if None.
+    """
+    assert x.ndim == 2
+    NUM_FP4_BLOCKS, BLOCK_SIZE = x.shape
+
+    if out_dtype is None:
+        out_dtype = x.dtype
+
+    amax = amax.float()  # Requires to be in float32
+    scale = amax / 6.0  # FP4 max representable value is 6.0
+
+    if quantize_block_scales:
+        from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl
+        from modelopt.torch.quantization.utils import reduce_amax
+
+        if global_amax is None:
+            global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True)
+
+        global_amax = global_amax.float()
+        scale_fp8_quant_amax = global_amax / 6.0
+        scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax)
+
+    x_flat = x.contiguous().view(-1)
+    y_flat = torch.empty_like(x_flat, dtype=out_dtype)
+    scale_flat = scale.view(NUM_FP4_BLOCKS).contiguous()
+
+    tl_out_dtype = _torch_dtype_to_tl(out_dtype)
+
+    grid = (NUM_FP4_BLOCKS,)
+
+    with torch.cuda.device(x.device):
+        static_blockwise_fp4_fake_quant_kernel[grid](
+            x_flat,
+            y_flat,
+            scale_flat,
+            NUM_FP4_BLOCKS,
+            BLOCK_SIZE,
+            OUT_DTYPE=tl_out_dtype,
+        )
+
+    return y_flat.view_as(x)

utils/position_embedding_utils.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Minimal temporal RoPE helpers used by multi-shot generation."""
+
+import torch
+
+
+def select_temporal_offset_for_sample(
+    temporal_offset,
+    sample_idx: int,
+    f: int,
+    start_frame: int = 0,
+):
+    """Select the offset slice that applies to one sample.
+
+    ``temporal_offset`` accepts a scalar, ``[B]`` per-sample constants,
+    ``[F]`` shared per-frame offsets, or ``[B, F]`` per-sample per-frame
+    offsets.  The returned value is still interpreted by
+    ``compute_temporal_freqs`` so full-length and local slices both work.
+    """
+    if temporal_offset is None:
+        return 0.0
+    if torch.is_tensor(temporal_offset):
+        if temporal_offset.ndim == 0:
+            return temporal_offset
+        if temporal_offset.ndim == 1:
+            # Usually this is a shared [F] vector. If it is too short to cover
+            # the requested frame range, treat it as [B] constants.
+            if temporal_offset.numel() == f or temporal_offset.numel() >= start_frame + f:
+                return temporal_offset
+            return temporal_offset[sample_idx]
+        if temporal_offset.ndim == 2:
+            return temporal_offset[sample_idx]
+        raise ValueError(
+            "temporal_offset tensor must be scalar, [B], [F], or [B, F], "
+            f"got shape={tuple(temporal_offset.shape)}"
+        )
+    if isinstance(temporal_offset, (list, tuple)):
+        if not temporal_offset:
+            return 0.0
+        if isinstance(temporal_offset[0], (list, tuple)):
+            return torch.as_tensor(temporal_offset[sample_idx])
+        if len(temporal_offset) == f or len(temporal_offset) >= start_frame + f:
+            return torch.as_tensor(temporal_offset)
+        return temporal_offset[sample_idx]
+    return temporal_offset
+
+
+def compute_temporal_freqs(
+    freqs_t: torch.Tensor,
+    f: int,
+    start_frame: int,
+    t_scale: float,
+    device: torch.device,
+    method: str = "linear",
+    original_seq_len: int | None = None,
+    temporal_offset: float = 0.0,
+) -> torch.Tensor:
+    """Compute linear temporal RoPE freqs with an optional multi-shot offset."""
+    if method != "linear":
+        raise ValueError(f"Only linear temporal RoPE is supported in this release, got {method}.")
+    if original_seq_len is not None:
+        raise ValueError("original_seq_len is not used by the release linear RoPE path.")
+    if temporal_offset is None:
+        temporal_offset = 0.0
+    if (
+        t_scale == 1.0
+        and not torch.is_tensor(temporal_offset)
+        and float(temporal_offset) == 0.0
+    ):
+        return freqs_t[start_frame:start_frame + f]
+
+    base_angles = torch.angle(freqs_t[1]).to(torch.float64)
+    positions = torch.arange(f, device=device, dtype=torch.float64) + start_frame
+    if torch.is_tensor(temporal_offset):
+        offset = temporal_offset.to(device=device, dtype=torch.float64)
+        if offset.ndim == 0:
+            positions = positions + offset
+        elif offset.ndim == 1:
+            if offset.numel() == f:
+                positions = positions + offset
+            elif offset.numel() >= start_frame + f:
+                positions = positions + offset[start_frame:start_frame + f]
+            else:
+                raise ValueError(
+                    "temporal_offset length is too short for requested RoPE "
+                    f"range: len={offset.numel()}, start={start_frame}, f={f}"
+                )
+        else:
+            raise ValueError(
+                "compute_temporal_freqs expects a scalar or 1D temporal_offset "
+                f"after sample selection, got shape={tuple(offset.shape)}"
+            )
+    else:
+        positions = positions + float(temporal_offset)
+    positions = positions * t_scale
+    angles = positions.unsqueeze(-1) * base_angles.unsqueeze(0)
+    return torch.polar(torch.ones_like(angles), angles)

utils/quant.py

@@ -0,0 +1,818 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+import importlib
+import inspect
+from contextlib import nullcontext
+from copy import deepcopy
+from dataclasses import dataclass
+import re
+from typing import Any
+import warnings
+
+import torch
+import torch.nn as nn
+from fouroversix import (
+    DataType,
+    ModelQuantizationConfig,
+    QuantizationConfig,
+    QuantizedTensor,
+    RoundStyle,
+    ScaleRule,
+    quantize_model,
+    quantize_to_fp4,
+)
+from fouroversix.quantize.quantized_tensor import from_blocked
+
+from utils.nvfp4_kernel import fp4_dequantize
+
+_FUSED_KV_DEQUANT_DISABLED = False
+_FUSED_KV_DEQUANT_WARNED = False
+
+QUANTIZATION_TYPE = {
+    "weight": "weight",
+    "activation": "activation",
+    "kv": "kv",
+}
+
+DEFAULT_GENERATOR_FILTERED_MODULES = [
+    "text_embedding.0",
+    "text_embedding.2",
+    "patch_embedding",
+    "time_projection.1",
+    "time_embedding.0",
+    "time_embedding.2",
+    "head.head",
+    "head.modulation",
+    "re:.*norm_k$",
+    "re:.*norm_q$",
+    "re:.*norm1$",
+    "re:.*norm2$",
+    "re:.*norm3$"
+]
+
+DEFAULT_REAL_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES)
+DEFAULT_FAKE_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES)
+DEFAULT_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES)
+
+FILTER_PROFILE_ALIASES = {
+    "generator": "generator",
+    "student": "generator",
+    "real_score": "real_score",
+    "teacher": "real_score",
+    "fake_score": "fake_score",
+    "critic": "fake_score",
+}
+
+
+@dataclass
+class LongLiveQuantizationConfig(QuantizationConfig):
+    type: str = "weight"
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+
+        if not isinstance(self.type, str):
+            raise TypeError("Quantization type must be a string.")
+        if self.type not in QUANTIZATION_TYPE:
+            allowed = ", ".join(QUANTIZATION_TYPE.keys())
+            raise ValueError(f"Unknown quantization type '{self.type}'. Expected one of: {allowed}.")
+
+        self.type = QUANTIZATION_TYPE[self.type]
+
+
+def _resolve_modules_to_not_convert(
+    model: nn.Module,
+    filtered_modules: list[str] | None,
+) -> list[str]:
+    if not filtered_modules:
+        return []
+
+    exact_names = set()
+    regex_patterns = []
+    for pattern in filtered_modules:
+        if not isinstance(pattern, str):
+            raise TypeError("Each filtered module pattern must be a string.")
+        if pattern.startswith("re:"):
+            regex_patterns.append(re.compile(pattern[3:]))
+        else:
+            exact_names.add(pattern)
+
+    resolved = []
+    for module_name, _ in model.named_modules():
+        if not module_name:
+            continue
+        if module_name in exact_names or any(
+            regex.search(module_name) for regex in regex_patterns
+        ):
+            resolved.append(module_name)
+
+    return sorted(set(resolved))
+
+
+def _get_default_filtered_modules(filter_profile: str | None) -> list[str]:
+    if filter_profile is None:
+        return list(DEFAULT_FILTERED_MODULES)
+
+    normalized = FILTER_PROFILE_ALIASES.get(filter_profile, filter_profile)
+    if normalized == "generator":
+        return list(DEFAULT_GENERATOR_FILTERED_MODULES)
+    if normalized == "real_score":
+        return list(DEFAULT_REAL_SCORE_FILTERED_MODULES)
+    if normalized == "fake_score":
+        return list(DEFAULT_FAKE_SCORE_FILTERED_MODULES)
+
+    allowed = ", ".join(sorted(FILTER_PROFILE_ALIASES))
+    raise ValueError(
+        f"Unknown filter_profile '{filter_profile}'. Expected one of: {allowed}.",
+    )
+
+
+def _warn_for_te_config_mismatch(model_quant_config: ModelQuantizationConfig) -> None:
+    config_entries = [("default", model_quant_config)]
+    module_overrides = getattr(model_quant_config, "module_config_overrides", None) or {}
+    config_entries.extend(sorted(module_overrides.items()))
+
+    mismatched_rules = []
+    for module_name, module_config in config_entries:
+        if getattr(module_config, "dtype", DataType.nvfp4) != DataType.nvfp4:
+            raise NotImplementedError(
+                "TransformerEngine replacement currently only supports NVFP4."
+            )
+
+        for attr_name in (
+            "scale_rule",
+            "activation_scale_rule",
+            "weight_scale_rule",
+            "gradient_scale_rule",
+        ):
+            rule = getattr(module_config, attr_name, None)
+            if rule is not None and rule != ScaleRule.static_6:
+                mismatched_rules.append(f"{module_name}:{attr_name}={rule}")
+
+    if mismatched_rules:
+        preview = ", ".join(mismatched_rules[:8])
+        if len(mismatched_rules) > 8:
+            preview += ", ..."
+        warnings.warn(
+            "TransformerEngine NVFP4 path maps to `NVFP4BlockScaling` and does not "
+            "replicate FourOverSix non-`static_6` scale rules exactly. "
+            f"Mismatched config entries: {preview}",
+            stacklevel=3,
+        )
+
+
+def _build_te_recipe(module_config: Any, te_recipe_kwargs: dict[str, Any] | None = None):
+    recipe_module = importlib.import_module("transformer_engine.common.recipe")
+    NVFP4BlockScaling = recipe_module.NVFP4BlockScaling
+
+    recipe_kwargs = {
+        "disable_2d_quantization": not getattr(module_config, "weight_scale_2d", False),
+        "disable_stochastic_rounding": (
+            getattr(module_config, "gradient_round_style", RoundStyle.nearest)
+            != RoundStyle.stochastic
+        ),
+        # FourOverSix only uses RHT in specific training paths, so keep TE conservative
+        # by default and let callers override via `te_recipe_kwargs`.
+        "disable_rht": True,
+    }
+    if te_recipe_kwargs:
+        recipe_kwargs.update(te_recipe_kwargs)
+    return NVFP4BlockScaling(**recipe_kwargs)
+
+
+class TransformerEngineLinear(nn.Module):
+    """A lightweight wrapper that routes a linear layer through TransformerEngine."""
+
+    def __init__(
+        self,
+        module: nn.Linear,
+        module_name: str,
+        module_config: Any,
+        inference_only: bool = False,
+        low_precision_weights: bool = False,
+        te_recipe_kwargs: dict[str, Any] | None = None,
+        te_module_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        super().__init__()
+
+        try:
+            te = importlib.import_module("transformer_engine.pytorch")
+        except ImportError as exc:
+            raise ImportError(
+                "TransformerEngine is not installed, but `use_transformer_engine=True` "
+                "was requested."
+            ) from exc
+
+        if module.weight.device.type != "cuda":
+            raise ValueError(
+                "TransformerEngine replacement requires CUDA modules. "
+                f"Module `{module_name}` is on `{module.weight.device}`."
+            )
+
+        self.module_name = module_name
+        self.in_features = module.in_features
+        self.out_features = module.out_features
+        self.inference_only = inference_only
+        self.low_precision_weights = low_precision_weights
+        self._te = te
+        self._recipe = _build_te_recipe(
+            module_config=module_config,
+            te_recipe_kwargs=te_recipe_kwargs,
+        )
+
+        module_kwargs = dict(te_module_kwargs or {})
+        module_kwargs.setdefault("device", module.weight.device)
+        module_kwargs.setdefault("params_dtype", module.weight.dtype)
+        module_kwargs.setdefault("name", module_name)
+
+        fp8_model_init_fn = getattr(te, "fp8_model_init", None)
+        if self.low_precision_weights and fp8_model_init_fn is None:
+            warnings.warn(
+                "TransformerEngine low-precision parameter init requested, but "
+                "`fp8_model_init` is unavailable. Falling back to regular TE parameter "
+                "storage for this inference path.",
+                stacklevel=2,
+            )
+            self.low_precision_weights = False
+
+        model_init_context = (
+            fp8_model_init_fn(
+                enabled=True,
+                recipe=self._recipe,
+                preserve_high_precision_init_val=False,
+            )
+            if self.low_precision_weights
+            else nullcontext()
+        )
+        with model_init_context:
+            self.linear = te.Linear(
+                module.in_features,
+                module.out_features,
+                bias=module.bias is not None,
+                **module_kwargs,
+            )
+
+        self._load_from_linear(module)
+
+        if self.inference_only:
+            self.linear.requires_grad_(False)
+            self.train(False)
+        else:
+            self.linear.weight.requires_grad_(module.weight.requires_grad)
+            if self.linear.bias is not None and module.bias is not None:
+                self.linear.bias.requires_grad_(module.bias.requires_grad)
+            self.train(module.training)
+
+    def _copy_tensor_into_parameter(
+        self,
+        destination: torch.Tensor,
+        source: torch.Tensor,
+    ) -> None:
+        source = source.detach().to(device=destination.device)
+        try:
+            destination.copy_(source)
+            return
+        except Exception:
+            pass
+
+        destination.copy_(source.to(dtype=destination.dtype))
+
+    def _load_from_linear(self, module: nn.Linear) -> None:
+        with torch.no_grad():
+            try:
+                self._copy_tensor_into_parameter(self.linear.weight, module.weight)
+                if module.bias is not None and self.linear.bias is not None:
+                    self._copy_tensor_into_parameter(self.linear.bias, module.bias)
+                return
+            except Exception as copy_exc:
+                state_dict = {
+                    "weight": module.weight.detach().to(device=self.linear.weight.device),
+                }
+                if module.bias is not None:
+                    state_dict["bias"] = module.bias.detach().to(
+                        device=self.linear.weight.device,
+                    )
+                incompatible_keys = self.linear.load_state_dict(state_dict, strict=False)
+                missing_keys = [
+                    key
+                    for key in getattr(incompatible_keys, "missing_keys", [])
+                    if key != "_extra_state"
+                ]
+                unexpected_keys = list(getattr(incompatible_keys, "unexpected_keys", []))
+                if missing_keys or unexpected_keys:
+                    raise RuntimeError(
+                        "Failed to load weights into TransformerEngine linear "
+                        f"`{self.module_name}`. missing_keys={missing_keys}, "
+                        f"unexpected_keys={unexpected_keys}"
+                    ) from copy_exc
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.linear.weight
+
+    @property
+    def bias(self) -> torch.Tensor | None:
+        return self.linear.bias
+
+    def _autocast_context(self):
+        autocast_fn = getattr(self._te, "autocast", None)
+        if autocast_fn is not None:
+            return autocast_fn(enabled=True, recipe=self._recipe)
+        fp8_autocast_fn = getattr(self._te, "fp8_autocast", None)
+        if fp8_autocast_fn is None:
+            raise AttributeError(
+                "TransformerEngine does not expose `autocast` or `fp8_autocast`."
+            )
+        return fp8_autocast_fn(enabled=True, fp8_recipe=self._recipe)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        with self._autocast_context():
+            return self.linear(input)
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, "
+            f"out_features={self.out_features}, "
+            f"bias={self.bias is not None}, "
+            f"inference_only={self.inference_only}, "
+            f"low_precision_weights={self.low_precision_weights}, "
+            "backend=transformer_engine"
+        )
+
+
+def quantize_model_with_optional_te(
+    model: nn.Module,
+    model_quant_config: ModelQuantizationConfig,
+    *,
+    use_transformer_engine: bool = False,
+    te_inference_only: bool = False,
+    te_low_precision_weights: bool | None = None,
+    te_recipe_kwargs: dict[str, Any] | None = None,
+    te_module_kwargs: dict[str, Any] | None = None,
+    te_fallback_to_fouroversix: bool = False,
+    **kwargs,
+) -> list[str]:
+    """
+    Quantize a model with FourOverSix by default, or replace `nn.Linear` with
+    TransformerEngine wrappers when `use_transformer_engine=True`.
+    """
+    if not use_transformer_engine:
+        quantize_model(model, model_quant_config, **kwargs)
+        return []
+
+    if te_low_precision_weights is None:
+        te_low_precision_weights = te_inference_only
+
+    if kwargs:
+        if te_fallback_to_fouroversix:
+            warnings.warn(
+                "Additional kwargs passed to `quantize_model_with_optional_te` will "
+                "only be forwarded to the fallback FourOverSix pass after "
+                f"TransformerEngine replacement: {sorted(kwargs)}",
+                stacklevel=2,
+            )
+        else:
+            warnings.warn(
+                "Additional kwargs passed to `quantize_model` are ignored in the "
+                f"TransformerEngine path: {sorted(kwargs)}",
+                stacklevel=2,
+            )
+
+    _warn_for_te_config_mismatch(model_quant_config)
+
+    replaced_modules = []
+    for module_name, module in list(model.named_modules()):
+        if (
+            module_name == ""
+            or module_name in model_quant_config.modules_to_not_convert
+            or not isinstance(module, nn.Linear)
+        ):
+            continue
+
+        model.set_submodule(
+            module_name,
+            TransformerEngineLinear(
+                module=module,
+                module_name=module_name,
+                module_config=model_quant_config.get_module_config(module_name),
+                inference_only=te_inference_only,
+                low_precision_weights=te_low_precision_weights,
+                te_recipe_kwargs=te_recipe_kwargs,
+                te_module_kwargs=te_module_kwargs,
+            ),
+        )
+        replaced_modules.append(module_name)
+
+    if te_fallback_to_fouroversix:
+        quantize_model(model, model_quant_config, **kwargs)
+
+    return replaced_modules
+
+
+def _tensor_nbytes(tensor: torch.Tensor | None) -> int:
+    if tensor is None:
+        return 0
+    return tensor.numel() * tensor.element_size()
+
+
+def _materialize_transformer_engine_weights_for_inference(
+    model: nn.Module,
+    target_device: torch.device | str | None = None,
+    cache_transposed_weights: bool = False,
+) -> tuple[list[str], int, int]:
+    del cache_transposed_weights
+
+    materialized_modules = []
+    master_weight_bytes = 0
+    quantized_weight_bytes = 0
+
+    for module_name, module in model.named_modules():
+        if not isinstance(module, TransformerEngineLinear):
+            continue
+
+        if target_device is not None:
+            module.to(device=torch.device(target_device))
+
+        quantized_weight_bytes += _tensor_nbytes(module.weight)
+        quantized_weight_bytes += _tensor_nbytes(module.bias)
+        materialized_modules.append(module_name)
+
+    return materialized_modules, master_weight_bytes, quantized_weight_bytes
+
+
+def _materialize_mixed_quantized_weights_for_inference(
+    model: nn.Module,
+    target_device: torch.device | str | None = None,
+    cache_transposed_weights: bool = False,
+) -> tuple[list[str], int, int]:
+    te_modules, te_master_bytes, te_quantized_bytes = (
+        _materialize_transformer_engine_weights_for_inference(
+            model,
+            target_device=target_device,
+            cache_transposed_weights=cache_transposed_weights,
+        )
+    )
+    f46_modules, f46_master_bytes, f46_quantized_bytes = (
+        _materialize_quantized_weights_for_inference(
+            model,
+            target_device=target_device,
+            cache_transposed_weights=cache_transposed_weights,
+        )
+    )
+
+    return (
+        sorted(set(te_modules + f46_modules)),
+        te_master_bytes + f46_master_bytes,
+        te_quantized_bytes + f46_quantized_bytes,
+    )
+
+
+def _materialize_quantized_weights_for_inference(
+    model: nn.Module,
+    target_device: torch.device | str | None = None,
+    cache_transposed_weights: bool = False,
+) -> tuple[list[str], int, int]:
+    """
+    Materialize quantized weights and drop master weights.
+
+    Optionally cache an additional transposed quantized layout for training paths that
+    still require dgrad after the master weight is deleted (e.g. NVFP4 + LoRA).
+
+    This function expects modules replaced by `fouroversix.quantize_model`.
+    """
+    materialized_modules = []
+    master_weight_bytes = 0
+    quantized_weight_bytes = 0
+
+    for module_name, module in model.named_modules():
+        if not hasattr(module, "parameters_to_quantize") or not hasattr(
+            module, "get_quantized_parameters",
+        ):
+            continue
+
+        parameters_to_quantize = getattr(module, "parameters_to_quantize", ())
+        if callable(parameters_to_quantize):
+            parameters_to_quantize = parameters_to_quantize()
+        if not parameters_to_quantize:
+            continue
+
+        did_materialize = False
+        for parameter_name in parameters_to_quantize:
+            parameter = getattr(module, parameter_name, None)
+            if parameter is None:
+                continue
+
+            if isinstance(parameter, nn.Parameter):
+                parameter_tensor = parameter.data
+            elif isinstance(parameter, torch.Tensor):
+                parameter_tensor = parameter
+            else:
+                continue
+
+            master_weight_bytes += parameter_tensor.numel() * parameter_tensor.element_size()
+            get_quantized_parameters = module.get_quantized_parameters
+            if (
+                cache_transposed_weights
+                and "include_transposed" in inspect.signature(
+                    get_quantized_parameters,
+                ).parameters
+            ):
+                quantized_params = get_quantized_parameters(
+                    parameter_name,
+                    parameter_tensor,
+                    include_transposed=True,
+                )
+            else:
+                quantized_params = get_quantized_parameters(
+                    parameter_name,
+                    parameter_tensor,
+                )
+
+            for quantized_name, quantized_tensor in quantized_params.items():
+                if not isinstance(quantized_tensor, torch.Tensor):
+                    continue
+
+                existing = getattr(module, quantized_name, None)
+                dst_dtype = (
+                    existing.dtype
+                    if isinstance(existing, torch.Tensor)
+                    else quantized_tensor.dtype
+                )
+                if target_device is not None:
+                    dst_device = torch.device(target_device)
+                elif isinstance(existing, torch.Tensor):
+                    dst_device = existing.device
+                else:
+                    dst_device = quantized_tensor.device
+
+                quantized_tensor = quantized_tensor.to(
+                    device=dst_device,
+                    dtype=dst_dtype,
+                )
+                setattr(module, quantized_name, quantized_tensor)
+                quantized_weight_bytes += (
+                    quantized_tensor.numel() * quantized_tensor.element_size()
+                )
+
+            # Drop high-precision master weight once quantized weights are materialized.
+            if isinstance(getattr(module, parameter_name, None), nn.Parameter):
+                module.register_parameter(parameter_name, None)
+            else:
+                setattr(module, parameter_name, None)
+            did_materialize = True
+
+        if did_materialize:
+            if hasattr(module, "_quantized_weight"):
+                delattr(module, "_quantized_weight")
+            if hasattr(module, "_quantized_weight_transposed"):
+                delattr(module, "_quantized_weight_transposed")
+            if hasattr(module, "_quantized_weights"):
+                delattr(module, "_quantized_weights")
+            if hasattr(module, "config") and hasattr(module.config, "keep_master_weights"):
+                module.config.keep_master_weights = False
+            materialized_modules.append(module_name)
+
+    return materialized_modules, master_weight_bytes, quantized_weight_bytes
+
+
+def quantize_model_with_filter(
+    model: nn.Module,
+    quant_config: ModelQuantizationConfig | dict | None = None,
+    filtered_modules: list[str] | None = None,
+    filter_profile: str | None = None,
+    use_default_filtered_modules: bool = False,
+    cast_model_to_bf16: bool = True,
+    materialize_for_inference: bool = False,
+    materialize_target_device: torch.device | str | None = None,
+    use_transformer_engine: bool = False,
+    te_inference_only: bool = False,
+    te_low_precision_weights: bool | None = None,
+    te_recipe_kwargs: dict[str, Any] | None = None,
+    te_module_kwargs: dict[str, Any] | None = None,
+    te_fallback_to_fouroversix: bool = False,
+    verbose: bool = True,
+    **kwargs,
+) -> tuple[nn.Module, list[str]]:
+    """
+    Quantize model with FourOverSix and optionally skip selected modules.
+
+    `filtered_modules` supports:
+    - Exact module names, e.g. "head.head"
+    - Regex patterns prefixed with "re:", e.g. "re:.*norm1$"
+
+    `filter_profile` selects which built-in filtered module profile to use when
+    `use_default_filtered_modules=True`. Supported values:
+    "generator"/"student" and "real_score"/"teacher".
+    """
+    if quant_config is None:
+        model_quant_config = ModelQuantizationConfig()
+    elif isinstance(quant_config, dict):
+        model_quant_config = ModelQuantizationConfig(**quant_config)
+    elif isinstance(quant_config, ModelQuantizationConfig):
+        model_quant_config = deepcopy(quant_config)
+    else:
+        raise TypeError(
+            "quant_config must be ModelQuantizationConfig, dict, or None.",
+        )
+
+    patterns = list(filtered_modules or [])
+    if use_default_filtered_modules:
+        patterns = _get_default_filtered_modules(filter_profile) + patterns
+
+    matched_modules = _resolve_modules_to_not_convert(model, patterns)
+    modules_to_not_convert = set(model_quant_config.modules_to_not_convert or [])
+    modules_to_not_convert.update(matched_modules)
+    model_quant_config.modules_to_not_convert = sorted(modules_to_not_convert)
+
+    if cast_model_to_bf16:
+        model.to(torch.bfloat16)
+
+    resolved_te_low_precision_weights = (
+        te_inference_only if te_low_precision_weights is None else te_low_precision_weights
+    )
+
+    te_replaced_modules = quantize_model_with_optional_te(
+        model,
+        model_quant_config,
+        use_transformer_engine=use_transformer_engine,
+        te_inference_only=te_inference_only,
+        te_low_precision_weights=resolved_te_low_precision_weights,
+        te_recipe_kwargs=te_recipe_kwargs,
+        te_module_kwargs=te_module_kwargs,
+        te_fallback_to_fouroversix=te_fallback_to_fouroversix,
+        **kwargs,
+    )
+
+    if materialize_for_inference:
+        materialize_fn = _materialize_quantized_weights_for_inference
+        if use_transformer_engine and te_fallback_to_fouroversix:
+            materialize_fn = _materialize_mixed_quantized_weights_for_inference
+        elif use_transformer_engine:
+            materialize_fn = _materialize_transformer_engine_weights_for_inference
+
+        materialized_modules, master_bytes, quantized_bytes = materialize_fn(
+            model,
+            target_device=materialize_target_device,
+        )
+        if verbose:
+            print(
+                "[quantize_model_with_filter] "
+                f"materialized_modules={len(materialized_modules)}, "
+                f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, "
+                f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB",
+            )
+
+    if verbose:
+        profile_label = filter_profile or "default"
+        print(
+            "[quantize_model_with_filter] "
+            f"profile={profile_label}, "
+            f"matched={len(matched_modules)}, "
+            f"total_excluded={len(model_quant_config.modules_to_not_convert)}",
+        )
+        if use_transformer_engine:
+            print(
+                "[quantize_model_with_filter] "
+                f"transformer_engine_replaced={len(te_replaced_modules)}, "
+                f"inference_only={te_inference_only}, "
+                f"low_precision_weights={resolved_te_low_precision_weights}, "
+                f"fallback_to_fouroversix={te_fallback_to_fouroversix}",
+            )
+
+    return model, matched_modules
+
+
+def _dequantize_kv_cache_fused_cuda(kv_list, max_blocks, num_heads, block_token_size, dtype):
+    global _FUSED_KV_DEQUANT_DISABLED, _FUSED_KV_DEQUANT_WARNED
+
+    if _FUSED_KV_DEQUANT_DISABLED or max_blocks <= 0:
+        return None
+
+    first_qt = kv_list[0]
+    if first_qt.values.device.type != "cuda":
+        return None
+
+    try:
+        from utils.kernel.kv_dequant import dequantize_kv_cache_fp4
+
+        blocks = kv_list[:max_blocks]
+        values = [qt.values for qt in blocks]
+        scale_factors = [qt.scale_factors for qt in blocks]
+        amax = [qt.amax for qt in blocks]
+
+        return dequantize_kv_cache_fp4(
+            values,
+            scale_factors,
+            amax,
+            num_heads=num_heads,
+            block_token_size=block_token_size,
+            dtype=dtype,
+            scale_rule=first_qt.scale_rule,
+        )
+    except Exception as exc:  # pragma: no cover - exercised only when extension is stale/missing
+        _FUSED_KV_DEQUANT_DISABLED = True
+        if not _FUSED_KV_DEQUANT_WARNED:
+            warnings.warn(
+                "Fused CUDA KV-cache dequantization is unavailable; falling back to "
+                f"the Triton per-block path. Reason: {exc}",
+                stacklevel=2,
+            )
+            _FUSED_KV_DEQUANT_WARNED = True
+        return None
+
+
+def dequantize_kv_cache(kv_list, max_blocks, num_heads, block_token_size, dtype, device):
+    """
+    Dequantize list of QuantizedTensor to a contiguous bf16 tensor.
+    kv_list[block_idx] -> QuantizedTensor(block_token_size * num_heads, 128)
+    Returns: [1, max_blocks * block_token_size, num_heads, 128]
+    """
+    fused_result = _dequantize_kv_cache_fused_cuda(
+        kv_list, max_blocks, num_heads, block_token_size, dtype,
+    )
+    if fused_result is not None:
+        return fused_result
+
+    total_tokens = max_blocks * block_token_size
+    result = torch.zeros([1, total_tokens, num_heads, 128], dtype=dtype, device=device)
+    for block_idx in range(max_blocks):
+        t_start = block_idx * block_token_size
+        t_end = t_start + block_token_size
+        # deq = kv_list[block_idx].dequantize(dtype)
+        # triton fp4_dequantize
+        qt = kv_list[block_idx]
+        padded_shape = qt.padded_shape
+        scales_2d = from_blocked(
+            qt.scale_factors,
+            (padded_shape[0], padded_shape[1] // 16),
+        )
+        global_scale = qt.amax / (
+            qt.scale_rule.max_allowed_e2m1_value()
+            * qt.scale_rule.max_allowed_e4m3_value()
+        )
+        deq = fp4_dequantize(
+            kv_list[block_idx].values,
+            scales_2d,
+            global_scale,
+            block_size=16,
+            dtype=dtype,
+        )
+        result[0, t_start:t_end, :, :] = deq.view(block_token_size, num_heads, 128)
+    return result
+
+def clone_quantized_tensor(qt):
+    """Clone a QuantizedTensor by cloning its internal tensors."""
+    return QuantizedTensor(
+        values=qt.values.clone(),
+        scale_factors=qt.scale_factors.clone(),
+        amax=qt.amax.clone() if qt.amax is not None else None,
+        dtype=qt.dtype,
+        original_shape=qt.original_shape,
+        scale_rule=qt.scale_rule,
+        padded_shape=qt.padded_shape,
+    )
+
+
+def copy_quantized_into(slot: QuantizedTensor, src: QuantizedTensor) -> None:
+    """In-place copy a QuantizedTensor's data into a pre-allocated slot.
+
+    Keeps the slot's `values`/`scale_factors`/`amax` buffers persistent
+    (their addresses don't change) so cudagraph allocator does not see them
+    as fresh outputs that can be reused across step boundaries. Used by the
+    quantized KV cache rolling/insert paths.
+    """
+    slot.values.copy_(src.values)
+    slot.scale_factors.copy_(src.scale_factors)
+    if src.amax is not None and slot.amax is not None:
+        slot.amax.copy_(src.amax)
+
+
+def k_smooth(k: torch.Tensor) -> torch.Tensor:
+    return k - k.mean(dim=-1, keepdim=True)
+
+def quantize_kv(k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    B, S, H, D = k.shape
+    # B is always 1
+    # S is the number of tokens
+    # H is the number of heads
+    # D is the dimension of the key and value
+
+    config = QuantizationConfig(scale_rule="mse", backend="cuda")
+    # per head quantization
+    for head in range(H):
+        k_head = k[:, :, head, :]
+        v_head = v[:, :, head, :]
+        k_head = k_smooth(k_head)
+        v_head = v_head
+        k_head = quantize_to_fp4(k_head, config)
+        v_head = quantize_to_fp4(v_head, config)
+        k[:, :, head, :] = k_head
+        v[:, :, head, :] = v_head
+    return k, v

utils/rope_triton.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# SPDX-License-Identifier: Apache-2.0
+"""Triton RoPE kernel for causal_rope_apply.
+
+iter-42: replaces the complex<double> × view_as_complex × view_as_real chain
+(33 ms / 1.3% of profile + feeding elementwise muls) with a single Triton
+kernel. Internal precision is fp32 — bf16 outputs cannot resolve any precision
+loss from fp32 vs fp64 arithmetic at this stage. cos / sin lookup tables come
+from the complex128 freqs split into real / imag floats up-front (one-shot per
+freqs_i cache entry).
+"""
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+# iter-45 (REVERTED): tried @triton.autotune over (BLOCK_N∈{4,8,16},
+# num_warps∈{2,4,8}, num_stages∈{1,2,3}) = 27 configs. Result was FLAT vs
+# iter-42 fixed BLOCK_N=8: median tied (-0.1%), total +0.4% (autotune warmup
+# bled into p1/p2 p90). The original BLOCK_N=8 / default warps was already
+# near-optimal for the (N=24, D_half=64) shape, autotune found no better.
+# Reverted to fixed config — same kernel as iter-42.
+#
+# iter-46: kernel now accepts FULL x[i] of shape [S_total, N, D] and a
+# runtime `seq_len` — for rows s < seq_len it applies rotation, for
+# s >= seq_len it copies through. This subsumes the `torch.cat([rotated,
+# x[i, seq_len:]])` step (1 fewer kernel + 1 fewer alloc per call). Also
+# skips the upstream `.contiguous()` because we no longer slice x.
+@triton.jit
+def _rope_apply_kernel(
+    x_ptr,         # [S_total, N, D] bf16  (D is even, pairs are (a,b)=(2d, 2d+1))
+    cos_ptr,       # [seq_len, D/2] fp32 (valid only for s < seq_len)
+    sin_ptr,       # [seq_len, D/2] fp32
+    out_ptr,       # [S_total, N, D] bf16
+    SEQ_LEN, N, D_half,
+    x_stride_s, x_stride_n,
+    o_stride_s, o_stride_n,
+    cs_stride_s,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    pid_s = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_D)  # over the D/2 pairs
+
+    n_mask = offs_n < N
+    d_mask = offs_d < D_half
+
+    x_row_base = pid_s * x_stride_s
+    o_row_base = pid_s * o_stride_s
+    a_offs = x_row_base + offs_n[:, None] * x_stride_n + (2 * offs_d)[None, :]
+    b_offs = a_offs + 1
+    mask = n_mask[:, None] & d_mask[None, :]
+    a = tl.load(x_ptr + a_offs, mask=mask, other=0.0).to(tl.float32)
+    b = tl.load(x_ptr + b_offs, mask=mask, other=0.0).to(tl.float32)
+
+    a_out_offs = o_row_base + offs_n[:, None] * o_stride_n + (2 * offs_d)[None, :]
+    b_out_offs = a_out_offs + 1
+
+    if pid_s < SEQ_LEN:
+        cs_base = pid_s * cs_stride_s
+        cos = tl.load(cos_ptr + cs_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
+        sin = tl.load(sin_ptr + cs_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
+        # Rotate: (a + bi) * (cos + sin i) = (a*cos - b*sin) + (a*sin + b*cos) i
+        out_a = a * cos[None, :] - b * sin[None, :]
+        out_b = a * sin[None, :] + b * cos[None, :]
+        tl.store(out_ptr + a_out_offs, out_a, mask=mask)
+        tl.store(out_ptr + b_out_offs, out_b, mask=mask)
+    else:
+        # passthrough copy for the unrotated tail
+        tl.store(out_ptr + a_out_offs, a, mask=mask)
+        tl.store(out_ptr + b_out_offs, b, mask=mask)
+
+
+def _split_complex_to_cos_sin(freqs_complex: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert complex128 freqs to (cos_f32, sin_f32) — once per cache entry."""
+    # freqs_complex shape: (S, 1, D/2). Squeeze the middle 1.
+    if freqs_complex.dim() == 3 and freqs_complex.size(1) == 1:
+        freqs_complex = freqs_complex.squeeze(1)
+    cos = freqs_complex.real.to(torch.float32).contiguous()
+    sin = freqs_complex.imag.to(torch.float32).contiguous()
+    return cos, sin
+
+
+def rope_apply_triton(
+    x: torch.Tensor,           # [S_total, N, D] bf16 (or fp16/fp32)
+    cos_f32: torch.Tensor,     # [seq_len, D/2] fp32
+    sin_f32: torch.Tensor,     # [seq_len, D/2] fp32
+    seq_len: int | None = None,
+) -> torch.Tensor:
+    """Apply rotary embedding via Triton kernel.
+
+    iter-46: when `seq_len < x.size(0)`, the kernel rotates the first
+    `seq_len` rows and copies through rows `[seq_len:]`. This replaces the
+    `cat([rotated, x[i, seq_len:]])` pattern in `causal_rope_apply` with a
+    single kernel + single allocation. `seq_len=None` (default) means rotate
+    all rows (equivalent to iter-42 behavior).
+
+    Returns a tensor of the same shape and dtype as `x`.
+    """
+    assert x.dim() == 3, f"expected x.shape == (S, N, D), got {x.shape}"
+    S_total, N, D = x.shape
+    assert D % 2 == 0
+    D_half = D // 2
+    if seq_len is None:
+        seq_len = S_total
+    assert seq_len <= S_total
+    assert cos_f32.shape == (seq_len, D_half), \
+        f"cos_f32 expected ({seq_len},{D_half}), got {cos_f32.shape}"
+    assert sin_f32.shape == (seq_len, D_half)
+
+    out = torch.empty_like(x)
+
+    BLOCK_N = 8
+    BLOCK_D = triton.next_power_of_2(D_half)
+    grid = (S_total, triton.cdiv(N, BLOCK_N))
+
+    _rope_apply_kernel[grid](
+        x, cos_f32, sin_f32, out,
+        seq_len, N, D_half,
+        x.stride(0), x.stride(1),
+        out.stride(0), out.stride(1),
+        cos_f32.stride(0),
+        BLOCK_N=BLOCK_N, BLOCK_D=BLOCK_D,
+    )
+    return out

utils/scheduler.py

@@ -0,0 +1,194 @@
+from abc import abstractmethod, ABC
+import torch
+
+
+class SchedulerInterface(ABC):
+    """
+    Base class for diffusion noise schedule.
+    """
+    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule
+
+    @abstractmethod
+    def add_noise(
+        self, clean_latent: torch.Tensor,
+        noise: torch.Tensor, timestep: torch.Tensor
+    ):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        pass
+
+    def convert_x0_to_noise(
+        self, x0: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's x0 prediction to noise predidction.
+        x0: the predicted clean data with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = x0.dtype
+        x0, xt, alphas_cumprod = map(
+            lambda x: x.double().to(x0.device), [x0, xt,
+                                                 self.alphas_cumprod]
+        )
+
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        noise_pred = (xt - alpha_prod_t **
+                      (0.5) * x0) / beta_prod_t ** (0.5)
+        return noise_pred.to(original_dtype)
+
+    def convert_noise_to_x0(
+        self, noise: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's noise prediction to x0 predidction.
+        noise: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)
+        """
+        # use higher precision for calculations
+        original_dtype = noise.dtype
+        noise, xt, alphas_cumprod = map(
+            lambda x: x.double().to(noise.device), [noise, xt,
+                                                    self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (xt - beta_prod_t **
+                   (0.5) * noise) / alpha_prod_t ** (0.5)
+        return x0_pred.to(original_dtype)
+
+    def convert_velocity_to_x0(
+        self, velocity: torch.Tensor, xt: torch.Tensor,
+        timestep: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Convert the diffusion network's velocity prediction to x0 predidction.
+        velocity: the predicted noise with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        v = sqrt(alpha_t) * noise - sqrt(beta_t) x0
+        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)
+        given v, x_t, we have
+        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v
+        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac56
+        """
+        # use higher precision for calculations
+        original_dtype = velocity.dtype
+        velocity, xt, alphas_cumprod = map(
+            lambda x: x.double().to(velocity.device), [velocity, xt,
+                                                       self.alphas_cumprod]
+        )
+        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
+        beta_prod_t = 1 - alpha_prod_t
+
+        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity
+        return x0_pred.to(original_dtype)
+
+
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) /
+                          num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (
+                self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B*T, C, H, W]
+            - noise: the noise with shape [B*T, C, H, W]
+            - timestep: the timestep with shape [B*T]
+        Output: the corrupted latent with shape [B*T, C, H, W]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        """
+        Input:
+            - timestep: the timestep with shape [B*T]
+        Output: the corresponding weighting [B*T]
+        """
+        if timestep.ndim == 2:
+            timestep = timestep.flatten(0, 1)
+        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights

utils/torch_compile_utils.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
+#
+# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+import torch.distributed as dist
+
+
+def _is_main_process() -> bool:
+    return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
+
+
+def _log_once(message: str) -> None:
+    if _is_main_process():
+        print(message)
+
+
+class SafeCompiledCallable:
+    """Lazy torch.compile wrapper that falls back to eager on compile/runtime errors."""
+
+    def __init__(
+        self,
+        fn,
+        *,
+        name: str,
+        backend: str = "inductor",
+        mode: str | None = "max-autotune-no-cudagraphs",
+        fullgraph: bool = False,
+        dynamic: bool | None = False,
+        options: dict | None = None,
+        suppress_errors: bool = True,
+    ) -> None:
+        self.fn = fn
+        self.name = name
+        self.enabled = True
+        self.failed = False
+        self.failure_reason = None
+
+        if suppress_errors:
+            try:
+                import torch._dynamo as torch_dynamo
+
+                torch_dynamo.config.suppress_errors = True
+            except Exception as exc:
+                _log_once(f"[torch.compile] Could not enable suppress_errors: {exc}")
+
+        compile_kwargs = {
+            "backend": backend,
+            "fullgraph": fullgraph,
+            "dynamic": dynamic,
+        }
+        if mode:
+            compile_kwargs["mode"] = mode
+        if options:
+            compile_kwargs["options"] = options
+
+        _log_once(
+            "[torch.compile] Preparing "
+            f"{name}: backend={backend}, mode={mode}, "
+            f"fullgraph={fullgraph}, dynamic={dynamic}"
+        )
+        self.compiled_fn = torch.compile(fn, **compile_kwargs)
+
+    def __call__(self, *args, **kwargs):
+        if not self.enabled:
+            return self.fn(*args, **kwargs)
+
+        try:
+            return self.compiled_fn(*args, **kwargs)
+        except Exception as exc:
+            self.enabled = False
+            self.failed = True
+            self.failure_reason = repr(exc)
+            _log_once(
+                f"[torch.compile][warn] {self.name} failed; "
+                f"falling back to eager. reason={exc}"
+            )
+            return self.fn(*args, **kwargs)
+
+
+def configure_module_call_torch_compile(
+    module,
+    *,
+    name: str,
+    backend: str = "inductor",
+    mode: str | None = "max-autotune-no-cudagraphs",
+    fullgraph: bool = False,
+    dynamic: bool | None = False,
+    options: dict | None = None,
+    suppress_errors: bool = True,
+):
+    if not torch.cuda.is_available():
+        _log_once(f"[torch.compile] Skipping {name}: CUDA is not available")
+        return None
+
+    try:
+        return SafeCompiledCallable(
+            module,
+            name=name,
+            backend=backend,
+            mode=mode,
+            fullgraph=fullgraph,
+            dynamic=dynamic,
+            options=options,
+            suppress_errors=suppress_errors,
+        )
+    except Exception as exc:
+        _log_once(
+            f"[torch.compile][warn] Could not prepare {name}; "
+            f"continuing in eager mode. reason={exc}"
+        )
+        return None

utils/wan_5b_wrapper.py

@@ -0,0 +1,582 @@
+import types
+from typing import List, Optional
+import os
+import torch
+from torch import nn
+
+from utils.scheduler import SchedulerInterface, FlowMatchScheduler
+
+from wan_5b.modules.tokenizers import HuggingfaceTokenizer
+from wan_5b.modules.model import WanModel
+from wan_5b.modules.vae2_2 import _video_vae
+from wan_5b.modules.t5 import umt5_xxl
+from wan_5b.modules.causal_model import CausalWanModel
+
+
+class WanTextEncoder(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self.text_encoder = umt5_xxl(
+            encoder_only=True,
+            return_tokenizer=False,
+            dtype=torch.float32,
+            device=torch.device('cpu')
+        ).eval().requires_grad_(False)
+        self.text_encoder.load_state_dict(
+            torch.load("wan_models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth",
+                       map_location='cpu', weights_only=False)
+        )
+        
+        # Move text encoder to GPU if available
+        if torch.cuda.is_available():
+            self.text_encoder = self.text_encoder.cuda()
+
+        self.tokenizer = HuggingfaceTokenizer(
+            name="wan_models/Wan2.2-TI2V-5B/google/umt5-xxl/", seq_len=512, clean='whitespace')
+
+    @property
+    def device(self):
+        # Assume we are always on GPU
+        return torch.cuda.current_device()
+
+    def forward(self, text_prompts: List[str]) -> dict:
+        ids, mask = self.tokenizer(
+            text_prompts, return_mask=True, add_special_tokens=True)
+        ids = ids.to(self.device)
+        mask = mask.to(self.device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        context = self.text_encoder(ids, mask)
+        for u, v in zip(context, seq_lens):
+            u[v:] = 0.0  # set padding to 0.0
+
+        return {
+            "prompt_embeds": context
+        }
+
+
+class WanVAEWrapper(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        mean = [
+                -0.2289,
+                -0.0052,
+                -0.1323,
+                -0.2339,
+                -0.2799,
+                0.0174,
+                0.1838,
+                0.1557,
+                -0.1382,
+                0.0542,
+                0.2813,
+                0.0891,
+                0.1570,
+                -0.0098,
+                0.0375,
+                -0.1825,
+                -0.2246,
+                -0.1207,
+                -0.0698,
+                0.5109,
+                0.2665,
+                -0.2108,
+                -0.2158,
+                0.2502,
+                -0.2055,
+                -0.0322,
+                0.1109,
+                0.1567,
+                -0.0729,
+                0.0899,
+                -0.2799,
+                -0.1230,
+                -0.0313,
+                -0.1649,
+                0.0117,
+                0.0723,
+                -0.2839,
+                -0.2083,
+                -0.0520,
+                0.3748,
+                0.0152,
+                0.1957,
+                0.1433,
+                -0.2944,
+                0.3573,
+                -0.0548,
+                -0.1681,
+                -0.0667,
+            ]
+        std = [
+                0.4765,
+                1.0364,
+                0.4514,
+                1.1677,
+                0.5313,
+                0.4990,
+                0.4818,
+                0.5013,
+                0.8158,
+                1.0344,
+                0.5894,
+                1.0901,
+                0.6885,
+                0.6165,
+                0.8454,
+                0.4978,
+                0.5759,
+                0.3523,
+                0.7135,
+                0.6804,
+                0.5833,
+                1.4146,
+                0.8986,
+                0.5659,
+                0.7069,
+                0.5338,
+                0.4889,
+                0.4917,
+                0.4069,
+                0.4999,
+                0.6866,
+                0.4093,
+                0.5709,
+                0.6065,
+                0.6415,
+                0.4944,
+                0.5726,
+                1.2042,
+                0.5458,
+                1.6887,
+                0.3971,
+                1.0600,
+                0.3943,
+                0.5537,
+                0.5444,
+                0.4089,
+                0.7468,
+                0.7744,
+            ]
+        self.mean = torch.tensor(mean, dtype=torch.float32)
+        self.std = torch.tensor(std, dtype=torch.float32)
+
+        # init model
+        self.model = _video_vae(
+            pretrained_path="wan_models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth",
+        ).eval().requires_grad_(False)
+
+    def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:
+        # pixel: [batch_size, num_channels, num_frames, height, width]
+        device, dtype = pixel.device, pixel.dtype
+
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        output = [
+            self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)
+            for u in pixel
+        ]
+        output = torch.stack(output, dim=0)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+    def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:
+        # from [batch_size, num_frames, num_channels, height, width]
+        # to [batch_size, num_channels, num_frames, height, width]
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        if use_cache:
+            decode_function = self.model.cached_decode
+        else:
+            decode_function = self.model.decode
+
+        output = []
+        for u in zs:
+            output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))
+        output = torch.stack(output, dim=0)
+        # from [batch_size, num_channels, num_frames, height, width]
+        # to [batch_size, num_frames, num_channels, height, width]
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+    def decode_to_pixel_chunk(self, latent: torch.Tensor, use_cache: bool = False, chunk_size: int = 1) -> torch.Tensor:
+        """
+        Decode latent frames to pixel space.
+        
+        Args:
+            latent: Latent tensor with shape [batch_size, num_frames, num_channels, height, width]
+            use_cache: Whether to use cached decoding (for streaming)
+            chunk_size: Number of latent frames to decode at once (default 240 to avoid OOM)
+        
+        Returns:
+            Decoded video tensor with shape [batch_size, num_frames, num_channels, height, width]
+        """
+        # latent shape: [batch_size, num_frames, num_channels, height, width]
+        # zs shape after permute: [batch_size, num_channels, num_frames, height, width]
+        zs = latent.permute(0, 2, 1, 3, 4)
+        if use_cache:
+            assert latent.shape[0] == 1, "Batch size must be 1 when using cache"
+
+        device, dtype = latent.device, latent.dtype
+        scale = [self.mean.to(device=device, dtype=dtype),
+                 1.0 / self.std.to(device=device, dtype=dtype)]
+
+        if use_cache:
+            decode_function = self.model.cached_decode
+        else:
+            decode_function = self.model.decode
+
+        output = []
+        for u in zs:
+            num_frames = u.shape[1]
+            if num_frames <= chunk_size:
+                # Decode short clips in one pass.
+                if use_cache:
+                    # Start this segment from a clean cache.
+                    self.model.clear_cache()
+                decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+                decoded = decoded.cpu()
+                if use_cache:
+                    # Clear after this segment so it cannot affect the next video.
+                    self.model.clear_cache()
+            else:
+                # Decode longer clips in temporal chunks.
+                decoded_chunks = []
+                if use_cache:
+                    # Clear once at the segment start; later chunks share the
+                    # internal cache.
+                    self.model.clear_cache()
+                for start_idx in range(0, num_frames, chunk_size):
+                    end_idx = min(start_idx + chunk_size, num_frames)
+                    chunk = u[:, start_idx:end_idx, :, :]  # [C, chunk_frames, H, W]
+                    decoded_chunk = decode_function(chunk.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)
+                    decoded_chunks.append(decoded_chunk.cpu())
+
+                    del decoded_chunk
+                    torch.cuda.empty_cache()
+                decoded = torch.cat(decoded_chunks, dim=1)
+                if use_cache:
+                    # Clear the cache after the full segment.
+                    self.model.clear_cache()
+            output.append(decoded)
+
+        output = torch.stack(output, dim=0)
+        output = output.permute(0, 2, 1, 3, 4)
+        return output
+
+
+class WanDiffusionWrapper(torch.nn.Module):
+    def __init__(
+            self,
+            model_name="Wan2.2-TI2V-5B",
+            timestep_shift=8.0,
+            is_causal=False,
+            local_attn_size=-1,
+            sink_size=0,
+            num_frame_per_block=1,
+            t_scale=1.0,
+            rope_method="linear",
+            original_seq_len=None,
+    ):
+        super().__init__()
+
+        if is_causal:
+            self.model = CausalWanModel.from_pretrained(
+                f"wan_models/{model_name}/", local_attn_size=local_attn_size, sink_size=sink_size,
+                num_frame_per_block=num_frame_per_block)
+        else:
+            self.model = WanModel.from_pretrained(f"wan_models/{model_name}/")
+        self.model.eval()
+        self.model.t_scale = t_scale
+        self.model.rope_method = rope_method
+        self.model.original_seq_len = original_seq_len
+
+        # For non-causal diffusion, all frames share the same timestep
+        self.uniform_timestep = not is_causal
+
+        self.scheduler = FlowMatchScheduler(
+            shift=timestep_shift, sigma_min=0.0, extra_one_step=True
+        )
+        self.scheduler.set_timesteps(1000, training=True)
+
+        self.seq_len = 28160  # [1, 32, 48, 44, 80]
+    
+        self.post_init()
+        self._compiled_model_call = None
+
+    def enable_gradient_checkpointing(self) -> None:
+        self.model.enable_gradient_checkpointing()
+
+    def configure_torch_compile(
+        self,
+        *,
+        backend: str = "inductor",
+        mode: str | None = "max-autotune-no-cudagraphs",
+        fullgraph: bool = False,
+        dynamic: bool | None = False,
+        options: dict | None = None,
+        suppress_errors: bool = True,
+    ) -> bool:
+        from utils.torch_compile_utils import configure_module_call_torch_compile
+
+        self._compiled_model_call = configure_module_call_torch_compile(
+            self.model,
+            name="WanDiffusionWrapper5B.model",
+            backend=backend,
+            mode=mode,
+            fullgraph=fullgraph,
+            dynamic=dynamic,
+            options=options,
+            suppress_errors=suppress_errors,
+        )
+        return self._compiled_model_call is not None
+
+    def _call_model(self, *args, **kwargs):
+        # iter-39 v2: publish kv_cache scalars BEFORE entering the compiled
+        # graph. The earlier version (iter-39 v1) published them inside
+        # `_forward_inference`, but that function IS compiled, so each
+        # `.item()` triggered a graph break. Moving the reads to this eager
+        # wrapper keeps the dict lookups in the compiled attention forward
+        # free of `.item()` syncs without adding any graph break.
+        kv_cache = kwargs.get("kv_cache", None)
+        if kv_cache is not None and len(kv_cache) > 0:
+            try:
+                from wan_5b.modules.causal_model import _CURRENT_GRID_META
+                first_block_cache = kv_cache[0]
+                _CURRENT_GRID_META["global_end_index"] = int(
+                    first_block_cache["global_end_index"].item()
+                )
+                _CURRENT_GRID_META["local_end_index"] = int(
+                    first_block_cache["local_end_index"].item()
+                )
+                _ps = first_block_cache.get("pinned_start", None)
+                if _ps is not None and hasattr(_ps, "item"):
+                    _CURRENT_GRID_META["pinned_start"] = int(_ps.item())
+                    _CURRENT_GRID_META["pinned_len"] = int(
+                        first_block_cache["pinned_len"].item()
+                    )
+                else:
+                    _CURRENT_GRID_META["pinned_start"] = -1
+                    _CURRENT_GRID_META["pinned_len"] = 0
+            except (KeyError, AttributeError, ImportError):
+                pass
+        defer_kv_updates = (
+            os.environ.get("LLV2_DEFER_KV_UPDATES", "0") == "1"
+            and kv_cache is not None
+        )
+        if defer_kv_updates:
+            kwargs["defer_cache_updates"] = True
+
+        if self._compiled_model_call is not None:
+            # iter-25: signal cudagraph allocator that a new "step" starts.
+            # Required for mode=reduce-overhead when modules cache state
+            # (KV cache rolling buffers, fp4-quant scale tensors) so the
+            # cudagraph pool knows it can safely reuse step-N memory now
+            # that step-(N+1) is starting.
+            mark_step = getattr(torch.compiler, "cudagraph_mark_step_begin", None)
+            if mark_step is not None:
+                mark_step()
+            result = self._compiled_model_call(*args, **kwargs)
+        else:
+            result = self.model(*args, **kwargs)
+
+        if defer_kv_updates:
+            if not isinstance(result, tuple) or len(result) != 2:
+                raise RuntimeError(
+                    "LLV2_DEFER_KV_UPDATES expected model to return "
+                    "(output, cache_update_infos)."
+                )
+            output, cache_update_infos = result
+            if cache_update_infos:
+                self.model._apply_cache_updates(kv_cache, cache_update_infos)
+            return output
+        return result
+
+    def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert flow matching's prediction to x0 prediction.
+        flow_pred: the prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = noise - x0
+        x_t = (1-sigma_t) * x0 + sigma_t * noise
+        we have x0 = x_t - sigma_t * pred
+        see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e
+        """
+        # use higher precision for calculations
+        original_dtype = flow_pred.dtype
+        flow_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(flow_pred.device), [flow_pred, xt,
+                                                        self.scheduler.sigmas,
+                                                        self.scheduler.timesteps]
+        )
+
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        x0_pred = xt - sigma_t * flow_pred
+        return x0_pred.to(original_dtype)
+
+    @staticmethod
+    def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
+        """
+        Convert x0 prediction to flow matching's prediction.
+        x0_pred: the x0 prediction with shape [B, C, H, W]
+        xt: the input noisy data with shape [B, C, H, W]
+        timestep: the timestep with shape [B]
+
+        pred = (x_t - x_0) / sigma_t
+        """
+        # use higher precision for calculations
+        original_dtype = x0_pred.dtype
+        x0_pred, xt, sigmas, timesteps = map(
+            lambda x: x.double().to(x0_pred.device), [x0_pred, xt,
+                                                      scheduler.sigmas,
+                                                      scheduler.timesteps]
+        )
+        timestep_id = torch.argmin(
+            (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        flow_pred = (xt - x0_pred) / sigma_t
+        return flow_pred.to(original_dtype)
+
+    def forward(
+        self,
+        noisy_image_or_video: torch.Tensor, conditional_dict: dict,
+        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,
+        crossattn_cache: Optional[List[dict]] = None,
+        current_start: Optional[int] = None,
+        classify_mode: Optional[bool] = False,
+        concat_time_embeddings: Optional[bool] = False,
+        clean_x: Optional[torch.Tensor] = None,
+        aug_t: Optional[torch.Tensor] = None,
+        cache_start: Optional[int] = None,
+        rope_temporal_offset: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        prompt_embeds = conditional_dict["prompt_embeds"]
+
+        # [B, F] -> [B]
+        if self.uniform_timestep:
+            input_timestep = timestep[:, 0]
+        else:
+            input_timestep = timestep
+
+        logits = None
+        rope_offset_was_set = (
+            rope_temporal_offset is not None
+            and hasattr(self.model, "rope_temporal_offset")
+        )
+        if rope_offset_was_set:
+            prev_rope_temporal_offset = self.model.rope_temporal_offset
+            self.model.rope_temporal_offset = rope_temporal_offset
+
+        # X0 prediction
+        if kv_cache is not None:
+            flow_pred = self._call_model(
+                noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                t=input_timestep, context=prompt_embeds,
+                seq_len=self.seq_len,
+                kv_cache=kv_cache,
+                crossattn_cache=crossattn_cache,
+                current_start=current_start,
+                cache_start=cache_start
+            ).permute(0, 2, 1, 3, 4)
+        else:
+            if clean_x is not None:
+                # teacher forcing
+                flow_pred = self._call_model(
+                    noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                    t=input_timestep, context=prompt_embeds,
+                    seq_len=self.seq_len,
+                    clean_x=clean_x.permute(0, 2, 1, 3, 4),
+                    aug_t=aug_t,
+                ).permute(0, 2, 1, 3, 4)
+            else:
+                if classify_mode:
+                    flow_pred, logits = self._call_model(
+                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                        t=input_timestep, context=prompt_embeds,
+                        seq_len=self.seq_len,
+                        classify_mode=True,
+                        register_tokens=self._register_tokens,
+                        cls_pred_branch=self._cls_pred_branch,
+                        gan_ca_blocks=self._gan_ca_blocks,
+                        concat_time_embeddings=concat_time_embeddings
+                    )
+                    flow_pred = flow_pred.permute(0, 2, 1, 3, 4)
+                else:
+                    flow_pred = self._call_model(
+                        noisy_image_or_video.permute(0, 2, 1, 3, 4),
+                        t=input_timestep, context=prompt_embeds,
+                        seq_len=self.seq_len
+                    ).permute(0, 2, 1, 3, 4)
+
+        if rope_offset_was_set:
+            self.model.rope_temporal_offset = prev_rope_temporal_offset
+
+        pred_x0 = self._convert_flow_pred_to_x0(
+            flow_pred=flow_pred.flatten(0, 1),
+            xt=noisy_image_or_video.flatten(0, 1),
+            timestep=timestep.flatten(0, 1)
+        ).unflatten(0, flow_pred.shape[:2])
+
+        if logits is not None:
+            return flow_pred, pred_x0, logits
+
+        return flow_pred, pred_x0
+
+    def get_scheduler(self) -> SchedulerInterface:
+        """
+        Update the current scheduler with the interface's static method
+        """
+        scheduler = self.scheduler
+        scheduler.convert_x0_to_noise = types.MethodType(
+            SchedulerInterface.convert_x0_to_noise, scheduler)
+        scheduler.convert_noise_to_x0 = types.MethodType(
+            SchedulerInterface.convert_noise_to_x0, scheduler)
+        scheduler.convert_velocity_to_x0 = types.MethodType(
+            SchedulerInterface.convert_velocity_to_x0, scheduler)
+        self.scheduler = scheduler
+        return scheduler
+
+    def post_init(self):
+        """
+        A few custom initialization steps that should be called after the object is created.
+        Currently, the only one we have is to bind a few methods to scheduler.
+        We can gradually add more methods here if needed.
+        """
+        self.get_scheduler()
+
+
+_MG_LIGHTVAE_DEFAULT_PATHS = {
+    "mg_lightvae": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE.pth"),
+    "mg_lightvae_v2": os.path.join("wan_models", "Matrix-Game-3.0", "MG-LightVAE_v2.pth"),
+}
+
+
+def build_vae_5b(args):
+    """Return the 5B VAE wrapper requested by args.vae_type."""
+    vae_type = str(getattr(args, "vae_type", "wan")).lower().strip()
+
+    if vae_type in ("wan", "wan2.2", ""):
+        return WanVAEWrapper()
+
+    if vae_type in _MG_LIGHTVAE_DEFAULT_PATHS:
+        from utils.lightvae_5b_wrapper import LightVAE5BWrapper
+
+        return LightVAE5BWrapper(vae_path=_MG_LIGHTVAE_DEFAULT_PATHS[vae_type])
+
+    raise ValueError(
+        f"Unknown vae_type '{vae_type}'. "
+        "Expected one of: wan, mg_lightvae, mg_lightvae_v2."
+    )