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, 11:00 UTC
Add cutlass support for blackwell fp8 blockwise gemm (#14383)
Signed-off-by: Shu Wang <shuw@nvidia.com>
Shu Wang committed
May 8, 2025 at 17:09 UTC
376786fac1fc50e8d788a39a91fa28d1709ad48b
11 files changed
+332
-64
CMakeLists.txt
+1
@@ -418,6 +418,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
418
set(SRCS
419
"csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm100.cu"
420
"csrc/quantization/cutlass_w8a8/c3x/scaled_mm_sm100_fp8.cu"
421
+ "csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8.cu"
422
)
423
set_gencode_flags_for_srcs(
424
SRCS "${SRCS}"
csrc/cutlass_extensions/common.hpp
+10
@@ -59,3 +59,13 @@ struct enable_sm90_only : Kernel {
59
#endif
60
}
61
};
62
+
63
+template <typename Kernel>
64
+struct enable_sm100_only : Kernel {
65
+ template <typename... Args>
66
+ CUTLASS_DEVICE void operator()(Args&&... args) {
67
+#if defined __CUDA_ARCH__ && __CUDA_ARCH__ == 1000
68
+ Kernel::operator()(std::forward<Args>(args)...);
69
+#endif
70
+ }
71
+};
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8.cu
new
+27
@@ -0,0 +1,27 @@
1
+#include "scaled_mm_kernels.hpp"
2
+#include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh"
3
+#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
4
+
5
+namespace vllm {
6
+
7
+void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out,
8
+ torch::Tensor const& a,
9
+ torch::Tensor const& b,
10
+ torch::Tensor const& a_scales,
11
+ torch::Tensor const& b_scales) {
12
+ TORCH_CHECK(
13
+ a.size(0) % 4 == 0,
14
+ "Input tensor must have a number of rows that is a multiple of 4. ",
15
+ "but got: ", a.size(0), " rows.");
16
+ if (out.dtype() == torch::kBFloat16) {
17
+ cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::bfloat16_t>(
18
+ out, a, b, a_scales, b_scales);
19
+
20
+ } else {
21
+ TORCH_CHECK(out.dtype() == torch::kFloat16);
22
+ cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::half_t>(
23
+ out, a, b, a_scales, b_scales);
24
+ }
25
+}
26
+
27
+} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_blockwise_sm100_fp8_dispatch.cuh
new
+205
@@ -0,0 +1,205 @@
1
+#pragma once
2
+
3
+#include "cutlass/cutlass.h"
4
+#include "cutlass/numeric_types.h"
5
+
6
+#include "cute/tensor.hpp"
7
+#include "cutlass/tensor_ref.h"
8
+#include "cutlass/gemm/dispatch_policy.hpp"
9
+#include "cutlass/gemm/collective/collective_builder.hpp"
10
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
11
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
12
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
13
+#include "cutlass/epilogue/dispatch_policy.hpp"
14
+#include "cutlass/epilogue/collective/collective_builder.hpp"
15
+
16
+#include "cutlass_extensions/gemm/dispatch_policy.hpp"
17
+#include "cutlass_extensions/gemm/collective/collective_builder.hpp"
18
+
19
+#include "cutlass_gemm_caller.cuh"
20
+
21
+namespace vllm {
22
+
23
+using namespace cute;
24
+
25
+template <typename OutType, typename MmaTileShape, typename ScalesPerTile,
26
+ class ClusterShape, typename EpilogueScheduler,
27
+ typename MainloopScheduler>
28
+struct cutlass_3x_gemm_fp8_blockwise {
29
+ using ElementAB = cutlass::float_e4m3_t;
30
+
31
+ using ElementA = ElementAB;
32
+ using LayoutA = cutlass::layout::RowMajor;
33
+ static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
34
+
35
+ using ElementB = ElementAB;
36
+ using LayoutB = cutlass::layout::ColumnMajor;
37
+ static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
38
+
39
+ using ElementC = void;
40
+ using ElementD = OutType;
41
+ using LayoutD = cutlass::layout::RowMajor;
42
+ static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
43
+
44
+ using LayoutC = LayoutD;
45
+ static constexpr int AlignmentC = AlignmentD;
46
+
47
+ using ElementAccumulator = float;
48
+ using ElementCompute = float;
49
+ using ElementBlockScale = float;
50
+
51
+ // MMA and Cluster Tile Shapes
52
+ // Shape of the tile computed by tcgen05 MMA, could be across 2 SMs if Cluster
53
+ // Shape %2 == 0 using MmaTileShape_MNK = Shape<_128,_128,_128>;
54
+ static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{});
55
+ static constexpr int ScaleGranularityM =
56
+ size<0>(MmaTileShape{}) / ScaleMsPerTile;
57
+ static constexpr int ScaleGranularityN =
58
+ size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{});
59
+ static constexpr int ScaleGranularityK =
60
+ size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{});
61
+
62
+ // Shape of the threadblocks in a cluster
63
+ using ClusterShape_MNK = ClusterShape;
64
+
65
+ using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<
66
+ ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
67
+ cute::UMMA::Major::MN, cute::UMMA::Major::K>;
68
+ using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
69
+ using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
70
+
71
+ using ArchTag = cutlass::arch::Sm100;
72
+ using OperatorClass = cutlass::arch::OpClassTensorOp;
73
+
74
+ static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
75
+ using ElementScalar = float;
76
+ // clang-format off
77
+ using DefaultOperation = cutlass::epilogue::fusion::LinearCombination<ElementD, ElementCompute, ElementC, ElementScalar, RoundStyle>;
78
+ using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
79
+ ArchTag,
80
+ OperatorClass,
81
+ MmaTileShape,
82
+ ClusterShape,
83
+ cutlass::epilogue::collective::EpilogueTileAuto,
84
+ ElementAccumulator,
85
+ ElementCompute,
86
+ ElementC,
87
+ LayoutC,
88
+ AlignmentC,
89
+ ElementD,
90
+ LayoutD,
91
+ AlignmentD,
92
+ EpilogueScheduler,
93
+ DefaultOperation
94
+ >::CollectiveOp;
95
+
96
+ using StageCountType = cutlass::gemm::collective::StageCountAuto;
97
+ using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
98
+ ArchTag,
99
+ OperatorClass,
100
+ ElementA,
101
+ cute::tuple<LayoutA, LayoutSFA>,
102
+ AlignmentA,
103
+ ElementB,
104
+ cute::tuple<LayoutB, LayoutSFB>,
105
+ AlignmentB,
106
+ ElementAccumulator,
107
+ MmaTileShape,
108
+ ClusterShape,
109
+
110
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
111
+ MainloopScheduler
112
+ >::CollectiveOp;
113
+ // clang-format on
114
+
115
+ using KernelType = enable_sm100_only<cutlass::gemm::kernel::GemmUniversal<
116
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
117
+
118
+ struct GemmKernel : public KernelType {};
119
+};
120
+
121
+template <typename Gemm>
122
+void cutlass_gemm_caller_blockwise(torch::Tensor& out, torch::Tensor const& a,
123
+ torch::Tensor const& b,
124
+ torch::Tensor const& a_scales,
125
+ torch::Tensor const& b_scales) {
126
+ using GemmKernel = typename Gemm::GemmKernel;
127
+ using StrideA = typename Gemm::GemmKernel::StrideA;
128
+ using StrideB = typename Gemm::GemmKernel::StrideB;
129
+ using StrideD = typename Gemm::GemmKernel::StrideD;
130
+ using StrideC = typename Gemm::GemmKernel::StrideC;
131
+ using LayoutSFA = typename Gemm::LayoutSFA;
132
+ using LayoutSFB = typename Gemm::LayoutSFB;
133
+ using ScaleConfig = typename Gemm::ScaleConfig;
134
+
135
+ using ElementAB = typename Gemm::ElementAB;
136
+ using ElementD = typename Gemm::ElementD;
137
+
138
+ int32_t m = a.size(0), n = b.size(1), k = a.size(1);
139
+ auto prob_shape = cute::make_shape(m, n, k, 1);
140
+
141
+ StrideA a_stride;
142
+ StrideB b_stride;
143
+ StrideC c_stride;
144
+ a_stride =
145
+ cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
146
+ b_stride =
147
+ cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
148
+ c_stride =
149
+ cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
150
+
151
+ LayoutSFA layout_SFA =
152
+ ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
153
+ LayoutSFB layout_SFB =
154
+ ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
155
+
156
+ auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
157
+ auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
158
+ auto a_scales_ptr = static_cast<float*>(a_scales.data_ptr());
159
+ auto b_scales_ptr = static_cast<float*>(b_scales.data_ptr());
160
+
161
+ typename GemmKernel::MainloopArguments mainloop_args{
162
+ a_ptr, a_stride, b_ptr, b_stride,
163
+ a_scales_ptr, layout_SFA, b_scales_ptr, layout_SFB};
164
+
165
+ auto c_ptr = static_cast<ElementD*>(out.data_ptr());
166
+ typename GemmKernel::EpilogueArguments epilogue_args{
167
+ {}, c_ptr, c_stride, c_ptr, c_stride};
168
+ c3x::cutlass_gemm_caller<GemmKernel>(a.device(), prob_shape, mainloop_args,
169
+ epilogue_args);
170
+}
171
+
172
+template <typename OutType>
173
+void cutlass_gemm_blockwise_sm100_fp8_dispatch(torch::Tensor& out,
174
+ torch::Tensor const& a,
175
+ torch::Tensor const& b,
176
+ torch::Tensor const& a_scales,
177
+ torch::Tensor const& b_scales) {
178
+ auto m = a.size(0);
179
+ auto k = a.size(1);
180
+ auto n = b.size(1);
181
+ int sms;
182
+ cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device());
183
+
184
+ auto should_use_2sm = [&sms](int m, int n, int tile1SM = 128) {
185
+ return std::ceil(static_cast<float>(m) / tile1SM) *
186
+ std::ceil(static_cast<float>(n) / tile1SM) >=
187
+ sms;
188
+ };
189
+ bool use_2sm = should_use_2sm(m, n);
190
+ if (use_2sm) {
191
+ cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
192
+ OutType, Shape<_256, _128, _128>, Shape<_256, _1, _1>,
193
+ Shape<_2, _2, _1>, cutlass::epilogue::TmaWarpSpecialized2Sm,
194
+ cutlass::gemm::KernelTmaWarpSpecializedBlockwise2SmSm100>>(
195
+ out, a, b, a_scales, b_scales);
196
+ } else {
197
+ cutlass_gemm_caller_blockwise<cutlass_3x_gemm_fp8_blockwise<
198
+ OutType, Shape<_128, _128, _128>, Shape<_128, _1, _1>,
199
+ Shape<_1, _1, _1>, cutlass::epilogue::TmaWarpSpecialized1Sm,
200
+ cutlass::gemm::KernelTmaWarpSpecializedBlockwise1SmSm100>>(
201
+ out, a, b, a_scales, b_scales);
202
+ }
203
+}
204
+
205
+} // namespace vllm
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_helper.hpp
new
+57
@@ -0,0 +1,57 @@
1
+#include <torch/all.h>
2
+#include "cuda_utils.h"
3
+
4
+template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
5
+void dispatch_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
6
+ torch::Tensor const& b, torch::Tensor const& a_scales,
7
+ torch::Tensor const& b_scales,
8
+ std::optional<torch::Tensor> const& bias,
9
+ Fp8Func fp8_func, Int8Func int8_func,
10
+ BlockwiseFunc blockwise_func) {
11
+ TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
12
+ TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
13
+
14
+ int M = a.size(0), N = b.size(1), K = a.size(1);
15
+
16
+ if ((a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
17
+ (b_scales.numel() == 1 || b_scales.numel() == b.size(1))) {
18
+ // Standard per-tensor/per-token/per-channel scaling
19
+ TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
20
+ if (a.dtype() == torch::kFloat8_e4m3fn) {
21
+ fp8_func(c, a, b, a_scales, b_scales, bias);
22
+ } else {
23
+ TORCH_CHECK(a.dtype() == torch::kInt8);
24
+ if constexpr (!std::is_same_v<Int8Func, std::nullptr_t>) {
25
+ int8_func(c, a, b, a_scales, b_scales, bias);
26
+ } else {
27
+ TORCH_CHECK(false, "Int8 not supported for this architecture");
28
+ }
29
+ }
30
+ } else {
31
+ using GroupShape = std::array<int64_t, 2>;
32
+ auto make_group_shape = [](torch::Tensor const& x,
33
+ torch::Tensor const& s) -> GroupShape {
34
+ TORCH_CHECK(s.dim() == 2, "cutlass_scaled_mm group scales must be 2D");
35
+ return {cuda_utils::ceil_div(x.size(0), s.size(0)),
36
+ cuda_utils::ceil_div(x.size(1), s.size(1))};
37
+ };
38
+
39
+ GroupShape a_scale_group_shape = make_group_shape(a, a_scales);
40
+ GroupShape b_scale_group_shape = make_group_shape(b, b_scales);
41
+
42
+ // 1x128 per-token group scales for activations
43
+ // 128x128 blockwise scales for weights
44
+ TORCH_CHECK((a_scale_group_shape == GroupShape{1, 128} &&
45
+ b_scale_group_shape == GroupShape{128, 128} &&
46
+ a.dtype() == torch::kFloat8_e4m3fn &&
47
+ b.dtype() == torch::kFloat8_e4m3fn),
48
+ "cutlass_scaled_mm only supports datatype float8_e4m3fn.\n"
49
+ "a_scale_group_shape must be [1, 128]. Got: [",
50
+ a_scale_group_shape[0], ", ", a_scale_group_shape[1],
51
+ "]\n"
52
+ "b_scale_group_shape must be [128, 128]. Got: [",
53
+ b_scale_group_shape[0], ", ", b_scale_group_shape[1], "]");
54
+ TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm");
55
+ blockwise_func(c, a, b, a_scales, b_scales);
56
+ }
57
+}
csrc/quantization/cutlass_w8a8/c3x/scaled_mm_kernels.hpp
+5
@@ -36,4 +36,9 @@ 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_blockwise_sm100_fp8(torch::Tensor& out,
40
+ torch::Tensor const& a,
41
+ torch::Tensor const& b,
42
+ torch::Tensor const& a_scales,
43
+ torch::Tensor const& b_scales);
44
} // namespace vllm
csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm100.cu
+5
-17
@@ -1,8 +1,6 @@
1
-#include <cudaTypedefs.h>
1
+#include "c3x/scaled_mm_helper.hpp"
2
#include "c3x/scaled_mm_kernels.hpp"
3
4
-#include "cuda_utils.h"
5
-
4
/*
5
This file defines quantized GEMM operations using the CUTLASS 3.x API, for
6
NVIDIA GPUs with sm100 (Blackwell).
@@ -15,20 +13,10 @@ void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
13
torch::Tensor const& a_scales,
14
torch::Tensor const& b_scales,
15
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_sm100_fp8(c, a, b, a_scales, b_scales, bias);
16
+ dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
17
+ vllm::cutlass_scaled_mm_sm100_fp8,
18
+ nullptr, // int8 not supported on SM100
19
+ vllm::cutlass_scaled_mm_blockwise_sm100_fp8);
20
}
21
22
#endif
csrc/quantization/cutlass_w8a8/scaled_mm_c3x_sm90.cu
+5
-46
@@ -1,8 +1,6 @@
1
-#include <cudaTypedefs.h>
1
+#include "c3x/scaled_mm_helper.hpp"
2
#include "c3x/scaled_mm_kernels.hpp"
3
4
-#include "cuda_utils.h"
5
-
4
/*
5
This file defines quantized GEMM operations using the CUTLASS 3.x API, for
6
NVIDIA GPUs with sm90a (Hopper).
@@ -15,49 +13,10 @@ void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
13
torch::Tensor const& a_scales,
14
torch::Tensor const& b_scales,
15
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
-
23
- if ((a_scales.numel() == 1 || a_scales.numel() == a.size(0)) &&
24
- (b_scales.numel() == 1 || b_scales.numel() == b.size(1))) {
25
- // Standard per-tensor/per-token/per-channel scaling
26
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
27
- if (a.dtype() == torch::kFloat8_e4m3fn) {
28
- vllm::cutlass_scaled_mm_sm90_fp8(c, a, b, a_scales, b_scales, bias);
29
- } else {
30
- TORCH_CHECK(a.dtype() == torch::kInt8);
31
- vllm::cutlass_scaled_mm_sm90_int8(c, a, b, a_scales, b_scales, bias);
32
- }
33
- } else {
34
- using GroupShape = std::array<int64_t, 2>;
35
- auto make_group_shape = [](torch::Tensor const& x,
36
- torch::Tensor const& s) -> GroupShape {
37
- TORCH_CHECK(s.dim() == 2, "cutlass_scaled_mm group scales must be 2D");
38
- return {cuda_utils::ceil_div(x.size(0), s.size(0)),
39
- cuda_utils::ceil_div(x.size(1), s.size(1))};
40
- };
41
-
42
- GroupShape a_scale_group_shape = make_group_shape(a, a_scales);
43
- GroupShape b_scale_group_shape = make_group_shape(b, b_scales);
44
-
45
- // 1x128 per-token group scales for activations
46
- // 128x128 blockwise scales for weights
47
- TORCH_CHECK((a_scale_group_shape == GroupShape{1, 128} &&
48
- b_scale_group_shape == GroupShape{128, 128} &&
49
- a.dtype() == torch::kFloat8_e4m3fn &&
50
- b.dtype() == torch::kFloat8_e4m3fn),
51
- "cutlass_scaled_mm only supports datatype float8_e4m3fn.\n"
52
- "a_scale_group_shape must be [1, 128]. Got: [",
53
- a_scale_group_shape[0], ", ", a_scale_group_shape[1],
54
- "]\n"
55
- "b_scale_group_shape must be [128, 128]. Got: [",
56
- b_scale_group_shape[0], ", ", b_scale_group_shape[1], "]");
57
- TORCH_CHECK(!bias, "Bias not yet supported blockwise scaled_mm");
58
-
59
- vllm::cutlass_scaled_mm_blockwise_sm90_fp8(c, a, b, a_scales, b_scales);
60
- }
16
+ dispatch_scaled_mm(c, a, b, a_scales, b_scales, bias,
17
+ vllm::cutlass_scaled_mm_sm90_fp8,
18
+ vllm::cutlass_scaled_mm_sm90_int8,
19
+ vllm::cutlass_scaled_mm_blockwise_sm90_fp8);
20
}
21
22
void cutlass_scaled_mm_azp_sm90(torch::Tensor& out, torch::Tensor const& a,
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+2
@@ -110,6 +110,8 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
110
#if defined CUDA_VERSION
111
if (cuda_device_capability >= 90 && cuda_device_capability < 100) {
112
return CUDA_VERSION >= 12000;
113
+ } else if (cuda_device_capability >= 100) {
114
+ return CUDA_VERSION >= 12080;
115
}
116
#endif
117
tests/kernels/quantization/test_cutlass_scaled_mm.py
+3
-1
@@ -95,7 +95,7 @@ def cutlass_fp8_gemm_helper(m: int,
95
out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
96
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
97
98
- torch.testing.assert_close(out, baseline, rtol=1e-2, atol=5e-2)
98
+ torch.testing.assert_close(out, baseline, rtol=1e-2, atol=1.5e-1)
99
100
opcheck(torch.ops._C.cutlass_scaled_mm,
101
(out, a, b, scale_a, scale_b, bias))
@@ -161,6 +161,8 @@ def test_cutlass_fp8_blockwise_scale_gemm(m: int, n: int, k: int,
161
return
162
if m % a_scale_group_shape[0] != 0 or k % a_scale_group_shape[1] != 0:
163
return
164
+ if m % 4 != 0 and current_platform.has_device_capability(100):
165
+ return
166
cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape,
167
use_bias)
168
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+12
@@ -57,6 +57,16 @@ def apply_w8a8_block_fp8_linear(
57
or br not in (1, weight.shape[0])):
58
shape_supported_by_cutlass = False
59
if cutlass_block_fp8_supported and shape_supported_by_cutlass:
60
+ rows, cols = input_2d.shape
61
+ # Blackwell GPUs (SM100) require row dimensions to be multiple of 4 for
62
+ # optimal tensor core usage. Can be removed when targeting platforms
63
+ # without this constraint.
64
+ should_pad = current_platform.has_device_capability(
65
+ 100) and rows % 4 != 0
66
+ if should_pad:
67
+ input_2d = torch.nn.functional.pad(input_2d,
68
+ (0, 0, 0, 4 - (rows % 4)),
69
+ value=0).contiguous()
70
q_input, x_scale = per_token_group_quant_fp8(input_2d,
71
block_size[1],
72
column_major_scales=True)
@@ -65,6 +75,8 @@ def apply_w8a8_block_fp8_linear(
75
out_dtype=input.dtype,
76
scale_a=x_scale,
77
scale_b=weight_scale.T)
78
+ if should_pad:
79
+ output = output[:rows, :]
80
else:
81
q_input, x_scale = per_token_group_quant_fp8(input_2d,
82
block_size[1],