# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Tests for the DeepSeek V4 CuTe DSL HCA decode integration."""

import math

import pytest
import torch

import flashinfer.mla._core as mla_core
from flashinfer.mla import (
    batch_decode_sparse_mla_dsv4,
    trtllm_batch_decode_sparse_mla_dsv4,
)
from flashinfer.utils import get_compute_capability


def _cpu_hca_inputs():
    batch_size, q_len, num_heads, head_dim = 2, 2, 128, 512
    rows = batch_size * q_len
    query = torch.empty(
        (batch_size, q_len, num_heads, head_dim), dtype=torch.float8_e4m3fn
    )
    window_cache = torch.empty((16, 32, head_dim), dtype=torch.float8_e4m3fn)
    compressed_cache = torch.empty((5, 128, head_dim), dtype=torch.float8_e4m3fn)
    window_positions = torch.arange(128, dtype=torch.int32)
    window_bases = torch.arange(rows, dtype=torch.int32).unsqueeze(1) * 128
    window_shifts = torch.tensor([7, 31, 63, 95], dtype=torch.int32).unsqueeze(1)
    window_indices = window_bases + torch.remainder(
        window_positions.unsqueeze(0) + window_shifts, 128
    )
    return {
        "query": query,
        "swa_kv_cache": window_cache,
        "workspace_buffer": torch.empty(4 << 20, dtype=torch.uint8),
        "sparse_indices": None,
        "compressed_kv_cache": compressed_cache,
        "sparse_topk_lens": torch.tensor([128, 129, 256, 257], dtype=torch.int32),
        "swa_topk_lens": torch.tensor([127, 128, 128, 128], dtype=torch.int32),
        "hca_swa_indices": window_indices.contiguous(),
        "hca_compressed_block_tables": torch.zeros((rows, 2), dtype=torch.int32),
        "hca_seq_lens": torch.tensor([129, 257], dtype=torch.int32),
        "sinks": torch.zeros(num_heads, dtype=torch.float32),
        "out": torch.empty(
            (batch_size, q_len, num_heads, head_dim), dtype=torch.bfloat16
        ),
        "bmm1_scale": 1.0 / math.sqrt(head_dim),
        "bmm2_scale": 1.0,
        "backend": "cute-dsl",
        "hca_is_causal": True,
    }


def _cpu_compressed_page_aligned_sparse_hca_inputs():
    args = _cpu_hca_inputs()
    rows = args["query"].shape[0] * args["query"].shape[1]
    compressed_page_size = args["compressed_kv_cache"].shape[1]
    window_indices = args["hca_swa_indices"].clone()
    compressed_tables = torch.tensor([[0, 1]] * rows, dtype=torch.int32)

    def expand(tables, page_size):
        offsets = torch.arange(page_size, dtype=torch.int32)
        return (tables.unsqueeze(-1) * page_size + offsets).flatten(1)

    sparse_indices = torch.cat(
        (
            window_indices,
            expand(compressed_tables, compressed_page_size),
        ),
        dim=1,
    )
    window_valid_lens = args["swa_topk_lens"]
    sparse_topk_lens = args["sparse_topk_lens"]
    window_columns = torch.arange(128)
    sparse_indices[:, :128].masked_fill_(
        window_columns.unsqueeze(0) >= window_valid_lens.unsqueeze(1), -1
    )
    compressed_columns = torch.arange(sparse_indices.shape[1] - 128)
    sparse_indices[:, 128:].masked_fill_(
        compressed_columns.unsqueeze(0) >= (sparse_topk_lens - 128).unsqueeze(1),
        0,
    )

    args["sparse_indices"] = sparse_indices
    args["seq_lens"] = torch.tensor([128, 16512], dtype=torch.int32)
    args["hca_sparse_indices_format"] = "compressed-page-aligned"
    args["swa_topk_lens"] = None
    args["hca_swa_indices"] = None
    args["hca_compressed_block_tables"] = None
    args["hca_seq_lens"] = None
    expected = {
        "window_indices": torch.where(
            torch.arange(128).unsqueeze(0) < window_valid_lens.unsqueeze(1),
            window_indices,
            torch.zeros_like(window_indices),
        ),
        "compressed_block_tables": torch.tensor(
            [[0, 0], [0, 0], [0, 0], [0, 1]], dtype=torch.int32
        ),
        "hca_seq_lens": torch.tensor([129, 257], dtype=torch.int32),
        "window_valid_lens": window_valid_lens,
    }
    return args, expected


