ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-12 01:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

tempsave, trload+asyncload done

Browse Source
This commit is contained in:
aska-0096
2025-07-21 05:55:55 +00:00
parent afd96d8180
commit 1b468bac0b
6 changed files with 168 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -676,7 +676,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
ck_tile::FillUniformDistribution<VDataType>{0.f, 1.f, seed}(vnew_host);
ck_tile::FillUniformDistribution<BiasDataType>{0.f, 1.f, seed}(bias_host);
}
else if(init_method == "exp" || init_method == "99")
else if(init_method == "v1" || init_method == "97")
{
ck_tile::FillUniformDistribution<QDataType>{0.f, 1.f, seed}(q_host);
ck_tile::FillUniformDistribution<KDataType>{0.f, 1.f, seed}(k_host);
@@ -685,6 +685,24 @@ bool run(const ck_tile::ArgParser& arg_parser)
ck_tile::FillUniformDistribution<VDataType>{1.f, 1.f, seed}(vnew_host);
ck_tile::FillUniformDistribution<BiasDataType>{0.f, 1.f, seed}(bias_host);
}
else if(init_method == "kv1" || init_method == "98")
{
ck_tile::FillUniformDistribution<QDataType>{0.f, 1.f, seed}(q_host);
ck_tile::FillUniformDistribution<KDataType>{1.f, 1.f, seed}(k_host);
ck_tile::FillUniformDistribution<KDataType>{1.f, 1.f, seed}(knew_host);
ck_tile::FillUniformDistribution<VDataType>{1.f, 1.f, seed}(v_host);
ck_tile::FillUniformDistribution<VDataType>{1.f, 1.f, seed}(vnew_host);
ck_tile::FillUniformDistribution<BiasDataType>{0.f, 1.f, seed}(bias_host);
}
else if(init_method == "qkv1" || init_method == "99")
{
ck_tile::FillUniformDistribution<QDataType>{1.f, 1.f, seed}(q_host);
ck_tile::FillUniformDistribution<KDataType>{1.f, 1.f, seed}(k_host);
ck_tile::FillUniformDistribution<KDataType>{1.f, 1.f, seed}(knew_host);
ck_tile::FillUniformDistribution<VDataType>{1.f, 1.f, seed}(v_host);
ck_tile::FillUniformDistribution<VDataType>{1.f, 1.f, seed}(vnew_host);
ck_tile::FillUniformDistribution<BiasDataType>{0.f, 1.f, seed}(bias_host);
}
else if(init_method == "nf")
{
ck_tile::FillNormalDistribution<QDataType>{0.f, 3.f, seed}(q_host);