import unittest

import torch

from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
from sglang.kernels.ops.attention.fla.fused_recurrent import (
    fused_recurrent_kda_packed_decode,
)
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
    fused_sigmoid_gating_delta_rule_update,
)
from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
from sglang.kernels.ops.attention.fla.kda import (
    chunk_kda,
    fused_recurrent_kda,
    kda_gate_chunk_cumsum,
)
from sglang.srt.utils.common import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase

register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=12, stage="stage-b", runner_config="1-gpu-large-amd")


@unittest.skipIf(
    not (torch.cuda.is_available() or torch.xpu.is_available()),
    "Test requires CUDA or XPU",
)
class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):
    def setUp(self):
        self.device = get_device()
        self.token_num = 4
        self.query_start_loc = torch.tensor([0, 1, 2, 3, 4], device=self.device)
        self.cache_indices = torch.tensor([0, 2, 5, 8], device=self.device)
        self.local_num_heads = 8
        self.head_dim = 128
        self.cache_len = 64

        self.A_log = torch.randn(
            1, 1, self.local_num_heads, 1, dtype=torch.float32, device=self.device
        )
        self.a = torch.randn(
            1,
            self.token_num,
            self.local_num_heads * self.head_dim,
            dtype=torch.bfloat16,
            device=self.device,
        )
        self.dt_bias = torch.randn(
            self.local_num_heads * self.head_dim,
            dtype=torch.bfloat16,
            device=self.device,
        )
        self.softplus_beta = 1.0
        self.softplus_threshold = 20.0
        self.q = torch.randn(
            1,
            self.token_num,
            self.local_num_heads,
            self.head_dim,
            dtype=torch.bfloat16,
            device=self.device,
        )
        self.k = torch.randn(
            1,
            self.token_num,
            self.local_num_heads,
            self.head_dim,
            dtype=torch.bfloat16,
            device=self.device,
        )
        self.v = torch.randn(
            1,
            self.token_num,
            self.local_num_heads,
            self.head_dim,
            dtype=torch.bfloat16,
            device=self.device,
        )
        self.beta = torch.randn(
            1,
            self.token_num,
            self.local_num_heads,
            dtype=torch.bfloat16,
            device=self.device,
        )

        self.ssm_states = torch.zeros(
            self.cache_len,
            self.local_num_heads,
            self.head_dim,
            self.head_dim,
            dtype=torch.float32,
            device=self.device,
        )

    def run_fused(self):
        ssm_states = self.ssm_states.clone()
        core_attn_out = fused_sigmoid_gating_delta_rule_update(
            A_log=self.A_log,
            dt_bias=self.dt_bias,
            q=self.q,
            k=self.k,
            v=self.v,
            a=self.a,
            b=self.beta,
            initial_state_source=ssm_states,
            initial_state_indices=self.cache_indices,
            cu_seqlens=self.query_start_loc,
            use_qk_l2norm_in_kernel=True,
            softplus_beta=self.softplus_beta,
            softplus_threshold=self.softplus_threshold,
            is_kda=True,
        )
        return core_attn_out, ssm_states[self.cache_indices]

    def run_kda(self):
        b = self.beta.float().sigmoid()
        # Reference gate activation using torch ops:
        #   g = -exp(A_log) * softplus(raw_g + dt_bias)
        H, K = self.local_num_heads, self.head_dim
        raw_g = self.a.float()  # [1, T, H*K]
        if self.dt_bias is not None:
            raw_g = raw_g + self.dt_bias.float()
        g = -torch.exp(
            self.A_log.float().view(1, 1, H, 1)
        ) * torch.nn.functional.softplus(raw_g.view(1, -1, H, K))
        initial_state = self.ssm_states[self.cache_indices].clone()
        core_attn_out, last_state = fused_recurrent_kda(
            q=self.q,
            k=self.k,
            v=self.v,
            g=g,
            beta=b,
            initial_state=initial_state,
            use_qk_l2norm_in_kernel=True,
            cu_seqlens=self.query_start_loc,
        )
        return core_attn_out, last_state

    def test_kda_fused_sigmoid_gating_recurrent(self):
        core_attn_out, last_state = self.run_fused()
        core_attn_out_ref, last_state_ref = self.run_kda()
        abs_diff_out = (core_attn_out - core_attn_out_ref).abs().max()
        abs_diff_state = (last_state - last_state_ref).abs().max()
        print(f"{abs_diff_out=}, {abs_diff_state=}")
        self.assertTrue(
            torch.allclose(core_attn_out, core_attn_out_ref, rtol=1e-3, atol=1e-4)
        )
        self.assertTrue(torch.allclose(last_state, last_state_ref))


