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 9, 17:06 UTC
[Bug] Fix benchmark script `moe_permute_unpermute` (#32949)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jan 23, 2026 at 16:18 UTC
dfab5f37648d123480acb8a6246c6abbd80dca58
1 file changed
+3
-5
benchmarks/kernels/benchmark_moe_permute_unpermute.py
+3
-5
@@ -8,7 +8,7 @@ import ray
8
import torch
9
from transformers import AutoConfig
10
11
-from vllm.model_executor.layers.fused_moe.fused_moe import *
11
+from vllm.model_executor.layers.fused_moe import fused_topk
12
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
13
_moe_permute,
14
_moe_unpermute_and_reduce,
@@ -86,9 +86,7 @@ def benchmark_permute(
86
sorted_token_ids,
87
expert_ids,
88
inv_perm,
89
- ) = _moe_permute(
90
- qhidden_states, None, topk_ids, num_experts, None, align_block_size
91
- )
89
+ ) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16)
90
91
# JIT compilation & warmup
92
run()
@@ -182,7 +180,7 @@ def benchmark_unpermute(
180
expert_ids,
181
inv_perm,
182
) = _moe_permute(
185
- qhidden_states, None, topk_ids, num_experts, None, align_block_size
183
+ qhidden_states, None, topk_ids, num_experts, None, block_m=16
184
)
185
# convert to fp16/bf16 as gemm output
186
return (