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, 01:00 UTC
Revert "[Performance] Performance improvements in non-blockwise fp8 CUTLASS MoE (#20762) (#21334)
Signed-off-by: Ming Yang <minos.future@gmail.com>
Ming Yang committed
Jul 21, 2025 at 21:49 UTC
e7b204268132cb775c139574c1ff4ad7e15c8f66
6 files changed
+38
-174
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+1
-34
@@ -80,11 +80,6 @@ def bench_run(
80
a, score, topk, renormalize=False
81
)
82
83
- ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
84
- ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)
85
- c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)
86
- c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
87
-
83
def run_triton_moe(
84
a: torch.Tensor,
85
w1: torch.Tensor,
@@ -116,10 +111,6 @@ def bench_run(
111
w2: torch.Tensor,
112
w1_scale: torch.Tensor,
113
w2_scale: torch.Tensor,
119
- ab_strides1: torch.Tensor,
120
- ab_strides2: torch.Tensor,
121
- c_strides1: torch.Tensor,
122
- c_strides2: torch.Tensor,
114
topk_weights: torch.Tensor,
115
topk_ids: torch.Tensor,
116
per_act_token: bool,
@@ -134,10 +125,6 @@ def bench_run(
125
topk_ids,
126
w1_scale,
127
w2_scale,
137
- ab_strides1,
138
- ab_strides2,
139
- c_strides1,
140
- c_strides2,
128
per_act_token,
129
a1_scale=None,
130
)
@@ -149,10 +136,6 @@ def bench_run(
136
w2_q: torch.Tensor,
137
w1_scale: torch.Tensor,
138
w2_scale: torch.Tensor,
152
- ab_strides1: torch.Tensor,
153
- ab_strides2: torch.Tensor,
154
- c_strides1: torch.Tensor,
155
- c_strides2: torch.Tensor,
139
topk_weights: torch.Tensor,
140
topk_ids: torch.Tensor,
141
):
@@ -167,10 +150,6 @@ def bench_run(
150
topk_ids,
151
w1_scale,
152
w2_scale,
170
- ab_strides1,
171
- ab_strides2,
172
- c_strides1,
173
- c_strides2,
153
per_act_token,
154
a1_scale=None,
155
)
@@ -215,10 +194,6 @@ def bench_run(
194
w2_q,
195
w1_scale,
196
w2_scale,
218
- ab_strides1,
219
- ab_strides2,
220
- c_strides1,
221
- c_strides2,
197
topk_weights,
198
topk_ids,
199
)
@@ -256,10 +231,6 @@ def bench_run(
231
"w1_scale": w1_scale,
232
"w2_scale": w2_scale,
233
"per_act_token": per_act_token,
259
- "ab_strides1": ab_strides1,
260
- "ab_strides2": ab_strides2,
261
- "c_strides1": c_strides1,
262
- "c_strides2": c_strides2,
234
# cuda graph params
235
"cutlass_graph": cutlass_graph,
236
"triton_graph": triton_graph,
@@ -318,10 +289,6 @@ def bench_run(
289
w2_q,
290
w1_scale,
291
w2_scale,
321
- ab_strides1,
322
- ab_strides2,
323
- c_strides1,
324
- c_strides2,
292
topk_weights,
293
topk_ids,
294
per_act_token,
@@ -330,7 +297,7 @@ def bench_run(
297
298
results.append(
299
benchmark.Timer(
333
- stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, ab_strides1, ab_strides2, c_strides1, c_strides2, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
300
+ stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
301
globals=globals,
302
label=label,
303
sub_label=sub_label,
csrc/moe/moe_permute_unpermute_op.cu
+11
-42
@@ -160,30 +160,6 @@ __global__ void shuffleInputRowsKernel(const T* input,
160
}
161
}
162
163
-template <typename T>
164
-__global__ void shuffleInputRowsKernelSlow(const T* input,
165
- const int32_t* dst2src_map,
166
- T* output, int64_t num_src_rows,
167
- int64_t num_dst_rows,
168
- int64_t num_cols) {
169
- int64_t dest_row_idx = blockIdx.x;
170
- int64_t const source_row_idx = dst2src_map[dest_row_idx];
171
-
172
- if (blockIdx.x < num_dst_rows) {
173
- // Duplicate and permute rows
174
- auto const* source_row_ptr = input + source_row_idx * num_cols;
175
- auto* dest_row_ptr = output + dest_row_idx * num_cols;
176
-
177
- int64_t const start_offset = threadIdx.x;
178
- int64_t const stride = blockDim.x;
179
-
180
- for (int elem_index = start_offset; elem_index < num_cols;
181
- elem_index += stride) {
182
- dest_row_ptr[elem_index] = source_row_ptr[elem_index];
183
- }
184
- }
185
-}
186
-
163
void shuffle_rows(const torch::Tensor& input_tensor,
164
const torch::Tensor& dst2src_map,
165
torch::Tensor& output_tensor) {
@@ -197,24 +173,17 @@ void shuffle_rows(const torch::Tensor& input_tensor,
173
int64_t const num_src_rows = input_tensor.size(0);
174
int64_t const num_cols = input_tensor.size(1);
175
200
- if (num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)) {
201
- // use slow kernel if num_cols can't be aligned to 128 bits
202
- MOE_DISPATCH(input_tensor.scalar_type(), [&] {
203
- shuffleInputRowsKernelSlow<scalar_t><<<blocks, threads, 0, stream>>>(
204
- reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),
205
- dst2src_map.data_ptr<int32_t>(),
206
- reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,
207
- num_dest_rows, num_cols);
208
- });
209
- } else {
210
- MOE_DISPATCH(input_tensor.scalar_type(), [&] {
211
- shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
212
- reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),
213
- dst2src_map.data_ptr<int32_t>(),
214
- reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,
215
- num_dest_rows, num_cols);
216
- });
217
- }
176
+ TORCH_CHECK(!(num_cols % (128 / sizeof(input_tensor.scalar_type()) / 8)),
177
+ "num_cols must be divisible by 128 / "
178
+ "sizeof(input_tensor.scalar_type()) / 8");
179
+
180
+ MOE_DISPATCH(input_tensor.scalar_type(), [&] {
181
+ shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>(
182
+ reinterpret_cast<scalar_t*>(input_tensor.data_ptr()),
183
+ dst2src_map.data_ptr<int32_t>(),
184
+ reinterpret_cast<scalar_t*>(output_tensor.data_ptr()), num_src_rows,
185
+ num_dest_rows, num_cols);
186
+ });
187
}
188
189
#else
tests/kernels/moe/test_cutlass_moe.py
+2
-12
@@ -207,10 +207,6 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
207
'topk_ids': topk_ids,
208
'w1_scale': moe_tensors.w1_scale,
209
'w2_scale': moe_tensors.w2_scale,
210
- 'ab_strides1': moe_tensors.ab_strides1,
211
- 'ab_strides2': moe_tensors.ab_strides2,
212
- 'c_strides1': moe_tensors.c_strides1,
213
- 'c_strides2': moe_tensors.c_strides2,
210
'per_act_token': per_act_token,
211
'a1_scale': None #moe_tensors.a_scale
212
}
@@ -444,11 +440,6 @@ def test_run_cutlass_moe_fp8(
440
expert_map[start:end] = list(range(num_local_experts))
441
expert_map = torch.tensor(expert_map, dtype=torch.int32, device="cuda")
442
447
- ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
448
- ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
449
- c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
450
- c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
451
-
443
activation = lambda o, i: torch.ops._C.silu_and_mul(o, i)
444
a1q, a1q_scale = moe_kernel_quantize_input(mt.a, mt.a_scale,
445
torch.float8_e4m3fn,
@@ -457,9 +448,8 @@ def test_run_cutlass_moe_fp8(
448
func = lambda output: run_cutlass_moe_fp8(
449
output, a1q, mt.w1_q, mt.w2_q, topk_ids, activation,
450
global_num_experts, expert_map, mt.w1_scale, mt.w2_scale,
460
- a1q_scale, None, ab_strides1, ab_strides2, c_strides1, c_strides2,
461
- workspace13, workspace2, None, mt.a.dtype, per_act_token,
462
- per_out_channel, False)
451
+ a1q_scale, None, workspace13, workspace2, None, mt.a.dtype,
452
+ per_act_token, per_out_channel, False)
453
454
workspace13.random_()
455
output_random_workspace = torch.empty(output_shape,
tests/kernels/moe/test_pplx_cutlass_moe.py
-22
@@ -75,7 +75,6 @@ def pplx_cutlass_moe(
75
assert torch.cuda.current_device() == pgi.local_rank
76
77
num_tokens, hidden_dim = a.shape
78
- intermediate_dim = w2.shape[2]
78
num_experts = w1.shape[0]
79
block_size = hidden_dim # TODO support more cases
80
device = pgi.device
@@ -124,31 +123,10 @@ def pplx_cutlass_moe(
123
num_local_experts=num_local_experts,
124
num_dispatchers=num_dispatchers)
125
127
- ab_strides1 = torch.full((num_local_experts, ),
128
- hidden_dim,
129
- device="cuda",
130
- dtype=torch.int64)
131
- ab_strides2 = torch.full((num_local_experts, ),
132
- intermediate_dim,
133
- device="cuda",
134
- dtype=torch.int64)
135
- c_strides1 = torch.full((num_local_experts, ),
136
- 2 * intermediate_dim,
137
- device="cuda",
138
- dtype=torch.int64)
139
- c_strides2 = torch.full((num_local_experts, ),
140
- hidden_dim,
141
- device="cuda",
142
- dtype=torch.int64)
143
-
126
experts = CutlassExpertsFp8(num_local_experts,
127
out_dtype,
128
per_act_token,
129
per_out_ch,
148
- ab_strides1,
149
- ab_strides2,
150
- c_strides1,
151
- c_strides2,
130
num_dispatchers=num_dispatchers,
131
use_batched_format=True)
132
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+23
-39
@@ -13,7 +13,8 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
13
MoEPrepareAndFinalizeNoEP)
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
TopKWeightAndReduceDelegate)
16
-from vllm.model_executor.layers.fused_moe.utils import (_fp8_quantize,
16
+from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,
17
+ _fp8_quantize,
18
_resize_cache,
19
extract_required_args)
20
from vllm.scalar_type import scalar_types
@@ -34,10 +35,6 @@ def run_cutlass_moe_fp8(
35
w2_scale: Optional[torch.Tensor],
36
a1q_scale: Optional[torch.Tensor],
37
a2_scale: Optional[torch.Tensor],
37
- ab_strides1: torch.Tensor,
38
- ab_strides2: torch.Tensor,
39
- c_strides1: torch.Tensor,
40
- c_strides2: torch.Tensor,
38
workspace13: torch.Tensor,
39
workspace2: torch.Tensor,
40
expert_num_tokens: Optional[torch.Tensor],
@@ -156,11 +153,27 @@ def run_cutlass_moe_fp8(
153
problem_sizes1, problem_sizes2, a_map,
154
c_map, global_num_experts, N, K)
155
159
- a1q = ops.shuffle_rows(a1q, a_map)
160
- a1q_scale = (ops.shuffle_rows(a1q_scale, a_map)
161
- if per_act_token else a1q_scale)
156
+ a1q = _fp8_perm(a1q, a_map)
157
+ a1q_scale = a1q_scale[a_map] if per_act_token else a1q_scale
158
expert_offsets = expert_offsets[:-1]
159
160
+ ab_strides1 = torch.full((w1.size(0), ),
161
+ K,
162
+ device=device,
163
+ dtype=torch.int64)
164
+ c_strides1 = torch.full((w1.size(0), ),
165
+ 2 * N,
166
+ device=device,
167
+ dtype=torch.int64)
168
+ ab_strides2 = torch.full((w1.size(0), ),
169
+ N,
170
+ device=device,
171
+ dtype=torch.int64)
172
+ c_strides2 = torch.full((w1.size(0), ),
173
+ K,
174
+ device=device,
175
+ dtype=torch.int64)
176
+
177
if use_batched_format:
178
c1 = _resize_cache(workspace13, (local_E * padded_M, N * 2))
179
c2 = _resize_cache(workspace2, (local_E * padded_M, N))
@@ -197,8 +210,7 @@ def run_cutlass_moe_fp8(
210
else:
211
# We can't do this inplace because output may point to the same tensor
212
# as c3.
200
- output.copy_(ops.shuffle_rows(c3, c_map).view(M * topk, K),
201
- non_blocking=True)
213
+ output.copy_(c3[c_map].view(M * topk, K), non_blocking=True)
214
215
216
# TODO (bnell): split class batched vs. non-batched?
@@ -211,10 +223,6 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
223
out_dtype: Optional[torch.dtype],
224
per_act_token_quant: bool,
225
per_out_ch_quant: bool,
214
- ab_strides1: torch.Tensor,
215
- ab_strides2: torch.Tensor,
216
- c_strides1: torch.Tensor,
217
- c_strides2: torch.Tensor,
226
block_shape: Optional[list[int]] = None,
227
num_dispatchers: Optional[int] = None,
228
use_batched_format: bool = False,
@@ -231,10 +239,6 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
239
self.max_experts_per_worker = max_experts_per_worker
240
self.num_dispatchers = num_dispatchers
241
self.out_dtype = out_dtype
234
- self.ab_strides1 = ab_strides1
235
- self.ab_strides2 = ab_strides2
236
- self.c_strides1 = c_strides1
237
- self.c_strides2 = c_strides2
242
self.use_batched_format = use_batched_format
243
244
@property
@@ -314,8 +318,7 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
318
run_cutlass_moe_fp8(
319
output, hidden_states, w1, w2, topk_ids, activation_callable,
320
global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,
317
- a2_scale, self.ab_strides1, self.ab_strides2, self.c_strides1,
318
- self.c_strides2, workspace13, workspace2, expert_num_tokens,
321
+ a2_scale, workspace13, workspace2, expert_num_tokens,
322
self.out_dtype if self.out_dtype is not None else in_dtype,
323
self.per_act_token_quant, self.per_out_ch_quant,
324
self.use_batched_format)
@@ -329,10 +332,6 @@ def cutlass_moe_fp8(
332
topk_ids: torch.Tensor,
333
w1_scale: torch.Tensor,
334
w2_scale: torch.Tensor,
332
- ab_strides1: torch.Tensor,
333
- ab_strides2: torch.Tensor,
334
- c_strides1: torch.Tensor,
335
- c_strides2: torch.Tensor,
335
per_act_token: Optional[bool] = None,
336
activation: str = "silu",
337
a1_scale: Optional[torch.Tensor] = None,
@@ -360,17 +359,6 @@ def cutlass_moe_fp8(
359
Shape: [num_experts] or [num_experts, 2N]
360
- w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
361
Shape: [num_experts] or [num_experts, K]
363
- - ab_strides1 (torch.Tensor): The input/weight strides for the first gemm.
364
- Shape: [num_experts]
365
- - ab_strides2 (torch.Tensor): The input/weight strides for the second gemm.
366
- Shape: [num_experts]
367
- - c_strides1 (torch.Tensor): The output strides for the first gemm.
368
- Shape: [num_experts]
369
- - c_strides2 (torch.Tensor): The output strides for the second gemm.
370
- Shape: [num_experts]
371
- - per_act_token (Optional[bool]): Whether the scale is per-token or
372
- per-tensor.
373
- - activation (str): The activation function to use.
362
- a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
363
Shape: scalar or [M]
364
- a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
@@ -403,10 +391,6 @@ def cutlass_moe_fp8(
391
out_dtype=a.dtype,
392
per_act_token_quant=per_act_token,
393
per_out_ch_quant=per_out_ch,
406
- ab_strides1=ab_strides1,
407
- ab_strides2=ab_strides2,
408
- c_strides1=c_strides1,
409
- c_strides2=c_strides2,
394
use_batched_format=False,
395
),
396
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+1
-25
@@ -859,21 +859,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
859
layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
860
requires_grad=False)
861
862
- device = layer.w13_weight.device
863
- # ab_strides1 and c_strides2 are the same
864
- self.ab_strides1_c_strides2 = torch.full((layer.local_num_experts, ),
865
- layer.hidden_size,
866
- device=device,
867
- dtype=torch.int64)
868
- self.ab_strides2 = torch.full((layer.local_num_experts, ),
869
- layer.intermediate_size_per_partition,
870
- device=device,
871
- dtype=torch.int64)
872
- self.c_strides1 = torch.full((layer.local_num_experts, ),
873
- 2 * layer.intermediate_size_per_partition,
874
- device=device,
875
- dtype=torch.int64)
876
-
862
def select_gemm_impl(
863
self,
864
prepare_finalize: FusedMoEPrepareAndFinalize,
@@ -896,10 +881,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
881
moe.in_dtype,
882
self.input_quant.strategy == QuantizationStrategy.TOKEN,
883
self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
899
- ab_strides1=self.ab_strides1_c_strides2,
900
- ab_strides2=self.ab_strides2,
901
- c_strides1=self.c_strides1,
902
- c_strides2=self.ab_strides1_c_strides2,
884
num_dispatchers=num_dispatchers,
885
use_batched_format=use_batched_format,
886
)
@@ -946,8 +927,7 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
927
num_expert_group=num_expert_group,
928
custom_routing_function=custom_routing_function,
929
scoring_func=scoring_func,
949
- e_score_correction_bias=e_score_correction_bias,
950
- indices_type=self.topk_indices_dtype)
930
+ e_score_correction_bias=e_score_correction_bias)
931
932
per_act_token = (
933
self.input_quant.strategy == QuantizationStrategy.TOKEN)
@@ -968,10 +948,6 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
948
expert_map=None if self.disable_expert_map else expert_map,
949
w1_scale=layer.w13_weight_scale,
950
w2_scale=layer.w2_weight_scale,
971
- ab_strides1=self.ab_strides1_c_strides2,
972
- ab_strides2=self.ab_strides2,
973
- c_strides1=self.c_strides1,
974
- c_strides2=self.ab_strides1_c_strides2,
951
a1_scale=layer.w13_input_scale,
952
a2_scale=layer.w2_input_scale,
953
)