[perf] Add FA4 CuTe backward support for VSA-256 (#1639)

Co-authored-by: Hyunsung Lee <hyunsungl@sizigistudios.com>
Co-authored-by: alexzms <3036648523@qq.com>

Author: 3 people

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
+ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without

fastvideo-kernel/README.md

@@ -17,7 +17,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
 | Kernels | Where | Used when |
 |---|---|---|
 | Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
-| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
+| FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
 | VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
 
 ## What gets built where, and when
@@ -64,28 +64,49 @@ cd fastvideo-kernel
 ./build.sh --rocm
 ```
 
-### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath)
+### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)
 
-The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the
+The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the
 FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is
 an **optional** dependency: it is imported lazily, and `video_sparse_attn`
 transparently falls back to the Triton backend when it is absent (so the package is
 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
+The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`
+and the public/private forward-backward bridges in `flash_attn.cute.interface`) are provided upstream by
 [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+commit `14c377950125c70b7a9dabf9c561fca53715ac7d`, 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.
+have incompatible block-sparse forward/backward interfaces.
+
+Install it under its distribution name so its own runtime stack resolves with it.
+Do **not** pre-install `nvidia-cutlass-dsl` by hand: this revision pins
+`nvidia-cutlass-dsl==4.6.0.dev0` exactly, and a hand-installed 4.5.x floor either
+gets silently upgraded or, if something else holds it back, leaves the CuTe
+kernels broken.
 
 ```bash
-pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
+pip install torchvision
+pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"
 ```
 
-The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
-`tests/test_vsa256_forward*.py`.
+That resolves `nvidia-cutlass-dsl` to 4.6.0.dev0 and `quack-kernels` to 0.5.3, a
+combination this revision works with. A mismatched CuTe DSL only surfaces when the
+kernel JIT-compiles, so the error points at CuTe internals rather than at the
+install:
+
+| Error on first VSA-128/256 CuTe call | Cause |
+|---|---|
+| `TypeError: fmax() missing 1 required positional argument: 'b'` | `nvidia-cutlass-dsl` 4.5.x |
+| `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'` | `quack-kernels` older than 0.5.1 |
+| `ImportError: cannot import name 'alloc_reserved_mbarrier'` | `quack-kernels` 0.6.2 or newer |
+
+An environment whose `flash_attn.cute` came from a prebuilt flash-attn wheel rather
+than from this pin hits the first row; that is what the overlay step in
+`docker/Dockerfile` works around.
+
+The CuTe kernels JIT-compile on first use. Forward and backward are verified on
+Blackwell (sm_100) against `tests/test_vsa128_*.py` and `tests/test_vsa256_*.py`.
 
 ## Usage
 
@@ -142,6 +163,14 @@ After building/installing `fastvideo-kernel`, run:
 ```bash
 cd fastvideo-kernel
 python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
+
+# VSA-256 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 256 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20
+
+# VSA-128 FA4 CuTe forward/backward on Blackwell
+python benchmarks/bench_vsa.py --block_size 128 --use_cute \
+  --batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40
 ```
 
 ### TurboDiffusion Kernels

fastvideo-kernel/benchmarks/bench_vsa.py

@@ -2,8 +2,9 @@
 """
 Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.
 
-This script benchmarks the autograd-enabled wrapper:
-  - fastvideo_kernel.block_sparse_attn.block_sparse_attn
+This script benchmarks the autograd-enabled wrappers:
+  - 64-token TK/Triton: fastvideo_kernel.block_sparse_attn.block_sparse_attn
+  - 128/256-token Triton/CuTe: fastvideo_kernel.block_sparse_attn_256
 
 So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
 """
@@ -23,9 +24,6 @@
 except Exception as e:  # pragma: no cover
     raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
 
-BLOCK_M = 64
-BLOCK_N = 64
-
 
 def set_seed(seed: int = 42) -> None:
     random.seed(seed)
@@ -41,7 +39,11 @@ def parse_arguments() -> argparse.Namespace:
     p.add_argument("--num_heads", type=int, default=12)
     p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
     p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
-    p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
+    p.add_argument("--q_seq_lens",
+                   type=int,
+                   nargs="+",
+                   default=[49152],
+                   help="Q sequence lengths (must be divisible by --block_size)")
     p.add_argument("--kv_seq_lens",
                    type=int,
                    nargs="+",
@@ -51,9 +53,13 @@ def parse_arguments() -> argparse.Namespace:
     p.add_argument("--rep", type=int, default=20)
     p.add_argument("--seed", type=int, default=42)
     p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
+    p.add_argument("--block_size", type=int, default=64, choices=[64, 128, 256])
     p.add_argument("--force_triton",
                    action="store_true",
                    help="Force wrapper to use Triton path (if supported by shapes).")
+    p.add_argument("--use_cute",
+                   action="store_true",
+                   help="Use the optional FA4 CuTe forward/backward path (requires --block_size 128 or 256).")
     return p.parse_args()
 
 
@@ -84,18 +90,38 @@ def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
     return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)
 
 
+def _configure_backend(args: argparse.Namespace) -> None:
+    if args.use_cute and args.block_size not in (128, 256):
+        raise ValueError("--use_cute requires --block_size 128 or 256")
+    if args.use_cute and args.force_triton:
+        raise ValueError("--use_cute and --force_triton are mutually exclusive")
+
+    if args.force_triton:
+        os.environ.pop("FASTVIDEO_VSA_CUTEDSL", None)
+        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
+    elif args.use_cute:
+        os.environ.pop("FASTVIDEO_VSA_TRITON", None)
+        os.environ.pop("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", None)
+        os.environ["FASTVIDEO_VSA_CUTEDSL"] = "1"
+
+
 def main() -> None:
     args = parse_arguments()
     set_seed(args.seed)
+    _configure_backend(args)
 
     dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
 
-    if args.force_triton:
-        os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
-
     from fastvideo_kernel.block_sparse_attn import block_sparse_attn
+    from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128, block_sparse_attn_256
 
     bs, h, d = args.batch_size, args.num_heads, args.head_dim
+    block_size = args.block_size
+    attention = {
+        64: block_sparse_attn,
+        128: block_sparse_attn_128,
+        256: block_sparse_attn_256,
+    }[block_size]
     kv_seq_lens = args.kv_seq_lens
     if kv_seq_lens is None:
         kv_seq_lens = args.q_seq_lens
@@ -105,20 +131,22 @@ def main() -> None:
     print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
     print(f"device: {torch.cuda.get_device_name(0)}")
     print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
-    print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
+    print(f"block_size={block_size}")
     print("NOTE: timings include wrapper overhead (map->index + dispatch).")
-    if args.force_triton:
+    if args.use_cute:
+        print("dispatch: FA4 CuTe")
+    elif args.force_triton:
         print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
     else:
         print("dispatch: SM90 if available, else Triton")
 
     for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
-        if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
-            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
+        if q_len % block_size != 0 or kv_len % block_size != 0:
+            print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by {block_size}")
             continue
 
-        num_q_blocks = q_len // BLOCK_M
-        num_kv_blocks = kv_len // BLOCK_N
+        num_q_blocks = q_len // block_size
+        num_kv_blocks = kv_len // block_size
         topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
         topk = min(topk, num_kv_blocks)
 
@@ -129,11 +157,11 @@ def main() -> None:
         q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
         block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)
 
