vLLM/tests/kernels/attention/test_cache.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import random

import pytest
import torch

from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
from vllm import _custom_ops as ops
from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
from vllm.platforms import current_platform
from vllm.utils.torch_utils import nvfp4_split_data_scale, set_random_seed

pytestmark = pytest.mark.skip_global_cleanup

COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
DTYPES = [torch.bfloat16, torch.float]
NUM_TOKENS = [42]  # Arbitrary values for testing
NUM_LAYERS = [1]  # Arbitrary values for testing
NUM_HEADS = [8]  # Arbitrary values for testing
HEAD_SIZES = [64, 80, 256]
BLOCK_SIZES = [8, 16, 32]
CACHE_LAYOUTS = ["NHD", "HND"]
KV_SCALE_TYPES = ["tensor", "attn_head"]

# Parameters for MLA tests.
KV_LORA_RANKS = [256, 512]
QK_ROPE_HEAD_DIMS = [64]
NUM_TOKENS_MLA = [42]
BLOCK_SIZES_MLA = [16]
NUM_BLOCKS_MLA = [8]

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

NUM_MAPPINGS = [256]  # Arbitrary values for testing
SEEDS = [0]
CUDA_DEVICES = [
    f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
]

# We assume fp8 is always enabled for testing.
KV_CACHE_DTYPE = ["auto", "fp8"]

RESHAPE_FLASH_IMPLEMENTATIONS = ["cuda", "triton"]


@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(
    kv_cache_factory,
    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:
    if kv_cache_dtype == "fp8" and head_size % 16:
        pytest.skip()
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)
    # Create a random slot mapping.
    num_slots = block_size * num_blocks
    slot_mapping_lst = random.sample(range(num_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long)

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

    # Create the KV caches.
    key_caches, value_caches = kv_cache_factory(
        num_blocks,
        block_size,
        1,
        num_heads,
        head_size,
        kv_cache_dtype,
        dtype,
        seed,
        device,
    )
    key_cache, value_cache = key_caches[0], value_caches[0]

    # Using default kv_scale
    k_scale = (key.amax() / 64.0).to(torch.float32)
    v_scale = (value.amax() / 64.0).to(torch.float32)

    # Clone the KV caches.
    if kv_cache_dtype == "fp8":
        cloned_key_cache = torch.empty_like(key_cache, dtype=torch.float16)
        ops.convert_fp8(cloned_key_cache, key_cache, k_scale.item())
        cloned_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
        ops.convert_fp8(cloned_value_cache, value_cache, v_scale.item())
    else:
        cloned_key_cache = key_cache.clone()
        cloned_value_cache = value_cache.clone()

    # Call the reshape_and_cache kernel.
    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]),
    )
    ops.reshape_and_cache(
        key,
        value,
        key_cache,
        value_cache,
        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)
        ops.convert_fp8(result_key_cache, key_cache, k_scale.item())
        result_value_cache = torch.empty_like(value_cache, dtype=torch.float16)
        ops.convert_fp8(result_value_cache, value_cache, v_scale.item())

    # Run the reference implementation.
    reshaped_key = key.reshape(num_tokens, *key_cache[0, :, :, 0, :].shape)
    block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
    block_indices_lst = block_indices.cpu().tolist()
    block_offsets = slot_mapping % block_size
    block_offsets_lst = block_offsets.cpu().tolist()
    for i in range(num_tokens):
        block_idx = block_indices_lst[i]
        block_offset = block_offsets_lst[i]
        cloned_key_cache[block_idx, :, :, block_offset, :] = reshaped_key[i]
        cloned_value_cache[block_idx, :, :, block_offset] = value[i]

    if kv_cache_dtype == "fp8":
        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
        )
    else:
        torch.testing.assert_close(key_cache, cloned_key_cache)
        torch.testing.assert_close(value_cache, cloned_value_cache)


