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

change squant to quant type

Browse Source
This commit is contained in:
ltqin
2025-11-03 12:31:45 +00:00
parent 1c17bae816
commit f72822d3a0
2 changed files with 11 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
example/ck_tile/01_fmha/example_fmha_fwd.cpp
View File

@@ -50,13 +50,10 @@ auto create_args(int argc, char* argv[])
"scale factor of S. 0 means equal to 1/sqrt(hdim).\n"
"note when squant=1, this value will be modified")
.insert("logits_soft_cap", "0", "attention logits soft capping value.")
.insert("squant",
"auto",
"if using static quantization fusion or not. auto: fp8 will default use squant, "
"other will not\n"
"0: no static quant(not implemented) 1: apply scale_p and scale_o with respect to "
"P and O.\n"
"calculate scale_s, scale_p, scale_o auto")
.insert("quant",
"0",
"if using quantization or not.\n"
"0: no quant 1: per-tensor quant 2: block scaling.\n")
.insert("iperm",
"1",
"permute input\n"
@@ -162,12 +159,7 @@ auto run(const ck_tile::ArgParser& arg_parser)
std::string init_method = arg_parser.get_str("init");
uint32_t seed = arg_parser.get_uint32("seed");
bool squant = [&]() {
if(arg_parser.get_str("squant") == "auto")
return std::is_same_v<DataTypeConfig, FmhaFwdFp8>;
else
return arg_parser.get_bool("squant");
}();
int quant = arg_parser.get_bool("quant");
ck_tile::stream_config stream_config{nullptr,
true,
@@ -208,7 +200,7 @@ auto run(const ck_tile::ArgParser& arg_parser)
drop_offset,
drop_prefs,
mask_str,
squant,
quant,
is_rotary_interleaved,
num_splits,
init_method,

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

@@ -178,7 +178,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
uint64_t drop_offset,
bool drop_prefs,
std::string mask_str,
bool squant,
int quant,
bool is_rotary_interleaved,
ck_tile::index_t num_splits,
std::string init_method,
@@ -705,7 +705,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
float scale_p = 1.f;
float scale_o = 1.f;
if(squant)
if(quant)
{
float q_dtype_max = ck_tile::type_convert<float>(ck_tile::numeric<QDataType>::max());
float k_dtype_max = ck_tile::type_convert<float>(ck_tile::numeric<KDataType>::max());
@@ -816,7 +816,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
<< (seqlen_kpads[0] < 0 ? ""
: (std::string("(") + std::to_string(seqlen_kpads[0]) + ")"))
<< ", d:" << hdim_q << "/" << hdim_v << ", scale_s:" << scale_s << ", bias:" << bias
<< ", p_drop:" << p_drop << ", lse:" << lse << ", squant:" << squant
<< ", p_drop:" << p_drop << ", lse:" << lse << ", quant:" << quant
<< ", mask:" << mask << ", v:" << (is_v_rowmajor ? "r" : "c");
#if CK_TILE_FMHA_FWD_APPENDKV_API
if(0 < rotary_dim)
@@ -908,7 +908,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
traits.mask_type = mask.type;
traits.bias_type = bias.type;
traits.has_lse = lse;
traits.do_fp8_static_quant = squant;
traits.do_fp8_static_quant = quant == 1;
if constexpr(std::is_same_v<fmha_fwd_traits, std::decay_t<decltype(traits)>>)
{
@@ -1774,7 +1774,7 @@ fwd_result fmha_fwd_run(mode_enum mode,
scale_s,
p_drop,
lse,
squant,
quant,
bias.type == bias_enum::elementwise_bias
? "elementwise_bias"
: (bias.type == bias_enum::alibi ? "alibi" : "no_bias"),