import pytest
import torch
import torch.nn.functional as F
import random

import flashinfer
from flashinfer.mla import (
    MLALayerDimensions,
    deepseek_mla_dimensions,
    supported_mla_layer_dimensions,
    smaller_mla_dimensions,
)
from flashinfer.utils import (
    get_compute_capability,
    get_device_sm_count,
    get_trtllm_gen_multi_ctas_kv_counter_bytes,
)

global_workspace_buffer = None  # can.be empty initialized
global_trtllm_gen_fmha_workspace_buffer = None
workspace_size = 128 * 1024 * 1024

# Guard region we zero past the softmax slab so we can detect OOB writes.
TRTLLM_GEN_WORKSPACE_CHECK_BYTES = 1 * 1024 * 1024


def trtllm_gen_workspace_softmax_end_bytes_decode(
    num_qo_heads: int, batch_size: int, max_q_len: int
) -> int:
    """End offset of the softmax slab in the generation-mode workspace [softmax|scratch].

    The C++ launcher allocates ``sizeof(float2) * softmax_slots`` bytes, i.e. 8 bytes per
    slot, where ``softmax_slots = num_qo_heads * batch_size * round_up(max_q_len, 256)``.
    """
    rounded_max_q_len = ((max_q_len + 255) // 256) * 256
    softmax_slab = (
        8 * num_qo_heads * batch_size * rounded_max_q_len
    )  # sizeof(float2) == 8
    return softmax_slab


def generate_sparse_indices(
    batch_size: int,
    q_len_per_request: int,
    seq_lens: torch.Tensor,
    topk: int,
    page_size: int,
    block_tables: torch.Tensor,
    device: str,
    seed: int = 42,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Generate sparse attention indices for MLA.

    Returns:
        abs_indices: [batch_size, q_len_per_request, topk] - absolute positions in sequence
        indices_in_kvcache: [batch_size, q_len_per_request, topk] - positions in blocked KV cache
    """
    random.seed(seed)
    torch.manual_seed(seed)

    block_tables_cpu = block_tables.cpu()
    seq_lens_cpu = seq_lens.cpu()

    abs_indices = torch.empty(
        batch_size, q_len_per_request, topk, dtype=torch.int32, device="cpu"
    )
    indices_in_kvcache = torch.empty(
        batch_size, q_len_per_request, topk, dtype=torch.int32, device="cpu"
    )

    for i in range(batch_size):
        cur_seq_len = int(seq_lens_cpu[i].item())
        # Generate indices for each query position
        for j in range(q_len_per_request):
            # Randomly sample topk positions from the sequence
            if cur_seq_len > 0:
                # cur_abs_indices = torch.randperm(cur_seq_len, device="cpu")[:topk]
                cur_abs_indices = torch.arange(0, topk, device="cpu")
                # Convert to blocked indices
                cur_blocked_indices = block_tables_cpu[
                    i, cur_abs_indices // page_size
                ] * page_size + (cur_abs_indices % page_size)
            else:
                cur_abs_indices = torch.empty(0, dtype=torch.int32, device="cpu")
                cur_blocked_indices = torch.empty(0, dtype=torch.int32, device="cpu")

            # Pad with -1 if we don't have enough indices
            if len(cur_abs_indices) < topk:
                pad_len = topk - len(cur_abs_indices)
                cur_abs_indices = torch.cat(
                    [
                        cur_abs_indices,
                        torch.full((pad_len,), -1, device="cpu", dtype=torch.int32),
                    ]
                )
                cur_blocked_indices = torch.cat(
                    [
                        cur_blocked_indices,
                        torch.full((pad_len,), -1, device="cpu", dtype=torch.int32),
                    ]
                )

            # Randomly permute the indices
            # perm = torch.randperm(topk, device="cpu")
            perm = torch.arange(0, topk, device="cpu")
            cur_abs_indices = cur_abs_indices[perm]
            cur_blocked_indices = cur_blocked_indices[perm]

            abs_indices[i, j, :] = cur_abs_indices
            indices_in_kvcache[i, j, :] = cur_blocked_indices

    return abs_indices.to(device), indices_in_kvcache.to(device)


def sparse_mla_reference_torch(
    cache_seqlens: torch.Tensor,  # [batch_size]
    block_table: torch.Tensor,  # [batch_size, ?]
    q: torch.Tensor,  # [batch_size, s_q, h_q, d]
    blocked_k: torch.Tensor,  # [?, block_size, d]
    blocked_v: torch.Tensor,  # [?, block_size, dv]
    page_size: int,
    is_causal: bool,
    sm_scale: float,
    indices: torch.Tensor | None = None,  # [batch_size, s_q, topk]
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    A reference implementation in PyTorch for MLA attention.
    Based on FlashMLA's reference implementation.

    Args:
        cache_seqlens: Sequence lengths for each batch [batch_size]
        block_table: Block table mapping [batch_size, max_num_blocks]
        q: Query tensor [batch_size, s_q, h_q, d]
        blocked_k: Blocked key cache [num_blocks, block_size, d]
        blocked_v: Blocked value cache [num_blocks, block_size, dv]
        page_size: Size of each block/page
        is_causal: Whether to apply causal masking
        sm_scale: Softmax scale factor
        indices: Optional sparse indices [batch_size, s_q, topk]

    Returns:
        output: Attention output [batch_size, s_q, h_q, dv]
        lse: Log-sum-exp values [batch_size, h_q, s_q]
    """

    def get_topk_attn_mask(s_q: int, s_k: int, indices: torch.Tensor):
        """Create attention mask for top-k sparse attention."""
        mask = torch.zeros(s_q, s_k, dtype=torch.bool)
        for i in range(s_q):
            cur_indices = indices[i]
            valid_indices = cur_indices[cur_indices != -1]
            mask[i, valid_indices] = True
        return mask

    def scaled_dot_product_attention(
        batch_idx: int,
        query: torch.Tensor,  # [h_q, s_q, d]
        key: torch.Tensor,  # [s_k, d]
        value: torch.Tensor,  # [s_k, dv]
        is_causal: bool,
        sm_scale: float,
        indices: torch.Tensor | None,  # [s_q, topk]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute scaled dot-product attention."""
        h_q = query.size(0)
        s_q = query.shape[-2]
        s_k = key.shape[-2]
        dv = value.shape[-1]

        query = query.float()
        key = key.float()
        value = value.float()

        # Handle NaN values in KV
        key[key != key] = 0.0
        value[value != value] = 0.0

        # Compute attention weights: [h_q, s_q, s_k]
        attn_weight = query @ key.transpose(-2, -1)

        # Apply masking if needed
        if (is_causal and query.size(1) > 1) or indices is not None:
            mask = torch.ones(s_q, s_k, dtype=torch.bool)
            if is_causal:
                mask = mask.tril(diagonal=s_k - s_q)
            if indices is not None:
                mask &= get_topk_attn_mask(s_q, s_k, indices)
            attn_bias = torch.zeros(s_q, s_k, dtype=torch.float, device=query.device)
            mask = mask.to(device=query.device)
            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
            attn_weight += attn_bias.to(query.dtype)

        # Scale and softmax
        attn_weight *= sm_scale
        lse = attn_weight.logsumexp(dim=-1)  # [h_q, s_q]
        attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)

        # Compute output
        output = attn_weight @ value  # [h_q, s_q, dv]

        # Correct for query tokens which have no attendable keys
        lonely_q_mask = lse == float("-inf")
        output[lonely_q_mask.unsqueeze(-1).broadcast_to(h_q, s_q, dv)] = 0.0
        lse[lonely_q_mask] = float("+inf")

        return output, lse

    b, s_q, h_q, d = q.size()
    dv = blocked_v.size(2)
    cache_seqlens_cpu = cache_seqlens.cpu()

    out_ref = torch.empty(b, s_q, h_q, dv, dtype=torch.float32)
    lse_ref = torch.empty(b, h_q, s_q, dtype=torch.float32)

    for i in range(b):
        cur_len = int(cache_seqlens_cpu[i].item())
        cur_num_blocks = (cur_len + page_size - 1) // page_size
        cur_block_indices = block_table[i][0:cur_num_blocks]

        # Gather KV for this sequence
        cur_key = blocked_k[cur_block_indices].view(-1, d)[:cur_len, ...]
        cur_value = blocked_v[cur_block_indices].view(-1, dv)[:cur_len, ...]

        cur_out, cur_lse = scaled_dot_product_attention(
            i,
            q[i].transpose(0, 1),  # [h_q, s_q, d]
            cur_key,  # [s_k, d]
            cur_value,  # [s_k, dv]
            is_causal,
            sm_scale,
            indices[i] if indices is not None else None,
        )
        out_ref[i] = cur_out.transpose(0, 1)
        lse_ref[i] = cur_lse

    out_ref = out_ref.to(torch.bfloat16).to(q.device)
    return out_ref, lse_ref


def torch_reference_mla(
    query,
    kv_cache,
    block_tables,
    seq_lens,
    kv_lora_rank,
    qk_rope_head_dim,
    softmax_scale,
    output_scale,
    page_size,
):
    """PyTorch reference for MLA decode. Returns [B, q_len, H, kv_lora_rank]."""
    B, q_len, H, D_qk = query.shape
    kv_flat = kv_cache.reshape(-1, D_qk)
    c_latent = kv_flat[:, :kv_lora_rank]
    c_rope = kv_flat[:, kv_lora_rank:]
    q_nope = query[..., :kv_lora_rank]
    q_rope = query[..., kv_lora_rank:]

    outputs = []
    for b in range(B):
        seq_len = seq_lens[b].item()
        num_pages = (seq_len + page_size - 1) // page_size
        pages = block_tables[b, :num_pages]
        kv_indices = []
        for p in pages:
            start = p.item() * page_size
            kv_indices.extend(range(start, start + page_size))
        kv_indices = kv_indices[:seq_len]
        kv_idx_t = torch.tensor(kv_indices, device=query.device)

        k_lat = c_latent[kv_idx_t]  # [seq_len, kv_lora_rank]
        k_rope = c_rope[kv_idx_t]  # [seq_len, rope_dim]

        attn_lat = torch.einsum("qhd,kd->qhk", q_nope[b].float(), k_lat.float())
        attn_rope = torch.einsum("qhd,kd->qhk", q_rope[b].float(), k_rope.float())
        attn = (attn_lat + attn_rope) * softmax_scale
        attn = F.softmax(attn, dim=-1)
        out_b = torch.einsum("qhk,kd->qhd", attn, k_lat.float()) * output_scale
        outputs.append(out_b)

    return torch.stack(outputs, dim=0)  # [B, q_len, H, kv_lora_rank]


def trtllm_batch_decode_mla(
    layer_dimensions: MLALayerDimensions,
    batch_size: int,
    scale: float,
    dtype: torch.dtype,
    page_size: int,
    q_len_per_request: int,
    dynamic_scale: bool,
    enable_pdl: bool,
    backend: str,
    MAX_SEQ_LEN: int,
    skips_softmax: bool,
    uses_shared_paged_kv_idx: bool = True,
    use_cum_seq_lens_q: bool = False,
    max_q_len_exceeds_total_q: bool = False,
):
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if backend == "xqa":
        if compute_capability[0] != 12:
            pytest.skip("XQA MLA only supports SM120/SM121 GPUs")
        if q_len_per_request != 1 or dtype != torch.float8_e4m3fn:
            pytest.skip(
                "XQA MLA only supports q_len_per_request == 1 and dtype == torch.float8_e4m3fn"
            )
        if not uses_shared_paged_kv_idx:
            pytest.skip("xqa backend does not support separate KV page indices")
    if backend == "trtllm-gen":
        if compute_capability[0] != 10:
            pytest.skip("TRTLLM-GEN MLA only supports SM100 and SM103 GPUs")
    if backend == "cute-dsl":
        if compute_capability[0] not in (10, 11):
            pytest.skip("cute-dsl MLA requires SM100-SM110 (tcgen05)")
        from flashinfer.cute_dsl.utils import is_cute_dsl_arch_supported

        if not is_cute_dsl_arch_supported(*compute_capability):
            pytest.skip("installed CuTe DSL cannot target this device architecture")
        if dynamic_scale:
            pytest.skip("cute-dsl does not support dynamic_scale")
        if enable_pdl is not None:
            pytest.skip("cute-dsl does not support enable_pdl")
        if skips_softmax:
            pytest.skip("cute-dsl does not support skip_softmax")
        if not uses_shared_paged_kv_idx:
            pytest.skip("cute-dsl does not support separate KV page indices")
    if dynamic_scale and dtype != torch.float8_e4m3fn:
        pytest.skip("Dynamic scale is not supported for non-fp8 dtype")

    if skips_softmax and backend != "trtllm-gen":
        pytest.skip("skips_softmax is only supported for trtllm-gen backend")
    if use_cum_seq_lens_q and backend == "xqa":
        pytest.skip("XQA does not support cum_seq_lens_q")

    torch.manual_seed(42)
    device = "cuda:0"
    qk_head_dim = (
        layer_dimensions.head_dimensions.qk_rope_head_dim
        + layer_dimensions.head_dimensions.kv_lora_rank
    )
    # Initialize tensors
    query = torch.randn(
        batch_size,
        q_len_per_request,
        layer_dimensions.num_heads,
        qk_head_dim,
        device=device,
    ).to(dtype)
    if use_cum_seq_lens_q:
        q_lens = (
            torch.arange(batch_size, device=device, dtype=torch.int32)
            % q_len_per_request
        ) + 1
        if max_q_len_exceeds_total_q:
            q_lens[0] = 0
        q_lens[-1] = q_len_per_request
        cum_seq_lens_q = torch.empty(batch_size + 1, device=device, dtype=torch.int32)
        cum_seq_lens_q[0] = 0
        cum_seq_lens_q[1:] = torch.cumsum(q_lens, dim=0)
        query_input = torch.cat(
            [query[i, : int(q_lens[i].item())] for i in range(batch_size)],
            dim=0,
        )
        # Overestimate the longest segment for CUDA-graph-friendly capacity.
        # Focused tests can also make this exceed the current compact length.
        max_q_len = (
            q_len_per_request + 1
            if max_q_len_exceeds_total_q
            else min(q_len_per_request + 1, query_input.size(0))
        )
    else:
        query_input = query
        cum_seq_lens_q = None
        q_lens = torch.full(
            (batch_size,), q_len_per_request, device=device, dtype=torch.int32
        )
        max_q_len = None

    num_tokens = MAX_SEQ_LEN * batch_size
    num_blocks = (num_tokens + page_size - 1) // page_size

    # Sequence lengths and block tables
    seq_lens = [torch.randint(1, MAX_SEQ_LEN, (1,)).item() for _ in range(batch_size)]
    seq_lens[-1] = MAX_SEQ_LEN
    max_seq_len = max(seq_lens)
    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device)

    blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size
    max_num_blocks_per_seq = blocks_per_seq.max().item()

    # Generate random but unique block IDs for all sequences
    total_blocks_needed = int(blocks_per_seq.sum().item())
    all_block_ids = torch.randperm(
        total_blocks_needed, device=device
    )  # Random permutation

    # Generate unique block IDs for all sequences
    block_id = 0
    block_tables = torch.zeros(
        (batch_size, max_num_blocks_per_seq), dtype=torch.int, device=device
    )

    # Populate block tables and track block assignments
    block_id = 0
    for i in range(batch_size):
        num_blocks_needed = int(blocks_per_seq[i].item())
        block_tables[i, :num_blocks_needed] = all_block_ids[
            block_id : block_id + num_blocks_needed
        ]
        block_id += num_blocks_needed

    # For separate KV page indices, duplicate the page table rows since
    # MLA K and V share the same compressed representation.
    if not uses_shared_paged_kv_idx:
        block_tables_kernel = torch.stack([block_tables, block_tables], dim=1)
    else:
        block_tables_kernel = block_tables

    # Create interleaved KV cache
    # Allocate more than needed blocks, block_id is just enough, to mimick real-world cases
    kv_cache = torch.randn(
        size=(
            num_blocks,
            page_size,
            qk_head_dim,
        ),
        device=device,
    ).to(dtype)
    # (num_blocks, 1, page_size, kv_lora_rank + qk_rope_head_dim)

    # Allocate workspace buffer
    # todo(Yingyi): calculate the actual size of workspace buffer
    global global_workspace_buffer, global_trtllm_gen_fmha_workspace_buffer
    if global_workspace_buffer is None:
        global_workspace_buffer = torch.empty(
            workspace_size, dtype=torch.int8, device=device
        )
    if global_trtllm_gen_fmha_workspace_buffer is None:
        global_trtllm_gen_fmha_workspace_buffer = torch.empty(
            workspace_size, dtype=torch.int8, device=device
        )
    workspace_buffer = global_trtllm_gen_fmha_workspace_buffer
    workspace_buffer_ref = global_workspace_buffer

    # Using a tiny threshold should give the same output as standard attention
    skip_softmax_threshold_scale_factor = 1e-30 if skips_softmax else None

    def maybe_get_lse_guard_end(softmax_end: int) -> int | None:
        # The C++ launcher carves the workspace as
        #     [softmax_slab | guard (1 MB) | scratch],
        # where ``softmax_slab = 8 * num_heads * batch * round_up(max_q_len, 256)``
        # bytes. The LSE guard subcheck is only meaningful when
        # ``softmax_end + 1 MB`` still fits in the test workspace. For the largest parametrized shapes
        # (e.g. batch=1024 * num_heads=128 * q_len=2 -> softmax_slab > 256 MB)
        # the slab alone overruns the 128 MB workspace; in that case we skip
        # the LSE subcheck and only exercise the output path (which does not
        # allocate softmax_slab when ``lse=None``). TODO: remove this branch
        # once trtllm-gen workspace sizing is computed dynamically per-call.
        guard_end = softmax_end + TRTLLM_GEN_WORKSPACE_CHECK_BYTES
        if guard_end > workspace_size:
            return None
        return guard_end

    # Only the trtllm-gen MLA path supports LSE output; other backends raise NotImplementedError.
    check_lse = (
        backend == "trtllm-gen"
        and not skips_softmax
        and not use_cum_seq_lens_q
        and dtype != torch.float8_e4m3fn
    )
    softmax_end = None
    guard_end = None
    if check_lse:
        softmax_end = trtllm_gen_workspace_softmax_end_bytes_decode(
            layer_dimensions.num_heads, batch_size, q_len_per_request
        )
        guard_end = maybe_get_lse_guard_end(softmax_end)
        if guard_end is None:
            # Preserve output coverage on shapes too large for the LSE guard
            # subcheck (see ``maybe_get_lse_guard_end`` for the size analysis).
            check_lse = False

    if check_lse:
        provided_lse = torch.full(
            (batch_size * q_len_per_request, layer_dimensions.num_heads),
            float("nan"),
            device=device,
            dtype=torch.float32,
        )
        assert softmax_end is not None
        assert guard_end is not None
        workspace_buffer[softmax_end:guard_end].zero_()
    else:
        provided_lse = None

    # Run decode-MLA
    output_and_lse = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
        query=query_input,
        kv_cache=kv_cache.unsqueeze(1),
        workspace_buffer=workspace_buffer,
        qk_nope_head_dim=layer_dimensions.head_dimensions.qk_nope_head_dim,
        kv_lora_rank=layer_dimensions.head_dimensions.kv_lora_rank,
        qk_rope_head_dim=layer_dimensions.head_dimensions.qk_rope_head_dim,
        block_tables=block_tables_kernel,
        seq_lens=seq_lens_tensor,
        max_seq_len=max_seq_len,
        bmm1_scale=scale / ((128 + 64) ** 0.5),
        bmm2_scale=1.0,
        skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale_factor,
        enable_pdl=enable_pdl,
        backend=backend,
        uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
        lse=provided_lse,
        return_lse=check_lse,
        cum_seq_lens_q=cum_seq_lens_q,
        max_q_len=max_q_len,
    )
    if check_lse:
        output, lse_out = output_and_lse
        assert lse_out is provided_lse
        assert lse_out.dtype == torch.float32
        assert lse_out.shape == (
            batch_size * q_len_per_request,
            layer_dimensions.num_heads,
        )
        assert torch.isfinite(lse_out).all(), (
            "trtllm-gen MLA decode produced non-finite LSE"
        )
        assert softmax_end is not None
        assert guard_end is not None
        assert (workspace_buffer[softmax_end:guard_end].cpu().numpy() == 0).all(), (
            "trtllm-gen MLA decode wrote past the softmax slab"
        )
    else:
        output = output_and_lse

    # Run reference attention and align output
    sm_scale = scale / (
        (128 + 64) ** 0.5
    )  # use head dimension before matrix absorption
    wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
        workspace_buffer_ref,
        backend="fa2",
    )

    if dtype == torch.float8_e4m3fn:
        # convert query and kv_cache to bfloat16
        query = query.to(torch.bfloat16)
        kv_cache = kv_cache.to(torch.bfloat16)

    q_ref = query_input if use_cum_seq_lens_q else query
    if dtype == torch.float8_e4m3fn and use_cum_seq_lens_q:
        q_ref = q_ref.to(torch.bfloat16)

    q_indptr = torch.empty(batch_size + 1, device=device, dtype=torch.int32)
    q_indptr[0] = 0
    q_indptr[1:] = torch.cumsum(q_lens, dim=0)
    kv_indptr = torch.zeros_like(q_indptr)
    kv_indptr[1:] = torch.cumsum(blocks_per_seq, dim=0)
    kv_indices = all_block_ids.int()

    wrapper.plan(
        q_indptr,
        kv_indptr,
        kv_indices,
        seq_lens_tensor,
        layer_dimensions.num_heads,
        layer_dimensions.head_dimensions.kv_lora_rank,
        layer_dimensions.head_dimensions.qk_rope_head_dim,
        page_size,
        True,
        sm_scale,
        q_ref.dtype,
        kv_cache.dtype,
    )
    q_nope = q_ref[..., : layer_dimensions.head_dimensions.kv_lora_rank].reshape(
        -1,
        layer_dimensions.num_heads,
        layer_dimensions.head_dimensions.kv_lora_rank,
    )
    q_pe = q_ref[..., layer_dimensions.head_dimensions.kv_lora_rank :].reshape(
        -1,
        layer_dimensions.num_heads,
        layer_dimensions.head_dimensions.qk_rope_head_dim,
    )

    # todo: fix kv_cache
    ckv = kv_cache[..., : layer_dimensions.head_dimensions.kv_lora_rank]
    kpe = kv_cache[..., layer_dimensions.head_dimensions.kv_lora_rank :]

    o_ref = wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False)

    # cute-dsl fp8 kernel outputs fp8; cast to bf16 to match trtllm-gen / reference
    if backend == "cute-dsl" and output.dtype == torch.float8_e4m3fn:
        output = output.to(torch.bfloat16)

    if backend in ("trtllm-gen", "cute-dsl", "auto"):
        # check is nan
        assert not torch.isnan(o_ref).any(), "o_ref is nan"
        assert not torch.isnan(output).any(), "output is nan"

        if use_cum_seq_lens_q:
            output_view = output
            o_ref_view = o_ref
        else:
            output_view = output.reshape(
                batch_size, q_len_per_request, layer_dimensions.num_heads, -1
            )
            o_ref_view = o_ref.view(
                batch_size, q_len_per_request, layer_dimensions.num_heads, -1
            )

        if dtype == torch.float8_e4m3fn:
            rtol, atol = 1e-1, 1e-1
        else:
            rtol, atol = 1e-2, 1e-2

        try:
            torch.testing.assert_close(output_view, o_ref_view, rtol=rtol, atol=atol)
        except AssertionError as fa2_err:
            if backend == "cute-dsl":
                # fa2 reference may diverge from cute-dsl in some configs;
                # fall back to torch reference as ground truth.
                query_for_ref = (
                    query.to(torch.bfloat16) if dtype == torch.float8_e4m3fn else query
                )
                kv_for_ref = (
                    kv_cache.to(torch.bfloat16)
                    if dtype == torch.float8_e4m3fn
                    else kv_cache
                )
                o_torch_ref = torch_reference_mla(
                    query_for_ref,
                    kv_for_ref,
                    block_tables,
                    seq_lens_tensor,
                    layer_dimensions.head_dimensions.kv_lora_rank,
                    layer_dimensions.head_dimensions.qk_rope_head_dim,
                    softmax_scale=sm_scale,
                    output_scale=1.0,
                    page_size=page_size,
                ).to(output.dtype)
                torch.testing.assert_close(output, o_torch_ref, rtol=rtol, atol=atol)
            else:
                print("output:", output)
                print("o_ref:", o_ref)
                raise fa2_err
    elif backend == "xqa":
        atol = 0.05
        rtol = 0.05

        diff_abs = torch.abs(
            o_ref.view(batch_size, q_len_per_request, layer_dimensions.num_heads, -1)
            - output
        )
        diff_rel = diff_abs / (
            torch.abs(
                o_ref.view(
                    batch_size, q_len_per_request, layer_dimensions.num_heads, -1
                )
            )
            + 1e-8
        )

        within_tolerance = (diff_abs <= atol) | (diff_rel <= rtol)

        pass_ratio = within_tolerance.float().mean().item()

        required_ratio = 0.95
        assert pass_ratio >= required_ratio, (
            f"Total {o_ref.numel()} elements, only {pass_ratio:.1%} meet tolerance criteria, "
            f"require at least {required_ratio:.1%}"
        )


