DeepInfer/deepinfer/neuron_ops/vendor/deep_gemm/einsum.py

"""Batched GEMM / einsum entry points; mirrors ``DeepGEMM-Ascend/csrc/apis/einsum.hpp``.

Only the three contraction strings the upstream kernels accept are supported:

* ``"bhr,hdr->bhd"``: a [b, h, r], b [h, d, r] -> d [b, h, d]
* ``"bhd,hdr->bhr"``: a [b, h, d], b [h, d, r] -> d [b, h, r]
* ``"bhd,bhr->hdr"``: a [b, h, d], b [b, h, r] -> d [h, d, r]

For the FP8 variants each operand's scale factor has the same leading dims as its
data and is indexed by the K block along its last dim (``gran_m == gran_n == 1``,
which is what DeepInfer passes: ``recipe=(1, 1, gran_k)``).
"""

import torch

try:  # the NKI kernels live in deepinfer; standalone shim tests run without it
    from deepinfer.neuron_ops import nki_gemm
except ImportError:  # pragma: no cover
    nki_gemm = None

from ._quant import dequantize
from .layout import get_recipe

_EXPRS = ("bhr,hdr->bhd", "bhd,hdr->bhr", "bhd,bhr->hdr")


def _check_expr(expr: str) -> None:
    assert expr in _EXPRS, f"Unsupported einsum expression: {expr}"


def _accumulate(out: torch.Tensor, d: torch.Tensor, c: torch.Tensor | None) -> None:
    if c is not None:
        if c.data_ptr() == d.data_ptr() and c.shape == d.shape:
            out = out + d.float()
        else:
            out = out + c.float()
    d.copy_(out.to(d.dtype))


def bf16_bmm(a, b, d, c=None, compiled_dims="nk", epilogue=None) -> None:
    """D[i] = A[i] @ B[i]^T (+ C[i])."""
    assert epilogue is None, "epilogue classes are not supported"
    assert a.dtype == b.dtype == torch.bfloat16, f"{a.dtype=}, {b.dtype=}"
    if d.numel() == 0:
        return
    _accumulate(a.float() @ b.float().transpose(-1, -2), d, c)


def fp8_bmm(a, b, d, c=None, recipe=(1, 1, 32), compiled_dims="nk", epilogue=None) -> None:
    """D[i] = A[i] @ B[i]^T (+ C[i]) with FP8/FP4 operands and UE8M0 scales."""
    assert epilogue is None, "epilogue classes are not supported"
    if isinstance(d, tuple):
        raise NotImplementedError("FP8 output quantization (D, SFD) is not supported on Neuron")
    a_data, a_sf = a
    b_data, b_sf = b
    if d.numel() == 0:
        return
    gran_m, gran_n, gran_k = get_recipe(recipe, None, None)
    assert gran_m == gran_n == 1, f"only per-row scales are supported, got {recipe=}"
    a_f32 = dequantize(a_data, a_sf, gran_k=gran_k)
    b_f32 = dequantize(b_data, b_sf, gran_k=gran_k)
    _accumulate(a_f32 @ b_f32.transpose(-1, -2), d, c)


def _permuted_c(c: torch.Tensor | None, perm) -> torch.Tensor | None:
    return None if c is None else c.permute(*perm)


def einsum(expr: str, a, b, d, c=None, use_cublaslt: bool = False, epilogue=None) -> None:
    """BF16 einsum over the fixed set of contraction expressions."""
    _check_expr(expr)
    if use_cublaslt:
        raise NotImplementedError("einsum: use_cublaslt is not implemented")
    if expr == "bhr,hdr->bhd":
        bf16_bmm(a.permute(1, 0, 2), b, d.permute(1, 0, 2), _permuted_c(c, (1, 0, 2)), "nk", epilogue)
    elif expr == "bhd,hdr->bhr":
        bf16_bmm(a.permute(1, 0, 2), b.permute(0, 2, 1), d.permute(1, 0, 2), _permuted_c(c, (1, 0, 2)), "nk", epilogue)
    else:  # "bhd,bhr->hdr"
        bf16_bmm(a.permute(1, 2, 0), b.permute(1, 2, 0), d, c, "mn", epilogue)


def fp8_einsum(expr: str, a, b, d, c=None, recipe=(1, 1, 32), epilogue=None) -> None:
    """FP8/FP4 einsum over the fixed set of contraction expressions."""
    _check_expr(expr)
    if isinstance(d, tuple):
        raise NotImplementedError("FP8 einsum output quantization is not supported on Neuron")
    a_data, a_sf = a
    b_data, b_sf = b
    if (
        expr == "bhr,hdr->bhd"
        and c is None
        and nki_gemm is not None
        and nki_gemm.supported(a_data, b_data, b_data.shape[-1], b_data.shape[-2])
    ):
        # NEURON: one NKI block-quant GEMM per group (vllm-neuron#71's BlockQuantLinear.grouped).
        for h in range(b_data.shape[0]):
            nki_gemm.gemm(a_data[:, h], a_sf[:, h], b_data, d[:, h], group=h)
        return
    if expr == "bhr,hdr->bhd":
        fp8_bmm(
            (a_data.permute(1, 0, 2), a_sf.permute(1, 0, 2)),
            (b_data, b_sf),
            d.permute(1, 0, 2),
            _permuted_c(c, (1, 0, 2)),
            recipe,
            "nk",
            epilogue,
        )
    elif expr == "bhd,hdr->bhr":
        fp8_bmm(
            (a_data.permute(1, 0, 2), a_sf.permute(1, 0, 2)),
            (b_data.permute(0, 2, 1), b_sf.permute(0, 2, 1)),
            d.permute(1, 0, 2),
            _permuted_c(c, (1, 0, 2)),
            recipe,
            "nk",
            epilogue,
        )
    else:  # "bhd,bhr->hdr"
        fp8_bmm(
            (a_data.permute(1, 2, 0), a_sf.permute(1, 2, 0)),
            (b_data.permute(1, 2, 0), b_sf.permute(1, 2, 0)),
            d,
            c,
            recipe,
            "mn",
            epilogue,
        )