diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index eaf13184..4e531755 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -187,8 +187,8 @@ namespace sd { } copy->buffer = buffer; // Capability checks inspect alignment even before weights are resident. - copy->data = source->data ? source->data : reinterpret_cast(ggml_backend_get_alignment(external)); - copy->extra = nullptr; + copy->data = source->data ? source->data : reinterpret_cast(ggml_backend_get_alignment(external)); + copy->extra = nullptr; } } auto copy_graph = ggml_new_graph_custom(context, graph_size, false); diff --git a/src/core/cuda/affine.cu b/src/core/cuda/affine.cu index e410d40b..2cd660b7 100644 --- a/src/core/cuda/affine.cu +++ b/src/core/cuda/affine.cu @@ -129,15 +129,15 @@ namespace SD_CUDA_NAMESPACE { ggml_ext_gpu_context context; if (!cuda_context(execution->backend, &context)) return GGML_STATUS_FAILED; - auto stream = static_cast(context.stream); - const auto* node = nodes[0]; - const auto* a = node->src[0]; - const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1]; - auto* dst = nodes[count - 1]; - const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr; - const bool gate = residual != nullptr; - const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) || - (residual && overlaps(dst, residual) && dst->data != residual->data); + auto stream = static_cast(context.stream); + const auto* node = nodes[0]; + const auto* a = node->src[0]; + const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1]; + auto* dst = nodes[count - 1]; + const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr; + const bool gate = residual != nullptr; + const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) || + (residual && overlaps(dst, residual) && dst->data != residual->data); void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; const int64_t count4 = ggml_nelements(dst) / 4; const unsigned blocks = unsigned((count4 + 255) / 256); diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index c4ca30ff..cdafc998 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -209,8 +209,8 @@ namespace sd_fattn_cute { using Element = cutlass::half_t; static constexpr int Threads = Warps * 32; using Mma = TiledMMA, - Layout, _1, _1>>, - Tile, _16, _16>>; + Layout, _1, _1>>, + Tile, _16, _16>>; static constexpr int SmemWidth = D % 64 == 0 ? 64 : D % 32 == 0 ? 32 : 16; static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2 @@ -449,7 +449,7 @@ static __global__ void copy_to_half(const char* src, half* dst, int64_t count, i if (i < count) { const int64_t row = i / dim; const auto* p = reinterpret_cast(src + (row % rows) * nb1 + - ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3); + ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3); dst[i] = __float2half_rn(float(p[i % dim]) * scale); } } @@ -784,9 +784,9 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) { desc.preferred = nullptr; desc.optimize = nullptr; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); - desc.name = "sd.i8_rope_cute_attention.cuda.v1"; + desc.name = "sd.i8_rope_cute_attention.cuda.v1"; desc.op = GGML_OP_CUSTOM; desc.custom_op = sd_ggml_i8_op(); - desc.priority = 400; + desc.priority = 400; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index 84d23326..69f91d2a 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -197,8 +197,8 @@ namespace sd_i8_cutlass { using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal, cutlass::epilogue::thread::LinearCombination, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel; using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks; using Kernel = std::conditional_t, - cutlass::gemm::kernel::GemmWithEpilogueVisitor>; + cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk, + cutlass::gemm::kernel::GemmWithEpilogueVisitor>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; struct Plan { diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 46a12676..45f6ecdf 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -661,10 +661,10 @@ namespace SD_CUDA_NAMESPACE { #else constexpr auto compute_type = HIPBLAS_R_32I; #endif - const auto status = hipblasGemmEx(static_cast(handle), HIPBLAS_OP_T, HIPBLAS_OP_N, - int(p.n), int(p.padded), int(p.k), &alpha, - op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k), - &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT); + const auto status = hipblasGemmEx(static_cast(handle), HIPBLAS_OP_T, HIPBLAS_OP_N, + int(p.n), int(p.padded), int(p.k), &alpha, + op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k), + &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT); if (status != HIPBLAS_STATUS_SUCCESS) #else const auto status = cublasGemmEx(static_cast(handle), CUBLAS_OP_T, CUBLAS_OP_N, diff --git a/src/core/cuda/sage_attention.cu b/src/core/cuda/sage_attention.cu index 46b9a08a..bb42d5b9 100644 --- a/src/core/cuda/sage_attention.cu +++ b/src/core/cuda/sage_attention.cu @@ -287,9 +287,9 @@ namespace sd_cuda { constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D; auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D, - sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp, - sageattention::QuantGranularity::kPerWarp, float, false, half, - sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>; + sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp, + sageattention::QuantGranularity::kPerWarp, float, false, half, + sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>; kernel<<>>( qi, ki, static_cast(v->data), out, nullptr, qs, ks, nullptr, lq, lk, hq / hk, diff --git a/src/core/cuda/sageattention/attn.cuh b/src/core/cuda/sageattention/attn.cuh index 61b57580..540d5822 100644 --- a/src/core/cuda/sageattention/attn.cuh +++ b/src/core/cuda/sageattention/attn.cuh @@ -17,14 +17,14 @@ #ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH #define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH -#include "math.cuh" #include "attn_utils.cuh" +#include "math.cuh" namespace sageattention { -#define PACK_SIZE_QK 16 // as if it is int8 -#define PACK_SIZE_V 8 // fp16 -#define PACK_SIZE_O 8 // fp16 +#define PACK_SIZE_QK 16 // as if it is int8 +#define PACK_SIZE_V 8 // fp16 +#define PACK_SIZE_O 8 // fp16 // treat as if int8 tensor core #define MMA_QK_M 16 @@ -36,640 +36,546 @@ namespace sageattention { #define MMA_SV_N 16 #define MMA_SV_K 16 -template -__global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, half *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse, - float *__restrict__ Q_scale, float *__restrict__ K_scale, DTypeOut *__restrict__ V_mean, - const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, - const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, - const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, - const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, - const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, - float sm_scale) -{ - // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes. + template + __global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) { + // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes. #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 - // compile time check - static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); - static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); - static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); - static_assert(std::is_same::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float"); - static_assert(std::is_same::value || std::is_same::value, "DTypeSVAccum must be float or half"); - static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); - static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); - static_assert(!fuse_v_mean || std::is_same::value, "fuse_v_mean only supports half"); - static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation - - using DTypeOut2 = typename std::conditional::value, half2, nv_bfloat162>::type; - - constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; - constexpr uint32_t num_warps_k = CTA_K / WARP_K; - constexpr uint32_t num_warps = num_warps_q * num_warps_k; - constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; - constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; - constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); - constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; - - constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); - constexpr uint32_t O_SMEM_STRIDE = head_dim; - constexpr uint32_t V_SMEM_STRIDE = head_dim; - - extern __shared__ int8_t smem[]; - - const uint32_t lane_id = get_lane_id(); - const uint32_t warp_id = get_warp_id(); - - // maximize L2 hit rate - const uint32_t batch_id = blockIdx.z; - const uint32_t bx = blockIdx.x; - const uint32_t num_qo_heads = gridDim.y; - const uint32_t head_id = blockIdx.y; - - // transfer to base 2 instead of base e with better numerical efficiency - sm_scale *= math::log2e; - - // RS holds the fragment of S - int32_t RS[num_tiles_q][num_tiles_k][8]; - DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; - float m[num_tiles_q][2]; // max - float d[num_tiles_q][2]; // denominator - - uint32_t q_scale_idx, k_scale_idx; - - 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 * num_warps_q; - q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); - } - else if constexpr (Q_GRAN == QuantGranularity::kPerThread) - { - const uint32_t num_warp_block_q = gridDim.x * num_warps_q; - q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; - } - - if constexpr (K_GRAN == QuantGranularity::kPerBlock) - { - const uint32_t num_block_k = SAGE_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::kPerWarp) - { - const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); - k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); - } - else if constexpr (K_GRAN == QuantGranularity::kPerThread) - { - const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); - k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; - } - - constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; - - // initialize o, m, d -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - if constexpr (std::is_same::value) - { -#pragma unroll - for (uint32_t k = 0; k < 8; k++) - { - RO[fq][fv][k] = 0.0f; - } - } - else if constexpr (std::is_same::value) - { -#pragma unroll - for (uint32_t k = 0; k < 4; k++) - { - ((int32_t*)RO[fq][fv])[k] = 0; - } - } - } - } -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t k = 0; k < 2; k++) - { - m[fq][k] = -5000000.0f; - d[fq][k] = 1.0f; - } - } - - constexpr uint32_t K_smem_idx_offset = CTA_Q; - constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; - - constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_Q(smem); - smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); - constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); - constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_O(smem); - - constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; - constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; - constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8; - constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4; - constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; - constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; - - constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); - constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); - constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); - constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); - constexpr uint32_t V_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V); - constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); - constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); - - int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; - int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; - half *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; - uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); - uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); - uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); - - uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); - uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); - uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); - - // for causal masking - uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; - uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); - - // for loading - uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; - uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; - uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V; - - const uint32_t num_iterations = SAGE_DIV_CEIL( - mask_mode == MaskMode::kCausal - ? min(kv_len, (bx + 1) * CTA_Q) - : kv_len, - CTA_K); - - // load Q with predicate - load_global_to_share( - &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); - cp_async::commit_group(); - cp_async::wait_group<0>(); - __syncthreads(); - - // for num_tiles_qk_inner = 1, we load all Qs in register - uint32_t RQ[num_tiles_q][4]; - if constexpr (num_tiles_qk_inner == 1) - { -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); - Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); - } - } - - // load K with predicate - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); - cp_async::commit_group(); - - float q_scale = Q_scale[q_scale_idx]; - - float original_sm_scale = sm_scale; - float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; - - sm_scale = original_sm_scale * dequant_scale; - - // load V with predicate - load_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); - cp_async::commit_group(); - - K_load_idx_lane_base += CTA_K; - V_load_idx_lane_base += CTA_K; - -#pragma unroll - for (uint32_t iter = 1; iter < num_iterations - 1; iter++) - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - - 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]); - } - } - } - - // do not apply causal mask and out of bound mask for these iterations - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; - RS_32_to_16(RS_f32, RS_f16); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d(RS_f16, d); - } - - __syncthreads(); - - // load K - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); - cp_async::commit_group(); - - dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; - sm_scale = original_sm_scale * dequant_scale; - - // ensure V is ready - cp_async::wait_group<1>(); - __syncthreads(); - - if constexpr (!use_inst_buffer) - { - compute_fp16_sv_permuted( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - else - { - compute_fp16_sv_permuted_inst_buf( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - - __syncthreads(); - // load V - load_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V); - cp_async::commit_group(); - K_load_idx_lane_base += CTA_K; - V_load_idx_lane_base += CTA_K; - } - - // second last iter, apply causal mask - if (num_iterations > 1) - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - - 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]) * dequant_scale; - } - } - } - - if constexpr (mask_mode == MaskMode::kCausal) - { - apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); - } - // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; - RS_32_to_16(RS_f32, RS_f16); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d(RS_f16, d); - } - - __syncthreads(); - - // load K with predicate - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); - cp_async::commit_group(); - - dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; - sm_scale = original_sm_scale * dequant_scale; - - // ensure V is ready - cp_async::wait_group<1>(); - __syncthreads(); - - if constexpr (!use_inst_buffer) - { - compute_fp16_sv_permuted( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - else - { - compute_fp16_sv_permuted_inst_buf( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - - __syncthreads(); - // load V with predicate - load_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); - cp_async::commit_group(); - K_load_idx_lane_base += CTA_K; - V_load_idx_lane_base += CTA_K; - } - - // last iter, apply causal mask and out of bound mask - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - - 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]) * dequant_scale; - } - } - } - - if constexpr (mask_mode == MaskMode::kCausal) - { - apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); - } - // check out of bound in the last iter - apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; - RS_32_to_16(RS_f32, RS_f16); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d(RS_f16, d); - } - - // ensure V is ready - cp_async::wait_group<0>(); - __syncthreads(); - - if constexpr (!use_inst_buffer) - { - compute_fp16_sv_permuted( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - else - { - compute_fp16_sv_permuted_inst_buf( - smem_V, RS_f16, RO, d, V_smem_offset_mma); - } - - __syncthreads(); - - } - - // TODO: thread block sync mdo state for num_warps_k > 0 - - normalize_d(RO, m, d); - - // save the result - // if (get_warp_idx_k() == 0) - // { - - // convert half to bfloat16 - if constexpr (std::is_same::value && std::is_same::value) - { -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0])); - ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1])); - ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2])); - ((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3])); - } - } - } - - // add v_mean - if constexpr (fuse_v_mean) - { - DTypeOut2 v_mean[2]; - DTypeOut *V_mean_lane_ptr = V_mean + 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++) - { - v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16)); - v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16)); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]); - ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]); - ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]); - ((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]); - } - } - } - - // save the result to shared memory - uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); - - if constexpr (std::is_same::value) - { - // convert RO to half - uint32_t RO_f16[4]; -#pragma unroll - for (uint32_t k = 0; k < 4; k++) - { - if constexpr (std::is_same::value) - { - ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); - } - else if constexpr (std::is_same::value) - { - ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); - } + // compile time check + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); + static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(std::is_same::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float"); + static_assert(std::is_same::value || std::is_same::value, "DTypeSVAccum must be float or half"); + static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); + static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); + static_assert(!fuse_v_mean || std::is_same::value, "fuse_v_mean only supports half"); + static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation + + using DTypeOut2 = typename std::conditional::value, half2, nv_bfloat162>::type; + + constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; + constexpr uint32_t num_warps_k = CTA_K / WARP_K; + constexpr uint32_t num_warps = num_warps_q * num_warps_k; + constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; + constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; + constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); + constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; + + constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); + constexpr uint32_t O_SMEM_STRIDE = head_dim; + constexpr uint32_t V_SMEM_STRIDE = head_dim; + + extern __shared__ int8_t smem[]; + + const uint32_t lane_id = get_lane_id(); + const uint32_t warp_id = get_warp_id(); + + // maximize L2 hit rate + const uint32_t batch_id = blockIdx.z; + const uint32_t bx = blockIdx.x; + const uint32_t num_qo_heads = gridDim.y; + const uint32_t head_id = blockIdx.y; + + // transfer to base 2 instead of base e with better numerical efficiency + sm_scale *= math::log2e; + + // RS holds the fragment of S + int32_t RS[num_tiles_q][num_tiles_k][8]; + DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; + float m[num_tiles_q][2]; // max + float d[num_tiles_q][2]; // denominator + + uint32_t q_scale_idx, k_scale_idx; + + 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 * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); + } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; } - ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; - ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + if constexpr (K_GRAN == QuantGranularity::kPerBlock) { + const uint32_t num_block_k = SAGE_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::kPerWarp) { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); + } else if constexpr (K_GRAN == QuantGranularity::kPerThread) { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; + } - // ! permuted, make sure you know what you are doing - ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2]; - ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; - } - else if constexpr (std::is_same::value) - { - ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0]; - ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1]; + constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) + : (CTA_K / WARP_K) * 4; - // ! permuted, make sure you know what you are doing - ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[2]; - ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3]; - } - } - } + // initialize o, m, d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + if constexpr (std::is_same::value) { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) { + RO[fq][fv][k] = 0.0f; + } + } else if constexpr (std::is_same::value) { +#pragma unroll + for (uint32_t k = 0; k < 4; k++) { + ((int32_t*)RO[fq][fv])[k] = 0; + } + } + } + } +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) { + m[fq][k] = -5000000.0f; + d[fq][k] = 1.0f; + } + } - // ! do we need to sync here? - __syncwarp(); + constexpr uint32_t K_smem_idx_offset = CTA_Q; + constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; - // shared memory to global memory - DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; - uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); - uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B + : SwizzleMode::k128B; + smem_t smem_Q(smem); + smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_O(smem); + + constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 + : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 + : 4; + constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4; + constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; + + constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); + constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); + constexpr uint32_t V_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V); + constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); + constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); + + int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + half* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; + uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); + + uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); + uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); + + // for causal masking + uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); + + // for loading + uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V; + + const uint32_t num_iterations = SAGE_DIV_CEIL( + mask_mode == MaskMode::kCausal + ? min(kv_len, (bx + 1) * CTA_Q) + : kv_len, + CTA_K); + + // load Q with predicate + load_global_to_share( + &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + // for num_tiles_qk_inner = 1, we load all Qs in register + uint32_t RQ[num_tiles_q][4]; + if constexpr (num_tiles_qk_inner == 1) { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); + Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); + } + } + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + float q_scale = Q_scale[q_scale_idx]; + + float original_sm_scale = sm_scale; + float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; + + sm_scale = original_sm_scale * dequant_scale; + + // load V with predicate + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; #pragma unroll - for (uint32_t i = 0; i < O_smem_iters_col; i++) - { + for (uint32_t iter = 1; iter < num_iterations - 1; iter++) { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + #pragma unroll - for (uint32_t j = 0; j < O_smem_iters_row; j++) - { - if (O_load_idx_lane_base < qo_len) - { - smem_O.store_128b(offset_O, O_lane_ptr); - } - O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); - offset_O = smem_O.advance_offset_by_column(offset_O); - } + 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]); + } + } + } - offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); - O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); - O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; - } + // do not apply causal mask and out of bound mask for these iterations + K_idx_lane_base += CTA_K; - if constexpr (return_lse) - { - uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); - float *lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; - uint32_t fq = (lane_id % 4) / 2; - uint32_t k = (lane_id % 4) % 2; + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } - if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) - { - lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); - } - } + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } - // } + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; + } + + // second last iter, apply causal mask + if (num_iterations > 1) { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + 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]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V with predicate + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; + } + + // last iter, apply causal mask and out of bound mask + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + 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]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // check out of bound in the last iter + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + } + + // TODO: thread block sync mdo state for num_warps_k > 0 + + normalize_d(RO, m, d); + + // save the result + // if (get_warp_idx_k() == 0) + // { + + // convert half to bfloat16 + if constexpr (std::is_same::value && std::is_same::value) { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0])); + ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1])); + ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2])); + ((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3])); + } + } + } + + // add v_mean + if constexpr (fuse_v_mean) { + DTypeOut2 v_mean[2]; + DTypeOut* V_mean_lane_ptr = V_mean + 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++) { + v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16)); + v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]); + ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]); + ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]); + ((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]); + } + } + } + + // save the result to shared memory + uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); + + if constexpr (std::is_same::value) { + // convert RO to half + uint32_t RO_f16[4]; +#pragma unroll + for (uint32_t k = 0; k < 4; k++) { + if constexpr (std::is_same::value) { + ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); + } else if constexpr (std::is_same::value) { + ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); + } + } + + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + + // ! permuted, make sure you know what you are doing + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2]; + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; + } else if constexpr (std::is_same::value) { + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1]; + + // ! permuted, make sure you know what you are doing + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[2]; + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3]; + } + } + } + + // ! do we need to sync here? + __syncwarp(); + + // shared memory to global memory + DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; + uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); + uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + +#pragma unroll + for (uint32_t i = 0; i < O_smem_iters_col; i++) { +#pragma unroll + for (uint32_t j = 0; j < O_smem_iters_row; j++) { + if (O_load_idx_lane_base < qo_len) { + smem_O.store_128b(offset_O, O_lane_ptr); + } + O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); + offset_O = smem_O.advance_offset_by_column(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); + O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); + O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; + } + + if constexpr (return_lse) { + uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); + float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; + uint32_t fq = (lane_id % 4) / 2; + uint32_t k = (lane_id % 4) % 2; + + if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) { + lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); + } + } + + // } #else - RUNTIME_ASSERT("SageAttention requires SM80 or newer"); + RUNTIME_ASSERT("SageAttention requires SM80 or newer"); #endif -} - + } #undef PACK_SIZE_QK #undef PACK_SIZE_V @@ -681,6 +587,6 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest #undef MMA_SV_N #undef MMA_SV_K -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/attn_fp8.cuh b/src/core/cuda/sageattention/attn_fp8.cuh index 507f9e7f..838e4c66 100644 --- a/src/core/cuda/sageattention/attn_fp8.cuh +++ b/src/core/cuda/sageattention/attn_fp8.cuh @@ -17,14 +17,14 @@ #ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH #define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH -#include "math.cuh" #include "attn_utils.cuh" +#include "math.cuh" namespace sageattention { -#define PACK_SIZE_QK 16 // as if it is int8 -#define PACK_SIZE_V 16 // fp8 -#define PACK_SIZE_O 8 // fp16 +#define PACK_SIZE_QK 16 // as if it is int8 +#define PACK_SIZE_V 16 // fp8 +#define PACK_SIZE_O 8 // fp16 // treat as if int8 tensor core #define MMA_QK_M 16 @@ -36,672 +36,576 @@ namespace sageattention { #define MMA_SV_N 16 #define MMA_SV_K 32 -template -__global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, int8_t *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse, - float *__restrict__ Q_scale, float *__restrict__ K_scale, float *__restrict__ V_scale, float *__restrict__ V_mean, - const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, - const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, - const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, - const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, - const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, - float sm_scale) -{ - // Keep unsupported device passes from parsing architecture-specific intrinsics. + template + __global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) { + // Keep unsupported device passes from parsing architecture-specific intrinsics. #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890 - // compile time check - static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); - static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); - static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); - static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); - static_assert(std::is_same::value, "DTypeSVAccum must be float, half is WIP"); - static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); - static_assert(CTA_K % 64 == 0); - static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation - - constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; - constexpr uint32_t num_warps_k = CTA_K / WARP_K; - constexpr uint32_t num_warps = num_warps_q * num_warps_k; - constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; - constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; - constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); - constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; - - constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); - constexpr uint32_t O_SMEM_STRIDE = head_dim; - // for fp16: head_dim - constexpr uint32_t V_SMEM_STRIDE = CTA_K; - - extern __shared__ int8_t smem[]; - - const uint32_t lane_id = get_lane_id(); - const uint32_t warp_id = get_warp_id(); - - // maximize L2 hit rate - const uint32_t batch_id = blockIdx.z; - const uint32_t bx = blockIdx.x; - const uint32_t num_qo_heads = gridDim.y; - const uint32_t head_id = blockIdx.y; - - // transfer to base 2 instead of base e with better numerical efficiency - sm_scale *= math::log2e; - - // RS holds the fragment of S - int32_t RS[num_tiles_q][num_tiles_k][8]; - DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; - float m[num_tiles_q][2]; // max - float d[num_tiles_q][2]; // denominator - - uint32_t q_scale_idx, k_scale_idx; - - 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 * num_warps_q; - q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); - } - else if constexpr (Q_GRAN == QuantGranularity::kPerThread) - { - const uint32_t num_warp_block_q = gridDim.x * num_warps_q; - q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; - } - - if constexpr (K_GRAN == QuantGranularity::kPerBlock) - { - const uint32_t num_block_k = SAGE_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::kPerWarp) - { - const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); - k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); - } - else if constexpr (K_GRAN == QuantGranularity::kPerThread) - { - const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); - k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; - } - - constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; - - // initialize o, m, d -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - if constexpr (std::is_same::value) - { -#pragma unroll - for (uint32_t k = 0; k < 8; k++) - { - RO[fq][fv][k] = 0.0f; - } - } - else if constexpr (std::is_same::value) - { -#pragma unroll - for (uint32_t k = 0; k < 4; k++) - { - ((int32_t*)RO[fq][fv])[k] = 0; - } - } - } - } -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t k = 0; k < 2; k++) - { - m[fq][k] = -5000000.0f; - d[fq][k] = 1.0f; - } - } - - constexpr uint32_t K_smem_idx_offset = CTA_Q; - constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; - - constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_Q(smem); - smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); - // for fp16: 32 - constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); - constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; - smem_t smem_O(smem); - - constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; - constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; - // for fp16: 32 - constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8; - // for fp16: 32 - constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4; - constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; - constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; - - constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); - constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); - constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); - constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); - // for fp16: CTA_K - constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V); - constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); - constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); - - int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; - int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; - // for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v - int8_t *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; - uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); - uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); - uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); - - uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); - uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); - // for fp 16: - // uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); - uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2); - - // for causal masking - uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; - uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); - - // for loading - uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; - uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; - - const uint32_t num_iterations = SAGE_DIV_CEIL( - mask_mode == MaskMode::kCausal - ? min(kv_len, (bx + 1) * CTA_Q) - : kv_len, - CTA_K); - - // load Q with predicate - load_global_to_share( - &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); - cp_async::commit_group(); - cp_async::wait_group<0>(); - __syncthreads(); - - // for num_tiles_qk_inner = 1, we load all Qs in register - uint32_t RQ[num_tiles_q][4]; - if constexpr (num_tiles_qk_inner == 1) - { -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); - Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); - } - } - - // load K with predicate - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); - cp_async::commit_group(); - - float q_scale = Q_scale[q_scale_idx]; - - float original_sm_scale = sm_scale; - float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; - - sm_scale = original_sm_scale * dequant_scale; - - // load V - // ! we assume that V is padded. If not, there might be illegal memory access or nan issue. - // for fp16: - // load_global_to_share stride_seq_v - load_fp8_V_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); - cp_async::commit_group(); - - K_load_idx_lane_base += CTA_K; - -#pragma unroll - for (uint32_t iter = 1; iter < num_iterations - 1; iter++) - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - 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]); - } - } - } - - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; - RS_32_to_8(RS_f32, RS_f8); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d_f8(RS_f8, d); - } - - __syncthreads(); - - // load K without predicate - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); - cp_async::commit_group(); - - dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; - sm_scale = original_sm_scale * dequant_scale; - - // ensure V is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // for fp16: - // compute_fp16_sv_permuted( - // smem_V, RS_f16, RO, d, V_smem_offset_mma); - if constexpr (!use_inst_buffer) - { - compute_fp8_sv( - smem_V, RS_f8, RO, d); - } - else - { - if constexpr (!use_pv_fp16_accu){ - compute_fp8_sv_inst_buf( - smem_V, RS_f8, RO, d); - } - else{ - compute_fp8_sv_inst_buf_fp16_accu( - smem_V, RS_f8, RO, d); - } - } - __syncthreads(); - // load V - // for fp16: - // load_global_to_share stride_seq_v - load_fp8_V_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); - cp_async::commit_group(); - - K_load_idx_lane_base += CTA_K; - } - - // second last iter, apply causal mask - if (num_iterations > 1) - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - - 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]) * dequant_scale; - } - } - } - - if constexpr (mask_mode == MaskMode::kCausal) - { - apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); - } - // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; - RS_32_to_8(RS_f32, RS_f8); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d_f8(RS_f8, d); - } - - __syncthreads(); - - // load K with predicate - load_global_to_share( - &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); - cp_async::commit_group(); - - dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; - sm_scale = original_sm_scale * dequant_scale; - - // ensure V is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // for fp16: - // compute_fp16_sv_permuted( - // smem_V, RS_f16, RO, d, V_smem_offset_mma); - if constexpr (!use_inst_buffer) - { - compute_fp8_sv( - smem_V, RS_f8, RO, d); - } - else - { - if constexpr (!use_pv_fp16_accu){ - compute_fp8_sv_inst_buf( - smem_V, RS_f8, RO, d); - } - else{ - compute_fp8_sv_inst_buf_fp16_accu( - smem_V, RS_f8, RO, d); - } - } - - __syncthreads(); - // load V - // for fp16: - // load_global_to_share stride_seq_v - load_fp8_V_global_to_share( - &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); - cp_async::commit_group(); - K_load_idx_lane_base += CTA_K; - } - - // last iter, apply causal mask and out of bound mask - { - // ensure K is ready - cp_async::wait_group<1>(); - __syncthreads(); - - // compute QK^T - if constexpr (num_tiles_qk_inner == 1) - { - compute_int_qk( - smem_K, RS, RQ, K_smem_offset_mma); - } - else - { - compute_int_qk( - smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); - } - - 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]) * dequant_scale; - } - } - } - - if constexpr (mask_mode == MaskMode::kCausal) - { - apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); - } - apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); - K_idx_lane_base += CTA_K; - - if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - else if constexpr (std::is_same::value) - { - update_mdo(RS_f32, RO, m, d, original_sm_scale); - } - - if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) - { - accumulate_d(RS_f32, d); - } - - uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; - RS_32_to_8(RS_f32, RS_f8); - - if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) - { - accumulate_d_f8(RS_f8, d); - } - - // ensure V is ready - cp_async::wait_group<0>(); - __syncthreads(); - - // for fp16: - // compute_fp16_sv_permuted( - // smem_V, RS_f16, RO, d, V_smem_offset_mma); - if constexpr (!use_inst_buffer) - { - compute_fp8_sv( - smem_V, RS_f8, RO, d); - } - else - { - if constexpr (!use_pv_fp16_accu){ - compute_fp8_sv_inst_buf( - smem_V, RS_f8, RO, d); - } - else{ - compute_fp8_sv_inst_buf_fp16_accu( - smem_V, RS_f8, RO, d); - } - } - - __syncthreads(); - - } - - // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving. - - normalize_d(RO, m, d); - - // ! here we just implement the case for fp32 acumulation - 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]; - } - } - } - - if constexpr (fuse_v_mean) - { - float v_mean[4]; - float *V_mean_base_ptr = V_mean + 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_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16)); - ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8)); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - RO[fq][fv][0] += v_mean[0]; - RO[fq][fv][1] += v_mean[1]; - RO[fq][fv][2] += v_mean[0]; - RO[fq][fv][3] += v_mean[1]; - RO[fq][fv][4] += v_mean[2]; - RO[fq][fv][5] += v_mean[3]; - RO[fq][fv][6] += v_mean[2]; - RO[fq][fv][7] += v_mean[3]; - } - } - } - - // save the result to shared memory - uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); - - if constexpr (std::is_same::value) - { - // convert RO to half - uint32_t RO_f16[4]; -#pragma unroll - for (uint32_t k = 0; k < 4; k++) - { - if constexpr (std::is_same::value) - { - ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); - } - else - { - ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); - } + // compile time check + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); + static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); + static_assert(std::is_same::value, "DTypeSVAccum must be float, half is WIP"); + static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); + static_assert(CTA_K % 64 == 0); + static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation + + constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; + constexpr uint32_t num_warps_k = CTA_K / WARP_K; + constexpr uint32_t num_warps = num_warps_q * num_warps_k; + constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; + constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; + constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); + constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; + + constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); + constexpr uint32_t O_SMEM_STRIDE = head_dim; + // for fp16: head_dim + constexpr uint32_t V_SMEM_STRIDE = CTA_K; + + extern __shared__ int8_t smem[]; + + const uint32_t lane_id = get_lane_id(); + const uint32_t warp_id = get_warp_id(); + + // maximize L2 hit rate + const uint32_t batch_id = blockIdx.z; + const uint32_t bx = blockIdx.x; + const uint32_t num_qo_heads = gridDim.y; + const uint32_t head_id = blockIdx.y; + + // transfer to base 2 instead of base e with better numerical efficiency + sm_scale *= math::log2e; + + // RS holds the fragment of S + int32_t RS[num_tiles_q][num_tiles_k][8]; + DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; + float m[num_tiles_q][2]; // max + float d[num_tiles_q][2]; // denominator + + uint32_t q_scale_idx, k_scale_idx; + + 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 * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); + } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; } - ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; - ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + if constexpr (K_GRAN == QuantGranularity::kPerBlock) { + const uint32_t num_block_k = SAGE_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::kPerWarp) { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); + } else if constexpr (K_GRAN == QuantGranularity::kPerThread) { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; + } - offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1); - ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2]; - ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; - } - else if constexpr (std::is_same::value) - { - // TODO: not implement - } - } - } + constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) + : (CTA_K / WARP_K) * 4; - // ! do we need to sync here? - __syncwarp(); + // initialize o, m, d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + if constexpr (std::is_same::value) { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) { + RO[fq][fv][k] = 0.0f; + } + } else if constexpr (std::is_same::value) { +#pragma unroll + for (uint32_t k = 0; k < 4; k++) { + ((int32_t*)RO[fq][fv])[k] = 0; + } + } + } + } +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) { + m[fq][k] = -5000000.0f; + d[fq][k] = 1.0f; + } + } - // shared memory to global memory - DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; - uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); - uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + constexpr uint32_t K_smem_idx_offset = CTA_Q; + constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; + + constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B + : SwizzleMode::k128B; + smem_t smem_Q(smem); + smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); + // for fp16: 32 + constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_O(smem); + + constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 + : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 + : 4; + // for fp16: 32 + constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8; + // for fp16: 32 + constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4; + constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; + + constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); + constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); + // for fp16: CTA_K + constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V); + constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); + constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); + + int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + // for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v + int8_t* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; + uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); + + uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); + // for fp 16: + // uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); + uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2); + + // for causal masking + uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); + + // for loading + uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + + const uint32_t num_iterations = SAGE_DIV_CEIL( + mask_mode == MaskMode::kCausal + ? min(kv_len, (bx + 1) * CTA_Q) + : kv_len, + CTA_K); + + // load Q with predicate + load_global_to_share( + &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + // for num_tiles_qk_inner = 1, we load all Qs in register + uint32_t RQ[num_tiles_q][4]; + if constexpr (num_tiles_qk_inner == 1) { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); + Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); + } + } + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + float q_scale = Q_scale[q_scale_idx]; + + float original_sm_scale = sm_scale; + float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; + + sm_scale = original_sm_scale * dequant_scale; + + // load V + // ! we assume that V is padded. If not, there might be illegal memory access or nan issue. + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; #pragma unroll - for (uint32_t i = 0; i < O_smem_iters_col; i++) - { + for (uint32_t iter = 1; iter < num_iterations - 1; iter++) { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + float RS_f32[num_tiles_q][num_tiles_k][8]; + #pragma unroll - for (uint32_t j = 0; j < O_smem_iters_row; j++) - { - if (O_load_idx_lane_base < qo_len) - { - smem_O.store_128b(offset_O, O_lane_ptr); - } - O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); - offset_O = smem_O.advance_offset_by_column(offset_O); - } + 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]); + } + } + } - offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); - O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); - O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; - } + K_idx_lane_base += CTA_K; - if constexpr (return_lse) - { - // ! this only works for num_tiles_q = 2 - uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); - float *lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; - uint32_t fq = (lane_id % 4) / 2; - uint32_t k = (lane_id % 4) % 2; + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } - if (lse_idx < qo_len) - { - lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); - } - } + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K without predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; + } + + // second last iter, apply causal mask + if (num_iterations > 1) { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + 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]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + } + + // last iter, apply causal mask and out of bound mask + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + 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]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + + __syncthreads(); + } + + // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving. + + normalize_d(RO, m, d); + + // ! here we just implement the case for fp32 acumulation + 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]; + } + } + } + + if constexpr (fuse_v_mean) { + float v_mean[4]; + float* V_mean_base_ptr = V_mean + 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_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16)); + ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + RO[fq][fv][0] += v_mean[0]; + RO[fq][fv][1] += v_mean[1]; + RO[fq][fv][2] += v_mean[0]; + RO[fq][fv][3] += v_mean[1]; + RO[fq][fv][4] += v_mean[2]; + RO[fq][fv][5] += v_mean[3]; + RO[fq][fv][6] += v_mean[2]; + RO[fq][fv][7] += v_mean[3]; + } + } + } + + // save the result to shared memory + uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); + + if constexpr (std::is_same::value) { + // convert RO to half + uint32_t RO_f16[4]; +#pragma unroll + for (uint32_t k = 0; k < 4; k++) { + if constexpr (std::is_same::value) { + ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); + } else { + ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); + } + } + + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + + offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1); + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; + } else if constexpr (std::is_same::value) { + // TODO: not implement + } + } + } + + // ! do we need to sync here? + __syncwarp(); + + // shared memory to global memory + DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; + uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); + uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + +#pragma unroll + for (uint32_t i = 0; i < O_smem_iters_col; i++) { +#pragma unroll + for (uint32_t j = 0; j < O_smem_iters_row; j++) { + if (O_load_idx_lane_base < qo_len) { + smem_O.store_128b(offset_O, O_lane_ptr); + } + O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); + offset_O = smem_O.advance_offset_by_column(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); + O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); + O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; + } + + if constexpr (return_lse) { + // ! this only works for num_tiles_q = 2 + uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); + float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; + uint32_t fq = (lane_id % 4) / 2; + uint32_t k = (lane_id % 4) % 2; + + if (lse_idx < qo_len) { + lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); + } + } #else - RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer"); + RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer"); #endif -} + } #undef PACK_SIZE_QK #undef PACK_SIZE_V @@ -713,6 +617,6 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr #undef MMA_SV_N #undef MMA_SV_K -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/attn_utils.cuh b/src/core/cuda/sageattention/attn_utils.cuh index f112ef26..07fa4e5a 100644 --- a/src/core/cuda/sageattention/attn_utils.cuh +++ b/src/core/cuda/sageattention/attn_utils.cuh @@ -22,8 +22,8 @@ #include "cp_async.cuh" #include "math.cuh" #include "mma.cuh" -#include "permuted_smem.cuh" #include "numeric_conversion.cuh" +#include "permuted_smem.cuh" namespace sageattention { @@ -35,965 +35,780 @@ namespace sageattention { #define SAGE_DIV_CEIL(M, N) (((M) + (N)-1) / (N)) -enum class MaskMode { - kNone = 0, - kCausal = 1, -}; + enum class MaskMode { + kNone = 0, + kCausal = 1, + }; -enum class DataType { - kHalf, - kInt8, - kInt4, - kE4M3, - kE5M2, -}; + enum class DataType { + kHalf, + kInt8, + kInt4, + kE4M3, + kE5M2, + }; -enum class QuantGranularity { - kPerTensor = 0, - kPerBlock = 1, - kPerWarp = 2, - kPerThread = 3, -}; + enum class QuantGranularity { + kPerTensor = 0, + kPerBlock = 1, + kPerWarp = 2, + kPerThread = 3, + }; -enum class ComputeUnit { - kTensorCore, - kCudaCore, -}; + enum class ComputeUnit { + kTensorCore, + kCudaCore, + }; -__device__ __forceinline__ uint32_t get_warp_id() -{ - return threadIdx.y; -} - -__device__ __forceinline__ uint32_t get_lane_id() -{ - return threadIdx.x; -} - -template -__device__ __forceinline__ uint32_t get_warp_idx_q() -{ - return get_warp_id() / num_warps_k; -} - -template -__device__ __forceinline__ uint32_t get_warp_idx_k() -{ - return get_warp_id() % num_warps_k; -} - -template -__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset, - const uint32_t &gmem_stride, - const smem_t &smem) -{ - static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); - static_assert(std::is_same::value || std::is_same::value); - - constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; - -#pragma unroll - for (uint32_t i = 0; i < smem_iters_col; i++) - { -#pragma unroll - for (uint32_t j = 0; j < smem_iters_row; j++) - { - smem.load_128b_async(smem_offset, *lane_ptr); - *lane_ptr += (global_to_shared_line_lanes * pack_size); - smem_offset = smem.advance_offset_by_column(smem_offset); + __device__ __forceinline__ uint32_t get_warp_id() { + return threadIdx.y; } - smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); - *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); - } - smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); - *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; -} - -// with predicate -template -__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset, - const uint32_t &gmem_stride, - const smem_t &smem, uint32_t base_idx, uint32_t max_len) -{ - static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); - static_assert(std::is_same::value || std::is_same::value); - - constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; - -#pragma unroll - for (uint32_t i = 0; i < smem_iters_col; i++) - { -#pragma unroll - for (uint32_t j = 0; j < smem_iters_row; j++) - { - // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN. - smem.load_128b_async::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len); - *lane_ptr += (global_to_shared_line_lanes * pack_size); - smem_offset = smem.advance_offset_by_column(smem_offset); + __device__ __forceinline__ uint32_t get_lane_id() { + return threadIdx.x; } - smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); - *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); - base_idx += global_to_shared_copy_lines_per_warp_per_iter; - } - smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); - *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; -} - -template -__device__ __forceinline__ void load_fp8_V_global_to_share(int8_t **lane_ptr, uint32_t &smem_offset, - const uint32_t &gmem_stride, - const smem_t &smem) -{ - static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); - constexpr uint32_t pack_size_fp8 = 16; - -#pragma unroll - for (uint32_t i = 0; i < smem_iters_col; i++) - { -#pragma unroll - for (uint32_t j = 0; j < smem_iters_row; j++) - { - smem.load_128b_async(smem_offset, *lane_ptr); - *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8); - smem_offset = smem.advance_offset_by_column(smem_offset); + template + __device__ __forceinline__ uint32_t get_warp_idx_q() { + return get_warp_id() / num_warps_k; } - smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); - *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8)); - } - smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); - // for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; - *lane_ptr += CTA; // ! prevent underflow - *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; -} - -template -__device__ __forceinline__ void compute_int_qk(const smem_t &smem_Q, const smem_t &smem_K, int32_t RS[][num_tiles_k][8], uint32_t &offset_Q, uint32_t &offset_K) -{ - static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); - - uint32_t RQ[num_tiles_q][4]; - uint32_t RK[4]; - - // the first iteration, mma mode is kInit -#pragma unroll - for (uint32_t iter = 0; iter < 1; iter++) - { - // load RQ -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); - offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + template + __device__ __forceinline__ uint32_t get_warp_idx_k() { + return get_warp_id() % num_warps_k; } - // ! using permutation invariance - offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + + template + __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t& smem) { + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + static_assert(std::is_same::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - // load RK - smem_K.ldmatrix_m8n8x4(offset_K, RK); - offset_K = smem_K.advance_offset_by_row<16>(offset_K); - - // mma + for (uint32_t i = 0; i < smem_iters_col; i++) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (DTypeQK == DataType::kInt8) - { - mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + for (uint32_t j = 0; j < smem_iters_row; j++) { + smem.load_128b_async(smem_offset, *lane_ptr); + *lane_ptr += (global_to_shared_line_lanes * pack_size); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); } - else if constexpr (DTypeQK == DataType::kInt4) - { - mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); - } - } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; } - offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); - } - // following iteration, mma mode is kInplace + // with predicate + template + __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t& smem, uint32_t base_idx, uint32_t max_len) { + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + static_assert(std::is_same::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; + #pragma unroll - for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) - { - // load RQ + for (uint32_t i = 0; i < smem_iters_col; i++) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); - offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + for (uint32_t j = 0; j < smem_iters_row; j++) { + // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN. + smem.load_128b_async::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len); + *lane_ptr += (global_to_shared_line_lanes * pack_size); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); + base_idx += global_to_shared_copy_lines_per_warp_per_iter; + } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; } - offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + + template + __device__ __forceinline__ void load_fp8_V_global_to_share(int8_t** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t& smem) { + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + constexpr uint32_t pack_size_fp8 = 16; #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - // load RK - smem_K.ldmatrix_m8n8x4(offset_K, RK); - offset_K = smem_K.advance_offset_by_row<16>(offset_K); - - // mma + for (uint32_t i = 0; i < smem_iters_col; i++) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (DTypeQK == DataType::kInt8) - { - mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + for (uint32_t j = 0; j < smem_iters_row; j++) { + smem.load_128b_async(smem_offset, *lane_ptr); + *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8)); } - else if constexpr (DTypeQK == DataType::kInt4) - { - mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); - } - } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + // for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; + *lane_ptr += CTA; // ! prevent underflow + *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; } - offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); - } - offset_Q -= (2 * num_tiles_qk_inner); - offset_K -= (2 * num_tiles_qk_inner); -} + template + __device__ __forceinline__ void compute_int_qk(const smem_t& smem_Q, const smem_t& smem_K, int32_t RS[][num_tiles_k][8], uint32_t& offset_Q, uint32_t& offset_K) { + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); -// for case when num_tiles_qk_inner = 1 -template -__device__ __forceinline__ void compute_int_qk(const smem_t &smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) -{ - static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); - static_assert(num_tiles_qk_inner == 1); + uint32_t RQ[num_tiles_q][4]; + uint32_t RK[4]; - uint32_t RK[4]; - - // mma mode is kInit + // the first iteration, mma mode is kInit #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - // load RK - smem_K.ldmatrix_m8n8x4(offset_K, RK); - offset_K = smem_K.advance_offset_by_row<16>(offset_K); - - // mma + for (uint32_t iter = 0; iter < 1; iter++) { + // load RQ #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (DTypeQK == DataType::kInt8) - { - mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); - } - else if constexpr (DTypeQK == DataType::kInt4) - { - mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); - } + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); + offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + } + // ! using permutation invariance + offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (DTypeQK == DataType::kInt8) { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } + offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); + } + + // following iteration, mma mode is kInplace +#pragma unroll + for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) { + // load RQ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); + offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + } + offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (DTypeQK == DataType::kInt8) { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } + offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); + } + + offset_Q -= (2 * num_tiles_qk_inner); + offset_K -= (2 * num_tiles_qk_inner); } - } -} -template -__device__ __forceinline__ void apply_causal_mask(const uint32_t &Q_idx_lane_base, const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][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++) - { - const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2); - const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; - const bool out_of_boundary = (kv_idx > q_idx); + // for case when num_tiles_qk_inner = 1 + template + __device__ __forceinline__ void compute_int_qk(const smem_t& smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) { + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); + static_assert(num_tiles_qk_inner == 1); - if constexpr (std::is_same::value) - { - RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + uint32_t RK[4]; + + // mma mode is kInit +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (DTypeQK == DataType::kInt8) { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } } - else if constexpr (std::is_same::value) - { - RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); - } - } } - } -} -template -__device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t &kv_len) -{ + template + __device__ __forceinline__ void apply_causal_mask(const uint32_t& Q_idx_lane_base, const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { #pragma unroll - for (uint32_t k = 0; k < 8; k++) - { - const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; - const bool out_of_boundary = (kv_idx >= kv_len); + for (uint32_t k = 0; k < 8; k++) { + const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2); + const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; + const bool out_of_boundary = (kv_idx > q_idx); - if constexpr (std::is_same::value) - { - RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } else if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } } - else if constexpr (std::is_same::value) - { - RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); - } - } } - } -} -// for DTypeQKAccum float -template -__device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float &sm_scale) -{ - static_assert(std::is_same::value || (!use_half_o_scale)); + template + __device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t& kv_len) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t k = 0; k < 2; k++) - { - // assign the smallest value possible - float m_prev = m[fq][k]; - float m_temp = -5000000.0f; + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]), - max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5])); - m_temp = max(m_temp, m_local); - } + for (uint32_t k = 0; k < 8; k++) { + const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; + const bool out_of_boundary = (kv_idx >= kv_len); - if constexpr (!fuse_scale) - { - if constexpr (exp_offset) - { - m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET); + if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } else if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } } - else - { - m_temp *= sm_scale; - } - } - else if constexpr (exp_offset) - { - m_temp += (-S_FP8_OFFSET); - } - - // exchange element with the 4 threads in the row - m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1)); // 0 exchange with 1, 2 exchange with 3 - m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2)); // 0 exchange with 2, 1 exchange with 3 - - m[fq][k] = max(m[fq][k], m_temp); - - float o_scale = math::ptx_exp2(m_prev - m[fq][k]); - - // update denominator - d[fq][k] *= o_scale; - - half2 o_scale2; - if constexpr (use_half_o_scale) - { - o_scale2 = __floats2half2_rn(o_scale, o_scale); - } - - // update RO -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - if constexpr (std::is_same::value) - { - RO[fq][fv][k * 2 + 0] *= o_scale; - RO[fq][fv][k * 2 + 1] *= o_scale; - RO[fq][fv][k * 2 + 4] *= o_scale; - RO[fq][fv][k * 2 + 5] *= o_scale; - } - else if constexpr (std::is_same::value) - { - if constexpr (use_half_o_scale) - { - ((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2); - ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2); - } - else - { - RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale); - RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale); - RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale); - RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale); - } - } - } - - // raise RS to exponent - float negative_m = -m[fq][k]; -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - if constexpr (fuse_scale) - { - RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m); - RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m); - RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m); - RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m); - } - else - { - RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m)); - RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m)); - RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m)); - RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m)); - } - } } - } -} -template -__device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) -{ - static_assert(sizeof(T) == 4); + // for DTypeQKAccum float + template + __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float& sm_scale) { + static_assert(std::is_same::value || (!use_half_o_scale)); #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]); - ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]); - ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]); - ((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]); - } - } -} + for (uint32_t k = 0; k < 2; k++) { + // assign the smallest value possible + float m_prev = m[fq][k]; + float m_temp = -5000000.0f; +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]), + max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5])); + m_temp = max(m_temp, m_local); + } -template -__device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) -{ -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) - { - floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4); - floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6); - floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4); - floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6); - } - } -} + if constexpr (!fuse_scale) { + if constexpr (exp_offset) { + m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET); + } else { + m_temp *= sm_scale; + } + } else if constexpr (exp_offset) { + m_temp += (-S_FP8_OFFSET); + } -template -__device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) -{ -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) - { - halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2); - halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3); - halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2); - halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3); - } - } -} + // exchange element with the 4 threads in the row + m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1)); // 0 exchange with 1, 2 exchange with 3 + m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2)); // 0 exchange with 2, 1 exchange with 3 -template -__device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) -{ -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) - { - e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]); - e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1); - e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2); - e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3); - } - } -} + m[fq][k] = max(m[fq][k], m_temp); -template -__device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore)? 4 : 8], float d[][2]) -{ - // for compute unit cuda core, RS is float - // for compute unit tensor core, RS is packed half - static_assert((std::is_same::value && compute_unit == ComputeUnit::kCudaCore) || - (std::is_same::value && compute_unit == ComputeUnit::kTensorCore)); + float o_scale = math::ptx_exp2(m_prev - m[fq][k]); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { - if constexpr (compute_unit == ComputeUnit::kTensorCore) - { - // full accumulate with tensor core - mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk])); - } - else if constexpr (compute_unit == ComputeUnit::kCudaCore) - { - // partial accumulate with cuda core - d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5]; - d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7]; - } - } - } -} + // update denominator + d[fq][k] *= o_scale; -template -__device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) -{ -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { -#pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) - { - mma::rowsum_f8f8f32(d[fq], RS[fq][fk]); - } - } -} + half2 o_scale2; + if constexpr (use_half_o_scale) { + o_scale2 = __floats2half2_rn(o_scale, o_scale); + } -template -__device__ __forceinline__ void compute_fp16_sv(const smem_t &smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) -{ - uint32_t smem_V_row_base = get_warp_idx_k() * (num_tiles_k * 16) + get_lane_id() % 16; - uint32_t smem_V_col_base = get_lane_id() / 16; + // update RO #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + if constexpr (std::is_same::value) { + RO[fq][fv][k * 2 + 0] *= o_scale; + RO[fq][fv][k * 2 + 1] *= o_scale; + RO[fq][fv][k * 2 + 4] *= o_scale; + RO[fq][fv][k * 2 + 5] *= o_scale; + } else if constexpr (std::is_same::value) { + if constexpr (use_half_o_scale) { + ((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2); + ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2); + } else { + RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale); + RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale); + RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale); + RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale); + } + } + } + + // raise RS to exponent + float negative_m = -m[fq][k]; #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2); - smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV); + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + if constexpr (fuse_scale) { + RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m); + RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m); + RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m); + RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m); + } else { + RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m)); + } + } + } } - else if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV); - } - } } - } -} -template -__device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V) -{ - static_assert(sizeof(T) == 4); - - // ! be sure you know what you are doing + template + __device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) { + static_assert(sizeof(T) == 4); #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k; fk++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]); + ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]); + ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]); + ((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]); + } } - else if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); - } - } - - offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); } - offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); - } - // make offset_V their original value - offset_V -= (16 * num_tiles_k * stride); -} - -template -__device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V) -{ - static_assert(sizeof(T) == 4); - static_assert(std::is_same::value); - - uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4]; - - // ! be sure you know what you are doing + template + __device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) { #pragma unroll - for (uint32_t fk = 0; fk < 1; fk++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - { - mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) { + floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4); + floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6); + floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4); + floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6); + } } - } - - offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); } - offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); - } + template + __device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) { #pragma unroll - for (uint32_t fk = 1; fk < num_tiles_k; fk++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - { - mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) { + halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2); + halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3); + halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2); + halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3); + } } - } - - offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); } - offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); - } - // accumulate into RO + template + __device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x); - RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y); - RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x); - RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y); - RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x); - RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y); - RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x); - RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y); + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) { + e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3); + } + } } - } - // make offset_V their original value - offset_V -= (16 * num_tiles_k * stride); -} + template + __device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore) ? 4 : 8], float d[][2]) { + // for compute unit cuda core, RS is float + // for compute unit tensor core, RS is packed half + static_assert((std::is_same::value && compute_unit == ComputeUnit::kCudaCore) || + (std::is_same::value && compute_unit == ComputeUnit::kTensorCore)); -template -__device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) -{ - if constexpr (compute_unit == ComputeUnit::kCudaCore) - { - // accumulate_d performs partial accumulation with cuda core - // aggregate d #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t k = 0; k < 2; k++) - { - d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3 - d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3 - } + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { + if constexpr (compute_unit == ComputeUnit::kTensorCore) { + // full accumulate with tensor core + mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk])); + } else if constexpr (compute_unit == ComputeUnit::kCudaCore) { + // partial accumulate with cuda core + d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5]; + d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7]; + } + } + } } - } - // divide O by d - float d_rcp[num_tiles_q][2]; + template + __device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t k = 0; k < 2; k++) - { - // TODO: check m to prevent nan - d_rcp[fq][k] = math::ptx_rcp(d[fq][k]); + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) { + mma::rowsum_f8f8f32(d[fq], RS[fq][fk]); + } + } } - } + template + __device__ __forceinline__ void compute_fp16_sv(const smem_t& smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) { + uint32_t smem_V_row_base = get_warp_idx_k() * (num_tiles_k * 16) + get_lane_id() % 16; + uint32_t smem_V_col_base = get_lane_id() / 16; #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2); + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); #pragma unroll - for (uint32_t k = 0; k < 8; k++) - { - if constexpr (std::is_same::value) - { - RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2]; + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV); + } else if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV); + } + } + } } - else if constexpr (std::is_same::value) - { - RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]); - } - } } - } -} -template -__device__ __forceinline__ void compute_fp8_sv(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) -{ - uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; - // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; - uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + template + __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) { + static_assert(sizeof(T) == 4); + + // ! be sure you know what you are doing #pragma unroll - for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) - { - uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); + for (uint32_t fk = 0; fk < num_tiles_k; fk++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); - smem_V.ldmatrix_m8n8x4(offset_V, RV); + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV); + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } else if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); } - else if constexpr (std::is_same::value) - { - // ! Not Implemented - } - } - offset_V = smem_V.advance_offset_by_row<16>(offset_V); + + // make offset_V their original value + offset_V -= (16 * num_tiles_k * stride); } - } -} -template -__device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) -{ - uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; - // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; - uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + template + __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) { + static_assert(sizeof(T) == 4); + static_assert(std::is_same::value); - float RO_inst_buf[num_tiles_q][num_tiles_v][8]; + uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4]; + // ! be sure you know what you are doing #pragma unroll - for (uint32_t fk = 0; fk < 1; fk++) - { - uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); + for (uint32_t fk = 0; fk < 1; fk++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); - smem_V.ldmatrix_m8n8x4(offset_V, RV); + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); } - else if constexpr (std::is_same::value) - { - // ! Not Implemented + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k; fk++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); } - } - offset_V = smem_V.advance_offset_by_row<16>(offset_V); + + // accumulate into RO +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x); + RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y); + RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x); + RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y); + RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x); + RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y); + RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x); + RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y); + } + } + + // make offset_V their original value + offset_V -= (16 * num_tiles_k * stride); } - } + + template + __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) { + if constexpr (compute_unit == ComputeUnit::kCudaCore) { + // accumulate_d performs partial accumulation with cuda core + // aggregate d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) { + d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3 + d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3 + } + } + } + + // divide O by d + float d_rcp[num_tiles_q][2]; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) { + // TODO: check m to prevent nan + d_rcp[fq][k] = math::ptx_rcp(d[fq][k]); + } + } #pragma unroll - for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) - { - uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); - smem_V.ldmatrix_m8n8x4(offset_V, RV); + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + for (uint32_t k = 0; k < 8; k++) { + if constexpr (std::is_same::value) { + RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2]; + } else if constexpr (std::is_same::value) { + RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]); + } + } + } } - else if constexpr (std::is_same::value) - { - // ! Not Implemented - } - } - offset_V = smem_V.advance_offset_by_row<16>(offset_V); } - } + template + __device__ __forceinline__ void compute_fp8_sv(const smem_t& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) { + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; - RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; - RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; - RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; - RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; - RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; - RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; - RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; - } - } -} - -template -__device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) -{ - uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; - // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; - uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; - - uint32_t RO_int32[num_tiles_q][num_tiles_v][4]; - + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); #pragma unroll - for (uint32_t fk = 0; fk < 1; fk++) - { - uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); -#pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); - smem_V.ldmatrix_m8n8x4(offset_V, RV); -#pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - //mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); - mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } } - else if constexpr (std::is_same::value) - { - // ! Not Implemented - } - } - offset_V = smem_V.advance_offset_by_row<16>(offset_V); } - } + + template + __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) { + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + + float RO_inst_buf[num_tiles_q][num_tiles_v][8]; #pragma unroll - for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) - { - uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); + for (uint32_t fk = 0; fk < 1; fk++) { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); #pragma unroll - for (uint32_t fv = 0; fv < num_tiles_v; fv++) - { - // load RV - uint32_t RV[4]; - // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); - smem_V.ldmatrix_m8n8x4(offset_V, RV); + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); #pragma unroll - for (uint32_t fq = 0; fq < num_tiles_q; fq++) - { - if constexpr (std::is_same::value) - { - //mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); - mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } } - else if constexpr (std::is_same::value) - { - // ! Not Implemented + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } } - } - offset_V = smem_V.advance_offset_by_row<16>(offset_V); - } - } - float RO_tmp_float[2]; -#pragma unroll - for(int i = 0; i < num_tiles_q; i++){ -#pragma unroll - for(int j = 0; j < num_tiles_v; j++){ - #pragma unroll - for(int k = 0; k < 4; k++){ - unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]); - RO[i][j][k * 2 + 0] += RO_tmp_float[0]; - RO[i][j][k * 2 + 1] += RO_tmp_float[1]; - } - } - } -// #pragma unroll -// for (uint32_t fq = 0; fq < num_tiles_q; fq++) -// { -// #pragma unroll -// for (uint32_t fv = 0; fv < num_tiles_v; fv++) -// { -// RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; -// RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; -// RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; -// RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; -// RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; -// RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; -// RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; -// RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; -// } -// } -} +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; + RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; + RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; + RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; + RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; + RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; + RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; + RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; + } + } + } -} // namespace sageattention + template + __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) { + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + + uint32_t RO_int32[num_tiles_q][num_tiles_v][4]; + +#pragma unroll + for (uint32_t fk = 0; fk < 1; fk++) { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + // mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) { + if constexpr (std::is_same::value) { + // mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + float RO_tmp_float[2]; +#pragma unroll + for (int i = 0; i < num_tiles_q; i++) { +#pragma unroll + for (int j = 0; j < num_tiles_v; j++) { +#pragma unroll + for (int k = 0; k < 4; k++) { + unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]); + RO[i][j][k * 2 + 0] += RO_tmp_float[0]; + RO[i][j][k * 2 + 1] += RO_tmp_float[1]; + } + } + } + + // #pragma unroll + // for (uint32_t fq = 0; fq < num_tiles_q; fq++) + // { + // #pragma unroll + // for (uint32_t fv = 0; fv < num_tiles_v; fv++) + // { + // RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; + // RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; + // RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; + // RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; + // RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; + // RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; + // RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; + // RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; + // } + // } + } + +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/cp_async.cuh b/src/core/cuda/sageattention/cp_async.cuh index 51cf59d5..98ceca5b 100644 --- a/src/core/cuda/sageattention/cp_async.cuh +++ b/src/core/cuda/sageattention/cp_async.cuh @@ -24,17 +24,17 @@ namespace sageattention { -namespace cp_async { + namespace cp_async { -enum class SharedMemFillMode { - kFillZero, // Fill zero to shared memory when predicate is false - kNoFill // Do not fill zero to shared memory when predicate is false -}; + enum class SharedMemFillMode { + kFillZero, // Fill zero to shared memory when predicate is false + kNoFill // Do not fill zero to shared memory when predicate is false + }; -enum class PrefetchMode { - kNoPrefetch, // Do not fetch additional data from global memory to L2 - kPrefetch // Fetch additional data from global memory to L2 -}; + enum class PrefetchMode { + kNoPrefetch, // Do not fetch additional data from global memory to L2 + kPrefetch // Fetch additional data from global memory to L2 + }; #if (__CUDACC_VER_MAJOR__ >= 11) #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) @@ -42,107 +42,107 @@ enum class PrefetchMode { #endif #endif -/*! - * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted - * cp.async instructions to a group - */ -__device__ __forceinline__ void commit_group() { + /*! + * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted + * cp.async instructions to a group + */ + __device__ __forceinline__ void commit_group() { #ifdef CP_ASYNC_ENABLED - asm volatile("cp.async.commit_group;\n" ::); + asm volatile("cp.async.commit_group;\n" ::); #endif -} + } -/*! - * \brief Wrapper of PTX cp.async.wait_group instruction - * \tparam n Wait till most recent n groups are committed - */ -template -__device__ __forceinline__ void wait_group() { + /*! + * \brief Wrapper of PTX cp.async.wait_group instruction + * \tparam n Wait till most recent n groups are committed + */ + template + __device__ __forceinline__ void wait_group() { #ifdef CP_ASYNC_ENABLED - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); #endif -} + } -/*! - * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from - * global memory to shared memory - * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 - * \tparam T Data type - * \param smem_ptr Pointer to shared memory - * \param gmem_ptr Pointer to global memory - */ -template -__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { + /*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + */ + template + __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { #ifdef CP_ASYNC_ENABLED - uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { - asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), - "l"(gmem_ptr), "n"(16), "r"(16)); - } else { - asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), - "l"(gmem_ptr), "n"(16), "r"(16)); - } + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } #else - *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); #endif -} + } -/*! - * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from - * global memory to shared memory with predicate. - * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 - * \tparam fill_mode Whether to fill zero to shared memory when predicate is false - * \tparam T Data type - * \param smem_ptr Pointer to shared memory - * \param gmem_ptr Pointer to global memory - * \param predicate Predicate value - * \note fill zero is slower than not fill zero - */ -template -__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { + /*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory with predicate. + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam fill_mode Whether to fill zero to shared memory when predicate is false + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + * \param predicate Predicate value + * \note fill zero is slower than not fill zero + */ + template + __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { #ifdef CP_ASYNC_ENABLED - uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (fill_mode == SharedMemFillMode::kFillZero) { - int src_in_bytes = predicate ? 16 : 0; - if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { - asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), - "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); - } else { - asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), - "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); - } - } else { - if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" - "}\n" ::"r"((int)predicate), - "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); - } else { - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)predicate), - "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); - } - } + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + int src_in_bytes = predicate ? 16 : 0; + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } + } else { + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } else { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } + } #else - if (predicate) { - *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); - } else { - if constexpr (fill_mode == SharedMemFillMode::kFillZero) { - *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0); - } - } + if (predicate) { + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); + } else { + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0); + } + } #endif -} + } -} // namespace cp_async + } // namespace cp_async -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/math.cuh b/src/core/cuda/sageattention/math.cuh index 3233a03f..b92ac798 100644 --- a/src/core/cuda/sageattention/math.cuh +++ b/src/core/cuda/sageattention/math.cuh @@ -26,137 +26,141 @@ namespace sageattention { #ifndef USHORT_TYPE #define USHORT_TYPE -typedef unsigned short ushort; + typedef unsigned short ushort; #endif -namespace math { + namespace math { -// log2(e) -constexpr float log2e = 1.44269504088896340736f; -constexpr float log2e_recp = 1.0f / log2e; + // log2(e) + constexpr float log2e = 1.44269504088896340736f; + constexpr float log2e_recp = 1.0f / log2e; -__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; } + __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { + return *(half2*)&x; + } -__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; } + __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { + return *(uint32_t*)&x; + } -/*! - * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x - * \param x input - */ -__forceinline__ __device__ float ptx_exp2(float x) { - float y; - asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} + /*! + * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x + * \param x input + */ + __forceinline__ __device__ float ptx_exp2(float x) { + float y; + asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } -/*! - * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) - * \param x input - */ -__forceinline__ __device__ float ptx_log2(float x) { - float y; - asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} + /*! + * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) + * \param x input + */ + __forceinline__ __device__ float ptx_log2(float x) { + float y; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } -/*! - * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x - * \param x input - */ -__forceinline__ __device__ half2 ptx_exp2(half2 x) { - uint32_t y_u32; - uint32_t x_u32 = half2_as_uint32(x); - asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); - return uint32_as_half2(y_u32); -} + /*! + * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x + * \param x input + */ + __forceinline__ __device__ half2 ptx_exp2(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); + } -/*! - * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x - * \param x input - */ -__forceinline__ __device__ half ptx_exp2(half x) { - ushort y_u16; - asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); - return __ushort_as_half(y_u16); -} + /*! + * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x + * \param x input + */ + __forceinline__ __device__ half ptx_exp2(half x) { + ushort y_u16; + asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); + } -/*! - * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x - * \param x input - */ -__forceinline__ __device__ float ptx_rcp(float x) { - float y; - asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} + /*! + * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x + * \param x input + */ + __forceinline__ __device__ float ptx_rcp(float x) { + float y; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } -/*! - * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle - * between threads in a warp. - * \param x The value in the source lane - * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta] - */ -__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) { - float y; - asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;" - : "=f"(y) - : "f"(x), "r"(lane_mask)); - return y; -} + /*! + * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle + * between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta] + */ + __forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) { + float y; + asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;" + : "=f"(y) + : "f"(x), "r"(lane_mask)); + return y; + } -/*! - * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly - * shuffle between threads in a warp. - * \param x The value in the source lane - * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask] - */ -__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) { - return __shfl_xor_sync(0xffffffff, x, lane_mask); -} + /*! + * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly + * shuffle between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask] + */ + __forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) { + return __shfl_xor_sync(0xffffffff, x, lane_mask); + } -/*! - * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x) - * \param x input - */ -__forceinline__ __device__ float rsqrt(float x) { - float y; - asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} + /*! + * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x) + * \param x input + */ + __forceinline__ __device__ float rsqrt(float x) { + float y; + asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } -/*! - * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x) - * \param x input - */ -__forceinline__ __device__ float tanh(float x) { - float y; - asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -} + /*! + * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x) + * \param x input + */ + __forceinline__ __device__ float tanh(float x) { + float y; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; + } -/*! - * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x) - * \param x input - */ -__forceinline__ __device__ half2 tanh(half2 x) { - uint32_t y_u32; - uint32_t x_u32 = half2_as_uint32(x); - asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); - return uint32_as_half2(y_u32); -} + /*! + * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x) + * \param x input + */ + __forceinline__ __device__ half2 tanh(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); + } -/*! - * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x) - * \param x input - */ -__forceinline__ __device__ half tanh(half x) { - ushort y_u16; - asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); - return __ushort_as_half(y_u16); -} + /*! + * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x) + * \param x input + */ + __forceinline__ __device__ half tanh(half x) { + ushort y_u16; + asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); + } -} // namespace math + } // namespace math -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/mma.cuh b/src/core/cuda/sageattention/mma.cuh index 03c4e69e..4de48286 100644 --- a/src/core/cuda/sageattention/mma.cuh +++ b/src/core/cuda/sageattention/mma.cuh @@ -26,7 +26,7 @@ namespace sageattention { -namespace mma{ + namespace mma { #if (__CUDACC_VER_MAJOR__ >= 11) #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) @@ -62,668 +62,622 @@ namespace mma{ #define RUNTIME_ASSERT(x) assert(0 && x) #endif -enum class MMAMode { - kInit = 0U, - kInplaceUpdate = 1U, -}; + enum class MMAMode { + kInit = 0U, + kInplaceUpdate = 1U, + }; -/*! - * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory - * to fragment - * \tparam T data type of the fragment - * \param R pointer to the fragment - * \param smem_ptr pointer to the shared memory - */ -template -__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) { + /*! + * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ + template + __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) { #ifdef LDMATRIX_M8N8X2_ENABLED - uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" - : "=r"(R[0]), "=r"(R[1]) - : "r"(smem_int_ptr)); + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); #endif -} + } -/*! - * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory - * to fragment - * \tparam T data type of the fragment - * \param R pointer to the fragment - * \param smem_ptr pointer to the shared memory - */ -template -__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { + /*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ + template + __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { #ifdef LDMATRIX_M8N8X4_ENABLED - uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" - : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) - : "r"(smem_int_ptr)); + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); #endif -} + } -/*! - * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from - * shared memory to fragment and transposes the fragment - * \tparam T data type of the fragment - * \param R pointer to the fragment - * \param smem_ptr pointer to the shared memory - */ -template -__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) { + /*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from + * shared memory to fragment and transposes the fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ + template + __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) { #ifdef LDMATRIX_M8N8X4_ENABLED - uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" - : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) - : "r"(smem_int_ptr)); + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix - * multiplication, accumulated in f32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F16F16F32_M16N8K16_ENABLED - // ! only support half dtype now - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), - "f"(C[2]), "f"(C[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); - } + // ! only support half dtype now + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix - * multiplication, accumulated in f32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F16F16F32_M16N8K16_ENABLED - // ! only support half dtype now - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), - "f"(C[2]), "f"(C[3])); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), - "f"(C[6]), "f"(C[7])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); - } + // ! only support half dtype now + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix - * multiplication, accumulated in f16. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F16F16F16_M16N8K16_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[0]), "=r"(C[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[0]), "=r"(C[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix - * multiplication, accumulated in f16. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F16F16F16_M16N8K16_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[0]), "=r"(C[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[0]), "=r"(C[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix - * multiplication, accumulated in int32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_S8S8S32_M16N8K32_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), - "r"(C[2]), "r"(C[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), - "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), + "r"(C[2]), "r"(C[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix - * multiplication, accumulated in int32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_S8S8S32_M16N8K32_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), - "r"(C[2]), "r"(C[3])); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), - "r"(C[6]), "r"(C[7])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), - "r"(0), "r"(0)); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), - "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), + "r"(C[2]), "r"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), + "r"(C[6]), "r"(C[7])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), + "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix - * multiplication, accumulated in int32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_S4S4S32_M16N8K64_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), - "r"(C[2]), "r"(C[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), - "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), + "r"(C[2]), "r"(C[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix - * multiplication, accumulated in int32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) { #ifdef MMA_S4S4S32_M16N8K64_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), - "r"(C[2]), "r"(C[3])); - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), - "r"(C[6]), "r"(C[7])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), - "r"(0), "r"(0)); - asm volatile( - "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), - "r"(0), "r"(0)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), + "r"(C[2]), "r"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), + "r"(C[6]), "r"(C[7])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0), + "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix - * multiplication, accumulated in fp32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix + * multiplication, accumulated in fp32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F8F8F32_M16N8K16_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), - "f"(C[2]), "f"(C[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); - } + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix - * multiplication, accumulated in fp16. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, - uint32_t* B) { - //uint32_t* C_uint32 = reinterpret_cast(C); + /*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix + * multiplication, accumulated in fp16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) { + // uint32_t* C_uint32 = reinterpret_cast(C); #ifdef MMA_F8F8F16_M16N8K16_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C_uint32[0]), "=r"(C_uint32[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1])); + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[0]), "=r"(C_uint32[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1])); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C_uint32[2]), "=r"(C_uint32[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C_uint32[0]), "=r"(C_uint32[1]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[2]), "=r"(C_uint32[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[0]), "=r"(C_uint32[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " - "{%0, %1}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, %9};\n" - : "=r"(C_uint32[2]), "=r"(C_uint32[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); - } + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[2]), "=r"(C_uint32[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } - - -/*! - * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix - * multiplication, accumulated in fp32. - * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate - * \param C pointer to the accumulator - * \param A pointer to the fragment of matrix A - * \param B pointer to the fragment of matrix B - */ -template -__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, - uint32_t* B) { + /*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix + * multiplication, accumulated in fp32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ + template + __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) { #ifdef MMA_F8F8F32_M16N8K16_ENABLED - if constexpr (mma_mode == MMAMode::kInplaceUpdate) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), - "f"(C[2]), "f"(C[3])); + if constexpr (mma_mode == MMAMode::kInplaceUpdate) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), - "f"(C[6]), "f"(C[7])); - } - else if constexpr (mma_mode == MMAMode::kInit) - { - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } else if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) - : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), - "f"(0.f), "f"(0.f)); - } + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Use mma instructions to compute rowsum. - */ -__device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) { + /*! + * \brief Use mma instructions to compute rowsum. + */ + __device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) { #ifdef MMA_F16F16F32_M16N8K16_ENABLED - asm volatile( - "{\n" - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0, _, %1, _}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, 0., %9, 0.};\n" - "}\n" - : "=f"(d[0]), "=f"(d[1]) - : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320), // 1006648320 packs two 1.0f in half precision - "r"(1006648320), "f"(d[0]), "f"(d[1])); + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320), // 1006648320 packs two 1.0f in half precision + "r"(1006648320), "f"(d[0]), "f"(d[1])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -/*! - * \brief Use mma instructions to compute rowsum. - */ -__device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) { + /*! + * \brief Use mma instructions to compute rowsum. + */ + __device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) { #ifdef MMA_F8F8F32_M16N8K16_ENABLED - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0, _, %1, _}," - "{%2, %3, %4, %5}," - "{%6, %7}," - "{%8, 0., %9, 0.};\n" - : "=f"(d[0]), "=f"(d[1]) - : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504), // 943208504 packs four 1.0f in e4m3 - "f"(d[0]), "f"(d[1])); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504), // 943208504 packs four 1.0f in e4m3 + "f"(d[0]), "f"(d[1])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); #endif -} + } -} // namespace mma + } // namespace mma -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/numeric_conversion.cuh b/src/core/cuda/sageattention/numeric_conversion.cuh index 7049bacb..066f7fe0 100644 --- a/src/core/cuda/sageattention/numeric_conversion.cuh +++ b/src/core/cuda/sageattention/numeric_conversion.cuh @@ -38,118 +38,111 @@ namespace sageattention { #define RUNTIME_ASSERT(x) assert(0 && x) #endif -__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) { - uint16_t h0 = source & 0xFFFF; - uint16_t h1 = (source >> 16) & 0xFFFF; - asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0)); - asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1)); -} + __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) { + uint16_t h0 = source & 0xFFFF; + uint16_t h1 = (source >> 16) & 0xFFFF; + asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0)); + asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1)); + } -__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t *dest, float *source0, float *source1) -{ + __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo;\n" \ - ".reg .b16 hi;\n" \ - "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" \ - "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" \ - "mov.b32 %0, {lo, hi};\n" \ - "}" \ - : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1])); + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + "mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t *dest, float *source0, float *source1) -{ + __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo;\n" \ - ".reg .b16 hi;\n" \ - "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \ - "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \ - "mov.b32 %0, {lo, hi};\n" \ - "}" \ - : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1])); + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" + "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" + "mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) -{ + __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo;\n" \ - ".reg .b16 hi;\n" \ - "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \ - "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \ - "mov.b32 %0, {lo, hi};\n" \ - "}" \ - : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" + "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" + "mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) -{ + __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo;\n" \ - ".reg .b16 hi;\n" \ - "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \ - "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \ - "mov.b32 %0, {lo, hi};\n" \ - "}" \ - : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); + asm volatile( + "{\n" + ".reg .b16 lo;\n" + ".reg .b16 hi;\n" + "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" + "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" + "mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) -{ + __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo, hi;\n" \ - "mov.b32 {lo, hi}, %2;\n" \ - "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ - "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ - "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); + asm volatile( + "{\n" + ".reg .b16 lo, hi;\n" + "mov.b32 {lo, hi}, %2;\n" + "cvt.rn.f16x2.e4m3x2 %0, lo;\n" + "cvt.rn.f16x2.e4m3x2 %1, hi;\n" + "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) -{ + __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) { #ifdef FP8_CAST_ENABLED - asm volatile( \ - "{\n" \ - ".reg .b16 lo, hi;\n" \ - "mov.b32 {lo, hi}, %2;\n" \ - "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ - "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ - "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); + asm volatile( + "{\n" + ".reg .b16 lo, hi;\n" + "mov.b32 {lo, hi}, %2;\n" + "cvt.rn.f16x2.e5m2x2 %0, lo;\n" + "cvt.rn.f16x2.e5m2x2 %1, hi;\n" + "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); #else - RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); #endif -} + } -__device__ __forceinline__ int8_t float_to_int8_rn(float x) -{ - uint32_t dst; - asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); - return reinterpret_cast(dst); -} + __device__ __forceinline__ int8_t float_to_int8_rn(float x) { + uint32_t dst; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); + return reinterpret_cast(dst); + } -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/cuda/sageattention/permuted_smem.cuh b/src/core/cuda/sageattention/permuted_smem.cuh index f44be703..37ad7704 100644 --- a/src/core/cuda/sageattention/permuted_smem.cuh +++ b/src/core/cuda/sageattention/permuted_smem.cuh @@ -23,180 +23,179 @@ #include #include - #include "cp_async.cuh" #include "mma.cuh" namespace sageattention { -enum class SwizzleMode { - k32B, // for k32B mode, a line of shared memory must have 32B (16 half value) - k64B, // for k64B mode, a line of shared memory must have 64B (32 half value) - k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B. -}; + enum class SwizzleMode { + k32B, // for k32B mode, a line of shared memory must have 32B (16 half value) + k64B, // for k64B mode, a line of shared memory must have 64B (32 half value) + k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B. + }; -// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth -using b128_t = uint4; + // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth + using b128_t = uint4; -/*! - * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure - * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through - * permutation in the granularity of 128bit. - * - * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless, - * meaning it does not maintain any information about the current pointer position. The offset returnd by - * the struct can be used to access the shared memory through the provided interface. - * - * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j). - * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp. - * Under this assumption, we do not permute for k32B mode. - */ -template -struct smem_t { - // The base pointer. - b128_t* base; - // How many b128_t value a row contains - // uint32_t stride; + /*! + * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure + * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through + * permutation in the granularity of 128bit. + * + * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless, + * meaning it does not maintain any information about the current pointer position. The offset returnd by + * the struct can be used to access the shared memory through the provided interface. + * + * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j). + * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp. + * Under this assumption, we do not permute for k32B mode. + */ + template + struct smem_t { + // The base pointer. + b128_t* base; + // How many b128_t value a row contains + // uint32_t stride; - __device__ __forceinline__ smem_t() : base(nullptr) {} - template - __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) { - if constexpr (swizzle_mode == SwizzleMode::k128B) { - static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode"); - } else if constexpr (swizzle_mode == SwizzleMode::k64B) { - static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode"); - } else if constexpr (swizzle_mode == SwizzleMode::k32B) { - static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode"); - } else { - static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode"); - } - } + __device__ __forceinline__ smem_t() : base(nullptr) {} + template + __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode"); + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode"); + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode"); + } else { + static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode"); + } + } - /*! - * \brief Set the base pointer. - */ - template - __device__ __forceinline__ void set_base(T* new_base) { - base = (b128_t*)new_base; - } + /*! + * \brief Set the base pointer. + */ + template + __device__ __forceinline__ void set_base(T* new_base) { + base = (b128_t*)new_base; + } - /*! - * \brief Compute the element offset given coordinates in a permuted shared memory. - * \param i The row index. - * \param j The column index. - */ - static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t &i, const uint32_t &j) { - if constexpr (swizzle_mode == SwizzleMode::k128B) { - return i * stride + (j ^ (i % 8)); - } else if constexpr (swizzle_mode == SwizzleMode::k64B) { - return i * stride + (j ^ ((i / 2) % 4)); - } else if constexpr (swizzle_mode == SwizzleMode::k32B) { - return i * stride + j; - } - } + /*! + * \brief Compute the element offset given coordinates in a permuted shared memory. + * \param i The row index. + * \param j The column index. + */ + static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + return i * stride + (j ^ (i % 8)); + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + return i * stride + (j ^ ((i / 2) % 4)); + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return i * stride + j; + } + } - /*! - * \tparam step_size The step size to advance the offset in the permuted shared memory. - * \param offset The current offset. - */ - template - static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset) { - if constexpr (swizzle_mode == SwizzleMode::k128B) { - static_assert(step_size % 8 == 0, - "Unsupported step size"); - return offset + step_size; - } else if constexpr (swizzle_mode == SwizzleMode::k64B) { - static_assert(step_size == 4, "Unsupported step size"); - return offset + step_size; - } else if constexpr (swizzle_mode == SwizzleMode::k32B) { - static_assert(step_size == 2, "Unsupported step size"); - return offset + step_size; - } - } + /*! + * \tparam step_size The step size to advance the offset in the permuted shared memory. + * \param offset The current offset. + */ + template + static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size % 8 == 0, + "Unsupported step size"); + return offset + step_size; + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 4, "Unsupported step size"); + return offset + step_size; + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + static_assert(step_size == 2, "Unsupported step size"); + return offset + step_size; + } + } - // ! use with care - template - static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset, const uint32_t &step_idx) { - if constexpr (swizzle_mode == SwizzleMode::k128B) { - static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0, - "Unsupported step size"); - if constexpr (step_size == 2) { - return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8; - } else if constexpr (step_size == 4) { - return (offset ^ 0x4) + (step_idx % 2 == 1) * 8; - } else { - // step_size % 8 == 0 - return offset + step_size; - } - } else if constexpr (swizzle_mode == SwizzleMode::k64B) { - static_assert(step_size == 2 || step_size == 4, "Unsupported step size"); - if constexpr (step_size == 2) { - return (offset ^ 0x2) + (step_idx % 2 == 1) * 4; - } else { - return offset + step_size; - } - } else if constexpr (swizzle_mode == SwizzleMode::k32B) { - return offset + step_size; - } - } + // ! use with care + template + static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0, + "Unsupported step size"); + if constexpr (step_size == 2) { + return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8; + } else if constexpr (step_size == 4) { + return (offset ^ 0x4) + (step_idx % 2 == 1) * 8; + } else { + // step_size % 8 == 0 + return offset + step_size; + } + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 2 || step_size == 4, "Unsupported step size"); + if constexpr (step_size == 2) { + return (offset ^ 0x2) + (step_idx % 2 == 1) * 4; + } else { + return offset + step_size; + } + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return offset + step_size; + } + } - template - static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t &offset) { - if constexpr (swizzle_mode == SwizzleMode::k128B) { - static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); - if constexpr (step_size == 4) { - return (offset ^ 0x4) + step_size * stride; - } else { - // step_size % 8 == 0 - return offset + step_size * stride; - } - } else if constexpr (swizzle_mode == SwizzleMode::k64B) { - static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); - if constexpr (step_size == 4) { - return (offset ^ 0x2) + step_size * stride; - } else { - // step_size % 8 == 0 - return offset + step_size * stride; - } - } else if constexpr (swizzle_mode == SwizzleMode::k32B) { - return offset + step_size * stride; - } - } + template + static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x4) + step_size * stride; + } else { + // step_size % 8 == 0 + return offset + step_size * stride; + } + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x2) + step_size * stride; + } else { + // step_size % 8 == 0 + return offset + step_size * stride; + } + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return offset + step_size * stride; + } + } - __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t &offset, uint32_t* R) const { - b128_t* smem_ptr = base + offset; - mma::ldmatrix_m8n8x2(R, smem_ptr); - } + __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x2(R, smem_ptr); + } - __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t &offset, uint32_t* R) const { - b128_t* smem_ptr = base + offset; - mma::ldmatrix_m8n8x4(R, smem_ptr); - } + __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4(R, smem_ptr); + } - __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t &offset, uint32_t* R) const { - b128_t* smem_ptr = base + offset; - mma::ldmatrix_m8n8x4_trans(R, smem_ptr); - } + __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_trans(R, smem_ptr); + } - template - __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr, bool predicate) const { - b128_t* smem_ptr = base + offset; - cp_async::pred_load_128b( - smem_ptr, reinterpret_cast(gptr), predicate); - } + template + __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const { + b128_t* smem_ptr = base + offset; + cp_async::pred_load_128b( + smem_ptr, reinterpret_cast(gptr), predicate); + } - template - __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const { - b128_t* smem_ptr = base + offset; - cp_async::load_128b(smem_ptr, reinterpret_cast(gptr)); - } + template + __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const { + b128_t* smem_ptr = base + offset; + cp_async::load_128b(smem_ptr, reinterpret_cast(gptr)); + } - template - __device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const { - *reinterpret_cast(gptr) = *(base + offset); - } -}; + template + __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const { + *reinterpret_cast(gptr) = *(base + offset); + } + }; -} // namespace sageattention +} // namespace sageattention #endif diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index fe6b48f1..4a97704d 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -4,9 +4,9 @@ #include #include +#include "core/ggml_extensions.h" #include "ggml-backend.h" #include "ggml.h" -#include "core/ggml_extensions.h" #define EPS 1e-05f diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index 04505f13..28f426cd 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -613,7 +613,7 @@ namespace sd_vulkan_i8 { auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst; ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}}; auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows, - std::initializer_list constants, uint32_t row_offset, uint32_t dispatch_rows) { + std::initializer_list constants, uint32_t row_offset, uint32_t dispatch_rows) { if (!dispatch_rows) return GGML_STATUS_SUCCESS; const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n; diff --git a/src/core/vulkan/shaders/modulate_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp index 731bfd04..3e4b3a21 100644 --- a/src/core/vulkan/shaders/modulate_convrot.comp +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -2,16 +2,28 @@ #extension GL_EXT_control_flow_attributes : require -layout(constant_id = 0) const uint K = 4096; +layout(constant_id = 0) const uint K = 4096; layout(constant_id = 1) const bool AFFINE = false; layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer X { float data_x[]; }; -layout(binding = 1) readonly buffer M { float data_m[]; }; -layout(binding = 2) readonly buffer H { float data_h[]; }; -layout(binding = 3) writeonly buffer D { uint data_d[]; }; -layout(binding = 4) writeonly buffer S { float data_s[]; }; +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer M { + float data_m[]; +}; +layout(binding = 2) readonly buffer H { + float data_h[]; +}; +layout(binding = 3) writeonly buffer D { + uint data_d[]; +}; +layout(binding = 4) writeonly buffer S { + float data_s[]; +}; #ifdef RMS_AFFINE -layout(binding = 5) readonly buffer W { float data_w[]; }; +layout(binding = 5) readonly buffer W { + float data_w[]; +}; #endif layout(push_constant) uniform parameter { uint rows; @@ -21,7 +33,8 @@ layout(push_constant) uniform parameter { float eps; uint k; uint w_offset; -} p; +} +p; shared float values[K]; shared vec2 sums[512]; @@ -36,7 +49,7 @@ float rms_row(uint row, uint tid) { [[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) { float x = 0.0; if (col < p.k) { - x = data_x[p.x_offset + row * K + col]; + x = data_x[p.x_offset + row * K + col]; values[col] = x; } sum += x * x; @@ -63,30 +76,31 @@ void main() { for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) { #ifdef RMS_AFFINE float inv_std = rms_row(row, tid); - float mean = 0.0; + float mean = 0.0; #else vec2 sum = vec2(0.0); [[unroll]] for (uint col = tid; col < K; col += 512) { const float x = data_x[p.x_offset + row * K + col]; - values[col] = x; + values[col] = x; if (!AFFINE) { sum.x += x; sum.y = fma(x, x, sum.y); } } - float mean = 0.0; + float mean = 0.0; float inv_std = 1.0; if (!AFFINE) { sums[tid] = sum; barrier(); [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { - if (tid < stride) sums[tid] += sums[tid + stride]; + if (tid < stride) + sums[tid] += sums[tid + stride]; barrier(); } mean = sums[0].x / p.k; // Match NORM's contraction even when precise propagates from the quantized output. const float variance = fma(-mean, mean, sums[0].y / p.k); - inv_std = inversesqrt(variance + p.eps); + inv_std = inversesqrt(variance + p.eps); } #endif [[unroll]] for (uint col = tid; col < K; col += 512) { @@ -96,7 +110,7 @@ void main() { #endif if (AFFINE) { precise float modulated = values[col] * data_m[p.m_offset + col]; - value = values[col] + modulated; + value = values[col] + modulated; value += data_h[p.h_offset + col]; } else { value = (values[col] - mean) * inv_std; @@ -107,15 +121,15 @@ void main() { barrier(); [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) { for (uint i = tid; i < K / 4; i += 512) { - const uint group = i / 64; - const uint lane = i % 64; - const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride; - const float a = values[base]; - const float b = values[base + stride]; - const float c = values[base + 2 * stride]; - const float d = values[base + 3 * stride]; - values[base] = a + b + c - d; - values[base + stride] = a + b - c + d; + const uint group = i / 64; + const uint lane = i % 64; + const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride; + const float a = values[base]; + const float b = values[base + stride]; + const float c = values[base + 2 * stride]; + const float d = values[base + 3 * stride]; + values[base] = a + b + c - d; + values[base + stride] = a + b - c + d; values[base + 2 * stride] = a - b + c + d; values[base + 3 * stride] = -a + b + c + d; } @@ -133,16 +147,13 @@ void main() { } barrier(); } - const float scale = sums[0].x / 127.0; + const float scale = sums[0].x / 127.0; const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale; if (tid == 0) { data_s[row] = scale; } for (uint i = tid; i < K / 4; i += 512) { - data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) - | (quantize(values[i * 4 + 1], inv_scale) << 8) - | (quantize(values[i * 4 + 2], inv_scale) << 16) - | (quantize(values[i * 4 + 3], inv_scale) << 24); + data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24); } barrier(); } diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp index af9df1f3..7c0292e3 100644 --- a/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp @@ -5,43 +5,55 @@ layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform parameter { uint n; uint rows; uint k; uint has_bias; uint row_offset; -} p; +} +p; -layout(binding = 0, std430) readonly buffer W { int data_w[]; }; -layout(binding = 1, std430) readonly buffer A { int data_a[]; }; -layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; -layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; -layout(binding = 4, std430) readonly buffer B { float data_b[]; }; -layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; +layout(binding = 0, std430) readonly buffer W { + int data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 5, std430) writeonly buffer D { + float data_d[]; +}; shared int weights[16 * 8]; shared int activations[16 * 8]; void main() { const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x; - const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y; - const uint lx = gl_LocalInvocationID.x; - const uint ly = gl_LocalInvocationID.y; - const uint k_packed = p.k / 4; - int sum = 0; + const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y; + const uint lx = gl_LocalInvocationID.x; + const uint ly = gl_LocalInvocationID.y; + const uint k_packed = p.k / 4; + int sum = 0; for (uint k_base = 0; k_base < k_packed; k_base += 8) { if (ly < 8) { weights[lx * 8 + ly] = output_index < p.n - ? data_w[output_index * k_packed + k_base + ly] - : 0; + ? data_w[output_index * k_packed + k_base + ly] + : 0; } if (lx < 8) { activations[ly * 8 + lx] = row < p.rows - ? data_a[row * k_packed + k_base + lx] - : 0; + ? data_a[row * k_packed + k_base + lx] + : 0; } barrier(); diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp index 50bc297a..573e257d 100644 --- a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp @@ -9,35 +9,47 @@ #extension GL_KHR_shader_subgroup_basic : require layout(constant_id = 0) const uint BLOCK_SIZE = 32; -layout(constant_id = 1) const uint CM_M = 16; -layout(constant_id = 2) const uint CM_N = 16; -layout(constant_id = 3) const uint CM_K = 16; -const uint TILE_M = 4; -const uint TILE_N = 4; +layout(constant_id = 1) const uint CM_M = 16; +layout(constant_id = 2) const uint CM_N = 16; +layout(constant_id = 3) const uint CM_K = 16; +const uint TILE_M = 4; +const uint TILE_N = 4; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform parameter { uint n; uint rows; uint k; uint has_bias; uint row_offset; -} p; +} +p; -layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; -layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; -layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; -layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; -layout(binding = 4, std430) readonly buffer B { float data_b[]; }; -layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; +layout(binding = 0, std430) readonly buffer W { + int8_t data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int8_t data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 5, std430) writeonly buffer D { + float data_d[]; +}; shared int32_t accumulators[1024]; void main() { const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N; - const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M; + const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M; coopmat sums[TILE_M * TILE_N]; [[unroll]] for (uint i = 0; i < TILE_M * TILE_N; ++i) { @@ -64,7 +76,7 @@ void main() { [[unroll]] for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) { const uint index = tile_m * TILE_N + tile_n; - sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]); + sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]); } } } @@ -79,9 +91,9 @@ void main() { gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) { - const uint row = row_base + tile_m * CM_M + index / CM_N; + const uint row = row_base + tile_m * CM_M + index / CM_N; const uint output_index = output_base + tile_n * CM_N + index % CM_N; - float value = float(accumulators[index]) * data_as[row] * data_ws[output_index]; + float value = float(accumulators[index]) * data_as[row] * data_ws[output_index]; if (p.has_bias != 0) { value += data_b[output_index]; } diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp index d268f7ce..3cb4547c 100644 --- a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp @@ -7,15 +7,14 @@ #extension GL_NV_cooperative_matrix2 : require #extension GL_KHR_memory_scope_semantics : require -layout(constant_id = 0) const uint BN = 128; +layout(constant_id = 0) const uint BN = 128; layout(constant_id = 1) const uint ALIGNMENT = 0; -layout(constant_id = 2) const bool SWIGLU = false; +layout(constant_id = 2) const bool SWIGLU = false; layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; layout(constant_id = 3) const uint BM = 128; layout(constant_id = 4) const uint BK = 64; -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform parameter { uint n; uint rows; uint k; @@ -23,25 +22,45 @@ layout(push_constant) uniform parameter uint row_offset; float output_scale; uint group_cols; -} p; +} +p; -layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; -layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; -layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; -layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; -layout(binding = 4, std430) readonly buffer B { float data_b[]; }; +layout(binding = 0, std430) readonly buffer W { + int8_t data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int8_t data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; // Store FP32 bit patterns without changing the integer accumulator fragment layout. -layout(binding = 5, std430) writeonly buffer D { int32_t data_d[]; }; +layout(binding = 5, std430) writeonly buffer D { + int32_t data_d[]; +}; #ifdef SPLIT_SWIGLU -layout(binding = 6, std430) readonly buffer UW { int8_t data_uw[]; }; -layout(binding = 7, std430) readonly buffer US { float data_us[]; }; +layout(binding = 6, std430) readonly buffer UW { + int8_t data_uw[]; +}; +layout(binding = 7, std430) readonly buffer US { + float data_us[]; +}; #endif #ifdef RESIDUAL -layout(binding = 6, std430) readonly buffer G { float data_g[]; }; -layout(binding = 7, std430) readonly buffer R { float data_r[]; }; +layout(binding = 6, std430) readonly buffer G { + float data_g[]; +}; +layout(binding = 7, std430) readonly buffer R { + float data_r[]; +}; #endif -int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, - const in uint row_base, const in uint col_base) { +int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) { const uint r = row_base + col; const uint c = col_base + row; if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) { @@ -75,32 +94,32 @@ int32_t dequantize_swiglu(const in uint row, const in uint col, const in int32_t } void main() { - const uint k_size = p.k & ~255u; + const uint k_size = p.k & ~255u; tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); - a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size); - w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size); - d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n); - tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); + a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size); + w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size); + d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n); + tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); // Keep a weight panel in L2 while visiting the activation tiles. - const uint col_tiles = gl_NumWorkGroups.x; - const uint row_tiles = gl_NumWorkGroups.y; - const uint group_cols = clamp(p.group_cols, 1u, col_tiles); - const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x; - const uint group = linear / (group_cols * row_tiles); - const uint first_col = group * group_cols; - const uint cols = min(col_tiles - first_col, group_cols); - const uint within_group = linear - group * group_cols * row_tiles; - const uint row = p.row_offset + (within_group / cols) * BN; - const uint col = (first_col + within_group % cols) * BM; - coopmat sum = coopmat(0); + const uint col_tiles = gl_NumWorkGroups.x; + const uint row_tiles = gl_NumWorkGroups.y; + const uint group_cols = clamp(p.group_cols, 1u, col_tiles); + const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x; + const uint group = linear / (group_cols * row_tiles); + const uint first_col = group * group_cols; + const uint cols = min(col_tiles - first_col, group_cols); + const uint within_group = linear - group * group_cols * row_tiles; + const uint row = p.row_offset + (within_group / cols) * BN; + const uint col = (first_col + within_group % cols) * BM; + coopmat sum = coopmat(0); coopmat sum_up = coopmat(0); if (ALIGNMENT != 0) { tensorLayoutNV<2> a_full = createTensorLayoutNV(2); tensorLayoutNV<2> w_full = createTensorLayoutNV(2); - a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size); - w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size); + a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size); + w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size); for (uint k = 0; k < p.k; k += BK) { coopmat w; coopmat a; @@ -145,7 +164,7 @@ void main() { // The aligned dispatch contains only complete tiles and a 128-aligned output width. if (ALIGNMENT == 1) { tensorLayoutNV<2> d_full = createTensorLayoutNV(2); - d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u); + d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u); coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose); } else { coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose); diff --git a/src/core/vulkan/shaders/quantize_i8_convrot.comp b/src/core/vulkan/shaders/quantize_i8_convrot.comp index 1c69d9dd..20edd616 100644 --- a/src/core/vulkan/shaders/quantize_i8_convrot.comp +++ b/src/core/vulkan/shaders/quantize_i8_convrot.comp @@ -4,45 +4,51 @@ layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform parameter { uint k; uint rows; -} p; +} +p; -layout(binding = 0, std430) readonly buffer A { float data_a[]; }; -layout(binding = 1, std430) writeonly buffer D { uint data_d[]; }; -layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; +layout(binding = 0, std430) readonly buffer A { + float data_a[]; +}; +layout(binding = 1, std430) writeonly buffer D { + uint data_d[]; +}; +layout(binding = 2, std430) writeonly buffer S { + float data_s[]; +}; shared float values[256]; shared float maxima[64]; void transform_group(const uint row, const uint group) { - const uint tid = gl_LocalInvocationID.x; + const uint tid = gl_LocalInvocationID.x; const uint base_offset = row * p.k + group * 256; [[unroll]] for (uint i = 0; i < 4; ++i) { const uint index = tid + i * 64; - values[index] = data_a[base_offset + index] * (1.0 / 16.0); + values[index] = data_a[base_offset + index] * (1.0 / 16.0); } barrier(); [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) { const uint base = (tid / stride) * 4 * stride + tid % stride; - const uint i0 = base; - const uint i1 = i0 + stride; - const uint i2 = i1 + stride; - const uint i3 = i2 + stride; - const float a = values[i0]; - const float b = values[i1]; - const float c = values[i2]; - const float d = values[i3]; - values[i0] = a + b + c - d; - values[i1] = a + b - c + d; - values[i2] = a - b + c + d; - values[i3] = -a + b + c + d; + const uint i0 = base; + const uint i1 = i0 + stride; + const uint i2 = i1 + stride; + const uint i3 = i2 + stride; + const float a = values[i0]; + const float b = values[i1]; + const float c = values[i2]; + const float d = values[i3]; + values[i0] = a + b + c - d; + values[i1] = a + b - c + d; + values[i2] = a - b + c + d; + values[i3] = -a + b + c + d; barrier(); } } @@ -53,9 +59,9 @@ uint quantize_value(const float value, const float inv_scale) { } void quantize_row(const uint row) { - const uint tid = gl_LocalInvocationID.x; + const uint tid = gl_LocalInvocationID.x; const uint groups = p.k / 256; - float local_max = 0.0; + float local_max = 0.0; for (uint group = 0; group < groups; ++group) { transform_group(row, group); @@ -85,11 +91,8 @@ void quantize_row(const uint row) { const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; for (uint group = 0; group < groups; ++group) { transform_group(row, group); - const uint index = tid * 4; - const uint packed = quantize_value(values[index], inv_scale) - | (quantize_value(values[index + 1], inv_scale) << 8) - | (quantize_value(values[index + 2], inv_scale) << 16) - | (quantize_value(values[index + 3], inv_scale) << 24); + const uint index = tid * 4; + const uint packed = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24); data_d[(row * p.k + group * 256) / 4 + tid] = packed; barrier(); } diff --git a/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp b/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp index 7ac1883f..eb451f6b 100644 --- a/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp +++ b/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp @@ -8,12 +8,24 @@ #extension GL_EXT_control_flow_attributes : require layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -layout(push_constant) uniform parameter { uint k; uint rows; } p; -layout(binding = 0, std430) readonly buffer A { float data_a[]; }; -layout(binding = 1, std430) writeonly buffer D { int8_t data_d[]; }; -layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; +layout(push_constant) uniform parameter { + uint k; + uint rows; +} +p; +layout(binding = 0, std430) readonly buffer A { + float data_a[]; +}; +layout(binding = 1, std430) writeonly buffer D { + int8_t data_d[]; +}; +layout(binding = 2, std430) writeonly buffer S { + float data_s[]; +}; #if defined(SWIGLU) || defined(SIGMOID_MUL) -layout(binding = 3, std430) readonly buffer B { float data_b[]; }; +layout(binding = 3, std430) readonly buffer B { + float data_b[]; +}; #endif shared float maxima[8]; #ifdef SINGLE_PASS @@ -26,7 +38,7 @@ layout(constant_id = 0) const uint K = 4096; vec4 rotate(float a, float b, float c, float d) { precise vec4 r = vec4(a + b + c - d, a + b - c + d, - a - b + c + d, -a + b + c + d); + a - b + c + d, -a + b + c + d); return r; } @@ -37,12 +49,12 @@ void transform_group(uint row, uint group, out float v[8]) { float value = data_a[base + lane + j * 32]; #ifdef SWIGLU precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32]; - value = activated; + value = activated; #endif #ifdef SIGMOID_MUL - precise float gate = 1.0 / (1.0 + exp(-value)); + precise float gate = 1.0 / (1.0 + exp(-value)); precise float activated = data_b[base + lane + j * 32] * gate; - value = activated; + value = activated; #endif v[j] = value * (1.0 / 16.0); } @@ -51,18 +63,21 @@ void transform_group(uint row, uint group, out float v[8]) { [[unroll]] for (uint j = 0; j < 8; ++j) { vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride), subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride)); - v[j] = r[(lane / stride) & 3u]; + v[j] = r[(lane / stride) & 3u]; } } [[unroll]] for (uint j = 0; j < 8; j += 2) { - vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16), - subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16)); - v[j] = r[lane / 16]; + vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16), + subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16)); + v[j] = r[lane / 16]; v[j + 1] = r[2 + lane / 16]; } [[unroll]] for (uint j = 0; j < 2; ++j) { - vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]); - v[j] = r.x; v[j + 2] = r.y; v[j + 4] = r.z; v[j + 6] = r.w; + vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]); + v[j] = r.x; + v[j + 2] = r.y; + v[j + 4] = r.z; + v[j + 6] = r.w; } } @@ -72,26 +87,32 @@ void main() { #ifdef SINGLE_PASS float values[((K + 2047) / 2048) * 8]; [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { - if (gl_SubgroupID + g * 8 >= K / 256) continue; - float v[8]; transform_group(row, gl_SubgroupID + g * 8, v); + if (gl_SubgroupID + g * 8 >= K / 256) + continue; + float v[8]; + transform_group(row, gl_SubgroupID + g * 8, v); [[unroll]] for (uint j = 0; j < 8; ++j) { values[g * 8 + j] = v[j]; - local_max = max(local_max, abs(v[j])); + local_max = max(local_max, abs(v[j])); } } #else for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { - float v[8]; transform_group(row, group, v); - [[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j])); + float v[8]; + transform_group(row, group, v); + [[unroll]] for (uint j = 0; j < 8; ++j) + local_max = max(local_max, abs(v[j])); } #endif local_max = subgroupMax(local_max); - if (gl_SubgroupInvocationID == 0) maxima[gl_SubgroupID] = local_max; + if (gl_SubgroupInvocationID == 0) + maxima[gl_SubgroupID] = local_max; barrier(); if (gl_LocalInvocationID.x == 0) { float row_max = 0.0; - for (uint j = 0; j < 8; ++j) row_max = max(row_max, maxima[j]); - maxima[0] = row_max; + for (uint j = 0; j < 8; ++j) + row_max = max(row_max, maxima[j]); + maxima[0] = row_max; data_s[row] = row_max / 127.0; } barrier(); @@ -99,14 +120,16 @@ void main() { const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; #ifdef SINGLE_PASS [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { - if (gl_SubgroupID + g * 8 >= K / 256) continue; + if (gl_SubgroupID + g * 8 >= K / 256) + continue; [[unroll]] for (uint j = 0; j < 8; ++j) { data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127)); } } #else for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { - float v[8]; transform_group(row, group, v); + float v[8]; + transform_group(row, group, v); [[unroll]] for (uint j = 0; j < 8; ++j) { data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127)); } diff --git a/src/core/vulkan/shaders/rms_rope_apply.comp b/src/core/vulkan/shaders/rms_rope_apply.comp index 5477c584..7d08f741 100644 --- a/src/core/vulkan/shaders/rms_rope_apply.comp +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -5,16 +5,26 @@ #extension GL_EXT_control_flow_attributes : require layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer X { float data_x[]; }; -layout(binding = 1) readonly buffer PE { float data_pe[]; }; +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer PE { + float data_pe[]; +}; #ifdef CAST_F16 #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -layout(binding = 2) writeonly buffer D { float16_t data_d[]; }; +layout(binding = 2) writeonly buffer D { + float16_t data_d[]; +}; #else -layout(binding = 2) writeonly buffer D { float data_d[]; }; +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; #endif -layout(binding = 3) readonly buffer W { float data_w[]; }; +layout(binding = 3) readonly buffer W { + float data_w[]; +}; layout(push_constant) uniform parameter { uint half_dim; @@ -31,10 +41,11 @@ layout(push_constant) uniform parameter { float eps; uint w_offset; float output_scale; -} pc; +} +pc; float shuffle_channel(vec4 values, uint channel) { - const uint lane = channel & 31u; + const uint lane = channel & 31u; const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane), subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane)); return shuffled[channel / 32]; @@ -45,43 +56,43 @@ void main() { if (head >= pc.heads) { return; } - const uint lane = gl_SubgroupInvocationID; + const uint lane = gl_SubgroupInvocationID; const uint batches = pc.pairs / (64 * pc.heads * pc.tokens); for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) { for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) { - const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; - vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32], - data_x[base + lane + 64], data_x[base + lane + 96]); + const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; + vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32], + data_x[base + lane + 64], data_x[base + lane + 96]); precise vec4 squares = x * x; - precise float sum = (squares.x + squares.z) + (squares.y + squares.w); + precise float sum = (squares.x + squares.z) + (squares.y + squares.w); [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) { sum += subgroupShuffle(sum, lane ^ stride); } - sum = subgroupShuffle(sum, 0); - const float scale = inversesqrt(sum / 128.0 + pc.eps); + sum = subgroupShuffle(sum, 0); + const float scale = inversesqrt(sum / 128.0 + pc.eps); precise vec4 weighted = scale * x; [[unroll]] for (uint i = 0; i < 4; ++i) { weighted[i] *= data_w[pc.w_offset + lane + i * 32]; } [[unroll]] for (uint i = 0; i < 4; ++i) { - const uint col = lane + i * 32; + const uint col = lane + i * 32; precise float value = weighted[i]; if (pc.interleaved != 0) { const float a = subgroupShuffle(value, lane & ~1u); const float b = subgroupShuffle(value, lane | 1u); if (col / 2 < pc.rot_half) { const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1]; - value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]]; + value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]]; } } else { const uint pair = col % pc.rot_half; - const float a = shuffle_channel(weighted, pair); - const float b = shuffle_channel(weighted, pair + pc.rot_half); + const float a = shuffle_channel(weighted, pair); + const float b = shuffle_channel(weighted, pair + pc.rot_half); if (col < 2 * pc.rot_half) { const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3]; - const float c = data_pe[matrix]; - const float s = data_pe[matrix + pc.pe_stride[1]]; - value = col < pc.rot_half ? a * c - b * s : a * s + b * c; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; + value = col < pc.rot_half ? a * c - b * s : a * s + b * c; } } #ifdef CAST_F16 diff --git a/src/core/vulkan/shaders/rope_apply.comp b/src/core/vulkan/shaders/rope_apply.comp index 33d9c67d..ccfc6d63 100644 --- a/src/core/vulkan/shaders/rope_apply.comp +++ b/src/core/vulkan/shaders/rope_apply.comp @@ -2,9 +2,15 @@ layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer X { float data_x[]; }; -layout(binding = 1) readonly buffer PE { float data_pe[]; }; -layout(binding = 2) writeonly buffer D { float data_d[]; }; +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer PE { + float data_pe[]; +}; +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; layout(push_constant) uniform parameter { uint half_dim; @@ -18,41 +24,42 @@ layout(push_constant) uniform parameter { uint pe_offset; uint dst_offset; uint interleaved; -} pc; +} +pc; void main() { - const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x*gl_WorkGroupSize.x*gl_GlobalInvocationID.y; + const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y; if (index >= pc.pairs) { return; } - const uint j = index % pc.half_dim; - const uint row = index / pc.half_dim; - const uint token = row % pc.tokens; - const uint head = (row / pc.tokens) % pc.heads; - const uint batch = row / (pc.tokens*pc.heads); - const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j; - const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j + 1 : j + pc.rot_half; - const uint x_row = pc.x_offset + head*pc.x_stride[1] + token*pc.x_stride[2] + batch*pc.x_stride[3]; - const uint out_row = pc.dst_offset + row*pc.half_dim*2; - const float a = data_x[x_row + i0*pc.x_stride[0]]; - const float b = data_x[x_row + i1*pc.x_stride[0]]; + const uint j = index % pc.half_dim; + const uint row = index / pc.half_dim; + const uint token = row % pc.tokens; + const uint head = (row / pc.tokens) % pc.heads; + const uint batch = row / (pc.tokens * pc.heads); + const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j; + const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half; + const uint x_row = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; + const uint out_row = pc.dst_offset + row * pc.half_dim * 2; + const float a = data_x[x_row + i0 * pc.x_stride[0]]; + const float b = data_x[x_row + i1 * pc.x_stride[0]]; if (j >= pc.rot_half) { data_d[out_row + i0] = a; data_d[out_row + i1] = b; return; } - const uint matrix = pc.pe_offset + j*pc.pe_stride[2] + token*pc.pe_stride[3]; - const float c = data_pe[matrix]; - const float s = data_pe[matrix + pc.pe_stride[1]]; + const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3]; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; // Preserve the separate multiply/add rounding of the unfused graph. precise float y0; precise float y1; if (pc.interleaved != 0) { - y0 = a*c + b*data_pe[matrix + pc.pe_stride[0]]; - y1 = a*s + b*data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]]; + y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]]; + y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]]; } else { - y0 = a*c - b*s; - y1 = a*s + b*c; + y0 = a * c - b * s; + y1 = a * s + b * c; } data_d[out_row + i0] = y0; data_d[out_row + i1] = y1;