[kernel] [bugfix] [ci] bump v0.2.4. Fix STA output handling, TurboDiffusion CUDA norm dtypes for fastvideo-kernel unit tests. (#1020)

Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com>

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import Tuple
 
 import torch
@@ -23,6 +24,12 @@ def _is_sm90() -> bool:
     return major == 9 and minor == 0
 
 
+def _force_triton() -> bool:
+    # Force Triton even on SM90 and even if the compiled extension is available.
+    # Useful for CI / debugging / parity testing.
+    return os.environ.get("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", "0") == "1"
+
+
 def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Pure-torch (no triton) conversion:
@@ -281,7 +288,7 @@ def block_sparse_attn(
     - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today).
     """
     block_sparse_fwd, block_sparse_bwd = _get_sm90_ops()
-    if _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None):
+    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: generally assumes q/k/v share the same padded length
     if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]:

fastvideo-kernel/python/fastvideo_kernel/ops.py

@@ -46,14 +46,22 @@ def sliding_tile_attention(
     flag = shape_map[seq_shape]
 
     for head_idx, (t, h, w) in enumerate(window_size):
+        # Per-head slices are not contiguous in the batch dimension when batch>1
+        # (they keep the original head-stride). The TK kernel assumes contiguous
+        # [B, H, S, D] layout, so we materialize a contiguous [B,1,S,D] view.
+        q_h = q[:, head_idx:head_idx + 1].contiguous()
+        k_h = k[:, head_idx:head_idx + 1].contiguous()
+        v_h = v[:, head_idx:head_idx + 1].contiguous()
+        o_h = torch.empty_like(q_h)
         sta_fwd(
-            q[:, head_idx:head_idx + 1], k[:, head_idx:head_idx + 1],
-            v[:, head_idx:head_idx + 1], output[:, head_idx:head_idx + 1],
+            q_h, k_h,
+            v_h, o_h,
             t, h, w, text_length, False, has_text, flag
         )
+        output[:, head_idx:head_idx + 1] = o_h
 
     if has_text:
-        sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag)
+        sta_fwd(q.contiguous(), k.contiguous(), v.contiguous(), output, 3, 3, 3, text_length, True, True, flag)
 
     return output[:, :, :seq_length]
 

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"