-        # Variable block sizes: default full blocks (64 tokens per KV block)
-        variable_block_sizes = torch.full((num_kv_blocks, ), BLOCK_N, dtype=torch.int32, device="cuda")
+        # Variable block sizes: default full logical blocks.
+        variable_block_sizes = torch.full((num_kv_blocks, ), block_size, dtype=torch.int32, device="cuda")
 
         def _fwd():
-            return block_sparse_attn(q, k, v, block_map, variable_block_sizes)
+            return attention(q, k, v, block_map, variable_block_sizes)
 
         fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)
 
@@ -142,7 +170,7 @@ def _fwd():
         q_ = q.detach().requires_grad_(True)
         k_ = k.detach().requires_grad_(True)
         v_ = v.detach().requires_grad_(True)
-        o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
+        o_, _aux_ = attention(q_, k_, v_, block_map, variable_block_sizes)
         og = torch.randn_like(o_)
         loss = (o_ * og).sum()
 
@@ -156,7 +184,7 @@ def _fwd():
             rep=max(5, args.rep // 2),
         )
 
-        flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
+        flops = flops_sparse_attention(bs, h, d, q_len, topk, block_size)
         fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
         # Rough backward multiplier (attention backward typically ~2-3x forward)
         bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3

fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py

@@ -1,4 +1,4 @@
-"""VSA-256 block-sparse attention wrapper.
+"""VSA-128/256 block-sparse attention wrappers.
 
 The default 256-block path is Triton: it expands the logical 256-block map
 to the existing 64-block Triton kernel via a dense 4x4 expansion per logical
