fix: can not find pipeline & text preprocessor

Signed-off-by: Jinzhe Pan <eigensystem1318@gmail.com>

.github/workflows/fastvideo-kernel-publish.yml

@@ -67,9 +67,6 @@ jobs:
             - torch-version: '2.9.1'
               cuda-version: '12.8.0'
               torch-cuda-short: 'cu128'
-            # - torch-version: '2.10.0'
-            #   cuda-version: '12.8.0'
-            #   torch-cuda-short: 'cu128'
 
     steps:
       - name: Free up disk space
@@ -223,6 +220,3 @@ jobs:
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           packages-dir: fastvideo-kernel/dist/
-          # PyPI does not allow replacing an existing file with the same name.
-          # This makes re-runs idempotent by skipping files already uploaded.
-          skip-existing: true

README.md

@@ -3,7 +3,7 @@
 </div>
 
 <p align="center">
-     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://github.com/hao-ai-lab/FastVideo/discussions/1097" target="_blank"> <b> WeChat </b> </a> |
+     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 **FastVideo is a unified post-training and inference framework for accelerated video generation.**

examples/inference/basic/basic_ltx2.py

@@ -20,10 +20,10 @@ def main() -> None:
     # Uses FastVideo default sampling settings for LTX2 base.
     generator = VideoGenerator.from_pretrained(
         "Davids048/LTX2-Base-Diffusers",
-        num_gpus=8,
+        num_gpus=1,
     )
 
-    output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.4.mp4"
+    output_path = "outputs_video/ltx2_basic/output_ltx2_base_t2v_1088_1920_1.1.mp4"
     generator.generate_video(
         prompt=PROMPT,
         output_path=output_path,
@@ -31,12 +31,6 @@ def main() -> None:
         num_frames=121,
         height=1088,
         width=1920,
-        # LTX2 uses these parameters for multi-modal CFG instead of guidance_scale
-        # ltx2_cfg_scale_video=3.0,
-        # ltx2_cfg_scale_audio=7.0,
-        # ltx2_modality_scale_video=3.0,
-        # ltx2_modality_scale_audio=3.0,
-        # ltx2_rescale_scale=0.7,
     )
     generator.shutdown()
 

fastvideo-kernel/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "fastvideo-kernel"
-version = "0.2.6"
+version = "0.2.5"
 description = "Unified CUDA kernels for FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -287,8 +287,9 @@ def block_sparse_attn(
     block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
     if (not _force_triton()) and _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
         return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes)
-    # Triton path: supports q_seq_len != kv_seq_len as long as both are padded
-    # to a multiple of the block size (64 tokens).
+    # Triton path: generally assumes q/k/v share the same padded length
+    if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:
+        raise RuntimeError("Triton fallback requires q/k/v to have the same padded length.")
     return block_sparse_attn_triton(q, k, v, block_map, variable_block_sizes)
 
 

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -141,6 +141,12 @@ def video_sparse_attn(
         # Use autograd-enabled wrapper so backward works (and still uses SM90 kernel when available)
         out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
     else:
+        if q_seq_len != kv_seq_len:
+            raise RuntimeError(
+                "q/k have different lengths, but the compiled CUDA kernel (block_sparse_fwd) "
+                "is not available. The Triton fallback currently requires q and k/v to have "
+                "the same padded length."
+            )
         # Triton-only forward (kept for environments without the wrapper deps)
         out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes)
 

fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py

