diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 79fdf84b3..51ec02a47 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -318,6 +318,38 @@ if(BUILD_CXX_KERNELS) # Combined FastVideo Extension # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace + # --------------------------------------------------------------------------- + # VSA block-sparse attention forward, Blackwell (sm_100a) only. + # + # NOTE the "a" suffix: -arch=sm_100a is NOT enough -- it emits a plain sm_100 target and + # ptxas rejects every tcgen05 / setmaxnreg instruction. The explicit gencode spelling + # below is required, and matches the 10.0a entry in TORCH_CUDA_ARCH_LIST. + # + # Built for 64-token sparse blocks -- FastVideo's default (4,4,4) tiling, so no + # tile-size change and no top-k granularity change is needed. VSA_BHSD + # selects [B, H, S, D]. Both are compile-time; the Python is_supported() checks incoming + # tensors against them so callers fall back to Triton rather than getting a wrong answer. + # Read the ENVIRONMENT as well as the cache variable. When TORCH_CUDA_ARCH_LIST is + # exported (build.sh, and `pip install` with it set) the branch above only prints it -- + # the cmake variable stays empty, so testing that alone silently skips the kernel and + # leaves a build that succeeds with the op missing. + set(ENABLE_VSA_SM100A OFF) + set(_VSA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") + if(NOT _VSA_ARCH_LIST AND DEFINED ENV{TORCH_CUDA_ARCH_LIST}) + set(_VSA_ARCH_LIST "$ENV{TORCH_CUDA_ARCH_LIST}") + endif() + if(_VSA_ARCH_LIST MATCHES "(^|[; ,])(10\\.0a|100a|sm_100a)([; ,]|$)") + set(ENABLE_VSA_SM100A ON) + endif() + if(ENABLE_VSA_SM100A) + message(STATUS "fastvideo-kernel: building block_sparse_sm100a (Blackwell, 64- and 128-token blocks)") + list(APPEND EXTENSION_SOURCES csrc/attention/block_sparse_sm100a.cu + csrc/attention/block_sparse_blk128_sm100a.cu) + 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;-DVSA_BHSD=true") + endif() + Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${EXTENSION_SOURCES} ) @@ -333,10 +365,14 @@ if(BUILD_CXX_KERNELS) # Build compile definitions list set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) + if(ENABLE_VSA_SM100A) + list(APPEND COMPILE_DEFS TK_COMPILE_BLOCK_SPARSE_VSA_SM100A) + endif() if(ENABLE_TK_KERNELS) list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE) endif() + target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS}) target_compile_options(fastvideo_kernel_ops PRIVATE diff --git a/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu new file mode 100644 index 000000000..cb1c2eefe --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_blk128_sm100a.cu @@ -0,0 +1,7 @@ +// block_sparse_blk128_sm100a.cu -- the 128-token-block instantiation of the torch binding. +// +// Same source as block_sparse_sm100a.cu with VSA_BLK128 set: the kernel and launch land in +// namespace vsa_blk128 (distinct symbols, no ODR clash with the blk64 objects) and the +// exported entry point becomes block_sparse_sm100a_blk128_fwd. +#define VSA_BLK128 true +#include "block_sparse_sm100a.cu" diff --git a/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh new file mode 100644 index 000000000..3113b03b0 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_kernel_sm100a.cuh @@ -0,0 +1,1078 @@ +// block_sparse_kernel_sm100a.cuh -- VSA block-sparse FMHA forward (per-q-block top-k), sm_100a. +// Warp-specialized: load / MMA (tcgen05) / softmax / correction / epilogue / scheduler. +// Writes O and, when asked, the log-sum-exp the backward consumes. +// +// Generated (comments stripped). Do not edit by hand. +#ifndef BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH +#define BLOCK_SPARSE_VSA_KERNEL_SM100A_CUH + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "primitives.cuh" + +#ifndef VSA_BLK128 +#define VSA_BLK128 false +#endif +#ifndef VSA_BHSD +#define VSA_BHSD false +#endif + +// BLK128 is a file-scope constexpr, not a template parameter, so the blk64 and blk128 builds +// would instantiate the SAME kernel symbol with DIFFERENT bodies -- an ODR violation the linker +// resolves by silently keeping one. A config-named namespace keeps the two builds' symbols +// distinct so both can live in one extension. +#if VSA_BLK128 +#define VSA_NAMESPACE vsa_blk128 +#else +#define VSA_NAMESPACE vsa_blk64 +#endif +namespace VSA_NAMESPACE { + +constexpr bool BLK128 = VSA_BLK128; + +#ifndef VSA_DEFER_ROWSUM +#define VSA_DEFER_ROWSUM (!VSA_BLK128) +#endif +constexpr bool DEFER_ROWSUM = VSA_DEFER_ROWSUM; + +constexpr int BLOCK = BLK128 ? 128 : 64; +constexpr int M_TILE = BLOCK; +constexpr int M_TILES_PER_CTA = 2; +constexpr int K_TILE = 256 / (BLK128 ? 2 : 1); +constexpr int BLOCKS_PER_KTILE = K_TILE / BLOCK; +constexpr int KV_HALF = K_TILE / 2; +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 Q_SUBTILES = HEAD_DIM / SUB_COLS_BF16; +constexpr int K_SUBTILES = HEAD_DIM / SUB_COLS_BF16; +constexpr int V_SUBTILES = (BLK128 ? K_TILE : KV_HALF) / SUB_COLS_BF16; +constexpr int Q_SUB_COLS_BYTES = M_TILE * SUB_COLS_BYTES; +constexpr int K_SUB_COLS_BYTES = K_TILE * SUB_COLS_BYTES; +constexpr int Q_TILE_BYTES = Q_SUBTILES * Q_SUB_COLS_BYTES; + +constexpr int KV_RING_SLOT_BYTES = 32 * 1024; +constexpr int SLOTS_PER_KV_TILE = BLK128 ? 1 : 2; + +constexpr int NUM_KV_STAGES = BLK128 ? 3 : 4; + +constexpr int V_BLK_BYTES = HEAD_DIM * SUB_COLS_BYTES; + +constexpr int S_COLS = 128; +constexpr int K_ATOMS_PER_TILE = SUB_COLS_BF16 / 16; +constexpr int SPLIT_P_N = S_COLS / 4 * 3; +constexpr int SPLIT_P_ATOM = SPLIT_P_N / 16; +constexpr int SPLIT_P_COL = SPLIT_P_N / 2; +constexpr int EX2_FRG_PAIRS = 16; +constexpr int EX2_FRG_CNT = S_COLS / 32; +constexpr int EX2_FREQ = 16; +constexpr int EX2_RES = 4; + +constexpr int O_COLS = HEAD_DIM; +constexpr int TMEM_TOTAL = 512; + +constexpr int STATS = BLK128 ? M_TILE : 2 * M_TILE; +constexpr int STAT_REGIONS = BLK128 ? 2 : 3; + +constexpr int W_CORR0 = 8, W_MMA = 12, W_EPI = 13, W_LOAD = 14, W_SCHED = 15; +constexpr int N_WARPS = 16; +constexpr int CLC_STAGES = 2; + +extern __shared__ __align__(1024) uint8_t fmha_smem[]; + +union SmemDescPair { + uint64_t u64; + uint2 w; +}; + +__device__ __forceinline__ void desc_add_lo(SmemDescPair& 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.u64) + : "r"(inc)); +} + +__device__ __forceinline__ void tcgen05_mma_ws_f16_ss_1sm_lead(uint32_t lead, + 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, %4, 0;\n\t" + "@q tcgen05.mma.ws.cta_group::1.kind::f16 [%1], %2, %3, %5, p, 0;\n\t" + "}\n" + :: "r"(lead), "r"(tmem_d), "l"(desc_a), "l"(desc_b), + "r"(enable_input_d ? 1u : 0u), "r"(idesc)); +} +__device__ __forceinline__ void tcgen05_mma_ws_f16_ts_1sm_lead(uint32_t lead, + 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, %4, 0;\n\t" + "@q tcgen05.mma.ws.cta_group::1.kind::f16 [%1], [%2], %3, %5, p, 0;\n\t" + "}\n" + :: "r"(lead), "r"(tmem_d), "r"(tmem_a), "l"(desc_b), + "r"(enable_input_d ? 1u : 0u), "r"(idesc)); +} + +struct WorkItem { + int sample; + int head; + int mtile0; + int mtile1; + int global_mtile0; + int global_mtile1; + // Shared trip count for the CTA's q-tile pair: max of the two tiles' own + // counts, floored at 1. Every warp role derives its loop bounds from this + // one value, so the producer/consumer mbarrier pairing stays symmetric for + // ANY per-tile counts (including 0): a tile shorter than the pair max runs + // its extra iterations against clamped (valid, in-bounds) KV blocks and + // masks them to -inf via the per-tile count below, contributing nothing. + int num_kv_blocks; + // Each q-tile's OWN q2k_num. Used for (a) the q2k_idx window clamp in the + // load warp and (b) the vbs-threshold masking in the softmax warps. The + // pre-fix kernel used q2k_num[global_mtile0] for BOTH tiles, silently + // corrupting one tile of any pair whose rows have different counts. + int num_kv_blocks_mt[2]; +}; +template +__device__ __forceinline__ WorkItem decode_workitem( + int workitem_id, int num_heads, int num_blocks, int packed_mtiles_per_seq, + unsigned long long magic0, unsigned long long magic1, unsigned long long magic2, const int* q2k_num) { + WorkItem it; + const int per_sample = num_heads * packed_mtiles_per_seq; + it.sample = (int)fdiv((unsigned)workitem_id, magic0); + const int rem = workitem_id - it.sample * per_sample; + int p; + if constexpr (Q_RASTER) { + it.head = (int)fdiv((unsigned)rem, magic2); + p = rem - it.head * packed_mtiles_per_seq; + } else { + p = (int)fdiv((unsigned)rem, magic1); + it.head = rem - p * num_heads; + } + it.mtile0 = 2 * p; + it.mtile1 = 2 * p + 1; + it.global_mtile0 = (it.sample * num_heads + it.head) * num_blocks + it.mtile0; + it.global_mtile1 = it.global_mtile0 + 1; + it.num_kv_blocks_mt[0] = q2k_num[it.global_mtile0]; + it.num_kv_blocks_mt[1] = q2k_num[it.global_mtile1]; + // Floor of 1: a pair whose rows are BOTH empty still walks one K-tile so no + // mbarrier wait is left without its arrive (the load/MMA/softmax/correction + // pipeline has a fixed per-iteration handshake); the per-tile counts mask + // that K-tile completely, so such rows produce exactly-zero output. + it.num_kv_blocks = max(max(it.num_kv_blocks_mt[0], it.num_kv_blocks_mt[1]), 1); + return it; +} + +template +__global__ void __cluster_dims__(1, 1, 1) __launch_bounds__(N_WARPS * 32, 1) +fmha_context_bf16_gen_kernel(const __grid_constant__ CUtensorMap tmap_q, + const __grid_constant__ CUtensorMap tmap_k, const __grid_constant__ CUtensorMap tmap_v_t, + const __grid_constant__ CUtensorMap tmap_v, + const __grid_constant__ CUtensorMap tmap_o, int seqlen, + int num_heads, float scale_log2, int num_samples, int num_blocks, + int packed_mtiles_per_seq, int max_kv, + unsigned long long magic0, unsigned long long magic1, unsigned long long magic2, + const int* __restrict__ q2k_idx, const int* __restrict__ q2k_num, + const int* __restrict__ variable_block_sizes, + float* __restrict__ lse_out) { +// Multi-arch builds: torch's cmake appends -gencode for EVERY entry of +// TORCH_CUDA_ARCH_LIST to this TU on top of the pinned compute_100a pass, and +// tcgen05/setmaxnreg do not exist outside sm_100a -- ptxas rejects the sm_120a +// (or plain sm_100) pass outright. Keep the body only where it can compile: +// the host pass (no __CUDA_ARCH__, needed for launch plumbing) and the +// sm_100a device pass (arch 1000 WITH the family-specific feature set that the +// "a" suffix defines). Every other device pass gets an empty stub; the Python +// is_supported() / host launcher never dispatch here off sm_100, so the stub +// is unreachable at runtime. +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ == 1000 && defined(__CUDA_ARCH_FEAT_SM100_ALL)) + + const int total_workitems = num_samples * num_heads * packed_mtiles_per_seq; + + uint8_t* sQ0 = fmha_smem; + uint8_t* sQ1 = sQ0 + Q_TILE_BYTES; + uint8_t* sQ[2] = { sQ0, sQ1 }; + uint8_t* sKV = sQ1 + Q_TILE_BYTES; + __nv_bfloat16* sO0 = reinterpret_cast<__nv_bfloat16*>(sKV + NUM_KV_STAGES * KV_RING_SLOT_BYTES); + __nv_bfloat16* sO1 = sO0 + M_TILE * HEAD_DIM; + __nv_bfloat16* sO_bufs[2] = { sO0, sO1 }; + uint64_t* full_bar = reinterpret_cast(reinterpret_cast(sO1) + M_TILE * HEAD_DIM * sizeof(__nv_bfloat16)); + uint64_t* empty_bar= full_bar + NUM_KV_STAGES; + uint64_t* full_bar_q = empty_bar + NUM_KV_STAGES; + uint64_t* empty_bar_q = full_bar_q + 2; + uint64_t* full_bar_spo = empty_bar_q + 2; + uint64_t* empty_bar_spo = full_bar_spo + 2; + uint64_t* full_bar_o_acc = empty_bar_spo + 2; + uint64_t* full_bar_alpha = full_bar_o_acc + 2; + uint64_t* full_bar_l = full_bar_alpha + 2; + uint64_t* full_bar_p_last = full_bar_l + 2; + uint64_t* empty_bar_alpha_and_l = full_bar_p_last + 2; + uint64_t* full_bar_o_epi = empty_bar_alpha_and_l + 2; + uint64_t* empty_bar_o_epi = full_bar_o_epi + 2; + uint64_t* clc_full = empty_bar_o_epi + 2; + 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; + float* alpha_and_l_smem = reinterpret_cast(tmem_slot + 2); + + 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; + + if (tid == 0) { + #pragma unroll + for (int s = 0; s < NUM_KV_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&full_bar[s]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar[s]), 1); + } + for (int i = 0; i < 2; ++i) { + mbarrier_init(smem_ptr_u32(&full_bar_q[i]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_q[i]), 1); + mbarrier_init(smem_ptr_u32(&full_bar_l[i]), 128); + mbarrier_init(smem_ptr_u32(&full_bar_spo[i]), 1); + mbarrier_init(smem_ptr_u32(&empty_bar_spo[i]), 256); + mbarrier_init(smem_ptr_u32(&full_bar_o_acc[i]), 1); + mbarrier_init(smem_ptr_u32(&full_bar_alpha[i]), 128); + mbarrier_init(smem_ptr_u32(&empty_bar_alpha_and_l[i]), 128); + mbarrier_init(smem_ptr_u32(&full_bar_p_last[i]), 128); + mbarrier_init(smem_ptr_u32(&full_bar_o_epi[i]), 128); + mbarrier_init(smem_ptr_u32(&empty_bar_o_epi[i]), 1); + } + if constexpr (USE_CLC) { + #pragma unroll + for (int s = 0; s < CLC_STAGES; ++s) { + mbarrier_init(smem_ptr_u32(&clc_full[s]), 1); + mbarrier_init(smem_ptr_u32(&clc_empty[s]), N_WARPS); + } + #pragma unroll + for (int i = 0; i < CLC_STAGES * 4; ++i) + clc_response[i] = 0; + } + } + fence_mbarrier_init_release_cluster(); + __syncthreads(); + + if (warp_id == W_LOAD) { + setmaxnreg_dec<48>(); + + EmptyPhaseTracker kv_empty_ph; + EmptyPhaseTracker<1> q_empty_ph; + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (true) { + const WorkItem it = decode_workitem(workitem_id, num_heads, num_blocks, packed_mtiles_per_seq, magic0, magic1, magic2, q2k_num); + const int q_start = it.sample * seqlen; + const int k_start = it.sample * seqlen; + const int global_mtile[2] = { it.global_mtile0, it.global_mtile1 }; + const int mtile[2] = { it.mtile0, it.mtile1 }; + const int num_k_tiles = (it.num_kv_blocks + BLOCKS_PER_KTILE - 1) / BLOCKS_PER_KTILE; + + int window_start[2] = { -(1 << 30), -(1 << 30) }; + int kv_block_id_cache[2] = { 0, 0 }; + auto get_kv_block_id = [&](int mtile_idx, int j) -> int { + if (j >= window_start[mtile_idx] + 32) { + window_start[mtile_idx] = j & ~31; + // Clamp into THIS tile's own row of q2k_idx (the pair may run more + // iterations than this tile has blocks -- see WorkItem). Positions + // past the tile's count re-read its last valid entry; an empty tile + // never dereferences its row at all and loads block 0 instead. Both + // are fully masked by the per-tile vbs thresholds, so only the + // load address safety matters here. + const int cnt = it.num_kv_blocks_mt[mtile_idx]; + const int idx = max(0, min(window_start[mtile_idx] + lane, cnt - 1)); + kv_block_id_cache[mtile_idx] = + (cnt > 0) ? q2k_idx[global_mtile[mtile_idx] * max_kv + idx] : 0; + } + return __shfl_sync(0xffffffffu, kv_block_id_cache[mtile_idx], j & 31); + }; + + auto load_k_oneslot = [&](int mtile_idx, int ktile_idx, int s) { + const int kv_stage = kv_empty_ph.get_stage(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar[kv_stage]), kv_empty_ph.get_phase()); + kv_empty_ph.advance(); + + int kv_tok[BLOCKS_PER_KTILE]; + #pragma unroll + for (int blk = 0; blk < BLOCKS_PER_KTILE; ++blk) + kv_tok[blk] = k_start + get_kv_block_id(mtile_idx, ktile_idx * BLOCKS_PER_KTILE + blk) * BLOCK; + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar[kv_stage]), KV_RING_SLOT_BYTES); + if constexpr (BLK128) { + + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES)), &tmap_k, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[0] - k_start, 0, it.sample * num_heads + it.head); + else + tma_load_3d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES)), &tmap_k, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[0], it.head * K_SUBTILES); + } else { + #pragma unroll + for (int blk = 0; blk < BLOCKS_PER_KTILE; ++blk) + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES) + blk * BLOCK * SUB_COLS_BYTES), + &tmap_k, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[blk] - k_start, s, it.sample * num_heads + it.head); + else + tma_load_3d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES) + blk * BLOCK * SUB_COLS_BYTES), + &tmap_k, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[blk], it.head * K_SUBTILES + s); + } + } + }; + auto load_v_oneslot = [&](int mtile_idx, int ktile_idx, int p) { + const int kv_stage = kv_empty_ph.get_stage(); + + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar[kv_stage]), kv_empty_ph.get_phase()); + kv_empty_ph.advance(); + + if constexpr (BLK128) { + + const int kv_tok0 = k_start + get_kv_block_id(mtile_idx, ktile_idx) * BLOCK; + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar[kv_stage]), KV_RING_SLOT_BYTES); + + if constexpr (BHSD) + tma_load_4d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES)), &tmap_v, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok0 - k_start, 0, it.sample * num_heads + it.head); + else + tma_load_3d(smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES)), &tmap_v, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok0, it.head * K_SUBTILES); + } + } else { + int kv_tok[2]; + #pragma unroll + for (int h = 0; h < 2; ++h) + kv_tok[h] = k_start + get_kv_block_id(mtile_idx, ktile_idx * BLOCKS_PER_KTILE + 2 * h + p) * BLOCK; + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar[kv_stage]), KV_RING_SLOT_BYTES); + + #pragma unroll + for (int h = 0; h < 2; ++h) { + #pragma unroll + for (int s2 = 0; s2 < K_SUBTILES; ++s2) { + const uint32_t dst = smem_ptr_u32((sKV + kv_stage * KV_RING_SLOT_BYTES) + + h * V_BLK_BYTES + s2 * (BLOCK * SUB_COLS_BYTES)); + if constexpr (BHSD) + tma_load_4d(dst, &tmap_v, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[h] - k_start, s2, it.sample * num_heads + it.head); + else + tma_load_3d(dst, &tmap_v, smem_ptr_u32(&full_bar[kv_stage]), + 0, kv_tok[h], it.head * K_SUBTILES + s2); + } + } + } + } + }; + auto load_k = [&](int mtile_idx, int ktile_idx) { + load_k_oneslot(mtile_idx, ktile_idx, 0); + if constexpr (!BLK128) load_k_oneslot(mtile_idx, ktile_idx, 1); + }; + auto load_v = [&](int mtile_idx, int ktile_idx) { + load_v_oneslot(mtile_idx, ktile_idx, 0); + if constexpr (!BLK128) load_v_oneslot(mtile_idx, ktile_idx, 1); + }; + + load_k(0, 0); load_k(1, 0); + + #pragma unroll + for (int m = 0; m < M_TILES_PER_CTA; ++m) { + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_q[m]), q_empty_ph.get_phase()); + const int tok0 = q_start + mtile[m] * BLOCK; + if (elect_one_sync()) { + mbarrier_arrive_expect_tx(smem_ptr_u32(&full_bar_q[m]), Q_TILE_BYTES); + + tma_load_4d(smem_ptr_u32(sQ[m]), &tmap_q, smem_ptr_u32(&full_bar_q[m]), + 0, BHSD ? it.sample * num_heads + it.head : it.head, + BHSD ? tok0 - q_start : tok0, 0); + } + } + q_empty_ph.advance(); + for (int ktile_idx = 0; ktile_idx + 1 < num_k_tiles; ++ktile_idx) { + load_v(0, ktile_idx); load_k(0, ktile_idx + 1); + load_v(1, ktile_idx); load_k(1, ktile_idx + 1); + } + load_v(0, num_k_tiles - 1); load_v(1, num_k_tiles - 1); + + 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); + if (!next.valid) break; + workitem_id = next.n_tile; + } else { + workitem_id += gridDim.x; + if (workitem_id >= total_workitems) break; + } + } + } + else if (warp_id == W_MMA) { + setmaxnreg_dec<48>(); + + const uint32_t lead = elect_one_sync() ? 1u : 0u; + constexpr uint32_t DESC_SBO = 1024, DESC_LBO = 16; + + const uint32_t idesc_qk = make_idesc_bf16_f32(M_TILE, K_TILE, false, false); + const uint32_t idesc_pv = make_idesc_bf16_f32(M_TILE, BLK128 ? HEAD_DIM : 2 * HEAD_DIM, + false, true); + + constexpr uint32_t DESC_LBO_MN = (uint32_t)((BLK128 ? K_TILE : BLOCK) * SUB_COLS_BF16 * 2); + const uint64_t desc_v0 = build_smem_desc_blackwell(smem_ptr_u32(sKV), DESC_SBO, DESC_LBO_MN, + SmemSwizzleBlackwell::B128); + constexpr uint32_t PV_DESC_STEP = (uint32_t)((16 * SUB_COLS_BF16 * 2) >> 4); + const uint64_t desc_q0 = build_smem_desc_blackwell(smem_ptr_u32(sQ0), DESC_SBO, DESC_LBO, SmemSwizzleBlackwell::B128); + const uint64_t desc_kv0 = build_smem_desc_blackwell(smem_ptr_u32(sKV), DESC_SBO, DESC_LBO, SmemSwizzleBlackwell::B128); + + constexpr uint64_t KV_DESC_DELTA = KV_RING_SLOT_BYTES >> 4; + constexpr uint64_t Q_SUB_DELTA = Q_SUB_COLS_BYTES >> 4; + constexpr uint64_t Q_MTILE_DESC_DELTA = Q_TILE_BYTES >> 4; + + PhaseTracker kv_ph; + PhaseTracker<1> q_ph; + PhaseTracker<1> spo_ph; + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (true) { + const WorkItem it = decode_workitem(workitem_id, num_heads, num_blocks, packed_mtiles_per_seq, magic0, magic1, magic2, q2k_num); + const int num_k_tiles = (it.num_kv_blocks + BLOCKS_PER_KTILE - 1) / BLOCKS_PER_KTILE; + + auto bmm1 = [&](int i) { + const uint32_t s_tmem_addr = tmem_base + (uint32_t)(i * S_COLS); + const uint64_t da_base = desc_q0 + (uint64_t)i * Q_MTILE_DESC_DELTA; + + int slot = 0; + #pragma unroll + for (int s = 0; s < Q_SUBTILES; ++s) { + if (!BLK128 || s == 0) { + slot = kv_ph.get_stage(); + mbarrier_wait_parity(smem_ptr_u32(&full_bar[slot]), kv_ph.get_phase()); + kv_ph.advance(); + } + + const uint64_t da = da_base + (uint64_t)s * Q_SUB_DELTA; + const uint64_t db = desc_kv0 + (uint64_t)slot * KV_DESC_DELTA + + (BLK128 ? (uint64_t)s * (K_SUB_COLS_BYTES >> 4) : 0u); + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_TILE; ++ki) { + const bool enable_d = (s != 0) || (ki != 0); + if constexpr (BLK128) + tcgen05_mma_f16_ss_lead(lead, s_tmem_addr, da + 2 * ki, db + 2 * ki, idesc_qk, enable_d); + else + tcgen05_mma_ws_f16_ss_1sm_lead(lead, s_tmem_addr, da + 2 * ki, db + 2 * ki, idesc_qk, enable_d); + } + if (!BLK128 || s == Q_SUBTILES - 1) + tcgen05_commit1_lead(lead, smem_ptr_u32(&empty_bar[slot])); + } + + tcgen05_commit1_lead(lead, smem_ptr_u32(&full_bar_spo[i])); + }; + + auto bmm2 = [&](int i, bool first_ktile, auto last_c) { + constexpr bool last_ktile = decltype(last_c)::value; + const uint32_t s_tmem_addr = tmem_base + (uint32_t)(i * S_COLS); + const uint32_t o_tmem_addr = tmem_base + (uint32_t)(2 * S_COLS + i * O_COLS); + + int slot = 0; + #pragma unroll + for (int p = 0; p < V_SUBTILES; ++p) { + SmemDescPair desc_bv; + if (!BLK128 || p == 0) { + slot = kv_ph.get_stage(); + mbarrier_wait_parity(smem_ptr_u32(&full_bar[slot]), kv_ph.get_phase()); + kv_ph.advance(); + } + if (!BLK128 || p == 0) { + desc_bv.u64 = desc_v0; + desc_bv.w.x += (uint32_t)(slot * (int)KV_DESC_DELTA); + } + const uint64_t dbV = desc_kv0 + (uint64_t)slot * KV_DESC_DELTA + + (BLK128 ? (uint64_t)p * (V_BLK_BYTES >> 4) : 0u); + #pragma unroll + for (int ki = 0; ki < K_ATOMS_PER_TILE; ++ki) { + const int a = p * K_ATOMS_PER_TILE + ki; + if constexpr (SPLIT_P) if (a == SPLIT_P_ATOM) { + mbarrier_wait_parity(smem_ptr_u32(&full_bar_p_last[i]), spo_ph.get_phase()); + } + const bool accumulate = (!first_ktile) || (a != 0); + if constexpr (BLK128) { + tcgen05_mma_f16_ts_1sm_lead(lead, o_tmem_addr, s_tmem_addr + (uint32_t)(a * 8), desc_bv.u64, idesc_pv, accumulate); + desc_add_lo(desc_bv, PV_DESC_STEP); + } else { + tcgen05_mma_ws_f16_ts_1sm_lead(lead, o_tmem_addr, s_tmem_addr + (uint32_t)(a * 8), desc_bv.u64, idesc_pv, accumulate); + desc_add_lo(desc_bv, PV_DESC_STEP); + } + } + + if (!BLK128 || p == V_SUBTILES - 1) { + tcgen05_commit1_lead(lead, smem_ptr_u32(&empty_bar[slot])); + } + } + + if constexpr (last_ktile) { + tcgen05_commit1_lead(lead, smem_ptr_u32(&full_bar_o_acc[i])); + } + }; + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + mbarrier_wait_parity(smem_ptr_u32(&full_bar_q[i]), q_ph.get_phase()); + bmm1(i); + } + + for (int k = 0; k + 1 < num_k_tiles; ++k) { + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + mbarrier_wait_parity(smem_ptr_u32(&empty_bar_spo[i]), spo_ph.get_phase()); + bmm2(i, (k == 0), std::false_type{}); + bmm1(i); + } + + spo_ph.advance(); + } + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) + tcgen05_commit1_lead(lead, smem_ptr_u32(&empty_bar_q[i])); + q_ph.advance(); + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + mbarrier_wait_parity(smem_ptr_u32(&empty_bar_spo[i]), spo_ph.get_phase()); + bmm2(i, (num_k_tiles == 1), std::true_type{}); + } + + spo_ph.advance(); + + 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); + if (!next.valid) break; + workitem_id = next.n_tile; + } else { + workitem_id += gridDim.x; + if (workitem_id >= total_workitems) break; + } + } + } + else if (warp_id == W_EPI) { + setmaxnreg_dec<48>(); + + PhaseTracker<1> full_o_ph; + + if (elect_one_sync()) { + #pragma unroll + for (int m = 0; m < M_TILES_PER_CTA; ++m) + mbarrier_arrive(smem_ptr_u32(&empty_bar_o_epi[m])); + } + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (true) { + const WorkItem it = decode_workitem(workitem_id, num_heads, num_blocks, packed_mtiles_per_seq, magic0, magic1, magic2, q2k_num); + const int q_start = it.sample * seqlen; + const int mtile[2] = { it.mtile0, it.mtile1 }; + + #pragma unroll + for (int m = 0; m < M_TILES_PER_CTA; ++m) { + mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_o_epi[m]), full_o_ph.get_phase()); + + if (elect_one_sync()) { + const int tok0 = q_start + mtile[m] * BLOCK; + + tma_store_4d(&tmap_o, 0, BHSD ? it.sample * num_heads + it.head : it.head, + BHSD ? tok0 - q_start : tok0, 0, + smem_ptr_u32(reinterpret_cast(sO_bufs[m]))); + cp_async_bulk_commit_group(); + } + } + + if (elect_one_sync()) { + cp_async_bulk_wait_group_read<1>(); + mbarrier_arrive(smem_ptr_u32(&empty_bar_o_epi[0])); + cp_async_bulk_wait_group_read<0>(); + mbarrier_arrive(smem_ptr_u32(&empty_bar_o_epi[1])); + } + + full_o_ph.advance(); + 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); + if (!next.valid) break; + workitem_id = next.n_tile; + } else { + workitem_id += gridDim.x; + if (workitem_id >= total_workitems) break; + } + } + } + else if (warp_id == W_SCHED) { + setmaxnreg_dec<48>(); + + 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 next = 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 (!next.valid) break; + } + + for (int s = 0; s < CLC_STAGES; ++s) { + if (lane == 0) + mbarrier_wait_parity_suspend(smem_ptr_u32(&clc_empty[prod_stage]), prod_phase); + __syncwarp(); + advance_stage_phase(prod_stage, prod_phase); + } + } + } + else if (warp_id >= W_CORR0 && warp_id < W_MMA) { + setmaxnreg_dec<80>(); + + const int corr_warp_id = warp_id - W_CORR0; + const int corr_tid = corr_warp_id * 32 + lane; + const int corr_row = BLK128 ? corr_tid : (corr_tid & 63); + const bool kv_half0 = BLK128 || corr_tid < 64; + [[maybe_unused]] PhaseTracker<1> alpha_ph; + PhaseTracker<1> o_acc_ph; + PhaseTracker<1> o_epi_empty_ph; + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + mbarrier_arrive(smem_ptr_u32(&empty_bar_spo[i])); + mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + } + + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + int workitem_id = (int)blockIdx.x; + while (true) { + const WorkItem it = decode_workitem(workitem_id, num_heads, num_blocks, packed_mtiles_per_seq, magic0, magic1, magic2, q2k_num); + const int num_k_tiles = (it.num_kv_blocks + BLOCKS_PER_KTILE - 1) / BLOCKS_PER_KTILE; + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + if constexpr (FULL_NAMED_BAR) full_bar_wait(i, corr_warp_id); + else mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_alpha[i]), alpha_ph.get_phase()); + mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + } + if constexpr (!FULL_NAMED_BAR) alpha_ph.advance(); + + for (int k = 1; k < num_k_tiles; ++k) { + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + if constexpr (FULL_NAMED_BAR) full_bar_wait(i, corr_warp_id); + else mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_alpha[i]), alpha_ph.get_phase()); + + const float alpha = alpha_and_l_smem[(i * STAT_REGIONS + 0) * STATS + corr_tid]; + if constexpr (!SOFTMAX_THROTTLE) mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + + bool skip = __all_sync(0xffffffffu, alpha == 1.0f); + if (!skip) { + + const uint32_t o_tmem_addr = tmem_base + (uint32_t)(2 * S_COLS + i * O_COLS) + ((uint32_t)(corr_warp_id * 32) << 16); + + const float2 alpha2 = f32x2_splat(alpha); + + #pragma unroll + for (int c0 = 0; c0 < HEAD_DIM; c0 += 16) { + uint32_t o_regs[16]; + tcgen05_ld_32x32b_x16(o_tmem_addr + (uint32_t)c0, o_regs); + float2* o2 = reinterpret_cast(o_regs); + #pragma unroll + for (int e = 0; e < 8; ++e) o2[e] = fmul2(o2[e], alpha2); + tcgen05_st_32x32b_x16(o_tmem_addr + (uint32_t)c0, o_regs); + } + tcgen05_wait_st(); + + tcgen05_fence_before_thread_sync(); + } + if constexpr (SOFTMAX_THROTTLE) mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + mbarrier_arrive(smem_ptr_u32(&empty_bar_spo[i])); + } + if constexpr (!FULL_NAMED_BAR) alpha_ph.advance(); + } + + #pragma unroll + for (int i = 0; i < M_TILES_PER_CTA; ++i) { + mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_o_acc[i]), o_acc_ph.get_phase()); + if constexpr (FULL_NAMED_BAR) full_bar_wait(i, corr_warp_id); + else mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_l[i]), o_acc_ph.get_phase()); + + float scale_own; + if constexpr (BLK128) { + + const float l = alpha_and_l_smem[(i * STAT_REGIONS + 1) * STATS + corr_tid]; + mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + scale_own = (l > 0.f) ? rcp_approx_ftz_f32(l) : 0.f; + } else { + const float l_own = alpha_and_l_smem[(i * STAT_REGIONS + 1) * STATS + corr_tid]; + const float m_own = alpha_and_l_smem[(i * STAT_REGIONS + 2) * STATS + corr_tid]; + const float l_par = alpha_and_l_smem[(i * STAT_REGIONS + 1) * STATS + (corr_tid ^ 64)]; + const float m_par = alpha_and_l_smem[(i * STAT_REGIONS + 2) * STATS + (corr_tid ^ 64)]; + mbarrier_arrive(smem_ptr_u32(&empty_bar_alpha_and_l[i])); + const float d = (m_own - m_par) * scale_log2; + const float beta_lo = ex2_approx_f32(-fabsf(d)); + const float beta_own = (d >= 0.f) ? 1.f : beta_lo; + const float beta_par = (d >= 0.f) ? beta_lo : 1.f; + const float l_tot = beta_own * l_own + beta_par * l_par; + scale_own = (l_tot > 0.f) ? beta_own * rcp_approx_ftz_f32(l_tot) : 0.f; + + if (lse_out != nullptr && (corr_tid >> 6) == 0) { + const float m_tot = (d >= 0.f) ? m_own : m_par; + const int q_row = (i == 0 ? it.mtile0 : it.mtile1) * BLOCK + (corr_tid & 63); + const float l_safe = (l_tot > 0.f) ? l_tot : 1.0f; + lse_out[((long)it.sample * num_heads + it.head) * (long)seqlen + q_row] = + m_tot * scale_log2 + __log2f(l_safe); + } + } + const float2 scale2 = f32x2_splat(scale_own); + const uint32_t o_tmem_addr = tmem_base + (uint32_t)(2 * S_COLS + i * O_COLS) + ((uint32_t)(corr_warp_id * 32) << 16); + if constexpr (BLK128) { + + #pragma unroll + for (int c0 = 0; c0 < HEAD_DIM; c0 += 16) { + uint32_t o_regs[16]; + tcgen05_ld_32x32b_x16(o_tmem_addr + (uint32_t)c0, o_regs); + if (c0 == 0) mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_o_epi[i]), o_epi_empty_ph.get_phase()); + const float2* o2 = reinterpret_cast(o_regs); + const int s = c0 / SUB_COLS_BF16; + const int v_base = (c0 % SUB_COLS_BF16) / 8; + __nv_bfloat16* so_sub = sO_bufs[i] + s * (M_TILE * SUB_COLS_BF16); + #pragma unroll + for (int vv = 0; vv < 2; ++vv) { + const int v = v_base + vv; + const float2 r0 = fmul2(o2[vv * 4 + 0], scale2); + const float2 r1 = fmul2(o2[vv * 4 + 1], scale2); + const float2 r2 = fmul2(o2[vv * 4 + 2], scale2); + const float2 r3 = fmul2(o2[vv * 4 + 3], scale2); + 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(&so_sub[corr_row * SUB_COLS_BF16 + (v ^ (corr_row & 7)) * 8]) = packed; + } + } + } else { + + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_o_epi[i]), o_epi_empty_ph.get_phase()); + if (!kv_half0) { + #pragma unroll + for (int c0 = 0; c0 < HEAD_DIM; c0 += 16) { + uint32_t o_regs[16]; + tcgen05_ld_32x32b_x16(o_tmem_addr + (uint32_t)c0, o_regs); + float2* o2 = reinterpret_cast(o_regs); + #pragma unroll + for (int e = 0; e < 8; ++e) o2[e] = fmul2(o2[e], scale2); + const int s = c0 / SUB_COLS_BF16; + const int v_base = (c0 % SUB_COLS_BF16) / 8; + __nv_bfloat16* so_sub = sO_bufs[i] + s * (M_TILE * SUB_COLS_BF16); + #pragma unroll + for (int vv = 0; vv < 2; ++vv) { + const int v = v_base + vv; + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(o2[vv * 4 + 0].x, o2[vv * 4 + 0].y); + packed.y = cvt_f32x2_to_bf16x2(o2[vv * 4 + 1].x, o2[vv * 4 + 1].y); + packed.z = cvt_f32x2_to_bf16x2(o2[vv * 4 + 2].x, o2[vv * 4 + 2].y); + packed.w = cvt_f32x2_to_bf16x2(o2[vv * 4 + 3].x, o2[vv * 4 + 3].y); + *reinterpret_cast(&so_sub[corr_row * SUB_COLS_BF16 + (v ^ (corr_row & 7)) * 8]) = packed; + } + } + } + bar_sync<9>(128); + if (kv_half0) { + #pragma unroll + for (int c0 = 0; c0 < HEAD_DIM; c0 += 16) { + uint32_t o_regs[16]; + tcgen05_ld_32x32b_x16(o_tmem_addr + (uint32_t)c0, o_regs); + float2* o2 = reinterpret_cast(o_regs); + #pragma unroll + for (int e = 0; e < 8; ++e) o2[e] = fmul2(o2[e], scale2); + const int s = c0 / SUB_COLS_BF16; + const int v_base = (c0 % SUB_COLS_BF16) / 8; + __nv_bfloat16* so_sub = sO_bufs[i] + s * (M_TILE * SUB_COLS_BF16); + #pragma unroll + for (int vv = 0; vv < 2; ++vv) { + const int v = v_base + vv; + __nv_bfloat16* dst = &so_sub[corr_row * SUB_COLS_BF16 + (v ^ (corr_row & 7)) * 8]; + uint4 h1 = *reinterpret_cast(dst); + o2[vv * 4 + 0] = fadd2(o2[vv * 4 + 0], __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&h1.x))); + o2[vv * 4 + 1] = fadd2(o2[vv * 4 + 1], __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&h1.y))); + o2[vv * 4 + 2] = fadd2(o2[vv * 4 + 2], __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&h1.z))); + o2[vv * 4 + 3] = fadd2(o2[vv * 4 + 3], __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(&h1.w))); + uint4 packed; + packed.x = cvt_f32x2_to_bf16x2(o2[vv * 4 + 0].x, o2[vv * 4 + 0].y); + packed.y = cvt_f32x2_to_bf16x2(o2[vv * 4 + 1].x, o2[vv * 4 + 1].y); + packed.z = cvt_f32x2_to_bf16x2(o2[vv * 4 + 2].x, o2[vv * 4 + 2].y); + packed.w = cvt_f32x2_to_bf16x2(o2[vv * 4 + 3].x, o2[vv * 4 + 3].y); + *reinterpret_cast(dst) = packed; + } + } + } + bar_sync<9>(128); + } + + tcgen05_fence_before_thread_sync(); + + mbarrier_arrive(smem_ptr_u32(&empty_bar_spo[i])); + + fence_proxy_async_shared(); + + mbarrier_arrive(smem_ptr_u32(&full_bar_o_epi[i])); + } + o_acc_ph.advance(); + o_epi_empty_ph.advance(); + + 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); + if (!next.valid) break; + workitem_id = next.n_tile; + } else { + workitem_id += gridDim.x; + if (workitem_id >= total_workitems) break; + } + } + } + else { + setmaxnreg_inc<192>(); + + const int warp_id_u = __shfl_sync(0xffffffffu, warp_id, 0); + const int m_tile = warp_id_u < 4 ? 0 : 1; + const int warp_in_group = warp_id_u & 3; + + const int sm_tid = warp_in_group * 32 + lane; + const uint32_t alpha_slot_u32 = smem_ptr_u32(&alpha_and_l_smem[(m_tile * STAT_REGIONS + 0) * STATS + sm_tid]); + const uint32_t s_tmem_addr = tmem_base + (uint32_t)(m_tile * S_COLS) + ((uint32_t)(warp_in_group * 32) << 16); + PhaseTracker<1> spo_ph; + PhaseTracker<1> scale_empty_ph; + [[maybe_unused]] int clc_stage = 0; + [[maybe_unused]] uint32_t clc_phase = 0; + + int workitem_id = (int)blockIdx.x; + while (true) { + const WorkItem it = decode_workitem(workitem_id, num_heads, num_blocks, packed_mtiles_per_seq, magic0, magic1, magic2, q2k_num); + const int num_k_tiles = (it.num_kv_blocks + BLOCKS_PER_KTILE - 1) / BLOCKS_PER_KTILE; + + float m_run = -INFINITY, l_run = 0.f; + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_alpha_and_l[m_tile]), scale_empty_ph.get_phase()); + scale_empty_ph.advance(); + + int thr_window = -(1 << 30); + int thr_cache0 = BLOCK, thr_cache1 = BLOCK; + auto get_vbs_thresholds = [&](int k, int gqb_mt, int half, int& t0, int& t1) { + if (k >= thr_window + 32) { + thr_window = k & ~31; + const int kk = thr_window + lane; + // THIS tile's own count (this warp group serves exactly one m_tile). + // Positions at or past it -- the pair-max padding, and everything in + // an empty row -- get threshold 0, i.e. the whole 64/128-token block + // masks to -inf, regardless of which KV block the load warp fetched. + const int cnt = it.num_kv_blocks_mt[m_tile]; + if constexpr (BLK128) { + thr_cache0 = (kk < cnt) + ? variable_block_sizes[q2k_idx[gqb_mt * max_kv + kk]] : 0; + } else { + const int b0 = kk * BLOCKS_PER_KTILE + 2 * half; + thr_cache0 = (b0 < cnt) + ? variable_block_sizes[q2k_idx[gqb_mt * max_kv + b0]] : 0; + thr_cache1 = (b0 + 1 < cnt) + ? variable_block_sizes[q2k_idx[gqb_mt * max_kv + b0 + 1]] : 0; + } + } + t0 = __shfl_sync(0xffffffffu, thr_cache0, k & 31); + t1 = BLK128 ? 0 : __shfl_sync(0xffffffffu, thr_cache1, k & 31); + }; + auto softmax_step = [&](auto is_first_c, int k) { + constexpr bool IS_FIRST = decltype(is_first_c)::value; + const int gqb_mt_ = (m_tile == 0) ? it.global_mtile0 : it.global_mtile1; + int vbs_thr0_, vbs_thr1_; + get_vbs_thresholds(k, gqb_mt_, warp_in_group >> 1, vbs_thr0_, vbs_thr1_); + mbarrier_wait_parity_suspend(smem_ptr_u32(&full_bar_spo[m_tile]), spo_ph.get_phase()); + + uint32_t s_regs[S_COLS]; + #pragma unroll + for (int c0 = 0; c0 < S_COLS; c0 += S_LD_COLS) { + const uint32_t taddr = s_tmem_addr + (uint32_t)c0; + if constexpr (S_LD_COLS == 32) tcgen05_ld_32x32b_x32 (taddr, *reinterpret_cast(&s_regs[c0])); + else if constexpr (S_LD_COLS == 64) tcgen05_ld_32x32b_x64 (taddr, *reinterpret_cast(&s_regs[c0])); + else if constexpr (S_LD_COLS == 128) tcgen05_ld_32x32b_x128(taddr, *reinterpret_cast(&s_regs[c0])); + } + + float* scores = reinterpret_cast(s_regs); + float2* scores2 = reinterpret_cast(s_regs); + + tcgen05_fence_before_thread_sync(); + + if constexpr (BLK128) { + if (vbs_thr0_ < S_COLS) mask_s_row_r2p(scores, 0, 0, vbs_thr0_); + } else { + if (vbs_thr0_ < 64) mask_s_row_r2p(scores, 0, 0, vbs_thr0_); + if (vbs_thr1_ < 64) mask_s_row_r2p(scores + 64, 0, 0, vbs_thr1_); + } + + float rmax0 = m_run, rmax1 = -INFINITY, rmax2 = -INFINITY, rmax3 = -INFINITY; + #pragma unroll + for (int j = 0; j < S_COLS; j += 8) { + rmax0 = fmaxf(fmaxf(rmax0, scores[j + 0]), scores[j + 1]); + rmax1 = fmaxf(fmaxf(rmax1, scores[j + 2]), scores[j + 3]); + rmax2 = fmaxf(fmaxf(rmax2, scores[j + 4]), scores[j + 5]); + rmax3 = fmaxf(fmaxf(rmax3, scores[j + 6]), scores[j + 7]); + } + float new_m = fmaxf(fmaxf(rmax0, rmax1), fmaxf(rmax2, rmax3)); + + new_m = fmaxf(new_m, -FLT_MAX); + float alpha = 0.0f; + if constexpr (!IS_FIRST) { + + const float acc_scale_ = (m_run - new_m) * scale_log2; + if (acc_scale_ >= -(float)RESCALE_THRESHOLD) { new_m = m_run; alpha = 1.0f; } + else { alpha = ex2_approx_f32(acc_scale_); } + + sts_f32(alpha_slot_u32, alpha); + } + if constexpr (FULL_NAMED_BAR) full_bar_arrive(m_tile, warp_in_group); + else mbarrier_arrive(smem_ptr_u32(&full_bar_alpha[m_tile])); + + const float2 scale2 = f32x2_splat(scale_log2); + const float2 neg_m_scaled2 = f32x2_splat(-new_m * scale_log2); + uint32_t p_regs[S_COLS / 2]; + [[maybe_unused]] float2 lt2_live = make_float2(IS_FIRST ? 0.0f : l_run * alpha, 0.0f); + #pragma unroll + for (int c = 0; c < S_COLS / 2; ++c) { + const float2 a2 = ffma2(scores2[c], scale2, neg_m_scaled2); + + if constexpr (EX2_EMU) { + const int jj = c / EX2_FRG_PAIRS; + const int kk = 2 * (c % EX2_FRG_PAIRS); + const bool use_hw = (kk % EX2_FREQ < EX2_FREQ - EX2_RES) || (jj >= EX2_FRG_CNT - 1); + scores2[c] = use_hw ? make_float2(ex2_approx_f32(a2.x), ex2_approx_f32(a2.y)) + : ex2_emu_f32x2(a2.x, a2.y); + } else { + scores2[c] = make_float2(ex2_approx_f32(a2.x), ex2_approx_f32(a2.y)); + } + if constexpr (!DEFER_ROWSUM) lt2_live = fadd2(lt2_live, scores2[c]); + p_regs[c] = cvt_f32x2_to_bf16x2(scores2[c].x, scores2[c].y); + } + const uint32_t p_tmem_addr = s_tmem_addr; + + if constexpr (SPLIT_P) { + tcgen05_st_32x32b_x32(p_tmem_addr, *reinterpret_cast(&p_regs[0])); + tcgen05_st_32x32b_x16(p_tmem_addr + 32, *reinterpret_cast(&p_regs[32])); + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + mbarrier_arrive(smem_ptr_u32(&empty_bar_spo[m_tile])); + tcgen05_st_32x32b_x16(p_tmem_addr + SPLIT_P_COL, *reinterpret_cast(&p_regs[SPLIT_P_COL])); + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + mbarrier_arrive(smem_ptr_u32(&full_bar_p_last[m_tile])); + } else { + tcgen05_st_32x32b_x32(p_tmem_addr, *reinterpret_cast(&p_regs[0])); + tcgen05_st_32x32b_x32(p_tmem_addr + 32, *reinterpret_cast(&p_regs[32])); + tcgen05_wait_st(); + tcgen05_fence_before_thread_sync(); + mbarrier_arrive(smem_ptr_u32(&empty_bar_spo[m_tile])); + } + + spo_ph.advance(); + float2 lt2 = lt2_live; + if constexpr (DEFER_ROWSUM) { + + float2 lt2a = make_float2(IS_FIRST ? 0.0f : l_run * alpha, 0.0f), lt2b = make_float2(0.f, 0.f); + float2 lt2c = make_float2(0.f, 0.f), lt2d = make_float2(0.f, 0.f); + #pragma unroll + for (int c = 0; c < S_COLS / 2; c += 4) { + lt2a = fadd2(lt2a, scores2[c + 0]); + lt2b = fadd2(lt2b, scores2[c + 1]); + lt2c = fadd2(lt2c, scores2[c + 2]); + lt2d = fadd2(lt2d, scores2[c + 3]); + } + lt2 = fadd2(fadd2(lt2a, lt2b), fadd2(lt2c, lt2d)); + } + l_run = lt2.x + lt2.y; m_run = new_m; + mbarrier_wait_parity_suspend(smem_ptr_u32(&empty_bar_alpha_and_l[m_tile]), scale_empty_ph.get_phase()); + scale_empty_ph.advance(); + }; + if (num_k_tiles > 0) softmax_step(std::true_type{}, 0); + for (int k = 1; k < num_k_tiles; ++k) softmax_step(std::false_type{}, k); + + if constexpr (BLK128) { + if (lse_out != nullptr) { + const int q_row = (m_tile == 0 ? it.mtile0 : it.mtile1) * BLOCK + sm_tid; + const float l_safe = (l_run > 0.f) ? l_run : 1.0f; + lse_out[((long)it.sample * num_heads + it.head) * (long)seqlen + q_row] = + m_run * scale_log2 + __log2f(l_safe); + } + } + + alpha_and_l_smem[(m_tile * STAT_REGIONS + 1) * STATS + sm_tid] = l_run; + if constexpr (!BLK128) + alpha_and_l_smem[(m_tile * STAT_REGIONS + 2) * STATS + sm_tid] = m_run; + if constexpr (FULL_NAMED_BAR) full_bar_arrive(m_tile, warp_in_group); + else mbarrier_arrive(smem_ptr_u32(&full_bar_l[m_tile])); + + 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); + if (!next.valid) break; + workitem_id = next.n_tile; + } else { + workitem_id += gridDim.x; + if (workitem_id >= total_workitems) break; + } + } + } + __syncthreads(); + if (warp_id == 0) tcgen05_dealloc<1>(tmem_base, TMEM_TOTAL); +#endif // host pass or sm_100a device pass (multi-arch guard; see note at the top of the body) +} + +} // namespace VSA_NAMESPACE + +#endif diff --git a/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh b/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh new file mode 100644 index 000000000..7b30dfb41 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_launch_sm100a.cuh @@ -0,0 +1,201 @@ +#ifndef BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH +#define BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH + +// Launch surface for the sm_100a VSA block-sparse FMHA forward. +// +// Everything a caller needs: a POD argument struct, a predicate saying whether this build can +// run those arguments, and one launch entry point. The benchmark in +// block_sparse_bench_sm100a.cu and the torch binding both go through here, so there is +// one tensormap construction and one launch configuration rather than two that can drift. +// +// Two compile-time knobs select the four builds: +// VSA_BLK128 false -> 64-token sparse blocks, true -> 128-token +// VSA_BHSD false -> [token][head][dim] (BSHD), true -> [batch][head][token][dim] (BHSD) + +#include "block_sparse_kernel_sm100a.cuh" + +namespace VSA_NAMESPACE { + +struct BlockSparseVsaArgs { + const __nv_bfloat16* q; + const __nv_bfloat16* k; + const __nv_bfloat16* v; // natural layout; only blk128 reads it (blk64 still needs v_t) + const __nv_bfloat16* v_t; // unused: kept so the bench's V_T buffer still binds + __nv_bfloat16* o; + float* lse; // [batch, num_heads, seqlen] fp32, or nullptr + + const int* q2k_idx; // [batch*num_heads*num_blocks, max_kv] int32 + const int* q2k_num; // [batch*num_heads*num_blocks] int32 + const int* variable_block_sizes; // [num_blocks] int32, valid tokens per block + + int batch; + int num_heads; + int seqlen; + int head_dim; + int num_blocks; + int max_kv; + float sm_scale; +}; + +// cudaSuccess iff this build can run `a`. Deliberately conservative: the caller is expected +// to fall back to its own implementation rather than get a wrong answer. +__host__ inline cudaError_t block_sparse_supported(const BlockSparseVsaArgs& a) { + if (a.head_dim != HEAD_DIM) return cudaErrorInvalidValue; // compile-time in the kernel + if (a.num_blocks % 2 != 0) return cudaErrorInvalidValue; // a CTA owns an adjacent pair + if (a.seqlen != a.num_blocks * BLOCK) return cudaErrorInvalidValue; + if (a.max_kv < 1 || a.num_blocks < 1) return cudaErrorInvalidValue; + if (a.q == nullptr || a.k == nullptr || a.o == nullptr) return cudaErrorInvalidValue; + if (a.q2k_idx == nullptr || a.q2k_num == nullptr) return cudaErrorInvalidValue; + // FastVideo always supplies this; without it padded keys would be attended as real zeros. + if (a.variable_block_sizes == nullptr) return cudaErrorInvalidValue; + // V is read MN-major at BOTH block sizes now, so no pre-transposed V_T is ever needed. + if (a.v == nullptr) return cudaErrorInvalidValue; + return cudaSuccess; +} + +__host__ inline cudaError_t launch_block_sparse_sm100a(const BlockSparseVsaArgs& a, + cudaStream_t stream) { + const cudaError_t sup = block_sparse_supported(a); + if (sup != cudaSuccess) return sup; + + const int B = a.batch, H = a.num_heads, S = a.seqlen, hd = a.head_dim; + const int num_blocks = a.num_blocks, max_kv = a.max_kv; + const long tq = (long)B * S; + const int packed_mtiles_per_seq = num_blocks / 2; + const int total_work = B * H * packed_mtiles_per_seq; + constexpr bool BHSD = VSA_BHSD; + + CUtensorMap tq_, tk_, tvt_, tv_, to_; + { + uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16, BHSD ? (uint64_t)((long)B * H) : (uint64_t)H, + BHSD ? (uint64_t)S : (uint64_t)tq, (uint64_t)Q_SUBTILES }; + uint64_t gs[3] = { BHSD ? (uint64_t)((long)S * hd) * 2u : (uint64_t)hd * 2u, + BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u, + (uint64_t)SUB_COLS_BF16 * 2u }; + uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, 1u, (uint32_t)M_TILE, (uint32_t)Q_SUBTILES }; + uint32_t es[4] = { 1u, 1u, 1u, 1u }; + if (cuTensorMapEncodeTiled(&tq_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, + const_cast<__nv_bfloat16*>(a.q), 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) != CUDA_SUCCESS) + return cudaErrorInvalidValue; + if (cuTensorMapEncodeTiled(&to_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, a.o, 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) != CUDA_SUCCESS) + return cudaErrorInvalidValue; + } + { + uint64_t gd[4] = { (uint64_t)SUB_COLS_BF16, + BHSD ? (uint64_t)S : (uint64_t)tq, + BHSD ? (uint64_t)(hd / SUB_COLS_BF16) + : (uint64_t)((long)H * hd / SUB_COLS_BF16), + (uint64_t)((long)B * H) }; + uint64_t gs[3] = { BHSD ? (uint64_t)hd * 2u : (uint64_t)((long)H * hd) * 2u, + (uint64_t)SUB_COLS_BF16 * 2u, + (uint64_t)((long)S * hd) * 2u }; + uint32_t bd[4] = { (uint32_t)SUB_COLS_BF16, (uint32_t)BLOCK, + BLK128 ? (uint32_t)K_SUBTILES : 1u, 1u }; + uint32_t es[4] = { 1u, 1u, 1u, 1u }; + if (cuTensorMapEncodeTiled(&tk_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3, + const_cast<__nv_bfloat16*>(a.k), 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) != CUDA_SUCCESS) + return cudaErrorInvalidValue; + // V map is byte-for-byte the K map over a.v: MN-major V needs no transpose (blk128). + const __nv_bfloat16* vbase = a.v ? a.v : a.k; + if (cuTensorMapEncodeTiled(&tv_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, BHSD ? 4 : 3, + const_cast<__nv_bfloat16*>(vbase), 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) != CUDA_SUCCESS) + return cudaErrorInvalidValue; + } + // V_T map: blk64 only. Unused at blk128 but must still be a valid tensormap to pass by value. + { + const __nv_bfloat16* vt = a.v_t ? a.v_t : a.k; + if constexpr (BLK128) { + uint64_t gd[3] = { (uint64_t)SUB_COLS_BF16, (uint64_t)((long)H * hd), + (uint64_t)((long)tq / SUB_COLS_BF16) }; + uint64_t gs[2] = { (uint64_t)tq * 2u, (uint64_t)SUB_COLS_BF16 * 2u }; + uint32_t bd[3] = { (uint32_t)SUB_COLS_BF16, (uint32_t)hd, (uint32_t)V_SUBTILES }; + uint32_t es[3] = { 1u, 1u, 1u }; + if (cuTensorMapEncodeTiled(&tvt_, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, + const_cast<__nv_bfloat16*>(vt), 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) != CUDA_SUCCESS) + return cudaErrorInvalidValue; + } else { + if (make_tma_2d_tiled(&tvt_, const_cast<__nv_bfloat16*>(vt), (long)H * hd, (int)tq, hd, + SUB_COLS_BF16, 2, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + CU_TENSOR_MAP_SWIZZLE_128B) != cudaSuccess) + return cudaErrorInvalidValue; + } + } + + const size_t smem = + (size_t)2 * Q_TILE_BYTES + NUM_KV_STAGES * KV_RING_SLOT_BYTES + + (size_t)2 * M_TILE * HEAD_DIM * sizeof(__nv_bfloat16) + + (2 * NUM_KV_STAGES + 22) * 8 + + (size_t)CLC_STAGES * (2 * 8 + 16) + 16 + + 8 + + (size_t)2 * STAT_REGIONS * STATS * sizeof(float) + + 256; + +#ifndef VSA_NAMED_BAR +#define VSA_NAMED_BAR false +#endif +#ifndef VSA_THROTTLE +#define VSA_THROTTLE false +#endif +#ifndef VSA_USE_CLC +#define VSA_USE_CLC true +#endif + constexpr bool FULL_NAMED_BAR = VSA_NAMED_BAR, EX2_EMU = true, SPLIT_P = true, + SOFTMAX_THROTTLE = VSA_THROTTLE, USE_CLC = VSA_USE_CLC, + Q_RASTER = true, MHA = true; + auto kfn = &fmha_context_bf16_gen_kernel<32, FULL_NAMED_BAR, EX2_EMU, SPLIT_P, + SOFTMAX_THROTTLE, USE_CLC, Q_RASTER, MHA, + /*RESCALE_THRESHOLD=*/8, /*BHSD=*/VSA_BHSD>; + cudaError_t e = cudaFuncSetAttribute(kfn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem); + if (e != cudaSuccess) return e; + + const unsigned long long magic0 = make_magic((unsigned)(H * packed_mtiles_per_seq)); + const unsigned long long magic1 = make_magic((unsigned)H); + const unsigned long long magic2 = make_magic((unsigned)packed_mtiles_per_seq); + const float scale_log2 = a.sm_scale * (float)M_LOG2E; + + int numSM = 0; + e = cudaDeviceGetAttribute(&numSM, cudaDevAttrMultiProcessorCount, 0); + if (e != cudaSuccess) return e; + const int num_ctas = USE_CLC ? total_work : (total_work < numSM ? total_work : numSM); + dim3 grid(num_ctas, 1, 1), block(N_WARPS * 32, 1, 1); + + if (USE_CLC) { + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = grid; cfg.blockDim = block; cfg.dynamicSmemBytes = smem; cfg.stream = stream; + cudaLaunchAttribute cfgAttr[1]; + cfgAttr[0].id = cudaLaunchAttributeClusterDimension; + cfgAttr[0].val.clusterDim.x = 1; cfgAttr[0].val.clusterDim.y = 1; + cfgAttr[0].val.clusterDim.z = 1; + cfg.attrs = cfgAttr; cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kfn, tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B, + num_blocks, packed_mtiles_per_seq, max_kv, magic0, magic1, magic2, + a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse); + } + kfn<<>>(tq_, tk_, tvt_, tv_, to_, S, H, scale_log2, B, num_blocks, + packed_mtiles_per_seq, max_kv, magic0, magic1, magic2, + a.q2k_idx, a.q2k_num, a.variable_block_sizes, a.lse); + return cudaGetLastError(); +} + +} // namespace VSA_NAMESPACE + +// Callers (the bench, the torch binding) keep using unqualified names; each translation unit +// only ever sees the one configuration its VSA_BLK128 selected. +using namespace VSA_NAMESPACE; + +#endif // BLOCK_SPARSE_VSA_LAUNCH_SM100A_CUH diff --git a/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu b/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu new file mode 100644 index 000000000..c7bee3420 --- /dev/null +++ b/fastvideo-kernel/csrc/attention/block_sparse_sm100a.cu @@ -0,0 +1,114 @@ +// block_sparse_sm100a.cu -- torch binding for the sm_100a VSA block-sparse FMHA forward. +// +// Forward only: returns (out, lse) so FastVideo's existing Triton backward keeps working +// unchanged. lse is exactly the M tensor triton_block_sparse_attn_forward writes -- +// max(qk * qk_scale) + log2(l), [B, H, S] fp32 -- which is what lets +// block_sparse_attn_backward_triton run against our forward untouched. +// +// The build is fixed at compile time by two flags, so one extension carries one configuration: +// VSA_BLK128 false -> 64-token sparse blocks, true -> 128-token +// VSA_BHSD false -> [B, S, H, D], true -> [B, H, S, D] +#include + +#include +#include + +#include "block_sparse_launch_sm100a.cuh" + +namespace { + +void check_qkv(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"); +} + +} // namespace + +// The exported symbol carries the block size: block_sparse_sm100a_fwd is the 64-token build, +// block_sparse_sm100a_blk128_fwd the 128-token one (block_sparse_blk128_sm100a.cu re-includes +// this file with VSA_BLK128 set). The python backend picks by the metadata's block size. +#if VSA_BLK128 +#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_blk128_fwd +#else +#define BLOCK_SPARSE_SM100A_FWD block_sparse_sm100a_fwd +#endif + +// Returns {out} or {out, lse}. Layout of out matches the inputs. +std::vector BLOCK_SPARSE_SM100A_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) { + const at::cuda::OptionalCUDAGuard guard(device_of(q)); + + 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_qkv(q, "q", B, H, S, D); + check_qkv(k, "k", B, H, S, D); + check_qkv(v, "v", B, H, S, D); + check_index(q2k_idx, "q2k_idx"); + check_index(q2k_num, "q2k_num"); + check_index(variable_block_sizes, "variable_block_sizes"); + + const int64_t num_blocks = variable_block_sizes.numel(); + const int64_t max_kv = q2k_idx.size(-1); + TORCH_CHECK(S == num_blocks * BLOCK, "seqlen ", S, " must equal num_blocks (", num_blocks, + ") * ", BLOCK, "; FastVideo pads the sequence up to whole blocks"); + + auto out = torch::empty_like(q); + torch::Tensor lse; + if (need_lse) lse = torch::empty({B, H, S}, q.options().dtype(torch::kFloat32)); + + BlockSparseVsaArgs a{}; + a.q = reinterpret_cast(q.data_ptr()); + a.k = reinterpret_cast(k.data_ptr()); + a.v = reinterpret_cast(v.data_ptr()); + a.v_t = v_t.has_value() ? reinterpret_cast(v_t->data_ptr()) : nullptr; + a.o = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + a.lse = need_lse ? lse.data_ptr() : nullptr; + a.q2k_idx = q2k_idx.data_ptr(); + a.q2k_num = q2k_num.data_ptr(); + a.variable_block_sizes = variable_block_sizes.data_ptr(); + a.batch = (int)B; + a.num_heads = (int)H; + a.seqlen = (int)S; + a.head_dim = (int)D; + a.num_blocks = (int)num_blocks; + a.max_kv = (int)max_kv; + a.sm_scale = (float)sm_scale; + + // Report an unsupported regime loudly rather than returning plausible-looking wrong values. + TORCH_CHECK(block_sparse_supported(a) == cudaSuccess, + "block_sparse_sm100a: unsupported configuration -- requires head_dim==", + HEAD_DIM, ", an even num_blocks, seqlen == num_blocks*", BLOCK, + ", and a variable_block_sizes tensor. Got head_dim=", D, " num_blocks=", + num_blocks, " seqlen=", S); + + const cudaError_t err = launch_block_sparse_sm100a(a, at::cuda::getCurrentCUDAStream()); + TORCH_CHECK(err == cudaSuccess, + "block_sparse_sm100a launch failed: ", cudaGetErrorString(err)); + + if (need_lse) return {out, lse}; + return {out}; +} diff --git a/fastvideo-kernel/csrc/attention/primitives.cuh b/fastvideo-kernel/csrc/attention/primitives.cuh new file mode 100644 index 000000000..ac0cef77f --- /dev/null +++ b/fastvideo-kernel/csrc/attention/primitives.cuh @@ -0,0 +1,877 @@ +// primitives.cuh -- device primitives for the sm_100a VSA block-sparse attention +// forward: tcgen05 (alloc / mma / ld / st / commit / wait / fence), TMA load / store / +// tensormap, mbarrier, cluster launch control, setmaxnreg, fast math, and the FMHA helpers. +// +// Generated and pruned to what the kernel reaches -- do not edit by hand. +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef CUDA_CHECK +#define CUDA_CHECK(stmt) do { \ + cudaError_t _e = (stmt); \ + if (_e != cudaSuccess) { \ + fprintf(stderr, "CUDA error %s:%d: %s -> %s\n", \ + __FILE__, __LINE__, #stmt, cudaGetErrorString(_e)); \ + std::exit(1); \ + } \ + } while (0) +#endif + +__device__ __forceinline__ +uint64_t mbarrier_arrive(uint32_t mbar_smem) { + uint64_t state; + asm volatile("mbarrier.arrive.shared::cta.b64 %0, [%1];\n" + : "=l"(state) : "r"(mbar_smem) : "memory"); + return state; +} + +__device__ __forceinline__ +void mbarrier_arrive_nostate(uint32_t mbar_smem) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];\n" + :: "r"(mbar_smem) : "memory"); +} + +__device__ __forceinline__ +void mbarrier_arrive_cluster_default(uint32_t cluster_smem_addr) { + asm volatile("mbarrier.arrive.shared::cluster.b64 _, [%0];\n" + :: "r"(cluster_smem_addr) : "memory"); +} + +__device__ __forceinline__ +void mbarrier_arrive_expect_tx(uint32_t mbar_smem, uint32_t expected_bytes) { + asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n" + :: "r"(mbar_smem), "r"(expected_bytes) : "memory"); +} + +__device__ __forceinline__ +void mbarrier_wait_parity_suspend(uint32_t mbar_smem, uint32_t phase_parity) { + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, 10000000;\n" + "@!P1 bra.uni LAB_WAIT;\n" + "}\n" + :: "r"(mbar_smem), "r"(phase_parity) : "memory"); +} + +__device__ __forceinline__ +void mbarrier_wait_parity(uint32_t mbar_smem, uint32_t phase_parity) { + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT_HOT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@!P1 bra.uni LAB_WAIT_HOT;\n" + "}\n" + :: "r"(mbar_smem), "r"(phase_parity) : "memory"); +} + +__device__ __forceinline__ +void fence_proxy_async_shared_cta() { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); +} + +__device__ __forceinline__ +void fence_proxy_async_shared() { + asm volatile("fence.proxy.async.shared::cta;\n" ::: "memory"); +} + +__device__ __forceinline__ void clc_try_cancel_async( + uint32_t smem_dst, uint32_t mbar_smem) { + asm volatile( + "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.b128" + " [%0], [%1];\n" + :: "r"(smem_dst), "r"(mbar_smem) : "memory"); +} + +__device__ __forceinline__ void clc_load_response( + uint32_t smem_slot, uint32_t& r0, uint32_t& r1, + uint32_t& r2, uint32_t& r3) { + asm volatile("ld.shared::cta.v4.b32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(smem_slot)); +} + +template +struct MbarrierPhaseTracker { + + uint32_t phase[NUM_STAGES]; + int idx; + + __device__ __forceinline__ + void init() { + for (int i = 0; i < NUM_STAGES; ++i) phase[i] = 0; + idx = 0; + } + + __device__ __forceinline__ + uint32_t current_phase() const { return phase[idx]; } + + __device__ __forceinline__ + void advance() { + phase[idx] ^= 1u; + idx = (idx + 1) % NUM_STAGES; + } + + __device__ __forceinline__ + int stage() const { return idx; } +}; + +template +struct PhaseTracker { + int stage; + uint32_t phase; + + __device__ __forceinline__ + PhaseTracker() : stage(0), phase(0) {} + + __device__ __forceinline__ + void advance() { + stage++; + if (stage == NUM_STAGES) { + stage = 0; + phase ^= 1; + } + } + + __device__ __forceinline__ + int get_stage() const { return stage; } + + __device__ __forceinline__ + uint32_t get_phase() const { return phase; } +}; + +template +struct EmptyPhaseTracker { + int stage; + uint32_t phase; + + __device__ __forceinline__ + EmptyPhaseTracker() : stage(0), phase(1) {} + + __device__ __forceinline__ + void advance() { + stage++; + if (stage == NUM_STAGES) { + stage = 0; + phase ^= 1; + } + } + + __device__ __forceinline__ + int get_stage() const { return stage; } + + __device__ __forceinline__ + uint32_t get_phase() const { return phase; } +}; + +template +__device__ __forceinline__ +void advance_stage_phase(int& stage, uint32_t& phase) { + ++stage; + if (stage == STAGES) { + stage = 0; + phase ^= 1u; + } +} + +static constexpr uint32_t SM100_CLC_PEER_MASK = 0xFEFFFFFF; + +struct ClcTileInfo { + int m_tile; + int n_tile; + bool valid; +}; + +enum class ClcRasterOrder { AlongN, AlongM }; + +__device__ __forceinline__ +void clc_arrive_expect_tx_cta(uint32_t clc_full_local_addr, uint32_t tx_bytes) { + if ((threadIdx.x & 31) == 0) { + mbarrier_arrive_expect_tx(clc_full_local_addr, tx_bytes); + } +} + +__device__ __forceinline__ +void clc_consumer_release(uint32_t clc_empty_local_addr) { + uint32_t peer0_addr = clc_empty_local_addr & SM100_CLC_PEER_MASK; + mbarrier_arrive_cluster_default(peer0_addr); +} + +__device__ __forceinline__ +void clc_consumer_release_cta(uint32_t clc_empty_local_addr) { + mbarrier_arrive_nostate(clc_empty_local_addr); +} + +template +__device__ __forceinline__ +ClcTileInfo clc_parse_response(uint32_t resp_smem_addr) { + uint32_t d0, d1, d2, d3; + fence_proxy_async_shared_cta(); + clc_load_response(resp_smem_addr, d0, d1, d2, d3); + const int ctaid_x = static_cast(d0); + const int ctaid_y = static_cast(d1 & 0xFFFFu); + const bool valid = (d2 & 1u) != 0u; + (void)d3; + + ClcTileInfo info; + info.valid = valid; + if constexpr (ORDER == ClcRasterOrder::AlongN) { + info.m_tile = ctaid_y / CLUSTER_SHAPE_M; + info.n_tile = ctaid_x / CLUSTER_SHAPE_N; + } else { + info.m_tile = ctaid_x / CLUSTER_SHAPE_M; + info.n_tile = ctaid_y / CLUSTER_SHAPE_N; + } + return info; +} + +template +__device__ __forceinline__ +ClcTileInfo clc_fetch_next_tile( + uint64_t* clc_full_bar, uint64_t* clc_empty_bar, uint32_t* clc_response, + int clc_cons_stage, uint32_t clc_cons_phase, bool do_release) { + uint32_t full_addr = static_cast( + __cvta_generic_to_shared(&clc_full_bar[clc_cons_stage])); + if constexpr (SUSPEND) mbarrier_wait_parity_suspend(full_addr, clc_cons_phase); + else mbarrier_wait_parity(full_addr, clc_cons_phase); + uint32_t resp_addr = static_cast( + __cvta_generic_to_shared(&clc_response[clc_cons_stage * 4])); + ClcTileInfo t = clc_parse_response< + CLUSTER_SHAPE_M, CLUSTER_SHAPE_N, ORDER>(resp_addr); + if (do_release) { + uint32_t empty_local = static_cast( + __cvta_generic_to_shared(&clc_empty_bar[clc_cons_stage])); + if constexpr (CTA_GROUP == 1) { + clc_consumer_release_cta(empty_local); + } else { + clc_consumer_release(empty_local); + } + } + return t; +} + +template +__device__ __forceinline__ +void clc_fetch_next_tile_advance(int& clc_cons_stage, + uint32_t& clc_cons_phase) { + advance_stage_phase(clc_cons_stage, clc_cons_phase); +} + +__device__ __forceinline__ unsigned fdiv(unsigned n, unsigned long long pk) { + unsigned M = (unsigned)pk; + if (M == 0u) return n; + return __umulhi(n, M) >> (unsigned)(pk >> 32); +} +__host__ inline unsigned long long make_magic(unsigned d) { + if (d <= 1u) return 0ULL; + unsigned l = 0; while ((1u << (l + 1)) <= d) ++l; + unsigned p = 31u + l; + unsigned long long m = ((1ull << p) + (unsigned long long)d - 1ull) / d; + return (m & 0xffffffffULL) | ((unsigned long long)(p - 32u) << 32); +} + +template +__device__ __forceinline__ +void bar_sync(uint32_t thread_count) { + static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15, + "bar.sync: BARRIER_ID must be in [0, 15]"); + asm volatile("bar.sync %0, %1;\n" + :: "n"(BARRIER_ID), "r"(thread_count) : "memory"); +} + +template +__device__ __forceinline__ +void bar_arrive(uint32_t thread_count) { + static_assert(BARRIER_ID >= 0 && BARRIER_ID <= 15, + "bar.arrive: BARRIER_ID must be in [0, 15]"); + asm volatile("bar.arrive %0, %1;\n" + :: "n"(BARRIER_ID), "r"(thread_count) : "memory"); +} + +__device__ __forceinline__ void full_bar_arrive(int m_tile, int band) { + switch (1 + m_tile * 4 + band) { + case 1: bar_arrive<1>(64); break; + case 2: bar_arrive<2>(64); break; + case 3: bar_arrive<3>(64); break; + case 4: bar_arrive<4>(64); break; + case 5: bar_arrive<5>(64); break; + case 6: bar_arrive<6>(64); break; + case 7: bar_arrive<7>(64); break; + case 8: bar_arrive<8>(64); break; + } +} +__device__ __forceinline__ void full_bar_wait(int m_tile, int band) { + switch (1 + m_tile * 4 + band) { + case 1: bar_sync<1>(64); break; + case 2: bar_sync<2>(64); break; + case 3: bar_sync<3>(64); break; + case 4: bar_sync<4>(64); break; + case 5: bar_sync<5>(64); break; + case 6: bar_sync<6>(64); break; + case 7: bar_sync<7>(64); break; + case 8: bar_sync<8>(64); break; + } +} + +template +__device__ __forceinline__ void mask_s_row_r2p(float* scores, int k_offset, int q_pos, int seqlen_k) { + int n_keep = seqlen_k - k_offset; + if constexpr (IS_CAUSAL) { + const int causal = q_pos - k_offset + 1; + n_keep = n_keep < causal ? n_keep : causal; + } + #pragma unroll + for (int s = 0; s < K_TILE / 32; ++s) { + int m = (s + 1) * 32 - n_keep; + m = m < 0 ? 0 : (m > 32 ? 32 : m); + const uint32_t keep = (m >= 32) ? 0u : (0xFFFFFFFFu >> m); + #pragma unroll + for (int i = 0; i < 32; ++i) + if (!(keep & (1u << i))) scores[s * 32 + i] = -INFINITY; + } +} + +template +__device__ __forceinline__ void tcgen05_alloc(uint32_t smem_dst_ptr, + uint32_t n_cols) { + static_assert(CTA_GROUP == 1 || CTA_GROUP == 2, + "tcgen05_alloc: CTA_GROUP must be 1 or 2"); + if constexpr (CTA_GROUP == 1) { + asm volatile( + "tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "r"(smem_dst_ptr), "r"(n_cols)); + } else { + asm volatile( + "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;\n" + :: "r"(smem_dst_ptr), "r"(n_cols)); + } +} + +__device__ __forceinline__ void tcgen05_st_32x32b_x16( + uint32_t tmem_addr, const uint32_t (&r)[16]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x16.b32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10," + "%11,%12,%13,%14,%15,%16};\n" + :: "r"(tmem_addr), + "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]), + "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]), + "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]), + "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15])); +} + +__device__ __forceinline__ void tcgen05_st_32x32b_x32( + uint32_t tmem_addr, const uint32_t (&r)[32]) { + asm volatile("tcgen05.st.sync.aligned.32x32b.x32.b32 " + "[%0], {%1,%2,%3,%4,%5,%6,%7,%8,%9,%10," + "%11,%12,%13,%14,%15,%16,%17,%18,%19,%20," + "%21,%22,%23,%24,%25,%26,%27,%28,%29,%30," + "%31,%32};\n" + :: "r"(tmem_addr), + "r"(r[0]),"r"(r[1]),"r"(r[2]),"r"(r[3]), + "r"(r[4]),"r"(r[5]),"r"(r[6]),"r"(r[7]), + "r"(r[8]),"r"(r[9]),"r"(r[10]),"r"(r[11]), + "r"(r[12]),"r"(r[13]),"r"(r[14]),"r"(r[15]), + "r"(r[16]),"r"(r[17]),"r"(r[18]),"r"(r[19]), + "r"(r[20]),"r"(r[21]),"r"(r[22]),"r"(r[23]), + "r"(r[24]),"r"(r[25]),"r"(r[26]),"r"(r[27]), + "r"(r[28]),"r"(r[29]),"r"(r[30]),"r"(r[31])); +} + +__device__ __forceinline__ void tcgen05_commit1_lead(uint32_t lead, uint32_t mbar_smem_addr) { + asm volatile( + "{\n\t" + ".reg .pred q;\n\t" + "setp.ne.b32 q, %0, 0;\n\t" + "@q tcgen05.commit.cta_group::1.mbarrier::arrive::one.b64 [%1];\n\t" + "}\n" + :: "r"(lead), "r"(mbar_smem_addr)); +} + +__device__ __forceinline__ void tcgen05_wait_st() { + asm volatile("tcgen05.wait::st.sync.aligned;\n" ::: "memory"); +} + +__device__ __forceinline__ void tcgen05_fence_before_thread_sync() { + asm volatile("tcgen05.fence::before_thread_sync;\n" ::: "memory"); +} + +__device__ __forceinline__ +void tma_load_2d(uint32_t smem_dst, const void* tensormap_ptr, + uint32_t mbar_smem, int coord_x, int coord_y) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4}], [%2];\n" + :: "r"(smem_dst), "l"(tensormap_ptr), + "r"(mbar_smem), "r"(coord_x), "r"(coord_y) + : "memory"); +} + +__device__ __forceinline__ +void tma_load_3d(uint32_t smem_dst, const void* tensormap_ptr, + uint32_t mbar_smem, int c0, int c1, int c2) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];\n" + :: "r"(smem_dst), "l"(tensormap_ptr), + "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2) + : "memory"); +} + +__device__ __forceinline__ +void tma_load_4d(uint32_t smem_dst, const void* tensormap_ptr, + uint32_t mbar_smem, int c0, int c1, int c2, int c3) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];\n" + :: "r"(smem_dst), "l"(tensormap_ptr), + "r"(mbar_smem), "r"(c0), "r"(c1), "r"(c2), "r"(c3) + : "memory"); +} + +template +__device__ __forceinline__ void tcgen05_dealloc(uint32_t tmem_addr, + uint32_t n_cols) { + static_assert(CTA_GROUP == 1 || CTA_GROUP == 2, + "tcgen05_dealloc: CTA_GROUP must be 1 or 2"); + if constexpr (CTA_GROUP == 1) { + asm volatile( + "tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;\n" + :: "r"(tmem_addr), "r"(n_cols)); + } else { + asm volatile( + "tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;\n" + :: "r"(tmem_addr), "r"(n_cols)); + } +} + +__device__ __forceinline__ +void tma_store_2d(const void* tensormap_ptr, int coord_x, int coord_y, + uint32_t smem_src) { + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2}], [%3];\n" + :: "l"(tensormap_ptr), "r"(coord_x), "r"(coord_y), + "r"(smem_src) + : "memory"); +} + +__device__ __forceinline__ +void tma_store_3d(const void* tensormap_ptr, int c0, int c1, int c2, + uint32_t smem_src) { + asm volatile( + "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3}], [%4];\n" + :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(smem_src) + : "memory"); +} + +__device__ __forceinline__ +void tma_store_4d(const void* tensormap_ptr, int c0, int c1, int c2, int c3, + uint32_t smem_src) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];\n" + :: "l"(tensormap_ptr), "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(smem_src) + : "memory"); +} + +inline cudaError_t make_tma_2d_tiled( + CUtensorMap* out, + const void* ptr, int rows, int cols, int box_rows, int box_cols, + int elem_bytes, CUtensorMapDataType dtype, + CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B, + CUtensorMapL2promotion l2 = CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + CUtensorMapFloatOOBfill oob = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE) { + uint64_t globalDim[2] = { (uint64_t)cols, (uint64_t)rows }; + uint64_t globalStrides[1] = { (uint64_t)cols * (uint64_t)elem_bytes }; + uint32_t boxDim[2] = { (uint32_t)box_cols, (uint32_t)box_rows }; + uint32_t elemStrides[2] = { 1u, 1u }; + + CUresult r = cuTensorMapEncodeTiled( + out, dtype, 2, + const_cast(ptr), globalDim, globalStrides, + boxDim, elemStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, l2, oob); + return (r == CUDA_SUCCESS) ? cudaSuccess : cudaErrorInvalidValue; +} + +__device__ __forceinline__ +void cp_async_bulk_commit_group() { + asm volatile("cp.async.bulk.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ +void cp_async_bulk_wait_group_read() { + asm volatile("cp.async.bulk.wait_group.read %0;\n" :: "n"(N) : "memory"); +} + +__device__ __forceinline__ +void mbarrier_init(uint32_t mbar_smem, uint32_t arrive_count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(mbar_smem), "r"(arrive_count) : "memory"); +} + +template +__device__ __forceinline__ void tcgen05_relinquish_alloc_permit() { + static_assert(CTA_GROUP == 1 || CTA_GROUP == 2, + "tcgen05_relinquish_alloc_permit: CTA_GROUP must be 1 or 2"); + if constexpr (CTA_GROUP == 1) { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;\n" ::); + } else { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;\n" ::); + } +} + +__device__ __forceinline__ +void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;\n" ::: "memory"); +} + +__device__ __forceinline__ void tcgen05_mma_f16_ss_lead(uint32_t lead, + uint32_t tmem_c, 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.cta_group::1.kind::f16 [%1], %2, %3, %4, {%6, %7, %8, %9}, p;\n\t" + "}\n" + :: "r"(lead), "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u)); +} + +__device__ __forceinline__ void tcgen05_mma_f16_ts_1sm_lead(uint32_t lead, + uint32_t tmem_c, 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.cta_group::1.kind::f16 [%1], [%2], %3, %4, {%6, %7, %8, %9}, p;\n\t" + "}\n" + :: "r"(lead), "r"(tmem_c), "r"(tmem_a), "l"(desc_b), "r"(idesc), + "r"(enable_input_d ? 1u : 0u), "r"(0u), "r"(0u), "r"(0u), "r"(0u)); +} + +enum class SmemSwizzleBlackwell : uint32_t { + None = 0, + B128_32atom = 1, + B128 = 2, + B64 = 4, + B32 = 6, +}; + +__device__ __host__ __forceinline__ uint64_t build_smem_desc_blackwell( + uint32_t smem_addr, + uint32_t stride_byte_offset, + uint32_t leading_byte_offset, + SmemSwizzleBlackwell swizzle = SmemSwizzleBlackwell::B128, + uint32_t base_offset = 0) { + uint64_t d = 0; + d |= static_cast((smem_addr >> 4) & 0x3FFF); + d |= static_cast((leading_byte_offset >> 4) & 0x3FFF) << 16; + d |= static_cast((stride_byte_offset >> 4) & 0x3FFF) << 32; + d |= static_cast(1) << 46; + d |= (static_cast(base_offset) & 0x7) << 49; + d |= static_cast(static_cast(swizzle) & 0x7) << 61; + return d; +} + +__device__ __forceinline__ +uint32_t elect_one_sync() { + uint32_t elected; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "elect.sync %0|p, 0xffffffff;\n\t" + "selp.b32 %0, 1, 0, p;\n\t" + "}\n" + : "=r"(elected)); + return elected; +} + +__device__ __forceinline__ +uint32_t elect_one_sync(uint32_t membermask) { + uint32_t elected; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "elect.sync %0|p, %1;\n\t" + "selp.b32 %0, 1, 0, p;\n\t" + "}\n" + : "=r"(elected) : "r"(membermask)); + return elected; +} + +template +__device__ __forceinline__ +void setmaxnreg_dec() { + static_assert(N >= 24 && N <= 256 && N % 8 == 0, + "setmaxnreg_dec: N must be in [24, 256] and a multiple of 8"); + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" + :: "n"(N) : "memory"); +} + +template +__device__ __forceinline__ +void setmaxnreg_inc() { + static_assert(N >= 24 && N <= 256 && N % 8 == 0, + "setmaxnreg_inc: N must be in [24, 256] and a multiple of 8"); + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" + :: "n"(N) : "memory"); +} + +__device__ __forceinline__ +uint32_t cvt_f32x2_to_bf16x2(float a, float b) { + uint32_t r; + asm volatile("cvt.rn.bf16x2.f32 %0, %2, %1;\n" + : "=r"(r) : "f"(a), "f"(b)); + return r; +} + +namespace { +__device__ __forceinline__ uint64_t f32x2_bits(float2 v) { + uint64_t b; __builtin_memcpy(&b, &v, 8); return b; +} +__device__ __forceinline__ float2 f32x2_make(uint64_t b) { + float2 v; __builtin_memcpy(&v, &b, 8); return v; +} +} + +__device__ __forceinline__ float2 fmul2(float2 a, float2 b) { + uint64_t d; + asm volatile("mul.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b))); + return f32x2_make(d); +} + +__device__ __forceinline__ float2 fadd2(float2 a, float2 b) { + uint64_t d; + asm volatile("add.f32x2 %0, %1, %2;\n" : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b))); + return f32x2_make(d); +} + +__device__ __forceinline__ float2 ffma2(float2 a, float2 b, float2 c) { + uint64_t d; + asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" + : "=l"(d) : "l"(f32x2_bits(a)), "l"(f32x2_bits(b)), "l"(f32x2_bits(c))); + return f32x2_make(d); +} + +__device__ __forceinline__ float2 f32x2_splat(float s) { return make_float2(s, s); } + +__device__ __forceinline__ float ex2_approx_f32(float z) { + float d; + asm volatile("ex2.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(z)); + return d; +} + +__device__ __forceinline__ float2 ex2_emu_f32x2(float x, float y) { + uint32_t ox, oy; + asm volatile( + "{\n\t" + ".reg .f32 f1,f2,f3,f4,f5,f6,f7;\n\t" + ".reg .b64 l1,l2,l3,l4,l5,l6,l7,l8,l9,l10;\n\t" + ".reg .s32 r1,r2,r3,r4,r5,r6,r7,r8;\n\t" + "max.f32 f1, %2, 0fC2FE0000;\n\t" + "max.f32 f2, %3, 0fC2FE0000;\n\t" + "mov.b64 l1, {f1, f2};\n\t" + "mov.f32 f3, 0f4B400000;\n\t" + "mov.b64 l2, {f3, f3};\n\t" + "add.rm.f32x2 l7, l1, l2;\n\t" + "sub.rn.f32x2 l8, l7, l2;\n\t" + "sub.rn.f32x2 l9, l1, l8;\n\t" + "mov.f32 f7, 0f3D9DF09D;\n\t" + "mov.b64 l6, {f7, f7};\n\t" + "mov.f32 f6, 0f3E6906A4;\n\t" + "mov.b64 l5, {f6, f6};\n\t" + "mov.f32 f5, 0f3F31F519;\n\t" + "mov.b64 l4, {f5, f5};\n\t" + "mov.f32 f4, 0f3F800000;\n\t" + "mov.b64 l3, {f4, f4};\n\t" + "fma.rn.f32x2 l10, l9, l6, l5;\n\t" + "fma.rn.f32x2 l10, l10, l9, l4;\n\t" + "fma.rn.f32x2 l10, l10, l9, l3;\n\t" + "mov.b64 {r1, r2}, l7;\n\t" + "mov.b64 {r3, r4}, l10;\n\t" + "shl.b32 r5, r1, 23;\n\t" + "add.s32 r7, r5, r3;\n\t" + "shl.b32 r6, r2, 23;\n\t" + "add.s32 r8, r6, r4;\n\t" + "mov.b32 %0, r7;\n\t" + "mov.b32 %1, r8;\n\t" + "}\n" + : "=r"(ox), "=r"(oy) : "f"(x), "f"(y)); + float2 r; __builtin_memcpy(&r.x, &ox, 4); __builtin_memcpy(&r.y, &oy, 4); + return r; +} + +__device__ __forceinline__ float rcp_approx_ftz_f32(float x) { + float d; + asm("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(d) : "f"(x)); + return d; +} + +__device__ __forceinline__ uint32_t make_idesc_table44( + int M, int N, + uint32_t dtype, uint32_t atype, uint32_t btype, + bool transpose_a = false, bool transpose_b = false, + bool negate_a = false, bool negate_b = false) { + uint32_t idesc = 0; + idesc |= (dtype & 0x3) << 4; + idesc |= (atype & 0x7) << 7; + idesc |= (btype & 0x7) << 10; + idesc |= (negate_a ? 1u : 0u) << 13; + idesc |= (negate_b ? 1u : 0u) << 14; + idesc |= (transpose_a ? 1u : 0u) << 15; + idesc |= (transpose_b ? 1u : 0u) << 16; + idesc |= ((static_cast(N) >> 3) & 0x3F) << 17; + idesc |= ((static_cast(M) >> 4) & 0x1F) << 24; + return idesc; +} + +__device__ __forceinline__ uint32_t make_idesc_bf16_f32( + int M, int N, bool ta = false, bool tb = false) { + return make_idesc_table44(M, N, 1, + 1, 1, ta, tb); +} + +__device__ __forceinline__ void tcgen05_ld_32x32b_x16( + uint32_t tmem_addr, uint32_t (&r)[16]) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x16.b32 " + "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9," + "%10,%11,%12,%13,%14,%15}, [%16];\n" + : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]), + "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]), + "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]), + "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]) + : "r"(tmem_addr)); +} + +__device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t tmem_addr, uint32_t (&r)[32]) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9," + "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29," + "%30,%31}, [%32];\n" + : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]), + "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]), + "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]), + "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]), + "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]), + "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]), + "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]), + "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31]) + : "r"(tmem_addr)); +} + +__device__ __forceinline__ void tcgen05_ld_32x32b_x64( + uint32_t tmem_addr, uint32_t (&r)[64]) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x64.b32 " + "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9," + "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29," + "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39," + "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49," + "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59," + "%60,%61,%62,%63}, [%64];\n" + : "=r"(r[0]),"=r"(r[1]),"=r"(r[2]),"=r"(r[3]), + "=r"(r[4]),"=r"(r[5]),"=r"(r[6]),"=r"(r[7]), + "=r"(r[8]),"=r"(r[9]),"=r"(r[10]),"=r"(r[11]), + "=r"(r[12]),"=r"(r[13]),"=r"(r[14]),"=r"(r[15]), + "=r"(r[16]),"=r"(r[17]),"=r"(r[18]),"=r"(r[19]), + "=r"(r[20]),"=r"(r[21]),"=r"(r[22]),"=r"(r[23]), + "=r"(r[24]),"=r"(r[25]),"=r"(r[26]),"=r"(r[27]), + "=r"(r[28]),"=r"(r[29]),"=r"(r[30]),"=r"(r[31]), + "=r"(r[32]),"=r"(r[33]),"=r"(r[34]),"=r"(r[35]), + "=r"(r[36]),"=r"(r[37]),"=r"(r[38]),"=r"(r[39]), + "=r"(r[40]),"=r"(r[41]),"=r"(r[42]),"=r"(r[43]), + "=r"(r[44]),"=r"(r[45]),"=r"(r[46]),"=r"(r[47]), + "=r"(r[48]),"=r"(r[49]),"=r"(r[50]),"=r"(r[51]), + "=r"(r[52]),"=r"(r[53]),"=r"(r[54]),"=r"(r[55]), + "=r"(r[56]),"=r"(r[57]),"=r"(r[58]),"=r"(r[59]), + "=r"(r[60]),"=r"(r[61]),"=r"(r[62]),"=r"(r[63]) + : "r"(tmem_addr)); +} + +__device__ __forceinline__ void tcgen05_ld_32x32b_x128( + uint32_t tmem_addr, uint32_t (&r)[128]) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x128.b32 " + "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9," + "%10,%11,%12,%13,%14,%15,%16,%17,%18,%19," + "%20,%21,%22,%23,%24,%25,%26,%27,%28,%29," + "%30,%31,%32,%33,%34,%35,%36,%37,%38,%39," + "%40,%41,%42,%43,%44,%45,%46,%47,%48,%49," + "%50,%51,%52,%53,%54,%55,%56,%57,%58,%59," + "%60,%61,%62,%63,%64,%65,%66,%67,%68,%69," + "%70,%71,%72,%73,%74,%75,%76,%77,%78,%79," + "%80,%81,%82,%83,%84,%85,%86,%87,%88,%89," + "%90,%91,%92,%93,%94,%95,%96,%97,%98,%99," + "%100,%101,%102,%103,%104,%105,%106,%107,%108,%109," + "%110,%111,%112,%113,%114,%115,%116,%117,%118,%119," + "%120,%121,%122,%123,%124,%125,%126,%127}, [%128];\n" + : "=r"(r[ 0]),"=r"(r[ 1]),"=r"(r[ 2]),"=r"(r[ 3]), + "=r"(r[ 4]),"=r"(r[ 5]),"=r"(r[ 6]),"=r"(r[ 7]), + "=r"(r[ 8]),"=r"(r[ 9]),"=r"(r[ 10]),"=r"(r[ 11]), + "=r"(r[ 12]),"=r"(r[ 13]),"=r"(r[ 14]),"=r"(r[ 15]), + "=r"(r[ 16]),"=r"(r[ 17]),"=r"(r[ 18]),"=r"(r[ 19]), + "=r"(r[ 20]),"=r"(r[ 21]),"=r"(r[ 22]),"=r"(r[ 23]), + "=r"(r[ 24]),"=r"(r[ 25]),"=r"(r[ 26]),"=r"(r[ 27]), + "=r"(r[ 28]),"=r"(r[ 29]),"=r"(r[ 30]),"=r"(r[ 31]), + "=r"(r[ 32]),"=r"(r[ 33]),"=r"(r[ 34]),"=r"(r[ 35]), + "=r"(r[ 36]),"=r"(r[ 37]),"=r"(r[ 38]),"=r"(r[ 39]), + "=r"(r[ 40]),"=r"(r[ 41]),"=r"(r[ 42]),"=r"(r[ 43]), + "=r"(r[ 44]),"=r"(r[ 45]),"=r"(r[ 46]),"=r"(r[ 47]), + "=r"(r[ 48]),"=r"(r[ 49]),"=r"(r[ 50]),"=r"(r[ 51]), + "=r"(r[ 52]),"=r"(r[ 53]),"=r"(r[ 54]),"=r"(r[ 55]), + "=r"(r[ 56]),"=r"(r[ 57]),"=r"(r[ 58]),"=r"(r[ 59]), + "=r"(r[ 60]),"=r"(r[ 61]),"=r"(r[ 62]),"=r"(r[ 63]), + "=r"(r[ 64]),"=r"(r[ 65]),"=r"(r[ 66]),"=r"(r[ 67]), + "=r"(r[ 68]),"=r"(r[ 69]),"=r"(r[ 70]),"=r"(r[ 71]), + "=r"(r[ 72]),"=r"(r[ 73]),"=r"(r[ 74]),"=r"(r[ 75]), + "=r"(r[ 76]),"=r"(r[ 77]),"=r"(r[ 78]),"=r"(r[ 79]), + "=r"(r[ 80]),"=r"(r[ 81]),"=r"(r[ 82]),"=r"(r[ 83]), + "=r"(r[ 84]),"=r"(r[ 85]),"=r"(r[ 86]),"=r"(r[ 87]), + "=r"(r[ 88]),"=r"(r[ 89]),"=r"(r[ 90]),"=r"(r[ 91]), + "=r"(r[ 92]),"=r"(r[ 93]),"=r"(r[ 94]),"=r"(r[ 95]), + "=r"(r[ 96]),"=r"(r[ 97]),"=r"(r[ 98]),"=r"(r[ 99]), + "=r"(r[100]),"=r"(r[101]),"=r"(r[102]),"=r"(r[103]), + "=r"(r[104]),"=r"(r[105]),"=r"(r[106]),"=r"(r[107]), + "=r"(r[108]),"=r"(r[109]),"=r"(r[110]),"=r"(r[111]), + "=r"(r[112]),"=r"(r[113]),"=r"(r[114]),"=r"(r[115]), + "=r"(r[116]),"=r"(r[117]),"=r"(r[118]),"=r"(r[119]), + "=r"(r[120]),"=r"(r[121]),"=r"(r[122]),"=r"(r[123]), + "=r"(r[124]),"=r"(r[125]),"=r"(r[126]),"=r"(r[127]) + : "r"(tmem_addr)); +} + +__device__ __forceinline__ +uint32_t smem_ptr_u32(const void* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +__device__ __forceinline__ +void sts_f32(uint32_t smem_addr, float val) { + asm volatile("st.shared.f32 [%0], %1;" :: "r"(smem_addr), "f"(val) : "memory"); +} + diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index ad27bb26b..2c97e9dad 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -28,10 +28,31 @@ void register_rms_norm(pybind11::module_ &); void register_layer_norm(pybind11::module_ &); void register_gemm(pybind11::module_ &); +#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A +extern std::vector block_sparse_sm100a_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_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); +#endif + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "FastVideo CUDA Kernels"; +#ifdef TK_COMPILE_BLOCK_SPARSE_VSA_SM100A + m.def("block_sparse_sm100a_fwd", + torch::wrap_pybind_function(block_sparse_sm100a_fwd), + "VSA block-sparse attention forward, 64-token blocks (Blackwell sm100a)"); + 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)"); +#endif + #ifdef TK_COMPILE_ST_ATTN + m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)"); #endif diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py new file mode 100644 index 000000000..1ea406efc --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_sm100a.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +"""sm_100a (Blackwell) CUDA block-sparse VSA forward. + +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. + +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 +tensors and picks the op. Anything else 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. + # `import fastvideo_kernel._C as _C` resolves to the namespace package, whose __init__ is + # empty, so hasattr() fails on a wheel install and the caller silently falls back with the + # kernel built and present. + from fastvideo_kernel._C import fastvideo_kernel_ops as _C + _FWD_BY_BLOCK = { + 64: getattr(_C, "block_sparse_sm100a_fwd", None), + 128: getattr(_C, "block_sparse_sm100a_blk128_fwd", None), + } + _HAS_VSA_SM100A = any(_FWD_BY_BLOCK.values()) +except ImportError: # pragma: no cover - extension not built + _C = None + _FWD_BY_BLOCK = {} + _HAS_VSA_SM100A = False + +_SM100 = (10, 0) +HEAD_DIM = 128 +# Must match the -DVSA_BHSD the extension was compiled with (see CMakeLists). +BHSD = True + + +def _block_size(q: torch.Tensor, variable_block_sizes: torch.Tensor) -> int: + num_blocks = variable_block_sizes.numel() + seqlen = q.shape[2] if BHSD else q.shape[1] + return 0 if num_blocks == 0 or seqlen % num_blocks else seqlen // num_blocks + + +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. Deliberately NO reads of tensor + contents: the previous ``int(variable_block_sizes.min())`` was a GPU->CPU sync on every + call, and the kernel no longer needs it (see below). This predicate must stay cheap + enough to sit on a per-layer dispatch path. + + What the kernel accepts (and is tested to handle): + * q/k/v: contiguous 4-D bf16 CUDA tensors on an sm_100 device, head_dim 128, laid out + as compiled (BHSD here); seqlen == num_blocks * block with an EVEN num_blocks (a CTA + owns an adjacent pair of query blocks) and a 64- or 128-token build present. + * q2k_num: any per-row counts in [0, max_kv], NON-uniform across rows included. Rows + with count 0 produce exactly-zero output rows (and a finite LSE sentinel) rather + than attending anywhere -- so no ``.min()`` floor is required of the caller. + * q2k_idx: rows only need valid entries (in [0, num_blocks)) BELOW that row's count; + padding past the count (e.g. map_to_index's -1 fill) is never dereferenced. max_kv + (= q2k_idx.shape[-1]) must be >= 1, which the host launcher re-checks. + * variable_block_sizes: per-KV-block valid-token counts in [0, block]; keys at or past + a block's count are masked. Integer metadata is converted to int32/contiguous by + ``block_sparse_attn_sm100a`` itself, so int64 inputs merely cost a cast. + """ + if not _HAS_VSA_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 + if _FWD_BY_BLOCK.get(_block_size(q, variable_block_sizes)) is None: + return False + # A CTA owns an adjacent pair of query blocks. + if variable_block_sizes.numel() % 2 != 0: + 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 + return True + + +def block_sparse_attn_sm100a( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, + need_lse: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass. Returns ``(out, lse)``; ``out`` has q's layout.""" + fwd = _FWD_BY_BLOCK[_block_size(q, variable_block_sizes)] + idx = q2k_idx.to(torch.int32).contiguous() + num = q2k_num.to(torch.int32).contiguous() + vbs = variable_block_sizes.to(torch.int32).contiguous() + sm_scale = 1.0 / (q.shape[-1]**0.5) + res = fwd(q.contiguous(), k.contiguous(), v.contiguous(), None, + idx, num, vbs, sm_scale, need_lse) + return (res[0], res[1]) if need_lse else (res[0], None) diff --git a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py index d0d647cfe..1c6f3e8a6 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vsa_utils.py @@ -14,7 +14,9 @@ import math import torch VSA_TILE_SIZE = (4, 4, 4) -_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 256) +# 128 is served by the sm_100a CUDA backend (block_sparse_attn_sm100a); 64 and 256 by +# Triton and the CuTe-DSL path. A volume here only needs a backend that accepts it. +_SUPPORTED_VSA_BLOCK_VOLUMES = (64, 128, 256) def _canonicalize_device(device: torch.device | str) -> torch.device: diff --git a/tests/bench_block_sparse_sm100a.py b/tests/bench_block_sparse_sm100a.py new file mode 100644 index 000000000..6ced276dc --- /dev/null +++ b/tests/bench_block_sparse_sm100a.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +"""In-framework benchmark of the sm_100a VSA backend at real Wan shapes. + +Not the powers-of-two bench grid: this builds the metadata FastVideo actually builds, at the +deployed latent shapes, and drives the same block_sparse_attn_from_indices entry point the +model calls -- so it measures the path that will run, including the dispatch and the index +tensors as constructed rather than synthesised. + + PYTHONPATH=fastvideo-kernel/python python tests/bench_block_sparse_sm100a.py +""" + +import time + +import torch + +from fastvideo_kernel import block_sparse_attn_sm100a as vsa +from fastvideo_kernel.block_sparse_attn import block_sparse_attn_triton +from fastvideo_kernel.triton_kernels.index import map_to_index +from fastvideo_kernel.vsa_utils import build_vsa_metadata + +HEAD_DIM = 128 + +# (label, latent (T,H,W), tile, heads). The latents are Wan's; heads is per-rank. +CASES = [ + ("480P", (21, 30, 52), (4, 4, 4), 40), + ("720P", (21, 45, 80), (4, 4, 4), 40), +] + + +def timed(fn, iters=20, warmup=5): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters * 1e3 + + +def main(): + if torch.cuda.get_device_capability() != (10, 0): + print("not Blackwell; skipping") + return + print(f"{'case':<8} {'S_pad':>7} {'blocks':>7} {'topk':>5} {'sm100a ms':>10} " + f"{'triton ms':>10} {'speedup':>8} selected") + for label, latent, tile, heads in CASES: + meta = build_vsa_metadata(latent, tile_size=tile, device="cuda") + vbs = meta["variable_block_sizes"].to(torch.int32) + nb = vbs.numel() + block = int(meta["max_block_size"]) + S = nb * block + topk = max(1, int(0.1 * nb)) # sparsity 0.9, as deployed + + q, k, v = (torch.randn(1, heads, S, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + for _ in range(3)) + # index tensors exactly as FastVideo builds them, from a top-k bool map + scores = torch.randn(1, heads, nb, nb, device="cuda") + keep = torch.zeros_like(scores, dtype=torch.bool) + keep.scatter_(-1, scores.topk(topk, dim=-1).indices, True) + idx, num = map_to_index(keep) + idx, num = idx.to(torch.int32).contiguous(), num.to(torch.int32).contiguous() + + ok = vsa.is_supported(q, vbs) + ours = timed(lambda: vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs)) if ok else float("nan") + + # Triton needs 64-token blocks: expand each 128-block into its two halves. + if block == 128: # Triton is 64-granular; expand only when we run 128 + keep64 = keep.repeat_interleave(2, dim=-1).repeat_interleave(2, dim=-2) + vbs64 = torch.stack([vbs.clamp(max=64), (vbs - 64).clamp(min=0)], dim=-1).flatten() + i64, n64 = map_to_index(keep64) + i64, n64 = i64.to(torch.int32).contiguous(), n64.to(torch.int32).contiguous() + vbs64 = vbs64.to(torch.int32).contiguous() + else: + keep64, i64, n64, vbs64 = keep, idx, num, vbs + tri = timed(lambda: block_sparse_attn_triton(q, k, v, i64, n64, vbs64)) + + print(f"{label:<8} {S:>7} {nb:>7} {topk:>5} {ours:>10.3f} {tri:>10.3f} " + f"{tri / ours:>7.2f}x {'sm100a' if ok else 'FALLBACK'}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_block_sparse_sm100a.py b/tests/test_block_sparse_sm100a.py new file mode 100644 index 000000000..6f90f4017 --- /dev/null +++ b/tests/test_block_sparse_sm100a.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Correctness tests for the sm_100a CUDA block-sparse VSA forward. + +Compared against an explicit PyTorch reference rather than the Triton kernel: Triton's +block-sparse forward is hardcoded to 64-token blocks (BLOCK_M = BLOCK_N = 64) while this +extension also carries a 128-token build, so a direct comparison would be comparing two +different sparsity granularities. The reference below applies exactly the semantics the +kernel is supposed to implement -- selected blocks only, keys past variable_block_sizes +masked. Every case runs at both block sizes. + +Run with: python -m pytest tests/test_block_sparse_sm100a.py -v +""" + +import os +import subprocess +import sys + +import pytest +import torch + +from fastvideo_kernel import block_sparse_attn_sm100a as vsa + +HEAD_DIM = 128 +BLOCK_SIZES = [64, 128] + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0) + or not vsa._HAS_VSA_SM100A, + reason="requires Blackwell (sm_100a) and a built fastvideo_kernel extension", +) + + +def make_case(block, num_blocks=8, topk=4, heads=4, batch=1, ragged=False, seed=0): + torch.manual_seed(seed) + S = num_blocks * block + shape = (batch, heads, S, HEAD_DIM) if vsa.BHSD else (batch, S, heads, HEAD_DIM) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + + idx = torch.empty((batch * heads * num_blocks, topk), dtype=torch.int32, device="cuda") + for r in range(idx.shape[0]): + idx[r] = torch.randperm(num_blocks, device="cuda")[:topk].to(torch.int32).sort().values + 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, idx, num, vbs + + +def reference(q, k, v, idx, num, vbs, block): + """Dense attention restricted to the selected blocks, with padded keys masked.""" + if not vsa.BHSD: + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) # -> [B, H, S, D] + B, H, S, D = q.shape + num_blocks = vbs.numel() + scale = 1.0 / (D**0.5) + + 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): + row = (b * H + h) * num_blocks + qb + for j in range(int(num[row])): + kb = int(idx[row, j]) + valid = int(vbs[kb]) + keep[b, h, qb * block:(qb + 1) * block, + kb * block:kb * block + valid] = True + + scores = (q.float() @ k.float().transpose(-1, -2)) * scale + scores = scores.masked_fill(~keep, float("-inf")) + p = torch.softmax(scores, dim=-1) + out = p @ v.float() + lse = torch.logsumexp(scores, dim=-1) * 1.4426950408889634 + return out, lse + + +def run_and_compare(block, ragged, num_blocks=8, topk=4, heads=4, atol=0.02): + q, k, v, idx, num, vbs = make_case(block, num_blocks=num_blocks, topk=topk, heads=heads, + ragged=ragged) + assert vsa.is_supported(q, vbs) + got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + ref, ref_lse = reference(q, k, v, idx, num, vbs, block) + + got_o = got if vsa.BHSD else got.transpose(1, 2) + diff = (got_o.float() - ref).abs().max().item() + assert diff < atol, f"out: max |diff| = {diff:.5f}" + lse_diff = (got_lse.float() - ref_lse).abs().max().item() + assert lse_diff < 0.05, f"lse: max |diff| = {lse_diff:.5f}" + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_forward_matches_reference(block): + run_and_compare(block, ragged=False) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_forward_matches_reference_ragged(block): + """variable_block_sizes is what FastVideo always passes; padded keys must be masked.""" + run_and_compare(block, ragged=True) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +@pytest.mark.parametrize("topk", [1, 2, 3, 5, 7]) +def test_topk_not_a_multiple_of_the_group(block, topk): + """The kernel groups selected blocks; a ragged final group must still be correct.""" + run_and_compare(block, ragged=True, num_blocks=8, topk=topk) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +@pytest.mark.parametrize("num_blocks", [4, 8, 16]) +def test_sequence_lengths(block, num_blocks): + run_and_compare(block, ragged=True, num_blocks=num_blocks, topk=3) + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_lse_is_not_vacuous(block): + """Guards the lse assertion: a wrong lse must actually fail the comparison.""" + q, k, v, idx, num, vbs = make_case(block, ragged=True) + _, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + _, ref_lse = reference(q, k, v, idx, num, vbs, block) + assert (got_lse.float() - (ref_lse + 1.0)).abs().max().item() > 0.5 + + +def test_unsupported_is_rejected(): + q, _, _, _, _, vbs = make_case(64) + assert not vsa.is_supported(q.float(), vbs) # wrong dtype + assert not vsa.is_supported(q[..., :64].contiguous(), vbs) # wrong head_dim + odd = torch.full((7, ), 64, dtype=torch.int32, device="cuda") + assert not vsa.is_supported(q, odd) # seqlen/blocks mismatch + thirty_two = torch.full((16, ), 32, dtype=torch.int32, device="cuda") + assert not vsa.is_supported(q, thirty_two) # block size with no build + + +# --------------------------------------------------------------------------- +# Per-q-tile q2k_num regressions. A CTA owns an ADJACENT PAIR of query blocks; +# the kernel must honor each row's own count -- not the even row's -- in both +# the q2k_idx window clamp and the vbs-threshold masking. These cases pin the +# two failure modes of a pair-shared count: silent corruption of the odd tile +# when the pair's rows differ, and a hang when an even row's count is zero. +# --------------------------------------------------------------------------- + + +def make_two_class_case(block, num_blocks=16, prefix=3, topk=4, heads=4, seed=0, zero_rows=()): + """Two row classes, like a packed multimodal layout: `prefix` DENSE rows + (count = num_blocks) followed by rows at a uniform smaller count + (prefix + topk). With an ODD `prefix`, pair (prefix-1, prefix) straddles + the classes, so the two rows of one CTA carry different counts.""" + torch.manual_seed(seed) + 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)) + vbs = torch.randint(block // 2, block + 1, (num_blocks, ), dtype=torch.int32, device="cuda") + + rows = heads * num_blocks + # Pad past each row's count with 0 -- a VALID block id, like real metadata + # padding -- so a regression to the pair-shared count reads plausible + # in-bounds garbage and fails by WRONG VALUES rather than by luck. + idx = torch.zeros((rows, num_blocks), dtype=torch.int32, device="cuda") + num = torch.zeros((rows, ), dtype=torch.int32, device="cuda") + g = torch.Generator().manual_seed(seed + 1) + for h in range(heads): + for t in range(num_blocks): + r = h * num_blocks + t + if t < prefix: + sel = torch.arange(num_blocks, dtype=torch.int32) + else: + vid = torch.randperm(num_blocks - prefix, generator=g)[:topk] + prefix + sel = torch.cat([torch.arange(prefix), vid.sort().values]).to(torch.int32) + if t in zero_rows: + sel = sel[:0] + idx[r, :sel.numel()] = sel.cuda() + num[r] = sel.numel() + return q, k, v, idx, num, vbs + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_nonuniform_counts_straddling_a_pair(block): + """A dense row paired with a top-k row must both be exact. Regression: the + pair-shared count computed the odd tile with the even row's count (max + |out diff| ~0.6 on this exact case), leaving every other tile correct.""" + q, k, v, idx, num, vbs = make_two_class_case(block) + assert vsa.is_supported(q, vbs) + got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + ref, ref_lse = reference(q, k, v, idx, num, vbs, block) + got_o = got if vsa.BHSD else got.transpose(1, 2) + diff = (got_o.float() - ref).abs().max().item() + assert diff < 0.02, f"out: max |diff| = {diff:.5f}" + lse_diff = (got_lse.float() - ref_lse).abs().max().item() + assert lse_diff < 0.05, f"lse: max |diff| = {lse_diff:.5f}" + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_same_seed_determinism(block): + """Two same-seed runs must be bitwise identical. Also documents that the + uniform-count path is unchanged by the per-tile fix: with cnt0 == cnt1 >= 1 + the pair trip count max(cnt0, cnt1, 1) and the per-tile clamps reduce to + the original expressions.""" + runs = [] + for _ in range(2): + q, k, v, idx, num, vbs = make_case(block, ragged=True, seed=3) + o, m = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + torch.cuda.synchronize() + runs.append((o, m)) + assert torch.equal(runs[0][0], runs[1][0]), "out differs across same-seed runs" + assert torch.equal(runs[0][1], runs[1][1]), "lse differs across same-seed runs" + + +def _zero_count_case_main(block): + """Body of test_zero_count_rows, run in a subprocess (see the test).""" + zero_rows = (5, 6, 8, 9) # odd member; even member with nonzero sibling 7; a whole pair + q, k, v, idx, num, vbs = make_two_class_case(block, zero_rows=zero_rows) + assert vsa.is_supported(q, vbs) + got, got_lse = vsa.block_sparse_attn_sm100a(q, k, v, idx, num, vbs) + torch.cuda.synchronize() + ref, ref_lse = reference(q, k, v, idx, num, vbs, block) + got_o = (got if vsa.BHSD else got.transpose(1, 2)).float() + + empty = torch.zeros(ref.shape[:-1], dtype=torch.bool, device="cuda") # [B, H, S] + for t in zero_rows: + rows = slice(t * block, (t + 1) * block) + empty[:, :, rows] = True + zmax = got_o[:, :, rows].abs().max().item() + assert zmax == 0.0, f"zero-count tile {t}: expected exact zeros, got max |out| = {zmax}" + assert torch.isfinite(got_lse).all(), "lse must stay finite on zero-count rows" + # every non-empty row -- including tile 7, whose pair sibling is empty -- is exact + diff = (got_o - ref).abs().amax(dim=-1) + assert diff[~empty].max().item() < 0.02 + lse_diff = (got_lse.float() - ref_lse).abs() + assert lse_diff[~empty].max().item() < 0.05 + print("ZERO_COUNT_CASE_OK") + + +@pytest.mark.parametrize("block", BLOCK_SIZES) +def test_zero_count_rows(block): + """q2k_num == 0 rows must produce exactly-zero output rows (finite lse) and + leave every other row intact. Runs in a subprocess with a timeout because + the failure mode this pins is a HANG (a zero count on an even row starved + the softmax/correction mbarrier handshake); a regression must fail the + suite, not wedge it.""" + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--zero-count-case", str(block)], + capture_output=True, + text=True, + timeout=240, + ) + assert proc.returncode == 0 and "ZERO_COUNT_CASE_OK" in proc.stdout, ( + f"rc={proc.returncode}\nstdout: {proc.stdout[-2000:]}\nstderr: {proc.stderr[-2000:]}") + + +if __name__ == "__main__": + if len(sys.argv) == 3 and sys.argv[1] == "--zero-count-case": + _zero_count_case_main(int(sys.argv[2]))