@@ -7,8 +7,8 @@
 The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is
 *opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to
 :mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates
-on 128-token KV blocks (this wrapper expands the logical 256-block map /
-sizes into that physical 128-block representation). The CuTe kernel
+on 128-token Q/KV blocks (the 256 wrapper expands its logical KV map and
+sizes into that physical representation). The CuTe kernel
 (``flash_attn.cute`` with block-sparsity) is an optional dependency,
 imported lazily only when this fastpath is selected.
 
@@ -35,7 +35,7 @@
 
 
 def _resolve_backend() -> str:
-    """Pick the backend for the 256-block VSA path.
+    """Pick the backend for the 128/256-block VSA paths.
 
     Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in
     via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe
@@ -49,6 +49,26 @@ def _resolve_backend() -> str:
     return "triton"
 
 
+def _expand_mask_and_sizes_128_to_64(
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Expand a [B, H, Qb128, KVb128] map to 64-token Triton tiles."""
+    expanded_mask = logical_mask_128.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
+    sizes_i32 = logical_kv_sizes_128.to(torch.int32)
+    offsets = torch.tensor(
+        [0, _KV_BLOCK_TRITON],
+        dtype=torch.int32,
+        device=sizes_i32.device,
+    )
+    expanded_sizes = torch.clamp(
+        sizes_i32[:, None] - offsets[None, :],
+        min=0,
+        max=_KV_BLOCK_TRITON,
+    ).reshape(-1)
+    return expanded_mask, expanded_sizes
+
+
 def _expand_mask_and_sizes_256_to_128(
     logical_mask_256: torch.Tensor,
     logical_kv_sizes_256: torch.Tensor,
@@ -112,6 +132,63 @@ def _triton_via_route_a(
     return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
 
 
+def _triton_via_route_a_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_mask_128: torch.Tensor,
+    logical_kv_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    from .triton_kernels.index import map_to_index as triton_map_to_index
+
+    mask_64, sizes_64 = _expand_mask_and_sizes_128_to_64(logical_mask_128, logical_kv_sizes_128)
+    q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
+    return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
+
+
+def block_sparse_attn_128(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, H, S, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        return _triton_via_route_a_128(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
+    return block_sparse_attn_cute_fwd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
+def block_sparse_attn_128_bshd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    logical_block_map_128: torch.Tensor,
+    logical_variable_block_sizes_128: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """VSA-128 sparse-branch entrypoint for [B, S, H, D] inputs."""
+    if logical_block_map_128.dim() == 3:
+        logical_block_map_128 = logical_block_map_128.unsqueeze(0)
+
+    if _resolve_backend() == "triton":
+        out_bhsd, aux = _triton_via_route_a_128(
+            q.transpose(1, 2).contiguous(),
+            k.transpose(1, 2).contiguous(),
+            v.transpose(1, 2).contiguous(),
+            logical_block_map_128,
+            logical_variable_block_sizes_128,
+        )
+        return out_bhsd.transpose(1, 2).contiguous(), aux
+
+    from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
+    return block_sparse_attn_cute_fwd_bshd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
+
+
 def block_sparse_attn_256(
     q: torch.Tensor,
     k: torch.Tensor,