@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 209 lines 6.1 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_flash,
15 set_random_seed,
16 )
17 from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
18 triton_reshape_and_cache_flash,
19 )
20
21 logger = init_logger(__name__)
22
23
24 @torch.inference_mode()
25 def run_benchmark(
26 num_tokens: int,
27 num_heads: int,
28 head_size: int,
29 block_size: int,
30 num_blocks: int,
31 dtype: torch.dtype,
32 kv_cache_dtype: str,
33 kv_cache_layout: str,
34 num_iters: int,
35 implementation: str,
36 benchmark_mode: str,
37 device: str = "cuda",
38 ) -> float:
39 """Return latency (seconds) for given num_tokens."""
40 if kv_cache_dtype == "fp8" and head_size % 16:
41 raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.")
42
43 if implementation not in ("cuda", "triton"):
44 raise ValueError(
45 f"Unsupported implementation: {implementation}. "
46 "Only 'cuda' and 'triton' are supported."
47 )
48 if implementation == "triton" and kv_cache_layout == "HND":
49 return float("nan") # Triton does not support HND layout yet.
50
51 set_random_seed(42)
52 torch.set_default_device(device)
53
54 # create random key / value tensors [T, H, D].
55 key = torch.randn(num_tokens, num_heads, head_size, dtype=dtype, device=device)
56 value = torch.randn_like(key)
57
58 # prepare the slot mapping.
59 # each token is assigned a unique slot in the KV-cache.
60 num_slots = block_size * num_blocks
61 if num_tokens > num_slots:
62 raise ValueError("num_tokens cannot exceed the total number of cache slots")
63 slot_mapping_lst = random.sample(range(num_slots), num_tokens)
64 slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
65
66 key_caches, value_caches = create_kv_caches_with_random_flash(
67 num_blocks,
68 block_size,
69 1, # num_layers
70 num_heads,
71 head_size,
72 kv_cache_dtype,
73 dtype,
74 device=device,
75 cache_layout=kv_cache_layout,
76 )
77 key_cache, value_cache = key_caches[0], value_caches[0]
78 # to free unused memory
79 del key_caches, value_caches
80
81 # compute per-kernel scaling factors for fp8 conversion (if used).
82 k_scale = (key.amax() / 64.0).to(torch.float32)
83 v_scale = (value.amax() / 64.0).to(torch.float32)
84
85 if implementation == "cuda":
86 function_under_test = lambda: ops.reshape_and_cache_flash(
87 key, # noqa: F821
88 value, # noqa: F821
89 key_cache, # noqa: F821
90 value_cache, # noqa: F821
91 slot_mapping, # noqa: F821
92 kv_cache_dtype,
93 k_scale,
94 v_scale,
95 )
96 else:
97 function_under_test = lambda: triton_reshape_and_cache_flash(
98 key, # noqa: F821
99 value, # noqa: F821
100 key_cache, # noqa: F821
101 value_cache, # noqa: F821
102 slot_mapping, # noqa: F821
103 kv_cache_dtype,
104 k_scale,
105 v_scale,
106 )
107 if benchmark_mode == "cudagraph":
108 g = torch.cuda.CUDAGraph()
109 with torch.cuda.graph(g):
110 function_under_test()
111 torch.accelerator.synchronize()
112 function_under_test = lambda: g.replay()
113
114 def run_cuda_benchmark(n_iters: int) -> float:
115 nonlocal key, value, key_cache, value_cache, slot_mapping
116 torch.accelerator.synchronize()
117 start = time.perf_counter()
118 for _ in range(n_iters):
119 function_under_test()
120 torch.accelerator.synchronize()
121 end = time.perf_counter()
122 return (end - start) / n_iters
123
124 # warm-up
125 run_cuda_benchmark(3)
126
127 lat = run_cuda_benchmark(num_iters)
128
129 # free tensors to mitigate OOM when sweeping
130 del key, value, key_cache, value_cache, slot_mapping
131 torch.accelerator.empty_cache()
132
133 return lat
134
135
136 def main(args):
137 rows = []
138 for layout in ["NHD", "HND"]:
139 for exp in range(1, 17):
140 n_tok = 2**exp
141 lat = run_benchmark(
142 num_tokens=n_tok,
143 num_heads=args.num_heads,
144 head_size=args.head_size,
145 block_size=args.block_size,
146 num_blocks=args.num_blocks,
147 dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
148 kv_cache_dtype=args.kv_cache_dtype,
149 kv_cache_layout=layout,
150 num_iters=args.iters,
151 implementation=args.implementation,
152 benchmark_mode=args.mode,
153 device="cuda",
154 )
155 rows.append([n_tok, layout, f"{lat * 1e6:.3f}"])
156
157 print(
158 f"Benchmark results for implementation {args.implementation}"
159 f" (measuring with {args.mode}):"
160 )
161 print(tabulate(rows, headers=["num_tokens", "layout", "latency (µs)"]))
162
163
164 if __name__ == "__main__":
165 parser = FlexibleArgumentParser()
166
167 parser.add_argument("--num-heads", type=int, default=128)
168 parser.add_argument(
169 "--head-size",
170 type=int,
171 choices=[64, 80, 96, 112, 120, 128, 192, 256],
172 default=128,
173 )
174 parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
175 parser.add_argument("--num-blocks", type=int, default=128 * 512)
176
177 parser.add_argument(
178 "--dtype",
179 type=str,
180 choices=["half", "bfloat16", "float"],
181 default="bfloat16",
182 )
183
184 parser.add_argument(
185 "--kv-cache-dtype",
186 type=str,
187 choices=["auto", "fp8"],
188 default="auto",
189 )
190
191 parser.add_argument("--iters", type=int, default=100)
192
193 parser.add_argument(
194 "--implementation",
195 type=str,
196 choices=["cuda", "triton"],
197 default="cuda",
198 )
199
200 parser.add_argument(
201 "--mode",
202 type=str,
203 choices=["cudagraph", "no_graph"],
204 default="cudagraph",
205 )
206
207 args = parser.parse_args()
208
209 main(args)