# SPDX-License-Identifier: Apache-2.0
# Standard
from typing import List
import random

# Third Party
import pytest
import torch

# First Party
from lmcache.v1.memory_allocators.pin_memory_allocator import PinMemoryAllocator

if not torch.cuda.is_available():
    pytest.skip(
        "CUDA is not available, skipping the test",
        allow_module_level=True,
    )

# First Party
if torch.cuda.is_available():
    try:
        # First Party
        import lmcache.c_ops as lmc_ops
    except ImportError:
        lmc_ops = None
else:
    lmc_ops = None

# Mock c_ops when not available
if lmc_ops is None:

    class MockEngineKVFormat:
        NL_X_TWO_NB_BS_NH_HS = 0
        NL_X_NB_TWO_BS_NH_HS = 1
        NL_X_NB_BS_HS = 2
        NL_X_TWO_NB_NH_BS_HS = 3
        NL_X_NB_TWO_NH_BS_HS = 4

    class MockTransferDirection:
        H2D = 0
        D2H = 1

    class MockCOps:
        EngineKVFormat = MockEngineKVFormat
        GPUKVFormat = MockEngineKVFormat
        TransferDirection = MockTransferDirection

    lmc_ops = MockCOps()

# Local
from .utils import (
    check_mem_obj_equal,
    check_paged_kv_cache_equal,
    generate_kv_cache_paged,
    generate_kv_cache_paged_list_tensors,
)


def _tuple_kv_to_blob(
    kv_tensors,
) -> torch.Tensor:
    k_temp = []
    v_temp = []
    for kv_layer in kv_tensors:
        k_temp.append(kv_layer[0])
        v_temp.append(kv_layer[1])
    k_tensor_blob = torch.stack(k_temp)
    v_tensor_blob = torch.stack(v_temp)

    # kv_tensors: [num_layer, 2, num_tok, num_kv_head, head_size]
    kv_tensors_flatten = torch.stack((k_tensor_blob, v_tensor_blob))
    kv_tensors_flatten = kv_tensors_flatten.permute([1, 0, 2, 3, 4])

    return kv_tensors_flatten


def _slice_kv_at(
    start_idx: int,
    kv_tensors: torch.Tensor,
    chunk_size: int,
) -> List[torch.Tensor]:
    return [
        x.contiguous()
        for x in list(
            torch.split(
                kv_tensors[:, :, start_idx:, ...],
                chunk_size,
                dim=2,
            )
        )
    ]


@pytest.mark.parametrize("num_tokens", [256, 500, 1024, 8000])
def test_extract_and_load_back(num_tokens):
    device = "cuda"

    num_blocks = 1000
    block_size = 16
    num_heads = 8
    head_size = 128
    num_layers = 32
    dtype = torch.bfloat16
    kv_cache = generate_kv_cache_paged(num_blocks, device, block_size, dtype)

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

    pinned_cpu_size = 4 * 1024 * 1024 * 1024  # 4GB
    mem_allocator = PinMemoryAllocator(pinned_cpu_size)

    # Old extract
    kv_tuple_list = []
    memory_obj_old_list = []
    chunk_size = 256
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    for layer_id in range(num_layers):
        key_cache = kv_cache[layer_id][0].reshape(-1, num_heads, head_size)
        value_cache = kv_cache[layer_id][1].reshape(-1, num_heads, head_size)
        kv_tuple_list.append((key_cache[slot_mapping], value_cache[slot_mapping]))
    kv_blob = _tuple_kv_to_blob(kv_tuple_list)
    kv_chunked = _slice_kv_at(0, kv_blob, chunk_size)
    for chunk_id, chunk in enumerate(kv_chunked):
        mem_obj_shape = torch.Size(
            [2, num_layers, chunk.shape[2], num_heads * head_size]
        )

        memory_obj_old = mem_allocator.allocate(mem_obj_shape, dtype)
        chunk = chunk.contiguous()
        for layer_id in range(num_layers):
            memory_obj_old.tensor[0, layer_id].copy_(
                chunk[layer_id, 0].reshape(-1, 1024)
            )
            memory_obj_old.tensor[1, layer_id].copy_(
                chunk[layer_id, 1].reshape(-1, 1024)
            )
        memory_obj_old_list.append(memory_obj_old)
    end_event.record()
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("Old extract time: ", elapsed_time_ms / 1000)

    # New extract (zero-copy kernels)
    memory_obj_new_list = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        mem_obj_shape = torch.Size(
            [2, num_layers, len(slot_mapping_temp), num_heads * head_size]
        )

        memory_obj_new = mem_allocator.allocate(mem_obj_shape, dtype)
        for layer_id in range(num_layers):
            lmc_ops.load_and_reshape_flash(
                memory_obj_new.tensor,
                kv_cache[layer_id][0],
                kv_cache[layer_id][1],
                slot_mapping_temp,
                layer_id,
            )
        memory_obj_new_list.append(memory_obj_new)
    end_event.record()
    # wait for all the operations to finish
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("New extract time: ", elapsed_time_ms / 1000)
    check_mem_obj_equal(
        memory_obj_old_list,
        memory_obj_new_list,
    )

    # Generate new paged kv_cache
    kv_cache_new = generate_kv_cache_paged(num_blocks, device, block_size, dtype)

    # New load back (zero-copy kernels)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        memory_obj_new = memory_obj_new_list[chunk_id]
        for layer_id in range(num_layers):
            lmc_ops.reshape_and_cache_back_flash(
                memory_obj_new.tensor,
                kv_cache_new[layer_id][0],
                kv_cache_new[layer_id][1],
                slot_mapping_temp,
                layer_id,
            )
    check_paged_kv_cache_equal(
        kv_cache,
        kv_cache_new,
        slot_mapping,
    )

    mem_allocator.close()