def _cpu_compressed_page_aligned_conversion_inputs(
    window_page_size, compressed_page_size, extra_padding=0
):
    rows, head_dim = 4, 512
    seq_lens = torch.tensor([128, 16512], dtype=torch.int32)
    sparse_topk_lens = torch.tensor([128, 129, 256, 257], dtype=torch.int32)
    window_valid_lens = torch.tensor([127, 128, 128, 128], dtype=torch.int32)
    window_page_count = 128 // window_page_size
    compressed_table_capacity = 256
    compressed_input_capacity = 132
    compressed_page_count = compressed_table_capacity // compressed_page_size
    window_tables = (
        torch.arange(rows * window_page_count, dtype=torch.int32).reshape(rows, -1) + 3
    ) % 32
    compressed_tables = (
        torch.arange(rows * compressed_page_count, dtype=torch.int32).reshape(rows, -1)
        + 5
    ) % 32

    def expand(tables, page_size):
        offsets = torch.arange(page_size, dtype=torch.int32)
        return (tables.unsqueeze(-1) * page_size + offsets).flatten(1)

    window_indices = expand(window_tables, window_page_size)
    shifts = torch.tensor([7, 31, 63, 95], dtype=torch.int64).unsqueeze(1)
    gather_columns = torch.remainder(
        torch.arange(128, dtype=torch.int64).unsqueeze(0) + shifts, 128
    )
    window_indices = window_indices.gather(1, gather_columns)
    compressed_indices = expand(compressed_tables, compressed_page_size)[
        :, :compressed_input_capacity
    ]
    window_positions = torch.arange(128)
    compressed_positions = torch.arange(compressed_input_capacity)
    window_indices.masked_fill_(
        window_positions.unsqueeze(0) >= window_valid_lens.unsqueeze(1),
        torch.iinfo(torch.int32).max,
    )
    active_compressed_lens = sparse_topk_lens - 128
    compressed_indices.masked_fill_(
        compressed_positions.unsqueeze(0) >= active_compressed_lens.unsqueeze(1),
        torch.iinfo(torch.int32).max,
    )
    if extra_padding:
        compressed_indices = torch.cat(
            (
                compressed_indices,
                torch.full(
                    (rows, extra_padding),
                    torch.iinfo(torch.int32).max,
                    dtype=torch.int32,
                ),
            ),
            dim=1,
        )

    sparse_indices = torch.cat((window_indices, compressed_indices), dim=1)
    expected_window = torch.where(
        torch.arange(128).unsqueeze(0) < window_valid_lens.unsqueeze(1),
        window_indices,
        torch.zeros_like(window_indices),
    )
    expected_compressed = torch.where(
        (torch.arange(compressed_page_count) * compressed_page_size).unsqueeze(0)
        < active_compressed_lens.unsqueeze(1),
        compressed_tables,
        0,
    )
    return {
        "sparse_indices": sparse_indices,
        "sparse_topk_lens": sparse_topk_lens,
        "seq_lens": seq_lens,
        "swa_kv_cache": torch.empty(
            (32, window_page_size, head_dim), dtype=torch.float8_e4m3fn
        ),
        "compressed_kv_cache": torch.empty(
            (32, compressed_page_size, head_dim), dtype=torch.float8_e4m3fn
        ),
        "q_len": 2,
        "expected_window": expected_window,
        "expected_compressed": expected_compressed,
        "expected_window_lens": window_valid_lens,
    }


def test_dsv4_backend_resolution_is_explicit(monkeypatch):
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    assert (
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "auto")
        == "trtllm-gen"
    )
    assert (
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "cute-dsl")
        == "cute-dsl"
    )

    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (12, 0))
    assert (
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "auto")
        == "sparse"
    )
    with pytest.raises(ValueError, match="requires SM100/SM103"):
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "cute-dsl")

    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 1))
    with pytest.raises(ValueError, match="supports SM100/SM103"):
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "auto")
    with pytest.raises(ValueError, match="requires SM100/SM103"):
        mla_core._resolve_dsv4_sparse_mla_backend(torch.device("cpu"), "cute-dsl")


def test_dsv4_backend_neutral_alias_is_backward_compatible():
    assert batch_decode_sparse_mla_dsv4 is trtllm_batch_decode_sparse_mla_dsv4
    assert (
        trtllm_batch_decode_sparse_mla_dsv4.__name__
        == "trtllm_batch_decode_sparse_mla_dsv4"
    )


@pytest.mark.parametrize("kv_layout", ["HND", "NHD"])
def test_dsv4_hca_public_api_forwards_gather_indices(monkeypatch, kv_layout):
    pytest.importorskip("cutlass")
    from flashinfer.cute_dsl.attention.wrappers import batch_hca

    args = _cpu_hca_inputs()
    args["kv_layout"] = kv_layout
    args["hca_use_persistent"] = True
    if kv_layout == "NHD":
        args["swa_kv_cache"] = args["swa_kv_cache"].unsqueeze(2)
        args["compressed_kv_cache"] = args["compressed_kv_cache"].unsqueeze(2)
    captured = {}

    def fake_hca_decode(**kwargs):
        captured.update(kwargs)
        return kwargs["out"]

    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    monkeypatch.setattr(batch_hca, "cute_dsl_hca_decode", fake_hca_decode)

    result = trtllm_batch_decode_sparse_mla_dsv4(**args)
    assert result is args["out"]
    assert captured["window_indices"] is args["hca_swa_indices"]
    assert captured["compressed_block_tables"] is args["hca_compressed_block_tables"]
    assert captured["hca_seq_lens"] is args["hca_seq_lens"]
    assert captured["window_valid_lens"] is args["swa_topk_lens"]
    assert captured["is_persistent"] is True
    assert captured["window_kv_cache"].shape == (16 * 32, 512)
    assert captured["compressed_kv_cache"].shape == (5, 128, 512)
    assert captured["window_kv_cache"].data_ptr() == args["swa_kv_cache"].data_ptr()
    assert (
        captured["compressed_kv_cache"].data_ptr()
        == args["compressed_kv_cache"].data_ptr()
    )