def trtllm_batch_decode_mla_sparse(
    batch_size: int,
    scale: float,
    dtype: torch.dtype,
    q_len_per_request: int,
    topk: int,
    is_varlen: bool,
    enable_pdl: bool,
    backend: str,
    qk_nope_head_dim: int,
    num_attn_heads: int,
    use_cum_seq_lens_q: bool = False,
):
    compute_capability = get_compute_capability(torch.device(device="cuda"))
    if backend == "trtllm-gen":
        if compute_capability[0] != 10:
            pytest.skip("TRTLLM-GEN MLA only supports SM100 and SM103 GPUs")

    torch.manual_seed(42)
    device = "cuda:0"

    # Deepseek/GLM-5 attention config (decode-MLA)
    num_q_heads = num_attn_heads
    qk_rope_head_dim = 64
    kv_lora_rank = 512

    # Fixed or variable sequence lengths
    if is_varlen:
        # Variable sequence lengths
        MAX_SEQ_LEN = 4096
        seq_lens = [
            max(
                topk,
                int(
                    torch.distributions.Normal(MAX_SEQ_LEN, MAX_SEQ_LEN / 2)
                    .sample()
                    .item()
                ),
            )
            for _ in range(batch_size)
        ]
        seq_lens[-1] = MAX_SEQ_LEN  # Ensure at least one max length
        seq_lens = [min(s, MAX_SEQ_LEN) for s in seq_lens]
    else:
        # Fixed sequence length
        MAX_SEQ_LEN = 4096
        seq_lens = [MAX_SEQ_LEN] * batch_size

    max_seq_len = max(seq_lens)
    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int, device=device)

    # Initialize query tensors
    query = torch.randn(
        batch_size,
        q_len_per_request,
        num_q_heads,
        kv_lora_rank + qk_rope_head_dim,
        device=device,
    )
    query.clamp_(min=-1.0, max=1.0)
    query = query.to(dtype)

    # Calculate blocks needed
    page_size = 32
    blocks_per_seq = (seq_lens_tensor + page_size - 1) // page_size
    max_num_blocks_per_seq = blocks_per_seq.max().item()
    total_blocks_needed = int(blocks_per_seq.sum().item())

    # Generate random but unique block IDs
    all_block_ids = torch.randperm(total_blocks_needed, device=device)

    # Create block tables
    block_tables = torch.zeros(
        (batch_size, max_num_blocks_per_seq), dtype=torch.int, device=device
    )
    block_id = 0
    for i in range(batch_size):
        num_blocks_needed = int(blocks_per_seq[i].item())
        block_tables[i, :num_blocks_needed] = all_block_ids[
            block_id : block_id + num_blocks_needed
        ]
        block_id += num_blocks_needed

    # Create KV cache
    num_blocks = total_blocks_needed
    kv_cache = torch.randn(
        size=(num_blocks, page_size, kv_lora_rank + qk_rope_head_dim),
        device=device,
    )
    kv_cache.clamp_(min=-1.0, max=1.0)
    kv_cache = kv_cache.to(dtype)

    # Generate sparse indices
    abs_indices, indices_in_kvcache = generate_sparse_indices(
        batch_size,
        q_len_per_request,
        seq_lens_tensor,
        topk,
        page_size,
        block_tables,
        device,
    )

    # Mask unused KV cache entries with NaN for correctness checking
    kv_cache_ref = kv_cache.clone()
    if dtype == torch.float8_e4m3fn:
        kv_cache_ref = kv_cache_ref.to(torch.bfloat16)

    # Mark all positions as NaN initially
    all_indices = indices_in_kvcache.flatten().tolist()
    all_indices = list(set(all_indices))
    if -1 in all_indices:
        all_indices.remove(-1)

    # Only used indices should be valid
    kv_cache_flat = kv_cache_ref.view(-1, kv_lora_rank + qk_rope_head_dim)
    used_mask = torch.zeros(kv_cache_flat.size(0), dtype=torch.bool, device="cpu")
    used_mask[torch.tensor(all_indices, dtype=torch.int64, device="cpu")] = True
    kv_cache_flat[~used_mask] = float("0")

    # Allocate workspace buffers
    global global_workspace_buffer, global_trtllm_gen_fmha_workspace_buffer
    if global_workspace_buffer is None:
        global_workspace_buffer = torch.empty(
            workspace_size, dtype=torch.int8, device=device
        )
    if global_trtllm_gen_fmha_workspace_buffer is None:
        global_trtllm_gen_fmha_workspace_buffer = torch.empty(
            workspace_size, dtype=torch.int8, device=device
        )
    workspace_buffer = global_trtllm_gen_fmha_workspace_buffer
    # workspace_buffer_ref = global_workspace_buffer

    # Run sparse decode-MLA
    cum_seq_lens_q = None
    max_q_len = None
    block_tables_input = indices_in_kvcache
    expected_shape = (batch_size, q_len_per_request, num_q_heads, kv_lora_rank)
    if use_cum_seq_lens_q:
        q_lens = (
            torch.arange(batch_size, device=device, dtype=torch.int32)
            % q_len_per_request
        ) + 1
        q_lens[-1] = q_len_per_request
        cum_seq_lens_q = torch.empty(batch_size + 1, device=device, dtype=torch.int32)
        cum_seq_lens_q[0] = 0
        cum_seq_lens_q[1:] = torch.cumsum(q_lens, dim=0)
        query_input = torch.cat(
            [query[i, : int(q_lens[i].item())] for i in range(batch_size)],
            dim=0,
        )
        block_tables_input = torch.cat(
            [indices_in_kvcache[i, : int(q_lens[i].item())] for i in range(batch_size)],
            dim=0,
        )
        max_q_len = q_len_per_request
        expected_shape = (query_input.size(0), num_q_heads, kv_lora_rank)
    else:
        query_input = query.clone()

    output = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
        query=query_input,
        kv_cache=kv_cache.unsqueeze(1),
        workspace_buffer=workspace_buffer,
        qk_nope_head_dim=qk_nope_head_dim,
        kv_lora_rank=kv_lora_rank,
        qk_rope_head_dim=qk_rope_head_dim,
        block_tables=block_tables_input,
        seq_lens=seq_lens_tensor,
        max_seq_len=max_seq_len,
        sparse_mla_top_k=topk,
        bmm1_scale=scale / ((qk_nope_head_dim + qk_rope_head_dim) ** 0.5),
        bmm2_scale=1.0,
        enable_pdl=enable_pdl,
        backend=backend,
        cum_seq_lens_q=cum_seq_lens_q,
        max_q_len=max_q_len,
    )

    # For now, just check that output has correct shape and no NaNs
    assert output.shape == expected_shape, (
        f"Output shape {output.shape} != {expected_shape}"
    )

    # Check for NaNs
    if dtype != torch.float8_e4m3fn:
        assert not torch.isnan(output).any(), "Output contains NaN values"

    # Generate reference output using PyTorch implementation
    query_ref = query.clone()
    if dtype == torch.float8_e4m3fn:
        query_ref = query_ref.to(torch.bfloat16)

    # Split kv_cache into K and V components
    # K uses full dimension (kv_lora_rank + qk_rope_head_dim)
    # V uses only kv_lora_rank dimension
    blocked_k = kv_cache_ref  # [num_blocks, page_size, kv_lora_rank + qk_rope_head_dim]
    blocked_v = kv_cache_ref[
        ..., :kv_lora_rank
    ]  # [num_blocks, page_size, kv_lora_rank]

    sm_scale = scale / ((qk_nope_head_dim + qk_rope_head_dim) ** 0.5)

    out_ref, _ = sparse_mla_reference_torch(
        cache_seqlens=seq_lens_tensor,
        block_table=block_tables,
        q=query_ref,
        blocked_k=blocked_k,
        blocked_v=blocked_v,
        page_size=page_size,
        is_causal=True,  # Cover cases where number of attendable kv values are less than topk
        sm_scale=sm_scale,
        indices=abs_indices,
    )
    if use_cum_seq_lens_q:
        out_ref = torch.cat(
            [out_ref[i, : int(q_lens[i].item())] for i in range(batch_size)],
            dim=0,
        )

    # Compare outputs
    assert not torch.isnan(output).any(), "Kernel output contains NaN values"
    assert not torch.isnan(out_ref).any(), "Reference output contains NaN values"

    if dtype == torch.float8_e4m3fn:
        # FP8 has lower precision, use more relaxed tolerances
        try:
            torch.testing.assert_close(
                output.float(),
                out_ref.float(),
                rtol=1e-1,
                atol=1e-1,
            )
        except AssertionError as e:
            # Calculate element-wise differences for debugging
            diff = torch.abs(output.float() - out_ref.float())
            max_diff = diff.max().item()
            mean_diff = diff.mean().item()
            print(f"Max difference: {max_diff}, Mean difference: {mean_diff}")
            output_sample = (
                output[0, 0, :8] if output.ndim == 3 else output[0, 0, 0, :8]
            )
            ref_sample = (
                out_ref[0, 0, :8] if out_ref.ndim == 3 else out_ref[0, 0, 0, :8]
            )
            print(f"Output sample: {output_sample}")
            print(f"Reference sample: {ref_sample}")
            raise e
    else:
        # BF16 should have better precision
        try:
            torch.testing.assert_close(
                output.float(),
                out_ref.float(),
                rtol=2e-2,
                atol=8e-4,
            )
        except AssertionError as e:
            # Calculate element-wise differences for debugging
            diff = torch.abs(output.float() - out_ref.float())
            max_diff = diff.max().item()
            mean_diff = diff.mean().item()
            print(f"Max difference: {max_diff}, Mean difference: {mean_diff}")
            if output.ndim == 3:
                print(f"Output sample: {output[0, 0, :8]}")
                print(f"Output sample: {output[1, 0, :8]}")
                print(f"Reference sample: {out_ref[0, 0, :8]}")
                print(f"Reference sample: {out_ref[1, 0, :8]}")
            else:
                print(f"Output sample: {output[0, 0, 0, :8]}")
                print(f"Output sample: {output[0, 1, 0, :8]}")
                print(f"Reference sample: {out_ref[0, 0, 0, :8]}")
                print(f"Reference sample: {out_ref[0, 1, 0, :8]}")
            raise e

    print(
        f"Sparse MLA test passed: batch_size={batch_size}, topk={topk}, "
        f"q_len={q_len_per_request}, varlen={is_varlen}, dtype={dtype}"
    )


