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/csrc/dispatch_utils.h b/csrc/dispatch_utils.h index 6d798e68..17fc199f 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 == 256) { \ + constexpr int HEAD_DIM = 256; \ + __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 new file mode 100644 index 00000000..347b0244 --- /dev/null +++ b/csrc/qattn/qk_int_sv_f8_cuda_dsk_sm90.cu @@ -0,0 +1,975 @@ +#include "../utils.cuh" +#include +#include +#include +#include + +#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) { + 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_[]; + + 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]; + 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(sK, &tensorMapK, &barrier_K, 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]; + 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]); + } + } + } + } +} + +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) +{ + 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, 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); + 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.size(0); + 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_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); + + if (tensor_layout == 0) + { + 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_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); + 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.size(1) == head_dim); + assert(value.size(2) == num_kv_heads); + } + else + { + 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_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); + 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.size(2) == head_dim); + assert(value.size(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()); + } + + torch::Tensor lse = torch::empty({0}); + if (return_lse) + { + 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.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_PYTORCH_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.size(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_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_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); + + 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_ptr()), + reinterpret_cast(key_scale.data_ptr()), + nullptr, + reinterpret_cast(output.data_ptr()), + stride_bz_o, stride_h_o, stride_seq_o, + qo_len, kv_len, num_kv_groups, sm_scale); + }); + }); + }); + }); + + 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) +{ + 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, 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); + 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.size(0); + 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_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); + + if (tensor_layout == 0) + { + 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_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); + 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.size(1) == head_dim); + assert(value.size(2) == num_kv_heads); + } + else + { + 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_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); + 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.size(2) == head_dim); + assert(value.size(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()); + } + + torch::Tensor lse = torch::empty({1}); + if (return_lse) + { + 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.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_PYTORCH_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.size(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_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_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_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); + + 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_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); + }); + }); + }); + }); + + return lse; +} \ No newline at end of file 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..a39a8a4d 100644 --- a/sageattention/core.py +++ b/sageattention/core.py @@ -901,4 +901,91 @@ def sageattn_qk_int8_pv_fp8_cuda_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 + + +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_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 + + 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 < 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." + + 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, 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_causal, _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: + return o + 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,