@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestKDAGateChunkCumsum(unittest.TestCase):
    """Test kda_gate_chunk_cumsum against torch reference (gate activation + cumsum)."""

    CHUNK_SIZE = 64

    def _ref_gate_cumsum(self, raw_g, A_log, dt_bias, cu_seqlens, chunk_size):
        """Reference: torch gate activation then chunk_local_cumsum."""
        B, T, H, K = raw_g.shape
        g = raw_g.float()
        if dt_bias is not None:
            g = g + dt_bias.float().view(1, 1, H, K)
        g = -torch.exp(A_log.float().view(1, 1, H, 1)) * torch.nn.functional.softplus(g)
        chunk_indices = (
            prepare_chunk_indices(cu_seqlens, chunk_size)
            if cu_seqlens is not None
            else None
        )
        return chunk_local_cumsum(
            g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
        )

    def _run_case(self, B, T_per_seq, H, K, use_bias, use_varlen):
        T = B * T_per_seq
        torch.manual_seed(42)
        raw_g = torch.randn(1, T, H, K, dtype=torch.bfloat16, device="cuda")
        A_log = torch.randn(H, dtype=torch.float32, device="cuda") * 0.5
        dt_bias = (
            torch.randn(H * K, dtype=torch.float32, device="cuda") * 0.1
            if use_bias
            else None
        )
        cu_seqlens = (
            torch.arange(
                0, (B + 1) * T_per_seq, T_per_seq, dtype=torch.long, device="cuda"
            )
            if use_varlen
            else None
        )

        out_fused = kda_gate_chunk_cumsum(
            raw_g,
            A_log=A_log,
            chunk_size=self.CHUNK_SIZE,
            dt_bias=dt_bias,
            cu_seqlens=cu_seqlens,
        )
        out_ref = self._ref_gate_cumsum(
            raw_g, A_log, dt_bias, cu_seqlens, self.CHUNK_SIZE
        )

        max_diff = (out_fused - out_ref).abs().max().item()
        rel_diff = max_diff / (out_ref.abs().mean().item() + 1e-8)
        return max_diff, rel_diff

    def test_varlen_with_bias(self):
        max_diff, rel_diff = self._run_case(
            B=4, T_per_seq=256, H=16, K=128, use_bias=True, use_varlen=True
        )
        self.assertLess(
            max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}"
        )

    def test_varlen_no_bias(self):
        max_diff, rel_diff = self._run_case(
            B=4, T_per_seq=256, H=16, K=128, use_bias=False, use_varlen=True
        )
        self.assertLess(
            max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}"
        )

    def test_fixed_len_with_bias(self):
        max_diff, rel_diff = self._run_case(
            B=4, T_per_seq=256, H=16, K=128, use_bias=True, use_varlen=False
        )
        self.assertLess(
            max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}"
        )

    def test_single_seq_long(self):
        max_diff, rel_diff = self._run_case(
            B=1, T_per_seq=2048, H=16, K=128, use_bias=True, use_varlen=True
        )
        self.assertLess(
            max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}"
        )

    def test_small_head_dim(self):
        max_diff, rel_diff = self._run_case(
            B=4, T_per_seq=128, H=8, K=64, use_bias=True, use_varlen=True
        )
        self.assertLess(
            max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}"
        )