def test_dsv4_hca_generates_metadata_from_compressed_page_aligned_sparse_indices(
    monkeypatch,
):
    pytest.importorskip("cutlass")
    from flashinfer.cute_dsl.attention.wrappers import batch_hca

    args, expected = _cpu_compressed_page_aligned_sparse_hca_inputs()
    captured = {}

    def fake_hca_decode(**kwargs):
        captured.update(kwargs)
        return kwargs["out"]

    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    monkeypatch.setattr(batch_hca, "cute_dsl_hca_decode", fake_hca_decode)

    result = trtllm_batch_decode_sparse_mla_dsv4(**args)
    assert result is args["out"]
    for name, value in expected.items():
        torch.testing.assert_close(captured[name], value)
    assert captured["is_causal"] is True


def test_dsv4_hca_rejects_out_of_range_swa_index(monkeypatch):
    args, _ = _cpu_compressed_page_aligned_sparse_hca_inputs()
    args["sparse_indices"][1, 0] = (
        args["swa_kv_cache"].shape[0] * args["swa_kv_cache"].shape[1]
    )
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    with pytest.raises(ValueError, match="active SWA sparse_indices"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)


@pytest.mark.parametrize(
    ("window_page_size", "compressed_page_size"),
    [(32, 128), (16, 32), (32, 16)],
)
def test_dsv4_hca_compressed_page_aligned_conversion_ignores_inactive_padding(
    window_page_size, compressed_page_size
):
    minimal = _cpu_compressed_page_aligned_conversion_inputs(
        window_page_size, compressed_page_size
    )
    padded = _cpu_compressed_page_aligned_conversion_inputs(
        window_page_size, compressed_page_size, extra_padding=512
    )

    def convert(inputs):
        return mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            inputs["sparse_indices"],
            inputs["sparse_topk_lens"],
            inputs["seq_lens"],
            inputs["swa_kv_cache"],
            inputs["compressed_kv_cache"],
            q_len=inputs["q_len"],
        )

    minimal_metadata = convert(minimal)
    padded_metadata = convert(padded)
    for field in (
        "hca_swa_indices",
        "hca_compressed_block_tables",
        "hca_seq_lens",
        "swa_topk_lens",
    ):
        minimal_value = getattr(minimal_metadata, field)
        padded_value = getattr(padded_metadata, field)
        torch.testing.assert_close(minimal_value, padded_value)

    torch.testing.assert_close(
        minimal_metadata.hca_swa_indices, minimal["expected_window"]
    )
    torch.testing.assert_close(
        minimal_metadata.hca_compressed_block_tables,
        minimal["expected_compressed"],
    )
    torch.testing.assert_close(
        minimal_metadata.swa_topk_lens, minimal["expected_window_lens"]
    )
    assert minimal_metadata.hca_compressed_block_tables.shape[1] == (
        256 // compressed_page_size
    )


def test_dsv4_hca_compressed_page_aligned_conversion_rejects_in_page_reordering():
    inputs = _cpu_compressed_page_aligned_conversion_inputs(32, 16)
    inputs["sparse_indices"][3, 129] += 16
    with pytest.raises(ValueError, match="canonical compressed-page expansion"):
        mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            inputs["sparse_indices"],
            inputs["sparse_topk_lens"],
            inputs["seq_lens"],
            inputs["swa_kv_cache"],
            inputs["compressed_kv_cache"],
            q_len=inputs["q_len"],
        )


def test_dsv4_hca_compressed_page_aligned_conversion_checks_across_chunks(monkeypatch):
    monkeypatch.setattr(mla_core, "_DSV4_HCA_CANONICAL_CHECK_MAX_ELEMENTS", 8)
    inputs = _cpu_compressed_page_aligned_conversion_inputs(32, 16)

    mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
        inputs["sparse_indices"],
        inputs["sparse_topk_lens"],
        inputs["seq_lens"],
        inputs["swa_kv_cache"],
        inputs["compressed_kv_cache"],
        q_len=inputs["q_len"],
    )
    inputs["sparse_indices"][3, 128 + 17] += 16
    with pytest.raises(ValueError, match="canonical compressed-page expansion"):
        mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            inputs["sparse_indices"],
            inputs["sparse_topk_lens"],
            inputs["seq_lens"],
            inputs["swa_kv_cache"],
            inputs["compressed_kv_cache"],
            q_len=inputs["q_len"],
        )


def test_dsv4_hca_compressed_page_aligned_conversion_fills_fully_inactive_pages():
    window_page_size = 32
    offsets = torch.arange(window_page_size, dtype=torch.int32)
    sparse_indices = (
        torch.arange(4, dtype=torch.int32).unsqueeze(-1) * window_page_size + offsets
    ).reshape(1, 128)
    sparse_indices[:, 1:] = torch.iinfo(torch.int32).max
    metadata = mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
        sparse_indices,
        torch.tensor([128], dtype=torch.int32),
        torch.tensor([1], dtype=torch.int32),
        torch.empty((4, window_page_size, 512), dtype=torch.float8_e4m3fn),
        torch.empty((1, 128, 512), dtype=torch.float8_e4m3fn),
        q_len=1,
    )
    torch.testing.assert_close(
        metadata.hca_swa_indices,
        torch.zeros((1, 128), dtype=torch.int32),
    )
    torch.testing.assert_close(
        metadata.hca_compressed_block_tables,
        torch.tensor([[0]], dtype=torch.int32),
    )
    torch.testing.assert_close(
        metadata.hca_seq_lens, torch.tensor([128], dtype=torch.int32)
    )
    torch.testing.assert_close(
        metadata.swa_topk_lens, torch.tensor([1], dtype=torch.int32)
    )


