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

hot fix expert over flow

Browse Source
This commit is contained in:
coderfeli
2025-02-20 13:54:44 +00:00
parent 6cb9c0e67c
commit 70aef2ed33
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp
View File

@@ -1174,7 +1174,9 @@ struct GridwiseMoeGemm
const index_t max_token_id = __builtin_amdgcn_readfirstlane(p_max_token_id[0]);
// constexpr int expert_tile_cnt[8] = {2, 1, 1, 2, 2, 2, 1, 2};
// const index_t b_block_id = blockIdx.x % problem.NBlock;
const index_t expert_block_id = blockIdx.x / problem.NBlock;
const index_t expert_block_id = NSwizzle ? blockIdx.x / problem.NBlock : blockIdx.x;
if (expert_block_id * MPerBlock >= max_token_id)
return;
const index_t expert_id = NSwizzle ? __builtin_amdgcn_readfirstlane(p_sorted_expert_ids[blockIdx.x / problem.NBlock]) :
__builtin_amdgcn_readfirstlane(p_sorted_expert_ids[blockIdx.y]);
const auto block_mn = [&]() -> std::pair<int, int> {
@@ -1219,7 +1221,7 @@ struct GridwiseMoeGemm
constexpr auto AMRepeats = MPerBlock / AMThreads;
const index_t token_pos = block_m_id * MPerBlock + threadIdx.x / AKThreads * AMRepeats;
if(token_pos >= max_token_id || expert_block_id * MPerBlock >= max_token_id || token0 >= problem.NumTokens)
if(token_pos >= max_token_id || token0 >= problem.NumTokens)
return;
StaticallyIndexedArray<index_t, AMRepeats> gather_offsets; //= p_sorted_token_ids[token_pos];
static_for<0, AMRepeats, 1>{}([&](auto m0) {