@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestKDAChunkExponentDomain(CustomTestCase):
    """Guard KDA prefill against mixing natural-log gates with exp2 kernels."""

    @staticmethod
    def _naive_recurrent(q, k, v, g, beta, initial_state, lengths):
        q, k, v, g, beta = (tensor.float() for tensor in (q, k, v, g, beta))
        scale = q.shape[-1] ** -0.5
        output = torch.empty_like(v)
        final_state = initial_state.float().clone()

        offset = 0
        for sequence_index, length in enumerate(lengths):
            state = final_state[sequence_index]
            for token_index in range(offset, offset + length):
                state = state * g[0, token_index].exp().unsqueeze(-2)
                residual = v[0, token_index] - torch.einsum(
                    "hvk,hk->hv", state, k[0, token_index]
                )
                state = state + torch.einsum(
                    "hv,hk->hvk",
                    residual * beta[0, token_index, :, None],
                    k[0, token_index],
                )
                output[0, token_index] = (
                    torch.einsum("hvk,hk->hv", state, q[0, token_index]) * scale
                )
            final_state[sequence_index] = state
            offset += length
        return output, final_state

    @staticmethod
    def _relative_rmse(actual, expected):
        error = (actual.float() - expected.float()).square().mean().sqrt()
        baseline = expected.float().square().mean().sqrt().clamp_min(1e-8)
        return (error / baseline).item()

    @torch.inference_mode()
    def test_chunk_prefill_matches_natural_exp_recurrence(self):
        device = get_device()
        dtype = torch.bfloat16
        num_heads, head_dim = 2, 64

        cases = (
            ([129], False, False),
            ([15, 16, 17, 63, 65], True, True),
            # 129 chunks x 2 heads = 258 CTAs > 256 -> _small_grid=False: exercises
            # the standalone (non-fused) diagonal and recompute kernels.
            ([2] * 129, True, False),
        )
        for lengths, use_varlen, fuse_gate in cases:
            with self.subTest(
                lengths=lengths, use_varlen=use_varlen, fuse_gate=fuse_gate
            ):
                torch.manual_seed(42)
                total_tokens = sum(lengths)
                shape = (1, total_tokens, num_heads, head_dim)
                q = torch.nn.functional.normalize(
                    torch.randn(shape, dtype=torch.float32, device=device), dim=-1
                ).to(dtype)
                k = torch.nn.functional.normalize(
                    torch.randn(shape, dtype=torch.float32, device=device), dim=-1
                ).to(dtype)
                v = torch.randn(shape, dtype=dtype, device=device) * 0.1
                raw_gate = (
                    torch.randn(shape, dtype=torch.float32, device=device) * 0.5 - 2.0
                ).to(dtype)
                A_log = torch.randn(num_heads, dtype=torch.float32, device=device) * 0.1
                dt_bias = (
                    torch.randn(
                        num_heads * head_dim,
                        dtype=torch.float32,
                        device=device,
                    )
                    * 0.1
                )
                activated_gate = -torch.exp(
                    A_log.view(1, 1, num_heads, 1)
                ) * torch.nn.functional.softplus(
                    raw_gate.float() + dt_bias.view(1, 1, num_heads, head_dim)
                )
                kernel_gate = raw_gate if fuse_gate else activated_gate.to(dtype)
                reference_gate = activated_gate if fuse_gate else kernel_gate.float()
                beta = torch.rand(
                    1, total_tokens, num_heads, dtype=dtype, device=device
                ).sigmoid()
                initial_state = (
                    torch.randn(
                        len(lengths),
                        num_heads,
                        head_dim,
                        head_dim,
                        dtype=torch.float32,
                        device=device,
                    )
                    * 0.05
                )

                expected_output, expected_state = self._naive_recurrent(
                    q=q,
                    k=k,
                    v=v,
                    g=reference_gate,
                    beta=beta,
                    initial_state=initial_state,
                    lengths=lengths,
                )
                actual_state = initial_state.clone()
                cu_seqlens = None
                if use_varlen:
                    cu_seqlens = torch.tensor(
                        [0, *torch.tensor(lengths).cumsum(0).tolist()],
                        dtype=torch.int32,
                        device=device,
                    )
                actual_output = chunk_kda(
                    q=q.clone(),
                    k=k.clone(),
                    v=v.clone(),
                    g=kernel_gate.clone(),
                    beta=beta.clone(),
                    initial_state=actual_state,
                    initial_state_indices=torch.arange(
                        len(lengths), dtype=torch.int32, device=device
                    ),
                    cu_seqlens=cu_seqlens,
                    A_log=A_log if fuse_gate else None,
                    dt_bias=dt_bias if fuse_gate else None,
                )

                output_error = self._relative_rmse(actual_output, expected_output)
                state_error = self._relative_rmse(actual_state, expected_state)
                self.assertLess(output_error, 1e-2, f"output error={output_error:.3%}")
                self.assertLess(state_error, 1e-2, f"state error={state_error:.3%}")


