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 11, 03:45 UTC
[MoE Refactor][3/N] Deprecate cutlass block quant fp8 (b200) (#30990)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Dec 19, 2025 at 16:09 UTC
83a317f650f210b86572b13b8198b7d38aaacb7e
8 files changed
+3
-704
CMakeLists.txt
-18
@@ -799,24 +799,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
799
else()
800
cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
801
endif()
802
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
803
- set(SRCS "csrc/quantization/w8a8/cutlass/moe/blockwise_scaled_group_mm_sm100.cu")
804
- set_gencode_flags_for_srcs(
805
- SRCS "${SRCS}"
806
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
807
- list(APPEND VLLM_EXT_SRC "${SRCS}")
808
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
809
- message(STATUS "Building blockwise_scaled_group_mm_sm100 for archs: ${SCALED_MM_ARCHS}")
810
- else()
811
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
812
- message(STATUS "Not building blockwise_scaled_group_mm_sm100 kernels as CUDA Compiler version is "
813
- "not >= 12.8, we recommend upgrading to CUDA 12.8 or later "
814
- "if you intend on running FP8 quantized MoE models on Blackwell.")
815
- else()
816
- message(STATUS "Not building blockwise_scaled_group_mm_sm100 as no compatible archs found "
817
- "in CUDA target architectures")
818
- endif()
819
- endif()
802
803
#
804
# Machete kernels
csrc/quantization/w8a8/cutlass/moe/blockwise_scaled_group_mm_sm100.cu
deleted
-373
@@ -1,373 +0,0 @@
1
-#include "core/registration.h"
2
-
3
-#include <torch/all.h>
4
-#include <cutlass/arch/arch.h>
5
-
6
-#include <ATen/cuda/CUDAContext.h>
7
-#include <c10/cuda/CUDAGuard.h>
8
-#include <c10/cuda/CUDAStream.h>
9
-
10
-#include "cute/tensor.hpp"
11
-#include "cutlass/tensor_ref.h"
12
-#include "cutlass/epilogue/collective/default_epilogue.hpp"
13
-#include "cutlass/epilogue/thread/linear_combination.h"
14
-#include "cutlass/gemm/dispatch_policy.hpp"
15
-#include "cutlass/gemm/group_array_problem_shape.hpp"
16
-#include "cutlass/gemm/collective/collective_builder.hpp"
17
-#include "cutlass/epilogue/collective/collective_builder.hpp"
18
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
19
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
20
-
21
-#include "cutlass/util/command_line.h"
22
-#include "cutlass/util/distribution.h"
23
-#include "cutlass/util/host_tensor.h"
24
-#include "cutlass/util/packed_stride.hpp"
25
-#include "cutlass/util/tensor_view_io.h"
26
-#include "cutlass/util/reference/device/gemm.h"
27
-#include "cutlass/util/reference/device/tensor_compare.h"
28
-#include "cutlass/util/reference/host/tensor_fill.h"
29
-#include "cutlass/util/reference/host/gett.hpp"
30
-#include "cutlass/util/reference/host/tensor_norm.h"
31
-#include "cutlass/util/reference/host/tensor_compare.h"
32
-#include <cassert>
33
-
34
-using namespace cute;
35
-
36
-template <typename ElementAB, typename ElementC, typename ElementAccumulator,
37
- typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
38
-__global__ void get_ggemm_starts(
39
- int32_t* expert_offsets, ElementAB** a_offsets, ElementAB** b_offsets,
40
- ElementC** out_offsets, ElementAccumulator** a_scale_offsets,
41
- ElementAccumulator** b_scale_offsets, ElementAB* a_base_as_int,
42
- ElementAB* b_base_as_int, ElementC* out_base_as_int,
43
- ElementAccumulator* a_scale_base_as_int,
44
- ElementAccumulator* b_scale_base_as_int, LayoutSFA* layout_sfa_base_as_int,
45
- LayoutSFB* layout_sfb_base_as_int, int* problem_sizes) {
46
- int expert_id = threadIdx.x;
47
-
48
- if (expert_id >= gridDim.x * blockDim.x) {
49
- return;
50
- }
51
-
52
- int m = problem_sizes[expert_id * 3];
53
- int n = problem_sizes[expert_id * 3 + 1];
54
- int k = problem_sizes[expert_id * 3 + 2];
55
-
56
- int32_t expert_offset = expert_offsets[expert_id];
57
- int a_stride = expert_offset * k;
58
- int b_stride = expert_id * k * n;
59
- int a_scale_stride = expert_offset * k / 128;
60
- int b_scale_stride = expert_id * k * n / 128 / 128;
61
-
62
- a_offsets[expert_id] = a_base_as_int + a_stride;
63
- b_offsets[expert_id] = b_base_as_int + b_stride;
64
- out_offsets[expert_id] = out_base_as_int + expert_offset * n;
65
- a_scale_offsets[expert_id] = a_scale_base_as_int + a_scale_stride;
66
- b_scale_offsets[expert_id] = b_scale_base_as_int + b_scale_stride;
67
-
68
- LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
69
- LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
70
-
71
- *layout_sfa_ptr =
72
- ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(m, n, k, 1));
73
- *layout_sfb_ptr =
74
- ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(m, n, k, 1));
75
-}
76
-
77
-#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB, \
78
- ScaleConfig) \
79
- else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
80
- get_ggemm_starts<cutlass::float_e4m3_t, C_TYPE, float, LayoutSFA, \
81
- LayoutSFB, ScaleConfig><<<1, num_experts, 0, stream>>>( \
82
- static_cast<int32_t*>(expert_offsets.data_ptr()), \
83
- static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \
84
- static_cast<cutlass::float_e4m3_t**>(b_ptrs.data_ptr()), \
85
- static_cast<C_TYPE**>(out_ptrs.data_ptr()), \
86
- static_cast<float**>(a_scales_ptrs.data_ptr()), \
87
- static_cast<float**>(b_scales_ptrs.data_ptr()), \
88
- static_cast<cutlass::float_e4m3_t*>(a_tensors.data_ptr()), \
89
- static_cast<cutlass::float_e4m3_t*>(b_tensors.data_ptr()), \
90
- static_cast<C_TYPE*>(out_tensors.data_ptr()), \
91
- static_cast<float*>(a_scales.data_ptr()), \
92
- static_cast<float*>(b_scales.data_ptr()), \
93
- reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
94
- reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
95
- static_cast<int*>(problem_sizes.data_ptr())); \
96
- }
97
-
98
-template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
99
-void run_get_ggemm_starts(
100
- torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
101
- torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
102
- torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
103
- torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
104
- torch::Tensor out_tensors, torch::Tensor const& a_scales,
105
- torch::Tensor const& b_scales, torch::Tensor const& layout_sfa,
106
- torch::Tensor const& layout_sfb, torch::Tensor const& problem_sizes) {
107
- TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
108
- TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);
109
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
110
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
111
- TORCH_CHECK(out_tensors.size(1) % 128 == 0 or out_tensors.size(0) % 128 == 0);
112
- TORCH_CHECK(a_tensors.size(1) % 128 == 0 or a_tensors.size(0) % 128 == 0);
113
-
114
- int num_experts = (int)expert_offsets.size(0);
115
- auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
116
-
117
- if (false) {
118
- }
119
- __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t, LayoutSFA,
120
- LayoutSFB, ScaleConfig)
121
- __CALL_GET_STARTS_KERNEL(torch::kFloat16, cutlass::half_t, LayoutSFA,
122
- LayoutSFB, ScaleConfig)
123
- else {
124
- TORCH_CHECK(false, "Unsupported output tensor type");
125
- }
126
-}
127
-
128
-template <typename OutType, typename ScheduleConfig, typename LayoutD>
129
-void run_blockwise_scaled_group_mm(
130
- torch::Tensor& out_ptrs, const torch::Tensor& a_ptrs,
131
- const torch::Tensor& b_ptrs, const torch::Tensor& a_scales_ptrs,
132
- const torch::Tensor& b_scales_ptrs, const torch::Tensor& stride_a,
133
- const torch::Tensor& stride_b, const torch::Tensor& stride_c,
134
- const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
135
- const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets) {
136
- using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;
137
-
138
- // Types
139
- using ElementA = cutlass::float_e4m3_t;
140
- using ElementB = cutlass::float_e4m3_t;
141
- using ElementC = OutType;
142
- using ElementD = ElementC;
143
- using ElementAccumulator = float;
144
- using LayoutA = cutlass::layout::RowMajor;
145
- using LayoutB = cutlass::layout::ColumnMajor;
146
- using LayoutC = LayoutD;
147
-
148
- // Alignments
149
- static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
150
- static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
151
- static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
152
-
153
- using ArchTag = cutlass::arch::Sm100;
154
- using OperatorClass = cutlass::arch::OpClassTensorOp;
155
-
156
- using CollectiveEpilogue =
157
- typename cutlass::epilogue::collective::CollectiveBuilder<
158
- ArchTag, OperatorClass, typename ScheduleConfig::MmaTileShape,
159
- typename ScheduleConfig::ClusterShape,
160
- cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
161
- ElementAccumulator, void, LayoutC*, AlignmentC, ElementD, LayoutC*,
162
- AlignmentC, typename ScheduleConfig::EpilogueSchedule>::CollectiveOp;
163
-
164
- using CollectiveMainloop =
165
- typename cutlass::gemm::collective::CollectiveBuilder<
166
- ArchTag, OperatorClass, ElementA,
167
- cute::tuple<LayoutA*, typename ScheduleConfig::LayoutSFA*>,
168
- AlignmentA, ElementB,
169
- cute::tuple<LayoutB*, typename ScheduleConfig::LayoutSFB*>,
170
- AlignmentB, ElementAccumulator, typename ScheduleConfig::MmaTileShape,
171
- typename ScheduleConfig::ClusterShape,
172
- cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
173
- sizeof(typename CollectiveEpilogue::SharedStorage))>,
174
- typename ScheduleConfig::KernelSchedule>::CollectiveOp;
175
-
176
- using GemmKernel =
177
- cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
178
- CollectiveEpilogue, void>;
179
-
180
- using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
181
- using StrideA = typename Gemm::GemmKernel::InternalStrideA;
182
- using StrideB = typename Gemm::GemmKernel::InternalStrideB;
183
- using StrideC = typename Gemm::GemmKernel::InternalStrideC;
184
- using StrideD = typename Gemm::GemmKernel::InternalStrideD;
185
-
186
- using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
187
- int num_experts = (int)expert_offsets.size(0);
188
-
189
- Gemm gemm_op;
190
-
191
- // Mainloop Arguments
192
- typename GemmKernel::MainloopArguments mainloop_args{
193
- static_cast<const ElementA**>(a_ptrs.data_ptr()),
194
- static_cast<StrideA*>(stride_a.data_ptr()),
195
- static_cast<const ElementB**>(b_ptrs.data_ptr()),
196
- static_cast<StrideB*>(stride_b.data_ptr()),
197
- static_cast<const ElementAccumulator**>(a_scales_ptrs.data_ptr()),
198
- reinterpret_cast<typename ScheduleConfig::LayoutSFA*>(
199
- layout_sfa.data_ptr()),
200
- static_cast<const ElementAccumulator**>(b_scales_ptrs.data_ptr()),
201
- reinterpret_cast<typename ScheduleConfig::LayoutSFB*>(
202
- layout_sfb.data_ptr())};
203
-
204
- int device_id = a_ptrs.device().index();
205
- static const cutlass::KernelHardwareInfo hw_info{
206
- device_id, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
207
- device_id)};
208
-
209
- // Epilogue Arguments
210
- typename GemmKernel::EpilogueArguments epilogue_args{
211
- {}, // epilogue.thread
212
- nullptr,
213
- static_cast<StrideC*>(stride_c.data_ptr()),
214
- static_cast<ElementD**>(out_ptrs.data_ptr()),
215
- static_cast<StrideC*>(stride_c.data_ptr())};
216
-
217
- UnderlyingProblemShape* problem_sizes_as_shapes =
218
- static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
219
-
220
- // Gemm Arguments
221
- typename GemmKernel::Arguments args{
222
- cutlass::gemm::GemmUniversalMode::kGrouped,
223
- {num_experts, problem_sizes_as_shapes, nullptr},
224
- mainloop_args,
225
- epilogue_args,
226
- hw_info};
227
-
228
- at::cuda::CUDAGuard device_guard{(char)a_ptrs.device().index()};
229
- const cudaStream_t stream =
230
- at::cuda::getCurrentCUDAStream(a_ptrs.get_device());
231
-
232
- auto can_implement_status = gemm_op.can_implement(args);
233
- TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
234
- "Failed to implement GEMM");
235
-
236
- size_t workspace_size = gemm_op.get_workspace_size(args);
237
- auto const workspace_options =
238
- torch::TensorOptions().dtype(torch::kUInt8).device(a_ptrs.device());
239
- auto workspace = torch::empty(workspace_size, workspace_options);
240
-
241
- auto status = gemm_op.initialize(args, workspace.data_ptr(), stream);
242
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
243
-
244
- status = gemm_op.run(stream);
245
- TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
246
-}
247
-
248
-template <typename OutType>
249
-void blockwise_scaled_group_mm_dispatch_shape(
250
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
251
- const torch::Tensor& scales_a, const torch::Tensor& scales_b,
252
- const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets) {
253
- struct MmaConfig {
254
- using ElementA = cutlass::float_e4m3_t;
255
- using KernelSchedule =
256
- cutlass::gemm::KernelPtrArrayTmaWarpSpecializedBlockwise1SmSm100;
257
- using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
258
- using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<
259
- 1, 128, 128, cute::UMMA::Major::K, cute::UMMA::Major::K>;
260
- using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
261
- using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
262
- using LayoutC = cutlass::layout::RowMajor;
263
- using MmaTileShape = Shape<_128, _128, _128>;
264
- using ClusterShape = Shape<_1, _1, _1>;
265
- };
266
-
267
- int num_experts = (int)expert_offsets.size(0);
268
-
269
- auto a_ptrs = torch::empty(
270
- {num_experts},
271
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
272
- auto b_ptrs = torch::empty(
273
- {num_experts},
274
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
275
- auto out_ptrs = torch::empty(
276
- {num_experts},
277
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
278
- auto a_scales_ptrs = torch::empty(
279
- {num_experts},
280
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
281
- auto b_scales_ptrs = torch::empty(
282
- {num_experts},
283
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
284
-
285
- auto layout_sfa = torch::empty(
286
- {num_experts, 5},
287
- torch::TensorOptions().dtype(torch::kInt32).device(a.device()));
288
- auto layout_sfb = torch::empty(
289
- {num_experts, 5},
290
- torch::TensorOptions().dtype(torch::kInt32).device(a.device()));
291
-
292
- auto stride_a = torch::full(
293
- {num_experts}, a.size(1),
294
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
295
- auto stride_b = torch::full(
296
- {num_experts}, a.size(1),
297
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
298
- auto stride_c = torch::full(
299
- {num_experts}, output.size(1),
300
- torch::TensorOptions().dtype(torch::kInt64).device(a.device()));
301
-
302
- torch::TensorOptions options_int =
303
- torch::TensorOptions().dtype(torch::kInt64).device(a.device());
304
-
305
- run_get_ggemm_starts<typename MmaConfig::LayoutSFA,
306
- typename MmaConfig::LayoutSFB,
307
- typename MmaConfig::ScaleConfig>(
308
- expert_offsets, a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, a,
309
- b, output, scales_a, scales_b, layout_sfa, layout_sfb, problem_sizes);
310
-
311
- run_blockwise_scaled_group_mm<OutType, MmaConfig,
312
- typename MmaConfig::LayoutC>(
313
- out_ptrs, a_ptrs, b_ptrs, a_scales_ptrs, b_scales_ptrs, stride_a,
314
- stride_b, stride_c, layout_sfa, layout_sfb, problem_sizes,
315
- expert_offsets);
316
-}
317
-
318
-void cutlass_blockwise_scaled_grouped_mm(
319
- torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
320
- const torch::Tensor& scales_a, const torch::Tensor& scales_b,
321
- const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets) {
322
- TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
323
- TORCH_CHECK(problem_sizes.size(1) == 3,
324
- "problem_sizes must have shape (num_experts, 3)");
325
- TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
326
- "Number of experts in problem_sizes must match expert_offsets");
327
- TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
328
- "problem_sizes must be int32");
329
- TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn,
330
- "a must be kFloat8_e4m3fn");
331
- TORCH_CHECK(b.scalar_type() == torch::kFloat8_e4m3fn,
332
- "b must be kFloat8_e4m3fn");
333
- TORCH_CHECK(output.scalar_type() == torch::kBFloat16 ||
334
- output.scalar_type() == torch::kHalf,
335
- "output must be bfloat16 or half");
336
- TORCH_CHECK(scales_a.scalar_type() == torch::kFloat32,
337
- "scales_a must be float32");
338
- TORCH_CHECK(scales_b.scalar_type() == torch::kFloat32,
339
- "scales_b must be float32");
340
- TORCH_CHECK(expert_offsets.scalar_type() == torch::kInt32,
341
- "expert_offsets must be int32");
342
-
343
- TORCH_CHECK(output.dim() == 2, "output must be 2D tensor");
344
- TORCH_CHECK(a.dim() == 2, "a must be 2D tensor");
345
- TORCH_CHECK(b.dim() == 3, "b must be 3D tensor");
346
- TORCH_CHECK(scales_a.dim() == 2, "scales_a must be 2D tensor");
347
- TORCH_CHECK(scales_b.dim() == 3, "scales_b must be 3D tensor");
348
- TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
349
- TORCH_CHECK(problem_sizes.size(1) == 3,
350
- "problem_sizes must have shape (num_experts, 3)");
351
- TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
352
- "Number of experts in problem_sizes must match expert_offsets");
353
- TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
354
- "problem_sizes must be int32");
355
- TORCH_CHECK(expert_offsets.dim() == 1, "expert_offsets must be 1D tensor");
356
-
357
-#if defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100
358
- if (output.scalar_type() == torch::kBFloat16) {
359
- blockwise_scaled_group_mm_dispatch_shape<cutlass::bfloat16_t>(
360
- output, a, b, scales_a, scales_b, problem_sizes, expert_offsets);
361
- } else if (output.scalar_type() == torch::kFloat16) {
362
- blockwise_scaled_group_mm_dispatch_shape<cutlass::half_t>(
363
- output, a, b, scales_a, scales_b, problem_sizes, expert_offsets);
364
- } else {
365
- TORCH_CHECK(false, "Unsupported output tensor type");
366
- }
367
-#endif
368
-}
369
-
370
-TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
371
- m.impl("cutlass_blockwise_scaled_grouped_mm",
372
- &cutlass_blockwise_scaled_grouped_mm);
373
-}
csrc/torch_bindings.cpp
-7
@@ -416,13 +416,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
416
" Tensor alpha) -> ()");
417
ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
418
419
- // cutlass blockwise scaledgroup GEMM
420
- ops.def(
421
- "cutlass_blockwise_scaled_grouped_mm(Tensor! output, Tensor a, Tensor b, "
422
- "Tensor scales_a, Tensor scales_b, "
423
- "Tensor problem_sizes, Tensor expert_offsets) -> ()");
424
- // conditionally compiled so impl registration is in source file
425
-
419
// cutlass nvfp4 block scaled group GEMM
420
ops.def(
421
"cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
tests/kernels/moe/test_cutlass_grouped_gemm.py
deleted
-92
@@ -1,92 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# DeepGEMM Style Cutlass Grouped GEMM Test
5
-# See https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py
6
-
7
-import random
8
-
9
-import pytest
10
-import torch
11
-
12
-from tests.kernels.moe.utils import per_token_cast_to_fp8
13
-from tests.kernels.utils import baseline_scaled_mm
14
-from vllm import _custom_ops as ops
15
-from vllm.platforms import current_platform
16
-from vllm.utils.deep_gemm import per_block_cast_to_fp8
17
-from vllm.utils.math_utils import cdiv
18
-
19
-
20
-@pytest.mark.parametrize(
21
- "num_groups, expected_m_per_group, k, n",
22
- [
23
- (4, 8192, 7168, 4096),
24
- (4, 8192, 2048, 7168),
25
- (8, 4096, 7168, 4096),
26
- (8, 4096, 2048, 7168),
27
- (32, 1024, 7168, 4096),
28
- (32, 1024, 2048, 7168),
29
- ],
30
-)
31
-@pytest.mark.parametrize("out_dtype", [torch.float16])
32
-@pytest.mark.skipif(
33
- (lambda x: x is None or x.to_int() != 100)(
34
- current_platform.get_device_capability()
35
- ),
36
- reason="Block Scaled Grouped GEMM is only supported on SM100.",
37
-)
38
-def test_cutlass_grouped_gemm(
39
- num_groups: int,
40
- expected_m_per_group: int,
41
- k: int,
42
- n: int,
43
- out_dtype: torch.dtype,
44
-):
45
- device = "cuda"
46
- alignment = 128
47
- group_ms = [
48
- int(expected_m_per_group * random.uniform(0.7, 1.3)) for _ in range(num_groups)
49
- ]
50
- m = sum([cdiv(m, alignment) * alignment for m in group_ms])
51
-
52
- x = torch.randn((m, k), device=device, dtype=out_dtype)
53
- y = torch.randn((num_groups, n, k), device=device, dtype=out_dtype)
54
- out = torch.empty((m, n), device=device, dtype=out_dtype)
55
- ref_out = torch.randn((m, n), device=device, dtype=out_dtype)
56
-
57
- ep_offset = [0] + [sum(group_ms[:i]) for i in range(1, num_groups)] + [m]
58
- pb_size = []
59
- for i in range(num_groups):
60
- pb_size.append([ep_offset[i + 1] - ep_offset[i], n, k])
61
- problem_sizes = torch.tensor(pb_size, device=device, dtype=torch.int32)
62
- expert_offsets = torch.tensor(ep_offset, device=device, dtype=torch.int32)
63
-
64
- x_fp8 = per_token_cast_to_fp8(x)
65
- y_fp8 = (
66
- torch.empty_like(y, dtype=torch.float8_e4m3fn),
67
- torch.empty(
68
- (num_groups, cdiv(n, 128), k // 128), device=device, dtype=torch.float
69
- ),
70
- )
71
- for i in range(num_groups):
72
- y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], [128, 128])
73
-
74
- for i in range(num_groups):
75
- a = x_fp8[0][ep_offset[i] : ep_offset[i + 1]]
76
- a_scale = x_fp8[1][ep_offset[i] : ep_offset[i + 1]]
77
- b = y_fp8[0][i].t()
78
- b_scale = y_fp8[1][i].t()
79
- baseline = baseline_scaled_mm(a, b, a_scale, b_scale, out_dtype)
80
- ref_out[ep_offset[i] : ep_offset[i + 1]] = baseline
81
-
82
- ops.cutlass_blockwise_scaled_grouped_mm(
83
- out,
84
- x_fp8[0],
85
- y_fp8[0],
86
- x_fp8[1],
87
- y_fp8[1],
88
- problem_sizes,
89
- expert_offsets[:-1],
90
- )
91
-
92
- torch.testing.assert_close(ref_out, out, atol=5e-1, rtol=1e-3)
vllm/_custom_ops.py
-14
@@ -788,20 +788,6 @@ def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
788
return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
789
790
791
-def cutlass_blockwise_scaled_grouped_mm(
792
- output: torch.Tensor,
793
- a: torch.Tensor,
794
- b: torch.Tensor,
795
- scales_a: torch.Tensor,
796
- scales_b: torch.Tensor,
797
- problem_sizes: torch.Tensor,
798
- expert_offsets: torch.Tensor,
799
-):
800
- torch.ops._C.cutlass_blockwise_scaled_grouped_mm(
801
- output, a, b, scales_a, scales_b, problem_sizes, expert_offsets
802
- )
803
-
804
-
791
def cutlass_scaled_fp4_mm(
792
a: torch.Tensor,
793
b: torch.Tensor,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+1
-157
@@ -21,7 +21,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
21
TopKWeightAndReduceDelegate,
22
TopKWeightAndReduceNoOP,
23
)
24
-from vllm.model_executor.layers.fused_moe.utils import _fp8_quantize, _resize_cache
24
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
25
from vllm.scalar_type import scalar_types
26
27
logger = init_logger(__name__)
@@ -896,162 +896,6 @@ def cutlass_moe_fp4(
896
)
897
898
899
-def _valid_cutlass_block_scaled_grouped_gemm(
900
- w1: torch.Tensor,
901
- w2: torch.Tensor,
902
- inplace: bool,
903
- activation: str,
904
- apply_router_weight_on_input: bool,
905
- expert_map: torch.Tensor | None,
906
-) -> bool:
907
- def _valid_cutlass_block_scaled_grouped_gemm_shape(N: int, K: int):
908
- return N % 128 == 0 and K % 128 == 0
909
-
910
- _, K, N = w2.size()
911
- if not _valid_cutlass_block_scaled_grouped_gemm_shape(N, K):
912
- logger.debug_once(
913
- "CutlassBlockScaledGroupedGemm disabled: unaligned problem size. "
914
- "N: %s, K: %s",
915
- N,
916
- K,
917
- )
918
- return False
919
-
920
- if w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn:
921
- logger.debug_once(
922
- "CutlassBlockScaledGroupedGemm disabled: invalid weight dtype(s). "
923
- "w1.dtype: %s, w2.dtype: %s",
924
- w1.dtype,
925
- w2.dtype,
926
- )
927
- return False
928
-
929
- if expert_map is not None:
930
- logger.debug_once(
931
- "CutlassBlockScaledGroupedGemm disabled: expert_parallel is not supported."
932
- )
933
- return False
934
-
935
- if activation != "silu":
936
- logger.debug_once(
937
- "CutlassBlockScaledGroupedGemm disabled: only activation silu is supported."
938
- )
939
- return False
940
-
941
- if apply_router_weight_on_input:
942
- logger.debug_once(
943
- "CutlassBlockScaledGroupedGemm disabled:"
944
- " apply_router_weight_on_input is not supported."
945
- )
946
- return False
947
-
948
- if inplace:
949
- logger.debug_once(
950
- "CutlassBlockScaledGroupedGemm disabled: inplace is not supported."
951
- )
952
- return False
953
-
954
- return True
955
-
956
-
957
-# TODO(bnell): would be nice combine/integrate with regular cutlass_fp8.
958
-def run_cutlass_block_scaled_fused_experts(
959
- a: torch.Tensor,
960
- w1: torch.Tensor,
961
- w2: torch.Tensor,
962
- w1_scale: torch.Tensor,
963
- w2_scale: torch.Tensor,
964
- topk_weights: torch.Tensor,
965
- topk_ids: torch.Tensor,
966
-) -> torch.Tensor:
967
- w1_q = w1.transpose(1, 2)
968
- w2_q = w2.transpose(1, 2)
969
- w1_scale = w1_scale.transpose(1, 2)
970
- w2_scale = w2_scale.transpose(1, 2)
971
-
972
- assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
973
- assert a.shape[0] == topk_ids.shape[0], (
974
- "a and topk_ids must have the same batch size"
975
- )
976
- assert w1_q.dtype == torch.float8_e4m3fn, "w1_q must be float8_e4m3fn"
977
- assert w2_q.dtype == torch.float8_e4m3fn, "w2_q must be float8_e4m3fn"
978
- assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
979
- assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
980
- assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
981
- assert w1_q.shape[0] == w1_scale.shape[0], "w1_scale expert number mismatch"
982
- assert w1_q.shape[0] == w2_scale.shape[0], "w2_scale expert number mismatch"
983
- assert a.dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
984
-
985
- out_dtype = a.dtype
986
- num_experts = w1_q.size(0)
987
- m = a.size(0)
988
- k = w1_q.size(1)
989
- n = w2_q.size(1)
990
-
991
- topk = topk_ids.size(1)
992
-
993
- a_q, a1_scale = _fp8_quantize(
994
- a, A_scale=None, per_act_token=False, block_shape=[128, 128]
995
- )
996
- device = a_q.device
997
-
998
- expert_offsets = torch.empty((num_experts + 1,), dtype=torch.int32, device=device)
999
- problem_sizes1 = torch.empty((num_experts, 3), dtype=torch.int32, device=device)
1000
- problem_sizes2 = torch.empty((num_experts, 3), dtype=torch.int32, device=device)
1001
-
1002
- a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
1003
- c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
1004
-
1005
- ops.get_cutlass_moe_mm_data(
1006
- topk_ids,
1007
- expert_offsets,
1008
- problem_sizes1,
1009
- problem_sizes2,
1010
- a_map,
1011
- c_map,
1012
- num_experts,
1013
- n,
1014
- k,
1015
- )
1016
-
1017
- rep_a_q = a_q.view(dtype=torch.uint8)[a_map].view(dtype=a_q.dtype)
1018
- rep_a1_scales = a1_scale[a_map]
1019
-
1020
- c1 = torch.empty((m * topk, n * 2), dtype=out_dtype, device=device)
1021
- c2 = torch.empty((m * topk, k), dtype=out_dtype, device=device)
1022
-
1023
- ops.cutlass_blockwise_scaled_grouped_mm(
1024
- c1,
1025
- rep_a_q,
1026
- w1_q,
1027
- rep_a1_scales,
1028
- w1_scale,
1029
- problem_sizes1,
1030
- expert_offsets[:-1],
1031
- )
1032
-
1033
- intermediate = torch.empty((m * topk, n), dtype=out_dtype, device=device)
1034
- torch.ops._C.silu_and_mul(intermediate, c1)
1035
-
1036
- intermediate_q, a2_scale = _fp8_quantize(
1037
- intermediate, A_scale=None, per_act_token=False, block_shape=[128, 128]
1038
- )
1039
-
1040
- ops.cutlass_blockwise_scaled_grouped_mm(
1041
- c2,
1042
- intermediate_q,
1043
- w2_q,
1044
- a2_scale,
1045
- w2_scale,
1046
- problem_sizes2,
1047
- expert_offsets[:-1],
1048
- )
1049
-
1050
- return (
1051
- c2[c_map].view(m, topk, k) * topk_weights.view(m, topk, 1).to(out_dtype)
1052
- ).sum(dim=1)
1053
-
1054
-
899
# W4A8
900
def run_cutlass_moe_w4a8_fp8(
901
output: torch.Tensor,
vllm/model_executor/layers/fused_moe/fused_moe.py
-23
@@ -25,10 +25,6 @@ from vllm.model_executor.layers.fused_moe.config import (
25
FusedMoEQuantConfig,
26
_get_config_dtype_str,
27
)
28
-from vllm.model_executor.layers.fused_moe.cutlass_moe import (
29
- _valid_cutlass_block_scaled_grouped_gemm,
30
- run_cutlass_block_scaled_fused_experts,
31
-)
28
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
29
_valid_deep_gemm,
30
deep_gemm_moe_fp8,
@@ -1678,11 +1674,9 @@ def fused_experts(
1674
expert_map: torch.Tensor | None = None,
1675
quant_config: FusedMoEQuantConfig | None = None,
1676
allow_deep_gemm: bool = False,
1681
- allow_cutlass_block_scaled_grouped_gemm: bool = False,
1677
) -> torch.Tensor:
1678
if quant_config is None:
1679
quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
1685
- use_fp8_w8a8 = quant_config.use_fp8_w8a8
1680
1681
# For now, disable DeepGemm for small N (<= 512) until better
1682
# permute/unpermute ops are available.
@@ -1712,23 +1706,6 @@ def fused_experts(
1706
a2_scale=quant_config.a2_scale,
1707
apply_router_weight_on_input=apply_router_weight_on_input,
1708
)
1715
- elif (
1716
- allow_cutlass_block_scaled_grouped_gemm
1717
- and use_fp8_w8a8
1718
- and _valid_cutlass_block_scaled_grouped_gemm(
1719
- w1, w2, inplace, activation, apply_router_weight_on_input, expert_map
1720
- )
1721
- ):
1722
- assert quant_config is not None
1723
- return run_cutlass_block_scaled_fused_experts(
1724
- a=hidden_states,
1725
- w1=w1,
1726
- w2=w2,
1727
- w1_scale=quant_config.w1_scale,
1728
- w2_scale=quant_config.w2_scale,
1729
- topk_weights=topk_weights,
1730
- topk_ids=topk_ids,
1731
- )
1709
else:
1710
return dispatch_fused_experts_func(inplace)(
1711
hidden_states=hidden_states,
vllm/model_executor/layers/quantization/fp8.py
+2
-20
@@ -118,9 +118,8 @@ class Fp8MoeBackend(Enum):
118
FLASHINFER_TRTLLM = 1
119
FLASHINFER_CUTLASS = 2
120
DEEPGEMM = 3
121
- CUTLASS_BLOCK_SCALED_GROUPED_GEMM = 4
122
- MARLIN = 5
123
- TRITON = 6
121
+ MARLIN = 4
122
+ TRITON = 5
123
124
125
def get_fp8_moe_backend(
@@ -191,17 +190,6 @@ def get_fp8_moe_backend(
190
logger.info_once("Using DeepGEMM backend for FP8 MoE", scope="local")
191
return Fp8MoeBackend.DEEPGEMM
192
194
- # CUTLASS BlockScaled GroupedGemm on SM100 with block-quantized weights
195
- if (
196
- current_platform.is_cuda()
197
- and current_platform.is_device_capability_family(100)
198
- and block_quant
199
- ):
200
- logger.info_once(
201
- "Using Cutlass BlockScaled GroupedGemm backend for FP8 MoE", scope="local"
202
- )
203
- return Fp8MoeBackend.CUTLASS_BLOCK_SCALED_GROUPED_GEMM
204
-
193
# default to Triton
194
logger.info_once("Using Triton backend for FP8 MoE")
195
return Fp8MoeBackend.TRITON
@@ -752,9 +740,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
740
)
741
742
self.allow_deep_gemm = self.fp8_backend == Fp8MoeBackend.DEEPGEMM
755
- self.allow_cutlass_block_scaled_grouped_gemm = (
756
- self.fp8_backend == Fp8MoeBackend.CUTLASS_BLOCK_SCALED_GROUPED_GEMM
757
- )
743
744
def create_weights(
745
self,
@@ -1316,9 +1301,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1301
expert_map=layer.expert_map,
1302
quant_config=self.moe_quant_config,
1303
allow_deep_gemm=self.allow_deep_gemm,
1319
- allow_cutlass_block_scaled_grouped_gemm=(
1320
- self.allow_cutlass_block_scaled_grouped_gemm
1321
- ),
1304
)
1305
1306
if layer.zero_expert_num != 0 and layer.zero_expert_type is not None: