onnx--onnx/onnx/reference/ops/op_linear_attention.py

# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import numpy as np

from onnx.reference.op_run import OpRun


def _unpack_3d_to_4d(x: np.ndarray, num_heads: int) -> np.ndarray:
    """Reshape (B, T, H*D) -> (B, H, T, D)."""
    b, t, hd = x.shape
    if hd % num_heads != 0:
        raise ValueError(
            f"Last dim {hd} not divisible by num_heads {num_heads} for shape {x.shape}."
        )
    d = hd // num_heads
    return x.reshape(b, t, num_heads, d).transpose(0, 2, 1, 3)


class LinearAttention(OpRun):
    def _run(
        self,
        query,
        key,
        value,
        past_state=None,
        decay=None,
        beta=None,
        chunk_size=None,  # noqa: ARG002 — tuning hint, no effect on output
        kv_num_heads=None,
        q_num_heads=None,
        scale=None,
        update_rule=None,
    ):
        # --- Step 1: defaults and validation ---
        if update_rule is None:
            update_rule = "gated_delta"
        if update_rule not in ("linear", "gated", "delta", "gated_delta"):
            raise ValueError(
                f"Unsupported update_rule '{update_rule}'. "
                "Expected one of: 'linear', 'gated', 'delta', 'gated_delta'."
            )
        if q_num_heads is None or kv_num_heads is None:
            raise ValueError("q_num_heads and kv_num_heads are required attributes.")
        if q_num_heads <= 0 or kv_num_heads <= 0 or q_num_heads % kv_num_heads != 0:
            raise ValueError(
                f"q_num_heads ({q_num_heads}) must be a positive multiple of "
                f"kv_num_heads ({kv_num_heads})."
            )

        gating = update_rule in ("gated", "gated_delta")
        delta_correction = update_rule in ("delta", "gated_delta")

        if gating and decay is None:
            raise ValueError(f"update_rule '{update_rule}' requires decay input.")
        if not gating and decay is not None:
            raise ValueError(f"update_rule '{update_rule}' forbids decay input.")
        if delta_correction and beta is None:
            raise ValueError(f"update_rule '{update_rule}' requires beta input.")
        if not delta_correction and beta is not None:
            raise ValueError(f"update_rule '{update_rule}' forbids beta input.")

        for name, arr in (("query", query), ("key", key), ("value", value)):
            if arr.ndim != 3:
                raise ValueError(
                    f"{name} must be rank 3 (B, T, H*D), got shape {arr.shape}."
                )

        # --- Step 2: unpack Q/K/V to 4D (B, H, T, D) ---
        out_dtype = query.dtype
        b, t, _ = query.shape
        d_k = query.shape[-1] // q_num_heads
        d_v = value.shape[-1] // kv_num_heads
        group_size = q_num_heads // kv_num_heads

        q4 = _unpack_3d_to_4d(query, q_num_heads).astype(np.float32)
        k4 = _unpack_3d_to_4d(key, kv_num_heads).astype(np.float32)
        v4 = _unpack_3d_to_4d(value, kv_num_heads).astype(np.float32)

        # --- Step 3: unpack decay (broadcastable to (B, H_kv, T, d_k)) ---
        if decay is not None:
            if decay.ndim != 3:
                raise ValueError(f"decay must be rank 3, got shape {decay.shape}.")
            decay_last = decay.shape[-1]
            if decay_last == kv_num_heads:
                # Per-head scalar: (B, T, H_kv) -> (B, H_kv, T, 1)
                decay4 = decay.reshape(b, t, kv_num_heads, 1).transpose(0, 2, 1, 3)
            elif decay_last == kv_num_heads * d_k:
                # Per-key-dim: (B, T, H_kv*d_k) -> (B, H_kv, T, d_k)
                decay4 = _unpack_3d_to_4d(decay, kv_num_heads)
            else:
                raise ValueError(
                    f"decay last dim {decay_last} must equal kv_num_heads "
                    f"({kv_num_heads}) or kv_num_heads*d_k ({kv_num_heads * d_k})."
                )
            decay4 = decay4.astype(np.float32)

        # --- Step 4: unpack beta (broadcastable to (B, H_kv, T, 1)) ---
        if beta is not None:
            if beta.ndim != 3:
                raise ValueError(f"beta must be rank 3, got shape {beta.shape}.")
            beta_last = beta.shape[-1]
            if beta_last not in (kv_num_heads, 1):
                raise ValueError(
                    f"beta last dim {beta_last} must be kv_num_heads "
                    f"({kv_num_heads}) or 1."
                )
            # (B, T, H_kv_or_1) -> (B, H_kv_or_1, T, 1)
            beta4 = beta.reshape(b, t, beta_last, 1).transpose(0, 2, 1, 3)
            beta4 = beta4.astype(np.float32)

        # --- Step 5: initialize state in float32 ---
        # TODO(review): The proposal allows S != T (e.g., float32 state with
        # float16/bfloat16 activations). We accumulate internally in float32
        # regardless, then cast `present_state` back to `past_state.dtype` (or
        # `query.dtype` when `past_state` is omitted, since there is no S
        # anchor in that case). A cleaner contract would propagate S
        # explicitly — possibly via a new attribute or by inferring S from a
        # zero-shape sentinel — once the spec resolves how S is signalled
        # when past_state is absent. Mirrors the same TODO in the C++
        # function-body builder in onnx/defs/nn/defs.cc.
        if past_state is not None:
            if past_state.shape != (b, kv_num_heads, d_k, d_v):
                raise ValueError(
                    f"past_state shape {past_state.shape} does not match "
                    f"({b}, {kv_num_heads}, {d_k}, {d_v})."
                )
            state_in_dtype = past_state.dtype
            state = past_state.astype(np.float32).copy()
        else:
            state_in_dtype = out_dtype
            state = np.zeros((b, kv_num_heads, d_k, d_v), dtype=np.float32)

        # --- Step 6: scale ---
        if scale is None or scale == 0.0:
            scale_val = 1.0 / np.sqrt(d_k)
        else:
            scale_val = float(scale)

        # --- Step 7+8: recurrence with GQA expansion at read time ---
        outputs = np.zeros((b, q_num_heads, t, d_v), dtype=np.float32)
        for i in range(t):
            q_t = q4[:, :, i, :]  # (B, H_q, d_k)
            k_t = k4[:, :, i, :]  # (B, H_kv, d_k)
            v_t = v4[:, :, i, :]  # (B, H_kv, d_v)

            # Decay: state *= exp(g_t)
            if gating:
                g_t = decay4[:, :, i, :]  # (B, H_kv, 1) or (B, H_kv, d_k)
                state = state * np.exp(g_t)[..., None]  # broadcast over d_v

            # Delta correction: v_t <- beta_t * (v_t - S^T @ k_t)
            if delta_correction:
                # retrieved[b, h, m] = sum_d state[b, h, d, m] * k_t[b, h, d]
                retrieved = np.einsum("bhdm,bhd->bhm", state, k_t)
                v_t = beta4[:, :, i, :] * (v_t - retrieved)

            # Write: state += k_t ⊗ v_t  (outer product over last dims)
            state = state + k_t[..., :, None] * v_t[..., None, :]

            # Read with GQA: replicate KV-head state across query heads.
            if group_size == 1:
                state_for_read = state
            else:
                # (B, H_kv, d_k, d_v) -> (B, H_q, d_k, d_v) by interleave-repeat
                state_for_read = np.repeat(state, group_size, axis=1)
            # o_t[b, h, m] = scale * sum_d q_t[b, h, d] * state_for_read[b, h, d, m]
            outputs[:, :, i, :] = scale_val * np.einsum(
                "bhd,bhdm->bhm", q_t, state_for_read
            )

        # --- Step 9: repack output (B, H_q, T, d_v) -> (B, T, H_q*d_v) ---
        output = outputs.transpose(0, 2, 1, 3).reshape(b, t, q_num_heads * d_v)
        output = output.astype(out_dtype)

        # --- Step 10: present_state in same dtype as past_state (or query) ---
        present_state = state.astype(state_in_dtype)

        return (output, present_state)