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, 22:15 UTC
[NVIDIA] Support nvfp4 cutlass gemm (#13571)
Kaixi Hou committed
Feb 22, 2025 at 05:24 UTC
e109e598c796ff4a911f92e24003a0a776bdd6e0
7 files changed
+494
-1
CMakeLists.txt
+3
-1
@@ -229,7 +229,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
229
230
# Set CUTLASS_REVISION manually -- its revision detection doesn't work in this case.
231
# Please keep this in sync with FetchContent_Declare line below.
232
- set(CUTLASS_REVISION "v3.7.0" CACHE STRING "CUTLASS revision to use")
232
+ set(CUTLASS_REVISION "v3.8.0" CACHE STRING "CUTLASS revision to use")
233
234
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
235
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
@@ -267,6 +267,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
267
"csrc/permute_cols.cu"
268
"csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
269
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
270
+ "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
271
"csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
272
"csrc/cutlass_extensions/common.cpp")
273
@@ -383,6 +384,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
384
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND FP4_ARCHS)
385
set(SRCS
386
"csrc/quantization/fp4/nvfp4_quant_kernels.cu"
387
+ "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
388
)
389
set_gencode_flags_for_srcs(
390
SRCS "${SRCS}"
csrc/ops.h
+5
@@ -152,6 +152,11 @@ torch::Tensor ggml_mul_mat_a8(torch::Tensor W, torch::Tensor X, int64_t type,
152
int64_t row);
153
154
#ifndef USE_ROCM
155
+void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
156
+ torch::Tensor const& B, torch::Tensor const& A_sf,
157
+ torch::Tensor const& B_sf,
158
+ torch::Tensor const& alpha);
159
+
160
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
161
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
162
csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu
new
+37
@@ -0,0 +1,37 @@
1
+/*
2
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
+ *
4
+ * Licensed under the Apache License, Version 2.0 (the "License");
5
+ * you may not use this file except in compliance with the License.
6
+ * You may obtain a copy of the License at
7
+ *
8
+ * http://www.apache.org/licenses/LICENSE-2.0
9
+ *
10
+ * Unless required by applicable law or agreed to in writing, software
11
+ * distributed under the License is distributed on an "AS IS" BASIS,
12
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ * See the License for the specific language governing permissions and
14
+ * limitations under the License.
15
+ */
16
+
17
+#include <torch/all.h>
18
+
19
+#if defined ENABLE_NVFP4 && ENABLE_NVFP4
20
+void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
21
+ torch::Tensor const& B,
22
+ torch::Tensor const& A_sf,
23
+ torch::Tensor const& B_sf,
24
+ torch::Tensor const& alpha);
25
+#endif
26
+
27
+void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
28
+ torch::Tensor const& B, torch::Tensor const& A_sf,
29
+ torch::Tensor const& B_sf,
30
+ torch::Tensor const& alpha) {
31
+#if defined ENABLE_NVFP4 && ENABLE_NVFP4
32
+ return cutlass_scaled_fp4_mm_sm100a(D, A, B, A_sf, B_sf, alpha);
33
+#endif
34
+ TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 mm kernel, vLLM should "
35
+ "be compiled using CUDA 12.8 and target "
36
+ "compute capability 100 or above.");
37
+}
csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu
new
+280
@@ -0,0 +1,280 @@
1
+/*
2
+ * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3
+ *
4
+ * Licensed under the Apache License, Version 2.0 (the "License");
5
+ * you may not use this file except in compliance with the License.
6
+ * You may obtain a copy of the License at
7
+ *
8
+ * http://www.apache.org/licenses/LICENSE-2.0
9
+ *
10
+ * Unless required by applicable law or agreed to in writing, software
11
+ * distributed under the License is distributed on an "AS IS" BASIS,
12
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ * See the License for the specific language governing permissions and
14
+ * limitations under the License.
15
+ */
16
+
17
+#include <torch/all.h>
18
+
19
+#include <ATen/cuda/CUDAContext.h>
20
+#include <c10/cuda/CUDAGuard.h>
21
+
22
+#include "cutlass_extensions/common.hpp"
23
+
24
+#include "cutlass/cutlass.h"
25
+
26
+#include "cutlass/gemm/collective/collective_builder.hpp"
27
+#include "cutlass/epilogue/collective/collective_builder.hpp"
28
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
29
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
30
+
31
+#include "cutlass/util/packed_stride.hpp"
32
+
33
+using namespace cute;
34
+
35
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
36
+// Kernel Perf config
37
+template <typename T>
38
+struct KernelTraits;
39
+
40
+template <>
41
+struct KernelTraits<float> {
42
+ using MmaTileShape = Shape<_128, _128, _256>;
43
+ using ClusterShape = Shape<_1, _1, _1>;
44
+ using PerSmTileShape_MNK = Shape<_128, _128, _256>;
45
+};
46
+
47
+template <>
48
+struct KernelTraits<cutlass::half_t> {
49
+ using MmaTileShape = Shape<_256, _256, _256>;
50
+ using ClusterShape = Shape<_4, _4, _1>;
51
+ using PerSmTileShape_MNK = Shape<_128, _256, _256>;
52
+};
53
+
54
+template <>
55
+struct KernelTraits<cutlass::bfloat16_t> {
56
+ using MmaTileShape = Shape<_256, _256, _256>;
57
+ using ClusterShape = Shape<_4, _4, _1>;
58
+ using PerSmTileShape_MNK = Shape<_128, _256, _256>;
59
+};
60
+
61
+template <typename T>
62
+struct Fp4GemmSm100 {
63
+ // A matrix configuration
64
+ using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
65
+ using LayoutATag = cutlass::layout::RowMajor;
66
+ static constexpr int AlignmentA = 32;
67
+
68
+ // B matrix configuration
69
+ using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
70
+ using LayoutBTag = cutlass::layout::ColumnMajor;
71
+ static constexpr int AlignmentB = 32;
72
+
73
+ // C/D matrix configuration
74
+ using ElementD = T;
75
+ using ElementC = T;
76
+ using LayoutCTag = cutlass::layout::RowMajor;
77
+ using LayoutDTag = cutlass::layout::RowMajor;
78
+ static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
79
+ static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
80
+ // Kernel functional config
81
+ using ElementAccumulator = float;
82
+ using ArchTag = cutlass::arch::Sm100;
83
+ using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
84
+
85
+ // Kernel Perf config
86
+ using MmaTileShape = typename KernelTraits<T>::MmaTileShape;
87
+ using ClusterShape = typename KernelTraits<T>::ClusterShape;
88
+ using PerSmTileShape_MNK = typename KernelTraits<T>::PerSmTileShape_MNK;
89
+
90
+ using CollectiveEpilogue =
91
+ typename cutlass::epilogue::collective::CollectiveBuilder<
92
+ ArchTag, OperatorClass, PerSmTileShape_MNK, ClusterShape,
93
+ cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
94
+ ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
95
+ LayoutDTag, AlignmentD,
96
+ cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
97
+
98
+ using CollectiveMainloop =
99
+ typename cutlass::gemm::collective::CollectiveBuilder<
100
+ ArchTag, OperatorClass, ElementA, LayoutATag, AlignmentA, ElementB,
101
+ LayoutBTag, AlignmentB, ElementAccumulator, MmaTileShape,
102
+ ClusterShape,
103
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
104
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
105
+ cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
106
+
107
+ using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
108
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
109
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
110
+ using StrideA = typename Gemm::GemmKernel::StrideA;
111
+ using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
112
+ using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
113
+ using StrideB = typename Gemm::GemmKernel::StrideB;
114
+ using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{}));
115
+ using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
116
+ using StrideC = typename Gemm::GemmKernel::StrideC;
117
+ using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{}));
118
+ using StrideD = typename Gemm::GemmKernel::StrideD;
119
+ using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
120
+};
121
+
122
+template <typename T>
123
+typename T::Gemm::Arguments args_from_options(
124
+ at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
125
+ at::Tensor const& A_sf, at::Tensor const& B_sf, at::Tensor const& alpha,
126
+ int64_t M, int64_t N, int64_t K) {
127
+ using ElementA = typename T::Gemm::ElementA;
128
+ using ElementB = typename T::Gemm::ElementB;
129
+ using ElementSFA = cutlass::float_ue4m3_t;
130
+ using ElementSFB = cutlass::float_ue4m3_t;
131
+ using ElementD = typename T::Gemm::ElementD;
132
+ using ElementCompute = float;
133
+ using StrideA = typename T::StrideA;
134
+ using StrideB = typename T::StrideB;
135
+ using StrideD = typename T::StrideD;
136
+ using Sm100BlkScaledConfig =
137
+ typename T::Gemm::GemmKernel::CollectiveMainloop::Sm100BlkScaledConfig;
138
+
139
+ int m = static_cast<int>(M);
140
+ int n = static_cast<int>(N);
141
+ int k = static_cast<int>(K);
142
+ auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
143
+ auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
144
+ auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
145
+
146
+ auto layout_SFA = Sm100BlkScaledConfig::tile_atom_to_shape_SFA(
147
+ cute::make_shape(m, n, k, 1));
148
+ auto layout_SFB = Sm100BlkScaledConfig::tile_atom_to_shape_SFB(
149
+ cute::make_shape(m, n, k, 1));
150
+
151
+ typename T::Gemm::Arguments arguments{
152
+ cutlass::gemm::GemmUniversalMode::kGemm,
153
+ {m, n, k, 1},
154
+ {// Mainloop arguments
155
+ static_cast<ElementA const*>(A.data_ptr()), stride_A,
156
+ static_cast<ElementB const*>(B.data_ptr()), stride_B,
157
+ static_cast<ElementSFA const*>(A_sf.data_ptr()), layout_SFA,
158
+ static_cast<ElementSFB const*>(B_sf.data_ptr()), layout_SFB},
159
+ { // Epilogue arguments
160
+ {}, // epilogue.thread
161
+ static_cast<ElementD const*>(D.data_ptr()),
162
+ stride_D,
163
+ static_cast<ElementD*>(D.data_ptr()),
164
+ stride_D}};
165
+ auto& fusion_args = arguments.epilogue.thread;
166
+ fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr());
167
+ return arguments;
168
+}
169
+
170
+template <typename T>
171
+void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
172
+ at::Tensor const& A_sf, at::Tensor const& B_sf,
173
+ at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
174
+ cudaStream_t stream) {
175
+ typename Fp4GemmSm100<T>::Gemm gemm;
176
+
177
+ auto arguments =
178
+ args_from_options<Fp4GemmSm100<T>>(D, A, B, A_sf, B_sf, alpha, m, n, k);
179
+
180
+ size_t workspace_size = Fp4GemmSm100<T>::Gemm::get_workspace_size(arguments);
181
+ auto const workspace_options =
182
+ torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
183
+ auto workspace = torch::empty(workspace_size, workspace_options);
184
+
185
+ CUTLASS_CHECK(gemm.can_implement(arguments));
186
+
187
+ CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
188
+
189
+ CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
190
+}
191
+#else
192
+template <typename T>
193
+void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
194
+ at::Tensor const& A_sf, at::Tensor const& B_sf,
195
+ at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
196
+ cudaStream_t stream) {
197
+ TORCH_CHECK(false, "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
198
+ "a CUTLASS 3.8 source directory to enable support.");
199
+}
200
+#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
201
+
202
+#define CHECK_TYPE(x, st, m) \
203
+ TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type:", m)
204
+#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
205
+#define CHECK_CONTIGUOUS(x, m) \
206
+ TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
207
+#define CHECK_INPUT(x, st, m) \
208
+ CHECK_TH_CUDA(x, m); \
209
+ CHECK_CONTIGUOUS(x, m); \
210
+ CHECK_TYPE(x, st, m)
211
+
212
+constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
213
+constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
214
+
215
+void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
216
+ torch::Tensor const& B,
217
+ torch::Tensor const& A_sf,
218
+ torch::Tensor const& B_sf,
219
+ torch::Tensor const& alpha) {
220
+ CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
221
+ CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
222
+
223
+ CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
224
+ CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
225
+
226
+ CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
227
+
228
+ TORCH_CHECK(A.dim() == 2, "a must be a matrix");
229
+ TORCH_CHECK(B.dim() == 2, "b must be a matrix");
230
+ TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
231
+ "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
232
+ A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
233
+
234
+ auto const m = A.sizes()[0];
235
+ auto const n = B.sizes()[0];
236
+ auto const k = A.sizes()[1] * 2;
237
+
238
+ constexpr int alignment = 32;
239
+ TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
240
+ ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
241
+ "), k: ", k, ".");
242
+ TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
243
+ ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
244
+
245
+ auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
246
+ int rounded_m = round_up(m, 128);
247
+ int rounded_n = round_up(n, 128);
248
+ // Since k is divisible by 32 (alignment), k / 16 is guaranteed to be an
249
+ // integer.
250
+ int rounded_k = round_up(k / 16, 4);
251
+
252
+ TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
253
+ TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
254
+ TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
255
+ "scale_a and scale_b shapes cannot be multiplied (",
256
+ A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
257
+ "x", B_sf.sizes()[1], ")");
258
+ TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
259
+ "scale_a must be padded and swizzled to a shape (", rounded_m,
260
+ "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
261
+ A_sf.sizes()[1], ")");
262
+ TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
263
+ "scale_b must be padded and swizzled to a shape (", rounded_n,
264
+ "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
265
+ B_sf.sizes()[1], ")");
266
+
267
+ auto out_dtype = D.dtype();
268
+ at::cuda::CUDAGuard device_guard{(char)A.get_device()};
269
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
270
+
271
+ if (out_dtype == at::ScalarType::Half) {
272
+ runGemm<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
273
+ } else if (out_dtype == at::ScalarType::BFloat16) {
274
+ runGemm<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
275
+ } else if (out_dtype == at::ScalarType::Float) {
276
+ runGemm<float>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
277
+ } else {
278
+ TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm");
279
+ }
280
+}
csrc/torch_bindings.cpp
+7
@@ -302,6 +302,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
302
"SymInt size_k) -> Tensor");
303
// conditionally compiled so impl registration is in source file
304
305
+ // CUTLASS nvfp4 block scaled GEMM
306
+ ops.def(
307
+ "cutlass_scaled_fp4_mm(Tensor! out, Tensor a, Tensor b,"
308
+ " Tensor block_scale_a, Tensor block_scale_b,"
309
+ " Tensor alpha) -> ()");
310
+ ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
311
+
312
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
313
// quantization, as well as bias
314
ops.def(
tests/kernels/test_nvfp4_scaled_mm.py
new
+150
@@ -0,0 +1,150 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+import torch
4
+
5
+from vllm import _custom_ops as ops
6
+from vllm.platforms import current_platform
7
+from vllm.scalar_type import scalar_types
8
+
9
+if not current_platform.has_device_capability(100):
10
+ pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
11
+ allow_module_level=True)
12
+
13
+DTYPES = [torch.float16, torch.bfloat16]
14
+# m, n, k
15
+SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
16
+PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]
17
+SHAPES.extend(PAD_SHAPES)
18
+
19
+SEEDS = [42]
20
+CUDA_DEVICES = ['cuda:0']
21
+
22
+FLOAT4_E2M1_MAX = scalar_types.float4_e2m1fn.max()
23
+FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
24
+
25
+kE2M1ToFloatArray = [
26
+ 0.,
27
+ 0.5,
28
+ 1.,
29
+ 1.5,
30
+ 2.,
31
+ 3.,
32
+ 4.,
33
+ 6.,
34
+]
35
+
36
+
37
+def e2m1_to_fp32(int4_value):
38
+ signBit = (int4_value & 0x8)
39
+ int4_absValue = int4_value & 0x7
40
+ float_result = kE2M1ToFloatArray[int4_absValue]
41
+ if (signBit):
42
+ float_result = -float_result
43
+ return float_result
44
+
45
+
46
+def break_fp4_bytes(a, dtype):
47
+ assert (a.dtype == torch.uint8)
48
+ m, n = a.shape
49
+ a = a.flatten()
50
+ # Get upper 4 bits
51
+ highHalfByte = (a & 0xF0) >> 4
52
+ # Get lower 4 bits
53
+ lowHalfByte = a & 0x0F
54
+ fH = torch.tensor([e2m1_to_fp32(x) for x in highHalfByte]).to(a.device)
55
+ fL = torch.tensor([e2m1_to_fp32(x) for x in lowHalfByte]).to(a.device)
56
+ # [0xAB, 0xCD] -> [0xB, 0xA, 0xD, 0xC]
57
+ out = torch.stack((fL, fH), dim=-1).reshape(m, n * 2)
58
+ return out
59
+
60
+
61
+def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
62
+ sf_m, sf_k = a_sf_swizzled.shape
63
+ m_tiles = (m + 128 - 1) // 128
64
+ f = block_size * 4
65
+ k_tiles = (k + f - 1) // f
66
+ tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
67
+ tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
68
+ out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
69
+ return out[0:m, 0:k]
70
+
71
+
72
+def dequantize_to_dtype(tensor_fp4,
73
+ tensor_sf,
74
+ global_scale,
75
+ dtype,
76
+ device,
77
+ block_size=16):
78
+ """Dequantize the fp4 tensor back to high precision."""
79
+ # Two fp4 values are packed into one uint8.
80
+ assert tensor_fp4.dtype == torch.uint8
81
+ m, packed_k = tensor_fp4.shape
82
+ k = packed_k * 2
83
+ tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
84
+ tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
85
+ tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
86
+ tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
87
+ tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
88
+
89
+ # scale the tensor
90
+ out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
91
+ return out
92
+
93
+
94
+def get_ref_results(a_fp4, b_fp4, a_sf, b_sf, a_global_scale, b_global_scale,
95
+ m, n, dtype, block_size, device):
96
+ _, m_k = a_fp4.shape
97
+ _, n_k = b_fp4.shape
98
+ assert (m_k == n_k)
99
+ a_in_dtype = dequantize_to_dtype(a_fp4,
100
+ a_sf,
101
+ a_global_scale,
102
+ dtype=dtype,
103
+ device=device,
104
+ block_size=block_size)
105
+ b_in_dtype = dequantize_to_dtype(b_fp4,
106
+ b_sf,
107
+ b_global_scale,
108
+ dtype=dtype,
109
+ device=device,
110
+ block_size=block_size)
111
+ return torch.matmul(a_in_dtype, b_in_dtype.t())
112
+
113
+
114
+@pytest.mark.parametrize("dtype", DTYPES)
115
+@pytest.mark.parametrize("shape", SHAPES)
116
+@pytest.mark.parametrize("seed", SEEDS)
117
+@pytest.mark.parametrize("device", CUDA_DEVICES)
118
+@torch.inference_mode()
119
+def test_nvfp4_gemm(
120
+ dtype: torch.dtype,
121
+ shape: tuple[int, int, int],
122
+ seed: int,
123
+ device: str,
124
+) -> None:
125
+ current_platform.seed_everything(seed)
126
+ m, n, packed_k = shape
127
+ k = packed_k * 2
128
+ block_size = 16
129
+ a_dtype = torch.randn((m, k), dtype=dtype, device=device)
130
+ b_dtype = torch.randn((n, k), dtype=dtype, device=device)
131
+
132
+ a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
133
+ torch.amax(a_dtype.flatten(), dim=-1)).to(torch.float32)
134
+ b_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
135
+ torch.amax(b_dtype.flatten(), dim=-1)).to(torch.float32)
136
+ alpha = 1. / (a_global_scale * b_global_scale)
137
+ a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale)
138
+ b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
139
+
140
+ expected_out = get_ref_results(a_fp4, b_fp4, a_scale_interleaved,
141
+ b_scale_interleaved, a_global_scale,
142
+ b_global_scale, m, n, dtype, block_size,
143
+ device)
144
+ out = ops.cutlass_scaled_fp4_mm(a_fp4, b_fp4, a_scale_interleaved,
145
+ b_scale_interleaved, alpha, dtype)
146
+
147
+ torch.testing.assert_close(out,
148
+ expected_out.to(dtype=dtype),
149
+ atol=1e-1,
150
+ rtol=1e-1)
vllm/_custom_ops.py
+12
@@ -433,6 +433,18 @@ if hasattr(torch.ops._C, "ggml_dequantize"):
433
434
435
# cutlass
436
+def cutlass_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
437
+ block_scale_a: torch.Tensor,
438
+ block_scale_b: torch.Tensor, alpha: torch.Tensor,
439
+ out_dtype: torch.dtype) -> torch.Tensor:
440
+ assert a.ndim == 2 and b.ndim == 2
441
+ m, n = a.shape[0], b.shape[0]
442
+ out = torch.empty((m, n), dtype=out_dtype, device=a.device)
443
+ torch.ops._C.cutlass_scaled_fp4_mm(out, a, b, block_scale_a, block_scale_b,
444
+ alpha)
445
+ return out
446
+
447
+
448
def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool:
449
return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
450