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, 20:45 UTC
Add missing import of fused_topk to benchmark_moe (#32784)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Jan 21, 2026 at 20:30 UTC
f99953986945be69b20cc395fc5b74a5b78d4585
1 file changed
+20
-16
benchmarks/kernels/benchmark_moe.py
+20
-16
@@ -15,6 +15,7 @@ import ray
15
import torch
16
from ray.experimental.tqdm_ray import tqdm
17
18
+from vllm.model_executor.layers.fused_moe import fused_topk
19
from vllm.model_executor.layers.fused_moe.config import (
20
FusedMoEConfig,
21
FusedMoEParallelConfig,
@@ -200,23 +201,26 @@ def benchmark_config(
201
block_shape=block_quant_shape,
202
)
203
203
- deep_gemm_experts = mk.FusedMoEModularKernel(
204
- prepare_finalize=MoEPrepareAndFinalizeNoEP(),
205
- fused_experts=TritonOrDeepGemmExperts(
206
- moe_config=FusedMoEConfig(
207
- num_experts=num_experts,
208
- experts_per_token=topk,
209
- hidden_dim=hidden_size,
210
- intermediate_size_per_partition=shard_intermediate_size,
211
- num_local_experts=num_experts,
212
- activation="silu",
213
- parallel_config=FusedMoEParallelConfig.make_no_parallel(),
214
- in_dtype=init_dtype,
215
- routing_method=RoutingMethodType.TopK,
204
+ deep_gemm_experts = None
205
+ if use_deep_gemm:
206
+ deep_gemm_experts = mk.FusedMoEModularKernel(
207
+ prepare_finalize=MoEPrepareAndFinalizeNoEP(),
208
+ fused_experts=TritonOrDeepGemmExperts(
209
+ moe_config=FusedMoEConfig(
210
+ num_experts=num_experts,
211
+ experts_per_token=topk,
212
+ hidden_dim=hidden_size,
213
+ intermediate_size_per_partition=shard_intermediate_size,
214
+ num_local_experts=num_experts,
215
+ activation="silu",
216
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
217
+ in_dtype=init_dtype,
218
+ routing_method=RoutingMethodType.TopK,
219
+ device="cuda",
220
+ ),
221
+ quant_config=quant_config,
222
),
217
- quant_config=quant_config,
218
- ),
219
- )
223
+ )
224
225
with override_config(config):
226
topk_weights, topk_ids, token_expert_indices = fused_topk(