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

add sorting params

Browse Source
This commit is contained in:
coderfeli
2025-03-17 07:44:49 +00:00
parent 7dbdff9f9f
commit 04dc9908d0
7 changed files with 10 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
example/ck_tile/13_moe_sorting/moe_sorting.cpp
View File

@@ -175,7 +175,8 @@ bool test_moe_sorting(ck_tile::ArgParser args)
unit_size,
num_experts,
topk,
static_cast<ck_tile::index_t>(moe_buf_size * sizeof(float))};
static_cast<ck_tile::index_t>(moe_buf_size * sizeof(float)),
false};
ck_tile::stream_config sc{nullptr,
true,

2
example/ck_tile/13_moe_sorting/moe_sorting_api.cpp
View File

@@ -153,7 +153,7 @@ float moe_sorting(moe_sorting_trait t, moe_sorting_args a, ck_tile::stream_confi
}
}
#else
if(moe_sorting_get_workspace_size(a.tokens, a.num_experts) != 0)
if(!a.force_one_shoot && moe_sorting_get_workspace_size(a.tokens, a.num_experts) != 0)
{
return moe_sorting_mp(t, a, s);
}

1
example/ck_tile/15_fused_moe/fused_moe.hpp
View File

@@ -36,6 +36,7 @@ struct fused_moe_args
ck_tile::index_t topk; // need this?
ck_tile::index_t stride_token; // for input/output, stride for each row, should >= hidden_size
bool force_one_shoot;
};
// This is the public API, will be generated by script

3
example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp
View File

@@ -32,7 +32,8 @@ float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_conf
a.block_m, // index_t unit_size;
a.num_experts, // index_t num_experts;
a.topk, // index_t topk;
static_cast<ck::long_index_t>(a.num_tokens) * a.stride_token * o_data_bytes // index_t moe_buf_bytes;
static_cast<ck::long_index_t>(a.num_tokens) * a.stride_token * o_data_bytes, // index_t moe_buf_bytes;
a.force_one_shoot
};
auto t1 = fused_moegemm_traits{t.prec_i,

3
example/ck_tile/15_fused_moe/main.cpp
View File

@@ -413,7 +413,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
tokens,
experts,
topk,
stride};
stride,
false};
float ave_time = fused_moe(
traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});