From a1df170f04cb993585ab2b92bbd5f4a6a10a2556 Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Wed, 19 Feb 2025 19:25:42 +0800 Subject: [PATCH 1/6] upload file --- csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu | 952 +++++++++++++++++++++++ 1 file changed, 952 insertions(+) create mode 100644 csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu diff --git a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu new file mode 100644 index 00000000..dd843da2 --- /dev/null +++ b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu @@ -0,0 +1,952 @@ +#include +#include +#include "paddle/extension.h" + +#include "sageattn_utils.cuh" + +template +CUtensorMap create_tensor_map_4D(T* gmem_ptr, int d1, int d2, int d3, int d4, int stride1, int stride2, int stride3) { + constexpr int smem_stride = BlockMinorSize * sizeof(T); + static_assert(sizeof(T) == 2 || sizeof(T) == 1); + static_assert(smem_stride == 32 || smem_stride == 64 || smem_stride == 128); + + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)d4, (uint64_t)d3, (uint64_t)d2, (uint64_t)d1, 1}; + uint64_t gmem_prob_stride[5] = {(uint64_t)stride3 * sizeof(T), (uint64_t)stride2 * sizeof(T), (uint64_t)stride1 * sizeof(T), 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, (sizeof(T) == 2) ? CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 : CU_TENSOR_MAP_DATA_TYPE_UINT8, 4, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + (swizzle == false) ? CU_TENSOR_MAP_SWIZZLE_NONE : (smem_stride == 128) ? CU_TENSOR_MAP_SWIZZLE_128B : (smem_stride == 64) ? CU_TENSOR_MAP_SWIZZLE_64B : CU_TENSOR_MAP_SWIZZLE_32B, + promotion_mode, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + + return tma_map; +} + +__device__ __forceinline__ void init_barrier(uint64_t* bar, int thread_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count) + ); +} + +template +__device__ __forceinline__ void expect_bytes(uint64_t* bar) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "n"(bytes)); +} + +template +__device__ __forceinline__ void load_async_4D(T *dst, void const* const src_tma_map, uint64_t* bar, int s0, int s1, int s2, int s3) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "r"(s0), "r"(s1), "r"(s2), "r"(s3) + : "memory" + ); +} + +template +__device__ __forceinline__ void store_async_4D(void const* dst_tma_map, T *src, int global_token_idx, int global_head_idx, int global_batch_idx) { + uint64_t tma_ptr = reinterpret_cast(dst_tma_map); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(src)); + + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(global_token_idx), "r"(global_head_idx), "r"(global_batch_idx) + : "memory" + ); +} + +__device__ __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +template +__device__ __forceinline__ void arrive(uint64_t* bar) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "n"(count) + : "memory" + ); +} + +// +// ======= kernel impl ======= +// + +template +__global__ void qk_int8_sv_f8_attn_dsk_kernel(const __grid_constant__ CUtensorMap tensorMapQ, + const __grid_constant__ CUtensorMap tensorMapK, + const __grid_constant__ CUtensorMap tensorMapQ_pe, + const __grid_constant__ CUtensorMap tensorMapK_pe, + const __grid_constant__ CUtensorMap tensorMapV, + float *__restrict__ Q_scale, float *__restrict__ K_scale, float *__restrict__ V_scale, + DTypeOut* O, uint32_t stride_bz_o, uint32_t stride_h_o, uint32_t stride_seq_o, + const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, + float sm_scale) +{ + static_assert(NUM_THREADS == 128); + static_assert(CTA_Q <= CTA_K); + + const uint32_t warp_idx = (threadIdx.x % 128) / 32; + const uint32_t lane_id = threadIdx.x % 32; + + constexpr uint32_t num_tiles_q = CTA_Q / 64; + constexpr uint32_t num_tiles_k = CTA_K / 16; + constexpr uint32_t num_tiles_qk_inner = head_dim / 32; + constexpr uint32_t num_tiles_qk_pe_inner = head_dim_pe / 32; + constexpr uint32_t num_tiles_v = head_dim / 16; + constexpr uint32_t num_tiles_pv_inner = CTA_K / 32; + + const uint32_t batch_id = blockIdx.z; + const uint32_t bx = blockIdx.x; + const uint32_t head_id = blockIdx.y; + const uint32_t num_qo_heads = gridDim.y; + const uint32_t kv_head_id = head_id / num_kv_groups; + + sm_scale *= math::log2e; + + extern __shared__ __align__(128) int8_t smem_[]; + + /* // original: + * int8_t *sQ = (int8_t*)smem_; + * int8_t *sK = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t)); + * int8_t *sV = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t)); + * half *sO = (half*)smem_; + */ + + int8_t *sQ = (int8_t*)smem_; // 0 + int8_t *sQ_pe = (int8_t*)(smem_ + CTA_Q * (head_dim) * sizeof(int8_t)); // 0 + head_dim + + int8_t *sK = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t)); // 0 + head_dim + pe + int8_t *sK_pe = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t) + CTA_K * (head_dim) * sizeof(int8_t)); // 0 + head_dim + pe + head_dim + int8_t *sV = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t) + CTA_K * (head_dim + head_dim_pe) * sizeof(int8_t)); + half *sO = (half*)smem_; + + int32_t RS[num_tiles_q][num_tiles_k][8]; + int32_t RS_pe[num_tiles_q][num_tiles_k][8]; + float RO[num_tiles_q][num_tiles_v][8]; + float m[num_tiles_q][2]; + float d[num_tiles_q][2]; + + uint32_t q_scale_idx, k_scale_idx; + + // scale shape: (b, h_qo, (qo_len + BLKQ - 1) // BLKQ) + if constexpr (Q_GRAN == QuantGranularity::kPerBlock) + { + const uint32_t num_block_q = gridDim.x; + q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx; + } + else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_warp_block_q = gridDim.x * 4; + q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * 4 + warp_idx; + } + else if constexpr (Q_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_warp_block_q = gridDim.x * 4; + q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (4 * 8) + warp_idx * 8 + lane_id / 4; + } + + if constexpr (K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_block_k = div_ceil(kv_len, CTA_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k; + } + else if constexpr (K_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_block_k = div_ceil(kv_len, CTA_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_block_k * 4) + (head_id / num_kv_groups) * (num_block_k * 4) + lane_id % 4; + } + + constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp) ? 1 : 4; + + uint32_t Q_idx_lane_base = bx * CTA_Q + warp_idx * 16 + lane_id / 4; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + m[fq][0] = -5000000.0f; + m[fq][1] = -5000000.0f; + d[fq][0] = 1.0f; + d[fq][1] = 1.0f; + } + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RO[fq][fv][k] = 0.0f; + } + } + } + + __shared__ __align__(8) uint64_t barrier_Q; + __shared__ __align__(8) uint64_t barrier_K; + __shared__ __align__(8) uint64_t barrier_Q_pe; + __shared__ __align__(8) uint64_t barrier_K_pe; + __shared__ __align__(8) uint64_t barrier_V; + + if (threadIdx.x == 0) { + init_barrier(&barrier_Q, 1); + init_barrier(&barrier_K, 1); + init_barrier(&barrier_Q_pe, 1); + init_barrier(&barrier_K_pe, 1); + init_barrier(&barrier_V, 1); + } + + __syncthreads(); + + // load Q, K, V + if (threadIdx.x == 0) + { + expect_bytes<(CTA_Q * (head_dim)) * sizeof(int8_t)>(&barrier_Q); + expect_bytes<(CTA_K * (head_dim)) * sizeof(int8_t)>(&barrier_K); + expect_bytes<(CTA_Q * (head_dim_pe)) * sizeof(int8_t)>(&barrier_Q_pe); + expect_bytes<(CTA_K * (head_dim_pe)) * sizeof(int8_t)>(&barrier_K_pe); + expect_bytes<(CTA_K * (head_dim)) * sizeof(int8_t)>(&barrier_V); + + load_async_4D(sQ, &tensorMapQ, &barrier_Q, 0, bx * CTA_Q, head_id, batch_id); + load_async_4D(sQ_pe, &tensorMapQ_pe, &barrier_Q_pe, 0, bx * CTA_Q, head_id, batch_id); + load_async_4D(sK, &tensorMapK, &barrier_K, 0, 0, kv_head_id, batch_id); + load_async_4D(sK_pe, &tensorMapK_pe, &barrier_K_pe, 0, 0, kv_head_id, batch_id); + load_async_4D(sV, &tensorMapV, &barrier_V, 0, 0, kv_head_id, batch_id); + } + + float q_scale = Q_scale[q_scale_idx]; + float original_sm_scale = sm_scale; + + // wait for Q + wait(&barrier_Q, 0); + wait(&barrier_Q_pe, 0); + + const uint32_t num_iterations = div_ceil( + mask_mode == MaskMode::kCausal + ? min(kv_len, (bx + 1) * CTA_Q) + : kv_len, + CTA_K); + + int p = 1; + for (uint32_t iter = 1; iter < num_iterations; iter++) + { + p ^= 1; + + float dequant_scale = q_scale * K_scale[k_scale_idx + (iter - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // wait for K + wait(&barrier_K, p); + wait(&barrier_K_pe, p); + + // compute QK^T + wgmma::warpgroup_arrive(); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + int8_t *sQ_local = sQ + fq * 64 * head_dim; + int8_t *sQ_local_pe = sQ_pe + fq * 64 * head_dim_pe; + wgmma::wgmma_s8s8s32(RS[fq], sQ_local, sK); + wgmma::wgmma_s8s8s32(RS_pe[fq], sQ_local_pe, sK_pe); // add one line +#pragma unroll + for (int k_it = 1; k_it < num_tiles_qk_inner; k_it++) + { + wgmma::wgmma_s8s8s32(RS[fq], &sQ_local[k_it*32], &sK[k_it*32]); + if (k_it < num_tiles_qk_pe_inner) { + wgmma::wgmma_s8s8s32(RS_pe[fq], &sQ_local_pe[k_it*32], &sK_pe[k_it*32]); // add one line + } + } + } + wgmma::warpgroup_commit_batch(); + wgmma::warpgroup_wait<0>(); + + // load K + if (threadIdx.x == 0) + { + expect_bytes<(CTA_K * head_dim) * sizeof(int8_t)>(&barrier_K); + expect_bytes<(CTA_K * head_dim_pe) * sizeof(int8_t)>(&barrier_K_pe); // add one line + load_async_4D(sK, &tensorMapK, &barrier_K, 0, iter * CTA_K, kv_head_id, batch_id); + load_async_4D(sK_pe, &tensorMapK_pe, &barrier_K_pe, 0, iter * CTA_K, kv_head_id, batch_id); // add one line + } + + // convert RS to float + float RS_f32[num_tiles_q][num_tiles_k][8]; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k] + RS_pe[fq][fk][k]); // add one line + } + } + } + + update_mdo(RS_f32, RO, m, d, sm_scale); + + // accumulate d on thread basis +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unrol + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + d[fq][0] += (RS_f32[fq][fk][0] + RS_f32[fq][fk][1] + RS_f32[fq][fk][4] + RS_f32[fq][fk][5]); + d[fq][1] += (RS_f32[fq][fk][2] + RS_f32[fq][fk][3] + RS_f32[fq][fk][6] + RS_f32[fq][fk][7]); + } + } + + uint32_t RS_f8[num_tiles_q][num_tiles_pv_inner][4]; + RS_32_to_8(RS_f32, RS_f8); + + // wait for V + wait(&barrier_V, p); + + float RO_temp[num_tiles_q][num_tiles_v][8]; + wgmma::warpgroup_arrive(); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + wgmma::wgmma_f8f8f32(RO_temp[fq], RS_f8[fq][0], &sV[0]); +#pragma unroll + for (uint32_t v_it = 1; v_it < num_tiles_pv_inner; v_it++) + { + wgmma::wgmma_f8f8f32(RO_temp[fq], RS_f8[fq][v_it], &sV[v_it * 32]); + } + } + + wgmma::warpgroup_commit_batch(); + wgmma::warpgroup_wait<0>(); + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RO[fq][fv][k] += RO_temp[fq][fv][k]; + } + } + } + + // load V + if (threadIdx.x == 0) + { + expect_bytes<(CTA_K * head_dim) * sizeof(int8_t)>(&barrier_V); + load_async_4D(sV, &tensorMapV, &barrier_V, iter * CTA_K, 0, kv_head_id, batch_id); + } + } + + { + p ^= 1; + + float dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale; + + // wait for K + wait(&barrier_K, p); + wait(&barrier_K_pe, p); + + // compute QK^T + wgmma::warpgroup_arrive(); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + int8_t *sQ_local = sQ + fq * 64 * head_dim; + int8_t *sQ_local_pe = sQ_pe + fq * 64 * head_dim_pe; + wgmma::wgmma_s8s8s32(RS[fq], sQ_local, sK); + wgmma::wgmma_s8s8s32(RS_pe[fq], sQ_local_pe, sK_pe); +#pragma unroll + for (int k_it = 1; k_it < num_tiles_qk_inner; k_it++) + { + wgmma::wgmma_s8s8s32(RS[fq], &sQ_local[k_it*32], &sK[k_it*32]); + if (k_it < num_tiles_qk_pe_inner) { + wgmma::wgmma_s8s8s32(RS_pe[fq], &sQ_local_pe[k_it*32], &sK_pe[k_it*32]); // add one line + } + } + } + wgmma::warpgroup_commit_batch(); + wgmma::warpgroup_wait<0>(); + + // convert RS to float + float RS_f32[num_tiles_q][num_tiles_k][8]; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k] + RS_pe[fq][fk][k]) * dequant_scale; + } + } + } + + // masking +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + const uint32_t q_idx = Q_idx_lane_base + fq * 64 + 8 * ((k % 4) / 2); + const uint32_t k_idx = (num_iterations - 1) * CTA_K + fk * 16 + 2 * (lane_id % 4) + 8 * (k / 4) + k % 2; + + bool is_out_of_bounds; + + if constexpr (mask_mode == MaskMode::kCausal) + { + is_out_of_bounds = (k_idx > q_idx) || (k_idx >= kv_len); + } + else + { + is_out_of_bounds = (k_idx >= kv_len); + } + + if (is_out_of_bounds) + { + RS_f32[fq][fk][k] = -5000000.0f; + } + } + } + } + + update_mdo(RS_f32, RO, m, d, sm_scale); + + // accumulate d on thread basis +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unrol + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + d[fq][0] += (RS_f32[fq][fk][0] + RS_f32[fq][fk][1] + RS_f32[fq][fk][4] + RS_f32[fq][fk][5]); + d[fq][1] += (RS_f32[fq][fk][2] + RS_f32[fq][fk][3] + RS_f32[fq][fk][6] + RS_f32[fq][fk][7]); + } + } + + uint32_t RS_f8[num_tiles_q][num_tiles_pv_inner][4]; + RS_32_to_8(RS_f32, RS_f8); + + // wait for V + wait(&barrier_V, p); + + float RO_temp[num_tiles_q][num_tiles_v][8]; + wgmma::warpgroup_arrive(); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + wgmma::wgmma_f8f8f32(RO_temp[fq], RS_f8[fq][0], &sV[0]); +#pragma unroll + for (uint32_t v_it = 1; v_it < num_tiles_pv_inner; v_it++) + { + wgmma::wgmma_f8f8f32(RO_temp[fq], RS_f8[fq][v_it], &sV[v_it * 32]); + } + } + + wgmma::warpgroup_commit_batch(); + wgmma::warpgroup_wait<0>(); + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RO[fq][fv][k] += RO_temp[fq][fv][k]; + } + } + } + } + + normalize_d(RO, m, d); + + if constexpr (fuse_v_scale) + { + float v_scale[4]; + float *V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4 ) * 2; + #pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + ((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16)); + ((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8)); + + #pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + RO[fq][fv][0] *= v_scale[0]; + RO[fq][fv][1] *= v_scale[1]; + RO[fq][fv][2] *= v_scale[0]; + RO[fq][fv][3] *= v_scale[1]; + RO[fq][fv][4] *= v_scale[2]; + RO[fq][fv][5] *= v_scale[3]; + RO[fq][fv][6] *= v_scale[2]; + RO[fq][fv][7] *= v_scale[3]; + } + } + } + + DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + warp_idx * 16 + (lane_id / 4)) * stride_seq_o + (lane_id % 4) * 2 ; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < head_dim/16; fv++) + { + if (Q_idx_lane_base + fq * 64 < qo_len) + { + if constexpr (std::is_same::value) + { + ((half2*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16))[0] = __float22half2_rn(((float2*)(RO[fq][fv]))[0]); + ((half2*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8))[0] = __float22half2_rn(((float2*)(RO[fq][fv]))[2]); + } + else + { + ((nv_bfloat162*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16))[0] = __float22bfloat162_rn(((float2*)(RO[fq][fv]))[0]); + ((nv_bfloat162*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8))[0] = __float22bfloat162_rn(((float2*)(RO[fq][fv]))[2]); + } + } + + if (Q_idx_lane_base + fq * 64 + 8 < qo_len) + { + if constexpr (std::is_same::value) + { + ((half2*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8 * stride_seq_o))[0] = __float22half2_rn(((float2*)(RO[fq][fv]))[1]); + ((half2*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8 + 8 * stride_seq_o))[0] = __float22half2_rn(((float2*)(RO[fq][fv]))[3]); + } + else + { + ((nv_bfloat162*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8 * stride_seq_o))[0] = __float22bfloat162_rn(((float2*)(RO[fq][fv]))[1]); + ((nv_bfloat162*)(O_lane_ptr + fq * 64 * stride_seq_o + fv * 16 + 8 + 8 * stride_seq_o))[0] = __float22bfloat162_rn(((float2*)(RO[fq][fv]))[3]); + } + } + } + } +} + +std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( + paddle::Tensor& query, + paddle::Tensor& key, + paddle::Tensor& query_pe, + paddle::Tensor& key_pe, + paddle::Tensor& value, + paddle::Tensor& output, + paddle::Tensor& query_scale, + paddle::Tensor& key_scale, + int tensor_layout, + int is_causal, + int qk_quant_gran, + float sm_scale, + int return_lse) +{ + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(query_pe); + CHECK_CUDA(key_pe); + CHECK_CUDA(value); + CHECK_CUDA(output); + CHECK_CUDA(query_scale); + CHECK_CUDA(key_scale); + + CHECK_LASTDIM_CONTIGUOUS(query); + CHECK_LASTDIM_CONTIGUOUS(key); + CHECK_LASTDIM_CONTIGUOUS(query_pe); + CHECK_LASTDIM_CONTIGUOUS(key_pe); + CHECK_LASTDIM_CONTIGUOUS(value); + CHECK_LASTDIM_CONTIGUOUS(output); + CHECK_CONTIGUOUS(query_scale); + CHECK_CONTIGUOUS(key_scale); + + CHECK_DTYPE(query, paddle::DataType::INT8); + CHECK_DTYPE(key, paddle::DataType::INT8); + CHECK_DTYPE(query_pe, paddle::DataType::INT8); + CHECK_DTYPE(key_pe, paddle::DataType::INT8); + CHECK_DTYPE(value, paddle::DataType::FLOAT8_E4M3FN); + CHECK_DTYPE(query_scale, paddle::DataType::FLOAT32); + CHECK_DTYPE(key_scale, paddle::DataType::FLOAT32); + + CHECK_DIMS(query, 4); + CHECK_DIMS(key, 4); + CHECK_DIMS(query_pe, 4); + CHECK_DIMS(key_pe, 4); + CHECK_DIMS(value, 4); + CHECK_DIMS(output, 4); + CHECK_DIMS(query_scale, 3); + CHECK_DIMS(key_scale, 3); + + const int batch_size = query.shape()[0]; + const int head_dim = query.shape()[3]; // 现在query是正常的128， 多出来的64在query_pe里面，所以这样做没什么问题 + + int stride_bz_q = query.strides()[0]; + int stride_bz_k = key.strides()[0]; + int stride_bz_v = value.strides()[0]; + int stride_bz_o = output.strides()[0]; + + int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; + int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; + + assert(value.shape()[0] == batch_size); + + if (tensor_layout == 0) + { + qo_len = query.shape()[1]; + kv_len = key.shape()[1]; + num_qo_heads = query.shape()[2]; + num_kv_heads = key.shape()[2]; + + stride_seq_q = query.strides()[1]; + stride_h_q = query.strides()[2]; + stride_seq_k = key.strides()[1]; + stride_h_k = key.strides()[2]; + stride_h_v = value.strides()[2]; + stride_d_v = value.strides()[1]; + stride_seq_o = output.strides()[1]; + stride_h_o = output.strides()[2]; + + CHECK_SHAPE(key, batch_size, kv_len, num_kv_heads, head_dim); + CHECK_SHAPE(output, batch_size, qo_len, num_qo_heads, head_dim); + assert(value.shape()[1] == head_dim); + assert(value.shape()[2] == num_kv_heads); + } + else + { + qo_len = query.shape()[2]; + kv_len = key.shape()[2]; + num_qo_heads = query.shape()[1]; + num_kv_heads = key.shape()[1]; + + stride_seq_q = query.strides()[2]; + stride_h_q = query.strides()[1]; + stride_seq_k = key.strides()[2]; + stride_h_k = key.strides()[1]; + stride_h_v = value.strides()[1]; + stride_d_v = value.strides()[2]; + stride_seq_o = output.strides()[2]; + stride_h_o = output.strides()[1]; + + CHECK_SHAPE(key, batch_size, num_kv_heads, kv_len, head_dim); + CHECK_SHAPE(output, batch_size, num_qo_heads, qo_len, head_dim); + assert(value.shape()[2] == head_dim); + assert(value.shape()[1] == num_kv_heads); + } + + if (num_qo_heads % num_kv_heads != 0) { + std::ostringstream err_msg; + err_msg << "num_qo_heads (" << num_qo_heads << ") must be divisible by num_kv_heads (" << num_kv_heads << ")"; + throw std::invalid_argument(err_msg.str()); + } + + paddle::Tensor lse = paddle::empty({0}, paddle::DataType::FLOAT32); + if (return_lse) + { + lse = paddle::empty({batch_size, num_qo_heads, qo_len}, paddle::DataType::FLOAT32); + } + + const int num_kv_groups = num_qo_heads / num_kv_heads; + + auto output_type = output.dtype(); + + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + DISPATCH_CAUSAL(is_causal, IS_CAUSAL, { + DISPATCH_QK_QUANT_GRAN(qk_quant_gran, QK_QUANT_GRAN, { + DISPATCH_PADDLE_DTYPE_TO_CTYPE_FP16(output_type, DTypeOut, { + constexpr int CTA_Q = 64; + constexpr int CTA_K = 128; + constexpr int NUM_THREADS = 128; + constexpr int HEAD_DIM_PE = 64; + + constexpr MaskMode mask_mode = IS_CAUSAL ? MaskMode::kCausal : MaskMode::kNone; + + assert(value.shape()[3] >= div_ceil(kv_len, CTA_K) * CTA_K); + + if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp)) + { + CHECK_SHAPE(query_scale, batch_size, num_qo_heads, static_cast(div_ceil(qo_len, CTA_Q) * (NUM_THREADS / 32))); + CHECK_SHAPE(key_scale, batch_size, num_kv_heads, static_cast(div_ceil(kv_len, CTA_K))); + } + else if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerThread)) + { + CHECK_SHAPE(query_scale, batch_size, num_qo_heads, static_cast(div_ceil(qo_len, CTA_Q) * (NUM_THREADS / 32) * 8)); + CHECK_SHAPE(key_scale, batch_size, num_kv_heads, static_cast(div_ceil(kv_len, CTA_K) * 4)); + } + else + { + static_assert(QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp) || QK_QUANT_GRAN == static_cast(QuantGranularity::kPerThread), "Unsupported quantization granularity"); + } + + CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + + CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data()), batch_size, num_kv_heads, HEAD_DIM, value.shape()[3], stride_bz_v, stride_h_v, stride_d_v); + + auto* kernel = qk_int8_sv_f8_attn_dsk_kernel(QK_QUANT_GRAN), static_cast(QK_QUANT_GRAN), DTypeOut, mask_mode, false>; + size_t sMemSize = CTA_Q * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t); + sMemSize += CTA_Q * HEAD_DIM_PE * sizeof(int8_t) + CTA_K * HEAD_DIM_PE * sizeof(int8_t); // add extra space for qk pe + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize); + + dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size); + kernel<<>>( + tma_map_Q, + tma_map_K, + tma_map_Q_pe, + tma_map_K_pe, + tma_map_V, + reinterpret_cast(query_scale.data()), + reinterpret_cast(key_scale.data()), + nullptr, + reinterpret_cast(output.data()), + stride_bz_o, stride_h_o, stride_seq_o, + qo_len, kv_len, num_kv_groups, sm_scale); + }); + }); + }); + }); + + return {lse}; +} + +std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90_fwd( + paddle::Tensor& query, + paddle::Tensor& key, + paddle::Tensor& query_pe, + paddle::Tensor& key_pe, + paddle::Tensor& value, + paddle::Tensor& output, + paddle::Tensor& query_scale, + paddle::Tensor& key_scale, + paddle::Tensor& value_scale, + int tensor_layout, + int is_causal, + int qk_quant_gran, + float sm_scale, + int return_lse) +{ + CHECK_CUDA(query); + CHECK_CUDA(key); + CHECK_CUDA(query_pe); + CHECK_CUDA(key_pe); + CHECK_CUDA(value); + CHECK_CUDA(output); + CHECK_CUDA(query_scale); + CHECK_CUDA(key_scale); + CHECK_CUDA(value_scale); + + CHECK_LASTDIM_CONTIGUOUS(query); + CHECK_LASTDIM_CONTIGUOUS(key); + CHECK_LASTDIM_CONTIGUOUS(query_pe); + CHECK_LASTDIM_CONTIGUOUS(key_pe); + CHECK_LASTDIM_CONTIGUOUS(value); + CHECK_LASTDIM_CONTIGUOUS(output); + CHECK_CONTIGUOUS(query_scale); + CHECK_CONTIGUOUS(key_scale); + CHECK_CONTIGUOUS(value_scale); + + CHECK_DTYPE(query, paddle::DataType::INT8); + CHECK_DTYPE(key, paddle::DataType::INT8); + CHECK_DTYPE(query_pe, paddle::DataType::INT8); + CHECK_DTYPE(key_pe, paddle::DataType::INT8); + CHECK_DTYPE(value, paddle::DataType::FLOAT8_E4M3FN); + CHECK_DTYPE(query_scale, paddle::DataType::FLOAT32); + CHECK_DTYPE(key_scale, paddle::DataType::FLOAT32); + CHECK_DTYPE(value_scale, paddle::DataType::FLOAT32); + + CHECK_DIMS(query, 4); + CHECK_DIMS(key, 4); + CHECK_DIMS(query_pe, 4); + CHECK_DIMS(key_pe, 4); + CHECK_DIMS(value, 4); + CHECK_DIMS(output, 4); + CHECK_DIMS(query_scale, 3); + CHECK_DIMS(key_scale, 3); + CHECK_DIMS(value_scale, 3); + + const int batch_size = query.shape()[0]; + const int head_dim = query.shape()[3]; + + int stride_bz_q = query.strides()[0]; + int stride_bz_k = key.strides()[0]; + int stride_bz_v = value.strides()[0]; + int stride_bz_o = output.strides()[0]; + + int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; + int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; + + assert(value.shape()[0] == batch_size); + + if (tensor_layout == 0) + { + qo_len = query.shape()[1]; + kv_len = key.shape()[1]; + num_qo_heads = query.shape()[2]; + num_kv_heads = key.shape()[2]; + + stride_seq_q = query.strides()[1]; + stride_h_q = query.strides()[2]; + stride_seq_k = key.strides()[1]; + stride_h_k = key.strides()[2]; + stride_h_v = value.strides()[2]; + stride_d_v = value.strides()[1]; + stride_seq_o = output.strides()[1]; + stride_h_o = output.strides()[2]; + + CHECK_SHAPE(key, batch_size, kv_len, num_kv_heads, head_dim); + CHECK_SHAPE(output, batch_size, qo_len, num_qo_heads, head_dim); + + assert(value.shape()[1] == head_dim); + assert(value.shape()[2] == num_kv_heads); + } + else + { + qo_len = query.shape()[2]; + kv_len = key.shape()[2]; + num_qo_heads = query.shape()[1]; + num_kv_heads = key.shape()[1]; + + stride_seq_q = query.strides()[2]; + stride_h_q = query.strides()[1]; + stride_seq_k = key.strides()[2]; + stride_h_k = key.strides()[1]; + stride_h_v = value.strides()[1]; + stride_d_v = value.strides()[2]; + stride_seq_o = output.strides()[2]; + stride_h_o = output.strides()[1]; + + CHECK_SHAPE(key, batch_size, num_kv_heads, kv_len, head_dim); + CHECK_SHAPE(output, batch_size, num_qo_heads, qo_len, head_dim); + assert(value.shape()[2] == head_dim); + assert(value.shape()[1] == num_kv_heads); + } + + if (num_qo_heads % num_kv_heads != 0) { + std::ostringstream err_msg; + err_msg << "num_qo_heads (" << num_qo_heads << ") must be divisible by num_kv_heads (" << num_kv_heads << ")"; + throw std::invalid_argument(err_msg.str()); + } + + paddle::Tensor lse = paddle::empty({1}, paddle::DataType::FLOAT32); + if (return_lse) + { + lse = paddle::empty({batch_size, num_qo_heads, qo_len}, paddle::DataType::FLOAT32); + } + + const int num_kv_groups = num_qo_heads / num_kv_heads; + + auto output_dtype = output.dtype(); + + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + DISPATCH_CAUSAL(is_causal, IS_CAUSAL, { + DISPATCH_QK_QUANT_GRAN(qk_quant_gran, QK_QUANT_GRAN, { + DISPATCH_PADDLE_DTYPE_TO_CTYPE_FP16(output_dtype, DTypeOut, { + constexpr int CTA_Q = 64; + constexpr int CTA_K = 128; + constexpr int NUM_THREADS = 128; + constexpr int HEAD_DIM_PE = 64; + + constexpr MaskMode mask_mode = IS_CAUSAL ? MaskMode::kCausal : MaskMode::kNone; + + assert(value.shape()[3] >= div_ceil(kv_len, CTA_K) * CTA_K); + + if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp)) + { + CHECK_SHAPE(query_scale, batch_size, num_qo_heads, static_cast(div_ceil(qo_len, CTA_Q) * (NUM_THREADS / 32))); + CHECK_SHAPE(key_scale, batch_size, num_kv_heads, static_cast(div_ceil(kv_len, CTA_K))); + } + else if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerThread)) + { + CHECK_SHAPE(query_scale, batch_size, num_qo_heads, static_cast(div_ceil(qo_len, CTA_Q) * (NUM_THREADS / 32) * 8)); + CHECK_SHAPE(key_scale, batch_size, num_kv_heads, static_cast(div_ceil(kv_len, CTA_K) * 4)); + } + else + { + static_assert(QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp) || QK_QUANT_GRAN == static_cast(QuantGranularity::kPerThread), "Unsupported quantization granularity"); + } + + CHECK_SHAPE(value_scale, batch_size, num_kv_heads, HEAD_DIM); + CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + + CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data()), batch_size, num_kv_heads, HEAD_DIM, value.shape()[3], stride_bz_v, stride_h_v, stride_d_v); + + auto* kernel = qk_int8_sv_f8_attn_dsk_kernel(QK_QUANT_GRAN), static_cast(QK_QUANT_GRAN), DTypeOut, mask_mode, true>; + size_t sMemSize = CTA_Q * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t); + sMemSize += CTA_Q * HEAD_DIM_PE * sizeof(int8_t) + CTA_K * HEAD_DIM_PE * sizeof(int8_t); + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize); + + dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size); + kernel<<>>( + tma_map_Q, + tma_map_K, + tma_map_Q_pe, + tma_map_K_pe, + tma_map_V, + reinterpret_cast(query_scale.data()), + reinterpret_cast(key_scale.data()), + reinterpret_cast(value_scale.data()), + reinterpret_cast(output.data()), + stride_bz_o, stride_h_o, stride_seq_o, + qo_len, kv_len, num_kv_groups, sm_scale); + }); + }); + }); + }); + + return {lse}; +} \ No newline at end of file From 36230f4149a730ffa0849172f20dc8b356455581 Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Wed, 19 Feb 2025 19:46:00 +0800 Subject: [PATCH 2/6] convert to torch --- csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu index dd843da2..85e5970d 100644 --- a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu +++ b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu @@ -1,8 +1,14 @@ -#include -#include -#include "paddle/extension.h" +#include "../utils.cuh" +#include +#include +#include +#include -#include "sageattn_utils.cuh" +#include "../wgmma.cuh" +#include "../math.cuh" +#include "../dispatch_utils.h" + +#include "attn_utils.cuh" template CUtensorMap create_tensor_map_4D(T* gmem_ptr, int d1, int d2, int d3, int d4, int stride1, int stride2, int stride3) { From e207ba61650bb84526736029c228ee9cc1932dfc Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Thu, 20 Feb 2025 00:45:46 +0800 Subject: [PATCH 3/6] complete patch --- csrc/dispatch_utils.h | 16 ++ csrc/fused/fused.cu | 4 +- csrc/qattn/attn_cuda_sm90.h | 31 +++ csrc/qattn/pybind_sm90.cpp | 2 + csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu | 274 +++++++++++------------ patch_test/test_dsk.py | 66 ++++++ sageattention/__init__.py | 2 +- sageattention/core.py | 87 +++++++ setup.py | 1 + 9 files changed, 343 insertions(+), 140 deletions(-) create mode 100644 patch_test/test_dsk.py diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index 6d798e68..fc41af7e 100644 --- a/csrc/dispatch_utils.h +++ b/csrc/dispatch_utils.h @@ -33,6 +33,22 @@ throw std::invalid_argument(err_msg.str()); \ } +#define DISPATCH_HEAD_DIM_QK(head_dim, HEAD_DIM, ...) \ + if (head_dim == 64) { \ + constexpr int HEAD_DIM = 64; \ + __VA_ARGS__ \ + } else if (head_dim == 128) { \ + constexpr int HEAD_DIM = 128; \ + __VA_ARGS__ \ + } else if (head_dim == 192) { \ + constexpr int HEAD_DIM = 192; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported head dim: " << int(head_dim); \ + throw std::invalid_argument(err_msg.str()); \ + } + #define DISPATCH_CAUSAL(is_causal, IS_CAUSAL, ...) \ if (is_causal == 1) { \ constexpr bool IS_CAUSAL = true; \ diff --git a/csrc/fused/fused.cu b/csrc/fused/fused.cu index fb8b9f15..4c10df0e 100644 --- a/csrc/fused/fused.cu +++ b/csrc/fused/fused.cu @@ -652,7 +652,7 @@ void quant_per_block_int8_fuse_sub_mean_cuda( DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input_dtype, c_type, { DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, { - DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + DISPATCH_HEAD_DIM_QK(head_dim, HEAD_DIM, { CHECK_SHAPE(mean, batch_size, num_heads, head_dim); CHECK_SHAPE(output, input.size(0), input.size(1), input.size(2), input.size(3)); @@ -738,7 +738,7 @@ void quant_per_warp_int8_cuda( DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input_dtype, c_type, { DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, { DISPATCH_WARP_BLOCK_SIZE(warp_block_size, WARP_BLOCK_SIZE, { - DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + DISPATCH_HEAD_DIM_QK(head_dim, HEAD_DIM, { CHECK_SHAPE(output, input.size(0), input.size(1), input.size(2), input.size(3)); CHECK_SHAPE(scale, batch_size, num_heads, (num_tokens + BLOCK_SIZE - 1) / BLOCK_SIZE * (BLOCK_SIZE / WARP_BLOCK_SIZE)); diff --git a/csrc/qattn/attn_cuda_sm90.h b/csrc/qattn/attn_cuda_sm90.h index 822ab7f4..e10c7d29 100644 --- a/csrc/qattn/attn_cuda_sm90.h +++ b/csrc/qattn/attn_cuda_sm90.h @@ -29,6 +29,21 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf( float sm_scale, int return_lse); +torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( + torch::Tensor query, + torch::Tensor key, + torch::Tensor query_pe, + torch::Tensor key_pe, + torch::Tensor value, + torch::Tensor output, + torch::Tensor query_scale, + torch::Tensor key_scale, + int tensor_layout, + int is_causal, + int qk_quant_gran, + float sm_scale, + int return_lse); + torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf( torch::Tensor query, torch::Tensor key, @@ -41,4 +56,20 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf( int is_causal, int qk_quant_gran, float sm_scale, + int return_lse); + +torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( + torch::Tensor query, + torch::Tensor key, + torch::Tensor query_pe, + torch::Tensor key_pe, + torch::Tensor value, + torch::Tensor output, + torch::Tensor query_scale, + torch::Tensor key_scale, + torch::Tensor value_scale, + int tensor_layout, + int is_causal, + int qk_quant_gran, + float sm_scale, int return_lse); \ No newline at end of file diff --git a/csrc/qattn/pybind_sm90.cpp b/csrc/qattn/pybind_sm90.cpp index 8900b641..0d6b905d 100644 --- a/csrc/qattn/pybind_sm90.cpp +++ b/csrc/qattn/pybind_sm90.cpp @@ -22,4 +22,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qk_int8_sv_f8_accum_f32_attn_inst_buf", &qk_int8_sv_f8_accum_f32_attn_inst_buf); m.def("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf", &qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf); + m.def("qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90", &qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90); + m.def("qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90", &qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90); } \ No newline at end of file diff --git a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu index 85e5970d..b3809d21 100644 --- a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu +++ b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu @@ -580,15 +580,15 @@ __global__ void qk_int8_sv_f8_attn_dsk_kernel(const __grid_constant__ CUtensorMa } } -std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( - paddle::Tensor& query, - paddle::Tensor& key, - paddle::Tensor& query_pe, - paddle::Tensor& key_pe, - paddle::Tensor& value, - paddle::Tensor& output, - paddle::Tensor& query_scale, - paddle::Tensor& key_scale, +torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( + torch::Tensor query, + torch::Tensor key, + torch::Tensor query_pe, + torch::Tensor key_pe, + torch::Tensor value, + torch::Tensor output, + torch::Tensor query_scale, + torch::Tensor key_scale, int tensor_layout, int is_causal, int qk_quant_gran, @@ -613,13 +613,13 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( CHECK_CONTIGUOUS(query_scale); CHECK_CONTIGUOUS(key_scale); - CHECK_DTYPE(query, paddle::DataType::INT8); - CHECK_DTYPE(key, paddle::DataType::INT8); - CHECK_DTYPE(query_pe, paddle::DataType::INT8); - CHECK_DTYPE(key_pe, paddle::DataType::INT8); - CHECK_DTYPE(value, paddle::DataType::FLOAT8_E4M3FN); - CHECK_DTYPE(query_scale, paddle::DataType::FLOAT32); - CHECK_DTYPE(key_scale, paddle::DataType::FLOAT32); + CHECK_DTYPE(query, torch::kInt8); + CHECK_DTYPE(key, torch::kInt8); + CHECK_DTYPE(query_pe, torch::kInt8); + CHECK_DTYPE(key_pe, torch::kInt8); + CHECK_DTYPE(value, at::ScalarType::Float8_e4m3fn); + CHECK_DTYPE(query_scale, torch::kFloat32); + CHECK_DTYPE(key_scale, torch::kFloat32); CHECK_DIMS(query, 4); CHECK_DIMS(key, 4); @@ -630,60 +630,60 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( CHECK_DIMS(query_scale, 3); CHECK_DIMS(key_scale, 3); - const int batch_size = query.shape()[0]; - const int head_dim = query.shape()[3]; // 现在query是正常的128， 多出来的64在query_pe里面，所以这样做没什么问题 + const int batch_size = query.size(0); + const int head_dim = query.size(3); // 现在query是正常的128， 多出来的64在query_pe里面，所以这样做没什么问题 - int stride_bz_q = query.strides()[0]; - int stride_bz_k = key.strides()[0]; - int stride_bz_v = value.strides()[0]; - int stride_bz_o = output.strides()[0]; + int stride_bz_q = query.stride(0); + int stride_bz_k = key.stride(0); + int stride_bz_v = value.stride(0); + int stride_bz_o = output.stride(0); int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; - assert(value.shape()[0] == batch_size); + assert(value.size(0) == batch_size); if (tensor_layout == 0) { - qo_len = query.shape()[1]; - kv_len = key.shape()[1]; - num_qo_heads = query.shape()[2]; - num_kv_heads = key.shape()[2]; - - stride_seq_q = query.strides()[1]; - stride_h_q = query.strides()[2]; - stride_seq_k = key.strides()[1]; - stride_h_k = key.strides()[2]; - stride_h_v = value.strides()[2]; - stride_d_v = value.strides()[1]; - stride_seq_o = output.strides()[1]; - stride_h_o = output.strides()[2]; + qo_len = query.size(1); + kv_len = key.size(1); + num_qo_heads = query.size(2); + num_kv_heads = key.size(2); + + stride_seq_q = query.stride(1); + stride_h_q = query.stride(2); + stride_seq_k = key.stride(1); + stride_h_k = key.stride(2); + stride_h_v = value.stride(2); + stride_d_v = value.stride(1); + stride_seq_o = output.stride(1); + stride_h_o = output.stride(2); CHECK_SHAPE(key, batch_size, kv_len, num_kv_heads, head_dim); CHECK_SHAPE(output, batch_size, qo_len, num_qo_heads, head_dim); - assert(value.shape()[1] == head_dim); - assert(value.shape()[2] == num_kv_heads); + assert(value.size(1) == head_dim); + assert(value.size(2) == num_kv_heads); } else { - qo_len = query.shape()[2]; - kv_len = key.shape()[2]; - num_qo_heads = query.shape()[1]; - num_kv_heads = key.shape()[1]; - - stride_seq_q = query.strides()[2]; - stride_h_q = query.strides()[1]; - stride_seq_k = key.strides()[2]; - stride_h_k = key.strides()[1]; - stride_h_v = value.strides()[1]; - stride_d_v = value.strides()[2]; - stride_seq_o = output.strides()[2]; - stride_h_o = output.strides()[1]; + qo_len = query.size(2); + kv_len = key.size(2); + num_qo_heads = query.size(1); + num_kv_heads = key.size(1); + + stride_seq_q = query.stride(2); + stride_h_q = query.stride(1); + stride_seq_k = key.stride(2); + stride_h_k = key.stride(1); + stride_h_v = value.stride(1); + stride_d_v = value.stride(2); + stride_seq_o = output.stride(2); + stride_h_o = output.stride(1); CHECK_SHAPE(key, batch_size, num_kv_heads, kv_len, head_dim); CHECK_SHAPE(output, batch_size, num_qo_heads, qo_len, head_dim); - assert(value.shape()[2] == head_dim); - assert(value.shape()[1] == num_kv_heads); + assert(value.size(2) == head_dim); + assert(value.size(1) == num_kv_heads); } if (num_qo_heads % num_kv_heads != 0) { @@ -692,20 +692,20 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( throw std::invalid_argument(err_msg.str()); } - paddle::Tensor lse = paddle::empty({0}, paddle::DataType::FLOAT32); + torch::Tensor lse = torch::empty({0}); if (return_lse) { - lse = paddle::empty({batch_size, num_qo_heads, qo_len}, paddle::DataType::FLOAT32); + lse = torch::empty({batch_size, num_qo_heads, qo_len}, query.options().dtype(torch::kFloat32)); } const int num_kv_groups = num_qo_heads / num_kv_heads; - auto output_type = output.dtype(); + auto output_type = output.scalar_type(); DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { DISPATCH_CAUSAL(is_causal, IS_CAUSAL, { DISPATCH_QK_QUANT_GRAN(qk_quant_gran, QK_QUANT_GRAN, { - DISPATCH_PADDLE_DTYPE_TO_CTYPE_FP16(output_type, DTypeOut, { + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(output_type, DTypeOut, { constexpr int CTA_Q = 64; constexpr int CTA_K = 128; constexpr int NUM_THREADS = 128; @@ -713,7 +713,7 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( constexpr MaskMode mask_mode = IS_CAUSAL ? MaskMode::kCausal : MaskMode::kNone; - assert(value.shape()[3] >= div_ceil(kv_len, CTA_K) * CTA_K); + assert(value.size(3) >= div_ceil(kv_len, CTA_K) * CTA_K); if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp)) { @@ -730,12 +730,12 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( static_assert(QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp) || QK_QUANT_GRAN == static_cast(QuantGranularity::kPerThread), "Unsupported quantization granularity"); } - CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data()), batch_size, num_kv_heads, HEAD_DIM, value.shape()[3], stride_bz_v, stride_h_v, stride_d_v); + CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data_ptr()), batch_size, num_kv_heads, HEAD_DIM, value.size(3), stride_bz_v, stride_h_v, stride_d_v); auto* kernel = qk_int8_sv_f8_attn_dsk_kernel(QK_QUANT_GRAN), static_cast(QK_QUANT_GRAN), DTypeOut, mask_mode, false>; size_t sMemSize = CTA_Q * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t); @@ -751,10 +751,10 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( tma_map_Q_pe, tma_map_K_pe, tma_map_V, - reinterpret_cast(query_scale.data()), - reinterpret_cast(key_scale.data()), + reinterpret_cast(query_scale.data_ptr()), + reinterpret_cast(key_scale.data_ptr()), nullptr, - reinterpret_cast(output.data()), + reinterpret_cast(output.data_ptr()), stride_bz_o, stride_h_o, stride_seq_o, qo_len, kv_len, num_kv_groups, sm_scale); }); @@ -762,19 +762,19 @@ std::vector qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90_fwd( }); }); - return {lse}; + return lse; } -std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90_fwd( - paddle::Tensor& query, - paddle::Tensor& key, - paddle::Tensor& query_pe, - paddle::Tensor& key_pe, - paddle::Tensor& value, - paddle::Tensor& output, - paddle::Tensor& query_scale, - paddle::Tensor& key_scale, - paddle::Tensor& value_scale, +torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( + torch::Tensor query, + torch::Tensor key, + torch::Tensor query_pe, + torch::Tensor key_pe, + torch::Tensor value, + torch::Tensor output, + torch::Tensor query_scale, + torch::Tensor key_scale, + torch::Tensor value_scale, int tensor_layout, int is_causal, int qk_quant_gran, @@ -801,14 +801,14 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d CHECK_CONTIGUOUS(key_scale); CHECK_CONTIGUOUS(value_scale); - CHECK_DTYPE(query, paddle::DataType::INT8); - CHECK_DTYPE(key, paddle::DataType::INT8); - CHECK_DTYPE(query_pe, paddle::DataType::INT8); - CHECK_DTYPE(key_pe, paddle::DataType::INT8); - CHECK_DTYPE(value, paddle::DataType::FLOAT8_E4M3FN); - CHECK_DTYPE(query_scale, paddle::DataType::FLOAT32); - CHECK_DTYPE(key_scale, paddle::DataType::FLOAT32); - CHECK_DTYPE(value_scale, paddle::DataType::FLOAT32); + CHECK_DTYPE(query, torch::kInt8); + CHECK_DTYPE(key, torch::kInt8); + CHECK_DTYPE(query_pe, torch::kInt8); + CHECK_DTYPE(key_pe, torch::kInt8); + CHECK_DTYPE(value, at::ScalarType::Float8_e4m3fn); + CHECK_DTYPE(query_scale, torch::kFloat32); + CHECK_DTYPE(key_scale, torch::kFloat32); + CHECK_DTYPE(value_scale, torch::kFloat32); CHECK_DIMS(query, 4); CHECK_DIMS(key, 4); @@ -820,61 +820,61 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d CHECK_DIMS(key_scale, 3); CHECK_DIMS(value_scale, 3); - const int batch_size = query.shape()[0]; - const int head_dim = query.shape()[3]; + const int batch_size = query.size(0); + const int head_dim = query.size(3); - int stride_bz_q = query.strides()[0]; - int stride_bz_k = key.strides()[0]; - int stride_bz_v = value.strides()[0]; - int stride_bz_o = output.strides()[0]; + int stride_bz_q = query.stride(0); + int stride_bz_k = key.stride(0); + int stride_bz_v = value.stride(0); + int stride_bz_o = output.stride(0); int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; - assert(value.shape()[0] == batch_size); + assert(value.size(0) == batch_size); if (tensor_layout == 0) { - qo_len = query.shape()[1]; - kv_len = key.shape()[1]; - num_qo_heads = query.shape()[2]; - num_kv_heads = key.shape()[2]; - - stride_seq_q = query.strides()[1]; - stride_h_q = query.strides()[2]; - stride_seq_k = key.strides()[1]; - stride_h_k = key.strides()[2]; - stride_h_v = value.strides()[2]; - stride_d_v = value.strides()[1]; - stride_seq_o = output.strides()[1]; - stride_h_o = output.strides()[2]; + qo_len = query.size(1); + kv_len = key.size(1); + num_qo_heads = query.size(2); + num_kv_heads = key.size(2); + + stride_seq_q = query.stride(1); + stride_h_q = query.stride(2); + stride_seq_k = key.stride(1); + stride_h_k = key.stride(2); + stride_h_v = value.stride(2); + stride_d_v = value.stride(1); + stride_seq_o = output.stride(1); + stride_h_o = output.stride(2); CHECK_SHAPE(key, batch_size, kv_len, num_kv_heads, head_dim); CHECK_SHAPE(output, batch_size, qo_len, num_qo_heads, head_dim); - assert(value.shape()[1] == head_dim); - assert(value.shape()[2] == num_kv_heads); + assert(value.size(1) == head_dim); + assert(value.size(2) == num_kv_heads); } else { - qo_len = query.shape()[2]; - kv_len = key.shape()[2]; - num_qo_heads = query.shape()[1]; - num_kv_heads = key.shape()[1]; - - stride_seq_q = query.strides()[2]; - stride_h_q = query.strides()[1]; - stride_seq_k = key.strides()[2]; - stride_h_k = key.strides()[1]; - stride_h_v = value.strides()[1]; - stride_d_v = value.strides()[2]; - stride_seq_o = output.strides()[2]; - stride_h_o = output.strides()[1]; + qo_len = query.size(2); + kv_len = key.size(2); + num_qo_heads = query.size(1); + num_kv_heads = key.size(1); + + stride_seq_q = query.stride(2); + stride_h_q = query.stride(1); + stride_seq_k = key.stride(2); + stride_h_k = key.stride(1); + stride_h_v = value.stride(1); + stride_d_v = value.stride(2); + stride_seq_o = output.stride(2); + stride_h_o = output.stride(1); CHECK_SHAPE(key, batch_size, num_kv_heads, kv_len, head_dim); CHECK_SHAPE(output, batch_size, num_qo_heads, qo_len, head_dim); - assert(value.shape()[2] == head_dim); - assert(value.shape()[1] == num_kv_heads); + assert(value.size(2) == head_dim); + assert(value.size(1) == num_kv_heads); } if (num_qo_heads % num_kv_heads != 0) { @@ -883,20 +883,20 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d throw std::invalid_argument(err_msg.str()); } - paddle::Tensor lse = paddle::empty({1}, paddle::DataType::FLOAT32); + torch::Tensor lse = torch::empty({1}); if (return_lse) { - lse = paddle::empty({batch_size, num_qo_heads, qo_len}, paddle::DataType::FLOAT32); + lse = torch::empty({batch_size, num_qo_heads, qo_len}, query.options().dtype(torch::kFloat32)); } const int num_kv_groups = num_qo_heads / num_kv_heads; - auto output_dtype = output.dtype(); + auto output_dtype = output.scalar_type(); DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { DISPATCH_CAUSAL(is_causal, IS_CAUSAL, { DISPATCH_QK_QUANT_GRAN(qk_quant_gran, QK_QUANT_GRAN, { - DISPATCH_PADDLE_DTYPE_TO_CTYPE_FP16(output_dtype, DTypeOut, { + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(output_dtype, DTypeOut, { constexpr int CTA_Q = 64; constexpr int CTA_K = 128; constexpr int NUM_THREADS = 128; @@ -904,7 +904,7 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d constexpr MaskMode mask_mode = IS_CAUSAL ? MaskMode::kCausal : MaskMode::kNone; - assert(value.shape()[3] >= div_ceil(kv_len, CTA_K) * CTA_K); + assert(value.size(3) >= div_ceil(kv_len, CTA_K) * CTA_K); if constexpr (QK_QUANT_GRAN == static_cast(QuantGranularity::kPerWarp)) { @@ -922,12 +922,12 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d } CHECK_SHAPE(value_scale, batch_size, num_kv_heads, HEAD_DIM); - CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data()), batch_size, num_kv_heads, HEAD_DIM, value.shape()[3], stride_bz_v, stride_h_v, stride_d_v); + CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data_ptr()), batch_size, num_kv_heads, HEAD_DIM, value.size(3), stride_bz_v, stride_h_v, stride_d_v); auto* kernel = qk_int8_sv_f8_attn_dsk_kernel(QK_QUANT_GRAN), static_cast(QK_QUANT_GRAN), DTypeOut, mask_mode, true>; size_t sMemSize = CTA_Q * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t); @@ -943,10 +943,10 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d tma_map_Q_pe, tma_map_K_pe, tma_map_V, - reinterpret_cast(query_scale.data()), - reinterpret_cast(key_scale.data()), - reinterpret_cast(value_scale.data()), - reinterpret_cast(output.data()), + reinterpret_cast(query_scale.data_ptr()), + reinterpret_cast(key_scale.data_ptr()), + reinterpret_cast(value_scale.data_ptr()), + reinterpret_cast(output.data_ptr()), stride_bz_o, stride_h_o, stride_seq_o, qo_len, kv_len, num_kv_groups, sm_scale); }); @@ -954,5 +954,5 @@ std::vector qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_d }); }); - return {lse}; + return lse; } \ No newline at end of file diff --git a/patch_test/test_dsk.py b/patch_test/test_dsk.py new file mode 100644 index 00000000..05c3f8f6 --- /dev/null +++ b/patch_test/test_dsk.py @@ -0,0 +1,66 @@ +from sageattention import sageattn_qk_int8_pv_fp8_cuda_dsk_sm90 as sageattn_qk_int8_pv_fp8_cuda_sm90a +from torch.nn.functional import scaled_dot_product_attention as sdpa + +import torch +import numpy as np +import os +import argparse + + +def precision_cmp_torch(t1: torch.Tensor, t2: torch.Tensor): + x, xx = t1.to(dtype=torch.float32), t2.to(dtype=torch.float32) + # 重塑张量并计算余弦相似度 + x_reshaped = torch.reshape(x, [1, -1]) + xx_reshaped = torch.reshape(xx, [1, -1]) + sim = torch.nn.functional.cosine_similarity(x_reshaped, xx_reshaped).item() + + # 计算 L1 误差 + l1 = (torch.abs(x - xx).sum() / torch.abs(xx).sum()).item() + + max_diff = torch.max(x - xx) + + return sim, l1, max_diff + +bsz = 2 +seq_len = 1026 +num_heads = 128 +head_dim_qk = 128 + 64 +head_dim_v = 128 + +tensor_layout = "NHD" +is_causal = True +return_lse = False + +torch.backends.cuda.enable_flash_sdp(True) + +q = torch.randn((bsz, seq_len, num_heads, head_dim_qk), dtype=torch.float16).cuda() +k = torch.randn((bsz, seq_len, num_heads, head_dim_qk), dtype=torch.float16).cuda() +v = torch.randn((bsz, seq_len, num_heads, head_dim_v), dtype=torch.float16).cuda() + +# to HND for sdpa input +q = q.transpose(2, 1) +k = k.transpose(2, 1) +v = v.transpose(2, 1) + +o_torch_fa2 = sdpa(q, k, v, is_causal=is_causal) +torch.cuda.synchronize() + +torch.backends.cuda.enable_flash_sdp(False) +o_torch_sdpa = sdpa(q, k, v, is_causal=is_causal) +torch.cuda.synchronize() + +head_dim_og = head_dim_qk +sm_scale = head_dim_og**-0.5 + +# to NHD for sageattn input +q = q.transpose(2, 1) +k = k.transpose(2, 1) +v = v.transpose(2, 1) +o_sa = sageattn_qk_int8_pv_fp8_cuda_sm90a(q, k, v, tensor_layout=tensor_layout, is_causal=is_causal, qk_quant_gran="per_warp", return_lse=return_lse, pv_accum_dtype="fp32+fp32") +torch.cuda.synchronize() + +sim, l1, max_diff = precision_cmp_torch(o_torch_fa2.transpose(2, 1), o_sa) +print(f"Sim and Diff of Sage Attn: {sim}, {max_diff}") + +sim, l1, max_diff = precision_cmp_torch(o_torch_fa2, o_torch_sdpa) +print(f"Sim and Diff of Flash Attn: {sim}, {max_diff}") diff --git a/sageattention/__init__.py b/sageattention/__init__.py index 73b0256d..5ca940b0 100644 --- a/sageattention/__init__.py +++ b/sageattention/__init__.py @@ -2,4 +2,4 @@ from .core import sageattn_qk_int8_pv_fp16_triton from .core import sageattn_qk_int8_pv_fp16_cuda from .core import sageattn_qk_int8_pv_fp8_cuda -from .core import sageattn_qk_int8_pv_fp8_cuda_sm90 \ No newline at end of file +from .core import sageattn_qk_int8_pv_fp8_cuda_sm90, sageattn_qk_int8_pv_fp8_cuda_dsk_sm90 \ No newline at end of file diff --git a/sageattention/core.py b/sageattention/core.py index 7d464ac8..25cce1fc 100644 --- a/sageattention/core.py +++ b/sageattention/core.py @@ -898,6 +898,93 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90( o = o[..., :head_dim_og] + if return_lse: + return o, lse / 1.44269504 + lse_correction * sm_scale if smooth_k else lse / 1.44269504 + else: + return o + + +def sageattn_qk_int8_pv_fp8_cuda_dsk_sm90( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + tensor_layout: str = "HND", + is_causal: bool = False, + qk_quant_gran: str = "per_warp", + sm_scale: Optional[float] = None, + pv_accum_dtype: str = "fp32+fp32", + smooth_k: bool = True, + return_lse: bool = False, + **kwargs: Any, +) -> torch.Tensor: + dtype = q.dtype + assert dtype in [torch.float16, torch.bfloat16], "Input tensors must be in dtype of torch.float16 or torch.bfloat16" + assert qk_quant_gran in ["per_warp", "per_thread"], "qk_quant_gran must be either 'per_warp' or 'per_thread'." + assert q.dtype == k.dtype == v.dtype, "All tensors must have the same dtype." + + torch.cuda.set_device(v.device) + + _tensor_layout = 0 if tensor_layout == "NHD" else 1 + _is_caual = 1 if is_causal else 0 + _qk_quant_gran = 3 if qk_quant_gran == "per_thread" else 2 + _return_lse = 1 if return_lse else 0 + + head_dim_og = q.size(-1) + + if head_dim_og < 64: + q = torch.nn.functional.pad(q, (0, 64 - head_dim_og)) + k = torch.nn.functional.pad(k, (0, 64 - head_dim_og)) + v = torch.nn.functional.pad(v, (0, 64 - head_dim_og)) + elif head_dim_og > 64 and head_dim_og < 128: + q = torch.nn.functional.pad(q, (0, 128 - head_dim_og)) + k = torch.nn.functional.pad(k, (0, 128 - head_dim_og)) + v = torch.nn.functional.pad(v, (0, 128 - head_dim_og)) + elif head_dim_og > 128 and head_dim_og < 192: + q = torch.nn.functional.pad(q, (0, 192 - head_dim_og)) + k = torch.nn.functional.pad(k, (0, 192 - head_dim_og)) + # v = torch.nn.functional.pad(v, (0, 128 - head_dim_og)) + elif head_dim_og > 192: + raise ValueError(f"Unsupported head_dim: {head_dim_og}") + + assert q.stride(-1) == 1 and k.stride(-1) == 1 and v.stride(-1) == 1, "Last dim of qkv must be contiguous." + + if sm_scale is None: + sm_scale = head_dim_og**-0.5 + + seq_dim = 1 if _tensor_layout == 0 else 2 + + if smooth_k: + km = k.mean(dim=seq_dim, keepdim=True) + if return_lse: + if tensor_layout == "NHD": + lse_correction = torch.matmul(q.transpose(1, 2), km.transpose(1, 2).transpose(2, 3)).squeeze(-1).to(torch.float32) + else: + lse_correction = torch.matmul(q, km.transpose(2, 3)).squeeze(-1).to(torch.float32) + else: + km = None + + if qk_quant_gran == "per_warp": + q_int8, q_scale, k_int8, k_scale = per_warp_int8_cuda(q, k, km, tensor_layout=tensor_layout, BLKQ=64, WARPQ=16, BLKK=128) + + o = torch.empty(v.size(), dtype=dtype, device=q.device) + + kv_len = k.size(seq_dim) + v_pad_len = 128 - (kv_len % 128) if kv_len % 128 != 0 else 0 + if v_pad_len > 0: + if tensor_layout == "HND": + v = torch.cat([v, torch.zeros(v.size(0), v.size(1), v_pad_len, v.size(3), dtype=v.dtype, device=v.device)], dim=2) + else: + v = torch.cat([v, torch.zeros(v.size(0), v_pad_len, v.size(2), v.size(3), dtype=v.dtype, device=v.device)], dim=1) + + v_fp8, v_scale, _ = per_channel_fp8(v, tensor_layout=tensor_layout, smooth_v=False) + q_int8_nope, q_int8_pe = torch.split(q_int8, [128, 64], dim=-1) + k_int8_nope, k_int8_pe = torch.split(k_int8, [128, 64], dim=-1) + + lse = _qattn_sm90.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90(q_int8_nope, k_int8_nope, q_int8_pe, k_int8_pe, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + + head_dim_og = v.shape[-1] + o = o[..., :head_dim_og] + if return_lse: return o, lse / 1.44269504 + lse_correction * sm_scale if smooth_k else lse / 1.44269504 else: diff --git a/setup.py b/setup.py index 06a75eb9..662af0d4 100644 --- a/setup.py +++ b/setup.py @@ -152,6 +152,7 @@ def get_nvcc_cuda_version(cuda_dir: str) -> Version: sources=[ "csrc/qattn/pybind_sm90.cpp", "csrc/qattn/qk_int_sv_f8_cuda_sm90.cu", + "csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu", ], extra_compile_args={ "cxx": CXX_FLAGS, From f5f8ed1a45f3e94edefaa0afbddca17bd0498e16 Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Fri, 21 Feb 2025 00:07:31 +0800 Subject: [PATCH 4/6] debug --- csrc/dispatch_utils.h | 4 ++-- patch_test/test_dsk.py | 4 ++-- sageattention/core.py | 17 ++++++++--------- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index fc41af7e..17fc199f 100644 --- a/csrc/dispatch_utils.h +++ b/csrc/dispatch_utils.h @@ -40,8 +40,8 @@ } else if (head_dim == 128) { \ constexpr int HEAD_DIM = 128; \ __VA_ARGS__ \ - } else if (head_dim == 192) { \ - constexpr int HEAD_DIM = 192; \ + } else if (head_dim == 256) { \ + constexpr int HEAD_DIM = 256; \ __VA_ARGS__ \ } else { \ std::ostringstream err_msg; \ diff --git a/patch_test/test_dsk.py b/patch_test/test_dsk.py index 05c3f8f6..31099eb2 100644 --- a/patch_test/test_dsk.py +++ b/patch_test/test_dsk.py @@ -60,7 +60,7 @@ def precision_cmp_torch(t1: torch.Tensor, t2: torch.Tensor): torch.cuda.synchronize() sim, l1, max_diff = precision_cmp_torch(o_torch_fa2.transpose(2, 1), o_sa) -print(f"Sim and Diff of Sage Attn: {sim}, {max_diff}") +print(f"Sim and Diff of Sage Attn & torch SDPA: {sim}, {max_diff}") sim, l1, max_diff = precision_cmp_torch(o_torch_fa2, o_torch_sdpa) -print(f"Sim and Diff of Flash Attn: {sim}, {max_diff}") +print(f"Sim and Diff of Flash Attn & torch SDPA: {sim}, {max_diff}") diff --git a/sageattention/core.py b/sageattention/core.py index 25cce1fc..49a2e92c 100644 --- a/sageattention/core.py +++ b/sageattention/core.py @@ -925,7 +925,7 @@ def sageattn_qk_int8_pv_fp8_cuda_dsk_sm90( torch.cuda.set_device(v.device) _tensor_layout = 0 if tensor_layout == "NHD" else 1 - _is_caual = 1 if is_causal else 0 + _is_causal = 1 if is_causal else 0 _qk_quant_gran = 3 if qk_quant_gran == "per_thread" else 2 _return_lse = 1 if return_lse else 0 @@ -939,11 +939,10 @@ def sageattn_qk_int8_pv_fp8_cuda_dsk_sm90( q = torch.nn.functional.pad(q, (0, 128 - head_dim_og)) k = torch.nn.functional.pad(k, (0, 128 - head_dim_og)) v = torch.nn.functional.pad(v, (0, 128 - head_dim_og)) - elif head_dim_og > 128 and head_dim_og < 192: - q = torch.nn.functional.pad(q, (0, 192 - head_dim_og)) - k = torch.nn.functional.pad(k, (0, 192 - head_dim_og)) - # v = torch.nn.functional.pad(v, (0, 128 - head_dim_og)) - elif head_dim_og > 192: + elif head_dim_og > 128 and head_dim_og < 256: + q = torch.nn.functional.pad(q, (0, 256 - head_dim_og)) + k = torch.nn.functional.pad(k, (0, 256 - head_dim_og)) + elif head_dim_og > 256: raise ValueError(f"Unsupported head_dim: {head_dim_og}") assert q.stride(-1) == 1 and k.stride(-1) == 1 and v.stride(-1) == 1, "Last dim of qkv must be contiguous." @@ -977,10 +976,10 @@ def sageattn_qk_int8_pv_fp8_cuda_dsk_sm90( v = torch.cat([v, torch.zeros(v.size(0), v_pad_len, v.size(2), v.size(3), dtype=v.dtype, device=v.device)], dim=1) v_fp8, v_scale, _ = per_channel_fp8(v, tensor_layout=tensor_layout, smooth_v=False) - q_int8_nope, q_int8_pe = torch.split(q_int8, [128, 64], dim=-1) - k_int8_nope, k_int8_pe = torch.split(k_int8, [128, 64], dim=-1) + q_int8_nope, q_int8_pe, _ = torch.split(q_int8, [128, 64, 64], dim=-1) + k_int8_nope, k_int8_pe, _ = torch.split(k_int8, [128, 64, 64], dim=-1) - lse = _qattn_sm90.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90(q_int8_nope, k_int8_nope, q_int8_pe, k_int8_pe, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_caual, _qk_quant_gran, sm_scale, _return_lse) + lse = _qattn_sm90.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90(q_int8_nope, k_int8_nope, q_int8_pe, k_int8_pe, v_fp8, o, q_scale, k_scale, v_scale, _tensor_layout, _is_causal, _qk_quant_gran, sm_scale, _return_lse) head_dim_og = v.shape[-1] o = o[..., :head_dim_og] From 825eea8615dc62be587ebef84ee0a3d350ab3594 Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Mon, 3 Mar 2025 22:15:32 +0800 Subject: [PATCH 5/6] bug fix for stride of q_pe and k_pe --- csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu | 61 +++++++++++++++--------- patch_test/test_dsk.py | 3 ++ sageattention/core.py | 3 +- 3 files changed, 44 insertions(+), 23 deletions(-) diff --git a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu index b3809d21..347b0244 100644 --- a/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu +++ b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu @@ -146,19 +146,12 @@ __global__ void qk_int8_sv_f8_attn_dsk_kernel(const __grid_constant__ CUtensorMa extern __shared__ __align__(128) int8_t smem_[]; - /* // original: - * int8_t *sQ = (int8_t*)smem_; - * int8_t *sK = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t)); - * int8_t *sV = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t)); - * half *sO = (half*)smem_; - */ - - int8_t *sQ = (int8_t*)smem_; // 0 - int8_t *sQ_pe = (int8_t*)(smem_ + CTA_Q * (head_dim) * sizeof(int8_t)); // 0 + head_dim - - int8_t *sK = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t)); // 0 + head_dim + pe - int8_t *sK_pe = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t) + CTA_K * (head_dim) * sizeof(int8_t)); // 0 + head_dim + pe + head_dim - int8_t *sV = (int8_t*)(smem_ + CTA_Q * (head_dim + head_dim_pe) * sizeof(int8_t) + CTA_K * (head_dim + head_dim_pe) * sizeof(int8_t)); + int8_t *sQ = (int8_t*)smem_; + int8_t *sK = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t)); + int8_t *sV = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t)); + int8_t *sQ_pe = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t)); + int8_t *sK_pe = (int8_t*)(smem_ + CTA_Q * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t) + CTA_K * head_dim * sizeof(int8_t) + CTA_Q * head_dim_pe * sizeof(int8_t)); + half *sO = (half*)smem_; int32_t RS[num_tiles_q][num_tiles_k][8]; @@ -250,10 +243,10 @@ __global__ void qk_int8_sv_f8_attn_dsk_kernel(const __grid_constant__ CUtensorMa expect_bytes<(CTA_K * (head_dim)) * sizeof(int8_t)>(&barrier_V); load_async_4D(sQ, &tensorMapQ, &barrier_Q, 0, bx * CTA_Q, head_id, batch_id); - load_async_4D(sQ_pe, &tensorMapQ_pe, &barrier_Q_pe, 0, bx * CTA_Q, head_id, batch_id); load_async_4D(sK, &tensorMapK, &barrier_K, 0, 0, kv_head_id, batch_id); - load_async_4D(sK_pe, &tensorMapK_pe, &barrier_K_pe, 0, 0, kv_head_id, batch_id); load_async_4D(sV, &tensorMapV, &barrier_V, 0, 0, kv_head_id, batch_id); + load_async_4D(sQ_pe, &tensorMapQ_pe, &barrier_Q_pe, 0, bx * CTA_Q, head_id, batch_id); + load_async_4D(sK_pe, &tensorMapK_pe, &barrier_K_pe, 0, 0, kv_head_id, batch_id); } float q_scale = Q_scale[q_scale_idx]; @@ -634,12 +627,16 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( const int head_dim = query.size(3); // 现在query是正常的128， 多出来的64在query_pe里面，所以这样做没什么问题 int stride_bz_q = query.stride(0); + int stride_bz_q_pe = query_pe.stride(0); int stride_bz_k = key.stride(0); + int stride_bz_k_pe = key_pe.stride(0); int stride_bz_v = value.stride(0); int stride_bz_o = output.stride(0); int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; - int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; + int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, + stride_seq_q_pe, stride_h_q_pe, stride_seq_k_pe, stride_h_k_pe, + stride_h_v, stride_d_v, stride_seq_o, stride_h_o; assert(value.size(0) == batch_size); @@ -652,8 +649,12 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( stride_seq_q = query.stride(1); stride_h_q = query.stride(2); + stride_seq_q_pe = query_pe.stride(1); + stride_h_q_pe = query_pe.stride(2); stride_seq_k = key.stride(1); stride_h_k = key.stride(2); + stride_seq_k_pe = key_pe.stride(1); + stride_h_k_pe = key_pe.stride(2); stride_h_v = value.stride(2); stride_d_v = value.stride(1); stride_seq_o = output.stride(1); @@ -673,8 +674,12 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( stride_seq_q = query.stride(2); stride_h_q = query.stride(1); + stride_seq_q_pe = query_pe.stride(2); + stride_h_q_pe = query_pe.stride(1); stride_seq_k = key.stride(2); stride_h_k = key.stride(1); + stride_seq_k_pe = key_pe.stride(2); + stride_h_k_pe = key_pe.stride(1); stride_h_v = value.stride(1); stride_d_v = value.stride(2); stride_seq_o = output.stride(2); @@ -732,8 +737,8 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90( CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q_pe, stride_h_q_pe, stride_seq_q_pe); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k_pe, stride_h_k_pe, stride_seq_k_pe); CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data_ptr()), batch_size, num_kv_heads, HEAD_DIM, value.size(3), stride_bz_v, stride_h_v, stride_d_v); @@ -824,12 +829,16 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( const int head_dim = query.size(3); int stride_bz_q = query.stride(0); + int stride_bz_q_pe = query_pe.stride(0); int stride_bz_k = key.stride(0); + int stride_bz_k_pe = key_pe.stride(0); int stride_bz_v = value.stride(0); int stride_bz_o = output.stride(0); int qo_len, kv_len, padded_kv_len, num_qo_heads, num_kv_heads; - int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, stride_h_v, stride_d_v, stride_seq_o, stride_h_o; + int stride_seq_q, stride_h_q, stride_seq_k, stride_h_k, + stride_seq_q_pe, stride_h_q_pe, stride_seq_k_pe, stride_h_k_pe, + stride_h_v, stride_d_v, stride_seq_o, stride_h_o; assert(value.size(0) == batch_size); @@ -842,8 +851,12 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( stride_seq_q = query.stride(1); stride_h_q = query.stride(2); + stride_seq_q_pe = query_pe.stride(1); + stride_h_q_pe = query_pe.stride(2); stride_seq_k = key.stride(1); stride_h_k = key.stride(2); + stride_seq_k_pe = key_pe.stride(1); + stride_h_k_pe = key_pe.stride(2); stride_h_v = value.stride(2); stride_d_v = value.stride(1); stride_seq_o = output.stride(1); @@ -864,8 +877,12 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( stride_seq_q = query.stride(2); stride_h_q = query.stride(1); + stride_seq_q_pe = query_pe.stride(2); + stride_h_q_pe = query_pe.stride(1); stride_seq_k = key.stride(2); stride_h_k = key.stride(1); + stride_seq_k_pe = key_pe.stride(2); + stride_h_k_pe = key_pe.stride(1); stride_h_v = value.stride(1); stride_d_v = value.stride(2); stride_seq_o = output.stride(2); @@ -923,9 +940,9 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( CHECK_SHAPE(value_scale, batch_size, num_kv_heads, HEAD_DIM); CUtensorMap tma_map_Q = create_tensor_map_4D(reinterpret_cast(query.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM, stride_bz_q, stride_h_q, stride_seq_q); - CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q, stride_h_q, stride_seq_q); + CUtensorMap tma_map_Q_pe = create_tensor_map_4D(reinterpret_cast(query_pe.data_ptr()), batch_size, num_qo_heads, qo_len, HEAD_DIM_PE, stride_bz_q_pe, stride_h_q_pe, stride_seq_q_pe); CUtensorMap tma_map_K = create_tensor_map_4D(reinterpret_cast(key.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM, stride_bz_k, stride_h_k, stride_seq_k); - CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k, stride_h_k, stride_seq_k); + CUtensorMap tma_map_K_pe = create_tensor_map_4D(reinterpret_cast(key_pe.data_ptr()), batch_size, num_kv_heads, kv_len, HEAD_DIM_PE, stride_bz_k_pe, stride_h_k_pe, stride_seq_k_pe); CUtensorMap tma_map_V = create_tensor_map_4D(reinterpret_cast(value.data_ptr()), batch_size, num_kv_heads, HEAD_DIM, value.size(3), stride_bz_v, stride_h_v, stride_d_v); @@ -935,7 +952,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf_dsk_sm90( cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize); - + dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size); kernel<<>>( tma_map_Q, diff --git a/patch_test/test_dsk.py b/patch_test/test_dsk.py index 31099eb2..6f1c192d 100644 --- a/patch_test/test_dsk.py +++ b/patch_test/test_dsk.py @@ -62,5 +62,8 @@ def precision_cmp_torch(t1: torch.Tensor, t2: torch.Tensor): sim, l1, max_diff = precision_cmp_torch(o_torch_fa2.transpose(2, 1), o_sa) print(f"Sim and Diff of Sage Attn & torch SDPA: {sim}, {max_diff}") +sim, l1, max_diff = precision_cmp_torch(o_torch_sdpa.transpose(2, 1), o_sa) +print(f"Sim and Diff of Sage Attn & Flash Attn: {sim}, {max_diff}") + sim, l1, max_diff = precision_cmp_torch(o_torch_fa2, o_torch_sdpa) print(f"Sim and Diff of Flash Attn & torch SDPA: {sim}, {max_diff}") diff --git a/sageattention/core.py b/sageattention/core.py index 49a2e92c..a39a8a4d 100644 --- a/sageattention/core.py +++ b/sageattention/core.py @@ -987,4 +987,5 @@ def sageattn_qk_int8_pv_fp8_cuda_dsk_sm90( if return_lse: return o, lse / 1.44269504 + lse_correction * sm_scale if smooth_k else lse / 1.44269504 else: - return o \ No newline at end of file + return o + From b580f8cae78d78a8efa3a904411b208f6986677a Mon Sep 17 00:00:00 2001 From: l1cacheDell Date: Tue, 4 Mar 2025 11:18:24 +0800 Subject: [PATCH 6/6] offer a benchmark code & remove patch_test --- bench/bench_qk_int8_pv_fp8_cuda_dsk_sm90.py | 89 +++++++++++++++++++++ patch_test/test_dsk.py | 69 ---------------- 2 files changed, 89 insertions(+), 69 deletions(-) create mode 100644 bench/bench_qk_int8_pv_fp8_cuda_dsk_sm90.py delete mode 100644 patch_test/test_dsk.py diff --git a/bench/bench_qk_int8_pv_fp8_cuda_dsk_sm90.py b/bench/bench_qk_int8_pv_fp8_cuda_dsk_sm90.py new file mode 100644 index 00000000..91e4818e --- /dev/null +++ b/bench/bench_qk_int8_pv_fp8_cuda_dsk_sm90.py @@ -0,0 +1,89 @@ +import torch +from flash_attn.utils.benchmark import benchmark_forward +import sageattention._qattn_sm90 as qattn +import argparse + +parser = argparse.ArgumentParser(description='Benchmark QK Int8 PV FP8 SM90') +parser.add_argument('--batch_size', type=int, default=4, help='Batch size') +parser.add_argument('--num_heads', type=int, default=32, help='Number of heads') +# parser.add_argument('--head_dim_qk', type=int, default=192, help='Head dimension') +# parser.add_argument('--head_dim_v', type=int, default=128, help='Head dimension') +parser.add_argument('--pv_accum_dtype', type=str, default='fp32+fp32', choices=['fp32', 'fp32+fp32']) +parser.add_argument('--quant_gran', type=str, default='per_warp', choices=['per_warp', 'per_thread'], help='Quantization granularity') +args = parser.parse_args() + +head = args.num_heads +batch = args.batch_size +# head_dim_qk = args.head_dim_qk +# head_dim_v = args.head_dim_v +head_dim_qk = 192 +head_dim_v = 128 + +assert args.pv_accum_dtype == 'fp32+fp32', "pure fp32 accumulator is not supported for now" + +print(f"CUDA QK Int8 PV FP8 SM90 Mixed HeadDim=192 Benchmark") +print(f"batch: {batch}, head: {head}, head_dim_qk: {head_dim_qk}, head_dim_v: {head_dim_v}") + +WARP_Q = 32 +WARP_K = 64 + +kernel = qattn.qk_int8_sv_f8_accum_f32_attn_inst_buf_dsk_sm90 + +_qk_quant_gran = 3 if args.quant_gran == 'per_thread' else 2 + +is_causal = False +_is_causal = 1 if is_causal else 0 +print(f"is_causal: {is_causal}") +for seq_len in {1024, 2048, 4096, 8192, 16384, 32768}: + flops = 2 * head * batch * (head_dim_qk + head_dim_v) * seq_len * seq_len / (2 if is_causal else 1) + + q = torch.randint(-95, 95, (batch, head, seq_len, head_dim_qk), dtype=torch.int8, device="cuda") + k = torch.randint(-95, 95, (batch, head, seq_len, head_dim_qk), dtype=torch.int8, device="cuda") + q_nope, q_pe = torch.split(q, [128, 64], dim=-1) + k_nope, k_pe = torch.split(k, [128, 64], dim=-1) + o = torch.empty(batch, head, seq_len, head_dim_v, dtype=torch.float16, device="cuda") + + v_scale = torch.randn(batch, head, head_dim_v, dtype=torch.float, device="cuda") + + if args.quant_gran == 'per_warp': + q_scale = torch.randn(batch, head, seq_len // 64 * 4, dtype=torch.float, device="cuda") + k_scale = torch.randn(batch, head, seq_len // 128, dtype=torch.float, device="cuda") + elif args.quant_gran == 'per_thread': + q_scale = torch.randn(batch, head, seq_len // 64 * 4 * 8, dtype=torch.float, device="cuda") + k_scale = torch.randn(batch, head, seq_len // 128 * 4, dtype=torch.float, device="cuda") + + v = torch.randn(batch, head, head_dim_v, seq_len, dtype=torch.float16, device="cuda").to(torch.float8_e4m3fn) + sm_scale = 1 / (head_dim_qk ** 0.5) + for i in range(5): kernel(q_nope, k_nope, q_pe, k_pe, v, o, q_scale, k_scale, 1, _is_causal, _qk_quant_gran, sm_scale, 0) + torch.cuda.synchronize() + _, time = benchmark_forward(kernel, q_nope, k_nope, q_pe, k_pe, v, o, q_scale, k_scale, 1, _is_causal, _qk_quant_gran, sm_scale, 0, repeats=100, verbose=False, desc='Triton') + print(f'{seq_len} flops:{flops/time.mean*1e-12}') + + +is_causal = True +_is_causal = 1 if is_causal else 0 +print(f"is_causal: {is_causal}") +for seq_len in {1024, 2048, 4096, 8192, 16384, 32768}: + flops = 2 * head * batch * (head_dim_qk + head_dim_v) * seq_len * seq_len / (2 if is_causal else 1) + + q = torch.randint(-95, 95, (batch, head, seq_len, head_dim_qk), dtype=torch.int8, device="cuda") + k = torch.randint(-95, 95, (batch, head, seq_len, head_dim_qk), dtype=torch.int8, device="cuda") + q_nope, q_pe = torch.split(q, [128, 64], dim=-1) + k_nope, k_pe = torch.split(k, [128, 64], dim=-1) + o = torch.empty(batch, head, seq_len, head_dim_v, dtype=torch.float16, device="cuda") + + v_scale = torch.randn(batch, head, head_dim_v, dtype=torch.float, device="cuda") + + if args.quant_gran == 'per_warp': + q_scale = torch.randn(batch, head, seq_len // 64 * 4, dtype=torch.float, device="cuda") + k_scale = torch.randn(batch, head, seq_len // 128, dtype=torch.float, device="cuda") + elif args.quant_gran == 'per_thread': + q_scale = torch.randn(batch, head, seq_len // 64 * 4 * 8, dtype=torch.float, device="cuda") + k_scale = torch.randn(batch, head, seq_len // 128 * 4, dtype=torch.float, device="cuda") + + v = torch.randn(batch, head, head_dim_v, seq_len, dtype=torch.float16, device="cuda").to(torch.float8_e4m3fn) + sm_scale = 1 / (head_dim_qk ** 0.5) + for i in range(5): kernel(q_nope, k_nope, q_pe, k_pe, v, o, q_scale, k_scale, 1, _is_causal, _qk_quant_gran, sm_scale, 0) + torch.cuda.synchronize() + _, time = benchmark_forward(kernel, q_nope, k_nope, q_pe, k_pe, v, o, q_scale, k_scale, 1, _is_causal, _qk_quant_gran, sm_scale, 0, repeats=100, verbose=False, desc='Triton') + print(f'{seq_len} flops:{flops/time.mean*1e-12}') \ No newline at end of file diff --git a/patch_test/test_dsk.py b/patch_test/test_dsk.py deleted file mode 100644 index 6f1c192d..00000000 --- a/patch_test/test_dsk.py +++ /dev/null @@ -1,69 +0,0 @@ -from sageattention import sageattn_qk_int8_pv_fp8_cuda_dsk_sm90 as sageattn_qk_int8_pv_fp8_cuda_sm90a -from torch.nn.functional import scaled_dot_product_attention as sdpa - -import torch -import numpy as np -import os -import argparse - - -def precision_cmp_torch(t1: torch.Tensor, t2: torch.Tensor): - x, xx = t1.to(dtype=torch.float32), t2.to(dtype=torch.float32) - # 重塑张量并计算余弦相似度 - x_reshaped = torch.reshape(x, [1, -1]) - xx_reshaped = torch.reshape(xx, [1, -1]) - sim = torch.nn.functional.cosine_similarity(x_reshaped, xx_reshaped).item() - - # 计算 L1 误差 - l1 = (torch.abs(x - xx).sum() / torch.abs(xx).sum()).item() - - max_diff = torch.max(x - xx) - - return sim, l1, max_diff - -bsz = 2 -seq_len = 1026 -num_heads = 128 -head_dim_qk = 128 + 64 -head_dim_v = 128 - -tensor_layout = "NHD" -is_causal = True -return_lse = False - -torch.backends.cuda.enable_flash_sdp(True) - -q = torch.randn((bsz, seq_len, num_heads, head_dim_qk), dtype=torch.float16).cuda() -k = torch.randn((bsz, seq_len, num_heads, head_dim_qk), dtype=torch.float16).cuda() -v = torch.randn((bsz, seq_len, num_heads, head_dim_v), dtype=torch.float16).cuda() - -# to HND for sdpa input -q = q.transpose(2, 1) -k = k.transpose(2, 1) -v = v.transpose(2, 1) - -o_torch_fa2 = sdpa(q, k, v, is_causal=is_causal) -torch.cuda.synchronize() - -torch.backends.cuda.enable_flash_sdp(False) -o_torch_sdpa = sdpa(q, k, v, is_causal=is_causal) -torch.cuda.synchronize() - -head_dim_og = head_dim_qk -sm_scale = head_dim_og**-0.5 - -# to NHD for sageattn input -q = q.transpose(2, 1) -k = k.transpose(2, 1) -v = v.transpose(2, 1) -o_sa = sageattn_qk_int8_pv_fp8_cuda_sm90a(q, k, v, tensor_layout=tensor_layout, is_causal=is_causal, qk_quant_gran="per_warp", return_lse=return_lse, pv_accum_dtype="fp32+fp32") -torch.cuda.synchronize() - -sim, l1, max_diff = precision_cmp_torch(o_torch_fa2.transpose(2, 1), o_sa) -print(f"Sim and Diff of Sage Attn & torch SDPA: {sim}, {max_diff}") - -sim, l1, max_diff = precision_cmp_torch(o_torch_sdpa.transpose(2, 1), o_sa) -print(f"Sim and Diff of Sage Attn & Flash Attn: {sim}, {max_diff}") - -sim, l1, max_diff = precision_cmp_torch(o_torch_fa2, o_torch_sdpa) -print(f"Sim and Diff of Flash Attn & torch SDPA: {sim}, {max_diff}")