diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index ac4b30aea7..a7f9682428 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -411,6 +411,12 @@ if(BUILD_CXX_KERNELS) set_source_files_properties(csrc/attention/block_sparse_sm100a.cu csrc/attention/block_sparse_blk128_sm100a.cu PROPERTIES COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-gencode;arch=compute_103a,code=sm_103a;-DVSA_BHSD=true") + # VSA block-sparse attention BACKWARD, 64-token blocks. sm_100a only for now: validated + # on GB200, not yet on B300/GB300, so no sm_103a image is built and the Python side keeps + # the Triton backward for sm_103a devices. + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_bwd_sm100a.cu) + set_source_files_properties(csrc/attention/block_sparse_bwd_sm100a.cu PROPERTIES + COMPILE_OPTIONS "-gencode;arch=compute_100a,code=sm_100a;-DVSA_BHSD=true") endif() Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh new file mode 100644 index 0000000000..86d738cacd --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_kernel_sm100a.cuh @@ -0,0 +1,1087 @@ +// block_sparse_bwd_kernel_sm100a.cuh -- VSA block-sparse attention BACKWARD (blk64 +// one-pass), sm_100a. Warp-specialized: load / MMA (tcgen05) / softmax (P^T, dS^T) / +// epilogue (dQ drain) / scheduler. Three kernels: preprocess (Delta, Q^T, dO^T, dqaccum +// zero), main (dK, dV, dQ partials), postprocess (dQ unscramble + scale). +#ifndef BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_CUH +#define BLOCK_SPARSE_VSA_BWD_KERNEL_SM100A_CUH + +#include +#include +#include +#include +#include +#include +#include "primitives.cuh" + +namespace vsa_bwd_blk64 { + +constexpr int BLOCK = 64; +constexpr int KV_TILE = BLOCK; +constexpr int QBLOCKS_PER_QUAD = 4; +constexpr int Q_QUAD = QBLOCKS_PER_QUAD * BLOCK; +constexpr int M_TILE = KV_TILE; +[[maybe_unused]] constexpr int K_TILE = Q_QUAD; +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 KV_SUBTILES = HEAD_DIM / SUB_COLS_BF16; +constexpr int KV_SUB_COLS_BYTES = KV_TILE * SUB_COLS_BYTES; +constexpr int KV_TILE_BYTES = KV_SUBTILES * KV_SUB_COLS_BYTES; +constexpr int Q_BLK_BYTES = HEAD_DIM * SUB_COLS_BYTES; +constexpr int QBLOCKS_PER_SLOT = 2; +constexpr int Q_RING_SLOT_BYTES = QBLOCKS_PER_SLOT * Q_BLK_BYTES; +constexpr int SLOTS_PER_QUAD = QBLOCKS_PER_QUAD / QBLOCKS_PER_SLOT; +constexpr int PRE_QBLOCKS = 2; +constexpr int PRE_TOKENS = PRE_QBLOCKS * BLOCK; +constexpr int NUM_Q_STAGES = 2 * SLOTS_PER_QUAD; +constexpr int DOT_SLOT0 = 0; +constexpr int QT_SLOT0 = SLOTS_PER_QUAD; +constexpr int DST_TILE_BYTES = KV_TILE * BLOCK * (int)sizeof(__nv_bfloat16); +constexpr int DST_TILES = 4; +constexpr int DST_BYTES = DST_TILES * DST_TILE_BYTES; +constexpr int MMA_K = 16; +constexpr int K_ATOMS_PER_SUBTILE = SUB_COLS_BF16 / MMA_K; +constexpr int K_ATOMS_PER_QBLOCK = BLOCK / MMA_K; +constexpr int K_ATOMS_PER_KV_TILE = KV_TILE / MMA_K; +constexpr int BF16X2_COLS_PER_K16 = MMA_K / 2; + +template +struct DQConfig { + static constexpr int COLS = sizeof(DQ_DTYPE) == 2 ? 64 : 32; + static constexpr int DQ_ONE_PUSH_BYTES = BLOCK * COLS * (int)sizeof(DQ_DTYPE); + static constexpr int WARP_ELEMS = 32 * COLS; + static constexpr int DQ_QBLOCKS_PER_STAGE = 2; + static constexpr int DQ_STAGE_BYTES = DQ_QBLOCKS_PER_STAGE * DQ_ONE_PUSH_BYTES; + static constexpr int DQ_STAGE_BUFFERS = 2; + static constexpr int DQ_BLOCK_ELEMS = BLOCK * HEAD_DIM; +}; +static_assert(DQConfig::DQ_STAGE_BYTES == DQConfig::DQ_STAGE_BYTES, + "f16 and fp32 pushes must match"); + +constexpr int N_WARPS = 16; +constexpr int W_EPI0 = 0, W_SOFTMAX0 = 4, W_MMA = 12, W_LOAD = 13, W_SCHED = 14; +constexpr int CLC_STAGES = 2; +constexpr int CLC_ARRIVALS = 15; + +constexpr int NUM_BARS = 2 * NUM_Q_STAGES + 14 + 2 * CLC_STAGES; +constexpr int SMEM_TOTAL = 2 * KV_TILE_BYTES + NUM_Q_STAGES * Q_RING_SLOT_BYTES + DST_BYTES + + DQConfig<>::DQ_STAGE_BUFFERS * DQConfig<>::DQ_STAGE_BYTES + + 2 * Q_QUAD * (int)sizeof(float) + NUM_BARS * 8 + CLC_STAGES * 16 + 48; + +constexpr int ST_COLS = Q_QUAD / 2; +constexpr int ST_QBLOCK_COLS = BLOCK; +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_QBLOCK_COLS, "two q blocks per lane-half"); +static_assert(TMEM_TOTAL == 512, "TMEM map must fill exactly 512 columns"); + +extern __shared__ __align__(1024) uint8_t bwd_smem[]; + +struct WorkItem { + int batch; + int head; + int kv_block_id_in_seq; + const int* local_k2q_idx; + int local_k2q_num; + int num_quads; +}; + +__device__ __forceinline__ WorkItem decode_workitem( + int workitem_id, const int* __restrict__ workitem_remap, 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; + const int real_item_id = workitem_remap ? workitem_remap[workitem_id] : workitem_id; + const int batch_head = real_item_id / num_kv_blocks_per_seq; + it.batch = batch_head / num_heads; + it.head = batch_head % num_heads; + it.kv_block_id_in_seq = real_item_id % num_kv_blocks_per_seq; + it.local_k2q_idx = k2q_idx + (size_t)real_item_id * (size_t)max_q_blocks; + it.local_k2q_num = k2q_num[real_item_id]; + static_assert(QBLOCKS_PER_QUAD == 4, "num_quads uses >> 2"); + it.num_quads = (it.local_k2q_num + 3) >> 2; + return it; +} + +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; +} + +template +__global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) vsa_bwd_main_kernel( + const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v, + const __grid_constant__ CUtensorMap tmap_qt, const __grid_constant__ CUtensorMap tmap_dot, + const __grid_constant__ CUtensorMap tmap_dk, const __grid_constant__ CUtensorMap tmap_dv, + DQ_DTYPE* __restrict__ dqaccum, const float* __restrict__ lse_rows, + const float* __restrict__ delta_rows, const int* __restrict__ k2q_idx, + const int* __restrict__ k2q_num, const int* __restrict__ workitem_remap, + const int* __restrict__ variable_block_sizes, int max_q_blocks, 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)) + using DQ = DQConfig; + const int num_kv_blocks_per_seq = seqlen / BLOCK; + + uint8_t* sK = bwd_smem; + uint8_t* sV = sK + KV_TILE_BYTES; + uint8_t* sRING = sV + KV_TILE_BYTES; + uint8_t* sDST = sRING + NUM_Q_STAGES * Q_RING_SLOT_BYTES; + uint8_t* sDQ_STAGE_bytes = sDST + DST_BYTES; + DQ_DTYPE* sDQ_STAGE[2] = {reinterpret_cast(sDQ_STAGE_bytes), + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BYTES)}; + __nv_bfloat16* sDV_STAGE = reinterpret_cast<__nv_bfloat16*>(sRING); + __nv_bfloat16* sDK_STAGE = reinterpret_cast<__nv_bfloat16*>(sRING + KV_TILE_BYTES); + float* sLSE = + reinterpret_cast(sDQ_STAGE_bytes + DQ::DQ_STAGE_BUFFERS * DQ::DQ_STAGE_BYTES); + float* sDelta = sLSE + Q_QUAD; + + uint64_t* full_bar_ring = reinterpret_cast(sDelta + Q_QUAD); + uint64_t* empty_bar_ring = full_bar_ring + NUM_Q_STAGES; + uint64_t* full_bar_lse = empty_bar_ring + NUM_Q_STAGES; + uint64_t* empty_bar_lse = full_bar_lse + 1; + uint64_t* full_bar_delta = empty_bar_lse + 1; + 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 (warp_id == 0) { + tcgen05_alloc<1>(smem_ptr_u32(tmem_slot), TMEM_TOTAL); + tcgen05_relinquish_alloc_permit<1>(); + } + __syncthreads(); + + 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; + + if (tid == 0) { + #pragma unroll + for (int s = 0; s < NUM_Q_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar_ring[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_ring[s]), 1); + } + mbarrier_init(smem_ptr_u32(full_bar_lse), 1); + mbarrier_init(smem_ptr_u32(empty_bar_lse), 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); + mbarrier_init(smem_ptr_u32(empty_bar_kv), 1); + mbarrier_init(smem_ptr_u32(empty_bar_epi), 256); + if constexpr (USE_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(); + + [[maybe_unused]] const int total_workitems = num_samples * num_heads * num_kv_blocks_per_seq; + + if (warp_id == W_LOAD) { + setmaxnreg_dec<88>(); + + EmptyPhaseTracker ring_empty_ph; + EmptyPhaseTracker<1> lse_empty_ph, delta_empty_ph, kv_empty_ph, epi_empty_ph; + + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (workitem_id >= 0) { + const WorkItem it = decode_workitem(workitem_id, workitem_remap, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + + if (it.local_k2q_num != 0) { + auto get_global_qblock_id = [&](int quad_idx, int qblock_id_in_quad) { + return it.local_k2q_idx[min(QBLOCKS_PER_QUAD * quad_idx + qblock_id_in_quad, + it.local_k2q_num - 1)]; + }; + + auto load_kv_tile = [&](uint8_t* dst, const CUtensorMap* map, uint64_t* full_bar) { + #pragma unroll + for (int s = 0; s < KV_SUBTILES; ++s) { + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32(dst + s * KV_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), 0, + it.kv_block_id_in_seq * KV_TILE, s, it.batch * num_heads + it.head); + else + tma_load_3d(smem_ptr_u32(dst + s * KV_SUB_COLS_BYTES), map, smem_ptr_u32(full_bar), 0, + it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * KV_SUBTILES + s); + } + }; + + auto load_quad = [&](const CUtensorMap* map, int quad_idx, bool with_kv) { + if (with_kv) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_kv), kv_empty_ph.get_phase()); + kv_empty_ph.advance(); + } + + for (int pair_idx = 0; pair_idx < SLOTS_PER_QUAD; ++pair_idx) { + const int slot = ring_empty_ph.get_stage(); + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_ring[slot]), + ring_empty_ph.get_phase()); + ring_empty_ph.advance(); + + const int qblock_in_quad = QBLOCKS_PER_SLOT * pair_idx; + const int global_qblock_id0 = get_global_qblock_id(quad_idx, qblock_in_quad); + const int global_qblock_id1 = get_global_qblock_id(quad_idx, qblock_in_quad + 1); + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_ring[slot]), + Q_RING_SLOT_BYTES + (with_kv ? KV_TILE_BYTES : 0)); + tma_load_2d(smem_ptr_u32(sRING + slot * Q_RING_SLOT_BYTES), map, + smem_ptr_u32(&full_bar_ring[slot]), + it.batch * seqlen + global_qblock_id0 * BLOCK, it.head * HEAD_DIM); + tma_load_2d(smem_ptr_u32(sRING + slot * Q_RING_SLOT_BYTES + Q_BLK_BYTES), map, + smem_ptr_u32(&full_bar_ring[slot]), + it.batch * seqlen + global_qblock_id1 * BLOCK, it.head * HEAD_DIM); + if (with_kv) { + if (pair_idx == 0) + load_kv_tile(sK, &tmap_k, &full_bar_ring[slot]); + else + load_kv_tile(sV, &tmap_v, &full_bar_ring[slot]); + } + } + } + }; + + auto load_lse_and_delta = [&](int quad_idx) { + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_lse), lse_empty_ph.get_phase()); + lse_empty_ph.advance(); + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(full_bar_lse), Q_QUAD * 4); + #pragma unroll + for (int i = 0; i < QBLOCKS_PER_QUAD; ++i) + cpasync_bulk_load_mbarrier(smem_ptr_u32(sLSE + i * BLOCK), + lse_rows + + (size_t)(it.batch * num_heads + it.head) * seqlen + + (size_t)get_global_qblock_id(quad_idx, i) * BLOCK, + BLOCK * 4, smem_ptr_u32(full_bar_lse)); + } + + 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_QUAD * 4); + #pragma unroll + for (int i = 0; i < QBLOCKS_PER_QUAD; ++i) + cpasync_bulk_load_mbarrier(smem_ptr_u32(sDelta + i * BLOCK), + delta_rows + + (size_t)(it.batch * num_heads + it.head) * seqlen + + (size_t)get_global_qblock_id(quad_idx, i) * BLOCK, + BLOCK * 4, smem_ptr_u32(full_bar_delta)); + } + }; + + mbarrier_wait_parity_suspend(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + load_quad(&tmap_dot, 0, false); + load_quad(&tmap_qt, 0, true); + load_lse_and_delta(0); + for (int j = 1; j < it.num_quads; ++j) { + load_quad(&tmap_dot, j, false); + load_lse_and_delta(j); + load_quad(&tmap_qt, j, false); + } + } + + if constexpr (USE_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 { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } + return; + } else if (warp_id == W_MMA) { + setmaxnreg_dec<88>(); + + const uint32_t lead = elect_one_sync() ? 1u : 0u; + const uint32_t idesc_st_dpt = make_idesc_bf16_f32(M_TILE, 2 * BLOCK, false, true); + const uint32_t idesc_dv_dk = make_idesc_bf16_f32(M_TILE, 2 * HEAD_DIM, false, false); + const uint32_t idesc_dq = make_idesc_bf16_f32(Q_QUAD / 2, HEAD_DIM, true, true); + + constexpr uint32_t DESC_SBO = 1024, DESC_LBO = 16; + auto make_smem_desc = [](const uint8_t* 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_ring = make_smem_desc(sRING, DESC_LBO); + const uint64_t desc_ring_mn = make_smem_desc(sRING, Q_BLK_BYTES); + const uint64_t desc_k_mn = make_smem_desc(sK, KV_SUB_COLS_BYTES); + const uint64_t desc_dst0 = make_smem_desc(sDST, DST_TILE_BYTES); + const uint64_t desc_dst1 = make_smem_desc(sDST + 2 * DST_TILE_BYTES, DST_TILE_BYTES); + + constexpr uint32_t K16_ROWS_DELTA = (MMA_K * SUB_COLS_BYTES) >> 4; + constexpr uint64_t K16_COLS_DELTA = (MMA_K * (int)sizeof(__nv_bfloat16)) >> 4; + constexpr uint64_t KV_SUB_COLS_DELTA = KV_SUB_COLS_BYTES >> 4; + constexpr uint64_t RING_DELTA = Q_RING_SLOT_BYTES >> 4; + + PhaseTracker<1> pt_ph, dst_ph; + PhaseTracker<1> dq_empty_ph; + EmptyPhaseTracker<1> epi_empty_ph; + PhaseTracker<1> ring_full_ph; + + auto gemm12_st_dpt = [&](auto is_st_const) { + 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; + uint64_t* commit_bar = is_st ? full_bar_st : full_bar_dpt; + #pragma unroll + for (int u = 0; u < SLOTS_PER_QUAD; ++u) { + const int slot = (is_st ? QT_SLOT0 : DOT_SLOT0) + u; + mbarrier_wait_parity(smem_ptr_u32(&full_bar_ring[slot]), ring_full_ph.get_phase()); + #pragma unroll + for (int s = 0; s < KV_SUBTILES; ++s) { + const uint64_t da = da_base + (uint64_t)s * KV_SUB_COLS_DELTA; + const uint64_t db = desc_ring_mn + (uint64_t)slot * RING_DELTA + + (uint64_t)s * K_ATOMS_PER_SUBTILE * K16_ROWS_DELTA; + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_SUBTILE; ++ki) { + const bool enable_d = (s != 0) || (ki != 0); + tcgen05_mma_ws_f16_ss_1sm_predicated( + lead, tmem_acc + (uint32_t)(BLOCK * u), da + ki * K16_COLS_DELTA, + db + (uint64_t)ki * K16_ROWS_DELTA, idesc_st_dpt, enable_d); + } + } + } + tcgen05_commit1_lead(lead, smem_ptr_u32(commit_bar)); + }; + + auto gemm35_dv_dk = [&](auto is_dv_const, 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; + #pragma unroll + for (int p = 0; p < SLOTS_PER_QUAD; ++p) { + const int slot = (is_dv ? DOT_SLOT0 : QT_SLOT0) + p; + const uint64_t db = desc_ring + (uint64_t)slot * RING_DELTA; + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_QBLOCK; ++ki) { + const int a = p * K_ATOMS_PER_QBLOCK + ki; + const uint32_t tmem_a = + tmem_a_base + (uint32_t)(p * ST_QBLOCK_COLS + ki * BF16X2_COLS_PER_K16); + const bool accumulate = (!first) || (a != 0); + tcgen05_mma_ws_f16_ts_1sm_predicated(lead, tmem_acc, tmem_a, db + ki * K16_COLS_DELTA, + idesc_dv_dk, accumulate); + } + tcgen05_commit1_lead(lead, smem_ptr_u32(&empty_bar_ring[slot])); + } + }; + + auto gemm4_dq_half = [&](uint64_t desc_dst_h) { + uint64_t adst = desc_dst_h; + 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)); + }; + + auto gemm4_dq = [&]() { + gemm4_dq_half(desc_dst0); + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + gemm4_dq_half(desc_dst1); + }; + + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (workitem_id >= 0) { + const WorkItem it = decode_workitem(workitem_id, workitem_remap, 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{}); + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + gemm12_st_dpt(std::false_type{}); + ring_full_ph.advance(); + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + mbarrier_wait_parity(smem_ptr_u32(empty_bar_epi), epi_empty_ph.get_phase()); + epi_empty_ph.advance(); + gemm35_dv_dk(std::true_type{}, true); + + for (int j = 0; j < it.num_quads - 1; ++j) { + mbarrier_wait_parity(smem_ptr_u32(full_bar_dst), dst_ph.get_phase()); + dst_ph.advance(); + gemm35_dv_dk(std::false_type{}, j == 0); + gemm4_dq(); + gemm12_st_dpt(std::true_type{}); + mbarrier_wait_parity(smem_ptr_u32(empty_bar_dq), dq_empty_ph.get_phase()); + dq_empty_ph.advance(); + gemm12_st_dpt(std::false_type{}); + ring_full_ph.advance(); + mbarrier_wait_parity(smem_ptr_u32(full_bar_pt), pt_ph.get_phase()); + pt_ph.advance(); + gemm35_dv_dk(std::true_type{}, false); + } + + 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{}, it.num_quads == 1); + tcgen05_commit1_lead(lead, smem_ptr_u32(full_bar_dk)); + gemm4_dq(); + tcgen05_commit1_lead(lead, smem_ptr_u32(empty_bar_kv)); + } + + if constexpr (USE_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 { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } + + bar_sync<10>(416); + tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL); + return; + } else if (warp_id == W_SCHED) { + setmaxnreg_dec<88>(); + + if constexpr (USE_CLC) { + 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); + } + } + return; + } else if (warp_id >= W_SOFTMAX0 && warp_id < W_MMA) { + setmaxnreg_inc<136>(); + + PhaseTracker<1> st_ph, dpt_ph, lse_ph, delta_ph, dv_ph, dk_ph; + constexpr int HALF_COLS = SUB_COLS_BF16; + static_assert(ST_COLS == 2 * HALF_COLS && DV_COLS == 2 * HALF_COLS, "two column halves"); + + 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 q_half = lane_group >> 1; + const int kv_row = (lane_group & 1) * 32 + lane; + const int qblock_in_quad = 2 * col_half + q_half; + + 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_QBLOCK_COLS); + const uint32_t tmem_bf16x2_offset = tmem_f32_offset; + + const float2* lse2 = reinterpret_cast(sLSE + qblock_in_quad * BLOCK); + const float2* delta2 = reinterpret_cast(sDelta + qblock_in_quad * BLOCK); + + constexpr int CHUNK_BF16 = 16 / (int)sizeof(__nv_bfloat16); + constexpr int CHUNKS_PER_ROW = SUB_COLS_BF16 / CHUNK_BF16; + __nv_bfloat16* sdst_row = + reinterpret_cast<__nv_bfloat16*>(sDST + (size_t)(2 * q_half + col_half) * DST_TILE_BYTES) + + kv_row * SUB_COLS_BF16; + + auto merge_lane_halves_and_store = [&](uint32_t tmem_acc, auto apply_sm_scale_const, + const CUtensorMap* map, __nv_bfloat16* stage_tile, + const WorkItem& it) { + constexpr bool apply_sm_scale = decltype(apply_sm_scale_const)::value; + __nv_bfloat16* stage_row = + stage_tile + (size_t)col_half * (KV_TILE * SUB_COLS_BF16) + kv_row * SUB_COLS_BF16; + uint32_t acc_regs[HALF_COLS]; + tcgen05_ld_32x32b_x64(tmem_acc + tmem_f32_offset, acc_regs); + tcgen05_wait_ld(); + tcgen05_fence_before_thread_sync(); + const float2* acc2 = reinterpret_cast(acc_regs); + const float2 scale2 = f32x2_splat(sm_scale); + + if (q_half == 1) { + #pragma unroll + for (int v = 0; v < CHUNKS_PER_ROW; ++v) { + const float2* a2 = acc2 + v * (CHUNK_BF16 / 2); + const float2 r0 = apply_sm_scale ? fmul2(a2[0], scale2) : a2[0]; + const float2 r1 = apply_sm_scale ? fmul2(a2[1], scale2) : a2[1]; + const float2 r2 = apply_sm_scale ? fmul2(a2[2], scale2) : a2[2]; + const float2 r3 = apply_sm_scale ? fmul2(a2[3], scale2) : a2[3]; + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(r0.x, r0.y); + packed.y = cvt_f32x2_to_bf16x2(r1.x, r1.y); + packed.z = cvt_f32x2_to_bf16x2(r2.x, r2.y); + packed.w = cvt_f32x2_to_bf16x2(r3.x, r3.y); + *reinterpret_cast(stage_row + (v ^ (kv_row & 7)) * CHUNK_BF16) = packed; + } + } + bar_sync<14>(256); + + if (q_half == 0) { + #pragma unroll + for (int v = 0; v < CHUNKS_PER_ROW; ++v) { + const float2* a2 = acc2 + v * (CHUNK_BF16 / 2); + uint4* chunk_ptr = reinterpret_cast(stage_row + (v ^ (kv_row & 7)) * CHUNK_BF16); + const uint4 staged = *chunk_ptr; + const __nv_bfloat162* staged_pairs = reinterpret_cast(&staged); + const float2 s0 = __bfloat1622float2(staged_pairs[0]); + const float2 s1 = __bfloat1622float2(staged_pairs[1]); + const float2 s2 = __bfloat1622float2(staged_pairs[2]); + const float2 s3 = __bfloat1622float2(staged_pairs[3]); + const float2 r0 = apply_sm_scale ? ffma2(a2[0], scale2, s0) : fadd2(a2[0], s0); + const float2 r1 = apply_sm_scale ? ffma2(a2[1], scale2, s1) : fadd2(a2[1], s1); + const float2 r2 = apply_sm_scale ? ffma2(a2[2], scale2, s2) : fadd2(a2[2], s2); + const float2 r3 = apply_sm_scale ? ffma2(a2[3], scale2, s3) : fadd2(a2[3], s3); + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(r0.x, r0.y); + packed.y = cvt_f32x2_to_bf16x2(r1.x, r1.y); + packed.z = cvt_f32x2_to_bf16x2(r2.x, r2.y); + packed.w = cvt_f32x2_to_bf16x2(r3.x, r3.y); + *chunk_ptr = packed; + } + } + fence_proxy_async_shared(); + bar_sync<14>(256); + + if (softmax_warp_id == 0 && elect_one_sync()) { + #pragma unroll + for (int s = 0; s < KV_SUBTILES; ++s) { + const uint32_t src = smem_ptr_u32(stage_tile + (size_t)s * KV_TILE * SUB_COLS_BF16); + if constexpr (BHSD) + tma_store_4d(map, 0, it.kv_block_id_in_seq * KV_TILE, s, it.batch * num_heads + it.head, + src); + else + tma_store_3d(map, 0, it.batch * seqlen + it.kv_block_id_in_seq * KV_TILE, + it.head * KV_SUBTILES + s, src); + } + cp_async_bulk_commit_group(); + } + bar_sync<14>(256); + }; + + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (workitem_id >= 0) { + const WorkItem it = decode_workitem(workitem_id, workitem_remap, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + + if (it.local_k2q_num != 0) { + const bool kv_row_valid = kv_row < variable_block_sizes[it.kv_block_id_in_seq]; + + for (int j = 0; j < it.num_quads; ++j) { + const bool p_valid = + kv_row_valid && (qblock_in_quad < it.local_k2q_num - QBLOCKS_PER_QUAD * j); + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_lse), 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[HALF_COLS]; + tcgen05_ld_32x32b_x64(tmem_st + tmem_f32_offset, st_regs); + tcgen05_wait_ld(); + tcgen05_fence_before_thread_sync(); + + float2* pt_fp32 = reinterpret_cast(st_regs); + uint32_t pt_bf16x2[HALF_COLS / 2]; + const float2 scale2 = f32x2_splat(scale_log2); + #pragma unroll + for (int c = 0; c < 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 (!p_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)); + } + + 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[HALF_COLS]; + tcgen05_ld_32x32b_x64(tmem_dpt + tmem_f32_offset, dpt_regs); + tcgen05_wait_ld(); + tcgen05_fence_before_thread_sync(); + + const float2* dpt2 = reinterpret_cast(dpt_regs); + uint32_t (&dst_bf16x2)[HALF_COLS / 2] = + reinterpret_cast(st_regs); + #pragma unroll + for (int c = 0; c < HALF_COLS / 2; ++c) { + const float2 ds = + fmul2(pt_fp32[c], fadd2(dpt2[c], make_float2(-delta2[c].x, -delta2[c].y))); + dst_bf16x2[c] = cvt_f32x2_to_bf16x2(ds.x, ds.y); + } + + 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(); + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(full_bar_dst)); + mbarrier_arrive(smem_ptr_u32(empty_bar_delta)); + } + } + + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dv), dv_ph.get_phase()); + dv_ph.advance(); + merge_lane_halves_and_store(tmem_dv, std::false_type{}, &tmap_dv, sDV_STAGE, it); + mbarrier_wait_parity_suspend(smem_ptr_u32(full_bar_dk), dk_ph.get_phase()); + dk_ph.advance(); + merge_lane_halves_and_store(tmem_dk, std::true_type{}, &tmap_dk, sDK_STAGE, it); + + if (softmax_warp_id == 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 (USE_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 { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } + + bar_sync<10>(416); + return; + } else if (warp_id < W_SOFTMAX0) { + setmaxnreg_inc<152>(); + + const int epi_warp_id = warp_id - W_EPI0; + const uint32_t tmem_lane_base = (uint32_t)(epi_warp_id * 32) << 16; + const bool is_leader = epi_warp_id == 0 && lane == 0; + + uint64_t dqaccum_l2_policy = 0; + if constexpr (DQ_L2_KEEP) { + dqaccum_l2_policy = make_l2cache_policy_fractional_evict_last_unchanged(0.25f); + } + + auto reduce_add_push = [&](DQ_DTYPE* dst, uint32_t src_smem) { + if constexpr (sizeof(DQ_DTYPE) == 4) { + if constexpr (DQ_L2_KEEP) { + cpasync_reduce_bulk_add_f32_l2hint(dst, src_smem, DQ::DQ_ONE_PUSH_BYTES, + dqaccum_l2_policy); + } else { + cpasync_reduce_bulk_add_f32(dst, src_smem, DQ::DQ_ONE_PUSH_BYTES); + } + } else { + if constexpr (DQ_L2_KEEP) { + cpasync_reduce_bulk_add_f16_l2hint(dst, src_smem, DQ::DQ_ONE_PUSH_BYTES, + dqaccum_l2_policy); + } else { + cpasync_reduce_bulk_add_f16(dst, src_smem, DQ::DQ_ONE_PUSH_BYTES); + } + } + }; + + PhaseTracker<1> dq_full_ph; + if (elect_one_sync()) { + mbarrier_arrive(smem_ptr_u32(empty_bar_dq)); + } + + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (workitem_id >= 0) { + const WorkItem it = decode_workitem(workitem_id, workitem_remap, k2q_idx, k2q_num, + max_q_blocks, num_heads, num_kv_blocks_per_seq); + + if (it.local_k2q_num != 0) { + DQ_DTYPE* dqaccum_head = dqaccum + (size_t)(it.batch * num_heads + it.head) * + num_kv_blocks_per_seq * DQ::DQ_BLOCK_ELEMS; + + for (int j = 0; j < it.num_quads; ++j) { + #pragma unroll 1 + for (int h = 0; h < 2; ++h) { + const int qblock_lo = + it.local_k2q_idx[min(QBLOCKS_PER_QUAD * j + h, it.local_k2q_num - 1)]; + const int qblock_hi = + it.local_k2q_idx[min(QBLOCKS_PER_QUAD * j + h + 2, it.local_k2q_num - 1)]; + DQ_DTYPE* dqaccum_lo = dqaccum_head + (size_t)qblock_lo * DQ::DQ_BLOCK_ELEMS; + DQ_DTYPE* dqaccum_hi = dqaccum_head + (size_t)qblock_hi * 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])); + } + // The loads are asynchronous; the arrive below hands tmem_dq back to the MMA warp, + // so it must sit behind their completion (the register consumers alone are + // scoreboarded, the arrive is not). + 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); + DQ_DTYPE* stage_row = sDQ_STAGE[stage_buf] + epi_warp_id * DQ::WARP_ELEMS + lane * 4; + #pragma unroll + for (int v4 = 0; v4 < DQ::COLS / 4; ++v4) { + const float4 v = dq_row4[v4]; + if constexpr (sizeof(DQ_DTYPE) == 2) { + uint2 packed; + packed.x = cvt_f32x2_to_f16x2(v.x, v.y); + packed.y = cvt_f32x2_to_f16x2(v.z, v.w); + *reinterpret_cast(stage_row + v4 * 32 * 4) = packed; + } else { + *reinterpret_cast(stage_row + v4 * 32 * 4) = v; + } + } + fence_proxy_async_shared(); + bar_sync<11>(128); + + if (is_leader) { + const size_t slice_offset = (size_t)hd_slice * BLOCK * DQ::COLS; + const uint32_t stage_lo = smem_ptr_u32(sDQ_STAGE[stage_buf]); + reduce_add_push(dqaccum_lo + slice_offset, stage_lo); + reduce_add_push(dqaccum_hi + slice_offset, stage_lo + DQ::DQ_ONE_PUSH_BYTES); + cp_async_bulk_commit_group(); + cp_async_bulk_wait_group_read<1>(); + } + bar_sync<11>(128); + } + } + } + } + + if constexpr (USE_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 { + workitem_id += (int)gridDim.x; + if (workitem_id >= total_workitems) workitem_id = -1; + } + } + + if (is_leader) { + cp_async_bulk_wait_group_read<0>(); + } + bar_sync<11>(128); + bar_sync<10>(416); + return; + } else { + setmaxnreg_dec<24>(); + return; + } +#endif +} + +constexpr int ORDER_THREADS = 1024; +constexpr int ORDER_SMEM_MAX = 227 * 1024; +constexpr int ORDER_MAX_BLOCKS = ORDER_SMEM_MAX / (2 * (int)sizeof(int)); + +__global__ void __launch_bounds__(ORDER_THREADS, 1) + vsa_bwd_order_kernel(const int* __restrict__ k2q_idx, const int* __restrict__ k2q_num, + int max_q_blocks, int num_kv_blocks_per_seq, int order_bin, bool snake, + int* __restrict__ order_out) { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) + extern __shared__ int order_smem[]; + int* sbin = order_smem; + int* smid = order_smem + num_kv_blocks_per_seq; + const int base = (int)blockIdx.x * num_kv_blocks_per_seq; + + for (int i = (int)threadIdx.x; i < num_kv_blocks_per_seq; i += ORDER_THREADS) { + const int count = k2q_num[base + i]; + sbin[i] = count / order_bin; + smid[i] = snake ? (count > 0 ? k2q_idx[(size_t)(base + i) * max_q_blocks + count / 2] : -1) : 0; + } + __syncthreads(); + + const int i = (int)blockIdx.y * ORDER_THREADS + (int)threadIdx.x; + if (i < num_kv_blocks_per_seq) { + const int bin = sbin[i], mid = smid[i]; + const bool descending = snake && (bin & 1); + int rank = 0; + + for (int j = 0; j < num_kv_blocks_per_seq; ++j) { + const int bin_j = sbin[j]; + if (bin_j < bin) { + ++rank; + } else if (bin_j == bin) { + const int mid_j = smid[j]; + const bool before = (mid_j == mid) ? (j < i) : (descending ? (mid_j > mid) : (mid_j < mid)); + if (before) ++rank; + } + } + + order_out[base + rank] = base + i; + } +#endif +} + +template +__global__ void __launch_bounds__(256, 1) + vsa_bwd_preprocess_kernel(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ o, + const __nv_bfloat16* __restrict__ dout, + float* __restrict__ delta_rows, DQ_DTYPE* __restrict__ dqaccum, + __nv_bfloat16* __restrict__ qt, __nv_bfloat16* __restrict__ dot, + __nv_bfloat16* __restrict__ dk, __nv_bfloat16* __restrict__ dv, + const int* __restrict__ k2q_num, int num_samples, int num_heads, + int seqlen) { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) + __shared__ __align__(128) __nv_bfloat16 tile[PRE_TOKENS][SUB_COLS_BF16 + 4]; + const int num_kv_blocks_per_seq = seqlen / BLOCK; + const int token_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 size_t total_tokens = (size_t)num_samples * seqlen; + const int token_begin = token_block_id * PRE_TOKENS; + + uint4* dqaccum_zero_destination = + reinterpret_cast(dqaccum + (size_t)batch_head * seqlen * HEAD_DIM + + (size_t)token_block_id * PRE_TOKENS * HEAD_DIM); + const uint4 zero_uint4 = make_uint4(0u, 0u, 0u, 0u); + constexpr int DQ_ZERO_CHUNKS = (PRE_TOKENS * HEAD_DIM * (int)sizeof(DQ_DTYPE) / 16) / 256; + #pragma unroll + for (int chunk = 0; chunk < DQ_ZERO_CHUNKS; ++chunk) + dqaccum_zero_destination[chunk * 256 + threadIdx.x] = zero_uint4; + + static_assert(PRE_TOKENS == 2 * BLOCK, "one preprocess CTA covers two kv64 blocks"); + constexpr int KV_ZERO_CHUNKS = (BLOCK * HEAD_DIM / 8) / 256; + #pragma unroll + for (int kv_block_in_cta = 0; kv_block_in_cta < PRE_QBLOCKS; ++kv_block_in_cta) { + const int kv_block_id_in_seq = token_block_id * PRE_QBLOCKS + kv_block_in_cta; + if (k2q_num[batch_head * num_kv_blocks_per_seq + kv_block_id_in_seq] == 0) { + #pragma unroll + for (int chunk = 0; chunk < KV_ZERO_CHUNKS; ++chunk) { + const int vector_index = chunk * 256 + (int)threadIdx.x; + const int row = vector_index >> 4; + const int column = (vector_index & 15) * 8; + const size_t token_element_offset = + token_offset(batch, head, num_heads, seqlen, kv_block_id_in_seq * BLOCK + row) + + column; + *reinterpret_cast(dk + token_element_offset) = zero_uint4; + *reinterpret_cast(dv + token_element_offset) = zero_uint4; + } + } + } + + const int row_base = (int)threadIdx.x >> 4; + const int dimension_begin = ((int)threadIdx.x & 15) * 4; + #pragma unroll + for (int dimension_block = 0; dimension_block < HEAD_DIM; dimension_block += 64) { + #pragma unroll + for (int row_pass = 0; row_pass < 8; ++row_pass) { + const int row = row_base + row_pass * 16; + *reinterpret_cast(&tile[row][dimension_begin]) = *reinterpret_cast( + q + token_offset(batch, head, num_heads, seqlen, token_begin + row) + + dimension_block + dimension_begin); + } + __syncthreads(); + + const int dimension = (int)threadIdx.x >> 2; + const int token_group_begin = ((int)threadIdx.x & 3) * 4; + #pragma unroll + for (int row_pass = 0; row_pass < 8; ++row_pass) { + const int row = token_group_begin + row_pass * 16; + uint2 packed_tokens; + uint16_t* packed_halves = reinterpret_cast(&packed_tokens); + #pragma unroll + for (int token_in_group = 0; token_in_group < 4; ++token_in_group) + packed_halves[token_in_group] = + *reinterpret_cast(&tile[row + token_in_group][dimension]); + __nv_bfloat16* transposed_row = + qt + ((size_t)head * HEAD_DIM + dimension_block + dimension) * total_tokens; + *reinterpret_cast(transposed_row + (size_t)batch * seqlen + token_begin + row) = + packed_tokens; + } + __syncthreads(); + } + + float delta_accumulator[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + #pragma unroll + for (int dimension_block = 0; dimension_block < HEAD_DIM; dimension_block += 64) { + #pragma unroll + for (int row_pass = 0; row_pass < 8; ++row_pass) { + const int row = row_base + row_pass * 16; + const size_t token_element_offset = + token_offset(batch, head, num_heads, seqlen, token_begin + row) + dimension_block + + dimension_begin; + const uint2 dout_vector = *reinterpret_cast(dout + token_element_offset); + *reinterpret_cast(&tile[row][dimension_begin]) = dout_vector; + + const uint2 o_vector = *reinterpret_cast(o + token_element_offset); + const __nv_bfloat162* dout_pairs = reinterpret_cast(&dout_vector); + const __nv_bfloat162* o_pairs = reinterpret_cast(&o_vector); + #pragma unroll + for (int pair = 0; pair < 2; ++pair) { + const float2 dout_pair_as_float = __bfloat1622float2(dout_pairs[pair]); + const float2 o_pair_as_float = __bfloat1622float2(o_pairs[pair]); + delta_accumulator[row_pass] += + o_pair_as_float.x * dout_pair_as_float.x + o_pair_as_float.y * dout_pair_as_float.y; + } + } + __syncthreads(); + + const int dimension = (int)threadIdx.x >> 2; + const int token_group_begin = ((int)threadIdx.x & 3) * 4; + #pragma unroll + for (int row_pass = 0; row_pass < 8; ++row_pass) { + const int row = token_group_begin + row_pass * 16; + uint2 packed_tokens; + uint16_t* packed_halves = reinterpret_cast(&packed_tokens); + #pragma unroll + for (int token_in_group = 0; token_in_group < 4; ++token_in_group) + packed_halves[token_in_group] = + *reinterpret_cast(&tile[row + token_in_group][dimension]); + __nv_bfloat16* transposed_row = + dot + ((size_t)head * HEAD_DIM + dimension_block + dimension) * total_tokens; + *reinterpret_cast(transposed_row + (size_t)batch * seqlen + token_begin + row) = + packed_tokens; + } + __syncthreads(); + } + + #pragma unroll + for (int row_pass = 0; row_pass < 8; ++row_pass) { + float delta_sum = delta_accumulator[row_pass]; + #pragma unroll + for (int shuffle_offset = 8; shuffle_offset > 0; shuffle_offset >>= 1) + delta_sum += __shfl_down_sync(0xffffffffu, delta_sum, shuffle_offset, 16); + if (((int)threadIdx.x & 15) == 0) + delta_rows[(size_t)batch_head * seqlen + token_begin + row_base + row_pass * 16] = delta_sum; + } +#endif +} + +template +__global__ void __launch_bounds__(128, 1) + vsa_bwd_postprocess_kernel(const DQ_DTYPE* __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)) + 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 DQ_DTYPE* dqaccum_block = dqaccum + ((size_t)batch_head * (seqlen / BLOCK) + q_block_id) * + DQConfig::DQ_BLOCK_ELEMS; + const int row = (int)threadIdx.x & 63; + const int dimension_half = (int)threadIdx.x >> 6; + + uint32_t dq_packed_bf16_pairs[32]; + #pragma unroll + for (int dimension_in_half = 0; dimension_in_half < 64; dimension_in_half += 4) { + const int dimension = dimension_half * 64 + dimension_in_half; + const DQ_DTYPE* dqaccum_element = + dqaccum_block + + (dimension / DQConfig::COLS) * (BLOCK * DQConfig::COLS) + + ((row & 32) >> 5) * DQConfig::WARP_ELEMS + + ((dimension % DQConfig::COLS) >> 2) * 128 + (row & 31) * 4; + float2 dq_pair_low, dq_pair_high; + if constexpr (sizeof(DQ_DTYPE) == 2) { + const uint2 dq_four_halves = *reinterpret_cast(dqaccum_element); + dq_pair_low = __half22float2(*reinterpret_cast(&dq_four_halves.x)); + dq_pair_high = __half22float2(*reinterpret_cast(&dq_four_halves.y)); + } else { + const float4 dq_four_floats = *reinterpret_cast(dqaccum_element); + dq_pair_low = make_float2(dq_four_floats.x, dq_four_floats.y); + dq_pair_high = make_float2(dq_four_floats.z, dq_four_floats.w); + } + dq_packed_bf16_pairs[dimension_in_half / 2 + 0] = + cvt_f32x2_to_bf16x2(dq_pair_low.x * sm_scale, dq_pair_low.y * sm_scale); + dq_packed_bf16_pairs[dimension_in_half / 2 + 1] = + cvt_f32x2_to_bf16x2(dq_pair_high.x * sm_scale, dq_pair_high.y * sm_scale); + } + + const int token = q_block_id * BLOCK + row; + uint4* dq_destination = reinterpret_cast( + dq + token_offset(batch, head, num_heads, seqlen, token) + dimension_half * 64); + const uint4* dq_packed_uint4 = reinterpret_cast(dq_packed_bf16_pairs); + #pragma unroll + for (int vector_index = 0; vector_index < 8; ++vector_index) + dq_destination[vector_index] = dq_packed_uint4[vector_index]; +#endif +} + +} + +#endif diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh new file mode 100644 index 0000000000..559152eac3 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_launch_sm100a.cuh @@ -0,0 +1,330 @@ +// block_sparse_bwd_launch_sm100a.cuh -- host surface of the VSA block-sparse backward drop: +// argument struct, workspace sizes, the support predicate and the stream-chained launch +// (preprocess -> order -> main -> postprocess). Tensor maps are encoded per call (no static +// cache: a torch caller hands us fresh pointers every time). +#ifndef BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH +#define BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH + +#include +#include +#include "block_sparse_bwd_kernel_sm100a.cuh" + +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif +#ifndef VSA_BWD_DQ_F16 +#define VSA_BWD_DQ_F16 false +#endif +#ifndef VSA_BWD_USE_CLC +#define VSA_BWD_USE_CLC true +#endif + +namespace vsa_bwd_blk64 { + +#if VSA_BWD_DQ_F16 +using dq_accum_t = uint16_t; +#else +using dq_accum_t = float; +#endif + +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 / 64. + + // 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. + // [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). + const int* k2q_num; + // [nb] int32: valid kv tokens per block (<= 64); kv rows at or past the count are masked. + const int* variable_block_sizes; + + // Work order: which (batch, head, kv block) item each CTA processes. + // [B*H*nb] int32 work id -> item ((b*H + h)*nb + kv). nullptr: identity order below + // ORDER_MIN_KV_BLOCKS, else the launch computes the length-binned order into order_workspace. + const int* workitem_remap; + // [B*H*nb] int32; required when workitem_remap is nullptr and nb >= ORDER_MIN_KV_BLOCKS. + int* order_workspace; + + // Outputs, inputs' layout; dk/dv rows of unselected kv blocks are zeroed by the preprocess. + __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] dq_accum_t, drain-native; preprocess zeroes, main reduce-adds, postprocess reads. + dq_accum_t* dqaccum; + // [H*128, B*S] Q^T, written by the preprocess. + __nv_bfloat16* qt; + // [H*128, B*S] dO^T, written by the preprocess. + __nv_bfloat16* dot; + // [B*H*S] fp32 rowsum(bf16(o) * dout), written by the preprocess. + float* delta; + + int batch; + int num_heads; + // S; a multiple of 128 (the preprocess works in 128-token blocks). + int seqlen; + // Must be 128. + int head_dim; + // nb = seqlen / 64. + 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_order_bytes(int batch, int num_heads, + int num_kv_blocks_per_seq) { + return (size_t)batch * num_heads * num_kv_blocks_per_seq * sizeof(int); +} +__host__ inline size_t block_sparse_bwd_transposed_bytes(int batch, int num_heads, int seqlen) { + return (size_t)num_heads * HEAD_DIM * (size_t)batch * seqlen * sizeof(__nv_bfloat16); +} +__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); +} + +// Below this many kv blocks the identity order is as fast and the order kernel is not free. +constexpr int ORDER_MIN_KV_BLOCKS = 1024; + +// One head's dQ accumulator (S x HEAD_DIM x sizeof) beyond which chunked DQ_L2_KEEP launches pay. +constexpr size_t L2_RESIDENT_DQ_ACCUM_BYTES_PER_HEAD = size_t{128} << 20; + +// Safety guard: chunked launches index the order's sub-array, so the order kernel must run there. +static_assert(L2_RESIDENT_DQ_ACCUM_BYTES_PER_HEAD / (HEAD_DIM * sizeof(dq_accum_t)) >= + (size_t)ORDER_MIN_KV_BLOCKS * BLOCK, + "the L2 transition (chunked launches) must not sit below the order-kernel threshold"); + +__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 || args.num_kv_blocks_per_seq % PRE_QBLOCKS != 0) { + 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; + } + 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.qt || !args.dot || !args.delta) { + return cudaErrorInvalidValue; + } + if (!args.workitem_remap && args.num_kv_blocks_per_seq >= ORDER_MIN_KV_BLOCKS && + (!args.order_workspace || args.num_kv_blocks_per_seq > ORDER_MAX_BLOCKS)) { + return cudaErrorInvalidValue; + } + return cudaSuccess; +} + +// K, V, dK, dV tensor maps, one 64-token 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_kv_units(CUtensorMap* map, const __nv_bfloat16* ptr, int B, + int H, int S) { + CUresult r; + if (VSA_BHSD) { + uint64_t gd[4] = {(uint64_t)SUB_COLS_BF16, (uint64_t)S, (uint64_t)KV_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)BLOCK, 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 * KV_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)BLOCK, 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; +} + +template +__host__ inline cudaError_t launch_main(const BlockSparseVsaBwdArgs& args, const int* work_remap, + const CUtensorMap& tk, const CUtensorMap& tv, + const CUtensorMap& tqt, const CUtensorMap& tdot, + const CUtensorMap& tdk, const CUtensorMap& tdv, int sms, + cudaStream_t stream) { + auto kernel = vsa_bwd_main_kernel; + cudaError_t e = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_TOTAL); + if (e != cudaSuccess) { + return e; + } + const float scale_log2 = args.sm_scale * 1.4426950408889634f; + const int B = args.batch, H = args.num_heads, S = args.seqlen; + const int total = B * H * args.num_kv_blocks_per_seq; + if constexpr (USE_CLC) { + // Above the L2 transition, one SM-wide launch at a time keeps each list neighbourhood + // resident; below it one launch of every item lets CLC steal freely. + const int chunk = DQ_L2_KEEP ? std::min(sms, total) : total; + cudaLaunchConfig_t cfg = {}; + cfg.blockDim = dim3(N_WARPS * 32, 1, 1); + cfg.dynamicSmemBytes = SMEM_TOTAL; + cfg.stream = stream; + cudaLaunchAttribute at[1]; + at[0].id = cudaLaunchAttributeClusterDimension; + at[0].val.clusterDim.x = 1; + at[0].val.clusterDim.y = 1; + at[0].val.clusterDim.z = 1; + cfg.attrs = at; + cfg.numAttrs = 1; + if (work_remap == nullptr && chunk != total) { + return cudaErrorInvalidValue; + } + for (int base = 0; base < total; base += chunk) { + const int count = std::min(chunk, total - base); + cfg.gridDim = dim3((unsigned)count, 1, 1); + // A chunk starts at work id `base`: it gets the order's sub-array. + e = cudaLaunchKernelEx(&cfg, kernel, tk, tv, tqt, tdot, tdk, tdv, args.dqaccum, args.lse, + args.delta, args.k2q_idx, args.k2q_num, + work_remap ? work_remap + base : nullptr, + args.variable_block_sizes, args.max_q_blocks, B, H, S, scale_log2, + args.sm_scale); + if (e != cudaSuccess) { + return e; + } + } + return cudaSuccess; + } else { + const int grid = std::min(total, sms); + kernel<<>>( + tk, tv, tqt, tdot, tdk, tdv, args.dqaccum, args.lse, args.delta, args.k2q_idx, args.k2q_num, + work_remap, args.variable_block_sizes, args.max_q_blocks, B, H, S, scale_log2, + args.sm_scale); + return cudaGetLastError(); + } +} + +__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; + const long n_tokens = (long)B * S; + + CUtensorMap tk, tv, tqt, tdot, tdk, tdv; + if (make_tma_kv_units(&tk, args.k, B, H, S) != cudaSuccess) { + return cudaErrorInvalidValue; + } + if (make_tma_kv_units(&tv, args.v, B, H, S) != cudaSuccess) { + return cudaErrorInvalidValue; + } + if (make_tma_kv_units(&tdk, args.dk, B, H, S) != cudaSuccess) { + return cudaErrorInvalidValue; + } + if (make_tma_kv_units(&tdv, args.dv, B, H, S) != cudaSuccess) { + return cudaErrorInvalidValue; + } + // Q^T, dO^T ([H*hd rows, B*S cols], token contiguous): box [hd rows, BLOCK cols] = one q64 + // block per TMA. + if (make_tma_2d_tiled(&tqt, args.qt, H * HEAD_DIM, (int)n_tokens, HEAD_DIM, BLOCK, 2, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16) != cudaSuccess) { + return cudaErrorInvalidValue; + } + if (make_tma_2d_tiled(&tdot, args.dot, H * HEAD_DIM, (int)n_tokens, HEAD_DIM, BLOCK, 2, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16) != cudaSuccess) { + return cudaErrorInvalidValue; + } + + int dev = 0, sms = 0; + cudaError_t e = cudaGetDevice(&dev); + if (e != cudaSuccess) { + return e; + } + e = cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev); + if (e != cudaSuccess) { + return e; + } + + vsa_bwd_preprocess_kernel + <<>>( + args.q, args.o, args.dout, args.delta, args.dqaccum, args.qt, args.dot, args.dk, args.dv, + args.k2q_num, B, H, S); + e = cudaGetLastError(); + if (e != cudaSuccess) { + return e; + } + + const bool keep_dq_l2 = + (size_t)S * HEAD_DIM * sizeof(dq_accum_t) >= L2_RESIDENT_DQ_ACCUM_BYTES_PER_HEAD; + + // Work-item order: the caller's, else identity below ORDER_MIN_KV_BLOCKS, else the device order. + const int* work_remap = args.workitem_remap; + if (work_remap == nullptr && args.num_kv_blocks_per_seq >= ORDER_MIN_KV_BLOCKS) { + const int order_smem = 2 * args.num_kv_blocks_per_seq * (int)sizeof(int); + if (order_smem > 48 * 1024) { + e = cudaFuncSetAttribute(vsa_bwd_order_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + order_smem); + if (e != cudaSuccess) { + return e; + } + } + const unsigned item_chunks = + (unsigned)((args.num_kv_blocks_per_seq + ORDER_THREADS - 1) / ORDER_THREADS); + vsa_bwd_order_kernel<<>>(args.k2q_idx, args.k2q_num, args.max_q_blocks, + args.num_kv_blocks_per_seq, keep_dq_l2 ? 12 : 8, + keep_dq_l2, args.order_workspace); + e = cudaGetLastError(); + if (e != cudaSuccess) { + return e; + } + work_remap = args.order_workspace; + } + + e = keep_dq_l2 ? launch_main(args, work_remap, tk, tv, tqt, tdot, + tdk, tdv, sms, stream) + : launch_main(args, work_remap, tk, tv, tqt, + tdot, tdk, tdv, sms, stream); + if (e != cudaSuccess) { + return e; + } + + vsa_bwd_postprocess_kernel + <<>>( + args.dqaccum, args.dq, H, S, args.sm_scale); + return cudaGetLastError(); +} + +} // namespace vsa_bwd_blk64 + +using namespace vsa_bwd_blk64; + +#endif // BLOCK_SPARSE_VSA_BWD_LAUNCH_SM100A_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu new file mode 100644 index 0000000000..929979377d --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_bwd_sm100a.cu @@ -0,0 +1,164 @@ +// block_sparse_bwd_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA backward. +// +// Pairs with block_sparse_sm100a_fwd. Inputs are the forward's operands plus its output o and +// its lse; lse is the Triton "M format" tensor the forward returns -- [B, H, S] fp32, +// M = max(qk * sm_scale * log2e) + log2(l) -- and is consumed as-is. Sparsity arrives as +// FastVideo's k2q metadata (fastvideo_kernel.triton_kernels.index.invert_indices): for every +// (batch, head, kv block) the LOCAL q64 block ids that selected it, padded to max_q_blocks, +// plus a count; entries past the count are never read. Returns {dq, dk, dv} in bf16 with the +// inputs' layout and the Triton backward's scaling: dk and dq carry sm_scale, dv does not. +// +// 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]). +#include + +#include +#include + +#include + +#include "block_sparse_bwd_launch_sm100a.cuh" + +namespace { + +void check_activation(const torch::Tensor& t, const char* name, int64_t B, int64_t H, int64_t S, + int64_t D) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == at::kBFloat16, name, " must be bfloat16, got ", t.scalar_type()); + TORCH_CHECK(t.dim() == 4, name, " must be 4-D, got ", t.dim(), " dims"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + if (VSA_BHSD) { + TORCH_CHECK(t.size(0) == B && t.size(1) == H && t.size(2) == S && t.size(3) == D, name, + " has shape ", t.sizes(), ", expected [", B, ",", H, ",", S, ",", D, "]"); + } else { + TORCH_CHECK(t.size(0) == B && t.size(1) == S && t.size(2) == H && t.size(3) == D, name, + " has shape ", t.sizes(), ", expected [", B, ",", S, ",", H, ",", D, "]"); + } +} + +void check_index(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == at::kInt, name, " must be int32, got ", t.scalar_type()); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +__nv_bfloat16* bf16_ptr(const torch::Tensor& t) { + return reinterpret_cast<__nv_bfloat16*>(t.data_ptr()); +} + +} // 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, + 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) { + const c10::cuda::OptionalCUDAGuard guard(device_of(q)); + + TORCH_CHECK(q.dim() == 4, "q must be 4-D, got ", q.dim(), " dims"); + const int64_t B = q.size(0); + const int64_t H = VSA_BHSD ? q.size(1) : q.size(2); + const int64_t S = VSA_BHSD ? q.size(2) : q.size(1); + const int64_t D = q.size(3); + + check_activation(q, "q", B, H, S, D); + check_activation(k, "k", B, H, S, D); + check_activation(v, "v", B, H, S, D); + check_activation(o, "o", B, H, S, D); + check_activation(grad_o, "grad_o", B, H, S, D); + + TORCH_CHECK(lse.is_cuda(), "lse must be a CUDA tensor"); + TORCH_CHECK(lse.scalar_type() == at::kFloat, "lse must be float32, got ", lse.scalar_type()); + TORCH_CHECK(lse.is_contiguous(), "lse must be contiguous"); + TORCH_CHECK(lse.numel() == B * H * S, "lse must hold [B, H, S] = ", B * H * S, + " values (Triton M format), got ", lse.numel()); + + check_index(k2q_idx, "k2q_idx"); + check_index(k2q_num, "k2q_num"); + check_index(variable_block_sizes, "variable_block_sizes"); + + const int64_t num_kv_blocks_per_seq = variable_block_sizes.numel(); + TORCH_CHECK(S == num_kv_blocks_per_seq * BLOCK, "seqlen ", S, + " must equal num_kv_blocks_per_seq (", num_kv_blocks_per_seq, ") * ", BLOCK, + "; FastVideo pads the sequence up to whole blocks"); + + const int64_t num_items = B * H * num_kv_blocks_per_seq; + TORCH_CHECK(k2q_idx.dim() == 4 || k2q_idx.dim() == 2, + "k2q_idx must be [B, H, num_kv_blocks_per_seq, max_q_blocks] or " + "[B*H*num_kv_blocks_per_seq, max_q_blocks], got shape ", + k2q_idx.sizes()); + const int64_t max_q_blocks = k2q_idx.size(-1); + const bool k2q_dims_ok = k2q_idx.dim() == 2 || (k2q_idx.size(0) == B && k2q_idx.size(1) == H && + k2q_idx.size(2) == num_kv_blocks_per_seq); + TORCH_CHECK(k2q_dims_ok && k2q_idx.numel() == num_items * max_q_blocks, "k2q_idx has shape ", + k2q_idx.sizes(), ", expected [", B, ",", H, ",", num_kv_blocks_per_seq, + ",max_q_blocks] or [", num_items, ",max_q_blocks]"); + TORCH_CHECK(k2q_num.numel() == num_items, + "k2q_num must hold one count per (batch, head, kv " + "block) = ", + num_items, " values, got ", k2q_num.numel()); + + auto dq = torch::empty_like(q); + auto dk = torch::empty_like(k); + 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. + 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); + 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); + } + + BlockSparseVsaBwdArgs a{}; + a.q = bf16_ptr(q); + a.k = bf16_ptr(k); + a.v = bf16_ptr(v); + a.o = bf16_ptr(o); + a.dout = bf16_ptr(grad_o); + a.dq = bf16_ptr(dq); + a.dk = bf16_ptr(dk); + a.dv = bf16_ptr(dv); + a.lse = lse.data_ptr(); + 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()); + a.batch = b; + a.num_heads = h; + a.seqlen = s; + a.head_dim = (int)D; + a.num_kv_blocks_per_seq = (int)num_kv_blocks_per_seq; + a.max_q_blocks = (int)max_q_blocks; + a.sm_scale = (float)sm_scale; + + // 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, + ", seqlen == num_kv_blocks_per_seq*", BLOCK, + " with seqlen % 128 == 0, " + "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)); + + return {dq, dk, dv}; +} diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh index 8075bb908d..545e2e3e61 100644 --- a/fastvideo-kernel/csrc/attention/primitives.cuh +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -874,3 +874,114 @@ __device__ __forceinline__ void sts_f32(uint32_t smem_addr, float val) { asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory"); } + +__device__ __forceinline__ void tcgen05_mma_ws_f16_ss_1sm_predicated( + uint32_t issue, uint32_t tmem_d, uint64_t desc_a, uint64_t desc_b, + uint32_t idesc, bool enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p, q;\n\t" + "setp.ne.b32 q, %0, 0;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "@q tcgen05.mma.ws.cta_group::1.kind::f16 " + "[%1], %2, %3, %4, p, 0;\n\t" + "}\n" + :: "r"(issue), "r"(tmem_d), "l"(desc_a), "l"(desc_b), + "r"(idesc), "r"(enable_input_d ? 1u : 0u)); +} + +__device__ __forceinline__ void tcgen05_mma_ws_f16_ts_1sm_predicated( + uint32_t issue, uint32_t tmem_d, uint32_t tmem_a, uint64_t desc_b, + uint32_t idesc, bool enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p, q;\n\t" + "setp.ne.b32 q, %0, 0;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "@q tcgen05.mma.ws.cta_group::1.kind::f16 " + "[%1], [%2], %3, %4, p, 0;\n\t" + "}\n" + :: "r"(issue), "r"(tmem_d), "r"(tmem_a), "l"(desc_b), + "r"(idesc), "r"(enable_input_d ? 1u : 0u)); +} + +__device__ __forceinline__ void tcgen05_wait_ld() { + asm volatile("tcgen05.wait::ld.sync.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ +void cpasync_bulk_load_mbarrier(uint32_t smem_dst, const void* gmem_src, + uint32_t bytes, uint32_t mbar_smem) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];\n" + :: "r"(smem_dst), "l"(gmem_src), "r"(bytes), "r"(mbar_smem) + : "memory"); +} + +__device__ __forceinline__ +void cpasync_reduce_bulk_add_f32(float* global_dst, uint32_t smem_src, + uint32_t bytes) { + asm volatile( + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;\n" + :: "l"(global_dst), "r"(smem_src), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ +void cpasync_reduce_bulk_add_f16(uint16_t* global_dst, uint32_t smem_src, + uint32_t bytes) { + asm volatile( + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.f16" + " [%0], [%1], %2;\n" + :: "l"(global_dst), "r"(smem_src), "r"(bytes) + : "memory"); +} + +__device__ __forceinline__ +void cpasync_reduce_bulk_add_f32_l2hint(float* global_dst, uint32_t smem_src, + uint32_t bytes, uint64_t cache_policy) { + asm volatile( + "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32" + " [%0], [%1], %2, %3;\n" + :: "l"(global_dst), "r"(smem_src), "r"(bytes), "l"(cache_policy) + : "memory"); +} + +__device__ __forceinline__ +void cpasync_reduce_bulk_add_f16_l2hint(uint16_t* global_dst, uint32_t smem_src, + uint32_t bytes, uint64_t cache_policy) { + asm volatile( + "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.noftz.f16" + " [%0], [%1], %2, %3;\n" + :: "l"(global_dst), "r"(smem_src), "r"(bytes), "l"(cache_policy) + : "memory"); +} + +__device__ __forceinline__ void smem_desc_add_lo(uint64_t& d, uint32_t inc) { + asm volatile("{\n\t" + ".reg .b32 lo, hi;\n\t" + "mov.b64 {lo, hi}, %0;\n\t" + "add.u32 lo, lo, %1;\n\t" + "mov.b64 %0, {lo, hi};\n\t" + "}" : "+l"(d) : "r"(inc)); +} + +__device__ __forceinline__ +uint32_t cvt_f32x2_to_f16x2(float a, float b) { + uint32_t r; + asm volatile("cvt.rn.f16x2.f32 %0, %2, %1;\n" + : "=r"(r) : "f"(a), "f"(b)); + return r; +} + +__device__ __forceinline__ +uint64_t make_l2cache_policy_fractional_evict_last_unchanged(float fraction_f32) { + uint64_t policy; + asm("createpolicy.fractional.L2::evict_last.L2::evict_unchanged.b64" + " %0, %1;\n" + : "=l"(policy) + : "f"(fraction_f32)); + return policy; +} diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 0243b70480..f58a397117 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -42,6 +42,10 @@ extern std::vector block_sparse_sm100a_blk128_fwd( torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional v_t, torch::Tensor q2k_idx, torch::Tensor q2k_num, torch::Tensor variable_block_sizes, double sm_scale, bool need_lse); +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); #endif PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -54,6 +58,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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 diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py index 34b975e4a0..cb5127d749 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn.py @@ -372,13 +372,65 @@ def _backward_sm90(ctx, grad_o, grad_lse): # --------------------------------------------------------------------------- # Data-center Blackwell backend custom op (index-native; legacy sm100a API name) # -# Forward runs the sm_100a/sm_103a CUDA extension; backward reuses the Triton kernels. +# 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 the pairing needs no conversion. The Triton backward is +# log2(l)), so either pairing needs no conversion. Both backwards are # hardcoded to 64-token blocks, hence the block-size assert below. # --------------------------------------------------------------------------- +@torch.library.custom_op( + "fastvideo_kernel::block_sparse_attn_backward_sm100a", + mutates_args=(), + device_types="cuda", +) +def block_sparse_attn_backward_sm100a( + grad_o: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from fastvideo_kernel.block_sparse_attn_bwd_sm100a import ( + block_sparse_attn_backward_sm100a_from_k2q, ) + + num_kv_blocks = variable_block_sizes.numel() + k2q_idx, k2q_num = _invert_indices_for_backward(q2k_idx, q2k_num, num_kv_blocks) + dq, dk, dv = block_sparse_attn_backward_sm100a_from_k2q( + grad_o.contiguous(), q.contiguous(), k.contiguous(), v.contiguous(), o.contiguous(), + lse.contiguous(), k2q_idx, k2q_num, variable_block_sizes) + return dq, dk, dv + + +@torch.library.register_fake("fastvideo_kernel::block_sparse_attn_backward_sm100a") +def _block_sparse_attn_backward_sm100a_fake( + grad_o: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + + +def _sm100a_backward_is_supported(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> bool: + try: + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd_sm100a + except ImportError: # pragma: no cover - extension not built + return False + return vsa_bwd_sm100a.is_supported(q, variable_block_sizes) + + @torch.library.custom_op( "fastvideo_kernel::block_sparse_attn_sm100a", mutates_args=(), @@ -428,11 +480,15 @@ def _backward_sm100a(ctx, grad_o, grad_M): 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 the " - f"Triton backward, which is hardcoded to 64-token blocks; got {block}. " + "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.") - dq, dk, dv = block_sparse_attn_backward_triton(grad_o, q, k, v, o, M, q2k_idx, - q2k_num, variable_block_sizes) + 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: + 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 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 new file mode 100644 index 0000000000..76d338cd63 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_bwd_sm100a.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""sm_100a (Blackwell) CUDA block-sparse VSA backward. + +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``. +""" + +from typing import Tuple + +import torch + +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 +except ImportError: # pragma: no cover - extension not built + _C = None + _BWD = None + _HAS_VSA_BWD_SM100A = False + +_SM100 = (10, 0) +HEAD_DIM = 128 +BLOCK = 64 +# 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``. + + A standalone build of ``block_sparse_bwd_sm100a.cu`` (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 + _C = module + _BWD = getattr(module, "block_sparse_sm100a_bwd", None) + _HAS_VSA_BWD_SM100A = _BWD is not None + + +def _seqlen(q: torch.Tensor) -> int: + return q.shape[2] if BHSD else q.shape[1] + + +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. + """ + if not _HAS_VSA_BWD_SM100A or not q.is_cuda: + return False + if torch.cuda.get_device_capability(q.device) != _SM100: + return False + if q.dtype != torch.bfloat16 or q.dim() != 4 or q.shape[-1] != HEAD_DIM: + return False + if not q.is_contiguous(): + 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: + return False + if _seqlen(q) != BLOCK * num_blocks: + return False + return True + + +def block_sparse_attn_backward_sm100a_from_k2q( + 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: torch.Tensor, +) -> 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. + """ + 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) + return res[0], res[1], res[2] + + +def block_sparse_attn_backward_sm100a( + grad_o: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + lse: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Backward pass from the forward's q2k metadata. Returns ``(dq, dk, dv)``. + + Mirrors ``block_sparse_attn_backward_triton``: the k2q inversion is recomputed here with + FastVideo's Triton ``invert_indices`` rather than saved by the forward. + """ + from fastvideo_kernel.triton_kernels.index import invert_indices + + num_kv_blocks = variable_block_sizes.numel() + batch = q.shape[0] + heads = q.shape[1] if BHSD else q.shape[2] + idx = q2k_idx.to(torch.int32).contiguous() + num = q2k_num.to(torch.int32).contiguous() + if idx.dim() != 4: + idx = idx.view(batch, heads, -1, idx.shape[-1]) + if num.dim() != 3: + num = num.view(batch, heads, -1) + k2q_idx, k2q_num = invert_indices(idx, num, num_kv_blocks) + return block_sparse_attn_backward_sm100a_from_k2q(grad_o, q, k, v, o, lse, k2q_idx, k2q_num, + variable_block_sizes) 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 e8347450a5..84872f52d8 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -4,10 +4,12 @@ The historical ``sm100a`` module and symbol names are retained for compatibility, but the extension carries native sm_100a and sm_103a images and supports both device generations. -A third backend behind the same VSA op as the Triton and CuTe-DSL paths. Forward only: 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 -``block_sparse_attn_backward_triton`` runs against it unchanged. +A third backend behind the same VSA op as the Triton and CuTe-DSL paths. This module is the +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 +unchanged. The extension carries TWO instantiations of the kernel, for 64- and 128-token sparse blocks (tile volumes 64 and 128 in ``build_vsa_metadata``); the block size is inferred from the diff --git a/tests/test_block_sparse_bwd_sm100a.py b/tests/test_block_sparse_bwd_sm100a.py new file mode 100644 index 0000000000..a2a7c13518 --- /dev/null +++ b/tests/test_block_sparse_bwd_sm100a.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Correctness tests for the sm_100a CUDA block-sparse VSA backward. + +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. + +Run with: python -m pytest tests/test_block_sparse_bwd_sm100a.py -v +""" + +import itertools +import os + +import pytest +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. +REL_MAX_TOL = 1e-2 # max|got - ref| / max|ref| +MEAN_ABS_TOL = 1e-3 # mean|got - ref| + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0) + or not bwd._HAS_VSA_BWD_SM100A, + reason="requires a compute-capability (10, 0) GPU (sm_100a) and a fastvideo_kernel " + "extension built with block_sparse_sm100a_bwd", +) + + +def make_case(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 + 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)) + + pool = torch.arange(num_blocks) if kv_pool is None else torch.as_tensor(list(kv_pool)) + assert topk <= pool.numel() + rows = batch * heads * num_blocks + idx = torch.empty((rows, topk), dtype=torch.int32) + for r in range(rows): + idx[r] = pool[torch.randperm(pool.numel())[:topk]].sort().values.to(torch.int32) + idx = idx.view(batch, heads, num_blocks, topk).cuda() + 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, + device="cuda") + else: + vbs = torch.full((num_blocks, ), BLOCK, dtype=torch.int32, device="cuda") + return q, k, v, grad_o, idx, num, vbs + + +def invert_indices_torch(q2k_idx, q2k_num, num_kv_blocks, pad_value=0): + """k2q from q2k without Triton: a stable sort, so each row lists its q blocks ascending. + + Returns (k2q_idx [B, H, num_kv_blocks, Nq] int32, k2q_num [B, H, num_kv_blocks] int32), the + layout fastvideo_kernel.triton_kernels.index.invert_indices produces. Entries past a row's + count hold ``pad_value`` -- a VALID block id, so an over-read fails by wrong values rather + than by luck. + """ + B, H, Nq, Mk = q2k_idx.shape + device = q2k_idx.device + valid = torch.arange(Mk, device=device).view(1, 1, 1, Mk) < q2k_num.view(B, H, Nq, 1) + row = (torch.arange(B * H, device=device).view(B, H, 1, 1) * num_kv_blocks + + q2k_idx.long())[valid] + qblock = torch.arange(Nq, device=device).view(1, 1, Nq, 1).expand(B, H, Nq, Mk)[valid] + order = torch.sort(row * Nq + qblock).indices + row, qblock = row[order], qblock[order] + counts = torch.bincount(row, minlength=B * H * num_kv_blocks) + starts = torch.cumsum(counts, 0) - counts + slot = torch.arange(row.numel(), device=device) - starts[row] + k2q_idx = torch.full((B * H * num_kv_blocks, Nq), pad_value, dtype=torch.int32, + device=device) + k2q_idx[row, slot] = qblock.to(torch.int32) + return (k2q_idx.view(B, H, num_kv_blocks, Nq), + counts.to(torch.int32).view(B, H, num_kv_blocks)) + + +def reference(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 + q's layout. + """ + if not bwd.BHSD: + q, k, v, grad_o = (t.transpose(1, 2) for t in (q, k, v, grad_o)) # -> [B, H, S, D] + B, H, S, D = q.shape + num_blocks = vbs.numel() + scale = 1.0 / (D**0.5) + + idx, num, vbs = idx.cpu(), num.cpu(), vbs.cpu() + keep = torch.zeros((B, H, S, S), dtype=torch.bool, device=q.device) + for b in range(B): + for h in range(H): + for qb in range(num_blocks): + 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 + + q32, k32, v32 = (t.detach().float().requires_grad_(True) for t in (q, k, v)) + scores = (q32 @ k32.transpose(-1, -2)) * scale + scores = scores.masked_fill(~keep, float("-inf")) + p = torch.softmax(scores, dim=-1) + out = p @ v32 + loss = (out * grad_o.float()).sum() + dq, dk, dv = torch.autograd.grad(loss, (q32, k32, v32)) + lse = (torch.logsumexp(scores, dim=-1) * LOG2E).detach().contiguous() + o = out.detach().to(torch.bfloat16) + + if not bwd.BHSD: + o, dq, dk, dv = (t.transpose(1, 2).contiguous() for t in (o, dq, dk, dv)) + return o, lse, dq.detach(), dk.detach(), dv.detach() + + +def check_close(name, got, ref): + got, ref = got.float(), ref.float() + assert got.shape == ref.shape, f"{name}: shape {tuple(got.shape)} vs {tuple(ref.shape)}" + assert torch.isfinite(got).all(), f"{name}: non-finite values" + diff = (got - ref).abs() + rel_max = diff.max().item() / max(ref.abs().max().item(), 1e-6) + mean_abs = diff.mean().item() + if os.environ.get("VSA_BWD_TEST_VERBOSE"): + print(f"{name}: rel_max={rel_max:.3e} mean_abs={mean_abs:.3e} " + f"mean|ref|={ref.abs().mean().item():.3e}") + assert rel_max <= REL_MAX_TOL and mean_abs <= MEAN_ABS_TOL, ( + f"{name}: max|diff|/max|ref| = {rel_max:.3e} (tol {REL_MAX_TOL:.1e}), " + f"mean|diff| = {mean_abs:.3e} (tol {MEAN_ABS_TOL:.1e}), " + 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) + assert bwd.is_supported(q, vbs) + o, lse, ref_dq, ref_dk, ref_dv = reference(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) + torch.cuda.synchronize() + for name, got, ref in (("dq", dq, ref_dq), ("dk", dk, ref_dk), ("dv", dv, ref_dv)): + assert got.dtype == torch.bfloat16, f"{name}: dtype {got.dtype}" + check_close(name, got, ref) + return (dq, dk, dv), (ref_dq, ref_dk, ref_dv), vbs + + +def _to_bhsd(t): + return t if bwd.BHSD else t.transpose(1, 2) + + +def test_backward_matches_reference(): + run_and_compare() + + +def test_batch_two(): + run_and_compare(batch=2) + + +def test_ragged_block_sizes(): + """variable_block_sizes is what FastVideo always passes; padded keys must be masked.""" + run_and_compare(ragged=True) + + +def test_ragged_padded_key_rows_are_zero(): + """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) + dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() + checked = 0 + for kb in range(vbs.numel()): + if int(vbs[kb]) == BLOCK: + continue # a full block has no padded rows + 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 + assert checked > 0, "the ragged draw produced no padded kv block; change the seed" + + +@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) + + +@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)) + + +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.""" + 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, + kv_pool=pool) + dk, dv = _to_bhsd(dk).float(), _to_bhsd(dv).float() + for kb in excluded: + 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}" + + +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) + 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) + assert int(k2q_num.sum()) == B * H * Nq * Mk + for b, h, qb, j in itertools.product(range(B), range(H), range(Nq), range(Mk)): + kb = int(idx[b, h, qb, j]) + listed = k2q_idx[b, h, kb, :int(k2q_num[b, h, kb])].tolist() + assert listed.count(qb) == 1 + assert listed == sorted(listed) + + +def test_unsupported_is_rejected(): + q, k, v, grad_o, idx, num, vbs = make_case() + 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 + 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) + + # 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, + device="cuda") + with pytest.raises(RuntimeError): + bwd.block_sparse_attn_backward_sm100a_from_k2q(grad_o.float(), q.float(), k.float(), + v.float(), o.float(), lse, k2q_idx, + k2q_num, vbs) diff --git a/tests/test_block_sparse_sm100a_dispatch.py b/tests/test_block_sparse_sm100a_dispatch.py index 5697b488d2..9bfdc719c4 100644 --- a/tests/test_block_sparse_sm100a_dispatch.py +++ b/tests/test_block_sparse_sm100a_dispatch.py @@ -2,14 +2,18 @@ """Routing tests for the opt-in sm_100a/sm_103a dispatch in ``block_sparse_attn_from_indices``. ``FASTVIDEO_VSA_SM100A=1`` routes the forward to the sm_100a extension when -``block_sparse_attn_sm100a.is_supported`` passes, pairing it with the Triton -backward (the sm_100a lse is already in Triton's M format). Everything else -- -env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- must keep -the pre-existing selection, bit-for-bit. +``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. +Everything else -- env unset, unsupported input, ``FASTVIDEO_VSA_TRITON`` override -- +must keep the pre-existing selection, bit-for-bit. Run with: python -m pytest tests/test_block_sparse_sm100a_dispatch.py -v """ +import importlib + import pytest import torch @@ -107,27 +111,125 @@ def test_force_triton_overrides_sm100a(monkeypatch): assert not torch.equal(got[0], sm100a_o) -def test_backward_runs_triton_and_matches(monkeypatch): - """sm_100a forward + Triton backward: grads match the all-Triton path.""" - monkeypatch.setenv(ENV, "1") - q, k, v, idx, num, vbs = make_case(64, requires_grad=True) - out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) - out.float().square().sum().backward() - got = [t.grad.float().clone() for t in (q, k, v)] +def _grads_sm100a_route_vs_triton(monkeypatch, vbs=None): + """Grads of (out**2).sum() through the sm_100a route vs the all-Triton route. + + On a device where the sm_100a backward is built and supported, the route must not enter + the Triton backward at all (it is monkeypatched to fail); elsewhere the route pairs the + sm_100a forward with the Triton backward, as before. + """ + # The package exports a FUNCTION named block_sparse_attn; the module needs importlib. + dispatch = importlib.import_module("fastvideo_kernel.block_sparse_attn") + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + + q, k, v, idx, num, default_vbs = make_case(64, requires_grad=True) + vbs = default_vbs if vbs is None else vbs q2, k2, v2 = (t.detach().clone().requires_grad_(True) for t in (q, k, v)) monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") out2, _ = block_sparse_attn_from_indices(q2, k2, v2, idx, num, vbs) out2.float().square().sum().backward() ref = [t.grad.float() for t in (q2, k2, v2)] + monkeypatch.delenv("FASTVIDEO_VSA_TRITON") + + monkeypatch.setenv(ENV, "1") + sm100a_backward = vsa_bwd.is_supported(q, vbs) + if sm100a_backward: + + def no_triton_backward(*args, **kwargs): + raise AssertionError("Triton backward entered on a supported sm_100a input") + + monkeypatch.setattr(dispatch, "block_sparse_attn_backward_triton", no_triton_backward) + out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs) + out.float().square().sum().backward() + got = [t.grad.float() for t in (q, k, v)] + return got, ref, sm100a_backward + +def _assert_grads_close(got, ref): + # Two bf16 kernels against each other (not an fp32 reference): twice the tolerances the + # sm_100a backward holds against fp32 in tests/test_block_sparse_bwd_sm100a.py. for g, r, name in zip(got, ref, "qkv"): - assert torch.allclose(g, r, atol=5e-2, rtol=5e-2), \ - f"d{name} max|diff|={(g - r).abs().max().item()}" + diff = (g - r).abs() + rel_max = diff.max().item() / max(r.abs().max().item(), 1e-6) + mean_abs = diff.mean().item() + assert rel_max <= 2e-2 and mean_abs <= 2e-3, \ + f"d{name}: max|diff|/max|ref|={rel_max:.3e} mean|diff|={mean_abs:.3e}" + + +def test_backward_matches_triton(monkeypatch): + """sm_100a route (CUDA backward where built, Triton backward otherwise) vs all-Triton.""" + got, ref, _ = _grads_sm100a_route_vs_triton(monkeypatch) + _assert_grads_close(got, ref) + + +def test_backward_ragged_block_sizes_matches_triton(monkeypatch): + """variable_block_sizes below 64 (padded kv rows) through the same two routes.""" + torch.manual_seed(1) + vbs = torch.randint(32, 65, (8, ), dtype=torch.int32, device="cuda") + got, ref, _ = _grads_sm100a_route_vs_triton(monkeypatch, vbs=vbs) + _assert_grads_close(got, ref) + + +def test_backward_uses_sm100a_kernel_when_built(monkeypatch): + """Guards against a silent Triton fallback: with the op built on an sm_100a device the + CUDA backward must be the one that runs.""" + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + if not vsa_bwd._HAS_VSA_BWD_SM100A or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("sm_100a backward not built for this device") + _, _, sm100a_backward = _grads_sm100a_route_vs_triton(monkeypatch) + assert sm100a_backward + + +def test_backward_large_seq_matches_triton_and_is_deterministic(monkeypatch): + """65536 tokens (1024 kv blocks, 12.5% density): the device-computed work order (nb >= 1024) + and the sequence regime where a TMEM ordering bug corrupted dk/dv while every <= 16-block + test stayed green. sm_100a route vs all-Triton, and two sm_100a runs must agree to within + summation-order noise: invert_indices compacts each kv row's q list with atomics, so the + kernel sums quads in a different order per call (measured run-to-run delta on GB200: + rel_max 5e-3, mean|diff| 7e-6). The TMEM race this guards against gave rel_max 0.6-1.0 + and mean|diff| 4e-2, an order of magnitude past the bounds below on both metrics.""" + from fastvideo_kernel import block_sparse_attn_bwd_sm100a as vsa_bwd + + torch.manual_seed(0) + block, num_blocks, heads, topk = 64, 1024, 4, 128 + S = num_blocks * block + shape = (1, heads, S, HEAD_DIM) if vsa.BHSD else (1, S, heads, HEAD_DIM) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + scores = torch.rand(1, heads, num_blocks, num_blocks, device="cuda") + idx = scores.topk(topk, dim=-1).indices.sort(dim=-1).values.to(torch.int32) + del scores + num = torch.full((1, heads, num_blocks), topk, dtype=torch.int32, device="cuda") + vbs = torch.full((num_blocks, ), block, dtype=torch.int32, device="cuda") + if not vsa_bwd.is_supported(q, vbs): + pytest.skip("sm_100a backward not built for this device") + + def grads(route): + qq, kk, vv = (t.clone().requires_grad_(True) for t in (q, k, v)) + if route == "triton": + monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1") + else: + monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False) + monkeypatch.setenv(ENV, "1") + out, _ = block_sparse_attn_from_indices(qq, kk, vv, idx, num, vbs) + out.float().square().sum().backward() + return [t.grad.float() for t in (qq, kk, vv)] + + ref = grads("triton") + got = grads("sm100a") + again = grads("sm100a") + _assert_grads_close(got, ref) + for g1, g2, name in zip(got, again, "qkv"): + diff = (g1 - g2).abs() + rel_max = diff.max().item() / max(g1.abs().max().item(), 1e-6) + mean_abs = diff.mean().item() + assert rel_max <= 2e-2 and mean_abs <= 1e-4, \ + f"d{name}: two sm_100a runs differ beyond summation-order noise: " \ + 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 (Triton bwd is 64-block only).""" + """128-token blocks: forward runs, backward refuses (both backwards are 64-block only).""" monkeypatch.setenv(ENV, "1") q, k, v, idx, num, vbs = make_case(128, requires_grad=True) out, _ = block_sparse_attn_from_indices(q, k, v, idx, num, vbs)