@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 196 lines 5.96 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import random
5 import time
6
7 import torch
8
9 from vllm import _custom_ops as ops
10 from vllm.logger import init_logger
11 from vllm.platforms import current_platform
12 from vllm.utils.argparse_utils import FlexibleArgumentParser
13 from vllm.utils.torch_utils import (
14 STR_DTYPE_TO_TORCH_DTYPE,
15 create_kv_caches_with_random,
16 set_random_seed,
17 )
18
19 logger = init_logger(__name__)
20
21 NUM_BLOCKS = 128 * 1024
22 PARTITION_SIZE_ROCM = 256
23
24
25 @torch.inference_mode()
26 def main(
27 num_seqs: int,
28 seq_len: int,
29 num_query_heads: int,
30 num_kv_heads: int,
31 head_size: int,
32 use_alibi: bool,
33 block_size: int,
34 dtype: torch.dtype,
35 seed: int,
36 do_profile: bool,
37 device: str = "cuda",
38 kv_cache_dtype: str | None = None,
39 ) -> None:
40 set_random_seed(seed)
41
42 scale = float(1.0 / (head_size**0.5))
43 query = torch.empty(
44 num_seqs, num_query_heads, head_size, dtype=dtype, device=device
45 )
46 query.uniform_(-scale, scale)
47
48 assert num_query_heads % num_kv_heads == 0
49 alibi_slopes = None
50 if use_alibi:
51 alibi_slopes = torch.randn(num_query_heads, dtype=torch.float, device=device)
52
53 seq_lens = [seq_len for _ in range(num_seqs)]
54 max_seq_len = max(seq_lens)
55 seq_lens = torch.tensor(seq_lens, dtype=torch.int, device=device)
56
57 # Create the block tables.
58 max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
59 block_tables_lst: list[list[int]] = []
60 for _ in range(num_seqs):
61 block_table = [
62 random.randint(0, NUM_BLOCKS - 1) for _ in range(max_num_blocks_per_seq)
63 ]
64 block_tables_lst.append(block_table)
65
66 block_tables = torch.tensor(block_tables_lst, dtype=torch.int, device=device)
67
68 # Create the KV cache.
69 key_caches, value_caches = create_kv_caches_with_random(
70 NUM_BLOCKS,
71 block_size,
72 1,
73 num_kv_heads,
74 head_size,
75 kv_cache_dtype,
76 dtype,
77 device=device,
78 )
79 key_cache, value_cache = key_caches[0], value_caches[0]
80
81 # Prepare for the paged attention kernel.
82 output = torch.empty_like(query)
83 num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM
84 tmp_output = torch.empty(
85 size=(num_seqs, num_query_heads, num_partitions, head_size),
86 dtype=output.dtype,
87 device=output.device,
88 )
89 exp_sums = torch.empty(
90 size=(num_seqs, num_query_heads, num_partitions),
91 dtype=torch.float32,
92 device=output.device,
93 )
94 max_logits = torch.empty_like(exp_sums)
95
96 def run_benchmark(num_iters: int, profile: bool = False) -> float:
97 torch.accelerator.synchronize()
98 if profile:
99 torch.cuda.cudart().cudaProfilerStart()
100 start_time = time.perf_counter()
101
102 # Using default kv_scale
103 k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
104
105 for _ in range(num_iters):
106 ops.paged_attention_rocm(
107 output,
108 exp_sums,
109 max_logits,
110 tmp_output,
111 query,
112 key_cache,
113 value_cache,
114 num_kv_heads,
115 scale,
116 block_tables,
117 seq_lens,
118 None,
119 block_size,
120 max_seq_len,
121 alibi_slopes,
122 kv_cache_dtype,
123 k_scale,
124 v_scale,
125 )
126 torch.accelerator.synchronize()
127
128 end_time = time.perf_counter()
129 if profile:
130 torch.cuda.cudart().cudaProfilerStop()
131 return (end_time - start_time) / num_iters
132
133 # Warmup.
134 print("Warming up...")
135 run_benchmark(num_iters=3, profile=False)
136
137 # Benchmark.
138 if do_profile:
139 latency = run_benchmark(num_iters=1, profile=True)
140 else:
141 latency = run_benchmark(num_iters=100, profile=False)
142 print(f"Kernel running time: {latency * 1000000:.3f} us")
143
144
145 if __name__ == "__main__":
146 logger.warning(
147 "This script benchmarks the ROCm paged attention kernel. "
148 "By default this is no longer used in vLLM inference."
149 )
150 if not current_platform.is_rocm():
151 raise RuntimeError("This benchmark requires the ROCm platform.")
152
153 parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
154 parser.add_argument("--batch-size", type=int, default=8)
155 parser.add_argument("--seq-len", type=int, default=4096)
156 parser.add_argument("--num-query-heads", type=int, default=64)
157 parser.add_argument("--num-kv-heads", type=int, default=8)
158 parser.add_argument(
159 "--head-size",
160 type=int,
161 choices=[64, 128],
162 default=128,
163 )
164 parser.add_argument("--block-size", type=int, choices=[16, 32], default=16)
165 parser.add_argument("--use-alibi", action="store_true")
166 parser.add_argument(
167 "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"
168 )
169 parser.add_argument("--seed", type=int, default=0)
170 parser.add_argument("--profile", action="store_true")
171 parser.add_argument(
172 "--kv-cache-dtype",
173 type=str,
174 choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"],
175 default="auto",
176 help="Data type for kv cache storage. If 'auto', will use model "
177 "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)",
178 )
179 args = parser.parse_args()
180 print(args)
181
182 if args.num_query_heads % args.num_kv_heads != 0:
183 raise ValueError("num_query_heads must be divisible by num_kv_heads")
184 main(
185 num_seqs=args.batch_size,
186 seq_len=args.seq_len,
187 num_query_heads=args.num_query_heads,
188 num_kv_heads=args.num_kv_heads,
189 head_size=args.head_size,
190 block_size=args.block_size,
191 use_alibi=args.use_alibi,
192 dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
193 seed=args.seed,
194 do_profile=args.profile,
195 kv_cache_dtype=args.kv_cache_dtype,
196 )