@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 169 lines 4.77 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import argparse
4 import itertools
5
6 import torch
7
8 import vllm._custom_ops as ops
9 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
10 moe_align_block_size,
11 )
12 from vllm.triton_utils import triton
13 from vllm.utils.torch_utils import set_random_seed
14
15
16 def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:
17 return torch.stack(
18 [
19 torch.randperm(num_experts, dtype=torch.int32, device="cuda")[:topk]
20 for _ in range(num_tokens)
21 ]
22 )
23
24
25 # test configurations
26 num_tokens_range = [1, 16, 256, 4096]
27 num_experts_range = [16, 64, 224, 256, 280, 512]
28 topk_range = [1, 2, 8]
29 ep_size_range = [1, 8]
30 configs = list(
31 itertools.product(num_tokens_range, num_experts_range, topk_range, ep_size_range)
32 )
33
34
35 @triton.testing.perf_report(
36 triton.testing.Benchmark(
37 x_names=["num_tokens", "num_experts", "topk", "ep_size"],
38 x_vals=configs,
39 line_arg="provider",
40 line_vals=["vllm"],
41 line_names=["vLLM"],
42 plot_name="moe-align-block-size-performance",
43 args={},
44 )
45 )
46 def benchmark(num_tokens, num_experts, topk, ep_size, provider):
47 """Benchmark function for Triton."""
48 block_size = 256
49 set_random_seed(0)
50 topk_ids = get_topk_ids(num_tokens, num_experts, topk)
51
52 e_map = None
53 if ep_size != 1:
54 local_e = num_experts // ep_size
55 e_ids = torch.randperm(num_experts, device="cuda", dtype=torch.int32)[:local_e]
56 e_map = torch.full((num_experts,), -1, device="cuda", dtype=torch.int32)
57 e_map[e_ids] = torch.arange(local_e, device="cuda", dtype=torch.int32)
58
59 quantiles = [0.5, 0.2, 0.8]
60
61 if provider == "vllm":
62 ms, min_ms, max_ms = triton.testing.do_bench(
63 lambda: moe_align_block_size(
64 topk_ids, block_size, num_experts, e_map, ignore_invalid_experts=True
65 ),
66 quantiles=quantiles,
67 )
68
69 return 1000 * ms, 1000 * max_ms, 1000 * min_ms
70
71
72 batched_configs = list(
73 itertools.product(
74 [8, 32, 64, 256],
75 [16, 256, 512, 2048],
76 [8, 32, 48, 64],
77 ["uniform", "skewed", "mostly_empty"],
78 )
79 )
80
81
82 @triton.testing.perf_report(
83 triton.testing.Benchmark(
84 x_names=[
85 "num_batches",
86 "max_tokens_per_batch",
87 "block_size",
88 "distribution",
89 ],
90 x_vals=batched_configs,
91 line_arg="provider",
92 line_vals=["vllm"],
93 line_names=["vLLM"],
94 plot_name="batched-moe-align-block-size-performance",
95 args={},
96 )
97 )
98 def benchmark_batched(
99 num_batches,
100 max_tokens_per_batch,
101 block_size,
102 distribution,
103 provider,
104 ):
105 set_random_seed(0)
106 if distribution == "uniform":
107 expert_num_tokens = torch.full(
108 (num_batches,),
109 max_tokens_per_batch,
110 dtype=torch.int32,
111 device="cuda",
112 )
113 elif distribution == "skewed":
114 expert_num_tokens = torch.ones((num_batches,), dtype=torch.int32, device="cuda")
115 expert_num_tokens[0] = max_tokens_per_batch
116 else:
117 expert_num_tokens = torch.zeros(
118 (num_batches,), dtype=torch.int32, device="cuda"
119 )
120 expert_num_tokens[::8] = max_tokens_per_batch
121
122 max_num_tokens_padded = (
123 num_batches
124 * ((max_tokens_per_batch + block_size - 1) // block_size)
125 * block_size
126 )
127 sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
128 expert_ids = torch.empty(
129 (max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda"
130 )
131 num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda")
132
133 quantiles = [0.5, 0.2, 0.8]
134 if provider == "vllm":
135 ms, min_ms, max_ms = triton.testing.do_bench(
136 lambda: ops.batched_moe_align_block_size(
137 max_tokens_per_batch,
138 block_size,
139 expert_num_tokens,
140 sorted_ids,
141 expert_ids,
142 num_tokens_post_pad,
143 ),
144 quantiles=quantiles,
145 )
146
147 return 1000 * ms, 1000 * max_ms, 1000 * min_ms
148
149
150 if __name__ == "__main__":
151 parser = argparse.ArgumentParser()
152 parser.add_argument("--batched", action="store_true")
153 parser.add_argument(
154 "--num_experts",
155 type=int,
156 default=64,
157 choices=[8, 16, 32, 64, 128, 256],
158 )
159 parser.add_argument(
160 "--topk",
161 type=int,
162 default=8,
163 choices=[2, 4, 8],
164 help="Top-k value for correctness check.",
165 )
166 args = parser.parse_args()
167
168 selected_benchmark = benchmark_batched if args.batched else benchmark
169 selected_benchmark.run(print_data=True, show_plots=True)