From d85c034977c6cec7faadfd028622d1d31dc01112 Mon Sep 17 00:00:00 2001 From: coderfeli Date: Thu, 13 Mar 2025 02:30:07 +0000 Subject: [PATCH] fix2 --- include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp index 81aedf8e80..e6fad5c857 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_moe_gemm.hpp @@ -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);