ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 17:48:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

add qkv scale all

Browse Source
This commit is contained in:
ltqin
2025-11-11 13:56:25 +00:00
parent e1bc48d5f3
commit 100dcc9ea2
5 changed files with 150 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp
View File

@@ -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,

47
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp
View File

@@ -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