ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-04 05:31:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

use int64_t as expert stride to avoid overflow

Browse Source
This commit is contained in:
Feng Shijie
2025-08-21 06:58:55 +00:00
parent 9fbcc8f8a4
commit 85976b0b87
3 changed files with 19 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
include/ck_tile/ops/moe_flatmm/kernel/moe_flatmm_kernel.hpp
View File

@@ -644,7 +644,8 @@ struct MoeFlatmmKernel
});
const SplitKBatchOffset splitk_batch_offset(kargs);
const index_t expert_stride = __builtin_amdgcn_readfirstlane(kargs.N * kargs.K);
const long_index_t expert_stride =
__builtin_amdgcn_readfirstlane(long_index_t(kargs.N) * kargs.K);
const ADataType* a_ptr =
static_cast<const ADataType*>(kargs.a_ptr) + splitk_batch_offset.a_k_split_offset;