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

from __future__ import annotations

import math
from typing import Optional, Union

import torch

from vllm.triton_utils import HAS_TRITON

if HAS_TRITON:
    import triton
    import triton.language as tl


if HAS_TRITON:

    @triton.autotune(
        configs=[
            triton.Config(
                {"BLOCK_Q": bq, "BLOCK_K": bk},
                num_warps=num_warps,
                num_stages=num_stages,
            ) for bq in [32, 64] for bk in [32, 64] for num_warps in [4, 8]
            for num_stages in [3, 4]
        ],
        key=["QUERY_GROUP_SIZE", "D", "ROWS_MAX"],
    )
    @triton.jit
    def _lse_and_store_logits_kernel(
        Q,
        K,
        cu_q,
        cu_k,
        w_b,
        out_m,
        out_S,
        LOGITS,
        sm_scale,
        QUERY_GROUP_SIZE: tl.constexpr,
        D: tl.constexpr,
        STRIDE_Q_NQ,
        STRIDE_Q_HQ,
        STRIDE_K_NK,
        STRIDE_K_HK,
        STRIDE_M_B,
        STRIDE_M_H,
        STRIDE_M_R,
        STRIDE_S_B,
        STRIDE_S_H,
        STRIDE_S_R,
        STRIDE_LG_NK,
        STRIDE_LG_HK,
        STRIDE_LG_R,
        BLOCK_Q: tl.constexpr,
        BLOCK_K: tl.constexpr,
        ROWS_MAX,
    ):
        b = tl.program_id(0)
        hk = tl.program_id(1)
        rid = tl.program_id(2)

        q_end = tl.load(cu_q + b + 1)
        k_beg = tl.load(cu_k + b)
        k_end = tl.load(cu_k + b + 1)
        win = tl.load(w_b + b)

        q_win_beg = q_end - win
        k_eff_end = k_end - win
        if (win <= 0) or (k_eff_end <= k_beg):
            return

        rows_b = win * QUERY_GROUP_SIZE
        row0 = rid * BLOCK_Q
        if row0 >= rows_b:
            return

        qk_scale = sm_scale * 1.4426950408889634  # exp -> exp2.

        offs_qrow = row0 + tl.arange(0, BLOCK_Q)
        row_mask = offs_qrow < rows_b

        hq_local = offs_qrow % QUERY_GROUP_SIZE
        q_off = offs_qrow // QUERY_GROUP_SIZE
        q_idx = q_win_beg + q_off
        hq_glob = hk * QUERY_GROUP_SIZE + hq_local

        offs_d = tl.arange(0, D)
        q_ptrs = (Q + q_idx[:, None] * STRIDE_Q_NQ +
                  hq_glob[:, None] * STRIDE_Q_HQ + offs_d[None, :])
        q_rows = tl.load(q_ptrs, mask=row_mask[:, None], other=0.0)

        m = tl.zeros([BLOCK_Q], dtype=tl.float32) + (-float("inf"))
        S = tl.zeros([BLOCK_Q], dtype=tl.float32)

        for ks in tl.range(k_beg, k_eff_end, BLOCK_K):
            nk = ks + tl.arange(0, BLOCK_K)
            kmask = nk < k_eff_end

            k_ptrs = (K + nk[:, None] * STRIDE_K_NK +
                      hk * STRIDE_K_HK + offs_d[None, :])
            k_blk = tl.load(k_ptrs, mask=kmask[:, None], other=0.0)

            s = tl.dot(q_rows, k_blk.T) * qk_scale
            s = tl.where(kmask[None, :], s, -float("inf"))

            log_ptrs = (LOGITS + nk[:, None] * STRIDE_LG_NK +
                        hk * STRIDE_LG_HK +
                        (row0 + tl.arange(0, BLOCK_Q))[None, :] * STRIDE_LG_R)
            tl.store(log_ptrs, s.T, mask=kmask[:, None] & row_mask[None, :])

            cur_max = tl.max(s, 1)
            n_m = tl.maximum(m, cur_max)
            rescale = tl.math.exp2(m - n_m)
            S = S * rescale + tl.sum(tl.math.exp2(s - n_m[:, None]), 1)
            m = n_m

        m_base = out_m + b * STRIDE_M_B + hk * STRIDE_M_H + row0 * STRIDE_M_R
        S_base = out_S + b * STRIDE_S_B + hk * STRIDE_S_H + row0 * STRIDE_S_R
        tl.store(m_base + tl.arange(0, BLOCK_Q) * STRIDE_M_R, m, mask=row_mask)
        tl.store(S_base + tl.arange(0, BLOCK_Q) * STRIDE_S_R, S, mask=row_mask)

    @triton.jit
    def _lse_and_store_logits_kernel_rocm_safe(
        Q,
        K,
        cu_q,
        cu_k,
        w_b,
        out_m,
        out_S,
        LOGITS,
        sm_scale,
        QUERY_GROUP_SIZE: tl.constexpr,
        D: tl.constexpr,
        STRIDE_Q_NQ,
        STRIDE_Q_HQ,
        STRIDE_K_NK,
        STRIDE_K_HK,
        STRIDE_M_B,
        STRIDE_M_H,
        STRIDE_M_R,
        STRIDE_S_B,
        STRIDE_S_H,
        STRIDE_S_R,
        STRIDE_LG_NK,
        STRIDE_LG_HK,
        STRIDE_LG_R,
        BLOCK_Q: tl.constexpr,
        BLOCK_K: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        """ROCm-safe variant of `_lse_and_store_logits_kernel`.

        On some ROCm + Triton (HIP) stacks we have observed memory corruption
        from the tl.dot-based implementation. This variant avoids `tl.dot` and
        instead computes dot-products via explicit outer-product accumulation.
        """
        b = tl.program_id(0)
        hk = tl.program_id(1)
        rid = tl.program_id(2)

        q_end = tl.load(cu_q + b + 1)
        k_beg = tl.load(cu_k + b)
        k_end = tl.load(cu_k + b + 1)
        win = tl.load(w_b + b)

        q_win_beg = q_end - win
        k_eff_end = k_end - win
        if (win <= 0) or (k_eff_end <= k_beg):
            return

        rows_b = win * QUERY_GROUP_SIZE
        row0 = rid * BLOCK_Q
        if row0 >= rows_b:
            return

        qk_scale = sm_scale * 1.4426950408889634  # exp -> exp2.

        offs_qrow = row0 + tl.arange(0, BLOCK_Q)
        row_mask = offs_qrow < rows_b

        hq_local = offs_qrow % QUERY_GROUP_SIZE
        q_off = offs_qrow // QUERY_GROUP_SIZE
        q_idx = q_win_beg + q_off
        hq_glob = hk * QUERY_GROUP_SIZE + hq_local

        m = tl.zeros([BLOCK_Q], dtype=tl.float32) + (-float("inf"))
        S = tl.zeros([BLOCK_Q], dtype=tl.float32)

        for ks in tl.range(k_beg, k_eff_end, BLOCK_K):
            nk = ks + tl.arange(0, BLOCK_K)
            kmask = nk < k_eff_end

            # Accumulate s = Q @ K^T in fp32 via outer products.
            s = tl.zeros([BLOCK_Q, BLOCK_K], dtype=tl.float32)

            for ds in tl.static_range(0, D, BLOCK_D):
                offs_d = ds + tl.arange(0, BLOCK_D)
                dmask = offs_d < D

                q_ptrs = (Q + q_idx[:, None] * STRIDE_Q_NQ +
                          hq_glob[:, None] * STRIDE_Q_HQ + offs_d[None, :])
                q_chunk = tl.load(
                    q_ptrs,
                    mask=row_mask[:, None] & dmask[None, :],
                    other=0.0,
                ).to(tl.float32)  # [BQ, BD]

                k_ptrs = (K + nk[:, None] * STRIDE_K_NK +
                          hk * STRIDE_K_HK + offs_d[None, :])
                k_chunk = tl.load(
                    k_ptrs,
                    mask=kmask[:, None] & dmask[None, :],
                    other=0.0,
                ).to(tl.float32)  # [BK, BD]

                s += tl.sum(q_chunk[:, None, :] * k_chunk[None, :, :], axis=2)

            s = s * qk_scale
            s = tl.where(kmask[None, :], s, -float("inf"))

            log_ptrs = (LOGITS + nk[:, None] * STRIDE_LG_NK +
                        hk * STRIDE_LG_HK +
                        (row0 + tl.arange(0, BLOCK_Q))[None, :] * STRIDE_LG_R)
            tl.store(log_ptrs, s.T, mask=kmask[:, None] & row_mask[None, :])

            cur_max = tl.max(s, 1)
            n_m = tl.maximum(m, cur_max)
            rescale = tl.math.exp2(m - n_m)
            S = S * rescale + tl.sum(tl.math.exp2(s - n_m[:, None]), 1)
            m = n_m

        m_base = out_m + b * STRIDE_M_B + hk * STRIDE_M_H + row0 * STRIDE_M_R
        S_base = out_S + b * STRIDE_S_B + hk * STRIDE_S_H + row0 * STRIDE_S_R
        tl.store(m_base + tl.arange(0, BLOCK_Q) * STRIDE_M_R,
                 m,
                 mask=row_mask)
        tl.store(S_base + tl.arange(0, BLOCK_Q) * STRIDE_S_R,
                 S,
                 mask=row_mask)

    @triton.autotune(
        configs=[
            triton.Config({"BLOCK_Q": bq, "BLOCK_K": bk})
            for bq in [16, 32, 64] for bk in [32, 64, 128]
        ],
        key=["HK", "HQ"],
    )
    @triton.jit
    def _scores_from_logits_kernel(
        cu_k,
        w_b,
        in_m,
        in_S,
        LOGITS,
        OUT,
        QUERY_GROUP_SIZE: tl.constexpr,
        STRIDE_M_B,
        STRIDE_M_H,
        STRIDE_M_R,
        STRIDE_S_B,
        STRIDE_S_H,
        STRIDE_S_R,
        STRIDE_LG_NK,
        STRIDE_LG_HK,
        STRIDE_LG_R,
        STRIDE_OUT_NK,
        STRIDE_OUT_HK,
        BLOCK_Q: tl.constexpr,
        BLOCK_K: tl.constexpr,
        DO_POOL: tl.constexpr,
        KPOOL: tl.constexpr,
        PROTECT_LAST: tl.constexpr,
    ):
        b = tl.program_id(0)
        hk = tl.program_id(1)

        k_beg = tl.load(cu_k + b)
        k_end = tl.load(cu_k + b + 1)
        win = tl.load(w_b + b)

        k_eff_end = k_end - win
        if (win <= 0) or (k_eff_end <= k_beg):
            return

        rows_b = win * QUERY_GROUP_SIZE

        for ks in tl.range(k_beg, k_eff_end, BLOCK_K):
            nk = ks + tl.arange(0, BLOCK_K)
            kmask = nk < k_eff_end

            scores = tl.zeros([BLOCK_K], dtype=tl.float32)
            for row0 in tl.range(0, rows_b, BLOCK_Q):
                r_idx = row0 + tl.arange(0, BLOCK_Q)
                rmask = r_idx < rows_b

                m_ptr = (in_m + b * STRIDE_M_B + hk * STRIDE_M_H +
                         row0 * STRIDE_M_R)
                S_ptr = (in_S + b * STRIDE_S_B + hk * STRIDE_S_H +
                         row0 * STRIDE_S_R)
                m = tl.load(m_ptr + tl.arange(0, BLOCK_Q) * STRIDE_M_R,
                            mask=rmask,
                            other=-float("inf"))
                S = tl.load(S_ptr + tl.arange(0, BLOCK_Q) * STRIDE_S_R,
                            mask=rmask,
                            other=0.0)

                valid_row = S > 0
                m = tl.where(valid_row, m, 0.0)
                S = tl.where(valid_row, S, 1.0)

                log_ptrs = (LOGITS + nk[:, None] * STRIDE_LG_NK +
                            hk * STRIDE_LG_HK +
                            (row0 + tl.arange(0, BLOCK_Q))[None, :] *
                            STRIDE_LG_R)
                s_T = tl.load(log_ptrs,
                              mask=kmask[:, None] & rmask[None, :],
                              other=-float("inf"))

                probs_T = tl.math.exp2(s_T - m[None, :]) / S[None, :]
                probs_T = tl.where(valid_row[None, :], probs_T, 0.0)
                scores += tl.sum(probs_T, 1)

            if DO_POOL and (KPOOL > 1):
                i = tl.arange(0, BLOCK_K)[:, None]
                j = tl.arange(0, BLOCK_K)[None, :]
                band = (j <= i) & ((i - j) < KPOOL)
                band = band & kmask[None, :]
                sums = tl.sum(tl.where(band, scores[None, :], 0.0), 1)
                denom = tl.sum(band, 1).to(tl.float32)
                denom = tl.where(denom > 0, denom, 1.0)
                scores = sums / denom

            out_ptrs = OUT + nk * STRIDE_OUT_NK + hk * STRIDE_OUT_HK
            tl.store(out_ptrs, scores, mask=kmask)

        if PROTECT_LAST:
            pad_beg = k_eff_end
            pad_end = k_end
            if pad_end > pad_beg:
                for ks in tl.range(pad_beg, pad_end, BLOCK_K):
                    nk = ks + tl.arange(0, BLOCK_K)
                    kmask = nk < pad_end
                    out_ptrs = OUT + nk * STRIDE_OUT_NK + hk * STRIDE_OUT_HK
                    tl.store(out_ptrs,
                             tl.full([BLOCK_K], float("inf"), dtype=tl.float32),
                             mask=kmask)

    @triton.autotune(
        configs=[triton.Config({"BLOCK_K": bk}) for bk in [32, 64, 128]],
        key=["HK"],
    )
    @triton.jit
    def _zscore_per_batch_epilogue(
        OUT,
        cu_k,
        w_b,
        STRIDE_OUT_NK,
        STRIDE_OUT_HK,
        HK: tl.constexpr,
        EPS: tl.constexpr,
        BLOCK_K: tl.constexpr,
    ):
        b = tl.program_id(0)

        k_beg = tl.load(cu_k + b)
        k_end = tl.load(cu_k + b + 1)
        win = tl.load(w_b + b)

        k_eff_end = k_end - win
        if k_eff_end <= k_beg:
            return

        sumv = tl.zeros([], dtype=tl.float32)
        sumsq = tl.zeros([], dtype=tl.float32)
        count = ((k_eff_end - k_beg) * HK).to(tl.float32)

        for ks in tl.range(k_beg, k_eff_end, BLOCK_K):
            nk = ks + tl.arange(0, BLOCK_K)
            kmask = nk < k_eff_end
            for h in tl.range(0, HK):
                ptrs = OUT + nk * STRIDE_OUT_NK + h * STRIDE_OUT_HK
                vals = tl.load(ptrs, mask=kmask, other=0.0).to(tl.float32)
                sumv += tl.sum(vals, 0)
                sumsq += tl.sum(vals * vals, 0)

        mean = sumv / count
        var = tl.maximum(sumsq / count - mean * mean, 0.0)
        invstd = 1.0 / tl.sqrt(var + EPS)

        for ks in tl.range(k_beg, k_eff_end, BLOCK_K):
            nk = ks + tl.arange(0, BLOCK_K)
            kmask = nk < k_eff_end
            for h in tl.range(0, HK):
                ptrs = OUT + nk * STRIDE_OUT_NK + h * STRIDE_OUT_HK
                vals = tl.load(ptrs, mask=kmask, other=0.0).to(tl.float32)
                vals = (vals - mean) * invstd
                tl.store(ptrs, vals, mask=kmask)


def query_aware_key_scores(
    q: torch.Tensor,  # [N_q, Hq, D]
    k: torch.Tensor,  # [N_k, Hk, D]
    cu_seqlens_q: torch.Tensor,  # [B+1] int32
    cu_seqlens_k: torch.Tensor,  # [B+1] int32
    w: Union[int, torch.Tensor],  # [B] int32 or scalar
    sm_scale: Optional[float] = None,
    *,
    pool: bool = True,
    kpool: int = 5,
    protect_last: bool = True,
    normalize: bool = False,
) -> torch.Tensor:
    """SnapKV query-aware key scores (Triton), returns [N_k, Hk] float32."""
    if not HAS_TRITON:
        raise RuntimeError("Triton is not available.")
    if q.device.type != "cuda" or k.device.type != "cuda":
        raise RuntimeError("Triton SnapKV requires CUDA/ROCm tensors.")
    if q.ndim != 3 or k.ndim != 3:
        raise ValueError("q and k must be 3D tensors.")
    if q.stride(-1) != 1 or k.stride(-1) != 1:
        raise ValueError("Last dim must be contiguous for Triton SnapKV.")

    device = q.device
    N_q, Hq, D = q.shape
    N_k, Hk, Dk = k.shape
    if D != Dk:
        raise ValueError("q and k must have the same head size.")
    if (Hq % Hk) != 0:
        raise ValueError("Hq must be a multiple of Hk.")
    if sm_scale is None:
        sm_scale = 1.0 / math.sqrt(D)

    B = int(cu_seqlens_q.numel() - 1)
    if B != int(cu_seqlens_k.numel() - 1):
        raise ValueError("cu_seqlens_q and cu_seqlens_k must match.")

    G = Hq // Hk
    if isinstance(w, int):
        max_w = int(w)
        w = torch.full((B, ),
                       fill_value=max_w,
                       device=device,
                       dtype=torch.int32)
    else:
        if w.numel() != B:
            raise ValueError("w must have shape [B].")
        w = w.to(device=device, dtype=torch.int32)
        max_w = int(w.max().item())

    rows_max = max_w * G
    if rows_max <= 0:
        return torch.zeros((N_k, Hk), dtype=torch.float32, device=device)

    if kpool < 1:
        raise ValueError("kpool must be >= 1.")

    out = torch.zeros((N_k, Hk), dtype=torch.float32, device=device)
    m_scratch = torch.empty((B, Hk, rows_max),
                            dtype=torch.float32,
                            device=device)
    S_scratch = torch.empty((B, Hk, rows_max),
                            dtype=torch.float32,
                            device=device)
    logits_buf = torch.empty((N_k, Hk, rows_max),
                             dtype=torch.float32,
                             device=device)

    STRIDE_Q_NQ, STRIDE_Q_HQ, _ = q.stride()
    STRIDE_K_NK, STRIDE_K_HK, _ = k.stride()
    STRIDE_M_B, STRIDE_M_H, STRIDE_M_R = m_scratch.stride()
    STRIDE_S_B, STRIDE_S_H, STRIDE_S_R = S_scratch.stride()
    STRIDE_LG_NK, STRIDE_LG_HK, STRIDE_LG_R = logits_buf.stride()
    STRIDE_OUT_NK, STRIDE_OUT_HK = out.stride()

    def grid(meta):
        return B, Hk, triton.cdiv(rows_max, meta["BLOCK_Q"])

    cu_q = cu_seqlens_q.to(device=device, dtype=torch.int32)
    cu_k = cu_seqlens_k.to(device=device, dtype=torch.int32)

    # NOTE: On ROCm/HIP, we prefer a dot-free kernel variant to avoid known
    # correctness issues (silent memory corruption) observed with the tl.dot
    # implementation on some stacks.
    is_rocm = getattr(torch.version, "hip", None) is not None
    if is_rocm:
        _lse_and_store_logits_kernel_rocm_safe[grid](
            q,
            k,
            cu_q,
            cu_k,
            w,
            m_scratch,
            S_scratch,
            logits_buf,
            sm_scale,
            QUERY_GROUP_SIZE=G,
            D=D,
            STRIDE_Q_NQ=STRIDE_Q_NQ,
            STRIDE_Q_HQ=STRIDE_Q_HQ,
            STRIDE_K_NK=STRIDE_K_NK,
            STRIDE_K_HK=STRIDE_K_HK,
            STRIDE_M_B=STRIDE_M_B,
            STRIDE_M_H=STRIDE_M_H,
            STRIDE_M_R=STRIDE_M_R,
            STRIDE_S_B=STRIDE_S_B,
            STRIDE_S_H=STRIDE_S_H,
            STRIDE_S_R=STRIDE_S_R,
            STRIDE_LG_NK=STRIDE_LG_NK,
            STRIDE_LG_HK=STRIDE_LG_HK,
            STRIDE_LG_R=STRIDE_LG_R,
            BLOCK_Q=32,
            BLOCK_K=32,
            BLOCK_D=16,
            num_warps=4,
            num_stages=1,
        )
    else:
        _lse_and_store_logits_kernel[grid](
            q,
            k,
            cu_q,
            cu_k,
            w,
            m_scratch,
            S_scratch,
            logits_buf,
            sm_scale,
            QUERY_GROUP_SIZE=G,
            D=D,
            STRIDE_Q_NQ=STRIDE_Q_NQ,
            STRIDE_Q_HQ=STRIDE_Q_HQ,
            STRIDE_K_NK=STRIDE_K_NK,
            STRIDE_K_HK=STRIDE_K_HK,
            STRIDE_M_B=STRIDE_M_B,
            STRIDE_M_H=STRIDE_M_H,
            STRIDE_M_R=STRIDE_M_R,
            STRIDE_S_B=STRIDE_S_B,
            STRIDE_S_H=STRIDE_S_H,
            STRIDE_S_R=STRIDE_S_R,
            STRIDE_LG_NK=STRIDE_LG_NK,
            STRIDE_LG_HK=STRIDE_LG_HK,
            STRIDE_LG_R=STRIDE_LG_R,
            ROWS_MAX=rows_max,
        )

    _scores_from_logits_kernel[(B, Hk)](
        cu_k,
        w,
        m_scratch,
        S_scratch,
        logits_buf,
        out,
        QUERY_GROUP_SIZE=G,
        STRIDE_M_B=STRIDE_M_B,
        STRIDE_M_H=STRIDE_M_H,
        STRIDE_M_R=STRIDE_M_R,
        STRIDE_S_B=STRIDE_S_B,
        STRIDE_S_H=STRIDE_S_H,
        STRIDE_S_R=STRIDE_S_R,
        STRIDE_LG_NK=STRIDE_LG_NK,
        STRIDE_LG_HK=STRIDE_LG_HK,
        STRIDE_LG_R=STRIDE_LG_R,
        STRIDE_OUT_NK=STRIDE_OUT_NK,
        STRIDE_OUT_HK=STRIDE_OUT_HK,
        DO_POOL=pool,
        KPOOL=kpool,
        PROTECT_LAST=protect_last,
    )

    if normalize:
        _zscore_per_batch_epilogue[(B, )](
            out,
            cu_k,
            w,
            STRIDE_OUT_NK,
            STRIDE_OUT_HK,
            HK=Hk,
            EPS=1e-12,
        )

    return out
