From 100dcc9ea2c24b5af084e0fdbbe3c7c38a0a3246 Mon Sep 17 00:00:00 2001 From: ltqin Date: Tue, 11 Nov 2025 13:56:25 +0000 Subject: [PATCH] add qkv scale all --- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 117 ++++++++++++------ include/ck_tile/core/utility/functional.hpp | 9 ++ .../host/reference/reference_batched_gemm.hpp | 41 ++++++ .../ops/fmha/kernel/fmha_fwd_kernel.hpp | 21 +--- .../pipeline/block_fmha_pipeline_qr_ks_vs.hpp | 47 +++---- 5 files changed, 150 insertions(+), 85 deletions(-) diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index d44cca458e..c545699989 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -172,7 +172,7 @@ class BlockQuantizer << " num_blocks_: " << num_blocks_ << std::endl; std::random_device rd; std::mt19937 gen(rd()); - std::uniform_real_distribution dis(2.0f, 2.0f); + std::uniform_real_distribution dis(0.5f, 2.0f); for(size_t b = 0; b < batch; ++b) { for(size_t h = 0; h < head; ++h) @@ -214,7 +214,7 @@ class BlockQuantizer } } // save scale to tensor - block_scale(b, h, block) = scale; + block_scale(b, h, block) = 1.0f / scale; std::cout << "block: " << block << " scale: " << scale << " max_value: " << max_value << " block_scale: " << block_scale << std::endl; @@ -252,7 +252,7 @@ class BlockQuantizer if(!i_perm) idx = {b, s, h, d}; float val = ck_tile::type_convert(in(idx)); - out(idx) = ck_tile::type_convert(val / scale); + out(idx) = ck_tile::type_convert(val * scale); } } } @@ -806,17 +806,32 @@ fwd_result fmha_fwd_run(mode_enum mode, float scale_o = 1.f; if(quant == 2) { - q_host.savetxt("./q_org.txt"); - k_host.savetxt("./k_org.txt"); - v_host.savetxt("./v_org.txt"); - BlockQuantizer quantizer(i_perm); - quantizer.quantize(q_host, q_host, q_scale, block_scale_m_); - quantizer.quantize(k_host, k_host, k_scale, block_scale_n_); - // quantizer.quantize(v_host, v_host, v_scale, block_scale_n_); - // scale_p = quantizer.scale_p(); - q_host.savetxt("./q_quant.txt"); - k_host.savetxt("./k_quant.txt"); - v_host.savetxt("./v_quant.txt"); + ck_tile::FillUniformDistributionIntegerValue{1.f, 10.f, next_seed()}(q_scale); + ck_tile::FillUniformDistributionIntegerValue{1.f, 10.f, next_seed()}(k_scale); + ck_tile::FillUniformDistributionIntegerValue{1.f, 10.f, next_seed()}(v_scale); + + { //debug info + std::cout << "q_scale: " << q_scale << " k_scale: " << k_scale + << " v_scale: " << v_scale << std::endl; + + ck_tile::HostTensor q_host_deq( + get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q)); + ck_tile::HostTensor k_host_deq( + 0 < page_block_size + ? get_lengths(i_perm, max_num_page_blocks, nhead_k, page_block_size, hdim_q) + : get_lengths(i_perm, shape_batch, nhead_k, shape_seqlen_k, hdim_q)); + ck_tile::HostTensor v_host_deq( + 0 < page_block_size + ? get_lengths(i_perm, max_num_page_blocks, nhead_k, page_block_size, hdim_q) + : get_lengths(i_perm, shape_batch, nhead_k, shape_seqlen_k, hdim_q)); + BlockQuantizer quantizer(i_perm); + quantizer.dequantize(q_host, q_host_deq, q_scale, block_scale_m_); + quantizer.dequantize(k_host, k_host_deq, k_scale, block_scale_n_); + quantizer.dequantize(v_host, v_host_deq, v_scale, block_scale_n_); + q_host_deq.savetxt("./q_deq.txt"); + k_host_deq.savetxt("./k_deq.txt"); + v_host_deq.savetxt("./v_deq.txt"); + } } else if(quant == 1) { @@ -1525,18 +1540,6 @@ fwd_result fmha_fwd_run(mode_enum mode, uint8_t(std::floor(p_undrop * std::numeric_limits::max())); float rp_undrop = 1.0 / p_undrop; - if(quant == 2) - { - // dequant data for host - BlockQuantizer quantizer(i_perm); - quantizer.dequantize(q_host, q_host, q_scale, block_scale_m_); - quantizer.dequantize(k_host, k_host, k_scale, block_scale_n_); - // quantizer.dequantize(v_host, v_host, v_scale, block_scale_n_); - q_host.savetxt("./q_dequant.txt"); - k_host.savetxt("./k_dequant.txt"); - v_host.savetxt("./v_dequant.txt"); - // scale_s = scale_s / 48.0 / 48.0; - } for(ck_tile::index_t wb = 0; wb < batch; ++wb) { ck_tile::index_t real_seqlen_q = seqstart_q_host[wb + 1] - seqstart_q_host[wb]; @@ -1723,14 +1726,34 @@ fwd_result fmha_fwd_run(mode_enum mode, #endif // reference - ck_tile:: - reference_batched_gemm( + if(quant == 2) + { + ck_tile::reference_batched_quant_gemm( q_host_ref, k_host_ref, s_host_ref, - ck_tile::identity{}, - ck_tile::identity{}, - ck_tile::scales(scale_s)); + ck_tile::idx_identity{}, + ck_tile::idx_identity{}, + [&q_scale, &k_scale, scale_s, wb](auto idx, auto value) { + return value * scale_s * + q_scale(wb, std::get<0>(idx), std::get<1>(idx) / 128) * + k_scale(wb, std::get<0>(idx), std::get<2>(idx) / 128); + }); + } + else + { + ck_tile:: + reference_batched_gemm( + q_host_ref, + k_host_ref, + s_host_ref, + ck_tile::identity{}, + ck_tile::identity{}, + ck_tile::scales(scale_s)); + } if(0.f < logits_soft_cap) { @@ -1888,13 +1911,31 @@ fwd_result fmha_fwd_run(mode_enum mode, } } - ck_tile::reference_batched_gemm( - p_host_ref, - v_host_ref, - o_host_ref, - ck_tile::identity{}, - ck_tile::identity{}, - oacc_element_func); + if(quant == 2) + { + ck_tile:: + reference_batched_quant_gemm( + p_host_ref, + v_host_ref, + o_host_ref, + ck_tile::idx_identity{}, + [&v_scale, wb](auto idx, auto value) { + // idx: b, m, n, k --> h, sq, d, sk + return ck_tile::type_convert(value) * + v_scale(wb, std::get<0>(idx), std::get<2>(idx) / 128); + }, + ck_tile::idx_identity{}); + } + else + { + ck_tile::reference_batched_gemm( + p_host_ref, + v_host_ref, + o_host_ref, + ck_tile::identity{}, + ck_tile::identity{}, + oacc_element_func); + } ck_tile::HostTensor o_host_result({nhead, real_seqlen_q, hdim_v}); // clang-format off diff --git a/include/ck_tile/core/utility/functional.hpp b/include/ck_tile/core/utility/functional.hpp index fd0252d3ca..b1f9193036 100644 --- a/include/ck_tile/core/utility/functional.hpp +++ b/include/ck_tile/core/utility/functional.hpp @@ -91,6 +91,15 @@ struct identity } }; +struct idx_identity +{ + template + CK_TILE_HOST_DEVICE constexpr T&& operator()(auto, T&& arg) const noexcept + { + return std::forward(arg); + } +}; + namespace detail { // RemainLengths: sequence<...> diff --git a/include/ck_tile/host/reference/reference_batched_gemm.hpp b/include/ck_tile/host/reference/reference_batched_gemm.hpp index 826358de30..c588803029 100644 --- a/include/ck_tile/host/reference/reference_batched_gemm.hpp +++ b/include/ck_tile/host/reference/reference_batched_gemm.hpp @@ -47,4 +47,45 @@ CK_TILE_HOST void reference_batched_gemm(const HostTensor& a_b_m_k, make_ParallelTensorFunctor(f, c_b_m_n.mDesc.get_lengths()[0], c_b_m_n.mDesc.get_lengths()[1])( std::thread::hardware_concurrency()); } + +template +CK_TILE_HOST void reference_batched_quant_gemm(const HostTensor& a_b_m_k, + const HostTensor& b_b_n_k, + HostTensor& c_b_m_n, + const AElementOp& a_element_op = {}, + const BElementOp& b_element_op = {}, + const ACCElementOp& acc_element_op = {}) +{ + const int N = b_b_n_k.mDesc.get_lengths()[1]; + const int K = b_b_n_k.mDesc.get_lengths()[2]; + + auto f = [&](auto batch, auto m) { + for(int n = 0; n < N; ++n) + { + AccDataType v_acc = 0; + + for(int k = 0; k < K; ++k) + { + AccDataType v_a = ck_tile::type_convert( + a_element_op(std::make_tuple(batch, m, k), a_b_m_k(batch, m, k))); + AccDataType v_b = ck_tile::type_convert( + b_element_op(std::make_tuple(batch, n, k), b_b_n_k(batch, n, k))); + + v_acc += v_a * v_b; + } + + c_b_m_n(batch, m, n) = ck_tile::type_convert(acc_element_op(std::make_tuple(batch, m, n), v_acc)); + } + }; + + make_ParallelTensorFunctor(f, c_b_m_n.mDesc.get_lengths()[0], c_b_m_n.mDesc.get_lengths()[1])( + std::thread::hardware_concurrency()); +} + } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index 07ee2e3b3d..910bab8b26 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -1403,7 +1403,6 @@ struct FmhaFwdKernel const float* k_scale_ptr = nullptr; const float* v_scale_ptr = nullptr; float q_scale = 1; - float v_scale = 1; if constexpr(kDoFp8StaticQuant) { if(kargs.q_scale_ptr) @@ -1422,24 +1421,6 @@ struct FmhaFwdKernel size_t idx = i_m0 / kargs.block_scale_m; q_scale = q_scale_ptr[idx]; - v_scale = v_scale_ptr[idx]; - } - - if(get_block_1d_id() == 0 && get_thread_local_1d_id() == 0) - { - size_t idx = i_m0 / kargs.block_scale_m; - printf("blockIdx.x: %d, blockIdx.y: %d,blockIdx.z: %d,i_batch: %d, i_nhead: " - "%d, i_nhead_k: %d, i_m0: %d, idx: %zu, q_scale: %f, v_scale: %f\n", - blockIdx.x, - blockIdx.y, - blockIdx.z, - i_batch, - i_nhead, - i_nhead / kargs.nhead_ratio_qk, - i_m0, - idx, - q_scale, - v_scale); } } @@ -1747,7 +1728,7 @@ struct FmhaFwdKernel o_acc_element_func, // o_acc_element_func mask, position_encoding, - kargs.scale_s / q_scale, + kargs.scale_s * q_scale, variant, variant_params, block_indices, diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp index 8bf529146c..9e3b4ae066 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp @@ -319,37 +319,20 @@ struct BlockFmhaPipelineQRKSVS static_assert(2 <= k0_loops); static_assert(1 <= k1_loops); - const float store_scale_s = scale_s; + // const float store_scale_s = scale_s; do { + float k_scale = 1.0f; + if constexpr(kDoFp8StaticQuant) { - scale_s = store_scale_s; - float k_scale = 1; - if constexpr(kDoFp8StaticQuant) + const auto k_origin = k_dram_block_window.get_window_origin(); + const auto row = k_origin.at(number<0>{}); + if(k_scale_ptr) { - const auto k_origin = k_dram_block_window.get_window_origin(); - const auto row = k_origin.at(number<0>{}); - if(k_scale_ptr) - { - const index_t idx = row / block_scale_n; - k_scale = k_scale_ptr[idx]; - scale_s = scale_s / k_scale; - if(get_block_1d_id() == 0 && get_thread_local_1d_id() == 0) - { - printf("blockIdx.x: %d, blockIdx.y: %d,blockIdx.z: %d, row: %d, idx: " - "%d, k_scale: %f " - "\n", - blockIdx.x, - blockIdx.y, - blockIdx.z, - row, - idx, - k_scale); - } - } + const index_t idx = row / block_scale_n; + k_scale = k_scale_ptr[idx]; } } - // STAGE 1, QK gemm auto k_dram_window = make_tile_window( k_dram_block_window.get_bottom_tensor_view(), @@ -419,6 +402,15 @@ struct BlockFmhaPipelineQRKSVS k_lds_window); schedule_gemm0(); } + // dequant + if constexpr(kDoFp8StaticQuant) + { + if(k_scale_ptr) + { + tile_elementwise_inout([k_scale](auto& x) { x = x * k_scale; }, s_acc); + } + } + // STAGE 2, scale_s, add bias, mask, softmax if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS) { @@ -690,13 +682,14 @@ struct BlockFmhaPipelineQRKSVS } // o_acc += o_acc_tmp; // o_acc += tile_elementwise_in(scale(1.0f / v_scale), o_acc_tmp); - ck_tile::ignore = v_scale; + // ck_tile::ignore = v_scale; sweep_tile_span(o_spans[number<0>{}], [&](auto idx0) { sweep_tile_span(o_spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); - o_acc(i_j_idx) += o_acc_tmp(i_j_idx); // / v_scale; + o_acc(i_j_idx) += o_acc_tmp(i_j_idx) * v_scale; }); }); + } while(++i_total_loops < num_total_loop); // store lse