nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[Bug Fix] Add autograd wrapper for block-sparse attention in fastvideo-kernel + fix CMake extension linking (#1015)

Browse Source
This commit is contained in:
alexzms
2026-01-16 21:16:43 -08:00
committed by GitHub
parent 0f53b51e6c
commit 560628709c
9 changed files with 558 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/attention/vsa/index.md
+3 -1
View File
@@ -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
)
```
fastvideo-kernel/CMakeLists.txt
+26
View File
@@ -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)
$<$<COMPILE_LANGUAGE:CUDA>:${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()
fastvideo-kernel/README.md
+1 -1
View File
@@ -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, ...)
fastvideo-kernel/csrc/attention/block_sparse_h100.cu
+174 -220
View File
@@ -639,7 +639,8 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> 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<D> g) {
tma::store_async_wait();
}
template<int D>
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<D>;
using q_tile = st_bf<K::qo_height, K::tile_width>;
using k_tile = st_bf<K::kv_height, K::tile_width>;
using v_tile = st_bf<K::kv_height, K::tile_width>;
using l_col_vec = col_vec<st_fl<K::qo_height, K::tile_width>>;
using o_tile = st_bf<K::qo_height, K::tile_width>;
using q_global = gl<bf16, -1, -1, -1, -1, q_tile>;
using k_global = gl<bf16, -1, -1, -1, -1, k_tile>;
using v_global = gl<bf16, -1, -1, -1, -1, v_tile>;
using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
using o_global = gl<bf16, -1, -1, -1, -1, o_tile>;
using globals = fwd_globals<D>;
q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
globals g{
qg_arg, kg_arg, vg_arg, lg_arg, og_arg,
static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(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<D>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
mem_size
);
fwd_attend_ker<D><<<grid, (128), mem_size, stream>>>(g);
}
template<int D>
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<D>;
using og_tile = st_bf<4*16, D>;
using o_tile = st_bf<4*16, D>;
using d_tile = col_vec<st_fl<4*16, D>>;
using og_global = gl<bf16, -1, -1, -1, -1, og_tile>;
using o_global = gl<bf16, -1, -1, -1, -1, o_tile>;
using d_global = gl<float, -1, -1, -1, -1, d_tile>;
using prep_globals = bwd_prep_globals<D>;
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<bf16>*4), qo_heads, batch);
og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
o_global prep_o_arg {d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
d_global prep_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg};
cudaFuncSetAttribute(
bwd_attend_prep_ker<D>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
mem_size_prep
);
bwd_attend_prep_ker<D><<<grid_bwd_prep, threads_prep, mem_size_prep, stream>>>(bwd_g);
using bwd_q_tile = st_bf<G::tile_h_qo, G::tile_width>;
using bwd_k_tile = st_bf<G::tile_h, G::tile_width>;
using bwd_v_tile = st_bf<G::tile_h, G::tile_width>;
using bwd_og_tile = st_bf<G::tile_h_qo, G::tile_width>;
using bwd_qg_tile = st_fl<G::tile_h_qo, G::tile_width>;
using bwd_kg_tile = st_fl<G::tile_h, G::tile_width>;
using bwd_vg_tile = st_fl<G::tile_h, G::tile_width>;
using bwd_l_tile = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
using bwd_d_tile = row_vec<st_fl<G::tile_h_qo, G::tile_h>>;
using bwd_q_global = gl<bf16, -1, -1, -1, -1, bwd_q_tile>;
using bwd_k_global = gl<bf16, -1, -1, -1, -1, bwd_k_tile>;
using bwd_v_global = gl<bf16, -1, -1, -1, -1, bwd_v_tile>;
using bwd_og_global = gl<bf16, -1, -1, -1, -1, bwd_og_tile>;
using bwd_qg_global = gl<float, -1, -1, -1, -1, bwd_qg_tile>;
using bwd_kg_global = gl<float, -1, -1, -1, -1, bwd_kg_tile>;
using bwd_vg_global = gl<float, -1, -1, -1, -1, bwd_vg_tile>;
using bwd_l_global = gl<float, -1, -1, -1, -1, bwd_l_tile>;
using bwd_d_global = gl<float, -1, -1, -1, -1, bwd_d_tile>;
using bwd_global_args = bwd_globals<D>;
bwd_q_global bwd_q_arg {d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_k_global bwd_k_arg {d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_v_global bwd_v_arg {d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), static_cast<uint32_t>(D)};
bwd_l_global bwd_l_arg {d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
bwd_d_global bwd_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(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<int>(seq_len), static_cast<int>(hr), static_cast<int>(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<D>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
bwd_mem_size
);
bwd_attend_ker<D><<<grid_bwd_main, threads_main, bwd_mem_size, stream>>>(bwd_global);
}
#include "pyutils/torch_helpers.cuh"
#include <ATen/cuda/CUDAContext.h>
#include <iostream>
@@ -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<const uint>(batch),
static_cast<const uint>(qo_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(q_seq_len),
static_cast<const uint>(head_dim)}, v.options());
torch::Tensor l_vec = torch::empty({static_cast<const uint>(batch),
static_cast<const uint>(qo_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(q_seq_len),
static_cast<const uint>(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<float*>(l_vec.data_ptr<float>());
float* d_l = reinterpret_cast<float*>(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<int32_t*>(q2k_block_sparse_index.data_ptr()),
using q_tile = st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
using k_tile = st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
using v_tile = st_bf<fwd_attend_ker_tile_dims<64>::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>;
using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>>;
using o_tile = st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
using q_global = gl<bf16, -1, -1, -1, -1, q_tile>;
using k_global = gl<bf16, -1, -1, -1, -1, k_tile>;
using v_global = gl<bf16, -1, -1, -1, -1, v_tile>;
using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
using o_global = gl<bf16, -1, -1, -1, -1, o_tile>;
using globals = fwd_globals<64>;
q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(q_seq_len)};
o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
globals g{
qg_arg,
kg_arg,
vg_arg,
lg_arg,
og_arg,
static_cast<int>(q_seq_len),
static_cast<int>(hr),
static_cast<int>(max_kv_blocks_per_q),
reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
reinterpret_cast<int32_t*>(block_size.data_ptr()),
stream
reinterpret_cast<int32_t*>(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<int32_t*>(q2k_block_sparse_index.data_ptr()),
fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
CHECK_CUDA_ERROR(cudaGetLastError());
// cudaStreamSynchronize(stream);
}
if (head_dim == 128) {
using q_tile = st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
using k_tile = st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
using v_tile = st_bf<fwd_attend_ker_tile_dims<128>::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>;
using l_col_vec = col_vec<st_fl<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>;
using o_tile = st_bf<fwd_attend_ker_tile_dims<128>::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>;
using q_global = gl<bf16, -1, -1, -1, -1, q_tile>;
using k_global = gl<bf16, -1, -1, -1, -1, k_tile>;
using v_global = gl<bf16, -1, -1, -1, -1, v_tile>;
using l_global = gl<float, -1, -1, -1, -1, l_col_vec>;
using o_global = gl<bf16, -1, -1, -1, -1, o_tile>;
using globals = fwd_globals<128>;
q_global qg_arg{d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
k_global kg_arg{d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
v_global vg_arg{d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
l_global lg_arg{d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(q_seq_len)};
o_global og_arg{d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
globals g{
qg_arg,
kg_arg,
vg_arg,
lg_arg,
og_arg,
static_cast<int>(q_seq_len),
static_cast<int>(hr),
static_cast<int>(max_kv_blocks_per_q),
reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()),
reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr()),
reinterpret_cast<int32_t*>(block_size.data_ptr()),
stream
reinterpret_cast<int32_t*>(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><<<grid, (128), mem_size, stream>>>(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<const uint>(batch),
static_cast<const uint>(qo_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(q_seq_len),
static_cast<const uint>(head_dim)}, l_vec.options());
torch::Tensor kg = torch::zeros({static_cast<const uint>(batch),
static_cast<const uint>(kv_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(kv_seq_len),
static_cast<const uint>(head_dim)}, l_vec.options());
torch::Tensor vg = torch::zeros({static_cast<const uint>(batch),
static_cast<const uint>(kv_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(kv_seq_len),
static_cast<const uint>(head_dim)}, l_vec.options());
torch::Tensor d_vec = torch::empty({static_cast<const uint>(batch),
static_cast<const uint>(qo_heads),
static_cast<const uint>(seq_len),
static_cast<const uint>(q_seq_len),
static_cast<const uint>(1)}, l_vec.options());
float* qg_ptr = qg.data_ptr<float>();
@@ -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<bf16>*4), qo_heads, batch);
dim3 grid_bwd(q_seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM<bf16>*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<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
o_global prep_o_arg {d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
d_global prep_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
o_global prep_o_arg {d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
d_global prep_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(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<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_k_global bwd_k_arg {d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_v_global bwd_v_arg {d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 64U};
bwd_l_global bwd_l_arg {d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
bwd_d_global bwd_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
bwd_q_global bwd_q_arg {d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
bwd_k_global bwd_k_arg {d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
bwd_v_global bwd_v_arg {d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 64U};
bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 64U};
bwd_l_global bwd_l_arg {d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(q_seq_len)};
bwd_d_global bwd_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(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<int>(seq_len),
static_cast<int>(kv_seq_len), // N is not used in the kernel
static_cast<int>(hr),
static_cast<int>(max_q_blocks_per_kv),
reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
reinterpret_cast<int32_t*>(block_size.data_ptr())};
reinterpret_cast<int32_t*>(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<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
o_global prep_o_arg {d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
d_global prep_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
og_global prep_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
o_global prep_o_arg {d_o, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
d_global prep_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(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<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_k_global bwd_k_arg {d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_v_global bwd_v_arg {d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(seq_len), 128U};
bwd_l_global bwd_l_arg {d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
bwd_d_global bwd_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(seq_len)};
bwd_q_global bwd_q_arg {d_q, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
bwd_k_global bwd_k_arg {d_k, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
bwd_v_global bwd_v_arg {d_v, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
bwd_og_global bwd_og_arg{d_og, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
bwd_qg_global bwd_qg_arg{d_qg, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), static_cast<unsigned int>(q_seq_len), 128U};
bwd_kg_global bwd_kg_arg{d_kg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
bwd_vg_global bwd_vg_arg{d_vg, static_cast<unsigned int>(batch), static_cast<unsigned int>(kv_heads), static_cast<unsigned int>(kv_seq_len), 128U};
bwd_l_global bwd_l_arg {d_l, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(q_seq_len)};
bwd_d_global bwd_d_arg {d_d, static_cast<unsigned int>(batch), static_cast<unsigned int>(qo_heads), 1U, static_cast<unsigned int>(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<int>(seq_len),
static_cast<int>(kv_seq_len), // N is not used in the kernel
static_cast<int>(hr),
static_cast<int>(max_q_blocks_per_kv),
reinterpret_cast<int32_t*>(k2q_block_sparse_index.data_ptr()),
reinterpret_cast<int32_t*>(k2q_block_sparse_num.data_ptr()),
reinterpret_cast<int32_t*>(block_size.data_ptr())};
reinterpret_cast<int32_t*>(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();
}
}
fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py
+291
View File
@@ -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)
fastvideo-kernel/python/fastvideo_kernel/ops.py
+47 -12
View File
@@ -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<Tensor>
# 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
fastvideo-kernel/tests/test_vsa.py
+6 -31
View File
@@ -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
fastvideo-kernel/tests/test_vsa_forward.py
+7 -20
View File
@@ -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"
fastvideo/attention/backends/video_sparse_attn.py
+3 -2
View File
@@ -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)