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

fix other device gemm from preshuffle

Browse Source
This commit is contained in:
qin letao
2025-03-19 08:00:45 +00:00
parent 2d30c55a83
commit 6ce257371f
5 changed files with 18 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8.cpp
View File

@@ -356,7 +356,9 @@ int main(int argc, char* argv[])
KBatch,
a_element_op,
b_element_op,
cde_element_op);
cde_element_op,
N,
K);
if(!device_op.IsSupportedArgument(argument))
{

4
example/65_gemm_multiply_multiply/moe_gemm1_xdl_pk_i4.cpp
View File

@@ -438,7 +438,9 @@ int main(int argc, char* argv[])
KBatch,
a_element_op,
b_element_op,
cde_element_op);
cde_element_op,
N,
K);
if(!device_op.IsSupportedArgument(argument))
{

4
example/65_gemm_multiply_multiply/moe_gemm2_xdl_fp8.cpp
View File

@@ -365,7 +365,9 @@ int main(int argc, char* argv[])
KBatch,
a_element_op,
b_element_op,
cde_element_op);
cde_element_op,
N,
K);
if(!device_op.IsSupportedArgument(argument))
{

4
example/65_gemm_multiply_multiply/moe_gemm2_xdl_pk_i4.cpp
View File

@@ -405,7 +405,9 @@ int main(int argc, char* argv[])
KBatch,
a_element_op,
b_element_op,
cde_element_op);
cde_element_op,
N,
K);
if(!device_op.IsSupportedArgument(argument))
{

8
include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
View File

@@ -391,7 +391,9 @@ struct DeviceMoeGemm : public DeviceGemmMultipleDSplitKBPreShuffle<ALayout,
index_t KBatch,
AElementwiseOperation a_element_op,
BElementwiseOperation b_element_op,
CElementwiseOperation c_element_op)
CElementwiseOperation c_element_op,
int,
int)
{
return Argument{static_cast<const index_t*>(p_sorted_token_ids),
static_cast<const index_t*>(p_sorted_expert_ids),
@@ -432,7 +434,9 @@ struct DeviceMoeGemm : public DeviceGemmMultipleDSplitKBPreShuffle<ALayout,
index_t KBatch,
AElementwiseOperation a_element_op,
BElementwiseOperation b_element_op,
CElementwiseOperation c_element_op) override
CElementwiseOperation c_element_op,
int,
int) override
{
return std::make_unique<Argument>(nullptr,
nullptr,