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 15, 06:00 UTC
[Bugfix] Fix benchmark_moe.py (#19016)
Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Tianyu Guo committed
Jun 10, 2025 at 09:04 UTC
4589b940326d1480effc66ee33d30404ca31ed40
1 file changed
+4
-10
benchmarks/kernels/benchmark_moe.py
+4
-10
@@ -7,7 +7,6 @@ import time
7
from contextlib import nullcontext
8
from datetime import datetime
9
from itertools import product
10
-from types import SimpleNamespace
10
from typing import Any, TypedDict
11
12
import ray
@@ -43,7 +42,7 @@ def benchmark_config(
42
use_fp8_w8a8: bool,
43
use_int8_w8a16: bool,
44
num_iters: int = 100,
46
- block_quant_shape: List[int] = None,
45
+ block_quant_shape: list[int] = None,
46
use_deep_gemm: bool = False,
47
) -> float:
48
init_dtype = torch.float16 if use_fp8_w8a8 else dtype
@@ -400,7 +399,7 @@ class BenchmarkWorker:
399
dtype: torch.dtype,
400
use_fp8_w8a8: bool,
401
use_int8_w8a16: bool,
403
- block_quant_shape: List[int] = None,
402
+ block_quant_shape: list[int] = None,
403
use_deep_gemm: bool = False,
404
) -> tuple[dict[str, int], float]:
405
current_platform.seed_everything(self.seed)
@@ -532,7 +531,7 @@ def save_configs(
531
dtype: torch.dtype,
532
use_fp8_w8a8: bool,
533
use_int8_w8a16: bool,
535
- block_quant_shape: List[int],
534
+ block_quant_shape: list[int],
535
) -> None:
536
dtype_str = get_config_dtype_str(
537
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
@@ -563,7 +562,6 @@ def main(args: argparse.Namespace):
562
config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
563
if args.model_prefix:
564
config = getattr(config, args.model_prefix)
566
- config = SimpleNamespace(**config)
565
566
if config.architectures[0] == "DbrxForCausalLM":
567
E = config.ffn_config.moe_num_experts
@@ -595,11 +593,7 @@ def main(args: argparse.Namespace):
593
shard_intermediate_size = 2 * intermediate_size // args.tp_size
594
595
hidden_size = config.hidden_size
598
- dtype = (
599
- torch.float16
600
- if current_platform.is_rocm()
601
- else getattr(torch, config.torch_dtype)
602
- )
596
+ dtype = torch.float16 if current_platform.is_rocm() else config.torch_dtype
597
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
598
use_int8_w8a16 = args.dtype == "int8_w8a16"
599
block_quant_shape = get_weight_block_size_safety(config)