@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 + ["nvfp4"])
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@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,
    kv_cache_layout: str,
    kv_scale_type: str,
    implementation: str,
) -> None:
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)
    assert implementation in ["cuda", "triton"]
    if implementation == "triton" and kv_cache_layout == "HND":
        pytest.skip("Triton implementation only supports NHD layout.")

    if kv_scale_type == "attn_head" and implementation != "cuda":
        pytest.skip("Only CUDA implementation supports attn_head scaling.")

    if kv_cache_dtype == "nvfp4":
        if not current_platform.has_device_capability(100):
            pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
        if implementation != "cuda":
            pytest.skip("NVFP4 only supports CUDA implementation.")
        if kv_scale_type != "tensor":
            pytest.skip("NVFP4 only supports per-tensor scaling.")
        if head_size % 16 != 0:
            pytest.skip("NVFP4 requires head_size divisible by 16.")
        if (head_size // 16) % 4 != 0:
            pytest.skip(
                "NVFP4 requires (head_size // 16) divisible by 4 "
                "for 4x4 block scale swizzle."
            )
        if block_size % 4 != 0:
            pytest.skip("NVFP4 requires block_size divisible by 4.")
        if dtype not in (torch.float16, torch.bfloat16):
            pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")

    # fp8 conversion requires continugous memory buffer. Reduce the number of
    # blocks and tokens to consume less memory.
    num_tokens = num_tokens // 2
    num_blocks = num_blocks // 2
    # Create a random slot mapping.
    num_slots = block_size * num_blocks
    slot_mapping_lst = random.sample(range(num_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
    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,
        device=device,
        cache_layout=kv_cache_layout,
    )
    key_cache, value_cache = key_caches[0], value_caches[0]
    del key_caches
    del value_caches

    # For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
    # Split views are still needed for dequant verification.
    key_scale_cache = None
    value_scale_cache = None
    nvfp4_key_data = None
    nvfp4_value_data = None
    if kv_cache_dtype == "nvfp4":
        nvfp4_key_data, key_scale_cache = nvfp4_split_data_scale(key_cache)
        nvfp4_value_data, value_scale_cache = nvfp4_split_data_scale(value_cache)

    if kv_cache_dtype == "nvfp4":
        # Global scale = amax / 448 (per-tensor)
        k_scale = (key.abs().amax() / 448.0).to(torch.float32)
        v_scale = (value.abs().amax() / 448.0).to(torch.float32)
    elif kv_scale_type == "tensor":
        k_scale = (key.amax() / 64.0).to(torch.float32)
        v_scale = (value.amax() / 64.0).to(torch.float32)
    else:  # "attn_head"
        k_scale = (key.amax(dim=(0, 2)) / 64.0).to(torch.float32)
        v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32)

    def permute_and_compact(x):
        y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
        return y.contiguous()

    if kv_cache_dtype != "nvfp4":
        key_cache_compact = permute_and_compact(key_cache)
        value_cache_compact = permute_and_compact(value_cache)

    def convert_fp8_local(output, input, scale, kv_dtype):
        fp8_input = input.view(current_platform.fp8_dtype())
        if scale.numel() == 1:  # per-tensor
            result = scaled_dequantize(
                fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype
            ).reshape(*input.shape)
        else:  # per-head: broadcast scale along the head dimension
            # Original code uses dim 2 for NHD, dim 1 for HND
            if kv_cache_layout == "NHD":
                result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1)
            else:
                result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
        output.copy_(result)

    # Clone the KV caches (for non-nvfp4, used as reference baseline).
    if kv_cache_dtype == "fp8":
        cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
        convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
        cloned_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
        convert_fp8_local(
            cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
        )
    elif kv_cache_dtype != "nvfp4":
        cloned_key_cache = key_cache_compact.clone()
        cloned_value_cache = value_cache_compact.clone()

    # Call the reshape_and_cache kernel.
    if implementation == "cuda":
        if kv_cache_dtype != "nvfp4":
            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]),
            )
        ops.reshape_and_cache_flash(
            key,
            value,
            key_cache,
            value_cache,
            slot_mapping,
            kv_cache_dtype,
            k_scale,
            v_scale,
        )
    elif implementation == "triton":
        from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
            triton_reshape_and_cache_flash,
        )

        triton_reshape_and_cache_flash(
            key,
            value,
            key_cache,
            value_cache,
            slot_mapping,
            kv_cache_dtype,
            k_scale,
            v_scale,
        )

    if kv_cache_dtype == "nvfp4":
        # Verify NVFP4 by dequantizing the entire cache and comparing
        # the written positions against original bf16 values.
        # Same pattern as FP8: dequant whole cache, then extract and compare.
        from tests.kernels.quantization.nvfp4_utils import (
            dequant_nvfp4_kv_cache,
        )

        def dequant_nvfp4_cache(data_cache, scale_cache, global_scale, swizzled_scales):
            # data_cache:  [B, N, H, data_dim]  logical view (layout in strides)
            # scale_cache: [B, N, H, scale_dim] logical view (layout in strides)
            # Permute to [B, H, N, dim] for the dequant utility.
            data_hnd = data_cache.permute(0, 2, 1, 3)
            scale_hnd = scale_cache.permute(0, 2, 1, 3)
            result_hnd = dequant_nvfp4_kv_cache(
                data_hnd,
                scale_hnd,
                global_scale,
                head_size,
                block_size,
                swizzled_scales=swizzled_scales,
            )
            return result_hnd.permute(0, 2, 1, 3)  # back to [B, N, H, dim]

        # The kernel writes K scales linearly on every arch and V scales in
        # the 4x4 swizzle of the SM100 trtllm-gen reader. The FlashInfer
        # reader used on SM12x expects linear V scales as well.
        v_scales_swizzled = not current_platform.is_device_capability_family(120)
        result_key_cache = dequant_nvfp4_cache(
            nvfp4_key_data, key_scale_cache, k_scale.item(), swizzled_scales=False
        )
        result_value_cache = dequant_nvfp4_cache(
            nvfp4_value_data,
            value_scale_cache,
            v_scale.item(),
            swizzled_scales=v_scales_swizzled,
        )

        # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
        num_slots = num_blocks * block_size
        result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
        result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)

        torch.testing.assert_close(
            result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
        )
        torch.testing.assert_close(
            result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
        )
        return

    key_cache_compact = permute_and_compact(key_cache)
    value_cache_compact = permute_and_compact(value_cache)

    if kv_cache_dtype == "fp8":
        result_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
        convert_fp8_local(result_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
        result_value_cache = torch.empty_like(value_cache_compact, dtype=torch.float16)
        convert_fp8_local(
            result_value_cache,
            value_cache_compact,
            v_scale,
            kv_cache_dtype,
        )

    # Run the reference implementation.
    block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
    block_indices_lst = block_indices.cpu().tolist()
    block_offsets = slot_mapping % block_size
    block_offsets_lst = block_offsets.cpu().tolist()
    for i in range(num_tokens):
        block_idx = block_indices_lst[i]
        block_offset = block_offsets_lst[i]
        if kv_cache_layout == "NHD":
            cloned_key_cache[block_idx, block_offset, :, :] = key[i]
            cloned_value_cache[block_idx, block_offset, :, :] = value[i]
        else:
            cloned_key_cache[block_idx, :, block_offset, :] = key[i]
            cloned_value_cache[block_idx, :, block_offset, :] = value[i]

    if kv_cache_dtype == "fp8":
        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
        )
    else:
        torch.testing.assert_close(key_cache_compact, cloned_key_cache)
        torch.testing.assert_close(value_cache_compact, cloned_value_cache)


