benchmark_paged_attention.py 8.89 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
7
8
import random
import time

import torch

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
15
    STR_DTYPE_TO_TORCH_DTYPE,
    create_kv_caches_with_random,
16
    set_random_seed,
17
)
18
import vllm.envs as envs
19

20
21
logger = init_logger(__name__)

22
NUM_BLOCKS = 128 * 1024
23
PARTITION_SIZE = 512
24
PARTITION_SIZE_ROCM = 256
25
26
27
28
29
30


@torch.inference_mode()
def main(
    version: str,
    num_seqs: int,
31
    seq_len: int,
32
33
34
35
36
37
38
39
    num_query_heads: int,
    num_kv_heads: int,
    head_size: int,
    use_alibi: bool,
    block_size: int,
    dtype: torch.dtype,
    seed: int,
    do_profile: bool,
40
    device: str = "cuda",
41
    kv_cache_dtype: str | None = None,
42
) -> None:
43
    set_random_seed(seed)
44
45

    scale = float(1.0 / (head_size**0.5))
46
47
48
    query = torch.empty(
        num_seqs, num_query_heads, head_size, dtype=dtype, device=device
    )
49
50
51
52
53
    query.uniform_(-scale, scale)

    assert num_query_heads % num_kv_heads == 0
    alibi_slopes = None
    if use_alibi:
54
        alibi_slopes = torch.randn(num_query_heads, dtype=torch.float, device=device)
55

56
57
58
    seq_lens = [seq_len for _ in range(num_seqs)]
    max_seq_len = max(seq_lens)
    seq_lens = torch.tensor(seq_lens, dtype=torch.int, device=device)
59
60

    # Create the block tables.
61
    max_num_blocks_per_seq = (max_seq_len + block_size - 1) // block_size
62
    block_tables_lst: list[list[int]] = []
63
64
    for _ in range(num_seqs):
        block_table = [
65
            random.randint(0, NUM_BLOCKS - 1) for _ in range(max_num_blocks_per_seq)
66
        ]
67
68
        block_tables_lst.append(block_table)

69
    block_tables = torch.tensor(block_tables_lst, dtype=torch.int, device=device)
70
71

    # Create the KV cache.
72
73
74
75
76
77
78
79
80
81
    key_caches, value_caches = create_kv_caches_with_random(
        NUM_BLOCKS,
        block_size,
        1,
        num_kv_heads,
        head_size,
        kv_cache_dtype,
        dtype,
        device=device,
    )
82
    key_cache, value_cache = key_caches[0], value_caches[0]
83
84
85
86

    # Prepare for the paged attention kernel.
    output = torch.empty_like(query)
    if version == "v2":
87
88
        if current_platform.is_rocm():
            global PARTITION_SIZE
89
            if not args.custom_paged_attn and not current_platform.is_navi():
90
91
92
                PARTITION_SIZE = 1024
            else:
                PARTITION_SIZE = PARTITION_SIZE_ROCM
93
        num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE
94
95
96
97
98
99
100
101
102
103
104
        tmp_output = torch.empty(
            size=(num_seqs, num_query_heads, num_partitions, head_size),
            dtype=output.dtype,
            device=output.device,
        )
        exp_sums = torch.empty(
            size=(num_seqs, num_query_heads, num_partitions),
            dtype=torch.float32,
            device=output.device,
        )
        max_logits = torch.empty_like(exp_sums)
zhuwenwen's avatar
zhuwenwen committed
105
106
107
108
        
    if version == "v12":
        sliding_window = ((-1, -1))
        logits_soft_cap = 0.0
109

110
    def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float:
111
112
113
114
115
        torch.cuda.synchronize()
        if profile:
            torch.cuda.cudart().cudaProfilerStart()
        start_time = time.perf_counter()

116
        # Using default kv_scale
117
        k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
118

119
120
        for _ in range(num_iters):
            if version == "v1":
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
                ops.paged_attention_v1(
                output,
                query,
                key_cache,
                value_cache,
                num_kv_heads,
                scale,
                block_tables,
                seq_lens,
                block_size,
                max_seq_len,
                alibi_slopes,
                kv_cache_dtype,
                k_scale,
                v_scale,
            )
137
            elif version == "v2":
zhuwenwen's avatar
zhuwenwen committed
138
                if not args.custom_paged_attn:   
