test_cache.py 16.2 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
import random
2
from typing import List, Tuple
Woosuk Kwon's avatar
Woosuk Kwon committed
3

4
import pytest
Woosuk Kwon's avatar
Woosuk Kwon committed
5
6
import torch

7
from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
8
from vllm import _custom_ops as ops
9
from vllm.platforms import current_platform
Woosuk Kwon's avatar
Woosuk Kwon committed
10

Vladimir's avatar
Vladimir committed
11
COPYING_DIRECTION = [('cuda', 'cpu'), ('cuda', 'cuda'), ('cpu', 'cuda')]
12
DTYPES = [torch.half, torch.bfloat16, torch.float]
Simon Mo's avatar
Simon Mo committed
13
NUM_TOKENS = [42]  # Arbitrary values for testing
14
NUM_LAYERS = [1]  # Arbitrary values for testing
15
NUM_HEADS = [8]  # Arbitrary values for testing
16
HEAD_SIZES = [64, 80, 120, 256]
17
BLOCK_SIZES = [8, 16, 32]
18
19
20
21
22

# Arbitrary values for testing
# don't make it too large. e.g. [1024, 36000] will OOM
NUM_BLOCKS = [1024, 10000]

23
NUM_MAPPINGS = [256]  # Arbitrary values for testing
24
SEEDS = [0]
25
26
27
CUDA_DEVICES = [
    f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
]
28
29

# We assume fp8 is always enabled for testing.
30
KV_CACHE_DTYPE = ["auto", "fp8"]
31
32
33
34
35
36
37
38
39
40