def test_dsv4_hca_compressed_page_aligned_conversion_rejects_out_of_range_page():
    inputs = _cpu_compressed_page_aligned_conversion_inputs(32, 16)
    inputs["sparse_indices"][3, 128] = 32 * 16
    with pytest.raises(ValueError, match="valid physical page IDs"):
        mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            inputs["sparse_indices"],
            inputs["sparse_topk_lens"],
            inputs["seq_lens"],
            inputs["swa_kv_cache"],
            inputs["compressed_kv_cache"],
            q_len=inputs["q_len"],
        )


def test_dsv4_hca_compressed_page_aligned_tag_rejects_query_batch_mismatch(monkeypatch):
    args, _ = _cpu_compressed_page_aligned_sparse_hca_inputs()
    args["seq_lens"] = torch.tensor([16512], dtype=torch.int32)
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    with pytest.raises(ValueError, match="seq_lens with 2 entries"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)


@pytest.mark.arch_blackwell
def test_dsv4_hca_compressed_page_aligned_conversion_rejects_mixed_devices():
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    inputs = _cpu_compressed_page_aligned_conversion_inputs(32, 128)
    with pytest.raises(ValueError, match="swa_kv_cache must be on cpu"):
        mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            inputs["sparse_indices"],
            inputs["sparse_topk_lens"],
            inputs["seq_lens"],
            inputs["swa_kv_cache"].cuda(),
            inputs["compressed_kv_cache"],
            q_len=inputs["q_len"],
        )


def test_dsv4_hca_rejects_legacy_token_indices(monkeypatch):
    args = _cpu_hca_inputs()
    args["sparse_indices"] = torch.zeros((4, 256), dtype=torch.int32)
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    with pytest.raises(ValueError, match="sparse_indices must be None"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)


def test_dsv4_hca_rejects_ambiguous_page_aligned_format(monkeypatch):
    args = _cpu_hca_inputs()
    args["hca_sparse_indices_format"] = "page-aligned"
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    with pytest.raises(ValueError, match="compressed-page-aligned"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)


def test_dsv4_hca_rejects_backend_specific_metadata(monkeypatch):
    args = _cpu_hca_inputs()
    args["hca_is_causal"] = False
    monkeypatch.setattr(mla_core, "get_compute_capability", lambda _device: (10, 0))
    with pytest.raises(ValueError, match="supports causal HCA only"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)

    args = _cpu_hca_inputs()
    args["seq_lens"] = torch.tensor([129, 257], dtype=torch.int32)
    with pytest.raises(ValueError, match="uses hca_seq_lens"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)

    args = _cpu_hca_inputs()
    args["backend"] = "trtllm-gen"
    args["sparse_indices"] = torch.zeros((4, 256), dtype=torch.int32)
    with pytest.raises(ValueError, match="does not accept hca_swa_indices"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)

    args = _cpu_hca_inputs()
    args["backend"] = "trtllm-gen"
    args["sparse_indices"] = torch.zeros((4, 256), dtype=torch.int32)
    args["hca_swa_indices"] = None
    args["hca_compressed_block_tables"] = None
    args["hca_seq_lens"] = None
    args["hca_is_causal"] = False
    with pytest.raises(ValueError, match="does not accept hca_is_causal"):
        trtllm_batch_decode_sparse_mla_dsv4(**args)


def test_hca_rejects_nonpersistent_grid_overflow():
    pytest.importorskip("cutlass")
    from flashinfer.cute_dsl.attention.wrappers.batch_hca import (
        _check_nonpersistent_grid,
    )

    _check_nonpersistent_grid(batch_size=1, q_len=8192)
    with pytest.raises(ValueError, match=r"batch_size \* q_len <= 65535"):
        _check_nonpersistent_grid(batch_size=8, q_len=8192)