def _run_trtllm_batch_decode_mla_head_case(
    num_heads: int,
    batch_size: int,
    dtype: torch.dtype,
    max_seq_len: int,
) -> None:
    trtllm_batch_decode_mla(
        MLALayerDimensions(
            head_dimensions=deepseek_mla_dimensions,
            num_heads=num_heads,
        ),
        batch_size=batch_size,
        scale=1.0,
        dtype=dtype,
        page_size=64,
        q_len_per_request=1,
        dynamic_scale=False,
        enable_pdl=False,
        backend="trtllm-gen",
        MAX_SEQ_LEN=max_seq_len,
        skips_softmax=False,
    )


@pytest.mark.parametrize("num_heads", [6, 12, 24, 48])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
@pytest.mark.parametrize("max_seq_len", [128, 1024])
def test_trtllm_batch_decode_mla_non_power_of_two_heads(
    num_heads: int,
    dtype: torch.dtype,
    max_seq_len: int,
) -> None:
    _run_trtllm_batch_decode_mla_head_case(
        num_heads,
        1,
        dtype,
        max_seq_len,
    )


@pytest.mark.parametrize(
    "num_heads,dtype,max_seq_len",
    [
        (3, torch.bfloat16, 128),
        (9, torch.float8_e4m3fn, 1024),
        (17, torch.bfloat16, 1024),
        (33, torch.float8_e4m3fn, 128),
        (63, torch.bfloat16, 1024),
        (192, torch.bfloat16, 128),
    ],
)
def test_trtllm_batch_decode_mla_generic_non_power_of_two_heads(
    num_heads: int,
    dtype: torch.dtype,
    max_seq_len: int,
) -> None:
    _run_trtllm_batch_decode_mla_head_case(
        num_heads,
        1,
        dtype,
        max_seq_len,
    )


