rootandClaude Fable 5.1
749a9c5d97
Native GB200 (sm_100a) backward pairing with the sm_100a forward: one warp-specialized tcgen05 kernel per (batch, head, kv64 block) computing dK, dV and the dQ partials (fp32 accumulator, cp.reduce.async.bulk), plus a preprocess (Delta, Q^T, dO^T, dqaccum zero, exact-zero dK/dV rows for unselected kv blocks) and a postprocess (dQ unscramble + sm_scale). Consumes the forward's lse (Triton M format) and invert_indices' k2q metadata unchanged; returns (dq, dk, dv) in bf16 with the Triton backward's scaling. Wiring: with FASTVIDEO_VSA_SM100A=1 the autograd backward of the sm_100a forward op runs this kernel when block_sparse_attn_bwd_sm100a.is_supported passes (bf16, head_dim 128, 64-token blocks, even block count, sm_100a device, op built) and the Triton backward otherwise. No new environment variable. primitives.cuh gains only the backward's primitives (additions only; the forward's definitions and SASS are untouched). Built for sm_100a only: the kernel is validated on GB200, not yet on B300/GB300, so sm_103a devices keep the Triton backward. The H3 backend's grad gate is unchanged (separate PR). Semantics: per-row k2q counts may differ arbitrarily (one work item is one kv row; count 0 is skipped identically by every warp and gets zero dK/dV rows); variable_block_sizes masks padded kv rows per lane; is_supported reads no tensor contents. Work-item order: a device kernel sorts by list length from 1024 kv blocks per sequence on; below that a cached identity array is passed. Tests (GB200, extension built from this branch): tests/test_block_sparse_bwd_sm100a.py 15 passed (fp32 masked-dense autograd reference; ragged vbs, zero-count kv blocks, top-k 1/2/3/5/7, batch 2), tests/test_block_sparse_sm100a_dispatch.py 10 passed (sm_100a route vs all-Triton grads, ragged vbs, Triton backward monkeypatched to fail on a supported input), tests/test_block_sparse_sm100a.py 37 passed (forward regression). Perf (kernels only, same window, B=1 H=8 D=128, 25% density, fp32 accumulator, TFLOPS on selected blocks) Triton backward vs this kernel: 4k 198/381 (1.93x), 8k 313/575 (1.84x), 16k 389/715 (1.84x), 32k 433/794 (1.84x), 65k 448/821 (1.83x), 131k 452/736 (1.63x), 262k 452/784 (1.73x), 524k 437/771 (1.76x). FastVideo autograd path incl. invert_indices: 1.6-1.7x at Wan 480P/720P-like shapes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
87 lines
3.5 KiB
C++
87 lines
3.5 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_SM100A
|
|
extern std::vector<torch::Tensor> block_sparse_sm100a_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_sm100a_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_sm100a_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
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.doc() = "FastVideo CUDA Kernels";
|
|
|
|
#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A
|
|
m.def("block_sparse_sm100a_fwd",
|
|
torch::wrap_pybind_function(block_sparse_sm100a_fwd),
|
|
"VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a/sm103a)");
|
|
m.def("block_sparse_sm100a_blk128_fwd",
|
|
torch::wrap_pybind_function(block_sparse_sm100a_blk128_fwd),
|
|
"VSA block-sparse attention forward, 128-token blocks (Blackwell sm100a/sm103a)");
|
|
m.def("block_sparse_sm100a_bwd",
|
|
torch::wrap_pybind_function(block_sparse_sm100a_bwd),
|
|
"VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a)");
|
|
#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);
|
|
}
|