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 7, 11:30 UTC
[Kernel] Integrate SM100 MXFP8 blockscaled grouped MM and quant kernels (#34448)
Signed-off-by: EdalatiAli <aliedalati@cohere.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
EdalatiAli committed
Mar 2, 2026 at 02:31 UTC
cb21972a976b01e5ee1d899ab744207826435684
10 files changed
+1331
CMakeLists.txt
+27
@@ -771,6 +771,33 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
771
endif()
772
endif()
773
774
+ # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
775
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
776
+ cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
777
+ else()
778
+ cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
779
+ endif()
780
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
781
+ set(SRCS
782
+ "csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
783
+ "csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu")
784
+ set_gencode_flags_for_srcs(
785
+ SRCS "${SRCS}"
786
+ CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
787
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
788
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
789
+ message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
790
+ else()
791
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
792
+ AND ES_MXFP8_GROUPED_MM_ARCHS)
793
+ message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
794
+ "not >= 12.8.")
795
+ else()
796
+ message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
797
+ "in CUDA target architectures.")
798
+ endif()
799
+ endif()
800
+
801
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
802
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
803
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu
new
+60
@@ -0,0 +1,60 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
5
+
6
+#include <torch/all.h>
7
+
8
+#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
9
+
10
+void cutlass_mxfp8_grouped_mm(const torch::Tensor& a, const torch::Tensor& b,
11
+ const torch::Tensor& sfa,
12
+ const torch::Tensor& sfb, torch::Tensor& d,
13
+ const torch::Tensor& problem_sizes,
14
+ const torch::Tensor& expert_offsets,
15
+ const torch::Tensor& blockscale_offsets) {
16
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
17
+ TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
18
+ TORCH_CHECK(problem_sizes.size(1) == 3,
19
+ "problem_sizes must have shape (num_experts, 3)");
20
+ TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
21
+ "Number of experts in problem_sizes must match expert_offsets");
22
+ TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
23
+ "problem_sizes must be int32");
24
+ TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
25
+ "expert_offsets must be int32");
26
+ TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
27
+ "blockscale_offsets must be int32");
28
+ TORCH_CHECK(a.dim() == 2, "a must be a 2D tensor of shape (num_tokens, k)");
29
+ TORCH_CHECK(b.dim() == 3,
30
+ "b must be a 3D tensor of shape (num_experts, k, n)");
31
+ TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
32
+ "k should align 128");
33
+ TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
34
+ TORCH_CHECK(a.strides()[1] == 1, "a must be row major");
35
+ TORCH_CHECK(b.strides()[1] == 1, "b must be column major");
36
+
37
+ auto stream = at::cuda::getCurrentCUDAStream();
38
+ if (d.dtype() == torch::kBFloat16) {
39
+ expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
40
+ cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
41
+ blockscale_offsets, stream);
42
+ } else if (d.dtype() == torch::kFloat16) {
43
+ expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
44
+ cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
45
+ blockscale_offsets, stream);
46
+ } else {
47
+ TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
48
+ }
49
+#else
50
+ TORCH_CHECK(false,
51
+ "No implemented cutlass_mxfp8_grouped_mm for "
52
+ "current device");
53
+#endif
54
+}
55
+
56
+#include "core/registration.h"
57
+
58
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
59
+ m.impl("cutlass_mxfp8_grouped_mm", cutlass_mxfp8_grouped_mm);
60
+}
\ No newline at end of file
csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh
new
+141
@@ -0,0 +1,141 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_functor.cuh
5
+
6
+#pragma once
7
+#include <cuda.h>
8
+
9
+#include "cute/tensor.hpp"
10
+#include "cutlass/util/packed_stride.hpp"
11
+#include "cutlass_mxfp8_grouped_mm_traits.cuh"
12
+
13
+namespace expert_specialization {
14
+
15
+using namespace cute;
16
+
17
+template <typename GemmTraits>
18
+struct CutlassMxfp8GroupedMmOffsetFunctor {
19
+ using Gemm = typename GemmTraits::Gemm;
20
+ using ElementA = typename Gemm::ElementA;
21
+ using ElementB = typename Gemm::ElementB;
22
+ using ElementSF = typename GemmTraits::ElementSF;
23
+ using ElementD = typename GemmTraits::ElementOutput;
24
+ // Input
25
+ int* expert_offsets{nullptr};
26
+ int* blockscale_offsets{nullptr};
27
+ // Output
28
+ ElementA* a_base{nullptr};
29
+ ElementB* b_base{nullptr};
30
+ ElementSF* sfa_base{nullptr};
31
+ ElementSF* sfb_base{nullptr};
32
+ ElementD* d_base{nullptr};
33
+ ElementA** a_offsets{nullptr};
34
+ ElementB** b_offsets{nullptr};
35
+ ElementSF** sfa_offsets{nullptr};
36
+ ElementSF** sfb_offsets{nullptr};
37
+ ElementD** d_offsets{nullptr};
38
+
39
+ CutlassMxfp8GroupedMmOffsetFunctor() = default;
40
+ CutlassMxfp8GroupedMmOffsetFunctor(
41
+ int* _expert_offsets, int* _blockscale_offsets, ElementA* _a_base,
42
+ ElementB* _b_base, ElementSF* _sfa_base, ElementSF* _sfb_base,
43
+ ElementD* _d_base, ElementA** _a_offsets, ElementB** _b_offsets,
44
+ ElementSF** _sfa_offsets, ElementSF** _sfb_offsets, ElementD** _d_offsets)
45
+ : expert_offsets{_expert_offsets},
46
+ blockscale_offsets{_blockscale_offsets},
47
+ a_base(_a_base),
48
+ b_base(_b_base),
49
+ sfa_base(_sfa_base),
50
+ sfb_base(_sfb_base),
51
+ d_base(_d_base),
52
+ a_offsets(_a_offsets),
53
+ b_offsets(_b_offsets),
54
+ sfa_offsets(_sfa_offsets),
55
+ sfb_offsets(_sfb_offsets),
56
+ d_offsets(_d_offsets) {}
57
+
58
+ void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
59
+ int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
60
+ int64_t blockscale_offset =
61
+ static_cast<int64_t>(blockscale_offsets[expert_id]);
62
+ int64_t a_stride = expert_offset * k;
63
+ int64_t b_stride = expert_id * k * n;
64
+ int64_t d_stride = expert_offset * n;
65
+ int64_t sfa_stride = blockscale_offset * (k / 32);
66
+ int64_t sfb_stride = expert_id * n * (k / 32);
67
+
68
+ a_offsets[expert_id] = a_base + a_stride;
69
+ b_offsets[expert_id] = b_base + b_stride;
70
+ sfa_offsets[expert_id] = sfa_base + sfa_stride;
71
+ sfb_offsets[expert_id] = sfb_base + sfb_stride;
72
+ d_offsets[expert_id] = d_base + d_stride;
73
+ }
74
+};
75
+
76
+template <typename GemmTraits>
77
+struct CutlassMxfp8GroupedMmLayoutFunctor {
78
+ using Sm1xxBlkScaledConfig = typename GemmTraits::Sm1xxBlkScaledConfig;
79
+ using LayoutSFA = typename GemmTraits::LayoutSFA;
80
+ using LayoutSFB = typename GemmTraits::LayoutSFB;
81
+ LayoutSFA* layout_sfa_base{nullptr};
82
+ LayoutSFB* layout_sfb_base{nullptr};
83
+
84
+ CutlassMxfp8GroupedMmLayoutFunctor() = default;
85
+ CutlassMxfp8GroupedMmLayoutFunctor(LayoutSFA* _layout_sfa_base,
86
+ LayoutSFB* _layout_sfb_base)
87
+ : layout_sfa_base(_layout_sfa_base), layout_sfb_base(_layout_sfb_base) {}
88
+
89
+ void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
90
+ LayoutSFA* layout_sfa_ptr = layout_sfa_base + expert_id;
91
+ LayoutSFB* layout_sfb_ptr = layout_sfb_base + expert_id;
92
+ *layout_sfa_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(
93
+ cute::make_shape(m, n, k, 1));
94
+ *layout_sfb_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
95
+ cute::make_shape(m, n, k, 1));
96
+ }
97
+};
98
+
99
+template <typename GemmTraits>
100
+struct CutlassMxfp8GroupedMmStrideFunctor {
101
+ using StrideA = typename GemmTraits::StrideA;
102
+ using StrideB = typename GemmTraits::StrideB;
103
+ using StrideD = typename GemmTraits::StrideD;
104
+ StrideA* stride_A_base{nullptr};
105
+ StrideB* stride_B_base{nullptr};
106
+ StrideD* stride_D_base{nullptr};
107
+
108
+ CutlassMxfp8GroupedMmStrideFunctor() = default;
109
+ CutlassMxfp8GroupedMmStrideFunctor(StrideA* _stride_A_base,
110
+ StrideB* _stride_B_base,
111
+ StrideD* _stride_D_base)
112
+ : stride_A_base(_stride_A_base),
113
+ stride_B_base(_stride_B_base),
114
+ stride_D_base(_stride_D_base) {}
115
+
116
+ void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
117
+ StrideA* stride_A = stride_A_base + expert_id;
118
+ StrideB* stride_B = stride_B_base + expert_id;
119
+ StrideD* stride_D = stride_D_base + expert_id;
120
+ *stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
121
+ *stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
122
+ *stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
123
+ }
124
+};
125
+
126
+template <typename OffsetFunctor, typename LayoutFunctor,
127
+ typename StrideFunctor>
128
+__global__ void cutlassMxfp8GroupedMmPreComputeKernel(
129
+ int* problem_sizes, OffsetFunctor offset_functor,
130
+ LayoutFunctor layout_functor, StrideFunctor stride_functor) {
131
+ int64_t expert_id = static_cast<int64_t>(threadIdx.x);
132
+ int m = problem_sizes[expert_id * 3 + 0];
133
+ int n = problem_sizes[expert_id * 3 + 1];
134
+ int k = problem_sizes[expert_id * 3 + 2];
135
+
136
+ offset_functor(expert_id, m, n, k);
137
+ layout_functor(expert_id, m, n, k);
138
+ stride_functor(expert_id, m, n, k);
139
+}
140
+
141
+} // namespace expert_specialization
\ No newline at end of file
csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh
new
+179
@@ -0,0 +1,179 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
5
+
6
+#pragma once
7
+#include <ATen/cuda/CUDAContext.h>
8
+#include <c10/cuda/CUDAGuard.h>
9
+#include <torch/all.h>
10
+
11
+#include <cassert>
12
+#include <iostream>
13
+#include <string>
14
+
15
+#include "cute/tensor.hpp"
16
+#include "cutlass_mxfp8_grouped_mm_functor.cuh"
17
+#include "cutlass_mxfp8_grouped_mm_traits.cuh"
18
+
19
+namespace expert_specialization {
20
+
21
+template <typename GemmTraits>
22
+void cutlass_mxfp8_grouped_mm_pre_compute(
23
+ torch::Tensor& a_ptrs, torch::Tensor& b_ptrs, torch::Tensor& sfa_ptrs,
24
+ torch::Tensor& sfb_ptrs, torch::Tensor& d_ptrs, torch::Tensor& stride_a,
25
+ torch::Tensor& stride_b, torch::Tensor& stride_d, torch::Tensor& layout_sfa,
26
+ torch::Tensor& layout_sfb, const torch::Tensor& a, const torch::Tensor& b,
27
+ const torch::Tensor& sfa, const torch::Tensor& sfb, const torch::Tensor& d,
28
+ const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
29
+ const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
30
+ using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
31
+ using ElementA = typename OffsetFunctor::ElementA;
32
+ using ElementB = typename OffsetFunctor::ElementB;
33
+ using ElementSF = typename OffsetFunctor::ElementSF;
34
+ using ElementD = typename OffsetFunctor::ElementD;
35
+
36
+ using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor<GemmTraits>;
37
+ using LayoutSFA = typename LayoutFunctor::LayoutSFA;
38
+ using LayoutSFB = typename LayoutFunctor::LayoutSFB;
39
+
40
+ using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor<GemmTraits>;
41
+ using StrideA = typename StrideFunctor::StrideA;
42
+ using StrideB = typename StrideFunctor::StrideB;
43
+ using StrideD = typename StrideFunctor::StrideD;
44
+
45
+ int num_experts = (int)expert_offsets.size(0);
46
+ TORCH_CHECK(num_experts <= 1024,
47
+ "Number of experts cannot exceed 1024, the maximum number of "
48
+ "threads per block.");
49
+
50
+ OffsetFunctor offset_functor(
51
+ reinterpret_cast<int*>(expert_offsets.data_ptr()),
52
+ reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
53
+ reinterpret_cast<ElementA*>(a.data_ptr()),
54
+ reinterpret_cast<ElementB*>(b.data_ptr()),
55
+ reinterpret_cast<ElementSF*>(sfa.data_ptr()),
56
+ reinterpret_cast<ElementSF*>(sfb.data_ptr()),
57
+ reinterpret_cast<ElementD*>(d.data_ptr()),
58
+ reinterpret_cast<ElementA**>(a_ptrs.data_ptr()),
59
+ reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
60
+ reinterpret_cast<ElementSF**>(sfa_ptrs.data_ptr()),
61
+ reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
62
+ reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
63
+ LayoutFunctor layout_functor(
64
+ reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
65
+ reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()));
66
+ StrideFunctor stride_functor(reinterpret_cast<StrideA*>(stride_a.data_ptr()),
67
+ reinterpret_cast<StrideB*>(stride_b.data_ptr()),
68
+ reinterpret_cast<StrideD*>(stride_d.data_ptr()));
69
+ cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>(
70
+ static_cast<int*>(problem_sizes.data_ptr()), offset_functor,
71
+ layout_functor, stride_functor);
72
+}
73
+
74
+template <typename GemmTraits>
75
+void cutlass_mxfp8_grouped_mm(
76
+ const torch::Tensor& a_ptrs, const torch::Tensor& b_ptrs,
77
+ const torch::Tensor& sfa_ptrs, const torch::Tensor& sfb_ptrs,
78
+ const torch::Tensor& d_ptrs, const torch::Tensor& stride_a,
79
+ const torch::Tensor& stride_b, const torch::Tensor& stride_d,
80
+ const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
81
+ const torch::Tensor& problem_sizes, cudaStream_t stream) {
82
+ using Gemm = typename GemmTraits::Gemm;
83
+ using ElementA = typename Gemm::ElementA;
84
+ using ElementB = typename Gemm::ElementB;
85
+ using ElementSF = typename GemmTraits::ElementSF;
86
+ using ElementD = typename GemmTraits::ElementOutput;
87
+ using StrideA = typename GemmTraits::StrideA;
88
+ using StrideB = typename GemmTraits::StrideB;
89
+ using StrideD = typename GemmTraits::StrideD;
90
+ using LayoutSFA = typename GemmTraits::LayoutSFA;
91
+ using LayoutSFB = typename GemmTraits::LayoutSFB;
92
+ using UnderlyingProblemShape =
93
+ typename GemmTraits::ProblemShape::UnderlyingProblemShape;
94
+
95
+ cutlass::KernelHardwareInfo hw_info;
96
+ hw_info.device_id = c10::cuda::current_device();
97
+ hw_info.sm_count =
98
+ at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
99
+ hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
100
+ hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
101
+
102
+ int num_experts = (int)problem_sizes.size(0);
103
+
104
+ UnderlyingProblemShape* underlying_problem_shape =
105
+ reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
106
+
107
+ typename Gemm::Arguments arguments = {
108
+ cutlass::gemm::GemmUniversalMode::kGrouped,
109
+ {num_experts, underlying_problem_shape, nullptr},
110
+ {reinterpret_cast<const ElementA**>(a_ptrs.data_ptr()),
111
+ reinterpret_cast<StrideA*>(stride_a.data_ptr()),
112
+ reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
113
+ reinterpret_cast<StrideB*>(stride_b.data_ptr()),
114
+ reinterpret_cast<const ElementSF**>(sfa_ptrs.data_ptr()),
115
+ reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
116
+ reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr()),
117
+ reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())},
118
+ {{},
119
+ nullptr,
120
+ nullptr,
121
+ reinterpret_cast<ElementD**>(d_ptrs.data_ptr()),
122
+ reinterpret_cast<StrideD*>(stride_d.data_ptr())},
123
+ hw_info,
124
+ {} // Scheduler
125
+ };
126
+
127
+ Gemm gemm;
128
+
129
+ auto can_implement_status = gemm.can_implement(arguments);
130
+ TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
131
+ "Failed to implement GEMM");
132
+
133
+ torch::TensorOptions options_uint8 =
134
+ torch::TensorOptions().dtype(torch::kUInt8).device(d_ptrs.device());
135
+ size_t workspace_size = gemm.get_workspace_size(arguments);
136
+ torch::Tensor workspace = torch::empty(workspace_size, options_uint8);
137
+
138
+ auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
139
+ TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
140
+
141
+ status = gemm.run(stream, nullptr, true); // Enable PDL
142
+ TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
143
+}
144
+
145
+template <typename OutType>
146
+void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
147
+ const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& sfa,
148
+ const torch::Tensor& sfb, torch::Tensor& d,
149
+ const torch::Tensor& problem_sizes, const torch::Tensor& expert_offsets,
150
+ const torch::Tensor& blockscale_offsets, cudaStream_t stream) {
151
+ int num_experts = (int)problem_sizes.size(0);
152
+ torch::TensorOptions options_int64 =
153
+ torch::TensorOptions().dtype(torch::kInt64).device(a.device());
154
+ torch::TensorOptions options_int32 =
155
+ torch::TensorOptions().dtype(torch::kInt32).device(a.device());
156
+
157
+ torch::Tensor a_ptrs = torch::empty(num_experts, options_int64);
158
+ torch::Tensor b_ptrs = torch::empty(num_experts, options_int64);
159
+ torch::Tensor sfa_ptrs = torch::empty(num_experts, options_int64);
160
+ torch::Tensor sfb_ptrs = torch::empty(num_experts, options_int64);
161
+ torch::Tensor d_ptrs = torch::empty(num_experts, options_int64);
162
+
163
+ torch::Tensor stride_a = torch::empty(num_experts, options_int64);
164
+ torch::Tensor stride_b = torch::empty(num_experts, options_int64);
165
+ torch::Tensor stride_d = torch::empty(num_experts, options_int64);
166
+ torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int32);
167
+ torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int32);
168
+
169
+ using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
170
+ cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
171
+ a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
172
+ layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets,
173
+ blockscale_offsets, stream);
174
+ cutlass_mxfp8_grouped_mm<GemmTraits>(
175
+ a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
176
+ layout_sfa, layout_sfb, problem_sizes, stream);
177
+}
178
+
179
+} // namespace expert_specialization
\ No newline at end of file
csrc/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh
new
+127
@@ -0,0 +1,127 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_traits.cuh
5
+
6
+#pragma once
7
+
8
+// Misc
9
+#include "cute/tensor.hpp"
10
+#include "cutlass/arch/arch.h"
11
+#include "cutlass/arch/mma.h"
12
+#include "cutlass/cutlass.h"
13
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
14
+#include "cutlass/epilogue/dispatch_policy.hpp"
15
+#include "cutlass/gemm/dispatch_policy.hpp"
16
+#include "cutlass/gemm/group_array_problem_shape.hpp"
17
+#include "cutlass/layout/layout.h"
18
+#include "cutlass/numeric_conversion.h"
19
+#include "cutlass/numeric_size.h"
20
+
21
+// Collective Builder
22
+#include "cutlass/epilogue/collective/collective_builder.hpp"
23
+#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
24
+#include "cutlass/epilogue/thread/activation.h"
25
+#include "cutlass/gemm/collective/collective_builder.hpp"
26
+
27
+// Integration
28
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
29
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
30
+
31
+namespace expert_specialization {
32
+
33
+using namespace cute;
34
+
35
+// Different configs for 1SM and 2SM MMA kernel
36
+struct MMA1SMConfig {
37
+ using MmaTileShape = Shape<_128, _128, _128>;
38
+ using KernelSchedule =
39
+ cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100;
40
+ using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
41
+ const static dim3 preferred_cluster;
42
+ const static dim3 fallback_cluster;
43
+};
44
+const dim3 MMA1SMConfig::preferred_cluster(1, 4, 1);
45
+const dim3 MMA1SMConfig::fallback_cluster(1, 2, 1);
46
+
47
+template <typename _MMAConfig, typename OutputDtype>
48
+struct CutlassMxfp8GroupedMmGemmTraits {
49
+ using MMAConfig = _MMAConfig;
50
+ using ElementInput = cutlass::float_e4m3_t;
51
+ using ElementOutput = OutputDtype;
52
+ using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;
53
+
54
+ // A matrix configuration
55
+ using ElementA = cutlass::mx_float8_t<ElementInput>;
56
+ using LayoutA = cutlass::layout::RowMajor;
57
+ constexpr static int AlignmentA = 32;
58
+
59
+ // B matrix configuration
60
+ using ElementB = cutlass::mx_float8_t<ElementInput>;
61
+ using LayoutB = cutlass::layout::ColumnMajor;
62
+ constexpr static int AlignmentB = 32;
63
+
64
+ // C/D matrix configuration
65
+ using ElementC = void;
66
+ using ElementD = ElementOutput;
67
+ using LayoutC = cutlass::layout::RowMajor;
68
+ using LayoutD = cutlass::layout::RowMajor;
69
+ constexpr static int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
70
+ constexpr static int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
71
+ using ElementAccumulator = float;
72
+
73
+ static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
74
+ using CustomEVTIdentity = // acc
75
+ cutlass::epilogue::fusion::Sm90EVT<
76
+ cutlass::epilogue::fusion::Sm90Compute<
77
+ cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator,
78
+ RoundStyle>,
79
+ cutlass::epilogue::fusion::Sm90AccFetch>;
80
+
81
+ // Core kernel configurations
82
+ using ArchTag = cutlass::arch::Sm100;
83
+ using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
84
+ using StageCountType = cutlass::gemm::collective::StageCountAuto;
85
+
86
+ // Runtime Cluster Shape
87
+ using ClusterShape = Shape<int32_t, int32_t, _1>;
88
+
89
+ // Define Epilogue
90
+ using CollectiveEpilogue =
91
+ typename cutlass::epilogue::collective::CollectiveBuilder<
92
+ ArchTag, OperatorClass, typename MMAConfig::MmaTileShape,
93
+ ClusterShape, Shape<_64, _64>, ElementAccumulator, ElementAccumulator,
94
+ ElementC, LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD,
95
+ typename MMAConfig::EpilogueSchedule,
96
+ CustomEVTIdentity>::CollectiveOp;
97
+
98
+ // Define Mainloop
99
+ using CollectiveMainloop =
100
+ typename cutlass::gemm::collective::CollectiveBuilder<
101
+ ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
102
+ LayoutB*, AlignmentB, ElementAccumulator,
103
+ typename MMAConfig::MmaTileShape, ClusterShape,
104
+ cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
105
+ sizeof(typename CollectiveEpilogue::SharedStorage))>,
106
+ typename MMAConfig::KernelSchedule>::CollectiveOp;
107
+
108
+ // Define GemmKernel
109
+ using GemmKernel =
110
+ cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
111
+ CollectiveEpilogue>;
112
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
113
+
114
+ using ElementSF = typename Gemm::GemmKernel::ElementSF;
115
+ using StrideA = typename Gemm::GemmKernel::InternalStrideA;
116
+ using StrideB = typename Gemm::GemmKernel::InternalStrideB;
117
+ using StrideC = typename Gemm::GemmKernel::InternalStrideC;
118
+ using StrideD = typename Gemm::GemmKernel::InternalStrideD;
119
+ using LayoutSFA =
120
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
121
+ using LayoutSFB =
122
+ typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
123
+ using Sm1xxBlkScaledConfig =
124
+ typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
125
+};
126
+
127
+} // namespace expert_specialization
\ No newline at end of file
csrc/moe/mxfp8_moe/mxfp8_experts_quant.cu
new
+60
@@ -0,0 +1,60 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
5
+
6
+#include <torch/all.h>
7
+
8
+#include "mxfp8_experts_quant.cuh"
9
+
10
+void mxfp8_experts_quant(const torch::Tensor& input,
11
+ const torch::Tensor& problem_sizes,
12
+ const torch::Tensor& expert_offsets,
13
+ const torch::Tensor& blockscale_offsets,
14
+ torch::Tensor& quant_output,
15
+ torch::Tensor& scale_factor) {
16
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
17
+ TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
18
+ TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
19
+ TORCH_CHECK(input.strides()[1] == 1, "input must be row major");
20
+ TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
21
+ TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
22
+ "problem_sizes must be int32");
23
+ TORCH_CHECK(expert_offsets.dtype() == torch::kInt32,
24
+ "expert_offsets must be int32");
25
+ TORCH_CHECK(blockscale_offsets.dtype() == torch::kInt32,
26
+ "blockscale_offsets must be int32");
27
+
28
+ auto groups = problem_sizes.size(0);
29
+ TORCH_CHECK(
30
+ expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
31
+ "expert_offsets must be 1D and have size equal to the number of groups");
32
+ TORCH_CHECK(
33
+ blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
34
+ "blockscale_offsets must be 1D and have size equal to the number of "
35
+ "groups");
36
+
37
+ auto stream = at::cuda::getCurrentCUDAStream();
38
+ if (input.dtype() == torch::kBFloat16) {
39
+ expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
40
+ input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
41
+ scale_factor);
42
+ } else if (input.dtype() == torch::kFloat16) {
43
+ expert_specialization::launch_mxfp8_experts_quant<__half>(
44
+ input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
45
+ scale_factor);
46
+ } else {
47
+ TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
48
+ }
49
+#else
50
+ TORCH_CHECK(false,
51
+ "No implemented mxfp8_experts_quant for "
52
+ "current device");
53
+#endif
54
+}
55
+
56
+#include "core/registration.h"
57
+
58
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
59
+ m.impl("mxfp8_experts_quant", mxfp8_experts_quant);
60
+}
\ No newline at end of file
csrc/moe/mxfp8_moe/mxfp8_experts_quant.cuh
new
+414
@@ -0,0 +1,414 @@
1
+// SPDX-License-Identifier: Apache-2.0
2
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+// Adapted from SGLang:
4
+// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
5
+
6
+#pragma once
7
+#include <ATen/cuda/CUDAContext.h>
8
+#include <c10/cuda/CUDAGuard.h>
9
+#include <cuda.h>
10
+#include <cuda_bf16.h>
11
+#include <cuda_fp16.h>
12
+#include <torch/all.h>
13
+
14
+#include <cuda/ptx>
15
+
16
+#include "cute/tensor.hpp"
17
+
18
+namespace expert_specialization {
19
+
20
+using namespace cute;
21
+
22
+constexpr uint32_t THREAD_BLOCK_SIZE = 128;
23
+constexpr uint32_t WARP_SIZE = 32;
24
+constexpr int BLOCK_M = 128;
25
+constexpr int BLOCK_K = 128;
26
+using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
27
+using ValLayout = Layout<Shape<_1, _16>>;
28
+using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
29
+using SfR2SValLayout = Layout<Shape<_1, _1>>;
30
+using ScaleFactorTileLayout =
31
+ Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
32
+
33
+// Fast reciprocal.
34
+inline __device__ float reciprocal_approximate_ftz(float a) {
35
+ float b;
36
+ asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
37
+ return b;
38
+}
39
+
40
+// Some code references TRT-LLM:
41
+// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
42
+template <typename FragmentS, typename FragmentD>
43
+__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s,
44
+ FragmentD& fragment_d) {
45
+ using FragmentSLayout = typename FragmentS::layout_type;
46
+ using FragmentDLayout = typename FragmentD::layout_type;
47
+ FragmentSLayout fragment_s_layout;
48
+ FragmentDLayout fragment_d_layout;
49
+ static_assert(is_static<FragmentSLayout>::value &&
50
+ size(fragment_s_layout) == 16);
51
+ static_assert(is_static<FragmentDLayout>::value &&
52
+ size(fragment_d_layout) == 16);
53
+
54
+ constexpr int eles_per_thr = 16;
55
+ using ValType = typename FragmentS::element_type;
56
+ using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>,
57
+ __nv_bfloat162, __half2>;
58
+ VecType vec[8];
59
+ // Assign vals
60
+ vec[0].x = fragment_s(Int<0>{});
61
+ vec[0].y = fragment_s(Int<1>{});
62
+ vec[1].x = fragment_s(Int<2>{});
63
+ vec[1].y = fragment_s(Int<3>{});
64
+ vec[2].x = fragment_s(Int<4>{});
65
+ vec[2].y = fragment_s(Int<5>{});
66
+ vec[3].x = fragment_s(Int<6>{});
67
+ vec[3].y = fragment_s(Int<7>{});
68
+ vec[4].x = fragment_s(Int<8>{});
69
+ vec[4].y = fragment_s(Int<9>{});
70
+ vec[5].x = fragment_s(Int<10>{});
71
+ vec[5].y = fragment_s(Int<11>{});
72
+ vec[6].x = fragment_s(Int<12>{});
73
+ vec[6].y = fragment_s(Int<13>{});
74
+ vec[7].x = fragment_s(Int<14>{});
75
+ vec[7].y = fragment_s(Int<15>{});
76
+
77
+ auto local_max = __habs2(vec[0]);
78
+ for (int i = 1; i < eles_per_thr / 2; i++) {
79
+ local_max = __hmax2(__habs2(vec[i]), local_max);
80
+ }
81
+ local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
82
+
83
+ // Get the final absolute maximum values.
84
+ float block_max(0.0f);
85
+ if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
86
+ block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
87
+ } else {
88
+ block_max = __half2float(__hmax(local_max.x, local_max.y));
89
+ }
90
+ // Get the SF (max value of the vector / max value of mxfp8).
91
+ float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
92
+ // 8 bits representation of the SF.
93
+ uint8_t fp8_sf_val;
94
+
95
+ __nv_fp8_e8m0 tmp_sf_val;
96
+ tmp_sf_val.__x =
97
+ __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
98
+ sf_val = static_cast<float>(tmp_sf_val);
99
+ fp8_sf_val = tmp_sf_val.__x;
100
+ // Get the output scale (reciprocal of the SFValue).
101
+ float output_scale =
102
+ block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
103
+
104
+ // Convert the input to float.
105
+ float2 fp2_vals[eles_per_thr / 2];
106
+
107
+#pragma unroll
108
+ for (int i = 0; i < eles_per_thr / 2; i++) {
109
+ if constexpr (std::is_same_v<ValType, __half>) {
110
+ fp2_vals[i] = __half22float2(vec[i]);
111
+ } else {
112
+ fp2_vals[i] = __bfloat1622float2(vec[i]);
113
+ }
114
+ fp2_vals[i].x *= output_scale;
115
+ fp2_vals[i].y *= output_scale;
116
+ }
117
+ union {
118
+ uint8_t bytes[16];
119
+ __nv_fp8x2_e4m3 elts[8];
120
+ } u;
121
+ u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
122
+ u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
123
+ u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
124
+ u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
125
+ u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
126
+ u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
127
+ u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
128
+ u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
129
+ fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
130
+ fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
131
+ fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
132
+ fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
133
+ fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
134
+ fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
135
+ fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
136
+ fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
137
+ fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
138
+ fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
139
+ fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
140
+ fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
141
+ fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
142
+ fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
143
+ fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
144
+ fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
145
+ return fp8_sf_val;
146
+}
147
+
148
+template <typename TensorS, typename TensorP, typename TensorD,
149
+ typename TensorSharedSF, typename TensorSF, typename TiledCopyG2R,
150
+ typename TiledCopyR2G, typename TiledCopyR2S>
151
+__inline__ __device__ void mxfp8_experts_quant_tile(
152
+ TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d,
153
+ TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m,
154
+ TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g,
155
+ TiledCopyR2S& tiled_copy_r2s) {
156
+ static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 &&
157
+ size(get<1>(typename TensorS::layout_type{})) == 128 &&
158
+ stride(get<1>(typename TensorS::layout_type{})) == 1);
159
+ static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 &&
160
+ size(get<1>(typename TensorD::layout_type{})) == 128 &&
161
+ stride(get<1>(typename TensorD::layout_type{})) == 1);
162
+ static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 &&
163
+ size(get<1>(typename TensorP::layout_type{})) == 128);
164
+ static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
165
+ size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
166
+ static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 &&
167
+ size(get<1>(typename TensorSF::layout_type{})) == 4);
168
+
169
+ using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
170
+ auto tiler_mn = Tiler_MN{};
171
+ static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
172
+
173
+ auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
174
+ auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
175
+ auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
176
+ static_assert(size<2>(tiled_tensor_s) == 1);
177
+ static_assert(size<2>(tiled_tensor_p) == 1);
178
+ static_assert(size<2>(tiled_tensor_d) == 1);
179
+ auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
180
+ auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
181
+ auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
182
+
183
+ using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
184
+ auto sf_tiler_mn = SF_Tiler_MN{};
185
+ static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
186
+
187
+ auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
188
+ auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
189
+ auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
190
+ auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
191
+
192
+ constexpr int tile_loop_count = size<1>(tiled_tensor_s);
193
+ constexpr int rows_in_tile = 16;
194
+ // We don't need to clear shared memory
195
+ // clear(squeeze_tiled_tensor_shared_sf);
196
+#pragma unroll 4
197
+ for (int t = 0; t < tile_loop_count; t++) {
198
+ if (t * rows_in_tile >= m) {
199
+ break;
200
+ }
201
+ auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
202
+ auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
203
+ auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
204
+ auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
205
+ auto current_copy_tile_shared_sf =
206
+ tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
207
+
208
+ // Global to Register copy
209
+ auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
210
+ auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
211
+ auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
212
+ auto input_fragment = make_fragment_like(thr_tile_g2r_s);
213
+
214
+ // Register to Global copy
215
+ auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
216
+ auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
217
+ auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
218
+ auto output_fragment = make_fragment_like(thr_tile_r2g_d);
219
+
220
+ // Register to Shared copy
221
+ auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
222
+ auto thr_tile_r2s_shared_sf =
223
+ thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
224
+ auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
225
+
226
+ // CopyG2R & convert & CopyR2G
227
+ copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
228
+ uint8_t fp8_sf_val =
229
+ cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
230
+ copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
231
+ shared_sf_fragment[0] = fp8_sf_val;
232
+
233
+ // Before first copy r2s, clear shared memory and wait previous group
234
+ if (t == 0 && threadIdx.x == 0) {
235
+ // Wait for the group to have completed reading from shared memory.
236
+ cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
237
+ }
238
+ __syncthreads();
239
+
240
+ if (threadIdx.x % 2 == 0) {
241
+ copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
242
+ }
243
+ __syncthreads();
244
+ }
245
+
246
+ // Wait for shared memory writes to be visible to TMA engine.
247
+ cuda::ptx::fence_proxy_async(cuda::ptx::space_shared); // b)
248
+ __syncthreads();
249
+
250
+ if (threadIdx.x == 0) {
251
+ cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared,
252
+ squeeze_tiled_tensor_sf.data().get(),
253
+ squeeze_tiled_tensor_shared_sf.data().get(), 512);
254
+ // Wait for TMA transfer to have finished reading shared memory.
255
+ // Create a "bulk async-group" out of the previous bulk copy operation.
256
+ cuda::ptx::cp_async_bulk_commit_group();
257
+ }
258
+ __syncthreads();
259
+}
260
+
261
+template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G,
262
+ typename TiledCopyR2S>
263
+__global__ void mxfp8_experts_quant_kernel(
264
+ const T_IN* input, const int* problem_sizes, const int* expert_offsets,
265
+ const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output,
266
+ uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r,
267
+ TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) {
268
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
269
+ __shared__ __align__(512) uint8_t shared_memory[512];
270
+ ScaleFactorTileLayout scale_factor_tile_layout{};
271
+ auto scale_factor_shared =
272
+ make_tensor(make_smem_ptr(shared_memory),
273
+ scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1)
274
+ // TODO: Transform Groupwise Schedule into a more efficient Schedule
275
+ for (int g = 0; g < groups; g++) {
276
+ int m = problem_sizes[g * 3 + 0];
277
+ int k = problem_sizes[g * 3 + 2];
278
+ int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
279
+ int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
280
+
281
+ auto input_tensor = make_tensor(
282
+ make_gmem_ptr(input + expert_offset * k),
283
+ make_layout(make_shape(m, k),
284
+ LayoutRight{})); // (M, K):(K, 1) half_t/bfloat16_t
285
+
286
+ auto quant_output_tensor = make_tensor(
287
+ make_gmem_ptr(quant_output + expert_offset * k),
288
+ make_layout(make_shape(m, k),
289
+ LayoutRight{})); // (M, K):(K, 1) cutlass::float_e4m3_t
290
+
291
+ auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
292
+ auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout,
293
+ scale_factor_shape, LayoutRight{});
294
+ // layout<0>(layout<0>(scale_factor_layout)) (_32,_4):(_16,_4) -- static
295
+ // layout<1>(layout<0>(scale_factor_layout)) M_align_128 / 128 -- dynamic
296
+ // shape dynamic stride layout<0>(layout<1>(scale_factor_layout)) _4:_1 --
297
+ // static layout<1>(layout<1>(scale_factor_layout)) (K / 32) / 4 : _512 --
298
+ // dynamic shape static stride
299
+
300
+ // Reshape to zipped layout for 1D indexing
301
+ auto zipped_scale_factor_layout = make_layout(
302
+ make_layout(layout<0>(layout<0>(scale_factor_layout)),
303
+ layout<0>(layout<1>(scale_factor_layout))),
304
+ make_layout(
305
+ layout<1>(layout<0>(scale_factor_layout)),
306
+ layout<1>(layout<1>(
307
+ scale_factor_layout)))); // (((_32,_4),_4),(M_align_128 /
308
+ // 128,(K / 32) /
309
+ // 4)):(((_16,_4),_1),(?,_512))
310
+
311
+ auto scale_factor_tensor =
312
+ make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)),
313
+ zipped_scale_factor_layout);
314
+
315
+ // Used for cases where M is not divisible by 128 (most scenarios).
316
+ auto input_shape = shape(input_tensor); // (M, K):(K, 1)
317
+ auto identity_tensor = make_identity_tensor(input_shape);
318
+ auto predict_tensor = cute::lazy::transform(
319
+ identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
320
+
321
+ // (_128, _128)
322
+ auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
323
+
324
+ auto tiled_input_tensor = zipped_divide(
325
+ input_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
326
+ auto tiled_quant_output_tensor =
327
+ zipped_divide(quant_output_tensor,
328
+ tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
329
+ auto tiled_predict_tensor = zipped_divide(
330
+ predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
331
+
332
+ auto total_tiles =
333
+ size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128)
334
+ decltype(total_tiles) blk_offset = blockIdx.x;
335
+ while (blk_offset < total_tiles) {
336
+ auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
337
+ auto current_quant_output_tile =
338
+ tensor<0>(tiled_quant_output_tensor(_, blk_offset));
339
+ auto current_predict_tile =
340
+ tensor<0>(tiled_predict_tensor(_, blk_offset));
341
+ auto current_scale_factor_tile =
342
+ tensor<0>(scale_factor_tensor(_, blk_offset));
343
+
344
+ mxfp8_experts_quant_tile<
345
+ decltype(current_input_tile), decltype(current_predict_tile),
346
+ decltype(current_quant_output_tile), decltype(scale_factor_shared),
347
+ decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G,
348
+ TiledCopyR2S>(current_input_tile, current_predict_tile,
349
+ current_quant_output_tile, scale_factor_shared,
350
+ current_scale_factor_tile, m, tiled_copy_g2r,
351
+ tiled_copy_r2g, tiled_copy_r2s);
352
+ blk_offset += gridDim.x;
353
+ }
354
+ }
355
+#endif
356
+}
357
+
358
+template <typename T_IN>
359
+void launch_mxfp8_experts_quant(const torch::Tensor& input,
360
+ const torch::Tensor& problem_sizes,
361
+ const torch::Tensor& expert_offsets,
362
+ const torch::Tensor& blockscale_offsets,
363
+ torch::Tensor& quant_output,
364
+ torch::Tensor& scale_factor) {
365
+ ThrLayout thr_layout{};
366
+ ValLayout val_layout{};
367
+ SfR2SThrLayout r2s_thr_layout{};
368
+ SfR2SValLayout r2s_val_layout{};
369
+
370
+ using CopyOpG2R =
371
+ UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
372
+ using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
373
+ auto tiled_copy_g2r = cute::make_tiled_copy(
374
+ CopyAtomG2R{}, thr_layout, val_layout); // Tiler_MN: (16, 128)
375
+
376
+ using CopyOpR2G = UniversalCopy<
377
+ cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
378
+ using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
379
+ auto tiled_copy_r2g = cute::make_tiled_copy(
380
+ CopyAtomR2G{}, thr_layout, val_layout); // Tiler_MN: (16, 128)
381
+
382
+ using CopyOpR2S =
383
+ UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
384
+ using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
385
+ auto tiled_copy_r2s = cute::make_tiled_copy(
386
+ CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4)
387
+
388
+ int max_active_blocks_per_sm = -1;
389
+ AT_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
390
+ &max_active_blocks_per_sm,
391
+ mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
392
+ decltype(tiled_copy_r2g),
393
+ decltype(tiled_copy_r2s)>,
394
+ THREAD_BLOCK_SIZE, 0));
395
+
396
+ dim3 grid(at::cuda::getCurrentDeviceProperties()->multiProcessorCount *
397
+ max_active_blocks_per_sm,
398
+ 1, 1);
399
+ dim3 block(THREAD_BLOCK_SIZE, 1, 1);
400
+ int num_experts = (int)problem_sizes.size(0);
401
+ auto stream = at::cuda::getCurrentCUDAStream();
402
+ mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
403
+ decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
404
+ <<<grid, block, 0, stream>>>(
405
+ reinterpret_cast<const T_IN*>(input.data_ptr()),
406
+ reinterpret_cast<const int*>(problem_sizes.data_ptr()),
407
+ reinterpret_cast<const int*>(expert_offsets.data_ptr()),
408
+ reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
409
+ reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
410
+ reinterpret_cast<uint8_t*>(scale_factor.data_ptr()), num_experts,
411
+ tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s);
412
+}
413
+
414
+} // namespace expert_specialization
\ No newline at end of file
csrc/torch_bindings.cpp
+16
@@ -426,6 +426,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
426
" Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
427
// conditionally compiled so impl registration is in source file
428
429
+ // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
430
+ ops.def(
431
+ "mxfp8_experts_quant("
432
+ " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
433
+ " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
434
+ " -> ()");
435
+ // conditionally compiled so impl registration is in source file
436
+
437
+ // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
438
+ ops.def(
439
+ "cutlass_mxfp8_grouped_mm("
440
+ " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
441
+ " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
442
+ " -> ()");
443
+ // conditionally compiled so impl registration is in source file
444
+
445
// CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
446
// quantization, as well as bias
447
ops.def(
tests/kernels/moe/test_cutlass_mxfp8_grouped_mm.py
new
+237
@@ -0,0 +1,237 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Adapted from SGLang:
4
+# https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/tests/test_es_fp8_blockwise_moe.py
5
+
6
+"""Tests for SM100 CUTLASS MXFP8 grouped MoE kernels."""
7
+
8
+import random
9
+
10
+import pytest
11
+import torch
12
+
13
+from tests.kernels.utils import torch_moe_single
14
+from vllm import _custom_ops as ops
15
+from vllm.platforms import current_platform
16
+from vllm.utils.torch_utils import set_random_seed
17
+
18
+random.seed(42)
19
+set_random_seed(42)
20
+
21
+
22
+def align(val: int, alignment: int = 128) -> int:
23
+ return int((val + alignment - 1) // alignment * alignment)
24
+
25
+
26
+# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
27
+def calc_diff(x, y):
28
+ x, y = x.double(), y.double()
29
+ denominator = (x * x + y * y).sum()
30
+ sim = 2 * (x * y).sum() / denominator
31
+ return 1 - sim
32
+
33
+
34
+def is_sm100_supported() -> bool:
35
+ return current_platform.is_cuda() and current_platform.is_device_capability_family(
36
+ 100
37
+ )
38
+
39
+
40
+def compute_ref_output(
41
+ input_tensor: torch.Tensor,
42
+ weight_list: list[torch.Tensor],
43
+ expert_offsets: list[int],
44
+ expert_offset: int,
45
+ num_experts: int,
46
+) -> torch.Tensor:
47
+ # Build a top-1 routing score so each token maps to its owning expert.
48
+ score = torch.full(
49
+ (expert_offset, num_experts),
50
+ -1e9,
51
+ device=input_tensor.device,
52
+ dtype=torch.float32,
53
+ )
54
+ for g in range(num_experts):
55
+ start = expert_offsets[g]
56
+ end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
57
+ score[start:end, g] = 0.0
58
+
59
+ return torch_moe_single(
60
+ input_tensor, torch.stack(weight_list, dim=0), score, topk=1
61
+ )
62
+
63
+
64
+def compute_kernel_output(
65
+ input_tensor: torch.Tensor,
66
+ weight_tensor: torch.Tensor,
67
+ problem_sizes: list[list[int]],
68
+ aux_problem_sizes: list[list[int]],
69
+ expert_offsets: list[int],
70
+ aux_expert_offsets: list[int],
71
+ input_blockscale_offsets: list[int],
72
+ weight_blockscale_offsets: list[int],
73
+ input_blockscale_offset: int,
74
+ n_g: int,
75
+ k_g: int,
76
+ num_experts: int,
77
+ expert_offset: int,
78
+ out_dtype: torch.dtype,
79
+) -> torch.Tensor:
80
+ device = input_tensor.device
81
+ _problem_sizes = torch.tensor(problem_sizes).to(device=device, dtype=torch.int32)
82
+ _aux_problem_sizes = torch.tensor(aux_problem_sizes).to(
83
+ device=device, dtype=torch.int32
84
+ )
85
+ _expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32)
86
+ _aux_expert_offsets = torch.tensor(aux_expert_offsets).to(
87
+ device=device, dtype=torch.int32
88
+ )
89
+ _input_blockscale_offsets = torch.tensor(input_blockscale_offsets).to(
90
+ device=device, dtype=torch.int32
91
+ )
92
+ _weight_blockscale_offsets = torch.tensor(weight_blockscale_offsets).to(
93
+ device=device, dtype=torch.int32
94
+ )
95
+
96
+ input_quant = torch.zeros_like(
97
+ input_tensor, dtype=torch.float8_e4m3fn, device=device
98
+ )
99
+ input_scale_factor = torch.zeros(
100
+ (input_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device
101
+ )
102
+
103
+ weight_quant = torch.zeros_like(
104
+ weight_tensor, dtype=torch.float8_e4m3fn, device=device
105
+ )
106
+ weight_scale_factor = torch.zeros(
107
+ (num_experts, n_g, k_g // 32), dtype=torch.uint8, device=device
108
+ )
109
+
110
+ ops.mxfp8_experts_quant(
111
+ input_tensor,
112
+ _problem_sizes,
113
+ _expert_offsets,
114
+ _input_blockscale_offsets,
115
+ input_quant,
116
+ input_scale_factor,
117
+ )
118
+
119
+ ops.mxfp8_experts_quant(
120
+ weight_tensor,
121
+ _aux_problem_sizes,
122
+ _aux_expert_offsets,
123
+ _weight_blockscale_offsets,
124
+ weight_quant,
125
+ weight_scale_factor,
126
+ )
127
+ weight_quant = weight_quant.view(num_experts, n_g, k_g).transpose(1, 2)
128
+ weight_scale_factor = weight_scale_factor.view(
129
+ num_experts, n_g, k_g // 32
130
+ ).transpose(1, 2)
131
+
132
+ output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
133
+ ops.cutlass_mxfp8_grouped_mm(
134
+ input_quant,
135
+ weight_quant,
136
+ input_scale_factor,
137
+ weight_scale_factor,
138
+ output,
139
+ _problem_sizes,
140
+ _expert_offsets,
141
+ _input_blockscale_offsets,
142
+ )
143
+ return output
144
+
145
+
146
+@pytest.mark.skipif(
147
+ not is_sm100_supported(),
148
+ reason=(
149
+ "cutlass_mxfp8_grouped_mm and mxfp8_experts_quant "
150
+ "are only supported on CUDA SM100"
151
+ ),
152
+)
153
+@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
154
+@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16])
155
+def test_cutlass_mxfp8_grouped_mm(num_experts, out_dtype):
156
+ device = "cuda"
157
+ alignment = 128
158
+ n_g = random.randint(1, 64) * alignment
159
+ k_g = random.randint(1, 64) * alignment
160
+
161
+ expert_offset = 0
162
+ expert_offsets = []
163
+ aux_expert_offset = 0
164
+ aux_expert_offsets = []
165
+ input_blockscale_offset = 0
166
+ input_blockscale_offsets = []
167
+ weight_blockscale_offset = 0
168
+ weight_blockscale_offsets = []
169
+ problem_sizes = []
170
+ aux_problem_sizes = []
171
+ input_list = []
172
+ weight_list = []
173
+
174
+ for g in range(num_experts):
175
+ m_g = random.randint(1, 512)
176
+ expert_offsets.append(expert_offset)
177
+ expert_offset += m_g
178
+ aux_expert_offsets.append(aux_expert_offset)
179
+ aux_expert_offset += n_g
180
+ input_blockscale_offsets.append(input_blockscale_offset)
181
+ input_blockscale_offset += align(m_g, 128)
182
+ weight_blockscale_offsets.append(weight_blockscale_offset)
183
+ weight_blockscale_offset += n_g # n_g already align to 128
184
+ problem_sizes.append([m_g, n_g, k_g])
185
+ aux_problem_sizes.append([n_g, m_g, k_g])
186
+
187
+ input_tensor = torch.normal(
188
+ 0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype
189
+ ) # (M, K):(K, 1)
190
+ weight_tensor = torch.normal(
191
+ 0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype
192
+ ) # (N, K):(K, 1)
193
+
194
+ input_list.append(input_tensor)
195
+ weight_list.append(weight_tensor)
196
+ input_tensor = torch.concat(input_list, dim=0)
197
+ weight_tensor = torch.concat(weight_list, dim=0)
198
+
199
+ ref_output = compute_ref_output(
200
+ input_tensor=input_tensor,
201
+ weight_list=weight_list,
202
+ expert_offsets=expert_offsets,
203
+ expert_offset=expert_offset,
204
+ num_experts=num_experts,
205
+ )
206
+ output = compute_kernel_output(
207
+ input_tensor=input_tensor,
208
+ weight_tensor=weight_tensor,
209
+ problem_sizes=problem_sizes,
210
+ aux_problem_sizes=aux_problem_sizes,
211
+ expert_offsets=expert_offsets,
212
+ aux_expert_offsets=aux_expert_offsets,
213
+ input_blockscale_offsets=input_blockscale_offsets,
214
+ weight_blockscale_offsets=weight_blockscale_offsets,
215
+ input_blockscale_offset=input_blockscale_offset,
216
+ n_g=n_g,
217
+ k_g=k_g,
218
+ num_experts=num_experts,
219
+ expert_offset=expert_offset,
220
+ out_dtype=out_dtype,
221
+ )
222
+
223
+ for g in range(num_experts):
224
+ baseline = ref_output[
225
+ expert_offsets[g] : (expert_offsets[g] + problem_sizes[g][0])
226
+ ]
227
+ actual = output[expert_offsets[g] : (expert_offsets[g] + problem_sizes[g][0])]
228
+ diff = calc_diff(actual, baseline)
229
+ assert diff < 0.001
230
+ print(
231
+ f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, "
232
+ f"out_dtype={out_dtype}, diff={diff:.5f}: OK"
233
+ )
234
+
235
+
236
+if __name__ == "__main__":
237
+ pytest.main([__file__])
vllm/_custom_ops.py
+70
@@ -1102,6 +1102,76 @@ def cutlass_fp4_moe_mm(
1102
)
1103
1104
1105
+def mxfp8_experts_quant(
1106
+ input_tensor: torch.Tensor,
1107
+ problem_sizes: torch.Tensor,
1108
+ expert_offsets: torch.Tensor,
1109
+ blockscale_offsets: torch.Tensor,
1110
+ quant_output: torch.Tensor,
1111
+ scale_factor: torch.Tensor,
1112
+) -> None:
1113
+ torch.ops._C.mxfp8_experts_quant(
1114
+ input_tensor,
1115
+ problem_sizes,
1116
+ expert_offsets,
1117
+ blockscale_offsets,
1118
+ quant_output,
1119
+ scale_factor,
1120
+ )
1121
+
1122
+
1123
+def cutlass_mxfp8_grouped_mm(
1124
+ a_tensors: torch.Tensor,
1125
+ b_tensors: torch.Tensor,
1126
+ a_scales: torch.Tensor,
1127
+ b_scales: torch.Tensor,
1128
+ out_tensors: torch.Tensor,
1129
+ problem_sizes: torch.Tensor,
1130
+ expert_offsets: torch.Tensor,
1131
+ blockscale_offsets: torch.Tensor,
1132
+) -> None:
1133
+ torch.ops._C.cutlass_mxfp8_grouped_mm(
1134
+ a_tensors,
1135
+ b_tensors,
1136
+ a_scales,
1137
+ b_scales,
1138
+ out_tensors,
1139
+ problem_sizes,
1140
+ expert_offsets,
1141
+ blockscale_offsets,
1142
+ )
1143
+
1144
+
1145
+if hasattr(torch.ops._C, "mxfp8_experts_quant"):
1146
+
1147
+ @register_fake("_C::mxfp8_experts_quant")
1148
+ def _mxfp8_experts_quant_fake(
1149
+ input_tensor: torch.Tensor,
1150
+ problem_sizes: torch.Tensor,
1151
+ expert_offsets: torch.Tensor,
1152
+ blockscale_offsets: torch.Tensor,
1153
+ quant_output: torch.Tensor,
1154
+ scale_factor: torch.Tensor,
1155
+ ) -> None:
1156
+ return None
1157
+
1158
+
1159
+if hasattr(torch.ops._C, "cutlass_mxfp8_grouped_mm"):
1160
+
1161
+ @register_fake("_C::cutlass_mxfp8_grouped_mm")
1162
+ def _cutlass_mxfp8_grouped_mm_fake(
1163
+ a_tensors: torch.Tensor,
1164
+ b_tensors: torch.Tensor,
1165
+ a_scales: torch.Tensor,
1166
+ b_scales: torch.Tensor,
1167
+ out_tensors: torch.Tensor,
1168
+ problem_sizes: torch.Tensor,
1169
+ expert_offsets: torch.Tensor,
1170
+ blockscale_offsets: torch.Tensor,
1171
+ ) -> None:
1172
+ return None
1173
+
1174
+
1175
# gptq_marlin
1176
def gptq_marlin_repack(
1177
b_q_weight: torch.Tensor,