exp7: fastvideo-kernel 0.3.4 CUDA VSA w/ triton fallback

requirements.txt

@@ -5,6 +5,10 @@ diffusers==0.40.0
 # is what compiles the vendored VSA block-sparse kernels. torchvision is not needed anywhere in this Space.
 transformers
 accelerate
+# FastVideo's published CUDA kernels: the block-sparse VSA forward, sm_100a source with sm_120a cubins in the
+# same fatbin (this pool's GPU). Only needs torch>=2.5, so the platform torch is untouched. If the op refuses
+# to run, `vsa_h3` falls back to the vendored Triton kernels below.
+fastvideo-kernel==0.3.4
 # PyAV muxes the generated soundtrack onto the frames (`diffusers.utils.encode_video`).
 av
 pillow

vsa_h3.py

@@ -33,6 +33,7 @@ from __future__ import annotations
 
 import functools
 import math
+import os
 from dataclasses import dataclass
 
 import torch
@@ -279,6 +280,36 @@ def _block_mask(scores: torch.Tensor, geometry: VSAGeometry) -> torch.Tensor:
     return mask
 
 
+# The sparse-attention kernel, resolved once per process. FastVideo's published CUDA wheel carries the
+# block-sparse forward as sm_100a source with sm_120a cubins in the same fatbin — this pool's GPU is sm_120 —
+# but its `is_supported` gate hard-checks capability == (10, 0) and would never select it here. We call the
+# public `block_sparse_attn_sm100a` entry directly (need_lse=False, the production inference path) and treat
+# the dense-equivalence `/selftest` as the numerical gate. Any import or launch failure falls back to the
+# vendored Triton kernels; `H3_VSA_CUDA=0` forces Triton up front.
+_CUDA_SPARSE: dict = {"resolved": False, "op": None}
+
+
+def _resolve_cuda_sparse_op():
+    if _CUDA_SPARSE["resolved"]:
+        return _CUDA_SPARSE["op"]
+    _CUDA_SPARSE["resolved"] = True
+    if os.environ.get("H3_VSA_CUDA", "1") != "1":
+        return None
+    try:
+        from fastvideo_kernel import block_sparse_attn_sm100a as sm100a
+
+        op = sm100a.block_sparse_attn_sm100a
+        probe = sm100a._FWD_BY_BLOCK  # noqa: F841 - import-time check that the extension loaded
+        if not sm100a._HAS_VSA_SM100A:
+            raise ImportError("fastvideo_kernel compiled without the sm100a VSA forward")
+        _CUDA_SPARSE["op"] = op
+        print("[vsa] CUDA sparse-attention kernel: fastvideo_kernel (sm100a fatbin on sm120)", flush=True)
+    except Exception as error:  # noqa: BLE001 - any failure means Triton
+        print(f"[vsa] fastvideo_kernel unavailable ({type(error).__name__}: {error}); using vendored Triton", flush=True)
+        _CUDA_SPARSE["op"] = None
+    return _CUDA_SPARSE["op"]
+
+
 def sparse_attention(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -298,14 +329,34 @@ def sparse_attention(
 
     mask = _block_mask(scores, geometry)
     q2k_index, q2k_num = map_to_index(mask)
-    out_tiled, _ = triton_block_sparse_attn_forward(
-        query_tiled,
-        key_tiled,
-        value_tiled,
-        q2k_index,
-        q2k_num,
-        geometry.variable_block_sizes,
-    )
+    # The CUDA path only handles an even tile count (one CTA owns an adjacent pair of query blocks); anything
+    # else — or any launch failure — falls back to the vendored Triton kernels for the rest of the process.
+    out_tiled = None
+    cuda_op = _resolve_cuda_sparse_op()
+    if cuda_op is not None and geometry.n_tiles % 2 == 0:
+        try:
+            out_tiled, _ = cuda_op(
+                query_tiled,
+                key_tiled,
+                value_tiled,
+                q2k_index,
+                q2k_num,
+                geometry.variable_block_sizes,
+                need_lse=False,
+            )
+        except Exception as error:  # noqa: BLE001 - a refused launch must not kill the request
+            print(f"[vsa] CUDA kernel failed ({type(error).__name__}: {error}); falling back to Triton", flush=True)
+            _CUDA_SPARSE["op"] = None
+            out_tiled = None
+    if out_tiled is None:
+        out_tiled, _ = triton_block_sparse_attn_forward(
+            query_tiled,
+            key_tiled,
+            value_tiled,
+            q2k_index,
+            q2k_num,
+            geometry.variable_block_sizes,
+        )
     out = out_tiled.index_select(2, geometry.untile_index)
 
     if gate_compress is not None: