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 14, 22:35 UTC
[Misc] Minor enhancement of benchmark_moe (#22068)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Aug 2, 2025 at 01:35 UTC
8d705996dffbb2299750b7b2b50bbcd5ccb4a5ad
1 file changed
+8
-1
benchmarks/kernels/benchmark_moe.py
+8
-1
@@ -22,6 +22,13 @@ from vllm.utils import FlexibleArgumentParser
22
FP8_DTYPE = current_platform.fp8_dtype()
23
24
25
+def ensure_divisibility(numerator, denominator):
26
+ """Ensure that numerator is divisible by the denominator."""
27
+ assert numerator % denominator == 0, (
28
+ "intermediate_size {} is not divisible by tp {}.".format(numerator, denominator)
29
+ )
30
+
31
+
32
class BenchmarkConfig(TypedDict):
33
BLOCK_SIZE_M: int
34
BLOCK_SIZE_N: int
@@ -603,7 +610,7 @@ def main(args: argparse.Namespace):
610
topk = config.num_experts_per_tok
611
intermediate_size = config.intermediate_size
612
shard_intermediate_size = 2 * intermediate_size // args.tp_size
606
-
613
+ ensure_divisibility(intermediate_size, args.tp_size)
614
hidden_size = config.hidden_size
615
dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
616
use_fp8_w8a8 = args.dtype == "fp8_w8a8"