@pytest.mark.parametrize("num_tokens", [256, 500, 1024, 8000])
@pytest.mark.parametrize(
    "engine_kv_format",
    [
        lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS,  # vllm non-MLA flash attention
        lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS,  # vllm non-MLA flash infer
    ],
)
def test_multi_layer_kernel(num_tokens, engine_kv_format):
    device = "cuda"

    num_blocks = 1000
    block_size = 16
    num_heads = 8
    head_size = 128
    chunk_size = 256
    num_layers = 32
    dtype = torch.bfloat16
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype, engine_kv_format=engine_kv_format
    )
    # old deprecated kernels only handle vllm non-MLA flash attention format
    if engine_kv_format == lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS:
        kv_cache_force_old = [
            kv_layer.clone().permute(1, 0, 2, 3, 4).contiguous()
            for kv_layer in kv_cache
        ]
    else:
        kv_cache_force_old = kv_cache
    page_buffer_size = num_blocks * block_size

    slot_mapping = random.sample(range(0, page_buffer_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    pinned_cpu_size = 4 * 1024 * 1024 * 1024  # 4GB
    mem_allocator = PinMemoryAllocator(pinned_cpu_size)

    # layer by layer extract
    memory_obj_old_list = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        mem_obj_shape = torch.Size(
            [2, num_layers, len(slot_mapping_temp), num_heads * head_size]
        )

        memory_obj_old = mem_allocator.allocate(mem_obj_shape, dtype)
        for layer_id in range(num_layers):
            lmc_ops.load_and_reshape_flash(
                memory_obj_old.tensor,
                kv_cache_force_old[layer_id][0],
                kv_cache_force_old[layer_id][1],
                slot_mapping_temp,
                layer_id,
            )
        memory_obj_old_list.append(memory_obj_old)
    end_event.record()
    # wait for all the operations to finish
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("Old extract time: ", elapsed_time_ms / 1000)

    # New extract with multi layer kernel
    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_cache[i].data_ptr()

    memory_obj_new_list = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        mem_obj_shape = torch.Size(
            [2, num_layers, len(slot_mapping_temp), num_heads * head_size]
        )

        memory_obj_new = mem_allocator.allocate(mem_obj_shape, dtype)
        lmc_ops.multi_layer_kv_transfer(
            memory_obj_new.tensor,
            kv_cache_pointers,
            slot_mapping_temp,
            kv_cache[0].device,
            page_buffer_size,
            lmc_ops.TransferDirection.D2H,
            engine_kv_format,
            block_size,
        )
        memory_obj_new_list.append(memory_obj_new)

    end_event.record()
    # wait for all the operations to finish
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("New extract time: ", elapsed_time_ms / 1000)

    check_mem_obj_equal(
        memory_obj_old_list,
        memory_obj_new_list,
    )

    # Generate new paged kv_cache
    kv_cache_new = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype, engine_kv_format=engine_kv_format
    )

    kv_cache_pointers_new = torch.empty(
        num_layers, dtype=torch.int64, device="cpu", pin_memory=True
    )
    for i in range(num_layers):
        kv_cache_pointers_new[i] = kv_cache_new[i].data_ptr()

    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        memory_obj_new = memory_obj_new_list[chunk_id]
        lmc_ops.multi_layer_kv_transfer(
            memory_obj_new.tensor,
            kv_cache_pointers_new,
            slot_mapping_temp,
            kv_cache_new[0].device,
            page_buffer_size,
            lmc_ops.TransferDirection.H2D,
            engine_kv_format,
            block_size,
        )

    check_paged_kv_cache_equal(
        kv_cache,
        kv_cache_new,
        slot_mapping,
        engine_kv_format=engine_kv_format,
    )

    mem_allocator.close()


@pytest.mark.parametrize("num_tokens", [256, 500, 1024, 8000])
@pytest.mark.parametrize("head_size", [576, 66])  # Use 68 for dsv32 (132x int8)
@pytest.mark.parametrize(
    "engine_kv_format",
    [lmc_ops.EngineKVFormat.NL_X_NB_BS_HS],  # vllm MLA
)
def test_multi_layer_kernel_use_mla(num_tokens, head_size, engine_kv_format):
    device = "cuda"

    num_blocks = 1000
    block_size = 64
    chunk_size = 256
    dtype = torch.bfloat16
    num_layers = 32
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers,
        head_size,
        engine_kv_format=engine_kv_format,
    )

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

    pinned_cpu_size = 4 * 1024 * 1024 * 1024  # 4GB
    mem_allocator = PinMemoryAllocator(pinned_cpu_size)

    # layer by layer extract
    memory_obj_old_list = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        mem_obj_shape = torch.Size([1, num_layers, len(slot_mapping_temp), head_size])
        memory_obj_old = mem_allocator.allocate(mem_obj_shape, dtype)

        for layer_id in range(num_layers):
            for token_idx, slot_idx in enumerate(slot_mapping_temp):
                slot_idx = slot_idx.item()

                block_idx = slot_idx // block_size
                block_offset = slot_idx % block_size

                memory_obj_old.tensor[0][layer_id][token_idx] = kv_cache[layer_id][
                    block_idx
                ][block_offset]

        memory_obj_old_list.append(memory_obj_old)
    end_event.record()
    # wait for all the operations to finish
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("Old extract time: ", elapsed_time_ms / 1000)

    # New extract with multi layer kernel
    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_cache[i].data_ptr()

    memory_obj_new_list = []
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        mem_obj_shape = torch.Size([1, num_layers, len(slot_mapping_temp), head_size])

        memory_obj_new = mem_allocator.allocate(mem_obj_shape, dtype)
        lmc_ops.multi_layer_kv_transfer(
            memory_obj_new.tensor,
            kv_cache_pointers,
            slot_mapping_temp,
            kv_cache[0].device,
            0,
            lmc_ops.TransferDirection.D2H,
            engine_kv_format,
            block_size,
        )
        memory_obj_new_list.append(memory_obj_new)

    end_event.record()
    # wait for all the operations to finish
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print("New extract time: ", elapsed_time_ms / 1000)

    for left_mem_obj, right_mem_obj in zip(
        memory_obj_old_list, memory_obj_new_list, strict=False
    ):
        left_kv, right_kv = left_mem_obj.tensor[0], right_mem_obj.tensor[0]
        right_kv = right_kv.to(left_kv.device)

        assert len(left_kv.shape) == 3
        assert len(right_kv.shape) == 3

        assert (left_kv[:, :, :] == right_kv[:, :, :]).all()

    # Generate new paged kv_cache
    kv_cache_new = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers,
        head_size,
        engine_kv_format=engine_kv_format,
    )

    kv_cache_pointers_new = torch.empty(
        num_layers, dtype=torch.int64, device="cpu", pin_memory=True
    )
    for i in range(num_layers):
        kv_cache_pointers_new[i] = kv_cache_new[i].data_ptr()

    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        memory_obj_new = memory_obj_new_list[chunk_id]
        lmc_ops.multi_layer_kv_transfer(
            memory_obj_new.tensor,
            kv_cache_pointers_new,
            slot_mapping_temp,
            kv_cache_new[0].device,
            0,
            lmc_ops.TransferDirection.H2D,
            engine_kv_format,
            block_size,
        )

    for left_kv, right_kv in zip(kv_cache, kv_cache_new, strict=False):
        assert len(left_kv.shape) == 3
        assert len(right_kv.shape) == 3

        left_reshaped = left_kv.reshape(
            left_kv.shape[0] * left_kv.shape[1], left_kv.shape[2]
        )
        right_reshaped = right_kv.reshape(
            right_kv.shape[0] * right_kv.shape[1], right_kv.shape[2]
        )

        assert (left_reshaped[slot_mapping, :] == right_reshaped[slot_mapping, :]).all()

    mem_allocator.close()