def test_hca_opt_in_metadata_value_validation(monkeypatch):
    pytest.importorskip("cutlass")
    from flashinfer.cute_dsl.attention.wrappers.batch_hca import (
        _validate_hca_values,
    )

    args = _cpu_hca_inputs()
    validation_args = {
        "window_kv_cache": args["swa_kv_cache"].reshape(-1, 512),
        "compressed_kv_cache": args["compressed_kv_cache"],
        "window_indices": args["hca_swa_indices"],
        "compressed_block_tables": args["hca_compressed_block_tables"],
        "hca_seq_lens": args["hca_seq_lens"],
        "sparse_topk_lens": args["sparse_topk_lens"],
        "window_valid_lens": args["swa_topk_lens"],
        "q_len": 2,
        "is_causal": True,
    }
    monkeypatch.setenv("FLASHINFER_VALIDATE_INPUTS", "1")

    # Every gather coordinate must be legal, while only the tile-rounded
    # compressed footprint is validated; later page columns may use -1.
    validation_args["compressed_block_tables"] = torch.tensor(
        [[0, -1], [0, -1], [0, 1], [0, 1]], dtype=torch.int32
    )
    _validate_hca_values(**validation_args)

    validation_args["hca_seq_lens"] = torch.tensor([129, 385], dtype=torch.int32)
    with pytest.raises(ValueError, match="block-table capacity 384"):
        _validate_hca_values(**validation_args)

    validation_args["hca_seq_lens"] = args["hca_seq_lens"]
    validation_args["compressed_block_tables"][2, 1] = -1
    with pytest.raises(ValueError, match="compressed_block_tables active values"):
        _validate_hca_values(**validation_args)

    validation_args["compressed_block_tables"][2, 1] = 1
    validation_args["window_valid_lens"][0] = 0
    validation_args["window_indices"][0, 3] = -1
    with pytest.raises(ValueError, match="window_indices values"):
        _validate_hca_values(**validation_args)


def test_hca_static_contract_accepts_fp8_to_bf16():
    pytest.importorskip("cutlass")
    import cutlass

    from flashinfer.cute_dsl.attention.dsa.hca_fp8 import (
        MAX_SPLITS,
        BlackwellHeavilyCompressedAttentionForwardFP8,
    )

    can_implement = BlackwellHeavilyCompressedAttentionForwardFP8.can_implement
    config = {
        "B": 2,
        "S": 2,
        "K": 384,
        "H": 128,
        "L": 512,
        "in_dtype": cutlass.Float8E4M3FN,
        "out_dtype": cutlass.BFloat16,
        "acc_dtype": cutlass.Float32,
        "lse_dtype": cutlass.Float32,
        "mma_qk_tiler_mn": (128, 128),
        "mma_pv_tiler_mn": (128, 256),
        "split_kv": 1,
        "is_persistent": False,
        "is_var_seq": True,
        "is_var_split_kv": False,
        "page_size_cmp": 128,
    }

    def supports(**overrides):
        return can_implement(**{**config, **overrides})

    assert supports()
    assert not supports(page_size_cmp=0)
    assert not supports(page_size_cmp=48)
    assert not supports(split_kv=0)
    assert not supports(split_kv=MAX_SPLITS + 1)
    assert not supports(B=8, S=8192)
    assert supports(B=8, S=8192, is_persistent=True)


def test_hca_static_contract_distributes_swa_gather():
    pytest.importorskip("cutlass")
    import cutlass

    from flashinfer.cute_dsl.attention.dsa.hca_fp8 import (
        BlackwellHeavilyCompressedAttentionForwardFP8,
    )

    kernel = BlackwellHeavilyCompressedAttentionForwardFP8(
        cutlass.Float32,
        cutlass.Float32,
        (128, 128),
        (128, 256),
        1,
        128,
        0.0,
        True,
        False,
        False,
    )
    kernel._setup_attributes()

    assert kernel.swa_load_warp_ids == (12, 13, 14, 15)
    assert kernel.swa_k_iters_per_warp == 1
    assert kernel.swa_v_iters_per_warp == 1
    assert kernel.load_k_stage == 3
    assert kernel.load_v_stage == 2


def _reference_hca(
    query,
    window_cache,
    compressed_cache,
    window_indices,
    compressed_tables,
    sparse_topk_lens,
    window_valid_lens,
    sinks,
    softmax_scale,
    output_scale,
):
    batch_size, q_len, _, _ = query.shape
    flat_window_cache = window_cache.reshape(-1, query.shape[-1])
    output = torch.empty_like(query, dtype=torch.float32)
    for batch_idx in range(batch_size):
        for query_idx in range(q_len):
            row = batch_idx * q_len + query_idx
            window = flat_window_cache.index_select(0, window_indices[row])
            window = window[: int(window_valid_lens[row].item())]
            compressed = compressed_cache.index_select(
                0, compressed_tables[row]
            ).reshape(-1, query.shape[-1])
            compressed_len = int(sparse_topk_lens[row].item()) - 128
            compressed = compressed[:compressed_len]
            kv = torch.cat((window, compressed), dim=0).float()
            scores = torch.einsum("hd,kd->hk", query[batch_idx, query_idx].float(), kv)
            scores *= softmax_scale
            log_norm = torch.logaddexp(torch.logsumexp(scores, dim=-1), sinks.float())
            probabilities = torch.exp(scores - log_norm[:, None])
            output[batch_idx, query_idx] = (
                torch.einsum("hk,kd->hd", probabilities, kv) * output_scale
            )
    return output


