A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 16:30 UTC
[cpu][bench] Add CPU paged attention benchmarks (#31720)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Jan 6, 2026 at 10:57 UTC
799b5721f6bb89ad03e47ba79131e0ce65023635
1 file changed
+272
benchmarks/kernels/cpu/benchmark_cpu_attn.py
new
+272
@@ -0,0 +1,272 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import functools
5
+import time
6
+
7
+import numpy as np
8
+import torch
9
+
10
+from vllm._custom_ops import (
11
+ cpu_attention_with_kv_cache,
12
+ cpu_attn_get_scheduler_metadata,
13
+ cpu_attn_reshape_and_cache,
14
+)
15
+from vllm.platforms import CpuArchEnum, current_platform
16
+from vllm.utils.argparse_utils import FlexibleArgumentParser
17
+from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
18
+from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend, _get_attn_isa
19
+
20
+
21
+def get_attn_isa(
22
+ block_size: int | None = None,
23
+ dtype: torch.dtype | None = None,
24
+):
25
+ if block_size and dtype:
26
+ return _get_attn_isa(dtype, block_size)
27
+ else:
28
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
29
+ return "neon"
30
+ elif torch._C._cpu._is_amx_tile_supported():
31
+ return "amx"
32
+ else:
33
+ return "vec"
34
+
35
+
36
+# rand number generation takes too much time, cache rand tensors
37
+@functools.lru_cache(maxsize=128, typed=False)
38
+def tensor_cache(
39
+ elem_num: int,
40
+ dtype: torch.dtype,
41
+) -> torch.Tensor:
42
+ tensor = torch.randn(elem_num, dtype=dtype)
43
+ return tensor
44
+
45
+
46
+@torch.inference_mode()
47
+def main(
48
+ seq_lens: list[tuple[int, int]],
49
+ num_heads: tuple[int, int],
50
+ head_size: int,
51
+ sliding_window: int = None,
52
+ dtype: torch.dtype = torch.bfloat16,
53
+ block_size: int = 128,
54
+ num_blocks: int = 4096,
55
+ use_sink: bool = False,
56
+ enable_kv_split: bool = False,
57
+ isa: str | None = None,
58
+ seed: int = 0,
59
+ iters: int = 20,
60
+) -> None:
61
+ current_platform.seed_everything(seed)
62
+ num_seqs = len(seq_lens)
63
+ query_lens = [x[0] for x in seq_lens]
64
+ kv_lens = [x[1] for x in seq_lens]
65
+ num_query_heads = num_heads[0]
66
+ num_kv_heads = num_heads[1]
67
+ assert num_query_heads % num_kv_heads == 0
68
+ max_kv_len = max(kv_lens)
69
+ window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
70
+ scale = head_size**-0.5
71
+ token_num = sum(query_lens)
72
+
73
+ if isa is None:
74
+ isa = get_attn_isa(block_size, dtype)
75
+
76
+ s_aux = (
77
+ 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
78
+ )
79
+
80
+ query = tensor_cache(
81
+ elem_num=token_num * num_query_heads * head_size,
82
+ dtype=dtype,
83
+ )
84
+ query = query.view(
85
+ token_num,
86
+ num_query_heads,
87
+ head_size,
88
+ )
89
+
90
+ key_value = tensor_cache(
91
+ elem_num=2 * num_blocks * num_kv_heads * block_size * head_size,
92
+ dtype=dtype,
93
+ )
94
+ key_value = key_value.view(
95
+ 2,
96
+ num_blocks,
97
+ block_size,
98
+ num_kv_heads,
99
+ head_size,
100
+ )
101
+ key_cache, value_cache = key_value.unbind(0)
102
+
103
+ # KV cache for CPU attention
104
+ packed_key_cache = torch.empty(
105
+ num_blocks, num_kv_heads, block_size, head_size, dtype=dtype
106
+ )
107
+ packed_value_cache = torch.empty_like(packed_key_cache)
108
+
109
+ cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
110
+ dim=0, dtype=torch.int32
111
+ )
112
+ kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
113
+ max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
114
+ block_tables = torch.randint(
115
+ 0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
116
+ )
117
+
118
+ # use reshape_and_cache to pack key_cache and value_cache
119
+ slot_mapping = torch.arange(0, num_blocks * block_size, dtype=torch.int64)
120
+ cpu_attn_reshape_and_cache(
121
+ key=key_cache.view(-1, num_kv_heads, head_size),
122
+ value=value_cache.view(-1, num_kv_heads, head_size),
123
+ key_cache=packed_key_cache,
124
+ value_cache=packed_value_cache,
125
+ slot_mapping=slot_mapping,
126
+ isa=isa,
127
+ )
128
+
129
+ metadata = cpu_attn_get_scheduler_metadata(
130
+ num_reqs=num_seqs,
131
+ num_heads=num_query_heads,
132
+ num_kv_heads=num_kv_heads,
133
+ head_dim=head_size,
134
+ seq_lens=kv_lens_tensor,
135
+ dtype=dtype,
136
+ query_start_loc=cu_query_lens,
137
+ causal=True,
138
+ sliding_window_size=sliding_window if sliding_window is not None else -1,
139
+ isa=isa,
140
+ enable_kv_split=enable_kv_split,
141
+ )
142
+
143
+ out_with_split = torch.empty_like(query)
144
+
145
+ def run_benchmark(iters: int) -> list[float]:
146
+ times = []
147
+ for _ in range(iters):
148
+ start_time = time.perf_counter_ns()
149
+ cpu_attention_with_kv_cache(
150
+ query=query,
151
+ key_cache=packed_key_cache,
152
+ value_cache=packed_value_cache,
153
+ output=out_with_split,
154
+ query_start_loc=cu_query_lens,
155
+ seq_lens=kv_lens_tensor,
156
+ scale=scale,
157
+ causal=True,
158
+ alibi_slopes=None,
159
+ sliding_window=window_size,
160
+ block_table=block_tables,
161
+ softcap=0,
162
+ scheduler_metadata=metadata,
163
+ s_aux=s_aux,
164
+ )
165
+ end_time = time.perf_counter_ns()
166
+ times.append((end_time - start_time) / 1e6)
167
+ return times
168
+
169
+ # warmup
170
+ run_benchmark(5)
171
+ # benchmark
172
+ times = run_benchmark(iters)
173
+
174
+ time_min = min(times)
175
+ time_max = max(times)
176
+ time_mean = np.mean(times)
177
+ time_std = np.std(times)
178
+
179
+ print("\tmin (ms) = ", time_min)
180
+ print("\tmax (ms) = ", time_max)
181
+ print("\tmean (ms) = ", time_mean)
182
+ print("\tstd = ", time_std)
183
+ print("\tmedian (ms) = ", np.median(times))
184
+
185
+
186
+def generate_seq_lens(
187
+ batch_size: int,
188
+ q_len_min: int,
189
+ q_len_max: int,
190
+ kv_len_min: int,
191
+ kv_len_max: int,
192
+ seed: int = 0,
193
+) -> list[tuple[int, int]]:
194
+ assert 1 <= q_len_min <= q_len_max
195
+ assert 1 <= kv_len_min <= kv_len_max
196
+ assert kv_len_max >= q_len_min
197
+
198
+ g = torch.Generator(device="cpu").manual_seed(seed)
199
+
200
+ def rint(lo: int, hi: int) -> int:
201
+ return torch.randint(lo, hi + 1, (1,), generator=g).item()
202
+
203
+ seq_lens: list[tuple[int, int]] = []
204
+ for _ in range(batch_size):
205
+ # ensure q <= kv
206
+ kv = rint(max(kv_len_min, q_len_min), kv_len_max)
207
+ q = rint(q_len_min, min(q_len_max, kv))
208
+ seq_lens.append((q, kv))
209
+
210
+ return seq_lens
211
+
212
+
213
+if __name__ == "__main__":
214
+ parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
215
+ parser.add_argument("--batch-size", type=int, default=64)
216
+ parser.add_argument("--q-len-min", type=int, default=512)
217
+ parser.add_argument("--q-len-max", type=int, default=512)
218
+ parser.add_argument("--kv-len-min", type=int, default=512)
219
+ parser.add_argument("--kv-len-max", type=int, default=512)
220
+ parser.add_argument("--num-blocks", type=int, default=4096)
221
+
222
+ parser.add_argument("--sliding-window", type=int, default=None)
223
+ parser.add_argument("--num-query-heads", type=int, default=32)
224
+ parser.add_argument("--num-kv-heads", type=int, default=8)
225
+ parser.add_argument(
226
+ "--head-size",
227
+ type=int,
228
+ choices=CPUAttentionBackend.get_supported_head_sizes(),
229
+ default=128,
230
+ )
231
+ parser.add_argument("--enable-kv-split", action="store_true")
232
+ parser.add_argument("--block-size", type=int, choices=[32, 64, 128], default=128)
233
+ parser.add_argument(
234
+ "--dtype", type=str, choices=["half", "bfloat16", "float"], default="bfloat16"
235
+ )
236
+ parser.add_argument("--use-sink", action="store_true")
237
+ parser.add_argument(
238
+ "--isa", type=str, choices=["vec", "neon", "amx", "vec16"], default=None
239
+ )
240
+ parser.add_argument("--seed", type=int, default=0)
241
+ parser.add_argument("--iters", type=int, default=20)
242
+
243
+ args = parser.parse_args()
244
+ print(args)
245
+
246
+ seq_lens = generate_seq_lens(
247
+ args.batch_size,
248
+ args.q_len_min,
249
+ args.q_len_max,
250
+ args.kv_len_min,
251
+ args.kv_len_max,
252
+ args.seed,
253
+ )
254
+
255
+ print("batch (query len, kv len) = ", seq_lens)
256
+
257
+ main(
258
+ seq_lens=seq_lens,
259
+ num_heads=(args.num_query_heads, args.num_kv_heads),
260
+ head_size=args.head_size,
261
+ sliding_window=args.sliding_window,
262
+ dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
263
+ block_size=args.block_size,
264
+ num_blocks=args.num_blocks,
265
+ use_sink=args.use_sink,
266
+ enable_kv_split=args.enable_kv_split,
267
+ isa=args.isa
268
+ if args.isa is not None
269
+ else get_attn_isa(args.block_size, STR_DTYPE_TO_TORCH_DTYPE[args.dtype]),
270
+ seed=args.seed,
271
+ iters=args.iters,
272
+ )