"""Small fused decode kernels, preserving the reference BF16 rounding points."""
import torch
import triton
import triton.language as tl


@triton.jit
def _rmsnorm(x, weight, output, H: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    data = tl.load(x + row * H + cols, cols < H, other=0).to(tl.float32)
    scale = tl.rsqrt(tl.sum(data * data, 0) / H + EPS)
    # LlamaRMSNorm rounds normalized activations before applying its weight.
    normalized = (data * scale).to(tl.bfloat16).to(tl.float32)
    w = tl.load(weight + cols, cols < H, other=0).to(tl.float32)
    tl.store(output + row * H + cols, normalized * w, cols < H)


def rmsnorm(states, weight, eps):
    assert states.is_contiguous() and states.dtype == torch.bfloat16
    width = states.shape[-1]
    result = torch.empty_like(states)
    _rmsnorm[(states.numel() // width,)](states, weight, result, width, eps,
        triton.next_power_of_2(width), num_warps=4, enable_fp_fusion=False)
    return result


@triton.jit
def _rotary_pair(Q, K, COS, SIN, POS, OQ, OK, HQ: tl.constexpr, D: tl.constexpr,
                 R: tl.constexpr, BLOCK: tl.constexpr):
    head = tl.program_id(0)
    dims = tl.arange(0, BLOCK)
    position = tl.load(POS)
    source = tl.where(head < HQ, Q + head * D, K + (head - HQ) * D)
    destination = tl.where(head < HQ, OQ + head * D, OK + (head - HQ) * D)
    values = tl.load(source + dims, dims < D, other=0).to(tl.float32)
    partner = (dims + R // 2) % R
    other = tl.load(source + partner, dims < R, other=0).to(tl.float32)
    other = tl.where(dims < R // 2, -other, other)
    cos = tl.load(COS + position * R + dims, dims < R, other=1).to(tl.float32)
    sin = tl.load(SIN + position * R + dims, dims < R, other=0).to(tl.float32)
    rotated = (values * cos).to(tl.bfloat16).to(tl.float32) + (other * sin).to(tl.bfloat16).to(tl.float32)
    tl.store(destination + dims, tl.where(dims < R, rotated, values), dims < D)


def rotary_pair(query, key, position, cos, sin):
    dimension = query.shape[-1]
    assert query.shape[0] == query.shape[2] == 1 and query.stride(1) == dimension
    oq = torch.empty_like(query, memory_format=torch.contiguous_format)
    if key is None:
        key, ok, key_heads = query, oq, 0
    else:
        assert key.shape[-1] == dimension and key.stride(1) == dimension
        ok, key_heads = torch.empty_like(key, memory_format=torch.contiguous_format), key.shape[1]
    _rotary_pair[(query.shape[1] + key_heads,)](query, key, cos, sin, position, oq, ok,
        query.shape[1], dimension, cos.shape[-1], triton.next_power_of_2(dimension),
        num_warps=4, enable_fp_fusion=False)
    return oq, ok if key_heads else None


@triton.jit
def _small_attention(Q, K, V, MASK, SINK, OUT,
        L: tl.constexpr, D: tl.constexpr, DV: tl.constexpr, GROUP: tl.constexpr,
        KS: tl.constexpr, VS: tl.constexpr, SCALE: tl.constexpr, VALUE_SCALE: tl.constexpr,
        HAS_SINK: tl.constexpr, BN: tl.constexpr, BD: tl.constexpr, BV: tl.constexpr):
    head = tl.program_id(0)
    kv = head // GROUP
    rows, dims = tl.arange(0, BN), tl.arange(0, BD)
    q = tl.load(Q + head * D + dims, dims < D, other=0).to(tl.float32)
    k = tl.load(K + kv * KS + rows[:, None] * D + dims[None, :],
                (rows[:, None] < L) & (dims[None, :] < D), other=0).to(tl.float32)
    logits = (tl.sum(k * q[None, :], 1).to(tl.bfloat16).to(tl.float32) * SCALE).to(tl.bfloat16).to(tl.float32)
    valid = tl.load(MASK + rows, rows < L, other=False)
    logits = tl.where(valid & (rows < L), logits, -float('inf'))
    maximum = tl.max(logits, 0)
    sink = -float('inf')
    if HAS_SINK:
        sink = tl.load(SINK + head).to(tl.float32)
        maximum = tl.maximum(maximum, sink)
    exps = tl.exp(logits - maximum)
    denominator = tl.sum(exps, 0)
    if HAS_SINK:
        denominator += tl.exp(sink - maximum)
    probabilities = (exps / denominator).to(tl.bfloat16).to(tl.float32)
    vd = tl.arange(0, BV)
    values = tl.load(V + kv * VS + rows[:, None] * DV + vd[None, :],
                    (rows[:, None] < L) & (vd[None, :] < DV), other=0).to(tl.float32)
    values = (values * VALUE_SCALE).to(tl.bfloat16).to(tl.float32)
    result = tl.sum(probabilities[:, None] * values, 0)
    tl.store(OUT + head * DV + vd, result, vd < DV)


@triton.jit
def _attention_logits(Q, K, MASK, LOGITS, L: tl.constexpr, D: tl.constexpr,
        GROUP: tl.constexpr, KS: tl.constexpr, SCALE: tl.constexpr,
        BN: tl.constexpr, BD: tl.constexpr):
    head, block = tl.program_id(0), tl.program_id(1)
    rows, dims = block * BN + tl.arange(0, BN), tl.arange(0, BD)
    q = tl.load(Q + head * D + dims, dims < D, other=0).to(tl.float32)
    k = tl.load(K + (head // GROUP) * KS + rows[:, None] * D + dims[None, :],
                (rows[:, None] < L) & (dims[None, :] < D), other=0).to(tl.float32)
    logits = (tl.sum(k * q[None, :], 1).to(tl.bfloat16).to(tl.float32) * SCALE).to(tl.bfloat16)
    valid = tl.load(MASK + rows, rows < L, other=False)
    tl.store(LOGITS + head * L + rows, tl.where(valid, logits, -float('inf')), rows < L)


@triton.jit
def _attention_probabilities(LOGITS, SINK, PROBS, L: tl.constexpr, HAS_SINK: tl.constexpr, BN: tl.constexpr):
    head = tl.program_id(0)
    rows = tl.arange(0, BN)
    logits = tl.load(LOGITS + head * L + rows, rows < L, other=-float('inf')).to(tl.float32)
    maximum = tl.max(logits, 0)
    sink = -float('inf')
    if HAS_SINK:
        sink = tl.load(SINK + head).to(tl.float32)
        maximum = tl.maximum(maximum, sink)
    exps = tl.exp(logits - maximum)
    denominator = tl.sum(exps, 0)
    if HAS_SINK:
        denominator += tl.exp(sink - maximum)
    tl.store(PROBS + head * L + rows, exps / denominator, rows < L)


@triton.jit
def _attention_values(PROBS, V, OUT, L: tl.constexpr, DV: tl.constexpr,
        GROUP: tl.constexpr, VS: tl.constexpr, VALUE_SCALE: tl.constexpr,
        BN: tl.constexpr, BV: tl.constexpr):
    head = tl.program_id(0)
    dims = tl.arange(0, BV)
    result = tl.full((BV,), 0, tl.float32)
    for block in range(tl.cdiv(L, BN)):
        rows = block * BN + tl.arange(0, BN)
        probabilities = tl.load(PROBS + head * L + rows, rows < L, other=0).to(tl.float32)
        values = tl.load(V + (head // GROUP) * VS + rows[:, None] * DV + dims[None, :],
                        (rows[:, None] < L) & (dims[None, :] < DV), other=0).to(tl.float32)
        values = (values * VALUE_SCALE).to(tl.bfloat16).to(tl.float32)
        result += tl.sum(probabilities[:, None] * values, 0)
    tl.store(OUT + head * DV + dims, result, dims < DV)


def decode_attention(query, keys, values, allowed, sink, value_scale):
    heads, dimension = query.shape[1], query.shape[-1]
    length, value_dimension = keys.shape[-2], values.shape[-1]
    assert query.shape[0] == query.shape[2] == 1
    assert keys.stride(-1) == values.stride(-1) == 1
    assert keys.stride(-2) == dimension and values.stride(-2) == value_dimension
    result = torch.empty((1, heads, 1, value_dimension), device=query.device, dtype=query.dtype)
    sink_pointer = sink if sink is not None else result
    group = heads // keys.shape[1]
    scale = dimension ** -.5
    value_scale = 1. if value_scale is None else value_scale
    if length <= 128:
        _small_attention[(heads,)](query, keys, values, allowed, sink_pointer, result,
            length, dimension, value_dimension, group, keys.stride(1), values.stride(1),
            scale, value_scale, sink is not None, triton.next_power_of_2(length),
            triton.next_power_of_2(dimension), triton.next_power_of_2(value_dimension),
            num_warps=8, enable_fp_fusion=False)
    else:
        logits = torch.empty((heads, length), dtype=torch.bfloat16, device=query.device)
        probabilities = torch.empty_like(logits)
        _attention_logits[(heads, triton.cdiv(length, 64))](query, keys, allowed, logits,
            length, dimension, group, keys.stride(1), scale, 64, triton.next_power_of_2(dimension),
            num_warps=4, enable_fp_fusion=False)
        _attention_probabilities[(heads,)](logits, sink_pointer, probabilities, length,
            sink is not None, triton.next_power_of_2(length), num_warps=4)
        _attention_values[(heads,)](probabilities, values, result, length, value_dimension,
            group, values.stride(1), value_scale, 64, triton.next_power_of_2(value_dimension),
            num_warps=4, enable_fp_fusion=False)
    return result