diff --git a/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh index 1549ed96aa2be..24564efbd21be 100644 --- a/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -29,26 +29,12 @@ struct sm100_fp8_config_default { template typename Epilogue> struct sm100_fp8_config_M256 { - // M in (128, 256] + // M in (64, 256] static_assert(std::is_same()); using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; using TileShape = Shape<_128, _128, _128>; - using ClusterShape = Shape<_2, _2, _1>; - using Cutlass3xGemm = - cutlass_3x_gemm_sm100; -}; - -template typename Epilogue> -struct sm100_fp8_config_M128 { - // M in (64, 128] - static_assert(std::is_same()); - using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; - using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; - using TileShape = Shape<_128, _128, _256>; - using ClusterShape = Shape<_2, _4, _1>; + using ClusterShape = Shape<_2, _1, _1>; using Cutlass3xGemm = cutlass_3x_gemm_sm100; @@ -57,12 +43,26 @@ struct sm100_fp8_config_M128 { template typename Epilogue> struct sm100_fp8_config_M64 { - // M in [1, 64] + // M in (16, 64] static_assert(std::is_same()); using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; - using TileShape = Shape<_64, _64, _256>; - using ClusterShape = Shape<_1, _8, _1>; + using TileShape = Shape<_64, _64, _128>; + using ClusterShape = Shape<_1, _1, _1>; + using Cutlass3xGemm = + cutlass_3x_gemm_sm100; +}; + +template typename Epilogue> +struct sm100_fp8_config_M16 { + // M in [1, 16] + static_assert(std::is_same()); + using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; + using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; + using TileShape = Shape<_64, _64, _128>; + using ClusterShape = Shape<_1, _4, _1>; using Cutlass3xGemm = cutlass_3x_gemm_sm100; @@ -82,27 +82,27 @@ inline void cutlass_gemm_sm100_fp8_dispatch(torch::Tensor& out, using Cutlass3xGemmDefault = typename sm100_fp8_config_default::Cutlass3xGemm; + using Cutlass3xGemmM16 = + typename sm100_fp8_config_M16::Cutlass3xGemm; using Cutlass3xGemmM64 = typename sm100_fp8_config_M64::Cutlass3xGemm; - using Cutlass3xGemmM128 = - typename sm100_fp8_config_M128::Cutlass3xGemm; using Cutlass3xGemmM256 = typename sm100_fp8_config_M256::Cutlass3xGemm; uint32_t const m = a.size(0); uint32_t const mp2 = - std::max(static_cast(64), next_pow_2(m)); // next power of 2 + std::max(static_cast(16), next_pow_2(m)); // next power of 2 - if (mp2 <= 64) { - // m in [1, 64] + if (mp2 <= 16) { + // m in [1, 16] + return cutlass_gemm_caller( + out, a, b, std::forward(args)...); + } else if (mp2 <= 64) { + // m in (16, 64] return cutlass_gemm_caller( out, a, b, std::forward(args)...); - } else if (mp2 <= 128) { - // m in (64, 128] - return cutlass_gemm_caller( - out, a, b, std::forward(args)...); } else if (mp2 <= 256) { - // m in (128, 256] + // m in (64, 256] return cutlass_gemm_caller( out, a, b, std::forward(args)...); } else {