@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 197 lines 6.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Benchmark the TP > 1 input-embedding path of VocabParallelEmbedding.
4
5 Two providers, both computing one rank's partial embedding (the all-reduce
6 that follows is identical for each and is not measured):
7
8 unfused mask + id shift + int64 cast + F.embedding + masked_fill_
9 fused the single fused CUDA kernel, _C::vocab_parallel_embedding
10
11 With --use-cuda-graph each provider is captured into a CUDA graph and the
12 replay is timed, which is what a decode step actually costs.
13 """
14
15 import itertools
16
17 import torch
18
19 from vllm import _custom_ops as vllm_ops
20 from vllm.model_executor.layers.vocab_parallel_embedding import (
21 VocabParallelEmbedding,
22 get_masked_input_and_mask,
23 pad_vocab_size,
24 )
25 from vllm.triton_utils import triton
26
27
28 def get_shard_indices(vocab_size, org_vocab_size, tp_rank, tp_size):
29 return VocabParallelEmbedding._get_indices(
30 pad_vocab_size(vocab_size),
31 pad_vocab_size(org_vocab_size),
32 vocab_size,
33 org_vocab_size,
34 tp_rank,
35 tp_size,
36 )
37
38
39 def embedding_reference(input_ids, weight, shard_indices):
40 masked_input, input_mask = get_masked_input_and_mask(
41 input_ids,
42 shard_indices.org_vocab_start_index,
43 shard_indices.org_vocab_end_index,
44 shard_indices.num_org_vocab_padding,
45 shard_indices.added_vocab_start_index,
46 shard_indices.added_vocab_end_index,
47 )
48 output = torch.nn.functional.embedding(masked_input.long(), weight)
49 return output.masked_fill_(input_mask.unsqueeze(-1), 0)
50
51
52 def embedding_fused(input_ids, weight, shard_indices):
53 return vllm_ops.vocab_parallel_embedding(
54 input_ids,
55 weight,
56 shard_indices.org_vocab_start_index,
57 shard_indices.org_vocab_end_index,
58 shard_indices.num_org_vocab_padding,
59 shard_indices.added_vocab_start_index,
60 shard_indices.added_vocab_end_index,
61 )
62
63
64 PROVIDERS = {
65 "unfused": embedding_reference,
66 "fused": embedding_fused,
67 }
68
69
70 def make_inputs(vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank=0):
71 shard_indices = get_shard_indices(vocab_size, vocab_size, tp_rank, tp_size)
72 weight = torch.randn(
73 pad_vocab_size(vocab_size) // tp_size,
74 hidden_size,
75 dtype=dtype,
76 device="cuda",
77 )
78 input_ids = torch.randint(
79 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
80 )
81 return input_ids, weight, shard_indices
82
83
84 def graph_replay(fn):
85 """Capture fn into a CUDA graph and return its replay."""
86 stream = torch.cuda.Stream()
87 stream.wait_stream(torch.cuda.current_stream())
88 with torch.cuda.stream(stream):
89 for _ in range(3):
90 fn()
91 torch.cuda.current_stream().wait_stream(stream)
92 graph = torch.cuda.CUDAGraph()
93 with torch.cuda.graph(graph):
94 fn()
95 return graph.replay
96
97
98 def calculate_diff(vocab_size, hidden_size, tp_size, num_tokens):
99 """The fused kernel must be bit-exact with the path it replaces, on every
100 rank of the shard."""
101 dtype = torch.bfloat16
102 for tp_rank in range(tp_size):
103 input_ids, weight, shard_indices = make_inputs(
104 vocab_size, hidden_size, tp_size, num_tokens, dtype, tp_rank
105 )
106 output_reference = embedding_reference(input_ids, weight, shard_indices)
107 output_fused = embedding_fused(input_ids, weight, shard_indices)
108 if not torch.equal(output_reference, output_fused):
109 print(f"❌ rank {tp_rank}: fused output differs")
110 return
111 print(f"✅ fused matches the reference path on all {tp_size} ranks")
112
113
114 tp_size_range = [2, 4, 8]
115 num_tokens_range = [1, 8, 64, 512, 2048, 8192, 32768]
116 configs = list(itertools.product(tp_size_range, num_tokens_range))
117
118
119 def get_benchmark(vocab_size, hidden_size, dtype, use_cuda_graph):
120 @triton.testing.perf_report(
121 triton.testing.Benchmark(
122 x_names=["tp_size", "num_tokens"],
123 x_vals=[list(_) for _ in configs],
124 line_arg="provider",
125 line_vals=list(PROVIDERS),
126 line_names=["Unfused PyTorch", "Fused kernel"],
127 styles=[("green", "-"), ("red", "-")],
128 ylabel="us",
129 plot_name=(
130 f"vocab-parallel-embedding-perf-vocab{vocab_size}-"
131 f"hidden{hidden_size}{'-cudagraph' if use_cuda_graph else ''}"
132 ),
133 args={},
134 )
135 )
136 def benchmark(tp_size, num_tokens, provider):
137 input_ids, weight, shard_indices = make_inputs(
138 vocab_size, hidden_size, tp_size, num_tokens, dtype
139 )
140 fn = PROVIDERS[provider]
141
142 def run():
143 return fn(input_ids, weight, shard_indices)
144
145 if use_cuda_graph:
146 run = graph_replay(run)
147
148 ms, min_ms, max_ms = triton.testing.do_bench(run, quantiles=[0.5, 0.2, 0.8])
149 return 1000 * ms, 1000 * max_ms, 1000 * min_ms
150
151 return benchmark
152
153
154 if __name__ == "__main__":
155 import argparse
156
157 parser = argparse.ArgumentParser()
158 parser.add_argument(
159 "--vocab-size", type=int, default=128256, help="Vocabulary size"
160 )
161 parser.add_argument(
162 "--hidden-size", type=int, default=4096, help="Embedding dimension"
163 )
164 parser.add_argument(
165 "--dtype", type=str, default="bfloat16", help="Embedding table dtype"
166 )
167 parser.add_argument(
168 "--tp-size",
169 type=int,
170 default=4,
171 help="TP size used for the correctness check",
172 )
173 parser.add_argument(
174 "--use-cuda-graph",
175 action="store_true",
176 help="Time CUDA graph replay instead of eager launches",
177 )
178 parser.add_argument(
179 "--save-path",
180 type=str,
181 default=None,
182 help="Optional path to save benchmark results",
183 )
184 args = parser.parse_args()
185 dtype = getattr(torch, args.dtype)
186
187 calculate_diff(
188 vocab_size=args.vocab_size,
189 hidden_size=args.hidden_size,
190 tp_size=args.tp_size,
191 num_tokens=1024,
192 )
193
194 benchmark = get_benchmark(
195 args.vocab_size, args.hidden_size, dtype, args.use_cuda_graph
196 )
197 benchmark.run(print_data=True, save_path=args.save_path)