@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestKDAPackedDecode(unittest.TestCase):
    """Verify ``fused_recurrent_kda_packed_decode`` matches the existing decode
    path (split + unflatten + ``fused_sigmoid_gating_delta_rule_update``)."""

    @staticmethod
    def _make_inputs(B, H, HV, K, V, pool_size, dtype, device, seed=42):
        torch.manual_seed(seed)
        qkv_dim = 2 * H * K + HV * V
        mixed_qkv = (
            torch.randn(B, qkv_dim, dtype=dtype, device=device) * 0.1
        ).contiguous()
        a = (
            torch.randn(B, HV * K, dtype=dtype, device=device) * 0.5 - 1.0
        ).contiguous()
        b = (torch.randn(B, HV, dtype=dtype, device=device) * 0.5).contiguous()
        A_log = torch.randn(HV, dtype=torch.float32, device=device) * 0.2
        dt_bias = torch.randn(HV * K, dtype=torch.float32, device=device) * 0.1
        ssm_states = (
            torch.randn(pool_size, HV, V, K, dtype=dtype, device=device) * 0.01
        ).contiguous()
        cache_indices = torch.arange(B, device=device, dtype=torch.int32)
        return mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices

    @staticmethod
    def _run_baseline(
        mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, H, HV, K, V
    ):
        B = mixed_qkv.shape[0]
        q_flat, k_flat, v_flat = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1)
        q = q_flat.view(1, B, H, K)
        k = k_flat.view(1, B, H, K)
        v = v_flat.view(1, B, HV, V)
        # The real backend passes query_start_loc = [0, 1, ..., B] so that
        # each of the B tokens becomes its own length-1 sequence with an
        # independent state; without this the kernel would share state.
        cu_seqlens = torch.arange(B + 1, device=mixed_qkv.device, dtype=torch.int32)
        return fused_sigmoid_gating_delta_rule_update(
            A_log=A_log,
            dt_bias=dt_bias,
            softplus_beta=1.0,
            softplus_threshold=20.0,
            q=q,
            k=k,
            v=v,
            a=a,
            b=b,
            initial_state_source=ssm_states,
            initial_state_indices=cache_indices,
            cu_seqlens=cu_seqlens,
            scale=K**-0.5,
            use_qk_l2norm_in_kernel=True,
            is_kda=True,
        )

    @staticmethod
    def _run_packed(
        mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, HV, K, V
    ):
        B = mixed_qkv.shape[0]
        out = mixed_qkv.new_empty(B, 1, HV, V)
        fused_recurrent_kda_packed_decode(
            mixed_qkv=mixed_qkv,
            a=a,
            b=b,
            A_log=A_log,
            dt_bias=dt_bias,
            scale=K**-0.5,
            initial_state=ssm_states,
            out=out,
            ssm_state_indices=cache_indices,
            use_qk_l2norm_in_kernel=True,
        )
        return out.transpose(0, 1)

    def _check(self, B, H, HV, K, V):
        device = get_device()
        dtype = torch.bfloat16
        pool_size = B + 4
        mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices = self._make_inputs(
            B, H, HV, K, V, pool_size, dtype, device
        )
        s_packed = ssm_states.clone()
        s_baseline = ssm_states.clone()

        o_packed = self._run_packed(
            mixed_qkv, a, b, A_log, dt_bias, s_packed, cache_indices, HV, K, V
        )
        o_baseline = self._run_baseline(
            mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V
        )

        torch.testing.assert_close(
            o_packed.float(), o_baseline.float(), atol=2e-2, rtol=1e-2
        )
        torch.testing.assert_close(
            s_packed[cache_indices].float(),
            s_baseline[cache_indices].float(),
            atol=2e-2,
            rtol=1e-2,
        )

    def test_b1(self):
        self._check(B=1, H=16, HV=16, K=128, V=128)

    def test_b4(self):
        self._check(B=4, H=16, HV=16, K=128, V=128)

    def test_b32(self):
        self._check(B=32, H=16, HV=16, K=128, V=128)

    def test_b128(self):
        self._check(B=128, H=16, HV=16, K=128, V=128)

    def test_asymmetric_heads(self):
        # Common KDA config with HV > H (grouped query).
        self._check(B=8, H=8, HV=16, K=128, V=128)

    def test_pad_slot(self):
        """Entries with state_idx == -1 must produce zero output and skip state writeback."""
        device = get_device()
        dtype = torch.bfloat16
        B, H, HV, K, V = 8, 16, 16, 128, 128
        pool_size = B + 4
        mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices = self._make_inputs(
            B, H, HV, K, V, pool_size, dtype, device
        )
        # Mark every other request as padded.
        cache_indices = cache_indices.clone()
        cache_indices[::2] = -1

        s_packed = ssm_states.clone()
        s_baseline = ssm_states.clone()
        o_packed = self._run_packed(
            mixed_qkv, a, b, A_log, dt_bias, s_packed, cache_indices, HV, K, V
        )
        o_baseline = self._run_baseline(
            mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V
        )
        torch.testing.assert_close(
            o_packed.float(), o_baseline.float(), atol=2e-2, rtol=1e-2
        )

    def test_production_shapes_through_dispatcher(self):
        """Go through ``TritonKDAKernel.packed_decode`` with the exact tensor
        shapes the KimiDeltaAttention model produces at decode time, so the
        a/b/A_log/dt_bias reshape-normalization is unit-tested (not only E2E).

        Production decode shapes (see kimi_linear.py forward + __init__):
          - a (forget_gate): [B, HV*K]   (2D, not unflattened in decode)
          - b (beta):        [1, B, HV]  (unsqueeze(0), pre-sigmoid)
          - A_log:           [1, 1, HV, 1]
          - dt_bias:         [HV*K]
        """
        from sglang.srt.layers.attention.linear.kernels.kda_triton import (
            TritonKDAKernel,
        )

        device = get_device()
        dtype = torch.bfloat16
        B, H, HV, K, V = 4, 16, 16, 128, 128
        pool_size = B + 4
        mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices = self._make_inputs(
            B, H, HV, K, V, pool_size, dtype, device
        )

        # Reshape the flat reference tensors into the production layouts.
        b_prod = b.unsqueeze(0)  # [B, HV] -> [1, B, HV]
        A_log_prod = A_log.view(1, 1, HV, 1)  # [HV] -> [1, 1, HV, 1]

        kernel = TritonKDAKernel()
        self.assertTrue(kernel.supports_packed_decode)

        s_packed = ssm_states.clone()
        out = kernel.packed_decode(
            mixed_qkv,
            a,
            b_prod,
            A_log=A_log_prod,
            dt_bias=dt_bias,
            scale=K**-0.5,
            ssm_states=s_packed,
            cache_indices=cache_indices,
            num_v_heads=HV,
            head_v_dim=V,
        )

        s_baseline = ssm_states.clone()
        o_baseline = self._run_baseline(
            mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V
        )

        # Dispatcher returns [1, B, HV, V], same layout as the baseline.
        torch.testing.assert_close(
            out.float(), o_baseline.float(), atol=2e-2, rtol=1e-2
        )
        torch.testing.assert_close(
            s_packed[cache_indices].float(),
            s_baseline[cache_indices].float(),
            atol=2e-2,
            rtol=1e-2,
        )


if __name__ == "__main__":
    unittest.main()
