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, 13:30 UTC
[Kernel] Add expert_map support to Cutlass FP8 MOE (#16861)
Signed-off-by: varun sundar rabindranath <vsundarr@redhat.com> Co-authored-by: varun sundar rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Apr 21, 2025 at 23:44 UTC
7b8a2ab76fc3438b1bd6e3caa86165cc0058c462
5 files changed
+333
-173
csrc/quantization/cutlass_w8a8/moe/moe_data.cu
+15
-2
@@ -46,14 +46,26 @@ __global__ void compute_expert_offsets(
46
}
47
48
__global__ void compute_arg_sorts(const int* __restrict__ topk_ids,
49
+ const int32_t* __restrict__ expert_offsets,
50
int32_t* input_permutation,
51
int32_t* output_permutation,
52
int32_t* atomic_buffer, const int topk_length,
53
const int topk) {
53
- int expert_id = blockIdx.x;
54
+ int const blk_expert_id = blockIdx.x;
55
+ int const num_experts = gridDim.x;
56
+ int32_t const num_tokens = expert_offsets[num_experts];
57
58
for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
56
- if (topk_ids[i] == expert_id) {
59
+ int const expert_id = topk_ids[i];
60
+ if (expert_id == -1 && blockIdx.x == 0) {
61
+ // output_permutation is used to re-order the moe outputs. It is
62
+ // used as c2 = c2[c_map], where c2 is a torch.tensor that is the
63
+ // output of the cutlass kernels and c_map is the output_permutation.
64
+ // c2 is initialized to zeros, therefore by setting the output_permutation
65
+ // to num_tokens, we are guaranteed to fill the moe outputs to zero
66
+ // for "invalid" topk_ids.
67
+ output_permutation[i] = num_tokens;
68
+ } else if (expert_id == blk_expert_id) {
69
int start = atomicAdd(&atomic_buffer[expert_id], 1);
70
input_permutation[start] = i / topk;
71
output_permutation[i] = start;
@@ -83,6 +95,7 @@ void get_cutlass_moe_mm_data_caller(
95
static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
96
compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
97
static_cast<const int32_t*>(topk_ids.data_ptr()),
98
+ static_cast<const int32_t*>(expert_offsets.data_ptr()),
99
static_cast<int32_t*>(input_permutation.data_ptr()),
100
static_cast<int32_t*>(output_permutation.data_ptr()),
101
static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
tests/kernels/test_cutlass_moe.py
+280
-160
@@ -1,4 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import dataclasses
3
+from typing import Optional
4
+
5
import pytest
6
import torch
7
@@ -12,32 +15,204 @@ from vllm.platforms import current_platform
15
NUM_EXPERTS = [40, 64]
16
TOP_KS = [6, 8]
17
18
+MNK_FACTORS = [
19
+ (2, 1024, 1024),
20
+ (2, 1024, 1536),
21
+ (2, 3072, 1024),
22
+ (2, 3072, 1536),
23
+ (64, 1024, 1024),
24
+ (64, 1024, 1536),
25
+ (64, 3072, 1024),
26
+ (64, 3072, 1536),
27
+ (224, 1024, 1024),
28
+ (224, 1024, 1536),
29
+ (224, 3072, 1024),
30
+ (224, 3072, 1536),
31
+]
32
+
33
+
34
+@dataclasses.dataclass
35
+class MOETensors:
36
+ a: torch.Tensor
37
+ w1: torch.Tensor
38
+ w2: torch.Tensor
39
+ ab_strides1: torch.Tensor
40
+ c_strides1: torch.Tensor
41
+ ab_strides2: torch.Tensor
42
+ c_strides2: torch.Tensor
43
+
44
+ @staticmethod
45
+ def make_moe_tensors(m: int, k: int, n: int, e: int,
46
+ dtype: torch.dtype) -> "MOETensors":
47
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
48
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
49
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
50
+ ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
51
+ c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
52
+ ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
53
+ c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
54
+ return MOETensors(a=a,
55
+ w1=w1,
56
+ w2=w2,
57
+ ab_strides1=ab_strides1,
58
+ c_strides1=c_strides1,
59
+ ab_strides2=ab_strides2,
60
+ c_strides2=c_strides2)
61
+
62
+
63
+@dataclasses.dataclass
64
+class MOETensors8Bit(MOETensors):
65
+ # quantized
66
+ a_q: Optional[torch.Tensor] = None # a -> a_q
67
+ w1_q: Optional[torch.Tensor] = None # w1 -> w1_q
68
+ w2_q: Optional[torch.Tensor] = None # w2 -> w2_q
69
+ a_scale: Optional[torch.Tensor] = None
70
+ w1_scale: Optional[torch.Tensor] = None
71
+ w2_scale: Optional[torch.Tensor] = None
72
+ # dequantized
73
+ a_d: Optional[torch.Tensor] = None # a -> a_q -> a_d
74
+ w1_d: Optional[torch.Tensor] = None # w1 -> w1_q -> w1_d
75
+ w2_d: Optional[torch.Tensor] = None # w2 -> w2_q -> w2_d
76
+
77
+ @staticmethod
78
+ def make_moe_tensors_8bit(m: int, k: int, n: int, e: int,
79
+ per_act_token: bool,
80
+ per_out_channel: bool) -> "MOETensors8Bit":
81
+ dtype = torch.half
82
+ q_dtype = torch.float8_e4m3fn
83
16
-def run(a: torch.Tensor, a_scale: torch.Tensor, w1_q: torch.Tensor,
17
- w2_q: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor,
18
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
19
- ab_strides1: torch.Tensor, c_strides1: torch.Tensor,
20
- ab_strides2: torch.Tensor, c_strides2: torch.Tensor):
21
- with set_current_vllm_config(
22
- VllmConfig(parallel_config=ParallelConfig(
23
- pipeline_parallel_size=1))):
24
- return cutlass_moe_fp8(a,
25
- w1_q,
26
- w2_q,
27
- w1_scale,
28
- w2_scale,
29
- topk_weights,
30
- topk_ids,
31
- ab_strides1,
32
- c_strides1,
33
- ab_strides2,
34
- c_strides2,
35
- a1_scale=a_scale)
36
-
37
-
38
-@pytest.mark.parametrize("m", [2, 64, 224])
39
-@pytest.mark.parametrize("n", [1024, 3072])
40
-@pytest.mark.parametrize("k", [1024, 1536])
84
+ moe_tensors_fp16 = MOETensors.make_moe_tensors(m, k, n, e, dtype)
85
+
86
+ # a -> a_q, w1 -> w1_q, w2 -> w2_q
87
+ n_b_scales = 2 * n if per_out_channel else 1
88
+ k_b_scales = k if per_out_channel else 1
89
+ # Get the right scale for tests.
90
+ _, a_scale = ops.scaled_fp8_quant(
91
+ moe_tensors_fp16.a, use_per_token_if_dynamic=per_act_token)
92
+ a_q, _ = ops.scaled_fp8_quant(moe_tensors_fp16.a,
93
+ a_scale,
94
+ use_per_token_if_dynamic=per_act_token)
95
+ w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
96
+ w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
97
+
98
+ w1_scale = torch.empty((e, n_b_scales, 1),
99
+ device="cuda",
100
+ dtype=torch.float32)
101
+ w2_scale = torch.empty((e, k_b_scales, 1),
102
+ device="cuda",
103
+ dtype=torch.float32)
104
+ for expert in range(e):
105
+ w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
106
+ moe_tensors_fp16.w1[expert],
107
+ use_per_token_if_dynamic=per_out_channel)
108
+ w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
109
+ moe_tensors_fp16.w2[expert],
110
+ use_per_token_if_dynamic=per_out_channel)
111
+
112
+ # a_q -> a_d, w1_q -> w1_d, w2_q -> w2_d
113
+ a_d = a_q.float().mul(a_scale).to(dtype)
114
+ w1_d = torch.empty_like(moe_tensors_fp16.w1)
115
+ w2_d = torch.empty_like(moe_tensors_fp16.w2)
116
+ for expert in range(e):
117
+ w1_d[expert] = (w1_q[expert].float() * w1_scale[expert]).half()
118
+ w2_d[expert] = (w2_q[expert].float() * w2_scale[expert]).half()
119
+
120
+ return MOETensors8Bit(a=moe_tensors_fp16.a,
121
+ w1=moe_tensors_fp16.w1,
122
+ w2=moe_tensors_fp16.w2,
123
+ ab_strides1=moe_tensors_fp16.ab_strides1,
124
+ c_strides1=moe_tensors_fp16.c_strides1,
125
+ ab_strides2=moe_tensors_fp16.ab_strides2,
126
+ c_strides2=moe_tensors_fp16.c_strides2,
127
+ a_q=a_q,
128
+ w1_q=w1_q,
129
+ w2_q=w2_q,
130
+ a_scale=a_scale,
131
+ w1_scale=w1_scale,
132
+ w2_scale=w2_scale,
133
+ a_d=a_d,
134
+ w1_d=w1_d,
135
+ w2_d=w2_d)
136
+
137
+
138
+def run_with_expert_maps(num_experts: int, num_local_experts: int,
139
+ **cutlass_moe_kwargs):
140
+
141
+ def slice_experts():
142
+ slice_params = [
143
+ "w1_q", "w2_q", "ab_strides1", "ab_strides2", "c_strides1",
144
+ "c_strides2", "w1_scale", "w2_scale"
145
+ ]
146
+ full_tensors = {
147
+ k: v
148
+ for k, v in cutlass_moe_kwargs.items()
149
+ if k in slice_params and k in cutlass_moe_kwargs
150
+ }
151
+
152
+ for i in range(0, num_experts, num_local_experts):
153
+ s, e = i, i + num_local_experts
154
+
155
+ # make expert map
156
+ expert_map = [-1] * num_experts
157
+ expert_map[s:e] = list(range(num_local_experts))
158
+ expert_map = torch.tensor(expert_map,
159
+ dtype=torch.int32,
160
+ device="cuda")
161
+
162
+ # update cutlass moe arg with expert_map
163
+ cutlass_moe_kwargs["expert_map"] = expert_map
164
+ # update cutlass moe arg tensors
165
+ for k, t in full_tensors.items():
166
+ cutlass_moe_kwargs[k] = t[s:e]
167
+
168
+ yield cutlass_moe_kwargs
169
+
170
+ out_tensor = torch.zeros_like(cutlass_moe_kwargs["a"])
171
+ for kwargs in slice_experts():
172
+ out_tensor = out_tensor + cutlass_moe_fp8(**kwargs)
173
+
174
+ return out_tensor
175
+
176
+
177
+def run_8_bit(moe_tensors: MOETensors8Bit,
178
+ topk_weights: torch.Tensor,
179
+ topk_ids: torch.Tensor,
180
+ num_local_experts: Optional[int] = None) -> torch.Tensor:
181
+ assert not any([
182
+ t is None for t in [
183
+ moe_tensors.w1_q, moe_tensors.w2_q, moe_tensors.w1_scale,
184
+ moe_tensors.w2_scale, moe_tensors.a_scale
185
+ ]
186
+ ])
187
+
188
+ kwargs = {
189
+ 'a': moe_tensors.a,
190
+ 'w1_q': moe_tensors.w1_q.transpose(1, 2), # type: ignore[union-attr]
191
+ 'w2_q': moe_tensors.w2_q.transpose(1, 2), # type: ignore[union-attr]
192
+ 'topk_weights': topk_weights,
193
+ 'topk_ids_': topk_ids,
194
+ 'ab_strides1': moe_tensors.ab_strides1,
195
+ 'c_strides1': moe_tensors.c_strides1,
196
+ 'ab_strides2': moe_tensors.ab_strides2,
197
+ 'c_strides2': moe_tensors.c_strides2,
198
+ 'w1_scale': moe_tensors.w1_scale,
199
+ 'w2_scale': moe_tensors.w2_scale,
200
+ 'a1_scale': moe_tensors.a_scale
201
+ }
202
+
203
+ num_experts = moe_tensors.w1.size(0)
204
+ with_ep = num_local_experts is not None or num_local_experts == num_experts
205
+ if not with_ep:
206
+ return cutlass_moe_fp8(**kwargs)
207
+
208
+ assert num_local_experts is not None
209
+ return run_with_expert_maps(
210
+ num_experts,
211
+ num_local_experts, # type: ignore[arg-type]
212
+ **kwargs)
213
+
214
+
215
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
216
@pytest.mark.parametrize("e", NUM_EXPERTS)
217
@pytest.mark.parametrize("topk", TOP_KS)
218
@pytest.mark.parametrize("per_act_token", [True, False])
@@ -46,7 +221,7 @@ def run(a: torch.Tensor, a_scale: torch.Tensor, w1_q: torch.Tensor,
221
(lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
222
current_platform.get_device_capability()),
223
reason="Grouped gemm is not supported on this GPU type.")
49
-def test_cutlass_moe_no_graph(
224
+def test_cutlass_moe_8_bit_no_graph(
225
m: int,
226
n: int,
227
k: int,
@@ -60,80 +235,21 @@ def test_cutlass_moe_no_graph(
235
VllmConfig(parallel_config=ParallelConfig(
236
pipeline_parallel_size=1))):
237
63
- dtype = torch.half
238
+ mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
239
+ per_out_ch)
240
65
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
66
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
67
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
68
-
69
- # Get the right scale for tests.
70
- _, a_scale1 = ops.scaled_fp8_quant(
71
- a, use_per_token_if_dynamic=per_act_token)
72
- a_q, _ = ops.scaled_fp8_quant(a,
73
- a_scale1,
74
- use_per_token_if_dynamic=per_act_token)
241
+ score = torch.randn((m, e), device="cuda", dtype=torch.half)
242
+ topk_weights, topk_ids = fused_topk(mt.a,
243
+ score,
244
+ topk,
245
+ renormalize=False)
246
76
- a_d = a_q.float().mul(a_scale1).to(dtype)
247
+ # Note that we are using the dequantized versions of the tensors.
248
+ # Using a, w1 and w2 directly results in minor output differences.
249
+ triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
250
+ topk_ids)
251
78
- n_b_scales = 2 * n if per_out_ch else 1
79
- k_b_scales = k if per_out_ch else 1
80
-
81
- w1_q = torch.empty((e, 2 * n, k),
82
- device="cuda",
83
- dtype=torch.float8_e4m3fn)
84
- w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
85
- w1_scale = torch.empty((e, n_b_scales, 1),
86
- device="cuda",
87
- dtype=torch.float32)
88
- w2_scale = torch.empty((e, k_b_scales, 1),
89
- device="cuda",
90
- dtype=torch.float32)
91
-
92
- ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
93
- c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
94
- ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
95
- c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
96
-
97
- for expert in range(e):
98
- w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
99
- w1[expert], use_per_token_if_dynamic=per_out_ch)
100
- w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
101
- w2[expert], use_per_token_if_dynamic=per_out_ch)
102
- w1_q = w1_q.transpose(1, 2)
103
- w2_q = w2_q.transpose(1, 2)
104
-
105
- ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
106
- c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
107
- ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
108
- c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
109
-
110
- w1_d = torch.empty_like(w1)
111
- w2_d = torch.empty_like(w2)
112
- for expert in range(e):
113
- w1_d[expert] = (w1_q[expert].t().float() * w1_scale[expert]).half()
114
- w2_d[expert] = (w2_q[expert].t().float() * w2_scale[expert]).half()
115
-
116
- score = torch.randn((m, e), device="cuda", dtype=dtype)
117
- topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
118
-
119
- triton_output = fused_experts(a_d, w1_d, w2_d, topk_weights, topk_ids)
120
-
121
- cutlass_output = cutlass_moe_fp8(a,
122
- w1_q,
123
- w2_q,
124
- w1_scale,
125
- w2_scale,
126
- topk_weights,
127
- topk_ids,
128
- ab_strides1,
129
- c_strides1,
130
- ab_strides2,
131
- c_strides2,
132
- a1_scale=a_scale1)
133
-
134
- #print(triton_output)
135
- #print(cutlass_output)
136
- #print("*")
252
+ cutlass_output = run_8_bit(mt, topk_weights, topk_ids)
253
254
torch.testing.assert_close(triton_output,
255
cutlass_output,
@@ -141,9 +257,7 @@ def test_cutlass_moe_no_graph(
257
rtol=1e-2)
258
259
144
-@pytest.mark.parametrize("m", [2, 64, 224])
145
-@pytest.mark.parametrize("n", [1024, 3072])
146
-@pytest.mark.parametrize("k", [1024, 1536])
260
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
261
@pytest.mark.parametrize("e", NUM_EXPERTS)
262
@pytest.mark.parametrize("topk", TOP_KS)
263
@pytest.mark.parametrize("per_act_token", [True, False])
@@ -152,7 +266,7 @@ def test_cutlass_moe_no_graph(
266
(lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
267
current_platform.get_device_capability()),
268
reason="Grouped gemm is not supported on this GPU type.")
155
-def test_cutlass_moe_cuda_graph(
269
+def test_cutlass_moe_8_bit_cuda_graph(
270
m: int,
271
n: int,
272
k: int,
@@ -168,77 +282,83 @@ def test_cutlass_moe_cuda_graph(
282
283
dtype = torch.half
284
171
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
172
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
173
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
174
-
175
- # Get the right scale for tests.
176
- _, a_scale1 = ops.scaled_fp8_quant(
177
- a, use_per_token_if_dynamic=per_act_token)
178
- a_q, _ = ops.scaled_fp8_quant(a,
179
- a_scale1,
180
- use_per_token_if_dynamic=per_act_token)
181
-
182
- a_d = a_q.float().mul(a_scale1).to(dtype)
183
-
184
- n_b_scales = 2 * n if per_out_ch else 1
185
- k_b_scales = k if per_out_ch else 1
186
-
187
- w1_q = torch.empty((e, 2 * n, k),
188
- device="cuda",
189
- dtype=torch.float8_e4m3fn)
190
- w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
191
- w1_scale = torch.empty((e, n_b_scales, 1),
192
- device="cuda",
193
- dtype=torch.float32)
194
- w2_scale = torch.empty((e, k_b_scales, 1),
195
- device="cuda",
196
- dtype=torch.float32)
197
-
198
- ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
199
- c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
200
- ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
201
- c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
202
-
203
- for expert in range(e):
204
- w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
205
- w1[expert], use_per_token_if_dynamic=per_out_ch)
206
- w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
207
- w2[expert], use_per_token_if_dynamic=per_out_ch)
208
- w1_q = w1_q.transpose(1, 2)
209
- w2_q = w2_q.transpose(1, 2)
210
-
211
- ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
212
- c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
213
- ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
214
- c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
215
-
216
- w1_d = torch.empty_like(w1)
217
- w2_d = torch.empty_like(w2)
218
- for expert in range(e):
219
- w1_d[expert] = (w1_q[expert].t().float() * w1_scale[expert]).half()
220
- w2_d[expert] = (w2_q[expert].t().float() * w2_scale[expert]).half()
285
+ mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
286
+ per_out_ch)
287
288
score = torch.randn((m, e), device="cuda", dtype=dtype)
223
- topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
289
+ topk_weights, topk_ids = fused_topk(mt.a,
290
+ score,
291
+ topk,
292
+ renormalize=False)
293
225
- triton_output = fused_experts(a_d, w1_d, w2_d, topk_weights, topk_ids)
294
+ # Note that we are using the dequantized versions of the tensors.
295
+ # Using a, w1 and w2 directly results in minor output differences.
296
+ triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
297
+ topk_ids)
298
299
stream = torch.cuda.Stream()
300
graph = torch.cuda.CUDAGraph()
301
with torch.cuda.graph(graph, stream=stream):
230
- cutlass_output = run(a, a_scale1, w1_q, w2_q, w1_scale, w2_scale,
231
- topk_weights, topk_ids, ab_strides1,
232
- c_strides1, ab_strides2, c_strides2)
302
+ cutlass_output = run_8_bit(mt, topk_weights, topk_ids)
303
+
304
torch.cuda.synchronize()
305
graph.replay()
306
torch.cuda.synchronize()
307
237
- #print(triton_output)
238
- #print(cutlass_output)
239
- #print("*")
240
-
308
torch.testing.assert_close(triton_output,
309
cutlass_output,
310
atol=9e-2,
311
rtol=1e-2)
312
+
313
+
314
+@pytest.mark.parametrize("m", [64])
315
+@pytest.mark.parametrize("n", [1024])
316
+@pytest.mark.parametrize("k", [4096])
317
+@pytest.mark.parametrize("e", [16])
318
+@pytest.mark.parametrize("topk", [1, 8])
319
+@pytest.mark.parametrize("per_act_token", [True])
320
+@pytest.mark.parametrize("per_out_channel", [True])
321
+@pytest.mark.parametrize("ep_size", [1, 2, 4, 8, 16])
322
+@pytest.mark.skipif(
323
+ (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
324
+ current_platform.get_device_capability()),
325
+ reason="Grouped gemm is not supported on this GPU type.")
326
+def test_cutlass_moe_8_bit_EP(
327
+ m: int,
328
+ n: int,
329
+ k: int,
330
+ e: int,
331
+ topk: int,
332
+ per_act_token: bool,
333
+ per_out_channel: bool,
334
+ ep_size: int,
335
+):
336
+ current_platform.seed_everything(7)
337
+ with set_current_vllm_config(
338
+ VllmConfig(parallel_config=ParallelConfig(
339
+ pipeline_parallel_size=1))):
340
+
341
+ mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
342
+ per_out_channel)
343
+
344
+ score = torch.randn((m, e), device="cuda", dtype=torch.half)
345
+ topk_weights, topk_ids = fused_topk(mt.a,
346
+ score,
347
+ topk,
348
+ renormalize=False)
349
+
350
+ # Note that we are using the dequantized versions of the tensors.
351
+ # Using a, w1 and w2 directly results in minor output differences.
352
+ triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
353
+ topk_ids)
354
+
355
+ assert e % ep_size == 0, "Cannot distribute experts evenly"
356
+ cutlass_output = run_8_bit(mt,
357
+ topk_weights,
358
+ topk_ids,
359
+ num_local_experts=e // ep_size)
360
+
361
+ torch.testing.assert_close(triton_output,
362
+ cutlass_output,
363
+ atol=5e-2,
364
+ rtol=1e-2)
vllm/config.py
+1
@@ -1693,6 +1693,7 @@ class ParallelConfig:
1693
factors: list[Any] = []
1694
factors.append(self.pipeline_parallel_size)
1695
factors.append(self.tensor_parallel_size)
1696
+ factors.append(self.enable_expert_parallel)
1697
return hashlib.sha256(str(factors).encode()).hexdigest()
1698
1699
def __post_init__(self) -> None:
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+35
-8
@@ -15,7 +15,7 @@ def cutlass_moe_fp8(
15
w1_scale: torch.Tensor,
16
w2_scale: torch.Tensor,
17
topk_weights: torch.Tensor,
18
- topk_ids: torch.Tensor,
18
+ topk_ids_: torch.Tensor,
19
ab_strides1: torch.Tensor,
20
c_strides1: torch.Tensor,
21
ab_strides2: torch.Tensor,
@@ -23,6 +23,7 @@ def cutlass_moe_fp8(
23
a1_scale: Optional[torch.Tensor] = None,
24
a2_scale: Optional[torch.Tensor] = None,
25
out_dtype: torch.dtype = torch.half,
26
+ expert_map: Optional[torch.Tensor] = None,
27
apply_router_weight_on_input: bool = False,
28
) -> torch.Tensor:
29
"""
@@ -57,12 +58,19 @@ def cutlass_moe_fp8(
58
quantize the intermediate result between the gemms.
59
Shape: scalar or [M]
60
- out_dtype (torch.Tensor): The output tensor type.
61
+ - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
62
+ every Rank is responsible for a subset of experts. expert_map is a
63
+ mapping from global expert-id to local expert-id. When expert_map[i]
64
+ is -1, it means that this Rank is not responsible for global
65
+ expert-id i.
66
+ - apply_router_weight_on_input (bool): When true, the topk weights are
67
+ applied directly on the inputs. This is only applicable when topk is 1.
68
69
Returns:
70
- torch.Tensor: The fp16 output tensor after applying the MoE layer.
71
"""
72
65
- assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
73
+ assert topk_weights.shape == topk_ids_.shape, "topk shape mismatch"
74
assert w1_q.dtype == torch.float8_e4m3fn
75
assert w2_q.dtype == torch.float8_e4m3fn
76
assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
@@ -96,7 +104,13 @@ def cutlass_moe_fp8(
104
k = w1_q.size(1)
105
n = w2_q.size(1)
106
99
- topk = topk_ids.size(1)
107
+ local_topk_ids = topk_ids_
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
+
113
+ topk = local_topk_ids.size(1)
114
115
per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
116
a2_scale.numel() != 1 if a2_scale is not None else False)
@@ -120,10 +134,23 @@ def cutlass_moe_fp8(
134
dtype=torch.int32,
135
device=device)
136
123
- a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
124
- c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
125
-
126
- ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
137
+ a_map_initializer = torch.empty
138
+ c2_initializer = torch.empty
139
+ if expert_map is not None:
140
+ # With expert_map each Rank processes only a subset of experts. As
141
+ # a result not all of a_map and c2 tensors are filled. We fill it
142
+ # zeros for correctness.
143
+ a_map_initializer = torch.zeros
144
+ c2_initializer = torch.zeros
145
+
146
+ a_map = a_map_initializer((local_topk_ids.numel()),
147
+ dtype=torch.int32,
148
+ device=device)
149
+ c_map = torch.empty((local_topk_ids.numel()),
150
+ dtype=torch.int32,
151
+ device=device)
152
+
153
+ ops.get_cutlass_moe_mm_data(local_topk_ids, expert_offsets, problem_sizes1,
154
problem_sizes2, a_map, c_map, num_experts, n,
155
k)
156
@@ -131,7 +158,7 @@ def cutlass_moe_fp8(
158
rep_a1_scales = a1_scale[a_map] if per_act_token else a1_scale
159
160
c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
134
- c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
161
+ c2 = c2_initializer((m * topk, k), device=device, dtype=out_dtype)
162
163
ops.cutlass_moe_mm(c1, rep_a_q, w1_q, rep_a1_scales, w1_scale,
164
expert_offsets[:-1], problem_sizes1, ab_strides1,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+2
-3
@@ -67,7 +67,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
67
else:
68
return CompressedTensorsWNA16MarlinMoEMethod(quant_config)
69
elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
70
- and layer.activation == "silu" and layer.expert_map is None):
70
+ and layer.activation == "silu"):
71
return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
72
elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
73
return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
@@ -510,8 +510,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
510
) -> torch.Tensor:
511
512
assert activation == "silu"
513
- assert global_num_experts == layer.w13_weight.shape[0]
514
- assert expert_map is None
513
514
topk_weights, topk_ids = FusedMoE.select_experts(
515
hidden_states=x,
@@ -542,6 +540,7 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
540
a1_scale=layer.w13_input_scale,
541
a2_scale=layer.w2_input_scale,
542
out_dtype=x.dtype,
543
+ expert_map=expert_map,
544
apply_router_weight_on_input=apply_router_weight_on_input,
545
)
546