@torch.inference_mode()
def test_nvfp4_4over6_selects_lower_error_scale(
    kv_cache_factory_flashinfer,
) -> None:
    if not current_platform.has_device_capability(100):
        pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")

    device = CUDA_DEVICES[0]
    block_size = 16
    head_size = 64
    source = torch.tensor(
        [6.0] + [4.5] * 15,
        dtype=torch.bfloat16,
        device=device,
    ).repeat(1, 1, head_size // 16)
    slot_mapping = torch.zeros(1, dtype=torch.long, device=device)
    scale = torch.ones(1, dtype=torch.float32, device=device)

    def quantize_and_dequantize(cache_dtype: str) -> torch.Tensor:
        key_caches, value_caches = kv_cache_factory_flashinfer(
            1,
            block_size,
            1,
            1,
            head_size,
            cache_dtype,
            source.dtype,
            device=device,
            cache_layout="NHD",
        )
        key_cache, value_cache = key_caches[0], value_caches[0]
        key_cache.zero_()
        value_cache.zero_()

        ops.reshape_and_cache_flash(
            source,
            source,
            key_cache,
            value_cache,
            slot_mapping,
            cache_dtype,
            scale,
            scale,
        )

        from tests.kernels.quantization.nvfp4_utils import (
            dequant_nvfp4_kv_cache,
        )

        data, block_scales = nvfp4_split_data_scale(key_cache)
        return dequant_nvfp4_kv_cache(
            data.permute(0, 2, 1, 3),
            block_scales.permute(0, 2, 1, 3),
            scale.item(),
            head_size,
            block_size,
            swizzled_scales=False,
        )[0, 0, 0]

    default = quantize_and_dequantize("nvfp4")
    four_over_six = quantize_and_dequantize("nvfp4_4over6")
    source_f32 = source[0, 0].float()

    default_mse = (default - source_f32).square().mean()
    four_over_six_mse = (four_over_six - source_f32).square().mean()
    assert four_over_six_mse < default_mse


@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
@pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@torch.inference_mode()
def test_reshape_and_cache_flash_unaligned_rows(
    kv_cache_factory_flashinfer,
    dtype: torch.dtype,
    kv_cache_dtype: str,
    kv_cache_layout: str,
    implementation: str,
) -> None:
    """Regression test for https://github.com/vllm-project/vllm/issues/41257.

    head_size=46 with num_heads=13 places KV-cache rows at byte offsets
    that are not a multiple of the vector width (NHD row pitch
    13*46*itemsize, HND head pitch 46*itemsize), unlike HEAD_SIZES above
    which are all 16-byte multiples. The CUDA kernel used to issue
    vectorized stores to those rows -> CUDA misaligned address.
    """
    test_reshape_and_cache_flash(
        kv_cache_factory_flashinfer,
        num_tokens=42,
        num_heads=13,
        head_size=46,
        block_size=16,
        num_blocks=128,
        dtype=dtype,
        seed=0,
        device=CUDA_DEVICES[0],
        kv_cache_dtype=kv_cache_dtype,
        kv_cache_layout=kv_cache_layout,
        kv_scale_type="tensor",
        implementation=implementation,
    )


@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)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
@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,
    device: str,
    kv_cache_dtype: str,
) -> None:
    if kv_cache_dtype == "fp8" and "cpu" in direction:
        pytest.skip()
    if kv_cache_dtype == "fp8" and head_size % 16:
        pytest.skip()

    set_random_seed(seed)

    src_device = device if direction[0] == "cuda" else "cpu"
    dst_device = device if direction[1] == "cuda" else "cpu"

    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)

    block_mapping = list(zip(src_blocks, dst_blocks))
    block_mapping_tensor = torch.tensor(
        block_mapping, dtype=torch.int64, device="cpu"
    ).view(-1, 2)

    # Create the KV caches on the first device.
    src_key_caches, src_value_caches = kv_cache_factory(
        num_blocks,
        block_size,
        1,
        num_heads,
        head_size,
        kv_cache_dtype,
        dtype,
        seed,
        src_device,
    )

    # Create the KV caches on the second device.
    dist_key_caches, dist_value_caches = kv_cache_factory(
        num_blocks,
        block_size,
        1,
        num_heads,
        head_size,
        kv_cache_dtype,
        dtype,
        seed,
        dst_device,
    )

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

    # Call the swap_blocks kernel.
    do_opcheck = head_size == HEAD_SIZES[0]
    src_cache = src_key_caches[0]
    block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
        (
            src_key_caches[0],
            dist_key_caches[0],
            block_size_in_bytes,
            block_mapping_tensor,
        ),
        cond=do_opcheck,
    )
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
        (
            src_value_caches[0],
            dist_value_caches[0],
            block_size_in_bytes,
            block_mapping_tensor,
        ),
        cond=do_opcheck,
    )

    ops.swap_blocks(
        src_key_caches[0],
        dist_key_caches[0],
        block_size_in_bytes,
        block_mapping_tensor,
    )
    ops.swap_blocks(
        src_value_caches[0],
        dist_value_caches[0],
        block_size_in_bytes,
        block_mapping_tensor,
    )

    src_indices = block_mapping_tensor[:, 0].to(src_key_caches_clone.device)
    dst_indices = block_mapping_tensor[:, 1].to(dist_key_caches[0].device)

    torch.testing.assert_close(
        src_key_caches_clone.index_select(0, src_indices).cpu(),
        dist_key_caches[0].index_select(0, dst_indices).cpu(),
    )
    torch.testing.assert_close(
        src_value_caches_clone.index_select(0, src_indices).cpu(),
        dist_value_caches[0].index_select(0, dst_indices).cpu(),
    )


@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()
def test_fp8_e4m3_conversion(
    num_heads: int,
    head_size: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
) -> None:
    set_random_seed(seed)

    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)
    ops.convert_fp8(cache_fp8, cache)

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

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


def _create_mla_cache(
    num_blocks: int,
    block_size: int,
    entry_size: int,
    dtype: torch.dtype,
    kv_cache_dtype: str,
    device: str,
) -> torch.Tensor:
    cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype
    return torch.zeros(
        num_blocks, block_size, entry_size, dtype=cache_dtype, device=device
    )


def _fill_mla_cache(cache: torch.Tensor, kv_cache_dtype: str):
    rand_dtype = torch.float16 if kv_cache_dtype == "fp8" else cache.dtype

    vals = torch.randn(*cache.shape, device=cache.device, dtype=rand_dtype)
    if kv_cache_dtype == "fp8":
        temp = torch.zeros_like(cache)
        ops.convert_fp8(temp, vals, 1.0, kv_dtype=kv_cache_dtype)
        vals = temp
    cache.copy_(vals)