def test_trtllm_batch_decode_mla_non_power_of_two_heads_cum_seq_lens_q() -> None:
    trtllm_batch_decode_mla(
        MLALayerDimensions(
            head_dimensions=deepseek_mla_dimensions,
            num_heads=24,
        ),
        batch_size=2,
        scale=1.0,
        dtype=torch.bfloat16,
        page_size=64,
        q_len_per_request=2,
        dynamic_scale=False,
        enable_pdl=False,
        backend="trtllm-gen",
        MAX_SEQ_LEN=1024,
        skips_softmax=False,
        use_cum_seq_lens_q=True,
    )


def _run_trtllm_batch_decode_sparse_mla_head_case(
    num_heads: int,
    batch_size: int,
    dtype: torch.dtype,
    topk: int,
) -> None:
    trtllm_batch_decode_mla_sparse(
        batch_size=batch_size,
        scale=1.0,
        dtype=dtype,
        q_len_per_request=1,
        topk=topk,
        is_varlen=False,
        enable_pdl=False,
        backend="trtllm-gen",
        qk_nope_head_dim=128,
        num_attn_heads=num_heads,
        use_cum_seq_lens_q=False,
    )


@pytest.mark.parametrize(
    "num_heads,dtype",
    [
        (6, torch.bfloat16),
        (12, torch.bfloat16),
        (24, torch.bfloat16),
        (48, torch.bfloat16),
        (24, torch.float8_e4m3fn),
    ],
)
@pytest.mark.parametrize("topk", [128, 2048])
def test_trtllm_batch_decode_sparse_mla_non_power_of_two_heads(
    num_heads: int,
    dtype: torch.dtype,
    topk: int,
) -> None:
    _run_trtllm_batch_decode_sparse_mla_head_case(
        num_heads,
        1,
        dtype,
        topk,
    )


