Upload backend snapshot for sfp4 checkpoint-750

backend_snapshot/README.md

@@ -0,0 +1,35 @@
+# Backend snapshot for checkpoint-750
+
+This directory is the code snapshot for the training backend used by:
+
+`sfp4_v4_sparse09_hpo_on_ours_p_init2050_1n_interactive/checkpoint-750`
+
+Key runtime settings:
+
+- `FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_OURS_P_ATTN`
+- `FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O=1`
+- `VSA_SPARSITY=0.9`
+- `VSA_INIT_SPARSITY=0.9`
+- `VSA_WARMUP_STEPS=0`
+- tile size: `4 x 4 x 4 = 64` video tokens
+
+Important files:
+
+- `fastvideo/attention/backends/sparse_fp4_ours_p_attn.py`: Python attention backend, Q/K/V fake quantization, top-k block map, tile mean setup.
+- `fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ours_p.py`: PyTorch custom op and autograd wrapper.
+- `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton_ours_p.py`: Triton forward/backward kernel.
+- `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py`: FP4 quant/dequant utilities used by the kernel.
+- `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py`: Q/K/V fake quant kernels.
+- `fastvideo/attention/backends/video_sparse_attn.py`: VSA metadata and tile-size helper.
+- `fastvideo/platforms/interface.py` and `fastvideo/platforms/cuda.py`: backend enum and CUDA backend selection wiring.
+- `fastvideo/training/training_pipeline.py` and `fastvideo/training/wan_training_pipeline.py`: legacy SFT training path used by the launch script.
+- `scripts/training/run_sparse_fp4_train_v4_1n_sparse09_hpo_on_ours_p_init2050_interactive.sh`: exact Slurm wrapper for this run.
+- `scripts/training/run_sparse_fp4_train_v4_common.sh`: common SFT launch/resume script.
+
+Source repo HEAD when staged:
+
+`3f818d0fc532ec6494b465967d5f485150917d0c`
+
+Note: several backend files were uncommitted or locally modified when this
+snapshot was staged, so the files here are the authoritative copy for this
+checkpoint rather than the clean git commit alone.

