From 560628709c41a9a6fefdacaecff55d9b2ba83cbd Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Fri, 16 Jan 2026 21:16:43 -0800 Subject: [PATCH] [Bug Fix] Add autograd wrapper for block-sparse attention in fastvideo-kernel + fix CMake extension linking (#1015) --- docs/attention/vsa/index.md | 4 +- fastvideo-kernel/CMakeLists.txt | 26 ++ fastvideo-kernel/README.md | 2 +- .../csrc/attention/block_sparse_h100.cu | 394 ++++++++---------- .../fastvideo_kernel/block_sparse_attn.py | 291 +++++++++++++ .../python/fastvideo_kernel/ops.py | 59 ++- fastvideo-kernel/tests/test_vsa.py | 37 +- fastvideo-kernel/tests/test_vsa_forward.py | 27 +- .../attention/backends/video_sparse_attn.py | 5 +- 9 files changed, 558 insertions(+), 287 deletions(-) create mode 100644 fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py diff --git a/docs/attention/vsa/index.md b/docs/attention/vsa/index.md index 6cc8b5e6f..5dc5589f0 100644 --- a/docs/attention/vsa/index.md +++ b/docs/attention/vsa/index.md @@ -13,11 +13,13 @@ from fastvideo_kernel import video_sparse_attn # q, k, v: [batch_size, num_heads, seq_len, head_dim] # variable_block_sizes: Number of valid tokens per block +# q_variable_block_sizes: Number of valid tokens per q block (can differ from KV for q/k of different lengths) # topk: Number of blocks to attend output = video_sparse_attn( q, k, v, - variable_block_sizes=block_sizes, + block_sizes, + block_sizes, topk=32 ) ``` diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 6c433b936..ca8b67c5a 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -10,6 +10,8 @@ if(GPU_BACKEND STREQUAL "ROCM") enable_language(HIP) else() enable_language(CUDA) + # Ensure CUDA toolkit targets (CUDA::cudart, CUDA::cuda_driver, etc.) are available. + find_package(CUDAToolkit REQUIRED) endif() # Import common utils if needed, but we keep it simple for now @@ -153,6 +155,30 @@ if(BUILD_CXX_KERNELS) $<$:${CUDA_FLAGS}> ) + # Link against Torch libraries to avoid undefined symbols at import time + # (e.g., torch::autograd vtables) when loading the extension module. + target_link_libraries(fastvideo_kernel_ops PRIVATE ${TORCH_LIBRARIES}) + + # Also link against libtorch_python to satisfy Python-binding symbols + # (e.g., torch::PyWarningHandler) required by torch/extension.h. + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import torch; from pathlib import Path; p=Path(torch.__file__).parent/'lib'; m=sorted(p.glob('libtorch_python*')); print(str(m[0]) if m else '')" + OUTPUT_VARIABLE TORCH_PYTHON_LIBRARY_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + ) + if(TORCH_PYTHON_LIBRARY_PATH) + message(STATUS "TORCH_PYTHON_LIBRARY_PATH: ${TORCH_PYTHON_LIBRARY_PATH}") + target_link_libraries(fastvideo_kernel_ops PRIVATE "${TORCH_PYTHON_LIBRARY_PATH}") + else() + message(WARNING "Could not locate libtorch_python; fastvideo_kernel_ops may fail to import.") + endif() + + # Link CUDA runtime + driver explicitly (fixes missing symbols like cuGetErrorString at import time) + if(NOT GPU_BACKEND STREQUAL "ROCM") + target_link_libraries(fastvideo_kernel_ops PRIVATE CUDA::cudart CUDA::cuda_driver) + endif() + # We install it to fastvideo_kernel/_C so we can load it to register the ops install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C) endif() diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 262788f0b..a894ea733 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -34,7 +34,7 @@ from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_att out = sliding_tile_attention(q, k, v, window_sizes, text_len) # Example: Video Sparse Attention (with Triton fallback) -out = video_sparse_attn(q, k, v, block_sizes, topk=5) +out = video_sparse_attn(q, k, v, block_sizes, block_sizes, topk=5) # Example: VMoBA out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) diff --git a/fastvideo-kernel/csrc/attention/block_sparse_h100.cu b/fastvideo-kernel/csrc/attention/block_sparse_h100.cu index 58031322f..995bc43b2 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_h100.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_h100.cu @@ -639,7 +639,8 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { // store kq and vq - // ensuring all writes are finished + // ! the following two line seems unnecessary. + // tma::store_async_wait(); // ensure qg is finished __syncthreads(); warpgroup::store(kg_smem[0], kg_reg); @@ -660,145 +661,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { tma::store_async_wait(); } - -template -void block_sparse_attention_forward_impl( - bf16* d_q, bf16* d_k, bf16* d_v, float* d_l, bf16* d_o, - int batch, int qo_heads, int kv_heads, int seq_len, int hr, - int max_kv_blocks_per_q, - int32_t* q2k_block_sparse_index_ptr, - int32_t* q2k_block_sparse_num_ptr, - int32_t* block_size_ptr, - cudaStream_t stream -) { - using K = fwd_attend_ker_tile_dims; - using q_tile = st_bf; - using k_tile = st_bf; - using v_tile = st_bf; - using l_col_vec = col_vec>; - using o_tile = st_bf; - - using q_global = gl; - using k_global = gl; - using v_global = gl; - using l_global = gl; - using o_global = gl; - - using globals = fwd_globals; - - q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - 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), static_cast(D)}; - - 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), - q2k_block_sparse_index_ptr, q2k_block_sparse_num_ptr, block_size_ptr - }; - - // Shared memory size for the kernel - // 54000 bytes is calibrated for H100 shared memory constraints for these tile sizes - constexpr int mem_size = 54000; - - dim3 grid(seq_len/(64), qo_heads, batch); - - cudaFuncSetAttribute( - fwd_attend_ker, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - - fwd_attend_ker<<>>(g); -} - -template -void block_sparse_attention_backward_impl( - bf16* d_q, bf16* d_k, bf16* d_v, bf16* d_o, bf16* d_og, float* d_l, float* d_d, float* d_qg, float* d_kg, float* d_vg, - int batch, int qo_heads, int kv_heads, int seq_len, int hr, int max_q_blocks_per_kv, - int32_t* k2q_block_sparse_index_ptr, - int32_t* k2q_block_sparse_num_ptr, - int32_t* block_size_ptr, - cudaStream_t stream -) { - using G = bwd_attend_ker_tile_dims; - using og_tile = st_bf<4*16, D>; - using o_tile = st_bf<4*16, D>; - using d_tile = col_vec>; - - using og_global = gl; - using o_global = gl; - using d_global = gl; - - using prep_globals = bwd_prep_globals; - - constexpr int mem_size_prep = kittens::MAX_SHARED_MEMORY; - int threads_prep = PREP_NUM_WARPS * kittens::WARP_THREADS; - dim3 grid_bwd_prep(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); - - og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; - - prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; - - cudaFuncSetAttribute( - bwd_attend_prep_ker, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size_prep - ); - bwd_attend_prep_ker<<>>(bwd_g); - - using bwd_q_tile = st_bf; - using bwd_k_tile = st_bf; - using bwd_v_tile = st_bf; - using bwd_og_tile = st_bf; - using bwd_qg_tile = st_fl; - using bwd_kg_tile = st_fl; - using bwd_vg_tile = st_fl; - using bwd_l_tile = row_vec>; - using bwd_d_tile = row_vec>; - - using bwd_q_global = gl; - using bwd_k_global = gl; - using bwd_v_global = gl; - using bwd_og_global = gl; - using bwd_qg_global = gl; - using bwd_kg_global = gl; - using bwd_vg_global = gl; - using bwd_l_global = gl; - using bwd_d_global = gl; - - using bwd_global_args = bwd_globals; - - bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; - bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; - 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_global_args bwd_global{bwd_q_arg, bwd_k_arg, bwd_v_arg, bwd_og_arg, bwd_qg_arg, bwd_kg_arg, bwd_vg_arg, bwd_l_arg, bwd_d_arg, - static_cast(seq_len), static_cast(hr), static_cast(max_q_blocks_per_kv), - k2q_block_sparse_index_ptr, k2q_block_sparse_num_ptr, block_size_ptr}; - - dim3 grid_bwd_main(seq_len/64, qo_heads, batch); - int threads_main = 128; - // Calibrated shared memory sizes for different head dimensions - int bwd_mem_size = (D == 64) ? 72000 : 113000; - - cudaFuncSetAttribute( - bwd_attend_ker, - cudaFuncAttributeMaxDynamicSharedMemorySize, - bwd_mem_size - ); - bwd_attend_ker<<>>(bwd_global); -} - #include "pyutils/torch_helpers.cuh" #include #include @@ -810,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"); @@ -834,11 +705,9 @@ 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"); TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); @@ -864,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)); @@ -880,32 +749,110 @@ block_sparse_attention_forward( float* l_ptr = reinterpret_cast(l_vec.data_ptr()); float* d_l = reinterpret_cast(l_ptr); + //cudadevicesynchronize(); const c10::cuda::OptionalCUDAGuard device_guard(q.device()); const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - // Temporated implementation to avoid code duplication between head_dim=64 and 128 if (head_dim == 64) { - block_sparse_attention_forward_impl<64>( - d_q, d_k, d_v, d_l, d_o, - batch, qo_heads, kv_heads, seq_len, hr, - max_kv_blocks_per_q, - reinterpret_cast(q2k_block_sparse_index.data_ptr()), + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>; + + using q_global = gl; + using k_global = gl; + using v_global = gl; + using l_global = gl; + using o_global = gl; + + using globals = fwd_globals<64>; + + 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, + kg_arg, + vg_arg, + lg_arg, + og_arg, + 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()), - stream + reinterpret_cast(kv_block_size.data_ptr()) + }; + + constexpr int mem_size = 54000; + + dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); + + cudaFuncSetAttribute( + fwd_attend_ker<64>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size ); - } else if (head_dim == 128) { - block_sparse_attention_forward_impl<128>( - d_q, d_k, d_v, d_l, d_o, - batch, qo_heads, kv_heads, seq_len, hr, - max_kv_blocks_per_q, - reinterpret_cast(q2k_block_sparse_index.data_ptr()), + + fwd_attend_ker<64><<>>(g); + + CHECK_CUDA_ERROR(cudaGetLastError()); + // cudaStreamSynchronize(stream); + } + + if (head_dim == 128) { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + + using q_global = gl; + using k_global = gl; + using v_global = gl; + using l_global = gl; + using o_global = gl; + + using globals = fwd_globals<128>; + + 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, + kg_arg, + vg_arg, + lg_arg, + og_arg, + 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()), - stream + reinterpret_cast(kv_block_size.data_ptr()) + }; + + constexpr int mem_size = 54000; + + dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); + + cudaFuncSetAttribute( + fwd_attend_ker<128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size ); - } else { - TORCH_CHECK(false, "Unsupported head_dim: ", head_dim, ". Only 64 and 128 are supported."); + + fwd_attend_ker<128><<>>(g); + + CHECK_CUDA_ERROR(cudaGetLastError()); + // cudaStreamSynchronize(stream); } return {o, l_vec}; @@ -921,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); @@ -930,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"); @@ -945,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); @@ -988,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(); @@ -1030,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>; @@ -1043,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}; @@ -1082,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, @@ -1101,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(); @@ -1147,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}; @@ -1186,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, @@ -1205,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(); @@ -1233,4 +1187,4 @@ block_sparse_attention_backward(torch::Tensor q, return {qg, kg, vg}; //cudadevicesynchronize(); -} +} \ No newline at end of file diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py new file mode 100644 index 000000000..0ed935dbd --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +from typing import Tuple + +import torch + + +def _get_sm90_ops(): + try: + from fastvideo_kernel._C import fastvideo_kernel_ops # type: ignore + except Exception: + return None, None + return ( + getattr(fastvideo_kernel_ops, "block_sparse_fwd", None), + getattr(fastvideo_kernel_ops, "block_sparse_bwd", None), + ) + + +def _is_sm90() -> bool: + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability(0) + return major == 9 and minor == 0 + + +def _map_to_index_torch(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Pure-torch (no triton) conversion: + block_map: [B, H, Q, KV] bool (or [H, Q, KV] which will be treated as B=1) + returns: + index: [B, H, Q, KV] int32 (packed KV indices, -1 padding) + num: [B, H, Q] int32 (#kv blocks per q block) + """ + if block_map.dim() == 3: + block_map = block_map.unsqueeze(0) + if block_map.dim() != 4: + raise ValueError(f"block_map must be [B,H,Q,KV] (or [H,Q,KV]), got shape={tuple(block_map.shape)}") + if block_map.dtype != torch.bool: + block_map = block_map.to(torch.bool) + + B, H, Q, KV = block_map.shape + index = torch.full((B, H, Q, KV), -1, dtype=torch.int32, device=block_map.device) + num = torch.zeros((B, H, Q), dtype=torch.int32, device=block_map.device) + + # Small sizes in practice (B=1, H<=16, Q/KV<=64), so a Python loop is fine. + for b in range(B): + for h in range(H): + for q in range(Q): + kv_idx = torch.nonzero(block_map[b, h, q], as_tuple=False).flatten().to(torch.int32) + n = int(kv_idx.numel()) + if n: + index[b, h, q, :n] = kv_idx + num[b, h, q] = n + return index, num + + +@torch.library.custom_op( + "fastvideo_kernel::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.to(torch.bool) + q2k_idx, q2k_num = _map_to_index_torch(block_map) + + from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import + triton_block_sparse_attn_forward, + ) + + o, M = triton_block_sparse_attn_forward(q, k, v, q2k_idx, q2k_num, variable_block_sizes) + return o, M + + +@torch.library.register_fake("fastvideo_kernel::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]: + 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( + "fastvideo_kernel::block_sparse_attn_backward_triton", + mutates_args=(), + device_types="cuda", +) +def block_sparse_attn_backward_triton( + grad_output: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + M: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grad_output = grad_output.contiguous() + block_map = block_map.to(torch.bool) + q2k_idx, q2k_num = _map_to_index_torch(block_map) + k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous()) + + from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import ( # local import + triton_block_sparse_attn_backward, + ) + + dq, dk, dv = triton_block_sparse_attn_backward( + grad_output, q, k, v, o, M, q2k_idx, q2k_num, k2q_idx, k2q_num, variable_block_sizes + ) + return dq, dk, dv + + +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_triton") +def _block_sparse_attn_backward_triton_fake( + grad_output: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + M: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + return dq, dk, dv + + +def _backward_triton(ctx, grad_o, grad_M): + q, k, v, o, M, block_map, variable_block_sizes = ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, block_map, variable_block_sizes) + return dq, dk, dv, None, None + + +def _setup_context_triton(ctx, inputs, output): + q, k, v, block_map, variable_block_sizes = inputs + o, M = output + ctx.save_for_backward(q, k, v, o, M, block_map, variable_block_sizes) + + +block_sparse_attn_triton.register_autograd(_backward_triton, setup_context=_setup_context_triton) + + +@torch.library.custom_op( + "fastvideo_kernel::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]: + block_sparse_fwd, _ = _get_sm90_ops() + if block_sparse_fwd is None: + raise ImportError("fastvideo_kernel_ops.block_sparse_fwd is not available") + + q_padded = q_padded.contiguous() + k_padded = k_padded.contiguous() + v_padded = v_padded.contiguous() + block_map = block_map.to(torch.bool) + q2k_idx, q2k_num = _map_to_index_torch(block_map) + + o_padded, lse_padded = block_sparse_fwd( + q_padded, k_padded, v_padded, q2k_idx, q2k_num, variable_block_sizes.int() + ) + return o_padded, lse_padded + + +@torch.library.register_fake("fastvideo_kernel::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]: + o = torch.empty_like(q_padded) + lse = torch.empty((q_padded.shape[0], q_padded.shape[1], q_padded.shape[2], 1), device=q_padded.device, dtype=torch.float32) + return o, lse + + +@torch.library.custom_op( + "fastvideo_kernel::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]: + _, block_sparse_bwd = _get_sm90_ops() + if block_sparse_bwd is None: + raise ImportError("fastvideo_kernel_ops.block_sparse_bwd is not available") + + grad_output_padded = grad_output_padded.contiguous() + block_map = block_map.to(torch.bool) + k2q_idx, k2q_num = _map_to_index_torch(block_map.transpose(-1, -2).contiguous()) + + dq, dk, dv = block_sparse_bwd( + q_padded, + k_padded, + v_padded, + o_padded, + lse_padded, + grad_output_padded, + k2q_idx, + k2q_num, + variable_block_sizes.int(), + ) + # C++ kernel returns fp32 grads; cast back to match PyTorch convention if needed + return dq.to(grad_output_padded.dtype), dk.to(grad_output_padded.dtype), dv.to(grad_output_padded.dtype) + + +@torch.library.register_fake("fastvideo_kernel::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]: + dq = torch.empty_like(q_padded) + dk = torch.empty_like(k_padded) + dv = torch.empty_like(v_padded) + return dq, dk, dv + + +def _backward_sm90(ctx, grad_o, grad_lse): + q, k, v, o, lse, block_map, variable_block_sizes = ctx.saved_tensors + dq, dk, dv = block_sparse_attn_backward_sm90( + grad_o, q, k, v, o, lse, block_map, variable_block_sizes + ) + return dq, dk, dv, None, None + + +def _setup_context_sm90(ctx, inputs, output): + q, k, v, block_map, variable_block_sizes = inputs + o, lse = output + ctx.save_for_backward(q, k, v, o, lse, block_map, variable_block_sizes) + + +block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_context_sm90) + + +def block_sparse_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Unified block-sparse attention op with autograd support. + - On SM90 with compiled extension present: uses fastvideo_kernel_ops.block_sparse_fwd/bwd. + - Otherwise: uses Triton implementation (requires q/k/v to have same padded length today). + """ + block_sparse_fwd, block_sparse_bwd = _get_sm90_ops() + if _is_sm90() and (block_sparse_fwd is not None) and (block_sparse_bwd is not None): + return block_sparse_attn_sm90(q, k, v, block_map, variable_block_sizes) + # Triton path: generally assumes q/k/v share the same padded length + if q.shape[2] != k.shape[2] or q.shape[2] != v.shape[2]: + raise RuntimeError("Triton fallback requires q/k/v to have the same padded length.") + return block_sparse_attn_triton(q, k, v, block_map, variable_block_sizes) + + diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index 5dfe31e92..f3910c9b2 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -1,5 +1,6 @@ import math import torch +from .block_sparse_attn import block_sparse_attn from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward from .triton_kernels.st_attn_triton import sliding_tile_attention_triton from .triton_kernels.index import map_to_index @@ -62,6 +63,7 @@ def video_sparse_attn( k: torch.Tensor, v: torch.Tensor, variable_block_sizes: torch.Tensor, + q_variable_block_sizes: torch.Tensor, topk: int, block_size: int | tuple = 64, compress_attn_weight: torch.Tensor = None, @@ -70,14 +72,42 @@ def video_sparse_attn( block_size = (block_size, block_size, block_size) block_elements = block_size[0] * block_size[1] * block_size[2] - batch, heads, seq_len, dim = q.shape + batch, heads, q_seq_len, dim = q.shape + kv_seq_len = k.shape[2] + if v.shape[2] != kv_seq_len: + raise ValueError( + f"Expected k and v to have the same sequence length, got " + f"k.shape[2]={kv_seq_len}, v.shape[2]={v.shape[2]}" + ) + if k.shape[0] != batch or v.shape[0] != batch or k.shape[1] != heads or v.shape[1] != heads: + raise ValueError("Expected q/k/v to have the same batch and head dimensions.") + + if q_seq_len % block_elements != 0 or kv_seq_len % block_elements != 0: + raise ValueError( + f"q_seq_len and kv_seq_len must be divisible by block_elements={block_elements}, " + f"got q_seq_len={q_seq_len}, kv_seq_len={kv_seq_len}" + ) + q_num_blocks = q_seq_len // block_elements + kv_num_blocks = kv_seq_len // block_elements + + if variable_block_sizes.numel() != kv_num_blocks: + raise ValueError( + f"variable_block_sizes must have length kv_num_blocks={kv_num_blocks}, " + f"got {variable_block_sizes.numel()}" + ) + + if q_variable_block_sizes.numel() != q_num_blocks: + raise ValueError( + f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, " + f"got {q_variable_block_sizes.numel()}" + ) # Compression branch - q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim) - k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim) - v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim) + q_c = q.view(batch, heads, q_num_blocks, block_elements, dim) + k_c = k.view(batch, heads, kv_num_blocks, block_elements, dim) + v_c = v.view(batch, heads, kv_num_blocks, block_elements, dim) - q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to( + q_c = (q_c.float().sum(dim=3) / q_variable_block_sizes.view(1, 1, -1, 1)).to( q.dtype) k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to( k.dtype) @@ -88,9 +118,9 @@ def video_sparse_attn( attn = torch.softmax(scores, dim=-1) out_c = torch.matmul(attn, v_c) - out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim) + out_c = out_c.view(batch, heads, q_num_blocks, 1, dim) out_c = out_c.repeat(1, 1, 1, block_elements, - 1).view(batch, heads, seq_len, dim) + 1).view(batch, heads, q_seq_len, dim) # Sparse branch topk_idx = torch.topk(scores, topk, dim=-1).indices @@ -100,12 +130,17 @@ def video_sparse_attn( idx, num = map_to_index(mask) if block_sparse_fwd is not None: - out_s = block_sparse_fwd( - q, k, v, idx, num, variable_block_sizes.int() - )[0] # block_sparse_fwd returns vector + # Use autograd-enabled wrapper so backward works (and still uses SM90 kernel when available) + out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] else: - out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, - variable_block_sizes) + if q_seq_len != kv_seq_len: + raise RuntimeError( + "q/k have different lengths, but the compiled CUDA kernel (block_sparse_fwd) " + "is not available. The Triton fallback currently requires q and k/v to have " + "the same padded length." + ) + # Triton-only forward (kept for environments without the wrapper deps) + out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes) if compress_attn_weight is not None: return out_c * compress_attn_weight + out_s diff --git a/fastvideo-kernel/tests/test_vsa.py b/fastvideo-kernel/tests/test_vsa.py index b0f7c4626..0f7dd1a69 100644 --- a/fastvideo-kernel/tests/test_vsa.py +++ b/fastvideo-kernel/tests/test_vsa.py @@ -42,37 +42,13 @@ def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q 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) - # Use raw kernel or triton - try: - from fastvideo_kernel._C import fastvideo_kernel_ops - raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None) - except ImportError: - raw_kernel = None + # Use autograd-enabled wrapper (internally dispatches to SM90 kernel or Triton) + from fastvideo_kernel.block_sparse_attn import block_sparse_attn + output_padded, _aux = block_sparse_attn( + q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes + ) - from fastvideo_kernel.triton_kernels.index import map_to_index - - - # Convert mask to indices - # block_sparse_mask is [H, M, N] bool - # We need to map it to index. - # block_sparse_mask needs to be expanded/reshaped? - # generate_block_sparse_mask_for_function returns [H, NumBlocksQ, NumBlocksKV] - - # Ops.py logic: - # mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) - # idx, num = map_to_index(mask) - - idx, num = map_to_index(block_sparse_mask.unsqueeze(0)) # Add batch dim [1, H, M, N] - - if raw_kernel: - out_s = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int()) - output = out_s[0] - else: - # Fallback to triton testing if C++ not available - from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward - output, _ = triton_block_sparse_attn_forward(q_padded, k_padded, v_padded, idx, num, variable_block_sizes) - - output = output[:, :, q_non_pad_index, :] + output = output_padded[:, :, q_non_pad_index, :] output.backward(dO) return output, Q.grad, K.grad, V.grad @@ -264,7 +240,6 @@ def generate_error_graphs_qkdiff(h, d, error_mode='all'): print("-" * 150) -@pytest.mark.skip() def test_video_sparse_attention_backward(): if not torch.cuda.is_available(): return diff --git a/fastvideo-kernel/tests/test_vsa_forward.py b/fastvideo-kernel/tests/test_vsa_forward.py index f93fc779c..55ce8d14c 100644 --- a/fastvideo-kernel/tests/test_vsa_forward.py +++ b/fastvideo-kernel/tests/test_vsa_forward.py @@ -3,6 +3,7 @@ import sys from typing import Tuple import torch +import pytest from .utils import ( generate_block_sparse_mask_for_function, @@ -57,23 +58,14 @@ def block_sparse_forward_test( 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) - # Use raw kernel or triton + # Use autograd-enabled wrapper (internally dispatches SM90 C++ vs Triton) + from fastvideo_kernel.block_sparse_attn import block_sparse_attn try: - from fastvideo_kernel._C import fastvideo_kernel_ops - raw_kernel = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None) - except ImportError: - raw_kernel = None - - from fastvideo_kernel.triton_kernels.index import map_to_index - idx, num = map_to_index(block_sparse_mask) - - if raw_kernel: - out_padded = raw_kernel(q_padded, k_padded, v_padded, idx, num, variable_block_sizes.int())[0] - else: - from fastvideo_kernel.triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward - out_padded, _ = triton_block_sparse_attn_forward( - q_padded, k_padded, v_padded, idx, num, variable_block_sizes + out_padded, _aux = block_sparse_attn( + q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes ) + except RuntimeError as e: + pytest.skip(str(e)) # Remove padding on the query side out = out_padded[:, :, q_non_pad_index, :] @@ -156,11 +148,6 @@ def run_forward_qk_diff( ) -> 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" diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index 0adc56724..5e7209c41 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -276,8 +276,9 @@ class VideoSparseAttentionImpl(AttentionImpl): query, key, value, - variable_block_sizes=attn_metadata.variable_block_sizes, - topk=cur_topk, + attn_metadata.variable_block_sizes, + attn_metadata.variable_block_sizes, + cur_topk, block_size=VSA_TILE_SIZE, compress_attn_weight=gate_compress).transpose(1, 2)