[kernel] Bump fastvideo-kernel pin to 0.3.1 and version the FA4 tile_mn port as 0.3.2 (#1539)

apps/dreamverse/docker/README.md

@@ -44,11 +44,10 @@ CUDA_TAG=12.6.3-cudnn-devel-ubuntu22.04 apps/dreamverse/docker/docker_build.sh
 ```
 
 Do not set `CUDA_TAG` and `CUDA_VERSION` together. The image installs FastVideo
-from this checkout with the `dreamverse` extra, including the FA4
-flash-attention fork and FlashInfer for NVFP4 quantization, and builds native
-FFmpeg.
+from this checkout with the `dreamverse` extra, including FA4 flash-attention
+and FlashInfer for NVFP4 quantization, and builds native FFmpeg.
 
-FastVideo's pinned `fastvideo-kernel==0.3.0` package is installed by default.
+FastVideo's pinned `fastvideo-kernel==0.3.1` package is installed by default.
 To rebuild `fastvideo-kernel` from this checkout during the image build, set:
 
 ```bash

fastvideo-kernel/README.md

@@ -39,13 +39,13 @@ fully usable without it).
 The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
 `flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
 [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `c19cd20e`: the wrapper targets that revision's `_flash_attn_fwd` signature
-(`m_block_size` / `n_block_size`); later upstream revisions reshaped it into a
-`tile_mn` tuple and are not drop-in compatible.
+commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
+have an incompatible `_flash_attn_fwd` signature.
 
 ```bash
 pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@c19cd20e#subdirectory=flash_attn/cute"
+pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
 ```
 
 The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against

fastvideo-kernel/pyproject.toml

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

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py

@@ -46,8 +46,7 @@ def _load_fa4_cute():
     return BlockSparseTensorsTorch, _flash_attn_fwd
 
 
-# FA4 BSA fwd uses (m_block_size, n_block_size); m_block_size=128 is the
-# Q-side tile, kv_block_size comes from the caller's VSA logical KV block.
+# Q-side tile size; kv_block_size comes from the caller's VSA logical KV block.
 _M_BLOCK_SIZE_DEFAULT = 128
 
 
@@ -182,9 +181,7 @@ def _cute_forward(
         block_size=(q_sparse_block_size, kv_block_size),
     )
 
-    # Upstream FA4 cute (the cutlass-4.5-safe ref this repo pins) folds the
-    # fork's m_block_size/n_block_size pair into tile_mn, and its main path
-    # returns (out, lse, p, row_max) -- slice the first two.
+    # _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
     out, lse = _flash_attn_fwd(
         q_bshd,
         k_bshd,

fastvideo-kernel/python/fastvideo_kernel/version.py

@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"

pyproject.toml

@@ -71,7 +71,7 @@ dependencies = [
     "remote-pdb",
 
     # Kernel & Packaging
-    "fastvideo-kernel==0.2.6",
+    "fastvideo-kernel==0.3.1",
     "wheel",
 
     # Training Dependencies
@@ -97,12 +97,11 @@ prerelease = "allow"
 members = ["apps/dreamverse"]
 
 [tool.uv.sources]
-# No prebuilt aarch64 wheel for fastvideo-kernel exists on PyPI, so on ARM
-# (e.g. the NVIDIA DGX Spark / GB10) build it from the in-tree source instead;
-# x86_64 keeps the published PyPI wheel. This table is uv-only: plain pip ignores
-# it and it is not part of the published wheel metadata, so PyPI consumers are
-# unaffected. Pass --no-sources to bypass it (e.g. when the kernel is already
-# built/installed separately, as in the Spark install guide).
+# PyPI's aarch64 fastvideo-kernel wheels do not cover the NVIDIA DGX Spark /
+# GB10 (sm_121), so on ARM build from the in-tree source instead; x86_64 keeps
+# the published PyPI wheel. This table is uv-only: plain pip ignores it. Pass
+# --no-sources to bypass it (e.g. when the kernel is already installed, as in
+# the Spark install guide).
 fastvideo-kernel = [
   { path = "fastvideo-kernel", marker = "platform_machine == 'aarch64'" },
 ]
@@ -113,8 +112,8 @@ fastvideo-kernel = [
 # uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The
 # supported backends (cu126, cu130) both provide torch 2.12.0.
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
-# FA4 cute: upstream Dao-AILab (cutlass-4.5-safe; uses cute.ThrMma). torch.compile
-# support comes from FastVideo's own custom_op wrappers, so no fork is needed.
+# FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
+# comes from FastVideo's own custom_op wrappers.
 flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "940cd9680f3315f2f06b43ab5bea2c2cf2d96806", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]