@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 84166fee97
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

[Kernel] Integrate CUTLASS MoE kernel with PPLX (#18762)

Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

ElizaWszola committed Jun 7, 2025 at 03:26 UTC 84166fee9770e6fba71a96978b3e7d149392fb28
26 files changed +925 -416
CMakeLists.txt
+2 -2
@@ -543,8 +543,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
543 # CUTLASS MoE kernels
544
545 # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and only works
546 - # on Hopper). get_cutlass_moe_mm_data should only be compiled if it's possible
547 - # to compile MoE kernels that use its output.
546 + # on Hopper). get_cutlass_(pplx_)moe_mm_data should only be compiled
547 + # if it's possible to compile MoE kernels that use its output.
548 cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;10.0a" "${CUDA_ARCHS}")
549 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
550 set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu"
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+11 -50
@@ -7,8 +7,8 @@ from benchmark_shapes import WEIGHT_SHAPES_MOE
7
8 from vllm import _custom_ops as ops
9 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
10 +from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
11 from vllm.model_executor.layers.fused_moe.fused_moe import (
11 - cutlass_moe_fp8,
12 fused_experts,
13 fused_topk,
14 )
@@ -70,18 +70,9 @@ def bench_run(
70 w1_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32)
71 w2_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32)
72
73 - ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
74 - c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)
75 - ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)
76 - c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
77 -
73 for expert in range(num_experts):
74 w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(w1[expert])
75 w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(w2[expert])
81 - w1_q_notransp = w1_q.clone()
82 - w2_q_notransp = w2_q.clone()
83 - w1_q = w1_q.transpose(1, 2)
84 - w2_q = w2_q.transpose(1, 2)
76
77 score = torch.randn((m, num_experts), device="cuda", dtype=dtype)
78
@@ -122,10 +113,6 @@ def bench_run(
113 w2_scale: torch.Tensor,
114 topk_weights: torch.Tensor,
115 topk_ids: torch.Tensor,
125 - ab_strides1: torch.Tensor,
126 - c_strides1: torch.Tensor,
127 - ab_strides2: torch.Tensor,
128 - c_strides2: torch.Tensor,
116 num_repeats: int,
117 ):
118 for _ in range(num_repeats):
@@ -133,14 +120,10 @@ def bench_run(
120 a,
121 w1,
122 w2,
136 - w1_scale,
137 - w2_scale,
123 topk_weights,
124 topk_ids,
140 - ab_strides1,
141 - c_strides1,
142 - ab_strides2,
143 - c_strides2,
125 + w1_scale,
126 + w2_scale,
127 a1_scale=a_scale,
128 )
129
@@ -153,10 +136,6 @@ def bench_run(
136 w2_scale: torch.Tensor,
137 topk_weights: torch.Tensor,
138 topk_ids: torch.Tensor,
156 - ab_strides1: torch.Tensor,
157 - c_strides1: torch.Tensor,
158 - ab_strides2: torch.Tensor,
159 - c_strides2: torch.Tensor,
139 ):
140 with set_current_vllm_config(
141 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
@@ -165,14 +144,10 @@ def bench_run(
144 a,
145 w1_q,
146 w2_q,
168 - w1_scale,
169 - w2_scale,
147 topk_weights,
148 topk_ids,
172 - ab_strides1,
173 - c_strides1,
174 - ab_strides2,
175 - c_strides2,
149 + w1_scale,
150 + w2_scale,
151 a1_scale=a_scale,
152 )
153
@@ -218,10 +193,6 @@ def bench_run(
193 w2_scale,
194 topk_weights,
195 topk_ids,
221 - ab_strides1,
222 - c_strides1,
223 - ab_strides2,
224 - c_strides2,
196 )
197 torch.cuda.synchronize()
198
@@ -230,8 +201,8 @@ def bench_run(
201 with torch.cuda.graph(triton_graph, stream=triton_stream):
202 run_triton_from_graph(
203 a,
233 - w1_q_notransp,
234 - w2_q_notransp,
204 + w1_q,
205 + w2_q,
206 topk_weights,
207 topk_ids,
208 w1_scale,
@@ -250,18 +221,12 @@ def bench_run(
221 "w2": w2,
222 "score": score,
223 "topk": topk,
253 - "w1_q_notransp": w1_q_notransp,
254 - "w2_q_notransp": w2_q_notransp,
224 # Cutlass params
225 "a_scale": a_scale,
226 "w1_q": w1_q,
227 "w2_q": w2_q,
228 "w1_scale": w1_scale,
229 "w2_scale": w2_scale,
261 - "ab_strides1": ab_strides1,
262 - "c_strides1": c_strides1,
263 - "ab_strides2": ab_strides2,
264 - "c_strides2": c_strides2,
230 # cuda graph params
231 "cutlass_graph": cutlass_graph,
232 "triton_graph": triton_graph,
@@ -279,8 +244,8 @@ def bench_run(
244 # Warmup
245 run_triton_moe(
246 a,
282 - w1_q_notransp,
283 - w2_q_notransp,
247 + w1_q,
248 + w2_q,
249 topk_weights,
250 topk_ids,
251 w1_scale,
@@ -291,7 +256,7 @@ def bench_run(
256
257 results.append(
258 benchmark.Timer(
294 - stmt="run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501
259 + stmt="run_triton_moe(a, w1_q, w2_q, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501
260 globals=globals,
261 label=label,
262 sub_label=sub_label,
@@ -322,16 +287,12 @@ def bench_run(
287 w2_scale,
288 topk_weights,
289 topk_ids,
325 - ab_strides1,
326 - c_strides1,
327 - ab_strides2,
328 - c_strides2,
290 num_warmup,
291 )
292
293 results.append(
294 benchmark.Timer(
334 - stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2, num_runs)", # noqa: E501
295 + stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, num_runs)", # noqa: E501
296 globals=globals,
297 label=label,
298 sub_label=sub_label,
csrc/ops.h
+10 -1
@@ -236,7 +236,8 @@ void cutlass_moe_mm(
236 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
237 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
238 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
239 - torch::Tensor const& b_strides, torch::Tensor const& c_strides);
239 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
240 + bool per_act_token, bool per_out_ch);
241
242 void cutlass_fp4_group_mm(
243 torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
@@ -251,6 +252,14 @@ void get_cutlass_moe_mm_data(
252 const int64_t num_experts, const int64_t n, const int64_t k,
253 const std::optional<torch::Tensor>& blockscale_offsets);
254
255 +void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
256 + torch::Tensor& problem_sizes1,
257 + torch::Tensor& problem_sizes2,
258 + const torch::Tensor& expert_num_tokens,
259 + const int64_t num_local_experts,
260 + const int64_t padded_m, const int64_t n,
261 + const int64_t k);
262 +
263 void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
264 torch::Tensor const& b,
265 torch::Tensor const& a_scales,
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu
+19 -10
@@ -84,7 +84,8 @@ void run_cutlass_moe_mm_sm90(
84 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
85 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
86 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
87 - torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
87 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
88 + bool per_act_token, bool per_out_ch) {
89 TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
90 TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
91 TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
@@ -113,19 +114,23 @@ void run_cutlass_moe_mm_sm90(
114 if (n >= 8192) {
115 cutlass_group_gemm_caller<Cutlass3xGemmN8192>(
116 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
116 - problem_sizes, a_strides, b_strides, c_strides);
117 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
118 + per_out_ch);
119 } else if (k >= 8192) {
120 cutlass_group_gemm_caller<Cutlass3xGemmK8192>(
121 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
120 - problem_sizes, a_strides, b_strides, c_strides);
122 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
123 + per_out_ch);
124 } else if (m <= 16) {
125 cutlass_group_gemm_caller<Cutlass3xGemmM16>(
126 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
124 - problem_sizes, a_strides, b_strides, c_strides);
127 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
128 + per_out_ch);
129 } else {
130 cutlass_group_gemm_caller<Cutlass3xGemmDefault>(
131 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
128 - problem_sizes, a_strides, b_strides, c_strides);
132 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
133 + per_out_ch);
134 }
135 }
136
@@ -134,15 +139,18 @@ void dispatch_moe_mm_sm90(
139 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
140 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
141 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
137 - torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
142 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
143 + bool per_act_token, bool per_out_ch) {
144 if (out_tensors.dtype() == torch::kBFloat16) {
145 run_cutlass_moe_mm_sm90<cutlass::float_e4m3_t, cutlass::bfloat16_t>(
146 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
141 - problem_sizes, a_strides, b_strides, c_strides);
147 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
148 + per_out_ch);
149 } else {
150 run_cutlass_moe_mm_sm90<cutlass::float_e4m3_t, cutlass::half_t>(
151 out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
145 - problem_sizes, a_strides, b_strides, c_strides);
152 + problem_sizes, a_strides, b_strides, c_strides, per_act_token,
153 + per_out_ch);
154 }
155 }
156
@@ -153,8 +161,9 @@ void cutlass_moe_mm_sm90(
161 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
162 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
163 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
156 - torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
164 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
165 + bool per_act_token, bool per_out_ch) {
166 dispatch_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
167 expert_offsets, problem_sizes, a_strides, b_strides,
159 - c_strides);
168 + c_strides, per_act_token, per_out_ch);
169 }
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cuh
+2 -4
@@ -76,7 +76,8 @@ void cutlass_group_gemm_caller(
76 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
77 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
78 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
79 - torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
79 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
80 + bool per_act_token, bool per_out_ch) {
81 using ElementAB = typename Gemm::ElementAB;
82 using ElementD = typename Gemm::ElementD;
83
@@ -84,9 +85,6 @@ void cutlass_group_gemm_caller(
85 int k_size = a_tensors.size(1);
86 int n_size = out_tensors.size(1);
87
87 - bool per_act_token = a_scales.numel() != 1;
88 - bool per_out_ch = b_scales.numel() != num_experts;
89 -
88 auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
89
90 auto options_int =
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+38 -4
@@ -7,7 +7,7 @@
7
8 constexpr uint64_t THREADS_PER_EXPERT = 512;
9
10 -__global__ void compute_problem_sizes(const int* __restrict__ topk_ids,
10 +__global__ void compute_problem_sizes(const uint32_t* __restrict__ topk_ids,
11 int32_t* problem_sizes1,
12 int32_t* problem_sizes2,
13 int32_t* atomic_buffer,
@@ -62,7 +62,7 @@ __global__ void compute_expert_blockscale_offsets(
62 }
63 }
64
65 -__global__ void compute_arg_sorts(const int* __restrict__ topk_ids,
65 +__global__ void compute_arg_sorts(const uint32_t* __restrict__ topk_ids,
66 const int32_t* __restrict__ expert_offsets,
67 int32_t* input_permutation,
68 int32_t* output_permutation,
@@ -103,7 +103,7 @@ void get_cutlass_moe_mm_data_caller(
103
104 int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
105 compute_problem_sizes<<<num_experts, num_threads, 0, stream>>>(
106 - static_cast<const int32_t*>(topk_ids.data_ptr()),
106 + static_cast<const uint32_t*>(topk_ids.data_ptr()),
107 static_cast<int32_t*>(problem_sizes1.data_ptr()),
108 static_cast<int32_t*>(problem_sizes2.data_ptr()),
109 static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n, k);
@@ -120,10 +120,44 @@ void get_cutlass_moe_mm_data_caller(
120 static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
121 }
122 compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
123 - static_cast<const int32_t*>(topk_ids.data_ptr()),
123 + static_cast<const uint32_t*>(topk_ids.data_ptr()),
124 static_cast<const int32_t*>(expert_offsets.data_ptr()),
125 static_cast<int32_t*>(input_permutation.data_ptr()),
126 static_cast<int32_t*>(output_permutation.data_ptr()),
127 static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
128 topk_ids.size(1));
129 }
130 +
131 +__global__ void compute_pplx_data(int32_t* expert_offsets,
132 + int32_t* problem_sizes1,
133 + int32_t* problem_sizes2,
134 + const int32_t* __restrict__ expert_num_tokens,
135 + const int padded_m, const int n,
136 + const int k) {
137 + int expert_idx = threadIdx.x;
138 +
139 + expert_offsets[expert_idx] = expert_idx * padded_m;
140 + problem_sizes1[expert_idx * 3] = expert_num_tokens[expert_idx];
141 + problem_sizes1[expert_idx * 3 + 1] = 2 * n;
142 + problem_sizes1[expert_idx * 3 + 2] = k;
143 + problem_sizes2[expert_idx * 3] = expert_num_tokens[expert_idx];
144 + problem_sizes2[expert_idx * 3 + 1] = k;
145 + problem_sizes2[expert_idx * 3 + 2] = n;
146 +}
147 +
148 +void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
149 + torch::Tensor& problem_sizes1,
150 + torch::Tensor& problem_sizes2,
151 + const torch::Tensor& expert_num_tokens,
152 + const int64_t num_local_experts,
153 + const int64_t padded_m,
154 + const int64_t n, const int64_t k) {
155 + auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index());
156 +
157 + compute_pplx_data<<<1, num_local_experts, 0, stream>>>(
158 + static_cast<int32_t*>(expert_offsets.data_ptr()),
159 + static_cast<int32_t*>(problem_sizes1.data_ptr()),
160 + static_cast<int32_t*>(problem_sizes2.data_ptr()),
161 + static_cast<const int32_t*>(expert_num_tokens.data_ptr()), padded_m, n,
162 + k);
163 +}
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+36 -3
@@ -36,7 +36,8 @@ void cutlass_moe_mm_sm90(
36 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
37 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
38 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
39 - torch::Tensor const& b_strides, torch::Tensor const& c_strides);
39 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
40 + bool per_act_token, bool per_out_ch);
41
42 #endif
43
@@ -56,6 +57,14 @@ void get_cutlass_moe_mm_data_caller(
57 torch::Tensor& input_permutation, torch::Tensor& output_permutation,
58 const int64_t num_experts, const int64_t n, const int64_t k,
59 const std::optional<torch::Tensor>& blockscale_offsets);
60 +
61 +void get_cutlass_pplx_moe_mm_data_caller(torch::Tensor& expert_offsets,
62 + torch::Tensor& problem_sizes1,
63 + torch::Tensor& problem_sizes2,
64 + const torch::Tensor& expert_num_tokens,
65 + const int64_t num_local_experts,
66 + const int64_t padded_m,
67 + const int64_t n, const int64_t k);
68 #endif
69
70 void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a,
@@ -207,12 +216,13 @@ void cutlass_moe_mm(
216 torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
217 torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
218 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
210 - torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
219 + torch::Tensor const& b_strides, torch::Tensor const& c_strides,
220 + bool per_act_token, bool per_out_ch) {
221 int32_t version_num = get_sm_version_num();
222 #if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
223 cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
224 expert_offsets, problem_sizes, a_strides, b_strides,
215 - c_strides);
225 + c_strides, per_act_token, per_out_ch);
226 return;
227 #endif
228 TORCH_CHECK_NOT_IMPLEMENTED(
@@ -245,6 +255,29 @@ void get_cutlass_moe_mm_data(
255 version_num, ". Required capability: 90");
256 }
257
258 +void get_cutlass_pplx_moe_mm_data(torch::Tensor& expert_offsets,
259 + torch::Tensor& problem_sizes1,
260 + torch::Tensor& problem_sizes2,
261 + const torch::Tensor& expert_num_tokens,
262 + const int64_t num_local_experts,
263 + const int64_t padded_m, const int64_t n,
264 + const int64_t k) {
265 + // This function currently gets compiled only if we have a valid cutlass moe
266 + // mm to run it for.
267 + int32_t version_num = get_sm_version_num();
268 +#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
269 + get_cutlass_pplx_moe_mm_data_caller(expert_offsets, problem_sizes1,
270 + problem_sizes2, expert_num_tokens,
271 + num_local_experts, padded_m, n, k);
272 + return;
273 +#endif
274 + TORCH_CHECK_NOT_IMPLEMENTED(
275 + false,
276 + "No compiled get_cutlass_pplx_moe_mm_data: no cutlass_scaled_mm kernel "
277 + "for CUDA device capability: ",
278 + version_num, ". Required capability: 90");
279 +}
280 +
281 void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
282 torch::Tensor const& b,
283 torch::Tensor const& a_scales,
csrc/torch_bindings.cpp
+18 -1
@@ -435,7 +435,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
435 "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, "
436 " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
437 " Tensor problem_sizes, Tensor a_strides, "
438 - " Tensor b_strides, Tensor c_strides) -> ()",
438 + " Tensor b_strides, Tensor c_strides, bool per_act_token, "
439 + " bool per_out_ch) -> ()",
440 {stride_tag});
441 ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);
442
@@ -454,6 +455,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
455 {stride_tag});
456 ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
457
458 + // A function that computes data required to run fused MoE with w8a8 grouped
459 + // GEMM and PPLX. It takes expert_num_tokens and non_zero_expert_idxs
460 + // as an input, and computes expert_offsets (token start indices of each
461 + // expert). In addition to this, it computes problem sizes for each expert's
462 + // multiplication used by the two mms called from fused MoE operation.
463 + ops.def(
464 + "get_cutlass_pplx_moe_mm_data(Tensor! expert_offsets, "
465 + " Tensor! problem_sizes1, "
466 + " Tensor! problem_sizes2, "
467 + " Tensor expert_num_tokens, "
468 + " int num_local_experts, int padded_m, "
469 + " int n, int k) -> ()",
470 + {stride_tag});
471 + ops.impl("get_cutlass_pplx_moe_mm_data", torch::kCUDA,
472 + &get_cutlass_pplx_moe_mm_data);
473 +
474 // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
475 ops.def(
476 "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
tests/kernels/moe/test_cutlass_moe.py
+2 -6
@@ -193,14 +193,10 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
193
194 kwargs = {
195 'a': moe_tensors.a,
196 - 'w1_q': moe_tensors.w1_q.transpose(1, 2), # type: ignore[union-attr]
197 - 'w2_q': moe_tensors.w2_q.transpose(1, 2), # type: ignore[union-attr]
196 + 'w1_q': moe_tensors.w1_q, # type: ignore[union-attr]
197 + 'w2_q': moe_tensors.w2_q, # type: ignore[union-attr]
198 'topk_weights': topk_weights,
199 'topk_ids': topk_ids,
200 - 'ab_strides1': moe_tensors.ab_strides1,
201 - 'c_strides1': moe_tensors.c_strides1,
202 - 'ab_strides2': moe_tensors.ab_strides2,
203 - 'c_strides2': moe_tensors.c_strides2,
200 'w1_scale': moe_tensors.w1_scale,
201 'w2_scale': moe_tensors.w2_scale,
202 'a1_scale': moe_tensors.a_scale
tests/kernels/moe/test_pplx_cutlass_moe.py new
+287
@@ -0,0 +1,287 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from tests.pplx_utils import ProcessGroupInfo, parallel_launch
8 +from vllm import _custom_ops as ops
9 +from vllm.config import VllmConfig, set_current_vllm_config
10 +from vllm.model_executor.layers.activation import SiluAndMul
11 +from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
12 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
13 +from vllm.model_executor.layers.fused_moe.modular_kernel import (
14 + FusedMoEModularKernel)
15 +from vllm.platforms import current_platform
16 +
17 +try:
18 + from pplx_kernels import AllToAll
19 + from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
20 + nvshmem_finalize, nvshmem_get_unique_id,
21 + nvshmem_init)
22 + has_pplx = True
23 +except ImportError:
24 + has_pplx = False
25 +
26 +requires_pplx = pytest.mark.skipif(
27 + not has_pplx,
28 + reason="Requires PPLX kernels",
29 +)
30 +
31 +NUM_EXPERTS = [40, 64]
32 +TOP_KS = [6, 8]
33 +
34 +
35 +def rank_chunk(num, r, w):
36 + rem = num % w
37 + return (num // w) + (1 if r < rem else 0)
38 +
39 +
40 +def chunk_by_rank(t, r, w):
41 + num = t.shape[0]
42 + chunk = rank_chunk(num, r, w)
43 + rem = num % w
44 + if rem == 0 or r < rem:
45 + return t[(r * chunk):(r + 1) * chunk].contiguous()
46 + else:
47 + long_chunks = (num // w + 1) * rem
48 + short_chunks = (r - rem) * chunk
49 + start = long_chunks + short_chunks
50 + return t[start:start + chunk].contiguous()
51 +
52 +
53 +def pplx_cutlass_moe(
54 + pgi: ProcessGroupInfo,
55 + dp_size: int,
56 + a: torch.Tensor,
57 + w1: torch.Tensor,
58 + w2: torch.Tensor,
59 + w1_scale: torch.Tensor,
60 + w2_scale: torch.Tensor,
61 + topk_weights: torch.Tensor,
62 + topk_ids: torch.Tensor,
63 + a1_scale: torch.Tensor,
64 + out_dtype,
65 + per_act_token: bool,
66 + per_out_ch: bool,
67 +):
68 + from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
69 + PplxPrepareAndFinalize)
70 + assert torch.cuda.current_device() == pgi.local_rank
71 +
72 + num_tokens, hidden_dim = a.shape
73 + num_experts = w1.shape[0]
74 + block_size = hidden_dim # TODO support more cases
75 + device = pgi.device
76 + rank = pgi.rank
77 + world_size = pgi.world_size
78 + rank_num_tokens = rank_chunk(num_tokens, rank, world_size)
79 + max_num_tokens = rank_chunk(num_tokens, 0, world_size)
80 + topk = topk_ids.shape[1]
81 +
82 + if block_size == hidden_dim:
83 + scale_elems = 4 # hack to circumvent pplx data format requirements
84 + else:
85 + scale_elems = (hidden_dim + block_size - 1) // block_size
86 +
87 + ata = AllToAll.internode(
88 + max_num_tokens=max_num_tokens,
89 + num_experts=num_experts,
90 + experts_per_token=topk,
91 + rank=rank,
92 + world_size=pgi.world_size,
93 + dp_size=dp_size,
94 + hidden_dim=hidden_dim,
95 + hidden_dim_bytes=hidden_dim, # because a.dtype.itemsize == 1
96 + hidden_dim_scale_bytes=scale_elems * torch.float32.itemsize,
97 + )
98 +
99 + w1 = w1.to(device)
100 + w2 = w2.to(device)
101 + w1_scale = w1_scale.to(device)
102 + w2_scale = w2_scale.to(device)
103 + a1_scale = a1_scale.to(device)
104 +
105 + prepare_finalize = PplxPrepareAndFinalize(
106 + ata,
107 + max_num_tokens,
108 + pgi.world_size,
109 + rank,
110 + dp_size,
111 + quant_dtype=torch.float8_e4m3fn,
112 + per_act_token=per_act_token,
113 + )
114 +
115 + experts = CutlassExpertsFp8((num_experts + world_size - 1) // world_size,
116 + out_dtype, per_act_token, per_out_ch)
117 +
118 + fused_cutlass_experts = FusedMoEModularKernel(
119 + prepare_finalize,
120 + experts,
121 + )
122 +
123 + a_chunk = chunk_by_rank(a, rank, world_size).to(device)
124 + chunk_topk_weight = chunk_by_rank(topk_weights, rank,
125 + world_size).to(device)
126 + chunk_topk_ids = chunk_by_rank(topk_ids, rank,
127 + world_size).to(torch.uint32).to(device)
128 +
129 + out = fused_cutlass_experts(
130 + a_chunk,
131 + chunk_by_rank(w1, rank, world_size),
132 + chunk_by_rank(w2, rank, world_size),
133 + chunk_topk_weight,
134 + chunk_topk_ids,
135 + global_num_experts=num_experts,
136 + expert_map=None, #TODO
137 + w1_scale=chunk_by_rank(w1_scale, rank, world_size),
138 + w2_scale=chunk_by_rank(w2_scale, rank, world_size),
139 + a1_scale=chunk_by_rank(a1_scale, rank, world_size)
140 + if per_act_token else a1_scale[rank])
141 +
142 + torch.cuda.synchronize()
143 +
144 + ata.destroy()
145 +
146 + return out[:rank_num_tokens]
147 +
148 +
149 +vllm_config = VllmConfig()
150 +vllm_config.scheduler_config.max_num_seqs = 128
151 +vllm_config.scheduler_config.max_model_len = 8192
152 +
153 +
154 +def torch_moe2(a, w1, w2, topk_weight, topk_ids):
155 + M, K = a.shape
156 + topk = topk_ids.shape[1]
157 + a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
158 + out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
159 + num_experts = w1.shape[0]
160 + for i in range(num_experts):
161 + mask = (topk_ids == i).view(-1)
162 + if mask.sum():
163 + out[mask] = SiluAndMul()(
164 + a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
165 +
166 + return (out.view(M, -1, w2.shape[1]) *
167 + topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
168 +
169 +
170 +def _pplx_moe(
171 + pgi: ProcessGroupInfo,
172 + dp_size: int,
173 + a: torch.Tensor,
174 + w1: torch.Tensor,
175 + w2: torch.Tensor,
176 + w1_scale: torch.Tensor,
177 + w2_scale: torch.Tensor,
178 + topk_weights: torch.Tensor,
179 + topk_ids: torch.Tensor,
180 + a1_scale: torch.Tensor,
181 + out_dtype,
182 + a_full: torch.Tensor,
183 + w1_full: torch.Tensor,
184 + w2_full: torch.Tensor,
185 + per_act_token: bool,
186 + per_out_ch: bool,
187 +):
188 + uid = nvshmem_get_unique_id(
189 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
190 + torch.distributed.broadcast(uid, src=0)
191 + nvshmem_init(uid, pgi.rank, pgi.world_size)
192 +
193 + with set_current_vllm_config(vllm_config):
194 + torch_output = torch_moe2(a_full, w1_full, w2_full, topk_weights,
195 + topk_ids)
196 + pplx_output = pplx_cutlass_moe(pgi, dp_size, a, w1, w2, w1_scale,
197 + w2_scale, topk_weights, topk_ids,
198 + a1_scale, out_dtype, per_act_token,
199 + per_out_ch)
200 +
201 + torch_output = chunk_by_rank(torch_output, pgi.rank,
202 + pgi.world_size).to(pplx_output.device)
203 +
204 + # Uncomment if more debugging is needed
205 + # print("PPLX OUT:", pplx_output)
206 + # print("TORCH OUT:", torch_output)
207 +
208 + torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0)
209 +
210 + nvshmem_finalize()
211 +
212 +
213 +@pytest.mark.parametrize("m", [2, 224])
214 +@pytest.mark.parametrize("n", [3072])
215 +@pytest.mark.parametrize("k", [1536])
216 +@pytest.mark.parametrize("e", NUM_EXPERTS)
217 +@pytest.mark.parametrize("topk", TOP_KS)
218 +@pytest.mark.parametrize("per_act_token", [True, False])
219 +@pytest.mark.parametrize("per_out_ch", [True, False])
220 +@pytest.mark.parametrize("world_dp_size", [[2, 1]]) #, [4, 2]])
221 +@pytest.mark.skipif(
222 + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
223 + current_platform.get_device_capability()),
224 + reason="Grouped gemm is not supported on this GPU type.")
225 +@requires_pplx
226 +def test_cutlass_moe_pplx(
227 + m: int,
228 + n: int,
229 + k: int,
230 + e: int,
231 + topk: int,
232 + per_act_token: bool,
233 + per_out_ch: bool,
234 + world_dp_size: tuple[int, int],
235 +):
236 + current_platform.seed_everything(7)
237 +
238 + with set_current_vllm_config(vllm_config):
239 +
240 + dtype = torch.half
241 +
242 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10.0
243 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10.0
244 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10.0
245 +
246 + n_b_scales = 2 * n if per_out_ch else 1
247 + k_b_scales = k if per_out_ch else 1
248 +
249 + w1_q = torch.empty((e, 2 * n, k),
250 + device="cuda",
251 + dtype=torch.float8_e4m3fn)
252 + w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
253 + w1_scale = torch.empty((e, n_b_scales, 1),
254 + device="cuda",
255 + dtype=torch.float32)
256 + w2_scale = torch.empty((e, k_b_scales, 1),
257 + device="cuda",
258 + dtype=torch.float32)
259 +
260 + for expert in range(e):
261 + w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
262 + w1[expert], use_per_token_if_dynamic=per_out_ch)
263 + w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
264 + w2[expert], use_per_token_if_dynamic=per_out_ch)
265 +
266 + w1_d = torch.empty_like(w1)
267 + w2_d = torch.empty_like(w2)
268 + for expert in range(e):
269 + w1_d[expert] = (w1_q[expert].float() * w1_scale[expert]).half()
270 + w2_d[expert] = (w2_q[expert].float() * w2_scale[expert]).half()
271 +
272 + score = torch.randn((m, e), device="cuda", dtype=dtype)
273 + topk_weights, topk_ids, _ = fused_topk(a,
274 + score,
275 + topk,
276 + renormalize=False)
277 +
278 + world_size, dp_size = world_dp_size
279 + a_scale1 = torch.randn(
280 + (m if per_act_token else 1, 1), device="cuda",
281 + dtype=torch.float32) / 10.0
282 + if not per_act_token:
283 + a_scale1 = a_scale1.repeat(world_size, 1)
284 +
285 + parallel_launch(world_size, _pplx_moe, dp_size, a, w1_q, w2_q,
286 + w1_scale, w2_scale, topk_weights, topk_ids, a_scale1,
287 + dtype, a, w1_d, w2_d, per_act_token, per_out_ch)
tests/kernels/moe/test_pplx_moe.py
+7 -124
@@ -4,10 +4,7 @@
4
5 Run `pytest tests/kernels/test_pplx_moe.py`.
6 """
7 -import dataclasses
8 -import os
9 -import traceback
10 -from typing import Callable, Optional
7 +from typing import Optional
8
9 import pytest
10 import torch
@@ -21,10 +18,7 @@ try:
18 except ImportError:
19 has_pplx = False
20
24 -from torch.multiprocessing import (
25 - spawn) # pyright: ignore[reportPrivateImportUsage]
26 -from typing_extensions import Concatenate, ParamSpec
27 -
21 +from tests.pplx_utils import ProcessGroupInfo, parallel_launch
22 from vllm.config import VllmConfig, set_current_vllm_config
23 from vllm.model_executor.layers.activation import SiluAndMul
24 from vllm.model_executor.layers.fused_moe import override_config
@@ -36,6 +30,11 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
30 FusedMoEModularKernel)
31 from vllm.platforms import current_platform
32
33 +requires_pplx = pytest.mark.skipif(
34 + not has_pplx,
35 + reason="Requires PPLX kernels",
36 +)
37 +
38 PPLX_PREPARE_COMBOS = [(4, 128, 128), (32, 1024, 512), (64, 1024, 512),
39 (222, 2048, 1024)]
40
@@ -57,122 +56,6 @@ vllm_config = VllmConfig()
56 vllm_config.scheduler_config.max_num_seqs = 128
57 vllm_config.scheduler_config.max_model_len = 8192
58
60 -P = ParamSpec("P")
61 -
62 -requires_pplx = pytest.mark.skipif(
63 - not has_pplx,
64 - reason="Requires PPLX kernels",
65 -)
66 -
67 -
68 -@dataclasses.dataclass
69 -class ProcessGroupInfo:
70 - world_size: int
71 - world_local_size: int
72 - rank: int
73 - node_rank: int
74 - local_rank: int
75 - device: torch.device
76 -
77 -
78 -def _worker_parallel_launch(
79 - local_rank: int,
80 - world_size: int,
81 - world_local_size: int,
82 - node_rank: int,
83 - init_method: str,
84 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
85 - *args: P.args,
86 - **kwargs: P.kwargs,
87 -) -> None:
88 - rank = node_rank * world_local_size + local_rank
89 - torch.cuda.set_device(local_rank)
90 - device = torch.device("cuda", local_rank)
91 - torch.distributed.init_process_group(
92 - backend="cpu:gloo,cuda:nccl",
93 - init_method=init_method,
94 - rank=rank,
95 - world_size=world_size,
96 - device_id=device,
97 - )
98 - barrier = torch.tensor([rank], device=device)
99 - torch.distributed.all_reduce(barrier)
100 -
101 - try:
102 - worker(
103 - ProcessGroupInfo(
104 - world_size=world_size,
105 - world_local_size=world_local_size,
106 - rank=rank,
107 - node_rank=node_rank,
108 - local_rank=local_rank,
109 - device=device,
110 - ),
111 - *args,
112 - **kwargs,
113 - )
114 - except Exception as ex:
115 - print(ex)
116 - traceback.print_exc()
117 - raise
118 - finally:
119 - torch.distributed.destroy_process_group()
120 -
121 -
122 -def parallel_launch(
123 - world_size: int,
124 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
125 - *args: P.args,
126 - **kwargs: P.kwargs,
127 -) -> None:
128 - assert not kwargs
129 - spawn(
130 - _worker_parallel_launch,
131 - args=(
132 - world_size,
133 - world_size,
134 - 0,
135 - "tcp://localhost:29500",
136 - worker,
137 - ) + args,
138 - nprocs=world_size,
139 - join=True,
140 - )
141 -
142 -
143 -def parallel_launch_from_env(
144 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
145 - *args: P.args,
146 - **kwargs: P.kwargs,
147 -) -> None:
148 - """
149 - Launches a worker function in parallel across all processes in the current
150 - environment. The environment must have the following variables set:
151 - - WORLD_SIZE: The total number of processes.
152 - - WORLD_LOCAL_SIZE: The number of processes on the current node.
153 - - NODE_RANK: The rank of the current
154 - - MASTER_ADDR: The address of the master process.
155 - - MASTER_PORT: The port of the master process.
156 - """
157 - assert not kwargs
158 - world_size = int(os.environ["WORLD_SIZE"])
159 - world_local_size = int(os.environ["WORLD_LOCAL_SIZE"])
160 - node_rank = int(os.environ["NODE_RANK"])
161 - assert "MASTER_ADDR" in os.environ
162 - assert "MASTER_PORT" in os.environ
163 - spawn(
164 - _worker_parallel_launch,
165 - args=(
166 - world_size,
167 - world_local_size,
168 - node_rank,
169 - "env://",
170 - worker,
171 - ) + args,
172 - nprocs=world_local_size,
173 - join=True,
174 - )
175 -
59
60 def torch_prepare(
61 a: torch.Tensor,
tests/kernels/quantization/test_cutlass_scaled_mm.py
+2 -1
@@ -632,7 +632,8 @@ def test_cutlass_fp8_group_gemm(num_experts: int, per_act_token: bool,
632 ops.cutlass_moe_mm(out_tensors_stacked, a_tensors_stacked,
633 b_tensors_stacked, a_scales_tensors_stacked,
634 b_scales_tensors_stacked, expert_offsets[:-1],
635 - problem_sizes, ab_strides, ab_strides, c_strides)
635 + problem_sizes, ab_strides, ab_strides, c_strides,
636 + per_act_token, per_out_ch)
637
638 # Validate each group's result against the baseline
639 for g in range(num_experts):
tests/pplx_utils.py new
+123
@@ -0,0 +1,123 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import dataclasses
5 +import os
6 +import traceback
7 +from typing import Callable
8 +
9 +import torch
10 +from torch.multiprocessing import (
11 + spawn) # pyright: ignore[reportPrivateImportUsage]
12 +from typing_extensions import Concatenate, ParamSpec
13 +
14 +P = ParamSpec("P")
15 +
16 +
17 +@dataclasses.dataclass
18 +class ProcessGroupInfo:
19 + world_size: int
20 + world_local_size: int
21 + rank: int
22 + node_rank: int
23 + local_rank: int
24 + device: torch.device
25 +
26 +
27 +def _worker_parallel_launch(
28 + local_rank: int,
29 + world_size: int,
30 + world_local_size: int,
31 + node_rank: int,
32 + init_method: str,
33 + worker: Callable[Concatenate[ProcessGroupInfo, P], None],
34 + *args: P.args,
35 + **kwargs: P.kwargs,
36 +) -> None:
37 + rank = node_rank * world_local_size + local_rank
38 + torch.cuda.set_device(local_rank)
39 + device = torch.device("cuda", local_rank)
40 + torch.distributed.init_process_group(
41 + backend="cpu:gloo,cuda:nccl",
42 + init_method=init_method,
43 + rank=rank,
44 + world_size=world_size,
45 + device_id=device,
46 + )
47 + barrier = torch.tensor([rank], device=device)
48 + torch.distributed.all_reduce(barrier)
49 +
50 + try:
51 + worker(
52 + ProcessGroupInfo(
53 + world_size=world_size,
54 + world_local_size=world_local_size,
55 + rank=rank,
56 + node_rank=node_rank,
57 + local_rank=local_rank,
58 + device=device,
59 + ),
60 + *args,
61 + **kwargs,
62 + )
63 + except Exception as ex:
64 + print(ex)
65 + traceback.print_exc()
66 + raise
67 + finally:
68 + torch.distributed.destroy_process_group()
69 +
70 +
71 +def parallel_launch(
72 + world_size: int,
73 + worker: Callable[Concatenate[ProcessGroupInfo, P], None],
74 + *args: P.args,
75 + **kwargs: P.kwargs,
76 +) -> None:
77 + assert not kwargs
78 + spawn(
79 + _worker_parallel_launch,
80 + args=(
81 + world_size,
82 + world_size,
83 + 0,
84 + "tcp://localhost:29500",
85 + worker,
86 + ) + args,
87 + nprocs=world_size,
88 + join=True,
89 + )
90 +
91 +
92 +def parallel_launch_from_env(
93 + worker: Callable[Concatenate[ProcessGroupInfo, P], None],
94 + *args: P.args,
95 + **kwargs: P.kwargs,
96 +) -> None:
97 + """
98 + Launches a worker function in parallel across all processes in the current
99 + environment. The environment must have the following variables set:
100 + - WORLD_SIZE: The total number of processes.
101 + - WORLD_LOCAL_SIZE: The number of processes on the current node.
102 + - NODE_RANK: The rank of the current
103 + - MASTER_ADDR: The address of the master process.
104 + - MASTER_PORT: The port of the master process.
105 + """
106 + assert not kwargs
107 + world_size = int(os.environ["WORLD_SIZE"])
108 + world_local_size = int(os.environ["WORLD_LOCAL_SIZE"])
109 + node_rank = int(os.environ["NODE_RANK"])
110 + assert "MASTER_ADDR" in os.environ
111 + assert "MASTER_PORT" in os.environ
112 + spawn(
113 + _worker_parallel_launch,
114 + args=(
115 + world_size,
116 + world_local_size,
117 + node_rank,
118 + "env://",
119 + worker,
120 + ) + args,
121 + nprocs=world_local_size,
122 + join=True,
123 + )
vllm/_custom_ops.py
+27 -2
@@ -899,11 +899,36 @@ def shuffle_rows(input_tensor: torch.Tensor, dst2src_map: torch.Tensor):
899 return output_tensor
900
901
902 +def get_cutlass_pplx_moe_mm_data(expert_offsets: torch.Tensor,
903 + problem_sizes1: torch.Tensor,
904 + problem_sizes2: torch.Tensor,
905 + expert_num_tokens: torch.Tensor,
906 + num_local_experts: int, padded_m: int, n: int,
907 + k: int):
908 + """
909 + Prepare data necessary to perform CUTLASS grouped matrix multiplications
910 + used in CUTLASS-based fused MoE.
911 +
912 + The function takes in expert_num_tokens (token count per expert) and
913 + non_zero_expert_idxs (consecutive indices of experts with non-zero token
914 + counts) and uses them to compute:
915 + - expert_offsets: Indices that mark at which token index each expert begins
916 + its computation.
917 + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
918 + multiplication in two grouped MMs used in
919 + the fused MoE operation.
920 + """
921 + return torch.ops._C.get_cutlass_pplx_moe_mm_data(
922 + expert_offsets, problem_sizes1, problem_sizes2, expert_num_tokens,
923 + num_local_experts, padded_m, n, k)
924 +
925 +
926 def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
927 b_tensors: torch.Tensor, a_scales: torch.Tensor,
928 b_scales: torch.Tensor, expert_offsets: torch.Tensor,
929 problem_sizes: torch.Tensor, a_strides: torch.Tensor,
906 - b_strides: torch.Tensor, c_strides: torch.Tensor):
930 + b_strides: torch.Tensor, c_strides: torch.Tensor,
931 + per_act_token: bool, per_out_ch: bool):
932 """
933 A single grouped matrix multiplication used in CUTLASS-based fused MoE.
934 The function executes fp8-quantized OUT = AB matrix multiplication.
@@ -918,7 +943,7 @@ def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
943 return torch.ops._C.cutlass_moe_mm(out_tensors, a_tensors, b_tensors,
944 a_scales, b_scales, expert_offsets,
945 problem_sizes, a_strides, b_strides,
921 - c_strides)
946 + c_strides, per_act_token, per_out_ch)
947
948
949 def cutlass_fp4_moe_mm(a_tensors: torch.Tensor, b_tensors: torch.Tensor,
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+1
@@ -39,6 +39,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
39 def workspace_shapes(
40 self,
41 a: torch.Tensor,
42 + aq: torch.Tensor,
43 M: int,
44 N: int,
45 K: int,
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+3 -2
@@ -67,6 +67,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
67 def workspace_shapes(
68 self,
69 a: torch.Tensor,
70 + aq: torch.Tensor,
71 M: int,
72 N: int,
73 K: int,
@@ -78,11 +79,11 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
79 # even if we fall back to triton later, e.g. if expert maps are set.
80 if self.allow_deep_gemm and self.batched_deep_gemm_experts is not None:
81 return self.batched_deep_gemm_experts.workspace_shapes(
81 - a, M, N, K, topk, num_experts)
82 + a, aq, M, N, K, topk, num_experts)
83 else:
84 assert self.batched_triton_experts is not None
85 return self.batched_triton_experts.workspace_shapes(
85 - a, M, N, K, topk, num_experts)
86 + a, aq, M, N, K, topk, num_experts)
87
88 def apply(
89 self,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+220 -143
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ CUTLASS based Fused MoE kernels."""
4 -from typing import Optional
4 +from typing import Callable, Optional
5
6 import torch
7
@@ -13,110 +13,109 @@ from vllm.model_executor.layers.fused_moe.utils import _fp8_perm, _resize_cache
13 from vllm.scalar_type import scalar_types
14
15
16 -class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
17 -
18 - def __init__(
19 - self,
20 - ab_strides1: torch.Tensor,
21 - c_strides1: torch.Tensor,
22 - ab_strides2: torch.Tensor,
23 - c_strides2: torch.Tensor,
24 - out_dtype: torch.dtype,
25 - ):
26 - super().__init__()
27 - self.ab_strides1 = ab_strides1
28 - self.c_strides1 = c_strides1
29 - self.ab_strides2 = ab_strides2
30 - self.c_strides2 = c_strides2
31 - self.out_dtype = out_dtype
32 -
33 - def workspace_shapes(
34 - self,
35 - a: torch.Tensor,
36 - M: int,
37 - N: int,
38 - K: int,
39 - topk: int,
40 - num_experts: int,
41 - ) -> tuple[int, int, torch.dtype]:
42 - # Note that K, N are transposed
43 - N, K = K, N
44 - workspace1 = M * topk * max(2 * N, K)
45 - workspace2 = M * topk * N
46 - return (workspace1, workspace2, self.out_dtype)
47 -
48 - def apply(
49 - self,
50 - hidden_states: torch.Tensor,
51 - w1: torch.Tensor,
52 - w2: torch.Tensor,
53 - topk_ids: torch.Tensor,
54 - activation: str,
55 - global_num_experts: int,
56 - expert_map: Optional[torch.Tensor],
57 - w1_scale: Optional[torch.Tensor],
58 - w2_scale: Optional[torch.Tensor],
59 - w1_zp: Optional[torch.Tensor],
60 - w2_zp: Optional[torch.Tensor],
61 - a1q_scale: Optional[torch.Tensor],
62 - a2_scale: Optional[torch.Tensor],
63 - workspace13: torch.Tensor,
64 - workspace2: torch.Tensor,
65 - expert_num_tokens: Optional[torch.Tensor],
66 - ) -> torch.Tensor:
67 - a1q = hidden_states
68 -
69 - assert w1_scale is not None
70 - assert w2_scale is not None
71 - assert w1.dtype == torch.float8_e4m3fn
72 - assert w2.dtype == torch.float8_e4m3fn
73 - assert a1q.shape[1] == w1.shape[1], "Hidden size mismatch w1"
74 - assert w1.shape[2] == w2.shape[1] * 2, "Hidden size mismatch w2"
75 - assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
76 - assert a1q_scale is None or a1q_scale.dim(
77 - ) == 0 or a1q_scale.shape[0] == 1 or a1q_scale.shape[0] == a1q.shape[
78 - 0], "Input scale shape mismatch"
79 - assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
80 - 1] == w1.shape[2], "W1 scale shape mismatch"
81 - assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
82 - 1] == w2.shape[2], "W2 scale shape mismatch"
83 - assert w1.shape[0] == w2.shape[0], "Weights expert number mismatch"
84 - assert w1.shape[0] == w1_scale.shape[
85 - 0], "w1 scales expert number mismatch"
86 - assert w1.shape[0] == w2_scale.shape[
87 - 0], "w2 scales expert number mismatch"
88 - assert a2_scale is None or a1q_scale is None or a2_scale.shape == a1q_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
89 - assert self.ab_strides1.shape[0] == w1.shape[
90 - 0], "AB Strides 1 expert number mismatch"
91 - assert self.c_strides1.shape[0] == w1.shape[
92 - 0], "C Strides 1 expert number mismatch"
93 - assert self.ab_strides2.shape[0] == w2.shape[
94 - 0], "AB Strides 2 expert number mismatch"
95 - assert self.c_strides2.shape[0] == w2.shape[
96 - 0], "C Strides 2 expert number mismatch"
97 - assert self.out_dtype in [torch.half,
98 - torch.bfloat16], "Invalid output dtype"
99 -
100 - M = a1q.shape[0]
101 - _, N, K = w2.shape # because w1 + w2 are transposed
102 - device = a1q.device
103 -
104 - assert w1.shape[1] == K
105 - assert global_num_experts != -1
106 - assert a1q_scale is not None
107 -
108 - if expert_map is not None:
109 - "Translate info from expert_map to topk_ids"
110 - local_topk_ids = torch.where(expert_map[topk_ids] != -1,
111 - expert_map[topk_ids], -1)
112 - else:
113 - local_topk_ids = topk_ids
16 +def run_cutlass_moe_fp8(
17 + hidden_states: torch.Tensor,
18 + w1: torch.Tensor,
19 + w2: torch.Tensor,
20 + topk_ids: torch.Tensor,
21 + activation_callable: Callable,
22 + global_num_experts: int,
23 + expert_map: Optional[torch.Tensor],
24 + w1_scale: Optional[torch.Tensor],
25 + w2_scale: Optional[torch.Tensor],
26 + a1q_scale: Optional[torch.Tensor],
27 + a2_scale: Optional[torch.Tensor],
28 + workspace13: torch.Tensor,
29 + workspace2: torch.Tensor,
30 + expert_num_tokens: Optional[torch.Tensor],
31 + out_dtype: torch.dtype,
32 + per_act_token: bool,
33 + per_out_ch: bool,
34 +) -> torch.Tensor:
35 + a1q = hidden_states
36 +
37 + assert w1_scale is not None
38 + assert w2_scale is not None
39 + assert w1.dtype == torch.float8_e4m3fn
40 + assert w2.dtype == torch.float8_e4m3fn
41 + if expert_num_tokens is None:
42 + assert a1q.shape[1] == w1.shape[2], "Hidden size mismatch w1"
43 + else:
44 + assert a1q.shape[2] == w1.shape[2], "Hidden size mismatch w1"
45 + assert w1.shape[1] == w2.shape[2] * 2, "Hidden size mismatch w2"
46 + assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
47 + 1] == w1.shape[1], "W1 scale shape mismatch"
48 + assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
49 + 1] == w2.shape[1], "W2 scale shape mismatch"
50 + assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
51 + assert a1q_scale is None or a1q_scale.dim(
52 + ) == 0 or a1q_scale.shape[0] == 1 or a1q_scale.shape[0] == a1q.shape[
53 + 0], "Input scale shape mismatch"
54 + assert w1.shape[0] == w2.shape[0], "Weights expert number mismatch"
55 + assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
56 + assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
57 + assert a2_scale is None or a2_scale.dim(
58 + ) == 0 or a2_scale.shape[0] == 1 or a2_scale.shape[0] == a1q.shape[
59 + 0], "Intermediate scale shape mismatch"
60 + assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
61 + if expert_map is not None:
62 + assert expert_num_tokens is None
63 +
64 + # We have two modes: PPLX and non-PPLX. We differentiate them by checking
65 + # if expert_num_tokens is None (expert_num_tokens is a tensor which PPLX
66 + # uses to track the number of tokens per expert).
67 + # In the non-PPLX mode, the input tokens are not padded: thus, the shape
68 + # of the input is [total_num_tokens, hidden_size]. The input and output
69 + # require shuffling by a_map and c_map such that the tokens assigned to
70 + # each expert are contiguous.
71 + # In the PPLX mode, the input tokens are padded per expert to ensure that
72 + # the PPLX dispatch and combine functions work correctly: thus, the shape
73 + # of the input is [num_experts, max_num_tokens_per_expert, hidden_size].
74 + # The PPLX input and output require no shuffling by a_map and c_map since
75 + # their tokens are already contiguous for each expert as a result of
76 + # the dispatch function.
77 + is_pplx = expert_num_tokens is not None
78 +
79 + M = a1q.shape[0] # no pplx
80 + padded_M = a1q.shape[1] # pplx
81 + _, K, N = w2.shape
82 + device = a1q.device
83 +
84 + assert w1.shape[2] == K
85 + assert global_num_experts != -1
86 + assert a1q_scale is not None
87 +
88 + if expert_map is not None:
89 + "Translate info from expert_map to topk_ids"
90 + local_topk_ids = torch.where(expert_map[topk_ids] != -1,
91 + expert_map[topk_ids], -1)
92 + else:
93 + local_topk_ids = topk_ids
94 +
95 + topk = local_topk_ids.shape[1]
96 + local_E = w1.shape[0]
97 +
98 + if is_pplx:
99 + expert_offsets = torch.empty((local_E),
100 + dtype=torch.int32,
101 + device=device)
102 + problem_sizes1 = torch.empty((local_E, 3),
103 + dtype=torch.int32,
104 + device=device)
105 + problem_sizes2 = torch.empty((local_E, 3),
106 + dtype=torch.int32,
107 + device=device)
108
115 - topk = local_topk_ids.shape[1]
109 + ops.get_cutlass_pplx_moe_mm_data(expert_offsets, problem_sizes1,
110 + problem_sizes2, expert_num_tokens,
111 + local_E, padded_M, N, K)
112
117 - per_act_token = a1q_scale.numel() != 1 if a1q_scale is not None else (
118 - a2_scale.numel() != 1 if a2_scale is not None else False)
113 + w1_scale = w1_scale.reshape(w1_scale.shape[0], -1)
114 + w2_scale = w2_scale.reshape(w2_scale.shape[0], -1)
115 + a1q = a1q.reshape(-1, a1q.shape[2])
116 + a1q_scale = a1q_scale.reshape(-1, a1q_scale.shape[2]).contiguous()
117
118 + else:
119 expert_offsets = torch.empty((global_num_experts + 1),
120 dtype=torch.int32,
121 device=device)
@@ -149,50 +148,130 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
148
149 a1q = _fp8_perm(a1q, a_map)
150 a1q_scale = a1q_scale[a_map] if per_act_token else a1q_scale
152 -
151 + expert_offsets = expert_offsets[:-1]
152 +
153 + ab_strides1 = torch.full((w1.shape[0], ),
154 + K,
155 + device=device,
156 + dtype=torch.int64)
157 + c_strides1 = torch.full((w1.shape[0], ),
158 + 2 * N,
159 + device=device,
160 + dtype=torch.int64)
161 + ab_strides2 = torch.full((w1.shape[0], ),
162 + N,
163 + device=device,
164 + dtype=torch.int64)
165 + c_strides2 = torch.full((w1.shape[0], ),
166 + K,
167 + device=device,
168 + dtype=torch.int64)
169 +
170 + if is_pplx:
171 + c1 = _resize_cache(workspace13, (local_E * padded_M, N * 2))
172 + c2 = _resize_cache(workspace2, (local_E * padded_M, N))
173 + c3 = _resize_cache(workspace13, (local_E * padded_M, K))
174 + else:
175 c1 = _resize_cache(workspace13, (M * topk, N * 2))
176 c2 = _resize_cache(workspace2, (M * topk, N))
177 c3 = _resize_cache(workspace13, (M * topk, K))
178
157 - ops.cutlass_moe_mm(c1, a1q, w1, a1q_scale, w1_scale,
158 - expert_offsets[:-1], problem_sizes1,
159 - self.ab_strides1, self.ab_strides1, self.c_strides1)
179 + ops.cutlass_moe_mm(c1, a1q, w1, a1q_scale, w1_scale, expert_offsets,
180 + problem_sizes1, ab_strides1, ab_strides1, c_strides1,
181 + per_act_token, per_out_ch)
182
161 - self.activation(activation, c2, c1)
183 + activation_callable(c2, c1)
184
163 - a2q, a2q_scale = ops.scaled_fp8_quant(
164 - c2, a2_scale, use_per_token_if_dynamic=per_act_token)
185 + a2q, a2q_scale = ops.scaled_fp8_quant(
186 + c2, a2_scale, use_per_token_if_dynamic=per_act_token)
187
166 - if expert_map is not None:
167 - c3.fill_(0)
188 + if expert_map is not None:
189 + c3.fill_(0)
190 +
191 + ops.cutlass_moe_mm(c3, a2q, w2, a2q_scale, w2_scale, expert_offsets,
192 + problem_sizes2, ab_strides2, ab_strides2, c_strides2,
193 + per_act_token, per_out_ch)
194
169 - ops.cutlass_moe_mm(c3, a2q, w2, a2q_scale, w2_scale,
170 - expert_offsets[:-1], problem_sizes2,
171 - self.ab_strides2, self.ab_strides2, self.c_strides2)
195 + if is_pplx:
196 + return c3.reshape(local_E, padded_M, K)
197 + else:
198 + return c3[c_map].view(M, topk, K)
199
173 - c3 = c3[c_map]
200
175 - return c3
201 +class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
202 +
203 + def __init__(
204 + self,
205 + max_experts_per_worker: int,
206 + out_dtype: torch.dtype,
207 + per_act_token: bool,
208 + per_out_ch: bool,
209 + ):
210 + super().__init__()
211 + self.max_experts_per_worker = max_experts_per_worker
212 + self.out_dtype = out_dtype
213 + self.per_act_token = per_act_token
214 + self.per_out_ch = per_out_ch
215 +
216 + def workspace_shapes(
217 + self,
218 + a: torch.Tensor,
219 + aq: torch.Tensor,
220 + M: int,
221 + N: int,
222 + K: int,
223 + topk: int,
224 + num_experts: int,
225 + ) -> tuple[int, int, torch.dtype]:
226 + padded_M = aq.shape[1]
227 + workspace1 = self.max_experts_per_worker * padded_M * max(N, K)
228 + workspace2 = self.max_experts_per_worker * padded_M * (N // 2)
229 + return (workspace1, workspace2, self.out_dtype)
230 +
231 + def apply(
232 + self,
233 + hidden_states: torch.Tensor,
234 + w1: torch.Tensor,
235 + w2: torch.Tensor,
236 + topk_ids: torch.Tensor,
237 + activation: str,
238 + global_num_experts: int,
239 + expert_map: Optional[torch.Tensor],
240 + w1_scale: Optional[torch.Tensor],
241 + w2_scale: Optional[torch.Tensor],
242 + w1_zp: Optional[torch.Tensor],
243 + w2_zp: Optional[torch.Tensor],
244 + a1q_scale: Optional[torch.Tensor],
245 + a2_scale: Optional[torch.Tensor],
246 + workspace13: torch.Tensor,
247 + workspace2: torch.Tensor,
248 + expert_num_tokens: Optional[torch.Tensor],
249 + ) -> torch.Tensor:
250 + assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
251 + assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
252 + activation_callable = lambda i, o: self.activation(activation, i, o)
253 + return run_cutlass_moe_fp8(hidden_states, w1, w2, topk_ids,
254 + activation_callable, global_num_experts,
255 + expert_map, w1_scale, w2_scale, a1q_scale,
256 + a2_scale, workspace13, workspace2,
257 + expert_num_tokens, self.out_dtype,
258 + self.per_act_token, self.per_out_ch)
259
260
178 -#TODO make the grouped gemm kernel consistent with scaled gemm kernel
261 def cutlass_moe_fp8(
262 a: torch.Tensor,
263 w1_q: torch.Tensor,
264 w2_q: torch.Tensor,
183 - w1_scale: torch.Tensor,
184 - w2_scale: torch.Tensor,
265 topk_weights: torch.Tensor,
266 topk_ids: torch.Tensor,
187 - ab_strides1: torch.Tensor,
188 - c_strides1: torch.Tensor,
189 - ab_strides2: torch.Tensor,
190 - c_strides2: torch.Tensor,
267 + w1_scale: torch.Tensor,
268 + w2_scale: torch.Tensor,
269 + activation: str = "silu",
270 a1_scale: Optional[torch.Tensor] = None,
271 a2_scale: Optional[torch.Tensor] = None,
193 - out_dtype: torch.dtype = torch.half,
272 expert_map: Optional[torch.Tensor] = None,
273 apply_router_weight_on_input: bool = False,
274 + global_num_experts: int = -1,
275 ) -> torch.Tensor:
276 """
277 This function computes a a8w8-quantized Mixture of Experts (MoE) layer
@@ -207,25 +286,17 @@ def cutlass_moe_fp8(
286 Shape: [num_experts, K, 2N] (the weights are passed transposed)
287 - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
288 Shape: [num_experts, N, K] (the weights are passed transposed)
289 + - topk_weights (torch.Tensor): The weights of each token->expert mapping.
290 + - topk_ids (torch.Tensor): The token->expert mappings.
291 - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
292 Shape: [num_experts] or [num_experts, 2N]
293 - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
294 Shape: [num_experts] or [num_experts, K]
214 - - gating_output (torch.Tensor): The output of the gating operation
215 - (before softmax).
216 - - topk_weights (torch.Tensor): The weights of each token->expert mapping.
217 - - ab_strides1 (torch.Tensor): The input and weights strides of the first
218 - grouped gemm.
219 - - c_strides1 (torch.Tensor): The output strides of the first grouped gemm.
220 - - ab_strides2 (torch.Tensor): The input and weights strides of the second
221 - grouped gemm.
222 - - c_strides2 (torch.Tensor): The output strides of the second grouped gemm.
295 - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
296 Shape: scalar or [M]
297 - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
298 quantize the intermediate result between the gemms.
299 Shape: scalar or [M]
228 - - out_dtype (torch.dtype): The output tensor type.
300 - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
301 every Rank is responsible for a subset of experts. expert_map is a
302 mapping from global expert-id to local expert-id. When expert_map[i]
@@ -233,24 +304,27 @@ def cutlass_moe_fp8(
304 expert-id i.
305 - apply_router_weight_on_input (bool): When true, the topk weights are
306 applied directly on the inputs. This is only applicable when topk is 1.
307 + - global_num_experts (int): The total number of experts.
308
309 Returns:
310 - torch.Tensor: The fp16 output tensor after applying the MoE layer.
311 """
312 per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
313 a2_scale.numel() != 1 if a2_scale is not None else False)
314 + per_out_ch = w1_scale.numel() != w1_q.shape[0]
315 +
316 + out_dtype = a.dtype
317
318 fn = mk.FusedMoEModularKernel(
319 MoEPrepareAndFinalizeNoEP(
245 - per_channel_quant=per_act_token,
320 quant_dtype=torch.float8_e4m3fn,
321 + per_channel_quant=per_act_token,
322 ),
323 CutlassExpertsFp8(
249 - ab_strides1,
250 - c_strides1,
251 - ab_strides2,
252 - c_strides2,
253 - out_dtype,
324 + max_experts_per_worker=global_num_experts,
325 + out_dtype=out_dtype,
326 + per_act_token=per_act_token,
327 + per_out_ch=per_out_ch,
328 ),
329 )
330
@@ -260,9 +334,12 @@ def cutlass_moe_fp8(
334 w2_q,
335 topk_weights,
336 topk_ids,
263 - expert_map=expert_map,
264 - w1_scale=w1_scale,
265 - w2_scale=w2_scale,
337 + False,
338 + activation,
339 + global_num_experts if global_num_experts != -1 else w1_q.size(0),
340 + expert_map,
341 + w1_scale,
342 + w2_scale,
343 a1_scale=a1_scale,
344 a2_scale=a2_scale,
345 apply_router_weight_on_input=apply_router_weight_on_input,
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+1
@@ -73,6 +73,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
73 def workspace_shapes(
74 self,
75 a: torch.Tensor,
76 + aq: torch.Tensor,
77 M: int,
78 N: int,
79 K: int,
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+2
@@ -521,6 +521,7 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
521 def workspace_shapes(
522 self,
523 a: torch.Tensor,
524 + aq: torch.Tensor,
525 M: int,
526 N: int,
527 K: int,
@@ -632,6 +633,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
633 def workspace_shapes(
634 self,
635 a: torch.Tensor,
636 + aq: torch.Tensor,
637 M: int,
638 N: int,
639 K: int,
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
@@ -1545,6 +1545,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1545 def workspace_shapes(
1546 self,
1547 a: torch.Tensor,
1548 + aq: torch.Tensor,
1549 M: int,
1550 N: int,
1551 K: int,
vllm/model_executor/layers/fused_moe/layer.py
+47 -11
@@ -9,6 +9,9 @@ from typing import Callable, Optional, Union
9
10 import torch
11 import torch.nn.functional as F
12 +from compressed_tensors.quantization import (QuantizationArgs,
13 + QuantizationStrategy,
14 + QuantizationType)
15 from torch.nn.parameter import UninitializedParameter
16
17 import vllm.envs as envs
@@ -210,6 +213,7 @@ class MoEConfig:
213 moe_parallel_config: FusedMoEParallelConfig
214
215 in_dtype: torch.dtype # The activation type.
216 + quant_dtype: torch.dtype = None
217
218 # TODO: add more quantization params, blocked, per-token, etc.
219 block_size: int = 128
@@ -264,8 +268,22 @@ class FusedMoeWeightScaleSupported(Enum):
268 BLOCK = "block"
269
270
271 +def get_quant_config_input_activations(
272 + quant_config: Optional[QuantizationConfig]
273 +) -> Optional[QuantizationArgs]:
274 + if (quant_config is not None and hasattr(quant_config, 'target_scheme_map')
275 + and "Linear" in quant_config.target_scheme_map and
276 + "input_activations" in quant_config.target_scheme_map["Linear"]):
277 + return quant_config.target_scheme_map["Linear"].get(
278 + "input_activations")
279 + else:
280 + return None
281 +
282 +
283 class FusedMoEMethodBase(QuantizeMethodBase):
284
285 + moe: MoEConfig
286 +
287 @abstractmethod
288 def create_weights(self, layer: torch.nn.Module, num_experts: int,
289 hidden_size: int, intermediate_size_per_partition: int,
@@ -277,6 +295,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
295 all2all_manager = get_ep_group().device_communicator.all2all_manager
296 assert all2all_manager is not None
297
298 + self.moe = moe
299 quant_dtype = None
300 act_quant_block_size = None
301 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
@@ -297,13 +316,14 @@ class FusedMoEMethodBase(QuantizeMethodBase):
316 # dp_size actually means tp_size, bug in pplx kernels
317 dp_size=all2all_manager.tp_group.world_size,
318 hidden_dim=moe.hidden_dim,
300 - hidden_dim_bytes=moe.hidden_dim * moe.in_dtype.itemsize,
319 + hidden_dim_bytes=moe.hidden_dim * moe.quant_dtype.itemsize,
320 # For blocked per token: set to
321 # ceil_div(hidden_dim, block_size) * sizeof(float32)
322 # For per-token: set to sizeof(float32)
304 - hidden_dim_scale_bytes=(0 if moe.in_dtype.itemsize != 1 else (
305 - (moe.hidden_dim + moe.block_size - 1) // moe.block_size *
306 - torch.float32.itemsize)),
323 + hidden_dim_scale_bytes=(
324 + 0 if moe.quant_dtype.itemsize != 1 else
325 + ((moe.hidden_dim + moe.block_size - 1) // moe.block_size *
326 + torch.float32.itemsize)),
327 )
328
329 # Intranode pplx a2a takes a group name while internode does not.
@@ -313,6 +333,9 @@ class FusedMoEMethodBase(QuantizeMethodBase):
333
334 handle = all2all_manager.get_handle(all_to_all_args)
335
336 + input_activations = get_quant_config_input_activations(
337 + quant_config)
338 +
339 prepare_finalize = PplxPrepareAndFinalize(
340 handle,
341 max_num_tokens=moe.max_num_tokens,
@@ -320,7 +343,10 @@ class FusedMoEMethodBase(QuantizeMethodBase):
343 rank=all2all_manager.rank,
344 # dp_size actually means tp_size, bug in pplx kernels
345 dp_size=all2all_manager.tp_group.world_size,
323 - quant_dtype=moe.in_dtype,
346 + quant_dtype=moe.quant_dtype,
347 + per_act_token=(input_activations.strategy
348 + == QuantizationStrategy.TOKEN
349 + if input_activations is not None else False),
350 )
351 elif moe.use_deepep_ht_kernels:
352 assert moe.dp_size == all2all_manager.dp_world_size
@@ -365,15 +391,15 @@ class FusedMoEMethodBase(QuantizeMethodBase):
391 self.topk_indices_dtype = None
392 if prepare_finalize is not None:
393 self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
368 - experts = self.select_gemm_impl(prepare_finalize)
394 + experts = self.select_gemm_impl(prepare_finalize, moe)
395 self.fused_experts = FusedMoEModularKernel(
396 prepare_finalize,
397 experts,
398 )
399
400 def select_gemm_impl(
375 - self, prepare_finalize: FusedMoEPrepareAndFinalize
376 - ) -> FusedMoEPermuteExpertsUnpermute:
401 + self, prepare_finalize: FusedMoEPrepareAndFinalize,
402 + moe: Optional[MoEConfig]) -> FusedMoEPermuteExpertsUnpermute:
403 # based on the all2all implementation, select the appropriate
404 # gemm implementation
405 raise NotImplementedError(
@@ -419,7 +445,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
445 else:
446 self.rocm_aiter_fused_experts = None # type: ignore
447
422 - def select_gemm_impl(self, prepare_finalize: FusedMoEPrepareAndFinalize):
448 + def select_gemm_impl(self, prepare_finalize: FusedMoEPrepareAndFinalize,
449 + moe: Optional[MoEConfig]):
450
451 assert self.fused_experts == fused_experts
452
@@ -809,7 +836,6 @@ class FusedMoE(torch.nn.Module):
836 activation: str = "silu",
837 ):
838 super().__init__()
812 -
839 if params_dtype is None:
840 params_dtype = torch.get_default_dtype()
841 self.params_dtype = params_dtype
@@ -869,14 +895,24 @@ class FusedMoE(torch.nn.Module):
895 from vllm_hpu_extension.ops import DynamicFusedMOE
896 self.hpu_fused_moe = DynamicFusedMOE(self.global_num_experts)
897
898 + # Only support float8 for now.
899 + quant_dtype = params_dtype
900 + if quant_config is not None:
901 + input_activations = get_quant_config_input_activations(
902 + quant_config)
903 + if (input_activations is not None
904 + and input_activations.num_bits == 8
905 + and input_activations.type == QuantizationType.FLOAT):
906 + quant_dtype = torch.float8_e4m3fn
907 +
908 moe = MoEConfig(
909 num_experts=self.global_num_experts,
910 experts_per_token=top_k,
911 hidden_dim=hidden_size,
912 num_local_experts=self.local_num_experts,
913 moe_parallel_config=self.moe_parallel_config,
878 - # TODO (bnell): this needs to be fixed for quantized types.
914 in_dtype=params_dtype,
915 + quant_dtype=quant_dtype,
916 max_num_tokens=MOE_DP_CHUNK_SIZE,
917 )
918 self.moe_config = moe
vllm/model_executor/layers/fused_moe/modular_kernel.py
+2 -1
@@ -175,6 +175,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
175 def workspace_shapes(
176 self,
177 a: torch.Tensor,
178 + aq: torch.Tensor,
179 M: int,
180 N: int,
181 K: int,
@@ -309,7 +310,7 @@ class FusedMoEModularKernel(torch.nn.Module):
310
311 # Use a1 here to decipher the correct workspace datatype
312 workspace13_shape, workspace2_shape, workspace_dtype = (
312 - self.fused_experts.workspace_shapes(a1, M, N, K, top_k,
313 + self.fused_experts.workspace_shapes(a1, a1q, M, N, K, top_k,
314 global_num_experts))
315
316 # We can reuse the memory between cache1 and cache3 because by the time
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+13 -8
@@ -21,7 +21,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
21 rank: int,
22 dp_size: int,
23 quant_dtype: Optional[torch.dtype] = None,
24 - block_shape: Optional[list[int]] = None):
24 + block_shape: Optional[list[int]] = None,
25 + per_act_token: bool = False):
26 super().__init__()
27 assert max_num_tokens > 0
28 self.a2a = a2a
@@ -31,6 +32,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
32 self.rank = rank
33 self.dp_size = dp_size
34 self.quant_dtype = quant_dtype
35 + self.per_act_token = per_act_token
36
37 def max_num_tokens_per_rank(self) -> Optional[int]:
38 return self.max_num_tokens
@@ -66,13 +68,14 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
68 "apply_router_weight_on_input is only implemented for topk=1")
69 a1 = a1 * rank_topk_weights.to(a1.dtype)
70
69 - per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
70 - a2_scale.numel() != 1 if a2_scale is not None else False)
71 + repeat_cols = 4
72 + repeat_rows = 1 if self.per_act_token else a1.shape[0]
73 + a1q, a1q_scale = moe_kernel_quantize_input(
74 + a1, (None if self.per_act_token else a1_scale), self.quant_dtype,
75 + self.per_act_token, self.block_shape)
76
72 - a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
73 - self.quant_dtype,
74 - per_act_token,
75 - self.block_shape)
77 + if a1q_scale is not None:
78 + a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols)
79
80 # rem_experts need to be 0 for pplx to work properly.
81 rem_experts = num_experts % self.world_size
@@ -100,7 +103,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
103 else 1) * float32_size
104 expert_x_scale = torch.empty(
105 (
103 - num_experts,
106 + num_local_experts,
107 expert_x.size(1),
108 (expert_x.size(2) + block_size - 1) // block_size,
109 ),
@@ -121,6 +124,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
124 indices=rank_topk_ids,
125 bound_m=bound_m,
126 )
127 + if expert_x_scale is not None:
128 + expert_x_scale = expert_x_scale[:, :, 0:1]
129
130 return expert_x, expert_x_scale, expert_num_tokens, None, None
131
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+3 -2
@@ -37,6 +37,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
37 def workspace_shapes(
38 self,
39 a: torch.Tensor,
40 + aq: torch.Tensor,
41 M: int,
42 N: int,
43 K: int,
@@ -49,9 +50,9 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
50 if self.allow_deep_gemm and _valid_deep_gemm_shape(M, N, K):
51 assert self.deep_gemm_expert is not None
52 return self.deep_gemm_expert.workspace_shapes(
52 - a, M, N, K, topk, num_experts)
53 + a, aq, M, N, K, topk, num_experts)
54 else:
54 - return self.triton_expert.workspace_shapes(a, M, N, K, topk,
55 + return self.triton_expert.workspace_shapes(a, aq, M, N, K, topk,
56 num_experts)
57
58 def apply(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+47 -40
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import enum
5 +import importlib
6 from enum import Enum
7 from typing import Callable, Optional
8
@@ -11,7 +12,6 @@ from compressed_tensors.quantization import (ActivationOrdering,
12 QuantizationStrategy)
13
14 import vllm.envs as envs
14 -import vllm.model_executor.layers.fused_moe # noqa
15 from vllm import _custom_ops as ops
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
@@ -30,6 +30,15 @@ from vllm.model_executor.utils import set_weight_attrs
30 from vllm.platforms import current_platform
31 from vllm.scalar_type import scalar_types
32
33 +has_pplx = importlib.util.find_spec("pplx_kernels") is not None
34 +
35 +if current_platform.is_cuda_alike():
36 + from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
37 + BatchedPrepareAndFinalize)
38 + if has_pplx:
39 + from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
40 + PplxPrepareAndFinalize)
41 +
42 logger = init_logger(__name__)
43
44
@@ -77,8 +86,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
86 else:
87 logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
88 return CompressedTensorsWNA16MarlinMoEMethod(quant_config)
80 - elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
81 - and layer.activation == "silu"):
89 + elif quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant):
90 return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
91 elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
92 return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
@@ -421,6 +429,11 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
429 "For FP8 Fused MoE layer, we require either per tensor or "
430 "channelwise, dynamic per token quantization.")
431
432 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
433 + cutlass_moe_fp8)
434 + self.fused_experts = cutlass_moe_fp8 # type: ignore
435 + self.disable_expert_map = False
436 +
437 def create_weights(self, layer: torch.nn.Module, num_experts: int,
438 hidden_size: int, intermediate_size_per_partition: int,
439 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -499,25 +512,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
512 layer.w13_input_scale = None
513 layer.w2_input_scale = None
514
502 - device = w13_weight.device
503 - # TODO strides can be shared across multiple layers
504 - self.ab_strides1 = torch.full((num_experts, ),
505 - hidden_size,
506 - device=device,
507 - dtype=torch.int64)
508 - self.c_strides1 = torch.full((num_experts, ),
509 - 2 * intermediate_size_per_partition,
510 - device=device,
511 - dtype=torch.int64)
512 - self.ab_strides2 = torch.full((num_experts, ),
513 - intermediate_size_per_partition,
514 - device=device,
515 - dtype=torch.int64)
516 - self.c_strides2 = torch.full((num_experts, ),
517 - hidden_size,
518 - device=device,
519 - dtype=torch.int64)
520 -
515 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
516 # Fp8 moe kernels require a single activation scale.
517 # We take the max of all the scales in case they differ.
@@ -558,6 +552,27 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
552 layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
553 requires_grad=False)
554
555 + def select_gemm_impl(self, prepare_finalize, moe):
556 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
557 + CutlassExpertsFp8)
558 +
559 + assert moe is not None
560 +
561 + max_experts_per_worker = (
562 + (moe.num_experts + prepare_finalize.world_size - 1) //
563 + prepare_finalize.world_size)
564 + experts = CutlassExpertsFp8(
565 + max_experts_per_worker, moe.in_dtype,
566 + self.input_quant.strategy == QuantizationStrategy.TOKEN,
567 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
568 +
569 + if has_pplx and isinstance(
570 + prepare_finalize,
571 + (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
572 + # no expert_map support in this case
573 + self.disable_expert_map = True
574 + return experts
575 +
576 def apply(
577 self,
578 layer: torch.nn.Module,
@@ -577,9 +592,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
592 activation: str = "silu",
593 ) -> torch.Tensor:
594
580 - assert activation == "silu", (
581 - f"{activation} not supported for Cutlass MoE.")
582 -
595 topk_weights, topk_ids = FusedMoE.select_experts(
596 hidden_states=x,
597 router_logits=router_logits,
@@ -590,27 +602,22 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
602 num_expert_group=num_expert_group,
603 custom_routing_function=custom_routing_function,
604 scoring_func=scoring_func,
593 - e_score_correction_bias=e_score_correction_bias)
594 -
595 - from vllm.model_executor.layers.fused_moe import cutlass_moe_fp8
605 + e_score_correction_bias=e_score_correction_bias,
606 + indices_type=torch.uint32)
607
597 - return cutlass_moe_fp8(
608 + return self.fused_experts(
609 x,
599 - layer.w13_weight.transpose(1, 2),
600 - layer.w2_weight.transpose(1, 2),
601 - layer.w13_weight_scale,
602 - layer.w2_weight_scale,
610 + layer.w13_weight,
611 + layer.w2_weight,
612 topk_weights,
613 topk_ids,
605 - self.ab_strides1,
606 - self.c_strides1,
607 - self.ab_strides2,
608 - self.c_strides2,
614 + activation=activation,
615 + global_num_experts=global_num_experts,
616 + expert_map=None if self.disable_expert_map else expert_map,
617 + w1_scale=layer.w13_weight_scale,
618 + w2_scale=layer.w2_weight_scale,
619 a1_scale=layer.w13_input_scale,
620 a2_scale=layer.w2_input_scale,
611 - out_dtype=x.dtype,
612 - expert_map=expert_map,
613 - apply_router_weight_on_input=apply_router_weight_on_input,
621 )
622
623
vllm/model_executor/layers/quantization/fp8.py
+1 -1
@@ -769,7 +769,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
769 del layer.w13_input_scale
770 del layer.w2_input_scale
771
772 - def select_gemm_impl(self, prepare_finalize):
772 + def select_gemm_impl(self, prepare_finalize, moe):
773
774 from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
775 BatchedTritonOrDeepGemmExperts)