"""Row-invariant decode attention (Triton), usable inside CUDA/HIP graphs.

One program per (query row, head). Row i attends to the logical keys t in [0, pos[i]]
of its layer; `pos` is read from device memory (graph-capturable). Optional tree
addressing: keys t < cl[i] are read at cache index t, keys t >= cl[i] at
anc[i, t - cl[i]] (the slots of the row's own path in a draft tree; its last entry
is the row's own slot). The summation order (blocks of BLOCK_N logical keys) does not
depend on the addresses, so a chain (cl = pos + 1) is bitwise unchanged and a tree
node gets exactly the AR arithmetic of its own path.
"""

from __future__ import annotations

import torch
import triton
import triton.language as tl


@triton.jit
def _row_attention_kernel(Q, K, V, O, pos_ptr, cl_ptr, anc_ptr, sm_scale,
                          stride_qn, stride_qh, stride_kh, stride_kt,
                          stride_on, stride_oh,
                          GROUP: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr,
                          MAXD: tl.constexpr):
    i = tl.program_id(0)
    h = tl.program_id(1)
    kvh = h // GROUP
    length = tl.load(pos_ptr + i).to(tl.int32) + 1
    cl = tl.load(cl_ptr + i).to(tl.int32)
    d = tl.arange(0, HEAD_DIM)
    q = tl.load(Q + i * stride_qn + h * stride_qh + d).to(tl.float32) * sm_scale
    m_i = tl.full((), -1.0e30, tl.float32)
    l_i = tl.full((), 0.0, tl.float32)
    acc = tl.zeros((HEAD_DIM,), tl.float32)
    k_base = K + kvh * stride_kh
    v_base = V + kvh * stride_kh
    for s in range(0, length, BLOCK_N):
        t = s + tl.arange(0, BLOCK_N)
        mask = t < length
        e = t - cl
        tree = mask & (e >= 0)
        a = tl.where(tree, tl.load(anc_ptr + i * MAXD + e, mask=tree, other=0).to(tl.int32), t)
        k = tl.load(k_base + a[:, None] * stride_kt + d[None, :], mask=mask[:, None], other=0.0).to(tl.float32)
        scores = tl.sum(k * q[None, :], axis=1)
        scores = tl.where(mask, scores, -1.0e30)
        m_new = tl.maximum(m_i, tl.max(scores, axis=0))
        alpha = tl.exp(m_i - m_new)
        p = tl.exp(scores - m_new)
        p = tl.where(mask, p, 0.0)
        v = tl.load(v_base + a[:, None] * stride_kt + d[None, :], mask=mask[:, None], other=0.0).to(tl.float32)
        l_i = l_i * alpha + tl.sum(p, axis=0)
        acc = acc * alpha + tl.sum(p[:, None] * v, axis=0)
        m_i = m_new
    out = acc / l_i
    tl.store(O + i * stride_on + h * stride_oh + d, out.to(O.dtype.element_ty))


def row_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
                  pos: torch.Tensor, scale: float, out: torch.Tensor | None = None,
                  block_n: int = 64, cl: torch.Tensor | None = None,
                  anc: torch.Tensor | None = None) -> torch.Tensor:
    """q: (n, H, hd); k_cache/v_cache: (Hkv, max_len, hd); pos: int tensor (>= n,) on device.

    Returns (n, H, hd): row i attends to positions [0, pos[i]].
    """
    n, H, hd = q.shape
    Hkv = k_cache.shape[0]
    assert k_cache.stride(-1) == 1 and q.stride(-1) == 1
    if out is None:
        out = torch.empty_like(q)
    if cl is None:  # plain chains: every key at its own position
        cl, anc = pos + 1, pos[:, None]
    _row_attention_kernel[(n, H)](
        q, k_cache, v_cache, out, pos, cl, anc, scale,
        q.stride(0), q.stride(1), k_cache.stride(0), k_cache.stride(1),
        out.stride(0), out.stride(1),
        GROUP=H // Hkv, HEAD_DIM=hd, BLOCK_N=block_n, MAXD=anc.shape[1],
    )
    return out