From e1f390479919e156b30b8611d856ea340c0577dd Mon Sep 17 00:00:00 2001 From: lpc0220 Date: Tue, 29 Sep 2026 16:57:22 -0700 Subject: [PATCH] [kernel] sm_100a CUDA backward for VSA block-sparse attention, 128-token blocks (#1866) Co-authored-by: Pengcheng Li Co-authored-by: Claude Fable 5.1 --- fastvideo-kernel/CMakeLists.txt | 8 +- .../block_sparse_bwd_blk128_sm100a.cu | 7 + .../block_sparse_bwd_kernel_sm100a_blk128.cuh | 1006 +++++++++++++++++ .../block_sparse_bwd_launch_sm100a_blk128.cuh | 231 ++++ .../csrc/attention/block_sparse_bwd_sm100a.cu | 44 +- .../csrc/attention/primitives.cuh | 27 + fastvideo-kernel/csrc/common_extension.cpp | 7 + .../fastvideo_kernel/block_sparse_attn.py | 43 +- .../block_sparse_attn_bwd_sm100a.py | 73 +- .../block_sparse_attn_sm100a.py | 5 +- .../tests}/test_block_sparse_bwd_sm100a.py | 118 +- .../test_block_sparse_sm100a_dispatch.py | 53 +- .../backends/video_sparse_attn_h3.py | 2 +- .../attention/test_vsa_h3_sm100a_route.py | 2 +- 14 files changed, 1508 insertions(+), 118 deletions(-) create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh create mode 100644 fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh rename {tests => fastvideo-kernel/tests}/test_block_sparse_bwd_sm100a.py (69%) rename {tests => fastvideo-kernel/tests}/test_block_sparse_sm100a_dispatch.py (80%) diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 1fc50a3f1..ff09788c1 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -464,10 +464,12 @@ if(BUILD_CXX_KERNELS) csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() - # VSA block-sparse attention BACKWARD, 64-token blocks: the same images as the forward. + # VSA block-sparse attention BACKWARD, 64- and 128-token blocks: the same images as the forward. if(_VSA_FWD_GENCODE) - list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu) - set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu + csrc/attention/block_sparse_bwd_blk128_sm100a.cu) + set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu + csrc/attention/block_sparse_bwd_blk128_sm100a.cu PROPERTIES COMPILE_OPTIONS "${_VSA_FWD_GENCODE};-DVSA_BHSD=true") endif() diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu new file mode 100644 index 000000000..2d04613a1 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_blk128_sm100a.cu @@ -0,0 +1,7 @@ +// block_sparse_bwd_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding. +// +// Same source as block_sparse_bwd_sm100a.cu with VSA_BLK128 set: the kernel and launch land in +// namespace vsa_bwd_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the +// exported entry point becomes block_sparse_sm100a_blk128_bwd. +#define VSA_BLK128 true +#include "block_sparse_bwd_sm100a.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh new file mode 100644 index 000000000..1ca9bcf4e --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a_blk128.cuh @@ -0,0 +1,1006 @@ +// block_sparse_bwd_kernel_sm100a_blk128.cuh -- VSA block-sparse attention BACKWARD, +// 128-token blocks, sm_100a (one CTA per kv block). Warp-specialized: load / +// MMA (tcgen05) / softmax (P^T, dS^T) / epilogue (dQ drain). Three kernels: preprocess +// (Delta, dqaccum zero), main (dK, dV, dQ partials), postprocess (dQ unscramble + scale). +#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_BLK128_CUH + +#include +#include +#include +#include +#include +#include +#include "primitives.cuh" + +namespace vsa_bwd_blk128 { + +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif + +enum class Sched { NON_PERSISTENT, STATIC_PERSISTENT, CLC }; +#ifndef KERNEL_SCHED +#define KERNEL_SCHED Sched::NON_PERSISTENT +#endif + +#ifndef KERNEL_PDL +#define KERNEL_PDL true +#endif + +constexpr int BLOCK = 128; +constexpr int KV_TILE = BLOCK; +constexpr int Q_TILE = BLOCK; +constexpr int HEAD_DIM = 128; +constexpr int SUB_COLS_BF16 = 64; +constexpr int SUB_COLS_BYTES = SUB_COLS_BF16 * (int)sizeof(__nv_bfloat16); +constexpr int HD_SUBTILES = HEAD_DIM / SUB_COLS_BF16; +constexpr int KV_SUB_COLS_BYTES = KV_TILE * SUB_COLS_BYTES; +constexpr int KV_TILE_BYTES = HD_SUBTILES * KV_SUB_COLS_BYTES; +constexpr int Q_SUB_COLS_BYTES = Q_TILE * SUB_COLS_BYTES; +constexpr int Q_TILE_BYTES = HD_SUBTILES * Q_SUB_COLS_BYTES; +constexpr int NUM_Q_STAGES = 2; +constexpr int DST_TILE_BYTES = KV_TILE * SUB_COLS_BYTES; +constexpr int DST_BYTES = 2 * DST_TILE_BYTES; +constexpr int MMA_K = 16; +constexpr int K_ATOMS_PER_SUBTILE = SUB_COLS_BF16 / MMA_K; +constexpr int K_ATOMS_PER_Q_HALF = SUB_COLS_BF16 / MMA_K; +constexpr int K_ATOMS_PER_KV_TILE = KV_TILE / MMA_K; +constexpr int BF16X2_COLS_PER_K16 = MMA_K / 2; + +struct DQConfig { + static constexpr int COLS = 32; + + static constexpr int DQ_ONE_PUSH_BYTES = Q_TILE * COLS * (int)sizeof(float); + static constexpr int DQ_STAGE_BYTES = DQ_ONE_PUSH_BYTES; + static constexpr int DQ_STAGE_BUFFERS = 2; + static constexpr int DQ_BLOCK_ELEMS = Q_TILE * HEAD_DIM; +}; +using DQ = DQConfig; + +constexpr int N_WARPS = 16; +constexpr int W_EPI0 = 0, W_SOFTMAX0 = 4, W_MMA = 12, W_LOAD = 13; +constexpr int W_SCHED = 14; +constexpr int CLC_STAGES = 2; +constexpr int CLC_ARRIVALS = 15; + +constexpr int NUM_BARS = 4 * NUM_Q_STAGES + 14 + 2 * CLC_STAGES; +constexpr int SMEM_TOTAL = 2 * KV_TILE_BYTES + NUM_Q_STAGES * Q_TILE_BYTES + Q_TILE_BYTES + + DST_BYTES + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES + + (NUM_Q_STAGES + 1) * Q_TILE * (int)sizeof(float) + NUM_BARS * 8 + + CLC_STAGES * 16 + 48; + +constexpr int ST_COLS = Q_TILE; +constexpr int ST_HALF_COLS = SUB_COLS_BF16; +constexpr int DV_COLS = HEAD_DIM; +constexpr int DK_COLS = HEAD_DIM; +constexpr int TMEM_TOTAL = ST_COLS + DV_COLS + ST_COLS + DK_COLS; +static_assert(ST_COLS == 2 * ST_HALF_COLS, "two q-halves per tile"); +static_assert(TMEM_TOTAL == 512, "TMEM map must fill exactly 512 columns"); + +extern __shared__ __align__(1024) uint8_t bwd_smem[]; + +template +__device__ __forceinline__ size_t token_offset(int batch, int head, int num_heads, int seqlen, + int t) { + if constexpr (BHSD) + return ((size_t)(batch * num_heads + head) * seqlen + t) * HEAD_DIM; + else + return ((size_t)(batch * seqlen + t) * num_heads + head) * HEAD_DIM; +} + +struct WorkItem { + int batch; + int head; + int kv_block_id_in_seq; + int batch_head; + const int* local_k2q_idx; + + int local_k2q_num; +}; + +__device__ __forceinline__ WorkItem decode_workitem(int batch_head, int kv_block_id_in_seq, + const int* __restrict__ k2q_idx, + const int* __restrict__ k2q_num, + int max_q_blocks, int num_heads, + int num_kv_blocks_per_seq) { + WorkItem it; + it.batch_head = batch_head; + it.batch = batch_head / num_heads; + it.head = batch_head % num_heads; + it.kv_block_id_in_seq = kv_block_id_in_seq; + const int item = batch_head * num_kv_blocks_per_seq + kv_block_id_in_seq; + it.local_k2q_idx = k2q_idx + (size_t)item * (size_t)max_q_blocks; + it.local_k2q_num = k2q_num[item]; + return it; +} + +template +__global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) + vsa_bwd_main_kernel(const __grid_constant__ CUtensorMap tmap_q, + const __grid_constant__ CUtensorMap tmap_k, + const __grid_constant__ CUtensorMap tmap_v, + const __grid_constant__ CUtensorMap tmap_do, + const __grid_constant__ CUtensorMap tmap_dk, + const __grid_constant__ CUtensorMap tmap_dv, float* __restrict__ dqaccum, + const float* __restrict__ lse_rows, const float* __restrict__ delta_rows, + const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num, + int max_q_blocks, const int* __restrict__ workitem_remap, + const int* __restrict__ variable_block_sizes, int num_samples, + int num_heads, int seqlen, float scale_log2, float sm_scale) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + constexpr bool PERSISTENT = SCHED != Sched::NON_PERSISTENT; + constexpr bool CLC = SCHED == Sched::CLC; + const int num_kv_blocks_per_seq = seqlen / BLOCK; + [[maybe_unused]] const int total_workitems = num_samples * num_heads * num_kv_blocks_per_seq; + uint8_t* sK = bwd_smem; + uint8_t* sV = sK + KV_TILE_BYTES; + uint8_t* sQ[NUM_Q_STAGES] = {sV + KV_TILE_BYTES, sV + KV_TILE_BYTES + Q_TILE_BYTES}; + uint8_t* sDO = sQ[0] + NUM_Q_STAGES * Q_TILE_BYTES; + __nv_bfloat16* sDST = reinterpret_cast<__nv_bfloat16*>(sDO + Q_TILE_BYTES); + uint8_t* sDQ_STAGE_bytes = sDO + Q_TILE_BYTES + DST_BYTES; + float* sDQ_STAGE[DQ::DQ_STAGE_BUFFERS] = { + reinterpret_cast(sDQ_STAGE_bytes), + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BYTES)}; + float* sLSE = + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES); + float* sDelta = sLSE + NUM_Q_STAGES * Q_TILE; + + uint64_t* full_bar_q = reinterpret_cast(sDelta + Q_TILE); + uint64_t* empty_bar_q = full_bar_q + NUM_Q_STAGES; + uint64_t* full_bar_do = empty_bar_q + NUM_Q_STAGES; + uint64_t* empty_bar_do = full_bar_do + 1; + uint64_t* full_bar_lse = empty_bar_do + 1; + uint64_t* empty_bar_lse = full_bar_lse + NUM_Q_STAGES; + uint64_t* full_bar_delta = empty_bar_lse + NUM_Q_STAGES; + uint64_t* empty_bar_delta = full_bar_delta + 1; + uint64_t* full_bar_st = empty_bar_delta + 1; + uint64_t* full_bar_dpt = full_bar_st + 1; + uint64_t* full_bar_pt = full_bar_dpt + 1; + uint64_t* full_bar_dst = full_bar_pt + 1; + uint64_t* full_bar_dq = full_bar_dst + 1; + uint64_t* empty_bar_dq = full_bar_dq + 1; + uint64_t* full_bar_dv = empty_bar_dq + 1; + uint64_t* full_bar_dk = full_bar_dv + 1; + + uint64_t* empty_bar_kv = full_bar_dk + 1; + uint64_t* empty_bar_epi = empty_bar_kv + 1; + uint64_t* clc_full = empty_bar_epi + 1; + uint64_t* clc_empty = clc_full + CLC_STAGES; + uint32_t* clc_response = reinterpret_cast( + (reinterpret_cast(clc_empty + CLC_STAGES) + 15u) & ~uintptr_t(15u)); + uint32_t* tmem_slot = clc_response + CLC_STAGES * 4; + + const int tid = threadIdx.x, warp_id = tid >> 5, lane = tid & 31; + + if (tid == 0) { + #pragma unroll + for (int s = 0; s < NUM_Q_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar_q[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_q[s]), 1); + } + mbarrier_init(smem_ptr_u32(full_bar_do), 1); + mbarrier_init(smem_ptr_u32(empty_bar_do), 1); + #pragma unroll + for (int s = 0; s < NUM_Q_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar_lse[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_lse[s]), 8); + } + mbarrier_init(smem_ptr_u32(full_bar_delta), 1); + mbarrier_init(smem_ptr_u32(empty_bar_delta), 8); + mbarrier_init(smem_ptr_u32(full_bar_st), 1); + mbarrier_init(smem_ptr_u32(full_bar_dpt), 1); + mbarrier_init(smem_ptr_u32(full_bar_pt), 8); + mbarrier_init(smem_ptr_u32(full_bar_dst), 8); + mbarrier_init(smem_ptr_u32(full_bar_dq), 1); + mbarrier_init(smem_ptr_u32(empty_bar_dq), 4); + mbarrier_init(smem_ptr_u32(full_bar_dv), 1); + mbarrier_init(smem_ptr_u32(full_bar_dk), 1); + if constexpr (PERSISTENT) { + mbarrier_init(smem_ptr_u32(empty_bar_kv), 1); + mbarrier_init(smem_ptr_u32(empty_bar_epi), 8 * 32); + } + if constexpr (CLC) { + #pragma unroll + for (int st = 0; st < CLC_STAGES; ++st) { + mbarrier_init(smem_ptr_u32(&clc_full[st]), 1); + mbarrier_init(smem_ptr_u32(&clc_empty[st]), CLC_ARRIVALS); + } + #pragma unroll + for (int i = 0; i < CLC_STAGES * 4; ++i) clc_response[i] = 0; + } + } + fence_mbarrier_init_release_cluster(); + __syncthreads(); + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + + if (warp_id == W_LOAD) { + setmaxnreg_dec<88>(); + + EmptyPhaseTracker q_empty_ph, lse_empty_ph; + EmptyPhaseTracker<1> do_empty_ph, delta_empty_ph; + [[maybe_unused]] EmptyPhaseTracker<1> epi_empty_ph, kv_empty_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + if (it.local_k2q_num != 0) { + auto load_tile = [&](uint8_t* dst, const CUtensorMap* map, uint64_t* full_bar, + int token_begin) { + static_assert(Q_SUB_COLS_BYTES == KV_SUB_COLS_BYTES, + "one subtile stride for Q/dO and K/V"); + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32(dst + s * Q_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), + 0, token_begin, s, it.batch_head); + else + tma_load_3d(smem_ptr_u32(dst + s * Q_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), + 0, it.batch * seqlen + token_begin, it.head * HD_SUBTILES + s); + } + }; + auto load_lse = [&](int qblock) { + const int stage = lse_empty_ph.get_stage(); + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_lse[stage]), + lse_empty_ph.get_phase()); + lse_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_lse[stage]), + Q_TILE * (int)sizeof(float)); + cpasync_bulk_load_mbarrier( + smem_ptr_u32(sLSE + stage * Q_TILE), + lse_rows + (size_t)it.batch_head * seqlen + (size_t)qblock * Q_TILE, + Q_TILE * sizeof(float), smem_ptr_u32(&full_bar_lse[stage])); + } + }; + auto load_delta = [&](int qblock) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_delta), delta_empty_ph.get_phase()); + delta_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(full_bar_delta), Q_TILE * (int)sizeof(float)); + cpasync_bulk_load_mbarrier( + smem_ptr_u32(sDelta), + delta_rows + (size_t)it.batch_head * seqlen + (size_t)qblock * Q_TILE, + Q_TILE * sizeof(float), smem_ptr_u32(full_bar_delta)); + } + }; + + auto load_q = [&](int qblock, auto with_kv_const) { + constexpr bool with_kv = decltype(with_kv_const)::value; + const int stage = q_empty_ph.get_stage(); + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_q[stage]), q_empty_ph.get_phase()); + q_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_q[stage]), + Q_TILE_BYTES + (with_kv ? KV_TILE_BYTES : 0)); + load_tile(sQ[stage], &tmap_q, &full_bar_q[stage], qblock * Q_TILE); + if constexpr (with_kv) + load_tile(sK, &tmap_k, &full_bar_q[stage], it.kv_block_id_in_seq * KV_TILE); + } + }; + + auto load_do = [&](int qblock, auto with_kv_const) { + constexpr bool with_kv = decltype(with_kv_const)::value; + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_do), do_empty_ph.get_phase()); + do_empty_ph.advance(); + + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(full_bar_do), + Q_TILE_BYTES + (with_kv ? KV_TILE_BYTES : 0)); + load_tile(sDO, &tmap_do, full_bar_do, qblock * Q_TILE); + if constexpr (with_kv) + load_tile(sV, &tmap_v, full_bar_do, it.kv_block_id_in_seq * KV_TILE); + } + }; + + if constexpr (PERSISTENT) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_kv), kv_empty_ph.get_phase()); + kv_empty_ph.advance(); + } + + load_q(it.local_k2q_idx[0], std::true_type{}); + load_lse(it.local_k2q_idx[0]); + load_do(it.local_k2q_idx[0], std::true_type{}); + load_delta(it.local_k2q_idx[0]); + for (int j = 1; j < it.local_k2q_num; ++j) { + + const int qblock = it.local_k2q_idx[j]; + load_q(qblock, std::false_type{}); + load_lse(qblock); + load_do(qblock, std::false_type{}); + load_delta(qblock); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + return; + } + else if (warp_id == W_MMA) { + setmaxnreg_dec<88>(); + + tcgen05_alloc<1>(smem_ptr_u32(tmem_slot), TMEM_TOTAL); + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_st = tmem_base; + const uint32_t tmem_dv = tmem_st + ST_COLS; + const uint32_t tmem_dpt = tmem_dv + DV_COLS; + const uint32_t tmem_dk = tmem_dpt + ST_COLS; + const uint32_t tmem_pt_bf16 = tmem_st, tmem_dst_bf16 = tmem_dpt, tmem_dq = tmem_dpt; + const uint32_t lead = elect_one_sync() ? 1u : 0u; + + constexpr uint32_t DESC_SBO = 1024, DESC_LBO = 16; + auto make_smem_desc = [](const void* smem, uint32_t leading_byte_offset) { + return build_smem_desc_blackwell(smem_ptr_u32(smem), DESC_SBO, leading_byte_offset, + SmemSwizzleBlackwell::B128); + }; + const uint64_t desc_k = make_smem_desc(sK, DESC_LBO); + const uint64_t desc_v = make_smem_desc(sV, DESC_LBO); + const uint64_t desc_q0 = make_smem_desc(sQ[0], DESC_LBO); + const uint64_t desc_do = make_smem_desc(sDO, DESC_LBO); + const uint64_t desc_q0_mn = make_smem_desc(sQ[0], Q_SUB_COLS_BYTES); + const uint64_t desc_do_mn = make_smem_desc(sDO, Q_SUB_COLS_BYTES); + const uint64_t desc_k_mn = make_smem_desc(sK, KV_SUB_COLS_BYTES); + const uint64_t desc_dst_mn = make_smem_desc(sDST, DST_TILE_BYTES); + + constexpr uint64_t K16_COLS_DELTA = (MMA_K * (int)sizeof(__nv_bfloat16)) >> 4; + constexpr uint32_t K16_ROWS_DELTA = (MMA_K * SUB_COLS_BYTES) >> 4; + constexpr uint64_t Q_SUB_DELTA = Q_SUB_COLS_BYTES >> 4; + constexpr uint64_t KV_SUB_DELTA = KV_SUB_COLS_BYTES >> 4; + constexpr uint64_t Q_STAGE_DELTA = Q_TILE_BYTES >> 4; + + const uint32_t idesc_st_dpt = make_idesc_bf16_f32(KV_TILE, Q_TILE, false, false); + const uint32_t idesc_dv_dk = make_idesc_bf16_f32(KV_TILE, HEAD_DIM, false, true); + const uint32_t idesc_dq = make_idesc_bf16_f32(Q_TILE, HEAD_DIM, true, true); + + PhaseTracker q_full_ph; + PhaseTracker<1> do_full_ph, pt_ph, dst_ph, dq_empty_ph; + [[maybe_unused]] EmptyPhaseTracker<1> epi_empty_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + auto gemm12_st_dpt = [&](auto is_st_const, int stage) { + constexpr bool is_st = decltype(is_st_const)::value; + const uint32_t tmem_acc = is_st ? tmem_st : tmem_dpt; + const uint64_t da_base = is_st ? desc_k : desc_v; + const uint64_t db_base = is_st ? desc_q0 + (uint64_t)stage * Q_STAGE_DELTA : desc_do; + uint64_t* commit_bar = is_st ? full_bar_st : full_bar_dpt; + if constexpr (is_st) { + mbarrier_wait_parity(smem_ptr_u32(&full_bar_q[stage]), q_full_ph.get_phase()); + } else { + mbarrier_wait_parity(smem_ptr_u32(full_bar_do), do_full_ph.get_phase()); + do_full_ph.advance(); + } + + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_SUBTILE; ++ki) { + const bool accumulate = (s | ki) != 0; + tcgen05_mma_f16_ss_lead(lead, tmem_acc, da_base + s * KV_SUB_DELTA + ki * K16_COLS_DELTA, + db_base + s * Q_SUB_DELTA + ki * K16_COLS_DELTA, idesc_st_dpt, + accumulate); + } + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(commit_bar)); + }; + + auto gemm35_dv_dk = [&](auto is_dv_const, int stage, bool first) { + constexpr bool is_dv = decltype(is_dv_const)::value; + const uint32_t tmem_acc = is_dv ? tmem_dv : tmem_dk; + const uint32_t tmem_a_base = is_dv ? tmem_pt_bf16 : tmem_dst_bf16; + uint64_t db = is_dv ? desc_do_mn : desc_q0_mn + (uint64_t)stage * Q_STAGE_DELTA; + uint64_t* commit_bar = is_dv ? empty_bar_do : &empty_bar_q[stage]; + #pragma unroll + for (int q_half = 0; q_half < 2; ++q_half) { + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_Q_HALF; ++ki) { + const uint32_t tmem_a = tmem_a_base + (uint32_t)(q_half * ST_HALF_COLS + + ki * BF16X2_COLS_PER_K16); + const bool accumulate = !(first && (q_half | ki) == 0); + tcgen05_mma_f16_ts_1sm_lead(lead, tmem_acc, tmem_a, db, idesc_dv_dk, accumulate); + smem_desc_add_lo(db, K16_ROWS_DELTA); + } + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(commit_bar)); + }; + + auto gemm4_dq = [&]() { + uint64_t adst = desc_dst_mn; + uint64_t bk = desc_k_mn; + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_KV_TILE; ++ki) { + tcgen05_mma_f16_ss_lead(lead, tmem_dq, adst, bk, idesc_dq, ki != 0); + smem_desc_add_lo(adst, K16_ROWS_DELTA); + smem_desc_add_lo(bk, K16_ROWS_DELTA); + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dq)); + }; + + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + if (it.local_k2q_num != 0) { + + gemm12_st_dpt(std::true_type{}, q_full_ph.get_stage()); + + if constexpr (PERSISTENT) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + } + + gemm12_st_dpt(std::false_type{}, 0); + + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + + if constexpr (PERSISTENT) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + } + + gemm35_dv_dk(std::true_type{}, 0, true); + + for (int j = 0; j < it.local_k2q_num; ++j) { + + const int stage = q_full_ph.get_stage(); + q_full_ph.advance(); + const int next_stage = q_full_ph.get_stage(); + const bool last = j + 1 == it.local_k2q_num; + + if (!last) gemm12_st_dpt(std::true_type{}, next_stage); + + if (last) { + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dv)); + } + mbarrier_wait_parity(smem_ptr_u32(full_bar_dst), dst_ph.get_phase()); + dst_ph.advance(); + gemm35_dv_dk(std::false_type{}, stage, j == 0); + if (last) { + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dk)); + } + gemm4_dq(); + if (!last) { + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + gemm12_st_dpt(std::false_type{}, 0); + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + gemm35_dv_dk(std::true_type{}, 0, false); + } + } + + if constexpr (PERSISTENT) { + tcgen05_commit1_lead(lead, smem_ptr_u32(empty_bar_kv)); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + tcgen05_relinquish_alloc_permit<1>(); + bar_sync<10>(416); + tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL); + return; + } + else if (warp_id == W_SCHED) { + + if constexpr (CLC) { + setmaxnreg_dec<88>(); + int prod_stage = 0; uint32_t prod_phase = 1; + int cons_stage = 0; uint32_t cons_phase = 0; + while (true) { + if (lane == 0) + mbarrier_wait_parity_suspend(smem_ptr_u32(&clc_empty[prod_stage]), prod_phase); + __syncwarp(); + clc_arrive_expect_tx_cta(smem_ptr_u32(&clc_full[prod_stage]), 16); + if (lane == 0) + clc_try_cancel_async(smem_ptr_u32(&clc_response[prod_stage * 4]), + smem_ptr_u32(&clc_full[prod_stage])); + advance_stage_phase(prod_stage, prod_phase); + ClcTileInfo n = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, cons_stage, cons_phase, elect_one_sync()); + clc_fetch_next_tile_advance(cons_stage, cons_phase); + if (!n.valid) break; + } + + #pragma unroll + for (int st = 0; st < CLC_STAGES; ++st) { + if (lane == 0) + mbarrier_wait_parity_suspend(smem_ptr_u32(&clc_empty[prod_stage]), prod_phase); + __syncwarp(); + advance_stage_phase(prod_stage, prod_phase); + } + } else { + setmaxnreg_dec<24>(); + } + return; + } + else if (warp_id >= W_SOFTMAX0 && warp_id < W_MMA) { + setmaxnreg_inc<136>(); + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_st = tmem_base; + const uint32_t tmem_dv = tmem_st + ST_COLS; + const uint32_t tmem_dpt = tmem_dv + DV_COLS; + const uint32_t tmem_dk = tmem_dpt + ST_COLS; + const uint32_t tmem_pt_bf16 = tmem_st, tmem_dst_bf16 = tmem_dpt; + + const int softmax_warp_id = warp_id - W_SOFTMAX0; + const int lane_group = softmax_warp_id & 3; + const int col_half = softmax_warp_id >> 2; + const int kv_row = lane_group * 32 + lane; + + const uint32_t tmem_lane_base = (uint32_t)(lane_group * 32) << 16; + const uint32_t tmem_f32_offset = tmem_lane_base + (uint32_t)(col_half * ST_HALF_COLS); + const uint32_t tmem_bf16x2_offset = tmem_f32_offset; + + constexpr int CHUNK_BF16 = 16 / (int)sizeof(__nv_bfloat16); + constexpr int CHUNKS_PER_ROW = SUB_COLS_BF16 / CHUNK_BF16; + + __nv_bfloat16* sdst_row = + sDST + (size_t)col_half * KV_TILE * SUB_COLS_BF16 + kv_row * SUB_COLS_BF16; + + PhaseTracker<1> st_ph, dpt_ph, delta_ph, dv_ph, dk_ph; + PhaseTracker lse_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + do { + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + + auto store_dv_dk_tile = [&](uint32_t tmem_acc, auto apply_sm_scale_const, uint8_t* bounce, + const CUtensorMap* map) { + constexpr bool apply_sm_scale = decltype(apply_sm_scale_const)::value; + uint8_t* bounce_subtile = bounce + col_half * KV_SUB_COLS_BYTES; + __nv_bfloat16* stage_row = + reinterpret_cast<__nv_bfloat16*>(bounce_subtile) + kv_row * SUB_COLS_BF16; + #pragma unroll + for (int c0 = 0; c0 < ST_HALF_COLS; c0 += 32) { + uint32_t acc_regs[32]; + tcgen05_ld_32x32b_x32(tmem_acc + tmem_f32_offset + (uint32_t)c0, acc_regs); + tcgen05_fence_before_thread_sync(); + const float* acc = reinterpret_cast(acc_regs); + #pragma unroll + for (int v = 0; v < 32 / CHUNK_BF16; ++v) { + float value[CHUNK_BF16]; + #pragma unroll + for (int e = 0; e < CHUNK_BF16; ++e) + value[e] = apply_sm_scale ? acc[v * CHUNK_BF16 + e] * sm_scale + : acc[v * CHUNK_BF16 + e]; + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(value[0], value[1]); + packed.y = cvt_f32x2_to_bf16x2(value[2], value[3]); + packed.z = cvt_f32x2_to_bf16x2(value[4], value[5]); + packed.w = cvt_f32x2_to_bf16x2(value[6], value[7]); + const int chunk = c0 / CHUNK_BF16 + v; + *reinterpret_cast(stage_row + ((chunk ^ (kv_row & 7)) * CHUNK_BF16)) = packed; + } + } + fence_proxy_async_shared_cta(); + if (col_half == 0) bar_sync<12>(128); else bar_sync<13>(128); + if (lane_group == 0 && elect_one_sync()) { + if constexpr (BHSD) + tma_store_4d(map, 0, it.kv_block_id_in_seq * KV_TILE, col_half, it.batch_head, + smem_ptr_u32(bounce_subtile)); + else + tma_store_3d(map, 0, it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * HD_SUBTILES + col_half, smem_ptr_u32(bounce_subtile)); + cp_async_bulk_commit_group(); + } + }; + + const bool kv_row_valid = kv_row < variable_block_sizes[it.kv_block_id_in_seq]; + + for (int j = 0; j < it.local_k2q_num; ++j) { + const int lse_stage = lse_ph.get_stage(); + const float2* lse2 = reinterpret_cast(sLSE + lse_stage * Q_TILE + + col_half * ST_HALF_COLS); + const float2* delta2 = reinterpret_cast(sDelta + col_half * ST_HALF_COLS); + + mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_lse[lse_stage]), lse_ph.get_phase()); + lse_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_st), st_ph.get_phase()); + st_ph.advance(); + + uint32_t st_regs[ST_HALF_COLS]; + uint32_t pt_bf16x2[ST_HALF_COLS / 2]; + float2* pt_fp32 = reinterpret_cast(st_regs); + tcgen05_ld_32x32b_x64(tmem_st + tmem_f32_offset, st_regs); + tcgen05_fence_before_thread_sync(); + const float2 scale2 = f32x2_splat(scale_log2); + #pragma unroll + for (int c = 0; c < ST_HALF_COLS / 2; ++c) { + const float2 z = ffma2(pt_fp32[c], scale2, make_float2(-lse2[c].x, -lse2[c].y)); + float2 p = make_float2(ex2_approx_f32(z.x), ex2_approx_f32(z.y)); + if (!kv_row_valid) p = make_float2(0.f, 0.f); + pt_fp32[c] = p; + pt_bf16x2[c] = cvt_f32x2_to_bf16x2(p.x, p.y); + } + tcgen05_st_32x32b_x32(tmem_pt_bf16 + tmem_bf16x2_offset, pt_bf16x2); + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(full_bar_pt)); + mbarrier_arrive(smem_ptr_u32(&empty_bar_lse[lse_stage])); + } + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_delta), delta_ph.get_phase()); + delta_ph.advance(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dpt), dpt_ph.get_phase()); + dpt_ph.advance(); + + uint32_t dpt_regs[ST_HALF_COLS]; + tcgen05_ld_32x32b_x64(tmem_dpt + tmem_f32_offset, dpt_regs); + tcgen05_fence_before_thread_sync(); + const float2* dpt2 = reinterpret_cast(dpt_regs); + #pragma unroll + for (int c = 0; c < ST_HALF_COLS / 2; ++c) { + const float2 ds = + fmul2(pt_fp32[c], fadd2(dpt2[c], make_float2(-delta2[c].x, -delta2[c].y))); + st_regs[c] = cvt_f32x2_to_bf16x2(ds.x, ds.y); + } + uint32_t (&dst_bf16x2)[ST_HALF_COLS / 2] = + reinterpret_cast(st_regs); + tcgen05_st_32x32b_x32(tmem_dst_bf16 + tmem_bf16x2_offset, dst_bf16x2); + const uint4* dst_chunks = reinterpret_cast(dst_bf16x2); + #pragma unroll + for (int v = 0; v < CHUNKS_PER_ROW; ++v) + *reinterpret_cast(sdst_row + (v ^ (kv_row & 7)) * CHUNK_BF16) = dst_chunks[v]; + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + fence_proxy_async_shared_cta(); + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(full_bar_dst)); + mbarrier_arrive(smem_ptr_u32(empty_bar_delta)); + } + } + + if (it.local_k2q_num != 0) { + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dv), dv_ph.get_phase()); + dv_ph.advance(); + + store_dv_dk_tile(tmem_dv, std::false_type{}, sDO, &tmap_dv); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dk), dk_ph.get_phase()); + dk_ph.advance(); + + store_dv_dk_tile(tmem_dk, std::true_type{}, sQ[0], &tmap_dk); + + if constexpr (PERSISTENT) { + if (lane_group == 0 && elect_one_sync()) cp_async_bulk_wait_group_read<0>(); + bar_sync<14>(256); + mbarrier_arrive(smem_ptr_u32(empty_bar_epi)); + } + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + bar_sync<10>(416); + return; + } + else if (warp_id < W_SOFTMAX0) { + setmaxnreg_inc<152>(); + + bar_sync<10>(416); + const uint32_t tmem_base = *tmem_slot; + const uint32_t tmem_dq = tmem_base + ST_COLS + DV_COLS; + + const int epi_warp_id = warp_id - W_EPI0; + const uint32_t tmem_lane_base = (uint32_t)(epi_warp_id * 32) << 16; + const int q_row = epi_warp_id * 32 + lane; + const bool is_leader = epi_warp_id == 0 && lane == 0; + + PhaseTracker<1> dq_full_ph; + int clc_stage = 0; + uint32_t clc_phase = 0; + int workitem_id = PERSISTENT ? (int)blockIdx.x : 0; + + if constexpr (PERSISTENT) { + if (elect_one_sync()) mbarrier_arrive(smem_ptr_u32(empty_bar_dq)); + } + + do { + + int batch_head, kv_block_id_in_seq; + if constexpr (PERSISTENT) { + const int real_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + batch_head = real_id / num_kv_blocks_per_seq; + kv_block_id_in_seq = real_id % num_kv_blocks_per_seq; + } else { + batch_head = (int)blockIdx.y; + kv_block_id_in_seq = (int)blockIdx.x; + } + const WorkItem it = decode_workitem(batch_head, kv_block_id_in_seq, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + float* dqaccum_head = + dqaccum + (size_t)it.batch_head * num_kv_blocks_per_seq * DQ::DQ_BLOCK_ELEMS; + + for (int j = 0; j < it.local_k2q_num; ++j) { + float* dqaccum_block = dqaccum_head + (size_t)it.local_k2q_idx[j] * DQ::DQ_BLOCK_ELEMS; + + mbarrier_wait_parity(smem_ptr_u32(full_bar_dq), dq_full_ph.get_phase()); + dq_full_ph.advance(); + + uint32_t dq_regs[HEAD_DIM]; + #pragma unroll + for (int c = 0; c < HEAD_DIM / 64; ++c) + tcgen05_ld_32x32b_x64(tmem_dq + tmem_lane_base + (uint32_t)(c * 64), + reinterpret_cast(dq_regs[c * 64])); + + tcgen05_wait_ld(); + tcgen05_fence_before_thread_sync(); + if (elect_one_sync()) mbarrier_arrive(smem_ptr_u32(empty_bar_dq)); + + #pragma unroll + for (int hd_slice = 0; hd_slice < HEAD_DIM / DQ::COLS; ++hd_slice) { + const int stage_buf = hd_slice & 1; + const float4* dq_row4 = reinterpret_cast(dq_regs + hd_slice * DQ::COLS); + #pragma unroll + for (int v4 = 0; v4 < DQ::COLS / 4; ++v4) + *reinterpret_cast(sDQ_STAGE[stage_buf] + v4 * Q_TILE * 4 + q_row * 4) = + dq_row4[v4]; + fence_proxy_async_shared_cta(); + bar_sync<11>(128); + if (is_leader) { + const size_t slice_offset = (size_t)hd_slice * Q_TILE * DQ::COLS; + cpasync_reduce_bulk_add_f32(dqaccum_block + slice_offset, + smem_ptr_u32(sDQ_STAGE[stage_buf]), DQ::DQ_ONE_PUSH_BYTES); + cp_async_bulk_commit_group(); + cp_async_bulk_wait_group_read<1>(); + } + bar_sync<11>(128); + } + } + + if (it.local_k2q_num == 0) { + + if (is_leader) cp_async_bulk_wait_group_read<0>(); + bar_sync<11>(128); + static_assert(DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES == HD_SUBTILES * KV_SUB_COLS_BYTES, + "the dQ stage buffers hold one 128 kv x 128 hd bf16 tile"); + constexpr int ZERO_UINT4_PER_THREAD = + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES / (128 * 16); + uint4* zero_tile = reinterpret_cast(sDQ_STAGE_bytes) + epi_warp_id * 32 + lane; + #pragma unroll + for (int v = 0; v < ZERO_UINT4_PER_THREAD; ++v) + zero_tile[v * 128] = make_uint4(0u, 0u, 0u, 0u); + fence_proxy_async_shared_cta(); + bar_sync<11>(128); + if (is_leader) { + #pragma unroll + for (int which = 0; which < 2; ++which) { + const CUtensorMap* map = (which == 0) ? &tmap_dk : &tmap_dv; + #pragma unroll + for (int s = 0; s < HD_SUBTILES; ++s) { + const uint32_t src = smem_ptr_u32(sDQ_STAGE_bytes + (size_t)s * KV_SUB_COLS_BYTES); + if constexpr (BHSD) + tma_store_4d(map, 0, it.kv_block_id_in_seq * KV_TILE, s, it.batch_head, src); + else + tma_store_3d(map, 0, it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * HD_SUBTILES + s, src); + } + } + cp_async_bulk_commit_group(); + cp_async_bulk_wait_group_read<0>(); + } + bar_sync<11>(128); + } + + if constexpr (CLC) { + ClcTileInfo next = clc_fetch_next_tile<1, 1, ClcRasterOrder::AlongN, 1, true>( + clc_full, clc_empty, clc_response, clc_stage, clc_phase, elect_one_sync()); + clc_fetch_next_tile_advance(clc_stage, clc_phase); + workitem_id = next.valid ? (int)next.n_tile : -1; + } else if constexpr (PERSISTENT) { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } while (PERSISTENT && workitem_id >= 0); + + if (is_leader) cp_async_bulk_wait_group_read<0>(); + bar_sync<11>(128); + bar_sync<10>(416); + if constexpr (KERNEL_PDL) { + if (is_leader) griddepcontrol_launch_dependents(); + } + return; + } + else { + setmaxnreg_dec<24>(); + return; + } +#endif +} + +template +__global__ void __launch_bounds__(256, 1) + vsa_bwd_preprocess_kernel(const __nv_bfloat16* __restrict__ o, + const __nv_bfloat16* __restrict__ dout, + float* __restrict__ delta_rows, float* __restrict__ dqaccum, + int num_heads, int seqlen) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + const int q_block_id = (int)blockIdx.x; + const int batch_head = (int)blockIdx.y; + const int batch = batch_head / num_heads, head = batch_head % num_heads; + const int token_begin = q_block_id * Q_TILE; + + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + + float4* dqaccum_zero_destination = reinterpret_cast( + dqaccum + (size_t)batch_head * seqlen * HEAD_DIM + (size_t)q_block_id * DQ::DQ_BLOCK_ELEMS); + const float4 zero_float4 = make_float4(0.f, 0.f, 0.f, 0.f); + #pragma unroll + for (int i = 0; i < (DQ::DQ_BLOCK_ELEMS / 4) / 256; ++i) + dqaccum_zero_destination[i * 256 + threadIdx.x] = zero_float4; + + if constexpr (KERNEL_PDL) griddepcontrol_launch_dependents(); + + const int row_in_pass = (int)threadIdx.x / 16; + const int dimension_begin = ((int)threadIdx.x % 16) * 8; + #pragma unroll + for (int row_pass = 0; row_pass < Q_TILE / 16; ++row_pass) { + const int token = token_begin + row_pass * 16 + row_in_pass; + const size_t element = + token_offset(batch, head, num_heads, seqlen, token) + dimension_begin; + const uint4 o_vector = *reinterpret_cast(o + element); + const uint4 dout_vector = *reinterpret_cast(dout + element); + const __nv_bfloat162* o_pairs = reinterpret_cast(&o_vector); + const __nv_bfloat162* dout_pairs = reinterpret_cast(&dout_vector); + float delta_accumulator = 0.f; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 o_pair_as_float = __bfloat1622float2(o_pairs[i]); + const float2 dout_pair_as_float = __bfloat1622float2(dout_pairs[i]); + delta_accumulator += o_pair_as_float.x * dout_pair_as_float.x + + o_pair_as_float.y * dout_pair_as_float.y; + } + #pragma unroll + for (int shuffle_offset = 8; shuffle_offset > 0; shuffle_offset >>= 1) + delta_accumulator += __shfl_down_sync(0xffffffffu, delta_accumulator, shuffle_offset); + if (dimension_begin == 0) delta_rows[(size_t)batch_head * seqlen + token] = delta_accumulator; + } +#endif +} + +template +__global__ void __launch_bounds__(128, 1) + vsa_bwd_postprocess_kernel(const float* __restrict__ dqaccum, __nv_bfloat16* __restrict__ dq, + int num_heads, int seqlen, float sm_scale) { +#if !defined(__CUDA_ARCH__) || \ + ((__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) || \ + (__CUDA_ARCH__ == 1030 && defined(__CUDA_ARCH_FEAT_SM103_ALL))) + extern __shared__ __align__(16) float post_smem[]; + const int q_block_id = (int)blockIdx.x; + const int batch_head = (int)blockIdx.y; + const int batch = batch_head / num_heads, head = batch_head % num_heads; + const int token_begin = q_block_id * Q_TILE; + const int thread = (int)threadIdx.x; + constexpr int CHUNK_BF16 = 16 / (int)sizeof(__nv_bfloat16); + + const float4* dqaccum_block = reinterpret_cast( + dqaccum + (size_t)batch_head * seqlen * HEAD_DIM + (size_t)q_block_id * DQ::DQ_BLOCK_ELEMS); + if constexpr (KERNEL_PDL) griddepcontrol_wait(); + #pragma unroll + for (int i = 0; i < (DQ::DQ_BLOCK_ELEMS / 4) / 128; ++i) + cp_async_cg_16(smem_ptr_u32(post_smem + (i * 128 + thread) * 4), + dqaccum_block + i * 128 + thread); + cp_async_commit_group(); + cp_async_wait_group<0>(); + __syncthreads(); + + if constexpr (KERNEL_PDL) griddepcontrol_launch_dependents(); + + uint32_t dq_packed[HEAD_DIM / 2]; + #pragma unroll + for (int hd_slice = 0; hd_slice < HEAD_DIM / DQ::COLS; ++hd_slice) { + #pragma unroll + for (int v4 = 0; v4 < DQ::COLS / 4; ++v4) { + const float4 value = *reinterpret_cast( + post_smem + hd_slice * Q_TILE * DQ::COLS + v4 * Q_TILE * 4 + thread * 4); + const int pair_index = (hd_slice * DQ::COLS + v4 * 4) / 2; + dq_packed[pair_index + 0] = cvt_f32x2_to_bf16x2(value.x * sm_scale, value.y * sm_scale); + dq_packed[pair_index + 1] = cvt_f32x2_to_bf16x2(value.z * sm_scale, value.w * sm_scale); + } + } + __syncthreads(); + + __nv_bfloat16* dq_tile_bf16 = reinterpret_cast<__nv_bfloat16*>(post_smem); + const uint4* dq_chunks = reinterpret_cast(dq_packed); + __nv_bfloat16* dq_row_bf16 = dq_tile_bf16 + thread * HEAD_DIM; + #pragma unroll + for (int v = 0; v < HEAD_DIM / CHUNK_BF16; ++v) + *reinterpret_cast(dq_row_bf16 + (v / 8) * 64 + ((v % 8) ^ (thread & 7)) * CHUNK_BF16) = + dq_chunks[v]; + __syncthreads(); + + const int row_in_pass = thread / 16; + const int dimension_begin = (thread % 16) * 8; + #pragma unroll + for (int row_pass = 0; row_pass < Q_TILE / 8; ++row_pass) { + const int row = row_pass * 8 + row_in_pass; + const int half = dimension_begin / 64; + const int chunk = ((dimension_begin % 64) / CHUNK_BF16) ^ (row & 7); + const uint4 value = *reinterpret_cast( + dq_tile_bf16 + row * HEAD_DIM + half * 64 + chunk * CHUNK_BF16); + const size_t element = + token_offset(batch, head, num_heads, seqlen, token_begin + row) + dimension_begin; + *reinterpret_cast(dq + element) = value; + } +#endif +} + +} + +#endif diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh new file mode 100644 index 000000000..9cacbda24 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a_blk128.cuh @@ -0,0 +1,231 @@ +// block_sparse_bwd_launch_sm100a_blk128.cuh -- host surface of the 128-token VSA block-sparse +// backward drop: argument struct, workspace sizes, the support predicate and the stream-chained +// launch (preprocess -> main -> postprocess, each a programmatic dependent launch of the one +// before). Tensor maps are encoded per call (a torch caller hands us fresh pointers every time). +#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH +#define BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH + +#include +#include "block_sparse_bwd_kernel_sm100a_blk128.cuh" + +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif + +namespace vsa_bwd_blk128 { + +using dq_accum_t = float; + +struct BlockSparseVsaBwdArgs { + // Activations are bf16, contiguous, [B, H, S, 128] under VSA_BHSD, else [B*S, H, 128]. + // nb below = num_kv_blocks_per_seq = S / 128. + + // Forward operands and results. + const __nv_bfloat16* q; + const __nv_bfloat16* k; + const __nv_bfloat16* v; + const __nv_bfloat16* o; + // Gradient of the forward output. + const __nv_bfloat16* dout; + // [B, H, S] fp32 log-sum-exp in Triton's M form: max(qk * sm_scale * log2e) + log2(l). + const float* lse; + + // Sparsity metadata, FastVideo's invert_indices layout, in 128-token blocks. + // [B*H*nb, max_q_blocks] int32: q blocks selecting each kv block; entries past the count unread. + const int* k2q_idx; + // [B*H*nb] int32: valid entries per k2q_idx row (0 allowed: the kernel writes zero dk/dv rows). + const int* k2q_num; + // [nb] int32: valid kv tokens per block (<= 128); kv rows at or past the count are masked. + const int* variable_block_sizes; + + // Outputs, inputs' layout; every row is written, including the rows of unselected kv blocks. + __nv_bfloat16* dq; + __nv_bfloat16* dk; + __nv_bfloat16* dv; + + // Scratch, caller-allocated; byte sizes from the block_sparse_bwd_*_bytes helpers below. + // [B*H*S*128] fp32, drain-native; preprocess zeroes, main reduce-adds, postprocess reads. + dq_accum_t* dqaccum; + // [B*H*S] fp32 rowsum(bf16(o) * dout), written by the preprocess. + float* delta; + + int batch; + int num_heads; + // S = nb * 128. + int seqlen; + // Must be 128. + int head_dim; + // nb = seqlen / 128. + int num_kv_blocks_per_seq; + // k2q_idx row stride (FastVideo passes nb). + int max_q_blocks; + // Softmax scale; dq and dk carry it, dv does not. + float sm_scale; +}; + +__host__ inline size_t block_sparse_bwd_dqaccum_bytes(int batch, int num_heads, int seqlen) { + return (size_t)batch * num_heads * seqlen * HEAD_DIM * sizeof(dq_accum_t); +} +__host__ inline size_t block_sparse_bwd_delta_bytes(int batch, int num_heads, int seqlen) { + return (size_t)batch * num_heads * seqlen * sizeof(float); +} + +__host__ inline cudaError_t block_sparse_bwd_supported(const BlockSparseVsaBwdArgs& args) { + if (args.head_dim != HEAD_DIM) { + return cudaErrorInvalidValue; + } + if (args.num_kv_blocks_per_seq < 1) { + return cudaErrorInvalidValue; + } + if (args.seqlen != args.num_kv_blocks_per_seq * BLOCK) { + return cudaErrorInvalidValue; + } + if (args.batch < 1 || args.num_heads < 1 || args.max_q_blocks < 1) { + return cudaErrorInvalidValue; + } + // The main grid is (kv block, batch * head): gridDim.y is limited to 65535. + if ((long)args.batch * args.num_heads > 65535) { + return cudaErrorInvalidValue; + } + if (!std::isfinite(args.sm_scale)) { + return cudaErrorInvalidValue; + } + if (!args.q || !args.k || !args.v || !args.o || !args.dout || !args.lse) { + return cudaErrorInvalidValue; + } + if (!args.dq || !args.dk || !args.dv) { + return cudaErrorInvalidValue; + } + if (!args.k2q_idx || !args.k2q_num || !args.variable_block_sizes) { + return cudaErrorInvalidValue; + } + if (!args.dqaccum || !args.delta) { + return cudaErrorInvalidValue; + } + return cudaSuccess; +} + +// Q, K, V, dO, dK, dV tensor maps, one box_tokens x 64-hd box per TMA (two per tile): +// BSHD: 3D [64 hd, B*S tokens, H*2 hd units], strides {H*128*2, 128} bytes. +// BHSD: 4D [64 hd, S tokens, 2 hd units, B*H], strides {128*2, 128, S*128*2} bytes. +__host__ inline cudaError_t make_tma_tile_units(CUtensorMap* map, const __nv_bfloat16* ptr, int B, + int H, int S, int box_tokens) { + CUresult r; + if (VSA_BHSD) { + uint64_t gd[4] = {(uint64_t)SUB_COLS_BF16, (uint64_t)S, (uint64_t)HD_SUBTILES, (uint64_t)B * H}; + uint64_t gs[3] = {(uint64_t)HEAD_DIM * 2, (uint64_t)SUB_COLS_BYTES, (uint64_t)S * HEAD_DIM * 2}; + uint32_t bd[4] = {(uint32_t)SUB_COLS_BF16, (uint32_t)box_tokens, 1u, 1u}; + uint32_t es[4] = {1u, 1u, 1u, 1u}; + r = cuTensorMapEncodeTiled( + map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, const_cast<__nv_bfloat16*>(ptr), gd, gs, bd, es, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + } else { + uint64_t gd[3] = {(uint64_t)SUB_COLS_BF16, (uint64_t)B * S, (uint64_t)H * HD_SUBTILES}; + uint64_t gs[2] = {(uint64_t)H * HEAD_DIM * 2, (uint64_t)SUB_COLS_BYTES}; + uint32_t bd[3] = {(uint32_t)SUB_COLS_BF16, (uint32_t)box_tokens, 1u}; + uint32_t es[3] = {1u, 1u, 1u}; + r = cuTensorMapEncodeTiled( + map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, const_cast<__nv_bfloat16*>(ptr), gd, gs, bd, es, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + } + return (r == CUDA_SUCCESS) ? cudaSuccess : cudaErrorInvalidValue; +} + +// Every launch carries the programmatic-stream-serialization attribute when the kernels are +// built with KERNEL_PDL (their griddepcontrol.wait guards the predecessor's data). +__host__ inline cudaLaunchConfig_t pdl_launch_config(dim3 grid, dim3 block, size_t smem, + cudaStream_t stream, cudaLaunchAttribute* at) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = grid; + cfg.blockDim = block; + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + at->id = cudaLaunchAttributeProgrammaticStreamSerialization; + at->val.programmaticStreamSerializationAllowed = 1; + cfg.attrs = at; + cfg.numAttrs = KERNEL_PDL ? 1 : 0; + return cfg; +} + +__host__ inline cudaError_t launch_block_sparse_bwd_sm100a(const BlockSparseVsaBwdArgs& args, + cudaStream_t stream) { + const cudaError_t supported = block_sparse_bwd_supported(args); + if (supported != cudaSuccess) { + return supported; + } + const int B = args.batch, H = args.num_heads, S = args.seqlen; + + CUtensorMap tq, tk, tv, tdo, tdk, tdv; + if (make_tma_tile_units(&tq, args.q, B, H, S, Q_TILE) != cudaSuccess || + make_tma_tile_units(&tk, args.k, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tv, args.v, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tdo, args.dout, B, H, S, Q_TILE) != cudaSuccess || + make_tma_tile_units(&tdk, args.dk, B, H, S, KV_TILE) != cudaSuccess || + make_tma_tile_units(&tdv, args.dv, B, H, S, KV_TILE) != cudaSuccess) { + return cudaErrorInvalidValue; + } + + auto main_kernel = vsa_bwd_main_kernel; + auto postprocess_kernel = vsa_bwd_postprocess_kernel; + constexpr int POST_SMEM_BYTES = DQ::DQ_BLOCK_ELEMS * (int)sizeof(float); + cudaError_t e = + cudaFuncSetAttribute(main_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL); + if (e != cudaSuccess) { + return e; + } + e = cudaFuncSetAttribute(postprocess_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + POST_SMEM_BYTES); + if (e != cudaSuccess) { + return e; + } + + // Preprocess: dqaccum zero + Delta, one CTA per (q block, batch * head). + cudaLaunchAttribute pre_at[1]; + cudaLaunchConfig_t pre_cfg = pdl_launch_config( + dim3((unsigned)(S / Q_TILE), (unsigned)(B * H), 1), dim3(256, 1, 1), 0, stream, pre_at); + e = cudaLaunchKernelEx(&pre_cfg, vsa_bwd_preprocess_kernel, args.o, args.dout, + args.delta, args.dqaccum, H, S); + if (e != cudaSuccess) { + return e; + } + + // Main: one CTA per (kv block, batch * head). + const float scale_log2 = args.sm_scale * 1.4426950408889634f; + cudaLaunchAttribute main_at[2]; + cudaLaunchConfig_t main_cfg = {}; + main_cfg.gridDim = dim3((unsigned)args.num_kv_blocks_per_seq, (unsigned)(B * H), 1); + main_cfg.blockDim = dim3(N_WARPS * 32, 1, 1); + main_cfg.dynamicSmemBytes = SMEM_TOTAL; + main_cfg.stream = stream; + main_at[0].id = cudaLaunchAttributeClusterDimension; + main_at[0].val.clusterDim.x = 1; + main_at[0].val.clusterDim.y = 1; + main_at[0].val.clusterDim.z = 1; + main_at[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + main_at[1].val.programmaticStreamSerializationAllowed = 1; + main_cfg.attrs = main_at; + main_cfg.numAttrs = KERNEL_PDL ? 2 : 1; + const int* workitem_remap = nullptr; + e = cudaLaunchKernelEx(&main_cfg, main_kernel, tq, tk, tv, tdo, tdk, tdv, args.dqaccum, args.lse, + args.delta, args.k2q_idx, args.k2q_num, args.max_q_blocks, workitem_remap, + args.variable_block_sizes, B, H, S, scale_log2, args.sm_scale); + if (e != cudaSuccess) { + return e; + } + + // Postprocess: drain-native dqaccum -> bf16 dq, one CTA per (q block, batch * head). + cudaLaunchAttribute post_at[1]; + cudaLaunchConfig_t post_cfg = + pdl_launch_config(dim3((unsigned)(S / Q_TILE), (unsigned)(B * H), 1), dim3(128, 1, 1), + POST_SMEM_BYTES, stream, post_at); + return cudaLaunchKernelEx(&post_cfg, postprocess_kernel, args.dqaccum, args.dq, H, S, + args.sm_scale); +} + +} // namespace vsa_bwd_blk128 + +using namespace vsa_bwd_blk128; + +#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_BLK128_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu index 929979377..8397ebf56 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu @@ -10,6 +10,9 @@ // // The layout is fixed at compile time: VSA_BHSD true -> [B, H, S, 128] (FastVideo's build), // false -> [B, S, H, 128] (repo native; the kernel addresses it as [B*S tokens, H, 128]). +// VSA_BLK128 picks the kernel: false -> 64-token sparse blocks (block_sparse_sm100a_bwd), true -> +// 128-token (block_sparse_sm100a_blk128_bwd; block_sparse_bwd_blk128_sm100a.cu re-includes this +// file with it set). The python backend picks by the metadata's block size. #include #include @@ -17,7 +20,16 @@ #include +#ifndef VSA_BLK128 +#define VSA_BLK128 false +#endif +#if VSA_BLK128 +#include "block_sparse_bwd_launch_sm100a_blk128.cuh" +#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_blk128_bwd +#else #include "block_sparse_bwd_launch_sm100a.cuh" +#define BLOCK_SPARSE_SM100A_BWD block_sparse_sm100a_bwd +#endif namespace { @@ -49,7 +61,7 @@ __nv_bfloat16* bf16_ptr(const torch::Tensor& t) { } // namespace // Returns {dq, dk, dv}: bf16, each with the shape and layout of q, k, v respectively. -std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch::Tensor q, +std::vector 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, @@ -105,20 +117,23 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: auto dv = torch::empty_like(v); // Workspace: torch::empty is enough. The preprocess kernel zeroes dqaccum and fully writes - // qt, dot and delta before the main kernel reads them; it also zeroes the dk/dv rows of kv - // blocks that no q block selects, so the empty_like outputs above come back fully defined. + // delta (and, at 64-token blocks, qt and dot) before the main kernel reads them. The dk/dv rows + // of kv blocks that no q block selects are zeroed by the preprocess (64) or by the main kernel + // (128), so the empty_like outputs above come back fully defined. const auto bytes = q.options().dtype(at::kByte); const int b = (int)B, h = (int)H, s = (int)S; auto dqaccum = torch::empty({(int64_t)block_sparse_bwd_dqaccum_bytes(b, h, s)}, bytes); - auto qt = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); - auto dot = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); auto delta = torch::empty({(int64_t)block_sparse_bwd_delta_bytes(b, h, s)}, bytes); +#if !VSA_BLK128 + auto qt = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); + auto dot = torch::empty({(int64_t)block_sparse_bwd_transposed_bytes(b, h, s)}, bytes); torch::Tensor order; const bool device_order = num_kv_blocks_per_seq >= ORDER_MIN_KV_BLOCKS; if (device_order) { order = torch::empty({(int64_t)block_sparse_bwd_order_bytes(b, h, (int)num_kv_blocks_per_seq)}, bytes); } +#endif BlockSparseVsaBwdArgs a{}; a.q = bf16_ptr(q); @@ -133,12 +148,14 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: a.k2q_idx = k2q_idx.data_ptr(); a.k2q_num = k2q_num.data_ptr(); a.variable_block_sizes = variable_block_sizes.data_ptr(); - a.workitem_remap = nullptr; - a.order_workspace = device_order ? reinterpret_cast(order.data_ptr()) : nullptr; a.dqaccum = reinterpret_cast(dqaccum.data_ptr()); - a.qt = bf16_ptr(qt); - a.dot = bf16_ptr(dot); a.delta = reinterpret_cast(delta.data_ptr()); +#if !VSA_BLK128 + a.workitem_remap = nullptr; + a.order_workspace = device_order ? reinterpret_cast(order.data_ptr()) : nullptr; + a.qt = bf16_ptr(qt); + a.dot = bf16_ptr(dot); +#endif a.batch = b; a.num_heads = h; a.seqlen = s; @@ -149,16 +166,17 @@ std::vector block_sparse_sm100a_bwd(torch::Tensor grad_o, torch:: // Report an unsupported regime loudly rather than returning plausible-looking wrong values. TORCH_CHECK(block_sparse_bwd_supported(a) == cudaSuccess, - "block_sparse_sm100a_bwd: unsupported configuration -- requires head_dim==", HEAD_DIM, + "block_sparse_sm100a_bwd (", BLOCK, + "-token blocks): unsupported configuration -- requires head_dim==", HEAD_DIM, ", seqlen == num_kv_blocks_per_seq*", BLOCK, - " with seqlen % 128 == 0, " + " with seqlen % 128 == 0, batch*heads <= 65535, " "max_q_blocks >= 1 and a finite sm_scale. Got head_dim=", D, " num_kv_blocks_per_seq=", num_kv_blocks_per_seq, " seqlen=", S, " max_q_blocks=", max_q_blocks, " sm_scale=", sm_scale); const cudaError_t err = launch_block_sparse_bwd_sm100a(a, at::cuda::getCurrentCUDAStream()); - TORCH_CHECK(err == cudaSuccess, - "block_sparse_sm100a_bwd launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, "block_sparse_sm100a_bwd (", BLOCK, + "-token blocks) launch failed: ", cudaGetErrorString(err)); return {dq, dk, dv}; } diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh index 545e2e3e6..895519dda 100644 --- a/fastvideo-kernel/csrc/attention/primitives.cuh +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -985,3 +985,30 @@ uint64_t make_l2cache_policy_fractional_evict_last_unchanged(float fraction_f32) : "f"(fraction_f32)); return policy; } + +__device__ __forceinline__ +void cp_async_cg_16(uint32_t smem_dst, const void* gmem_src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + :: "r"(smem_dst), "l"(gmem_src) : "memory"); +} + +__device__ __forceinline__ +void cp_async_commit_group() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ +void cp_async_wait_group() { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N) : "memory"); +} + +__device__ __forceinline__ +void griddepcontrol_wait() { + asm volatile("griddepcontrol.wait;\n" ::: "memory"); +} + +__device__ __forceinline__ +void griddepcontrol_launch_dependents() { + asm volatile("griddepcontrol.launch_dependents;\n" ::: "memory"); +} diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 45d5744bb..9f839dd62 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -46,6 +46,10 @@ extern std::vector 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); +extern std::vector block_sparse_sm100a_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 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -61,6 +65,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("block_sparse_sm100a_bwd", torch::wrap_pybind_function(block_sparse_sm100a_bwd), "VSA block-sparse attention backward, 64-token blocks (Blackwell sm100a/sm103a)"); + m.def("block_sparse_sm100a_blk128_bwd", + torch::wrap_pybind_function(block_sparse_sm100a_blk128_bwd), + "VSA block-sparse attention backward, 128-token blocks (Blackwell sm100a/sm103a)"); #endif #ifdef TK_COMPILE_ST_ATTN diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 0b2c56df1..b463932ca 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -52,13 +52,13 @@ def _force_tk() -> bool: def _force_sm100a() -> bool: """True iff the data-center Blackwell forward is explicitly opted into. - Opt-in only (same legacy-named env the H3 backend honors): the extension is - forward-only, so this routing pairs it with the Triton backward -- its lse - is already in Triton's M format. Honored only when - ``block_sparse_attn_sm100a.is_supported`` passes. Unsupported 64-token - metadata falls through to the default selection; unsupported 128-token - metadata raises because Triton has no compatible fallback. - ``FASTVIDEO_VSA_TRITON`` still wins. + Opt-in only (same legacy-named env the H3 backend honors). Honored only when + ``block_sparse_attn_sm100a.is_supported`` passes; the backward then runs the + sm_100a/sm_103a CUDA backward when ``block_sparse_attn_bwd_sm100a.is_supported`` + passes and Triton otherwise (the forward's lse is already in Triton's M + format). Unsupported 64-token metadata falls through to the default + selection; unsupported 128-token metadata raises because Triton has no + compatible fallback. ``FASTVIDEO_VSA_TRITON`` still wins. """ return _env_bool("FASTVIDEO_VSA_SM100A") @@ -389,11 +389,12 @@ block_sparse_attn_sm90.register_autograd(_backward_sm90, setup_context=_setup_co # Data-center Blackwell backend custom op (index-native; legacy sm100a API name) # # Forward runs the sm_100a/sm_103a CUDA extension. Backward runs the sm_100a CUDA -# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64-token blocks, -# sm_100a device, extension built with the op) and the Triton kernels otherwise. -# The native forward emits lse in exactly Triton's M format (max*log2e + -# log2(l)), so either pairing needs no conversion. Both backwards are -# hardcoded to 64-token blocks, hence the block-size assert below. +# backward when block_sparse_attn_bwd_sm100a.is_supported passes (64- or 128-token +# blocks, data-center Blackwell device, extension built with that block's op) and +# the Triton kernels otherwise. The native forward emits lse in exactly Triton's M +# format (max*log2e + log2(l)), so either pairing needs no conversion. The Triton +# backward is hardcoded to 64-token blocks, so unsupported 128-token metadata +# raises instead of falling back. # --------------------------------------------------------------------------- @@ -493,16 +494,17 @@ def _setup_context_sm100a(ctx, inputs, output): def _backward_sm100a(ctx, grad_o, grad_M): q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes = ctx.saved_tensors - block = q.shape[2] // variable_block_sizes.numel() - if block != 64: - raise RuntimeError( - "block_sparse_attn_sm100a backward pairs the sm_100a/sm_103a forward with a " - f"backward that is hardcoded to 64-token blocks; got {block}. " - "Run 128-token-block metadata without grad, or use the Triton forward.") if _sm100a_backward_is_supported(q, variable_block_sizes): dq, dk, dv = block_sparse_attn_backward_sm100a(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) else: + block = q.shape[2] // variable_block_sizes.numel() + if block != 64: + raise RuntimeError( + "block_sparse_attn_sm100a backward: no sm_100a/sm_103a backward for " + f"{block}-token blocks on this build/device, and the Triton backward is " + "hardcoded to 64-token blocks. Run this metadata without grad, or build " + "the extension with block_sparse_sm100a_blk128_bwd.") dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, q2k_num, variable_block_sizes) return dq, dk, dv, None, None, None @@ -536,8 +538,9 @@ def block_sparse_attn_from_indices( # Backend resolution: # - FASTVIDEO_VSA_TRITON forces Triton everywhere. - # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward - # (Triton backward). The environment name is retained for compatibility. + # - FASTVIDEO_VSA_SM100A opts into the data-center Blackwell forward (CUDA + # backward when supported, else Triton). The environment name is retained + # for compatibility. # Unsupported 64-token metadata falls through; unsupported 128-token # metadata raises because Triton cannot consume it. # - FASTVIDEO_VSA_TK requests sm_90 TK; honored only when it's actually diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py index ea073d868..cb2b1c7ca 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -4,8 +4,10 @@ Companion of ``block_sparse_attn_sm100a`` (the forward): consumes the forward's ``lse`` in the Triton "M format" (``max(qk * sm_scale * log2e) + log2(l)``, ``[B, H, S]`` fp32) unchanged and FastVideo's k2q index metadata, returns ``(dq, dk, dv)`` in bf16 with the inputs' layout and the -Triton backward's scaling (dq and dk carry sm_scale, dv does not). 64-token blocks only; any -other configuration falls back to Triton via ``is_supported``. +Triton backward's scaling (dq and dk carry sm_scale, dv does not). Two kernels, picked by the +metadata's block size: 64-token blocks (``block_sparse_sm100a_bwd``) and 128-token blocks +(``block_sparse_sm100a_blk128_bwd``); any other configuration falls back to Triton via +``is_supported``. """ from typing import Tuple @@ -16,45 +18,64 @@ try: # The pybind symbols live on fastvideo_kernel_ops, NOT on the _C package that contains it # (its __init__ is empty, so hasattr on the package fails with the kernel built and present). from fastvideo_kernel._C import fastvideo_kernel_ops as _C - _BWD = getattr(_C, "block_sparse_sm100a_bwd", None) - _HAS_VSA_BWD_SM100A = _BWD is not None + _BWD_BY_BLOCK = { + 64: getattr(_C, "block_sparse_sm100a_bwd", None), + 128: getattr(_C, "block_sparse_sm100a_blk128_bwd", None), + } + _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) except ImportError: # pragma: no cover - extension not built _C = None - _BWD = None + _BWD_BY_BLOCK = {} _HAS_VSA_BWD_SM100A = False _SUPPORTED_COMPUTE_CAPABILITIES = {(10, 0), (10, 3)} HEAD_DIM = 128 -BLOCK = 64 +BLOCKS = (64, 128) # Must match the -DVSA_BHSD the extension was compiled with (FastVideo builds with true). BHSD = True def set_extension(module) -> None: - """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd``. + """Use an already-loaded extension module exposing ``block_sparse_sm100a_bwd`` and / or + ``block_sparse_sm100a_blk128_bwd``. - A standalone build of ``block_sparse_bwd_sm100a.cu`` (for example through + A standalone build of the binding .cu files (for example through ``torch.utils.cpp_extension.load`` with a ten-line pybind wrapper) can be injected here, so the backend can be exercised without rebuilding the fastvideo_kernel wheel. """ - global _C, _BWD, _HAS_VSA_BWD_SM100A + global _C, _BWD_BY_BLOCK, _HAS_VSA_BWD_SM100A _C = module - _BWD = getattr(module, "block_sparse_sm100a_bwd", None) - _HAS_VSA_BWD_SM100A = _BWD is not None + _BWD_BY_BLOCK = { + 64: getattr(module, "block_sparse_sm100a_bwd", None), + 128: getattr(module, "block_sparse_sm100a_blk128_bwd", None), + } + _HAS_VSA_BWD_SM100A = any(_BWD_BY_BLOCK.values()) def _seqlen(q: torch.Tensor) -> int: return q.shape[2] if BHSD else q.shape[1] +def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + """The block size the metadata implies (seqlen / num_blocks), 0 when it is not integral. + + A plain int: under dynamic-shape torch.compile the shapes are SymInts, which cannot key + ``_BWD_BY_BLOCK``; ``int()`` specializes the value with a shape guard. + """ + num_blocks = variable_block_sizes.numel() + seqlen = _seqlen(q) + return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks) + + def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: """True iff this build can run these tensors; otherwise the caller uses Triton. Static facts only (shapes, dtypes, arch, layout), never tensor contents, so it is cheap - enough for a per-layer dispatch path. The kernel is fixed at 64-token blocks with - head_dim 128 and needs seqlen == 64 * num_blocks with an even num_blocks (its preprocess - works in 128-token blocks). Per-row k2q counts may be anything in [0, num_q_blocks], - including 0: unselected kv blocks get exactly-zero dk/dv rows. + enough for a per-layer dispatch path. Both kernels take head_dim 128 and seqlen == block * + num_blocks; the 64-token kernel also needs an even num_blocks (its preprocess works in + 128-token blocks); batch * heads is the launch grid's y extent, so it must be <= 65535. + Per-row k2q counts may be anything in [0, num_q_blocks], including 0: unselected kv blocks + get exactly-zero dk/dv rows. """ if not _HAS_VSA_BWD_SM100A or not q.is_cuda: return False @@ -64,14 +85,16 @@ def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: return False if not q.is_contiguous(): return False + if q.shape[0] * (q.shape[1] if BHSD else q.shape[2]) > 65535: + return False # Metadata must be integer-typed so the wrapper's int32 conversion is value-preserving. if not variable_block_sizes.is_cuda or variable_block_sizes.dtype not in (torch.int32, torch.int64): return False - num_blocks = variable_block_sizes.numel() - if num_blocks == 0 or num_blocks % 2 != 0: + block = _block_size(q, variable_block_sizes) + if _BWD_BY_BLOCK.get(block) is None: return False - if _seqlen(q) != BLOCK * num_blocks: + if block == 64 and variable_block_sizes.numel() % 2 != 0: return False return True @@ -89,15 +112,21 @@ def block_sparse_attn_backward_sm100a_from_k2q( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Backward from k2q metadata already in hand (``invert_indices`` layout). - ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q64 - block ids, ``k2q_num`` ``[B, H, num_kv_blocks]``; entries past a row's count are never read. + ``k2q_idx`` is ``[B, H, num_kv_blocks, max_q_blocks]`` (or the flat 2-D view) of LOCAL q + block ids (64- or 128-token blocks, as the metadata's block size), ``k2q_num`` + ``[B, H, num_kv_blocks]``; entries past a row's count are never read. """ + block = _block_size(q, variable_block_sizes) + backward = _BWD_BY_BLOCK.get(block) + if backward is None: + raise RuntimeError(f"block_sparse_attn_backward_sm100a: no kernel for {block}-token " + f"blocks (built: {sorted(b for b, f in _BWD_BY_BLOCK.items() if f)})") sm_scale = 1.0 / (q.shape[-1]**0.5) idx = k2q_idx.to(torch.int32).contiguous() num = k2q_num.to(torch.int32).contiguous() vbs = variable_block_sizes.to(torch.int32).contiguous() - res = _BWD(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), - o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale) + res = backward(grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), + o.contiguous(), lse.contiguous(), idx, num, vbs, sm_scale) return res[0], res[1], res[2] diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py index 84872f52d..b2d2df6fe 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -8,7 +8,7 @@ A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This mo forward: it returns ``(out, lse)`` with ``lse`` in exactly the form ``triton_block_sparse_attn_forward`` writes -- ``max(qk * qk_scale) + log2(l)``, ``[B, H, S]`` fp32 -- so both ``block_sparse_attn_backward_triton`` and the sm_100a CUDA backward -(``block_sparse_attn_bwd_sm100a``, 64-token blocks, sm_100a devices only) run against it +(``block_sparse_attn_bwd_sm100a``, 64- and 128-token blocks, sm_100a / sm_103a) run against it unchanged. The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks @@ -43,9 +43,10 @@ BHSD = True def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + # A plain int: dynamic-shape torch.compile hands SymInt shapes, which cannot key _FWD_BY_BLOCK. num_blocks = variable_block_sizes.numel() seqlen = q.shape[2] if BHSD else q.shape[1] - return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks + return 0 if num_blocks == 0 or seqlen % num_blocks else int(seqlen // num_blocks) def is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: diff --git a/tests/test_block_sparse_bwd_sm100a.py b/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py similarity index 69% rename from tests/test_block_sparse_bwd_sm100a.py rename to fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py index b9daea240..832bef8df 100644 --- a/tests/test_block_sparse_bwd_sm100a.py +++ b/fastvideo-kernel/tests/test_block_sparse_bwd_sm100a.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -"""Correctness tests for the sm_100a CUDA block-sparse VSA backward. +"""Correctness tests for the sm_100a CUDA block-sparse VSA backward, 64- and 128-token blocks. Reference: fp32 dense attention restricted to the selected blocks, keys past variable_block_sizes masked to -inf, differentiated with torch autograd on fp32 copies of the bf16 inputs with loss = (out * grad_o).sum(). The kernel is fed the reference's lse in Triton's M format (logsumexp * log2e) and the reference output rounded to bf16 as ``o`` -- exactly what the sm_100a forward hands it in FastVideo. The k2q inversion is done here in torch (a stable -sort) so the tests do not need Triton. +sort) so the tests do not need Triton. Every case runs once per built block size. Run with: python -m pytest tests/test_block_sparse_bwd_sm100a.py -v """ @@ -20,13 +20,12 @@ import torch from fastvideo_kernel import block_sparse_attn_bwd_sm100a as bwd HEAD_DIM = 128 -BLOCK = 64 LOG2E = 1.4426950408889634 # Per-tensor tolerances on the bf16 outputs against the fp32 reference. Measured over all 15 -# cases on GB200 (2026-09-03, VSA_BWD_TEST_VERBOSE=1): max|diff|/max|ref| up to 3.0e-3 (dq) and -# 5.3e-3 (dk, dv); mean|diff| up to 2.8e-4 against mean|ref| of 0.06-0.12. The bounds below leave -# about 2x (rel max) and 3.5x (mean abs) headroom. +# blk64 cases on GB200 (2026-09-03, VSA_BWD_TEST_VERBOSE=1): max|diff|/max|ref| up to 3.0e-3 (dq) +# and 5.3e-3 (dk, dv); mean|diff| up to 2.8e-4 against mean|ref| of 0.06-0.12. The bounds below +# leave about 2x (rel max) and 3.5x (mean abs) headroom; the blk128 cases sit in the same range. REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref| MEAN_ABS_TOL = 1e-3 # mean|got - ref| @@ -34,18 +33,25 @@ pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in {(10, 0), (10, 3)} or not bwd._HAS_VSA_BWD_SM100A, reason="requires a data-center Blackwell GPU (sm_100a/sm_103a) and a fastvideo_kernel " - "extension built with block_sparse_sm100a_bwd", + "extension built with block_sparse_sm100a_bwd / block_sparse_sm100a_blk128_bwd", ) -def make_case(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): +@pytest.fixture(params=bwd.BLOCKS, ids=lambda b: f"blk{b}") +def block(request): + if bwd._BWD_BY_BLOCK.get(request.param) is None: + pytest.skip(f"no {request.param}-token backward in this build") + return request.param + + +def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): """Random bf16 q/k/v/grad_o plus q2k metadata [B, H, Nq, topk] / [B, H, Nq] and vbs. ``kv_pool`` restricts the kv blocks a q block may select (default: all), which is how the zero-count case leaves some kv blocks unselected by every q block. """ torch.manual_seed(seed) - S = num_blocks * BLOCK + S = num_blocks * block shape = (batch, heads, S, HEAD_DIM) if bwd.BHSD else (batch, S, heads, HEAD_DIM) q, k, v, grad_o = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(4)) @@ -59,10 +65,10 @@ def make_case(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_p num = torch.full((batch, heads, num_blocks), topk, dtype=torch.int32, device="cuda") if ragged: - vbs = torch.randint(BLOCK // 2, BLOCK + 1, (num_blocks, ), dtype=torch.int32, + vbs = torch.randint(block // 2, block + 1, (num_blocks, ), dtype=torch.int32, device="cuda") else: - vbs = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda") return q, k, v, grad_o, idx, num, vbs @@ -92,7 +98,7 @@ def invert_indices_torch(q2k_idx, q2k_num, num_kv_blocks, pad_value=0): counts.to(torch.int32).view(B, H, num_kv_blocks)) -def reference(q, k, v, grad_o, idx, num, vbs): +def reference(block, q, k, v, grad_o, idx, num, vbs): """fp32 masked-dense autograd reference. Returns (o bf16, lse fp32 [B, H, S] in M format, dq, dk, dv fp32) -- o and the grads in @@ -112,7 +118,7 @@ def reference(q, k, v, grad_o, idx, num, vbs): for j in range(int(num[b, h, qb])): kb = int(idx[b, h, qb, j]) valid = int(vbs[kb]) - keep[b, h, qb * BLOCK:(qb + 1) * BLOCK, kb * BLOCK:kb * BLOCK + valid] = True + keep[b, h, qb * block:(qb + 1) * block, kb * block:kb * block + valid] = True q32, k32, v32 = (t.detach().float().requires_grad_(True) for t in (q, k, v)) scores = (q32 @ k32.transpose(-1, -2)) * scale @@ -145,12 +151,13 @@ def check_close(name, got, ref): f"max|ref| = {ref.abs().max().item():.3e}, mean|ref| = {ref.abs().mean().item():.3e}") -def run_and_compare(num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, kv_pool=None): - q, k, v, grad_o, idx, num, vbs = make_case(num_blocks=num_blocks, topk=topk, heads=heads, - batch=batch, ragged=ragged, seed=seed, - kv_pool=kv_pool) +def run_and_compare(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0, + kv_pool=None): + q, k, v, grad_o, idx, num, vbs = make_case(block, num_blocks=num_blocks, topk=topk, + heads=heads, batch=batch, ragged=ragged, + seed=seed, kv_pool=kv_pool) assert bwd.is_supported(q, vbs) - o, lse, ref_dq, ref_dk, ref_dv = reference(q, k, v, grad_o, idx, num, vbs) + o, lse, ref_dq, ref_dk, ref_dv = reference(block, q, k, v, grad_o, idx, num, vbs) k2q_idx, k2q_num = invert_indices_torch(idx, num, num_blocks) dq, dk, dv = bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num, vbs) @@ -165,29 +172,29 @@ def _to_bhsd(t): return t if bwd.BHSD else t.transpose(1, 2) -def test_backward_matches_reference(): - run_and_compare() +def test_backward_matches_reference(block): + run_and_compare(block) -def test_batch_two(): - run_and_compare(batch=2) +def test_batch_two(block): + run_and_compare(block, batch=2) -def test_ragged_block_sizes(): +def test_ragged_block_sizes(block): """variable_block_sizes is what FastVideo always passes; padded keys must be masked.""" - run_and_compare(ragged=True) + run_and_compare(block, ragged=True) -def test_ragged_padded_key_rows_are_zero(): +def test_ragged_padded_key_rows_are_zero(block): """Keys at or past a block's count get P^T = 0 in-kernel, so their dk/dv rows are exactly 0 (Triton's backward stores zeros there as well).""" - (dq, dk, dv), _, vbs = run_and_compare(ragged=True, seed=1) + (dq, dk, dv), _, vbs = run_and_compare(block, ragged=True, seed=1) dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() checked = 0 for kb in range(vbs.numel()): - if int(vbs[kb]) == BLOCK: + if int(vbs[kb]) == block: continue # a full block has no padded rows - rows = slice(kb * BLOCK + int(vbs[kb]), (kb + 1) * BLOCK) + rows = slice(kb * block + int(vbs[kb]), (kb + 1) * block) assert dk[:, :, rows].abs().max().item() == 0.0, f"dk: padded rows of kv block {kb}" assert dv[:, :, rows].abs().max().item() == 0.0, f"dv: padded rows of kv block {kb}" checked += 1 @@ -195,27 +202,31 @@ def test_ragged_padded_key_rows_are_zero(): @pytest.mark.parametrize("topk", [1, 2, 3, 5, 7]) -def test_topk_not_a_multiple_of_the_quad(topk): - """The kernel walks each kv block's q list in quads of 4; a ragged tail quad must be exact.""" - run_and_compare(ragged=True, num_blocks=8, topk=topk) +def test_topk_not_a_multiple_of_the_quad(block, topk): + """blk64 walks each kv block's q list in quads of 4 (a ragged tail quad must be exact); + blk128 one q block per step (an odd list length must be exact).""" + run_and_compare(block, ragged=True, num_blocks=8, topk=topk) -@pytest.mark.parametrize("num_blocks", [2, 4, 8, 16]) -def test_sequence_lengths(num_blocks): - run_and_compare(ragged=True, num_blocks=num_blocks, topk=min(3, num_blocks)) +@pytest.mark.parametrize("num_blocks", [1, 2, 4, 8, 16]) +def test_sequence_lengths(block, num_blocks): + if block == 64 and num_blocks % 2: + pytest.skip("blk64 needs an even block count (128-token preprocess)") + run_and_compare(block, ragged=True, num_blocks=num_blocks, topk=min(3, num_blocks)) -def test_zero_count_kv_blocks(): - """kv blocks no q block selects: the main kernel skips them and the preprocess must write - their dk/dv rows as exact zeros (the outputs are empty_like), every other block exact.""" +def test_zero_count_kv_blocks(block): + """kv blocks no q block selects: their dk/dv rows must come back as exact zeros (the outputs + are empty_like; blk64's preprocess and blk128's main kernel write them), every other block + exact.""" num_blocks = 8 excluded = (2, 5) pool = [kb for kb in range(num_blocks) if kb not in excluded] - (dq, dk, dv), _, _ = run_and_compare(num_blocks=num_blocks, topk=4, ragged=True, + (dq, dk, dv), _, _ = run_and_compare(block, num_blocks=num_blocks, topk=4, ragged=True, kv_pool=pool) dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() for kb in excluded: - rows = slice(kb * BLOCK, (kb + 1) * BLOCK) + rows = slice(kb * block, (kb + 1) * block) assert dk[:, :, rows].abs().max().item() == 0.0, f"dk: unselected kv block {kb}" assert dv[:, :, rows].abs().max().item() == 0.0, f"dv: unselected kv block {kb}" @@ -223,7 +234,7 @@ def test_zero_count_kv_blocks(): def test_invert_indices_torch_matches_q2k(): """Guards the test's own k2q builder: every (q, kv) pair appears once, counts add up.""" num_blocks = 8 - _, _, _, _, idx, num, _ = make_case(num_blocks=num_blocks, topk=3, heads=2, batch=2) + _, _, _, _, idx, num, _ = make_case(64, num_blocks=num_blocks, topk=3, heads=2, batch=2) k2q_idx, k2q_num = invert_indices_torch(idx, num, num_blocks) B, H, Nq, Mk = idx.shape assert k2q_idx.shape == (B, H, num_blocks, Nq) and k2q_num.shape == (B, H, num_blocks) @@ -235,24 +246,31 @@ def test_invert_indices_torch_matches_q2k(): assert listed == sorted(listed) -def test_unsupported_is_rejected(): - q, k, v, grad_o, idx, num, vbs = make_case() +def test_unsupported_is_rejected(block): + q, k, v, grad_o, idx, num, vbs = make_case(block) assert not bwd.is_supported(q.float(), vbs) # wrong dtype assert not bwd.is_supported(q[..., :64].contiguous(), vbs) # wrong head_dim - seven = torch.full((7, ), 64, dtype=torch.int32, device="cuda") - assert not bwd.is_supported(q, seven) # seqlen != 64 * num_blocks + seven = torch.full((7, ), block, dtype=torch.int32, device="cuda") + assert not bwd.is_supported(q, seven) # seqlen != block * num_blocks q500 = (q[:, :, :500] if bwd.BHSD else q[:, :500]).contiguous() - assert not bwd.is_supported(q500, vbs) # S not a multiple of 64 - # An odd block count (S % 128 != 0) is refused statically too, so FastVideo falls back to - # Triton instead of tripping the binding's check. - q448 = (q[:, :, :448] if bwd.BHSD else q[:, :448]).contiguous() - assert not bwd.is_supported(q448, seven) + assert not bwd.is_supported(q500, vbs) # S not a multiple of the block + if block == 64: + # An odd block count (S % 128 != 0) is refused statically too, so FastVideo falls back + # to Triton instead of tripping the binding's check. + q448 = (q[:, :, :448] if bwd.BHSD else q[:, :448]).contiguous() + assert not bwd.is_supported(q448, seven) + # batch * heads is the launch grid's y extent: 65536 (batch, head) pairs are refused. + wide = (1, 65536, block, q.shape[-1]) if bwd.BHSD else (1, block, 65536, q.shape[-1]) + q_wide = torch.empty(wide, dtype=torch.bfloat16, device="cuda") + one = torch.full((1, ), block, dtype=torch.int32, device="cuda") + assert not bwd.is_supported(q_wide, one) + del q_wide # The binding itself refuses bad dtypes before touching the GPU. k2q_idx, k2q_num = invert_indices_torch(idx, num, vbs.numel()) o = torch.zeros_like(q) heads = q.shape[1] if bwd.BHSD else q.shape[2] - lse = torch.zeros((q.shape[0], heads, vbs.numel() * BLOCK), dtype=torch.float32, + lse = torch.zeros((q.shape[0], heads, vbs.numel() * block), dtype=torch.float32, device="cuda") with pytest.raises(RuntimeError): bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o.float(), q.float(), k.float(), diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py similarity index 80% rename from tests/test_block_sparse_sm100a_dispatch.py rename to fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py index 122aec2ea..73cca33fa 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/fastvideo-kernel/tests/test_block_sparse_sm100a_dispatch.py @@ -4,8 +4,9 @@ ``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when ``block_sparse_attn_sm100a.is_supported`` passes. Its backward is the sm_100a CUDA backward where that op is built and ``block_sparse_attn_bwd_sm100a.is_supported`` -passes (64-token blocks on an sm_100a device), the Triton backward otherwise; the -sm_100a lse is already in Triton's M format, so either pairing needs no conversion. +passes (64- or 128-token blocks on a data-center Blackwell device), the Triton backward +otherwise (64-token blocks only); the sm_100a lse is already in Triton's M format, so +either pairing needs no conversion. Everything else -- env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- must keep the pre-existing selection, bit-for-bit. @@ -229,10 +230,50 @@ def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): f"max|diff|/max|ref|={rel_max:.3e} mean|diff|={mean_abs:.3e}" -def test_blk128_backward_raises(monkeypatch): - """128-token blocks: forward runs, backward refuses (both backwards are 64-block only).""" +def _blk128_triton_reference_grads(q, k, v, idx, num, vbs, monkeypatch): + """All-Triton grads for 128-token metadata: the logical [B, H, Nq, Nk] map goes through + ``block_sparse_attn_128`` (Triton via the 128 -> 64 tile expansion, autograd-capable).""" + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128 + + B, H, Nq, Mk = idx.shape + rows = B * H * Nq + block_map = torch.zeros((rows, vbs.numel()), dtype=torch.bool, device="cuda") + valid = torch.arange(Mk, device="cuda").view(1, Mk) < num.view(rows, 1) + row_of = torch.arange(rows, device="cuda").view(rows, 1).expand(rows, Mk) + block_map[row_of[valid], idx.view(rows, Mk)[valid].long()] = True + block_map = block_map.view(B, H, Nq, vbs.numel()) + q2, k2, v2 = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False) + out2, _ = block_sparse_attn_128(q2, k2, v2, block_map, vbs) + out2.float().square().sum().backward() + monkeypatch.delenv("FASTVIDEO_VSA_TRITON") + return [t.grad.float() for t in (q2, k2, v2)] + + +def test_blk128_backward(monkeypatch): + """128-token blocks: with the 128-token backward built the sm_100a route runs it (the Triton + backward must not be entered) and its grads match the all-Triton 128 wrapper; without it the + backward raises instead of silently falling back to a 64-token kernel.""" + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + + dispatch = importlib.import_module("fastvideo_kernel.block_sparse_attn") monkeypatch.setenv(ENV, "1") q, k, v, idx, num, vbs = make_case(128, requires_grad=True) + if not vsa_bwd.is_supported(q, vbs): + out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) + with pytest.raises(RuntimeError, match="128-token"): + out.float().sum().backward() + return + + # The Triton reference first: the guard below disables the Triton backward for the route. + ref = _blk128_triton_reference_grads(q, k, v, idx, num, vbs, monkeypatch) + + def no_triton_backward(*args, **kwargs): + raise AssertionError("Triton backward entered on a supported 128-token input") + + monkeypatch.setattr(dispatch, "block_sparse_attn_backward_triton", no_triton_backward) out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) - with pytest.raises(RuntimeError, match="64-token"): - out.float().sum().backward() + out.float().square().sum().backward() + got = [t.grad.float() for t in (q, k, v)] + _assert_grads_close(got, ref) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index e37e6d5a1..a6f195ece 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -420,7 +420,7 @@ def _sm100a_unavailable_reason(sm100a_mod: Any, query_bhsd: torch.Tensor, variab if sm100a_mod is None: return "fastvideo_kernel.block_sparse_attn_sm100a is not installed" if grad_mode: - return "inputs require grad and the sm_100a/sm_103a kernel is forward-only; grad paths keep Triton" + return "inputs require grad; this backend keeps grad paths on Triton" if not sm100a_mod.is_supported(query_bhsd, variable_block_sizes): return ("block_sparse_attn_sm100a.is_supported returned False (needs an sm_100 or sm_103 device, a built " "extension, bf16, head_dim 128, an even tile count, and integer tile sizes)") diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 2deaed9f4..5b03fabe9 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -120,7 +120,7 @@ def test_reason_covers_every_precondition(): vbs = torch.full((2, ), 64, dtype=torch.long) assert "not installed" in _sm100a_unavailable_reason(None, q, vbs, grad_mode=False) ok = _FakeSm100a(supported=True) - assert "forward-only" in _sm100a_unavailable_reason(ok, q, vbs, grad_mode=True) + assert "require grad" in _sm100a_unavailable_reason(ok, q, vbs, grad_mode=True) bad = _FakeSm100a(supported=False) assert "is_supported" in _sm100a_unavailable_reason(bad, q, vbs, grad_mode=False) assert _sm100a_unavailable_reason(ok, q, vbs, grad_mode=False) is None