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 10, 16:30 UTC
[Perf] Optimize additional `fill(0)` in cutlass moe, 2.9% E2E throughput improvement, 10.8% TTFT improvement (#31754)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 5, 2026 at 21:01 UTC
ffec8154225f00a41447c6cd0617d5245972f0d5
1 file changed
+10
-10
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+10
-10
@@ -173,7 +173,7 @@ def run_cutlass_moe_fp8(
173
174
num_expert = global_num_experts if expert_map is None else expert_map.size(0)
175
# permuted a1q reuses workspace2
176
- a1q, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
176
+ a1q, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute(
177
a1q,
178
a1q_scale,
179
topk_ids,
@@ -182,7 +182,7 @@ def run_cutlass_moe_fp8(
182
expert_map,
183
permuted_hidden_states=a1q_perm,
184
)
185
- expert_offsets = expert_offsets[:-1]
185
+ expert_offsets = expert_first_token_offset[:-1]
186
187
ops.get_cutlass_moe_mm_problem_sizes(
188
local_topk_ids, problem_sizes1, problem_sizes2, global_num_experts, N, K
@@ -215,9 +215,6 @@ def run_cutlass_moe_fp8(
215
act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out
216
)
217
218
- if expert_map is not None:
219
- mm2_out.fill_(0)
220
-
218
ops.cutlass_moe_mm(
219
mm2_out,
220
a2q,
@@ -243,6 +240,9 @@ def run_cutlass_moe_fp8(
240
permuted_hidden_states=mm2_out,
241
topk_weights=topk_weights,
242
inv_permuted_idx=inv_perm,
243
+ expert_first_token_offset=(
244
+ expert_first_token_offset if expert_map is not None else None
245
+ ),
246
)
247
248
@@ -988,7 +988,7 @@ def run_cutlass_moe_w4a8_fp8(
988
989
num_expert = global_num_experts if expert_map is None else expert_map.size(0)
990
# permuted a1q reuses workspace2
991
- a1q, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
991
+ a1q, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute(
992
a1q,
993
a1q_scale,
994
topk_ids,
@@ -997,7 +997,7 @@ def run_cutlass_moe_w4a8_fp8(
997
expert_map,
998
permuted_hidden_states=a1q_perm,
999
)
1000
- expert_offsets = expert_offsets[:-1]
1000
+ expert_offsets = expert_first_token_offset[:-1]
1001
1002
# For RS gemm SwapAB is always enabled (swap logical M, N in the problem shape)
1003
ops.get_cutlass_moe_mm_problem_sizes(
@@ -1032,9 +1032,6 @@ def run_cutlass_moe_w4a8_fp8(
1032
act_out, a2_scale, use_per_token_if_dynamic=per_act_token, output=quant_out
1033
)
1034
1035
- if expert_map is not None:
1036
- mm2_out.fill_(0)
1037
-
1035
ops.cutlass_w4a8_moe_mm(
1036
mm2_out,
1037
a2q,
@@ -1058,6 +1055,9 @@ def run_cutlass_moe_w4a8_fp8(
1055
permuted_hidden_states=mm2_out,
1056
topk_weights=topk_weights,
1057
inv_permuted_idx=inv_perm,
1058
+ expert_first_token_offset=(
1059
+ expert_first_token_offset if expert_map is not None else None
1060
+ ),
1061
)
1062
1063