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, 21:00 UTC
add cutlass support for blackwell fp8 gemm (#13798)
kushanam committed
Mar 4, 2025 at 07:55 UTC
f89978ad7cf3732a4a482db8e784aba7e417abbd
11 files changed
+272
-65
CMakeLists.txt
+24
-8
@@ -31,7 +31,7 @@ set(ignoreMe "${VLLM_PYTHON_PATH}")
31
set(PYTHON_SUPPORTED_VERSIONS "3.9" "3.10" "3.11" "3.12")
32
33
# Supported NVIDIA architectures.
34
-set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
34
+set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0")
35
36
# Supported AMD GPU architectures.
37
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101")
@@ -297,7 +297,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
297
# Only build Marlin kernels if we are building for at least some compatible archs.
298
# Keep building Marlin for 9.0 as there are some group sizes and shapes that
299
# are not supported by Machete yet.
300
- cuda_archs_loose_intersection(MARLIN_ARCHS "8.0;8.6;8.7;8.9;9.0" "${CUDA_ARCHS}")
300
+ cuda_archs_loose_intersection(MARLIN_ARCHS "8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
301
if (MARLIN_ARCHS)
302
set(MARLIN_SRCS
303
"csrc/quantization/fp8/fp8_marlin.cu"
@@ -335,7 +335,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
335
336
# The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
337
# CUDA 12.0 or later (and only work on Hopper, 9.0a for now).
338
- cuda_archs_loose_intersection(SCALED_MM_3X_ARCHS "9.0a" "${CUDA_ARCHS}")
338
+ cuda_archs_loose_intersection(SCALED_MM_3X_ARCHS "9.0a;10.0a;10.1a;12.0a" "${CUDA_ARCHS}")
339
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.0 AND SCALED_MM_3X_ARCHS)
340
set(SRCS
341
"csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu"
@@ -369,7 +369,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
369
# For the cutlass_scaled_mm kernels we want to build the c2x (CUTLASS 2.x)
370
# kernels for the remaining archs that are not already built for 3x.
371
cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS
372
- "7.5;8.0;8.6;8.7;8.9;9.0" "${CUDA_ARCHS}")
372
+ "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
373
# subtract out the archs that are already built for 3x
374
list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS})
375
if (SCALED_MM_2X_ARCHS)
@@ -394,7 +394,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
394
# 2:4 Sparse Kernels
395
396
# The 2:4 sparse kernels cutlass_scaled_sparse_mm and cutlass_compressor
397
- # require CUDA 12.2 or later (and only work on Hopper, 9.0a for now).
397
+ # require CUDA 12.2 or later (and only work on Hopper and Blackwell).
398
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.2 AND SCALED_MM_3X_ARCHS)
399
set(SRCS "csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu")
400
set_gencode_flags_for_srcs(
@@ -419,8 +419,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
419
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND FP4_ARCHS)
420
set(SRCS
421
"csrc/quantization/fp4/nvfp4_quant_kernels.cu"
422
- "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
423
- )
422
+ "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu")
423
set_gencode_flags_for_srcs(
424
SRCS "${SRCS}"
425
CUDA_ARCHS "${FP4_ARCHS}")
@@ -433,6 +432,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
432
set(FP4_ARCHS)
433
endif()
434
435
+ # FP8 Blackwell Archs
436
+ cuda_archs_loose_intersection(BLACKWELL_ARCHS "10.0;10.1;12.0" "${CUDA_ARCHS}")
437
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND BLACKWELL_ARCHS)
438
+ set(SRCS
439
+ "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8.cu"
440
+ )
441
+ set_gencode_flags_for_srcs(
442
+ SRCS "${SRCS}"
443
+ CUDA_ARCHS "${BLACKWELL_ARCHS}")
444
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
445
+ message(STATUS "Building FP8 for archs: ${BLACKWELL_ARCHS}")
446
+ else()
447
+ # clear BLACKWELL_ARCHS
448
+ set(BLACKWELL_ARCHS)
449
+ endif()
450
+
451
#
452
# Machete kernels
453
@@ -514,6 +529,7 @@ define_gpu_extension_target(
529
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
530
ARCHITECTURES ${VLLM_GPU_ARCHES}
531
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
532
+ INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
533
USE_SABI 3
534
WITH_SOABI)
535
@@ -537,7 +553,7 @@ set_gencode_flags_for_srcs(
553
CUDA_ARCHS "${CUDA_ARCHS}")
554
555
if(VLLM_GPU_LANG STREQUAL "CUDA")
540
- cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0;8.6;8.7;8.9;9.0" "${CUDA_ARCHS}")
556
+ cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0" "${CUDA_ARCHS}")
557
if (MARLIN_MOE_ARCHS)
558
set(MARLIN_MOE_SRC
559
"csrc/moe/marlin_kernels/marlin_moe_kernel.h"
csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+23
-25
@@ -22,7 +22,7 @@ struct identity {
22
T operator()(T lhs) const { return lhs; }
23
};
24
25
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
25
+template <typename ElementAcc, typename ElementD, typename TileShape>
26
struct TrivialEpilogue {
27
private:
28
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
@@ -44,32 +44,30 @@ struct TrivialEpilogue {
44
* This class provides the common load descriptors for the
45
* ScaledEpilogue[...] classes
46
*/
47
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
47
+template <typename ElementAcc, typename ElementD, typename TileShape>
48
struct ScaledEpilogueBase {
49
protected:
50
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
51
52
template <typename T>
53
using ColOrScalarLoad = cutlass::epilogue::fusion::Sm90ColOrScalarBroadcast<
54
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
55
- Stride<Int<1>, Int<0>, Int<0>>>;
54
+ 0 /*Stages*/, TileShape, T, Stride<Int<1>, Int<0>, Int<0>>>;
55
56
template <typename T>
57
using RowOrScalarLoad = cutlass::epilogue::fusion::Sm90RowOrScalarBroadcast<
59
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
60
- Stride<Int<0>, Int<1>, Int<0>>>;
58
+ 0 /*Stages*/, TileShape, T, Stride<Int<0>, Int<1>, Int<0>>>;
59
60
// Don't want to support nullptr by default
61
template <typename T, bool EnableNullPtr = false>
62
using ColLoad = cutlass::epilogue::fusion::Sm90ColBroadcast<
65
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
66
- Stride<Int<1>, Int<0>, Int<0>>, 128 / sizeof_bits_v<T>, EnableNullPtr>;
63
+ 0 /*Stages*/, TileShape, T, T, Stride<Int<1>, Int<0>, Int<0>>,
64
+ 128 / sizeof_bits_v<T>, EnableNullPtr>;
65
66
// Don't want to support nullptr by default
67
template <typename T, bool EnableNullPtr = false>
68
using RowLoad = cutlass::epilogue::fusion::Sm90RowBroadcast<
71
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
72
- Stride<Int<0>, Int<1>, Int<0>>, 128 / sizeof_bits_v<T>, EnableNullPtr>;
69
+ 0 /*Stages*/, TileShape, T, T, Stride<Int<0>, Int<1>, Int<0>>,
70
+ 128 / sizeof_bits_v<T>, EnableNullPtr>;
71
72
// This utility function constructs the arguments for the load descriptors
73
// from a tensor. It can handle both row and column, as well as row/column or
@@ -116,11 +114,11 @@ struct ScaledEpilogueBase {
114
the A and B operands respectively. These scales may be either per-tensor or
115
per row or column.
116
*/
119
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
117
+template <typename ElementAcc, typename ElementD, typename TileShape>
118
struct ScaledEpilogue
121
- : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
119
+ : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
120
private:
123
- using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
121
+ using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
122
using Accum = typename SUPER::Accum;
123
using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
124
using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
@@ -160,11 +158,11 @@ struct ScaledEpilogue
158
* The bias tensor must be per-output channel.
159
* ScaleA and ScaleB can be per-tensor or per-token/per-channel.
160
*/
163
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
161
+template <typename ElementAcc, typename ElementD, typename TileShape>
162
struct ScaledEpilogueBias
165
- : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
163
+ : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
164
private:
167
- using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
165
+ using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
166
using Accum = typename SUPER::Accum;
167
using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
168
using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
@@ -203,11 +201,11 @@ struct ScaledEpilogueBias
201
* bias is a column vector instead of a row vector. Useful e.g. if we are
202
* computing a GEMM via C^T += B^T A^T. This happens in the 2:4 sparse kernels.
203
*/
206
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
204
+template <typename ElementAcc, typename ElementD, typename TileShape>
205
struct ScaledEpilogueColumnBias
208
- : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
206
+ : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
207
private:
210
- using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
208
+ using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
209
using Accum = typename SUPER::Accum;
210
using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
211
using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
@@ -249,11 +247,11 @@ struct ScaledEpilogueColumnBias
247
*
248
* This epilogue also supports bias, which remains per-channel.
249
*/
252
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
250
+template <typename ElementAcc, typename ElementD, typename TileShape>
251
struct ScaledEpilogueBiasAzp
254
- : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
252
+ : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
253
private:
256
- using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
254
+ using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
255
using Accum = typename SUPER::Accum;
256
using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
257
using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
@@ -314,11 +312,11 @@ struct ScaledEpilogueBiasAzp
312
*
313
* This epilogue also supports bias, which remains per-channel.
314
*/
317
-template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
315
+template <typename ElementAcc, typename ElementD, typename TileShape>
316
struct ScaledEpilogueBiasAzpToken
319
- : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
317
+ : private ScaledEpilogueBase<ElementAcc, ElementD, TileShape> {
318
private:
321
- using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
319
+ using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, TileShape>;
320
using Accum = typename SUPER::Accum;
321
using ScaleA = typename SUPER::template ColOrScalarLoad<float>;
322
using ScaleB = typename SUPER::template RowOrScalarLoad<float>;
csrc/quantization/cutlass_w8a8/c3x/cutlass_gemm_caller.cuh
+20
-12
@@ -16,6 +16,7 @@
16
#include "cutlass/gemm/kernel/gemm_universal.hpp"
17
#include "cutlass/epilogue/collective/collective_builder.hpp"
18
#include "cutlass/gemm/collective/collective_builder.hpp"
19
+#include "cutlass/util/packed_stride.hpp"
20
21
#include "core/math.hpp"
22
#include "cutlass_extensions/common.hpp"
@@ -64,22 +65,28 @@ void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
65
torch::Tensor const& b,
66
EpilogueArgs&&... epilogue_params) {
67
using ElementAB = typename Gemm::ElementAB;
68
+ using ElementC = typename Gemm::ElementC;
69
using ElementD = typename Gemm::ElementD;
70
using GemmKernel = typename Gemm::GemmKernel;
71
70
- int64_t lda = a.stride(0);
71
- int64_t ldb = b.stride(1);
72
- int64_t ldc = out.stride(0);
73
-
74
- using StrideA = cute::Stride<int64_t, cute::Int<1>, int64_t>;
75
- using StrideB = cute::Stride<int64_t, cute::Int<1>, int64_t>;
76
- using StrideC = typename Gemm::StrideC;
77
-
78
- StrideA a_stride{lda, cute::Int<1>{}, 0};
79
- StrideB b_stride{ldb, cute::Int<1>{}, 0};
80
- StrideC c_stride{ldc, cute::Int<1>{}, cute::Int<0>{}};
72
+ using StrideA = typename Gemm::GemmKernel::StrideA;
73
+ using StrideB = typename Gemm::GemmKernel::StrideB;
74
+ using StrideC = typename Gemm::GemmKernel::StrideC;
75
+ using StrideD = StrideC;
76
+ using StrideAux = StrideC;
77
78
typename GemmKernel::ProblemShape prob_shape = get_problem_shape(a, b);
79
+ auto [M, N, K, L] = prob_shape;
80
+
81
+ StrideA a_stride =
82
+ cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, L));
83
+ StrideB b_stride =
84
+ cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, L));
85
+ StrideC c_stride =
86
+ cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L));
87
+ StrideD d_stride =
88
+ cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L));
89
+ StrideAux aux_stride = d_stride;
90
91
auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
92
auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
@@ -87,10 +94,11 @@ void cutlass_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
94
b_stride};
95
96
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
97
+ // auto d_ptr = static_cast<ElementC*>(out.data_ptr());
98
typename GemmKernel::EpilogueArguments epilogue_args{
99
Gemm::Epilogue::prepare_args(
100
std::forward<EpilogueArgs>(epilogue_params)...),
93
- c_ptr, c_stride, c_ptr, c_stride};
101
+ c_ptr, c_stride, c_ptr, d_stride};
102
103
cutlass_gemm_caller<GemmKernel>(a.device(), prob_shape, mainloop_args,
104
epilogue_args);
csrc/quantization/cutlass_w8a8/c3x/scaled_mm.cuh
+62
-6
@@ -40,12 +40,7 @@ struct cutlass_3x_gemm {
40
typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
41
float>::type;
42
43
- using EpilogueDescriptor =
44
- cutlass::epilogue::collective::detail::EpilogueDescriptor<
45
- TileShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementD,
46
- ElementD, EpilogueSchedule>;
47
-
48
- using Epilogue = Epilogue_<ElementAcc, ElementD, EpilogueDescriptor>;
43
+ using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
44
45
using StrideD = Stride<int64_t, Int<1>, Int<0>>;
46
using ElementC = void;
@@ -88,4 +83,65 @@ struct cutlass_3x_gemm {
83
struct GemmKernel : public KernelType {};
84
};
85
86
+template <typename ElementAB_, typename ElementD_,
87
+ template <typename, typename, typename> typename Epilogue_,
88
+ typename TileShape, typename ClusterShape, typename KernelSchedule,
89
+ typename EpilogueSchedule>
90
+struct cutlass_3x_gemm_sm100 {
91
+ using ElementAB = ElementAB_;
92
+ using LayoutA = cutlass::layout::RowMajor;
93
+ static constexpr int AlignmentA =
94
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
95
+
96
+ using LayoutB = cutlass::layout::ColumnMajor;
97
+ static constexpr int AlignmentB =
98
+ 128 / cutlass::sizeof_bits<ElementAB>::value;
99
+
100
+ using ElementC = void;
101
+ using LayoutC = cutlass::layout::RowMajor;
102
+ static constexpr int AlignmentC =
103
+ 128 / cutlass::sizeof_bits<ElementD_>::value;
104
+
105
+ using ElementD = ElementD_;
106
+ using LayoutD = cutlass::layout::RowMajor;
107
+ static constexpr int AlignmentD = AlignmentC;
108
+
109
+ using ElementAcc =
110
+ typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
111
+ float>::type;
112
+ using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
113
+
114
+ // MMA type
115
+ using ElementAccumulator = float;
116
+
117
+ // Epilogue types
118
+ using ElementBias = cutlass::half_t;
119
+ using ElementCompute = float;
120
+ using ElementAux = ElementD;
121
+ using LayoutAux = LayoutD;
122
+ using ElementAmax = float;
123
+
124
+ using EVTCompute = typename Epilogue::EVTCompute;
125
+
126
+ using CollectiveEpilogue =
127
+ typename cutlass::epilogue::collective::CollectiveBuilder<
128
+ cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, TileShape,
129
+ ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
130
+ ElementAccumulator, ElementCompute, ElementC, LayoutC, AlignmentC,
131
+ ElementD, LayoutD, AlignmentD, EpilogueSchedule,
132
+ EVTCompute>::CollectiveOp;
133
+
134
+ using CollectiveMainloop =
135
+ typename cutlass::gemm::collective::CollectiveBuilder<
136
+ cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, ElementAB,
137
+ LayoutA, AlignmentA, ElementAB, LayoutB, AlignmentB,
138
+ ElementAccumulator, TileShape, ClusterShape,
139
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
140
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
141
+ KernelSchedule>::CollectiveOp;
142
+
143
+ using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
144
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
145
+};
146
+
147
} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
+6
@@ -30,4 +30,10 @@ void cutlass_scaled_mm_blockwise_sm90_fp8(torch::Tensor& out,
30
torch::Tensor const& a_scales,
31
torch::Tensor const& b_scales);
32
33
+void cutlass_scaled_mm_sm100_fp8(torch::Tensor& out, torch::Tensor const& a,
34
+ torch::Tensor const& b,
35
+ torch::Tensor const& a_scales,
36
+ torch::Tensor const& b_scales,
37
+ std::optional<torch::Tensor> const& bias);
38
+
39
} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8.cu
new
+24
@@ -0,0 +1,24 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_sm100_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_sm100_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_sm100_fp8_epilogue<c3x::ScaledEpilogueBias>(
17
+ out, a, b, a_scales, b_scales, *bias);
18
+ } else {
19
+ return cutlass_scaled_mm_sm100_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_sm100_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 SM100 (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 sm100_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<_256, _128, _64>;
22
+ using ClusterShape = Shape<_2, _2, _1>;
23
+ using Cutlass3xGemm =
24
+ cutlass_3x_gemm_sm100<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_sm100_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 sm100_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_sm100_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_sm100_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_sm100_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.cu
+25
@@ -71,3 +71,28 @@ void cutlass_scaled_mm_azp_sm90(torch::Tensor& out, torch::Tensor const& a,
71
vllm::cutlass_scaled_mm_azp_sm90_int8(out, a, b, a_scales, b_scales, azp_adj,
72
azp, bias);
73
}
74
+
75
+#if defined CUDA_VERSION && CUDA_VERSION >= 12080
76
+
77
+void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
78
+ torch::Tensor const& b,
79
+ torch::Tensor const& a_scales,
80
+ torch::Tensor const& b_scales,
81
+ std::optional<torch::Tensor> const& bias) {
82
+ TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
83
+ TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
84
+
85
+ int M = a.size(0), N = b.size(1), K = a.size(1);
86
+ TORCH_CHECK(
87
+ (a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
88
+ (b_scales.numel() == 1 || b_scales.numel() == b.size(1)),
89
+ "Currently, block scaled fp8 gemm is not implemented for Blackwell");
90
+
91
+ // Standard per-tensor/per-token/per-channel scaling
92
+ TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
93
+ TORCH_CHECK(a.dtype() == torch::kFloat8_e4m3fn,
94
+ "Currently, only fp8 gemm is implemented for Blackwell");
95
+ vllm::cutlass_scaled_mm_sm100_fp8(c, a, b, a_scales, b_scales, bias);
96
+}
97
+
98
+#endif
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+19
-2
@@ -29,6 +29,11 @@ void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
29
torch::Tensor const& a_scales,
30
torch::Tensor const& b_scales,
31
std::optional<torch::Tensor> const& bias);
32
+void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
33
+ torch::Tensor const& b,
34
+ torch::Tensor const& a_scales,
35
+ torch::Tensor const& b_scales,
36
+ std::optional<torch::Tensor> const& bias);
37
#endif
38
39
void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a,
@@ -86,7 +91,7 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
91
// and at least SM90 (Hopper)
92
93
#if defined CUDA_VERSION
89
- if (cuda_device_capability >= 90) {
94
+ if (cuda_device_capability >= 90 && cuda_device_capability < 100) {
95
return CUDA_VERSION >= 12000;
96
}
97
#endif
@@ -120,10 +125,22 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
125
126
// Guard against compilation issues for sm90 kernels
127
#if defined ENABLE_SCALED_MM_C3X && ENABLE_SCALED_MM_C3X
123
- if (version_num >= 90) {
128
+
129
+ #if defined CUDA_VERSION && CUDA_VERSION < 12080
130
+ if (version_num >= 90 && version_num < 100) {
131
+ cutlass_scaled_mm_sm90(c, a, b, a_scales, b_scales, bias);
132
+ return;
133
+ }
134
+ #else
135
+ if (version_num >= 90 && version_num < 100) {
136
cutlass_scaled_mm_sm90(c, a, b, a_scales, b_scales, bias);
137
return;
138
+ } else if (version_num >= 100) {
139
+ cutlass_scaled_mm_sm100(c, a, b, a_scales, b_scales, bias);
140
+ return;
141
}
142
+ #endif
143
+
144
#endif
145
146
#if defined ENABLE_SCALED_MM_C2X && ENABLE_SCALED_MM_C2X
csrc/quantization/machete/machete_mm_kernel.cuh
+1
-6
@@ -126,15 +126,10 @@ struct MacheteKernelTemplate {
126
std::is_same_v<ElementSChannel, ElementSToken>),
127
"Currently token and channel scales (if present) must be the same type");
128
129
- using EpilogueDescriptor =
130
- cutlass::epilogue::collective::detail::EpilogueDescriptor<
131
- TileShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementD,
132
- ElementD, EpilogueSchedule>;
133
-
129
// Currently only supports float scales
130
using ChTokScalesEpilogue =
131
typename vllm::c3x::ScaledEpilogue<ElementAccumulator, ElementD,
137
- EpilogueDescriptor>;
132
+ TileShape>;
133
static_assert((with_channel_scales || with_token_scales) ||
134
(std::is_same_v<ElementSChannel, float> &&
135
std::is_same_v<ElementSToken, float>),
csrc/sparse/cutlass/sparse_scaled_mm_c3x.cuh
+1
-6
@@ -65,12 +65,7 @@ struct cutlass_sparse_3x_gemm {
65
typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
66
float>::type;
67
68
- using EpilogueDescriptor =
69
- cutlass::epilogue::collective::detail::EpilogueDescriptor<
70
- TileShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementD,
71
- ElementD, EpilogueSchedule>;
72
-
73
- using Epilogue = Epilogue_<ElementAcc, ElementD, EpilogueDescriptor>;
68
+ using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
69
70
using ElementC = void;
71
using LayoutC = cutlass::layout::RowMajor;