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

only padding N in epilogue

Browse Source
This commit is contained in:
lalala-sh
2025-09-05 02:50:47 +00:00
parent 5c484a5672
commit adc8db73e0
1 changed files with 10 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -372,13 +372,13 @@ struct MoeFlatmmKernel
if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
{
if(kargs.N % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
{
std::cerr << "Can't support N that is not a multiple of NPerBlock"
" without padding!"
<< std::endl;
return false;
}
// if(kargs.N % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
// {
// std::cerr << "Can't support N that is not a multiple of NPerBlock"
// " without padding!"
// << std::endl;
// return false;
// }
if(kargs.N % FlatmmPipeline::GetVectorSizeB() != 0)
{
std::cerr << "N is not a multiple of vector load size for B tensor!" << std::endl;
@@ -441,7 +441,7 @@ struct MoeFlatmmKernel
if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
{
if(kargs.N % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
if(kargs.stride_C % TilePartitioner::NPerBlock != 0 && FlatmmPipeline::kPadN == false)
{
std::cerr << "Can't support N that is not a multiple of NPerBlock"
" without padding!"
@@ -525,7 +525,7 @@ struct MoeFlatmmKernel
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
e_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens,
IsGateUp ? kargs.N / 2 : kargs.N),
IsGateUp ? kargs.stride_C / 2 : kargs.stride_C),
make_tuple(kargs.stride_C, 1),
number<EpiloguePipeline::GetVectorSizeC()>{},
number<1>{});
@@ -535,7 +535,7 @@ struct MoeFlatmmKernel
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
e_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken,
IsGateUp ? kargs.N / 2 : kargs.N),
IsGateUp ? kargs.stride_C / 2 : kargs.stride_C),
make_tuple(1, kargs.stride_C),
number<1>{},
number<1>{});