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 14, 22:35 UTC
Support CUTLASS NVFP4 (w4a4) for Blackwell Geforce GPUs (SM120) (#21309)
Signed-off-by: LopezCastroRoberto <roberto.lopez.castro@udc.es>
Roberto L. Castro committed
Aug 3, 2025 at 09:54 UTC
789562c28c143201a1d2ca35f7adcdf54ef832e5
6 files changed
+329
-13
CMakeLists.txt
+20
-1
@@ -529,6 +529,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
529
endif()
530
endif()
531
532
+ # The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
533
+ # CUDA 12.8 or later
534
+ cuda_archs_loose_intersection(FP4_ARCHS "12.0;12.0a" "${CUDA_ARCHS}")
535
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
536
+ set(SRCS
537
+ "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
538
+ "csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu")
539
+ set_gencode_flags_for_srcs(
540
+ SRCS "${SRCS}"
541
+ CUDA_ARCHS "${FP4_ARCHS}")
542
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
543
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
544
+ message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
545
+ else()
546
+ message(STATUS "Not building NVFP4 as no compatible archs were found.")
547
+ # clear FP4_ARCHS
548
+ set(FP4_ARCHS)
549
+ endif()
550
+
551
# FP4 Archs and flags
552
cuda_archs_loose_intersection(FP4_ARCHS "10.0a" "${CUDA_ARCHS}")
553
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
@@ -541,7 +560,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
560
SRCS "${SRCS}"
561
CUDA_ARCHS "${FP4_ARCHS}")
562
list(APPEND VLLM_EXT_SRC "${SRCS}")
544
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4=1")
563
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
564
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
565
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
566
else()
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu
+3
-3
@@ -335,7 +335,7 @@ void run_fp4_blockwise_scaled_group_mm(
335
TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
336
}
337
338
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
338
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
339
constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
340
constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
341
#endif
@@ -356,7 +356,7 @@ void cutlass_fp4_group_mm(
356
const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
357
const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
358
const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
359
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
359
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
360
// Input validation
361
CHECK_INPUT(a, FLOAT4_E2M1X2, "a");
362
CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
@@ -398,7 +398,7 @@ void cutlass_fp4_group_mm(
398
TORCH_CHECK_NOT_IMPLEMENTED(
399
false,
400
"No compiled cutlass_fp4_group_mm kernel, vLLM must "
401
- "be compiled with ENABLE_NVFP4 for SM100+ and CUDA "
401
+ "be compiled with ENABLE_NVFP4_SM100 for SM100+ and CUDA "
402
"12.8 or above.");
403
#endif
404
}
csrc/quantization/fp4/nvfp4_quant_entry.cu
+8
-6
@@ -16,14 +16,15 @@
16
17
#include <torch/all.h>
18
19
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
20
-void scaled_fp4_quant_sm100a(torch::Tensor const& output,
19
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
20
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
21
+void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
22
torch::Tensor const& input,
23
torch::Tensor const& output_sf,
24
torch::Tensor const& input_sf);
25
#endif
26
26
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
27
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
28
void scaled_fp4_experts_quant_sm100a(
29
torch::Tensor& output, torch::Tensor& output_scale,
30
torch::Tensor const& input, torch::Tensor const& input_global_scale,
@@ -33,8 +34,9 @@ void scaled_fp4_experts_quant_sm100a(
34
35
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
36
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
36
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
37
- return scaled_fp4_quant_sm100a(output, input, output_sf, input_sf);
37
+#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
38
+ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
39
+ return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf);
40
#endif
41
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
42
}
@@ -44,7 +46,7 @@ void scaled_fp4_experts_quant(
46
torch::Tensor const& input, torch::Tensor const& input_global_scale,
47
torch::Tensor const& input_offset_by_experts,
48
torch::Tensor const& output_scale_offset_by_experts) {
47
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
49
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
50
return scaled_fp4_experts_quant_sm100a(
51
output, output_scale, input, input_global_scale, input_offset_by_experts,
52
output_scale_offset_by_experts);
csrc/quantization/fp4/nvfp4_quant_kernels.cu
+1
-1
@@ -332,7 +332,7 @@ template void invokeFP4Quantization(int m, int n, __nv_bfloat16 const* input,
332
int multiProcessorCount,
333
cudaStream_t stream);
334
335
-void scaled_fp4_quant_sm100a(torch::Tensor const& output,
335
+void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
336
torch::Tensor const& input,
337
torch::Tensor const& output_sf,
338
torch::Tensor const& input_sf) {
csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu
+12
-2
@@ -16,7 +16,7 @@
16
17
#include <torch/all.h>
18
19
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
19
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
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,
@@ -24,12 +24,22 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
24
torch::Tensor const& alpha);
25
#endif
26
27
+#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
28
+void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
29
+ torch::Tensor const& B,
30
+ torch::Tensor const& A_sf,
31
+ torch::Tensor const& B_sf,
32
+ torch::Tensor const& alpha);
33
+#endif
34
+
35
void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
36
torch::Tensor const& B, torch::Tensor const& A_sf,
37
torch::Tensor const& B_sf,
38
torch::Tensor const& alpha) {
31
-#if defined ENABLE_NVFP4 && ENABLE_NVFP4
39
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
40
return cutlass_scaled_fp4_mm_sm100a(D, A, B, A_sf, B_sf, alpha);
41
+#elif defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
42
+ return cutlass_scaled_fp4_mm_sm120a(D, A, B, A_sf, B_sf, alpha);
43
#endif
44
TORCH_CHECK_NOT_IMPLEMENTED(false,
45
"No compiled nvfp4 mm kernel, vLLM should "
csrc/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu
new
+285
@@ -0,0 +1,285 @@
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
+#include "core/math.hpp"
34
+
35
+using namespace cute;
36
+
37
+#define CHECK_TYPE(x, st, m) \
38
+ TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
39
+#define CHECK_TH_CUDA(x, m) \
40
+ TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
41
+#define CHECK_CONTIGUOUS(x, m) \
42
+ TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
43
+#define CHECK_INPUT(x, st, m) \
44
+ CHECK_TH_CUDA(x, m); \
45
+ CHECK_CONTIGUOUS(x, m); \
46
+ CHECK_TYPE(x, st, m)
47
+
48
+constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
49
+constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
50
+
51
+struct sm120_fp4_config_M256 {
52
+ using ClusterShape = Shape<_1, _1, _1>;
53
+ using MmaTileShape = Shape<_128, _128, _128>;
54
+ using PerSmTileShape_MNK = Shape<_128, _128, _128>;
55
+};
56
+
57
+struct sm120_fp4_config_default {
58
+ using ClusterShape = Shape<_1, _1, _1>;
59
+ using MmaTileShape = Shape<_256, _128, _128>;
60
+ using PerSmTileShape_MNK = Shape<_256, _128, _128>;
61
+};
62
+
63
+template <typename Config, typename OutType>
64
+struct Fp4GemmSm120 {
65
+ using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
66
+ using LayoutATag = cutlass::layout::RowMajor;
67
+ static constexpr int AlignmentA = 32;
68
+
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
+ using ElementD = OutType;
74
+ using ElementC = OutType;
75
+ using LayoutCTag = cutlass::layout::RowMajor;
76
+ using LayoutDTag = cutlass::layout::RowMajor;
77
+ static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
78
+ static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
79
+
80
+ using ElementAccumulator = float;
81
+ using ArchTag = cutlass::arch::Sm120;
82
+ using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
83
+
84
+ using MmaTileShape = typename Config::MmaTileShape;
85
+ using ClusterShape = typename Config::ClusterShape;
86
+ using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK;
87
+
88
+ using CollectiveEpilogue =
89
+ typename cutlass::epilogue::collective::CollectiveBuilder<
90
+ ArchTag, OperatorClass, PerSmTileShape_MNK, ClusterShape,
91
+ cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
92
+ ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
93
+ LayoutDTag, AlignmentD,
94
+ cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp;
95
+
96
+ using CollectiveMainloop =
97
+ typename cutlass::gemm::collective::CollectiveBuilder<
98
+ ArchTag, OperatorClass, ElementA, LayoutATag, AlignmentA, ElementB,
99
+ LayoutBTag, AlignmentB, ElementAccumulator, MmaTileShape,
100
+ ClusterShape,
101
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
102
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
103
+ cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
104
+
105
+ using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
106
+ Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
107
+
108
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
109
+};
110
+
111
+template <typename Gemm>
112
+typename Gemm::Arguments args_from_options(at::Tensor& D, at::Tensor const& A,
113
+ at::Tensor const& B,
114
+ at::Tensor const& A_sf,
115
+ at::Tensor const& B_sf,
116
+ torch::Tensor const& alpha, int M,
117
+ int N, int K) {
118
+ using ElementA = typename Gemm::ElementA;
119
+ using ElementB = typename Gemm::ElementB;
120
+ using ElementD = typename Gemm::ElementD;
121
+ using ElementSFA = cutlass::float_ue4m3_t;
122
+ using ElementSFB = cutlass::float_ue4m3_t;
123
+ using ElementCompute = float;
124
+
125
+ using StrideA = typename Gemm::GemmKernel::StrideA;
126
+ using StrideB = typename Gemm::GemmKernel::StrideB;
127
+ using StrideC = typename Gemm::GemmKernel::StrideC;
128
+ using StrideD = typename Gemm::GemmKernel::StrideD;
129
+
130
+ using Sm1xxBlkScaledConfig =
131
+ typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
132
+
133
+ auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
134
+ auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
135
+ auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
136
+
137
+ auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(
138
+ cute::make_shape(M, N, K, 1));
139
+ auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
140
+ cute::make_shape(M, N, K, 1));
141
+
142
+ typename Gemm::Arguments arguments{
143
+ cutlass::gemm::GemmUniversalMode::kGemm,
144
+ {M, N, K, 1},
145
+ {static_cast<ElementA const*>(A.data_ptr()), stride_A,
146
+ static_cast<ElementB const*>(B.data_ptr()), stride_B,
147
+ static_cast<ElementSFA const*>(A_sf.data_ptr()), layout_SFA,
148
+ static_cast<ElementSFB const*>(B_sf.data_ptr()), layout_SFB},
149
+ {{},
150
+ static_cast<ElementD const*>(D.data_ptr()),
151
+ stride_D,
152
+ static_cast<ElementD*>(D.data_ptr()),
153
+ stride_D}};
154
+ auto& fusion_args = arguments.epilogue.thread;
155
+ fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr());
156
+
157
+ return arguments;
158
+}
159
+
160
+template <typename Gemm>
161
+void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
162
+ at::Tensor const& A_sf, at::Tensor const& B_sf,
163
+ torch::Tensor const& alpha, int M, int N, int K,
164
+ cudaStream_t stream) {
165
+ Gemm gemm;
166
+
167
+ auto arguments = args_from_options<Gemm>(D, A, B, A_sf, B_sf, alpha, M, N, K);
168
+
169
+ size_t workspace_size = Gemm::get_workspace_size(arguments);
170
+ auto const workspace_options =
171
+ torch::TensorOptions().dtype(torch::kUInt8).device(A.device());
172
+ auto workspace = torch::empty(workspace_size, workspace_options);
173
+
174
+ CUTLASS_CHECK(gemm.can_implement(arguments));
175
+
176
+ CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
177
+
178
+ CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
179
+}
180
+
181
+void cutlass_fp4_bf16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
182
+ torch::Tensor const& B,
183
+ torch::Tensor const& A_sf,
184
+ torch::Tensor const& B_sf,
185
+ torch::Tensor const& alpha, int m, int n,
186
+ int k, cudaStream_t stream) {
187
+ uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
188
+ if (mp2 <= 256) {
189
+ runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::bfloat16_t>::Gemm>(
190
+ D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
191
+ } else {
192
+ runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::bfloat16_t>::Gemm>(
193
+ D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
194
+ }
195
+}
196
+
197
+void cutlass_fp4_f16_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
198
+ torch::Tensor const& B,
199
+ torch::Tensor const& A_sf,
200
+ torch::Tensor const& B_sf,
201
+ torch::Tensor const& alpha, int m, int n,
202
+ int k, cudaStream_t stream) {
203
+ uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
204
+ if (mp2 <= 256) {
205
+ runGemm<Fp4GemmSm120<sm120_fp4_config_M256, cutlass::half_t>::Gemm>(
206
+ D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
207
+ } else {
208
+ runGemm<Fp4GemmSm120<sm120_fp4_config_default, cutlass::half_t>::Gemm>(
209
+ D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
210
+ }
211
+}
212
+
213
+void cutlass_scaled_fp4_mm_sm120a(torch::Tensor& D, torch::Tensor const& A,
214
+ torch::Tensor const& B,
215
+ torch::Tensor const& A_sf,
216
+ torch::Tensor const& B_sf,
217
+ torch::Tensor const& alpha) {
218
+#if defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
219
+ CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
220
+ CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
221
+
222
+ CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
223
+ CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
224
+
225
+ CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
226
+
227
+ TORCH_CHECK(A.dim() == 2, "a must be a matrix");
228
+ TORCH_CHECK(B.dim() == 2, "b must be a matrix");
229
+ TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
230
+ "a and b shapes cannot be multiplied (", A.sizes()[0], "x",
231
+ A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
232
+
233
+ auto const m = A.sizes()[0];
234
+ auto const n = B.sizes()[0];
235
+ auto const k = A.sizes()[1] * 2;
236
+
237
+ constexpr int alignment = 32;
238
+ TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
239
+ ", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
240
+ "), k: ", k, ".");
241
+ TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
242
+ ", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
243
+
244
+ auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
245
+ int rounded_m = round_up(m, 128);
246
+ int rounded_n = round_up(n, 128);
247
+ // Since k is divisible by 32 (alignment), k / 16 is guaranteed to be an
248
+ // integer.
249
+ int rounded_k = round_up(k / 16, 4);
250
+
251
+ TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
252
+ TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
253
+ TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
254
+ "scale_a and scale_b shapes cannot be multiplied (",
255
+ A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
256
+ "x", B_sf.sizes()[1], ")");
257
+ TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
258
+ "scale_a must be padded and swizzled to a shape (", rounded_m,
259
+ "x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
260
+ A_sf.sizes()[1], ")");
261
+ TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
262
+ "scale_b must be padded and swizzled to a shape (", rounded_n,
263
+ "x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
264
+ B_sf.sizes()[1], ")");
265
+
266
+ auto out_dtype = D.dtype();
267
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
268
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
269
+
270
+ if (out_dtype == at::ScalarType::BFloat16) {
271
+ return cutlass_fp4_bf16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
272
+ stream);
273
+ } else if (out_dtype == at::ScalarType::Half) {
274
+ return cutlass_fp4_f16_gemm_dispatch(D, A, B, A_sf, B_sf, alpha, m, n, k,
275
+ stream);
276
+ } else {
277
+ TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm sm120 (",
278
+ out_dtype, ")");
279
+ }
280
+#else
281
+ TORCH_CHECK(false,
282
+ "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to "
283
+ "a CUTLASS 3.8 source directory to enable support.");
284
+#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
285
+}
\ No newline at end of file