From dbf3917bf4fa264d8aaa1d2dd8b2a7c1da40cde1 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Thu, 22 Jan 2026 11:35:02 -0800 Subject: [PATCH] [fastvideo-kernel] replace map to index with Triton implementation + add vsa benchmark (#1029) --- fastvideo-kernel/README.md | 11 ++ fastvideo-kernel/benchmarks/bench_vsa.py | 166 ++++++++++++++++++ fastvideo-kernel/pyproject.toml | 2 +- .../fastvideo_kernel/block_sparse_attn.py | 45 +++-- .../python/fastvideo_kernel/version.py | 2 +- pyproject.toml | 2 +- pyproject_other.toml | 2 +- 7 files changed, 202 insertions(+), 28 deletions(-) create mode 100644 fastvideo-kernel/benchmarks/bench_vsa.py diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index a894ea733..010779280 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -40,6 +40,17 @@ out = video_sparse_attn(q, k, v, block_sizes, block_sizes, topk=5) out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) ``` +## Benchmark + +### VSA (block-sparse) TFLOPs + +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 +``` + ### TurboDiffusion Kernels This package also includes kernels from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion), including INT8 GEMM, Quantization, RMSNorm and LayerNorm. diff --git a/fastvideo-kernel/benchmarks/bench_vsa.py b/fastvideo-kernel/benchmarks/bench_vsa.py new file mode 100644 index 000000000..28a9f6aef --- /dev/null +++ b/fastvideo-kernel/benchmarks/bench_vsa.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python3 +""" +Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs. + +This script benchmarks the autograd-enabled wrapper: + - fastvideo_kernel.block_sparse_attn.block_sparse_attn + +So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time. +""" + +from __future__ import annotations + +import argparse +import os +import random +from typing import Tuple, Callable + +import numpy as np +import torch + +try: + from triton.testing import do_bench +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) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def parse_arguments() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Benchmark FastVideo VSA block-sparse attention") + p.add_argument("--batch_size", type=int, default=1) + 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("--kv_seq_lens", type=int, nargs="+", default=None, help="KV sequence lengths (defaults to q_seq_len)") + p.add_argument("--warmup", type=int, default=5) + 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("--force_triton", action="store_true", help="Force wrapper to use Triton path (if supported by shapes).") + return p.parse_args() + + +def create_qkv(batch: int, heads: int, q_len: int, kv_len: int, d: int, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = torch.randn(batch, heads, q_len, d, dtype=dtype, device="cuda") + k = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda") + v = torch.randn(batch, heads, kv_len, d, dtype=dtype, device="cuda") + return q, k, v + + +def make_block_map(bs: int, h: int, num_q_blocks: int, num_kv_blocks: int, topk: int) -> torch.Tensor: + # block_map: [bs, h, num_q_blocks, num_kv_blocks] bool + scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device="cuda") + topk = min(max(1, topk), num_kv_blocks) + idx = torch.topk(scores, topk, dim=-1).indices + block_map = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device="cuda") + block_map.scatter_(-1, idx, True) + return block_map + + +def flops_sparse_attention(bs: int, h: int, d: int, q_len: int, topk_blocks: int, block_n: int) -> float: + # Approx: QK^T + PV, each is ~2*bs*h*q_len*(topk_blocks*block_n)*d + return 4.0 * bs * h * d * q_len * (topk_blocks * block_n) + + +def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float: + return do_bench(fn, warmup=warmup, rep=rep, quantiles=None) + + +def main() -> None: + args = parse_arguments() + set_seed(args.seed) + + 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 + + bs, h, d = args.batch_size, args.num_heads, args.head_dim + kv_seq_lens = args.kv_seq_lens + if kv_seq_lens is None: + kv_seq_lens = args.q_seq_lens + if len(kv_seq_lens) != len(args.q_seq_lens): + raise ValueError("kv_seq_lens must have the same number of entries as q_seq_lens (or be omitted).") + + 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("NOTE: timings include wrapper overhead (map->index + dispatch).") + if 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") + continue + + num_q_blocks = q_len // BLOCK_M + num_kv_blocks = kv_len // BLOCK_N + topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10) + topk = min(topk, num_kv_blocks) + + print("\n" + "=" * 80) + print(f"q_len={q_len}, kv_len={kv_len}, num_q_blocks={num_q_blocks}, num_kv_blocks={num_kv_blocks}, topk={topk}") + + 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") + + def _fwd(): + return block_sparse_attn(q, k, v, block_map, variable_block_sizes) + + fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep) + + # Backward benchmark (wrapper autograd). We build the graph once, then repeatedly run backward + # on the retained graph so bwd timing excludes the forward compute. + 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) + og = torch.randn_like(o_) + loss = (o_ * og).sum() + + for _ in range(max(1, args.warmup // 2)): + torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True) + torch.cuda.synchronize() + + bwd_ms = bench_ms( + lambda: torch.autograd.grad(loss, (q_, k_, v_), retain_graph=True), + warmup=0, + rep=max(5, args.rep // 2), + ) + + flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N) + 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 + + print(f"fwd(wrapper): {fwd_ms:.3f} ms | {fwd_tflops:.2f} TFLOPs (approx)") + print(f"bwd(wrapper): {bwd_ms:.3f} ms | {bwd_tflops:.2f} TFLOPs (approx)") + + +if __name__ == "__main__": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for this benchmark.") + main() + + diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml index d593c6f38..40fbafd12 100644 --- a/fastvideo-kernel/pyproject.toml +++ b/fastvideo-kernel/pyproject.toml @@ -9,7 +9,7 @@ build-backend = "scikit_build_core.build" [project] name = "fastvideo-kernel" -version = "0.2.4" +version = "0.2.5" description = "Unified CUDA kernels for FastVideo" readme = "README.md" requires-python = ">=3.10" diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 7206034fd..aeb25cf34 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -30,13 +30,12 @@ def _force_triton() -> bool: 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]: +def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ - Pure-torch (no triton) conversion: - block_map: [B, H, Q, KV] bool (or [H, Q, KV] which will be treated as B=1) - returns: - index: [B, H, Q, KV] int32 (packed KV indices, -1 padding) - num: [B, H, Q] int32 (#kv blocks per q block) + Preferred map->index conversion used by the wrapper. + + This wrapper **requires** the Triton implementation. + If Triton (or the Triton map_to_index module) is not available, it raises. """ if block_map.dim() == 3: block_map = block_map.unsqueeze(0) @@ -45,20 +44,17 @@ def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Te if block_map.dtype != torch.bool: block_map = block_map.to(torch.bool) - B, H, Q, KV = block_map.shape - index = torch.full((B, H, Q, KV), -1, dtype=torch.int32, device=block_map.device) - num = torch.zeros((B, H, Q), dtype=torch.int32, device=block_map.device) + if not block_map.is_cuda: + raise RuntimeError("block_map must be a CUDA tensor (Triton map_to_index required).") - # Small sizes in practice (B=1, H<=16, Q/KV<=64), so a Python loop is fine. - for b in range(B): - for h in range(H): - for q in range(Q): - kv_idx = torch.nonzero(block_map[b, h, q], as_tuple=False).flatten().to(torch.int32) - n = int(kv_idx.numel()) - if n: - index[b, h, q, :n] = kv_idx - num[b, h, q] = n - return index, num + try: + from fastvideo_kernel.triton_kernels.index import map_to_index as triton_map_to_index # local import + except Exception as e: + raise ImportError( + "Triton map_to_index is required but not available. " + "Ensure Triton is installed and fastvideo_kernel.triton_kernels.index is importable." + ) from e + return triton_map_to_index(block_map) @torch.library.custom_op( @@ -77,7 +73,7 @@ def block_sparse_attn_triton( k = k.contiguous() v = v.contiguous() block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index_torch(block_map) + q2k_idx, q2k_num = _map_to_index(block_map) from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import triton_block_sparse_attn_forward, @@ -87,6 +83,7 @@ def block_sparse_attn_triton( return o, M + @torch.library.register_fake("fastvideo_kernel::block_sparse_attn_triton") def _block_sparse_attn_triton_fake( q: torch.Tensor, @@ -117,8 +114,8 @@ def block_sparse_attn_backward_triton( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: grad_output = grad_output.contiguous() block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index_torch(block_map) - k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous()) + q2k_idx, q2k_num = _map_to_index(block_map) + k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous()) from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import triton_block_sparse_attn_backward, @@ -182,7 +179,7 @@ def block_sparse_attn_sm90( k_padded = k_padded.contiguous() v_padded = v_padded.contiguous() block_map = block_map.to(torch.bool) - q2k_idx, q2k_num = _map_to_index_torch(block_map) + q2k_idx, q2k_num = _map_to_index(block_map) o_padded, lse_padded = block_sparse_fwd( q_padded, k_padded, v_padded, q2k_idx, q2k_num, variable_block_sizes.int() @@ -224,7 +221,7 @@ def block_sparse_attn_backward_sm90( grad_output_padded = grad_output_padded.contiguous() block_map = block_map.to(torch.bool) - k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous()) + k2q_idx, k2q_num = _map_to_index(block_map.transpose(-1, -2).contiguous()) dq, dk, dv = block_sparse_bwd( q_padded, diff --git a/fastvideo-kernel/python/fastvideo_kernel/version.py b/fastvideo-kernel/python/fastvideo_kernel/version.py index 788da1fb3..fe404ae57 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/version.py +++ b/fastvideo-kernel/python/fastvideo_kernel/version.py @@ -1 +1 @@ -__version__ = "0.2.4" +__version__ = "0.2.5" diff --git a/pyproject.toml b/pyproject.toml index a704bd8dc..11bf2b2b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ dependencies = [ "remote-pdb", # Kernel & Packaging - "fastvideo-kernel==0.2.4", + "fastvideo-kernel==0.2.5", "wheel", # Training Dependencies diff --git a/pyproject_other.toml b/pyproject_other.toml index 746501a6d..68dc9cf8e 100644 --- a/pyproject_other.toml +++ b/pyproject_other.toml @@ -63,7 +63,7 @@ dependencies = [ "remote-pdb", # Kernel & Packaging - "fastvideo-kernel==0.2.4", + "fastvideo-kernel==0.2.5", "wheel", # Training Dependencies