#!/usr/bin/env python3 """ Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs. 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. """ 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 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 divisible by --block_size)") 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("--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() 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 _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 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 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_size={block_size}") print("NOTE: timings include wrapper overhead (map->index + dispatch).") 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_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_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) 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 logical blocks. variable_block_sizes = torch.full((num_kv_blocks, ), block_size, dtype=torch.int32, device="cuda") def _fwd(): return attention(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_ = attention(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_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 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()