# SPDX-License-Identifier: Apache-2.0
# From vllm-project/vllm-neuron PR #71 (contrib/models/deepseek_v41/src/kernels/compress_pool.py,
# commit 30d54d1351c80728fe5685478fac838b44c53d43, Zifan He / OswaldHe), unmodified.
# Copyright Amazon.com, Inc. or its affiliates. See vllm-neuron NOTICE.
"""The KV compressor's learned softmax pooling, as one NKI kernel.

    g0 = nl.load(cgt[r0:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
    g1 = nl.load(cgt[r0 + 1:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
    v0 = nl.load(ckv[r0:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
    v1 = nl.load(ckv[r0 + 1:r1:RATIO, 0:HEAD_DIM], dtype=nl.float32)
    mx = nl.maximum(g0, g1)
    e0 = nl.exp(nl.subtract(g0, mx))
    e1 = nl.exp(nl.subtract(g1, mx))
    nisa.reciprocal(dst=rz, data=nl.add(e0, e1))
    p0 = nl.multiply(e0, rz)
    p1 = nl.multiply(e1, rz)
    acc = nl.add(nl.multiply(v0, p0), nl.multiply(v1, p1))
"""

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

#: Partition width; one *group* per partition, 128 groups per tile.
PT = 128


def _chunk(dim):
    """Free-axis chunk for the pooling: a divisor of `dim` small enough that
    `ratio + 3` fp32 copies of it fit a partition comfortably.
    """
    for cand in (2048, 1024, 512, 256, 128):
        if dim % cand == 0:
            return cand
    return dim


def _recip(d, shape):
    """`1 / d` with one Newton step: `y * (2 - d * y)`."""
    y = nl.ndarray(shape, dtype=nl.float32, buffer=nl.sbuf)
    nisa.reciprocal(dst=y, data=d)
    return nl.multiply(y, nl.add(nl.multiply(nl.multiply(d, y), -1.0), 2.0))


@nki.jit
def softmax_pool(kv, score):
    """`(kv * score.softmax(dim=1)).sum(dim=1)`. See the module docstring."""
    groups, ratio, dim = kv.shape
    assert score.shape == (groups, ratio, dim), (
        f"score {score.shape} wants ({groups}, {ratio}, {dim})"
    )
    assert ratio >= 2, f"ratio={ratio} — a single-member group needs no pooling"

    kv2 = kv.reshape((groups * ratio, dim))
    sc2 = score.reshape((groups * ratio, dim))

    ch = _chunk(dim)
    nchunks = dim // ch

    out = nl.ndarray((groups, dim), dtype=nl.float32, buffer=nl.shared_hbm)

    def body(g0, ng):
        r0, rn = g0 * ratio, ng * ratio
        for j in range(nchunks):
            c0 = j * ch
            g = [
                nl.load(
                    sc2[r0 + m : r0 + rn : ratio, c0 : c0 + ch], dtype=nl.float32
                )
                for m in range(ratio)
            ]
            mx = g[0]
            for m in range(1, ratio):
                mx = nl.maximum(mx, g[m])
            e = [nl.exp(nl.subtract(g[m], mx)) for m in range(ratio)]
            den = e[0]
            for m in range(1, ratio):
                den = nl.add(den, e[m])
            rz = _recip(den, (ng, ch))

            acc = None
            for m in range(ratio):
                v = nl.load(
                    kv2[r0 + m : r0 + rn : ratio, c0 : c0 + ch], dtype=nl.float32
                )
                term = nl.multiply(v, nl.multiply(e[m], rz))
                acc = term if acc is None else nl.add(acc, term)
            nl.store(out[g0 : g0 + ng, c0 : c0 + ch], value=acc)

    full, tail = groups // PT, groups % PT
    for t in nl.affine_range(full):
        body(t * PT, PT)
    if tail:
        body(full * PT, tail)
    return out