"""
Copyright (c) 2023 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import pytest
import torch
from tests.test_helpers.jit_utils import (
    gen_decode_attention_modules,
    gen_prefill_attention_modules,
)
from tests.test_helpers.utils_fp4 import create_nvfp4_kv, nvfp4_to_float
from functools import partial
import flashinfer
from flashinfer.utils import get_compute_capability, has_flashinfer_jit_cache


def head_dim_512_supported() -> bool:
    # 16-bit FA2 head_dim > 256 uses the Ampere+ large-head path.
    return get_compute_capability(torch.device("cuda:0"))[0] >= 8


def skip_if_head_dim_unsupported(head_dim: int):
    if head_dim > 256 and not head_dim_512_supported():
        pytest.skip("16-bit FA2 head_dim > 256 is only supported on SM80 or newer")


def skip_if_head_dim_dtype_unsupported(head_dim: int, kv_dtype: torch.dtype):
    skip_if_head_dim_unsupported(head_dim)
    if (
        head_dim > 256
        and kv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
        and get_compute_capability(torch.device("cuda:0"))[0] < 10
    ):
        pytest.skip("head_dim > 256 with FP8 KV is only validated on SM100 or newer")


def skip_if_nvfp4_large_head_decode_unsupported(head_dim: int):
    if head_dim > 256 and get_compute_capability(torch.device("cuda:0"))[0] < 10:
        pytest.skip(
            "head_dim > 256 with NVFP4 KV decode is only validated on SM100 or newer"
        )


def skip_if_nvfp4_kv_unsupported():
    # Mirrors the guard in flashinfer.decode/prefill: no NVFP4-KV decode
    # kernels exist for SM107.
    if get_compute_capability(torch.device("cuda:0")) == (10, 7):
        pytest.skip("KV Cache NVFP4 is not supported on SM107")


@pytest.fixture(
    autouse=not has_flashinfer_jit_cache(),
    scope="module",
)
def warmup_jit():
    flashinfer.jit.build_jit_specs(
        gen_decode_attention_modules(
            [torch.float16],  # q_dtypes
            [
                torch.float16,
                torch.float8_e4m3fn,
            ],  # kv_dtypes
            [128, 256],  # head_dims
            [0, 1],  # pos_encoding_modes
            [False],  # use_sliding_windows
            [False],  # use_logits_soft_caps
        )
        + gen_prefill_attention_modules(
            [torch.float16],  # q_dtypes
            [
                torch.float16,
                torch.float8_e4m3fn,
            ],  # kv_dtypes
            [128, 256],  # head_dims
            [0, 1],  # pos_encoding_modes
            [False],  # use_sliding_windows
            [False],  # use_logits_soft_caps
            [False],  # use_fp16_qk_reductions
        ),
        verbose=False,
    )
    yield


@pytest.mark.parametrize("batch_size", [12, 17, 128])
@pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256, 512])
@pytest.mark.parametrize("kv_layout", ["NHD"])
@pytest.mark.parametrize("pos_encoding_mode", ["NONE", "ROPE_LLAMA"])
@pytest.mark.parametrize("logits_soft_cap", [0.0])
@pytest.mark.parametrize("return_lse", [True])
@pytest.mark.parametrize("q_dtype", [torch.float16])
@pytest.mark.parametrize("kv_dtype", [torch.float16, torch.float8_e4m3fn])
@pytest.mark.parametrize("contiguous_kv", [True])
def test_batch_decode_with_paged_kv_cache(
    batch_size,
    kv_len,
    page_size,
    num_kv_heads,
    num_qo_heads,
    head_dim,
    kv_layout,
    pos_encoding_mode,
    logits_soft_cap,
    return_lse,
    q_dtype,
    kv_dtype,
    contiguous_kv,
):
    skip_if_head_dim_dtype_unsupported(head_dim, kv_dtype)
    q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
    num_pages_per_seq = (kv_len + page_size - 1) // page_size
    total_num_pages = num_pages_per_seq * batch_size
    if kv_layout == "HND":
        kv_shape = [total_num_pages, 2, num_kv_heads, page_size, head_dim]
    else:
        kv_shape = [total_num_pages, 2, page_size, num_kv_heads, head_dim]
    if not contiguous_kv:
        tmp = [kv_shape[0]]
        for v in kv_shape[1:]:
            tmp.append(2)
            tmp.append(v)
        kv_shape = tmp
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
        kv_data = kv_data[:, 1, :, 1, :, 1, :, 1, :]
        kv_data_fp32 = kv_data_fp32[:, 1, :, 1, :, 1, :, 1, :]
        # actual data is stored in non-contiguous memory
        assert (
            kv_data.stride(-4)
            != kv_data.shape[-3] * kv_data.shape[-2] * kv_data.shape[-1]
        )
    else:
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
    kv_indptr = (
        torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32)
        * num_pages_per_seq
    )
    kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32)
    kv_last_page_len = torch.full(
        (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0"
    )

    workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda:0")
    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer, kv_layout
    )
    wrapper.plan(
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        logits_soft_cap=logits_soft_cap,
        pos_encoding_mode=pos_encoding_mode,
        data_type=kv_dtype,
        q_data_type=q_dtype,
    )
    if return_lse:
        o, _ = wrapper.run(q, kv_data, return_lse=True)
    else:
        o = wrapper.run(q, kv_data)

    for i in range(batch_size):
        perm_dims = [0, 2, 1, 3] if kv_layout == "HND" else [0, 1, 2, 3]
        perm_dims_last = [1, 0, 2] if kv_layout == "HND" else [0, 1, 2]
        qi = q[i]
        ki = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 0]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 0, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 0, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        vi = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 1]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 1, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 1, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
            qi,
            ki,
            vi,
            pos_encoding_mode=pos_encoding_mode,
            logits_soft_cap=logits_soft_cap,
        )
        torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3)

    # test user-allocated output
    o_buffer = torch.empty_like(o)
    wrapper.run(q, kv_data, out=o_buffer)
    torch.testing.assert_close(o, o_buffer, rtol=1e-3, atol=1e-3)


global_override_indptr_cpu = None
MAX_BATCH_SIZE = 128


@pytest.mark.parametrize("batch_size", [12, 17, 128])
@pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256])
@pytest.mark.parametrize("kv_layout", ["NHD"])
@pytest.mark.parametrize("pos_encoding_mode", ["NONE", "ROPE_LLAMA"])
@pytest.mark.parametrize("logits_soft_cap", [0.0])
@pytest.mark.parametrize("return_lse", [True])
@pytest.mark.parametrize("q_dtype", [torch.float16])
@pytest.mark.parametrize("kv_dtype", [torch.float16, torch.float8_e4m3fn])
@pytest.mark.parametrize("contiguous_kv", [True])
def test_batch_decode_with_paged_kv_cache_with_fast_plan(
    batch_size,
    kv_len,
    page_size,
    num_kv_heads,
    num_qo_heads,
    head_dim,
    kv_layout,
    pos_encoding_mode,
    logits_soft_cap,
    return_lse,
    q_dtype,
    kv_dtype,
    contiguous_kv,
):
    q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
    num_pages_per_seq = (kv_len + page_size - 1) // page_size
    total_num_pages = num_pages_per_seq * batch_size

    global global_override_indptr_cpu
    if global_override_indptr_cpu is None:
        global_override_indptr_cpu = torch.empty(MAX_BATCH_SIZE + 1, device="cpu")
    if global_override_indptr_cpu is not None:
        global_override_indptr_cpu = (
            torch.arange(0, batch_size + 1, device="cpu", dtype=torch.int32)
            * num_pages_per_seq
        )

    if kv_layout == "HND":
        kv_shape = [total_num_pages, 2, num_kv_heads, page_size, head_dim]
    else:
        kv_shape = [total_num_pages, 2, page_size, num_kv_heads, head_dim]
    if not contiguous_kv:
        tmp = [kv_shape[0]]
        for v in kv_shape[1:]:
            tmp.append(2)
            tmp.append(v)
        kv_shape = tmp
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
        kv_data = kv_data[:, 1, :, 1, :, 1, :, 1, :]
        kv_data_fp32 = kv_data_fp32[:, 1, :, 1, :, 1, :, 1, :]
        # actual data is stored in non-contiguous memory
        assert (
            kv_data.stride(-4)
            != kv_data.shape[-3] * kv_data.shape[-2] * kv_data.shape[-1]
        )
    else:
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
    kv_indptr = (
        torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32)
        * num_pages_per_seq
    )
    kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32)
    kv_last_page_len = torch.full(
        (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0"
    )

    workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda:0")
    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer, kv_layout
    )
    wrapper.plan(
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        logits_soft_cap=logits_soft_cap,
        pos_encoding_mode=pos_encoding_mode,
        data_type=kv_dtype,
        q_data_type=q_dtype,
    )
    wrapper.plan = partial(flashinfer.fast_decode_plan, wrapper)
    wrapper.plan(
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        logits_soft_cap=logits_soft_cap,
        pos_encoding_mode=pos_encoding_mode,
        data_type=kv_dtype,
        q_data_type=q_dtype,
        non_blocking=True,
        global_override_indptr_cpu=global_override_indptr_cpu,
    )
    if return_lse:
        o, _ = wrapper.run(q, kv_data, return_lse=True)
    else:
        o = wrapper.run(q, kv_data)

    for i in range(batch_size):
        perm_dims = [0, 2, 1, 3] if kv_layout == "HND" else [0, 1, 2, 3]
        perm_dims_last = [1, 0, 2] if kv_layout == "HND" else [0, 1, 2]
        qi = q[i]
        ki = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 0]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 0, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 0, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        vi = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 1]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 1, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 1, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
            qi,
            ki,
            vi,
            pos_encoding_mode=pos_encoding_mode,
            logits_soft_cap=logits_soft_cap,
        )
        torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3)

    # test user-allocated output
    o_buffer = torch.empty_like(o)
    wrapper.run(q, kv_data, out=o_buffer)
    torch.testing.assert_close(o, o_buffer, rtol=1e-3, atol=1e-3)


@pytest.mark.parametrize("batch_size", [12, 17, 128])
@pytest.mark.parametrize("kv_len", [54, 97, 512, 2048, 16384])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256])
@pytest.mark.parametrize("kv_layout", ["NHD"])
@pytest.mark.parametrize("pos_encoding_mode", ["NONE", "ROPE_LLAMA"])
@pytest.mark.parametrize("logits_soft_cap", [0.0])
@pytest.mark.parametrize("return_lse", [True])
@pytest.mark.parametrize("q_dtype", [torch.float16])
@pytest.mark.parametrize("kv_dtype", [torch.float16, torch.float8_e4m3fn])
@pytest.mark.parametrize("contiguous_kv", [True])
def test_batch_decode_with_tuple_paged_kv_cache(
    batch_size,
    kv_len,
    page_size,
    num_kv_heads,
    num_qo_heads,
    head_dim,
    kv_layout,
    pos_encoding_mode,
    logits_soft_cap,
    return_lse,
    q_dtype,
    kv_dtype,
    contiguous_kv,
):
    q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
    num_pages_per_seq = (kv_len + page_size - 1) // page_size
    total_num_pages = num_pages_per_seq * batch_size
    if kv_layout == "HND":
        kv_shape = [total_num_pages, num_kv_heads, page_size, head_dim]
    else:
        kv_shape = [total_num_pages, page_size, num_kv_heads, head_dim]
    if not contiguous_kv:
        tmp = [kv_shape[0]]
        for v in kv_shape[1:]:
            tmp.append(2)
            tmp.append(v)
        kv_shape = tmp
        kv_data_fp32 = [
            torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
            for _ in range(2)
        ]
        kv_data = [kv_data_fp32[i].to(kv_dtype) for i in range(2)]
        for i in range(2):
            kv_data_fp32[i] = kv_data_fp32[i][:, 1, :, 1, :, 1, :]
            kv_data[i] = kv_data[i][:, 1, :, 1, :, 1, :]
            # actual data is stored in non-contiguous memory
            assert (
                kv_data[i].stride(-4)
                != kv_data[i].shape[-3] * kv_data[i].shape[-2] * kv_data[i].shape[-1]
            )
    else:
        kv_data_fp32 = [
            torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
            for _ in range(2)
        ]
        kv_data = [kv_data_fp32[i].to(kv_dtype) for i in range(2)]
    kv_data = tuple(kv_data)
    kv_indptr = (
        torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32)
        * num_pages_per_seq
    )
    kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32)
    kv_last_page_len = torch.full(
        (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0"
    )

    workspace_buffer = torch.empty(32 * 1024 * 1024, dtype=torch.int8, device="cuda:0")
    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer, kv_layout
    )
    wrapper.plan(
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        logits_soft_cap=logits_soft_cap,
        pos_encoding_mode=pos_encoding_mode,
        data_type=kv_dtype,
        q_data_type=q_dtype,
    )
    if return_lse:
        o, _ = wrapper.run(q, kv_data, return_lse=True)
    else:
        o = wrapper.run(q, kv_data)

    k_cache, v_cache = kv_data_fp32
    for i in range(batch_size):
        perm_dims = [0, 2, 1, 3] if kv_layout == "HND" else [0, 1, 2, 3]
        perm_dims_last = [1, 0, 2] if kv_layout == "HND" else [0, 1, 2]
        qi = q[i]
        ki = torch.cat(
            [
                k_cache[kv_indptr[i] : kv_indptr[i + 1] - 1]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    k_cache[kv_indptr[i + 1] - 1, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else k_cache[kv_indptr[i + 1] - 1, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        vi = torch.cat(
            [
                v_cache[kv_indptr[i] : kv_indptr[i + 1] - 1]
                .to(torch.float32)  # torch.cat does not support some fp8 types
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    v_cache[kv_indptr[i + 1] - 1, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else v_cache[kv_indptr[i + 1] - 1, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
            qi,
            ki,
            vi,
            pos_encoding_mode=pos_encoding_mode,
            logits_soft_cap=logits_soft_cap,
        )
        torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3)


@pytest.mark.parametrize("batch_size", [12, 17, 128])
@pytest.mark.parametrize("kv_len", [54, 2048, 16384])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256])
@pytest.mark.parametrize("kv_layout", ["NHD"])
@pytest.mark.parametrize("pos_encoding_mode", ["NONE", "ROPE_LLAMA"])
@pytest.mark.parametrize("q_dtype", [torch.float16])
@pytest.mark.parametrize("kv_dtype", [torch.float16, torch.float8_e4m3fn])
@pytest.mark.parametrize("contiguous_kv", [True])
def test_cuda_graph_batch_decode_with_paged_kv_cache(
    batch_size,
    kv_len,
    page_size,
    num_kv_heads,
    num_qo_heads,
    head_dim,
    kv_layout,
    pos_encoding_mode,
    q_dtype,
    kv_dtype,
    contiguous_kv,
):
    q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)
    num_pages_per_seq = (kv_len + page_size - 1) // page_size
    total_num_pages = num_pages_per_seq * batch_size
    if kv_layout == "HND":
        kv_shape = [total_num_pages, 2, num_kv_heads, page_size, head_dim]
    else:
        kv_shape = [total_num_pages, 2, page_size, num_kv_heads, head_dim]
    if not contiguous_kv:
        tmp = [kv_shape[0]]
        for v in kv_shape[1:]:
            tmp.append(2)
            tmp.append(v)
        kv_shape = tmp
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
        kv_data_fp32 = kv_data_fp32[:, 1, :, 1, :, 1, :, 1, :]
        kv_data = kv_data[:, 1, :, 1, :, 1, :, 1, :]
        # actual data is stored in non-contiguous memory
        assert (
            kv_data.stride(-4)
            != kv_data.shape[-3] * kv_data.shape[-2] * kv_data.shape[-1]
        )
    else:
        kv_data_fp32 = torch.randn(*kv_shape, dtype=torch.float32, device="cuda:0")
        kv_data = kv_data_fp32.to(kv_dtype)
    kv_indptr_host_warmup = torch.arange(
        0, batch_size + 1, device="cuda:0", dtype=torch.int32
    )
    kv_indices_host_warmup = torch.arange(
        0, batch_size, device="cuda:0", dtype=torch.int32
    )
    kv_last_page_len_host_warmup = torch.full(
        (batch_size,), page_size, dtype=torch.int32
    )

    # NOTE(Zihao): allocate more space than needed for testing
    kv_indptr_device_buffer = torch.empty(
        batch_size + 1, device="cuda:0", dtype=torch.int32
    )
    kv_indices_device_buffer = torch.empty(
        total_num_pages, device="cuda:0", dtype=torch.int32
    )
    kv_last_page_device_buffer = torch.empty(
        batch_size, device="cuda:0", dtype=torch.int32
    )

    workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda:0")
    wrapper = flashinfer.decode.CUDAGraphBatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer,
        kv_indptr_device_buffer,
        kv_indices_device_buffer,
        kv_last_page_device_buffer,
        kv_layout,
    )
    wrapper.plan(
        kv_indptr_host_warmup,
        kv_indices_host_warmup,
        kv_last_page_len_host_warmup,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        data_type=kv_dtype,
        pos_encoding_mode=pos_encoding_mode,
        q_data_type=q_dtype,
    )
    # warmup
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(3):
            o = wrapper.run(q, kv_data)
    torch.cuda.current_stream().wait_stream(s)

    # capture
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        o = wrapper.run(q, kv_data)

    # replay multiple times
    for i in range(1, min(4, num_pages_per_seq)):
        kv_indptr_host = torch.arange(0, batch_size + 1).int() * i
        kv_indices_host = torch.arange(0, i * batch_size).int()
        kv_last_page_len_host = torch.full((batch_size,), page_size, dtype=torch.int32)

        wrapper.plan(
            kv_indptr_host,
            kv_indices_host,
            kv_last_page_len_host,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            data_type=kv_dtype,
            pos_encoding_mode=pos_encoding_mode,
            q_data_type=q_dtype,
        )
        g.replay()

    # replay again
    kv_indptr_host = torch.arange(0, batch_size + 1).int() * num_pages_per_seq
    kv_indices_host = torch.arange(0, total_num_pages).int()
    kv_last_page_len_host = torch.full(
        (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32
    )

    wrapper.plan(
        kv_indptr_host,
        kv_indices_host,
        kv_last_page_len_host,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        data_type=kv_dtype,
        pos_encoding_mode=pos_encoding_mode,
        q_data_type=q_dtype,
    )
    g.replay()

    # compute ground truth and compare
    kv_indptr = kv_indptr_host.to(0)
    kv_last_page_len = kv_last_page_len_host.to(0)
    for i in range(batch_size):
        perm_dims = [0, 2, 1, 3] if kv_layout == "HND" else [0, 1, 2, 3]
        perm_dims_last = [1, 0, 2] if kv_layout == "HND" else [0, 1, 2]
        qi = q[i]
        ki = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 0]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 0, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 0, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        vi = torch.cat(
            [
                kv_data_fp32[kv_indptr[i] : kv_indptr[i + 1] - 1, 1]
                .permute(*perm_dims)
                .reshape(-1, num_kv_heads, head_dim),
                (
                    kv_data_fp32[kv_indptr[i + 1] - 1, 1, :, : kv_last_page_len[i]]
                    if kv_layout == "HND"
                    else kv_data_fp32[kv_indptr[i + 1] - 1, 1, : kv_last_page_len[i], :]
                )
                .permute(*perm_dims_last)
                .reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        ).to(kv_dtype)
        o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
            qi, ki, vi, pos_encoding_mode=pos_encoding_mode
        )
        torch.testing.assert_close(o[i], o_ref_i, rtol=1e-3, atol=1e-3)


@pytest.mark.parametrize("batch_size", [12, 17, 128])
@pytest.mark.parametrize("kv_len", [54, 97, 512, 2048])
@pytest.mark.parametrize("page_size", [1, 8, 16])
@pytest.mark.parametrize("num_kv_heads", [4])
@pytest.mark.parametrize("num_qo_heads", [4, 32])
@pytest.mark.parametrize("head_dim", [128, 256])
@pytest.mark.parametrize("q_dtype", [torch.float16, torch.bfloat16])
def test_batch_decode_with_paged_kv_cache_nvfp4(
    batch_size,
    kv_len,
    page_size,
    num_kv_heads,
    num_qo_heads,
    head_dim,
    q_dtype,
):
    """Test BatchDecodeWithPagedKVCacheWrapper with NVFP4 KV cache.

    KV cache layout (NHD):
      kv_cache:    [num_pages, 2, page_size, num_kv_heads, head_dim//2]   uint8 (packed FP4x2)
      kv_cache_sf: [num_pages, 2, page_size, num_kv_heads, head_dim//16]  uint8 (FP8 SFs)

    Reference is computed by dequantizing the packed KV back to q_dtype and running
    single_decode_with_kv_cache per batch item.
    """
    skip_if_nvfp4_kv_unsupported()
    kv_layout = "NHD"
    torch.manual_seed(42)

    # --- query: one token per request ---
    q = torch.randn(batch_size, num_qo_heads, head_dim, device="cuda:0", dtype=q_dtype)

    # --- paged KV metadata ---
    num_pages_per_seq = (kv_len + page_size - 1) // page_size
    total_num_pages = num_pages_per_seq * batch_size
    kv_indptr = (
        torch.arange(0, batch_size + 1, device="cuda:0", dtype=torch.int32)
        * num_pages_per_seq
    )
    kv_indices = torch.arange(0, total_num_pages, device="cuda:0", dtype=torch.int32)
    kv_last_page_len = torch.full(
        (batch_size,), (kv_len - 1) % page_size + 1, dtype=torch.int32, device="cuda:0"
    )

    # --- create NVFP4 KV pages directly (NHD: [num_pages, page_size, num_kv_heads, head_dim//2]) ---
    kv_shape = (total_num_pages, page_size, num_kv_heads, head_dim // 2)
    k_packed, k_sf, k_global_scale = create_nvfp4_kv(kv_shape, "cuda:0")
    v_packed, v_sf, v_global_scale = create_nvfp4_kv(kv_shape, "cuda:0")

    # Dequantize for reference attention
    k_dq = nvfp4_to_float(k_packed, k_sf, k_global_scale).to(q_dtype)
    v_dq = nvfp4_to_float(v_packed, v_sf, v_global_scale).to(q_dtype)

    # Pack into combined tensors:
    #   kv_cache:    [num_pages, 2, page_size, num_kv_heads, head_dim//2]
    #   kv_cache_sf: [num_pages, 2, page_size, num_kv_heads, head_dim//16]
    kv_cache = torch.stack([k_packed, v_packed], dim=1)  # [P, 2, S, H, D//2]
    kv_cache_sf = torch.stack([k_sf, v_sf], dim=1)  # [P, 2, S, H, D//16]

    # --- run BatchDecodeWithPagedKVCacheWrapper ---
    # NVFP4 KV is only supported via the tensor-cores path (FA2 prefill kernel),
    # not the legacy FA2 decode kernel.
    workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda:0")

    # This actually initialize the smem buffers to 0x7F (NaN as FP8), to trigger the SF addressing issue.
    workspace_buffer.fill_(0x7F)

    wrapper = flashinfer.decode.BatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer, kv_layout, use_tensor_cores=True
    )
    wrapper.plan(
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        pos_encoding_mode="NONE",
        logits_soft_cap=0.0,
        kv_data_type=torch.uint8,
        q_data_type=q_dtype,
    )
    o = wrapper.run(
        q,
        kv_cache,
        k_scale=k_global_scale.item(),
        v_scale=v_global_scale.item(),
        kv_cache_sf=kv_cache_sf,
    )

    # --- reference: single_decode_with_kv_cache per batch item using dequantized KV ---
    kv_indptr_cpu = kv_indptr.cpu()
    kv_last_page_len_cpu = kv_last_page_len.cpu()
    for i in range(batch_size):
        qi = q[i]

        # Gather full (non-padded) KV for sequence i from pages
        full_pages_k = k_dq[
            kv_indptr_cpu[i] : kv_indptr_cpu[i + 1] - 1
        ]  # [p-1, S, H, D]
        last_page_k = k_dq[
            kv_indptr_cpu[i + 1] - 1, : kv_last_page_len_cpu[i]
        ]  # [l, H, D]
        ki = torch.cat(
            [
                full_pages_k.reshape(-1, num_kv_heads, head_dim),
                last_page_k.reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        )

        full_pages_v = v_dq[kv_indptr_cpu[i] : kv_indptr_cpu[i + 1] - 1]
        last_page_v = v_dq[kv_indptr_cpu[i + 1] - 1, : kv_last_page_len_cpu[i]]
        vi = torch.cat(
            [
                full_pages_v.reshape(-1, num_kv_heads, head_dim),
                last_page_v.reshape(-1, num_kv_heads, head_dim),
            ],
            dim=0,
        )

        o_ref_i = flashinfer.decode.single_decode_with_kv_cache(
            qi, ki, vi, pos_encoding_mode="NONE", logits_soft_cap=0.0
        )

        # NVFP4 is 4-bit; use relaxed tolerance
        torch.testing.assert_close(o[i], o_ref_i, rtol=1e-1, atol=1e-1)


def test_batch_decode_with_paged_kv_cache_nvfp4_large_head():
    skip_if_nvfp4_large_head_decode_unsupported(512)
    test_batch_decode_with_paged_kv_cache_nvfp4(
        batch_size=4,
        kv_len=128,
        page_size=16,
        num_kv_heads=1,
        num_qo_heads=1,
        head_dim=512,
        q_dtype=torch.float16,
    )


if __name__ == "__main__":
    test_batch_decode_with_paged_kv_cache(
        256,
        54,
        8,
        8,
        8,
        128,
        "NHD",
        "NONE",
        0.0,
        False,
        torch.float16,
        torch.float16,
        True,
    )
    test_batch_decode_with_tuple_paged_kv_cache(
        256,
        54,
        8,
        8,
        8,
        128,
        "NHD",
        "NONE",
        0.0,
        False,
        torch.float16,
        torch.float16,
        True,
    )
    test_batch_decode_with_paged_kv_cache(
        12,
        2048,
        8,
        8,
        8,
        128,
        "NHD",
        "NONE",
        0.0,
        False,
        torch.float16,
        torch.float16,
        True,
    )
    test_batch_decode_with_paged_kv_cache(
        12,
        54,
        1,
        8,
        8,
        128,
        "HND",
        "NONE",
        0.0,
        True,
        torch.float16,
        torch.float8_e5m2,
        True,
    )
    test_cuda_graph_batch_decode_with_paged_kv_cache(
        12, 2048, 8, 8, 8, 128, "NHD", "NONE", torch.float16, torch.float16, True
    )
    test_cuda_graph_batch_decode_with_paged_kv_cache(
        128, 54, 8, 8, 8, 128, "NHD", "NONE", torch.float16, torch.float16, True
    )
    test_batch_decode_with_paged_kv_cache(
        12,
        54,
        1,
        8,
        8,
        128,
        "HND",
        "NONE",
        0.0,
        True,
        torch.float16,
        torch.float8_e5m2,
        True,
    )
    test_cuda_graph_batch_decode_with_paged_kv_cache(
        12, 54, 8, 8, 8, 128, "HND", "NONE", torch.float16, torch.float8_e5m2, True
    )
    test_batch_decode_with_paged_kv_cache_nvfp4(4, 128, 64, 1, 1, 128, torch.float16)


def test_single_decode_torch_compile_cuda_graph():
    """Issue #541: single_decode_with_kv_cache with torch.compile(mode='reduce-overhead')
    should not raise RuntimeError about unaccounted cudagraph pool pointers.

    Runs as a subprocess because torch.library registration must happen at module
    level for torch.compile compatibility -- registering inside a pytest function
    causes dynamo tracing errors.
    """
    import subprocess
    import sys
    import textwrap

    script = textwrap.dedent("""\
        import torch
        from flashinfer import single_decode_with_kv_cache

        torch.library.define(
            "test541d::op", "(Tensor q, Tensor k, Tensor v) -> Tensor")

        @torch.library.impl("test541d::op", "cuda")
        def _impl(q, k, v):
            return single_decode_with_kv_cache(q, k, v)

        @torch.library.register_fake("test541d::op")
        def _fake(q, k, v):
            return torch.empty_like(q)

        compiled_fn = torch.compile(
            lambda q, k, v: torch.ops.test541d.op(q, k, v),
            mode="reduce-overhead", fullgraph=True)

        KV_LEN, QH, KH, D = 256, 8, 8, 128
        for i in range(3):
            q = torch.randn(QH, D, device="cuda", dtype=torch.float16)
            k = torch.randn(KV_LEN, KH, D, device="cuda", dtype=torch.float16)
            v = torch.randn(KV_LEN, KH, D, device="cuda", dtype=torch.float16)
            o = compiled_fn(q, k, v)
            assert o.shape == (QH, D)
        torch.cuda.synchronize()
        print("PASS")
    """)
    import gc
    import os

    # torch.compile's inductor calls getpass.getuser() for cache dir, which fails
    # in CI containers where the uid has no /etc/passwd entry. Setting USER avoids this.
    env = os.environ.copy()
    env.setdefault("USER", "ci")

    # The parent pytest process has already run thousands of decode cases in this
    # file. Release its cached blocks before the subprocess initializes
    # torch.compile/cudagraph state on memory-constrained A10G runners.
    torch.cuda.synchronize()
    gc.collect()
    torch.cuda.empty_cache()

    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True,
        text=True,
        timeout=300,
        env=env,
    )
    assert result.returncode == 0 and "PASS" in result.stdout, (
        f"Test failed:\nstdout: {result.stdout[-500:]}\nstderr: {result.stderr[-500:]}"
    )


@pytest.mark.parametrize("batch_size", [4, 8])
@pytest.mark.parametrize("q_len_per_req", [2, 3])
@pytest.mark.parametrize("num_kv_heads", [2, 8])
@pytest.mark.parametrize("gqa_group_size", [4, 8])
def test_cuda_graph_uniform_multi_token_decode_with_paged_kv_cache(
    batch_size,
    q_len_per_req,
    num_kv_heads,
    gqa_group_size,
):
    # Spec-decode verify shape: every request carries q_len_per_req tokens.
    # The tensor-core decode path must produce prefill-wrapper-identical
    # results eagerly, and stay correct when the plan is refreshed with
    # different kv lengths between CUDA graph replays.
    page_size = 16
    head_dim = 128
    dtype = torch.bfloat16
    num_qo_heads = num_kv_heads * gqa_group_size
    capture_kv_lens = [129] * batch_size
    # tightest legal kv under the plan contract kv_len >= q_len_per_req
    replay_kv_lens_sets = [
        [[33, 17, 65, 129][i % 4] for i in range(batch_size)],
        [[q_len_per_req, 128, 64, 100][i % 4] for i in range(batch_size)],
    ]

    def build_kv_layout(kv_lens):
        pages_per = [(length + page_size - 1) // page_size for length in kv_lens]
        indptr = torch.tensor(
            [0] + list(torch.cumsum(torch.tensor(pages_per), 0)),
            dtype=torch.int32,
            device="cuda:0",
        )
        indices = torch.arange(int(indptr[-1]), dtype=torch.int32, device="cuda:0")
        last_len = torch.tensor(
            [(length - 1) % page_size + 1 for length in kv_lens],
            dtype=torch.int32,
            device="cuda:0",
        )
        return indptr, indices, last_len

    indptr, indices, last_len = build_kv_layout(capture_kv_lens)
    total_pages = int(indptr[-1])
    kv_data = torch.randn(
        total_pages, 2, page_size, num_kv_heads, head_dim, dtype=dtype, device="cuda:0"
    )
    q = torch.randn(
        batch_size * q_len_per_req, num_qo_heads, head_dim, dtype=dtype, device="cuda:0"
    )

    def reference(ref_indptr, ref_indices, ref_last_len):
        ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
        prefill = flashinfer.BatchPrefillWithPagedKVCacheWrapper(ws, "NHD")
        qo_indptr = (
            torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda:0")
            * q_len_per_req
        )
        prefill.plan(
            qo_indptr,
            ref_indptr,
            ref_indices,
            ref_last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            causal=True,
            q_data_type=dtype,
            kv_data_type=dtype,
        )
        return prefill.run(q, kv_data)

    workspace_buffer = torch.empty(
        128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0"
    )
    wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        workspace_buffer,
        "NHD",
        use_cuda_graph=True,
        use_tensor_cores=True,
        paged_kv_indptr_buffer=torch.empty(
            batch_size + 1, dtype=torch.int32, device="cuda:0"
        ),
        paged_kv_indices_buffer=torch.empty(
            total_pages, dtype=torch.int32, device="cuda:0"
        ),
        paged_kv_last_page_len_buffer=torch.empty(
            batch_size, dtype=torch.int32, device="cuda:0"
        ),
    )

    def plan(plan_indptr, plan_indices, plan_last_len):
        wrapper.plan(
            plan_indptr,
            plan_indices,
            plan_last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=q_len_per_req,
        )

    plan(indptr, indices, last_len)
    s = torch.cuda.Stream()
    s.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(s):
        for _ in range(3):
            o = wrapper.run(q, kv_data)
    torch.cuda.current_stream().wait_stream(s)

    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        o = wrapper.run(q, kv_data, out=o)

    g.replay()
    torch.cuda.synchronize()
    torch.testing.assert_close(
        o, reference(indptr, indices, last_len), rtol=1e-2, atol=1e-2
    )

    for kv_lens in replay_kv_lens_sets:
        new_indptr, new_indices, new_last_len = build_kv_layout(kv_lens)
        plan(new_indptr, new_indices, new_last_len)
        torch.cuda.synchronize()
        g.replay()
        torch.cuda.synchronize()
        torch.testing.assert_close(
            o, reference(new_indptr, new_indices, new_last_len), rtol=1e-2, atol=1e-2
        )


def test_tensor_core_decode_rejects_mismatched_q_len():
    batch_size = 4
    page_size = 16
    num_qo_heads, num_kv_heads, head_dim = 8, 2, 128
    dtype = torch.bfloat16
    indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda:0")
    indices = torch.arange(batch_size, dtype=torch.int32, device="cuda:0")
    last_len = torch.full((batch_size,), page_size, dtype=torch.int32, device="cuda:0")
    kv_data = torch.randn(
        batch_size, 2, page_size, num_kv_heads, head_dim, dtype=dtype, device="cuda:0"
    )
    ws = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0")
    wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws, "NHD", use_tensor_cores=True
    )
    wrapper.plan(
        indptr,
        indices,
        last_len,
        num_qo_heads,
        num_kv_heads,
        head_dim,
        page_size,
        q_data_type=dtype,
        kv_data_type=dtype,
        q_len_per_req=2,
    )
    q = torch.randn(
        batch_size * 3, num_qo_heads, head_dim, dtype=dtype, device="cuda:0"
    )
    with pytest.raises(ValueError, match="q_len_per_req"):
        wrapper.run(q, kv_data)

    cuda_core_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws, "NHD", use_tensor_cores=False
    )
    with pytest.raises(ValueError, match="use_tensor_cores"):
        cuda_core_wrapper.plan(
            indptr,
            indices,
            last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=2,
        )

    fa3_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws, "NHD", use_tensor_cores=True, backend="fa3"
    )
    with pytest.raises(NotImplementedError, match="fa2"):
        fa3_wrapper.plan(
            indptr,
            indices,
            last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=2,
        )

    short_kv_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws, "NHD", use_tensor_cores=True
    )
    short_last_len = torch.ones(batch_size, dtype=torch.int32, device="cuda:0")
    with pytest.raises(ValueError, match="empty KV range"):
        short_kv_wrapper.plan(
            indptr,
            indices,
            short_last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=2,
        )

    graph_wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws,
        "NHD",
        use_cuda_graph=True,
        use_tensor_cores=True,
        paged_kv_indptr_buffer=torch.empty(
            batch_size + 1, dtype=torch.int32, device="cuda:0"
        ),
        paged_kv_indices_buffer=torch.empty(
            batch_size, dtype=torch.int32, device="cuda:0"
        ),
        paged_kv_last_page_len_buffer=torch.empty(
            batch_size, dtype=torch.int32, device="cuda:0"
        ),
    )

    def plan_graph(q_len):
        graph_wrapper.plan(
            indptr,
            indices,
            last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=q_len,
        )

    plan_graph(2)
    plan_graph(2)  # same value re-plans fine
    with pytest.raises(ValueError, match="frozen"):
        plan_graph(3)

    unplanned = flashinfer.BatchDecodeWithPagedKVCacheWrapper(
        ws, "NHD", use_tensor_cores=True
    )
    with pytest.raises(ValueError, match="prior plan"):
        flashinfer.decode.fast_decode_plan(
            unplanned,
            indptr,
            indices,
            last_len,
            num_qo_heads,
            num_kv_heads,
            head_dim,
            page_size,
            q_data_type=dtype,
            kv_data_type=dtype,
            q_len_per_req=2,
        )
