| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import itertools |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk |
| 9 | from vllm.triton_utils import triton |
| 10 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 11 | |
| 12 | num_tokens_range = [2**i for i in range(0, 8, 2)] |
| 13 | num_experts_range = [16, 32, 64, 128, 256, 512] |
| 14 | topk_range = [3, 4] |
| 15 | configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) |
| 16 | |
| 17 | |
| 18 | def torch_topk( |
| 19 | gating_output: torch.Tensor, |
| 20 | topk: int, |
| 21 | renormalize: bool, |
| 22 | scoring_func: str = "softmax", |
| 23 | ): |
| 24 | if scoring_func == "softmax": |
| 25 | scores = torch.softmax(gating_output.float(), dim=-1) |
| 26 | else: |
| 27 | scores = torch.sigmoid(gating_output.float()) |
| 28 | topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1) |
| 29 | |
| 30 | if renormalize: |
| 31 | topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) |
| 32 | |
| 33 | return topk_weights, topk_ids |
| 34 | |
| 35 | |
| 36 | def get_benchmark(scoring_func): |
| 37 | @triton.testing.perf_report( |
| 38 | triton.testing.Benchmark( |
| 39 | x_names=["num_tokens", "num_experts", "topk"], |
| 40 | x_vals=[list(_) for _ in configs], |
| 41 | line_arg="provider", |
| 42 | line_vals=["torch", "vllm"], |
| 43 | line_names=["Torch", "vLLM"], |
| 44 | styles=[("blue", "-"), ("red", "-")], |
| 45 | ylabel="us", |
| 46 | plot_name=f"fused-topk-perf-{scoring_func}", |
| 47 | args={}, |
| 48 | ) |
| 49 | ) |
| 50 | def benchmark(num_tokens, num_experts, topk, provider): |
| 51 | dtype = torch.bfloat16 |
| 52 | hidden_size = 1024 |
| 53 | renormalize = True |
| 54 | hidden_states = torch.randn( |
| 55 | (num_tokens, hidden_size), dtype=dtype, device="cuda" |
| 56 | ) |
| 57 | gating_output = torch.randn( |
| 58 | (num_tokens, num_experts), dtype=dtype, device="cuda" |
| 59 | ) |
| 60 | |
| 61 | quantiles = [0.5, 0.2, 0.8] |
| 62 | |
| 63 | if provider == "torch": |
| 64 | ms, min_ms, max_ms = triton.testing.do_bench( |
| 65 | lambda: torch_topk( |
| 66 | gating_output=gating_output, |
| 67 | topk=topk, |
| 68 | renormalize=renormalize, |
| 69 | scoring_func=scoring_func, |
| 70 | ), |
| 71 | quantiles=quantiles, |
| 72 | ) |
| 73 | else: |
| 74 | ms, min_ms, max_ms = triton.testing.do_bench( |
| 75 | lambda: fused_topk( |
| 76 | hidden_states=hidden_states, |
| 77 | gating_output=gating_output, |
| 78 | topk=topk, |
| 79 | renormalize=renormalize, |
| 80 | scoring_func=scoring_func, |
| 81 | ), |
| 82 | quantiles=quantiles, |
| 83 | ) |
| 84 | |
| 85 | return 1000 * ms, 1000 * max_ms, 1000 * min_ms |
| 86 | |
| 87 | return benchmark |
| 88 | |
| 89 | |
| 90 | if __name__ == "__main__": |
| 91 | parser = FlexibleArgumentParser(description="Benchmark the MoE topk kernel.") |
| 92 | parser.add_argument("--scoring-func", type=str, default="softmax") |
| 93 | parser.add_argument("--save-path", type=str, default="./configs/fused_topk/") |
| 94 | args = parser.parse_args() |
| 95 | |
| 96 | # Get the benchmark function |
| 97 | benchmark = get_benchmark(args.scoring_func) |
| 98 | # Run performance benchmark |
| 99 | benchmark.run(print_data=True, save_path=args.save_path) |