"""Triton paged-attention decode kernel for RTX PRO 6000 (SM120 Blackwell).

Single-query decode: one query per batch element attends over a paged KV cache.
KV cache layout (num_blocks, P, Hkv, 2*D) packs [K|V] on the last axis.

Grid = (B, Hkv): one program per (batch, kv_head), streaming that kv_head's K/V
once per token while handling all G = H//Hkv query heads that share it together
(GQA reuse). KV tiles use flash-attention online-softmax, writing straight to
output with zero extra traffic. On the memory-bound decode regime this kernel
saturates ~87% of DRAM bandwidth at a large grid.

On sm90, tl.dot requires M,N multiples of 16 (wgmma tile); G in {4,8} is padded
to the next multiple of 16 in the head dimension of both matmuls; the extra rows
are masked out and sliced off before writing. Decoding is bandwidth-bound, so
that padding overhead is cheap.
"""
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


def _cfg():
    return [triton.Config({"BLOCK_K": bk}) for bk in (32, 64)]


@triton.autotune(configs=_cfg(), key=["D", "G_pad", "P"])
@triton.heuristics({
    # sm90 wgmma tl.dot requires M,N multiples of 16; G in {4,8} -> pad to 16.
    "G_pad": lambda args: max(triton.next_power_of_2(args["G"]), 16),
})
@triton.jit
def _paged_attn(
    Q, KV, BT, SL, Out,
    stride_q_b, stride_q_h, stride_q_d,
    stride_kv_b, stride_kv_p, stride_kv_h, stride_kv_d,
    stride_bt_b, stride_bt_m,
    stride_o_b, stride_o_h, stride_o_d,
    B: tl.constexpr, H: tl.constexpr, Hkv: tl.constexpr, G: tl.constexpr,
    scale, P: tl.constexpr, max_blocks: tl.constexpr,
    D: tl.constexpr, G_pad: tl.constexpr, BLOCK_K: tl.constexpr,
):
    b = tl.program_id(0)
    kvh = tl.program_id(1)
    sl = tl.load(SL + b)

    offs_d = tl.arange(0, D)
    offs_g = tl.arange(0, G_pad)

    # Q for the G query heads sharing this kv_head (padded to G_pad rows), bf16.
    q_ptrs = (Q + b * stride_q_b
              + (kvh * G + offs_g[:, None]) * stride_q_h
              + offs_d[None, :] * stride_q_d)
    q = tl.load(q_ptrs, mask=(offs_g < G)[:, None], other=0.0)

    m = tl.full((G_pad,), -float("inf"), dtype=tl.float32)
    l = tl.full((G_pad,), 0.0, dtype=tl.float32)
    o = tl.zeros((G_pad, D), dtype=tl.float32)

    start = 0
    while start < sl:
        toks = start + tl.arange(0, BLOCK_K)
        mask = toks < sl

        page_idx = toks // P
        inpage = toks % P
        pages = tl.load(BT + b * stride_bt_b + page_idx * stride_bt_m,
                        mask=mask, other=0).to(tl.int32)

        # Offset of (page, inpage, kvh, :) in KV; K is the first D of the pack.
        base = (pages * stride_kv_b
                + inpage * stride_kv_p
                + kvh * stride_kv_h)

        k = tl.load(KV + base[:, None] + offs_d[None, :] * stride_kv_d,
                    mask=mask[:, None], other=0.0)
        v = tl.load(KV + base[:, None] + D * stride_kv_d
                    + offs_d[None, :] * stride_kv_d,
                    mask=mask[:, None], other=0.0)

        # QK^T: (G_pad, D) x (D, BLOCK_K) -> (G_pad, BLOCK_K); bf16 -> fp32 acc.
        scores = tl.dot(q, tl.trans(k))
        scores = scores * scale
        scores = tl.where(mask[None, :], scores, -1e30)

        rowmax = tl.max(scores, axis=1)
        m_new = tl.maximum(m, rowmax)
        p = tl.exp(scores - m_new[:, None])
        p = tl.where(mask[None, :], p, 0.0)
        l_sum = tl.sum(p, axis=1)
        l_new = tl.exp(m - m_new) * l + l_sum
        # PV: keep p in bf16 (tensorcore) against bf16 V, accumulate in fp32.
        o = tl.exp(m - m_new)[:, None] * o + tl.dot(p.to(tl.bfloat16), v)
        m = m_new
        l = l_new
        start += BLOCK_K

    o = o / l[:, None]
    o = o.to(tl.bfloat16)
    out_ptrs = (Out + b * stride_o_b
                + (kvh * G + offs_g[:, None]) * stride_o_h
                + offs_d[None, :] * stride_o_d)
    tl.store(out_ptrs, o, mask=(offs_g < G)[:, None])


def _launch(q, kv, bt, sl, out, scale, P, max_blocks):
    B, H, D = q.shape
    Hkv = kv.shape[2]
    G = H // Hkv

    grid = (B, Hkv)
    _paged_attn[grid](
        q, kv, bt, sl, out,
        q.stride(0), q.stride(1), q.stride(2),
        kv.stride(0), kv.stride(1), kv.stride(2), kv.stride(3),
        bt.stride(0), bt.stride(1),
        out.stride(0), out.stride(1), out.stride(2),
        B=B, H=H, Hkv=Hkv, G=G, scale=scale, P=P, max_blocks=max_blocks, D=D,
    )


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)

    def forward(self, query, kv_cache, block_table, seq_lens):
        out = torch.empty_like(query)
        max_blocks = block_table.shape[1]
        _launch(query, kv_cache, block_table, seq_lens, out,
                self.scale, self.page_size, max_blocks)
        return out


def get_inputs():
    B = BATCH
    H = NUM_HEADS
    Hkv = NUM_KV_HEADS
    D = HEAD_DIM
    L = SEQ_LEN
    P = 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.empty(total_pages, P, Hkv, 2 * D, dtype=torch.bfloat16)
    kv_cache.normal_().mul_(0.1)

    perm = torch.randperm(total_pages)[: B * pages_per_seq].reshape(B, pages_per_seq).int()
    block_table = perm.contiguous()
    seq_lens = torch.full((B,), L, dtype=torch.int32)

    return [query, kv_cache, block_table, seq_lens]


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