跳到内容

vllm_gaudi.extension.cache_ops

copy_blocks 

copy_blocks(key_caches, value_caches, block_mapping)
Source code in vllm_gaudi/extension/cache_ops.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def copy_blocks(key_caches, value_caches, block_mapping):
    if block_mapping.numel() == 0:
        return

    block_mapping = block_mapping.transpose(0, 1)
    src = block_mapping[0]
    dst = block_mapping[1]

    for key_cache, value_cache in zip(key_caches, value_caches):
        key_cache.index_copy_(0, dst, key_cache.index_select(0, src))
        value_cache.index_copy_(0, dst, value_cache.index_select(0, src))

    if key_caches[0].device.type == 'hpu':
        htorch.core.mark_step()

swap_blocks 

swap_blocks(src, dst, block_mapping)
Source code in vllm_gaudi/extension/cache_ops.py 
12
13
14
15
16
17
18
19
20
21
22
23
def swap_blocks(src, dst, block_mapping):
    if block_mapping.numel() == 0:
        return

    block_mapping = block_mapping.transpose(0, 1)
    src_indices = block_mapping[0]
    dst_indices = block_mapping[1]

    dst.index_put_(dst_indices, src.index_select(0, src_indices))

    htorch.core.mark_step()
    torch.hpu.synchronize()