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 15, 03:30 UTC
[NVIDIA] Support Cutlass w8a8 FP8 for Blackwell Geforce GPUs (sm120) (#17280)
Signed-off-by: kaln27 <liaojuncheng123@foxmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Joonchen Liau committed
Jul 2, 2025 at 20:47 UTC
9e5552aa13908fcc4a4e4c22ac28529b196001f3
7 files changed
+238
-1
CMakeLists.txt
+30
@@ -420,6 +420,36 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
420
endif()
421
endif()
422
423
+
424
+ # The cutlass_scaled_mm kernels for Geforce Blackwell SM120 (c3x, i.e. CUTLASS 3.x) require
425
+ # CUDA 12.8 or later
426
+ cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0;12.0a" "${CUDA_ARCHS}")
427
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND SCALED_MM_ARCHS)
428
+ set(SRCS
429
+ "csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm120.cu"
430
+ "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8.cu"
431
+ )
432
+ set_gencode_flags_for_srcs(
433
+ SRCS "${SRCS}"
434
+ CUDA_ARCHS "${SCALED_MM_ARCHS}")
435
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
436
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
437
+ # Let scaled_mm_c2x know it doesn't need to build these arches
438
+ list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
439
+ message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
440
+ else()
441
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND SCALED_MM_ARCHS)
442
+ message(STATUS "Not building scaled_mm_c3x_sm120 as CUDA Compiler version is "
443
+ "not >= 12.8, we recommend upgrading to CUDA 12.8 or "
444
+ "later if you intend on running FP8 quantized models on "
445
+ "Blackwell.")
446
+ else()
447
+ message(STATUS "Not building scaled_mm_c3x_120 as no compatible archs found "
448
+ "in CUDA target architectures")
449
+ endif()
450
+ endif()
451
+
452
+
453
# The cutlass_scaled_mm kernels for Blackwell SM100 (c3x, i.e. CUTLASS 3.x)
454
# require CUDA 12.8 or later
455
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a" "${CUDA_ARCHS}")
csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
+61
@@ -144,4 +144,65 @@ struct cutlass_3x_gemm_sm100 {
144
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
145
};
146
147
+template <typename ElementAB_, typename ElementD_,
148
+ template <typename, typename, typename> typename Epilogue_,
149
+ typename TileShape, typename ClusterShape, typename KernelSchedule,
150
+ typename EpilogueSchedule>
151
+struct cutlass_3x_gemm_sm120 {
152
+ using ElementAB = ElementAB_;
153
+ using LayoutA = cutlass::layout::RowMajor;
154
+ static constexpr int AlignmentA =
155
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
156
+
157
+ using LayoutB = cutlass::layout::ColumnMajor;
158
+ static constexpr int AlignmentB =
159
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
160
+
161
+ using ElementC = void;
162
+ using LayoutC = cutlass::layout::RowMajor;
163
+ static constexpr int AlignmentC =
164
+ 128 / cutlass::sizeof_bits<ElementD_>::value;
165
+
166
+ using ElementD = ElementD_;
167
+ using LayoutD = cutlass::layout::RowMajor;
168
+ static constexpr int AlignmentD = AlignmentC;
169
+
170
+ using ElementAcc =
171
+ typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
172
+ float>::type;
173
+ using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
174
+
175
+ // MMA type
176
+ using ElementAccumulator = float;
177
+
178
+ // Epilogue types
179
+ using ElementBias = cutlass::half_t;
180
+ using ElementCompute = float;
181
+ using ElementAux = ElementD;
182
+ using LayoutAux = LayoutD;
183
+ using ElementAmax = float;
184
+
185
+ using EVTCompute = typename Epilogue::EVTCompute;
186
+
187
+ using CollectiveEpilogue =
188
+ typename cutlass::epilogue::collective::CollectiveBuilder<
189
+ cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, TileShape,
190
+ ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
191
+ ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
192
+ ElementD, LayoutD, AlignmentD, EpilogueSchedule,
193
+ EVTCompute>::CollectiveOp;
194
+
195
+ using CollectiveMainloop =
196
+ typename cutlass::gemm::collective::CollectiveBuilder<
197
+ cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp, ElementAB,
198
+ LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
199
+ ElementAccumulator, TileShape, ClusterShape,
200
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
201
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
202
+ KernelSchedule>::CollectiveOp;
203
+
204
+ using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
205
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
206
+};
207
+
208
} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
+6
@@ -36,6 +36,12 @@ void cutlass_scaled_mm_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
36
torch::Tensor const& b_scales,
37
std::optional<torch::Tensor> const& bias);
38
39
+void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
40
+ torch::Tensor const& b,
41
+ torch::Tensor const& a_scales,
42
+ torch::Tensor const& b_scales,
43
+ std::optional<torch::Tensor> const& bias);
44
+
45
void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out,
46
torch::Tensor const& a,
47
torch::Tensor const& b,
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8.cu
new
+24
@@ -0,0 +1,24 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm120_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_sm120_fp8(torch::Tensor& out, torch::Tensor const& a,
8
+ torch::Tensor const& b,
9
+ torch::Tensor const& a_scales,
10
+ torch::Tensor const& b_scales,
11
+ std::optional<torch::Tensor> const& bias) {
12
+ TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
13
+ if (bias) {
14
+ TORCH_CHECK(bias->dtype() == out.dtype(),
15
+ "currently bias dtype must match output dtype ", out.dtype());
16
+ return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
17
+ out, a, b, a_scales, b_scales, *bias);
18
+ } else {
19
+ return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
20
+ out, a, b, a_scales, b_scales);
21
+ }
22
+}
23
+
24
+} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm120_fp8_dispatch.cuh
new
+67
@@ -0,0 +1,67 @@
1
+#pragma once
2
+
3
+#include "scaled_mm.cuh"
4
+#include "cutlass_gemm_caller.cuh"
5
+
6
+/**
7
+ * This file defines Gemm kernel configurations for SM120 (fp8) based on the
8
+ * Gemm shape.
9
+ */
10
+
11
+namespace vllm {
12
+
13
+using c3x::cutlass_gemm_caller;
14
+
15
+template <typename InType, typename OutType,
16
+ template <typename, typename, typename> typename Epilogue>
17
+struct sm120_fp8_config_default {
18
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
19
+ using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
20
+ using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
21
+ using TileShape = Shape<_128, _128, _128>;
22
+ using ClusterShape = Shape<_1, _1, _1>; // Only work with Shape<_1, _1, _1>
23
+ using Cutlass3xGemm =
24
+ cutlass_3x_gemm_sm120<InType, OutType, Epilogue, TileShape, ClusterShape,
25
+ KernelSchedule, EpilogueSchedule>;
26
+};
27
+
28
+template <typename InType, typename OutType,
29
+ template <typename, typename, typename> typename Epilogue,
30
+ typename... EpilogueArgs>
31
+inline void cutlass_gemm_sm120_fp8_dispatch(torch::Tensor& out,
32
+ torch::Tensor const& a,
33
+ torch::Tensor const& b,
34
+ EpilogueArgs&&... args) {
35
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
36
+ TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
37
+ TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
38
+
39
+ using Cutlass3xGemmDefault =
40
+ typename sm120_fp8_config_default<InType, OutType,
41
+ Epilogue>::Cutlass3xGemm;
42
+ return cutlass_gemm_caller<Cutlass3xGemmDefault>(
43
+ out, a, b, std::forward<EpilogueArgs>(args)...);
44
+}
45
+
46
+template <template <typename, typename, typename> typename Epilogue,
47
+ typename... EpilogueArgs>
48
+void cutlass_scaled_mm_sm120_fp8_epilogue(torch::Tensor& out,
49
+ torch::Tensor const& a,
50
+ torch::Tensor const& b,
51
+ EpilogueArgs&&... epilogue_args) {
52
+ TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn);
53
+ TORCH_CHECK(b.dtype() == torch::kFloat8_e4m3fn);
54
+
55
+ if (out.dtype() == torch::kBFloat16) {
56
+ return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
57
+ cutlass::bfloat16_t, Epilogue>(
58
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
59
+ } else {
60
+ TORCH_CHECK(out.dtype() == torch::kFloat16);
61
+ return cutlass_gemm_sm120_fp8_dispatch<cutlass::float_e4m3_t,
62
+ cutlass::half_t, Epilogue>(
63
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
64
+ }
65
+}
66
+
67
+} // namespace vllm
\ No newline at end of file
csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm120.cu
new
+34
@@ -0,0 +1,34 @@
1
+#include <cudaTypedefs.h>
2
+#include "c3x/scaled_mm_kernels.hpp"
3
+
4
+#include "cuda_utils.h"
5
+
6
+/*
7
+ This file defines quantized GEMM operations using the CUTLASS 3.x API, for
8
+ NVIDIA GPUs with sm120 (Blackwell Geforce).
9
+*/
10
+
11
+#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
12
+
13
+void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
14
+ torch::Tensor const& b,
15
+ torch::Tensor const& a_scales,
16
+ torch::Tensor const& b_scales,
17
+ std::optional<torch::Tensor> const& bias) {
18
+ TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
19
+ TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
20
+
21
+ int M = a.size(0), N = b.size(1), K = a.size(1);
22
+ TORCH_CHECK(
23
+ (a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
24
+ (b_scales.numel() == 1 || b_scales.numel() == b.size(1)),
25
+ "Currently, block scaled fp8 gemm is not implemented for Blackwell");
26
+
27
+ // Standard per-tensor/per-token/per-channel scaling
28
+ TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
29
+ TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn,
30
+ "Currently, only fp8 gemm is implemented for Blackwell");
31
+ vllm::cutlass_scaled_mm_sm120_fp8(c, a, b, a_scales, b_scales, bias);
32
+}
33
+
34
+#endif
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+16
-1
@@ -41,6 +41,14 @@ void cutlass_moe_mm_sm90(
41
42
#endif
43
44
+#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
45
+void cutlass_scaled_mm_sm120(torch::Tensor& c, torch::Tensor const& a,
46
+ torch::Tensor const& b,
47
+ torch::Tensor const& a_scales,
48
+ torch::Tensor const& b_scales,
49
+ std::optional<torch::Tensor> const& bias);
50
+#endif
51
+
52
#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
53
void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
54
torch::Tensor const& b,
@@ -168,8 +176,15 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
176
at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
177
int32_t version_num = get_sm_version_num();
178
179
+#if defined ENABLE_SCALED_MM_SM120 && ENABLE_SCALED_MM_SM120
180
+ if (version_num >= 120) {
181
+ cutlass_scaled_mm_sm120(c, a, b, a_scales, b_scales, bias);
182
+ return;
183
+ }
184
+#endif
185
+
186
#if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
172
- if (version_num >= 100) {
187
+ if (version_num >= 100 && version_num < 120) {
188
cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
189
return;
190
}