@pytest.mark.arch_blackwell
@pytest.mark.parametrize(
    "compare_trtllm", (False, True), ids=("pytorch-reference", "trtllm-parity")
)
def test_cute_dsl_hca_fp8_to_bf16_correctness(monkeypatch, compare_trtllm):
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    if get_compute_capability(torch.device("cuda")) not in ((10, 0), (10, 3)):
        pytest.skip("CuTe DSL HCA requires SM100/SM103")
    pytest.importorskip("cutlass")

    torch.manual_seed(17)
    device = torch.device("cuda")
    batch_size, q_len, num_heads, head_dim = 2, 2, 128, 512
    query = (
        torch.randn((batch_size, q_len, num_heads, head_dim), device=device).clamp_(
            -2, 2
        )
    ).to(torch.float8_e4m3fn)
    window_cache = torch.randn((16, 32, head_dim), device=device).clamp_(-2, 2)
    window_cache = window_cache.to(torch.float8_e4m3fn)
    compressed_cache = torch.randn((5, 128, head_dim), device=device).clamp_(-2, 2)
    compressed_cache = compressed_cache.to(torch.float8_e4m3fn)
    window_tables = torch.tensor(
        [
            [7, 0, 13, 2],
            [3, 15, 1, 9],
            [12, 5, 8, 4],
            [10, 14, 6, 11],
        ],
        dtype=torch.int32,
        device=device,
    )
    window_offsets = torch.arange(32, dtype=torch.int32, device=device)
    window_indices = (window_tables.unsqueeze(-1) * 32 + window_offsets).flatten(1)
    window_shifts = torch.tensor(
        [7, 31, 63, 95], dtype=torch.int64, device=device
    ).unsqueeze(1)
    window_indices = window_indices.gather(
        1,
        torch.remainder(
            torch.arange(128, dtype=torch.int64, device=device).unsqueeze(0)
            + window_shifts,
            128,
        ),
    ).contiguous()
    compressed_tables = torch.tensor(
        [[3, 4], [3, 4], [1, 0], [1, 0]], dtype=torch.int32, device=device
    )
    # Keep three backing tiles for every request while the effective lengths
    # remain partial. This catches any causal mask optimization that assumes
    # the scheduler drops backing-cache tiles beyond sparse_topk_lens.
    hca_seq_lens = torch.tensor([384, 384], dtype=torch.int32, device=device)
    sparse_topk_lens = torch.tensor(
        [128, 129, 256, 257], dtype=torch.int32, device=device
    )
    window_valid_lens = torch.tensor(
        [127, 128, 128, 128], dtype=torch.int32, device=device
    )
    workspace = torch.empty(8 << 20, dtype=torch.uint8, device=device)
    softmax_scale = 1.0 / math.sqrt(head_dim)
    output_scale = 1.0
    compressed_offsets = torch.arange(128, dtype=torch.int32, device=device)
    compressed_page_aligned_sparse_indices = torch.cat(
        (
            window_indices,
            (compressed_tables.unsqueeze(-1) * 128 + compressed_offsets).flatten(1),
        ),
        dim=1,
    )
    window_positions = torch.arange(128, device=device)
    compressed_positions = torch.arange(256, device=device)
    compressed_page_aligned_sparse_indices[:, :128].masked_fill_(
        window_positions.unsqueeze(0) >= window_valid_lens.unsqueeze(1), -1
    )
    compressed_page_aligned_sparse_indices[:, 128:].masked_fill_(
        compressed_positions.unsqueeze(0) >= (sparse_topk_lens - 128).unsqueeze(1),
        0,
    )
    raw_seq_lens = torch.tensor([128, 16512], dtype=torch.int32, device=device)
    precomputed_metadata = (
        mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
            compressed_page_aligned_sparse_indices,
            sparse_topk_lens,
            raw_seq_lens,
            window_cache,
            compressed_cache,
            q_len=q_len,
        )
    )

    for sinks in (
        torch.zeros(num_heads, dtype=torch.float32, device=device),
        torch.linspace(4.0, 6.0, num_heads, dtype=torch.float32, device=device),
    ):
        output = trtllm_batch_decode_sparse_mla_dsv4(
            query=query,
            swa_kv_cache=window_cache,
            workspace_buffer=workspace,
            sparse_indices=None,
            compressed_kv_cache=compressed_cache,
            sparse_topk_lens=sparse_topk_lens,
            out=None,
            bmm1_scale=softmax_scale,
            bmm2_scale=output_scale,
            sinks=sinks,
            swa_topk_lens=window_valid_lens,
            backend="cute-dsl",
            hca_swa_indices=window_indices,
            hca_compressed_block_tables=compressed_tables,
            hca_seq_lens=hca_seq_lens,
            hca_is_causal=True,
            hca_use_persistent=True,
        )
        tagged_output = trtllm_batch_decode_sparse_mla_dsv4(
            query=query,
            swa_kv_cache=window_cache,
            workspace_buffer=workspace,
            sparse_indices=compressed_page_aligned_sparse_indices,
            compressed_kv_cache=compressed_cache,
            sparse_topk_lens=sparse_topk_lens,
            seq_lens=raw_seq_lens,
            out=None,
            bmm1_scale=softmax_scale,
            bmm2_scale=output_scale,
            sinks=sinks,
            backend="cute-dsl",
            hca_sparse_indices_format="compressed-page-aligned",
            hca_is_causal=True,
            hca_use_persistent=True,
        )
        reference = _reference_hca(
            query,
            window_cache,
            compressed_cache,
            window_indices,
            compressed_tables,
            sparse_topk_lens,
            window_valid_lens,
            sinks,
            softmax_scale,
            output_scale,
        )
        trtllm_output = None
        if compare_trtllm:
            try:
                trtllm_output = trtllm_batch_decode_sparse_mla_dsv4(
                    query=query,
                    swa_kv_cache=window_cache,
                    workspace_buffer=workspace,
                    sparse_indices=compressed_page_aligned_sparse_indices,
                    compressed_kv_cache=compressed_cache,
                    sparse_topk_lens=sparse_topk_lens,
                    seq_lens=raw_seq_lens,
                    bmm1_scale=softmax_scale,
                    bmm2_scale=output_scale,
                    sinks=sinks,
                    backend="trtllm-gen",
                    enable_pdl=False,
                )
            except RuntimeError as error:
                message = str(error)
                if (
                    "trtllm_paged_attention_decode_sparse_mla_dsv4 is not available"
                    in message
                    or "Missing TRTLLM-GEN kernel (decode)" in message
                    or (
                        "Ninja build failed" in message
                        and "missing and no known rule" in message
                    )
                ):
                    pytest.skip("TRTLLM-GEN sparse MLA artifacts are unavailable")
                raise
        assert output.dtype == torch.bfloat16
        assert output.shape == query.shape
        assert torch.isfinite(output).all()
        torch.testing.assert_close(output.float(), reference, atol=0.13, rtol=1e-5)
        if trtllm_output is not None:
            assert trtllm_output.dtype == torch.bfloat16
            assert trtllm_output.shape == query.shape
            assert torch.isfinite(trtllm_output).all()
            torch.testing.assert_close(
                trtllm_output.float(), reference, atol=0.13, rtol=1e-5
            )
            torch.testing.assert_close(output, trtllm_output, atol=0.13, rtol=1e-5)
        torch.testing.assert_close(tagged_output, output, atol=0.00390625, rtol=0)

    explicit_output = torch.empty_like(query, dtype=torch.bfloat16)
    precomputed_output = torch.empty_like(explicit_output)
    trtllm_batch_decode_sparse_mla_dsv4(
        query=query,
        swa_kv_cache=window_cache,
        workspace_buffer=workspace,
        sparse_indices=None,
        compressed_kv_cache=compressed_cache,
        sparse_topk_lens=sparse_topk_lens,
        out=explicit_output,
        bmm1_scale=softmax_scale,
        swa_topk_lens=window_valid_lens,
        backend="cute-dsl",
        hca_swa_indices=window_indices,
        hca_compressed_block_tables=compressed_tables,
        hca_seq_lens=hca_seq_lens,
        hca_use_persistent=True,
    )

    def run_precomputed_metadata():
        return trtllm_batch_decode_sparse_mla_dsv4(
            query=query,
            swa_kv_cache=window_cache,
            workspace_buffer=workspace,
            sparse_indices=None,
            compressed_kv_cache=compressed_cache,
            sparse_topk_lens=sparse_topk_lens,
            out=precomputed_output,
            bmm1_scale=softmax_scale,
            swa_topk_lens=precomputed_metadata.swa_topk_lens,
            backend="cute-dsl",
            hca_swa_indices=precomputed_metadata.hca_swa_indices,
            hca_compressed_block_tables=(
                precomputed_metadata.hca_compressed_block_tables
            ),
            hca_seq_lens=precomputed_metadata.hca_seq_lens,
            hca_is_causal=precomputed_metadata.hca_is_causal,
            hca_use_persistent=True,
        )

    run_precomputed_metadata()
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        run_precomputed_metadata()
    graph.replay()
    torch.cuda.synchronize()
    torch.testing.assert_close(
        precomputed_output, explicit_output, atol=0.00390625, rtol=0
    )

    with monkeypatch.context() as capture_monkeypatch:
        capture_monkeypatch.setattr(
            torch.cuda, "is_current_stream_capturing", lambda: True
        )
        with pytest.raises(RuntimeError, match="not CUDA Graph capture safe"):
            mla_core.convert_compressed_page_aligned_sparse_indices_to_hca_metadata(
                compressed_page_aligned_sparse_indices,
                sparse_topk_lens,
                raw_seq_lens,
                window_cache,
                compressed_cache,
                q_len=q_len,
            )
        with pytest.raises(RuntimeError, match="not CUDA Graph capture safe"):
            trtllm_batch_decode_sparse_mla_dsv4(
                query=query,
                swa_kv_cache=window_cache,
                workspace_buffer=workspace,
                sparse_indices=compressed_page_aligned_sparse_indices,
                compressed_kv_cache=compressed_cache,
                sparse_topk_lens=sparse_topk_lens,
                seq_lens=raw_seq_lens,
                bmm1_scale=softmax_scale,
                backend="cute-dsl",
                hca_sparse_indices_format="compressed-page-aligned",
                hca_use_persistent=True,
            )

    # Capture the private LSE scratch/result to verify the all-empty split-K
    # reduction. The public DSV4 API intentionally returns output only.
    from flashinfer.cute_dsl.attention.wrappers import batch_hca

    real_compile_hca_kernel = batch_hca._compile_hca_kernel
    captured = {}

    def compile_and_capture_lse(*args, **kwargs):
        compiled_kernel = real_compile_hca_kernel(*args, **kwargs)

        def run_and_capture_lse(*kernel_args, **kernel_kwargs):
            captured["lse"] = kernel_args[6]
            return compiled_kernel(*kernel_args, **kernel_kwargs)

        return run_and_capture_lse

    monkeypatch.setattr(batch_hca, "_compile_hca_kernel", compile_and_capture_lse)
    sparse_topk_lens.fill_(128)
    window_valid_lens.zero_()
    output = trtllm_batch_decode_sparse_mla_dsv4(
        query=query,
        swa_kv_cache=window_cache,
        workspace_buffer=workspace,
        sparse_indices=None,
        compressed_kv_cache=compressed_cache,
        sparse_topk_lens=sparse_topk_lens,
        bmm1_scale=softmax_scale,
        bmm2_scale=output_scale,
        sinks=None,
        swa_topk_lens=window_valid_lens,
        backend="cute-dsl",
        hca_swa_indices=window_indices,
        hca_compressed_block_tables=compressed_tables,
        hca_seq_lens=hca_seq_lens,
        hca_is_causal=True,
        hca_use_persistent=True,
    )
    torch.testing.assert_close(output, torch.zeros_like(output), atol=0, rtol=0)
    assert torch.isneginf(captured["lse"]).all()