@pytest.mark.parametrize("num_tokens", [256, 500, 1024, 8000])
@pytest.mark.parametrize("token_major", [True, False])
@pytest.mark.parametrize(
    "engine_kv_format",
    [
        lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS,  # vllm non-MLA flash attention
        lmc_ops.EngineKVFormat.NL_X_NB_TWO_BS_NH_HS,  # vllm non-MLA flash infer
    ],
)
def test_single_layer_kernel(num_tokens, token_major, engine_kv_format):
    device = "cuda"

    num_layers = 32
    num_blocks = 1000
    block_size = 16
    num_heads = 8
    head_size = 128
    hidden_dim_size = num_heads * head_size
    dtype = torch.bfloat16
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype, engine_kv_format=engine_kv_format
    )
    kv_cache_new = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype, engine_kv_format=engine_kv_format
    )
    slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    if token_major:
        tmp_gpu_buffer = torch.empty(
            (num_tokens, 2, hidden_dim_size), dtype=dtype, device=device
        )
    else:
        tmp_gpu_buffer = torch.empty(
            (2, num_tokens, hidden_dim_size), dtype=dtype, device=device
        )

    for layer_id in range(num_layers):
        lmc_ops.single_layer_kv_transfer(
            tmp_gpu_buffer,
            kv_cache[layer_id],
            slot_mapping,
            lmc_ops.TransferDirection.D2H,
            engine_kv_format,
            token_major,
        )
        lmc_ops.single_layer_kv_transfer(
            tmp_gpu_buffer,
            kv_cache_new[layer_id],
            slot_mapping,
            lmc_ops.TransferDirection.H2D,
            engine_kv_format,
            token_major,
        )

    check_paged_kv_cache_equal(
        kv_cache,
        kv_cache_new,
        slot_mapping,
        engine_kv_format=engine_kv_format,
    )