backend_snapshot/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ours_p.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import os
+
+import torch
+
+
+def _use_high_prec_output_for_backward() -> bool:
+    value = os.environ.get("FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O", "1")
+    return value.lower() not in ("0", "false", "no", "off")
+
+
+def _map_to_index(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if block_map.dim() == 3:
+        block_map = block_map.unsqueeze(0)
+    if block_map.dim() != 4:
+        raise ValueError(
+            f"block_map must be [B,H,Q,KV] or [H,Q,KV], got {tuple(block_map.shape)}"
+        )
+    if block_map.dtype != torch.bool:
+        block_map = block_map.to(torch.bool)
+    if not block_map.is_cuda:
+        raise RuntimeError("block_map must be a CUDA tensor.")
+
+    try:
+        from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index
+    except Exception as e:
+        raise ImportError("Triton map_to_index is required for ours-P Sparse FP4.") from e
+    return triton_map_to_index(block_map)
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_ours_p_triton",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_ours_p_triton(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    block_map = block_map.to(torch.bool)
+    q2k_idx, q2k_num = _map_to_index(block_map)
+
+    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton_ours_p import (
+        triton_block_sparse_attn_forward,
+    )
+
+    return triton_block_sparse_attn_forward(
+        q, k, v, q2k_idx, q2k_num, variable_block_sizes, is_qat=True
+    )
+
+
+@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_ours_p_triton")
+def _block_sparse_attn_ours_p_triton_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    o = torch.empty_like(q)
+    high_prec_o = torch.empty_like(q)
+    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+    return o, M, high_prec_o
+
+
+@torch.library.custom_op(
+    "fastvideo_kernel::block_sparse_attn_ours_p_backward_triton",
+    mutates_args=(),
+    device_types="cuda",
+)
+def block_sparse_attn_ours_p_backward_triton(
+    grad_output: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    M: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    grad_output = grad_output.contiguous()
+    block_map = block_map.to(torch.bool)
+    q2k_idx, q2k_num = _map_to_index(block_map)
+    k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+
+    from fastvideo_kernel.triton_kernels.block_sparse_attn_triton_ours_p import (
+        triton_block_sparse_attn_backward,
+    )
+
+    return triton_block_sparse_attn_backward(
+        grad_output,
+        q,
+        k,
+        v,
+        o,
+        M,
+        q2k_idx,
+        q2k_num,
+        k2q_idx,
+        k2q_num,
+        variable_block_sizes,
+        is_qat=True,
+    )
+
+
+@torch.library.register_fake(
+    "fastvideo_kernel::block_sparse_attn_ours_p_backward_triton"
+)
+def _block_sparse_attn_ours_p_backward_triton_fake(
+    grad_output: torch.Tensor,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    M: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
+
+
+def _backward_triton(ctx, grad_o, grad_M, grad_high_prec_o):
+    q, k, v, o_for_bwd, M, block_map, variable_block_sizes = ctx.saved_tensors
+    dq, dk, dv = block_sparse_attn_ours_p_backward_triton(
+        grad_o, q, k, v, o_for_bwd, M, block_map, variable_block_sizes
+    )
+    return dq, dk, dv, None, None
+
+
+def _setup_context_triton(ctx, inputs, output):
+    q, k, v, block_map, variable_block_sizes = inputs
+    o, M, high_prec_o = output
+    o_for_bwd = high_prec_o if _use_high_prec_output_for_backward() else o
+    ctx.save_for_backward(q, k, v, o_for_bwd, M, block_map, variable_block_sizes)
+
+
+block_sparse_attn_ours_p_triton.register_autograd(
+    _backward_triton, setup_context=_setup_context_triton
+)
+
+
+class _BlockSparseAttnOursPTileComp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, q_mean, k_mean, v_mean, block_map, variable_block_sizes):
+        q = q.contiguous()
+        k = k.contiguous()
+        v = v.contiguous()
+        q_mean = q_mean.contiguous()
+        k_mean = k_mean.contiguous()
+        v_mean = v_mean.contiguous()
+        block_map = block_map.to(torch.bool)
+        dropped_block_map = torch.logical_not(block_map)
+
+        q2k_idx, q2k_num = _map_to_index(block_map)
+        dropped_q2k_idx, dropped_q2k_num = _map_to_index(dropped_block_map)
+
+        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton_ours_p import (
+            triton_block_sparse_attn_forward,
+        )
+
+        o, M, high_prec_o = triton_block_sparse_attn_forward(
+            q,
+            k,
+            v,
+            q2k_idx,
+            q2k_num,
+            variable_block_sizes,
+            is_qat=True,
+            q_mean=q_mean,
+            k_mean=k_mean,
+            v_mean=v_mean,
+            dropped_q2k_index=dropped_q2k_idx,
+            dropped_q2k_num=dropped_q2k_num,
+        )
+        o_for_bwd = high_prec_o if _use_high_prec_output_for_backward() else o
+        ctx.save_for_backward(
+            q,
+            k,
+            v,
+            q_mean,
+            k_mean,
+            v_mean,
+            o_for_bwd,
+            M,
+            block_map,
+            dropped_block_map,
+            variable_block_sizes,
+        )
+        return o, M
+
+    @staticmethod
+    def backward(ctx, grad_o, grad_M):
+        (
+            q,
+            k,
+            v,
+            q_mean,
+            k_mean,
+            v_mean,
+            o_for_bwd,
+            M,
+            block_map,
+            dropped_block_map,
+            variable_block_sizes,
+        ) = ctx.saved_tensors
+
+        q2k_idx, q2k_num = _map_to_index(block_map)
+        k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous())
+        dropped_q2k_idx, dropped_q2k_num = _map_to_index(dropped_block_map)
+        dropped_k2q_idx, dropped_k2q_num = _map_to_index(
+            dropped_block_map.transpose(-1, -2).contiguous()
+        )
+
+        from fastvideo_kernel.triton_kernels.block_sparse_attn_triton_ours_p import (
+            triton_block_sparse_attn_backward,
+        )
+
+        dq, dk, dv = triton_block_sparse_attn_backward(
+            grad_o.contiguous(),
+            q,
+            k,
+            v,
+            o_for_bwd,
+            M,
+            q2k_idx,
+            q2k_num,
+            k2q_idx,
+            k2q_num,
+            variable_block_sizes,
+            is_qat=True,
+            q_mean=q_mean,
+            k_mean=k_mean,
+            v_mean=v_mean,
+            dropped_q2k_index=dropped_q2k_idx,
+            dropped_q2k_num=dropped_q2k_num,
+            dropped_k2q_index=dropped_k2q_idx,
+            dropped_k2q_num=dropped_k2q_num,
+        )
+        return dq, dk, dv, None, None, None, None, None
+
+
+def block_sparse_attn_ours_p(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+    q_mean: torch.Tensor | None = None,
+    k_mean: torch.Tensor | None = None,
+    v_mean: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if (q_mean is not None) or (k_mean is not None) or (v_mean is not None):
+        if q_mean is None or k_mean is None or v_mean is None:
+            raise ValueError("q_mean, k_mean, and v_mean must be provided together")
+        return _BlockSparseAttnOursPTileComp.apply(
+            q, k, v, q_mean, k_mean, v_mean, block_map, variable_block_sizes
+        )
+
+    o, M, _ = block_sparse_attn_ours_p_triton(
+        q, k, v, block_map, variable_block_sizes
+    )
+    return o, M

backend_snapshot/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton_ours_p.py

@@ -0,0 +1,1155 @@
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+
+Credits: OpenAI kernel team
+"""
+
+import torch
+import triton
+import triton.language as tl
+from .quant_utils import fake_quantize
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+import math  # small utility needed by the sparse wrapper
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+configs = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [64]\
+    for BN in [64]\
+    for s in [3, 4, 7]\
+    for w in [4, 8]\
+]
+
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd_sparse(
+        Q,
+        K,
+        V,
+        QMean,
+        KMean,
+        VMean,
+        sm_scale,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,  #
+        dropped_q2k_index,
+        dropped_q2k_num,
+        max_dropped_kv_blks,  #
+        variable_block_sizes,
+        M,
+        Out,  #
+        HighPrecOut,  #
+        stride_qz,
+        stride_qh,
+        stride_qm,
+        stride_qk,
+        stride_kz,
+        stride_kh,
+        stride_kn,
+        stride_kk,
+        stride_vz,
+        stride_vh,
+        stride_vk,
+        stride_vn,
+        stride_oz,
+        stride_oh,
+        stride_om,
+        stride_on,
+        Z,
+        H,
+        N_CTX_Q,  #
+        N_CTX_KV,  #
+        HEAD_DIM: tl.constexpr,  #
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        STAGE: tl.constexpr,
+        IS_QAT: tl.constexpr = False,
+        USE_TILE_COMP: tl.constexpr = False):
+    """
+    64x64 block-sparse forward kernel for the independent "ours P quant" path.
+
+    P quantization is group-local: each selected KV tile quantizes
+    exp2(logit - tile_row_max), then applies exp2(tile_row_max - online_max)
+    after the FP4 PV GEMM. This intentionally differs from the QAT-style
+    backend, which quantizes exp2(logit - online_max) directly.
+    """
+
+    # ----- program-id mapping -----
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(1)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    q_tiles = N_CTX_Q // BLOCK_M
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    dropped_kv_blocks = tl.load(dropped_q2k_num + meta_base)
+    dropped_kv_ptr = dropped_q2k_index + meta_base * max_dropped_kv_blks
+
+    # ----- base pointers -----
+    q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh)
+
+    Q_ptr = tl.make_block_ptr(base=Q + q_off,
+                              shape=(N_CTX_Q, HEAD_DIM),
+                              strides=(stride_qm, stride_qk),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
+
+    K_base = tl.make_block_ptr(base=K + k_off,
+                               shape=(HEAD_DIM, N_CTX_KV),
+                               strides=(stride_kk, stride_kn),
+                               offsets=(0, 0),
+                               block_shape=(HEAD_DIM, BLOCK_N),
+                               order=(0, 1))
+
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1,
+                                                                              0)
+    V_base = tl.make_block_ptr(base=V + v_off,
+                               shape=(N_CTX_KV, HEAD_DIM),
+                               strides=(stride_vk, stride_vn),
+                               offsets=(0, 0),
+                               block_shape=(BLOCK_N, HEAD_DIM),
+                               order=v_order)
+
+    O_ptr = tl.make_block_ptr(base=Out + o_off,
+                              shape=(N_CTX_Q, HEAD_DIM),
+                              strides=(stride_om, stride_on),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
+    HPO_ptr = tl.make_block_ptr(base=HighPrecOut + o_off,
+                                shape=(N_CTX_Q, HEAD_DIM),
+                                strides=(stride_om, stride_on),
+                                offsets=(q_blk * BLOCK_M, 0),
+                                block_shape=(BLOCK_M, HEAD_DIM),
+                                order=(1, 0))
+
+    # ----- accumulators -----
+    offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    high_prec_acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    qk_scale = sm_scale * 1.44269504  # 1/ln2
+    q = tl.load(Q_ptr)
+    offs_d = tl.arange(0, HEAD_DIM)
+
+    # ----- sparse loop over valid K/V tiles -----
+    for i in range(0, kv_blocks):
+        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
+        block_size = tl.load(variable_block_sizes + kv_idx)
+        K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
+        V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
+
+        k = tl.load(K_ptr)
+        mask = tl.arange(0, BLOCK_N) < block_size
+        qk = tl.dot(q, k) * qk_scale
+        # mask out invalid columns
+        qk = tl.where(mask[None, :], qk, -float("inf"))
+        group_m = tl.max(qk, 1)
+        m_ij = tl.maximum(m_i, group_m)
+
+        p_local = tl.math.exp2(qk - group_m[:, None])
+        p_local = tl.where(mask[None, :], p_local, 0.0)
+        p_comp = tl.math.exp2(group_m - m_ij)
+        p_valid = mask[None, :] & (
+            tl.full(shape=p_local.shape, value=1.0,
+                    dtype=p_local.dtype) == 1.0
+        )
+        p_quant, high_prec_p = fake_quantize(
+            src_tensor=p_local, valid_src_mask=p_valid,
+            BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=BLOCK_N,
+            dst_dtype=tl.bfloat16, use_global_sf=False,
+        )
+        l_ij = tl.sum(high_prec_p, 1) * p_comp
+
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+        high_prec_acc = high_prec_acc * alpha[:, None]
+
+        v = tl.load(V_ptr)
+        acc = acc + tl.dot(
+            p_quant.to(tl.bfloat16),
+            v.to(tl.bfloat16),
+        ) * p_comp[:, None]
+        high_prec_acc = high_prec_acc + tl.dot(
+            high_prec_p.to(tl.bfloat16),
+            v.to(tl.bfloat16),
+        ) * p_comp[:, None]
+        m_i = m_ij
+
+    if USE_TILE_COMP:
+        q_mean_base = (off_hz * q_tiles + q_blk).to(tl.int64) * HEAD_DIM
+        q_mean = tl.load(QMean + q_mean_base + offs_d).to(tl.float32)
+        kv_tiles = N_CTX_KV // BLOCK_N
+
+        for i in range(0, dropped_kv_blocks):
+            kv_idx = tl.load(dropped_kv_ptr + i).to(tl.int32)
+            block_size = tl.load(variable_block_sizes + kv_idx).to(tl.float32)
+            kv_mean_base = (off_hz * kv_tiles + kv_idx).to(tl.int64) * HEAD_DIM
+            k_mean = tl.load(KMean + kv_mean_base + offs_d).to(tl.float32)
+            v_mean = tl.load(VMean + kv_mean_base + offs_d).to(tl.float32)
+
+            score = tl.sum(q_mean * k_mean, axis=0) * qk_scale
+            m_ij = tl.maximum(m_i, score)
+            alpha = tl.math.exp2(m_i - m_ij)
+            beta = tl.math.exp2(score - m_ij)
+
+            l_i = l_i * alpha + block_size * beta
+            comp = (block_size * beta)[:, None] * v_mean[None, :]
+            acc = acc * alpha[:, None] + comp
+            high_prec_acc = high_prec_acc * alpha[:, None] + comp
+            m_i = m_ij
+
+    # ----- epilogue -----
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    high_prec_acc = high_prec_acc / l_i[:, None]
+    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
+    tl.store(O_ptr, acc.to(Out.type.element_ty))
+    tl.store(HPO_ptr, high_prec_acc.to(HighPrecOut.type.element_ty))
+
+
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.jit
+def _attn_bwd_preprocess(
+        O,
+        DO,  #
+        Delta,  #
+        Z,
+        H,
+        N_CTX,  #
+        BLOCK_M: tl.constexpr,
+        HEAD_DIM: tl.constexpr  #
+):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
+                off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM +
+                 off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        QMean,
+        KMean,
+        VMean,
+        sm_scale,  #
+        DO,  #
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        dropped_k2q_index,
+        dropped_k2q_num,
+        max_dropped_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX_KV,
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,  #
+        # Filled in by the wrapper.
+    start_n,
+        start_m,
+        num_steps,
+        IS_QAT: tl.constexpr = False,
+        USE_TILE_COMP: tl.constexpr = False):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    step_m = BLOCK_M1
+    kv_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    kv_tiles = N_CTX_KV // BLOCK_N1
+    meta_base = ((b * H + h) * kv_tiles + kv_blk)
+
+    q_blocks = tl.load(k2q_num + meta_base)  # int32
+    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
+    dropped_q_blocks = tl.load(dropped_k2q_num + meta_base)
+    dropped_q_ptr = dropped_k2q_index + meta_base * max_dropped_q_blks
+    block_size = tl.load(variable_block_sizes + kv_blk)
+    block_size_f = block_size.to(tl.float32)
+
+    for blk_idx in range(q_blocks * 2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * step_m
+        qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k.to(tl.bfloat16), qT)
+        qkT = qkT * sm_scale * 1.44269504
+        mask = tl.arange(0, BLOCK_N1) < block_size
+        qkT = tl.where(mask[:, None], qkT, -float("inf"))
+        group_m = tl.max(qkT, 0)
+        pT = tl.math.exp2(qkT - m[None, :])
+        pT = tl.where(mask[:, None], pT, 0.0)
+
+        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
+        # Compute dV with group-local P quantization:
+        # quantize exp2(logit - tile_col_max), then multiply dO by
+        # exp2(tile_col_max - final_lse) to recover the final softmax scale.
+        p_local_T = tl.math.exp2(qkT - group_m[None, :])
+        p_local_T = tl.where(mask[:, None], p_local_T, 0.0)
+        p_comp = tl.math.exp2(group_m - m)
+        p_for_quant = tl.trans(p_local_T)
+        p_valid = mask[None, :] & (
+            tl.full(
+                shape=p_for_quant.shape,
+                value=1.0,
+                dtype=p_for_quant.dtype,
+            ) == 1.0
+        )
+        p_quant, _ = fake_quantize(
+            src_tensor=p_for_quant, valid_src_mask=p_valid,
+            BLOCK_SIZE_OUT_DIM=BLOCK_M1, BLOCK_SIZE_QUANT_DIM=BLOCK_N1,
+            dst_dtype=p_for_quant.dtype, use_global_sf=False,
+        )
+        dv += tl.dot(
+            tl.trans(p_quant.to(tl.bfloat16)),
+            (do * p_comp[:, None]).to(tl.bfloat16),
+        )
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.bfloat16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+
+    if USE_TILE_COMP:
+        k_mean = tl.load(KMean + kv_blk * HEAD_DIM + offs_k).to(tl.float32)
+        v_mean = tl.load(VMean + kv_blk * HEAD_DIM + offs_k).to(tl.float32)
+        qk_scale = sm_scale * 1.44269504
+
+        for blk_idx in range(dropped_q_blocks * 2):
+            q_blk_idx = tl.load(dropped_q_ptr + blk_idx // 2).to(tl.int32)
+            half = (blk_idx % 2).to(tl.int32)
+            block_sparse_offset = (q_blk_idx * 2 + half) * step_m
+            offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+            q_mean = tl.load(QMean + q_blk_idx * HEAD_DIM +
+                             offs_k).to(tl.float32)
+            m = tl.load(M + offs_m)
+            do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
+            Di = tl.load(D + offs_m)
+            q_block_size = tl.load(variable_block_sizes +
+                                   q_blk_idx).to(tl.float32)
+
+            score = tl.sum(q_mean * k_mean, axis=0) * qk_scale
+            p = tl.math.exp2(score - m)
+            dp = tl.sum(do.to(tl.float32) * v_mean[None, :], axis=1)
+            ds = block_size_f * p * (dp - Di)
+
+            dk_mean = tl.sum(ds[:, None] * q_mean[None, :],
+                             axis=0) / block_size_f
+            dv_mean = tl.sum(p[:, None] * do.to(tl.float32), axis=0)
+            dk += dk_mean[None, :]
+            dv += dv_mean[None, :]
+    return dk, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        QMean,
+        KMean,
+        VMean,
+        do,
+        m,
+        m_vec,
+        D,
+        # shared by Q/K/V/DO.
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        dropped_q2k_index,
+        dropped_q2k_num,
+        max_dropped_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        # Filled in by the wrapper.
+        start_m,
+        start_n,
+        num_steps,
+        sm_scale=1.0,
+        IS_QAT: tl.constexpr = False,
+        USE_TILE_COMP: tl.constexpr = False):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    step_n = BLOCK_N2
+
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    q_tiles = N_CTX // BLOCK_M2
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    dropped_kv_blocks = tl.load(dropped_q2k_num + meta_base)
+    dropped_kv_ptr = dropped_q2k_index + meta_base * max_dropped_kv_blks
+
+    for blk_idx in range(kv_blocks * 2):
+        kv_idx = tl.load(kv_ptr + blk_idx // 2).to(tl.int32)
+        # variable_block_sizes is defined per KV block (tile). Mask must therefore
+        # use kv_idx (not q_blk). Also, because we split each 64-token block into
+        # two 32-token halves, the mask must account for the half-block offset.
+        block_size = tl.load(variable_block_sizes + kv_idx).to(tl.int32)
+        half = (blk_idx % 2).to(tl.int32)
+        block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
+        kT = tl.load(kT_ptrs + block_sparse_offset)
+        vT = tl.load(vT_ptrs + block_sparse_offset)
+        qk = tl.dot(q, kT)
+        qk = qk * sm_scale * 1.44269504
+        p = tl.math.exp2(qk - m)
+        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
+        mask = offs_in_block < block_size
+        p = tl.where(mask[None, :], p, 0.0)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.bfloat16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+
+    if USE_TILE_COMP:
+        q_mean = tl.load(QMean + q_blk * HEAD_DIM + offs_k).to(tl.float32)
+        q_block_size = tl.load(variable_block_sizes + q_blk).to(tl.float32)
+        qk_scale = sm_scale * 1.44269504
+        dq_mean = tl.zeros([HEAD_DIM], dtype=tl.float32)
+
+        for blk_idx in range(dropped_kv_blocks):
+            kv_idx = tl.load(dropped_kv_ptr + blk_idx).to(tl.int32)
+            block_size = tl.load(variable_block_sizes + kv_idx).to(tl.float32)
+            k_mean = tl.load(KMean + kv_idx * HEAD_DIM +
+                             offs_k).to(tl.float32)
+            v_mean = tl.load(VMean + kv_idx * HEAD_DIM +
+                             offs_k).to(tl.float32)
+
+            score = tl.sum(q_mean * k_mean, axis=0) * qk_scale
+            p = tl.math.exp2(score - m_vec)
+            dp = tl.sum(do.to(tl.float32) * v_mean[None, :], axis=1)
+            ds = block_size * p * (dp - Di)
+            dq_mean = dq_mean + tl.sum(ds, axis=0) * k_mean
+
+        dq += dq_mean[None, :] / q_block_size
+    return dq
+
+
+@triton.jit
+def _attn_bwd(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        DO,  #
+        DQ,
+        DK,
+        DV,  #
+        M,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_z,
+        stride_h,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        IS_QAT: tl.constexpr = False):
+    LN2 = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    num_steps = N_CTX // BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(  #
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        DO,  #
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1,
+        BLOCK_N1,
+        HEAD_DIM,  #
+        start_n,
+        start_m,
+        num_steps,  #
+        IS_QAT=IS_QAT,
+        USE_TILE_COMP=False,
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = 0
+
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m_vec = tl.load(M + offs_m)
+    m = m_vec[:, None]
+
+    num_steps = N_CTX // BLOCK_N2
+    dq = _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        Q,
+        K,
+        V,
+        do,
+        m,
+        m_vec,
+        D,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2,
+        BLOCK_N2,
+        HEAD_DIM,  #
+        start_m,
+        end_n,
+        num_steps,  #
+        sm_scale=sm_scale,
+        IS_QAT=IS_QAT,
+        USE_TILE_COMP=False,
+    )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= sm_scale
+    tl.store(dq_ptrs, dq)
+
+
+@triton.jit
+def _attn_bwd_dkdv_kernel(
+        Q,
+        K,
+        V,
+        QMean,
+        KMean,
+        VMean,
+        sm_scale,  #
+        DO,  #
+        DK,
+        DV,  #
+        M,
+        D,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        dropped_k2q_index,
+        dropped_k2q_num,
+        max_dropped_q_blks,
+        variable_block_sizes,
+        # shared token/dim strides (assumed contiguous along token and dim)
+        stride_tok,
+        stride_d,  #
+        # batch/head strides (may differ between Q and KV)
+        stride_qz,
+        stride_qh,
+        stride_kz,
+        stride_kh,
+        stride_vz,
+        stride_vh,
+        stride_doz,
+        stride_doh,
+        stride_dkz,
+        stride_dkh,
+        stride_dvz,
+        stride_dvh,
+        H,
+        N_CTX_Q,
+        N_CTX_KV,
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        IS_QAT: tl.constexpr = False,
+        USE_TILE_COMP: tl.constexpr = False):
+    """
+    Backward kernel that computes dK and dV for each KV block (64 tokens).
+    Grid:
+      pid0: kv_blk in [0, N_CTX_KV/BLOCK_N1)
+      pid2: fused (batch, head) in [0, B*H)
+    """
+    bhid = tl.program_id(2)
+    b = bhid // H
+    h = bhid % H
+    kv_blk = tl.program_id(0)
+
+    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
+    dk_adj = (b.to(tl.int64) * stride_dkz + h.to(tl.int64) * stride_dkh)
+    dv_adj = (b.to(tl.int64) * stride_dvz + h.to(tl.int64) * stride_dvh)
+
+    Q = Q + q_adj
+    K = K + kv_adj_k
+    V = V + kv_adj_v
+    DO = DO + do_adj
+    DK = DK + dk_adj
+    DV = DV + dv_adj
+
+    q_tiles = N_CTX_Q // BLOCK_M1 // 2
+    kv_tiles = N_CTX_KV // BLOCK_N1
+    mean_q_adj = (bhid * q_tiles * HEAD_DIM).to(tl.int64)
+    mean_kv_adj = (bhid * kv_tiles * HEAD_DIM).to(tl.int64)
+    QMean = QMean + mean_q_adj
+    KMean = KMean + mean_kv_adj
+    VMean = VMean + mean_kv_adj
+
+    # M and D (delta) are always sized by Q length.
+    M = M + (bhid * N_CTX_Q).to(tl.int64)
+    D = D + (bhid * N_CTX_Q).to(tl.int64)
+
+    offs_k = tl.arange(0, HEAD_DIM)
+    start_n = kv_blk * BLOCK_N1
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    dv_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    num_steps = N_CTX_Q // BLOCK_M1
+    dk_acc, dv_acc = _attn_bwd_dkdv(
+        dk_acc,
+        dv_acc,
+        Q,
+        k,
+        v,
+        QMean,
+        KMean,
+        VMean,
+        sm_scale,
+        DO,
+        M,
+        D,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        dropped_k2q_index,
+        dropped_k2q_num,
+        max_dropped_q_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,
+        H,
+        N_CTX_KV,
+        BLOCK_M1=BLOCK_M1,
+        BLOCK_N1=BLOCK_N1,
+        HEAD_DIM=HEAD_DIM,
+        start_n=start_n,
+        start_m=0,
+        num_steps=num_steps,
+        IS_QAT=IS_QAT,
+        USE_TILE_COMP=USE_TILE_COMP,
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv_acc)
+
+    dk_acc *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk_acc)
+
+
+@triton.jit
+def _attn_bwd_dq_kernel(
+        Q,
+        K,
+        V,
+        QMean,
+        KMean,
+        VMean,
+        DO,  #
+        DQ,
+        M,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        dropped_q2k_index,
+        dropped_q2k_num,
+        max_dropped_kv_blks,
+        variable_block_sizes,
+        # shared token/dim strides (assumed contiguous along token and dim)
+        stride_tok,
+        stride_d,  #
+        # batch/head strides (may differ between Q and KV)
+        stride_qz,
+        stride_qh,
+        stride_kz,
+        stride_kh,
+        stride_vz,
+        stride_vh,
+        stride_doz,
+        stride_doh,
+        stride_dqz,
+        stride_dqh,
+        H,
+        N_CTX_Q,
+        sm_scale,
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        IS_QAT: tl.constexpr = False,
+        USE_TILE_COMP: tl.constexpr = False):
+    """
+    Backward kernel that computes dQ for each Q block (64 tokens).
+    Grid:
+      pid0: q_blk in [0, N_CTX_Q/BLOCK_M2)
+      pid2: fused (batch, head) in [0, B*H)
+    """
+    LN2 = 0.6931471824645996  # = ln(2)
+    bhid = tl.program_id(2)
+    b = bhid // H
+    h = bhid % H
+    q_blk = tl.program_id(0)
+
+    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
+    dq_adj = (b.to(tl.int64) * stride_dqz + h.to(tl.int64) * stride_dqh)
+
+    Q = Q + q_adj
+    K = K + kv_adj_k
+    V = V + kv_adj_v
+    DO = DO + do_adj
+    DQ = DQ + dq_adj
+
+    q_tiles = N_CTX_Q // BLOCK_M2
+    kv_tiles = N_CTX_Q // 64
+    mean_q_adj = (bhid * q_tiles * HEAD_DIM).to(tl.int64)
+    mean_kv_adj = (bhid * kv_tiles * HEAD_DIM).to(tl.int64)
+    QMean = QMean + mean_q_adj
+    KMean = KMean + mean_kv_adj
+    VMean = VMean + mean_kv_adj
+
+    M = M + (bhid * N_CTX_Q).to(tl.int64)
+    D = D + (bhid * N_CTX_Q).to(tl.int64)
+
+    offs_k = tl.arange(0, HEAD_DIM)
+    start_m = q_blk * BLOCK_M2
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    m_vec = tl.load(M + offs_m)
+    m = m_vec[:, None]
+
+    dq_acc = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    num_steps = 0  # unused in _attn_bwd_dq
+    dq_acc = _attn_bwd_dq(
+        dq_acc,
+        q,
+        K,
+        V,
+        QMean,
+        KMean,
+        VMean,
+        do,
+        m,
+        m_vec,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        dropped_q2k_index,
+        dropped_q2k_num,
+        max_dropped_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,
+        H,
+        N_CTX_Q,
+        BLOCK_M2=BLOCK_M2,
+        BLOCK_N2=BLOCK_N2,
+        HEAD_DIM=HEAD_DIM,
+        start_m=start_m,
+        start_n=0,
+        num_steps=num_steps,
+        sm_scale=sm_scale,
+        IS_QAT=IS_QAT,
+        USE_TILE_COMP=USE_TILE_COMP,
+    )
+
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq_acc *= sm_scale
+    tl.store(dq_ptrs, dq_acc)
+
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
+                                     variable_block_sizes, is_qat=False,
+                                     q_mean=None, k_mean=None, v_mean=None,
+                                     dropped_q2k_index=None,
+                                     dropped_q2k_num=None):
+    B, H, Tq, D = q.shape
+    Tkv = k.shape[2]
+    sm_scale = 1.0 / math.sqrt(D)
+    max_kv_blks = q2k_index.shape[-1]
+    use_tile_comp = q_mean is not None
+    if use_tile_comp:
+        assert k_mean is not None and v_mean is not None
+        assert dropped_q2k_index is not None and dropped_q2k_num is not None
+        q_mean = q_mean.contiguous()
+        k_mean = k_mean.contiguous()
+        v_mean = v_mean.contiguous()
+        max_dropped_kv_blks = dropped_q2k_index.shape[-1]
+    else:
+        q_mean = q
+        k_mean = k
+        v_mean = v
+        dropped_q2k_index = q2k_index
+        dropped_q2k_num = q2k_num
+        max_dropped_kv_blks = max_kv_blks
+    assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}"
+    assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}"
+    assert q2k_num.shape[
+        -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
+    assert variable_block_sizes.numel() == Tkv // 64, (
+        f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
+        f"got {variable_block_sizes.numel()}"
+    )
+    o = torch.empty_like(q)
+    high_prec_o = torch.empty_like(q)
+    M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
+
+    grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
+    _attn_fwd_sparse[grid](q,
+                           k,
+                           v,
+                           q_mean,
+                           k_mean,
+                           v_mean,
+                           sm_scale,
+                           q2k_index,
+                           q2k_num,
+                           max_kv_blks,
+                           dropped_q2k_index,
+                           dropped_q2k_num,
+                           max_dropped_kv_blks,
+                           variable_block_sizes,
+                           M,
+                           o,
+                           high_prec_o,
+                           q.stride(0),
+                           q.stride(1),
+                           q.stride(2),
+                           q.stride(3),
+                           k.stride(0),
+                           k.stride(1),
+                           k.stride(2),
+                           k.stride(3),
+                           v.stride(0),
+                           v.stride(1),
+                           v.stride(2),
+                           v.stride(3),
+                           o.stride(0),
+                           o.stride(1),
+                           o.stride(2),
+                           o.stride(3),
+                           B,
+                           H,
+                           Tq,
+                           Tkv,
+                           HEAD_DIM=D,
+                           STAGE=3,
+                           IS_QAT=is_qat,
+                           USE_TILE_COMP=use_tile_comp)
+
+    return o, M, high_prec_o
+
+
+def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
+                                      k2q_index, k2q_num, variable_block_sizes,
+                                      is_qat=False, q_mean=None, k_mean=None,
+                                      v_mean=None, dropped_q2k_index=None,
+                                      dropped_q2k_num=None,
+                                      dropped_k2q_index=None,
+                                      dropped_k2q_num=None):
+    assert do.is_contiguous()
+
+    B, H, Tq, D = q.shape
+    Tkv = k.shape[2]
+    sm_scale = 1.0 / math.sqrt(D)
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    BATCH, N_HEAD = q.shape[:2]
+    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+    RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+    # Ours-P mode keeps K unscaled and applies sm_scale inside the bwd kernels.
+    arg_k = k
+    PRE_BLOCK = 64
+    assert Tq % PRE_BLOCK == 0
+    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
+    delta = torch.empty_like(M)
+    _attn_bwd_preprocess[pre_grid](
+        o,
+        do,  #
+        delta,  #
+        BATCH,
+        N_HEAD,
+        Tq,  #
+        BLOCK_M=PRE_BLOCK,
+        HEAD_DIM=D  #
+    )
+
+    max_q_blks = k2q_index.shape[-1]
+    max_kv_blks = q2k_index.shape[-1]
+    use_tile_comp = q_mean is not None
+    if use_tile_comp:
+        assert k_mean is not None and v_mean is not None
+        assert dropped_q2k_index is not None and dropped_q2k_num is not None
+        assert dropped_k2q_index is not None and dropped_k2q_num is not None
+        q_mean = q_mean.contiguous()
+        k_mean = k_mean.contiguous()
+        v_mean = v_mean.contiguous()
+        max_dropped_kv_blks = dropped_q2k_index.shape[-1]
+        max_dropped_q_blks = dropped_k2q_index.shape[-1]
+    else:
+        q_mean = q
+        k_mean = k
+        v_mean = v
+        dropped_q2k_index = q2k_index
+        dropped_q2k_num = q2k_num
+        dropped_k2q_index = k2q_index
+        dropped_k2q_num = k2q_num
+        max_dropped_kv_blks = max_kv_blks
+        max_dropped_q_blks = max_q_blks
+
+    # dK/dV kernel: grid over KV blocks
+    grid_kv = (Tkv // BLOCK_N1, 1, BATCH * N_HEAD)
+    _attn_bwd_dkdv_kernel[grid_kv](
+        q,
+        arg_k,
+        v,
+        q_mean,
+        k_mean,
+        v_mean,
+        sm_scale,
+        do,
+        dk,
+        dv,
+        M,
+        delta,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        dropped_k2q_index,
+        dropped_k2q_num,
+        max_dropped_q_blks,
+        variable_block_sizes,
+        q.stride(2),
+        q.stride(3),
+        q.stride(0),
+        q.stride(1),
+        arg_k.stride(0),
+        arg_k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        do.stride(0),
+        do.stride(1),
+        dk.stride(0),
+        dk.stride(1),
+        dv.stride(0),
+        dv.stride(1),
+        N_HEAD,
+        Tq,
+        Tkv,
+        BLOCK_M1=BLOCK_M1,
+        BLOCK_N1=BLOCK_N1,
+        HEAD_DIM=D,
+        IS_QAT=is_qat,
+        USE_TILE_COMP=use_tile_comp,
+    )
+
+    # dQ kernel: grid over Q blocks
+    grid_q = (Tq // BLOCK_M2, 1, BATCH * N_HEAD)
+    _attn_bwd_dq_kernel[grid_q](
+        q,
+        arg_k,
+        v,
+        q_mean,
+        k_mean,
+        v_mean,
+        do,
+        dq,
+        M,
+        delta,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        dropped_q2k_index,
+        dropped_q2k_num,
+        max_dropped_kv_blks,
+        variable_block_sizes,
+        q.stride(2),
+        q.stride(3),
+        q.stride(0),
+        q.stride(1),
+        arg_k.stride(0),
+        arg_k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        do.stride(0),
+        do.stride(1),
+        dq.stride(0),
+        dq.stride(1),
+        N_HEAD,
+        Tq,
+        sm_scale,
+        BLOCK_M2=BLOCK_M2,
+        BLOCK_N2=BLOCK_N2,
+        HEAD_DIM=D,
+        IS_QAT=is_qat,
+        USE_TILE_COMP=use_tile_comp,
+    )
+
+    return dq, dk, dv

backend_snapshot/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/numerics_details/mxfp_details/_upcast_from_mxfp.py
+# and https://github.com/triton-lang/triton/blob/main/python/triton_kernels/triton_kernels/numerics_details/mxfp_details/_downcast_to_mxfp.py
+
+import triton
+import triton.language as tl
+try:
+    from triton.language.target_info import cuda_capability_geq
+    _HAS_CAPABILITY_CHECK = True
+except ImportError:
+    cuda_capability_geq = None
+    _HAS_CAPABILITY_CHECK = False
+
+MXFP_BLOCK_SIZE = tl.constexpr(16)
+
+@triton.jit
+def _compute_quant_and_scale(
+    src_tensor,
+    valid_src_mask,
+    mx_tensor_dtype: tl.constexpr = tl.uint8,
+    use_global_sf=True,
+    two_level_quant_P=False,
+    IS_BLACKWELL: tl.constexpr = False,
+):
+    BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0]
+    BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1]
+    BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // MXFP_BLOCK_SIZE
+    is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8
+
+    is_fp8e4: tl.constexpr = mx_tensor_dtype == tl.float8e4nv
+    is_fp8e5: tl.constexpr = mx_tensor_dtype == tl.float8e5
+    tl.static_assert(
+        is_fp4 or (is_fp8e4 or is_fp8e5),
+        "mx_tensor_dtype must be uint8, float8e4nv, or float8e5",
+    )
+
+    # Explicit cast to fp32 since most ops are not supported on bfloat16. We avoid needless conversions to and from bf16
+    f32_tensor = src_tensor.to(tl.float32)
+    abs_tensor = tl.abs(f32_tensor)
+    abs_tensor = tl.where(valid_src_mask, abs_tensor, -1.0)  # Don't consider padding tensors in scale computation
+    
+    if two_level_quant_P:
+        # row max from SageAttn3 paper
+        global_max_val = tl.max(f32_tensor, axis=1, keep_dims=True)  # (BLOCK_SIZE_OUT_DIM, 1)
+        global_max_val = tl.maximum(global_max_val, 1e-8)
+        s_enc = ((6 * 448) / global_max_val).reshape([BLOCK_SIZE_OUT_DIM, 1, 1])
+        s_dec = (1 / s_enc)
+
+    abs_tensor = tl.reshape(abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
+
+    if use_global_sf and not two_level_quant_P:
+        global_max_val = tl.max(abs_tensor)
+        # Avoid division by zero: if all values are padding (max is 0), use a default scale
+        global_max_val = tl.maximum(global_max_val, 1e-8)
+        s_enc = (6 * 448) / global_max_val
+        s_dec = (1 / s_enc)
+    elif not two_level_quant_P and not use_global_sf:
+        s_dec = 1.0
+        s_enc = 1.0
+    
+    max_val = tl.max(abs_tensor, axis=2, keep_dims=True)  # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)  # per block maxima
+    s_dec_b = max_val / 6  # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
+    s_dec_b_e4m3 = (s_dec_b * s_enc).to(tl.float8e4nv)  # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
+    s_enc_b = 1 / (s_dec_b_e4m3.to(tl.float32) * s_dec)  # (BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1)
+
+    f32_tensor = tl.reshape(f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
+    quant_tensor = f32_tensor * s_enc_b
+
+    # Reshape the tensors after scaling
+    quant_tensor = quant_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM])
+    # Set the invalid portions of the tensor to 0. This will ensure that any padding tensors are 0 in the mx format.
+    quant_tensor = tl.where(valid_src_mask, quant_tensor, 0.0)
+    dequant_scale = s_dec_b_e4m3.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE])
+
+    if is_fp4 and IS_BLACKWELL:
+        # Convert scaled values to two f32 lanes and use PTX cvt to e2m1x2 with two f32 operands.
+        pairs = tl.reshape(quant_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2])
+        lo_f, hi_f = tl.split(pairs)
+        lo_f32 = lo_f.to(tl.float32)
+        hi_f32 = hi_f.to(tl.float32)
+
+        # Inline PTX: cvt.rn.satfinite.e2m1x2.f32 takes two f32 sources and produces one .b8 packed e2m1x2.
+        out_tensor = tl.inline_asm_elementwise(
+            """
+            {
+                .reg .b8 r;
+                cvt.rn.satfinite.e2m1x2.f32 r, $1, $2;
+                mov.b32 $0, {r, r, r, r};
+            }
+            """,
+            constraints="=r,f,f",
+            args=[hi_f32, lo_f32],
+            dtype=tl.uint8,
+            is_pure=True,
+            pack=1,
+        )
+    elif is_fp4:
+        quant_tensor = quant_tensor.to(tl.uint32, bitcast=True)
+        signs = quant_tensor & 0x80000000
+        exponents = (quant_tensor >> 23) & 0xFF
+        mantissas_orig = (quant_tensor & 0x7FFFFF)
+
+        # For RTNE: 0.25 < x < 0.75 maps to 0.5 (denormal); exactly 0.25 maps to 0.0
+        E8_BIAS = 127
+        E2_BIAS = 1
+        # Move implicit bit 1 at the beginning to mantissa for denormals
+        is_subnormal = exponents < E8_BIAS
+        adjusted_exponents = tl.core.sub(E8_BIAS, exponents + 1, sanitize_overflow=False)
+        mantissas_pre = (0x400000 | (mantissas_orig >> 1))
+        mantissas = tl.where(is_subnormal, mantissas_pre >> adjusted_exponents, mantissas_orig)
+
+        # For normal numbers, we change the bias from 127 to 1, and for subnormals, we keep exponent as 0.
+        exponents = tl.maximum(exponents, E8_BIAS - E2_BIAS) - (E8_BIAS - E2_BIAS)
+
+        # Combine sign, exponent, and mantissa, while saturating
+        # Round to nearest, ties to even (RTNE): use guard/sticky and LSB to decide increment
+        m2bits = mantissas >> 21
+        lsb_keep = (m2bits >> 1) & 0x1
+        guard = m2bits & 0x1
+        IS_SRC_FP32: tl.constexpr = src_tensor.dtype == tl.float32
+        if IS_SRC_FP32:
+            bit0_dropped = (mantissas_orig & 0x1) != 0
+            mask = (1 << tl.minimum(adjusted_exponents, 31)) - 1
+            dropped_post = (mantissas_pre & mask) != 0
+            sticky = is_subnormal & (bit0_dropped | dropped_post)
+            sticky |= ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
+        else:
+            sticky = ((mantissas & 0x1FFFFF) != 0).to(tl.uint32)
+        round_inc = guard & (sticky | lsb_keep)
+        e2m1_tmp = tl.minimum((((exponents << 2) | m2bits) + round_inc) >> 1, 0x7)
+        e2m1_value = ((signs >> 28) | e2m1_tmp).to(tl.uint8)
+
+        e2m1_value = tl.reshape(e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2])
+        evens, odds = tl.split(e2m1_value)
+        out_tensor = evens | (odds << 4)
+    else:
+        out_tensor = quant_tensor.to(mx_tensor_dtype)
+
+    return out_tensor, dequant_scale, s_dec
+
+@triton.jit
+def _compute_dequant(
+    mx_tensor,
+    scale,
+    s_dec,
+    BLOCK_SIZE_OUT_DIM: tl.constexpr,
+    BLOCK_SIZE_QUANT_DIM: tl.constexpr,
+    dst_dtype: tl.constexpr,
+    IS_BLACKWELL: tl.constexpr = False,
+):
+    tl.static_assert(BLOCK_SIZE_QUANT_DIM % MXFP_BLOCK_SIZE == 0, f"Block size along quantization block must be a multiple of {MXFP_BLOCK_SIZE=}")
+    # uint8 signifies two fp4 e2m1 values packed into a single byte
+    mx_tensor_dtype: tl.constexpr = mx_tensor.dtype
+    _is_f16: tl.constexpr = dst_dtype == tl.float16
+    _is_bf16: tl.constexpr = dst_dtype == tl.bfloat16
+    _is_f32: tl.constexpr = dst_dtype == tl.float32
+    tl.static_assert(_is_f16 or (_is_bf16 or _is_f32))
+    _is_u8: tl.constexpr = mx_tensor_dtype == tl.uint8
+    _is_e4: tl.constexpr = mx_tensor_dtype == tl.float8e4nv
+    _is_e5: tl.constexpr = mx_tensor_dtype == tl.float8e5
+    _is_dst: tl.constexpr = mx_tensor_dtype == dst_dtype
+    tl.static_assert(
+        _is_u8 or ((_is_e4 or _is_e5) or _is_dst),
+        "mx_tensor_ptr must be uint8 or float8 or dst_dtype")
+    tl.static_assert(scale.dtype == tl.float8e4nv, "scale must be float8e4nv")
+
+    # Determine if we are dealing with fp8 types.
+    is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8
+    BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // MXFP_BLOCK_SIZE
+
+    # Upcast the scale to the destination type.
+    if dst_dtype == tl.bfloat16:
+        dst_scale = scale.to(tl.bfloat16)
+    else:
+        dst_scale = scale.to(tl.float32)
+        if dst_dtype == tl.float16:
+            dst_scale = dst_scale.to(tl.float16)
+
+    # Now upcast the tensor.
+    intermediate_dtype: tl.constexpr = tl.bfloat16 if dst_dtype == tl.float32 else dst_dtype
+    if IS_BLACKWELL:
+        assert is_fp4
+        packed_u32 = tl.inline_asm_elementwise(
+            asm="""
+            {
+            .reg .b8 in_8;
+            .reg .f16x2 out;
+            cvt.u8.u32 in_8, $1;
+            cvt.rn.f16x2.e2m1x2 out, in_8;
+            mov.b32 $0, out;
+            }
+            """,
+            constraints="=r,r",
+            args=[mx_tensor],  # tl.uint8 passed in as a 32-bit reg with value in low 8 bits
+            dtype=tl.uint32,
+            is_pure=True,
+            pack=1,
+        )
+        lo_u16 = (packed_u32 & 0xFFFF).to(tl.uint16)
+        hi_u16 = (packed_u32 >> 16).to(tl.uint16)
+        lo_f16 = lo_u16.to(tl.float16, bitcast=True)
+        hi_f16 = hi_u16.to(tl.float16, bitcast=True)
+
+        if intermediate_dtype == tl.float16:
+            x0, x1 = lo_f16, hi_f16
+        else:
+            x0 = lo_f16.to(intermediate_dtype)
+            x1 = hi_f16.to(intermediate_dtype)
+
+        dst_tensor = tl.interleave(x0, x1)
+
+    else:
+        assert is_fp4
+        dst_bias: tl.constexpr = 127 if intermediate_dtype == tl.bfloat16 else 15  # exponent bias
+        dst_0p5: tl.constexpr = 16128 if intermediate_dtype == tl.bfloat16 else 0x3800
+        dst_m_bits: tl.constexpr = 7 if intermediate_dtype == tl.bfloat16 else 10  # mantissa bits
+        # e2m1
+        em0 = mx_tensor & 0x07
+        em1 = mx_tensor & 0x70
+        x0 = (em0.to(tl.uint16) << (dst_m_bits - 1)) | ((mx_tensor & 0x08).to(tl.uint16) << 12)
+        x1 = (em1.to(tl.uint16) << (dst_m_bits - 5)) | ((mx_tensor & 0x80).to(tl.uint16) << 8)
+        # Three cases:
+        # 1) x is normal and non-zero: Correct bias
+        x0 = tl.where((em0 & 0x06) != 0, x0 + ((dst_bias - 1) << dst_m_bits), x0)
+        x1 = tl.where((em1 & 0x60) != 0, x1 + ((dst_bias - 1) << dst_m_bits), x1)
+        # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type
+        x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0)
+        x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1)
+        # 3) x is zero, do nothing
+        dst_tensor = tl.interleave(x0, x1).to(intermediate_dtype, bitcast=True)
+
+    dst_tensor = dst_tensor.to(dst_dtype)
+
+    # Reshape for proper broadcasting: the scale was stored with a 16‐sized “inner” grouping.
+    dst_tensor = dst_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, MXFP_BLOCK_SIZE])
+    dst_scale = dst_scale.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 1])
+    scale = scale.reshape(dst_scale.shape)
+
+    out_tensor = dst_tensor * dst_scale * s_dec  # NVFP4 has the additional global scale factor
+    if dst_dtype == tl.float32:
+        max_fin = 3.4028234663852886e+38
+    elif dst_dtype == tl.bfloat16:
+        max_fin = 3.3895313892515355e+38
+    else:
+        tl.static_assert(dst_dtype == tl.float16)
+        max_fin = 65504
+    out_tensor = tl.clamp(out_tensor, min=-max_fin, max=max_fin)
+    out_tensor = out_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM])
+    out_tensor = out_tensor.to(dst_dtype)
+    return out_tensor

backend_snapshot/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py

@@ -0,0 +1,80 @@
+import triton
+import triton.language as tl
+
+from .nvfp4_utils import _compute_quant_and_scale, _compute_dequant
+
+@triton.jit
+def fake_quantize(src_tensor, valid_src_mask, BLOCK_SIZE_OUT_DIM: tl.constexpr,
+                    BLOCK_SIZE_QUANT_DIM: tl.constexpr,
+                    dst_dtype: tl.constexpr,
+                    mx_tensor_dtype: tl.constexpr = tl.uint8,
+                    use_global_sf: tl.constexpr = True,
+                    two_level_quant_P: tl.constexpr = False):
+    high_prec_src_tensor = src_tensor
+    src_tensor, src_scale, src_s_dec = _compute_quant_and_scale(src_tensor=src_tensor,
+                                                                valid_src_mask=valid_src_mask,
+                                                                mx_tensor_dtype=mx_tensor_dtype,
+                                                                use_global_sf=use_global_sf,
+                                                                two_level_quant_P=two_level_quant_P)
+    src_tensor = _compute_dequant(mx_tensor=src_tensor, 
+                                  scale=src_scale, 
+                                  s_dec=src_s_dec, 
+                                  BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, 
+                                  BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, 
+                                  dst_dtype=dst_dtype)
+    return src_tensor, high_prec_src_tensor.to(src_tensor.dtype)
+
+@triton.jit
+def fake_quantize_q(Q, fake_Q, stride_z_q, stride_h_q, 
+                    stride_tok_q, stride_d_q,
+                    fake_stride_z_q, fake_stride_h_q,
+                    fake_stride_tok_q, fake_stride_d_q,
+                    H, N_CTX_Q,
+                    BLOCK_M: tl.constexpr,
+                    HEAD_DIM: tl.constexpr,
+                    use_global_sf: tl.constexpr = True):
+    bhid = tl.program_id(1)
+    adj_q = (stride_h_q * (bhid % H) + stride_z_q * (bhid // H))
+    fake_adj_q = (fake_stride_h_q * (bhid % H) + fake_stride_z_q * (bhid // H))
+    Q += adj_q
+    fake_Q += fake_adj_q
+
+    pid = tl.program_id(0)
+    start_m = pid * BLOCK_M
+    offs_m = start_m + tl.arange(0, BLOCK_M)
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    q_valid = offs_m < N_CTX_Q
+    q = tl.load(Q + offs_m[:, None] * stride_tok_q + offs_k[None, :] * stride_d_q, mask=q_valid[:, None], other=0.0)
+    q, _ = fake_quantize(src_tensor=q, valid_src_mask=q_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_M, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=q.dtype, use_global_sf=use_global_sf)
+    tl.store(fake_Q + offs_m[:, None] * fake_stride_tok_q + offs_k[None, :] * fake_stride_d_q, q, mask=q_valid[:, None])
+
+@triton.jit
+def fake_quantize_kv(K, V, fake_K, fake_V, stride_z_kv, stride_h_kv, 
+                    stride_tok_kv, stride_d_kv,
+                    fake_stride_z_kv, fake_stride_h_kv,
+                    fake_stride_tok_kv, fake_stride_d_kv,
+                    H, N_CTX_KV,
+                    BLOCK_N: tl.constexpr,
+                    HEAD_DIM: tl.constexpr,
+                    use_global_sf: tl.constexpr = True):
+    bhid = tl.program_id(1)
+    adj_kv = (stride_h_kv * (bhid % H) + stride_z_kv * (bhid // H))
+    fake_adj_kv = (fake_stride_h_kv * (bhid % H) + fake_stride_z_kv * (bhid // H))
+    K += adj_kv
+    V += adj_kv
+    fake_K += fake_adj_kv
+    fake_V += fake_adj_kv
+
+    pid = tl.program_id(0)
+    start_n = pid * BLOCK_N
+    offs_n = start_n + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, HEAD_DIM)
+    
+    kv_valid = offs_n < N_CTX_KV
+    k_block = tl.load(K + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
+    v_block = tl.load(V + offs_n[:, None] * stride_tok_kv + offs_k[None, :] * stride_d_kv, mask=kv_valid[:, None], other=0.0)
+    k, _ = fake_quantize(src_tensor=k_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=k_block.dtype, use_global_sf=use_global_sf)
+    v, _ = fake_quantize(src_tensor=v_block, valid_src_mask=kv_valid[:, None], BLOCK_SIZE_OUT_DIM=BLOCK_N, BLOCK_SIZE_QUANT_DIM=HEAD_DIM, dst_dtype=v_block.dtype, use_global_sf=use_global_sf)
+    tl.store(fake_K + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, k, mask=kv_valid[:, None])
+    tl.store(fake_V + offs_n[:, None] * fake_stride_tok_kv + offs_k[None, :] * fake_stride_d_kv, v, mask=kv_valid[:, None])

backend_snapshot/fastvideo/attention/backends/sparse_fp4_ours_p_attn.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Sparse FP4 Attention backend with the independent ours-P quant kernel."""
+
+import math
+
+import torch
+import torch.nn.functional as F
+import triton
+
+from fastvideo_kernel.triton_kernels.quant_utils import (
+    fake_quantize_q,
+    fake_quantize_kv,
+)
+from fastvideo_kernel.block_sparse_attn_ours_p import block_sparse_attn_ours_p
+from fastvideo.forward_context import get_forward_context
+
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder,
+)
+from fastvideo.attention.backends.video_sparse_attn import (
+    VideoSparseAttentionMetadata,
+    VideoSparseAttentionMetadataBuilder,
+    VSA_TILE_SIZE,
+)
+from fastvideo.distributed import get_sp_group
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _dense_sdpa_blhd(query, key, value):
+    q = query.transpose(1, 2)
+    k = key.transpose(1, 2)
+    v = value.transpose(1, 2)
+    out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+    return out.transpose(1, 2)
+
+
+def _quantize_qkv_bhld(q, k, v):
+    """FP4 fake quantize Q/K/V in BHLD layout, same as attn_qat_train."""
+    H = q.shape[1]
+    N_Q = q.shape[2]
+    N_KV = k.shape[2]
+    D = q.shape[3]
+    BLOCK = 32
+
+    fake_q = torch.empty_like(q)
+    fake_k = torch.empty_like(k)
+    fake_v = torch.empty_like(v)
+
+    grid_q = (triton.cdiv(N_Q, BLOCK), q.shape[0] * H, 1)
+    grid_kv = (triton.cdiv(N_KV, BLOCK), q.shape[0] * H, 1)
+
+    fake_quantize_q[grid_q](
+        q, fake_q,
+        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+        fake_q.stride(0), fake_q.stride(1), fake_q.stride(2), fake_q.stride(3),
+        H, N_Q, BLOCK_M=BLOCK, HEAD_DIM=D, use_global_sf=False,
+    )
+    fake_quantize_kv[grid_kv](
+        k, v, fake_k, fake_v,
+        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+        fake_k.stride(0), fake_k.stride(1), fake_k.stride(2), fake_k.stride(3),
+        H, N_KV, BLOCK_N=BLOCK, HEAD_DIM=D, use_global_sf=False,
+    )
+    return fake_q, fake_k, fake_v
+
+
+class SparseFP4OursPAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SPARSE_FP4_OURS_P_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["SparseFP4OursPAttentionImpl"]:
+        return SparseFP4OursPAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
+        return VideoSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
+        return VideoSparseAttentionMetadataBuilder
+
+
+class SparseFP4OursPAttentionImpl(AttentionImpl):
+
+    def __init__(self, num_heads, head_size, causal, softmax_scale,
+                 num_kv_heads=None, prefix="", **extra):
+        self.prefix = prefix
+        self.sp_size = get_sp_group().world_size
+
+    def tile(self, x, num_tiles, tile_partition_indices, non_pad_index):
+        t_p = num_tiles[0] * VSA_TILE_SIZE[0]
+        h_p = num_tiles[1] * VSA_TILE_SIZE[1]
+        w_p = num_tiles[2] * VSA_TILE_SIZE[2]
+        out = torch.zeros(
+            (x.shape[0], t_p * h_p * w_p, x.shape[-2], x.shape[-1]),
+            device=x.device, dtype=x.dtype,
+        )
+        out[:, non_pad_index] = x[:, tile_partition_indices]
+        return out
+
+    def untile(self, x, reverse_tile_partition_indices, non_pad_index):
+        return x[:, non_pad_index][:, reverse_tile_partition_indices]
+
+    def _is_force_dense(self) -> bool:
+        ctx = get_forward_context()
+        return ctx.force_dense
+
+    def preprocess_qkv(self, qkv, attn_metadata):
+        if attn_metadata is None or self._is_force_dense():
+            return qkv
+        return self.tile(qkv, attn_metadata.num_tiles,
+                         attn_metadata.tile_partition_indices,
+                         attn_metadata.non_pad_index)
+
+    def postprocess_output(self, output, attn_metadata):
+        if attn_metadata is None or self._is_force_dense():
+            return output
+        return self.untile(output,
+                           attn_metadata.reverse_tile_partition_indices,
+                           attn_metadata.non_pad_index)
+
+    def forward(self, query, key, value,
+                gate_compress_or_metadata=None, attn_metadata=None):
+        # Handle both call conventions
+        if attn_metadata is None and isinstance(
+                gate_compress_or_metadata, (VideoSparseAttentionMetadata, type(None))):
+            attn_metadata = gate_compress_or_metadata
+
+        # ── force_dense: true dense BF16 SDPA (for teacher in distillation) ──
+        ctx = get_forward_context()
+        if ctx.force_dense:
+            return _dense_sdpa_blhd(query, key, value)
+
+        is_cross = query.shape[1] != key.shape[1]
+
+        # ── Cross-attention/no metadata: keep dense. The sparse VSA metadata only
+        # applies to tiled video self-attention.
+        if attn_metadata is None or is_cross:
+            return _dense_sdpa_blhd(query, key, value)
+
+        # ── Self-attention: FP4 quant Q/K/V + block-sparse attention ──
+        # BLHD → BHLD
+        q = query.transpose(1, 2).contiguous()
+        k = key.transpose(1, 2).contiguous()
+        v = value.transpose(1, 2).contiguous()
+
+        # Step 1: FP4 fake quantize Q/K/V with STE (straight-through estimator)
+        with torch.no_grad():
+            fq, fk, fv = _quantize_qkv_bhld(q, k, v)
+        # STE: forward uses quantized values, backward passes gradient through as-is
+        fq = q + (fq - q).detach()
+        fk = k + (fk - k).detach()
+        fv = v + (fv - v).detach()
+
+        # Step 2: Build sparse block map
+        B, H, S, D = fq.shape
+        block_elements = math.prod(VSA_TILE_SIZE)
+        num_blocks = S // block_elements
+
+        VSA_sparsity = attn_metadata.VSA_sparsity
+        cur_topk = max(1, math.ceil((1 - VSA_sparsity) * num_blocks))
+        logger.info(f"[SFP4] S={S} num_blocks={num_blocks} sparsity={VSA_sparsity} topk={cur_topk}/{num_blocks}")
+
+        block_sizes = attn_metadata.variable_block_sizes.to(
+            device=fq.device, dtype=torch.float32).clamp_min(1)
+        block_sizes = block_sizes.view(1, 1, num_blocks, 1)
+        q_c = (fq.view(B, H, num_blocks, block_elements, D).float().sum(3) /
+               block_sizes).to(fq.dtype)
+        k_c = (fk.view(B, H, num_blocks, block_elements, D).float().sum(3) /
+               block_sizes).to(fk.dtype)
+        v_c = (fv.view(B, H, num_blocks, block_elements, D).float().sum(3) /
+               block_sizes).to(fv.dtype)
+        scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (D ** 0.5)
+        topk_idx = torch.topk(scores, cur_topk, dim=-1).indices
+        block_map = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True)
+
+        # Step 3: Block-sparse attention with independent group-local P quant.
+        out, _ = block_sparse_attn_ours_p(fq, fk, fv, block_map,
+                                          attn_metadata.variable_block_sizes,
+                                          q_c, k_c, v_c)
+
+        return out.transpose(1, 2)  # BHLD → BLHD