@@ -29,7 +29,7 @@ configs = [
 
 
 # ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
-@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
+@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"])
 @triton.jit
 def _attn_fwd_sparse(
         Q,
@@ -60,8 +60,7 @@ def _attn_fwd_sparse(
         stride_on,
         Z,
         H,
-        N_CTX_Q,  #
-        N_CTX_KV,  #
+        N_CTX,  #
         HEAD_DIM: tl.constexpr,  #
         BLOCK_M: tl.constexpr,
         BLOCK_N: tl.constexpr,
@@ -76,29 +75,24 @@ def _attn_fwd_sparse(
     off_hz = tl.program_id(1)  # fused (batch, head)
     b = off_hz // H
     h = off_hz % H
-    q_tiles = N_CTX_Q // BLOCK_M
+    q_tiles = N_CTX // 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
 
     # ----- base pointers -----
-    # Note: when q and kv have different sequence lengths, their per-(batch,head)
-    # strides differ, so we must compute separate base offsets.
-    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)
+    qvk_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
 
-    Q_ptr = tl.make_block_ptr(base=Q + q_off,
-                              shape=(N_CTX_Q, HEAD_DIM),
+    Q_ptr = tl.make_block_ptr(base=Q + qvk_off,
+                              shape=(N_CTX, 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),
+    K_base = tl.make_block_ptr(base=K + qvk_off,
+                               shape=(HEAD_DIM, N_CTX),
                                strides=(stride_kk, stride_kn),
                                offsets=(0, 0),
                                block_shape=(HEAD_DIM, BLOCK_N),
@@ -106,15 +100,15 @@ def _attn_fwd_sparse(
 
     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),
+    V_base = tl.make_block_ptr(base=V + qvk_off,
+                               shape=(N_CTX, 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),
+    O_ptr = tl.make_block_ptr(base=Out + qvk_off,
+                              shape=(N_CTX, HEAD_DIM),
                               strides=(stride_om, stride_on),
                               offsets=(q_blk * BLOCK_M, 0),
                               block_shape=(BLOCK_M, HEAD_DIM),
@@ -156,7 +150,7 @@ def _attn_fwd_sparse(
     # ----- epilogue -----
     m_i += tl.math.log2(l_i)
     acc = acc / l_i[:, None]
-    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
+    tl.store(M + off_hz * N_CTX + offs_m, m_i)
     tl.store(O_ptr, acc.to(Out.type.element_ty))
 
 
@@ -207,7 +201,7 @@ def _attn_bwd_dkdv(
         stride_tok,
         stride_d,  #
         H,
-        N_CTX_KV,
+        N_CTX,
         BLOCK_M1: tl.constexpr,  #
         BLOCK_N1: tl.constexpr,  #
         HEAD_DIM: tl.constexpr,  #
@@ -227,8 +221,8 @@ def _attn_bwd_dkdv(
     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_tiles = N_CTX // BLOCK_N1
+    meta_base = ((b * H + h) * q_tiles + kv_blk)
 
     q_blocks = tl.load(k2q_num + meta_base)  # int32
     q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
@@ -308,21 +302,16 @@ def _attn_bwd_dq(
 
     kv_blocks = tl.load(q2k_num + meta_base)  # int32
     kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+    block_size = tl.load(variable_block_sizes + q_blk)
 
     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
+        block_sparse_offset = (tl.load(kv_ptr + blk_idx // 2).to(tl.int32) * 2 +
+                               blk_idx % 2) * 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)
         p = tl.math.exp2(qk - m)
-        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
-        mask = offs_in_block < block_size
+        mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32)
         p = tl.where(mask[None, :], p, 0.0)
         # Compute dP and dS.
         dp = tl.dot(do, vT).to(tl.float32)
@@ -478,235 +467,19 @@ def _attn_bwd(
     tl.store(dq_ptrs, dq)
 
 
-@triton.jit
-def _attn_bwd_dkdv_kernel(
-        Q,
-        K,
-        V,
-        sm_scale,  #
-        DO,  #
-        DK,
-        DV,  #
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_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):
-    """
-    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
-
-    # 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,
-        sm_scale,
-        DO,
-        M,
-        D,
-        k2q_index,
-        k2q_num,
-        max_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,
-    )
-
-    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,
-        DO,  #
-        DQ,
-        M,
-        D,
-        q2k_index,
-        q2k_num,
-        max_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,
-        BLOCK_M2: tl.constexpr,  #
-        BLOCK_N2: tl.constexpr,  #
-        HEAD_DIM: tl.constexpr):
-    """
-    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
-
-    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 = tl.load(M + offs_m)[:, 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,
-        do,
-        m,
-        D,
-        q2k_index,
-        q2k_num,
-        max_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,
-    )
-
-    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
-    dq_acc *= LN2
-    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):
-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
     sm_scale = 1.0 / math.sqrt(D)
     max_kv_blks = q2k_index.shape[-1]
-    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 T % 64 == 0, f"T must be a multiple of 64, but got {T}"
     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()}"
-    )
+        -1] == T // 64, f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
     o = torch.empty_like(q)
-    M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
+    M = torch.empty((B, H, T), dtype=torch.float32, device=q.device)
 
-    grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
+    grid = lambda _: (triton.cdiv(T, 64), B * H, 1)
     _attn_fwd_sparse[grid](q,
                            k,
                            v,
@@ -735,8 +508,7 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
                            o.stride(3),
                            B,
                            H,
-                           Tq,
-                           Tkv,
+                           T,
                            HEAD_DIM=D,
                            STAGE=3)
 
@@ -746,21 +518,21 @@ def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num,
 def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
                                       k2q_index, k2q_num, variable_block_sizes):
     assert do.is_contiguous()
+    assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
 
-    B, H, Tq, D = q.shape
-    Tkv = k.shape[2]
+    B, H, T, D = q.shape
     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]
+    BATCH, N_HEAD, N_CTX = q.shape[:3]
     BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
     RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
     arg_k = k
     arg_k = arg_k * (sm_scale * RCP_LN2)
     PRE_BLOCK = 64
-    assert Tq % PRE_BLOCK == 0
-    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
+    assert N_CTX % PRE_BLOCK == 0
+    pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
     delta = torch.empty_like(M)
     _attn_bwd_preprocess[pre_grid](
         o,
@@ -768,7 +540,7 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
         delta,  #
         BATCH,
         N_HEAD,
-        Tq,  #
+        N_CTX,  #
         BLOCK_M=PRE_BLOCK,
         HEAD_DIM=D  #
     )
@@ -776,75 +548,36 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num,
     max_q_blks = k2q_index.shape[-1]
     max_kv_blks = q2k_index.shape[-1]
 
-    # dK/dV kernel: grid over KV blocks
-    grid_kv = (Tkv // BLOCK_N1, 1, BATCH * N_HEAD)
-    _attn_bwd_dkdv_kernel[grid_kv](
+    grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+    _attn_bwd[grid](
         q,
         arg_k,
         v,
         sm_scale,
         do,
+        dq,
         dk,
-        dv,
+        dv,  #
         M,
-        delta,
+        delta,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
         k2q_index,
         k2q_num,
         max_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),
+        q.stride(2),
+        q.stride(3),  #
         N_HEAD,
-        Tq,
-        Tkv,
+        N_CTX,  #
         BLOCK_M1=BLOCK_M1,
-        BLOCK_N1=BLOCK_N1,
-        HEAD_DIM=D,
-    )
-
-    # 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,
-        do,
-        dq,
-        M,
-        delta,
-        q2k_index,
-        q2k_num,
-        max_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,
+        BLOCK_N1=BLOCK_N1,  #
         BLOCK_M2=BLOCK_M2,
-        BLOCK_N2=BLOCK_N2,
-        HEAD_DIM=D,
+        BLOCK_N2=BLOCK_N2,  #
+        HEAD_DIM=D  #
     )
 
     return dq, dk, dv

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.6"
+__version__ = "0.2.5"

fastvideo/configs/pipelines/hunyuangamecraft.py

@@ -57,10 +57,17 @@ def llama_preprocess_text(prompt: str) -> str:
 def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
     """Extract hidden states from LLaMA output, skipping instruction tokens."""
     hidden_state_skip_layer = 2
-    assert outputs.hidden_states is not None
-    hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
-    last_hidden_state: torch.Tensor = hidden_states[-(hidden_state_skip_layer +
-                                                      1)]
+    hidden_states = outputs.hidden_states
+    if hidden_states is not None and len(hidden_states) > hidden_state_skip_layer:
+        last_hidden_state: torch.Tensor = hidden_states[
+            -(hidden_state_skip_layer + 1)]
+    elif outputs.last_hidden_state is not None:
+        # Fallback for encoder outputs without hidden_states.
+        last_hidden_state = outputs.last_hidden_state
+    else:
+        raise ValueError(
+            "LLaMA encoder output must contain hidden_states or last_hidden_state."
+        )
     crop_start = prompt_template_video.get("crop_start", -1)
     last_hidden_state = last_hidden_state[:, crop_start:]
     return last_hidden_state

fastvideo/configs/sample/ltx2.py

@@ -9,7 +9,6 @@ class LTX2BaseSamplingParam(SamplingParam):
     """Default sampling parameters for LTX-2 base one-stage T2V.
 
     Values follow the official LTX-2 one-stage defaults.
-    Multi-modal CFG params are read by ``LTX2DenoisingStage``.
     """
 
     seed: int = 10
@@ -38,12 +37,6 @@ class LTX2BaseSamplingParam(SamplingParam):
         "awkward pauses, incorrect timing, unnatural transitions, "
         "inconsistent framing, tilted camera, flat lighting, inconsistent "
         "tone, cinematic oversaturation, stylized filters, or AI artifacts.")
-    # Official LTX-2 multi-modal CFG defaults.
-    ltx2_cfg_scale_video: float = 3.0
-    ltx2_cfg_scale_audio: float = 7.0
-    ltx2_modality_scale_video: float = 3.0
-    ltx2_modality_scale_audio: float = 3.0
-    ltx2_rescale_scale: float = 0.7
 
 
 @dataclass

fastvideo/models/dits/ltx2.py

@@ -1534,7 +1534,6 @@ class BasicAVTransformerBlock(torch.nn.Module):
         audio: TransformerArgs | None,
         video_attention_mask: torch.Tensor | None = None,
         audio_attention_mask: torch.Tensor | None = None,
-        skip_cross_modal_attn: bool = False,
     ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
         """Forward pass for transformer block.
 
@@ -1543,8 +1542,6 @@ class BasicAVTransformerBlock(torch.nn.Module):
             audio: Audio transformer args
             video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
             audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
-            skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
-                attention (used for the modality-isolated CFG pass).
         """
         vx = video.x if video is not None else None
         ax = audio.x if audio is not None else None
@@ -1588,7 +1585,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 mask=audio.context_mask,
             )
 
-        if (run_a2v or run_v2a) and not skip_cross_modal_attn:
+        if run_a2v or run_v2a:
             vx_norm3 = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps)
             ax_norm3 = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps)
 
@@ -2009,7 +2006,6 @@ class LTXModel(torch.nn.Module):
         audio: TransformerArgs | None,
         video_attention_mask: torch.Tensor | None = None,
         audio_attention_mask: torch.Tensor | None = None,
-        skip_cross_modal_attn: bool = False,
     ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
         for block in self.transformer_blocks:
             video, audio = block(
@@ -2017,7 +2013,6 @@ class LTXModel(torch.nn.Module):
                 audio=audio,
                 video_attention_mask=video_attention_mask,
                 audio_attention_mask=audio_attention_mask,
-                skip_cross_modal_attn=skip_cross_modal_attn,
             )
         return video, audio
 
@@ -2044,7 +2039,6 @@ class LTXModel(torch.nn.Module):
         audio: Modality | None,
         video_attention_mask: torch.Tensor | None = None,
         audio_attention_mask: torch.Tensor | None = None,
-        skip_cross_modal_attn: bool = False,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Forward pass through the LTX model.
 
@@ -2053,9 +2047,6 @@ class LTXModel(torch.nn.Module):
             audio: Audio modality input
             video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
             audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
-            skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
-                attention in all transformer blocks (modality-isolated
-                CFG pass).
         """
         if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
             _debug_block_log_line(
@@ -2079,7 +2070,6 @@ class LTXModel(torch.nn.Module):
             audio_args,
             video_attention_mask=video_attention_mask,
             audio_attention_mask=audio_attention_mask,
-            skip_cross_modal_attn=skip_cross_modal_attn,
         )
 
         vx = (
@@ -2248,7 +2238,6 @@ class LTX2Transformer3DModel(CachableDiT):
         audio_encoder_hidden_states: torch.Tensor | None = None,
         audio_timestep: torch.Tensor | None = None,
         audio_encoder_attention_mask: torch.Tensor | None = None,
-        skip_cross_modal_attn: bool = False,
         **kwargs,
     ) -> torch.Tensor:
         if isinstance(encoder_hidden_states, list):
@@ -2405,7 +2394,6 @@ class LTX2Transformer3DModel(CachableDiT):
             audio=audio_modality,
             video_attention_mask=video_attention_mask,
             audio_attention_mask=audio_attention_mask,
-            skip_cross_modal_attn=skip_cross_modal_attn,
         )
 
         # Denoised prediction

fastvideo/models/registry.py

@@ -454,4 +454,4 @@ ModelRegistry = _ModelRegistry({
     )
     for model_arch, (component_name, mod_relname,
                      cls_name) in _FAST_VIDEO_MODELS.items()
-})
+})

fastvideo/pipelines/pipeline_batch_info.py

@@ -175,13 +175,6 @@ class ForwardBatch:
     eta: float = 0.0
     sigmas: list[float] | None = None
 
-    # LTX-2 multi-modal CFG parameters
-    ltx2_cfg_scale_video: float = 1.0
-    ltx2_cfg_scale_audio: float = 1.0
-    ltx2_modality_scale_video: float = 1.0
-    ltx2_modality_scale_audio: float = 1.0
-    ltx2_rescale_scale: float = 0.0
-
     n_tokens: int | None = None
 
     # Other parameters that may be needed by specific schedulers

fastvideo/pipelines/pipeline_registry.py

@@ -174,7 +174,7 @@ def _import_pipeline_classes_cached(
                     pipeline_module = importlib.import_module(module_name)
                     if not hasattr(pipeline_module, "EntryClass"):
                         continue
-                    entry_cls = pipeline_module.EntryClass
+                    entry_cls = getattr(pipeline_module, "EntryClass")
                     entry_cls_list = ([
                         entry_cls
                     ] if not isinstance(entry_cls, list) else entry_cls)

fastvideo/pipelines/stages/gamecraft_denoising.py

@@ -175,11 +175,11 @@ class GameCraftDenoisingStage(DenoisingStage):
 
         # Debug logging
         logger.debug("[GameCraft DEBUG] latents shape: %s, min/max: %.4f/%.4f",
-                     latents.shape, latents.min(), latents.max())
+                    latents.shape, latents.min(), latents.max())
         logger.debug("[GameCraft DEBUG] camera_states: %s",
-                     camera_states.shape if camera_states is not None else None)
+                    camera_states.shape if camera_states is not None else None)
         logger.debug("[GameCraft DEBUG] prompt_embeds[0] shape: %s",
-                     prompt_embeds[0].shape)
+                    prompt_embeds[0].shape)
 
         # Initialize lists for trajectory
         trajectory_timesteps: list[torch.Tensor] = []

fastvideo/pipelines/stages/ltx2_denoising.py

@@ -141,7 +141,6 @@ class LTX2DenoisingStage(PipelineStage):
                 video_shape)
         else:
             token_count = 1
-
         timestep_template = torch.ones(
             (latents.shape[0], token_count),
             device=latents.device,
@@ -210,27 +209,12 @@ class LTX2DenoisingStage(PipelineStage):
                 device=latents.device,
                 dtype=torch.float32,
             )
-        # Multi-modal CFG parameters (per-stream scales).
-        cfg_scale_video = batch.ltx2_cfg_scale_video
-        cfg_scale_audio = batch.ltx2_cfg_scale_audio
-        modality_scale_video = batch.ltx2_modality_scale_video
-        modality_scale_audio = batch.ltx2_modality_scale_audio
-        rescale_scale = batch.ltx2_rescale_scale
-
-        do_cfg = batch.do_classifier_free_guidance
-
         logger.info(
-            "[LTX2] Denoising start: steps=%d dtype=%s "
-            "cfg_video=%.1f cfg_audio=%.1f mod_video=%.1f "
-            "mod_audio=%.1f rescale=%.2f "
+            "[LTX2] Denoising start: steps=%d dtype=%s guidance=%s "
             "sigmas_shape=%s latents_shape=%s",
             batch.num_inference_steps,
             target_dtype,
-            cfg_scale_video,
-            cfg_scale_audio,
-            modality_scale_video,
-            modality_scale_audio,
-            rescale_scale,
+            batch.guidance_scale,
             tuple(sigmas.shape),
             tuple(latents.shape),
         )
@@ -251,7 +235,6 @@ class LTX2DenoisingStage(PipelineStage):
                     attn_metadata=None,
                     forward_batch=batch,
             ):
-                # Pass 1: Full conditioning (text + cross-modal)
                 pos_outputs = self.transformer(
                     hidden_states=latents.to(target_dtype),
                     encoder_hidden_states=prompt_embeds,
@@ -267,8 +250,8 @@ class LTX2DenoisingStage(PipelineStage):
                     pos_denoised = pos_outputs
                     pos_audio = None
 
-                if do_cfg:
-                    # Pass 2: Unconditioned text (negative prompt)
+                # Only run negative pass if CFG is enabled
+                if batch.do_classifier_free_guidance:
                     neg_outputs = self.transformer(
                         hidden_states=latents.to(target_dtype),
                         encoder_hidden_states=neg_prompt_embeds,
@@ -283,48 +266,11 @@ class LTX2DenoisingStage(PipelineStage):
                     else:
                         neg_denoised = neg_outputs
                         neg_audio = None
-
-                    # Pass 3: Modality-isolated (skip cross-modal attn)
-                    mod_outputs = self.transformer(
-                        hidden_states=latents.to(target_dtype),
-                        encoder_hidden_states=prompt_embeds,
-                        encoder_attention_mask=prompt_mask,
-                        timestep=timestep,
-                        audio_hidden_states=audio_latents,
-                        audio_encoder_hidden_states=audio_context_p,
-                        audio_timestep=audio_timestep,
-                        skip_cross_modal_attn=True,
-                    )
-                    if isinstance(mod_outputs, tuple):
-                        mod_denoised, mod_audio = mod_outputs
-                    else:
-                        mod_denoised = mod_outputs
-                        mod_audio = None
-
-                    # Multi-modal guidance formula per stream.
-                    vid = (pos_denoised + (cfg_scale_video - 1) *
-                           (pos_denoised - neg_denoised) +
-                           (modality_scale_video - 1) *
-                           (pos_denoised - mod_denoised))
-                    aud = None
-                    if pos_audio is not None:
-                        aud = (pos_audio + (cfg_scale_audio - 1) *
-                               (pos_audio - neg_audio) +
-                               (modality_scale_audio - 1) *
-                               (pos_audio - mod_audio))
-
-                    # Guidance rescaling (prevents saturation).
-                    if rescale_scale > 0:
-                        f_v = pos_denoised.std() / vid.std()
-                        f_v = rescale_scale * f_v + (1 - rescale_scale)
-                        vid = vid * f_v
-                        if aud is not None:
-                            f_a = pos_audio.std() / aud.std()
-                            f_a = rescale_scale * f_a + (1 - rescale_scale)
-                            aud = aud * f_a
-
-                    pos_denoised = vid
-                    pos_audio = aud
+                    pos_denoised = pos_denoised + (batch.guidance_scale - 1) * (
+                        pos_denoised - neg_denoised)
+                    if pos_audio is not None and neg_audio is not None:
+                        pos_audio = pos_audio + (batch.guidance_scale -
+                                                 1) * (pos_audio - neg_audio)
 
             sigma_value = sigma.to(torch.float32) if isinstance(
                 sigma, torch.Tensor) else torch.tensor(
@@ -335,7 +281,6 @@ class LTX2DenoisingStage(PipelineStage):
             dt = sigma_next - sigma
             velocity = ((latents.float() - pos_denoised.float()) /
                         sigma_value).to(latents.dtype)
-
             latents = (latents.float() + velocity.float() * dt).to(
                 latents.dtype)
             if pos_audio is not None and audio_latents is not None:

fastvideo/registry.py

@@ -356,7 +356,7 @@ def _register_configs() -> None:
         sampling_param_cls=HunyuanGameCraftSamplingParam,
         pipeline_config_cls=HunyuanGameCraftPipelineConfig,
         hf_model_paths=[
-            "FastVideo/HunyuanGameCraft-Diffusers",
+            "tencent/Hunyuan-GameCraft-1.0",
         ],
         model_detectors=[lambda path: "gamecraft" in path.lower()],
     )
@@ -679,4 +679,4 @@ __all__ = [
     "get_pipeline_config_cls_from_name",
     "get_sampling_param_cls_for_name",
     "get_pipeline_config_classes",
-]
+]

fastvideo/tests/ssim/test_causal_similarity.py

@@ -19,10 +19,6 @@ if "A40" in device_name:
     device_reference_folder = "A40" + device_reference_folder_suffix
 elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-elif "H200" in device_name:
-    device_reference_folder = "H200" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")

fastvideo/tests/ssim/test_inference_similarity.py

@@ -41,10 +41,6 @@ if "A40" in device_name:
     device_reference_folder = "A40" + device_reference_folder_suffix
 elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-elif "H200" in device_name:
-    device_reference_folder = "H200" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")

fastvideo/tests/ssim/test_lingbot_similarity.py

@@ -47,8 +47,6 @@ elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
 elif "H100" in device_name:
     device_reference_folder = "H100" + device_reference_folder_suffix
-elif "H200" in device_name:
-    device_reference_folder = "H200" + device_reference_folder_suffix
 else:
     device_reference_folder = None
     logger.warning("Unsupported device for ssim tests: %s", device_name)

fastvideo/tests/ssim/test_longcat_similarity.py

@@ -35,8 +35,6 @@ elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
 elif "H100" in device_name:
     device_reference_folder = "H100" + device_reference_folder_suffix
-elif "H200" in device_name:
-    device_reference_folder = "H200" + device_reference_folder_suffix
 else:
     logger.warning(f"Unsupported device for ssim tests: {device_name}")
 

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -22,10 +22,6 @@ if "A40" in device_name:
     device_reference_folder = "A40" + device_reference_folder_suffix
 elif "L40S" in device_name:
     device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-elif "H200" in device_name:
-    device_reference_folder = "H200" + device_reference_folder_suffix
 else:
     # device_reference_folder = "L40S" + device_reference_folder_suffix
     logger.warning(f"Unsupported device for ssim tests: {device_name}")

fastvideo/tests/ssim/test_sd35_similarity.py

@@ -22,8 +22,6 @@ def _device_reference_folder() -> str:
         return "L40S" + suffix
     if "H100" in device_name:
         return "H100" + suffix
-    if "H200" in device_name:
-        return "H200" + suffix
     if "RTX 4090" in device_name or "4090" in device_name:
         return "RTX4090" + suffix
 

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -121,9 +121,9 @@ def test_distributed_training():
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
-        'avg_step_time': 15,
+        'avg_step_time': 5,
         'grad_norm': 0.5,
-        'step_time': 15,
+        'step_time': 5,
         'train_loss': 0.04
     }
 

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -136,9 +136,9 @@ def test_distributed_training():
     wandb_summary = json.load(open(summary_file))
 
     fields_and_thresholds = {
-        'avg_step_time': 15.0,
+        'avg_step_time': 6.0,
         'grad_norm': 0.3,
-        'step_time': 15.0,
+        'step_time': 6.0,
         'train_loss': 0.0025
     }
 

fastvideo/training/training_pipeline.py

@@ -19,7 +19,6 @@ 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:
@@ -606,7 +605,6 @@ class TrainingPipeline(LoRAPipeline, ABC):
             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()

pyproject.toml

@@ -66,7 +66,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.6",
+    "fastvideo-kernel==0.2.5",
     "wheel",
 
     # Training Dependencies

pyproject_other.toml

@@ -65,7 +65,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.6",
+    "fastvideo-kernel==0.2.5",
     "wheel",
 
     # Training Dependencies