"""Paged attention decode kernel for H100 (Hopper SM90).

Single-query decode with grouped-query attention. The KV cache is laid out as
(num_blocks, page_size, num_kv_heads, head_dim * 2) with [K | V] packed on the
last dim.

Two-stage "flash decoding":
  1. `_paged_attn_split_kernel`: each program handles one (batch, kv_head,
     split) and computes a partial (acc, m, l) over a contiguous chunk of the
     sequence. The number of splits is chosen so the grid roughly fills
     the SMs.
  2. `_paged_attn_combine_kernel`: per (batch, kv_head), combines the
     per-split (acc, m, l) into the final output via a log-sum-exp reduction.
"""
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"]

LOG2E = 1.4426950408889634

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


# ============================================================================
# Stage 1a: single-pass attention (no split)
# ============================================================================
@triton.jit
def _paged_attn_kernel(
    Q_ptr,           # (B, H, D) bf16
    KV_ptr,          # (num_blocks, P, H_kv, 2*D) bf16, [K | V]
    BT_ptr,          # (B, max_blocks) int32
    SL_ptr,          # (B,) int32
    O_ptr,           # (B, H, D) bf16
    stride_qb, stride_qh, stride_qd,
    stride_kvn, stride_kvp, stride_kvh, stride_kvd,
    stride_btb, stride_btn,
    stride_ob, stride_oh, stride_od,
    scale_log2,
    G: tl.constexpr,
    D: tl.constexpr,
    P: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_hkv = tl.program_id(1)

    seq_len = tl.load(SL_ptr + pid_b)

    h_start = pid_hkv * G
    q_head_offs = h_start + tl.arange(0, G)
    d_offs = tl.arange(0, D)

    q_ptrs = (
        Q_ptr
        + pid_b * stride_qb
        + q_head_offs[:, None] * stride_qh
        + d_offs[None, :] * stride_qd
    )
    q = tl.load(q_ptrs)  # (G, D)

    m_i = tl.full((G,), -1.0e30, dtype=tl.float32)
    l_i = tl.zeros((G,), dtype=tl.float32)
    acc = tl.zeros((G, D), dtype=tl.float32)

    n_offs = tl.arange(0, BLOCK_N)

    for start in tl.range(0, seq_len, BLOCK_N):
        token_offs = start + n_offs
        page_offs = token_offs // P
        offset_in_page = token_offs % P
        valid = token_offs < seq_len

        page_ids = tl.load(
            BT_ptr + pid_b * stride_btb + page_offs * stride_btn,
            mask=valid, other=0,
        )

        k_offs = (
            page_ids * stride_kvn
            + offset_in_page * stride_kvp
            + pid_hkv * stride_kvh
        )
        k_ptrs = KV_ptr + k_offs[:, None] + d_offs[None, :] * stride_kvd
        v_ptrs = KV_ptr + k_offs[:, None] + (d_offs[None, :] + D) * stride_kvd

        k = tl.load(k_ptrs, mask=valid[:, None], other=0.0)
        v = tl.load(v_ptrs, mask=valid[:, None], other=0.0)

        s = tl.dot(q, tl.trans(k)) * scale_log2
        s = tl.where(valid[None, :], s, -1.0e30)

        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp2(m_i - m_new)
        p = tl.exp2(s - m_new[:, None])

        l_i = l_i * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m_i = m_new

    acc = acc / l_i[:, None]

    o_ptrs = (
        O_ptr
        + pid_b * stride_ob
        + q_head_offs[:, None] * stride_oh
        + d_offs[None, :] * stride_od
    )
    tl.store(o_ptrs, acc.to(tl.bfloat16))


# ============================================================================
# Stage 1b: split-K attention
# ============================================================================
@triton.jit
def _paged_attn_split_kernel(
    Q_ptr,           # (B, H, D) bf16
    KV_ptr,          # (num_blocks, P, H_kv, 2*D) bf16, [K | V]
    BT_ptr,          # (B, max_blocks) int32
    SL_ptr,          # (B,) int32
    O_partial_ptr,   # (B, H_kv, S, G, D) fp32
    M_partial_ptr,   # (B, H_kv, S, G) fp32  -- running max in log2 space
    L_partial_ptr,   # (B, H_kv, S, G) fp32  -- running unnormalized sum
    # Q strides
    stride_qb, stride_qh, stride_qd,
    # KV strides (in elements)
    stride_kvn, stride_kvp, stride_kvh, stride_kvd,
    # block_table strides
    stride_btb, stride_btn,
    # partial output strides
    stride_op_b, stride_op_h, stride_op_s, stride_op_g, stride_op_d,
    # partial m, l strides
    stride_mp_b, stride_mp_h, stride_mp_s, stride_mp_g,
    stride_lp_b, stride_lp_h, stride_lp_s, stride_lp_g,
    scale_log2,
    NUM_SPLITS: tl.constexpr,
    G: tl.constexpr,
    D: tl.constexpr,
    P: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_hkv = tl.program_id(1)
    pid_split = tl.program_id(2)

    seq_len = tl.load(SL_ptr + pid_b)

    chunk = (seq_len + NUM_SPLITS - 1) // NUM_SPLITS
    n_start = pid_split * chunk
    n_end = tl.minimum(n_start + chunk, seq_len)

    h_start = pid_hkv * G
    q_head_offs = h_start + tl.arange(0, G)        # absolute head index, for Q load
    d_offs = tl.arange(0, D)
    g_offs_l = tl.arange(0, G)                     # group index (0..G-1), for partial buffers

    if n_start >= seq_len:
        # write -inf m, 0 l so combine ignores this split
        m_ptrs = (
            M_partial_ptr
            + pid_b * stride_mp_b
            + pid_hkv * stride_mp_h
            + pid_split * stride_mp_s
            + g_offs_l * stride_mp_g
        )
        l_ptrs = (
            L_partial_ptr
            + pid_b * stride_lp_b
            + pid_hkv * stride_lp_h
            + pid_split * stride_lp_s
            + g_offs_l * stride_lp_g
        )
        tl.store(m_ptrs, tl.full((G,), -1.0e30, dtype=tl.float32))
        tl.store(l_ptrs, tl.zeros((G,), dtype=tl.float32))
        return

    # ----- load Q (G, D) -----
    q_ptrs = (
        Q_ptr
        + pid_b * stride_qb
        + q_head_offs[:, None] * stride_qh
        + d_offs[None, :] * stride_qd
    )
    q = tl.load(q_ptrs)  # (G, D)

    # ----- online-softmax accumulators -----
    m_i = tl.full((G,), -1.0e30, dtype=tl.float32)
    l_i = tl.zeros((G,), dtype=tl.float32)
    acc = tl.zeros((G, D), dtype=tl.float32)

    n_offs = tl.arange(0, BLOCK_N)

    for start in tl.range(n_start, n_end, BLOCK_N):
        token_offs = start + n_offs
        page_offs = token_offs // P
        offset_in_page = token_offs % P
        valid = token_offs < n_end

        page_ids = tl.load(
            BT_ptr + pid_b * stride_btb + page_offs * stride_btn,
            mask=valid, other=0,
        )

        k_offs = (
            page_ids * stride_kvn
            + offset_in_page * stride_kvp
            + pid_hkv * stride_kvh
        )
        k_ptrs = KV_ptr + k_offs[:, None] + d_offs[None, :] * stride_kvd
        v_ptrs = KV_ptr + k_offs[:, None] + (d_offs[None, :] + D) * stride_kvd

        k = tl.load(k_ptrs, mask=valid[:, None], other=0.0)  # (BLOCK_N, D)
        v = tl.load(v_ptrs, mask=valid[:, None], other=0.0)  # (BLOCK_N, D)

        s = tl.dot(q, tl.trans(k)) * scale_log2
        s = tl.where(valid[None, :], s, -1.0e30)

        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp2(m_i - m_new)
        p = tl.exp2(s - m_new[:, None])

        l_i = l_i * alpha + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
        m_i = m_new

    # write partial output and (m, l)
    o_ptrs = (
        O_partial_ptr
        + pid_b * stride_op_b
        + pid_hkv * stride_op_h
        + pid_split * stride_op_s
        + g_offs_l[:, None] * stride_op_g
        + d_offs[None, :] * stride_op_d
    )
    tl.store(o_ptrs, acc)

    m_ptrs = (
        M_partial_ptr
        + pid_b * stride_mp_b
        + pid_hkv * stride_mp_h
        + pid_split * stride_mp_s
        + g_offs_l * stride_mp_g
    )
    l_ptrs = (
        L_partial_ptr
        + pid_b * stride_lp_b
        + pid_hkv * stride_lp_h
        + pid_split * stride_lp_s
        + g_offs_l * stride_lp_g
    )
    tl.store(m_ptrs, m_i)
    tl.store(l_ptrs, l_i)


# ============================================================================
# Stage 2: combine partial outputs across splits
# ============================================================================
@triton.jit
def _paged_attn_combine_kernel(
    O_partial_ptr,   # (B, H_kv, S, G, D) fp32
    M_partial_ptr,   # (B, H_kv, S, G) fp32
    L_partial_ptr,   # (B, H_kv, S, G) fp32
    O_ptr,           # (B, H, D) bf16
    stride_op_b, stride_op_h, stride_op_s, stride_op_g, stride_op_d,
    stride_mp_b, stride_mp_h, stride_mp_s, stride_mp_g,
    stride_lp_b, stride_lp_h, stride_lp_s, stride_lp_g,
    stride_ob, stride_oh, stride_od,
    G: tl.constexpr,
    D: tl.constexpr,
    NUM_SPLITS: tl.constexpr,
    BLOCK_S: tl.constexpr,
):
    pid_b = tl.program_id(0)
    pid_hkv = tl.program_id(1)

    h_start = pid_hkv * G
    g_offs = h_start + tl.arange(0, G)        # absolute head index, for output
    g_offs_l = tl.arange(0, G)                # group index, for partial buffers
    d_offs = tl.arange(0, D)
    s_offs = tl.arange(0, BLOCK_S)
    s_mask = s_offs < NUM_SPLITS

    # Load m, l (G, S) — m in log2 space, l unnormalized in normal space
    m_ptrs = (
        M_partial_ptr
        + pid_b * stride_mp_b
        + pid_hkv * stride_mp_h
        + s_offs[None, :] * stride_mp_s
        + g_offs_l[:, None] * stride_mp_g
    )
    l_ptrs = (
        L_partial_ptr
        + pid_b * stride_lp_b
        + pid_hkv * stride_lp_h
        + s_offs[None, :] * stride_lp_s
        + g_offs_l[:, None] * stride_lp_g
    )
    m = tl.load(m_ptrs, mask=s_mask[None, :], other=-1.0e30)
    l = tl.load(l_ptrs, mask=s_mask[None, :], other=0.0)

    m_max = tl.max(m, axis=1)                          # (G,)
    alpha = tl.exp2(m - m_max[:, None])                # (G, S)
    alpha_l = alpha * l                                # (G, S)
    l_global = tl.sum(alpha_l, axis=1)                 # (G,)

    # Load O_partial (G, S, D)
    op_ptrs = (
        O_partial_ptr
        + pid_b * stride_op_b
        + pid_hkv * stride_op_h
        + s_offs[None, :, None] * stride_op_s
        + g_offs_l[:, None, None] * stride_op_g
        + d_offs[None, None, :] * stride_op_d
    )
    op = tl.load(op_ptrs, mask=s_mask[None, :, None], other=0.0)  # (G, S, D)

    acc = tl.sum(alpha[:, :, None] * op, axis=1)       # (G, D)
    out = acc / l_global[:, None]

    o_ptrs = (
        O_ptr
        + pid_b * stride_ob
        + g_offs[:, None] * stride_oh
        + d_offs[None, :] * stride_od
    )
    tl.store(o_ptrs, out.to(tl.bfloat16))


# ============================================================================
# Driver
# ============================================================================
def _pick_block_n(D: int) -> int:
    if D >= 128:
        return 64
    return 128


def _pick_num_warps(D: int) -> int:
    if D >= 64:
        return 4
    return 2


def _next_pow2(x: int) -> int:
    p = 1
    while p < x:
        p <<= 1
    return p


# H100 PCIe: 114 SMs. Aim for ~2-3x oversubscription of work-units.
def _pick_num_splits(B: int, H_kv: int, seq_len: int) -> int:
    base = B * H_kv
    if base >= 64:
        return 1
    # aim for ~256 program-units
    splits = max(1, 256 // base)
    # cap so each split has at least 128 tokens of work
    splits = min(splits, max(1, seq_len // 128))
    return splits


_WORKSPACE = {}


def _get_workspace(B: int, H_kv: int, G: int, D: int, S: int, device):
    key = (B, H_kv, G, D, S, str(device))
    ws = _WORKSPACE.get(key)
    if ws is None:
        op = torch.empty(B, H_kv, S, G, D, dtype=torch.float32, device=device)
        mp = torch.empty(B, H_kv, S, G, dtype=torch.float32, device=device)
        lp = torch.empty(B, H_kv, S, G, dtype=torch.float32, device=device)
        _WORKSPACE[key] = (op, mp, lp)
        return op, mp, lp
    return ws


class Model(nn.Module):
    """Single-query paged attention decode (Triton, split-K)."""

    def __init__(
        self,
        batch: int,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
        seq_len: int,
        page_size: int,
    ):
        super().__init__()
        assert num_heads % num_kv_heads == 0, "num_heads must be a multiple of num_kv_heads (GQA)"
        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_log2 = LOG2E / math.sqrt(head_dim)
        self.register_buffer(
            "_dummy", torch.zeros(1, dtype=torch.bfloat16), persistent=False
        )

    def forward(
        self,
        query: torch.Tensor,
        kv_cache: torch.Tensor,
        block_table: torch.Tensor,
        seq_lens: torch.Tensor,
    ) -> torch.Tensor:
        B, H, D = query.shape
        H_kv = self.num_kv_heads
        G = self.group_size
        P = self.page_size

        out = torch.empty(B, H, D, dtype=query.dtype, device=query.device)

        BLOCK_N = _pick_block_n(D)
        num_warps = _pick_num_warps(D)

        num_splits = _pick_num_splits(B, H_kv, int(seq_lens.max().item()))

        if num_splits == 1:
            grid = (B, H_kv)
            _paged_attn_kernel[grid](
                query, kv_cache, block_table, seq_lens, out,
                query.stride(0), query.stride(1), query.stride(2),
                kv_cache.stride(0), kv_cache.stride(1),
                kv_cache.stride(2), kv_cache.stride(3),
                block_table.stride(0), block_table.stride(1),
                out.stride(0), out.stride(1), out.stride(2),
                self.scale_log2,
                G=G, D=D, P=P, BLOCK_N=BLOCK_N,
                num_warps=num_warps, num_stages=2,
            )
            return out

        # split path
        num_splits_pow2 = max(1, _next_pow2(num_splits))
        op, mp, lp = _get_workspace(B, H_kv, G, D, num_splits, query.device)

        grid_split = (B, H_kv, num_splits)
        _paged_attn_split_kernel[grid_split](
            query, kv_cache, block_table, seq_lens, op, mp, lp,
            query.stride(0), query.stride(1), query.stride(2),
            kv_cache.stride(0), kv_cache.stride(1),
            kv_cache.stride(2), kv_cache.stride(3),
            block_table.stride(0), block_table.stride(1),
            op.stride(0), op.stride(1), op.stride(2), op.stride(3), op.stride(4),
            mp.stride(0), mp.stride(1), mp.stride(2), mp.stride(3),
            lp.stride(0), lp.stride(1), lp.stride(2), lp.stride(3),
            self.scale_log2,
            NUM_SPLITS=num_splits,
            G=G, D=D, P=P, BLOCK_N=BLOCK_N,
            num_warps=num_warps, num_stages=2,
        )

        grid_combine = (B, H_kv)
        _paged_attn_combine_kernel[grid_combine](
            op, mp, lp, out,
            op.stride(0), op.stride(1), op.stride(2), op.stride(3), op.stride(4),
            mp.stride(0), mp.stride(1), mp.stride(2), mp.stride(3),
            lp.stride(0), lp.stride(1), lp.stride(2), lp.stride(3),
            out.stride(0), out.stride(1), out.stride(2),
            G=G, D=D, NUM_SPLITS=num_splits, BLOCK_S=num_splits_pow2,
            num_warps=2, num_stages=1,
        )
        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.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()
    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]