@pytest.mark.parametrize("num_tokens", [256, 500, 1024])
@pytest.mark.parametrize(
    "engine_kv_format",
    [
        lmc_ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS,  # vllm HND flash attention
        lmc_ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS,  # vllm HND flash infer
    ],
)
def test_multi_layer_kernel_hnd(num_tokens, engine_kv_format):
    """Round-trip test for HND multi-layer kernel: D2H then H2D."""
    device = "cuda"

    num_blocks = 200
    block_size = 16
    num_heads = 8
    head_size = 128
    chunk_size = 256
    num_layers = 4
    dtype = torch.bfloat16

    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers=num_layers,
        head_size=head_size,
        engine_kv_format=engine_kv_format,
    )
    page_buffer_size = num_blocks * block_size

    slot_mapping = random.sample(range(0, page_buffer_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    pinned_cpu_size = 512 * 1024 * 1024  # 512MB
    mem_allocator = PinMemoryAllocator(pinned_cpu_size)

    # D2H: extract from paged cache into flat memory
    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_cache[i].data_ptr()

    memory_obj_list = []
    slot_mapping_chunked = torch.split(slot_mapping, chunk_size)
    for slot_mapping_temp in slot_mapping_chunked:
        mem_obj_shape = torch.Size(
            [2, num_layers, len(slot_mapping_temp), num_heads * head_size]
        )
        memory_obj = mem_allocator.allocate(mem_obj_shape, dtype)
        lmc_ops.multi_layer_kv_transfer(
            memory_obj.tensor,
            kv_cache_pointers,
            slot_mapping_temp,
            kv_cache[0].device,
            page_buffer_size,
            lmc_ops.TransferDirection.D2H,
            engine_kv_format,
            block_size,
            head_size=head_size,
        )
        memory_obj_list.append(memory_obj)
    torch.cuda.synchronize()

    # H2D: write back into a fresh paged cache
    kv_cache_new = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers=num_layers,
        head_size=head_size,
        engine_kv_format=engine_kv_format,
    )
    kv_cache_pointers_new = torch.empty(
        num_layers, dtype=torch.int64, device="cpu", pin_memory=True
    )
    for i in range(num_layers):
        kv_cache_pointers_new[i] = kv_cache_new[i].data_ptr()

    for chunk_id, slot_mapping_temp in enumerate(slot_mapping_chunked):
        lmc_ops.multi_layer_kv_transfer(
            memory_obj_list[chunk_id].tensor,
            kv_cache_pointers_new,
            slot_mapping_temp,
            kv_cache_new[0].device,
            page_buffer_size,
            lmc_ops.TransferDirection.H2D,
            engine_kv_format,
            block_size,
            head_size=head_size,
        )
    torch.cuda.synchronize()

    check_paged_kv_cache_equal(
        kv_cache,
        kv_cache_new,
        slot_mapping,
        engine_kv_format=engine_kv_format,
    )
    mem_allocator.close()


@pytest.mark.parametrize("num_tokens", [256, 500, 1024])
@pytest.mark.parametrize("token_major", [True, False])
@pytest.mark.parametrize(
    "engine_kv_format",
    [
        lmc_ops.EngineKVFormat.NL_X_TWO_NB_NH_BS_HS,  # vllm HND flash attention
        lmc_ops.EngineKVFormat.NL_X_NB_TWO_NH_BS_HS,  # vllm HND flash infer
    ],
)
def test_single_layer_kernel_hnd(num_tokens, token_major, engine_kv_format):
    """Round-trip test for HND single-layer kernel: D2H then H2D per layer."""
    device = "cuda"

    num_layers = 4
    num_blocks = 200
    block_size = 16
    num_heads = 8
    head_size = 128
    hidden_dim_size = num_heads * head_size
    dtype = torch.bfloat16

    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers=num_layers,
        head_size=head_size,
        engine_kv_format=engine_kv_format,
    )
    kv_cache_new = generate_kv_cache_paged_list_tensors(
        num_blocks,
        device,
        block_size,
        dtype,
        num_layers=num_layers,
        head_size=head_size,
        engine_kv_format=engine_kv_format,
    )
    slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    if token_major:
        tmp_gpu_buffer = torch.empty(
            (num_tokens, 2, hidden_dim_size), dtype=dtype, device=device
        )
    else:
        tmp_gpu_buffer = torch.empty(
            (2, num_tokens, hidden_dim_size), dtype=dtype, device=device
        )

    for layer_id in range(num_layers):
        lmc_ops.single_layer_kv_transfer(
            tmp_gpu_buffer,
            kv_cache[layer_id],
            slot_mapping,
            lmc_ops.TransferDirection.D2H,
            engine_kv_format,
            token_major,
        )
        lmc_ops.single_layer_kv_transfer(
            tmp_gpu_buffer,
            kv_cache_new[layer_id],
            slot_mapping,
            lmc_ops.TransferDirection.H2D,
            engine_kv_format,
            token_major,
        )

    check_paged_kv_cache_equal(
        kv_cache,
        kv_cache_new,
        slot_mapping,
        engine_kv_format=engine_kv_format,
    )


