SolitaryThinker
32d8a8765b
FastH3's tile-128 route runs the ptx_blackwell CUDA forward on sm_100/sm_103 and, on every other GPU,
splits each 128-token tile into two 64-token halves on the Triton kernels. On the DGX Spark that split
route is what OmniRef's BF16 attention runs today. This adds a 128-token forward built for sm_121a
(mma.sync + TMA; the GB10 has no tcgen05 or tensor memory) and routes the tile-128 call to it on a
GB10, ahead of the split.
- csrc/attention/block_sparse_kernel_blk128_sm121.cuh and block_sparse_blk128_sm121.cu: the kernel and
its binding block_sparse_sm121_blk128_fwd(q, k, v, q2k_idx, q2k_num, variable_block_sizes) -> {out},
inference only, no LSE.
- CMakeLists.txt: built when TORCH_CUDA_ARCH_LIST has 12.1a or 12.1 and CUDA >= 13, so both the
aarch64 wheel and the GB10 Docker image carry it.
- fastvideo_kernel.block_sparse_attn_sm121: is_supported and an opaque custom op with a fake. It is a
separate module because block_sparse_attn_ptx_blackwell.is_supported also gates an LSE-requiring
autograd path.
- video_sparse_attn_h3: the tile-128 route tries it after ptx_blackwell and before split_tile128.
- Tests: kernel correctness against an FP32 reference (ragged blocks, empty selections, odd block
counts, many waves, a NaN-filled allocator, determinism), and the H3 route (fakes, plus a real-GB10
oracle test).
106 lines
4.6 KiB
C++
106 lines
4.6 KiB
C++
#include <torch/extension.h>
|
|
#include <ATen/ATen.h>
|
|
#include <vector>
|
|
|
|
// Forward declarations
|
|
#ifdef TK_COMPILE_ST_ATTN
|
|
extern torch::Tensor sta_forward(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
|
|
int kernel_t_size, int kernel_w_size, int kernel_h_size,
|
|
int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
|
|
);
|
|
#endif
|
|
|
|
#ifdef TK_COMPILE_BLOCK_SPARSE
|
|
extern std::vector<torch::Tensor> block_sparse_attention_forward(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v,
|
|
torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size
|
|
);
|
|
extern std::vector<torch::Tensor> block_sparse_attention_backward(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og,
|
|
torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size
|
|
);
|
|
#endif
|
|
|
|
#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
|
|
// Ulysses sequence-parallel all-to-all (csrc/comm/)
|
|
void register_ulysses_a2a(pybind11::module_ &);
|
|
#endif
|
|
|
|
// TurboDiffusion kernels
|
|
void register_quant(pybind11::module_ &);
|
|
void register_rms_norm(pybind11::module_ &);
|
|
void register_layer_norm(pybind11::module_ &);
|
|
void register_gemm(pybind11::module_ &);
|
|
|
|
#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL
|
|
extern std::vector<torch::Tensor> block_sparse_ptx_blackwell_fwd(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
|
|
torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
|
|
double sm_scale, bool need_lse);
|
|
extern std::vector<torch::Tensor> block_sparse_ptx_blackwell_blk128_fwd(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional<torch::Tensor> v_t,
|
|
torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes,
|
|
double sm_scale, bool need_lse);
|
|
extern std::vector<torch::Tensor> block_sparse_ptx_blackwell_bwd(
|
|
torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
|
|
torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num,
|
|
torch::Tensor variable_block_sizes, double sm_scale);
|
|
extern std::vector<torch::Tensor> block_sparse_ptx_blackwell_blk128_bwd(
|
|
torch::Tensor grad_o, torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
|
|
torch::Tensor lse, torch::Tensor k2q_idx, torch::Tensor k2q_num,
|
|
torch::Tensor variable_block_sizes, double sm_scale);
|
|
#endif
|
|
|
|
#ifdef FASTVIDEO_KERNEL_COMPILE_BLOCK_SPARSE_VSA_SM121
|
|
extern std::vector<torch::Tensor> block_sparse_sm121_blk128_fwd(
|
|
torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_idx, torch::Tensor q2k_num,
|
|
torch::Tensor variable_block_sizes);
|
|
#endif
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.doc() = "FastVideo CUDA Kernels";
|
|
|
|
#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_PTX_BLACKWELL
|
|
m.def("block_sparse_ptx_blackwell_fwd",
|
|
torch::wrap_pybind_function(block_sparse_ptx_blackwell_fwd),
|
|
"VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a/sm103a)");
|
|
m.def("block_sparse_ptx_blackwell_blk128_fwd",
|
|
torch::wrap_pybind_function(block_sparse_ptx_blackwell_blk128_fwd),
|
|
"VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)");
|
|
m.def("block_sparse_ptx_blackwell_bwd",
|
|
torch::wrap_pybind_function(block_sparse_ptx_blackwell_bwd),
|
|
"VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)");
|
|
m.def("block_sparse_ptx_blackwell_blk128_bwd",
|
|
torch::wrap_pybind_function(block_sparse_ptx_blackwell_blk128_bwd),
|
|
"VSA block-sparse attention backward, 128-token blocks (Blackwell sm100a/sm103a)");
|
|
#endif
|
|
|
|
#ifdef FASTVIDEO_KERNEL_COMPILE_BLOCK_SPARSE_VSA_SM121
|
|
m.def("block_sparse_sm121_blk128_fwd",
|
|
torch::wrap_pybind_function(block_sparse_sm121_blk128_fwd),
|
|
"VSA block-sparse attention forward, 128-token blocks, inference only (GB10 / DGX Spark, sm121a)");
|
|
#endif
|
|
|
|
#ifdef TK_COMPILE_ST_ATTN
|
|
|
|
m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)");
|
|
#endif
|
|
|
|
#ifdef TK_COMPILE_BLOCK_SPARSE
|
|
m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention forward (Hopper)");
|
|
m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)");
|
|
#endif
|
|
|
|
#ifdef FASTVIDEO_KERNEL_COMPILE_ULYSSES_A2A
|
|
// Ulysses sequence-parallel all-to-all
|
|
register_ulysses_a2a(m);
|
|
#endif
|
|
|
|
// TurboDiffusion
|
|
register_quant(m);
|
|
register_rms_norm(m);
|
|
register_layer_norm(m);
|
|
register_gemm(m);
|
|
}
|