zhuwenwen's avatar
zhuwenwen committed
139
                    ops.paged_attention_v2(
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
                        output,
                        exp_sums,
                        max_logits,
                        tmp_output,
                        query,
                        key_cache,
                        value_cache,
                        num_kv_heads,
                        scale,
                        block_tables,
                        seq_lens,
                        block_size,
                        max_seq_len,
                        alibi_slopes,
                        kv_cache_dtype,
                        k_scale,
                        v_scale,
                    )
                else:
                    ops.paged_attention_rocm(
                        output,
                        exp_sums,
                        max_logits,
                        tmp_output,
                        query,
                        key_cache,
                        value_cache,
                        num_kv_heads,
                        scale,
                        block_tables,
                        seq_lens,
171
                        None,
172
173
174
175
176
177
178
                        block_size,
                        max_seq_len,
                        alibi_slopes,
                        kv_cache_dtype,
                        k_scale,
                        v_scale,
                    )
zhuwenwen's avatar
zhuwenwen committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
            elif version == "v12":
                from flash_attn import vllm_flash_attn_with_kvcache
                vllm_flash_attn_with_kvcache(
                    q=query.unsqueeze(1),  
                    k_cache=key_cache,  
                    v_cache=value_cache,  
                    cache_seqlens=seq_lens,  
                    block_table=block_tables, 
                    softmax_scale=scale,
                    causal=True,
                    window_size=sliding_window,  
                    softcap=logits_soft_cap,
                    alibi_slopes=alibi_slopes,
                    return_softmax_lse=False,
                    k_scale=k_scale,  
                    v_scale=v_scale, 
                    kv_cache_dtype=kv_cache_dtype,  
                ).squeeze(1) 
197
198
199
200
201
202
            else:
                raise ValueError(f"Invalid version: {version}")
        torch.cuda.synchronize()

        end_time = time.perf_counter()
        if profile:
203
            torch.cuda.cudart().cudaProfilerStop()
204
205
206
207
        return (end_time - start_time) / num_iters

    # Warmup.
    print("Warming up...")
208
    run_benchmark = run_cuda_benchmark
209
210
211
212
213
214
215
216
217
218
    run_benchmark(num_iters=3, profile=False)

    # Benchmark.
    if do_profile:
        latency = run_benchmark(num_iters=1, profile=True)
    else:
        latency = run_benchmark(num_iters=100, profile=False)
    print(f"Kernel running time: {latency * 1000000:.3f} us")


219
220
221
222
223
if __name__ == "__main__":
    logger.warning(
        "This script benchmarks the paged attention kernel. "
        "By default this is no longer used in vLLM inference."
    )
224

225
    parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.")
zhuwenwen's avatar
zhuwenwen committed
226
    parser.add_argument("--version", type=str, choices=["v1", "v2", "v12"], default="v12")
227
    parser.add_argument("--batch-size", type=int, default=8)
Allen.Dou's avatar
Allen.Dou committed
228
    parser.add_argument("--seq-len", type=int, default=4096)
229
230
    parser.add_argument("--num-query-heads", type=int, default=64)
    parser.add_argument("--num-kv-heads", type=int, default=8)
231
232
233
234
235
236
    parser.add_argument(
        "--head-size",
        type=int,
        choices=[64, 80, 96, 112, 120, 128, 192, 256],
        default=128,
    )
zhuwenwen's avatar
zhuwenwen committed
237
    parser.add_argument("--block-size", type=int, choices=[16, 32, 64], default=64)
238
    parser.add_argument("--use-alibi", action="store_true")
239
240
241
    parser.add_argument(
        "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"
    )
242
243
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--profile", action="store_true")
244
245
246
    parser.add_argument(
        "--kv-cache-dtype",
        type=str,
247
        choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"],
248
        default="auto",
249
250
        help="Data type for kv cache storage. If 'auto', will use model "
        "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "
zhuwenwen's avatar
zhuwenwen committed
251
        "ROCm (hcu) supports fp8 (=fp8_e4m3)")
252
253
254
    parser.add_argument(
        "--custom-paged-attn", action="store_true", help="Use custom paged attention"
    )
255
256
257
258
259
260
261
262
    args = parser.parse_args()
    print(args)

    if args.num_query_heads % args.num_kv_heads != 0:
        raise ValueError("num_query_heads must be divisible by num_kv_heads")
    main(
        version=args.version,
        num_seqs=args.batch_size,
263
        seq_len=args.seq_len,
264
265
266
267
268
        num_query_heads=args.num_query_heads,
        num_kv_heads=args.num_kv_heads,
        head_size=args.head_size,
        block_size=args.block_size,
        use_alibi=args.use_alibi,
269
        dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype],
270
271
        seed=args.seed,
        do_profile=args.profile,
272
        kv_cache_dtype=args.kv_cache_dtype,
273
    )