offloading_benchmarks.py

import time
import torch
import random
# import pytest
from typing import List
import lmc_ops
import argparse
import json
from vllm.attention.ops.paged_attn import PagedAttention
from typing import Optional

use_mla = False


def check_paged_kv_cache_equal(
    left, right, num_tokens, slot_mapping, num_heads=8, head_size=128
):
    """
    check whether two paged kv caches are the same at slot_mapping
    """
    token_dim = 0
    for left_kv, right_kv in zip(left, right, strict=False):
        left_k = left_kv[0].reshape(-1, num_heads, head_size)
        left_v = left_kv[1].reshape(-1, num_heads, head_size)
        right_k = right_kv[0].reshape(-1, num_heads, head_size)
        right_v = right_kv[1].reshape(-1, num_heads, head_size)

        assert len(left_k.shape) == 3
        assert len(left_v.shape) == 3
        assert len(right_k.shape) == 3
        assert len(right_v.shape) == 3

        assert left_k.shape[token_dim] >= num_tokens
        assert left_v.shape[token_dim] >= num_tokens
        assert right_k.shape[token_dim] >= num_tokens
        assert right_v.shape[token_dim] >= num_tokens

        assert (left_k[slot_mapping, :, :] == right_k[slot_mapping, :, :]).all()
        assert (left_v[slot_mapping, :, :] == right_v[slot_mapping, :, :]).all()

def generate_kv_cache_paged_list_tensors(num_blocks,
                                         block_size,
                                         num_layers,
                                         num_heads,
                                         head_size,
                                         device,
                                         dtype=torch.bfloat16):
    """
    Instead of Tuple[Tuple[Tensor, Tensor]], return List[Tensor]
    where KV are in the same tensor
    """
    ret = []
    shape = [2, num_blocks, block_size, num_heads, head_size]

    for i in range(num_layers):
        kv = torch.rand(shape, dtype=dtype, device=device)
        ret.append(kv)

    return ret

def initialize_pointers(kv_caches: List[torch.Tensor], num_layers):
    kv_cache_pointers = torch.empty(num_layers,
                                        dtype=torch.int64,
                                        device='cpu',
                                        pin_memory=True)
    for i in range(num_layers):
        kv_cache_pointers[i] = kv_caches[i].data_ptr()
    return kv_cache_pointers

def get_page_buffer_size(kv_caches):        # from gpu_connector, line 353
    page_buffer_size = 0
    if use_mla:                                                          
        # kv_caches[0].shape: [num_pages, page_size, head_size]
        assert kv_caches[0].dim() == 3
        page_buffer_size = kv_caches[0].shape[0] * kv_caches[0].shape[1]
    else:
        # kv_caches[0].shape: [2, num_pages, page_size, num_heads, head_size]
        assert kv_caches[0].dim() == 5
        page_buffer_size = kv_caches[0].shape[1] * kv_caches[0].shape[2]
    return page_buffer_size

def benchmark_offload_lmcache(block_size, chunk_size,num_blocks, num_layers, num_heads, head_size, device, num_tokens):
    hidden_dim = num_heads * head_size
    offloads_timings = []

    for i in range(31):
        print(f"begin iteration {i+1}:")
        gpu_kv_src = generate_kv_cache_paged_list_tensors(num_blocks=num_blocks, block_size=block_size, num_layers=num_layers, num_heads=num_heads, head_size=head_size, device=device)
        gpu_kv_dst = generate_kv_cache_paged_list_tensors(num_blocks=num_blocks, block_size=block_size, num_layers=num_layers, num_heads=num_heads, head_size=head_size, device=device)

        memlist = []
        for i,start in enumerate(range(0, num_tokens, chunk_size)):
            end = min(start + chunk_size, num_tokens)
            memory_obj = torch.empty(2, num_layers, end - start, hidden_dim, dtype=gpu_kv_dst[0][0].dtype, pin_memory=True)
            memlist.append(memory_obj)

        slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
        slot_mapping = torch.tensor(slot_mapping, device="cuda", dtype=torch.int64)

        pointers_src = initialize_pointers(gpu_kv_src, num_layers)
        pointers_dst = initialize_pointers(gpu_kv_dst, num_layers)
        page_buffer_size = get_page_buffer_size(gpu_kv_src)

        for i,start in enumerate(range(0, num_tokens, chunk_size)):
            end = min(start + chunk_size, num_tokens)
            lmc_ops.multi_layer_kv_transfer_Memcpy(
                    memlist[i],
                    pointers_src,
                    slot_mapping[start:end],
                    gpu_kv_src[0].device,
                    page_buffer_size,                   # page buffer size
                    True,                              # direction of transfer, True: GPU->CPU, False: CPU->GPU
                    use_mla,                            # use_mlu, Multi-Head Latent Attention
                )
            
        torch.cuda.synchronize()
        start_time = time.time()

        for i,start in enumerate(range(0, num_tokens, chunk_size)):
            end = min(start + chunk_size, num_tokens)
            lmc_ops.multi_layer_kv_transfer_Memcpy(
                    memlist[i],
                    pointers_dst,
                    slot_mapping[start:end],
                    gpu_kv_dst[0].device,
                    page_buffer_size,                   # page buffer size
                    False,                              # direction of transfer, True: GPU->CPU, False: CPU->GPU
                    use_mla,                            # use_mlu, Multi-Head Latent Attention
                )
            
        torch.cuda.synchronize()
        end_time = time.time()

        check_paged_kv_cache_equal(gpu_kv_src, gpu_kv_dst, num_tokens, slot_mapping, num_heads, head_size)

        if i>0:
            offloads_timings.append(end_time - start_time)
            print(f"Offload Time CPU->GPU: {end_time - start_time:.6f} sec")         

    return offloads_timings

