# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/sparse_attn.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""The attention arithmetic — QK score, masked softmax with the learned sink,
and the probability-weighted sum — as one NKI kernel.

  * the QK product takes bf16 operands into an fp32 psum accumulator, and the
    scale multiplies the *logits* rather than `q`, so it never passes through a
    bf16 rounding;
  * the denominator sums **fp32** probabilities while the PV product consumes
    **bf16** ones — that is the reference's own split (its staging buffer is
    bf16), which is why `_softmax` writes the exponentials twice;
  * the sink joins the denominator after the row max is final.
"""

import nki
import nki.isa as nisa
import nki.language as nl

#: Partition width of the tensor engine.
PT = 128
#: A psum bank holds 2 KiB per partition = 512 fp32, which is also `nc_matmul`'s
#: widest moving free axis.
PSUM_FP32 = 512

#: `gathered_attn_idx` takes the softmax scale in column 0 of its bias tensor
#: rather than as its own operand, and this is how much of the row comes before
#: the bias proper. 16 fp32 = 64 bytes, so the bias slice the score product
BIAS_PAD = 16


def _ktile(k):
    """Free-axis width for the score tile: as wide as a psum bank allows, and a
    divisor of `k` so no slots are dropped.
    """
    if k <= PSUM_FP32:
        return k
    for cand in (PSUM_FP32, 256, PT):
        if k % cand == 0:
            return cand
    return PT


def _align16(n):
    """Round an SBUF free extent up to 16 elements = 32 bytes.

        TDRV:dma_util_expand_desc_transpose  transpose dest offset ... must be
        32B aligned
    """
    return -(-n // 16) * 16


@nki.jit
def gathered_attn(q, g, bias, sink):
    """Multi-head attention over per-query *pre-gathered* slots.

      * `softmax_scale` arrives in column 0 of `bias` rather than as its own
        `(1, 1)` operand — see `BIAS_PAD`. As two separate broadcast tiles the
        scale and the bias hit `[NCC_IBIR297]` (two SBUF inputs of one
        instruction at different base partitions) at `tokens=512, topk=128`,
        in *this* kernel too, so the shipped form cannot be compiled standalone
        at every live shape and could not serve as the reference unchanged. The
        instruction and both operand *values* are untouched, and the tiny-shape
        run of `devcmp.py` compared the shipped form against the combined one
        and got identical bytes;
      * nothing else.
    """
    tq, d = q.shape
    tokens, kb_ = bias.shape
    k = kb_ - BIAS_PAD
    heads = tq // tokens

    assert tq == tokens * heads, f"q rows {tq} is not tokens({tokens}) * heads"
    assert g.shape == (tokens * k, d), (
        f"g {g.shape} does not match tokens({tokens}) * topk({k}) x head_dim({d})"
    )
    assert d % PT == 0, f"head_dim={d} must be a multiple of {PT}"
    assert k % PT == 0, f"topk={k} must be a multiple of {PT} — the caller pads"
    assert heads <= PT, f"heads={heads} exceeds the partition width {PT}"
    assert sink.shape == (heads, 1), f"sink {sink.shape} wants ({heads}, 1)"

    dchunks = d // PT
    ktiles = k // PT
    kt = _ktile(k)
    ntiles = k // kt
    hp = _align16(heads)

    out = nl.ndarray((tq, d), dtype=nl.bfloat16, buffer=nl.shared_hbm)

    sink_t = nl.load(sink[0:heads, 0:1], dtype=nl.float32)

    for t in nl.affine_range(tokens):
        qt = nl.ndarray((PT, dchunks, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
        gt = nl.ndarray((PT, dchunks, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        gn = nl.ndarray((PT, ktiles, d), dtype=nl.bfloat16, buffer=nl.sbuf)
        for c in nl.static_range(dchunks):
            nisa.dma_transpose(
                dst=qt[0:PT, c, 0:heads],
                src=q[t * heads : t * heads + heads, c * PT : (c + 1) * PT],
            )
            nisa.dma_transpose(
                dst=gt[0:PT, c, 0:k],
                src=g[t * k : (t + 1) * k, c * PT : (c + 1) * PT],
            )
        for kb in nl.static_range(ktiles):
            nisa.tensor_copy(
                dst=gn[0:PT, kb, 0:d],
                src=nl.load(
                    g[t * k + kb * PT : t * k + (kb + 1) * PT, 0:d],
                    dtype=nl.bfloat16,
                ),
            )

        bias_b = nl.broadcast_to(
            nl.load(bias[t : t + 1, 0 : BIAS_PAD + k]), (heads, BIAS_PAD + k)
        )
        scale_p = bias_b[0:heads, 0:1]
        sc = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
        for nb in nl.static_range(ntiles):
            acc = nl.ndarray((heads, kt), dtype=nl.float32, buffer=nl.psum)
            for c in nl.static_range(dchunks):
                nisa.nc_matmul(
                    dst=acc,
                    stationary=qt[0:PT, c, 0:heads],
                    moving=gt[0:PT, c, nb * kt : (nb + 1) * kt],
                    accumulate=(c > 0),
                )
            nisa.scalar_tensor_tensor(
                dst=sc[0:heads, nb * kt : (nb + 1) * kt],
                data=acc,
                op0=nl.multiply,
                operand0=scale_p,
                op1=nl.add,
                operand1=bias_b[
                    0:heads, BIAS_PAD + nb * kt : BIAS_PAD + (nb + 1) * kt
                ],
            )

        neg_max = nl.negative(nl.max(sc[0:heads, 0:k], axis=1, keepdims=True))
        p32 = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
        den = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.activation(
            dst=p32,
            op=nl.exp,
            data=sc[0:heads, 0:k],
            bias=neg_max,
            scale=1.0,
            reduce_op=nl.add,
            reduce_res=den,
        )
        sk = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.activation(dst=sk, op=nl.exp, data=sink_t, bias=neg_max)
        rcp = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.reciprocal(dst=rcp, data=nl.add(den, sk))
        p16 = nl.copy(p32, dtype=nl.bfloat16)

        pt = nl.ndarray((PT, ktiles, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
        for kb in nl.static_range(ktiles):
            ptps = nl.ndarray((PT, hp), dtype=nl.bfloat16, buffer=nl.psum)
            nisa.nc_transpose(
                dst=ptps[0:PT, 0:heads], data=p16[0:heads, kb * PT : (kb + 1) * PT]
            )
            nisa.tensor_copy(
                dst=pt[0:PT, kb, 0:heads], src=ptps[0:PT, 0:heads]
            )
        acc_o = nl.ndarray((heads, d), dtype=nl.float32, buffer=nl.psum)
        for kb in nl.static_range(ktiles):
            nisa.nc_matmul(
                dst=acc_o,
                stationary=pt[0:PT, kb, 0:heads],
                moving=gn[0:PT, kb, 0:d],
                accumulate=(kb > 0),
            )
        o16 = nl.ndarray((heads, d), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=o16, data=acc_o, op0=nl.multiply, operand0=rcp
        )
        nl.store(out[t * heads : t * heads + heads, 0:d], value=o16)
    return out


@nki.jit
def gathered_attn_idx(q, kv, idx, bias, sink):
    """`gathered_attn` with the gather inside it. **This is the live path.**

      * the logit of an empty slot is `q·kv[0] * scale + (-1e30)`. `|q·kv|` is
        bounded by a few times `head_dim` and one ulp of 1e30 is ~1.2e23, so
        the sum is **exactly** -1e30, which is what the zeroed row gives. The
        row max, and so every probability, is therefore identical;
      * the value product weights that row by `p = exp(-1e30 - max) = 0`, and
        `0 * finite` is exactly 0. `kv` is fp8/fp4-rounded activation, never
        inf or NaN, so the product cannot be NaN;
      * in the all-empty case (`p = 1` for every slot) the accumulator holds
        `sum(kv[0])` instead of 0, and the reciprocal of the denominator is
        exactly 0 because `exp(sink + 1e30)` is `+inf` — so the row is exactly
        zero either way, which is the reference's convention.
    """
    tq, d = q.shape
    tokens, kb_ = bias.shape
    k = kb_ - BIAS_PAD
    heads = tq // tokens
    nrows, kvd = kv.shape

    assert tq == tokens * heads, f"q rows {tq} is not tokens({tokens}) * heads"
    assert kvd == d, f"kv head_dim {kvd} does not match q head_dim {d}"
    assert idx.shape == (tokens, k), (
        f"idx {idx.shape} does not match tokens({tokens}) x topk({k})"
    )
    assert nrows > 0, "kv is empty"
    assert d % PT == 0, f"head_dim={d} must be a multiple of {PT}"
    assert k % PT == 0, f"topk={k} must be a multiple of {PT} — the caller pads"
    assert heads <= PT, f"heads={heads} exceeds the partition width {PT}"
    assert sink.shape == (heads, 1), f"sink {sink.shape} wants ({heads}, 1)"

    dchunks = d // PT
    ktiles = k // PT
    kt = _ktile(k)
    ntiles = k // kt
    hp = _align16(heads)

    out = nl.ndarray((tq, d), dtype=nl.bfloat16, buffer=nl.shared_hbm)

    sink_t = nl.load(sink[0:heads, 0:1], dtype=nl.float32)

    for t in nl.affine_range(tokens):
        qt = nl.ndarray((PT, dchunks, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
        gt = nl.ndarray((PT, dchunks, k), dtype=nl.bfloat16, buffer=nl.sbuf)
        gn = nl.ndarray((PT, ktiles, d), dtype=nl.bfloat16, buffer=nl.sbuf)

        ids = nl.ndarray((PT, ktiles), dtype=nl.int32, buffer=nl.sbuf)
        nisa.dma_copy(
            dst=ids, src=idx[t : t + 1, 0:k].ap(pattern=[[1, PT], [PT, ktiles]])
        )

        for c in nl.static_range(dchunks):
            nisa.dma_transpose(
                dst=qt[0:PT, c, 0:heads],
                src=q[t * heads : t * heads + heads, c * PT : (c + 1) * PT],
            )

        for kb in nl.static_range(ktiles):
            nisa.dma_copy(
                dst=gn[0:PT, kb, 0:d],
                src=kv.ap(
                    pattern=[[d, PT], [1, d]],
                    offset=0,
                    vector_offset=ids[0:PT, kb : kb + 1],
                    indirect_dim=0,
                ),
            )
            tps = nl.ndarray((PT, dchunks, PT), dtype=nl.bfloat16, buffer=nl.psum)
            for c in nl.static_range(dchunks):
                nisa.nc_transpose(
                    dst=tps[0:PT, c, 0:PT],
                    data=gn[0:PT, kb, c * PT : (c + 1) * PT],
                )
            nisa.tensor_copy(
                dst=gt[0:PT, 0:dchunks, kb * PT : (kb + 1) * PT], src=tps
            )

        bias_b = nl.broadcast_to(
            nl.load(bias[t : t + 1, 0 : BIAS_PAD + k]), (heads, BIAS_PAD + k)
        )
        scale_p = bias_b[0:heads, 0:1]
        sc = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
        for nb in nl.static_range(ntiles):
            acc = nl.ndarray((heads, kt), dtype=nl.float32, buffer=nl.psum)
            for c in nl.static_range(dchunks):
                nisa.nc_matmul(
                    dst=acc,
                    stationary=qt[0:PT, c, 0:heads],
                    moving=gt[0:PT, c, nb * kt : (nb + 1) * kt],
                    accumulate=(c > 0),
                )
            nisa.scalar_tensor_tensor(
                dst=sc[0:heads, nb * kt : (nb + 1) * kt],
                data=acc,
                op0=nl.multiply,
                operand0=scale_p,
                op1=nl.add,
                operand1=bias_b[
                    0:heads, BIAS_PAD + nb * kt : BIAS_PAD + (nb + 1) * kt
                ],
            )

        neg_max = nl.negative(nl.max(sc[0:heads, 0:k], axis=1, keepdims=True))
        p32 = nl.ndarray((heads, k), dtype=nl.float32, buffer=nl.sbuf)
        den = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.activation(
            dst=p32,
            op=nl.exp,
            data=sc[0:heads, 0:k],
            bias=neg_max,
            scale=1.0,
            reduce_op=nl.add,
            reduce_res=den,
        )
        sk = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.activation(dst=sk, op=nl.exp, data=sink_t, bias=neg_max)
        rcp = nl.ndarray((heads, 1), dtype=nl.float32, buffer=nl.sbuf)
        nisa.reciprocal(dst=rcp, data=nl.add(den, sk))
        p16 = nl.copy(p32, dtype=nl.bfloat16)

        pt = nl.ndarray((PT, ktiles, hp), dtype=nl.bfloat16, buffer=nl.sbuf)
        for kb in nl.static_range(ktiles):
            ptps = nl.ndarray((PT, hp), dtype=nl.bfloat16, buffer=nl.psum)
            nisa.nc_transpose(
                dst=ptps[0:PT, 0:heads], data=p16[0:heads, kb * PT : (kb + 1) * PT]
            )
            nisa.tensor_copy(
                dst=pt[0:PT, kb, 0:heads], src=ptps[0:PT, 0:heads]
            )
        acc_o = nl.ndarray((heads, d), dtype=nl.float32, buffer=nl.psum)
        for kb in nl.static_range(ktiles):
            nisa.nc_matmul(
                dst=acc_o,
                stationary=pt[0:PT, kb, 0:heads],
                moving=gn[0:PT, kb, 0:d],
                accumulate=(kb > 0),
            )
        o16 = nl.ndarray((heads, d), dtype=nl.bfloat16, buffer=nl.sbuf)
        nisa.tensor_scalar(
            dst=o16, data=acc_o, op0=nl.multiply, operand0=rcp
        )
        nl.store(out[t * heads : t * heads + heads, 0:d], value=o16)
    return out