@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
@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_concat_and_cache_mla(
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    num_tokens: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
    kv_cache_dtype: str,
) -> None:
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)

    total_slots = num_blocks * block_size
    slot_mapping_lst = random.sample(range(total_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)

    kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
    k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)
    entry_size = kv_lora_rank + qk_rope_head_dim

    scale = torch.tensor(0.1, dtype=torch.float32, device=device)
    kv_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)

    for i in range(num_tokens):
        slot = slot_mapping[i].item()
        block_idx = slot // block_size
        block_offset = slot % block_size
        ref_temp[block_idx, block_offset, :kv_lora_rank] = kv_c[i]
        ref_temp[block_idx, block_offset, kv_lora_rank:] = k_pe[i]

    if kv_cache_dtype == "fp8":
        ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
        ops.convert_fp8(ref_kv_cache, ref_temp, scale.item(), kv_dtype=kv_cache_dtype)
    else:
        ref_kv_cache = ref_temp

    opcheck(
        torch.ops._C_cache_ops.concat_and_cache_mla,
        (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)

    if kv_cache_dtype == "fp8":
        result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
        ops.convert_fp8(
            result_temp, kv_cache.contiguous(), scale.item(), kv_dtype=kv_cache_dtype
        )
        expected_temp = torch.empty_like(ref_kv_cache, dtype=torch.float16)
        ops.convert_fp8(
            expected_temp, ref_kv_cache, scale.item(), kv_dtype=kv_cache_dtype
        )
        torch.testing.assert_close(result_temp, expected_temp, atol=0.001, rtol=0.1)
    else:
        torch.testing.assert_close(kv_cache, ref_kv_cache)


@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("shared_slot_mapping", [False, True])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e5m2"])
@torch.inference_mode()
def test_concat_and_cache_mla_grouped(
    device: str,
    shared_slot_mapping: bool,
    kv_cache_dtype: str,
) -> None:
    if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
        pytest.skip("fp8_e5m2 KV cache is not supported on ROCm/HIP")

    set_random_seed(0)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)

    num_layers = 5
    num_tokens = 42
    num_blocks = 8
    block_size = 16
    kv_lora_rank = 512
    qk_rope_head_dim = 64
    entry_size = kv_lora_rank + qk_rope_head_dim
    total_slots = num_blocks * block_size

    kv_c = torch.randn(num_layers, num_tokens, kv_lora_rank, dtype=torch.bfloat16)
    k_pe = torch.randn(num_layers, num_tokens, qk_rope_head_dim, dtype=torch.bfloat16)
    kv_caches = torch.zeros(
        num_layers,
        num_blocks,
        block_size,
        entry_size,
        dtype=torch.bfloat16 if kv_cache_dtype == "auto" else torch.uint8,
    )
    reference = torch.zeros_like(kv_caches)
    scales = torch.linspace(0.05, 0.25, num_layers, dtype=torch.float32)

    slot_mapping = torch.stack(
        [torch.randperm(total_slots)[:num_tokens] for _ in range(num_layers)]
    )
    slot_mapping[:, -1] = -1
    if shared_slot_mapping:
        slot_mapping = slot_mapping[:1].expand(num_layers, -1)

    for layer_idx in range(num_layers):
        ops.concat_and_cache_mla(
            kv_c[layer_idx],
            k_pe[layer_idx],
            reference[layer_idx],
            slot_mapping[layer_idx],
            kv_cache_dtype,
            scales[layer_idx],
        )

    cache_ptrs = torch.tensor(
        [kv_caches[layer_idx].data_ptr() for layer_idx in range(num_layers)],
        dtype=torch.int64,
    )
    ref_cache = kv_caches[0]

    def run_grouped(kv_scales: torch.Tensor | None) -> None:
        ops.concat_and_cache_mla_grouped(
            kv_c,
            k_pe,
            cache_ptrs,
            slot_mapping,
            ref_cache.size(1),
            ref_cache.stride(0),
            ref_cache.stride(1),
            kv_scales,
            kv_cache_dtype,
        )

    run_grouped(None if kv_cache_dtype == "auto" else scales)

    torch.testing.assert_close(kv_caches, reference, rtol=0, atol=0)

    if kv_cache_dtype == "fp8" and not shared_slot_mapping:
        noncontiguous_scales = torch.ones(num_layers, 2)[:, 0]
        assert not noncontiguous_scales.is_contiguous()
        for invalid_scales, error in (
            (scales.cpu(), "same CUDA device"),
            (noncontiguous_scales, "must be contiguous"),
        ):
            with pytest.raises(RuntimeError, match=error):
                run_grouped(invalid_scales)