backend_snapshot/fastvideo/attention/backends/video_sparse_attn.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+import functools
+import math
+from dataclasses import dataclass
+
+import torch
+
+try:
+    from fastvideo_kernel import video_sparse_attn
+except ImportError:
+    video_sparse_attn = None
+
+from typing import Any
+
+from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata,
+                                                   AttentionMetadataBuilder)
+from fastvideo.distributed import get_sp_group
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+VSA_TILE_SIZE = (4, 4, 4)
+
+
+@functools.lru_cache(maxsize=10)
+def get_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    T, H, W = dit_seq_shape
+    ts, hs, ws = tile_size
+    indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
+    ls = []
+    for t in range(math.ceil(T / ts)):
+        for h in range(math.ceil(H / hs)):
+            for w in range(math.ceil(W / ws)):
+                ls.append(indices[t * ts:min(t * ts + ts, T), h * hs:min(h * hs + hs, H),
+                                  w * ws:min(w * ws + ws, W)].flatten())
+    index = torch.cat(ls, dim=0)
+    return index
+
+
+@functools.lru_cache(maxsize=10)
+def get_reverse_tile_partition_indices(
+    dit_seq_shape: tuple[int, int, int],
+    tile_size: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device))
+
+
+@functools.lru_cache(maxsize=10)
+def construct_variable_block_sizes(
+    dit_seq_shape: tuple[int, int, int],
+    num_tiles: tuple[int, int, int],
+    device: torch.device,
+) -> torch.LongTensor:
+    """
+    Compute the number of valid (non‑padded) tokens inside every
+    (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order
+    (t‑tile, h‑tile, w‑tile) that `rearrange` uses.
+
+    Returns
+    -------
+    torch.LongTensor  # shape: [∏ full_window_size]
+    """
+    # unpack
+    t, h, w = dit_seq_shape
+    ts_t, ts_h, ts_w = VSA_TILE_SIZE
+    n_t, n_h, n_w = num_tiles
+
+    def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor:
+        """Vector with the size of each tile along one dimension."""
+        sizes = torch.full((n_tiles, ), tile, dtype=torch.int, device=device)
+        # size of last (possibly partial) tile
+        remainder = dim_len - (n_tiles - 1) * tile
+        sizes[-1] = remainder if remainder > 0 else tile
+        return sizes
+
+    t_sizes = _sizes(t, ts_t, n_t)  # [n_t]
+    h_sizes = _sizes(h, ts_h, n_h)  # [n_h]
+    w_sizes = _sizes(w, ts_w, n_w)  # [n_w]
+
+    # broadcast‑multiply to get voxels per tile, then flatten
+    block_sizes = (
+        t_sizes[:, None, None]  # [n_t, 1,   1]
+        * h_sizes[None, :, None]  # [1,   n_h, 1]
+        * w_sizes[None, None, :]  # [1,   1,   n_w]
+    ).reshape(-1)  # [n_t * n_h * n_w]
+
+    return block_sizes
+
+
+@functools.lru_cache(maxsize=10)
+def get_non_pad_index(
+    variable_block_sizes: torch.LongTensor,
+    max_block_size: int,
+):
+    n_win = variable_block_sizes.shape[0]
+    device = variable_block_sizes.device
+    starts_pad = torch.arange(n_win, device=device) * max_block_size
+    index_pad = starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :]
+    index_mask = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
+    return index_pad[index_mask]
+
+
+class VideoSparseAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128]
+
+    @staticmethod
+    def get_name() -> str:
+        return "VIDEO_SPARSE_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
+        return VideoSparseAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
+        return VideoSparseAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
+        return VideoSparseAttentionMetadataBuilder
+
+
+@dataclass
+class VideoSparseAttentionMetadata(AttentionMetadata):
+    current_timestep: int
+    dit_seq_shape: list[int]
+    num_tiles: list[int]
+    total_seq_length: int
+    tile_partition_indices: torch.LongTensor
+    reverse_tile_partition_indices: torch.LongTensor
+    variable_block_sizes: torch.LongTensor
+    non_pad_index: torch.LongTensor
+
+
+class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
+
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(  # type: ignore
+        self,
+        current_timestep: int,
+        raw_latent_shape: tuple[int, int, int],
+        patch_size: tuple[int, int, int],
+        VSA_sparsity: float,
+        device: torch.device,
+        **kwargs: dict[str, Any],
+    ) -> VideoSparseAttentionMetadata:
+        patch_size = patch_size
+        dit_seq_shape = (raw_latent_shape[0] // patch_size[0], raw_latent_shape[1] // patch_size[1],
+                         raw_latent_shape[2] // patch_size[2])
+
+        num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]),
+                     math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2]))
+        total_seq_length = math.prod(dit_seq_shape)
+
+        tile_partition_indices = get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device)
+        reverse_tile_partition_indices = get_reverse_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device)
+        variable_block_sizes = construct_variable_block_sizes(dit_seq_shape, num_tiles, device)
+        non_pad_index = get_non_pad_index(variable_block_sizes, math.prod(VSA_TILE_SIZE))
+
+        return VideoSparseAttentionMetadata(
+            current_timestep=current_timestep,
+            dit_seq_shape=dit_seq_shape,  # type: ignore
+            VSA_sparsity=VSA_sparsity,  # type: ignore
+            num_tiles=num_tiles,  # type: ignore
+            total_seq_length=total_seq_length,  # type: ignore
+            tile_partition_indices=tile_partition_indices,  # type: ignore
+            reverse_tile_partition_indices=reverse_tile_partition_indices,
+            variable_block_sizes=variable_block_sizes,
+            non_pad_index=non_pad_index)
+
+
+class VideoSparseAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.prefix = prefix
+        sp_group = get_sp_group()
+        self.sp_size = sp_group.world_size
+
+    def tile(self, x: torch.Tensor, num_tiles: list[int], tile_partition_indices: torch.LongTensor,
+             non_pad_index: torch.LongTensor) -> torch.Tensor:
+        t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0]
+        h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1]
+        w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
+
+        x_padded = torch.zeros((x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1]),
+                               device=x.device,
+                               dtype=x.dtype)
+        x_padded[:, non_pad_index] = x[:, tile_partition_indices]
+        return x_padded
+
+    def untile(self, x: torch.Tensor, reverse_tile_partition_indices: torch.LongTensor,
+               non_pad_index: torch.LongTensor) -> torch.Tensor:
+        x = x[:, non_pad_index][:, reverse_tile_partition_indices]
+        return x
+
+    def preprocess_qkv(
+        self,
+        qkv: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.tile(qkv, attn_metadata.num_tiles, attn_metadata.tile_partition_indices,
+                         attn_metadata.non_pad_index)
+
+    def postprocess_output(
+        self,
+        output: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        return self.untile(output, attn_metadata.reverse_tile_partition_indices, attn_metadata.non_pad_index)
+
+    def forward(  # type: ignore[override]
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        gate_compress: torch.Tensor,
+        attn_metadata: VideoSparseAttentionMetadata,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+        gate_compress = gate_compress.transpose(1, 2).contiguous()
+
+        VSA_sparsity = attn_metadata.VSA_sparsity
+
+        cur_topk = math.ceil((1 - VSA_sparsity) * (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE)))
+
+        if video_sparse_attn is None:
+            raise NotImplementedError("video_sparse_attn is not installed")
+        hidden_states = video_sparse_attn(query,
+                                          key,
+                                          value,
+                                          attn_metadata.variable_block_sizes,
+                                          attn_metadata.variable_block_sizes,
+                                          cur_topk,
+                                          block_size=VSA_TILE_SIZE,
+                                          compress_attn_weight=gate_compress).transpose(1, 2)
+
+        return hidden_states

backend_snapshot/fastvideo/configs/models/dits/base.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+@dataclass
+class DiTArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=list)
+    _compile_conditions: list = field(default_factory=list)
+    param_names_mapping: dict = field(default_factory=dict)
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+    _supported_attention_backends: tuple[AttentionBackendEnum,
+                                         ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
+                                                 AttentionBackendEnum.TORCH_SDPA,
+                                                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+                                                 AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE,
+                                                 AttentionBackendEnum.ATTN_QAT_INFER,
+                                                 AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN,
+                                                 AttentionBackendEnum.SAGE_SLA_ATTN,
+                                                 AttentionBackendEnum.SPARSE_FP4_ATTN,
+                                                 AttentionBackendEnum.SPARSE_FP4_OURS_P_ATTN)
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+    in_channels: int | None = 0
+    out_channels: int | None = 0
+    patch_size: int | tuple[int, int, int] | None = None
+    expand_timesteps: bool = False
+    num_layers: int = 0
+    ffn_dim: int = 0
+    exclude_lora_layers: list[str] = field(default_factory=list)
+    boundary_ratio: float | None = None
+
+    def __post_init__(self) -> None:
+        if not self._compile_conditions:
+            self._compile_conditions = self._fsdp_shard_conditions.copy()
+
+
+@dataclass
+class DiTConfig(ModelConfig):
+    arch_config: DiTArchConfig = field(default_factory=DiTArchConfig)
+
+    # FastVideoDiT-specific parameters
+    prefix: str = ""
+    quant_config: QuantizationConfig | None = None
+    expand_timesteps: bool = False
+    boundary_ratio: float | None = None
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.arch_config.expand_timesteps = self.expand_timesteps
+        self.arch_config.boundary_ratio = self.boundary_ratio
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for DiTConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=DiTConfig.prefix,
+            help="Prefix for the DiT model",
+        )
+
+        parser.add_argument(
+            f"--{prefix}.quant-config",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.quant_config",
+            default=None,
+            help="Quantization configuration for the DiT model",
+        )
+
+        return parser

backend_snapshot/fastvideo/forward_context.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/forward_context.py
+
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from fastvideo.logger import init_logger
+
+if TYPE_CHECKING:
+    from fastvideo.attention import AttentionMetadata
+    from fastvideo.pipelines import ForwardBatch
+
+logger = init_logger(__name__)
+
+# TODO(will): check if this is needed
+# track_batchsize: bool = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL >= 0
+track_batchsize: bool = False
+last_logging_time: float = 0
+forward_start_time: float = 0
+# batchsize_logging_interval: float = envs.FASTVIDEO_LOG_BATCHSIZE_INTERVAL
+batchsize_logging_interval: float = 1000
+batchsize_forward_time: defaultdict = defaultdict(list)
+
+
+#
+@dataclass
+class ForwardContext:
+    current_timestep: int
+    # TODO(will): check this arg
+    # copy from vllm_config.compilation_config.static_forward_context
+    # attn_layers: Dict[str, Any]
+    # TODO: extend to support per-layer dynamic forward context
+    attn_metadata: "AttentionMetadata"  # set dynamically for each forward pass
+    forward_batch: Optional["ForwardBatch"] = None
+    force_dense: bool = False
+
+
+_forward_context: Optional["ForwardContext"] = None
+
+
+def get_forward_context() -> "ForwardContext":
+    """Get the current forward context."""
+    assert _forward_context is not None, ("Forward context is not set. "
+                                          "Please use `set_forward_context` to set the forward context.")
+    return _forward_context
+
+
+# TODO(will): finalize the interface
+@contextmanager
+def set_forward_context(current_timestep, attn_metadata, forward_batch: Optional["ForwardBatch"] = None, force_dense: bool = False):
+    """A context manager that stores the current forward context,
+    can be attention metadata, etc.
+    Here we can inject common logic for every model forward pass.
+    """
+    global forward_start_time
+    need_to_track_batchsize = track_batchsize and attn_metadata is not None
+    if need_to_track_batchsize:
+        forward_start_time = time.perf_counter()
+    global _forward_context
+    prev_context = _forward_context
+    _forward_context = ForwardContext(current_timestep=current_timestep,
+                                      attn_metadata=attn_metadata,
+                                      forward_batch=forward_batch,
+                                      force_dense=force_dense)
+
+    try:
+        yield
+    finally:
+        global last_logging_time, batchsize_logging_interval
+        if need_to_track_batchsize:
+            if hasattr(attn_metadata, "num_prefill_tokens"):
+                # for v0 attention backends
+                batchsize = attn_metadata.num_prefill_tokens + \
+                    attn_metadata.num_decode_tokens
+            else:
+                # for v1 attention backends
+                batchsize = attn_metadata.num_input_tokens
+            now = time.perf_counter()
+            # time measurement is in milliseconds
+            batchsize_forward_time[batchsize].append((now - forward_start_time) * 1000)
+            if now - last_logging_time > batchsize_logging_interval:
+                last_logging_time = now
+                forward_stats = []
+                for bs, times in batchsize_forward_time.items():
+                    if len(times) <= 1:
+                        # can be cudagraph / profiling run
+                        continue
+                    medium = torch.quantile(torch.tensor(times), q=0.5).item()
+                    medium = round(medium, 2)
+                    forward_stats.append((bs, len(times), medium))
+                forward_stats.sort(key=lambda x: x[1], reverse=True)
+                if forward_stats:
+                    logger.info(("Batchsize forward time stats "
+                                 "(batchsize, count, median_time(ms)): %s"), forward_stats)
+        _forward_context = prev_context

backend_snapshot/fastvideo/pipelines/stages/denoising.py