def test_lmcache_memcpy_async():
    # Register 4 memory regions and try to launch copy
    chunk_size = 1 << 26
    num_chunks = 4
    dtype = torch.bfloat16
    elements_per_chunk = chunk_size // dtype.itemsize

    big_cpu_tensor = torch.rand(
        elements_per_chunk * num_chunks, dtype=dtype, device="cpu"
    )
    big_gpu_tensor = torch.rand(
        elements_per_chunk * num_chunks, dtype=dtype, device="cuda"
    )

    def check_gpu_and_cpu_equal(
        gpu_tensor: torch.Tensor,
        cpu_tensor: torch.Tensor,
    ):
        cpu_tensor_from_gpu = gpu_tensor.to("cpu")
        assert torch.allclose(cpu_tensor_from_gpu, cpu_tensor, atol=1e-5)

    rt = torch.cuda.cudart()

    # Register the cpu memory
    ptr = big_cpu_tensor.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostRegister(ptr + i * chunk_size, chunk_size, 0)

    # Launch default cuda copy
    with pytest.raises(RuntimeError):
        big_gpu_tensor.copy_(big_cpu_tensor, non_blocking=True)
        torch.cuda.synchronize()

    # Launc default cuda copy for a small page
    big_gpu_tensor[: elements_per_chunk // 2].copy_(
        big_cpu_tensor[: elements_per_chunk // 2], non_blocking=True
    )
    torch.cuda.synchronize()

    check_gpu_and_cpu_equal(
        big_gpu_tensor[: elements_per_chunk // 2],
        big_cpu_tensor[: elements_per_chunk // 2],
    )

    # Positions to copy
    # - 0.25 chunk -- 0.75 chunk
    # - 0.75 chunk -- 1.25 chunk
    # - 1.75 chunk -- 3.25 chunk
    # - whole 4 chunks
    starts = [chunk_size // 4, (3 * chunk_size) // 4, (7 * chunk_size) // 4, 0]
    ends = [
        (3 * chunk_size) // 4,
        (5 * chunk_size) // 4,
        (13 * chunk_size) // 4,
        chunk_size * num_chunks,
    ]

    # H2D copy
    for start, end in zip(starts, ends, strict=False):
        # should not be equal before copy
        with pytest.raises(AssertionError):
            check_gpu_and_cpu_equal(
                big_gpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
                big_cpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
            )

        lmc_ops.lmcache_memcpy_async(
            big_gpu_tensor.data_ptr() + start,
            big_cpu_tensor.data_ptr() + start,
            end - start,
            lmc_ops.TransferDirection.H2D,
            start,
            chunk_size,
        )
        torch.cuda.synchronize()

        check_gpu_and_cpu_equal(
            big_gpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
            big_cpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
        )

    # Reset the data in gpu
    big_gpu_tensor = torch.rand(
        elements_per_chunk * num_chunks, dtype=dtype, device="cuda"
    )
    # D2H copy
    for start, end in zip(starts, ends, strict=False):
        # copy from gpu to cpu
        with pytest.raises(AssertionError):
            check_gpu_and_cpu_equal(
                big_gpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
                big_cpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
            )

        lmc_ops.lmcache_memcpy_async(
            big_cpu_tensor.data_ptr() + start,
            big_gpu_tensor.data_ptr() + start,
            end - start,
            lmc_ops.TransferDirection.D2H,
            start,
            chunk_size,
        )
        torch.cuda.synchronize()

        check_gpu_and_cpu_equal(
            big_gpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
            big_cpu_tensor[start // dtype.itemsize : end // dtype.itemsize],
        )

    # Unregister the cpu memory
    ptr = big_cpu_tensor.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostUnregister(ptr + i * chunk_size)


def test_lmcache_memcpy_async_int8_hidden132():
    """Test lmcache_memcpy_async with int8 dtype and hidden_size=132.

    hidden_size=132 means each token occupies 132 bytes, which is NOT a
    multiple of 4 (uint32 size). This tests that the kernel correctly handles
    non-uint32-aligned data sizes when splitting copies across registered
    memory chunk boundaries.

    The KV buffer shape is (2, num_layers, num_tokens, hidden_size), so the
    total bytes = 2 * num_layers * num_tokens * 132. We pick values such that
    the copy spans across at least one PIN_CHUNK_SIZE (64 MB) boundary.
    """
    # Use PIN_CHUNK_SIZE = 64 MB as the registration granularity
    chunk_size = 1 << 26  # 64 MB
    num_chunks = 2
    dtype = torch.int8
    hidden_size = 132  # NOT a multiple of 4 — tests uint32 alignment handling

    # Total bytes must span at least one chunk boundary.
    # We allocate num_chunks * chunk_size bytes worth of int8 elements.
    total_bytes = chunk_size * num_chunks
    total_elements = total_bytes  # int8: 1 byte per element

    cpu_tensor = torch.randint(-128, 127, (total_elements,), dtype=dtype, device="cpu")
    gpu_tensor = torch.zeros(total_elements, dtype=dtype, device="cuda")

    rt = torch.cuda.cudart()

    # Register the cpu memory in PIN_CHUNK_SIZE chunks
    ptr = cpu_tensor.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostRegister(ptr + i * chunk_size, chunk_size, 0)

    def check_equal(gpu_t: torch.Tensor, cpu_t: torch.Tensor):
        assert torch.equal(gpu_t.cpu(), cpu_t), "GPU and CPU tensors are not equal"

    # Test cases: copy ranges that cross chunk boundaries with non-4-byte-aligned
    # sizes, simulating real KV cache transfers where hidden_size=132 (int8).
    #
    # Each "KV row" is hidden_size=132 bytes. We pick copy sizes that are
    # multiples of 132 but NOT multiples of 4, so the copy boundary falls on
    # a non-uint32-aligned offset within a registered chunk.
    #
    # num_tokens=3: 3 * 132 = 396 bytes (396 % 4 == 0, but 132 % 4 != 0)
    # num_tokens=1: 1 * 132 = 132 bytes (132 % 4 == 0)
    # num_tokens=5: 5 * 132 = 660 bytes (660 % 4 == 0)
    # Use an offset that puts the copy right across the 64MB boundary.
    boundary = chunk_size  # 64 MB boundary in bytes

    # Place the copy so it straddles the chunk boundary:
    # start = boundary - 2 * hidden_size, length = 4 * hidden_size
    # This means 2 * hidden_size bytes before the boundary and 2 * hidden_size
    # bytes after — each segment is 264 bytes, not a multiple of 4*hidden_size.
    test_cases = [
        # (start_bytes, num_hidden_rows)
        (boundary - 2 * hidden_size, 4),  # straddles boundary, 528 bytes
        (boundary - hidden_size, 3),  # straddles boundary, 396 bytes
        (boundary - 3 * hidden_size, 6),  # straddles boundary, 792 bytes
        (0, 1),  # single row at start, 132 bytes
        (total_bytes - hidden_size, 1),  # single row at end, 132 bytes
    ]

    for start_bytes, num_rows in test_cases:
        nbytes = num_rows * hidden_size
        end_bytes = start_bytes + nbytes
        assert end_bytes <= total_bytes, (
            f"Test case out of bounds: start={start_bytes}, nbytes={nbytes}"
        )

        # Reset gpu_tensor slice to zeros so we can detect the copy
        gpu_tensor[start_bytes:end_bytes].zero_()

        # H2D copy
        lmc_ops.lmcache_memcpy_async(
            gpu_tensor.data_ptr() + start_bytes,
            cpu_tensor.data_ptr() + start_bytes,
            nbytes,
            lmc_ops.TransferDirection.H2D,
            start_bytes,
            chunk_size,
        )
        torch.cuda.synchronize()

        check_equal(
            gpu_tensor[start_bytes:end_bytes],
            cpu_tensor[start_bytes:end_bytes],
        )

    # D2H copy: write known values to GPU, copy back to CPU, verify
    gpu_src = torch.randint(-128, 127, (total_elements,), dtype=dtype, device="cuda")
    cpu_dst = torch.zeros(total_elements, dtype=dtype, device="cpu")

    # Register cpu_dst for D2H
    ptr_dst = cpu_dst.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostRegister(ptr_dst + i * chunk_size, chunk_size, 0)

    for start_bytes, num_rows in test_cases:
        nbytes = num_rows * hidden_size
        end_bytes = start_bytes + nbytes

        cpu_dst[start_bytes:end_bytes].zero_()

        lmc_ops.lmcache_memcpy_async(
            cpu_dst.data_ptr() + start_bytes,
            gpu_src.data_ptr() + start_bytes,
            nbytes,
            lmc_ops.TransferDirection.D2H,
            start_bytes,
            chunk_size,
        )
        torch.cuda.synchronize()

        check_equal(
            gpu_src[start_bytes:end_bytes],
            cpu_dst[start_bytes:end_bytes],
        )

    # Unregister memory
    ptr = cpu_tensor.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostUnregister(ptr + i * chunk_size)

    ptr_dst = cpu_dst.data_ptr()
    for i in range(num_chunks):
        rt.cudaHostUnregister(ptr_dst + i * chunk_size)