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

fix2

Browse Source
This commit is contained in:
coderfeli
2025-03-13 02:30:07 +00:00
parent 8b05fa935d
commit d85c034977
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1576,7 +1576,7 @@ struct GridwiseMoeGemm
static_for<0, EMRepeats, 1>{}([&](auto m0) {
const index_t fused_token = p_sorted_token_ids[c_token_pos + m0];
index_t token_offset = fused_token & 0xffffff;
float weight = token_offset < problem.NumTokens? p_sorted_weights_0[token_offset * problem.StrideDs[0]] : 1.0;
float weight = token_offset < problem.NumTokens? p_sorted_weights_0[token_offset * problem.StrideDs[0]] : 0.0;
if constexpr(IsInputGemm)
{
token_offset = token_offset * problem.TopK + (fused_token >> 24);
@@ -2084,7 +2084,8 @@ struct GridwiseMoeGemm
static_for<0, EMRepeats, 1>{}([&](auto m0) {
const index_t fused_token = p_sorted_token_ids[c_token_pos + m0];
index_t token_offset = fused_token & 0xffffff;
float weight = p_sorted_weights_0[token_offset * problem.StrideDs[0]];
float weight = token_offset < problem.NumTokens? p_sorted_weights_0[token_offset * problem.StrideDs[0]] : 0.0;
// float weight = p_sorted_weights_0[token_offset * problem.StrideDs[0]];
if constexpr(IsInputGemm)
{
token_offset = token_offset * problem.TopK + (fused_token >> 24);