"""Paged-attention decode: Triton GQA FlashDecoding + CUDA graphs.

Kernel: one CTA per (batch, kv_head, split); streams KV once for G query heads.
CUDA graph capture on static buffers removes launch overhead (critical under
the L2-flush timing protocol for small decode shapes).
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import triton
import triton.language as tl

OP_TYPE = "attention"
SUPPORTED_PRECISIONS = ["bf16"]
HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]

BATCH = 8
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
SEQ_LEN = 1024
PAGE_SIZE = 16


@triton.jit
def _kernel_nosplit(
    Q_ptr, KV_ptr, BT_ptr, SL_ptr, Out_ptr, sm_scale,
    stride_qb, stride_qh, stride_qd,
    stride_kv_b, stride_kv_p, stride_kv_h, stride_kv_d,
    stride_bt_b, stride_bt_n,
    stride_ob, stride_oh, stride_od,
    page_size,
    HEAD_DIM: tl.constexpr, PAGE_SIZE: tl.constexpr, G: tl.constexpr,
):
    batch_id = tl.program_id(0)
    kv_head = tl.program_id(1)
    seq_len = tl.maximum(tl.load(SL_ptr + batch_id), 0)
    num_blocks = (seq_len + page_size - 1) // page_size
    offs_d = tl.arange(0, HEAD_DIM)
    offs_g = tl.arange(0, G)
    offs_t = tl.arange(0, PAGE_SIZE)
    q_head0 = kv_head * G
    q = tl.load(
        Q_ptr + batch_id * stride_qb
        + (q_head0 + offs_g)[:, None] * stride_qh
        + offs_d[None, :] * stride_qd
    ).to(tl.float32) * sm_scale
    m_i = tl.full([G], float("-inf"), dtype=tl.float32)
    l_i = tl.zeros([G], dtype=tl.float32)
    acc = tl.zeros([G, HEAD_DIM], dtype=tl.float32)
    for logical_block in tl.range(0, num_blocks):
        phys = tl.load(BT_ptr + batch_id * stride_bt_b + logical_block * stride_bt_n)
        mask_t = (logical_block * PAGE_SIZE + offs_t) < seq_len
        base = (
            KV_ptr + phys * stride_kv_b
            + offs_t[:, None] * stride_kv_p
            + kv_head * stride_kv_h
            + offs_d[None, :] * stride_kv_d
        )
        k = tl.load(base, mask=mask_t[:, None], other=0.0).to(tl.float32)
        scores = tl.dot(q, tl.trans(k))
        scores = tl.where(mask_t[None, :], scores, float("-inf"))
        m_ij = tl.max(scores, axis=1)
        m_new = tl.maximum(m_i, m_ij)
        alpha = tl.exp(m_i - m_new)
        p = tl.where(mask_t[None, :], tl.exp(scores - m_new[:, None]), 0.0)
        l_i = l_i * alpha + tl.sum(p, axis=1)
        v = tl.load(base + HEAD_DIM * stride_kv_d, mask=mask_t[:, None], other=0.0).to(tl.float32)
        acc = acc * alpha[:, None] + tl.dot(p, v)
        m_i = m_new
    acc = tl.where(l_i[:, None] > 0, acc / l_i[:, None], 0.0)
    tl.store(
        Out_ptr + batch_id * stride_ob
        + (q_head0 + offs_g)[:, None] * stride_oh
        + offs_d[None, :] * stride_od,
        acc.to(tl.bfloat16),
    )


@triton.jit
def _kernel_split(
    Q_ptr, KV_ptr, BT_ptr, SL_ptr,
    PartialOut_ptr, PartialM_ptr, PartialL_ptr, sm_scale,
    stride_qb, stride_qh, stride_qd,
    stride_kv_b, stride_kv_p, stride_kv_h, stride_kv_d,
    stride_bt_b, stride_bt_n,
    stride_po_s, stride_po_b, stride_po_h, stride_po_d,
    stride_pm_s, stride_pm_b, stride_pm_h,
    page_size, max_blocks,
    HEAD_DIM: tl.constexpr, PAGE_SIZE: tl.constexpr,
    G: tl.constexpr, NUM_SPLITS: tl.constexpr,
):
    batch_id = tl.program_id(0)
    kv_head = tl.program_id(1)
    split_id = tl.program_id(2)
    seq_len = tl.maximum(tl.load(SL_ptr + batch_id), 0)
    num_blocks = (seq_len + page_size - 1) // page_size
    bps = (max_blocks + NUM_SPLITS - 1) // NUM_SPLITS
    block_lo = split_id * bps
    block_hi = tl.minimum(block_lo + bps, num_blocks)
    offs_d = tl.arange(0, HEAD_DIM)
    offs_g = tl.arange(0, G)
    offs_t = tl.arange(0, PAGE_SIZE)
    q_head0 = kv_head * G
    q = tl.load(
        Q_ptr + batch_id * stride_qb
        + (q_head0 + offs_g)[:, None] * stride_qh
        + offs_d[None, :] * stride_qd
    ).to(tl.float32) * sm_scale
    m_i = tl.full([G], float("-inf"), dtype=tl.float32)
    l_i = tl.zeros([G], dtype=tl.float32)
    acc = tl.zeros([G, HEAD_DIM], dtype=tl.float32)
    for logical_block in tl.range(block_lo, block_hi):
        phys = tl.load(BT_ptr + batch_id * stride_bt_b + logical_block * stride_bt_n)
        mask_t = (logical_block * PAGE_SIZE + offs_t) < seq_len
        base = (
            KV_ptr + phys * stride_kv_b
            + offs_t[:, None] * stride_kv_p
            + kv_head * stride_kv_h
            + offs_d[None, :] * stride_kv_d
        )
        k = tl.load(base, mask=mask_t[:, None], other=0.0).to(tl.float32)
        scores = tl.dot(q, tl.trans(k))
        scores = tl.where(mask_t[None, :], scores, float("-inf"))
        m_ij = tl.max(scores, axis=1)
        m_new = tl.maximum(m_i, m_ij)
        alpha = tl.exp(m_i - m_new)
        p = tl.where(mask_t[None, :], tl.exp(scores - m_new[:, None]), 0.0)
        l_i = l_i * alpha + tl.sum(p, axis=1)
        v = tl.load(base + HEAD_DIM * stride_kv_d, mask=mask_t[:, None], other=0.0).to(tl.float32)
        acc = acc * alpha[:, None] + tl.dot(p, v)
        m_i = m_new
    tl.store(
        PartialOut_ptr + split_id * stride_po_s + batch_id * stride_po_b
        + (q_head0 + offs_g)[:, None] * stride_po_h + offs_d[None, :] * stride_po_d,
        acc,
    )
    tl.store(
        PartialM_ptr + split_id * stride_pm_s + batch_id * stride_pm_b
        + (q_head0 + offs_g) * stride_pm_h,
        m_i,
    )
    tl.store(
        PartialL_ptr + split_id * stride_pm_s + batch_id * stride_pm_b
        + (q_head0 + offs_g) * stride_pm_h,
        l_i,
    )


@triton.jit
def _kernel_reduce(
    PartialOut_ptr, PartialM_ptr, PartialL_ptr, Out_ptr,
    stride_po_s, stride_po_b, stride_po_h, stride_po_d,
    stride_pm_s, stride_pm_b, stride_pm_h,
    stride_ob, stride_oh, stride_od,
    HEAD_DIM: tl.constexpr, NUM_SPLITS: tl.constexpr,
):
    batch_id = tl.program_id(0)
    head_id = tl.program_id(1)
    offs_d = tl.arange(0, HEAD_DIM)
    m = tl.full([], float("-inf"), dtype=tl.float32)
    for s in tl.static_range(0, NUM_SPLITS):
        m_s = tl.load(PartialM_ptr + s * stride_pm_s + batch_id * stride_pm_b + head_id * stride_pm_h)
        m = tl.maximum(m, m_s)
    l = 0.0
    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
    for s in tl.static_range(0, NUM_SPLITS):
        m_s = tl.load(PartialM_ptr + s * stride_pm_s + batch_id * stride_pm_b + head_id * stride_pm_h)
        l_s = tl.load(PartialL_ptr + s * stride_pm_s + batch_id * stride_pm_b + head_id * stride_pm_h)
        scale = tl.where(m_s > float("-inf"), tl.exp(m_s - m), 0.0)
        l += l_s * scale
        o_s = tl.load(
            PartialOut_ptr + s * stride_po_s + batch_id * stride_po_b
            + head_id * stride_po_h + offs_d * stride_po_d
        )
        acc += o_s * scale
    acc = tl.where(l > 0, acc / l, 0.0)
    tl.store(
        Out_ptr + batch_id * stride_ob + head_id * stride_oh + offs_d * stride_od,
        acc.to(tl.bfloat16),
    )


def _choose_num_splits(
    batch: int, num_kv_heads: int, seq_len: int, page_size: int, head_dim: int = 128
) -> int:
    sms = torch.cuda.get_device_properties(0).multi_processor_count
    base = batch * num_kv_heads
    pages = (seq_len + page_size - 1) // page_size
    if pages <= 2:
        return 1
    # Small head_dim has less work per page — need more CTAs for latency hiding
    mult = 6 if head_dim <= 64 else 4
    target = sms * mult
    if base >= target:
        return 1
    splits = (target + base - 1) // base
    splits = min(splits, max(1, pages // 2), 32)
    return max(1, int(splits))


_partial_cache: dict = {}


def _get_partials(num_splits: int, B: int, H: int, D: int, device):
    key = (num_splits, B, H, D, str(device))
    buf = _partial_cache.get(key)
    if buf is None:
        buf = (
            torch.empty(num_splits, B, H, D, device=device, dtype=torch.float32),
            torch.empty(num_splits, B, H, device=device, dtype=torch.float32),
            torch.empty(num_splits, B, H, device=device, dtype=torch.float32),
        )
        _partial_cache[key] = buf
    return buf


def paged_attention_decode(
    query: torch.Tensor,
    kv_cache: torch.Tensor,
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    num_kv_heads: int,
    page_size: int,
    sm_scale: float,
    out: torch.Tensor | None = None,
    seq_len_hint: int | None = None,
) -> torch.Tensor:
    """Core launch. If `out` is provided, write into it (for graph capture).

    `seq_len_hint` avoids host sync on seq_lens (required inside CUDA graphs).
    """
    B, H, D = query.shape
    G = H // num_kv_heads
    Hkv = num_kv_heads
    max_blocks = block_table.shape[1]
    if out is None:
        out = torch.empty(B, H, D, device=query.device, dtype=query.dtype)

    sqb, sqh, sqd = query.stride()
    skvb, skvp, skvh, skvd = kv_cache.stride()
    sbtb, sbtn = block_table.stride()
    sob, soh, sod = out.stride()

    if seq_len_hint is not None:
        max_seq = seq_len_hint
    else:
        max_seq = int(seq_lens.max().item()) if seq_lens.numel() else 0
    num_splits = _choose_num_splits(B, Hkv, max(max_seq, 1), page_size, D)
    num_splits = min(num_splits, max_blocks)
    num_splits = max(1, num_splits)

    # Tuned on SM120 with CUDA graphs + L2-flush timing
    num_warps, num_stages = 1, 3

    if num_splits == 1:
        _kernel_nosplit[(B, Hkv)](
            query, kv_cache, block_table, seq_lens, out, sm_scale,
            sqb, sqh, sqd, skvb, skvp, skvh, skvd, sbtb, sbtn, sob, soh, sod,
            page_size, HEAD_DIM=D, PAGE_SIZE=page_size, G=G,
            num_warps=num_warps, num_stages=num_stages,
        )
        return out

    partial_out, partial_m, partial_l = _get_partials(num_splits, B, H, D, query.device)
    spo_s, spo_b, spo_h, spo_d = partial_out.stride()
    spm_s, spm_b, spm_h = partial_m.stride()

    _kernel_split[(B, Hkv, num_splits)](
        query, kv_cache, block_table, seq_lens,
        partial_out, partial_m, partial_l, sm_scale,
        sqb, sqh, sqd, skvb, skvp, skvh, skvd, sbtb, sbtn,
        spo_s, spo_b, spo_h, spo_d, spm_s, spm_b, spm_h,
        page_size, max_blocks,
        HEAD_DIM=D, PAGE_SIZE=page_size, G=G, NUM_SPLITS=num_splits,
        num_warps=num_warps, num_stages=num_stages,
    )
    _kernel_reduce[(B, H)](
        partial_out, partial_m, partial_l, out,
        spo_s, spo_b, spo_h, spo_d, spm_s, spm_b, spm_h, sob, soh, sod,
        HEAD_DIM=D, NUM_SPLITS=num_splits, num_warps=4,
    )
    return out


class Model(nn.Module):
    def __init__(self, batch, num_heads, num_kv_heads, head_dim, seq_len, page_size):
        super().__init__()
        assert num_heads % num_kv_heads == 0
        self.batch = batch
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.seq_len = seq_len
        self.page_size = page_size
        self.group_size = num_heads // num_kv_heads
        self.scale = 1.0 / math.sqrt(head_dim)
        self.register_buffer("_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False)

        # CUDA graph: capture on the live input tensors (no extra KV copy).
        # Recaptures when tensor storage addresses change (new inputs).
        self._graph = None
        self._graph_key = None
        self._graph_out = None

    def forward(self, query, kv_cache, block_table, seq_lens):
        key = (
            query.data_ptr(),
            kv_cache.data_ptr(),
            block_table.data_ptr(),
            seq_lens.data_ptr(),
            query.shape,
            kv_cache.shape,
            block_table.shape,
        )
        try:
            if self._graph is None or self._graph_key != key:
                if (
                    self._graph_out is None
                    or self._graph_out.shape != query.shape
                    or self._graph_out.device != query.device
                ):
                    self._graph_out = torch.empty_like(query)
                # Warmup / compile before capture
                for _ in range(3):
                    paged_attention_decode(
                        query, kv_cache, block_table, seq_lens,
                        self.num_kv_heads, self.page_size, self.scale,
                        out=self._graph_out, seq_len_hint=self.seq_len,
                    )
                torch.cuda.synchronize()
                g = torch.cuda.CUDAGraph()
                with torch.cuda.graph(g):
                    paged_attention_decode(
                        query, kv_cache, block_table, seq_lens,
                        self.num_kv_heads, self.page_size, self.scale,
                        out=self._graph_out, seq_len_hint=self.seq_len,
                    )
                self._graph = g
                self._graph_key = key
            self._graph.replay()
            return self._graph_out
        except Exception:
            return paged_attention_decode(
                query, kv_cache, block_table, seq_lens,
                self.num_kv_heads, self.page_size, self.scale,
                seq_len_hint=self.seq_len,
            )


def get_inputs():
    B, H, Hkv, D, L, P = BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE
    pages_per_seq = (L + P - 1) // P
    total_pages = max(B * pages_per_seq + 8, 64)
    query = torch.randn(B, H, D, dtype=torch.bfloat16) * 0.1
    kv_cache = torch.randn(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16) * 0.1
    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
    return [query, kv_cache, perm.contiguous(), torch.full((B,), L, dtype=torch.int32)]


def get_init_inputs():
    return [BATCH, NUM_HEADS, NUM_KV_HEADS, HEAD_DIM, SEQ_LEN, PAGE_SIZE]