@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_concat_and_cache_ds_mla(
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    num_tokens: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
) -> None:
    if current_platform.is_rocm():
        pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
    if dtype.itemsize != 2:
        pytest.skip("ds_mla only supports 16-bit input")
    if kv_lora_rank != 512:
        pytest.skip("fp8_ds_mla requires kv_lora_rank == 512")
    kv_cache_dtype = "fp8_ds_mla"
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)

    total_slots = num_blocks * block_size
    slot_mapping_lst = random.sample(range(total_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)

    kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
    k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)
    entry_size = kv_lora_rank + (4 * 4) + (2 * qk_rope_head_dim)

    scale = torch.tensor(1.0, dtype=torch.float32, device=device)
    kv_cache = _create_mla_cache(
        num_blocks,
        block_size,
        entry_size,
        dtype=torch.uint8,
        kv_cache_dtype=kv_cache_dtype,
        device=device,
    )

    ref_cache = torch.zeros_like(kv_cache, dtype=kv_cache.dtype)
    tile_data = torch.zeros(128, dtype=dtype, device=device)

    for i in range(num_tokens):
        slot = slot_mapping[i].item()
        block_idx = slot // block_size
        block_offset = slot % block_size

        ref_cache_slice = ref_cache[block_idx, block_offset]
        ref_cache_16bit = ref_cache_slice.view(dtype)
        ref_cache_32bit = ref_cache_slice.view(torch.float32)

        kv_c_data = kv_c[i]
        num_tiles = kv_lora_rank // 128
        for tile_idx in range(num_tiles):
            tile_start = tile_idx * 128
            tile_end = (tile_idx + 1) * 128
            tile_data[:] = kv_c_data[tile_start:tile_end]

            # Using torch's amax() gives different results, so this must be
            # manually computed.
            tile_data_float = tile_data.to(torch.float32)
            manual_max = abs(tile_data_float[0])
            for j in range(1, 128):
                manual_max = max(manual_max, abs(tile_data_float[j]))
            raw_scale = torch.clamp(manual_max / 448.0, min=1e-4)
            tile_scale = torch.exp2(torch.ceil(torch.log2(raw_scale)))

            ref_cache_32bit[kv_lora_rank // 4 + tile_idx] = tile_scale

            ops.convert_fp8(
                ref_cache_slice[tile_start:tile_end],
                tile_data,
                tile_scale.item(),
                kv_dtype="fp8",
            )

        for j in range(qk_rope_head_dim):
            ref_cache_16bit[kv_lora_rank // 2 + 8 + j] = k_pe[i, j]

    opcheck(
        torch.ops._C_cache_ops.concat_and_cache_mla,
        (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)

    for i in range(num_tokens):
        slot = slot_mapping[i].item()
        block_idx = slot // block_size
        block_offset = slot % block_size
        kv_cache_slice = kv_cache[block_idx, block_offset]
        ref_cache_slice = ref_cache[block_idx, block_offset]

        kv_nope = kv_cache_slice[:kv_lora_rank]
        ref_nope = ref_cache_slice[:kv_lora_rank]
        kv_scales = kv_cache_slice.view(torch.float32)[
            kv_lora_rank // 4 : kv_lora_rank // 4 + 4
        ]
        ref_scales = ref_cache_slice.view(torch.float32)[
            kv_lora_rank // 4 : kv_lora_rank // 4 + 4
        ]
        kv_rope = kv_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :]
        ref_rope = ref_cache_slice.view(dtype)[kv_lora_rank // 2 + 8 :]

        torch.testing.assert_close(kv_nope, ref_nope, atol=0.001, rtol=0.1)
        torch.testing.assert_close(kv_scales, ref_scales, atol=0, rtol=0)
        torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)


# Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py).
NVFP4_DS_MLA_ENTRY_SIZE = 352

# e2m1 magnitude table; the sign lives in bit 3 of each 4-bit code.
E2M1_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]


def _unpack_e2m1(packed: torch.Tensor) -> torch.Tensor:
    """Unpack uint8 bytes holding two e2m1 codes each into float32 values.

    The low nibble holds the even element, the high nibble the odd one.
    """
    low = packed & 0x0F
    high = packed >> 4
    codes = torch.stack([low, high], dim=-1).reshape(-1).long()
    table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=packed.device)
    magnitude = table[codes & 0x7]
    sign = torch.where((codes & 0x8) != 0, -1.0, 1.0)
    return sign * magnitude


def _ref_nvfp4_sf_bytes(vals: torch.Tensor, divisor: float) -> torch.Tensor:
    """Reference e4m3 scale-factor bytes: e4m3(max(amax_per_16 / divisor, 2^-9)).

    In element-block order (byte s scales elements [16s, 16s + 16)).
    """
    tiles = vals.to(torch.float32).reshape(-1, 16)
    sf = torch.clamp(tiles.abs().amax(dim=1) / divisor, min=2.0**-9)
    return sf.to(torch.float8_e4m3fn).view(torch.uint8)