@pytest.mark.parametrize(
    "num_heads,dtype,topk",
    [
        (3, torch.bfloat16, 128),
        (9, torch.float8_e4m3fn, 128),
        (17, torch.bfloat16, 2048),
        (33, torch.float8_e4m3fn, 128),
        (63, torch.bfloat16, 2048),
    ],
)
def test_trtllm_batch_decode_sparse_mla_generic_non_power_of_two_heads(
    num_heads: int,
    dtype: torch.dtype,
    topk: int,
) -> None:
    _run_trtllm_batch_decode_sparse_mla_head_case(
        num_heads,
        1,
        dtype,
        topk,
    )


@pytest.mark.parametrize(
    "batch_size,dtype",
    [(1, torch.bfloat16), (2, torch.float8_e4m3fn)],
)
def test_trtllm_batch_decode_sparse_mla_power_of_two_heads(
    batch_size: int,
    dtype: torch.dtype,
) -> None:
    _run_trtllm_batch_decode_sparse_mla_head_case(
        16,
        batch_size,
        dtype,
        2048,
    )


@pytest.mark.parametrize(
    "layer_dimensions",
    supported_mla_layer_dimensions,
)
@pytest.mark.parametrize(
    "batch_size",
    [1, 2, 4, 16, 32, 64, 128, 256, 512, 768, 1024],
)
@pytest.mark.parametrize("scale", [1.0, 0.5])
@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
@pytest.mark.parametrize("page_size", [32, 64])
@pytest.mark.parametrize(
    "q_len_per_request", [1, 2]
)  # todo(Yingyi): verify larger q_len_per_request
@pytest.mark.parametrize("dynamic_scale", [False])
@pytest.mark.parametrize("enable_pdl", [True, False, None])
@pytest.mark.parametrize("backend", ["trtllm-gen", "xqa", "cute-dsl"])
@pytest.mark.parametrize("skips_softmax", [False, True])
@pytest.mark.parametrize("uses_shared_paged_kv_idx", [True, False])
def test_trtllm_batch_decode_mla(
    layer_dimensions: MLALayerDimensions,
    batch_size: int,
    scale: float,
    dtype: torch.dtype,
    page_size: int,
    q_len_per_request: int,
    dynamic_scale: bool,
    enable_pdl: bool,
    backend: str,
    skips_softmax: bool,
    uses_shared_paged_kv_idx: bool,
):
    if backend == "xqa" and layer_dimensions.head_dimensions == smaller_mla_dimensions:
        pytest.skip("XQA MLA does not support smaller MLA dimensions yet.")
    if backend == "xqa" and layer_dimensions.num_heads != 128:
        pytest.skip("XQA MLA only supports 128 query heads (head_group_ratio=128)")
    if (
        backend == "cute-dsl"
        and layer_dimensions.head_dimensions == smaller_mla_dimensions
    ):
        pytest.skip("cute-dsl MLA requires 512 latent dim and 64 rope dim")
    # Note: the cute-dsl branch runs the default implementation (auto →
    # monolithic, since this test passes no sinks).  We don't parametrize
    # cute_dsl_impl here because this test's purpose is public-API smoke
    # across all three backends; the modular vs monolithic matrix lives
    # in test_cute_dsl_mla_decode.py.
    trtllm_batch_decode_mla(
        layer_dimensions,
        batch_size,
        scale,
        dtype,
        page_size,
        q_len_per_request,
        dynamic_scale,
        enable_pdl,
        backend,
        1024,
        skips_softmax,
        uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
        use_cum_seq_lens_q=False,
    )


