ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

clear code

Browse Source
This commit is contained in:
ltqin
2025-11-12 06:10:03 +00:00
parent df55e264ad
commit 8f2e2886cd
2 changed files with 10 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -810,9 +810,6 @@ fwd_result fmha_fwd_run(mode_enum mode,
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_);
q_host.savetxt("./q_quant.txt");
k_host.savetxt("./k_quant.txt");
v_host.savetxt("./v_quant.txt");
}
else if(quant == 1)
{
@@ -1720,8 +1717,9 @@ fwd_result fmha_fwd_run(mode_enum mode,
ck_tile::idx_identity{},
[&](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) / nr, std::get<2>(idx) / 128);
q_scale(wb, std::get<0>(idx), std::get<1>(idx) / block_scale_m_) *
k_scale(
wb, std::get<0>(idx) / nr, std::get<2>(idx) / block_scale_n_);
});
}
else
@@ -1900,9 +1898,11 @@ fwd_result fmha_fwd_run(mode_enum mode,
v_host_ref,
o_host_ref,
ck_tile::idx_identity{},
[&v_scale, wb, nr](auto idx, auto value) {
[&v_scale, wb, nr, block_scale_n_](auto idx, auto value) {
return ck_tile::type_convert<float>(value) *
v_scale(wb, std::get<0>(idx) / nr, std::get<2>(idx) / 128);
v_scale(wb,
std::get<0>(idx) / nr,
std::get<2>(idx) / block_scale_n_);
},
ck_tile::idx_identity{});
}

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

@@ -408,7 +408,8 @@ struct BlockFmhaPipelineQRKSVS
{
if(k_scale_ptr)
{
tile_elementwise_inout([q_scale, k_scale](auto& x) { x = x * q_scale * k_scale; }, s_acc);
tile_elementwise_inout(
[q_scale, k_scale](auto& x) { x = x * q_scale * k_scale; }, s_acc);
}
}
@@ -646,15 +647,7 @@ struct BlockFmhaPipelineQRKSVS
{
auto acc0 = gemm_1(a, b);
tile_elementwise_inout(
[&v_scale](auto& o, auto o0) {
// asm volatile(";wrapper_gemm1\n\tv_mul_f32_e32 %0, %1, %2"
// : "=v"(o)
// : "s"(v_scale), "v"(o0)
// : "memory");
o += o0 * v_scale;
},
acc,
acc0);
[&v_scale](auto& o, auto o0) { o += o0 * v_scale; }, acc, acc0);
}
else
{
@@ -703,16 +696,6 @@ struct BlockFmhaPipelineQRKSVS
v_lds_window);
block_sync_lds();
}
// o_acc += o_acc_tmp;
// o_acc += tile_elementwise_in(scale(1.0f / v_scale), o_acc_tmp);
// 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;
// });
// });
} while(++i_total_loops < num_total_loop);
// store lse