ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 10:08:01 +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

117
example/ck_tile/01_fmha/fmha_fwd_runner.hpp
View File

@@ -172,7 +172,7 @@ class BlockQuantizer
<< " num_blocks_: " << num_blocks_ << std::endl;
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_real_distribution<float> dis(2.0f, 2.0f);
std::uniform_real_distribution<float> 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<float>(in(idx));
out(idx) = ck_tile::type_convert<OutDataType>(val / scale);
out(idx) = ck_tile::type_convert<OutDataType>(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<QDataType>();
q_host.savetxt("./q_quant.txt");
k_host.savetxt("./k_quant.txt");
v_host.savetxt("./v_quant.txt");
ck_tile::FillUniformDistributionIntegerValue<float>{1.f, 10.f, next_seed()}(q_scale);
ck_tile::FillUniformDistributionIntegerValue<float>{1.f, 10.f, next_seed()}(k_scale);
ck_tile::FillUniformDistributionIntegerValue<float>{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<float> q_host_deq(
get_lengths(i_perm, shape_batch, nhead, shape_seqlen_q, hdim_q));
ck_tile::HostTensor<float> 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<float> 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<uint8_t>::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<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
if(quant == 2)
{
ck_tile::reference_batched_quant_gemm<QDataType,
KDataType,
SaccDataType,
SMPLComputeDataType>(
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<QDataType, KDataType, SaccDataType, SMPLComputeDataType>(
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<PDataType, VDataType, OaccDataType, ODataType>(
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<PDataType, VDataType, OaccDataType, ODataType>(
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<float>(value) *
v_scale(wb, std::get<0>(idx), std::get<2>(idx) / 128);
},
ck_tile::idx_identity{});
}
else
{
ck_tile::reference_batched_gemm<PDataType, VDataType, OaccDataType, ODataType>(
p_host_ref,
v_host_ref,
o_host_ref,
ck_tile::identity{},
ck_tile::identity{},
oacc_element_func);
}
ck_tile::HostTensor<ODataType> o_host_result({nhead, real_seqlen_q, hdim_v});
// clang-format off

9
include/ck_tile/core/utility/functional.hpp
View File

@@ -91,6 +91,15 @@ struct identity
}
};
struct idx_identity
{
template <typename T>
CK_TILE_HOST_DEVICE constexpr T&& operator()(auto, T&& arg) const noexcept
{
return std::forward<T>(arg);
}
};
namespace detail {
// RemainLengths: sequence<...>

41
include/ck_tile/host/reference/reference_batched_gemm.hpp
View File

@@ -47,4 +47,45 @@ CK_TILE_HOST void reference_batched_gemm(const HostTensor<ADataType>& 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 <typename ADataType,
typename BDataType,
typename AccDataType,
typename CDataType,
typename AElementOp = ck_tile::identity,
typename BElementOp = ck_tile::identity,
typename ACCElementOp = ck_tile::identity>
CK_TILE_HOST void reference_batched_quant_gemm(const HostTensor<ADataType>& a_b_m_k,
const HostTensor<BDataType>& b_b_n_k,
HostTensor<CDataType>& 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<AccDataType>(
a_element_op(std::make_tuple(batch, m, k), a_b_m_k(batch, m, k)));
AccDataType v_b = ck_tile::type_convert<AccDataType>(
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<CDataType>(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

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