@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
@pytest.mark.parametrize("skips_softmax", [False, True])
@pytest.mark.parametrize("uses_shared_paged_kv_idx", [True, False])
def test_trtllm_batch_decode_mla_cum_seq_lens_q(
    dtype: torch.dtype,
    skips_softmax: bool,
    uses_shared_paged_kv_idx: bool,
):
    trtllm_batch_decode_mla(
        supported_mla_layer_dimensions[0],
        batch_size=2,
        scale=1.0,
        dtype=dtype,
        page_size=64,
        q_len_per_request=2,
        dynamic_scale=False,
        enable_pdl=False,
        backend="trtllm-gen",
        MAX_SEQ_LEN=1024,
        skips_softmax=skips_softmax,
        uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
        use_cum_seq_lens_q=True,
    )


@pytest.mark.parametrize("backend", ["trtllm-gen", "auto"])
def test_trtllm_batch_decode_mla_variable_q_capacity_overestimate(backend: str):
    """TRT accepts graph capacity above the current compact token count."""
    if get_compute_capability(torch.device("cuda"))[0] != 10:
        pytest.skip("TRTLLM-GEN MLA requires SM100 or SM103")
    trtllm_batch_decode_mla(
        supported_mla_layer_dimensions[0],
        batch_size=2,
        scale=1.0,
        dtype=torch.bfloat16,
        page_size=64,
        q_len_per_request=2,
        dynamic_scale=False,
        enable_pdl=False,
        backend=backend,
        MAX_SEQ_LEN=1024,
        skips_softmax=False,
        uses_shared_paged_kv_idx=False,
        use_cum_seq_lens_q=True,
        max_q_len_exceeds_total_q=True,
    )


