From 91364982dfb4fc1422fccbbc9866fc3d47f37666 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Wed, 17 Dec 2025 12:08:15 +0800 Subject: [PATCH] [Feature] Support for Variable Q/KV Sequence Lengths in VSA ThunderKittens kernel (#911) --- csrc/attn/tests/test_vsa.py | 109 +++++++- csrc/attn/tests/test_vsa_forward.py | 236 ++++++++++++++++++ csrc/attn/tests/utils.py | 48 ++-- .../vsa/block_sparse_h100.cu | 167 +++++++------ fastvideo/tests/modal/pr_test.py | 22 +- 5 files changed, 473 insertions(+), 109 deletions(-) create mode 100644 csrc/attn/tests/test_vsa_forward.py diff --git a/csrc/attn/tests/test_vsa.py b/csrc/attn/tests/test_vsa.py index 707d3e1ac..e2da3ec58 100644 --- a/csrc/attn/tests/test_vsa.py +++ b/csrc/attn/tests/test_vsa.py @@ -34,16 +34,16 @@ def pytorch_test(Q, K, V, block_sparse_mask, dO): ) -def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, non_pad_index, dO): +def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q_non_pad_index, kv_non_pad_index, q_num_blocks, kv_num_blocks, dO): Q = Q.detach().requires_grad_() K = K.detach().requires_grad_() V = V.detach().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) + q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M) + k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M) + v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M) output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes) - output = output[:, :, non_pad_index, :] + output = output[:, :, q_non_pad_index, :] output.backward(dO) return output, Q.grad, K.grad, V.grad @@ -64,7 +64,7 @@ def generate_tensor(shape, dtype, device): tensor = torch.randn(shape, dtype=dtype, device=device) return tensor -def generate_variable_block_sizes(num_blocks, min_size=32, max_size=64, device="cuda"): +def generate_variable_block_sizes(num_blocks, min_size=16, max_size=64, device="cuda"): return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32) @@ -86,19 +86,21 @@ def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all') 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) + block_mask = generate_block_sparse_mask_for_function(h, num_blocks, num_blocks, k, device) + full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, variable_block_sizes, device) for _ in range(num_iterations): Q = generate_tensor((1, h, S, d), torch.bfloat16, device) K = generate_tensor((1, h, S, d), torch.bfloat16, device) V = generate_tensor((1, h, S, d), torch.bfloat16, device) dO = generate_tensor((1, h, S, d), torch.bfloat16, device) + + # print(Q.shape, K.shape, V.shape, dO.shape) # 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) + 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, non_pad_index, num_blocks, num_blocks, 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 @@ -118,6 +120,60 @@ def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all') return results +def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, 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" + q_variable_block_sizes = generate_variable_block_sizes(num_q_blocks, device=device) + kv_variable_block_sizes = generate_variable_block_sizes(num_kv_blocks, device=device) + + S_q = int(q_variable_block_sizes.sum().item()) + S_kv = int(kv_variable_block_sizes.sum().item()) + + q_non_pad_index = get_non_pad_index(q_variable_block_sizes, num_q_blocks, BLOCK_M) + kv_non_pad_index = get_non_pad_index(kv_variable_block_sizes, num_kv_blocks, BLOCK_M) + + block_mask = generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device) + full_mask = create_full_mask_from_block_mask(block_mask, q_variable_block_sizes, kv_variable_block_sizes, device) + + for _ in range(num_iterations): + Q = generate_tensor((1, h, S_q, d), torch.bfloat16, device) + K = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) + V = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) + dO = generate_tensor((1, h, S_q, d), torch.bfloat16, device) + + # print(Q.shape, K.shape, V.shape, dO.shape) + + 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), kv_variable_block_sizes, q_non_pad_index, kv_non_pad_index, num_q_blocks, num_kv_blocks, 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() + rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt)) + results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item()) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + total_elements_q = h * S_q * d * num_iterations + total_elements_kv = h * S_kv * d * num_iterations + + for name, data in results.items(): + total_elements = total_elements_q if name in ['gQ', 'gO'] else total_elements_kv + 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, error_mode='all'): test_configs = [ {"num_blocks": 16, "k": 2, "description": "Small sequence"}, @@ -147,10 +203,43 @@ def generate_error_graphs(h, d, error_mode='all'): print("-" * 150) +def generate_error_graphs_qkdiff(h, d, error_mode='all'): + test_configs = [ + {"num_q_blocks": 16, "num_kv_blocks": 32, "k": 2, "description": "Small Q, Med KV"}, + {"num_q_blocks": 32, "num_kv_blocks": 16, "k": 4, "description": "Med Q, Small KV"}, + {"num_q_blocks": 53, "num_kv_blocks": 32, "k": 6, "description": "Large Q, Med KV"}, + {"num_q_blocks": 16, "num_kv_blocks": 48, "k": 2, "description": "Small Q, Large KV"}, + {"num_q_blocks": 48, "num_kv_blocks": 16, "k": 2, "description": "Large Q, Small KV"}, + ] + + print(f"\nError Analysis (QK Diff) for h={h}, d={d}, mode={error_mode}") + print("=" * 150) + print(f"{'Config':<20} {'Q Blks':<8} {'KV Blks':<8} {'K':<4} " + f"{'gQ Avg':<12} {'Rel gQ Max':<12} " + f"{'gK Avg':<12} {'Rel gK Max':<12} " + f"{'gV Avg':<12} {'Rel gV Max':<12} " + f"{'gO Avg':<12} {'Rel gO Max':<12}") + print("-" * 150) + + for config in test_configs: + num_q_blocks = config["num_q_blocks"] + num_kv_blocks = config["num_kv_blocks"] + k = config["k"] + description = config["description"] + results = check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, error_mode=error_mode) + print(f"{description:<20} {num_q_blocks:<8} {num_kv_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 print("Block Sparse Attention with Variable Block Sizes Analysis") print("=" * 60) for mode in ['backward']: generate_error_graphs(h, d, error_mode=mode) - print("\nAnalysis completed for all modes.") + generate_error_graphs_qkdiff(h, d, error_mode=mode) + print("\nAnalysis completed for all modes.") \ No newline at end of file diff --git a/csrc/attn/tests/test_vsa_forward.py b/csrc/attn/tests/test_vsa_forward.py new file mode 100644 index 000000000..d20055094 --- /dev/null +++ b/csrc/attn/tests/test_vsa_forward.py @@ -0,0 +1,236 @@ +import os +import sys +from typing import Tuple + +import torch + +# Make sure we can import from the project root (`vsa`, `tests.utils`, etc.) +CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) +PROJECT_ROOT = os.path.dirname(CURRENT_DIR) +if PROJECT_ROOT not in sys.path: + sys.path.append(PROJECT_ROOT) +if CURRENT_DIR not in sys.path: + sys.path.append(CURRENT_DIR) + +from tests.utils import ( + generate_block_sparse_mask_for_function, + create_full_mask_from_block_mask, +) +from vsa import block_sparse_attn, BLOCK_M +import test_vsa as ref # reuse helper functions from backward test + + +def pytorch_forward( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + block_sparse_mask: torch.Tensor, +) -> torch.Tensor: + """ + Dense PyTorch reference forward: + - Q: [1, h, S_q, d] + - K,V: [1, h, S_kv, d] + - block_sparse_mask: [h, S_q, S_kv] bool + """ + q = Q.clone().float() + k = K.clone().float() + v = V.clone().float() + + attn = torch.matmul(q, k.transpose(-2, -1)) # [1, h, S_q, S_kv] + attn = attn / (q.size(-1) ** 0.5) + attn = attn.masked_fill(~block_sparse_mask.unsqueeze(0), float("-inf")) + attn = torch.nn.functional.softmax(attn, dim=-1) + out = torch.matmul(attn, v) # [1, h, S_q, d] + return out.to(torch.bfloat16) + + +def block_sparse_forward_test( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + block_sparse_mask: torch.Tensor, + variable_block_sizes: torch.Tensor, + q_non_pad_index: torch.Tensor, + kv_non_pad_index: torch.Tensor, + q_num_blocks: int, + kv_num_blocks: int, +) -> torch.Tensor: + """ + Forward-only wrapper around `block_sparse_attn`, mirroring `block_sparse_kernel_test` + but without any backward / grad logic. + """ + Q = Q.detach() + K = K.detach() + V = V.detach() + + q_padded = ref.vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M) + k_padded = ref.vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M) + v_padded = ref.vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M) + + out_padded, _ = block_sparse_attn( + q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes + ) + # Remove padding on the query side + out = out_padded[:, :, q_non_pad_index, :] + return out + + +def run_forward_equal_qk( + h: int = 16, + d: int = 128, + num_blocks: int = 16, + k: int = 2, + num_iterations: int = 5, +) -> Tuple[float, float]: + """ + Forward-only correctness test for the case S_q == S_kv. + Mirrors `check_correctness` but only compares forward outputs. + """ + assert torch.cuda.is_available(), "VSA kernels require CUDA" + device = "cuda" + + variable_block_sizes = ref.generate_variable_block_sizes( + num_blocks, device=device + ) + S = int(variable_block_sizes.sum().item()) + non_pad_index = ref.get_non_pad_index( + variable_block_sizes, num_blocks, BLOCK_M + ) + + block_mask = generate_block_sparse_mask_for_function( + h, num_blocks, num_blocks, k, device + ) + full_mask = create_full_mask_from_block_mask( + block_mask, variable_block_sizes, variable_block_sizes, device + ) + print(f"[qkequal] h: {h}, d: {d}, num_blocks: {num_blocks}, k: {k}") + print(f"[qkequal] variable_block_sizes: {variable_block_sizes}, non_pad_index: {non_pad_index.shape}, block_mask: {block_mask.shape}, full_mask: {full_mask.shape}") + sum_diff = 0.0 + sum_abs = 0.0 + max_rel_diff = 0.0 + + for i in range(num_iterations): + Q = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) + K = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) + V = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) + + if i == 0: print(f"[qkequal] Q: {Q.shape}, K: {K.shape}, V: {V.shape}, full_mask: {full_mask.shape}") + if i == 0: print(f"[qkequal] block_mask: {block_mask.shape}") + + pt_o = pytorch_forward(Q, K, V, full_mask) + bs_o = block_sparse_forward_test( + Q, + K, + V, + block_mask.unsqueeze(0), + variable_block_sizes, + non_pad_index, + non_pad_index, + num_blocks, + num_blocks, + ) + + diff = (pt_o - bs_o).abs() + sum_diff += diff.sum().item() + sum_abs += pt_o.abs().sum().item() + rel_max = diff.max() / (pt_o.abs().mean() + 1e-6) + max_rel_diff = max(max_rel_diff, rel_max.item()) + + total_elems = h * S * d * num_iterations + avg_abs_err = sum_diff / total_elems + return avg_abs_err, max_rel_diff + + +def run_forward_qk_diff( + h: int = 16, + d: int = 128, + num_q_blocks: int = 16, + num_kv_blocks: int = 32, + k: int = 2, + num_iterations: int = 5, +) -> Tuple[float, float]: + """ + Forward-only correctness test for the case S_q != S_kv. + + NOTE: + - The Triton backend supports different Q/KV logical lengths via padding. + - The SM90 (H100) CUDA backend currently assumes the same number of blocks + for Q and KV, so we skip this test there. + """ + assert torch.cuda.is_available(), "VSA kernels require CUDA" + + device = "cuda" + + q_variable_block_sizes = ref.generate_variable_block_sizes( + num_q_blocks, device=device + ) + kv_variable_block_sizes = ref.generate_variable_block_sizes( + num_kv_blocks, device=device + ) + + S_q = int(q_variable_block_sizes.sum().item()) + S_kv = int(kv_variable_block_sizes.sum().item()) + + q_non_pad_index = ref.get_non_pad_index( + q_variable_block_sizes, num_q_blocks, BLOCK_M + ) + kv_non_pad_index = ref.get_non_pad_index( + kv_variable_block_sizes, num_kv_blocks, BLOCK_M + ) + + block_mask = generate_block_sparse_mask_for_function( + h, num_q_blocks, num_kv_blocks, k, device + ) + full_mask = create_full_mask_from_block_mask( + block_mask, q_variable_block_sizes, kv_variable_block_sizes, device + ) + + sum_diff = 0.0 + sum_abs = 0.0 + max_rel_diff = 0.0 + + for _ in range(num_iterations): + Q = ref.generate_tensor((1, h, S_q, d), torch.bfloat16, device) + K = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device) + V = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device) + + pt_o = pytorch_forward(Q, K, V, full_mask) + bs_o = block_sparse_forward_test( + Q, + K, + V, + block_mask.unsqueeze(0), + kv_variable_block_sizes, + q_non_pad_index, + kv_non_pad_index, + num_q_blocks, + num_kv_blocks, + ) + + diff = (pt_o - bs_o).abs() + sum_diff += diff.sum().item() + sum_abs += pt_o.abs().sum().item() + rel_max = diff.max() / (pt_o.abs().mean() + 1e-6) + max_rel_diff = max(max_rel_diff, rel_max.item()) + + total_elems = h * S_q * d * num_iterations + avg_abs_err = sum_diff / total_elems + return avg_abs_err, max_rel_diff + + +if __name__ == "__main__": + h, d = 16, 128 + print("Forward Block Sparse Attention Check (QK Equal)") + print("=" * 80) + avg_err_eq, max_rel_eq = run_forward_equal_qk(h, d, num_blocks=32, k=2) + print(f"QK equal: avg |ΔO| = {avg_err_eq:.6e}, max rel ΔO = {max_rel_eq:.6e}") + + print("\nForward Block Sparse Attention Check (QK Different)") + print("=" * 80) + avg_err_diff, max_rel_diff = run_forward_qk_diff( + h, d, num_q_blocks=32, num_kv_blocks=48, k=2 + ) + print( + f"QK diff: avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}" + ) + diff --git a/csrc/attn/tests/utils.py b/csrc/attn/tests/utils.py index 30e53deb0..3d9f15f45 100644 --- a/csrc/attn/tests/utils.py +++ b/csrc/attn/tests/utils.py @@ -1,54 +1,60 @@ import torch -def generate_block_sparse_mask_for_function(h, num_blocks, k, device="cuda"): +def generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device="cuda"): """ - Generate block sparse mask of shape [h, num_blocks, num_blocks]. + Generate block sparse mask of shape [h, num_q_blocks, num_kv_blocks]. Args: h: number of heads - num_blocks: number of blocks + 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: - block_sparse_mask: [h, num_blocks, num_blocks] bool tensor + block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor """ - k = min(k, num_blocks) - scores = torch.rand(h, num_blocks, num_blocks, device=device) + k = min(k, num_kv_blocks) + scores = torch.rand(h, num_q_blocks, num_kv_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 = torch.zeros(h, num_q_blocks, num_kv_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"): +def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, + kv_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 + block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor + q_variable_block_sizes: [num_q_blocks] tensor + kv_variable_block_sizes: [num_kv_blocks] tensor device: device to create tensors on Returns: - full_mask: [h, S, S] bool tensor where S = total sequence length + full_mask: [h, S_q, S_kv] 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]]) + h, num_q_blocks, num_kv_blocks = block_sparse_mask.shape + total_q_seq_len = q_variable_block_sizes.sum().item() + total_kv_seq_len = kv_variable_block_sizes.sum().item() + + q_cumsum = torch.cat([torch.tensor([0], device=device), q_variable_block_sizes.cumsum(dim=0)[:-1]]) + kv_cumsum = torch.cat([torch.tensor([0], device=device), kv_variable_block_sizes.cumsum(dim=0)[:-1]]) - full_mask = torch.zeros(h, total_seq_len, total_seq_len, dtype=torch.bool, device=device) + full_mask = torch.zeros(h, total_q_seq_len, total_kv_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 q_block in range(num_q_blocks): + q_start = q_cumsum[q_block] + q_end = q_start + q_variable_block_sizes[q_block] - for kv_block in range(num_blocks): + for kv_block in range(num_kv_blocks): if block_sparse_mask[head, q_block, kv_block]: - kv_start = cumsum[kv_block] - kv_end = kv_start + variable_block_sizes[kv_block] + kv_start = kv_cumsum[kv_block] + kv_end = kv_start + kv_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/video_sparse_attn/vsa/block_sparse_h100.cu b/csrc/attn/video_sparse_attn/vsa/block_sparse_h100.cu index dca9dd24a..f33d2719e 100644 --- a/csrc/attn/video_sparse_attn/vsa/block_sparse_h100.cu +++ b/csrc/attn/video_sparse_attn/vsa/block_sparse_h100.cu @@ -672,23 +672,32 @@ block_sparse_attention_forward( torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, - torch::Tensor block_size + torch::Tensor kv_block_size ) { CHECK_INPUT(q); CHECK_INPUT(k); CHECK_INPUT(v); + // q shape: (batch, qo_heads, q_seq_len, head_dim) + // k shape: (batch, kv_heads, kv_seq_len, head_dim) + // v shape: (batch, kv_heads, kv_seq_len, head_dim) + // q2k_block_sparse_index shape: (batch, qo_heads, num_q_blocks, max_kv_blocks_per_q) + // q2k_block_sparse_num shape: (batch, qo_heads, num_q_blocks) + // kv_block_size shape: (num_kv_blocks) This does not need other dimensions because across all batch/heads the padding is the same. + auto batch = q.size(0); - auto seq_len = q.size(2); + auto q_seq_len = q.size(2); + auto kv_seq_len = k.size(2); auto head_dim = q.size(3); 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); + auto num_q_blocks = q2k_block_sparse_index.size(2); + auto num_kv_blocks = kv_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"); + TORCH_CHECK(num_q_blocks * BLOCK_M == q_seq_len, "This kernel supports variable q block size, but it assumes the input sequence is properly padded."); + TORCH_CHECK(num_kv_blocks * BLOCK_M == kv_seq_len, "This kernel supports variable kv block size, but it assumes the input sequence is properly padded."); // 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"); @@ -696,11 +705,8 @@ block_sparse_attention_forward( TORCH_CHECK(q2k_block_sparse_index.size(0) == batch, "q2k_block_sparse_index batch dimension - idx 0 - must match for all inputs"); TORCH_CHECK(q2k_block_sparse_num.size(0) == batch, "q2k_block_sparse_num batch dimension - idx 0 - must match for all inputs"); - TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(q2k_block_sparse_index.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_index idx 2 - must match seq_len / BLOCK_M"); - TORCH_CHECK(q2k_block_sparse_num.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_num idx 2 - must match seq_len / BLOCK_M"); + TORCH_CHECK(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K inputs"); + TORCH_CHECK(q2k_block_sparse_num.size(2) == num_q_blocks, "q2k_block_sparse_num idx 2 - must match num_q_blocks"); TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); @@ -727,12 +733,12 @@ block_sparse_attention_forward( // for the returned outputs torch::Tensor o = torch::empty({static_cast(batch), static_cast(qo_heads), - static_cast(seq_len), + static_cast(q_seq_len), static_cast(head_dim)}, v.options()); torch::Tensor l_vec = torch::empty({static_cast(batch), static_cast(qo_heads), - static_cast(seq_len), + static_cast(q_seq_len), static_cast(1)}, torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous)); @@ -762,11 +768,11 @@ block_sparse_attention_forward( using globals = fwd_globals<64>; - q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - 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}; + q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; + o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; globals g{ qg_arg, @@ -774,17 +780,17 @@ block_sparse_attention_forward( vg_arg, lg_arg, og_arg, - static_cast(seq_len), + static_cast(q_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()) + reinterpret_cast(kv_block_size.data_ptr()) }; constexpr int mem_size = 54000; - dim3 grid(seq_len/(64), qo_heads, batch); + dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); cudaFuncSetAttribute( fwd_attend_ker<64>, @@ -813,11 +819,11 @@ block_sparse_attention_forward( using globals = fwd_globals<128>; - q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - 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}; + q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; + o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; globals g{ qg_arg, @@ -825,17 +831,17 @@ block_sparse_attention_forward( vg_arg, lg_arg, og_arg, - static_cast(seq_len), + static_cast(q_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()) + reinterpret_cast(kv_block_size.data_ptr()) }; constexpr int mem_size = 54000; - dim3 grid(seq_len/(64), qo_heads, batch); + dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); cudaFuncSetAttribute( fwd_attend_ker<128>, @@ -862,7 +868,7 @@ block_sparse_attention_backward(torch::Tensor q, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, - torch::Tensor block_size) + torch::Tensor kv_block_size) { CHECK_INPUT(q); CHECK_INPUT(k); @@ -871,11 +877,23 @@ block_sparse_attention_backward(torch::Tensor q, CHECK_INPUT(o); CHECK_INPUT(og); + // q: [batch, qo_heads, q_seq_len, head_dim] + // k: [batch, kv_heads, kv_seq_len, head_dim] + // v: [batch, kv_heads, kv_seq_len, head_dim] + // o: [batch, qo_heads, q_seq_len, head_dim] + // l_vec: [batch, qo_heads, q_seq_len, 1] + // og: [batch, qo_heads, q_seq_len, head_dim] + // k2q_block_sparse_index: [batch, kv_heads, num_kv_blocks, max_num_q_blocks] + // k2q_block_sparse_num: [batch, kv_heads, num_kv_blocks] + // kv_block_size: [num_kv_blocks] + auto batch = q.size(0); - auto seq_len = q.size(2); + auto q_seq_len = q.size(2); + auto kv_seq_len = k.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)"); + auto num_kv_blocks = kv_block_size.size(0); + TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index.size(2) must match num_kv_blocks (kv_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"); @@ -886,23 +904,18 @@ block_sparse_attention_backward(torch::Tensor q, TORCH_CHECK(k2q_block_sparse_index.size(0) == batch, "k2q_block_sparse_index batch dimension - idx 0 - must match for all inputs"); TORCH_CHECK(k2q_block_sparse_num.size(0) == batch, "k2q_block_sparse_num batch dimension - idx 0 - must match for all inputs"); - - TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(l_vec.size(2) == seq_len, "L sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(o.size(2) == seq_len, "O sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(og.size(2) == seq_len, "OG sequence length dimension - idx 2 - must match for all inputs"); - TORCH_CHECK(k2q_block_sparse_index.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_index idx 2 - must match seq_len / BLOCK_N"); - TORCH_CHECK(k2q_block_sparse_num.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_num idx 2 - must match seq_len / BLOCK_N"); + TORCH_CHECK(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K sequence length"); + TORCH_CHECK(l_vec.size(2) == q_seq_len, "L sequence length dimension - idx 2 - must match Q sequence length"); + TORCH_CHECK(o.size(2) == q_seq_len, "O sequence length dimension - idx 2 - must match Q sequence length"); + TORCH_CHECK(og.size(2) == q_seq_len, "OG sequence length dimension - idx 2 - must match Q sequence length"); + TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index idx 2 - must match num_kv_blocks (kv_block_size.size(0))"); + TORCH_CHECK(k2q_block_sparse_num.size(2) == num_kv_blocks, "k2q_block_sparse_num idx 2 - must match num_kv_blocks (kv_block_size.size(0))"); TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs"); TORCH_CHECK(o.size(3) == head_dim, "O head dimension - idx 3 - must match for all non-vector inputs"); TORCH_CHECK(og.size(3) == head_dim, "OG head dimension - idx 3 - must match for all non-vector inputs"); - - auto qo_heads = q.size(1); auto kv_heads = k.size(1); @@ -929,20 +942,20 @@ block_sparse_attention_backward(torch::Tensor q, torch::Tensor qg = torch::zeros({static_cast(batch), static_cast(qo_heads), - static_cast(seq_len), + static_cast(q_seq_len), static_cast(head_dim)}, l_vec.options()); torch::Tensor kg = torch::zeros({static_cast(batch), static_cast(kv_heads), - static_cast(seq_len), + static_cast(kv_seq_len), static_cast(head_dim)}, l_vec.options()); torch::Tensor vg = torch::zeros({static_cast(batch), static_cast(kv_heads), - static_cast(seq_len), + static_cast(kv_seq_len), static_cast(head_dim)}, l_vec.options()); torch::Tensor d_vec = torch::empty({static_cast(batch), static_cast(qo_heads), - static_cast(seq_len), + static_cast(q_seq_len), static_cast(1)}, l_vec.options()); float* qg_ptr = qg.data_ptr(); @@ -971,7 +984,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/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); + dim3 grid_bwd(q_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>; @@ -984,9 +997,9 @@ block_sparse_attention_backward(torch::Tensor q, using bwd_prep_globals = bwd_prep_globals<64>; - og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; @@ -1023,15 +1036,15 @@ block_sparse_attention_backward(torch::Tensor q, using bwd_global_args = bwd_globals<64>; - bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; - bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; - bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; - bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; + bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; + bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; + bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; bwd_global_args bwd_global{bwd_q_arg, bwd_k_arg, @@ -1042,14 +1055,14 @@ block_sparse_attention_backward(torch::Tensor q, bwd_vg_arg, bwd_l_arg, bwd_d_arg, - static_cast(seq_len), + static_cast(kv_seq_len), // N is not used in the kernel 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(block_size.data_ptr())}; + reinterpret_cast(kv_block_size.data_ptr())}; - dim3 grid_bwd_2(seq_len/64, qo_heads, batch); + dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch); threads = 128; //cudadevicesynchronize(); @@ -1088,9 +1101,9 @@ block_sparse_attention_backward(torch::Tensor q, using bwd_prep_globals = bwd_prep_globals<128>; - og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; @@ -1127,15 +1140,15 @@ block_sparse_attention_backward(torch::Tensor q, using bwd_global_args = bwd_globals<128>; - bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; - bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; - bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; - bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; + bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; + bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; + bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; bwd_global_args bwd_global{bwd_q_arg, bwd_k_arg, @@ -1146,14 +1159,14 @@ block_sparse_attention_backward(torch::Tensor q, bwd_vg_arg, bwd_l_arg, bwd_d_arg, - static_cast(seq_len), + static_cast(kv_seq_len), // N is not used in the kernel 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(block_size.data_ptr())}; + reinterpret_cast(kv_block_size.data_ptr())}; - dim3 grid_bwd_2(seq_len/64, qo_heads, batch); + dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch); threads = 128; //cudadevicesynchronize(); diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index b00089328..b7b1b987f 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -11,7 +11,21 @@ print(f"Using image: {image_tag}") image = ( modal.Image.from_registry(image_tag, add_python="3.12") .run_commands("rm -rf /FastVideo") - .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg") + .apt_install( + "cmake", + "pkg-config", + "build-essential", + "gcc-11", + "g++-11", + "clang-11", + "curl", + "libssl-dev", + "ffmpeg", + "git", + ) + .run_commands( + "update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11" + ) .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") .env({ @@ -52,7 +66,13 @@ def run_test(pytest_command: str): git clone {git_repo} /FastVideo && cd /FastVideo && {checkout_command} && + git submodule update --init --recursive && + export THUNDERKITTENS_ROOT=/FastVideo/csrc/attn/video_sparse_attn/tk && + export CC=gcc-11 && export CXX=g++-11 && uv pip install -e .[test] && + cd csrc/attn/video_sparse_attn && + python setup.py install && + cd ../../.. && {pytest_command} """