| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | |
| 7 | from vllm import _custom_ops as ops |
| 8 | from vllm.platforms import current_platform |
| 9 | from vllm.transformers_utils.config import get_config |
| 10 | from vllm.triton_utils import triton |
| 11 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 12 | |
| 13 | # Dimensions supported by the gpt-oss specialized kernel |
| 14 | GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128] |
| 15 | GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880] |
| 16 | |
| 17 | # Dimensions supported by the fp32 specialized kernel (MiniMax-M2) |
| 18 | FP32_SUPPORTED_NUM_EXPERTS = [256] |
| 19 | FP32_SUPPORTED_HIDDEN_SIZES = [3072] |
| 20 | FP32_MAX_TOKENS = 32 |
| 21 | |
| 22 | |
| 23 | def get_batch_size_range(max_batch_size): |
| 24 | return [2**x for x in range(14) if 2**x <= max_batch_size] |
| 25 | |
| 26 | |
| 27 | def get_model_params(config): |
| 28 | if config.architectures[0] in ( |
| 29 | "DeepseekV2ForCausalLM", |
| 30 | "DeepseekV3ForCausalLM", |
| 31 | "DeepseekV32ForCausalLM", |
| 32 | ): |
| 33 | num_experts = config.n_routed_experts |
| 34 | hidden_size = config.hidden_size |
| 35 | elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[ |
| 36 | 0 |
| 37 | ] in ("MiniMaxM2ForCausalLM",): |
| 38 | num_experts = config.num_local_experts |
| 39 | hidden_size = config.hidden_size |
| 40 | else: |
| 41 | raise ValueError(f"Unsupported architecture: {config.architectures}") |
| 42 | return num_experts, hidden_size |
| 43 | |
| 44 | |
| 45 | def get_benchmark(model, max_batch_size, trust_remote_code): |
| 46 | @triton.testing.perf_report( |
| 47 | triton.testing.Benchmark( |
| 48 | x_names=["batch_size"], |
| 49 | x_vals=get_batch_size_range(max_batch_size), |
| 50 | x_log=False, |
| 51 | line_arg="provider", |
| 52 | line_vals=[ |
| 53 | "torch", |
| 54 | "vllm", |
| 55 | ], |
| 56 | line_names=["PyTorch", "vLLM"], |
| 57 | styles=([("blue", "-"), ("red", "-")]), |
| 58 | ylabel="TFLOPs", |
| 59 | plot_name=f"{model} router gemm throughput", |
| 60 | args={}, |
| 61 | ) |
| 62 | ) |
| 63 | def benchmark(batch_size, provider): |
| 64 | config = get_config(model=model, trust_remote_code=trust_remote_code) |
| 65 | num_experts, hidden_size = get_model_params(config) |
| 66 | |
| 67 | is_hopper_or_blackwell = current_platform.is_device_capability( |
| 68 | 90 |
| 69 | ) or current_platform.is_device_capability_family(100) |
| 70 | allow_gpt_oss_router_gemm = ( |
| 71 | is_hopper_or_blackwell |
| 72 | and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS |
| 73 | and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES |
| 74 | ) |
| 75 | is_fp32_router_model = ( |
| 76 | is_hopper_or_blackwell |
| 77 | and num_experts in FP32_SUPPORTED_NUM_EXPERTS |
| 78 | and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES |
| 79 | ) |
| 80 | allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS |
| 81 | |
| 82 | # Weight dtype: fp32 kernel requires fp32 weights; others use bf16. |
| 83 | weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16 |
| 84 | mat_a = torch.randn( |
| 85 | (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda" |
| 86 | ).contiguous() |
| 87 | mat_b = torch.randn( |
| 88 | (num_experts, hidden_size), dtype=weight_dtype, device="cuda" |
| 89 | ).contiguous() |
| 90 | bias = torch.randn( |
| 91 | num_experts, dtype=torch.bfloat16, device="cuda" |
| 92 | ).contiguous() |
| 93 | |
| 94 | has_bias = allow_gpt_oss_router_gemm |
| 95 | |
| 96 | quantiles = [0.5, 0.2, 0.8] |
| 97 | |
| 98 | if provider == "torch": |
| 99 | |
| 100 | def runner(): |
| 101 | if allow_fp32_router_gemm: |
| 102 | F.linear(mat_a.float(), mat_b) |
| 103 | elif has_bias: |
| 104 | F.linear(mat_a, mat_b, bias) |
| 105 | else: |
| 106 | F.linear(mat_a, mat_b) |
| 107 | elif provider == "vllm": |
| 108 | |
| 109 | def runner(): |
| 110 | if allow_fp32_router_gemm: |
| 111 | ops.fp32_router_gemm(mat_a, mat_b) |
| 112 | elif allow_gpt_oss_router_gemm: |
| 113 | ops.gpt_oss_router_gemm(mat_a, mat_b, bias) |
| 114 | elif is_fp32_router_model: |
| 115 | # batch_size > FP32_MAX_TOKENS: fall back to F.linear |
| 116 | F.linear(mat_a.float(), mat_b) |
| 117 | else: |
| 118 | F.linear(mat_a, mat_b) |
| 119 | |
| 120 | ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( |
| 121 | runner, quantiles=quantiles |
| 122 | ) |
| 123 | |
| 124 | def tflops(t_ms): |
| 125 | flops = 2 * batch_size * hidden_size * num_experts |
| 126 | return flops / (t_ms * 1e-3) / 1e12 |
| 127 | |
| 128 | return tflops(ms), tflops(max_ms), tflops(min_ms) |
| 129 | |
| 130 | return benchmark |
| 131 | |
| 132 | |
| 133 | if __name__ == "__main__": |
| 134 | parser = FlexibleArgumentParser() |
| 135 | parser.add_argument("--model", type=str, default="openai/gpt-oss-20b") |
| 136 | parser.add_argument("--max-batch-size", default=16, type=int) |
| 137 | parser.add_argument("--trust-remote-code", action="store_true") |
| 138 | args = parser.parse_args() |
| 139 | |
| 140 | # Get the benchmark function |
| 141 | benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code) |
| 142 | # Run performance benchmark |
| 143 | benchmark.run(print_data=True) |