@pytest.mark.parametrize(
    "batch_size",
    [1, 2, 4, 16, 32, 64, 128],
)
@pytest.mark.parametrize("scale", [1.0])
@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
@pytest.mark.parametrize("q_len_per_request", [1, 2])
@pytest.mark.parametrize("topk", [128, 2048])
@pytest.mark.parametrize("is_varlen", [False, True])
@pytest.mark.parametrize("enable_pdl", [True, False, None])
@pytest.mark.parametrize("backend", ["trtllm-gen"])
@pytest.mark.parametrize("qk_nope_head_dim", [128, 192])
@pytest.mark.parametrize("num_attn_heads", [128, 64])
def test_trtllm_batch_decode_mla_sparse(
    batch_size: int,
    scale: float,
    dtype: torch.dtype,
    q_len_per_request: int,
    topk: int,
    is_varlen: bool,
    enable_pdl: bool,
    backend: str,
    qk_nope_head_dim: int,
    num_attn_heads: int,
):
    """
    Test sparse MLA decoding with top-k attention.
    Based on FlashMLA test patterns from:
    https://github.com/deepseek-ai/FlashMLA/blob/main/tests/test_flash_mla_decoding.py
    """
    trtllm_batch_decode_mla_sparse(
        batch_size,
        scale,
        dtype,
        q_len_per_request,
        topk,
        is_varlen,
        enable_pdl,
        backend,
        qk_nope_head_dim,
        num_attn_heads,
        use_cum_seq_lens_q=False,
    )


def test_trtllm_batch_decode_mla_sparse_cum_seq_lens_q():
    trtllm_batch_decode_mla_sparse(
        batch_size=2,
        scale=1.0,
        dtype=torch.bfloat16,
        q_len_per_request=2,
        topk=128,
        is_varlen=False,
        enable_pdl=False,
        backend="trtllm-gen",
        qk_nope_head_dim=128,
        num_attn_heads=128,
        use_cum_seq_lens_q=True,
    )


