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

moe sorting fix moebuf

Browse Source
This commit is contained in:
coderfeli
2025-03-17 06:20:57 +00:00
parent 5eaa36be18
commit 7dbdff9f9f
2 changed files with 11 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -32,7 +32,7 @@ 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;
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;
};
auto t1 = fused_moegemm_traits{t.prec_i,

20
include/ck_tile/ops/fused_moe/kernel/moe_sorting_kernel.hpp
View File

@@ -192,7 +192,7 @@ struct MoeSortingHostArgs
index_t unit_size; // this is the M_a of fused-moe kernel
index_t num_experts;
index_t topk;
index_t moe_buf_bytes; // byte size of p_moe_buf
long_index_t moe_buf_bytes; // byte size of p_moe_buf
};
template <typename Problem_>
@@ -219,7 +219,7 @@ struct MoeSortingKernel
void* p_moe_buf;
index_t tokens;
index_t num_experts;
index_t moe_buf_bytes;
long_index_t moe_buf_bytes;
index_t tokens_per_thread;
index_t smem_rows;
@@ -426,7 +426,7 @@ struct MoeSortingKernel
return row * total_col + col;
}
CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes) const
CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, long_index_t buf_bytes) const
{
const index_t offset = (blockIdx.x - 1) * blockDim.x + threadIdx.x;
if(offset < buf_bytes / 16)
@@ -1218,10 +1218,10 @@ CK_TILE_DEVICE void moe_sorting_wave_cumsum(data_t& thread_data)
}
template <index_t BLOCK_SIZE = 256>
CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes, index_t gid)
CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, long_index_t buf_bytes, index_t gid)
{
// const index_t offset = (blockIdx.x - 1) * BLOCK_SIZE + threadIdx.x;
index_t offset = gid * BLOCK_SIZE + threadIdx.x;
long_index_t offset = static_cast<long_index_t>(gid) * BLOCK_SIZE + threadIdx.x;
if(offset < buf_bytes / 16)
{
buf[offset] = uint8x16_t{0};
@@ -1233,6 +1233,10 @@ CK_TILE_DEVICE void moe_buf_set_zero_kernel(uint8x16_t* buf, index_t buf_bytes,
// prefer to run mp kernel if is not oneshot
CK_TILE_HOST bool moe_sorting_is_oneshot(int tokens_, int num_experts_)
{
if(tokens_ >= 65536 * 2)
{
return 0;
}
auto sub_token_ = moe_sorting_get_sub_token(tokens_, num_experts_);
bool is_sub_token_onshot = tokens_ <= sub_token_;
return is_sub_token_onshot;
@@ -1250,10 +1254,6 @@ CK_TILE_HOST index_t moe_sorting_mp_get_workspace_size(int tokens_, int num_expe
CK_TILE_HOST index_t moe_sorting_get_workspace_size(int tokens_, int num_experts_)
{
#if 1
if(tokens_ >= 65536 * 2)
{
return 0;
}
if(moe_sorting_is_oneshot(tokens_, num_experts_))
{
return 0;
@@ -1527,7 +1527,7 @@ struct MoeSortingMultiPhaseKernel_P2
index_t num_experts;
index_t mesh_stride; // mesh_stride for p_expert_mesh
mdiv unit_size_mdiv;
index_t moe_buf_bytes;
long_index_t moe_buf_bytes;
};
CK_TILE_HOST static constexpr auto MakeKargs(const Hargs& h)