A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 04:30 UTC
[kernel] Improve FP8 PTPC on Hopper for larger shapes (#28692)
Signed-off-by: czhu-cohere <conway.zhu@cohere.com>
czhu-cohere committed
Nov 14, 2025 at 12:59 UTC
cdd7025961cf79480f885804c21e7d60866fb33f
1 file changed
+27
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh
+27
@@ -116,6 +116,26 @@ struct sm90_fp8_config_default {
116
ClusterShape, KernelSchedule, EpilogueSchedule>>;
117
};
118
119
+template <typename InType, typename OutType, bool EnableBias>
120
+struct sm90_fp8_config_M8192_K6144 {
121
+ // M >= 8192, K >= 6144
122
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
123
+ using KernelSchedule =
124
+ cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
125
+ using EpilogueSchedule =
126
+ typename cutlass::epilogue::TmaWarpSpecializedCooperative;
127
+ using TileShape = Shape<_256, _128, _128>;
128
+ using ClusterShape = Shape<_2, _1, _1>;
129
+
130
+ using Cutlass3xGemm = conditional_t<
131
+ EnableBias,
132
+ cutlass_3x_gemm_sm90_fp8<InType, OutType, c3x::ScaledEpilogueBias,
133
+ TileShape, ClusterShape, KernelSchedule,
134
+ EpilogueSchedule>,
135
+ cutlass_3x_gemm_sm90_fp8<InType, OutType, c3x::ScaledEpilogue, TileShape,
136
+ ClusterShape, KernelSchedule, EpilogueSchedule>>;
137
+};
138
+
139
template <typename InType, typename OutType, bool EnableBias>
140
struct sm90_fp8_config_M128 {
141
// M in (64, 128]
@@ -273,6 +293,9 @@ inline void cutlass_gemm_sm90_fp8_dispatch(torch::Tensor& out,
293
using Cutlass3xGemmDefault =
294
typename sm90_fp8_config_default<InType, OutType,
295
EnableBias>::Cutlass3xGemm;
296
+ using Cutlass3xGemmM8192_K6144 =
297
+ typename sm90_fp8_config_M8192_K6144<InType, OutType,
298
+ EnableBias>::Cutlass3xGemm;
299
using Cutlass3xGemmM128 =
300
typename sm90_fp8_config_M128<InType, OutType, EnableBias>::Cutlass3xGemm;
301
@@ -291,6 +314,7 @@ inline void cutlass_gemm_sm90_fp8_dispatch(torch::Tensor& out,
314
315
uint32_t const m = a.size(0);
316
uint32_t const n = b.size(1);
317
+ uint32_t const k = a.size(1);
318
319
if (m <= 16) {
320
// m in [1, 16]
@@ -312,6 +336,9 @@ inline void cutlass_gemm_sm90_fp8_dispatch(torch::Tensor& out,
336
// m in (64, 128]
337
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM128>(
338
out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
339
+ } else if (m >= 8192 && k >= 6144) {
340
+ return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM8192_K6144>(
341
+ out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
342
} else {
343
// m in (128, inf)
344
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmDefault>(