@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 171 lines 4.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import random
4 import time
5
6 import torch
7 from tabulate import tabulate
8
9 from vllm import _custom_ops as ops
10 from vllm.logger import init_logger
11 from vllm.utils.argparse_utils import FlexibleArgumentParser
12 from vllm.utils.torch_utils import (
13 STR_DTYPE_TO_TORCH_DTYPE,
14 create_kv_caches_with_random,
15 set_random_seed,
16 )
17
18 logger = init_logger(__name__)
19
20
21 @torch.inference_mode()
22 def run_benchmark(
23 num_tokens: int,
24 num_heads: int,
25 head_size: int,
26 block_size: int,
27 num_blocks: int,
28 dtype: torch.dtype,
29 kv_cache_dtype: str,
30 num_iters: int,
31 benchmark_mode: str,
32 device: str = "cuda",
33 ) -> float:
34 """Return latency (seconds) for given num_tokens."""
35 if kv_cache_dtype == "fp8" and head_size % 16:
36 raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
37
38 set_random_seed(42)
39 torch.set_default_device(device)
40
41 # create random key / value tensors [T, H, D].
42 key = torch.randn(num_tokens, num_heads, head_size, dtype=dtype, device=device)
43 value = torch.randn_like(key)
44
45 # prepare the slot mapping.
46 # each token is assigned a unique slot in the KV-cache.
47 num_slots = block_size * num_blocks
48 if num_tokens > num_slots:
49 raise ValueError("num_tokens cannot exceed the total number of cache slots")
50 slot_mapping_lst = random.sample(range(num_slots), num_tokens)
51 slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
52
53 key_caches, value_caches = create_kv_caches_with_random(
54 num_blocks,
55 block_size,
56 1, # num_layers
57 num_heads,
58 head_size,
59 kv_cache_dtype,
60 dtype,
61 device=device,
62 )
63 key_cache, value_cache = key_caches[0], value_caches[0]
64 # to free unused memory
65 del key_caches, value_caches
66
67 # compute per-kernel scaling factors for fp8 conversion (if used).
68 k_scale = (key.amax() / 64.0).to(torch.float32)
69 v_scale = (value.amax() / 64.0).to(torch.float32)
70
71 function_under_test = lambda: ops.reshape_and_cache(
72 key, # noqa: F821
73 value, # noqa: F821
74 key_cache, # noqa: F821
75 value_cache, # noqa: F821
76 slot_mapping, # noqa: F821
77 kv_cache_dtype,
78 k_scale,
79 v_scale,
80 )
81
82 if benchmark_mode == "cudagraph":
83 g = torch.cuda.CUDAGraph()
84 with torch.cuda.graph(g):
85 function_under_test()
86 torch.accelerator.synchronize()
87 function_under_test = lambda: g.replay()
88
89 def run_cuda_benchmark(n_iters: int) -> float:
90 nonlocal key, value, key_cache, value_cache, slot_mapping
91 torch.accelerator.synchronize()
92 start = time.perf_counter()
93 for _ in range(n_iters):
94 function_under_test()
95 torch.accelerator.synchronize()
96 end = time.perf_counter()
97 return (end - start) / n_iters
98
99 # warm-up
100 run_cuda_benchmark(3)
101
102 lat = run_cuda_benchmark(num_iters)
103
104 # free tensors to mitigate OOM when sweeping
105 del key, value, key_cache, value_cache, slot_mapping
106 torch.accelerator.empty_cache()
107
108 return lat
109
110
111 def main(args):
112 rows = []
113 for exp in range(1, 17):
114 n_tok = 2**exp
115 lat = run_benchmark(
116 num_tokens=n_tok,
117 num_heads=args.num_heads,
118 head_size=args.head_size,
119 block_size=args.block_size,
120 num_blocks=args.num_blocks,
121 dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
122 kv_cache_dtype=args.kv_cache_dtype,
123 num_iters=args.iters,
124 benchmark_mode=args.mode,
125 device="cuda",
126 )
127 rows.append([n_tok, lat * 1e6]) # convert to microseconds
128
129 print(f"Benchmark results for implementation cuda (measuring with {args.mode}):")
130 print(tabulate(rows, headers=["num_tokens", "latency (µs)"], floatfmt=".3f"))
131
132
133 if __name__ == "__main__":
134 parser = FlexibleArgumentParser()
135
136 parser.add_argument("--num-heads", type=int, default=128)
137 parser.add_argument(
138 "--head-size",
139 type=int,
140 choices=[64, 80, 96, 112, 120, 128, 192, 256],
141 default=128,
142 )
143 parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
144 parser.add_argument("--num-blocks", type=int, default=128 * 128)
145
146 parser.add_argument(
147 "--dtype",
148 type=str,
149 choices=["half", "bfloat16", "float"],
150 default="bfloat16",
151 )
152
153 parser.add_argument(
154 "--kv-cache-dtype",
155 type=str,
156 choices=["auto", "fp8"],
157 default="auto",
158 )
159
160 parser.add_argument("--iters", type=int, default=200)
161
162 parser.add_argument(
163 "--mode",
164 type=str,
165 choices=["cudagraph", "no_graph"],
166 default="cudagraph",
167 )
168
169 args = parser.parse_args()
170
171 main(args)