From cbc3a6bb9dbd58b23db892d03dba249971f93f82 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Mon, 28 Jul 2025 20:14:40 -0700 Subject: [PATCH] [Feat] Support VSA with any resolution. (#650) --- .github/workflows/pr-test.yml | 2 +- csrc/attn/README.md | 29 +- csrc/attn/benchmarks/bench_vsa_hopper.py | 32 +- csrc/attn/tests/test_block_sparse.py | 289 ------------ .../tests/test_block_sparse_mask_to_index.py | 0 csrc/attn/tests/test_block_sparse_triton.py | 289 ------------ csrc/attn/tests/test_full_attn.py | 175 -------- csrc/attn/tests/test_triton.py | 183 -------- csrc/attn/tests/test_vsa.py | 159 +++++++ csrc/attn/tests/utils.py | 54 +++ csrc/attn/vsa.cpp | 4 +- csrc/attn/vsa/__init__.py | 237 +--------- csrc/attn/vsa/block_sparse_attn_triton.py | 424 ++++-------------- csrc/attn/vsa/block_sparse_h100.cu | 306 +++---------- csrc/attn/vsa/block_sparse_wrapper.py | 182 ++++++++ csrc/attn/vsa/index.py | 152 +++++++ csrc/attn/vsa/vsa.py | 47 -- examples/inference/basic/basic_dmd.py | 4 +- fastvideo/attention/backends/abstract.py | 9 +- .../attention/backends/sliding_tile_attn.py | 23 +- .../attention/backends/video_sparse_attn.py | 225 +++++++--- .../base_device_communicator.py | 47 +- fastvideo/pipelines/stages/denoising.py | 32 +- fastvideo/training/training_pipeline.py | 16 +- scripts/inference/v1_inference_wan_dmd.sh | 6 +- 25 files changed, 972 insertions(+), 1954 deletions(-) delete mode 100644 csrc/attn/tests/test_block_sparse.py delete mode 100644 csrc/attn/tests/test_block_sparse_mask_to_index.py delete mode 100644 csrc/attn/tests/test_block_sparse_triton.py delete mode 100644 csrc/attn/tests/test_full_attn.py delete mode 100644 csrc/attn/tests/test_triton.py create mode 100644 csrc/attn/tests/test_vsa.py create mode 100644 csrc/attn/tests/utils.py create mode 100644 csrc/attn/vsa/block_sparse_wrapper.py create mode 100644 csrc/attn/vsa/index.py delete mode 100644 csrc/attn/vsa/vsa.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index f3ae701bf..995d6b667 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -326,7 +326,7 @@ jobs: volume_size: 100 disk_size: 100 image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" - test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py" + test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_vsa.py" timeout_minutes: 30 secrets: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} diff --git a/csrc/attn/README.md b/csrc/attn/README.md index 87d94e556..dd6a12763 100644 --- a/csrc/attn/README.md +++ b/csrc/attn/README.md @@ -2,12 +2,7 @@ # Attention Kernel Used in FastVideo -## Sliding Tile Attention (STA) -We only support H100 for STA. -```bash -git submodule update --init --recursive -python setup_sta.py install -``` + ## Video Sparse Attention (VSA) We support H100 (via TK) and RTX 4090 (via triton) for VSA. @@ -16,6 +11,7 @@ git submodule update --init --recursive python setup_vsa.py install ``` + If you encounter error during installation, try below: Install C++20 for ThunderKittens: ```bash @@ -35,9 +31,17 @@ export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH ``` +## Sliding Tile Attention (STA) +We only support H100 for STA. +```bash +git submodule update --init --recursive +python setup_sta.py install +``` -## Usage -### STA + + + +### Usage End-2-end inference with FastVideo: ```bash bash scripts/inference/v1_inference_wan_STA.sh @@ -57,22 +61,19 @@ out = sliding_tile_attention(q, k, v, window_size, text_length) out = sliding_tile_attention(q, k, v, window_size, 0, False) ``` -### VSA -We do not officially supoort end-2-end inference with VSA in FastVideo yet. Stay tuned. - -## Test +### Test ```bash python tests/test_sta.py # test STA python tests/test_block_sparse.py # test VSA ``` -## Benchmark +### Benchmark ```bash python benchmarks/bench_sta.py ``` -## How Does STA Work? +### How Does STA Work? We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. diff --git a/csrc/attn/benchmarks/bench_vsa_hopper.py b/csrc/attn/benchmarks/bench_vsa_hopper.py index f3297f70e..aef0e4a48 100644 --- a/csrc/attn/benchmarks/bench_vsa_hopper.py +++ b/csrc/attn/benchmarks/bench_vsa_hopper.py @@ -1,9 +1,9 @@ import torch import argparse -from flash_attn.utils.benchmark import benchmark_forward +from triton.testing import do_bench from vsa import block_sparse_fwd, block_sparse_bwd from vsa import BLOCK_M, BLOCK_N - +import triton import numpy as np import random @@ -23,7 +23,7 @@ def parse_arguments(): parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') parser.add_argument('--batch_size', type=int, default=1, help='Batch size') parser.add_argument('--num_heads', type=int, default=12, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=64, help='Head dimension') + parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to') parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark') return parser.parse_args() @@ -134,15 +134,14 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_ torch.cuda.synchronize() # Benchmark forward - _, fwd_time = benchmark_forward( - block_sparse_fwd, - q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, - repeats=20, - verbose=False, - desc='Block Sparse Forward' + fwd_time = do_bench( + lambda: block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num), + warmup=5, + rep=20, + quantiles=None ) - sparse_tflops = flops / fwd_time.mean * 1e-12 + sparse_tflops = flops / fwd_time * 1e-12 * 1e3 print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}") # Backward pass @@ -154,16 +153,15 @@ def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_ torch.cuda.synchronize() # Benchmark backward - _, bwd_time = benchmark_forward( - block_sparse_bwd, - q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, - repeats=20, - verbose=False, - desc='Block Sparse Backward' + bwd_time = do_bench( + lambda: block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num), + warmup=5, + rep=20, + quantiles=None ) bwd_flops = 2.5 * flops # Approximation - sparse_bwd_tflops = bwd_flops / bwd_time.mean * 1e-12 + sparse_bwd_tflops = bwd_flops / bwd_time * 1e-12 * 1e3 print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}") return sparse_tflops, sparse_bwd_tflops diff --git a/csrc/attn/tests/test_block_sparse.py b/csrc/attn/tests/test_block_sparse.py deleted file mode 100644 index f6f85f99a..000000000 --- a/csrc/attn/tests/test_block_sparse.py +++ /dev/null @@ -1,289 +0,0 @@ -import torch -import argparse -from flash_attn.utils.benchmark import benchmark_forward -from flash_attn import flash_attn_func -from vsa import block_sparse_attn -from vsa import BLOCK_M, BLOCK_N - -import numpy as np -import random -import gc - -def set_seed(seed: int = 42): - # Python random module - random.seed(seed) - - # NumPy - np.random.seed(seed) - - # PyTorch - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) # if using multi-GPU - - -@torch.no_grad -def precision_metric(quant_o, fa2_o): - x, xx = quant_o.float(), fa2_o.float() - sim = torch.nn.functional.cosine_similarity(x.reshape(1, -1), xx.reshape(1, -1)).item() - l1 = ((x - xx).abs().sum() / xx.abs().sum() ).item() - rmse = torch.sqrt(torch.mean((x -xx) ** 2)).item() - - return sim, l1, rmse - -def create_input_tensors(batch, head, seq_len, headdim): - """Create random input tensors for attention.""" - q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - return q, k, v - -def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"): - """ - Generate a block sparse pattern where each q block attends to exactly k kv blocks. - - Args: - bs: batch size - h: number of heads - num_q_blocks: number of query blocks - num_kv_blocks: number of key-value blocks - k: number of kv blocks each q block attends to - device: device to create tensors on - - Returns: - q2k_block_sparse_index: [bs, h, num_q_blocks, k] - Contains the indices of kv blocks that each q block attends to. - q2k_block_sparse_num: [bs, h, num_q_blocks] - Contains the number of kv blocks that each q block attends to (all equal to k). - k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks] - Contains the indices of q blocks that attend to each kv block. - k2q_block_sparse_num: [bs, h, num_kv_blocks] - Contains the number of q blocks that attend to each kv block. - block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks] - Binary mask where 1 indicates attention connection. - """ - # Ensure k is not larger than num_kv_blocks - k = min(k, num_kv_blocks) - - # Create random scores for sampling - scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device) - - # Get top-k indices for each q block - _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1) - q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32) - - # sort q2k_block_sparse_index - q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1) - - # All q blocks attend to exactly k kv blocks - q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device) - - # Create the corresponding mask - block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device) - - # Fill in the mask based on the indices - for b in range(bs): - for head in range(h): - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx] - block_sparse_mask[b, head, q_idx, kv_indices] = True - - # Create the reverse mapping (k2q) - # First, initialize lists to collect q indices for each kv block - k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)] - - # Populate the lists based on q2k mapping - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist() - for kv_idx in kv_indices: - k2q_indices_list[flat_idx][kv_idx].append(q_idx) - - # Find the maximum number of q blocks that attend to any kv block - max_q_per_kv = 0 - for flat_idx in range(bs * h): - for kv_idx in range(num_kv_blocks): - max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx])) - - # Create tensors for k2q mapping - k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, - dtype=torch.int32, device=device) - k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), - dtype=torch.int32, device=device) - - # Fill the tensors - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for kv_idx in range(num_kv_blocks): - q_indices = k2q_indices_list[flat_idx][kv_idx] - num_q = len(q_indices) - k2q_block_sparse_num[b, head, kv_idx] = num_q - if num_q > 0: - k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor( - q_indices, dtype=torch.int32, device=device) - - return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask - -def main(args): - set_seed(42) - - # Extract parameters - batch = args.batch_size - head = args.num_heads - headdim = args.head_dim - num_iterations = args.num_iterations - - print(f"Block Sparse Attention Benchmark") - print(f"batch: {batch}, head: {head}, headdim: {headdim}, iterations: {num_iterations}") - - # Test with different sequence lengths - for seq_len in args.seq_lengths: - # Skip very long sequences if they might cause OOM - # if seq_len > 16384 and batch > 1: - # continue - - print("="*100) - print(f"\nSequence length: {seq_len}") - - # Collect metrics across iterations - forward_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_q_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_k_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_v_metrics = {'sim': [], 'l1': [], 'rmse': []} - - for iter_idx in range(num_iterations): - if num_iterations > 1: - print(f"\nIteration {iter_idx+1}/{num_iterations}") - - # Create input tensors - q, k, v = create_input_tensors(batch, head, seq_len, headdim) - - # Setup block sparse parameters - num_q_blocks = seq_len // BLOCK_M - num_kv_blocks = seq_len // BLOCK_N - - # Determine k value (number of kv blocks per q block) - topk = args.topk - if topk is None: - topk = num_kv_blocks // 10 # Default to ~90% sparsity if k is not specified - topk = max(1, topk) - if iter_idx == 0: # Only print this once - print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)") - - # Generate block sparse pattern - q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask = generate_block_sparse_pattern( - batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda") - - # expand block_sparse_mask to full mask - block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2) # [b, h, num_q_blocks, num_kv_blocks, 1, 1] - block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N) # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N] - full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len) - - q.requires_grad = True - k.requires_grad = True - v.requires_grad = True - - - # testing forward - o = block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) - del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded - grad_o = torch.randn_like(o) - o.backward(grad_o) - # clear memory - q_sdpa = q.detach().clone() - k_sdpa = k.detach().clone() - v_sdpa = v.detach().clone() - q_sdpa.requires_grad = True - k_sdpa.requires_grad = True - v_sdpa.requires_grad = True - q.data = torch.empty(0, device=q.device) - k.data = torch.empty(0, device=k.device) - v.data = torch.empty(0, device=v.device) - torch.cuda.empty_cache() - - o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask) - - - sim, l1, rmse = precision_metric(o, o_sdpa) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 8e-5, f"l1 too large: {l1}" - assert rmse < 2e-5, f"RMSE too large: {rmse}" - forward_metrics['sim'].append(sim) - forward_metrics['l1'].append(l1) - forward_metrics['rmse'].append(rmse) - - print(f"block_sparse_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - # test backward - o_sdpa.backward(grad_o) - - sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad) - # Error bounds collected on H100 - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 4e-3, f"l1 too large: {l1}" - assert rmse < 3e-4, f"RMSE too large: {rmse}" - grad_q_metrics['sim'].append(sim) - grad_q_metrics['l1'].append(l1) - grad_q_metrics['rmse'].append(rmse) - print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 4e-3, f"l1 too large: {l1}" - assert rmse < 2e-4, f"RMSE too large: {rmse}" - grad_k_metrics['sim'].append(sim) - grad_k_metrics['l1'].append(l1) - grad_k_metrics['rmse'].append(rmse) - print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 1e-4, f"l1 too large: {l1}" - assert rmse < 2e-5, f"RMSE too large: {rmse}" - grad_v_metrics['sim'].append(sim) - grad_v_metrics['l1'].append(l1) - grad_v_metrics['rmse'].append(rmse) - print(f"block_sparse_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa - gc.collect() - torch.cuda.empty_cache() - - # Print summary statistics if multiple iterations were run - if num_iterations > 1: - print("\n" + "="*50) - print(f"Summary Statistics (over {num_iterations} iterations):") - - print("\nForward metrics:") - print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}") - - print("\nGradient Q metrics:") - print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}") - - print("\nGradient K metrics:") - print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}") - - print("\nGradient V metrics:") - print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}") - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') - parser.add_argument('--batch_size', type=int, default=4, help='Batch size') - parser.add_argument('--num_heads', type=int, default=6, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') - parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to') - parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark') - parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run') - args = parser.parse_args() - main(args) \ No newline at end of file diff --git a/csrc/attn/tests/test_block_sparse_mask_to_index.py b/csrc/attn/tests/test_block_sparse_mask_to_index.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/csrc/attn/tests/test_block_sparse_triton.py b/csrc/attn/tests/test_block_sparse_triton.py deleted file mode 100644 index 256aa9954..000000000 --- a/csrc/attn/tests/test_block_sparse_triton.py +++ /dev/null @@ -1,289 +0,0 @@ -import torch -import argparse -from flash_attn.utils.benchmark import benchmark_forward -from flash_attn import flash_attn_func -from vsa import triton_attention_sparse -from vsa import BLOCK_M, BLOCK_N - -import numpy as np -import random -import gc - -def set_seed(seed: int = 42): - # Python random module - random.seed(seed) - - # NumPy - np.random.seed(seed) - - # PyTorch - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) # if using multi-GPU - - -@torch.no_grad -def precision_metric(quant_o, fa2_o): - x, xx = quant_o.float(), fa2_o.float() - sim = torch.nn.functional.cosine_similarity(x.reshape(1, -1), xx.reshape(1, -1)).item() - l1 = ((x - xx).abs().sum() / xx.abs().sum() ).item() - rmse = torch.sqrt(torch.mean((x -xx) ** 2)).item() - - return sim, l1, rmse - -def create_input_tensors(batch, head, seq_len, headdim): - """Create random input tensors for attention.""" - q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - return q, k, v - -def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"): - """ - Generate a block sparse pattern where each q block attends to exactly k kv blocks. - - Args: - bs: batch size - h: number of heads - num_q_blocks: number of query blocks - num_kv_blocks: number of key-value blocks - k: number of kv blocks each q block attends to - device: device to create tensors on - - Returns: - q2k_block_sparse_index: [bs, h, num_q_blocks, k] - Contains the indices of kv blocks that each q block attends to. - q2k_block_sparse_num: [bs, h, num_q_blocks] - Contains the number of kv blocks that each q block attends to (all equal to k). - k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks] - Contains the indices of q blocks that attend to each kv block. - k2q_block_sparse_num: [bs, h, num_kv_blocks] - Contains the number of q blocks that attend to each kv block. - block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks] - Binary mask where 1 indicates attention connection. - """ - # Ensure k is not larger than num_kv_blocks - k = min(k, num_kv_blocks) - - # Create random scores for sampling - scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device) - - # Get top-k indices for each q block - _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1) - q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32) - - # sort q2k_block_sparse_index - q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1) - - # All q blocks attend to exactly k kv blocks - q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device) - - # Create the corresponding mask - block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device) - - # Fill in the mask based on the indices - for b in range(bs): - for head in range(h): - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx] - block_sparse_mask[b, head, q_idx, kv_indices] = True - - # Create the reverse mapping (k2q) - # First, initialize lists to collect q indices for each kv block - k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)] - - # Populate the lists based on q2k mapping - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist() - for kv_idx in kv_indices: - k2q_indices_list[flat_idx][kv_idx].append(q_idx) - - # Find the maximum number of q blocks that attend to any kv block - max_q_per_kv = 0 - for flat_idx in range(bs * h): - for kv_idx in range(num_kv_blocks): - max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx])) - - # Create tensors for k2q mapping - k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, - dtype=torch.int32, device=device) - k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), - dtype=torch.int32, device=device) - - # Fill the tensors - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for kv_idx in range(num_kv_blocks): - q_indices = k2q_indices_list[flat_idx][kv_idx] - num_q = len(q_indices) - k2q_block_sparse_num[b, head, kv_idx] = num_q - if num_q > 0: - k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor( - q_indices, dtype=torch.int32, device=device) - - return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask - -def main(args): - set_seed(42) - - # Extract parameters - batch = args.batch_size - head = args.num_heads - headdim = args.head_dim - num_iterations = args.num_iterations - - print(f"Block Sparse Attention Benchmark") - print(f"batch: {batch}, head: {head}, headdim: {headdim}, iterations: {num_iterations}") - - # Test with different sequence lengths - for seq_len in args.seq_lengths: - # Skip very long sequences if they might cause OOM - # if seq_len > 16384 and batch > 1: - # continue - - print("="*100) - print(f"\nSequence length: {seq_len}") - - # Collect metrics across iterations - forward_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_q_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_k_metrics = {'sim': [], 'l1': [], 'rmse': []} - grad_v_metrics = {'sim': [], 'l1': [], 'rmse': []} - - for iter_idx in range(num_iterations): - if num_iterations > 1: - print(f"\nIteration {iter_idx+1}/{num_iterations}") - - # Create input tensors - q, k, v = create_input_tensors(batch, head, seq_len, headdim) - - # Setup block sparse parameters - num_q_blocks = seq_len // BLOCK_M - num_kv_blocks = seq_len // BLOCK_N - - # Determine k value (number of kv blocks per q block) - topk = args.topk - if topk is None: - topk = num_kv_blocks // 10 # Default to ~90% sparsity if k is not specified - topk = max(1, topk) - if iter_idx == 0: # Only print this once - print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)") - - # Generate block sparse pattern - q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask = generate_block_sparse_pattern( - batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda") - - # expand block_sparse_mask to full mask - block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2) # [b, h, num_q_blocks, num_kv_blocks, 1, 1] - block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N) # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N] - full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len) - - q.requires_grad = True - k.requires_grad = True - v.requires_grad = True - - - # testing forward - o = triton_attention_sparse(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) - del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded - grad_o = torch.randn_like(o) - o.backward(grad_o) - # clear memory - q_sdpa = q.detach().clone() - k_sdpa = k.detach().clone() - v_sdpa = v.detach().clone() - q_sdpa.requires_grad = True - k_sdpa.requires_grad = True - v_sdpa.requires_grad = True - q.data = torch.empty(0, device=q.device) - k.data = torch.empty(0, device=k.device) - v.data = torch.empty(0, device=v.device) - torch.cuda.empty_cache() - - o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask) - - - sim, l1, rmse = precision_metric(o, o_sdpa) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 8e-5, f"l1 too large: {l1}" - assert rmse < 5e-5, f"RMSE too large: {rmse}" - forward_metrics['sim'].append(sim) - forward_metrics['l1'].append(l1) - forward_metrics['rmse'].append(rmse) - - print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - # test backward - o_sdpa.backward(grad_o) - - sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad) - # Error bounds collected on H100 - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 4e-3, f"l1 too large: {l1}" - assert rmse < 5e-4, f"RMSE too large: {rmse}" - grad_q_metrics['sim'].append(sim) - grad_q_metrics['l1'].append(l1) - grad_q_metrics['rmse'].append(rmse) - print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 4e-3, f"l1 too large: {l1}" - assert rmse < 5e-4, f"RMSE too large: {rmse}" - grad_k_metrics['sim'].append(sim) - grad_k_metrics['l1'].append(l1) - grad_k_metrics['rmse'].append(rmse) - print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad) - assert sim > 0.9999, f"SSIM too low: {sim}" - assert l1 < 4e-3, f"l1 too large: {l1}" - assert rmse < 5e-4, f"RMSE too large: {rmse}" - grad_v_metrics['sim'].append(sim) - grad_v_metrics['l1'].append(l1) - grad_v_metrics['rmse'].append(rmse) - print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - - del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa - gc.collect() - torch.cuda.empty_cache() - - # Print summary statistics if multiple iterations were run - if num_iterations > 1: - print("\n" + "="*50) - print(f"Summary Statistics (over {num_iterations} iterations):") - - print("\nForward metrics:") - print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}") - - print("\nGradient Q metrics:") - print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}") - - print("\nGradient K metrics:") - print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}") - - print("\nGradient V metrics:") - print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}") - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') - parser.add_argument('--batch_size', type=int, default=4, help='Batch size') - parser.add_argument('--num_heads', type=int, default=6, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') - parser.add_argument('--topk', type=int, default=4, help='Number of kv blocks each q block attends to') - parser.add_argument('--seq_lengths', type=int, nargs='+', default=[4096], help='Sequence lengths to benchmark') - parser.add_argument('--num_iterations', type=int, default=10, help='Number of test iterations to run') - args = parser.parse_args() - main(args) \ No newline at end of file diff --git a/csrc/attn/tests/test_full_attn.py b/csrc/attn/tests/test_full_attn.py deleted file mode 100644 index 6dd0b7e3a..000000000 --- a/csrc/attn/tests/test_full_attn.py +++ /dev/null @@ -1,175 +0,0 @@ -import torch -from flash_attn_interface import flash_attn_func -from st_attn import mha_forward, mha_backward -import random -from tqdm import tqdm -import matplotlib.pyplot as plt -import numpy as np - -def pytorch_test(Q, K, V, dO): - q_ = Q.to(torch.float64).requires_grad_() - k_ = K.to(torch.float64).requires_grad_() - v_ = V.to(torch.float64).requires_grad_() - dO_ = dO.to(torch.float64) - - # manual pytorch implementation of scaled dot product attention - QK = torch.matmul(q_, k_.transpose(-2, -1)) - QK /= (q_.size(-1) ** 0.5) - - # Causal mask removed since causal is always false - - QK = torch.nn.functional.softmax(QK, dim=-1) - output = torch.matmul(QK, v_) - - output.backward(dO_) - - q_grad = q_.grad - k_grad = k_.grad - v_grad = v_.grad - - return output, q_grad, k_grad, v_grad - -def fa2_test(Q, K, V, dO): - Q.requires_grad = True - K.requires_grad = True - V.requires_grad = True - output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False) - output.backward(dO) - - return output, Q.grad, K.grad, V.grad - - -def mha_kernel_test(Q, K, V, dO, mode): - Q.requires_grad = True - K.requires_grad = True - V.requires_grad = True - - o, l_vec = mha_forward(Q, K, V) - - if mode == 'forward_only': - return o, None, None, None - else: # 'forward_backward' - qg, kg, vg = mha_backward(Q, K, V, o, l_vec, dO) - return o, qg, kg, vg - -def generate_tensor(shape, mean, std, dtype, device): - tensor = torch.randn(shape, dtype=dtype, device=device) - - magnitude = torch.norm(tensor, dim=-1, keepdim=True) - scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude - - return scaled_tensor.contiguous() - -def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'): - results = { - 'MHA vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0}, - 'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0}, - } - - for _ in range(num_iterations): - torch.manual_seed(0) - - Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - - pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO) - fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO) - - if test_mode == 'forward_only': - mha_o, _, _, _ = mha_kernel_test(Q, K, V, dO, 'forward_only') - tensors_mha_pt = [(pt_o, mha_o)] - tensors_fa2_pt = [(pt_o, fa2_o)] - else: # 'forward_backward' - mha_o, mha_qg, mha_kg, mha_vg = mha_kernel_test(Q, K, V, dO, 'forward_backward') - - if error_mode == 'output': - tensors_mha_pt = [(pt_o, mha_o)] - tensors_fa2_pt = [(pt_o, fa2_o)] - elif error_mode == 'backward': - tensors_mha_pt = [(pt_qg, mha_qg), - (pt_kg, mha_kg), - (pt_vg, mha_vg)] - tensors_fa2_pt = [(pt_qg, fa2_qg), - (pt_kg, fa2_kg), - (pt_vg, fa2_vg)] - else: # 'all' - tensors_mha_pt = [(pt_o, mha_o), - (pt_qg, mha_qg), - (pt_kg, mha_kg), - (pt_vg, mha_vg)] - tensors_fa2_pt = [(pt_o, fa2_o), - (pt_qg, fa2_qg), - (pt_kg, fa2_kg), - (pt_vg, fa2_vg)] - - for pt, mha in tensors_mha_pt: - diff = pt - mha - abs_diff = torch.abs(diff) - results['MHA vs PT']['sum_diff'] += torch.sum(abs_diff).item() - results['MHA vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item() - results['MHA vs PT']['max_diff'] = max(results['MHA vs PT']['max_diff'], torch.max(abs_diff).item()) - - for pt, fa2 in tensors_fa2_pt: - diff = pt - fa2 - abs_diff = torch.abs(diff) - results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item() - results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item() - results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item()) - - torch.cuda.empty_cache() - - # Calculate total elements based on test mode and error mode - if test_mode == 'forward_only': - total_elements = b * h * n * d * num_iterations - else: # 'forward_backward' - total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4) - - for name, data in results.items(): - avg_diff = data['sum_diff'] / total_elements - max_diff = data['max_diff'] - results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} - - return results - -def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'): - seq_lengths = [768 * (2**i) for i in range(1)] - - print(f"\n{'='*80}") - print(f"MHA ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})") - print(f"Mode: {error_mode}, Test: {test_mode}") - print(f"{'='*80}") - - # Print header - print(f"{'Seq Length':<12} | {'MHA vs PT Avg':<15} | {'MHA vs PT Max':<15} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}") - print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}") - - for n in seq_lengths: - results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode) - - mha_pt_avg = results['MHA vs PT']['avg_diff'] - mha_pt_max = results['MHA vs PT']['max_diff'] - fa2_pt_avg = results['FA2 vs PT']['avg_diff'] - fa2_pt_max = results['FA2 vs PT']['max_diff'] - - # Print row - print(f"{n:<12} | {mha_pt_avg:<15.6e} | {mha_pt_max:<15.6e} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}") - - print(f"{'='*80}\n") - -# fix random seed -torch.manual_seed(0) - -# Example usage -b, h, d = 2, 2, 64 -mean = 1e-1 -std = 10 - -# Test forward only -generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only') - -# Test forward and backward -generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward') - -print("MHA attention error comparison completed.") \ No newline at end of file diff --git a/csrc/attn/tests/test_triton.py b/csrc/attn/tests/test_triton.py deleted file mode 100644 index 27e1b0558..000000000 --- a/csrc/attn/tests/test_triton.py +++ /dev/null @@ -1,183 +0,0 @@ -import torch -from tqdm import tqdm -import matplotlib.pyplot as plt -import numpy as np - -from vsa import triton_attention - -def pytorch_test(Q, K, V, dO): - q_ = Q.to(torch.float64).requires_grad_() - k_ = K.to(torch.float64).requires_grad_() - v_ = V.to(torch.float64).requires_grad_() - q_.grad = None - k_.grad = None - v_.grad = None - dO_ = dO.to(torch.float64) - - # manual pytorch implementation of scaled dot product attention - QK = torch.matmul(q_, k_.transpose(-2, -1)) - QK /= (q_.size(-1) ** 0.5) - - # Causal mask removed since causal is always false - - QK = torch.nn.functional.softmax(QK, dim=-1) - output = torch.matmul(QK, v_) - - output.backward(dO_) - - q_grad = q_.grad - k_grad = k_.grad - v_grad = v_.grad - - return output, q_grad, k_grad, v_grad - -def fa2_test(Q, K, V, dO): - Q.requires_grad = True - K.requires_grad = True - V.requires_grad = True - Q.grad = None - K.grad = None - V.grad = None - - output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False) - output.backward(dO) - - return output, Q.grad, K.grad, V.grad - -def triton_test(Q, K, V, dO): - Q.requires_grad = True - K.requires_grad = True - V.requires_grad = True - Q.grad = None - K.grad = None - V.grad = None - - output = triton_attention(Q, K, V) - - output.backward(dO) - - q_grad = Q.grad - k_grad = K.grad - v_grad = V.grad - - return output.to(Q.dtype) if output is not None else None, q_grad, k_grad, v_grad - -def generate_tensor(shape, mean, std, dtype, device): - tensor = torch.randn(shape, dtype=dtype, device=device) - - magnitude = torch.norm(tensor, dim=-1, keepdim=True) - scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude - - return scaled_tensor.contiguous() - -def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'): - results = { - 'FA2 vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'Triton vs PT': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - } - - for _ in range(num_iterations): - torch.manual_seed(0) - - Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') - - pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO) - fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO) - triton_o, triton_qg, triton_kg, triton_vg = triton_test(Q, K, V, dO) - - if test_mode == 'forward_only': - tensors_fa2_pt = [(pt_o, fa2_o)] - tensors_triton_pt = [(pt_o, triton_o)] - else: # 'forward_backward' - if error_mode == 'output': - tensors_fa2_pt = [(pt_o, fa2_o)] - tensors_triton_pt = [(pt_o, triton_o)] - elif error_mode == 'backward': - tensors_fa2_pt = [(pt_qg, fa2_qg), - (pt_kg, fa2_kg), - (pt_vg, fa2_vg)] - tensors_triton_pt = [(pt_qg, triton_qg), - (pt_kg, triton_kg), - (pt_vg, triton_vg)] - else: # 'all' - tensors_fa2_pt = [(pt_o, fa2_o), - (pt_qg, fa2_qg), - (pt_kg, fa2_kg), - (pt_vg, fa2_vg)] - tensors_triton_pt = [(pt_o, triton_o), - (pt_qg, triton_qg), - (pt_kg, triton_kg), - (pt_vg, triton_vg)] - - for pt, fa2 in tensors_fa2_pt: - diff = pt - fa2 - abs_diff = torch.abs(diff) - results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item() - results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item() - results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item()) - - for pt, triton in tensors_triton_pt: - diff = pt - triton - abs_diff = torch.abs(diff) - results['Triton vs PT']['sum_diff'] += torch.sum(abs_diff).item() - results['Triton vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item() - results['Triton vs PT']['max_diff'] = max(results['Triton vs PT']['max_diff'], torch.max(abs_diff).item()) - - torch.cuda.empty_cache() - - # Calculate total elements based on test mode and error mode - if test_mode == 'forward_only': - total_elements = b * h * n * d * num_iterations - else: # 'forward_backward' - total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4) - - for name, data in results.items(): - avg_diff = data['sum_diff'] / total_elements - max_diff = data['max_diff'] - results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} - - return results - -def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'): - seq_lengths = [768 * (2**i) for i in range(1)] - - print(f"\n{'='*100}") - print(f"ATTENTION ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})") - print(f"Mode: {error_mode}, Test: {test_mode}") - print(f"{'='*100}") - - # Print header - print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15} | {'Triton vs PT Avg':<15} | {'Triton vs PT Max':<15}") - print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}") - - for n in seq_lengths: - results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode) - - fa2_pt_avg = results['FA2 vs PT']['avg_diff'] - fa2_pt_max = results['FA2 vs PT']['max_diff'] - triton_pt_avg = results['Triton vs PT']['avg_diff'] - triton_pt_max = results['Triton vs PT']['max_diff'] - - # Print row with both comparisons - print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e} | {triton_pt_avg:<15.6e} | {triton_pt_max:<15.6e}") - - print(f"{'='*100}\n") - -# fix random seed -torch.manual_seed(0) - -mean = 1e-1 -std = 10 -configs = [ - (4, 1, 128), # Larger batch, single head, larger dim - (2, 8, 64), # Medium batch, many heads, medium dim -] - -for b, h, d in configs: - print(f"\nConfiguration: batch={b}, heads={h}, dim={d}") - generate_error_tables(b, h, d, mean, std, error_mode='backward', test_mode='forward_backward') - -print("Attention error comparison completed.") \ No newline at end of file diff --git a/csrc/attn/tests/test_vsa.py b/csrc/attn/tests/test_vsa.py new file mode 100644 index 000000000..20a5d7e90 --- /dev/null +++ b/csrc/attn/tests/test_vsa.py @@ -0,0 +1,159 @@ +import torch +import sys +import os +import numpy as np +from tqdm import tqdm + +# Add the parent directory to the path to import block_sparse_attn +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from tests.utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask +from vsa import block_sparse_attn + +BLOCK_M = 64 +BLOCK_N = 64 + +def pytorch_test(Q, K, V, block_sparse_mask, dO): + q_ = Q.clone().requires_grad_() + k_ = K.clone().requires_grad_() + v_ = V.clone().requires_grad_() + + QK = torch.matmul(q_, k_.transpose(-2, -1)) + QK /= (q_.size(-1) ** 0.5) + QK = QK.masked_fill(~block_sparse_mask.unsqueeze(0), float('-inf')) + + QK = torch.nn.functional.softmax(QK, dim=-1) + output = torch.matmul(QK, v_) + + dO_ = dO + output.backward(dO_) + return ( + output.to(torch.bfloat16), + q_.grad.to(torch.bfloat16), + k_.grad.to(torch.bfloat16), + v_.grad.to(torch.bfloat16), + ) + + +def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, non_pad_index, dO): + Q = Q.clone().requires_grad_() + K = K.clone().requires_grad_() + V = V.clone().requires_grad_() + + q_padded = vsa_pad(Q, non_pad_index, variable_block_sizes.shape[0], BLOCK_M) + k_padded = vsa_pad(K, non_pad_index, variable_block_sizes.shape[0], BLOCK_M) + v_padded = vsa_pad(V, non_pad_index, variable_block_sizes.shape[0], BLOCK_M) + output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes) + output = output[:, :, non_pad_index, :] + output.backward(dO) + return output, Q.grad, K.grad, V.grad + + +def get_non_pad_index( + vid_len: torch.LongTensor, + n_win: int, + win_size: int, +): + device = vid_len.device + starts_pad = torch.arange(n_win, device=device) * win_size + index_pad = starts_pad[:, None] + torch.arange(win_size, device=device)[None, :] + index_mask = torch.arange(win_size, device=device)[None, :] < vid_len[:, None] + + return index_pad[index_mask] + +def generate_tensor(shape, mean, std, dtype, device): + tensor = torch.randn(shape, dtype=dtype, device=device) + magnitude = torch.norm(tensor, dim=-1, keepdim=True) + scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude + return scaled_tensor.contiguous() + +def generate_variable_block_sizes(num_blocks, min_size=32, max_size=64, device="cuda"): + return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32) + + +def vsa_pad(x, non_pad_index, num_blocks, block_size): + padded_x = torch.zeros((1, x.shape[1], num_blocks * BLOCK_M, x.shape[3]), device=x.device, dtype=x.dtype) + padded_x[:, :, non_pad_index, :] = x + return padded_x + +def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_mode='all'): + results = { + 'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + 'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + 'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + 'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, + } + + device = "cuda" if torch.cuda.is_available() else "cpu" + variable_block_sizes = generate_variable_block_sizes(num_blocks, device=device) + S = int(variable_block_sizes.sum().item()) + padded_S = num_blocks * BLOCK_M + non_pad_index = get_non_pad_index(variable_block_sizes, num_blocks, BLOCK_M) + block_mask = generate_block_sparse_mask_for_function(h, num_blocks, k, device) + full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, device) + for _ in range(num_iterations): + Q = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device) + K = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device) + V = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device) + dO = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device) + + # dO_padded = torch.zeros_like(dO_padded) + # dO_padded[:, :, non_pad_index, :] = dO + + pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) + bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes,non_pad_index, dO) + for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]): + if bs is not None: + diff = pt - bs + abs_diff = torch.abs(diff) + results[name]['sum_diff'] += torch.sum(abs_diff).item() + results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item() + results[name]['max_diff'] = max(results[name]['max_diff'], torch.max(abs_diff).item()) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + total_elements = h * S * d * num_iterations + for name, data in results.items(): + avg_diff = data['sum_diff'] / total_elements + max_diff = data['max_diff'] + results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} + + return results + +def generate_error_graphs(h, d, mean, std, error_mode='all'): + test_configs = [ + {"num_blocks": 16, "k": 2, "description": "Small sequence"}, + {"num_blocks": 32, "k": 4, "description": "Medium sequence"}, + {"num_blocks": 53, "k": 6, "description": "Large sequence"}, + ] + + print(f"\nError Analysis for h={h}, d={d}, mean={mean}, std={std}, mode={error_mode}") + print("=" * 150) + print(f"{'Config':<20} {'Blocks':<8} {'K':<4} " + f"{'gQ Avg':<12} {'gQ Max':<12} " + f"{'gK Avg':<12} {'gK Max':<12} " + f"{'gV Avg':<12} {'gV Max':<12} " + f"{'gO Avg':<12} {'gO Max':<12}") + print("-" * 150) + + for config in test_configs: + num_blocks = config["num_blocks"] + k = config["k"] + description = config["description"] + results = check_correctness(h, d, num_blocks, k, mean, std, error_mode=error_mode) + print(f"{description:<20} {num_blocks:<8} {k:<4} " + f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} " + f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} " + f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} " + f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}") + + print("-" * 150) + +if __name__ == "__main__": + h, d = 16, 128 + mean = 0.0 + std = 1 + print("Block Sparse Attention with Variable Block Sizes Analysis") + print("=" * 60) + for mode in ['backward']: + generate_error_graphs(h, d, mean, std, error_mode=mode) + print("\nAnalysis completed for all modes.") diff --git a/csrc/attn/tests/utils.py b/csrc/attn/tests/utils.py new file mode 100644 index 000000000..30e53deb0 --- /dev/null +++ b/csrc/attn/tests/utils.py @@ -0,0 +1,54 @@ +import torch + +def generate_block_sparse_mask_for_function(h, num_blocks, k, device="cuda"): + """ + Generate block sparse mask of shape [h, num_blocks, num_blocks]. + + Args: + h: number of heads + num_blocks: number of blocks + k: number of kv blocks each q block attends to + device: device to create tensors on + + Returns: + block_sparse_mask: [h, num_blocks, num_blocks] bool tensor + """ + k = min(k, num_blocks) + scores = torch.rand(h, num_blocks, num_blocks, device=device) + _, indices = torch.topk(scores, k, dim=-1) + block_sparse_mask = torch.zeros(h, num_blocks, num_blocks, dtype=torch.bool, device=device) + + block_sparse_mask = block_sparse_mask.scatter_(2, indices, 1).bool() + return block_sparse_mask + + +def create_full_mask_from_block_mask(block_sparse_mask, variable_block_sizes, device="cuda"): + """ + Convert block-level sparse mask to full attention mask. + + Args: + block_sparse_mask: [h, num_blocks, num_blocks] bool tensor + variable_block_sizes: [num_blocks] tensor + device: device to create tensors on + + Returns: + full_mask: [h, S, S] bool tensor where S = total sequence length + """ + h, num_blocks, _ = block_sparse_mask.shape + total_seq_len = variable_block_sizes.sum().item() + cumsum = torch.cat([torch.tensor([0], device=device), variable_block_sizes.cumsum(dim=0)[:-1]]) + + full_mask = torch.zeros(h, total_seq_len, total_seq_len, dtype=torch.bool, device=device) + + for head in range(h): + for q_block in range(num_blocks): + q_start = cumsum[q_block] + q_end = q_start + variable_block_sizes[q_block] + + for kv_block in range(num_blocks): + if block_sparse_mask[head, q_block, kv_block]: + kv_start = cumsum[kv_block] + kv_end = kv_start + variable_block_sizes[kv_block] + full_mask[head, q_start:q_end, kv_start:kv_end] = True + + return full_mask \ No newline at end of file diff --git a/csrc/attn/vsa.cpp b/csrc/attn/vsa.cpp index 149806167..8661103a0 100644 --- a/csrc/attn/vsa.cpp +++ b/csrc/attn/vsa.cpp @@ -9,10 +9,10 @@ #ifdef TK_COMPILE_BLOCK_SPARSE extern std::vector block_sparse_attention_forward( - torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size ); extern std::vector block_sparse_attention_backward( - torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size ); #endif diff --git a/csrc/attn/vsa/__init__.py b/csrc/attn/vsa/__init__.py index 789330c0c..2faa9e18d 100644 --- a/csrc/attn/vsa/__init__.py +++ b/csrc/attn/vsa/__init__.py @@ -1,12 +1,17 @@ import torch from typing import Tuple -from vsa.vsa import block_sparse_attn +block_sparse_attn=None + try: from vsa_cuda import block_sparse_fwd, block_sparse_bwd + from vsa.block_sparse_wrapper import block_sparse_attn_SM90 + block_sparse_attn = block_sparse_attn_SM90 except ImportError: + from vsa.block_sparse_wrapper import block_sparse_attn_triton block_sparse_fwd = None block_sparse_bwd = None -from vsa.block_sparse_attn_triton import attention as triton_attention, attention_sparse as triton_attention_sparse + block_sparse_attn = block_sparse_attn_triton + BLOCK_M = 64 BLOCK_N = 64 @@ -22,7 +27,7 @@ def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]: return output, QK -def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None): +def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight=None): """ q: [batch_size, num_heads, seq_len, head_dim] k: [batch_size, num_heads, seq_len, head_dim] @@ -32,7 +37,7 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None): video_shape: tuple of (T, H, W) compress_attn_weight: [batch_size, num_heads, seq_len, head_dim] select_attn_weight: [batch_size, num_heads, seq_len, head_dim] - + NOTE: We assume q, k, v is zero padded!! V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]] """ @@ -41,16 +46,16 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None): block_size = (block_size, block_size, block_size) block_elements = block_size[0] * block_size[1] * block_size[2] - assert block_elements % 64 == 0 and block_elements >= 64 + assert block_elements == 64 assert q.shape[2] % block_elements == 0 batch_size, num_heads, seq_len, head_dim = q.shape # compress attn - q_compress = q.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).mean(dim=3) - k_compress = k.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).mean(dim=3) - v_compress = v.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).mean(dim=3) + q_compress = (q.view(batch_size, num_heads, seq_len // block_elements, + block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype) + k_compress = (k.view(batch_size, num_heads, seq_len // block_elements, + block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype) + v_compress = (v.view(batch_size, num_heads, seq_len // block_elements, + block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype) output_compress, block_attn_score = torch_attention(q_compress, k_compress, v_compress) @@ -62,13 +67,9 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None): 1).view(batch_size, num_heads, seq_len, head_dim) - q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern( - block_attn_score, topk) - - output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index, - q2k_block_sparse_num, - k2q_block_sparse_index, - k2q_block_sparse_num) + topK_indices = torch.topk(block_attn_score, topk, dim=-1).indices + block_mask = torch.zeros_like(block_attn_score, dtype=torch.bool).scatter_(-1, topK_indices, True) + output_select, _ = block_sparse_attn(q, k, v, block_mask, variable_block_sizes) if compress_attn_weight is not None: final_output = output_compress * compress_attn_weight + output_select @@ -76,203 +77,3 @@ def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None): final_output = output_compress + output_select return final_output - -def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor, - topk: int): - """ - Generate a block sparse pattern where each q block attends to exactly topk kv blocks, - based on the provided attention scores. - - Args: - block_attn_score: [bs, h, num_q_blocks, num_kv_blocks] - Attention scores between query and key blocks - topk: int - Number of kv blocks each q block attends to - - Returns: - q2k_block_sparse_index: [bs, h, num_q_blocks, topk] - Contains the indices of kv blocks that each q block attends to. - q2k_block_sparse_num: [bs, h, num_q_blocks] - Contains the number of kv blocks that each q block attends to (all equal to topk). - k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv] - Contains the indices of q blocks that attend to each kv block. - k2q_block_sparse_num: [bs, h, num_kv_blocks] - Contains the number of q blocks that attend to each kv block. - """ - device = block_attn_score.device - # Extract dimensions from block_attn_score - bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape - - sorted_result = torch.sort(block_attn_score, dim=-1, descending=True) - - sorted_indice = sorted_result.indices - - q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk], - dim=-1) - q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32) - q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), - topk, - device=device, - dtype=torch.int32) - - block_map = topk_index_to_map(q2k_block_sparse_index, - num_kv_blocks, - transpose_map=True) - k2q_block_sparse_index, k2q_block_sparse_num = map_to_index( - block_map.transpose(2, 3)) - - return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num - - -## pytorch sdpa version of block sparse ## -import triton -import triton.language as tl - - -@triton.jit -def topk_index_to_map_kernel( - map_ptr, - index_ptr, - map_bs_stride, - map_h_stride, - map_q_stride, - map_kv_stride, - index_bs_stride, - index_h_stride, - index_q_stride, - index_kv_stride, - topk: tl.constexpr, -): - b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) - index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride - map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride - - for i in tl.static_range(topk): - index = tl.load(index_ptr_base + i * index_kv_stride) - tl.store(map_ptr_base + index * map_kv_stride, 1.0) - -@triton.jit -def map_to_index_kernel( - map_ptr, - index_ptr, - index_num_ptr, - map_bs_stride, - map_h_stride, - map_q_stride, - map_kv_stride, - index_bs_stride, - index_h_stride, - index_q_stride, - index_kv_stride, - index_num_bs_stride, - index_num_h_stride, - index_num_q_stride, - num_kv_blocks: tl.constexpr, -): - b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) - index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride - map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride - - num = 0 - for i in tl.static_range(num_kv_blocks): - map_entry = tl.load(map_ptr_base + i * map_kv_stride) - if map_entry: - tl.store(index_ptr_base + num * index_kv_stride, i) - num += 1 - - tl.store( - index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + - q * index_num_q_stride, num) - -def topk_index_to_map(index: torch.Tensor, - num_kv_blocks: int, - transpose_map: bool = False): - """ - Convert topk indices to a map. - - Args: - index: [bs, h, num_q_blocks, topk] - The topk indices tensor. - num_kv_blocks: int - The number of key-value blocks in the block_map returned - transpose_map: bool - If True, the block_map will be transposed on the final two dimensions. - - Returns: - block_map: [bs, h, num_q_blocks, num_kv_blocks] - A binary map where 1 indicates that the q block attends to the kv block. - """ - bs, h, num_q_blocks, topk = index.shape - - if transpose_map is False: - block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), - dtype=torch.bool, - device=index.device) - else: - block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), - dtype=torch.bool, - device=index.device) - block_map = block_map.transpose(2, 3) - - grid = (bs, h, num_q_blocks) - topk_index_to_map_kernel[grid]( - block_map, - index, - block_map.stride(0), - block_map.stride(1), - block_map.stride(2), - block_map.stride(3), - index.stride(0), - index.stride(1), - index.stride(2), - index.stride(3), - topk=topk, - ) - - return block_map - -def map_to_index(block_map: torch.Tensor): - """ - Convert a block map to indices and counts. - - Args: - block_map: [bs, h, num_q_blocks, num_kv_blocks] - The block map tensor. - - Returns: - index: [bs, h, num_q_blocks, num_kv_blocks] - The indices of the blocks. - index_num: [bs, h, num_q_blocks] - The number of blocks for each q block. - """ - bs, h, num_q_blocks, num_kv_blocks = block_map.shape - - index = torch.full((block_map.shape), - -1, - dtype=torch.int32, - device=block_map.device) - index_num = torch.empty((bs, h, num_q_blocks), - dtype=torch.int32, - device=block_map.device) - - grid = (bs, h, num_q_blocks) - map_to_index_kernel[grid]( - block_map, - index, - index_num, - block_map.stride(0), - block_map.stride(1), - block_map.stride(2), - block_map.stride(3), - index.stride(0), - index.stride(1), - index.stride(2), - index.stride(3), - index_num.stride(0), - index_num.stride(1), - index_num.stride(2), - num_kv_blocks=num_kv_blocks, - ) - - return index, index_num - diff --git a/csrc/attn/vsa/block_sparse_attn_triton.py b/csrc/attn/vsa/block_sparse_attn_triton.py index 1623a788e..db771dae3 100644 --- a/csrc/attn/vsa/block_sparse_attn_triton.py +++ b/csrc/attn/vsa/block_sparse_attn_triton.py @@ -18,43 +18,6 @@ import math # small utility needed by the sparse wrapper # ──────────────────────────── SPARSE ADDITION END ───────────────────────────── -@triton.jit -def _attn_fwd_inner(acc, l_i, m_i, q, # - K_block_ptr, V_block_ptr, # - start_m, qk_scale, # - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # - STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, # - N_CTX: tl.constexpr, fp8_v: tl.constexpr): - - # loop over k, v and update accumulator - for start_n in range(0, N_CTX, BLOCK_N): - # -- compute qk ---- - k = tl.load(K_block_ptr) - qk = tl.dot(q, k) - - m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) - qk = qk * qk_scale - m_ij[:, None] - p = tl.math.exp2(qk) - l_ij = tl.sum(p, 1) - # -- update m_i and l_i - alpha = tl.math.exp2(m_i - m_ij) - l_i = l_i * alpha + l_ij - # -- update output accumulator -- - acc = acc * alpha[:, None] - # update acc - v = tl.load(V_block_ptr) - if fp8_v: - p = p.to(tl.float8e5) - else: - p = p.to(tl.bfloat16) - acc = tl.dot(p, v, acc) - # update m_i and l_i - m_i = m_ij - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) - K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N)) - return acc, l_i, m_i - - # We don't run auto-tuning every time to keep the tutorial fast. Keeping # the code below and commenting out the equivalent parameters is convenient for # re-tuning. @@ -71,6 +34,7 @@ configs = [ @triton.jit def _attn_fwd_sparse(Q, K, V, sm_scale, # q2k_index, q2k_num, max_kv_blks, # + variable_block_sizes, M, Out, # stride_qz, stride_qh, stride_qm, stride_qk, stride_kz, stride_kh, stride_kn, stride_kk, @@ -136,12 +100,15 @@ def _attn_fwd_sparse(Q, K, V, sm_scale, # # ----- sparse loop over valid K/V tiles ----- for i in range(0, kv_blocks): kv_idx = tl.load(kv_ptr + i).to(tl.int32) - + block_size = tl.load(variable_block_sizes + kv_idx) K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N)) V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0)) k = tl.load(K_ptr) qk = tl.dot(q, k) + # mask out invalid columns + mask = tl.arange(0, BLOCK_N) < block_size + qk = tl.where(mask[None, :], qk, -float("inf")) m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) p = tl.math.exp2(qk * qk_scale - m_ij[:, None]) @@ -163,84 +130,6 @@ def _attn_fwd_sparse(Q, K, V, sm_scale, # # ──────────────────────────── SPARSE ADDITION END ───────────────────────────── -@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"]) -@triton.jit -def _attn_fwd(Q, K, V, sm_scale, M, Out, # - stride_qz, stride_qh, stride_qm, stride_qk, # - stride_kz, stride_kh, stride_kn, stride_kk, # - stride_vz, stride_vh, stride_vk, stride_vn, # - stride_oz, stride_oh, stride_om, stride_on, # - Z, H, N_CTX, # - HEAD_DIM: tl.constexpr, # - BLOCK_M: tl.constexpr, # - BLOCK_N: tl.constexpr, # - STAGE: tl.constexpr # - ): - tl.static_assert(BLOCK_N <= HEAD_DIM) - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh - - # block pointers - Q_block_ptr = tl.make_block_ptr( - base=Q + qvk_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_qm, stride_qk), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0) - V_block_ptr = tl.make_block_ptr( - base=V + qvk_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_vk, stride_vn), - offsets=(0, 0), - block_shape=(BLOCK_N, HEAD_DIM), - order=v_order, - ) - K_block_ptr = tl.make_block_ptr( - base=K + qvk_offset, - shape=(HEAD_DIM, N_CTX), - strides=(stride_kk, stride_kn), - offsets=(0, 0), - block_shape=(HEAD_DIM, BLOCK_N), - order=(0, 1), - ) - O_block_ptr = tl.make_block_ptr( - base=Out + qvk_offset, - shape=(N_CTX, HEAD_DIM), - strides=(stride_om, stride_on), - offsets=(start_m * BLOCK_M, 0), - block_shape=(BLOCK_M, HEAD_DIM), - order=(1, 0), - ) - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 - acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - # load scales - qk_scale = sm_scale - qk_scale *= 1.44269504 # 1/log(2) - # load q: it will stay in SRAM throughout - q = tl.load(Q_block_ptr) - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr, # - start_m, qk_scale, # - BLOCK_M, HEAD_DIM, BLOCK_N, # - 3, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5 # - ) - - # epilogue - m_i += tl.math.log2(l_i) - acc = acc / l_i[:, None] - m_ptrs = M + off_hz * N_CTX + offs_m - tl.store(m_ptrs, m_i) - tl.store(O_block_ptr, acc.to(Out.type.element_ty)) @triton.jit @@ -267,6 +156,7 @@ def _attn_bwd_dkdv(dk, dv, # DO, # M, D, # k2q_index, k2q_num, max_q_blks, + variable_block_sizes, # shared by Q/K/V/DO. stride_tok, stride_d, # H, N_CTX, BLOCK_M1: tl.constexpr, # @@ -291,6 +181,8 @@ def _attn_bwd_dkdv(dk, dv, # q_blocks = tl.load(k2q_num + meta_base) # int32 q_ptr = k2q_index + meta_base * max_q_blks # ptr to list + block_size = tl.load(variable_block_sizes + kv_blk) + for blk_idx in range(q_blocks*2): @@ -301,6 +193,9 @@ def _attn_bwd_dkdv(dk, dv, # m = tl.load(M + offs_m) qkT = tl.dot(k, qT) pT = tl.math.exp2(qkT - m[None, :]) + mask = tl.arange(0, BLOCK_N1) < block_size + pT = tl.where(mask[:, None], pT, 0.0) + do = tl.load(do_ptrs + block_sparse_offset * stride_tok) # Compute dV. ppT = pT @@ -324,6 +219,7 @@ def _attn_bwd_dq(dq, q, K, V, # do, m, D, # shared by Q/K/V/DO. q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, stride_tok, stride_d, # H, N_CTX, # BLOCK_M2: tl.constexpr, # @@ -355,10 +251,13 @@ def _attn_bwd_dq(dq, q, K, V, # for blk_idx in range(kv_blocks*2): block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok + block_size = tl.load(variable_block_sizes + blk_idx//2) - (blk_idx%2) * step_n kT = tl.load(kT_ptrs + block_sparse_offset) vT = tl.load(vT_ptrs + block_sparse_offset) qk = tl.dot(q, kT) p = tl.math.exp2(qk - m) + mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32) + p = tl.where(mask[None, :], p , 0.0) # Compute dP and dS. dp = tl.dot(do, vT).to(tl.float32) ds = p * (dp - Di[:, None]) @@ -378,6 +277,7 @@ def _attn_bwd(Q, K, V, sm_scale, # M, D, q2k_index, q2k_num, max_kv_blks, k2q_index, k2q_num, max_q_blks, + variable_block_sizes, # shared by Q/K/V/DO. stride_z, stride_h, stride_tok, stride_d, # H, N_CTX, # @@ -428,6 +328,7 @@ def _attn_bwd(Q, K, V, sm_scale, # DO, # M, D, # k2q_index, k2q_num, max_q_blks, + variable_block_sizes, stride_tok, stride_d, # H, N_CTX, # BLOCK_M1, BLOCK_N1, HEAD_DIM, # @@ -459,6 +360,7 @@ def _attn_bwd(Q, K, V, sm_scale, # dq = _attn_bwd_dq(dq, q, K, V, # do, m, D, # q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, stride_tok, stride_d, # H, N_CTX, # BLOCK_M2, BLOCK_N2, HEAD_DIM, # @@ -471,237 +373,77 @@ def _attn_bwd(Q, K, V, sm_scale, # -class _attention(torch.autograd.Function): - @staticmethod - def forward(ctx, q, k, v): - # shape constraints - sm_scale = 1.0 / (q.shape[-1] ** 0.5) - HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] - # when v is in float8_e5m2 it is transposed. - HEAD_DIM_V = v.shape[-1] - assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V - assert HEAD_DIM_K in {16, 32, 64, 128, 256} - o = torch.empty_like(q) - stage = 1 - extra_kern_args = {} - - - grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1) - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - _attn_fwd[grid]( - q, k, v, sm_scale, M, o, # - q.stride(0), q.stride(1), q.stride(2), q.stride(3), # - k.stride(0), k.stride(1), k.stride(2), k.stride(3), # - v.stride(0), v.stride(1), v.stride(2), v.stride(3), # - o.stride(0), o.stride(1), o.stride(2), o.stride(3), # - q.shape[0], q.shape[1], # - N_CTX=q.shape[2], # - HEAD_DIM=HEAD_DIM_K, # - STAGE=stage, # - **extra_kern_args) - - ctx.save_for_backward(q, k, v, o, M) - ctx.grid = grid - ctx.sm_scale = sm_scale - ctx.HEAD_DIM = HEAD_DIM_K - return o - - @staticmethod - def backward(ctx, do): - q, k, v, o, M = ctx.saved_tensors - assert do.is_contiguous() - assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - BATCH, N_HEAD, N_CTX = q.shape[:3] - PRE_BLOCK = 128 - BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32 - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) - arg_k = k - arg_k = arg_k * (ctx.sm_scale * RCP_LN2) - PRE_BLOCK = 128 - assert N_CTX % PRE_BLOCK == 0 - pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) - delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid]( - o, do, # - delta, # - BATCH, N_HEAD, N_CTX, # - BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM # - ) - grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD) - _attn_bwd[grid]( - q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, # - M, delta, # - q.stride(0), q.stride(1), q.stride(2), q.stride(3), # - N_HEAD, N_CTX, # - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, # - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, # - HEAD_DIM=ctx.HEAD_DIM # - ) - - return dq, dk, dv, None, None - - -attention = _attention.apply # ──────────────────────────── SPARSE ADDITION BEGIN ─────────────────────────── -class _attention_sparse(torch.autograd.Function): - """ - Thin autograd wrapper that uses the sparse forward kernel above and the - standard dense backward kernels defined earlier (no extra memory use). - """ +def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes): + B, H, T, D = q.shape + sm_scale = 1.0 / math.sqrt(D) + max_kv_blks = q2k_index.shape[-1] + assert T % 64 == 0, f"T must be a multiple of 64, but got {T}" + assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}" + o = torch.empty_like(q) + M = torch.empty((B, H, T), dtype=torch.float32, device=q.device) - @staticmethod - def forward(ctx, q, k, v, q2k_index, q2k_num, k2q_index, k2q_num): - B, H, T, D = q.shape - sm_scale = 1.0 / math.sqrt(D) - max_kv_blks = q2k_index.shape[-1] - assert T % 64 == 0, f"T must be a multiple of 64, but got {T}" - assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}" - o = torch.empty_like(q) - M = torch.empty((B, H, T), dtype=torch.float32, device=q.device) + grid = lambda _: (triton.cdiv(T, 64), B * H, 1) + _attn_fwd_sparse[grid]( + q, k, v, sm_scale, + q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, + M, o, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + v.stride(0), v.stride(1), v.stride(2), v.stride(3), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + B, H, T, + HEAD_DIM=D, STAGE=3 + ) - grid = lambda _: (triton.cdiv(T, 64), B * H, 1) - _attn_fwd_sparse[grid]( - q, k, v, sm_scale, - q2k_index, q2k_num, max_kv_blks, - M, o, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), - k.stride(0), k.stride(1), k.stride(2), k.stride(3), - v.stride(0), v.stride(1), v.stride(2), v.stride(3), - o.stride(0), o.stride(1), o.stride(2), o.stride(3), - B, H, T, - HEAD_DIM=D, STAGE=3 - ) + return o, M - ctx.save_for_backward(q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num) - ctx.grid = None - ctx.sm_scale = sm_scale - ctx.HEAD_DIM = D - return o +def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes): + assert do.is_contiguous() + assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() + + B, H, T, D = q.shape + sm_scale = 1.0 / math.sqrt(D) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + BATCH, N_HEAD, N_CTX = q.shape[:3] + BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32 + RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) + arg_k = k + arg_k = arg_k * (sm_scale * RCP_LN2) + PRE_BLOCK = 64 + assert N_CTX % PRE_BLOCK == 0 + pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) + delta = torch.empty_like(M) + _attn_bwd_preprocess[pre_grid]( + o, do, # + delta, # + BATCH, N_HEAD, N_CTX, # + BLOCK_M=PRE_BLOCK, HEAD_DIM=D # + ) + + + max_q_blks = k2q_index.shape[-1] + max_kv_blks = q2k_index.shape[-1] + + grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD) + _attn_bwd[grid]( + q, arg_k, v, sm_scale, do, dq, dk, dv, # + M, delta, # + q2k_index, q2k_num, max_kv_blks, + k2q_index, k2q_num, max_q_blks, + variable_block_sizes, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), # + N_HEAD, N_CTX, # + BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, # + BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, # + HEAD_DIM=D # + ) - @staticmethod - def backward(ctx, do): - q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num = ctx.saved_tensors - assert do.is_contiguous() - assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - BATCH, N_HEAD, N_CTX = q.shape[:3] - PRE_BLOCK = 128 - BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32 - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) - arg_k = k - arg_k = arg_k * (ctx.sm_scale * RCP_LN2) - PRE_BLOCK = 128 - assert N_CTX % PRE_BLOCK == 0 - pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) - delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid]( - o, do, # - delta, # - BATCH, N_HEAD, N_CTX, # - BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM # - ) - - - max_q_blks = k2q_index.shape[-1] - max_kv_blks = q2k_index.shape[-1] - - - grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD) - _attn_bwd[grid]( - q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, # - M, delta, # - q2k_index, q2k_num, max_kv_blks, - k2q_index, k2q_num, max_q_blks, - q.stride(0), q.stride(1), q.stride(2), q.stride(3), # - N_HEAD, N_CTX, # - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, # - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, # - HEAD_DIM=ctx.HEAD_DIM # - ) - - return dq, dk, dv, None, None, None, None - -attention_sparse = _attention_sparse.apply -# ──────────────────────────── SPARSE ADDITION END ───────────────────────────── + return dq, dk, dv -try: - from flash_attn.flash_attn_interface import \ - flash_attn_qkvpacked_func as flash_attn_func - HAS_FLASH = True -except BaseException: - HAS_FLASH = False - -TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2') -BATCH, N_HEADS, HEAD_DIM = 4, 32, 128 -# vary seq length for fixed head and batch=4 -configs = [] -for mode in ["fwd", "bwd"]: - configs.append( - triton.testing.Benchmark( - x_names=["N_CTX"], - x_vals=[2**i for i in range(10, 15)], - line_arg="provider", - line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) + - (["flash"] if HAS_FLASH else []), - line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) + - (["Flash-2"] if HAS_FLASH else []), - styles=[("red", "-"), ("blue", "-"), ("green", "-")], - ylabel="ms", - plot_name=f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}", - args={ - "H": N_HEADS, - "BATCH": BATCH, - "HEAD_DIM": HEAD_DIM, - "mode": mode, - }, - )) - - -@triton.testing.perf_report(configs) -def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, mode, provider, device="cuda"): - assert mode in ["fwd", "bwd"] - warmup = 25 - rep = 100 - dtype = torch.bfloat16 - if "triton" in provider: - q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - if mode == "fwd" and "fp8" in provider: - q = q.to(torch.float8_e5m2) - k = k.to(torch.float8_e5m2) - v = v.permute(0, 1, 3, 2).contiguous() - v = v.permute(0, 1, 3, 2) - v = v.to(torch.float8_e5m2) - fn = lambda: attention(q, k, v) - if mode == "bwd": - o = fn() - do = torch.randn_like(o) - fn = lambda: o.backward(do, retain_graph=True) - ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep) - if provider == "flash": - qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - fn = lambda: flash_attn_func(qkv, causal=False) - if mode == "bwd": - o = fn() - do = torch.randn_like(o) - fn = lambda: o.backward(do, retain_graph=True) - ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep) - flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM - total_flops = 2 * flops_per_matmul - if mode == "bwd": - total_flops *= 2.5 # 2.0(bwd) + 0.5(recompute) - return total_flops / ms * 1e-9 - - -if __name__ == "__main__": - # only works on post-Ampere GPUs right now - bench_flash_attention.run(save_path=".", print_data=True) diff --git a/csrc/attn/vsa/block_sparse_h100.cu b/csrc/attn/vsa/block_sparse_h100.cu index 2b74c7a8b..cf5ea8330 100644 --- a/csrc/attn/vsa/block_sparse_h100.cu +++ b/csrc/attn/vsa/block_sparse_h100.cu @@ -46,225 +46,9 @@ template struct fwd_globals { int32_t *__restrict__ q2k_block_sparse_index; int32_t *__restrict__ q2k_block_sparse_num; + int32_t *__restrict__ block_size; }; -template -__global__ __launch_bounds__(128, 3) // encourage compiler to reduce register usage so that an SM can hold 3 CTAs. Performance will drop from 391T to 353T if not specified explicitly. -void fwd_attend_ker_even(const __grid_constant__ fwd_globals g) { - extern __shared__ int __shm[]; - tma_swizzle_allocator al((int*)&__shm[0]); - - using K = fwd_attend_ker_tile_dims; - - using q_tile = st_bf<64, K::tile_width>; - using k_tile = st_bf<128, K::tile_width>; - using v_tile = st_bf<128, K::tile_width>; - using k_tile_half = st_bf<64, K::tile_width>; - using v_tile_half = st_bf<64, K::tile_width>; - using l_col_vec = col_vec>; - using o_tile = st_bf<64, K::tile_width>; - - q_tile (&q_smem)[1] = al.allocate(); - - k_tile_half (&k_smem_0)[1] = al.allocate(); - k_tile_half (&k_smem_1)[1] = al.allocate(); - k_tile (*k_smem) = reinterpret_cast(k_smem_0); - - v_tile_half (&v_smem_0)[1] = al.allocate(); - v_tile_half (&v_smem_1)[1] = al.allocate(); - v_tile (*v_smem) = reinterpret_cast(v_smem_0); - - l_col_vec (&l_smem)[1] = al.allocate(); - - auto (*o_smem) = reinterpret_cast(q_smem); - - int kv_head_idx = blockIdx.y / g.hr; - int seq_idx = blockIdx.x; - - - int32_t* q2k_block_sparse_index_ptr = g.q2k_block_sparse_index + blockIdx.z * gridDim.y * gridDim.x * g.max_kv_blocks_per_q + blockIdx.y * gridDim.x * g.max_kv_blocks_per_q + blockIdx.x * g.max_kv_blocks_per_q; - int32_t* q2k_block_sparse_num_ptr = g.q2k_block_sparse_num + blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x; - int32_t kv_blocks = q2k_block_sparse_num_ptr[0] / 2; // each iter load 2 kv blocks - __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived, v_smem_arrived; - if (threadIdx.x == 0) { - int32_t kv_block_index[2]; - reinterpret_cast(kv_block_index)[0] = reinterpret_cast(q2k_block_sparse_index_ptr)[0]; - - init_semaphore(qsmem_semaphore, 0, 1); - init_semaphore(k_smem_arrived, 0, 1); - init_semaphore(v_smem_arrived, 0, 1); - - // preload q block - coord q_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0}; - tma::expect_bytes(qsmem_semaphore, sizeof(q_smem)); - tma::load_async(q_smem[0], g.q, q_tile_idx, qsmem_semaphore); - - // preload the zeroth block of kv - tma::expect_bytes(k_smem_arrived, sizeof(k_tile)); - coord k_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0}; - coord k_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0}; - tma::load_async(k_smem_0[0], g.k, k_tile_idx_0, k_smem_arrived); - tma::load_async(k_smem_1[0], g.k, k_tile_idx_1, k_smem_arrived); - - tma::expect_bytes(v_smem_arrived, sizeof(v_tile)); - coord v_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0}; - coord v_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0}; - tma::load_async(v_smem_0[0], g.v, v_tile_idx_0, v_smem_arrived); - tma::load_async(v_smem_1[0], g.v, v_tile_idx_1, v_smem_arrived); - } - __syncthreads(); - - rt_fl<16, 128> att_block; - rt_bf<16, 128> att_block_mma; - rt_fl<16, K::tile_width> o_reg; - - col_vec> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled; - - neg_infty(max_vec); - zero(norm_vec); - zero(o_reg); - - // wait for q block - wait(qsmem_semaphore, 0); - - for (int kv_idx = 0; kv_idx < kv_blocks - 1; kv_idx++) { - // preload kv index - int32_t kv_block_index[2]; - reinterpret_cast(kv_block_index)[0] = reinterpret_cast(q2k_block_sparse_index_ptr)[kv_idx + 1]; - - // wait k - wait(k_smem_arrived, kv_idx % 2); - - // compute QK^T - warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); - - copy(max_vec_last_scaled, max_vec); - if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } - else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } - - warpgroup::mma_async_wait(); - - // load K - if (threadIdx.x == 0) { - tma::expect_bytes(k_smem_arrived, sizeof(k_tile)); - coord k_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0}; - coord k_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0}; - tma::load_async(k_smem_0[0], g.k, k_tile_idx_0, k_smem_arrived); - tma::load_async(k_smem_1[0], g.k, k_tile_idx_1, k_smem_arrived); - } - - // exp - row_max(max_vec, att_block, max_vec); - - if constexpr (D == 64) { - mul(att_block, att_block, 1.44269504089f*0.125f); - mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); - } - else { - mul(att_block, att_block, 1.44269504089f*0.08838834764f); - mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); - } - - sub_row(att_block, att_block, max_vec_scaled); - exp2(att_block, att_block); - sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); - exp2(max_vec_last_scaled, max_vec_last_scaled); - mul(norm_vec, norm_vec, max_vec_last_scaled); - row_sum(norm_vec, att_block, norm_vec); - add(att_block, att_block, 0.f); - copy(att_block_mma, att_block); - mul_row(o_reg, o_reg, max_vec_last_scaled); - - // wait v - wait(v_smem_arrived, kv_idx % 2); - - // compute SV - warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); - warpgroup::mma_async_wait(); - - // load V - if (threadIdx.x == 0) { - tma::expect_bytes(v_smem_arrived, sizeof(v_tile)); - // coord v_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0}; - // tma::load_async(v_smem[0], g.v, v_tile_idx, v_smem_arrived); - coord v_tile_idx_0 = {blockIdx.z, kv_head_idx, kv_block_index[0], 0}; - coord v_tile_idx_1 = {blockIdx.z, kv_head_idx, kv_block_index[1], 0}; - tma::load_async(v_smem_0[0], g.v, v_tile_idx_0, v_smem_arrived); - tma::load_async(v_smem_1[0], g.v, v_tile_idx_1, v_smem_arrived); - } - } - - // last iter - { - int kv_idx = kv_blocks - 1; - // wait k - wait(k_smem_arrived, kv_idx % 2); - - // compute QK^T - warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); - - copy(max_vec_last_scaled, max_vec); - if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } - else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } - - warpgroup::mma_async_wait(); - - // exp - row_max(max_vec, att_block, max_vec); - - if constexpr (D == 64) { - mul(att_block, att_block, 1.44269504089f*0.125f); - mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); - } - else { - mul(att_block, att_block, 1.44269504089f*0.08838834764f); - mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); - } - - sub_row(att_block, att_block, max_vec_scaled); - exp2(att_block, att_block); - sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); - exp2(max_vec_last_scaled, max_vec_last_scaled); - mul(norm_vec, norm_vec, max_vec_last_scaled); - row_sum(norm_vec, att_block, norm_vec); - add(att_block, att_block, 0.f); - copy(att_block_mma, att_block); - mul_row(o_reg, o_reg, max_vec_last_scaled); - - // wait v - wait(v_smem_arrived, kv_idx % 2); - - // compute SV - warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); - warpgroup::mma_async_wait(); - } - - div_row(o_reg, o_reg, norm_vec); - warpgroup::store(o_smem[0], o_reg); - __syncthreads(); - - // TK store_async internally calls syncwarp so we need to route on warp level - if (threadIdx.x / 32 == 0) { - coord o_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0}; - tma::store_async(g.o, o_smem[0], o_tile_idx); - } - - mul(max_vec_scaled, max_vec_scaled, 0.69314718056f); - log(norm_vec, norm_vec); - add(norm_vec, norm_vec, max_vec_scaled); - - if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); } - else { mul(norm_vec, norm_vec, -11.313708499f); } - - warpgroup::store(l_smem[0], norm_vec); - __syncthreads(); - - if (threadIdx.x / 32 == 0) { - coord tile_idx = {blockIdx.z, blockIdx.y, 0, seq_idx}; - tma::store_async(g.l, l_smem[0], tile_idx); - } - tma::store_async_wait(); -} template __global__ __launch_bounds__(128, 4) @@ -357,6 +141,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz } // exp + right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants::neg_infty()); row_max(max_vec, att_block, max_vec); if constexpr (D == 64) { @@ -409,6 +194,8 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz warpgroup::mma_async_wait(); // exp + right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants::neg_infty()); + row_max(max_vec, att_block, max_vec); if constexpr (D == 64) { @@ -484,8 +271,9 @@ struct bwd_prep_globals { d_gl d; }; +constexpr int PREP_NUM_WARPS = (1); template -__global__ __launch_bounds__(4*kittens::WARP_THREADS, (D == 64) ? 2 : 1) +__global__ __launch_bounds__(PREP_NUM_WARPS*kittens::WARP_THREADS, (D == 64) ? 6 / PREP_NUM_WARPS : 3 / PREP_NUM_WARPS) void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals g) { extern __shared__ int __shm[]; tma_swizzle_allocator al((int*)&__shm[0]); @@ -496,9 +284,9 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals g) { using o_tile = st_bf<4*16, D>; using d_tile = col_vec>; - og_tile (&og_smem)[4] = al.allocate(); - o_tile (&o_smem) [4] = al.allocate(); - d_tile (&d_smem) [4] = al.allocate(); + og_tile (&og_smem)[PREP_NUM_WARPS] = al.allocate(); + o_tile (&o_smem) [PREP_NUM_WARPS] = al.allocate(); + d_tile (&d_smem) [PREP_NUM_WARPS] = al.allocate(); rt_fl<4*16, D> og_reg, o_reg; col_vec> d_reg; @@ -507,13 +295,13 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals g) { if (threadIdx.x == 0) { init_semaphore(smem_semaphore, 0, 1); - tma::expect_bytes(smem_semaphore, sizeof(og_smem[0]) * 4 * 2); + tma::expect_bytes(smem_semaphore, sizeof(og_smem[0]) * PREP_NUM_WARPS * 2); } __syncthreads(); if (warpid == 0) { - for (int w = 0; w < 4; w++) { - coord tile_idx = {blockIdx.z, blockIdx.y, (blockIdx.x * 4) + w, 0}; + for (int w = 0; w < PREP_NUM_WARPS; w++) { + coord tile_idx = {blockIdx.z, blockIdx.y, (blockIdx.x * PREP_NUM_WARPS) + w, 0}; tma::load_async(o_smem[w], g.o, tile_idx, smem_semaphore); tma::load_async(og_smem[w], g.og, tile_idx, smem_semaphore); } @@ -528,8 +316,8 @@ void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals g) { __syncthreads(); if (warpid == 0) { - for (int w = 0; w < 4; w++) { - coord tile_idx = {blockIdx.z, blockIdx.y, 0, (blockIdx.x * 4) + w}; + for (int w = 0; w < PREP_NUM_WARPS; w++) { + coord tile_idx = {blockIdx.z, blockIdx.y, 0, (blockIdx.x * PREP_NUM_WARPS) + w}; tma::store_async(g.d, d_smem[w], tile_idx); } } @@ -591,6 +379,7 @@ struct bwd_globals { int32_t *__restrict__ k2q_block_sparse_index; int32_t *__restrict__ k2q_block_sparse_num; + int32_t *__restrict__ block_size; }; __device__ static inline void @@ -713,7 +502,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { // wait for kv wait(kv_b, 0); - + int fill_start = g.block_size[blockIdx.x] - 16 * kittens::warpid(); for (int qo_idx = 0; qo_idx < qo_blocks - 1; qo_idx++) { // preload q index store_qg_block_index = load_q_block_index; @@ -732,7 +521,8 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); } else { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); } - + + lower_fill(s_block_t, s_block_t, fill_start, base_types::constants::neg_infty()); exp2(s_block_t, s_block_t); // P_i copy(p_block_t, s_block_t); copy(p_block_t_mma, s_block_t); @@ -810,7 +600,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); } else { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); } - + lower_fill(s_block_t, s_block_t, fill_start, base_types::constants::neg_infty()); exp2(s_block_t, s_block_t); // P_i copy(p_block_t, s_block_t); copy(p_block_t_mma, s_block_t); @@ -876,7 +666,14 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { #include std::vector -block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num) +block_sparse_attention_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor q2k_block_sparse_index, + torch::Tensor q2k_block_sparse_num, + torch::Tensor block_size +) { CHECK_INPUT(q); CHECK_INPUT(k); @@ -888,6 +685,10 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v auto qo_heads = q.size(1); auto kv_heads = k.size(1); auto max_kv_blocks_per_q = q2k_block_sparse_index.size(3); + auto num_q_blocks = block_size.size(0); + TORCH_CHECK(batch==1, "Batch size dim will be removed in the future, please set batch to 1"); + TORCH_CHECK(num_q_blocks * 64 == seq_len, "This kernel supports variable block size, but it assumes the input sequence is properly padded."); + TORCH_CHECK(num_q_blocks == q2k_block_sparse_index.size(2), "Number of Q blocks does not match between q2k_block_sparse_index and block_size"); // check to see that these dimensions match for all inputs TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); @@ -967,7 +768,19 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(hr), static_cast(max_kv_blocks_per_q), reinterpret_cast(q2k_block_sparse_index.data_ptr()), reinterpret_cast(q2k_block_sparse_num.data_ptr())}; + globals g{ + qg_arg, + kg_arg, + vg_arg, + lg_arg, + og_arg, + static_cast(seq_len), + static_cast(hr), + static_cast(max_kv_blocks_per_q), + reinterpret_cast(q2k_block_sparse_index.data_ptr()), + reinterpret_cast(q2k_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr()) + }; constexpr int mem_size = 54000; @@ -1006,7 +819,19 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(hr), static_cast(max_kv_blocks_per_q), reinterpret_cast(q2k_block_sparse_index.data_ptr()), reinterpret_cast(q2k_block_sparse_num.data_ptr())}; + globals g{ + qg_arg, + kg_arg, + vg_arg, + lg_arg, + og_arg, + static_cast(seq_len), + static_cast(hr), + static_cast(max_kv_blocks_per_q), + reinterpret_cast(q2k_block_sparse_index.data_ptr()), + reinterpret_cast(q2k_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr()) + }; constexpr int mem_size = 54000; @@ -1036,7 +861,8 @@ block_sparse_attention_backward(torch::Tensor q, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, - torch::Tensor k2q_block_sparse_num) + torch::Tensor k2q_block_sparse_num, + torch::Tensor block_size) { CHECK_INPUT(q); CHECK_INPUT(k); @@ -1049,7 +875,7 @@ block_sparse_attention_backward(torch::Tensor q, auto seq_len = q.size(2); auto head_dim = q.size(3); auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3); - + TORCH_CHECK(k2q_block_sparse_index.size(2) == block_size.size(0), "k2q_block_sparse_index.size(2) must match block_size.size(0)"); // check to see that these dimensions match for all inputs TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); @@ -1136,7 +962,7 @@ block_sparse_attention_backward(torch::Tensor q, float* d_vg = reinterpret_cast(vg_ptr); constexpr int mem_size = kittens::MAX_SHARED_MEMORY; - int threads = 4 * kittens::WARP_THREADS; + int threads = PREP_NUM_WARPS * kittens::WARP_THREADS; //cudadevicesynchronize(); const c10::cuda::OptionalCUDAGuard device_guard(q.device()); @@ -1145,7 +971,7 @@ block_sparse_attention_backward(torch::Tensor q, // cudaStreamSynchronize(stream); // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256"); - dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM*4), qo_heads, batch); + dim3 grid_bwd(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); if (head_dim == 64) { using og_tile = st_bf<4*16, 64>; @@ -1220,8 +1046,8 @@ block_sparse_attention_backward(torch::Tensor q, static_cast(hr), static_cast(max_q_blocks_per_kv), reinterpret_cast(k2q_block_sparse_index.data_ptr()), - reinterpret_cast(k2q_block_sparse_num.data_ptr()) - }; + reinterpret_cast(k2q_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr())}; dim3 grid_bwd_2(seq_len/64, qo_heads, batch); threads = 128; @@ -1324,8 +1150,8 @@ block_sparse_attention_backward(torch::Tensor q, static_cast(hr), static_cast(max_q_blocks_per_kv), reinterpret_cast(k2q_block_sparse_index.data_ptr()), - reinterpret_cast(k2q_block_sparse_num.data_ptr()) - }; + reinterpret_cast(k2q_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr())}; dim3 grid_bwd_2(seq_len/64, qo_heads, batch); threads = 128; diff --git a/csrc/attn/vsa/block_sparse_wrapper.py b/csrc/attn/vsa/block_sparse_wrapper.py new file mode 100644 index 000000000..5a7418339 --- /dev/null +++ b/csrc/attn/vsa/block_sparse_wrapper.py @@ -0,0 +1,182 @@ +import torch +try: + from vsa_cuda import block_sparse_fwd, block_sparse_bwd +except ImportError: + block_sparse_fwd = None + block_sparse_bwd = None +from vsa.block_sparse_attn_triton import triton_block_sparse_attn_forward, triton_block_sparse_attn_backward +assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher" +from vsa.index import map_to_index +from typing import Tuple, Optional + + + +@torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda") +def block_sparse_attn_SM90( + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + )-> Tuple[torch.Tensor, torch.Tensor]: + q_padded = q_padded.contiguous() + k_padded = k_padded.contiguous() + v_padded = v_padded.contiguous() + q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) + variable_block_sizes = variable_block_sizes.int() + o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes) + return o_padded, lse_padded + + + + +@torch.library.register_fake("vsa::block_sparse_attn_SM90") +def _block_sparse_attn_SM90_fake( + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)] + B, H, S, D = q_padded.shape + o_padded = torch.empty_like(q_padded) + lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32) + return o_padded, lse_padded + + +@torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda") +def block_sparse_attn_backward_SM90( + grad_output_padded: torch.Tensor, + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + o_padded: torch.Tensor, + lse_padded: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grad_output_padded = grad_output_padded.contiguous() + k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2)) + grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd( + q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes + ) + grad_q_padded = grad_q_padded.to(grad_output_padded.dtype) + grad_k_padded = grad_k_padded.to(grad_output_padded.dtype) + grad_v_padded = grad_v_padded.to(grad_output_padded.dtype) + return grad_q_padded, grad_k_padded, grad_v_padded + +@torch.library.register_fake("vsa::block_sparse_attn_backward_SM90") +def _block_sparse_attn_backward_SM90_fake( + grad_output_padded: torch.Tensor, + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + o_padded: torch.Tensor, + lse_padded: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + torch._check(grad_output_padded.dtype == torch.bfloat16) + torch._check(lse_padded.dtype == torch.float32) + grad_output_padded = grad_output_padded.contiguous() + dq = torch.empty_like(grad_output_padded) + dk = torch.empty_like(grad_output_padded) + dv = torch.empty_like(grad_output_padded) + return dq, dk, dv + + +def backward_SM90(ctx, grad_output1, grad_output2): + q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes) + return dq, dk, dv, None, None + +def setup_context_SM90(ctx, inputs, output): + q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs + o_padded, lse_padded = output + ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes) + + +block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90) + + +@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda") +def block_sparse_attn_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + block_map = block_map.int() + q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) + o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes) + return o, M + + +@torch.library.register_fake("vsa::block_sparse_attn_triton") +def _block_sparse_attn_triton_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + o = torch.empty_like(q) + M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + return o, M + + +@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda") +def block_sparse_attn_backward_triton( + grad_output_padded: torch.Tensor, + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + o_padded: torch.Tensor, + M: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grad_output_padded = grad_output_padded.contiguous() + q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) + k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2)) + dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes) + return dq, dk, dv + +@torch.library.register_fake("vsa::block_sparse_attn_backward_triton") +def _block_sparse_attn_backward_triton_fake( + grad_output_padded: torch.Tensor, + q_padded: torch.Tensor, + k_padded: torch.Tensor, + v_padded: torch.Tensor, + o_padded: torch.Tensor, + M: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grad_output_padded = grad_output_padded.contiguous() + dq = torch.empty_like(grad_output_padded) + dk = torch.empty_like(grad_output_padded) + dv = torch.empty_like(grad_output_padded) + return dq, dk, dv + + +def backward_triton(ctx, grad_output1, grad_output2): + q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes) + return dq, dk, dv, None, None + + +def setup_context_triton(ctx, inputs, output): + q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs + o_padded, M = output + ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes) + +block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton) \ No newline at end of file diff --git a/csrc/attn/vsa/index.py b/csrc/attn/vsa/index.py new file mode 100644 index 000000000..6225f15b4 --- /dev/null +++ b/csrc/attn/vsa/index.py @@ -0,0 +1,152 @@ + +## pytorch sdpa version of block sparse ## +import triton +import triton.language as tl +import torch + +@triton.jit +def topk_index_to_map_kernel( + map_ptr, + index_ptr, + map_bs_stride, + map_h_stride, + map_q_stride, + map_kv_stride, + index_bs_stride, + index_h_stride, + index_q_stride, + index_kv_stride, + topk, +): + b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) + index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride + map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride + + for i in tl.static_range(topk): + index = tl.load(index_ptr_base + i * index_kv_stride) + tl.store(map_ptr_base + index * map_kv_stride, 1.0) + +@triton.jit +def map_to_index_kernel( + map_ptr, + index_ptr, + index_num_ptr, + map_bs_stride, + map_h_stride, + map_q_stride, + map_kv_stride, + index_bs_stride, + index_h_stride, + index_q_stride, + index_kv_stride, + index_num_bs_stride, + index_num_h_stride, + index_num_q_stride, + num_kv_blocks, +): + b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) + index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride + map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride + + num = 0 + for i in tl.range(num_kv_blocks): + map_entry = tl.load(map_ptr_base + i * map_kv_stride) + if map_entry: + tl.store(index_ptr_base + num * index_kv_stride, i) + num += 1 + + tl.store( + index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + + q * index_num_q_stride, num) + +def topk_index_to_map(index: torch.Tensor, + num_kv_blocks: int, + transpose_map: bool = False): + """ + Convert topk indices to a map. + + Args: + index: [bs, h, num_q_blocks, topk] + The topk indices tensor. + num_kv_blocks: int + The number of key-value blocks in the block_map returned + transpose_map: bool + If True, the block_map will be transposed on the final two dimensions. + + Returns: + block_map: [bs, h, num_q_blocks, num_kv_blocks] + A binary map where 1 indicates that the q block attends to the kv block. + """ + bs, h, num_q_blocks, topk = index.shape + + if transpose_map is False: + block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), + dtype=torch.bool, + device=index.device) + else: + block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), + dtype=torch.bool, + device=index.device) + block_map = block_map.transpose(2, 3) + + grid = (bs, h, num_q_blocks) + topk_index_to_map_kernel[grid]( + block_map, + index, + block_map.stride(0), + block_map.stride(1), + block_map.stride(2), + block_map.stride(3), + index.stride(0), + index.stride(1), + index.stride(2), + index.stride(3), + topk=topk, + ) + + return block_map + +def map_to_index(block_map: torch.Tensor): + """ + Convert a block map to indices and counts. + + Args: + block_map: [bs, h, num_q_blocks, num_kv_blocks] + The block map tensor. + + Returns: + index: [bs, h, num_q_blocks, num_kv_blocks] + The indices of the blocks. + index_num: [bs, h, num_q_blocks] + The number of blocks for each q block. + """ + bs, h, num_q_blocks, num_kv_blocks = block_map.shape + + index = torch.full((block_map.shape), + -1, + dtype=torch.int32, + device=block_map.device) + index_num = torch.empty((bs, h, num_q_blocks), + dtype=torch.int32, + device=block_map.device) + + grid = (bs, h, num_q_blocks) + map_to_index_kernel[grid]( + block_map, + index, + index_num, + block_map.stride(0), + block_map.stride(1), + block_map.stride(2), + block_map.stride(3), + index.stride(0), + index.stride(1), + index.stride(2), + index.stride(3), + index_num.stride(0), + index_num.stride(1), + index_num.stride(2), + num_kv_blocks=num_kv_blocks, + ) + + return index, index_num diff --git a/csrc/attn/vsa/vsa.py b/csrc/attn/vsa/vsa.py deleted file mode 100644 index 973df2ca2..000000000 --- a/csrc/attn/vsa/vsa.py +++ /dev/null @@ -1,47 +0,0 @@ -import torch -try: - from vsa_cuda import block_sparse_fwd, block_sparse_bwd -except ImportError: - block_sparse_fwd = None - block_sparse_bwd = None -from .block_sparse_attn_triton import attention_sparse as block_sparse_attn_triton - -class BlockSparseAttentionFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num): - o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num) - ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num) - return o - - @staticmethod - def backward(ctx, grad_output): - q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors - grad_q, grad_k, grad_v = block_sparse_bwd( - q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num - ) - return grad_q, grad_k, grad_v, None, None, None, None - - -@torch._dynamo.disable -def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num): - """ - Differentiable block sparse attention function. - - Args: - q: Query tensor [batch_size, num_heads, seq_len_q, head_dim] - k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim] - v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim] - q2k_block_sparse_index: Indices for query-to-key sparse blocks - q2k_block_sparse_num: Number of sparse blocks for each query block - k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass) - k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass) - - Returns: - output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim] - """ - if block_sparse_fwd is not None: - return BlockSparseAttentionFunction.apply( - q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num - ) - else: - return block_sparse_attn_triton(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) diff --git a/examples/inference/basic/basic_dmd.py b/examples/inference/basic/basic_dmd.py index 9bf7bc243..ca806257c 100644 --- a/examples/inference/basic/basic_dmd.py +++ b/examples/inference/basic/basic_dmd.py @@ -28,9 +28,7 @@ def main(): sampling_param = SamplingParam.from_pretrained(model_name) prompt = ( - "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " - "wide with interest. The playful yet serene atmosphere is complemented by soft " - "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + "A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. The puddles reflect glowing signs in kanji, advertising ramen, karaoke, and VR arcades. A woman in a translucent raincoat walks briskly with an LED umbrella. Steam rises from a street food cart, and a cat darts across the screen. Raindrops are visible on the camera lens, creating a cinematic bokeh effect." ) start_time = time.perf_counter() video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param) diff --git a/fastvideo/attention/backends/abstract.py b/fastvideo/attention/backends/abstract.py index f015ea82e..34a5bb665 100644 --- a/fastvideo/attention/backends/abstract.py +++ b/fastvideo/attention/backends/abstract.py @@ -6,8 +6,7 @@ from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar if TYPE_CHECKING: - from fastvideo.fastvideo_args import FastVideoArgs - from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + pass import torch @@ -87,10 +86,8 @@ class AttentionMetadataBuilder(ABC, Generic[T]): @abstractmethod def build( self, - current_timestep: int, - forward_batch: "ForwardBatch", - fastvideo_args: "FastVideoArgs", - ) -> T: + **kwargs: dict[str, Any], + ) -> AttentionMetadata: """Build attention metadata with on-device tensors.""" raise NotImplementedError diff --git a/fastvideo/attention/backends/sliding_tile_attn.py b/fastvideo/attention/backends/sliding_tile_attn.py index 1477de30c..77ab57a8b 100644 --- a/fastvideo/attention/backends/sliding_tile_attn.py +++ b/fastvideo/attention/backends/sliding_tile_attn.py @@ -13,10 +13,8 @@ from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionMetadata, AttentionMetadataBuilder) from fastvideo.distributed import get_sp_group -from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import ForwardContext, get_forward_context from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.utils import dict_to_3d_list logger = init_logger(__name__) @@ -76,13 +74,13 @@ class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder): def prepare(self): pass - def build( + def build( # type: ignore self, + STA_param: list[list[Any]], current_timestep: int, - forward_batch: ForwardBatch, - fastvideo_args: FastVideoArgs, + **kwargs: dict[str, Any], ) -> SlidingTileAttentionMetadata: - param = forward_batch.STA_param + param = STA_param if param is None: return SlidingTileAttentionMetadata( current_timestep=current_timestep, STA_param=[]) @@ -130,12 +128,6 @@ class SlidingTileAttentionImpl(AttentionImpl): } def tile(self, x: torch.Tensor) -> torch.Tensor: - x = rearrange(x, - "b (sp t h w) head d -> b (t sp h w) head d", - sp=self.sp_size, - t=self.dit_seq_shape_int[0] // self.sp_size, - h=self.dit_seq_shape_int[1], - w=self.dit_seq_shape_int[2]) return rearrange( x, "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", @@ -156,12 +148,7 @@ class SlidingTileAttentionImpl(AttentionImpl): ts_t=self.STA_base_tile_size[0], ts_h=self.STA_base_tile_size[1], ts_w=self.STA_base_tile_size[2]) - return rearrange(x, - "b (t sp h w) head d -> b (sp t h w) head d", - sp=self.sp_size, - t=self.dit_seq_shape_int[0] // self.sp_size, - h=self.dit_seq_shape_int[1], - w=self.dit_seq_shape_int[2]) + return x def preprocess_qkv( self, diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index acb516d79..15c1c6ba0 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -1,25 +1,114 @@ # SPDX-License-Identifier: Apache-2.0 +import functools import math from dataclasses import dataclass import torch -from einops import rearrange try: from vsa import video_sparse_attn except ImportError: video_sparse_attn = None +from typing import Any + from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, AttentionMetadataBuilder) from fastvideo.distributed import get_sp_group -from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger -from fastvideo.pipelines.pipeline_batch_info import ForwardBatch logger = init_logger(__name__) +VSA_TILE_SIZE = (4, 4, 4) + + +@functools.lru_cache(maxsize=10) +def get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + T, H, W = dit_seq_shape + ts, hs, ws = tile_size + indices = torch.arange(T * H * W, device=device, + dtype=torch.long).reshape(T, H, W) + ls = [] + for t in range(math.ceil(T / ts)): + for h in range(math.ceil(H / hs)): + for w in range(math.ceil(W / ws)): + ls.append(indices[t * ts:min(t * ts + ts, T), + h * hs:min(h * hs + hs, H), + w * ws:min(w * ws + ws, W)].flatten()) + index = torch.cat(ls, dim=0) + return index + + +@functools.lru_cache(maxsize=10) +def get_reverse_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + return torch.argsort( + get_tile_partition_indices(dit_seq_shape, tile_size, device)) + + +@functools.lru_cache(maxsize=10) +def construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + num_tiles: tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + """ + Compute the number of valid (non‑padded) tokens inside every + (ts_t × ts_h × ts_w) tile after padding ‑‑ flattened in the order + (t‑tile, h‑tile, w‑tile) that `rearrange` uses. + + Returns + ------- + torch.LongTensor # shape: [∏ full_window_size] + """ + # unpack + t, h, w = dit_seq_shape + ts_t, ts_h, ts_w = VSA_TILE_SIZE + n_t, n_h, n_w = num_tiles + + def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor: + """Vector with the size of each tile along one dimension.""" + sizes = torch.full((n_tiles, ), tile, dtype=torch.int, device=device) + # size of last (possibly partial) tile + remainder = dim_len - (n_tiles - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(t, ts_t, n_t) # [n_t] + h_sizes = _sizes(h, ts_h, n_h) # [n_h] + w_sizes = _sizes(w, ts_w, n_w) # [n_w] + + # broadcast‑multiply to get voxels per tile, then flatten + block_sizes = ( + t_sizes[:, None, None] # [n_t, 1, 1] + * h_sizes[None, :, None] # [1, n_h, 1] + * w_sizes[None, None, :] # [1, 1, n_w] + ).reshape(-1) # [n_t * n_h * n_w] + + return block_sizes + + +@functools.lru_cache(maxsize=10) +def get_non_pad_index( + variable_block_sizes: torch.LongTensor, + max_block_size: int, +): + n_win = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts_pad = torch.arange(n_win, device=device) * max_block_size + index_pad = starts_pad[:, None] + torch.arange(max_block_size, + device=device)[None, :] + index_mask = torch.arange( + max_block_size, device=device)[None, :] < variable_block_sizes[:, None] + return index_pad[index_mask] class VideoSparseAttentionBackend(AttentionBackend): @@ -52,6 +141,12 @@ class VideoSparseAttentionMetadata(AttentionMetadata): current_timestep: int dit_seq_shape: list[int] VSA_sparsity: float + num_tiles: list[int] + total_seq_length: int + tile_partition_indices: torch.LongTensor + reverse_tile_partition_indices: torch.LongTensor + variable_block_sizes: torch.LongTensor + non_pad_index: torch.LongTensor class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): @@ -62,30 +157,44 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder): def prepare(self): pass - def build( + def build( # type: ignore self, current_timestep: int, - forward_batch: ForwardBatch, - fastvideo_args: FastVideoArgs, + raw_latent_shape: tuple[int, int, int], + patch_size: tuple[int, int, int], + VSA_sparsity: float, + device: torch.device, + **kwargs: dict[str, Any], ) -> VideoSparseAttentionMetadata: - if forward_batch.latents is None: - raise ValueError("latents cannot be None") + patch_size = patch_size + dit_seq_shape = (raw_latent_shape[0] // patch_size[0], + raw_latent_shape[1] // patch_size[1], + raw_latent_shape[2] // patch_size[2]) - raw_latent_shape = forward_batch.raw_latent_shape - if raw_latent_shape is None: - raise ValueError("raw_latent_shape cannot be None") + num_tiles = (math.ceil(dit_seq_shape[0] / VSA_TILE_SIZE[0]), + math.ceil(dit_seq_shape[1] / VSA_TILE_SIZE[1]), + math.ceil(dit_seq_shape[2] / VSA_TILE_SIZE[2])) + total_seq_length = math.prod(dit_seq_shape) - patch_size = fastvideo_args.pipeline_config.dit_config.patch_size - dit_seq_shape = [ - raw_latent_shape[2] // patch_size[0], - raw_latent_shape[3] // patch_size[1], - raw_latent_shape[4] // patch_size[2] - ] - VSA_sparsity = forward_batch.VSA_sparsity + tile_partition_indices = get_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device) + reverse_tile_partition_indices = get_reverse_tile_partition_indices( + dit_seq_shape, VSA_TILE_SIZE, device) + variable_block_sizes = construct_variable_block_sizes( + dit_seq_shape, num_tiles, device) + non_pad_index = get_non_pad_index(variable_block_sizes, + math.prod(VSA_TILE_SIZE)) - return VideoSparseAttentionMetadata(current_timestep=current_timestep, - dit_seq_shape=dit_seq_shape, - VSA_sparsity=VSA_sparsity) + return VideoSparseAttentionMetadata( + current_timestep=current_timestep, + dit_seq_shape=dit_seq_shape, # type: ignore + VSA_sparsity=VSA_sparsity, # type: ignore + num_tiles=num_tiles, # type: ignore + total_seq_length=total_seq_length, # type: ignore + tile_partition_indices=tile_partition_indices, # type: ignore + reverse_tile_partition_indices=reverse_tile_partition_indices, + variable_block_sizes=variable_block_sizes, + non_pad_index=non_pad_index) class VideoSparseAttentionImpl(AttentionImpl): @@ -103,66 +212,44 @@ class VideoSparseAttentionImpl(AttentionImpl): self.prefix = prefix sp_group = get_sp_group() self.sp_size = sp_group.world_size - self.VSA_base_tile_size = [4, 4, 4] - self.dit_seq_shape: list[int] - self.full_window_size: list[int] - self.img_seq_length: int - def tile(self, x: torch.Tensor) -> torch.Tensor: - x = rearrange(x, - "b (sp t h w) head d -> b (t sp h w) head d", - sp=self.sp_size, - t=self.dit_seq_shape[0] // self.sp_size, - h=self.dit_seq_shape[1], - w=self.dit_seq_shape[2]) + def tile(self, x: torch.Tensor, num_tiles: list[int], + tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor) -> torch.Tensor: + t_padded_size = num_tiles[0] * VSA_TILE_SIZE[0] + h_padded_size = num_tiles[1] * VSA_TILE_SIZE[1] + w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2] - return rearrange( - x, - "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d", - n_t=self.full_window_size[0], - n_h=self.full_window_size[1], - n_w=self.full_window_size[2], - ts_t=self.VSA_base_tile_size[0], - ts_h=self.VSA_base_tile_size[1], - ts_w=self.VSA_base_tile_size[2]) + x_padded = torch.zeros( + (x.shape[0], t_padded_size * h_padded_size * w_padded_size, + x.shape[-2], x.shape[-1]), + device=x.device, + dtype=x.dtype) + x_padded[:, non_pad_index] = x[:, tile_partition_indices] + return x_padded - def untile(self, x: torch.Tensor) -> torch.Tensor: - x = rearrange( - x, - "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d", - n_t=self.full_window_size[0], - n_h=self.full_window_size[1], - n_w=self.full_window_size[2], - ts_t=self.VSA_base_tile_size[0], - ts_h=self.VSA_base_tile_size[1], - ts_w=self.VSA_base_tile_size[2]) - return rearrange(x, - "b (t sp h w) head d -> b (sp t h w) head d", - sp=self.sp_size, - t=self.dit_seq_shape[0] // self.sp_size, - h=self.dit_seq_shape[1], - w=self.dit_seq_shape[2]) + def untile(self, x: torch.Tensor, + reverse_tile_partition_indices: torch.LongTensor, + non_pad_index: torch.LongTensor) -> torch.Tensor: + x = x[:, non_pad_index][:, reverse_tile_partition_indices] + return x def preprocess_qkv( self, qkv: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: - self.dit_seq_shape = attn_metadata.dit_seq_shape - self.full_window_size = [ - self.dit_seq_shape[0] // self.VSA_base_tile_size[0], - self.dit_seq_shape[1] // self.VSA_base_tile_size[1], - self.dit_seq_shape[2] // self.VSA_base_tile_size[2] - ] - self.img_seq_length = math.prod(self.dit_seq_shape) - return self.tile(qkv) + return self.tile(qkv, attn_metadata.num_tiles, + attn_metadata.tile_partition_indices, + attn_metadata.non_pad_index) def postprocess_output( self, output: torch.Tensor, attn_metadata: VideoSparseAttentionMetadata, ) -> torch.Tensor: - return self.untile(output) + return self.untile(output, attn_metadata.reverse_tile_partition_indices, + attn_metadata.non_pad_index) def forward( # type: ignore[override] self, @@ -181,17 +268,17 @@ class VideoSparseAttentionImpl(AttentionImpl): cur_topk = math.ceil( (1 - VSA_sparsity) * - (self.img_seq_length / math.prod(self.VSA_base_tile_size))) + (attn_metadata.total_seq_length / math.prod(VSA_TILE_SIZE))) if video_sparse_attn is None: raise NotImplementedError("video_sparse_attn is not installed") - hidden_states = video_sparse_attn( query, key, value, + variable_block_sizes=attn_metadata.variable_block_sizes, topk=cur_topk, - block_size=(4, 4, 4), + block_size=VSA_TILE_SIZE, compress_attn_weight=gate_compress).transpose(1, 2) return hidden_states diff --git a/fastvideo/distributed/device_communicators/base_device_communicator.py b/fastvideo/distributed/device_communicators/base_device_communicator.py index 835fe826e..858e78e18 100644 --- a/fastvideo/distributed/device_communicators/base_device_communicator.py +++ b/fastvideo/distributed/device_communicators/base_device_communicator.py @@ -120,38 +120,43 @@ class DistributedAutograd: ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}" if scatter_dim == 2 and gather_dim == 1: - bs, shard_seqlen, hc, hs = input_.shape + bs, shard_seqlen, hn, hd = input_.shape seqlen = shard_seqlen * world_size - shard_hc = hc // world_size + shard_hn = hn // world_size - input_t = input_.reshape(bs, shard_seqlen, world_size, shard_hc, - hs).transpose(0, 2).contiguous() - output = torch.empty_like(input_t) + input_ = input_.transpose( + 0, 2).contiguous() # hn, shard_seqlen, bs, hd + output = torch.empty_like(input_) - dist.all_to_all_single(output, input_t, group=group) + dist.all_to_all_single(output, input_, + group=group) # hn, shard_seqlen, bs, hd - output = output.reshape(seqlen, bs, shard_hc, - hs).transpose(0, 1).contiguous() - output = output.reshape(bs, seqlen, shard_hc, hs) + output = torch.cat(output.split(shard_hn), + dim=1) # sharded hn, seqlen, bs, hd + + output = output.transpose( + 0, 2).contiguous() # bs, seqlen, sharded_hn, hd return output elif scatter_dim == 1 and gather_dim == 2: - bs, seqlen, shard_hc, hs = input_.shape - hc = shard_hc * world_size + bs, seqlen, shard_hn, hd = input_.shape + hn = shard_hn * world_size shard_seqlen = seqlen // world_size - input_t = input_.reshape(bs, world_size, shard_seqlen, shard_hc, - hs) - input_t = input_t.transpose(0, 3).transpose(0, 1).contiguous() - input_t = input_t.reshape(world_size, shard_hc, shard_seqlen, - bs, hs) + input_ = input_.transpose( + 0, 2).contiguous() # shard_hn, seqlen, bs, hd - output = torch.empty_like(input_t) - dist.all_to_all_single(output, input_t, group=group) + input_ = input_.reshape(shard_hn, world_size, shard_seqlen, bs, + hd).transpose(0, 1).reshape( + shard_hn * world_size, shard_seqlen, + bs, hd).contiguous() - output = output.reshape(hc, shard_seqlen, bs, hs) - output = output.transpose(0, 2).contiguous() - output = output.reshape(bs, shard_seqlen, hc, hs) + output = torch.empty_like(input_) + + dist.all_to_all_single(output, input_, group=group) + + output = output.transpose( + 0, 2).contiguous() # bs, seqlen, sharded_hn, hd return output else: diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index cbf5dfbb9..07d627c4a 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -228,10 +228,17 @@ class DenoisingStage(PipelineStage): self.attn_metadata_builder = self.attn_metadata_builder_cls( ) # TODO(will): clean this up - attn_metadata = self.attn_metadata_builder.build( - current_timestep=i, - forward_batch=batch, - fastvideo_args=fastvideo_args, + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=batch. + raw_latent_shape[2:5], # type: ignore + patch_size=fastvideo_args. + pipeline_config. # type: ignore + dit_config.patch_size, # type: ignore + STA_param=batch.STA_param, # type: ignore + VSA_sparsity=fastvideo_args. + VSA_sparsity, # type: ignore + device=get_local_torch_device(), ) assert attn_metadata is not None, "attn_metadata cannot be None" else: @@ -719,11 +726,18 @@ class DmdDenoisingStage(DenoisingStage): self.attn_metadata_builder = self.attn_metadata_builder_cls( ) # TODO(will): clean this up - attn_metadata = self.attn_metadata_builder.build( - current_timestep=i, - forward_batch=batch, - fastvideo_args=fastvideo_args, - ) + attn_metadata = self.attn_metadata_builder.build( # type: ignore + current_timestep=i, # type: ignore + raw_latent_shape=batch. + raw_latent_shape[2:5], # type: ignore + patch_size=fastvideo_args. + pipeline_config. # type: ignore + dit_config.patch_size, # type: ignore + STA_param=batch.STA_param, # type: ignore + VSA_sparsity=fastvideo_args. + VSA_sparsity, # type: ignore + device=get_local_torch_device(), # type: ignore + ) # type: ignore assert attn_metadata is not None, "attn_metadata cannot be None" else: attn_metadata = None diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 2be1f43ee..ad5f1093a 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -21,7 +21,7 @@ from tqdm.auto import tqdm import fastvideo.envs as envs from fastvideo.attention.backends.video_sparse_attn import ( - VideoSparseAttentionMetadata) + VideoSparseAttentionMetadataBuilder) from fastvideo.configs.sample import SamplingParam from fastvideo.dataset import build_parquet_map_style_dataloader from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v @@ -260,15 +260,13 @@ class TrainingPipeline(LoRAPipeline, ABC): assert latents_shape is not None assert training_batch.timesteps is not None if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN": - dit_seq_shape = [ - latents_shape[2] // patch_size[0], - latents_shape[3] // patch_size[1], - latents_shape[4] // patch_size[2] - ] - training_batch.attn_metadata = VideoSparseAttentionMetadata( + training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder( # type: ignore + ).build( # type: ignore + raw_latent_shape=latents_shape[2:5], current_timestep=training_batch.timesteps, - dit_seq_shape=dit_seq_shape, - VSA_sparsity=current_vsa_sparsity) + patch_size=patch_size, + VSA_sparsity=current_vsa_sparsity, + device=get_local_torch_device()) else: training_batch.attn_metadata = None diff --git a/scripts/inference/v1_inference_wan_dmd.sh b/scripts/inference/v1_inference_wan_dmd.sh index a38b5b77f..59b5048f2 100755 --- a/scripts/inference/v1_inference_wan_dmd.sh +++ b/scripts/inference/v1_inference_wan_dmd.sh @@ -10,9 +10,9 @@ fastvideo generate \ --sp-size $num_gpus \ --tp-size 1 \ --num-gpus $num_gpus \ - --height 448 \ - --width 832 \ - --num-frames 61 \ + --height 480 \ + --width 848 \ + --num-frames 81 \ --num-inference-steps 3 \ --fps 16 \ --prompt-txt assets/prompt.txt \