From dfab5f37648d123480acb8a6246c6abbd80dca58 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 23 Jan 2026 16:18:56 -0500 Subject: [PATCH] [Bug] Fix benchmark script `moe_permute_unpermute` (#32949) Signed-off-by: yewentao256 --- benchmarks/kernels/benchmark_moe_permute_unpermute.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe_permute_unpermute.py b/benchmarks/kernels/benchmark_moe_permute_unpermute.py index e70f5c0a1..8920e90fc 100644 --- a/benchmarks/kernels/benchmark_moe_permute_unpermute.py +++ b/benchmarks/kernels/benchmark_moe_permute_unpermute.py @@ -8,7 +8,7 @@ import ray import torch from transformers import AutoConfig -from vllm.model_executor.layers.fused_moe.fused_moe import * +from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import ( _moe_permute, _moe_unpermute_and_reduce, @@ -86,9 +86,7 @@ def benchmark_permute( sorted_token_ids, expert_ids, inv_perm, - ) = _moe_permute( - qhidden_states, None, topk_ids, num_experts, None, align_block_size - ) + ) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16) # JIT compilation & warmup run() @@ -182,7 +180,7 @@ def benchmark_unpermute( expert_ids, inv_perm, ) = _moe_permute( - qhidden_states, None, topk_ids, num_experts, None, align_block_size + qhidden_states, None, topk_ids, num_experts, None, block_m=16 ) # convert to fp16/bf16 as gemm output return (