@@ -0,0 +1,1184 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Denoising stage for diffusion pipelines.
+"""
+
+import inspect
+import weakref
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from tqdm.auto import tqdm
+
+from fastvideo.attention import get_attn_backend
+from fastvideo.distributed import (get_local_torch_device, get_world_group)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.loader.component_loader import TransformerLoader
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (FlowMatchEulerDiscreteScheduler)
+from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import StageValidators as V
+from fastvideo.pipelines.stages.validators import VerificationResult
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.utils import dict_to_3d_list, masks_like
+
+try:
+    from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend
+    from fastvideo.utils import is_vmoba_available
+    vmoba_attn_available = is_vmoba_available()
+except ImportError:
+    vmoba_attn_available = False
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+
+try:
+    from fastvideo.attention.backends.sparse_fp4_attn import (SparseFP4AttentionBackend)
+except ImportError:
+    SparseFP4AttentionBackend = None  # type: ignore[assignment]
+
+try:
+    from fastvideo.attention.backends.sparse_fp4_ours_p_attn import (SparseFP4OursPAttentionBackend)
+except ImportError:
+    SparseFP4OursPAttentionBackend = None  # type: ignore[assignment]
+
+sparse_fp4_backends = tuple(
+    backend for backend in (
+        SparseFP4AttentionBackend,
+        SparseFP4OursPAttentionBackend,
+    ) if backend is not None)
+sparse_fp4_available = bool(sparse_fp4_backends)
+
+logger = init_logger(__name__)
+
+
+class DenoisingStage(PipelineStage):
+    """
+    Stage for running the denoising loop in diffusion pipelines.
+    
+    This stage handles the iterative denoising process that transforms
+    the initial noise into the final output.
+    """
+
+    def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.transformer_2 = transformer_2
+        self.scheduler = scheduler
+        self.vae = vae
+        self.pipeline = weakref.ref(pipeline) if pipeline else None
+        attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
+        self.attn_backend = get_attn_backend(
+            head_size=attn_head_size,
+            dtype=torch.float16,  # TODO(will): hack
+            supported_attention_backends=(
+                AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN,
+                AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE,
+                AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_TRAIN,
+                AttentionBackendEnum.SPARSE_FP4_ATTN, AttentionBackendEnum.SPARSE_FP4_OURS_P_ATTN)  # hack
+        )
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run the denoising loop.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        pipeline = self.pipeline() if self.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        # Prepare extra step kwargs for scheduler
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        # Setup precision and autocast settings
+        # TODO(will): make the precision configurable for inference
+        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        # Get timesteps and calculate warmup steps
+        timesteps = batch.timesteps
+        # TODO(will): remove this once we add input/output validation for stages
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        # Prepare image latents and embeddings for I2V generation
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert not torch.isnan(image_embeds[0]).any(), "image_embeds contains nan"
+            image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]
+
+        image_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24)
+            },
+        )
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        neg_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_neg,
+                "encoder_attention_mask": batch.negative_attention_mask,
+            },
+        )
+
+        action_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "mouse_cond": batch.mouse_cond,
+                "keyboard_cond": batch.keyboard_cond,
+                "c2ws_plucker_emb": batch.c2ws_plucker_emb,
+            },
+        )
+
+        camera_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "camera_states": batch.camera_states,
+            },
+        )
+
+        # Get latents and embeddings
+        latents = batch.latents
+        prompt_embeds = batch.prompt_embeds
+        assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan"
+        if batch.do_classifier_free_guidance:
+            neg_prompt_embeds = batch.negative_prompt_embeds
+            assert neg_prompt_embeds is not None
+            assert not torch.isnan(neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
+
+        # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
+        boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
+        if batch.boundary_ratio is not None:
+            logger.info("Overriding boundary ratio from %s to %s", boundary_ratio, batch.boundary_ratio)
+            boundary_ratio = batch.boundary_ratio
+
+        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps if boundary_ratio is not None else None
+        latent_model_input = latents.to(target_dtype)
+        assert latent_model_input.shape[0] == 1, "only support batch size 1"
+
+        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            # TI2V directly replaces the first frame of the latent with
+            # the image latent instead of appending along the channel dim
+            assert batch.image_latent is None, "TI2V task should not have image latents"
+            assert self.vae is not None, "VAE is not provided for TI2V task"
+            z = self.vae.encode(batch.pil_image).mean.float()
+            if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
+                if isinstance(self.vae.shift_factor, torch.Tensor):
+                    z -= self.vae.shift_factor.to(z.device, z.dtype)
+                else:
+                    z -= self.vae.shift_factor
+
+            if isinstance(self.vae.scaling_factor, torch.Tensor):
+                z = z * self.vae.scaling_factor.to(z.device, z.dtype)
+            else:
+                z = z * self.vae.scaling_factor
+
+            latent_model_input = latent_model_input.squeeze(0)
+            _, mask2 = masks_like([latent_model_input], zero=True)
+
+            latent_model_input = (1. - mask2[0]) * z + mask2[0] * latent_model_input
+            # latent_model_input = latent_model_input.unsqueeze(0)
+            latent_model_input = latent_model_input.to(get_local_torch_device())
+            latents = latent_model_input
+            F = batch.num_frames
+            temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
+            spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            if not isinstance(patch_size, tuple):
+                raise ValueError(f"Expected 3D patch_size tuple for denoising, got {patch_size!r}")
+            seq_len = ((F - 1) // temporal_scale + 1) * (batch.height // spatial_scale) * (
+                batch.width // spatial_scale) // (patch_size[1] * patch_size[2])
+
+        # Initialize lists for ODE trajectory
+        trajectory_timesteps: list[torch.Tensor] = []
+        trajectory_latents: list[torch.Tensor] = []
+
+        # Run denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Skip if interrupted
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+
+                if boundary_timestep is None or t >= boundary_timestep:
+                    if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
+                            and self.transformer_2 is not None
+                            and next(self.transformer_2.parameters()).device.type == 'cuda'):
+                        self.transformer_2.to('cpu')
+                    current_model = self.transformer
+                    if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
+                            and not fastvideo_args.use_fsdp_inference and current_model is not None):
+                        transformer_device = next(current_model.parameters()).device.type
+                        if transformer_device == 'cpu':
+                            current_model.to(get_local_torch_device())
+                    current_guidance_scale = batch.guidance_scale
+                else:
+                    # low-noise stage in wan2.2
+                    if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
+                            and next(self.transformer.parameters()).device.type == 'cuda'):
+                        self.transformer.to('cpu')
+                    current_model = self.transformer_2
+                    if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload
+                            and not fastvideo_args.use_fsdp_inference and current_model is not None):
+                        transformer_2_device = next(current_model.parameters()).device.type
+                        if transformer_2_device == 'cpu':
+                            current_model.to(get_local_torch_device())
+                    current_guidance_scale = batch.guidance_scale_2
+                assert current_model is not None, "current_model is None"
+
+                # Expand latents for V2V/I2V
+                latent_model_input = latents.to(target_dtype)
+                if batch.video_latent is not None:
+                    latent_model_input = torch.cat([latent_model_input, batch.video_latent,
+                                                    torch.zeros_like(latents)],
+                                                   dim=1).to(target_dtype)
+                elif batch.image_latent is not None:
+                    assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
+                    latent_model_input = torch.cat([latent_model_input, batch.image_latent], dim=1).to(target_dtype)
+
+                assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan"
+                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                    timestep = torch.stack([t]).to(get_local_torch_device())
+                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
+                    temp_ts = torch.cat([temp_ts, temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep])
+                    timestep = temp_ts.unsqueeze(0)
+                    t_expand = timestep.repeat(latent_model_input.shape[0], 1)
+                else:
+                    t_expand = t.repeat(latent_model_input.shape[0])
+                t_expand = t_expand.to(get_local_torch_device())
+
+                use_meanflow = getattr(self.transformer.config, "use_meanflow", False)
+                if use_meanflow:
+                    if i == len(timesteps) - 1:
+                        timesteps_r = torch.tensor([0.0], device=get_local_torch_device())
+                    else:
+                        timesteps_r = timesteps[i + 1]
+                    timesteps_r = timesteps_r.repeat(latent_model_input.shape[0])
+                else:
+                    timesteps_r = None
+
+                timesteps_r_kwarg = self.prepare_extra_func_kwargs(
+                    self.transformer.forward,
+                    {
+                        "timestep_r": timesteps_r,
+                    },
+                )
+
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # Prepare inputs for transformer
+                guidance_expand = (torch.tensor(
+                    [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0],
+                    dtype=torch.float32,
+                    device=get_local_torch_device(),
+                ).to(target_dtype) * 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale is not None else None)
+
+                # Predict noise residual
+                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                    if (vsa_available and self.attn_backend == VideoSparseAttentionBackend) or \
+                       (sparse_fp4_available and self.attn_backend in sparse_fp4_backends):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls()
+
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls()
+                            # TODO(will): clean this up
+                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                                current_timestep=i,  # type: ignore
+                                raw_latent_shape=batch.raw_latent_shape[2:5],  # type: ignore
+                                patch_size=fastvideo_args.pipeline_config.  # type: ignore
+                                dit_config.patch_size,  # type: ignore
+                                VSA_sparsity=fastvideo_args.VSA_sparsity,  # type: ignore
+                                device=get_local_torch_device(),
+                            )
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    elif (vmoba_attn_available and self.attn_backend == VMOBAAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls()
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls()
+                            # Prepare V-MoBA parameters from config
+                            moba_params = fastvideo_args.moba_config.copy()
+                            assert batch.raw_latent_shape is not None, "raw_latent_shape must be set for V-MoBA"
+                            moba_params.update({
+                                "current_timestep": i,
+                                "raw_latent_shape": batch.raw_latent_shape[2:5],
+                                "patch_size": fastvideo_args.pipeline_config.dit_config.patch_size,
+                                "device": get_local_torch_device(),
+                            })
+                            attn_metadata = self.attn_metadata_builder.build(**moba_params)
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+                    # TODO(will): finalize the interface. vLLM uses this to
+                    # support torch dynamo compilation. They pass in
+                    # attn_metadata, vllm_config, and num_tokens. We can pass in
+                    # fastvideo_args or training_args, and attn_metadata.
+                    batch.is_cfg_negative = False
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            forward_batch=batch,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        noise_pred = current_model(
+                            latent_model_input,
+                            prompt_embeds,
+                            t_expand,
+                            guidance=guidance_expand,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                            **action_kwargs,
+                            **camera_kwargs,
+                            **timesteps_r_kwarg,
+                        )
+
+                    if batch.do_classifier_free_guidance:
+                        batch.is_cfg_negative = True
+                        with set_forward_context(
+                                current_timestep=i,
+                                attn_metadata=attn_metadata,
+                                forward_batch=batch,
+                        ):
+                            noise_pred_uncond = current_model(
+                                latent_model_input,
+                                neg_prompt_embeds,
+                                t_expand,
+                                guidance=guidance_expand,
+                                **image_kwargs,
+                                **neg_cond_kwargs,
+                                **action_kwargs,
+                                **camera_kwargs,
+                                **timesteps_r_kwarg,
+                            )
+
+                        noise_pred_text = noise_pred
+                        noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                        # Apply guidance rescale if needed
+                        if batch.guidance_rescale > 0.0:
+                            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                            noise_pred = self.rescale_noise_cfg(
+                                noise_pred,
+                                noise_pred_text,
+                                guidance_rescale=batch.guidance_rescale,
+                            )
+                    # Compute the previous noisy sample
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                        latents = latents.squeeze(0)
+                        latents = (1. - mask2[0]) * z + mask2[0] * latents
+                        # latents = latents.unsqueeze(0)
+
+                # save trajectory latents if needed
+                if batch.return_trajectory_latents:
+                    trajectory_timesteps.append(t)
+                    trajectory_latents.append(latents)
+
+                # Update progress bar
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and
+                                               (i + 1) % self.scheduler.order == 0 and progress_bar is not None):
+                    progress_bar.update()
+
+        trajectory_tensor: torch.Tensor | None = None
+        if trajectory_latents:
+            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
+            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0)
+        else:
+            trajectory_tensor = None
+            trajectory_timesteps_tensor = None
+
+        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
+            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
+
+        # Update batch with final latents
+        batch.latents = latents
+
+        if fastvideo_args.dit_layerwise_offload:
+            mgr = getattr(self.transformer, "_layerwise_offload_manager", None)
+            if mgr is not None and getattr(mgr, "enabled", False):
+                mgr.release_all()
+            if self.transformer_2 is not None:
+                mgr2 = getattr(self.transformer_2, "_layerwise_offload_manager", None)
+                if mgr2 is not None and getattr(mgr2, "enabled", False):
+                    mgr2.release_all()
+
+        # deallocate transformer if on mps
+        if torch.backends.mps.is_available():
+            logger.info("Memory before deallocating transformer: %s", torch.mps.current_allocated_memory())
+            del self.transformer
+            if pipeline is not None and "transformer" in pipeline.modules:
+                del pipeline.modules["transformer"]
+            fastvideo_args.model_loaded["transformer"] = False
+            logger.info("Memory after deallocating transformer: %s", torch.mps.current_allocated_memory())
+
+        return batch
+
+    def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, Any]:
+        """
+        Prepare extra kwargs for the scheduler step / denoise step.
+        
+        Args:
+            func: The function to prepare kwargs for.
+            kwargs: The kwargs to prepare.
+            
+        Returns:
+            The prepared kwargs.
+        """
+        extra_step_kwargs = {}
+        for k, v in kwargs.items():
+            accepts = k in set(inspect.signature(func).parameters.keys())
+            if accepts:
+                extra_step_kwargs[k] = v
+        return extra_step_kwargs
+
+    def progress_bar(self, iterable: Iterable | None = None, total: int | None = None) -> tqdm:
+        """
+        Create a progress bar for the denoising process.
+        
+        Args:
+            iterable: The iterable to iterate over.
+            total: The total number of items.
+            
+        Returns:
+            A tqdm progress bar.
+        """
+        local_rank = get_world_group().local_rank
+        if local_rank == 0:
+            return tqdm(iterable=iterable, total=total)
+        else:
+            return tqdm(iterable=iterable, total=total, disable=True)
+
+    def rescale_noise_cfg(self, noise_cfg, noise_pred_text, guidance_rescale=0.0) -> torch.Tensor:
+        """
+        Rescale noise prediction according to guidance_rescale.
+        
+        Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed"
+        (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4.
+        
+        Args:
+            noise_cfg: The noise prediction with guidance.
+            noise_pred_text: The text-conditioned noise prediction.
+            guidance_rescale: The guidance rescale factor.
+            
+        Returns:
+            The rescaled noise prediction.
+        """
+        std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+        std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+        # Rescale the results from guidance (fixes overexposure)
+        noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+        # Mix with the original results from guidance by factor guidance_rescale
+        noise_cfg = (guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg)
+        return noise_cfg
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check("image_embeds", batch.image_embeds, V.is_list)
+        result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5))
+        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
+        result.add_check("guidance_scale", batch.guidance_scale, V.positive_float)
+        result.add_check("eta", batch.eta, V.non_negative_float)
+        result.add_check("generator", batch.generator, V.generator_or_list_generators)
+        result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
+                         lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify denoising stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class CosmosDenoisingStage(DenoisingStage):
+    """
+    Denoising stage for Cosmos models using FlowMatchEulerDiscreteScheduler.
+    """
+
+    def __init__(self, transformer, scheduler, pipeline=None) -> None:
+        super().__init__(transformer, scheduler, pipeline)
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        pipeline = self.pipeline() if self.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        latents = batch.latents
+        num_inference_steps = batch.num_inference_steps
+        guidance_scale = batch.guidance_scale
+
+        sigma_max = 80.0
+        sigma_min = 0.002
+        sigma_data = 1.0
+        final_sigmas_type = "sigma_min"
+
+        if self.scheduler is not None:
+            self.scheduler.register_to_config(
+                sigma_max=sigma_max,
+                sigma_min=sigma_min,
+                sigma_data=sigma_data,
+                final_sigmas_type=final_sigmas_type,
+            )
+
+        self.scheduler.set_timesteps(num_inference_steps, device=latents.device)
+        timesteps = self.scheduler.timesteps
+
+        if (hasattr(self.scheduler.config, 'final_sigmas_type')
+                and self.scheduler.config.final_sigmas_type == "sigma_min" and len(self.scheduler.sigmas) > 1):
+            self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
+
+        conditioning_latents = getattr(batch, 'conditioning_latents', None)
+        unconditioning_latents = conditioning_latents
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+
+                current_sigma = self.scheduler.sigmas[i]
+                current_t = current_sigma / (current_sigma + 1)
+                c_in = 1 - current_t
+                c_skip = 1 - current_t
+                c_out = -current_t
+
+                timestep = current_t.view(1, 1, 1, 1, 1).expand(latents.size(0), -1, latents.size(2), -1,
+                                                                -1)  # [B, 1, T, 1, 1]
+
+                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+
+                    cond_latent = latents * c_in
+
+                    if hasattr(
+                            batch,
+                            'cond_indicator') and batch.cond_indicator is not None and conditioning_latents is not None:
+                        cond_latent = batch.cond_indicator * conditioning_latents + (1 -
+                                                                                     batch.cond_indicator) * cond_latent
+                    else:
+                        logger.warning(
+                            "Step %s: Missing conditioning data - cond_indicator: %s, conditioning_latents: %s", i,
+                            hasattr(batch, 'cond_indicator'), conditioning_latents is not None)
+
+                    cond_latent = cond_latent.to(target_dtype)
+
+                    cond_timestep = timestep
+                    if hasattr(batch, 'cond_indicator') and batch.cond_indicator is not None:
+                        sigma_conditioning = 0.0001
+                        t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
+                        cond_timestep = batch.cond_indicator * t_conditioning + (1 - batch.cond_indicator) * timestep
+                        cond_timestep = cond_timestep.to(target_dtype)
+
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=None,
+                            forward_batch=batch,
+                    ):
+                        # Use conditioning masks from CosmosLatentPreparationStage
+                        condition_mask = batch.cond_mask.to(target_dtype) if hasattr(batch, 'cond_mask') else None
+                        padding_mask = torch.zeros(1,
+                                                   1,
+                                                   batch.height,
+                                                   batch.width,
+                                                   device=cond_latent.device,
+                                                   dtype=target_dtype)
+
+                        # Fallback if masks not available
+                        if condition_mask is None:
+                            batch_size, num_channels, num_frames, height, width = cond_latent.shape
+                            condition_mask = torch.zeros(batch_size,
+                                                         1,
+                                                         num_frames,
+                                                         height,
+                                                         width,
+                                                         device=cond_latent.device,
+                                                         dtype=target_dtype)
+
+                        noise_pred = self.transformer(
+                            hidden_states=cond_latent,
+                            timestep=cond_timestep.to(target_dtype),
+                            encoder_hidden_states=batch.prompt_embeds[0].to(target_dtype),
+                            fps=24,  # TODO: get fps from batch or config
+                            condition_mask=condition_mask,
+                            padding_mask=padding_mask,
+                            return_dict=False,
+                        )[0]
+
+                    cond_pred = (c_skip * latents + c_out * noise_pred.float()).to(target_dtype)
+
+                    if hasattr(
+                            batch,
+                            'cond_indicator') and batch.cond_indicator is not None and conditioning_latents is not None:
+                        cond_pred = batch.cond_indicator * conditioning_latents + (1 - batch.cond_indicator) * cond_pred
+
+                    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds is not None:
+                        uncond_latent = latents * c_in
+
+                        if hasattr(batch, 'uncond_indicator'
+                                   ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
+                            uncond_latent = batch.uncond_indicator * unconditioning_latents + (
+                                1 - batch.uncond_indicator) * uncond_latent
+
+                        with set_forward_context(
+                                current_timestep=i,
+                                attn_metadata=None,
+                                forward_batch=batch,
+                        ):
+                            uncond_condition_mask = batch.uncond_mask.to(target_dtype) if hasattr(
+                                batch, 'uncond_mask') and batch.uncond_mask is not None else condition_mask
+
+                            uncond_timestep = timestep
+                            if hasattr(batch, 'uncond_indicator') and batch.uncond_indicator is not None:
+                                sigma_conditioning = 0.0001
+                                t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
+                                uncond_timestep = batch.uncond_indicator * t_conditioning + (
+                                    1 - batch.uncond_indicator) * timestep
+                                uncond_timestep = uncond_timestep.to(target_dtype)
+
+                            noise_pred_uncond = self.transformer(
+                                hidden_states=uncond_latent.to(target_dtype),
+                                timestep=uncond_timestep.to(target_dtype),
+                                encoder_hidden_states=batch.negative_prompt_embeds[0].to(target_dtype),
+                                fps=24,  # TODO: get fps from batch or config
+                                condition_mask=uncond_condition_mask,
+                                padding_mask=padding_mask,
+                                return_dict=False,
+                            )[0]
+
+                        uncond_pred = (c_skip * latents + c_out * noise_pred_uncond.float()).to(target_dtype)
+
+                        if hasattr(batch, 'uncond_indicator'
+                                   ) and batch.uncond_indicator is not None and unconditioning_latents is not None:
+                            uncond_pred = batch.uncond_indicator * unconditioning_latents + (
+                                1 - batch.uncond_indicator) * uncond_pred
+
+                        guidance_diff = cond_pred - uncond_pred
+                        final_pred = cond_pred + guidance_scale * guidance_diff
+                    else:
+                        final_pred = cond_pred
+
+                # Convert to noise for scheduler step
+                if current_sigma > 1e-8:
+                    noise_for_scheduler = (latents - final_pred) / current_sigma
+                else:
+                    logger.warning("Step %s: current_sigma too small (%s), using final_pred directly", i, current_sigma)
+                    noise_for_scheduler = final_pred
+
+                if torch.isnan(noise_for_scheduler).sum() > 0:
+                    logger.error("Step %s: NaN detected in noise_for_scheduler, sum: %s", i,
+                                 noise_for_scheduler.float().sum().item())
+                    logger.error("Step %s: latents sum: %s, final_pred sum: %s, current_sigma: %s", i,
+                                 latents.float().sum().item(),
+                                 final_pred.float().sum().item(), current_sigma)
+
+                latents = self.scheduler.step(noise_for_scheduler, t, latents, **extra_step_kwargs,
+                                              return_dict=False)[0]
+
+                progress_bar.update()
+
+        batch.latents = latents
+
+        return batch
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify Cosmos denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int)
+        result.add_check("guidance_scale", batch.guidance_scale, V.positive_float)
+        result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value)
+        result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds,
+                         lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x))
+        return result
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify Cosmos denoising stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
+        return result
+
+
+class Cosmos25DenoisingStage(CosmosDenoisingStage):
+    """Denoising stage for Cosmos 2.5 DiT (expects 1D/2D timestep, not 5D)."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        pipeline = self.pipeline() if self.pipeline else None
+        if not fastvideo_args.model_loaded["transformer"]:
+            loader = TransformerLoader()
+            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
+            if pipeline:
+                pipeline.add_module("transformer", self.transformer)
+            fastvideo_args.model_loaded["transformer"] = True
+
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step,
+            {
+                "generator": batch.generator,
+                "eta": batch.eta
+            },
+        )
+
+        if hasattr(self.transformer, 'module'):
+            transformer_dtype = next(self.transformer.module.parameters()).dtype
+        else:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+        target_dtype = transformer_dtype
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        latents = batch.latents
+        if latents is None:
+            raise ValueError("latents must be provided for Cosmos25DenoisingStage")
+        guidance_scale = batch.guidance_scale
+
+        if batch.timesteps is None:
+            self.scheduler.set_timesteps(batch.num_inference_steps, device=latents.device)
+            timesteps = self.scheduler.timesteps
+        else:
+            timesteps = batch.timesteps.to(latents.device)
+
+        cfg = fastvideo_args.pipeline_config
+
+        if batch.fps is None:
+            gen = batch.generator
+            if isinstance(gen, list) and len(gen) > 0:
+                gen = gen[0]
+            fps_tensor = torch.randint(
+                16,
+                32,
+                (1, ),
+                generator=gen if isinstance(gen, torch.Generator) else None,
+                device=latents.device,
+            ).float().to(dtype=target_dtype)
+        else:
+            fps_val = batch.fps
+            fps_tensor = torch.tensor(
+                [fps_val],
+                device=latents.device,
+                dtype=target_dtype,
+            )
+
+        latents_4d = latents[0]
+
+        # Masks are optional for T2W.
+        cond_mask = getattr(batch, "cond_mask", None)
+        condition_mask = cond_mask.to(target_dtype) if isinstance(cond_mask, torch.Tensor) else None
+        pad_mask = getattr(batch, "padding_mask", None)
+        padding_mask = pad_mask.to(target_dtype) if isinstance(pad_mask, torch.Tensor) else None
+
+        # Conditioning fields are attached by latent preparation stage.
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        cond_indicator = getattr(batch, "cond_indicator", None)
+        # Infer whether this is a conditioned run (V2W/I2W) purely from the presence
+        # of conditioning latents. Avoid carrying explicit mode flags on the batch.
+        is_conditioned = (conditioning_latents is not None)
+
+        init_noise_4d = latents_4d.clone()
+        if condition_mask is None:
+            _, t, h, w = latents_4d.shape
+            condition_mask = torch.zeros(1, 1, t, h, w, device=latents.device, dtype=target_dtype)
+        if padding_mask is None:
+            _, _, h, w = latents_4d.shape
+            padding_default = 0.0 if is_conditioned else 1.0
+            padding_mask = torch.full(
+                (1, 1, h, w),
+                float(padding_default),
+                device=latents.device,
+                dtype=target_dtype,
+            )
+
+        timestep_scale = 0.001
+
+        state_dtype = torch.float32
+
+        conditional_frame_timestep = 0.1
+        latents_4d = latents_4d.to(state_dtype)
+        init_noise_4d = init_noise_4d.to(state_dtype)
+
+        clamp_every_step = bool(getattr(cfg, "cosmos25_clamp_every_step", True)) if is_conditioned else False
+
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                t_val = float(t)
+                if is_conditioned:
+                    t_frames = int(latents_4d.shape[1])
+                    timestep = torch.full(
+                        (1, t_frames),
+                        float(t_val * timestep_scale),
+                        device=latents.device,
+                        dtype=torch.float32,
+                    )
+                    if cond_indicator is not None and t_frames > 0:
+                        cond_t = cond_indicator[0, 0, :t_frames, 0, 0]
+                        cond_mask_t = (cond_t > 0.5)
+                        if bool(cond_mask_t.any().item()):
+                            timestep[0, cond_mask_t] = float(conditional_frame_timestep)
+                else:
+                    timestep_val = t_val * timestep_scale
+                    timestep = torch.tensor(
+                        [[float(timestep_val)]],
+                        device=latents.device,
+                        dtype=target_dtype,
+                    )
+
+                # Conditioned runs: replace x_t with GT x0 on the conditioned frames.
+                if (is_conditioned and cond_indicator is not None and conditioning_latents is not None
+                        and (clamp_every_step or i == 0)):
+                    cond_ind_4d = cond_indicator[0].to(state_dtype)
+                    gt_x0 = conditioning_latents[0].to(state_dtype)
+                    latents_4d = gt_x0 * cond_ind_4d + latents_4d * (1 - cond_ind_4d)
+
+                model_hidden_states = latents_4d.unsqueeze(0)
+
+                with (
+                        set_forward_context(current_timestep=int(t_val), attn_metadata=None, forward_batch=batch),
+                        torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled),
+                ):
+                    cond_v = self.transformer(
+                        hidden_states=model_hidden_states.to(target_dtype),
+                        encoder_hidden_states=batch.prompt_embeds[0].to(target_dtype),
+                        timestep=timestep,
+                        fps=fps_tensor,
+                        condition_mask=condition_mask,
+                        padding_mask=padding_mask,
+                        return_dict=False,
+                    )[0]
+
+                    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
+                        uncond_v = self.transformer(
+                            hidden_states=model_hidden_states.to(target_dtype),
+                            encoder_hidden_states=batch.negative_prompt_embeds[0].to(target_dtype),
+                            timestep=timestep,
+                            fps=fps_tensor,
+                            condition_mask=condition_mask,
+                            padding_mask=padding_mask,
+                            return_dict=False,
+                        )[0]
+                        if is_conditioned:
+                            v = cond_v + guidance_scale * (cond_v - uncond_v)
+                        else:
+                            v = uncond_v + guidance_scale * (cond_v - uncond_v)
+                    else:
+                        v = cond_v
+
+                # Conditioned runs: replace velocity on conditioned frames with GT velocity.
+                if (is_conditioned and cond_indicator is not None and conditioning_latents is not None):
+                    cond_ind_4d = cond_indicator[0].to(state_dtype)
+                    gt_x0 = conditioning_latents[0].to(state_dtype)
+                    gt_v = init_noise_4d.to(state_dtype) - gt_x0
+                    v = cond_ind_4d * gt_v + (1 - cond_ind_4d) * v.to(state_dtype)
+
+                prev = self.scheduler.step(v.unsqueeze(0),
+                                           t,
+                                           latents_4d.unsqueeze(0),
+                                           **extra_step_kwargs,
+                                           return_dict=False)[0]
+                latents_4d = prev.squeeze(0)
+
+                progress_bar.update()
+
+        batch.latents = latents_4d.to(target_dtype).unsqueeze(0)
+        return batch
+
+
+class Cosmos25T2WDenoisingStage(Cosmos25DenoisingStage):
+    """Cosmos 2.5 Text2World denoising stage."""
+
+    _CONDITIONING_FIELDS = (
+        "conditioning_latents",
+        "cond_indicator",
+        "uncond_indicator",
+    )
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        for name in self._CONDITIONING_FIELDS:
+            if hasattr(batch, name):
+                setattr(batch, name, None)
+        return super().forward(batch, fastvideo_args)
+
+
+class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
+    """Cosmos 2.5 Video2World denoising stage."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        return super().forward(batch, fastvideo_args)
+
+
+class Cosmos25AutoDenoisingStage(PipelineStage):
+    """Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__()
+        self._t2w = Cosmos25T2WDenoisingStage(transformer=transformer, scheduler=scheduler)
+        self._v2w = Cosmos25V2WDenoisingStage(transformer=transformer, scheduler=scheduler)
+
+    def pipeline(self):
+        return self._v2w.pipeline() if self._v2w.pipeline else None
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.forward(batch, fastvideo_args)
+        return self._t2w.forward(batch, fastvideo_args)
+
+    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.verify_input(batch, fastvideo_args)
+        return self._t2w.verify_input(batch, fastvideo_args)
+
+    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
+        conditioning_latents = getattr(batch, "conditioning_latents", None)
+        if conditioning_latents is not None:
+            return self._v2w.verify_output(batch, fastvideo_args)
+        return self._t2w.verify_output(batch, fastvideo_args)
+
+
+class DmdDenoisingStage(DenoisingStage):
+    """
+    Denoising stage for DMD.
+    """
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__(transformer, scheduler)
+        self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Run the denoising loop.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with denoised latents.
+        """
+        # Setup precision and autocast settings
+        # TODO(will): make the precision configurable for inference
+        # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        # Get timesteps and calculate warmup steps
+        timesteps = batch.timesteps
+
+        # TODO(will): remove this once we add input/output validation for stages
+        if timesteps is None:
+            raise ValueError("Timesteps must be provided")
+        num_inference_steps = batch.num_inference_steps
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        # Prepare image latents and embeddings for I2V generation
+        image_embeds = batch.image_embeds
+        if len(image_embeds) > 0:
+            assert torch.isnan(image_embeds[0]).sum() == 0
+            image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds]
+
+        image_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_image": image_embeds,
+                "mask_strategy": dict_to_3d_list(None, t_max=50, l_max=60, h_max=24)
+            },
+        )
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        # Get latents and embeddings
+        assert batch.latents is not None, "latents must be provided"
+        latents = batch.latents
+
+        video_raw_latent_shape = latents.shape
+        prompt_embeds = batch.prompt_embeds
+        assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan"
+        timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps,
+                                 dtype=torch.long,
+                                 device=get_local_torch_device())
+
+        # Run denoising loop
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # Skip if interrupted
+                if hasattr(self, 'interrupt') and self.interrupt:
+                    continue
+                # Expand latents for I2V
+                noise_latents = latents.clone()
+                latent_model_input = latents.to(target_dtype)
+
+                if batch.image_latent is not None:
+                    latent_model_input = torch.cat(
+                        [latent_model_input, batch.image_latent.permute(0, 2, 1, 3, 4)], dim=2).to(target_dtype)
+                assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan"
+
+                # Prepare inputs for transformer
+                t_expand = t.repeat(latent_model_input.shape[0])
+                guidance_expand = (torch.tensor(
+                    [fastvideo_args.pipeline_config.embedded_cfg_scale] * latent_model_input.shape[0],
+                    dtype=torch.float32,
+                    device=get_local_torch_device(),
+                ).to(target_dtype) * 1000.0 if fastvideo_args.pipeline_config.embedded_cfg_scale is not None else None)
+
+                # Predict noise residual
+                with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                    if (vsa_available and self.attn_backend == VideoSparseAttentionBackend) or \
+                       (sparse_fp4_available and self.attn_backend in sparse_fp4_backends):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls()
+
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls()
+                            # TODO(will): clean this up
+                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                                current_timestep=i,  # type: ignore
+                                raw_latent_shape=batch.raw_latent_shape[2:5],  # type: ignore
+                                patch_size=fastvideo_args.pipeline_config.  # type: ignore
+                                dit_config.patch_size,  # type: ignore
+                                VSA_sparsity=fastvideo_args.VSA_sparsity,  # type: ignore
+                                device=get_local_torch_device(),  # type: ignore
+                            )  # type: ignore
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+
+                    batch.is_cfg_negative = False
+                    with set_forward_context(
+                            current_timestep=i,
+                            attn_metadata=attn_metadata,
+                            forward_batch=batch,
+                            # fastvideo_args=fastvideo_args
+                    ):
+                        # Run transformer
+                        pred_noise = self.transformer(
+                            latent_model_input.permute(0, 2, 1, 3, 4),
+                            prompt_embeds,
+                            t_expand,
+                            guidance=guidance_expand,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        ).permute(0, 2, 1, 3, 4)
+
+                    pred_video = pred_noise_to_pred_video(pred_noise=pred_noise.flatten(0, 1),
+                                                          noise_input_latent=noise_latents.flatten(0, 1),
+                                                          timestep=t_expand,
+                                                          scheduler=self.scheduler).unflatten(0, pred_noise.shape[:2])
+
+                    if i < len(timesteps) - 1:
+                        next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video.device)
+                        noise_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
+                        noise = torch.randn(video_raw_latent_shape, dtype=pred_video.dtype,
+                                            generator=noise_generator).to(self.device)
+                        latents = self.scheduler.add_noise(pred_video.flatten(0, 1), noise.flatten(0, 1),
+                                                           next_timestep).unflatten(0, pred_video.shape[:2])
+                    else:
+                        latents = pred_video
+
+                    # Update progress bar
+                    if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and
+                                                   (i + 1) % self.scheduler.order == 0 and progress_bar is not None):
+                        progress_bar.update()
+
+        # Gather results if using sequence parallelism
+        latents = latents.permute(0, 2, 1, 3, 4)
+        # Update batch with final latents
+        batch.latents = latents
+
+        return batch

backend_snapshot/fastvideo/platforms/cuda.py

@@ -0,0 +1,440 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/platforms/cuda.py
+"""Code inside this file can safely assume cuda platform, e.g. importing
+pynvml. However, it should not initialize cuda context.
+"""
+
+import os
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from typing import TypeVar
+
+import torch
+from typing_extensions import ParamSpec
+
+import fastvideo.envs as envs
+from fastvideo.logger import init_logger
+from fastvideo.platforms.interface import (AttentionBackendEnum, DeviceCapability, Platform, PlatformEnum)
+from fastvideo.utils import import_pynvml
+
+logger = init_logger(__name__)
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+pynvml = import_pynvml()  # type: ignore[no-untyped-call]
+
+# pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
+# see https://github.com/huggingface/diffusers/issues/9704 for details
+torch.backends.cuda.enable_cudnn_sdp(False)
+
+
+def device_id_to_physical_device_id(device_id: int) -> int:
+    if "CUDA_VISIBLE_DEVICES" in os.environ:
+        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
+        if device_ids == [""]:
+            msg = ("CUDA_VISIBLE_DEVICES is set to empty string, which means"
+                   " GPU support is disabled. If you are using ray, please unset"
+                   " the environment variable `CUDA_VISIBLE_DEVICES` inside the"
+                   " worker/actor. "
+                   "Check https://github.com/vllm-project/vllm/issues/8402 for"
+                   " more information.")
+            raise RuntimeError(msg)
+        physical_device_id = device_ids[device_id]
+        return int(physical_device_id)
+    else:
+        return device_id
+
+
+def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
+
+    @wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        pynvml.nvmlInit()
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            pynvml.nvmlShutdown()
+
+    return wrapper
+
+
+class CudaPlatformBase(Platform):
+    _enum = PlatformEnum.CUDA
+    device_name: str = "cuda"
+    device_type: str = "cuda"
+    dispatch_key: str = "CUDA"
+    ray_device_key: str = "GPU"
+    device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        if enforce_eager:
+            logger.warning("To see benefits of async output processing, enable CUDA "
+                           "graph. Since, enforce-eager is enabled, async output "
+                           "processor cannot be used")
+            return False
+        return True
+
+    @classmethod
+    def is_full_nvlink(cls, device_ids: list[int]) -> bool:
+        raise NotImplementedError
+
+    @classmethod
+    def log_warnings(cls) -> None:
+        pass
+
+    @classmethod
+    def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> float:
+        torch.cuda.reset_peak_memory_stats(device)
+        return float(torch.cuda.max_memory_allocated(device))
+
+    @classmethod
+    def get_torch_device(cls) -> object:
+        """
+        Return torch.cuda
+        """
+        return torch.cuda
+
+    @classmethod
+    def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
+                             dtype: torch.dtype) -> str:
+        # TODO(will): maybe come up with a more general interface for local attention
+        # if distributed is False, we always try to use Flash attn
+
+        logger.info("Trying FASTVIDEO_ATTENTION_BACKEND=%s", envs.FASTVIDEO_ATTENTION_BACKEND)
+        logger.info("Selected backend: %s", selected_backend)
+        if selected_backend == AttentionBackendEnum.SAGE_ATTN:
+            try:
+                from sageattention import sageattn  # noqa: F401
+
+                from fastvideo.attention.backends.sage_attn import (  # noqa: F401
+                    SageAttentionBackend)
+                logger.info("Using Sage Attention backend.")
+
+                return "fastvideo.attention.backends.sage_attn.SageAttentionBackend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info("Sage Attention backend is not installed. Fall back to Flash Attention.")
+        elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE:
+            try:
+                from sageattn3 import sageattn3_blackwell  # noqa: F401
+
+                from fastvideo.attention.backends.sage_attn3 import (  # noqa: F401
+                    SageAttention3Backend)
+                logger.info("Using Sage Attention 3 backend.")
+
+                return "fastvideo.attention.backends.sage_attn3.SageAttention3Backend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info("Sage Attention 3 backend is not installed. Fall back to Flash Attention.")
+        elif selected_backend == AttentionBackendEnum.ATTN_QAT_INFER:
+            try:
+                from fastvideo.attention.backends.attn_qat_infer import (  # noqa: F401
+                    AttnQatInferBackend, is_attn_qat_infer_available,
+                )
+                if not is_attn_qat_infer_available():
+                    raise ImportError("attn_qat_infer could not be imported.")
+                logger.info("Using attn_qat_infer backend.")
+
+                return "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info("attn_qat_infer backend is not installed. Fall back to Flash Attention.")
+        elif selected_backend == AttentionBackendEnum.ATTN_QAT_TRAIN:
+            try:
+                from fastvideo_kernel.triton_kernels.attn_qat_train import attention  # noqa: F401
+
+                from fastvideo.attention.backends.attn_qat_train import (  # noqa: F401
+                    AttnQatTrainBackend)
+                logger.info("Using attn_qat_train backend.")
+
+                return "fastvideo.attention.backends.attn_qat_train.AttnQatTrainBackend"
+            except ImportError as e:
+                logger.info(e)
+                logger.info("attn_qat_train backend is not installed. Fall back to Flash Attention.")
+        elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN:
+            try:
+                from fastvideo_kernel import video_sparse_attn  # noqa: F401
+
+                from fastvideo.attention.backends.video_sparse_attn import (  # noqa: F401
+                    VideoSparseAttentionBackend)
+                logger.info("Using Video Sparse Attention backend.")
+
+                return "fastvideo.attention.backends.video_sparse_attn.VideoSparseAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import Video Sparse Attention backend: %s", str(e))
+                raise ImportError("The Video Sparse Attention backend is not installed. "
+                                  "To install it, please follow the instructions at: "
+                                  "https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation ") from e
+        elif selected_backend == AttentionBackendEnum.SPARSE_FP4_ATTN:
+            try:
+                from fastvideo.attention.backends.sparse_fp4_attn import (  # noqa: F401
+                    SparseFP4AttentionBackend)
+                logger.info("Using Sparse FP4 Attention backend (FP4 quant + VSA).")
+                return "fastvideo.attention.backends.sparse_fp4_attn.SparseFP4AttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import Sparse FP4 Attention backend: %s", str(e))
+                raise ImportError("Sparse FP4 Attention backend is not available.") from e
+        elif selected_backend == AttentionBackendEnum.SPARSE_FP4_OURS_P_ATTN:
+            try:
+                from fastvideo.attention.backends.sparse_fp4_ours_p_attn import (  # noqa: F401
+                    SparseFP4OursPAttentionBackend)
+                logger.info(
+                    "Using Sparse FP4 Ours-P Attention backend (group-local P quant + VSA)."
+                )
+                return "fastvideo.attention.backends.sparse_fp4_ours_p_attn.SparseFP4OursPAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import Sparse FP4 Ours-P Attention backend: %s", str(e))
+                raise ImportError("Sparse FP4 Ours-P Attention backend is not available.") from e
+        elif selected_backend == AttentionBackendEnum.BSA_ATTN:
+            try:
+                from fastvideo.attention.backends.bsa_attn import (  # noqa: F401
+                    BSAAttentionBackend)
+                logger.info("Using BSA Attention backend.")
+
+                return "fastvideo.attention.backends.bsa_attn.BSAAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import BSA Attention backend: %s", str(e))
+                raise ImportError("The BSA Attention backend failed to import.") from e
+        elif selected_backend == AttentionBackendEnum.VMOBA_ATTN:
+            try:
+                from fastvideo_kernel import moba_attn_varlen  # noqa: F401
+                from fastvideo.attention.backends.vmoba import (  # noqa: F401
+                    VMOBAAttentionBackend)
+                logger.info("Using Video MOBA Attention backend.")
+
+                return "fastvideo.attention.backends.vmoba.VMOBAAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import Video MoBA Attention backend: %s", str(e))
+                raise ImportError("Video MoBA Attention backend is not installed. ") from e
+        elif selected_backend == AttentionBackendEnum.SLA_ATTN:
+            try:
+                from fastvideo.attention.backends.sla import (  # noqa: F401
+                    SLAAttentionBackend)
+                logger.info("Using SLA (Sparse-Linear Attention) backend.")
+
+                return "fastvideo.attention.backends.sla.SLAAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import SLA Attention backend: %s", str(e))
+                raise ImportError("SLA Attention backend is not available. ") from e
+        elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN:
+            try:
+                from fastvideo.attention.backends.sla import (  # noqa: F401
+                    SageSLAAttentionBackend)
+                logger.info("Using SageSLA (Quantized Sparse-Linear Attention) backend.")
+
+                return "fastvideo.attention.backends.sla.SageSLAAttentionBackend"
+            except ImportError as e:
+                logger.error("Failed to import SageSLA Attention backend: %s", str(e))
+                raise ImportError("SageSLA Attention backend requires spas_sage_attn. "
+                                  "Install with: pip install git+https://github.com/thu-ml/SpargeAttn.git") from e
+        elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend.")
+            return "fastvideo.attention.backends.sdpa.SDPABackend"
+        elif selected_backend == AttentionBackendEnum.FLASH_ATTN or selected_backend is None:
+            pass
+        elif selected_backend:
+            raise ValueError(f"Invalid attention backend for {cls.device_name}")
+
+        target_backend = AttentionBackendEnum.FLASH_ATTN
+        if not cls.has_device_capability(80):
+            logger.info("Cannot use FlashAttention-2 backend for Volta and Turing "
+                        "GPUs.")
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+        elif dtype not in (torch.float16, torch.bfloat16):
+            logger.info("Cannot use FlashAttention-2 backend for dtype other than "
+                        "torch.float16 or torch.bfloat16.")
+            target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        # FlashAttn is valid for the model, checking if the package is
+        # installed.
+        if target_backend == AttentionBackendEnum.FLASH_ATTN:
+            try:
+                import flash_attn  # noqa: F401
+
+                from fastvideo.attention.backends.flash_attn import (  # noqa: F401
+                    FlashAttentionBackend)
+
+                supported_sizes = \
+                    FlashAttentionBackend.get_supported_head_sizes()
+                if head_size not in supported_sizes:
+                    logger.info("Cannot use FlashAttention-2 backend for head size %d.", head_size)
+                    target_backend = AttentionBackendEnum.TORCH_SDPA
+            except ImportError:
+                logger.info("Cannot use FlashAttention-2 backend because the "
+                            "flash_attn package is not found. "
+                            "Make sure that flash_attn was built and installed "
+                            "(on by default).")
+                target_backend = AttentionBackendEnum.TORCH_SDPA
+
+        if target_backend == AttentionBackendEnum.TORCH_SDPA:
+            logger.info("Using Torch SDPA backend.")
+
+            return "fastvideo.attention.backends.sdpa.SDPABackend"
+
+        logger.info("Using Flash Attention backend.")
+
+        return "fastvideo.attention.backends.flash_attn.FlashAttentionBackend"
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "fastvideo.distributed.device_communicators.cuda_communicator.CudaCommunicator"  # noqa
+
+
+# NVML utils
+# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
+# all the related functions work on real physical device ids.
+# the major benefit of using NVML is that it will not initialize CUDA
+class NvmlCudaPlatform(CudaPlatformBase):
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
+        try:
+            physical_device_id = device_id_to_physical_device_id(device_id)
+            handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+            return DeviceCapability(major=major, minor=minor)
+        except RuntimeError:
+            return None
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        try:
+            return bool(super().has_device_capability(capability, device_id))
+        except RuntimeError:
+            return False
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_name(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        return cls._get_physical_device_name(physical_device_id)
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return str(pynvml.nvmlDeviceGetUUID(handle))
+
+    @classmethod
+    @lru_cache(maxsize=8)
+    @with_nvml_context
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        physical_device_id = device_id_to_physical_device_id(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(physical_device_id)
+        return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total)
+
+    @classmethod
+    @with_nvml_context
+    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
+        """
+        query if the set of gpus are fully connected by nvlink (1 hop)
+        """
+        handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids]
+        for i, handle in enumerate(handles):
+            for j, peer_handle in enumerate(handles):
+                if i < j:
+                    try:
+                        p2p_status = pynvml.nvmlDeviceGetP2PStatus(
+                            handle,
+                            peer_handle,
+                            pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
+                        )
+                        if p2p_status != pynvml.NVML_P2P_STATUS_OK:
+                            return False
+                    except pynvml.NVMLError:
+                        logger.exception("NVLink detection failed. This is normal if"
+                                         " your machine has no NVLink equipped.")
+                        return False
+        return True
+
+    @classmethod
+    def _get_physical_device_name(cls, device_id: int = 0) -> str:
+        handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
+        return str(pynvml.nvmlDeviceGetName(handle))
+
+    @classmethod
+    @with_nvml_context
+    def log_warnings(cls) -> None:
+        device_ids: int = pynvml.nvmlDeviceGetCount()
+        if device_ids > 1:
+            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
+            if (len(set(device_names)) > 1 and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID"):
+                logger.warning(
+                    "Detected different devices in the system: %s. Please"
+                    " make sure to set `CUDA_DEVICE_ORDER=PCI_BUS_ID` to "
+                    "avoid unexpected behavior.",
+                    ", ".join(device_names),
+                )
+
+
+class NonNvmlCudaPlatform(CudaPlatformBase):
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> DeviceCapability:
+        major, minor = torch.cuda.get_device_capability(device_id)
+        return DeviceCapability(major=major, minor=minor)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        return str(torch.cuda.get_device_name(device_id))
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        device_props = torch.cuda.get_device_properties(device_id)
+        return int(device_props.total_memory)
+
+    @classmethod
+    def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
+        logger.exception("NVLink detection not possible, as context support was"
+                         " not found. Assuming no NVLink available.")
+        return False
+
+
+# Autodetect either NVML-enabled or non-NVML platform
+# based on whether NVML is available.
+nvml_available = False
+try:
+    try:
+        pynvml.nvmlInit()
+        nvml_available = True
+    except Exception:
+        # On Jetson, NVML is not supported.
+        nvml_available = False
+finally:
+    if nvml_available:
+        pynvml.nvmlShutdown()
+
+CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
+
+try:
+    from sphinx.ext.autodoc.mock import _MockModule
+
+    if not isinstance(pynvml, _MockModule):
+        CudaPlatform.log_warnings()
+except ModuleNotFoundError:
+    CudaPlatform.log_warnings()

backend_snapshot/fastvideo/platforms/interface.py

@@ -0,0 +1,255 @@
+import enum
+import random
+from typing import Any, NamedTuple
+
+import numpy as np
+import torch
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class AttentionBackendEnum(enum.Enum):
+    FLASH_ATTN = enum.auto()
+    TORCH_SDPA = enum.auto()
+    SAGE_ATTN = enum.auto()
+    SAGE_ATTN_THREE = enum.auto()
+    ATTN_QAT_INFER = enum.auto()
+    ATTN_QAT_TRAIN = enum.auto()
+    VIDEO_SPARSE_ATTN = enum.auto()
+    BSA_ATTN = enum.auto()
+    VMOBA_ATTN = enum.auto()
+    SLA_ATTN = enum.auto()
+    SAGE_SLA_ATTN = enum.auto()
+    SPARSE_FP4_ATTN = enum.auto()
+    SPARSE_FP4_OURS_P_ATTN = enum.auto()
+    NO_ATTENTION = enum.auto()
+
+
+class PlatformEnum(enum.Enum):
+    CUDA = enum.auto()
+    ROCM = enum.auto()
+    TPU = enum.auto()
+    XPU = enum.auto()
+    CPU = enum.auto()
+    MPS = enum.auto()
+    OOT = enum.auto()
+    UNSPECIFIED = enum.auto()
+    NPU = enum.auto()
+
+
+class CpuArchEnum(enum.Enum):
+    X86 = enum.auto()
+    ARM = enum.auto()
+    UNSPECIFIED = enum.auto()
+
+
+class DeviceCapability(NamedTuple):
+    major: int
+    minor: int
+
+    def as_version_str(self) -> str:
+        return f"{self.major}.{self.minor}"
+
+    def to_int(self) -> int:
+        """
+        Express device capability as an integer ``<major><minor>``.
+
+        It is assumed that the minor version is always a single digit.
+        """
+        assert 0 <= self.minor < 10
+        return self.major * 10 + self.minor
+
+
+class Platform:
+    _enum: PlatformEnum
+    device_name: str
+    device_type: str
+
+    dispatch_key: str = "CPU"
+
+    # platform-agnostic way to specify the device control environment variable,
+    # .e.g. CUDA_VISIBLE_DEVICES for CUDA.
+    # hint: search for "get_visible_accelerator_ids_env_var" in
+    # https://github.com/ray-project/ray/tree/master/python/ray/_private/accelerators # noqa
+    device_control_env_var: str = "FASTVIDEO_DEVICE_CONTROL_ENV_VAR_PLACEHOLDER"
+
+    # available ray device keys:
+    # https://github.com/ray-project/ray/blob/10ba5adadcc49c60af2c358a33bb943fb491a171/python/ray/_private/ray_constants.py#L438 # noqa
+    # empty string means the device does not support ray
+    ray_device_key: str = ""
+    # The torch.compile backend for compiling simple and
+    # standalone functions. The default value is "inductor" to keep
+    # the same behavior as PyTorch.
+    # NOTE: for the forward part of the model, vLLM has another separate
+    # compilation strategy.
+    simple_compile_backend: str = "inductor"
+
+    supported_quantization: list[str] = []
+
+    additional_env_vars: list[str] = []
+
+    def is_cuda(self) -> bool:
+        return self._enum == PlatformEnum.CUDA
+
+    def is_rocm(self) -> bool:
+        return self._enum == PlatformEnum.ROCM
+
+    def is_tpu(self) -> bool:
+        return self._enum == PlatformEnum.TPU
+
+    def is_xpu(self) -> bool:
+        return self._enum == PlatformEnum.XPU
+
+    def is_cpu(self) -> bool:
+        return self._enum == PlatformEnum.CPU
+
+    def is_out_of_tree(self) -> bool:
+        return self._enum == PlatformEnum.OOT
+
+    def is_cuda_alike(self) -> bool:
+        """Stateless version of :func:`torch.cuda.is_available`."""
+        return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
+
+    def is_mps(self) -> bool:
+        return self._enum == PlatformEnum.MPS
+
+    def is_npu(self) -> bool:
+        return self._enum == PlatformEnum.NPU
+
+    @classmethod
+    def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int,
+                             dtype: torch.dtype) -> str:
+        """Get the attention backend class of a device."""
+        return ""
+
+    @classmethod
+    def get_device_capability(
+        cls,
+        device_id: int = 0,
+    ) -> DeviceCapability | None:
+        """Stateless version of :func:`torch.cuda.get_device_capability`."""
+        return None
+
+    @classmethod
+    def has_device_capability(
+        cls,
+        capability: tuple[int, int] | int,
+        device_id: int = 0,
+    ) -> bool:
+        """
+        Test whether this platform is compatible with a device capability.
+
+        The ``capability`` argument can either be:
+
+        - A tuple ``(major, minor)``.
+        - An integer ``<major><minor>``. (See :meth:`DeviceCapability.to_int`)
+        """
+        current_capability = cls.get_device_capability(device_id=device_id)
+        if current_capability is None:
+            return False
+
+        if isinstance(capability, tuple):
+            return current_capability >= capability
+
+        return current_capability.to_int() >= capability
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        """Get the name of a device."""
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_uuid(cls, device_id: int = 0) -> str:
+        """Get the uuid of a device, e.g. the PCI bus ID."""
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_total_memory(cls, device_id: int = 0) -> int:
+        """Get the total memory of a device in bytes."""
+        raise NotImplementedError
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: bool | None) -> bool:
+        """
+        Check if the current platform supports async output.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_torch_device(cls) -> Any:
+        """
+        Check if the current platform supports torch device.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def inference_mode(cls):
+        """A device-specific wrapper of `torch.inference_mode`.
+
+        This wrapper is recommended because some hardware backends such as TPU
+        do not support `torch.inference_mode`. In such a case, they will fall
+        back to `torch.no_grad` by overriding this method.
+        """
+        return torch.inference_mode(mode=True)
+
+    @classmethod
+    def seed_everything(cls, seed: int | None = None) -> None:
+        """
+        Set the seed of each random module.
+        `torch.manual_seed` will set seed on all devices.
+
+        Loosely based on: https://github.com/Lightning-AI/pytorch-lightning/blob/2.4.0/src/lightning/fabric/utilities/seed.py#L20
+        """
+        if seed is not None:
+            random.seed(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+
+    @classmethod
+    def verify_model_arch(cls, model_arch: str) -> None:
+        """
+        Verify whether the current platform supports the specified model
+        architecture.
+
+        - This will raise an Error or Warning based on the model support on
+        the current platform.
+        - By default all models are considered supported.
+        """
+        pass
+
+    @classmethod
+    def verify_quantization(cls, quant: str) -> None:
+        """
+        Verify whether the quantization is supported by the current platform.
+        """
+        if cls.supported_quantization and \
+            quant not in cls.supported_quantization:
+            raise ValueError(f"{quant} quantization is currently not supported in "
+                             f"{cls.device_name}.")
+
+    @classmethod
+    def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> float:
+        """
+        Return the memory usage in bytes.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        """
+        Get device specific communicator class for distributed communication.
+        """
+        return "fastvideo.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase"  # noqa
+
+    @classmethod
+    def get_cpu_architecture(cls) -> CpuArchEnum:
+        """Get the CPU architecture of the current platform."""
+        return CpuArchEnum.UNSPECIFIED
+
+
+class UnspecifiedPlatform(Platform):
+    _enum = PlatformEnum.UNSPECIFIED
+    device_type = ""