def _nvfp4_sf_element_order(wire_sf_bytes: torch.Tensor) -> torch.Tensor:
    """Stored NVFP4 scale-factor byte order -> element-block order.

    The 32 NoPE scale-factor bytes of an nvfp4_ds_mla entry are stored
    permuted (an 8x4 -> 4x8 transpose): the scale for element block s lives at
    byte 8 * (s & 3) + (s >> 2), so that the 8 scales one FlashMLA dequant
    thread needs are contiguous. This undoes that permutation.
    """
    num_sf = wire_sf_bytes.shape[-1]
    return wire_sf_bytes.unflatten(-1, (4, num_sf // 4)).transpose(-1, -2).flatten(-2)


@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_concat_and_cache_nvfp4_ds_mla(
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    num_tokens: int,
    block_size: int,
    num_blocks: int,
    seed: int,
    device: str,
) -> None:
    if current_platform.is_rocm():
        pytest.skip("concat_and_cache_mla doesn't support NVFP4 DS-MLA on ROCm")
    device_capability = current_platform.get_device_capability()
    if device_capability is None or device_capability.major != 10:
        pytest.skip("The NVFP4 DS-MLA kv-cache dtype requires SM 10.x")
    if kv_lora_rank != 512:
        pytest.skip("The NVFP4 DS-MLA layout requires kv_lora_rank == 512")
    dtype = torch.bfloat16
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)

    kv_cache_dtype = "nvfp4_ds_mla"
    entry_size = NVFP4_DS_MLA_ENTRY_SIZE

    total_slots = num_blocks * block_size
    slot_mapping_lst = random.sample(range(total_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)

    kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
    k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)

    scale = torch.tensor(1.0, dtype=torch.float32, device=device)
    kv_cache = _create_mla_cache(
        num_blocks,
        block_size,
        entry_size,
        dtype=torch.uint8,
        kv_cache_dtype=kv_cache_dtype,
        device=device,
    )

    opcheck(
        torch.ops._C_cache_ops.concat_and_cache_mla,
        (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)

    # Byte layout of a single cache entry:
    #   [0, 256)             512 e2m1 NoPE values packed 2/byte
    #   [256, 320)           64 unscaled e4m3 RoPE values
    #   [nope_sf_off, 352)   32 e4m3 NoPE SFs (one per 16 elements), stored
    #                        permuted (see _nvfp4_sf_element_order)
    nope_bytes = kv_lora_rank // 2
    rope_bytes = qk_rope_head_dim
    nope_sf_off = nope_bytes + rope_bytes
    num_nope_sf = kv_lora_rank // 16
    assert nope_sf_off + num_nope_sf == entry_size

    for i in range(num_tokens):
        slot = slot_mapping[i].item()
        block_idx = slot // block_size
        block_offset = slot % block_size
        entry = kv_cache[block_idx, block_offset]

        kv_c_ref = kv_c[i].to(torch.float32)
        k_pe_ref = k_pe[i].to(torch.float32)

        # Scale-factor bytes must match e4m3(max(amax_per_16 / divisor, 2^-9))
        # computed in python, allowing a 1-ulp e4m3 difference (all SFs are
        # positive, so e4m3 bit patterns are monotonic and adjacent codes
        # differ by 1).
        nope_sf_bytes = _nvfp4_sf_element_order(
            entry[nope_sf_off : nope_sf_off + num_nope_sf]
        )
        ref_nope_sf_bytes = _ref_nvfp4_sf_bytes(kv_c_ref, 6.0)
        assert (
            (nope_sf_bytes.to(torch.int16) - ref_nope_sf_bytes.to(torch.int16)).abs()
            <= 1
        ).all()

        # Dequantize with the *stored* scale factors: x = float(code) * float(sf).
        nope_sf = nope_sf_bytes.view(torch.float8_e4m3fn).to(torch.float32)
        nope_sf_per_elem = nope_sf.repeat_interleave(16)
        nope_vals = _unpack_e2m1(entry[:nope_bytes]) * nope_sf_per_elem
        # e2m1 error bound: the widest grid spacing is 2 (between codes 4 and
        # 6), so round-to-nearest is off by at most 1.0*sf; the e4m3 rounding
        # of the SF itself adds at most ~2^-4 relative, so 1.5*sf is safe for
        # either encode convention (quantized or unquantized SF).
        assert ((nope_vals - kv_c_ref).abs() <= 1.5 * nope_sf_per_elem).all()

        # Plain unscaled e4m3: ~2^-4 relative rounding, with a small absolute
        # floor for values in the subnormal range.
        rope_payload = entry[nope_bytes:nope_sf_off]
        rope_vals = rope_payload.view(torch.float8_e4m3fn).to(torch.float32)
        rope_tol = k_pe_ref.abs() * 2.0**-3 + 2.0**-9
        assert ((rope_vals - k_pe_ref).abs() <= rope_tol).all()


@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
@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_swap_blocks_mla(
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
    device: str,
    kv_cache_dtype: str,
) -> None:
    set_random_seed(seed)
    torch.set_default_device(device)
    torch.accelerator.set_device_index(device)

    entry_size = kv_lora_rank + qk_rope_head_dim

    src_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    dst_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )

    _fill_mla_cache(src_cache, kv_cache_dtype)
    _fill_mla_cache(dst_cache, kv_cache_dtype)

    src_cache_clone = src_cache.clone()

    num_mappings = min(2, num_blocks // 2)
    src_blocks = random.sample(range(num_blocks), num_mappings)
    remaining_blocks = list(set(range(num_blocks)) - set(src_blocks))
    dst_blocks = random.sample(remaining_blocks, num_mappings)
    block_mapping = list(zip(src_blocks, dst_blocks))
    block_mapping_tensor = torch.tensor(
        block_mapping, dtype=torch.int64, device="cpu"
    ).view(-1, 2)

    block_size_in_bytes = src_cache.element_size() * src_cache.stride(0)
    opcheck(
        torch.ops._C_cache_ops.swap_blocks,
        (src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.swap_blocks(src_cache, dst_cache, block_size_in_bytes, block_mapping_tensor)

    for src, dst in block_mapping:
        torch.testing.assert_close(
            src_cache_clone[src].cpu(),
            dst_cache[dst].cpu(),
            msg=f"Block {src} from src should have been swapped to block "
            f"{dst} in dst_cache.",
        )


@pytest.mark.parametrize("kv_lora_rank", [512])
@pytest.mark.parametrize("qk_rope_head_dim", [64])
@pytest.mark.parametrize("block_size", [16])
@pytest.mark.parametrize("num_blocks", [1024])
@pytest.mark.parametrize("max_seq_len", [512])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize("dtype", [torch.float32])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_gather_and_maybe_dequant_cache_mla(
    kv_lora_rank,
    qk_rope_head_dim,
    block_size,
    num_blocks,
    max_seq_len,
    batch_size,
    dtype,
    kv_cache_dtype,
    device,
):
    entry_size = kv_lora_rank + qk_rope_head_dim
    scale = torch.tensor(0.1, dtype=torch.float32, device=device)
    src_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    seq_len_tensor = torch.randint(
        max_seq_len, max_seq_len + 1, (batch_size,), device=device
    )

    total_tokens = seq_len_tensor.sum()
    cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device)
    cu_seq_lens[0] = 0
    cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32)
    token_to_seq = torch.arange(0, batch_size, dtype=torch.int32, device=device)
    token_to_seq = torch.repeat_interleave(token_to_seq, seq_len_tensor)
    print("seq_len_tensor", seq_len_tensor)

    tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size
    block_table = torch.empty(
        (batch_size, num_blocks), dtype=torch.int32, device=device
    )

    for b in range(batch_size):
        perm = torch.randperm(num_blocks, device=device)
        block_table[b, :] = perm

    dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device)

    expected_batches = []
    for b in range(batch_size):
        s = seq_len_tensor[b]
        if s == 0:
            continue
        tot = tot_blocks_tensor[b]
        blocks = block_table[b, :tot].tolist()

        gathered_rows = []
        for i in range(tot - 1):
            block_data = src_cache[blocks[i]]
            if kv_cache_dtype == "fp8":
                dequantized_block = torch.empty_like(block_data, dtype=dtype)
                ops.convert_fp8(dequantized_block, block_data, scale.item())
                gathered_rows.append(dequantized_block)
            else:
                gathered_rows.append(block_data)
        remaining = s - (tot - 1) * block_size
        last_block_data = src_cache[blocks[-1], :remaining, :]
        if kv_cache_dtype == "fp8":
            dequantized_last_block = torch.empty_like(last_block_data, dtype=dtype)
            ops.convert_fp8(dequantized_last_block, last_block_data, scale.item())
            gathered_rows.append(dequantized_last_block)
        else:
            gathered_rows.append(last_block_data)

        batch_expected = torch.cat(gathered_rows, dim=0)
        expected_batches.append(batch_expected)
    expected = torch.cat(expected_batches, dim=0)

    opcheck(
        torch.ops._C_cache_ops.gather_and_maybe_dequant_cache,
        (
            src_cache,
            dst,
            block_table,
            cu_seq_lens,
            token_to_seq,
            total_tokens,
            kv_cache_dtype,
            scale,
            None,
        ),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.gather_and_maybe_dequant_cache(
        src_cache,
        dst,
        block_table,
        cu_seq_lens,
        token_to_seq,
        total_tokens,
        kv_cache_dtype,
        scale,
        None,
    )
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_lora_rank", [512])
@pytest.mark.parametrize("qk_rope_head_dim", [64])
@pytest.mark.parametrize("block_size", [16])
@pytest.mark.parametrize("num_blocks", [128])
@pytest.mark.parametrize("dtype", [torch.float32])
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_gather_and_maybe_dequant_cache_mla_with_seq_starts(
    kv_lora_rank,
    qk_rope_head_dim,
    block_size,
    num_blocks,
    dtype,
    kv_cache_dtype,
    device,
):
    entry_size = kv_lora_rank + qk_rope_head_dim
    scale = torch.tensor(0.1, dtype=torch.float32, device=device)
    src_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    seq_starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
    seq_lens = torch.tensor([20, 10, 16], dtype=torch.int32, device=device)
    batch_size = seq_lens.shape[0]
    total_tokens = seq_lens.sum().item()
    cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device)
    cu_seq_lens[0] = 0
    cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
    token_to_seq = torch.repeat_interleave(
        torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens
    )

    block_table = torch.empty(
        (batch_size, num_blocks), dtype=torch.int32, device=device
    )
    for b in range(batch_size):
        block_table[b, :] = torch.randperm(num_blocks, device=device)

    if kv_cache_dtype == "fp8":
        dequant_src_cache = torch.empty_like(src_cache, dtype=dtype)
        ops.convert_fp8(dequant_src_cache, src_cache, scale.item())
    else:
        dequant_src_cache = src_cache

    expected_rows = []
    for b in range(batch_size):
        start = seq_starts[b].item()
        length = seq_lens[b].item()
        for offset in range(start, start + length):
            block_id = block_table[b, offset // block_size]
            slot = offset % block_size
            expected_rows.append(dequant_src_cache[block_id, slot])
    expected = torch.stack(expected_rows)

    dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device)
    ops.gather_and_maybe_dequant_cache(
        src_cache,
        dst,
        block_table,
        cu_seq_lens,
        token_to_seq,
        total_tokens,
        kv_cache_dtype,
        scale,
        seq_starts,
    )
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("use_seq_starts", [False, True])
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_gather_and_maybe_dequant_cache_mla_large_uneven_sequences(
    kv_cache_dtype,
    use_seq_starts,
    device,
):
    block_size = 64
    entry_size = 576
    num_blocks = 1024
    dtype = torch.bfloat16
    scale = torch.tensor(0.1, dtype=torch.float32, device=device)
    src_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
    seq_starts = starts if use_seq_starts else None
    seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device)
    batch_size = seq_lens.shape[0]
    total_tokens = seq_lens.sum().item()
    cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
    cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
    token_to_seq = torch.repeat_interleave(
        torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens
    )
    block_table = torch.stack(
        [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
    ).to(torch.int32)

    if kv_cache_dtype == "fp8":
        dequant_src_cache = torch.empty_like(src_cache, dtype=dtype)
        ops.convert_fp8(dequant_src_cache, src_cache, scale.item())
    else:
        dequant_src_cache = src_cache
    expected_batches = []
    for req_id in range(batch_size):
        start = starts[req_id].item() if use_seq_starts else 0
        source_tokens = torch.arange(
            start, start + seq_lens[req_id].item(), device=device
        )
        physical_blocks = block_table[req_id, source_tokens // block_size].long()
        expected_batches.append(
            dequant_src_cache[physical_blocks, source_tokens % block_size]
        )
    expected = torch.cat(expected_batches)
    dst = torch.empty_like(expected)

    ops.gather_and_maybe_dequant_cache(
        src_cache,
        dst,
        block_table,
        cu_seq_lens,
        token_to_seq,
        total_tokens,
        kv_cache_dtype,
        scale,
        seq_starts,
    )
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_lora_rank", [512])
@pytest.mark.parametrize("qk_rope_head_dim", [64])
@pytest.mark.parametrize("block_size", [16])
@pytest.mark.parametrize("num_blocks", [1024])
@pytest.mark.parametrize("max_seq_len", [512])
@pytest.mark.parametrize("batch_size", [8])
@pytest.mark.parametrize("dtype", [torch.float32])
@pytest.mark.parametrize(
    "kv_cache_dtype", ["auto"]
)  # You can also test "fp8" if needed.
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_cp_gather_cache_mla(
    kv_lora_rank,
    qk_rope_head_dim,
    block_size,
    num_blocks,
    max_seq_len,
    batch_size,
    dtype,
    kv_cache_dtype,
    device,
):
    entry_size = kv_lora_rank + qk_rope_head_dim
    src_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    seq_len_tensor = torch.randint(0, max_seq_len + 1, (batch_size,), device=device)

    total_tokens = seq_len_tensor.sum()
    cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device)
    cu_seq_lens[0] = 0
    cu_seq_lens[1:] = seq_len_tensor.cumsum(dim=0).to(dtype=torch.int32)
    print("seq_len_tensor", seq_len_tensor)

    tot_blocks_tensor = (seq_len_tensor + block_size - 1) // block_size
    block_table = torch.empty(
        (batch_size, num_blocks), dtype=torch.int32, device=device
    )

    for b in range(batch_size):
        perm = torch.randperm(num_blocks, device=device)
        block_table[b, :] = perm

    dst = torch.zeros((total_tokens, entry_size), dtype=src_cache.dtype, device=device)

    expected_batches = []
    for b in range(batch_size):
        s = seq_len_tensor[b]
        if s == 0:
            continue
        tot = tot_blocks_tensor[b]
        blocks = block_table[b, :tot].tolist()

        gathered_rows = []
        for i in range(tot - 1):
            gathered_rows.append(src_cache[blocks[i]])
        remaining = s - (tot - 1) * block_size
        gathered_rows.append(src_cache[blocks[-1], :remaining, :])

        batch_expected = torch.cat(gathered_rows, dim=0)
        expected_batches.append(batch_expected)
    expected = torch.cat(expected_batches, dim=0)

    opcheck(
        torch.ops._C_cache_ops.cp_gather_cache,
        (src_cache, dst, block_table, cu_seq_lens, batch_size, None),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.cp_gather_cache(src_cache, dst, block_table, cu_seq_lens, batch_size)
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("unaligned", [False, True])
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_cp_gather_cache_mla_with_seq_starts(
    kv_cache_dtype,
    unaligned,
    device,
):
    block_size = 16
    entry_size = 576
    num_blocks = 128
    src_storage = _create_mla_cache(
        num_blocks,
        block_size,
        entry_size + int(unaligned),
        torch.bfloat16,
        kv_cache_dtype,
        device,
    )
    src_cache = src_storage[..., 1:] if unaligned else src_storage
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    seq_starts = torch.tensor([3, 17, 5, 31], dtype=torch.int32, device=device)
    seq_lens = torch.tensor([0, 1, 63, 4], dtype=torch.int32, device=device)
    batch_size = seq_lens.shape[0]
    cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
    cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
    block_table = torch.stack(
        [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
    ).to(torch.int32)

    expected_rows = []
    for req_id in range(batch_size):
        start = seq_starts[req_id].item()
        for source_token in range(start, start + seq_lens[req_id].item()):
            block_id = block_table[req_id, source_token // block_size]
            expected_rows.append(src_cache[block_id, source_token % block_size])
    expected = torch.stack(expected_rows)
    if unaligned:
        dst_storage = torch.empty(
            (expected.shape[0], entry_size + 1),
            dtype=expected.dtype,
            device=device,
        )
        dst = dst_storage[:, 1:]
    else:
        dst = torch.empty_like(expected)

    ops.cp_gather_cache(
        src_cache,
        dst,
        block_table,
        cu_seq_lens,
        batch_size,
        seq_starts,
    )
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize("use_seq_starts", [False, True])
@pytest.mark.parametrize("device", CUDA_DEVICES)
@torch.inference_mode()
def test_cp_gather_cache_mla_large_uneven_sequences(
    kv_cache_dtype,
    use_seq_starts,
    device,
):
    block_size = 64
    entry_size = 576
    num_blocks = 1024
    src_cache = _create_mla_cache(
        num_blocks,
        block_size,
        entry_size,
        torch.bfloat16,
        kv_cache_dtype,
        device,
    )
    _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype)

    starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device)
    seq_starts = starts if use_seq_starts else None
    seq_lens = torch.tensor([17, 32_768, 71], dtype=torch.int32, device=device)
    batch_size = seq_lens.shape[0]
    cu_seq_lens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
    cu_seq_lens[1:] = seq_lens.cumsum(dim=0)
    block_table = torch.stack(
        [torch.randperm(num_blocks, device=device) for _ in range(batch_size)]
    ).to(torch.int32)

    expected_batches = []
    for req_id in range(batch_size):
        start = starts[req_id].item() if use_seq_starts else 0
        source_tokens = torch.arange(
            start, start + seq_lens[req_id].item(), device=device
        )
        physical_blocks = block_table[req_id, source_tokens // block_size].long()
        expected_batches.append(src_cache[physical_blocks, source_tokens % block_size])
    expected = torch.cat(expected_batches)
    dst = torch.empty_like(expected)

    ops.cp_gather_cache(
        src_cache,
        dst,
        block_table,
        cu_seq_lens,
        batch_size,
        seq_starts,
    )
    torch.testing.assert_close(dst, expected)


@pytest.mark.parametrize("kv_lora_rank", KV_LORA_RANKS)
@pytest.mark.parametrize("qk_rope_head_dim", QK_ROPE_HEAD_DIMS)
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_MLA)
@pytest.mark.parametrize("block_size", BLOCK_SIZES_MLA)
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS_MLA)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.cpu_model
@pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
@torch.inference_mode()
def test_concat_and_cache_mla_cpu(
    kv_lora_rank: int,
    qk_rope_head_dim: int,
    num_tokens: int,
    block_size: int,
    num_blocks: int,
    dtype: torch.dtype,
    seed: int,
) -> None:
    device = "cpu"
    kv_cache_dtype = "auto"
    set_random_seed(seed)
    torch.set_default_device(device)

    total_slots = num_blocks * block_size
    slot_mapping_lst = random.sample(range(total_slots), num_tokens)
    slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)

    kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype, device=device)
    k_pe = torch.randn(num_tokens, qk_rope_head_dim, dtype=dtype, device=device)
    entry_size = kv_lora_rank + qk_rope_head_dim

    scale = torch.tensor(0.1, dtype=torch.float32, device=device)
    kv_cache = _create_mla_cache(
        num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device
    )
    ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)

    for i in range(num_tokens):
        slot = slot_mapping[i].item()
        block_idx = slot // block_size
        block_offset = slot % block_size
        ref_temp[block_idx, block_offset, :kv_lora_rank] = kv_c[i]
        ref_temp[block_idx, block_offset, kv_lora_rank:] = k_pe[i]

    if kv_cache_dtype == "fp8":
        ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
        ops.convert_fp8(ref_kv_cache, ref_temp, scale.item(), kv_dtype=kv_cache_dtype)
    else:
        ref_kv_cache = ref_temp

    opcheck(
        torch.ops._C_cache_ops.concat_and_cache_mla,
        (kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale),
        test_utils=DEFAULT_OPCHECK_TEST_UTILS,
    )

    ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, kv_cache_dtype, scale)
    torch.testing.assert_close(kv_cache, ref_kv_cache)