@pytest.mark.parametrize("uses_shared_paged_kv_idx", [True, False])
def test_trtllm_batch_decode_mla_native_block_table_width(
    uses_shared_paged_kv_idx: bool,
):
    page_size = 32
    max_seq_len = 1025
    block_table_width = (max_seq_len + page_size - 1) // page_size
    assert block_table_width == 33
    assert block_table_width % (128 // page_size) != 0

    trtllm_batch_decode_mla(
        supported_mla_layer_dimensions[0],
        batch_size=1,
        scale=1.0,
        dtype=torch.bfloat16,
        page_size=page_size,
        q_len_per_request=1,
        dynamic_scale=False,
        enable_pdl=False,
        backend="trtllm-gen",
        MAX_SEQ_LEN=max_seq_len,
        skips_softmax=False,
        uses_shared_paged_kv_idx=uses_shared_paged_kv_idx,
    )


@pytest.mark.parametrize("q_len_per_request", [1, 2, 4])
@pytest.mark.parametrize("batch_size", [1, 4])
def test_trtllm_batch_decode_mla_preallocated_out(
    q_len_per_request: int,
    batch_size: int,
):
    """Issue #2856: pre-allocated out tensor rejected when q_len_per_req > 1.
    The shape check hardcoded 3D but query is 4D for multi-token generation."""
    cc = get_compute_capability(torch.device("cuda"))
    if cc[0] != 10:
        pytest.skip("trtllm-gen MLA requires SM100/SM103")

    device = "cuda:0"
    layer_dim = supported_mla_layer_dimensions[0]
    kv_lora_rank = layer_dim.head_dimensions.kv_lora_rank
    qk_nope_head_dim = layer_dim.head_dimensions.qk_nope_head_dim
    qk_rope_head_dim = layer_dim.head_dimensions.qk_rope_head_dim
    num_heads = layer_dim.num_heads
    head_dim_qk = kv_lora_rank + qk_rope_head_dim

    page_size = 64
    max_seq_len = 256
    num_pages_per_seq = (max_seq_len + page_size - 1) // page_size
    head_dim_ckv_kpe = kv_lora_rank + qk_rope_head_dim

    kv_cache = torch.randn(
        num_pages_per_seq * batch_size,
        1,
        page_size,
        head_dim_ckv_kpe,
        dtype=torch.bfloat16,
        device=device,
    )
    block_tables = torch.arange(
        num_pages_per_seq * batch_size,
        device=device,
        dtype=torch.int32,
    ).reshape(batch_size, num_pages_per_seq)
    seq_lens = torch.full((batch_size,), max_seq_len, device=device, dtype=torch.int32)

    query = torch.randn(
        batch_size,
        q_len_per_request,
        num_heads,
        head_dim_qk,
        dtype=torch.bfloat16,
        device=device,
    )

    global global_trtllm_gen_fmha_workspace_buffer
    if global_trtllm_gen_fmha_workspace_buffer is None:
        global_trtllm_gen_fmha_workspace_buffer = torch.empty(
            workspace_size,
            dtype=torch.int8,
            device=device,
        )
    workspace = global_trtllm_gen_fmha_workspace_buffer

    bmm1_scale = 1.0 / (head_dim_qk**0.5)
    counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes(
        batch_size, num_heads, get_device_sm_count(torch.device(device))
    )
    counter_buffer = torch.zeros(counter_bytes, dtype=torch.uint8, device=device)

    # out=None should work
    result_none = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
        query=query,
        kv_cache=kv_cache,
        workspace_buffer=workspace,
        qk_nope_head_dim=qk_nope_head_dim,
        kv_lora_rank=kv_lora_rank,
        qk_rope_head_dim=qk_rope_head_dim,
        block_tables=block_tables,
        seq_lens=seq_lens,
        max_seq_len=max_seq_len,
        bmm1_scale=bmm1_scale,
        bmm2_scale=1.0,
        backend="trtllm-gen",
        multi_ctas_kv_counter_buffer=counter_buffer,
    )
    expected_shape = (batch_size, q_len_per_request, num_heads, kv_lora_rank)
    assert result_none.shape == expected_shape
    assert torch.count_nonzero(counter_buffer).item() == 0

    # out=pre-allocated should also work (this was the bug)
    out = torch.empty(expected_shape, dtype=torch.bfloat16, device=device)
    result_pre = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
        query=query,
        kv_cache=kv_cache,
        workspace_buffer=workspace,
        qk_nope_head_dim=qk_nope_head_dim,
        kv_lora_rank=kv_lora_rank,
        qk_rope_head_dim=qk_rope_head_dim,
        block_tables=block_tables,
        seq_lens=seq_lens,
        max_seq_len=max_seq_len,
        out=out,
        bmm1_scale=bmm1_scale,
        bmm2_scale=1.0,
        backend="trtllm-gen",
        multi_ctas_kv_counter_buffer=counter_buffer,
    )
    assert result_pre.data_ptr() == out.data_ptr(), (
        "Expected kernel to write into provided out tensor"
    )
    assert result_pre.shape == expected_shape
    assert torch.count_nonzero(counter_buffer).item() == 0
    torch.testing.assert_close(result_none, result_pre, rtol=1e-3, atol=1e-3)

    # Buffer validation is shape-independent, so cover it once rather than for
    # every entry in the preallocated-output matrix.
    if batch_size == 1 and q_len_per_request == 1:
        with pytest.raises(
            ValueError, match="multi_ctas_kv_counter_buffer is too small"
        ):
            flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
                query=query,
                kv_cache=kv_cache,
                workspace_buffer=workspace,
                qk_nope_head_dim=qk_nope_head_dim,
                kv_lora_rank=kv_lora_rank,
                qk_rope_head_dim=qk_rope_head_dim,
                block_tables=block_tables,
                seq_lens=seq_lens,
                max_seq_len=max_seq_len,
                bmm1_scale=bmm1_scale,
                bmm2_scale=1.0,
                backend="trtllm-gen",
                multi_ctas_kv_counter_buffer=counter_buffer[:4],
            )

        strided_counter_buffer = torch.zeros(
            counter_bytes * 2, dtype=torch.uint8, device=device
        )[::2]
        assert strided_counter_buffer.numel() == counter_bytes
        assert not strided_counter_buffer.is_contiguous()
        with pytest.raises(
            ValueError, match="multi_ctas_kv_counter_buffer must be contiguous"
        ):
            flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
                query=query,
                kv_cache=kv_cache,
                workspace_buffer=workspace,
                qk_nope_head_dim=qk_nope_head_dim,
                kv_lora_rank=kv_lora_rank,
                qk_rope_head_dim=qk_rope_head_dim,
                block_tables=block_tables,
                seq_lens=seq_lens,
                max_seq_len=max_seq_len,
                bmm1_scale=bmm1_scale,
                bmm2_scale=1.0,
                backend="trtllm-gen",
                multi_ctas_kv_counter_buffer=strided_counter_buffer,
            )

        offset_counter_buffer = torch.zeros(
            counter_bytes + 1, dtype=torch.uint8, device=device
        )[1:]
        assert offset_counter_buffer.is_contiguous()
        assert offset_counter_buffer.data_ptr() % 16 != 0
        with pytest.raises(
            ValueError,
            match="multi_ctas_kv_counter_buffer must be 16-byte aligned",
        ):
            flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
                query=query,
                kv_cache=kv_cache,
                workspace_buffer=workspace,
                qk_nope_head_dim=qk_nope_head_dim,
                kv_lora_rank=kv_lora_rank,
                qk_rope_head_dim=qk_rope_head_dim,
                block_tables=block_tables,
                seq_lens=seq_lens,
                max_seq_len=max_seq_len,
                bmm1_scale=bmm1_scale,
                bmm2_scale=1.0,
                backend="trtllm-gen",
                multi_ctas_kv_counter_buffer=offset_counter_buffer,
            )