backend_snapshot/fastvideo/train/models/wan/wan.py

@@ -0,0 +1,680 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan model plugin (per-role instance)."""
+
+from __future__ import annotations
+
+import copy
+import gc
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.distributed import (
+    get_sp_group,
+    get_world_group,
+)
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.pipelines.basic.wan.wan_pipeline import (
+    WanPipeline, )
+from fastvideo.pipelines.pipeline_batch_info import (
+    ForwardBatch, )
+from fastvideo.training.activation_checkpoint import (
+    apply_activation_checkpointing, )
+from fastvideo.training.training_utils import (
+    compute_density_for_timestep_sampling,
+    get_sigmas,
+    normalize_dit_input,
+    shift_timestep,
+)
+from fastvideo.utils import (
+    is_vmoba_available,
+    is_vsa_available,
+)
+
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.module_state import (
+    apply_trainable, )
+from fastvideo.train.utils.moduleloader import (
+    load_module_from_path, )
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+VideoSparseAttentionMetadataBuilder: type[Any] | None
+VideoMobaAttentionMetadataBuilder: type[Any] | None
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionMetadataBuilder as _VideoSparseAttentionMetadataBuilder, )
+    from fastvideo.attention.backends.vmoba import (
+        VideoMobaAttentionMetadataBuilder as _VideoMobaAttentionMetadataBuilder, )
+    VideoSparseAttentionMetadataBuilder = _VideoSparseAttentionMetadataBuilder
+    VideoMobaAttentionMetadataBuilder = _VideoMobaAttentionMetadataBuilder
+except Exception:
+    VideoSparseAttentionMetadataBuilder = None
+    VideoMobaAttentionMetadataBuilder = None
+
+
+class WanModel(ModelBase):
+    """Wan per-role model: owns transformer + noise_scheduler."""
+
+    _transformer_cls_name: str = "WanTransformer3DModel"
+
+    def __init__(
+        self,
+        *,
+        init_from: str,
+        training_config: TrainingConfig,
+        trainable: bool = True,
+        disable_custom_init_weights: bool = False,
+        flow_shift: float = 3.0,
+        enable_gradient_checkpointing_type: str
+        | None = None,
+        transformer_override_safetensor: str
+        | None = None,
+    ) -> None:
+        self._init_from = str(init_from)
+        self._trainable = bool(trainable)
+
+        self.transformer = self._load_transformer(
+            init_from=self._init_from,
+            trainable=self._trainable,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+            training_config=training_config,
+            transformer_override_safetensor=(transformer_override_safetensor),
+        )
+
+        self.noise_scheduler = (FlowMatchEulerDiscreteScheduler(shift=float(flow_shift)))
+
+        # Filled by init_preprocessors (student only).
+        self.vae: Any = None
+        self.training_config: TrainingConfig = training_config
+        self.dataloader: Any = None
+        self.validator: Any = None
+        self.start_step: int = 0
+
+        self.world_group: Any = None
+        self.sp_group: Any = None
+
+        self.negative_prompt_embeds: (torch.Tensor | None) = None
+        self.negative_prompt_attention_mask: (torch.Tensor | None) = None
+
+        # Timestep mechanics.
+        self.timestep_shift: float = float(flow_shift)
+        self.num_train_timestep: int = int(self.noise_scheduler.num_train_timesteps)
+        self.min_timestep: int = 0
+        self.max_timestep: int = self.num_train_timestep
+
+    def _load_transformer(
+        self,
+        *,
+        init_from: str,
+        trainable: bool,
+        disable_custom_init_weights: bool,
+        enable_gradient_checkpointing_type: str | None,
+        training_config: TrainingConfig,
+        transformer_override_safetensor: str | None = None,
+    ) -> torch.nn.Module:
+        transformer = load_module_from_path(
+            model_path=init_from,
+            module_type="transformer",
+            training_config=training_config,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            override_transformer_cls_name=(self._transformer_cls_name),
+            transformer_override_safetensor=(transformer_override_safetensor),
+        )
+        transformer = apply_trainable(transformer, trainable=trainable)
+        # Fall back to training_config.model if not set on the
+        # model YAML section directly.
+        ckpt_type = (enable_gradient_checkpointing_type or getattr(
+            getattr(training_config, "model", None),
+            "enable_gradient_checkpointing_type",
+            None,
+        ))
+        if trainable and ckpt_type:
+            transformer = apply_activation_checkpointing(
+                transformer,
+                checkpointing_type=ckpt_type,
+            )
+        return transformer
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def init_preprocessors(self, training_config: TrainingConfig) -> None:
+        self.vae = load_module_from_path(
+            model_path=str(training_config.model_path),
+            module_type="vae",
+            training_config=training_config,
+        )
+
+        self.world_group = get_world_group()
+        self.sp_group = get_sp_group()
+
+        self._init_timestep_mechanics()
+
+        from fastvideo.dataset.dataloader.schema import (
+            pyarrow_schema_t2v, )
+        from fastvideo.train.utils.dataloader import (
+            build_parquet_t2v_train_dataloader, )
+
+        text_len = (
+            training_config.pipeline_config.text_encoder_configs[  # type: ignore[union-attr]
+                0].arch_config.text_len)
+        self.dataloader = build_parquet_t2v_train_dataloader(
+            training_config.data,
+            text_len=int(text_len),
+            parquet_schema=pyarrow_schema_t2v,
+        )
+        self.start_step = 0
+
+    @property
+    def num_train_timesteps(self) -> int:
+        return int(self.num_train_timestep)
+
+    def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
+        timestep = shift_timestep(
+            timestep,
+            self.timestep_shift,
+            self.num_train_timestep,
+        )
+        return timestep.clamp(self.min_timestep, self.max_timestep)
+
+    def on_train_start(self) -> None:
+        self.ensure_negative_conditioning()
+
+    # ------------------------------------------------------------------
+    # Runtime primitives
+    # ------------------------------------------------------------------
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        self.ensure_negative_conditioning()
+        assert self.training_config is not None
+        tc = self.training_config
+
+        dtype = self._get_training_dtype()
+        device = self.device
+
+        training_batch = TrainingBatch()
+        encoder_hidden_states = raw_batch["text_embedding"]
+        encoder_attention_mask = raw_batch["text_attention_mask"]
+        infos = raw_batch.get("info_list")
+
+        if latents_source == "zeros":
+            batch_size = encoder_hidden_states.shape[0]
+            vae_config = (
+                tc.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
+            )
+            num_channels = vae_config.z_dim
+            spatial_compression_ratio = (vae_config.spatial_compression_ratio)
+            latent_height = (tc.data.num_height // spatial_compression_ratio)
+            latent_width = (tc.data.num_width // spatial_compression_ratio)
+            latents = torch.zeros(
+                batch_size,
+                num_channels,
+                tc.data.num_latent_t,
+                latent_height,
+                latent_width,
+                device=device,
+                dtype=dtype,
+            )
+        elif latents_source == "data":
+            if "vae_latent" not in raw_batch:
+                raise ValueError("vae_latent not found in batch "
+                                 "and latents_source='data'")
+            latents = raw_batch["vae_latent"]
+            latents = latents[:, :, :tc.data.num_latent_t]
+            latents = latents.to(device, dtype=dtype)
+        else:
+            raise ValueError(f"Unknown latents_source: "
+                             f"{latents_source!r}")
+
+        training_batch.latents = latents
+        training_batch.encoder_hidden_states = (encoder_hidden_states.to(device, dtype=dtype))
+        training_batch.encoder_attention_mask = (encoder_attention_mask.to(device, dtype=dtype))
+        training_batch.infos = infos
+
+        training_batch.latents = normalize_dit_input("wan", training_batch.latents, self.vae)
+        training_batch = self._prepare_dit_inputs(training_batch, generator)
+        training_batch = self._build_attention_metadata(training_batch)
+
+        training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata)
+        if training_batch.attn_metadata is not None:
+            training_batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore[attr-defined]
+
+        return training_batch
+
+    def add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        b, t = clean_latents.shape[:2]
+        noisy = self.noise_scheduler.add_noise(
+            clean_latents.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep,
+        ).unflatten(0, (b, t))
+        return noisy
+
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+        force_dense: bool = False,
+    ) -> torch.Tensor:
+        device_type = self.device.type
+        dtype = noisy_latents.dtype
+        if conditional:
+            text_dict = batch.conditional_dict
+            if text_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+        else:
+            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
+
+        if attn_kind == "dense":
+            attn_metadata = batch.attn_metadata
+        elif attn_kind in ("vsa", "sparse_fp4"):
+            attn_metadata = batch.attn_metadata_vsa
+        else:
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+
+        with torch.autocast(device_type, dtype=dtype), set_forward_context(
+                current_timestep=batch.timesteps,
+                attn_metadata=attn_metadata,
+                force_dense=force_dense,
+        ):
+            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
+            transformer = self._get_transformer(timestep)
+            pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
+        return pred_noise
+
+    def backward(
+        self,
+        loss: torch.Tensor,
+        ctx: Any,
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        timesteps, attn_metadata = ctx
+        with set_forward_context(
+                current_timestep=timesteps,
+                attn_metadata=attn_metadata,
+        ):
+            (loss / max(1, int(grad_accum_rounds))).backward()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _get_training_dtype(self) -> torch.dtype:
+        return torch.bfloat16
+
+    def _init_timestep_mechanics(self) -> None:
+        assert self.training_config is not None
+        tc = self.training_config
+        flow_shift = tc.pipeline_config.flow_shift
+        self.timestep_shift = float(0.0 if flow_shift is None else flow_shift)
+        self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps)
+        # min/max timestep ratios now come from method_config;
+        # default to full range.
+        self.min_timestep = 0
+        self.max_timestep = self.num_train_timestep
+
+    def ensure_negative_conditioning(self) -> None:
+        if self.negative_prompt_embeds is not None:
+            return
+
+        assert self.training_config is not None
+        tc = self.training_config
+        world_group = self.world_group
+        device = self.device
+        dtype = self._get_training_dtype()
+
+        from fastvideo.train.utils.moduleloader import (
+            make_inference_args, )
+
+        neg_embeds: torch.Tensor | None = None
+        neg_mask: torch.Tensor | None = None
+
+        if world_group.rank_in_group == 0:
+            sampling_param = SamplingParam.from_pretrained(tc.model_path)
+            negative_prompt = sampling_param.negative_prompt
+
+            inference_args = make_inference_args(tc, model_path=tc.model_path)
+
+            prompt_pipeline = WanPipeline.from_pretrained(
+                tc.model_path,
+                args=inference_args,
+                inference_mode=True,
+                loaded_modules={"transformer": self.transformer},
+                tp_size=tc.distributed.tp_size,
+                sp_size=tc.distributed.sp_size,
+                num_gpus=tc.distributed.num_gpus,
+                pin_cpu_memory=(tc.distributed.pin_cpu_memory),
+                dit_cpu_offload=True,
+            )
+
+            batch_negative = ForwardBatch(
+                data_type="video",
+                prompt=negative_prompt,
+                prompt_embeds=[],
+                prompt_attention_mask=[],
+            )
+            result_batch = prompt_pipeline.prompt_encoding_stage(  # type: ignore[attr-defined]
+                batch_negative,
+                inference_args,
+            )
+
+            neg_embeds = result_batch.prompt_embeds[0].to(device=device, dtype=dtype)
+            neg_mask = (result_batch.prompt_attention_mask[0].to(device=device, dtype=dtype))
+
+            del prompt_pipeline
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        meta = torch.zeros((2, ), device=device, dtype=torch.int64)
+        if world_group.rank_in_group == 0:
+            assert neg_embeds is not None
+            assert neg_mask is not None
+            meta[0] = neg_embeds.ndim
+            meta[1] = neg_mask.ndim
+        world_group.broadcast(meta, src=0)
+        embed_ndim, mask_ndim = (
+            int(meta[0].item()),
+            int(meta[1].item()),
+        )
+
+        max_ndim = 8
+        embed_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
+        mask_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
+        if world_group.rank_in_group == 0:
+            assert neg_embeds is not None
+            assert neg_mask is not None
+            embed_shape[:embed_ndim] = torch.tensor(
+                list(neg_embeds.shape),
+                device=device,
+                dtype=torch.int64,
+            )
+            mask_shape[:mask_ndim] = torch.tensor(
+                list(neg_mask.shape),
+                device=device,
+                dtype=torch.int64,
+            )
+        world_group.broadcast(embed_shape, src=0)
+        world_group.broadcast(mask_shape, src=0)
+
+        embed_sizes = tuple(int(x) for x in embed_shape[:embed_ndim].tolist())
+        mask_sizes = tuple(int(x) for x in mask_shape[:mask_ndim].tolist())
+
+        if world_group.rank_in_group != 0:
+            neg_embeds = torch.empty(embed_sizes, device=device, dtype=dtype)
+            neg_mask = torch.empty(mask_sizes, device=device, dtype=dtype)
+        assert neg_embeds is not None
+        assert neg_mask is not None
+
+        world_group.broadcast(neg_embeds, src=0)
+        world_group.broadcast(neg_mask, src=0)
+
+        self.negative_prompt_embeds = neg_embeds
+        self.negative_prompt_attention_mask = neg_mask
+
+    def _sample_timesteps(
+        self,
+        batch_size: int,
+        device: torch.device,
+        generator: torch.Generator,
+    ) -> torch.Tensor:
+        assert self.training_config is not None
+        tc = self.training_config
+
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=tc.model.weighting_scheme,
+            batch_size=batch_size,
+            generator=generator,
+            device=device,
+            logit_mean=tc.model.logit_mean,
+            logit_std=tc.model.logit_std,
+            mode_scale=tc.model.mode_scale,
+        )
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices.cpu()].to(device=device)
+
+    def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+        latents_shape = training_batch.raw_latent_shape
+        patch_size = (
+            tc.pipeline_config.dit_config.patch_size  # type: ignore[union-attr]
+        )
+        assert latents_shape is not None
+        assert training_batch.timesteps is not None
+
+        if envs.FASTVIDEO_ATTENTION_BACKEND in (
+            "VIDEO_SPARSE_ATTN", "SPARSE_FP4_ATTN", "SPARSE_FP4_OURS_P_ATTN",
+        ):
+            if (not is_vsa_available() or VideoSparseAttentionMetadataBuilder is None):
+                raise ImportError(
+                    f"FASTVIDEO_ATTENTION_BACKEND is "
+                    f"{envs.FASTVIDEO_ATTENTION_BACKEND}, but "
+                    f"fastvideo_kernel is not correctly "
+                    f"installed or detected.")
+            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder().build(  # type: ignore[misc]
+                raw_latent_shape=latents_shape[2:5],
+                current_timestep=(training_batch.timesteps),
+                patch_size=patch_size,
+                VSA_sparsity=tc.vsa_sparsity,
+                device=self.device,
+            )
+        elif (envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN"):
+            if (not is_vmoba_available() or VideoMobaAttentionMetadataBuilder is None):
+                raise ImportError("FASTVIDEO_ATTENTION_BACKEND is "
+                                  "VMOBA_ATTN, but fastvideo_kernel "
+                                  "(or flash_attn>=2.7.4) is not "
+                                  "correctly installed.")
+            moba_params = tc.model.moba_config.copy()
+            assert training_batch.raw_latent_shape is not None
+            moba_params.update({
+                "current_timestep": (training_batch.timesteps),
+                "raw_latent_shape": (training_batch.raw_latent_shape[2:5]),
+                "patch_size": patch_size,
+                "device": self.device,
+            })
+            training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder().build(**
+                                                                                     moba_params)  # type: ignore[misc]
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _prepare_dit_inputs(
+        self,
+        training_batch: TrainingBatch,
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+        latents = training_batch.latents
+        assert isinstance(latents, torch.Tensor)
+        batch_size = latents.shape[0]
+
+        noise = torch.randn(
+            latents.shape,
+            generator=generator,
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+        timesteps = self._sample_timesteps(
+            batch_size,
+            latents.device,
+            generator,
+        )
+        if int(tc.distributed.sp_size or 1) > 1:
+            self.sp_group.broadcast(timesteps, src=0)
+
+        sigmas = get_sigmas(
+            self.noise_scheduler,
+            latents.device,
+            timesteps,
+            n_dim=latents.ndim,
+            dtype=latents.dtype,
+        )
+        noisy_model_input = ((1.0 - sigmas) * latents + sigmas * noise)
+
+        training_batch.noisy_model_input = (noisy_model_input)
+        training_batch.timesteps = timesteps
+        training_batch.sigmas = sigmas
+        training_batch.noise = noise
+        training_batch.raw_latent_shape = latents.shape
+
+        training_batch.conditional_dict = {
+            "encoder_hidden_states": (training_batch.encoder_hidden_states),
+            "encoder_attention_mask": (training_batch.encoder_attention_mask),
+        }
+
+        if (self.negative_prompt_embeds is not None and self.negative_prompt_attention_mask is not None):
+            neg_embeds = self.negative_prompt_embeds
+            neg_mask = (self.negative_prompt_attention_mask)
+            if (neg_embeds.shape[0] == 1 and batch_size > 1):
+                neg_embeds = neg_embeds.expand(batch_size, *neg_embeds.shape[1:]).contiguous()
+            if (neg_mask.shape[0] == 1 and batch_size > 1):
+                neg_mask = neg_mask.expand(batch_size, *neg_mask.shape[1:]).contiguous()
+            training_batch.unconditional_dict = {
+                "encoder_hidden_states": neg_embeds,
+                "encoder_attention_mask": neg_mask,
+            }
+
+        training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4))
+        return training_batch
+
+    def _build_distill_input_kwargs(
+        self,
+        noise_input: torch.Tensor,
+        timestep: torch.Tensor,
+        text_dict: dict[str, torch.Tensor] | None,
+    ) -> dict[str, Any]:
+        if text_dict is None:
+            raise ValueError("text_dict cannot be None for "
+                             "Wan distillation")
+        return {
+            "hidden_states": noise_input.permute(0, 2, 1, 3, 4),
+            "encoder_hidden_states": text_dict["encoder_hidden_states"],
+            "encoder_attention_mask": text_dict["encoder_attention_mask"],
+            "timestep": timestep,
+            "return_dict": False,
+        }
+
+    def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module:
+        return self.transformer
+
+    def _get_uncond_text_dict(
+        self,
+        batch: TrainingBatch,
+        *,
+        cfg_uncond: dict[str, Any] | None,
+    ) -> dict[str, torch.Tensor]:
+        if cfg_uncond is None:
+            text_dict = getattr(batch, "unconditional_dict", None)
+            if text_dict is None:
+                raise RuntimeError("Missing unconditional_dict; "
+                                   "ensure_negative_conditioning() "
+                                   "may have failed")
+            return text_dict
+
+        on_missing_raw = cfg_uncond.get("on_missing", "error")
+        if not isinstance(on_missing_raw, str):
+            raise ValueError("method_config.cfg_uncond.on_missing "
+                             "must be a string, got "
+                             f"{type(on_missing_raw).__name__}")
+        on_missing = on_missing_raw.strip().lower()
+        if on_missing not in {"error", "ignore"}:
+            raise ValueError("method_config.cfg_uncond.on_missing "
+                             "must be one of {error, ignore}, got "
+                             f"{on_missing_raw!r}")
+
+        for channel, policy_raw in cfg_uncond.items():
+            if channel in {"on_missing", "text"}:
+                continue
+            if policy_raw is None:
+                continue
+            if not isinstance(policy_raw, str):
+                raise ValueError("method_config.cfg_uncond values "
+                                 "must be strings, got "
+                                 f"{channel}="
+                                 f"{type(policy_raw).__name__}")
+            policy = policy_raw.strip().lower()
+            if policy == "keep":
+                continue
+            if on_missing == "ignore":
+                continue
+            raise ValueError("WanModel does not support "
+                             "cfg_uncond channel "
+                             f"{channel!r} (policy={policy!r}). "
+                             "Set cfg_uncond.on_missing=ignore or "
+                             "remove the channel.")
+
+        text_policy_raw = cfg_uncond.get("text", None)
+        if text_policy_raw is None:
+            text_policy = "negative_prompt"
+        elif not isinstance(text_policy_raw, str):
+            raise ValueError("method_config.cfg_uncond.text must be "
+                             "a string, got "
+                             f"{type(text_policy_raw).__name__}")
+        else:
+            text_policy = (text_policy_raw.strip().lower())
+
+        if text_policy in {"negative_prompt"}:
+            text_dict = getattr(batch, "unconditional_dict", None)
+            if text_dict is None:
+                raise RuntimeError("Missing unconditional_dict; "
+                                   "ensure_negative_conditioning() "
+                                   "may have failed")
+            return text_dict
+        if text_policy == "keep":
+            if batch.conditional_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+            return batch.conditional_dict
+        if text_policy == "zero":
+            if batch.conditional_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+            cond = batch.conditional_dict
+            enc = cond["encoder_hidden_states"]
+            mask = cond["encoder_attention_mask"]
+            if not torch.is_tensor(enc) or not torch.is_tensor(mask):
+                raise TypeError("conditional_dict must contain "
+                                "tensor text inputs")
+            return {
+                "encoder_hidden_states": (torch.zeros_like(enc)),
+                "encoder_attention_mask": (torch.zeros_like(mask)),
+            }
+        if text_policy == "drop":
+            raise ValueError("cfg_uncond.text=drop is not supported "
+                             "for Wan. Use "
+                             "{negative_prompt, keep, zero}.")
+        raise ValueError("cfg_uncond.text must be one of "
+                         "{negative_prompt, keep, zero, drop}, got "
+                         f"{text_policy_raw!r}")

backend_snapshot/fastvideo/training/training_pipeline.py

@@ -0,0 +1,1044 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import asdict
+from contextlib import AbstractContextManager, nullcontext
+import math
+import os
+import shutil
+import tempfile
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Iterator
+from typing import Any
+from fastvideo.profiler import profile_region
+import imageio
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from einops import rearrange
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm.auto import tqdm
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+import fastvideo.envs as envs
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionMetadataBuilder)
+    from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
+except Exception:
+    pass
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import build_parquet_map_style_dataloader
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.dataset.validation_dataset import ValidationDataset
+from fastvideo.distributed import (cleanup_dist_env_and_memory, get_local_torch_device, get_sp_group, get_world_group)
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.attention.selector import global_force_attn_backend_context_manager
+from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch, LoRAPipeline, TrainingBatch)
+from fastvideo.platforms import AttentionBackendEnum, current_platform
+from fastvideo.training.activation_checkpoint import (apply_activation_checkpointing)
+from fastvideo.training.trackers import (DummyTracker, TrackerType, initialize_trackers, Trackers)
+from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases,
+                                               compute_density_for_timestep_sampling, count_trainable, get_scheduler,
+                                               get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint,
+                                               swap_fp4_linear, traverse_swap_module)
+from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict)
+
+try:
+    vsa_available = is_vsa_available()
+    vmoba_available = is_vmoba_available()
+except Exception:
+    vsa_available = False
+    vmoba_available = False
+
+logger = init_logger(__name__)
+
+
+class TrainingPipeline(LoRAPipeline, ABC):
+    """
+    A pipeline for training a model. All training pipelines should inherit from this class.
+    All reusable components and code should be implemented in this class.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+    validation_pipeline: ComposedPipelineBase
+    train_dataloader: StatefulDataLoader
+    train_loader_iter: Iterator[dict[str, Any]]
+    current_epoch: int = 0
+    train_transformer_2: bool = False
+    tracker: TrackerType
+
+    def __init__(self,
+                 model_path: str,
+                 fastvideo_args: TrainingArgs,
+                 required_config_modules: list[str] | None = None,
+                 loaded_modules: dict[str, torch.nn.Module] | None = None) -> None:
+        fastvideo_args.inference_mode = False
+        self.lora_training = fastvideo_args.lora_training
+        if self.lora_training and fastvideo_args.lora_rank is None:
+            raise ValueError("lora rank must be set when using lora training")
+
+        set_random_seed(fastvideo_args.seed)  # for lora param init
+        super().__init__(model_path, fastvideo_args, required_config_modules, loaded_modules)  # type: ignore
+        self.tracker = DummyTracker()
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        raise RuntimeError("create_pipeline_stages should not be called for training pipeline")
+
+    @staticmethod
+    def _should_force_generator_attn_qat_train(fastvideo_args: FastVideoArgs) -> bool:
+        if not isinstance(fastvideo_args, TrainingArgs):
+            return False
+        return (fastvideo_args.generator_4bit_attn or envs.FASTVIDEO_ATTENTION_BACKEND == "ATTN_QAT_TRAIN")
+
+    def load_modules(self,
+                     fastvideo_args: FastVideoArgs,
+                     loaded_modules: dict[str, torch.nn.Module] | None = None) -> dict[str, Any]:
+        force_generator_qat = self._should_force_generator_attn_qat_train(fastvideo_args)
+        load_context: AbstractContextManager[None] = nullcontext()
+        if force_generator_qat:
+            logger.info("Forcing generator attention backend to ATTN_QAT_TRAIN during module loading")
+            load_context = global_force_attn_backend_context_manager(AttentionBackendEnum.ATTN_QAT_TRAIN)
+
+        with load_context:
+            return super().load_modules(fastvideo_args, loaded_modules)
+
+    def set_schemas(self) -> None:
+        self.train_dataset_schema = pyarrow_schema_t2v
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing training pipeline...")
+        self.device = get_local_torch_device()
+        self.training_args = training_args
+        world_group = get_world_group()
+        self.world_size = world_group.world_size
+        self.global_rank = world_group.rank
+        self.sp_group = get_sp_group()
+        self.rank_in_sp_group = self.sp_group.rank_in_group
+        self.sp_world_size = self.sp_group.world_size
+        self.local_rank = world_group.local_rank
+        self.transformer = self.get_module("transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
+        self.seed = training_args.seed
+        self.set_schemas()
+
+        # Set random seeds for deterministic training
+        assert self.seed is not None, "seed must be set"
+        set_random_seed(self.seed + self.global_rank)
+        self.transformer.train()
+        if training_args.enable_gradient_checkpointing_type is not None:
+            self.transformer = apply_activation_checkpointing(
+                self.transformer, checkpointing_type=training_args.enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2, checkpointing_type=training_args.enable_gradient_checkpointing_type)
+
+        if training_args.generator_4bit_linear:
+            num_swaps = traverse_swap_module(self.transformer, swap_fn=swap_fp4_linear)
+            logger.info("Swapped %s linear layers to the FP4 forward path in self.transformer", num_swaps)
+        noise_scheduler = self.modules["scheduler"]
+        self.set_trainable()
+        params_to_optimize = self.transformer.parameters()
+        params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
+        # Parse betas from string format "beta1,beta2"
+        betas_str = training_args.betas
+        betas = tuple(float(x.strip()) for x in betas_str.split(","))
+
+        self.optimizer = torch.optim.AdamW(
+            params_to_optimize,
+            lr=training_args.learning_rate,
+            betas=betas,
+            weight_decay=training_args.weight_decay,
+            eps=1e-8,
+        )
+
+        self.init_steps = 0
+        logger.info("optimizer: %s", self.optimizer)
+
+        self.lr_scheduler = get_scheduler(
+            training_args.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=training_args.lr_warmup_steps,
+            num_training_steps=training_args.max_train_steps,
+            num_cycles=training_args.lr_num_cycles,
+            power=training_args.lr_power,
+            min_lr_ratio=training_args.min_lr_ratio,
+            last_epoch=self.init_steps - 1,
+        )
+        if self.transformer_2 is not None:
+            # Ensure transformer_2 has trainable parameters before creating optimizer
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(filter(lambda p: p.requires_grad, params_to_optimize_2))
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=(0.9, 0.999),
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
+
+        self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
+            training_args.data_path,
+            training_args.train_batch_size,
+            parquet_schema=self.train_dataset_schema,
+            num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
+            drop_last=True,
+            text_padding_length=training_args.pipeline_config.text_encoder_configs[0].arch_config.
+            text_len,  # type: ignore[attr-defined]
+            seed=self.seed)
+
+        self.noise_scheduler = noise_scheduler
+        if self.training_args.boundary_ratio is not None:
+            self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        else:
+            self.boundary_timestep = None
+
+        logger.info("train_dataloader length: %s", len(self.train_dataloader))
+        logger.info("train_sp_batch_size: %s", training_args.train_sp_batch_size)
+        logger.info("gradient_accumulation_steps: %s", training_args.gradient_accumulation_steps)
+        logger.info("sp_size: %s", training_args.sp_size)
+
+        self.num_update_steps_per_epoch = math.ceil(
+            len(self.train_dataloader) / training_args.gradient_accumulation_steps * training_args.sp_size /
+            training_args.train_sp_batch_size)
+        self.num_train_epochs = math.ceil(training_args.max_train_steps / self.num_update_steps_per_epoch)
+
+        # TODO(will): is there a cleaner way to track epochs?
+        self.current_epoch = 0
+
+        trackers = list(training_args.trackers)
+        if not trackers and training_args.tracker_project_name:
+            trackers.append(Trackers.WANDB.value)
+        if self.global_rank != 0:
+            trackers = []
+
+        tracker_log_dir = training_args.output_dir or os.getcwd()
+        if trackers:
+            tracker_log_dir = os.path.join(tracker_log_dir, "tracker")
+
+        tracker_config = asdict(training_args) if trackers else None
+        tracker_run_name = training_args.wandb_run_name or None
+        project = training_args.tracker_project_name or "fastvideo"
+        self.tracker = initialize_trackers(
+            trackers,
+            experiment_name=project,
+            config=tracker_config,
+            log_dir=tracker_log_dir,
+            run_name=tracker_run_name,
+        )
+
+    @abstractmethod
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError("Training pipelines must implement this method")
+
+    def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
+        self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.optimizer_2.zero_grad()
+        training_batch.total_loss = 0.0
+        return training_batch
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        with self.tracker.timed("timing/get_next_batch"):
+            batch = next(self.train_loader_iter, None)  # type: ignore
+            if batch is None:
+                self.current_epoch += 1
+                logger.info("Starting epoch %s", self.current_epoch)
+                # Reset iterator for next epoch
+                self.train_loader_iter = iter(self.train_dataloader)
+                # Get first batch of new epoch
+                batch = next(self.train_loader_iter)
+
+            latents = batch['vae_latent']
+            latents = latents[:, :, :self.training_args.num_latent_t]
+            encoder_hidden_states = batch['text_embedding']
+            encoder_attention_mask = batch['text_attention_mask']
+            infos = batch['info_list']
+
+            training_batch.latents = latents.to(
+                get_local_torch_device(),
+                dtype=torch.bfloat16,
+                non_blocking=True,
+            )
+            training_batch.encoder_hidden_states = (encoder_hidden_states.to(
+                get_local_torch_device(),
+                dtype=torch.bfloat16,
+                non_blocking=True,
+            ))
+            training_batch.encoder_attention_mask = (encoder_attention_mask.to(
+                get_local_torch_device(),
+                dtype=torch.bfloat16,
+                non_blocking=True,
+            ))
+            training_batch.infos = infos
+
+        return training_batch
+
+    def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch:
+        # TODO(will): support other models
+        with self.tracker.timed("timing/normalize_input"):
+            training_batch.latents = normalize_dit_input(
+                'wan',
+                training_batch.latents,
+                self.get_module("vae"),
+            )
+        return training_batch
+
+    def _prepare_dit_inputs(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_args is not None, "training_args must be set"
+        with self.tracker.timed("timing/prepare_dit_inputs"):
+            latents = training_batch.latents
+            batch_size = latents.shape[0]
+            noise = torch.randn(latents.shape,
+                                generator=self.noise_gen_cuda,
+                                device=latents.device,
+                                dtype=latents.dtype)
+            timesteps = self._sample_timesteps(batch_size, latents.device)
+
+            if self.training_args.sp_size > 1:
+                # Make sure that the timesteps are the same across all sp processes.
+                sp_group = get_sp_group()
+                sp_group.broadcast(timesteps, src=0)
+                sp_group.broadcast(noise, src=0)
+            sigmas = get_sigmas(
+                self.noise_scheduler,
+                latents.device,
+                timesteps,
+                n_dim=latents.ndim,
+                dtype=latents.dtype,
+            )
+            noisy_model_input = (1.0 - sigmas) * training_batch.latents + sigmas * noise
+
+            training_batch.noisy_model_input = noisy_model_input
+            training_batch.timesteps = timesteps
+            training_batch.sigmas = sigmas
+            training_batch.noise = noise
+            training_batch.raw_latent_shape = training_batch.latents.shape
+
+        return training_batch
+
+    def _sample_timesteps(self, batch_size: int, device: torch.device) -> torch.Tensor:
+        # Determine which model to train based on the boundary timestep
+        if (self.transformer_2 is not None and self.boundary_timestep is not None
+                and torch.rand(1, generator=self.noise_random_generator).item() <= self.training_args.boundary_ratio):
+            self.train_transformer_2 = True
+        else:
+            self.train_transformer_2 = False
+
+        # Broadcast the decision to all processes
+        decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device)
+        dist.broadcast(decision, src=0)
+        self.train_transformer_2 = decision.item() == 1.0
+
+        # Sample u from the appropriate range
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+
+        boundary_ratio = self.training_args.boundary_ratio
+        if self.train_transformer_2:
+            u = (1 - boundary_ratio) + u * boundary_ratio  # min: 1 - boundary_ratio, max: 1
+        # elif self.transformer_2 is not None:
+        #     u = u * (1 - boundary_ratio)  # min: 0, max: 1 - boundary_ratio
+        # else:  # patch for now to align with non-MoE timestep logic
+        #     pass
+
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices].to(device=device)
+
+    def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch:
+        latents_shape = training_batch.raw_latent_shape
+        patch_size = self.training_args.pipeline_config.dit_config.patch_size
+        current_vsa_sparsity = training_batch.current_vsa_sparsity
+        assert latents_shape is not None
+        assert isinstance(patch_size, tuple), f"Expected tuple patch_size, got {patch_size!r}"
+        assert training_batch.timesteps is not None
+        if envs.FASTVIDEO_ATTENTION_BACKEND in (
+            "VIDEO_SPARSE_ATTN",
+            "SPARSE_FP4_ATTN",
+            "SPARSE_FP4_OURS_P_ATTN",
+        ):
+            if not vsa_available:
+                raise ImportError("FASTVIDEO_ATTENTION_BACKEND is set to VIDEO_SPARSE_ATTN, "
+                                  "but fastvideo_kernel is not correctly installed or detected. "
+                                  "Please ensure fastvideo-kernel is installed.")
+            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder(  # type: ignore
+            ).build(  # type: ignore
+                raw_latent_shape=latents_shape[2:5],
+                current_timestep=training_batch.timesteps,
+                patch_size=patch_size,
+                VSA_sparsity=current_vsa_sparsity,
+                device=get_local_torch_device())
+        elif envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
+            if not vmoba_available:
+                raise ImportError("FASTVIDEO_ATTENTION_BACKEND is set to VMOBA_ATTN, "
+                                  "but fastvideo_kernel (or flash_attn>=2.7.4) is not correctly installed.")
+            moba_params = self.training_args.moba_config.copy()
+            moba_params.update({
+                "current_timestep": training_batch.timesteps,
+                "raw_latent_shape": latents_shape[2:5],
+                "patch_size": self.training_args.pipeline_config.dit_config.patch_size,
+                "device": get_local_torch_device(),
+            })
+            training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder().build(**moba_params)
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _build_input_kwargs(self, training_batch: TrainingBatch) -> TrainingBatch:
+        training_batch.input_kwargs = {
+            "hidden_states": training_batch.noisy_model_input,
+            "encoder_hidden_states": training_batch.encoder_hidden_states,
+            "timestep": training_batch.timesteps.to(get_local_torch_device(), dtype=torch.bfloat16),
+            "encoder_attention_mask": training_batch.encoder_attention_mask,
+            "return_dict": False,
+        }
+        return training_batch
+
+    def _transformer_forward_and_compute_loss(self, training_batch: TrainingBatch) -> TrainingBatch:
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND in (
+            "VIDEO_SPARSE_ATTN",
+            "SPARSE_FP4_ATTN",
+            "SPARSE_FP4_OURS_P_ATTN",
+        ) or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
+            assert training_batch.attn_metadata is not None
+        else:
+            assert training_batch.attn_metadata is None
+        input_kwargs = training_batch.input_kwargs
+
+        # if 'hunyuan' in self.training_args.model_type:
+        #     input_kwargs["guidance"] = torch.tensor(
+        #         [1000.0],
+        #         device=training_batch.noisy_model_input.device,
+        #         dtype=torch.bfloat16)
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+
+        with self.tracker.timed("timing/forward_backward"), set_forward_context(
+                current_timestep=training_batch.current_timestep, attn_metadata=training_batch.attn_metadata):
+            model_pred = current_model(**input_kwargs)
+            if self.training_args.precondition_outputs:
+                assert training_batch.sigmas is not None
+                model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
+            assert training_batch.latents is not None
+            assert training_batch.noise is not None
+            target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+
+            # make sure no implicit broadcasting happens
+            assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+
+            loss = (torch.mean(
+                (model_pred.float() - target.float())**2) / self.training_args.gradient_accumulation_steps)
+
+            loss.backward()
+
+            avg_loss = loss.detach().clone()
+
+        # Reduce across ranks without forcing a CPU sync
+        with self.tracker.timed("timing/reduce_loss"):
+            world_group = get_world_group()
+            avg_loss = world_group.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        # Accumulate on GPU; materialize to CPU only once after
+        # all gradient-accumulation iterations (see train_one_step).
+        training_batch.total_loss += avg_loss
+
+        return training_batch
+
+    def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
+        max_grad_norm = self.training_args.max_grad_norm
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            with self.tracker.timed("timing/clip_grad_norm"):
+                # Only clip gradients for the model that is currently training
+                if self.train_transformer_2 and self.transformer_2 is not None:
+                    model_parts = [self.transformer_2]
+                else:
+                    model_parts = [self.transformer]
+
+                grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+                    [p for m in model_parts for p in m.parameters()],
+                    max_grad_norm,
+                    foreach=None,
+                )
+                assert grad_norm is not float('nan') or grad_norm is not float('inf')
+                grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+        training_batch.grad_norm = grad_norm
+        return training_batch
+
+    @profile_region("profiler_region_training_train_one_step")
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        training_batch = self._prepare_training(training_batch)
+
+        for _ in range(self.training_args.gradient_accumulation_steps):
+            training_batch = self._get_next_batch(training_batch)
+
+            # Normalize DIT input
+            training_batch = self._normalize_dit_input(training_batch)
+            # Create noisy model input
+            training_batch = self._prepare_dit_inputs(training_batch)
+            assert training_batch.latents is not None
+            assert training_batch.noisy_model_input is not None
+            assert training_batch.noise is not None
+
+            # old sharding code, need to shard latents and noise but not input
+            # Shard latents across sp groups
+            training_batch.latents = training_batch.latents[:, :, :self.training_args.num_latent_t]
+            # shard noisy_model_input to match
+            training_batch.noisy_model_input = training_batch.noisy_model_input[:, :, :self.training_args.num_latent_t]
+            # shard noise to match latents
+            training_batch.noise = training_batch.noise[:, :, :self.training_args.num_latent_t]
+
+            training_batch = self._build_attention_metadata(training_batch)
+            training_batch = self._build_input_kwargs(training_batch)
+
+            training_batch = self._transformer_forward_and_compute_loss(training_batch)
+
+        training_batch = self._clip_grad_norm(training_batch)
+
+        # Only step the optimizer and scheduler for the model that is currently training
+        with self.tracker.timed("timing/optimizer_step"):
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                self.optimizer_2.step()
+                self.lr_scheduler_2.step()
+            else:
+                self.optimizer.step()
+                self.lr_scheduler.step()
+
+        return training_batch
+
+    def _compute_current_sparsity(self, step: int) -> float:
+        """Compute the VSA sparsity for a given step using the decay schedule."""
+        vsa_sparsity = self.training_args.VSA_sparsity
+        vsa_decay_rate = self.training_args.VSA_decay_rate
+        vsa_decay_interval = self.training_args.VSA_decay_interval_steps
+        vsa_init = getattr(self.training_args, 'VSA_init_sparsity', 0.0)
+        vsa_warmup = getattr(self.training_args, 'VSA_warmup_steps', 0)
+        if step <= vsa_warmup:
+            return vsa_init
+        ramp_step = step - vsa_warmup
+        max_times = int((vsa_sparsity - vsa_init) / vsa_decay_rate) if vsa_decay_rate > 0 else 0
+        times = min(ramp_step // vsa_decay_interval, max_times)
+        return vsa_init + times * vsa_decay_rate
+
+    def _resolve_checkpoint_path(self, path: str) -> str | None:
+        """Resolve 'latest' to the most recent checkpoint in output_dir."""
+        import glob
+        if path == "latest":
+            output_dir = self.training_args.output_dir
+            ckpt_dirs = sorted(
+                glob.glob(os.path.join(output_dir, "checkpoint-*")),
+                key=lambda d: int(d.split("-")[-1]) if d.split("-")[-1].isdigit() else 0,
+            )
+            if ckpt_dirs:
+                latest = ckpt_dirs[-1]
+                logger.info("Auto-resolved 'latest' to %s", latest)
+                return latest
+            logger.info("No checkpoints found in %s, starting from scratch", output_dir)
+            return None
+        return path
+
+    def _resume_from_checkpoint(self) -> None:
+        ckpt_path = self._resolve_checkpoint_path(self.training_args.resume_from_checkpoint)
+        if ckpt_path is None:
+            logger.info("No checkpoint to resume from, starting from step 0")
+            return
+
+        safetensors_path = os.path.join(ckpt_path, "transformer", "diffusion_pytorch_model.safetensors")
+        step = int(os.path.basename(os.path.normpath(ckpt_path)).split('-')[-1])
+
+        resumed_step = load_checkpoint(self.transformer, self.global_rank, ckpt_path,
+                                       self.optimizer, self.train_dataloader,
+                                       self.lr_scheduler, self.noise_random_generator)
+        if resumed_step > 0 or step == 0:
+            self.init_steps = resumed_step
+            logger.info("Successfully resumed full training state from step %s", resumed_step)
+            return
+
+        if os.path.exists(safetensors_path):
+            self.init_steps = step
+            logger.warning("Distributed checkpoint resume failed; falling back to safetensors weights at step %s",
+                           step)
+            return
+
+        logger.warning("No usable checkpoint state found at %s; starting from step 0", ckpt_path)
+        self.init_steps = 0
+
+    @profile_region("profiler_region_training_train")
+    def train(self) -> None:
+        assert self.seed is not None, "seed must be set"
+        assert self.training_args is not None, "training_args must be set"
+        set_random_seed(self.seed + self.global_rank)
+        logger.info('rank: %s: start training', self.global_rank, local_main_process_only=False)
+        if not self.post_init_called:
+            self.post_init()
+        num_trainable_params = count_trainable(self.transformer)
+        logger.info("Starting training with %s B trainable parameters", round(num_trainable_params / 1e9, 3))
+
+        if getattr(self, "transformer_2", None) is not None:
+            num_trainable_params = count_trainable(self.transformer_2)
+            logger.info("Transformer 2: Starting training with %s B trainable parameters",
+                        round(num_trainable_params / 1e9, 3))
+
+        # Set random seeds for deterministic training
+        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(self.seed + self.global_rank)
+        self.noise_gen_cuda = torch.Generator(device=current_platform.device_name).manual_seed(self.seed +
+                                                                                               self.global_rank)
+        self.validation_random_generator = torch.Generator(device="cpu").manual_seed(self.seed + self.global_rank)
+        logger.info("Initialized random seeds with seed: %s", self.seed + self.global_rank)
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        if self.training_args.resume_from_checkpoint:
+            self._resume_from_checkpoint()
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        self._log_training_info()
+
+        # Validation at init uses the sparsity corresponding to init_steps
+        saved_sparsity = self.training_args.VSA_sparsity
+        self.training_args.VSA_sparsity = self._compute_current_sparsity(self.init_steps)
+        self._log_validation(self.transformer, self.training_args, self.init_steps)
+        self.training_args.VSA_sparsity = saved_sparsity
+
+        # Train!
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+        for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+            if vsa_available:
+                vsa_sparsity = self.training_args.VSA_sparsity
+                vsa_decay_rate = self.training_args.VSA_decay_rate
+                vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps
+                vsa_init_sparsity = getattr(self.training_args, 'VSA_init_sparsity', 0.0)
+                vsa_warmup_steps = getattr(self.training_args, 'VSA_warmup_steps', 0)
+                if step <= vsa_warmup_steps:
+                    current_vsa_sparsity = vsa_init_sparsity
+                else:
+                    ramp_step = step - vsa_warmup_steps
+                    max_decay_times = int((vsa_sparsity - vsa_init_sparsity) / vsa_decay_rate)
+                    current_decay_times = min(ramp_step // vsa_decay_interval_steps, max_decay_times)
+                    current_vsa_sparsity = vsa_init_sparsity + current_decay_times * vsa_decay_rate
+            elif vmoba_available:
+                #TODO: add vmoba sparsity scheduling here
+                current_vsa_sparsity = 0.0
+            else:
+                current_vsa_sparsity = 0.0
+
+            training_batch = TrainingBatch()
+            training_batch.current_timestep = step
+            training_batch.current_vsa_sparsity = current_vsa_sparsity
+            training_batch = self.train_one_step(training_batch)
+
+            loss = float(training_batch.total_loss)
+            grad_norm = training_batch.grad_norm
+
+            step_time = time.perf_counter() - start_time
+            step_times.append(step_time)
+            avg_step_time = sum(step_times) / len(step_times)
+
+            progress_bar.set_postfix({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.global_rank == 0:
+                metrics = {
+                    "train_loss": loss,
+                    "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                    "grad_norm": grad_norm,
+                    "vsa_sparsity": current_vsa_sparsity,
+                }
+                try:
+                    assert training_batch.raw_latent_shape is not None
+                    metrics["batch_size"] = int(training_batch.raw_latent_shape[0])
+
+                    patch_size = self.training_args.pipeline_config.dit_config.patch_size
+                    assert isinstance(patch_size, tuple), f"Expected tuple patch_size, got {patch_size!r}"
+                    patch_t, patch_h, patch_w = patch_size
+                    seq_len = (training_batch.raw_latent_shape[2] // patch_t) * (
+                        training_batch.raw_latent_shape[3] // patch_h) * (training_batch.raw_latent_shape[4] // patch_w)
+                    if training_batch.encoder_hidden_states is not None:
+                        context_len = int(training_batch.encoder_hidden_states.shape[1])
+                    else:
+                        context_len = 0
+
+                    metrics["dit_seq_len"] = int(seq_len)
+                    metrics["context_len"] = context_len
+
+                    arch_config = self.training_args.pipeline_config.dit_config.arch_config
+
+                    metrics["hidden_dim"] = arch_config.hidden_size
+                    metrics["num_layers"] = arch_config.num_layers
+                    metrics["ffn_dim"] = arch_config.ffn_dim
+                except Exception:
+                    pass
+
+                self.tracker.log(metrics, step)
+            if step % self.training_args.training_state_checkpointing_steps == 0:
+                with self.profiler_controller.region("profiler_region_training_save_checkpoint"):
+                    save_checkpoint(self.transformer, self.global_rank, self.training_args.output_dir, step,
+                                    self.optimizer, self.train_dataloader, self.lr_scheduler,
+                                    self.noise_random_generator,
+                                    self.training_args.checkpoints_total_limit)
+                self.transformer.train()
+                self.sp_group.barrier()
+
+            if self.training_args.log_visualization and step % self.training_args.visualization_steps == 0:
+                self.visualize_intermediate_latents(training_batch, self.training_args, step)
+
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                with self.profiler_controller.region("profiler_region_training_validation"):
+                    saved_sparsity = self.training_args.VSA_sparsity
+                    self.training_args.VSA_sparsity = current_vsa_sparsity
+                    self._log_validation(self.transformer, self.training_args, step)
+                    self.training_args.VSA_sparsity = saved_sparsity
+                    gpu_memory_usage = current_platform.get_torch_device().memory_allocated() / 1024**2
+                    trainable_params = round(count_trainable(self.transformer) / 1e9, 3)
+                    logger.info("GPU memory usage after validation: %s MB, trainable params: %sB", gpu_memory_usage,
+                                trainable_params)
+
+        self.tracker.finish()
+        save_checkpoint(self.transformer, self.global_rank, self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer, self.train_dataloader, self.lr_scheduler,
+                        self.noise_random_generator, self.training_args.checkpoints_total_limit)
+
+        if envs.FASTVIDEO_TORCH_PROFILER_DIR:
+            logger.info("Stopping profiler...")
+            self.profiler_controller.stop()
+            logger.info("Profiler stopped.")
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+    def _log_training_info(self) -> None:
+        assert self.training_args is not None, "training_args must be set"
+        total_batch_size = (self.world_size * self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size * self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(self.train_dataset))
+        logger.info("  Dataloader size = %s", len(self.train_dataloader))
+        logger.info("  Num Epochs = %s", self.num_train_epochs)
+        logger.info("  Resume training from step %s", self.init_steps)  # type: ignore
+        logger.info("  Instantaneous batch size per device = %s", self.training_args.train_batch_size)
+        logger.info("  Total train batch size (w. data & sequence parallel, accumulation) = %s", total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s", self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s", self.training_args.max_train_steps)
+        logger.info("  Total training parameters per FSDP shard = %s B",
+                    round(count_trainable(self.transformer) / 1e9, 3))
+        # print dtype
+        logger.info("  Master weight dtype: %s", self.transformer.parameters().__next__().dtype)
+
+        gpu_memory_usage = current_platform.get_torch_device().memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB", gpu_memory_usage)
+        logger.info("VSA validation sparsity: %s", self.training_args.VSA_sparsity)
+
+    def _prepare_validation_batch(self, sampling_param: SamplingParam, training_args: TrainingArgs,
+                                  validation_batch: dict[str, Any], num_inference_steps: int) -> ForwardBatch:
+        sampling_param.prompt = validation_batch['prompt']
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        if training_args.validation_guidance_scale:
+            sampling_param.guidance_scale = float(training_args.validation_guidance_scale)
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (training_args.num_latent_t - 1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=self.validation_random_generator,
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+
+        return batch
+
+    @torch.no_grad()
+    def _log_validation(self, transformer, training_args, global_step) -> None:
+        """
+        Generate a validation video and log it to the configured tracker to check the quality during training.
+        """
+        training_args.inference_mode = True
+        training_args.dit_cpu_offload = False
+        if not training_args.log_validation:
+            return
+        if self.validation_pipeline is None:
+            raise ValueError("Validation pipeline is not set")
+
+        logger.info("Starting validation")
+
+        # Create sampling parameters if not provided
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+
+        # Prepare validation prompts
+        logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
+                    self.global_rank,
+                    training_args.validation_dataset_file,
+                    local_main_process_only=False)
+        validation_dataset = ValidationDataset(training_args.validation_dataset_file)
+        validation_dataloader = DataLoader(validation_dataset, batch_size=None, num_workers=0)
+
+        self.transformer.eval()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.eval()
+
+        validation_steps = training_args.validation_sampling_steps.split(",")
+        validation_steps = [int(step) for step in validation_steps]
+        validation_steps = [step for step in validation_steps if step > 0]
+        # Log validation results for this step
+        world_group = get_world_group()
+        num_sp_groups = world_group.world_size // self.sp_group.world_size
+        one_prompt_per_rank = os.environ.get(
+            "FASTVIDEO_VALIDATION_ONE_PROMPT_PER_RANK",
+            "",
+        ).lower() in {"1", "true", "yes", "on"}
+
+        # Process each validation prompt for each validation step
+        for num_inference_steps in validation_steps:
+            logger.info("rank: %s: num_inference_steps: %s",
+                        self.global_rank,
+                        num_inference_steps,
+                        local_main_process_only=False)
+            step_videos: list[np.ndarray] = []
+            step_captions: list[str] = []
+
+            step_audio: list[np.ndarray | None] = []
+            step_sample_rates: list[int | None] = []
+
+            for prompt_idx, validation_batch in enumerate(validation_dataloader):
+                if one_prompt_per_rank and prompt_idx > 0:
+                    continue
+
+                batch = self._prepare_validation_batch(sampling_param, training_args, validation_batch,
+                                                       num_inference_steps)
+                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                            self.global_rank,
+                            self.rank_in_sp_group,
+                            batch.prompt,
+                            local_main_process_only=False)
+
+                assert batch.prompt is not None and isinstance(batch.prompt, str)
+                step_captions.append(batch.prompt)
+
+                # Run validation inference
+                output_batch = self.validation_pipeline.forward(batch, training_args)
+                samples = output_batch.output.cpu()
+
+                # Capture audio if available
+                audio = output_batch.extra.get("audio")
+                sample_rate = output_batch.extra.get("audio_sample_rate")
+
+                if audio is not None and torch.is_tensor(audio):
+                    audio = audio.detach().cpu().float().numpy()
+
+                step_audio.append(audio)
+                step_sample_rates.append(sample_rate)
+
+                if self.rank_in_sp_group != 0:
+                    continue
+
+                # Process outputs
+                video = rearrange(samples, "b c t h w -> t b c h w")
+                frames = []
+                for x in video:
+                    x = torchvision.utils.make_grid(x, nrow=6)
+                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                    frames.append((x * 255).numpy().astype(np.uint8))
+                step_videos.append(frames)
+
+            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
+            # results to global rank 0
+            if self.rank_in_sp_group == 0 and self.global_rank == 0:
+                # Global rank 0 collects results from all sp_group leaders
+                all_videos = step_videos  # Start with own results
+                all_captions = step_captions
+                all_audios = step_audio
+                all_sample_rates = step_sample_rates
+
+                # Receive from other sp_group leaders
+                for sp_group_idx in range(1, num_sp_groups):
+                    src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                    recv_videos = world_group.recv_object(src=src_rank)
+                    recv_captions = world_group.recv_object(src=src_rank)
+                    recv_audios = world_group.recv_object(src=src_rank)
+                    recv_sample_rates = world_group.recv_object(src=src_rank)
+
+                    all_videos.extend(recv_videos)
+                    all_captions.extend(recv_captions)
+                    all_audios.extend(recv_audios)
+                    all_sample_rates.extend(recv_sample_rates)
+
+                video_filenames = []
+                for i, (video, caption, audio, sample_rate) in enumerate(
+                        zip(all_videos, all_captions, all_audios, all_sample_rates, strict=True)):
+                    os.makedirs(training_args.output_dir, exist_ok=True)
+                    filename = os.path.join(
+                        training_args.output_dir,
+                        f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4")
+                    imageio.mimsave(filename, video, fps=sampling_param.fps)
+                    # Mux audio if available
+                    if (audio is not None and sample_rate is not None and not self._mux_audio(
+                            filename,
+                            audio,
+                            sample_rate,
+                    )):
+                        logger.warning("Audio mux failed for validation video %s; saved video without audio.", filename)
+                    video_filenames.append(filename)
+
+                artifacts = []
+                for filename, caption in zip(video_filenames, all_captions, strict=True):
+                    video_artifact = self.tracker.video(filename, caption=caption)
+                    if video_artifact is not None:
+                        artifacts.append(video_artifact)
+                if artifacts:
+                    logs = {f"validation_videos_{num_inference_steps}_steps": artifacts}
+                    self.tracker.log_artifacts(logs, global_step)
+            elif self.rank_in_sp_group == 0:
+                # Other sp_group leaders send their results to global rank 0
+                world_group.send_object(step_videos, dst=0)
+                world_group.send_object(step_captions, dst=0)
+                world_group.send_object(step_audio, dst=0)
+                world_group.send_object(step_sample_rates, dst=0)
+
+            world_group.barrier()
+
+        # Re-enable gradients for training
+        training_args.inference_mode = False
+        self.transformer.train()
+        if getattr(self, "transformer_2", None) is not None:
+            self.transformer_2.train()
+
+    @staticmethod
+    def _mux_audio(
+        video_path: str,
+        audio: torch.Tensor | np.ndarray,
+        sample_rate: int,
+    ) -> bool:
+        """Mux audio into video using PyAV."""
+        try:
+            import av
+        except ImportError:
+            logger.warning("PyAV not installed; cannot mux audio. "
+                           "Install with: pip install av")
+            return False
+
+        if torch.is_tensor(audio):
+            audio_np = audio.detach().cpu().float().numpy()
+        else:
+            audio_np = np.asarray(audio, dtype=np.float32)
+
+        if audio_np.ndim == 1:
+            audio_np = audio_np[:, None]
+        elif audio_np.ndim == 2:
+            if audio_np.shape[0] <= 8 and audio_np.shape[1] > audio_np.shape[0]:
+                audio_np = audio_np.T
+        else:
+            logger.warning("Unexpected audio shape %s; skipping mux.", audio_np.shape)
+            return False
+
+        audio_np = np.clip(audio_np, -1.0, 1.0)
+        audio_int16 = (audio_np * 32767.0).astype(np.int16)
+        num_channels = audio_int16.shape[1]
+        layout = "stereo" if num_channels == 2 else "mono"
+
+        try:
+            import wave
+            with tempfile.TemporaryDirectory() as tmpdir:
+                out_path = os.path.join(tmpdir, "muxed.mp4")
+                wav_path = os.path.join(tmpdir, "audio.wav")
+
+                # Write audio to WAV file
+                with wave.open(wav_path, "wb") as wav_file:
+                    wav_file.setnchannels(num_channels)
+                    wav_file.setsampwidth(2)
+                    wav_file.setframerate(sample_rate)
+                    wav_file.writeframes(audio_int16.tobytes())
+
+                # Open input video and audio
+                input_video = av.open(video_path)
+                input_audio = av.open(wav_path)
+
+                # Create output with both streams
+                output = av.open(out_path, mode="w")
+
+                # Add video stream (copy codec from input)
+                in_video_stream = input_video.streams.video[0]
+                out_video_stream = output.add_stream(
+                    codec_name=in_video_stream.codec_context.name,
+                    rate=in_video_stream.average_rate,
+                )
+                out_video_stream.width = in_video_stream.width
+                out_video_stream.height = in_video_stream.height
+                out_video_stream.pix_fmt = in_video_stream.pix_fmt
+
+                # Add audio stream (AAC)
+                out_audio_stream = output.add_stream("aac", rate=sample_rate)
+                out_audio_stream.layout = layout
+
+                # Remux video (decode and re-encode to be safe)
+                for frame in input_video.decode(video=0):
+                    for packet in out_video_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_video_stream.encode():
+                    output.mux(packet)
+
+                # Encode audio
+                for frame in input_audio.decode(audio=0):
+                    frame.pts = None  # Let encoder assign PTS
+                    for packet in out_audio_stream.encode(frame):
+                        output.mux(packet)
+                for packet in out_audio_stream.encode():
+                    output.mux(packet)
+
+                input_video.close()
+                input_audio.close()
+                output.close()
+                shutil.move(out_path, video_path)
+            return True
+        except Exception as e:
+            logger.warning("Audio mux failed: %s", e)
+            return False
+
+    def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int):
+        """Add visualization data to tracker logging and save frames to disk."""
+        raise NotImplementedError("Visualize intermediate latents is not implemented for training pipeline")

backend_snapshot/fastvideo/training/wan_training_pipeline.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+from copy import deepcopy
+
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (FlowUniPCMultistepScheduler)
+from fastvideo.pipelines.basic.wan.wan_pipeline import WanPipeline
+from fastvideo.training.training_pipeline import TrainingPipeline
+from fastvideo.utils import is_vsa_available
+
+try:
+    vsa_available = is_vsa_available()
+except Exception:
+    vsa_available = False
+
+logger = init_logger(__name__)
+
+
+class WanTrainingPipeline(TrainingPipeline):
+    """
+    A training pipeline for Wan.
+    """
+    _required_config_modules = ["scheduler", "transformer", "vae"]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_training_stages(self, training_args: TrainingArgs):
+        """
+        May be used in future refactors.
+        """
+        pass
+
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing validation pipeline...")
+        args_copy = deepcopy(training_args)
+
+        args_copy.inference_mode = True
+        validation_pipeline = WanPipeline.from_pretrained(
+            training_args.model_path,
+            args=args_copy,  # type: ignore
+            inference_mode=True,
+            loaded_modules={
+                "transformer": self.get_module("transformer"),
+            },
+            tp_size=training_args.tp_size,
+            sp_size=training_args.sp_size,
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
+
+        self.validation_pipeline = validation_pipeline
+
+
+def main(args) -> None:
+    logger.info("Starting training pipeline...")
+
+    pipeline = WanTrainingPipeline.from_pretrained(args.pretrained_model_name_or_path, args=args)
+    args = pipeline.training_args
+    pipeline.train()
+    logger.info("Training pipeline done")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    from fastvideo.fastvideo_args import TrainingArgs
+    from fastvideo.utils import FlexibleArgumentParser
+    parser = FlexibleArgumentParser()
+    parser = TrainingArgs.add_cli_args(parser)
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    args.dit_cpu_offload = False
+    main(args)

backend_snapshot/manifest.sha256

@@ -0,0 +1,17 @@
+45ff4b677a84fad92bd2ff596bf432cb1b9386c5923b6c0824f896074e7cfbc6  ./README.md
+9d1d8dc58aab529270fe31eb1735d6a1382c0c6d36fccca122a8dbffa1b714fd  ./fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_ours_p.py
+211c7f0445fbe9488250f01fa83457c6620e83bd6f3877db791fd155de93c08b  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton_ours_p.py
+3f3a407a88612ea17ad65e1b6b9cf6b7b02df56956d8301c4b13bffa92095016  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py
+56f17c602dede53c7c3677058f81274681530f1b83c086d9d1d44c6b51feefbb  ./fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py
+2b821b0e2e7bdb3581be6312ebbece42380a6ee28a7a982f0cf2dc71fab849c8  ./fastvideo/attention/backends/sparse_fp4_ours_p_attn.py
+a97adcc52d7558c49f418c09395fd1665e988ad290d2276b95f21dfca0f8eb7d  ./fastvideo/attention/backends/video_sparse_attn.py
+79ef6f38ec0f5bfe16b2b98327ad2ccd15f3c863dd87fd03affc5dbdaa0a8224  ./fastvideo/configs/models/dits/base.py
+ddcab6f4fd33c9813840571b6bf83bbbcea164b564166951ed4301297db6cef0  ./fastvideo/forward_context.py
+6cfd128e782b7787a27ddd28a5e2d50cb4b0e2e9425d51d9780f14c91e8206f0  ./fastvideo/pipelines/stages/denoising.py
+489388dbdd9e5e3ad24db3012bd9b108794509a9729891d7dd315a102abba828  ./fastvideo/platforms/cuda.py
+c046b1914041b59254bcdfe577aed20d6f007a72632ea1fe1ae92fa678eca760  ./fastvideo/platforms/interface.py
+2456d39ca28019e12bb7ab007774e86348f0582a017bf0e6c91e2a01d654a1a0  ./fastvideo/train/models/wan/wan.py
+bc46e84b732567de6c0325223405daecd1226c623e303be33c7be9b5b7fdec08  ./fastvideo/training/training_pipeline.py
+1d3898fa37e21029df6c37e05dc34ed7805a211c2f87de6642db890e5a8c6f2e  ./fastvideo/training/wan_training_pipeline.py
+5c982b64653fae83ebfdeb43fda8f29b3e2cb581fb4daee38cd3cf56aa9d73f5  ./scripts/training/run_sparse_fp4_train_v4_1n_sparse09_hpo_on_ours_p_init2050_interactive.sh
+5c1d5ce9ecc8b90e59ddfc2ddb3e2dae500bcd3acb90429c901444b1630f05fb  ./scripts/training/run_sparse_fp4_train_v4_common.sh

backend_snapshot/scripts/training/run_sparse_fp4_train_v4_1n_sparse09_hpo_on_ours_p_init2050_interactive.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+#SBATCH --job-name=sfp4-s09-oursp-i2050
+#SBATCH --account=nvr_elm_llm
+#SBATCH --partition=interactive
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:8
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --time=02:00:00
+#SBATCH --output=slurm_logs/sfp4_sparse09_ours_p_init2050_1n_interactive_%j.out
+#SBATCH --error=slurm_logs/sfp4_sparse09_ours_p_init2050_1n_interactive_%j.err
+
+export RUN_NAME="sfp4_v4_sparse09_hpo_on_ours_p_init2050_1n_interactive"
+export WANDB_RUN_ID="sfp4v4-sparse09-hpo-on-ours-p-init2050-1n-interactive"
+export FASTVIDEO_ATTENTION_BACKEND=SPARSE_FP4_OURS_P_ATTN
+export FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O=1
+export CHECKPOINT_LIMIT=5
+export SAVE_STEPS=50
+export EVAL_STEPS=50
+export VALIDATION_SAMPLING_STEPS=50
+export USE_SRUN=0
+
+export VSA_SPARSITY=0.9
+export VSA_INIT_SPARSITY=0.9
+export VSA_WARMUP_STEPS=0
+export VSA_DECAY_RATE=0.03
+export VSA_DECAY_INTERVAL_STEPS=50
+
+export INIT_WEIGHTS_FROM_SAFETENSORS="checkpoints/init/sfp4_v4_sparse06_hpo_on_ours_p_1n_interactive_v2_ckpt2050/transformer/diffusion_pytorch_model.safetensors"
+
+exec bash scripts/training/run_sparse_fp4_train_v4_common.sh

backend_snapshot/scripts/training/run_sparse_fp4_train_v4_common.sh

@@ -0,0 +1,199 @@
+#!/bin/bash
+
+set -euo pipefail
+set -x
+
+cd /lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/users/yitongl/code/FastVideo
+source .venv/bin/activate
+
+: "${RUN_NAME:?RUN_NAME must be set by the Slurm wrapper}"
+: "${WANDB_RUN_ID:?WANDB_RUN_ID must be set by the Slurm wrapper}"
+
+export PYTHONPATH=fastvideo-kernel/python:fastvideo-kernel:${PYTHONPATH:-}
+export FASTVIDEO_ATTENTION_BACKEND="${FASTVIDEO_ATTENTION_BACKEND:-SPARSE_FP4_ATTN}"
+export FASTVIDEO_VALIDATION_ONE_PROMPT_PER_RANK="${FASTVIDEO_VALIDATION_ONE_PROMPT_PER_RANK:-1}"
+export WANDB_MODE=online
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_RESUME=allow
+export WANDB_NAME="${RUN_NAME}"
+export TOKENIZERS_PARALLELISM=false
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export TRITON_CACHE_DIR="/tmp/triton_cache_${SLURM_JOB_ID:-manual}"
+
+if [[ -n "${SLURM_JOB_NODELIST:-}" ]]; then
+    MASTER_ADDR=$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1)
+else
+    MASTER_ADDR=127.0.0.1
+fi
+MASTER_PORT=$((20000 + (${SLURM_JOB_ID:-0} % 20000)))
+export MASTER_ADDR MASTER_PORT
+
+NUM_GPUS_PER_NODE=8
+NNODES=${SLURM_NNODES:-1}
+TOTAL_GPUS=$((NNODES * NUM_GPUS_PER_NODE))
+OUTPUT_DIR="checkpoints/${RUN_NAME}"
+CHECKPOINT_LIMIT="${CHECKPOINT_LIMIT:-5}"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/Wan-Syn_77x448x832_600k"
+VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json"
+SAVE_STEPS="${SAVE_STEPS:-50}"
+EVAL_STEPS="${EVAL_STEPS:-50}"
+VALIDATION_SAMPLING_STEPS="${VALIDATION_SAMPLING_STEPS:-50}"
+MAX_TRAIN_STEPS="${MAX_TRAIN_STEPS:-100000}"
+VSA_SPARSITY="${VSA_SPARSITY:-0.9}"
+VSA_INIT_SPARSITY="${VSA_INIT_SPARSITY:-0.6}"
+VSA_WARMUP_STEPS="${VSA_WARMUP_STEPS:-0}"
+VSA_DECAY_RATE="${VSA_DECAY_RATE:-0.03}"
+VSA_DECAY_INTERVAL_STEPS="${VSA_DECAY_INTERVAL_STEPS:-50}"
+INIT_WEIGHTS_FROM_SAFETENSORS="${INIT_WEIGHTS_FROM_SAFETENSORS:-}"
+
+mkdir -p slurm_logs "${OUTPUT_DIR}"
+
+find_latest_checkpoint() {
+    if [[ ! -d "${OUTPUT_DIR}" ]]; then
+        return 1
+    fi
+
+    mapfile -t checkpoint_steps < <(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name 'checkpoint-*' -printf '%f\n' \
+        | sed 's/checkpoint-//' \
+        | sort -nr)
+
+    local step
+    for step in "${checkpoint_steps[@]}"; do
+        if [[ -f "${OUTPUT_DIR}/checkpoint-${step}/transformer/diffusion_pytorch_model.safetensors" ]]; then
+            echo "${OUTPUT_DIR}/checkpoint-${step}"
+            return 0
+        fi
+    done
+
+    return 1
+}
+
+prune_checkpoints() {
+    if [[ ! -d "${OUTPUT_DIR}" ]]; then
+        return 0
+    fi
+
+    mapfile -t checkpoint_steps < <(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name 'checkpoint-*' -printf '%f\n' \
+        | sed 's/checkpoint-//' \
+        | sort -n)
+
+    local count=${#checkpoint_steps[@]}
+    if (( count <= CHECKPOINT_LIMIT )); then
+        return 0
+    fi
+
+    local remove_count=$((count - CHECKPOINT_LIMIT))
+    local step
+    for step in "${checkpoint_steps[@]:0:remove_count}"; do
+        rm -rf "${OUTPUT_DIR}/checkpoint-${step}"
+    done
+}
+
+RESUME_ARGS=()
+refresh_resume_args() {
+    RESUME_ARGS=()
+
+    local latest_ckpt
+    if latest_ckpt=$(find_latest_checkpoint); then
+        RESUME_ARGS=(
+            --resume_from_checkpoint latest
+            --init_weights_from_safetensors "${latest_ckpt}/transformer/diffusion_pytorch_model.safetensors"
+        )
+        echo "=== Resuming from ${latest_ckpt} ==="
+    fi
+}
+
+COMMON_ARGS=(
+    --tracker_project_name "wan_t2v_sparse_fp4"
+    --wandb_run_name "${RUN_NAME}"
+    --output_dir "${OUTPUT_DIR}"
+    --train_batch_size 1
+    --train_sp_batch_size 1
+    --gradient_accumulation_steps 1
+    --num_latent_t 20
+    --num_height 448
+    --num_width 832
+    --num_frames 77
+    --enable_gradient_checkpointing_type "full"
+    --num_gpus "${TOTAL_GPUS}"
+    --sp_size 1
+    --tp_size 1
+    --hsdp_replicate_dim "${TOTAL_GPUS}"
+    --hsdp_shard_dim 1
+    --model_path "${MODEL_PATH}"
+    --pretrained_model_name_or_path "${MODEL_PATH}"
+    --data_path "${DATA_DIR}"
+    --dataloader_num_workers 4
+    --log_validation
+    --validation_dataset_file "${VALIDATION_DATASET_FILE}"
+    --validation_steps "${EVAL_STEPS}"
+    --validation_sampling_steps "${VALIDATION_SAMPLING_STEPS}"
+    --validation_guidance_scale "5.0"
+    --learning_rate 1e-6
+    --mixed_precision "bf16"
+    --weight_only_checkpointing_steps "${SAVE_STEPS}"
+    --training_state_checkpointing_steps "${SAVE_STEPS}"
+    --weight_decay 0.01
+    --max_grad_norm 1.0
+    --inference_mode False
+    --checkpoints_total_limit "${CHECKPOINT_LIMIT}"
+    --training_cfg_rate 0.1
+    --dit_precision "fp32"
+    --ema_start_step 0
+    --flow_shift 1
+    --seed 1000
+    --VSA-sparsity "${VSA_SPARSITY}"
+    --VSA-init-sparsity "${VSA_INIT_SPARSITY}"
+    --VSA-warmup-steps "${VSA_WARMUP_STEPS}"
+    --VSA-decay-rate "${VSA_DECAY_RATE}"
+    --VSA-decay-interval-steps "${VSA_DECAY_INTERVAL_STEPS}"
+)
+
+if [[ -n "${INIT_WEIGHTS_FROM_SAFETENSORS}" ]]; then
+    COMMON_ARGS+=(
+        --init_weights_from_safetensors "${INIT_WEIGHTS_FROM_SAFETENSORS}"
+    )
+fi
+
+run_training() {
+    local max_steps=$1
+
+    local torchrun_cmd=(
+        torchrun
+        --nnodes="${NNODES}" \
+        --nproc_per_node="${NUM_GPUS_PER_NODE}" \
+        --rdzv_backend=c10d \
+        --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
+        fastvideo/training/wan_training_pipeline.py \
+        --max_train_steps "${max_steps}" \
+        "${COMMON_ARGS[@]}" \
+        "${RESUME_ARGS[@]}"
+    )
+
+    if [[ "${USE_SRUN:-1}" == "1" ]]; then
+        srun "${torchrun_cmd[@]}"
+    else
+        "${torchrun_cmd[@]}"
+    fi
+}
+
+echo "=== ${RUN_NAME} array=${SLURM_ARRAY_TASK_ID:-0} nodes=${NNODES} gpus=${TOTAL_GPUS} ==="
+echo "=== save_steps=${SAVE_STEPS} eval_steps=${EVAL_STEPS} checkpoint_limit=${CHECKPOINT_LIMIT} ==="
+echo "=== master=${MASTER_ADDR}:${MASTER_PORT} validation_one_prompt_per_rank=1 ==="
+
+if [[ ! -f "${OUTPUT_DIR}/checkpoint-0/transformer/diffusion_pytorch_model.safetensors" ]]; then
+    if ! find_latest_checkpoint >/dev/null; then
+        echo "=== Creating step-0 validation and checkpoint ==="
+        run_training 0
+        prune_checkpoints
+    fi
+fi
+
+refresh_resume_args
+run_training "${MAX_TRAIN_STEPS}"
+prune_checkpoints
+
+echo "=== Done arr=${SLURM_ARRAY_TASK_ID:-0} ==="