@pytest.mark.arch_blackwell
@pytest.mark.parametrize(
    ("num_heads", "window_page_size", "compressed_page_size"),
    ((64, 32, 128), (128, 32, 16), (128, 16, 32)),
)
def test_cute_dsl_hca_partial_head_tile_and_empty_attention(
    num_heads, window_page_size, compressed_page_size
):
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required")
    if get_compute_capability(torch.device("cuda")) not in ((10, 0), (10, 3)):
        pytest.skip("CuTe DSL HCA requires SM100/SM103")
    pytest.importorskip("cutlass")

    torch.manual_seed(23)
    device = torch.device("cuda")
    batch_size, q_len, head_dim = 1, 1, 512
    query = torch.randn((batch_size, q_len, num_heads, head_dim), device=device).clamp_(
        -2, 2
    )
    query = query.to(torch.float8_e4m3fn)
    active_window_pages = 128 // window_page_size
    active_compressed_pages = 256 // compressed_page_size
    window_table_pages = 2 * active_window_pages
    compressed_table_pages = 2 * active_compressed_pages
    window_cache = torch.randn(
        (window_table_pages, window_page_size, head_dim), device=device
    ).clamp_(-2, 2)
    window_cache = window_cache.to(torch.float8_e4m3fn)
    compressed_cache = torch.randn(
        (compressed_table_pages, compressed_page_size, head_dim), device=device
    ).clamp_(-2, 2)
    compressed_cache = compressed_cache.to(torch.float8_e4m3fn)
    window_indices = torch.remainder(
        torch.arange(128, dtype=torch.int32, device=device) + 17, 128
    ).unsqueeze(0)
    compressed_tables = torch.arange(
        compressed_table_pages, dtype=torch.int32, device=device
    ).unsqueeze(0)
    hca_seq_lens = torch.tensor([384], dtype=torch.int32, device=device)
    workspace = torch.empty(4 << 20, dtype=torch.uint8, device=device)
    softmax_scale = 1.0 / math.sqrt(head_dim)

    sparse_topk_lens = torch.tensor([384], dtype=torch.int32, device=device)
    window_valid_lens = torch.tensor([128], dtype=torch.int32, device=device)
    sinks = torch.linspace(2.0, 3.0, num_heads, dtype=torch.float32, device=device)
    output = trtllm_batch_decode_sparse_mla_dsv4(
        query=query,
        swa_kv_cache=window_cache,
        workspace_buffer=workspace,
        compressed_kv_cache=compressed_cache,
        sparse_topk_lens=sparse_topk_lens,
        bmm1_scale=softmax_scale,
        sinks=sinks,
        swa_topk_lens=window_valid_lens,
        backend="cute-dsl",
        hca_swa_indices=window_indices,
        hca_compressed_block_tables=compressed_tables,
        hca_seq_lens=hca_seq_lens,
    )
    reference = _reference_hca(
        query,
        window_cache,
        compressed_cache,
        window_indices,
        compressed_tables,
        sparse_topk_lens,
        window_valid_lens,
        sinks,
        softmax_scale,
        1.0,
    )
    torch.testing.assert_close(output.float(), reference, atol=0.13, rtol=1e-5)

    # With no visible KV entry and no attention sink, the defined contribution
    # is O=0 rather than NaN. This also exercises the split_kv=1 epilogue.
    sparse_topk_lens.zero_().add_(128)
    window_valid_lens.zero_()
    output = trtllm_batch_decode_sparse_mla_dsv4(
        query=query,
        swa_kv_cache=window_cache,
        workspace_buffer=workspace,
        compressed_kv_cache=compressed_cache,
        sparse_topk_lens=sparse_topk_lens,
        bmm1_scale=softmax_scale,
        sinks=None,
        swa_topk_lens=window_valid_lens,
        backend="cute-dsl",
        hca_swa_indices=window_indices,
        hca_compressed_block_tables=compressed_tables,
        hca_seq_lens=hca_seq_lens,
    )
    assert torch.isfinite(output).all()
    torch.testing.assert_close(output, torch.zeros_like(output), atol=0, rtol=0)