def benchmark_offload_vllm(block_size, num_blocks, num_layers, num_heads, head_size, num_tokens):
    offloads_timings = []
    hidden_dim = num_heads * head_size
    num_mapping = num_tokens // block_size

    for i in range(11):
        print(f"begin iteration {i+1}:")
        kv_src = generate_kv_cache_paged_list_tensors(num_blocks=num_blocks, block_size=block_size, num_layers=num_layers, num_heads=num_heads, head_size=head_size, device="cuda")
        kv_dst = generate_kv_cache_paged_list_tensors(num_blocks, block_size, num_layers, num_heads, head_size, device="cuda")
        
        src_blocks = random.sample(range(num_blocks), num_mapping)   
        dst_blocks = random.sample(range(num_blocks), num_mapping)
        block_mapping = list(zip(src_blocks, dst_blocks))
        block_mapping_tensor = torch.tensor(block_mapping,
                                            dtype=torch.int64,
                                            device="cpu").view(-1, 2)
        # Reverse the mapping: (src, dst) -> (dst, src)
        reverse_block_mapping_tensor = block_mapping_tensor[:, [1, 0]]

        memlist = []
        for _ in range(num_layers):
            memory_obj = torch.empty(size=kv_src[0].shape, dtype=kv_src[0][0].dtype, pin_memory=True)
            memlist.append(memory_obj)


        for layer in range(num_layers):
            PagedAttention.swap_blocks(kv_src[layer], memlist[layer], block_mapping_tensor) 

        torch.cuda.synchronize()
        start_time = time.time()

        for layer in range(num_layers):
            PagedAttention.swap_blocks(memlist[layer], kv_dst[layer], reverse_block_mapping_tensor)

        torch.cuda.synchronize()
        end_time = time.time()

        # check correctness
        for layer in range(num_layers):
            for j in src_blocks:
                torch.testing.assert_close(kv_src[layer][0][j].cpu(),
                                        kv_dst[layer][0][j].cpu())
                torch.testing.assert_close(kv_src[layer][1][j].cpu(),
                                        kv_dst[layer][1][j].cpu())

        if i > 0:
            print(f"Offload Time: {end_time - start_time:.6f} sec")
            offloads_timings.append(end_time - start_time)

    return offloads_timings

def save_results(offload_timings, parameters, block_size, chunk_size, num_layers, num_tokens):
    filename = f"bs_{block_size}_cs_{chunk_size}_l_{num_layers}"
    path = f"offload_results/version_cudaMemcpy/{filename}.json"
    with open(path, "w") as f:
        json.dump({"parameters": parameters, "data": offload_timings}, f, indent=2)
    print(f"\nSaved results to {path}\n")


def benchmark_offload():
    parser = argparse.ArgumentParser(description="Program with QPS argument")
    parser.add_argument("--case", type=str, default="vllm",
                            help="Case to test")
    args = parser.parse_args()
        
    # Access the qps value
    case=args.case

    num_blocks = 1280
    num_heads = 8
    head_size = 128
    device = "cuda"
    num_tokens = 10240
    
    for block_size in [8,16,32,64]: 
        for chunk_size in [128,256,512]:
            for num_layers in [8,16,24,32,40]:
                parameters = {
                        "num_blocks": num_blocks,
                        "num_layers": num_layers,
                        "num_heads": num_heads,
                        "head_size": head_size,
                        "num_tokens": num_tokens
                    }
                
                offload_timings=[]

                if case == "lmcache":
                    offload_timings = benchmark_offload_lmcache(block_size, chunk_size, num_blocks, num_layers=num_layers, num_heads=num_heads, head_size=head_size, device=device, num_tokens=num_tokens)
                elif case == "vllm":
                    offload_timings = benchmark_offload_vllm(block_size, num_blocks, num_layers=num_layers, num_heads=num_heads, head_size=head_size, num_tokens=num_tokens)

                save_results(offload_timings, parameters, block_size, chunk_size, num_layers, num_tokens)

if __name__ == "__main__":
    benchmark_offload()