@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("num_layers", NUM_LAYERS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
41
@pytest.mark.parametrize("device", CUDA_DEVICES)
42
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
43
@torch.inference_mode()
44
45
def test_copy_blocks(
    kv_cache_factory,
46
47
48
49
50
51
52
    num_mappings: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
53
    seed: int,
54
    kv_cache_dtype: str,
55
    device: str,
56
) -> None:
Joe's avatar
Joe committed
57
58
    if kv_cache_dtype == "fp8" and head_size % 16:
        pytest.skip()
59
    current_platform.seed_everything(seed)
60
    torch.set_default_device(device)
61
62
63
    # Generate random block mappings where each source block is mapped to two
    # destination blocks.
    assert 2 * num_mappings <= num_blocks
64
65
    src_blocks = random.sample(range(num_blocks), num_mappings)
    remainig_blocks = list(set(range(num_blocks)) - set(src_blocks))
66
    dst_blocks = random.sample(remainig_blocks, 2 * num_mappings)
67
    block_mapping: List[Tuple[int, int]] = []
68
    for i in range(num_mappings):
69
70
71
        src = src_blocks[i]
        dst1 = dst_blocks[2 * i]
        dst2 = dst_blocks[2 * i + 1]
72
73
        block_mapping.append((src, dst1))
        block_mapping.append((src, dst2))
74
75
76
77

    # Create the KV caches.
    key_caches, value_caches = kv_cache_factory(num_blocks, block_size,
                                                num_layers, num_heads,
78
                                                head_size, kv_cache_dtype,
79
                                                dtype, seed, device)
80
81
82
83

    # Clone the KV caches.
    cloned_key_caches = [key_cache.clone() for key_cache in key_caches]
    cloned_value_caches = [value_cache.clone() for value_cache in value_caches]
84
85

    # Call the copy blocks kernel.
86
87
88
    block_mapping_tensor = torch.tensor(block_mapping,
                                        dtype=torch.int64,
                                        device=device).view(-1, 2)
89
90
91
92
93

    opcheck(torch.ops._C_cache_ops.copy_blocks,
            (key_caches, value_caches, block_mapping_tensor),
            test_utils=DEFAULT_OPCHECK_TEST_UTILS,
            cond=(head_size == HEAD_SIZES[0]))
94
    ops.copy_blocks(key_caches, value_caches, block_mapping_tensor)
95

96
    # Run the reference implementation.
97
98
99
100
101
    for src, dst in block_mapping:
        for cloned_key_cache in cloned_key_caches:
            cloned_key_cache[dst].copy_(cloned_key_cache[src])
        for cloned_value_cache in cloned_value_caches:
            cloned_value_cache[dst].copy_(cloned_value_cache[src])
102
103
104

    # Compare the results.
    for key_cache, cloned_key_cache in zip(key_caches, cloned_key_caches):
105
        torch.testing.assert_close(key_cache, cloned_key_cache)
106
107
    for value_cache, cloned_value_cache in zip(value_caches,
                                               cloned_value_caches):
108
        torch.testing.assert_close(value_cache, cloned_value_cache)
109
110


111
112
113
114
115
116
117
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
118
@pytest.mark.parametrize("device", CUDA_DEVICES)
119
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
120
@torch.inference_mode()
121
122
def test_reshape_and_cache(
    kv_cache_factory,
Woosuk Kwon's avatar
Woosuk Kwon committed
123
124
125
126
127
128
    num_tokens: int,
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
129
    seed: int,
130
    device: str,
131
    kv_cache_dtype: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
132
) -> None:
Joe's avatar
Joe committed
133
134
    if kv_cache_dtype == "fp8" and head_size % 16:
        pytest.skip()
135
    current_platform.seed_everything(seed)
136
    torch.set_default_device(device)
137
    # Create a random slot mapping.
Woosuk Kwon's avatar
Woosuk Kwon committed
138
    num_slots = block_size * num_blocks
139
140
    slot_mapping_lst = random.sample(range(num_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long)
141
142

    qkv = torch.randn(num_tokens, 3, num_heads, head_size, dtype=dtype)
Woosuk Kwon's avatar
Woosuk Kwon committed
143
144
    _, key, value = qkv.unbind(dim=1)

145
146
    # Create the KV caches.
    key_caches, value_caches = kv_cache_factory(num_blocks, block_size, 1,
147
148
149
                                                num_heads, head_size,
                                                kv_cache_dtype, dtype, seed,
                                                device)
150
    key_cache, value_cache = key_caches[0], value_caches[0]
Woosuk Kwon's avatar
Woosuk Kwon committed
151

152
    # Clone the KV caches.
153
154
    if kv_cache_dtype == "fp8":
        cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
155
        ops.convert_fp8(cloned_key_cache, key_cache)
156
        cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
157
        ops.convert_fp8(cloned_value_cache, value_cache)
158
159
160
161
162
    else:
        cloned_key_cache = key_cache.clone()
        cloned_value_cache = value_cache.clone()

    # Using default kv_scale
163
    k_scale = v_scale = 1.0
Woosuk Kwon's avatar
Woosuk Kwon committed
164

165
    # Call the reshape_and_cache kernel.
166
167
168
169
    opcheck(torch.ops._C_cache_ops.reshape_and_cache,
            (key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype,
             k_scale, v_scale),
            cond=(head_size == HEAD_SIZES[0]))
170
    ops.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping,
171
                          kv_cache_dtype, k_scale, v_scale)
172
173
174

    if kv_cache_dtype == "fp8":
        result_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
175
        ops.convert_fp8(result_key_cache, key_cache)
176
        result_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
177
        ops.convert_fp8(result_value_cache, value_cache)
Woosuk Kwon's avatar
Woosuk Kwon committed
178

179
180
    # Run the reference implementation.
    reshaped_key = key.reshape(num_tokens, *key_cache[0, :, :, 0, :].shape)
181
    block_indicies = torch.div(slot_mapping, block_size, rounding_mode="floor")
182
    block_indicies_lst = block_indicies.cpu().tolist()
183
    block_offsets = slot_mapping % block_size
184
    block_offsets_lst = block_offsets.cpu().tolist()
Woosuk Kwon's avatar
Woosuk Kwon committed
185
    for i in range(num_tokens):
186
187
        block_idx = block_indicies_lst[i]
        block_offset = block_offsets_lst[i]
Woosuk Kwon's avatar
Woosuk Kwon committed
188
        cloned_key_cache[block_idx, :, :, block_offset, :] = reshaped_key[i]
189
        cloned_value_cache[block_idx, :, :, block_offset] = value[i]
Woosuk Kwon's avatar
Woosuk Kwon committed
190

191
    if kv_cache_dtype == "fp8":
192
193
194
195
196
197
198
199
        torch.testing.assert_close(result_key_cache,
                                   cloned_key_cache,
                                   atol=0.001,
                                   rtol=0.1)
        torch.testing.assert_close(result_value_cache,
                                   cloned_value_cache,
                                   atol=0.001,
                                   rtol=0.1)
200
    else:
201
202
        torch.testing.assert_close(key_cache, cloned_key_cache)
        torch.testing.assert_close(value_cache, cloned_value_cache)
Vladimir's avatar
Vladimir committed
203
204


205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@torch.inference_mode()
def test_reshape_and_cache_flash(
    kv_cache_factory_flashinfer,
    num_tokens: int,
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
    kv_cache_dtype: str,
) -> None:
227
    current_platform.seed_everything(seed)
228
    torch.set_default_device(device)
229
230
231

    # Create a random slot mapping.
    num_slots = block_size * num_blocks
232
233
234
235
    slot_mapping_lst = random.sample(range(num_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst,
                                dtype=torch.long,
                                device=device)
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253

    qkv = torch.randn(num_tokens,
                      3,
                      num_heads,
                      head_size,
                      dtype=dtype,
                      device=device)
    _, key, value = qkv.unbind(dim=1)

    # Create the KV caches.
    key_caches, value_caches = kv_cache_factory_flashinfer(
        num_blocks,
        block_size,
        1,
        num_heads,
        head_size,
        kv_cache_dtype,
        dtype,
254
        device=device,
255
    )
256
257
258
259
    key_cache, value_cache = key_caches[0].contiguous(
    ), value_caches[0].contiguous()
    del key_caches
    del value_caches
260

261
262
263
    k_scale = key.amax().item() / 256
    v_scale = value.amax().item() / 256

264
    # Clone the KV caches.
265
266
    if kv_cache_dtype == "fp8":
        cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
267
        ops.convert_fp8(cloned_key_cache, key_cache, k_scale, kv_cache_dtype)
268
        cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
269
270
        ops.convert_fp8(cloned_value_cache, value_cache, v_scale,
                        kv_cache_dtype)
271
272
273
274
    else:
        cloned_key_cache = key_cache.clone()
        cloned_value_cache = value_cache.clone()

275
    # Call the reshape_and_cache kernel.
276
277
278
279
    opcheck(torch.ops._C_cache_ops.reshape_and_cache_flash,
            (key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype,
             k_scale, v_scale),
            cond=(head_size == HEAD_SIZES[0]))
280
    ops.reshape_and_cache_flash(key, value, key_cache, value_cache,
281
282
283
284
                                slot_mapping, kv_cache_dtype, k_scale, v_scale)

    if kv_cache_dtype == "fp8":
        result_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
285
286
287
288
        ops.convert_fp8(result_key_cache,
                        key_cache,
                        k_scale,
                        kv_dtype=kv_cache_dtype)
289
        result_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
290
291
292
293
        ops.convert_fp8(result_value_cache,
                        value_cache,
                        v_scale,
                        kv_dtype=kv_cache_dtype)
294
295

    # Run the reference implementation.
296
297
    block_indicies = torch.div(slot_mapping, block_size, rounding_mode="floor")
    block_indicies_lst = block_indicies.cpu().tolist()
298
    block_offsets = slot_mapping % block_size
299
    block_offsets_lst = block_offsets.cpu().tolist()
300
    for i in range(num_tokens):
301
302
        block_idx = block_indicies_lst[i]
        block_offset = block_offsets_lst[i]
303
304
305
        cloned_key_cache[block_idx, block_offset, :, :] = key[i]
        cloned_value_cache[block_idx, block_offset, :, :] = value[i]

306
    if kv_cache_dtype == "fp8":
307
308
309
310
311
312
313
314
        torch.testing.assert_close(result_key_cache,
                                   cloned_key_cache,
                                   atol=0.001,
                                   rtol=0.1)
        torch.testing.assert_close(result_value_cache,
                                   cloned_value_cache,
                                   atol=0.001,
                                   rtol=0.1)
315
    else:
316
317
        torch.testing.assert_close(key_cache, cloned_key_cache)
        torch.testing.assert_close(value_cache, cloned_value_cache)
318
319


Vladimir's avatar
Vladimir committed
320
321
322
323
324
325
326
327
@pytest.mark.parametrize("direction", COPYING_DIRECTION)
@pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
328
@pytest.mark.parametrize("device", CUDA_DEVICES)
329
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
Vladimir's avatar
Vladimir committed
330
331
332
333
334
335
336
337
338
339
340
@torch.inference_mode()
def test_swap_blocks(
    kv_cache_factory,
    direction: Tuple[str, str],
    num_mappings: int,
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
341
    device: str,
342
    kv_cache_dtype: str,
Vladimir's avatar
Vladimir committed
343
) -> None:
344
345
    if kv_cache_dtype == "fp8" and "cpu" in direction:
        pytest.skip()
Joe's avatar
Joe committed
346
347
    if kv_cache_dtype == "fp8" and head_size % 16:
        pytest.skip()
348

349
    current_platform.seed_everything(seed)
350
351
352

    src_device = device if direction[0] == "cuda" else 'cpu'
    dst_device = device if direction[1] == "cuda" else 'cpu'
Vladimir's avatar
Vladimir committed
353
354
355
356
357
358
359
360
361

    src_blocks = random.sample(range(num_blocks), num_mappings)
    # For the same device, mapping must not overlap
    if src_device == dst_device:
        remaining_blocks = list(set(range(num_blocks)) - set(src_blocks))
        dst_blocks = random.sample(remaining_blocks, num_mappings)
    else:
        dst_blocks = random.sample(range(num_blocks), num_mappings)

362
363
364
365
    block_mapping = list(zip(src_blocks, dst_blocks))
    block_mapping_tensor = torch.tensor(block_mapping,
                                        dtype=torch.int64,
                                        device="cpu").view(-1, 2)
Vladimir's avatar
Vladimir committed
366
367
368

    # Create the KV caches on the first device.
    src_key_caches, src_value_caches = kv_cache_factory(
369
370
        num_blocks, block_size, 1, num_heads, head_size, kv_cache_dtype, dtype,
        seed, src_device)
Vladimir's avatar
Vladimir committed
371
372
373

    # Create the KV caches on the second device.
    dist_key_caches, dist_value_caches = kv_cache_factory(
374
375
        num_blocks, block_size, 1, num_heads, head_size, kv_cache_dtype, dtype,
        seed, dst_device)
Vladimir's avatar
Vladimir committed
376
377
378
379
380

    src_key_caches_clone = src_key_caches[0].clone()
    src_value_caches_clone = src_value_caches[0].clone()

    # Call the swap_blocks kernel.
381
382
383
384
385
386
387
388
    do_opcheck = (head_size == HEAD_SIZES[0])
    opcheck(torch.ops._C_cache_ops.swap_blocks,
            (src_key_caches[0], dist_key_caches[0], block_mapping_tensor),
            cond=do_opcheck)
    opcheck(torch.ops._C_cache_ops.swap_blocks,
            (src_value_caches[0], dist_value_caches[0], block_mapping_tensor),
            cond=do_opcheck)

389
390
391
392
    ops.swap_blocks(src_key_caches[0], dist_key_caches[0],
                    block_mapping_tensor)
    ops.swap_blocks(src_value_caches[0], dist_value_caches[0],
                    block_mapping_tensor)
Vladimir's avatar
Vladimir committed
393

394
    for src, dst in block_mapping:
395
396
397
398
        torch.testing.assert_close(src_key_caches_clone[src].cpu(),
                                   dist_key_caches[0][dst].cpu())
        torch.testing.assert_close(src_value_caches_clone[src].cpu(),
                                   dist_value_caches[0][dst].cpu())
399
400
401
402
403
404
405
406
407
408


@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("head_size", HEAD_SIZES)
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
409
def test_fp8_e4m3_conversion(
410
411
412
413
414
415
416
417
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
) -> None:
418
    current_platform.seed_everything(seed)
419
420
421
422
423
424
425
426

    low = -224.0
    high = 224.0
    shape = (num_blocks, num_heads, head_size, block_size)
    cache = torch.empty(shape, dtype=dtype, device=device)
    cache.uniform_(low, high)

    cache_fp8 = torch.empty_like(cache, dtype=torch.uint8)
427
    ops.convert_fp8(cache_fp8, cache)
428
429

    converted_cache = torch.empty_like(cache)
430
    ops.convert_fp8(converted_cache, cache_fp8)
431

432
    torch.testing.assert_close(cache, converted_cache, atol=0.001, rtol=0.1)