A high-throughput and memory-efficient inference and serving engine for LLMs
[Refactor] Remove dead cutlass mxfp8 code (#44681)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Wentao Ye committed
Jun 18, 2026 at 00:18 UTC
b4c80ec0fd19c13a53d89623bb5957cd5cd631bb
10 files changed
-1369
CMakeLists.txt
-29
@@ -363,35 +363,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
363
SRCS "${VLLM_EXT_SRC}"
364
CUDA_ARCHS "${CUDA_ARCHS}")
365
366
- # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
367
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
368
- cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
369
- else()
370
- cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
371
- endif()
372
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
373
- set(ES_MXFP8_GROUPED_MM_SRCS
374
- "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
375
- "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
376
- set_gencode_flags_for_srcs(
377
- SRCS "${ES_MXFP8_GROUPED_MM_SRCS}"
378
- CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
379
- list(APPEND VLLM_STABLE_EXT_SRC "${ES_MXFP8_GROUPED_MM_SRCS}")
380
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
381
- message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
382
- else()
383
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
384
- AND ES_MXFP8_GROUPED_MM_ARCHS)
385
- message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
386
- "not >= 12.8.")
387
- else()
388
- message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
389
- "in CUDA target architectures.")
390
- endif()
391
- endif()
392
-
393
-
394
-
366
# if CUDA endif
367
endif()
368
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu
deleted
-69
@@ -1,69 +0,0 @@
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/csrc/stable/library.h>
7
-#include <torch/csrc/stable/tensor.h>
8
-#include "libtorch_stable/torch_utils.h"
9
-
10
-#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
11
-
12
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a,
13
- const torch::stable::Tensor& b,
14
- const torch::stable::Tensor& sfa,
15
- const torch::stable::Tensor& sfb,
16
- torch::stable::Tensor& d,
17
- const torch::stable::Tensor& problem_sizes,
18
- const torch::stable::Tensor& expert_offsets,
19
- const torch::stable::Tensor& blockscale_offsets) {
20
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
21
- STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
22
- STD_TORCH_CHECK(problem_sizes.size(1) == 3,
23
- "problem_sizes must have shape (num_experts, 3)");
24
- STD_TORCH_CHECK(
25
- problem_sizes.size(0) == expert_offsets.size(0),
26
- "Number of experts in problem_sizes must match expert_offsets");
27
- STD_TORCH_CHECK(
28
- problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
29
- "problem_sizes must be int32");
30
- STD_TORCH_CHECK(
31
- expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
32
- "expert_offsets must be int32");
33
- STD_TORCH_CHECK(
34
- blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
35
- "blockscale_offsets must be int32");
36
- STD_TORCH_CHECK(a.dim() == 2,
37
- "a must be a 2D tensor of shape (num_tokens, k)");
38
- STD_TORCH_CHECK(b.dim() == 3,
39
- "b must be a 3D tensor of shape (num_experts, k, n)");
40
- STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
41
- "k should align 128");
42
- STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
43
- STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major");
44
- STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major");
45
-
46
- const torch::stable::accelerator::DeviceGuard device_guard(
47
- a.get_device_index());
48
- auto stream = get_current_cuda_stream(a.get_device_index());
49
- if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
50
- expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
51
- cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
52
- blockscale_offsets, stream);
53
- } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) {
54
- expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
55
- cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
56
- blockscale_offsets, stream);
57
- } else {
58
- STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
59
- }
60
-#else
61
- STD_TORCH_CHECK(false,
62
- "No implemented cutlass_mxfp8_grouped_mm for "
63
- "current device");
64
-#endif
65
-}
66
-
67
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
68
- m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm));
69
-}
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh
deleted
-141
@@ -1,141 +0,0 @@
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/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh
deleted
-198
@@ -1,198 +0,0 @@
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
-
8
-#include <torch/csrc/stable/tensor.h>
9
-#include <torch/headeronly/util/Exception.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
-#include "libtorch_stable/torch_utils.h"
19
-
20
-namespace expert_specialization {
21
-
22
-template <typename GemmTraits>
23
-void cutlass_mxfp8_grouped_mm_pre_compute(
24
- torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs,
25
- torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs,
26
- torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a,
27
- torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d,
28
- torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb,
29
- const torch::stable::Tensor& a, const torch::stable::Tensor& b,
30
- const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
31
- const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
32
- const torch::stable::Tensor& expert_offsets,
33
- const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
34
- using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
35
- using ElementA = typename OffsetFunctor::ElementA;
36
- using ElementB = typename OffsetFunctor::ElementB;
37
- using ElementSF = typename OffsetFunctor::ElementSF;
38
- using ElementD = typename OffsetFunctor::ElementD;
39
-
40
- using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor<GemmTraits>;
41
- using LayoutSFA = typename LayoutFunctor::LayoutSFA;
42
- using LayoutSFB = typename LayoutFunctor::LayoutSFB;
43
-
44
- using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor<GemmTraits>;
45
- using StrideA = typename StrideFunctor::StrideA;
46
- using StrideB = typename StrideFunctor::StrideB;
47
- using StrideD = typename StrideFunctor::StrideD;
48
-
49
- int num_experts = static_cast<int>(expert_offsets.size(0));
50
- STD_TORCH_CHECK(num_experts <= 1024,
51
- "Number of experts cannot exceed 1024, the maximum number of "
52
- "threads per block.");
53
-
54
- OffsetFunctor offset_functor(
55
- reinterpret_cast<int*>(expert_offsets.data_ptr()),
56
- reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
57
- reinterpret_cast<ElementA*>(a.data_ptr()),
58
- reinterpret_cast<ElementB*>(b.data_ptr()),
59
- reinterpret_cast<ElementSF*>(sfa.data_ptr()),
60
- reinterpret_cast<ElementSF*>(sfb.data_ptr()),
61
- reinterpret_cast<ElementD*>(d.data_ptr()),
62
- reinterpret_cast<ElementA**>(a_ptrs.data_ptr()),
63
- reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
64
- reinterpret_cast<ElementSF**>(sfa_ptrs.data_ptr()),
65
- reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
66
- reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
67
- LayoutFunctor layout_functor(
68
- reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
69
- reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()));
70
- StrideFunctor stride_functor(reinterpret_cast<StrideA*>(stride_a.data_ptr()),
71
- reinterpret_cast<StrideB*>(stride_b.data_ptr()),
72
- reinterpret_cast<StrideD*>(stride_d.data_ptr()));
73
- cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>(
74
- static_cast<int*>(problem_sizes.data_ptr()), offset_functor,
75
- layout_functor, stride_functor);
76
-}
77
-
78
-template <typename GemmTraits>
79
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
80
- const torch::stable::Tensor& b_ptrs,
81
- const torch::stable::Tensor& sfa_ptrs,
82
- const torch::stable::Tensor& sfb_ptrs,
83
- const torch::stable::Tensor& d_ptrs,
84
- const torch::stable::Tensor& stride_a,
85
- const torch::stable::Tensor& stride_b,
86
- const torch::stable::Tensor& stride_d,
87
- const torch::stable::Tensor& layout_sfa,
88
- const torch::stable::Tensor& layout_sfb,
89
- const torch::stable::Tensor& problem_sizes,
90
- cudaStream_t stream) {
91
- using Gemm = typename GemmTraits::Gemm;
92
- using ElementA = typename Gemm::ElementA;
93
- using ElementB = typename Gemm::ElementB;
94
- using ElementSF = typename GemmTraits::ElementSF;
95
- using ElementD = typename GemmTraits::ElementOutput;
96
- using StrideA = typename GemmTraits::StrideA;
97
- using StrideB = typename GemmTraits::StrideB;
98
- using StrideD = typename GemmTraits::StrideD;
99
- using LayoutSFA = typename GemmTraits::LayoutSFA;
100
- using LayoutSFB = typename GemmTraits::LayoutSFB;
101
- using UnderlyingProblemShape =
102
- typename GemmTraits::ProblemShape::UnderlyingProblemShape;
103
-
104
- cutlass::KernelHardwareInfo hw_info;
105
- hw_info.device_id = d_ptrs.get_device_index();
106
- hw_info.sm_count = get_device_prop()->multiProcessorCount;
107
- hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
108
- hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
109
-
110
- int num_experts = static_cast<int>(problem_sizes.size(0));
111
-
112
- UnderlyingProblemShape* underlying_problem_shape =
113
- reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
114
-
115
- typename Gemm::Arguments arguments = {
116
- cutlass::gemm::GemmUniversalMode::kGrouped,
117
- {num_experts, underlying_problem_shape, nullptr},
118
- {reinterpret_cast<const ElementA**>(a_ptrs.data_ptr()),
119
- reinterpret_cast<StrideA*>(stride_a.data_ptr()),
120
- reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
121
- reinterpret_cast<StrideB*>(stride_b.data_ptr()),
122
- reinterpret_cast<const ElementSF**>(sfa_ptrs.data_ptr()),
123
- reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
124
- reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr()),
125
- reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())},
126
- {{},
127
- nullptr,
128
- nullptr,
129
- reinterpret_cast<ElementD**>(d_ptrs.data_ptr()),
130
- reinterpret_cast<StrideD*>(stride_d.data_ptr())},
131
- hw_info,
132
- {} // Scheduler
133
- };
134
-
135
- Gemm gemm;
136
-
137
- auto can_implement_status = gemm.can_implement(arguments);
138
- STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
139
- "Failed to implement GEMM");
140
-
141
- size_t workspace_size = gemm.get_workspace_size(arguments);
142
- torch::stable::Tensor workspace = torch::stable::empty(
143
- {static_cast<int64_t>(workspace_size)},
144
- torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
145
-
146
- auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
147
- STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
148
- "Failed to initialize GEMM");
149
-
150
- status = gemm.run(stream, nullptr, true); // Enable PDL
151
- STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
152
-}
153
-
154
-template <typename OutType>
155
-void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
156
- const torch::stable::Tensor& a, const torch::stable::Tensor& b,
157
- const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
158
- torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
159
- const torch::stable::Tensor& expert_offsets,
160
- const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
161
- int num_experts = static_cast<int>(problem_sizes.size(0));
162
- auto device = a.device();
163
-
164
- torch::stable::Tensor a_ptrs = torch::stable::empty(
165
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
166
- torch::stable::Tensor b_ptrs = torch::stable::empty(
167
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
168
- torch::stable::Tensor sfa_ptrs = torch::stable::empty(
169
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
170
- torch::stable::Tensor sfb_ptrs = torch::stable::empty(
171
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
172
- torch::stable::Tensor d_ptrs = torch::stable::empty(
173
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
174
-
175
- torch::stable::Tensor stride_a = torch::stable::empty(
176
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
177
- torch::stable::Tensor stride_b = torch::stable::empty(
178
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
179
- torch::stable::Tensor stride_d = torch::stable::empty(
180
- num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
181
- torch::stable::Tensor layout_sfa =
182
- torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
183
- std::nullopt, device);
184
- torch::stable::Tensor layout_sfb =
185
- torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
186
- std::nullopt, device);
187
-
188
- using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
189
- cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
190
- a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
191
- layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets,
192
- blockscale_offsets, stream);
193
- cutlass_mxfp8_grouped_mm<GemmTraits>(
194
- a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
195
- layout_sfa, layout_sfb, problem_sizes, stream);
196
-}
197
-
198
-} // namespace expert_specialization
csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh
deleted
-127
@@ -1,127 +0,0 @@
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/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu
deleted
-66
@@ -1,66 +0,0 @@
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/csrc/stable/library.h>
7
-#include <torch/csrc/stable/tensor.h>
8
-#include "libtorch_stable/torch_utils.h"
9
-
10
-#include "mxfp8_experts_quant.cuh"
11
-
12
-void mxfp8_experts_quant(const torch::stable::Tensor& input,
13
- const torch::stable::Tensor& problem_sizes,
14
- const torch::stable::Tensor& expert_offsets,
15
- const torch::stable::Tensor& blockscale_offsets,
16
- torch::stable::Tensor& quant_output,
17
- torch::stable::Tensor& scale_factor) {
18
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
19
- STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
20
- STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
21
- STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major");
22
- STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
23
- STD_TORCH_CHECK(
24
- problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
25
- "problem_sizes must be int32");
26
- STD_TORCH_CHECK(
27
- expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
28
- "expert_offsets must be int32");
29
- STD_TORCH_CHECK(
30
- blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
31
- "blockscale_offsets must be int32");
32
-
33
- auto groups = problem_sizes.size(0);
34
- STD_TORCH_CHECK(
35
- expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
36
- "expert_offsets must be 1D and have size equal to the number of groups");
37
- STD_TORCH_CHECK(
38
- blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
39
- "blockscale_offsets must be 1D and have size equal to the number of "
40
- "groups");
41
-
42
- const torch::stable::accelerator::DeviceGuard device_guard(
43
- input.get_device_index());
44
- if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
45
- expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
46
- input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
47
- scale_factor);
48
- } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
49
- expert_specialization::launch_mxfp8_experts_quant<__half>(
50
- input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
51
- scale_factor);
52
- } else {
53
- STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
54
- }
55
-#else
56
- STD_TORCH_CHECK(false,
57
- "No implemented mxfp8_experts_quant for "
58
- "current device");
59
-#endif
60
-}
61
-
62
-// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM
63
-// is applied only under COMPILE_LANGUAGE:CUDA.
64
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
65
- m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant));
66
-}
csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh
deleted
-416
@@ -1,416 +0,0 @@
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 <cuda.h>
8
-#include <cuda_bf16.h>
9
-#include <cuda_fp16.h>
10
-
11
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
12
-#include <torch/csrc/stable/macros.h>
13
-#include <torch/csrc/stable/tensor.h>
14
-#include <torch/headeronly/util/Exception.h>
15
-
16
-#include <cuda/ptx>
17
-
18
-#include "cute/tensor.hpp"
19
-#include "libtorch_stable/torch_utils.h"
20
-
21
-namespace expert_specialization {
22
-
23
-using namespace cute;
24
-
25
-constexpr uint32_t THREAD_BLOCK_SIZE = 128;
26
-constexpr uint32_t WARP_SIZE = 32;
27
-constexpr int BLOCK_M = 128;
28
-constexpr int BLOCK_K = 128;
29
-using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
30
-using ValLayout = Layout<Shape<_1, _16>>;
31
-using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
32
-using SfR2SValLayout = Layout<Shape<_1, _1>>;
33
-using ScaleFactorTileLayout =
34
- Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
35
-
36
-// Fast reciprocal.
37
-inline __device__ float reciprocal_approximate_ftz(float a) {
38
- float b;
39
- asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
40
- return b;
41
-}
42
-
43
-// Some code references TRT-LLM:
44
-// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
45
-template <typename FragmentS, typename FragmentD>
46
-__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s,
47
- FragmentD& fragment_d) {
48
- using FragmentSLayout = typename FragmentS::layout_type;
49
- using FragmentDLayout = typename FragmentD::layout_type;
50
- FragmentSLayout fragment_s_layout;
51
- FragmentDLayout fragment_d_layout;
52
- static_assert(is_static<FragmentSLayout>::value &&
53
- size(fragment_s_layout) == 16);
54
- static_assert(is_static<FragmentDLayout>::value &&
55
- size(fragment_d_layout) == 16);
56
-
57
- constexpr int eles_per_thr = 16;
58
- using ValType = typename FragmentS::element_type;
59
- using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>,
60
- __nv_bfloat162, __half2>;
61
- VecType vec[8];
62
- // Assign vals
63
- vec[0].x = fragment_s(Int<0>{});
64
- vec[0].y = fragment_s(Int<1>{});
65
- vec[1].x = fragment_s(Int<2>{});
66
- vec[1].y = fragment_s(Int<3>{});
67
- vec[2].x = fragment_s(Int<4>{});
68
- vec[2].y = fragment_s(Int<5>{});
69
- vec[3].x = fragment_s(Int<6>{});
70
- vec[3].y = fragment_s(Int<7>{});
71
- vec[4].x = fragment_s(Int<8>{});
72
- vec[4].y = fragment_s(Int<9>{});
73
- vec[5].x = fragment_s(Int<10>{});
74
- vec[5].y = fragment_s(Int<11>{});
75
- vec[6].x = fragment_s(Int<12>{});
76
- vec[6].y = fragment_s(Int<13>{});
77
- vec[7].x = fragment_s(Int<14>{});
78
- vec[7].y = fragment_s(Int<15>{});
79
-
80
- auto local_max = __habs2(vec[0]);
81
- for (int i = 1; i < eles_per_thr / 2; i++) {
82
- local_max = __hmax2(__habs2(vec[i]), local_max);
83
- }
84
- local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
85
-
86
- // Get the final absolute maximum values.
87
- float block_max(0.0f);
88
- if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
89
- block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
90
- } else {
91
- block_max = __half2float(__hmax(local_max.x, local_max.y));
92
- }
93
- // Get the SF (max value of the vector / max value of mxfp8).
94
- float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
95
- // 8 bits representation of the SF.
96
- uint8_t fp8_sf_val;
97
-
98
- __nv_fp8_e8m0 tmp_sf_val;
99
- tmp_sf_val.__x =
100
- __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
101
- sf_val = static_cast<float>(tmp_sf_val);
102
- fp8_sf_val = tmp_sf_val.__x;
103
- // Get the output scale (reciprocal of the SFValue).
104
- float output_scale =
105
- block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
106
-
107
- // Convert the input to float.
108
- float2 fp2_vals[eles_per_thr / 2];
109
-
110
-#pragma unroll
111
- for (int i = 0; i < eles_per_thr / 2; i++) {
112
- if constexpr (std::is_same_v<ValType, __half>) {
113
- fp2_vals[i] = __half22float2(vec[i]);
114
- } else {
115
- fp2_vals[i] = __bfloat1622float2(vec[i]);
116
- }
117
- fp2_vals[i].x *= output_scale;
118
- fp2_vals[i].y *= output_scale;
119
- }
120
- union {
121
- uint8_t bytes[16];
122
- __nv_fp8x2_e4m3 elts[8];
123
- } u;
124
- u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
125
- u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
126
- u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
127
- u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
128
- u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
129
- u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
130
- u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
131
- u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
132
- fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
133
- fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
134
- fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
135
- fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
136
- fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
137
- fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
138
- fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
139
- fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
140
- fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
141
- fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
142
- fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
143
- fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
144
- fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
145
- fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
146
- fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
147
- fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
148
- return fp8_sf_val;
149
-}
150
-
151
-template <typename TensorS, typename TensorP, typename TensorD,
152
- typename TensorSharedSF, typename TensorSF, typename TiledCopyG2R,
153
- typename TiledCopyR2G, typename TiledCopyR2S>
154
-__inline__ __device__ void mxfp8_experts_quant_tile(
155
- TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d,
156
- TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m,
157
- TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g,
158
- TiledCopyR2S& tiled_copy_r2s) {
159
- static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 &&
160
- size(get<1>(typename TensorS::layout_type{})) == 128 &&
161
- stride(get<1>(typename TensorS::layout_type{})) == 1);
162
- static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 &&
163
- size(get<1>(typename TensorD::layout_type{})) == 128 &&
164
- stride(get<1>(typename TensorD::layout_type{})) == 1);
165
- static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 &&
166
- size(get<1>(typename TensorP::layout_type{})) == 128);
167
- static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
168
- size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
169
- static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 &&
170
- size(get<1>(typename TensorSF::layout_type{})) == 4);
171
-
172
- using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
173
- auto tiler_mn = Tiler_MN{};
174
- static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
175
-
176
- auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
177
- auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
178
- auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
179
- static_assert(size<2>(tiled_tensor_s) == 1);
180
- static_assert(size<2>(tiled_tensor_p) == 1);
181
- static_assert(size<2>(tiled_tensor_d) == 1);
182
- auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
183
- auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
184
- auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
185
-
186
- using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
187
- auto sf_tiler_mn = SF_Tiler_MN{};
188
- static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
189
-
190
- auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
191
- auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
192
- auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
193
- auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
194
-
195
- constexpr int tile_loop_count = size<1>(tiled_tensor_s);
196
- constexpr int rows_in_tile = 16;
197
- // We don't need to clear shared memory
198
- // clear(squeeze_tiled_tensor_shared_sf);
199
-#pragma unroll 4
200
- for (int t = 0; t < tile_loop_count; t++) {
201
- if (t * rows_in_tile >= m) {
202
- break;
203
- }
204
- auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
205
- auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
206
- auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
207
- auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
208
- auto current_copy_tile_shared_sf =
209
- tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
210
-
211
- // Global to Register copy
212
- auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
213
- auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
214
- auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
215
- auto input_fragment = make_fragment_like(thr_tile_g2r_s);
216
-
217
- // Register to Global copy
218
- auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
219
- auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
220
- auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
221
- auto output_fragment = make_fragment_like(thr_tile_r2g_d);
222
-
223
- // Register to Shared copy
224
- auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
225
- auto thr_tile_r2s_shared_sf =
226
- thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
227
- auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
228
-
229
- // CopyG2R & convert & CopyR2G
230
- copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
231
- uint8_t fp8_sf_val =
232
- cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
233
- copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
234
- shared_sf_fragment[0] = fp8_sf_val;
235
-
236
- // Before first copy r2s, clear shared memory and wait previous group
237
- if (t == 0 && threadIdx.x == 0) {
238
- // Wait for the group to have completed reading from shared memory.
239
- cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
240
- }
241
- __syncthreads();
242
-
243
- if (threadIdx.x % 2 == 0) {
244
- copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
245
- }
246
- __syncthreads();
247
- }
248
-
249
- // Wait for shared memory writes to be visible to TMA engine.
250
- cuda::ptx::fence_proxy_async(cuda::ptx::space_shared); // b)
251
- __syncthreads();
252
-
253
- if (threadIdx.x == 0) {
254
- cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared,
255
- squeeze_tiled_tensor_sf.data().get(),
256
- squeeze_tiled_tensor_shared_sf.data().get(), 512);
257
- // Wait for TMA transfer to have finished reading shared memory.
258
- // Create a "bulk async-group" out of the previous bulk copy operation.
259
- cuda::ptx::cp_async_bulk_commit_group();
260
- }
261
- __syncthreads();
262
-}
263
-
264
-template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G,
265
- typename TiledCopyR2S>
266
-__global__ void mxfp8_experts_quant_kernel(
267
- const T_IN* input, const int* problem_sizes, const int* expert_offsets,
268
- const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output,
269
- uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r,
270
- TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) {
271
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
272
- __shared__ __align__(512) uint8_t shared_memory[512];
273
- ScaleFactorTileLayout scale_factor_tile_layout{};
274
- auto scale_factor_shared =
275
- make_tensor(make_smem_ptr(shared_memory),
276
- scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1)
277
- // TODO: Transform Groupwise Schedule into a more efficient Schedule
278
- for (int g = 0; g < groups; g++) {
279
- int m = problem_sizes[g * 3 + 0];
280
- int k = problem_sizes[g * 3 + 2];
281
- int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
282
- int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
283
-
284
- auto input_tensor = make_tensor(
285
- make_gmem_ptr(input + expert_offset * k),
286
- make_layout(make_shape(m, k),
287
- LayoutRight{})); // (M, K):(K, 1) half_t/bfloat16_t
288
-
289
- auto quant_output_tensor = make_tensor(
290
- make_gmem_ptr(quant_output + expert_offset * k),
291
- make_layout(make_shape(m, k),
292
- LayoutRight{})); // (M, K):(K, 1) cutlass::float_e4m3_t
293
-
294
- auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
295
- auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout,
296
- scale_factor_shape, LayoutRight{});
297
- // layout<0>(layout<0>(scale_factor_layout)) (_32,_4):(_16,_4) -- static
298
- // layout<1>(layout<0>(scale_factor_layout)) M_align_128 / 128 -- dynamic
299
- // shape dynamic stride layout<0>(layout<1>(scale_factor_layout)) _4:_1 --
300
- // static layout<1>(layout<1>(scale_factor_layout)) (K / 32) / 4 : _512 --
301
- // dynamic shape static stride
302
-
303
- // Reshape to zipped layout for 1D indexing
304
- auto zipped_scale_factor_layout = make_layout(
305
- make_layout(layout<0>(layout<0>(scale_factor_layout)),
306
- layout<0>(layout<1>(scale_factor_layout))),
307
- make_layout(
308
- layout<1>(layout<0>(scale_factor_layout)),
309
- layout<1>(layout<1>(
310
- scale_factor_layout)))); // (((_32,_4),_4),(M_align_128 /
311
- // 128,(K / 32) /
312
- // 4)):(((_16,_4),_1),(?,_512))
313
-
314
- auto scale_factor_tensor =
315
- make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)),
316
- zipped_scale_factor_layout);
317
-
318
- // Used for cases where M is not divisible by 128 (most scenarios).
319
- auto input_shape = shape(input_tensor); // (M, K):(K, 1)
320
- auto identity_tensor = make_identity_tensor(input_shape);
321
- auto predict_tensor = cute::lazy::transform(
322
- identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
323
-
324
- // (_128, _128)
325
- auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
326
-
327
- auto tiled_input_tensor = zipped_divide(
328
- input_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
329
- auto tiled_quant_output_tensor =
330
- zipped_divide(quant_output_tensor,
331
- tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
332
- auto tiled_predict_tensor = zipped_divide(
333
- predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
334
-
335
- auto total_tiles =
336
- size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128)
337
- decltype(total_tiles) blk_offset = blockIdx.x;
338
- while (blk_offset < total_tiles) {
339
- auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
340
- auto current_quant_output_tile =
341
- tensor<0>(tiled_quant_output_tensor(_, blk_offset));
342
- auto current_predict_tile =
343
- tensor<0>(tiled_predict_tensor(_, blk_offset));
344
- auto current_scale_factor_tile =
345
- tensor<0>(scale_factor_tensor(_, blk_offset));
346
-
347
- mxfp8_experts_quant_tile<
348
- decltype(current_input_tile), decltype(current_predict_tile),
349
- decltype(current_quant_output_tile), decltype(scale_factor_shared),
350
- decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G,
351
- TiledCopyR2S>(current_input_tile, current_predict_tile,
352
- current_quant_output_tile, scale_factor_shared,
353
- current_scale_factor_tile, m, tiled_copy_g2r,
354
- tiled_copy_r2g, tiled_copy_r2s);
355
- blk_offset += gridDim.x;
356
- }
357
- }
358
-#endif
359
-}
360
-
361
-template <typename T_IN>
362
-void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
363
- const torch::stable::Tensor& problem_sizes,
364
- const torch::stable::Tensor& expert_offsets,
365
- const torch::stable::Tensor& blockscale_offsets,
366
- torch::stable::Tensor& quant_output,
367
- torch::stable::Tensor& scale_factor) {
368
- ThrLayout thr_layout{};
369
- ValLayout val_layout{};
370
- SfR2SThrLayout r2s_thr_layout{};
371
- SfR2SValLayout r2s_val_layout{};
372
-
373
- using CopyOpG2R =
374
- UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
375
- using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
376
- auto tiled_copy_g2r = cute::make_tiled_copy(
377
- CopyAtomG2R{}, thr_layout, val_layout); // Tiler_MN: (16, 128)
378
-
379
- using CopyOpR2G = UniversalCopy<
380
- cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
381
- using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
382
- auto tiled_copy_r2g = cute::make_tiled_copy(
383
- CopyAtomR2G{}, thr_layout, val_layout); // Tiler_MN: (16, 128)
384
-
385
- using CopyOpR2S =
386
- UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
387
- using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
388
- auto tiled_copy_r2s = cute::make_tiled_copy(
389
- CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4)
390
-
391
- int max_active_blocks_per_sm = -1;
392
- STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
393
- &max_active_blocks_per_sm,
394
- mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
395
- decltype(tiled_copy_r2g),
396
- decltype(tiled_copy_r2s)>,
397
- THREAD_BLOCK_SIZE, 0));
398
-
399
- dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
400
- 1, 1);
401
- dim3 block(THREAD_BLOCK_SIZE, 1, 1);
402
- int num_experts = static_cast<int>(problem_sizes.size(0));
403
- auto stream = get_current_cuda_stream(input.get_device_index());
404
- mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
405
- decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
406
- <<<grid, block, 0, stream>>>(
407
- reinterpret_cast<const T_IN*>(input.data_ptr()),
408
- reinterpret_cast<const int*>(problem_sizes.data_ptr()),
409
- reinterpret_cast<const int*>(expert_offsets.data_ptr()),
410
- reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
411
- reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
412
- reinterpret_cast<uint8_t*>(scale_factor.data_ptr()), num_experts,
413
- tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s);
414
-}
415
-
416
-} // namespace expert_specialization
\ No newline at end of file
csrc/libtorch_stable/torch_bindings.cpp
-16
@@ -308,22 +308,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
308
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
309
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
310
311
- // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
312
- ops.def(
313
- "mxfp8_experts_quant("
314
- " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
315
- " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
316
- " -> ()");
317
- // conditionally compiled so impl registration is in source file
318
-
319
- // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
320
- ops.def(
321
- "cutlass_mxfp8_grouped_mm("
322
- " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
323
- " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
324
- " -> ()");
325
- // conditionally compiled so impl registration is in source file
326
-
311
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
312
// conditionally compiled so impl registration is in source file
313
ops.def(
tests/kernels/moe/test_cutlass_mxfp8_grouped_mm.py
deleted
-237
@@ -1,237 +0,0 @@
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
@@ -1136,76 +1136,6 @@ def cutlass_mxfp4_moe_mm(
1136
)
1137
1138
1139
-def mxfp8_experts_quant(
1140
- input_tensor: torch.Tensor,
1141
- problem_sizes: torch.Tensor,
1142
- expert_offsets: torch.Tensor,
1143
- blockscale_offsets: torch.Tensor,
1144
- quant_output: torch.Tensor,
1145
- scale_factor: torch.Tensor,
1146
-) -> None:
1147
- torch.ops._C.mxfp8_experts_quant(
1148
- input_tensor,
1149
- problem_sizes,
1150
- expert_offsets,
1151
- blockscale_offsets,
1152
- quant_output,
1153
- scale_factor,
1154
- )
1155
-
1156
-
1157
-def cutlass_mxfp8_grouped_mm(
1158
- a_tensors: torch.Tensor,
1159
- b_tensors: torch.Tensor,
1160
- a_scales: torch.Tensor,
1161
- b_scales: torch.Tensor,
1162
- out_tensors: torch.Tensor,
1163
- problem_sizes: torch.Tensor,
1164
- expert_offsets: torch.Tensor,
1165
- blockscale_offsets: torch.Tensor,
1166
-) -> None:
1167
- torch.ops._C.cutlass_mxfp8_grouped_mm(
1168
- a_tensors,
1169
- b_tensors,
1170
- a_scales,
1171
- b_scales,
1172
- out_tensors,
1173
- problem_sizes,
1174
- expert_offsets,
1175
- blockscale_offsets,
1176
- )
1177
-
1178
-
1179
-if hasattr(torch.ops._C, "mxfp8_experts_quant"):
1180
-
1181
- @register_fake("_C::mxfp8_experts_quant")
1182
- def _mxfp8_experts_quant_fake(
1183
- input_tensor: torch.Tensor,
1184
- problem_sizes: torch.Tensor,
1185
- expert_offsets: torch.Tensor,
1186
- blockscale_offsets: torch.Tensor,
1187
- quant_output: torch.Tensor,
1188
- scale_factor: torch.Tensor,
1189
- ) -> None:
1190
- return None
1191
-
1192
-
1193
-if hasattr(torch.ops._C, "cutlass_mxfp8_grouped_mm"):
1194
-
1195
- @register_fake("_C::cutlass_mxfp8_grouped_mm")
1196
- def _cutlass_mxfp8_grouped_mm_fake(
1197
- a_tensors: torch.Tensor,
1198
- b_tensors: torch.Tensor,
1199
- a_scales: torch.Tensor,
1200
- b_scales: torch.Tensor,
1201
- out_tensors: torch.Tensor,
1202
- problem_sizes: torch.Tensor,
1203
- expert_offsets: torch.Tensor,
1204
- blockscale_offsets: torch.Tensor,
1205
- ) -> None:
1206
- return None
1207
-
1208
-
1139
# gptq_marlin
1140
def gptq_marlin_repack(
1141
b_q_weight: torch.Tensor,