DeepInfer/deepinfer/neuron_ops/nki_norm.py

"""RMSNorm (+ residual add, + fused fp8 cast) and the clamped SwiGLU through the vllm-neuron#71 NKI kernels
(issue #13).

The pure-torch shims (``tilelang_ops/neuron/norm.py``, ``tile_kernels.quant.swiglu_forward_kernel``) emulate the
Ascend kernels bit for bit: a 1-ulp compensated rstd built from error-free fma tricks, a 4x64-lane square-sum
order, ue8m0 exponents read out of the fp32 bits. Inside compiled region graphs those formulations crash
neuronx-cc; on the device the norm instead runs ``nki/rmsnorm.rms_norm`` (fp32 interior, rsqrt + one Newton step)
and the SwiGLU ``nki/swiglu.swiglu_fused``. Results agree with the shims to bf16 rounding, not bit for bit.

The fp8 cast after the norm stays in torch, but without dtype bitcasts: the ue8m0 exponent ``ceil(log2(v))`` comes
from ``log2`` corrected against exact ``exp2`` powers of two, and the int16 scale words are joined arithmetically.

``DEEPINFER_NEURON_NKI_NORM=0`` / ``DEEPINFER_NEURON_NKI_SWIGLU=0`` restore the shims.
"""

from __future__ import annotations

import os
from typing import Any

import torch

NORM_ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_NORM", "1") == "1"
SWIGLU_ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_SWIGLU", "1") == "1"

#: ``swiglu`` without a clamp: a limit no finite bf16 activation reaches.
_NO_CLAMP = 3.0e38


_KERNELS: dict = {}


def _kernel(module: str, name: str):  # noqa: ANN202
    """``wrap_nki(nki.<module>.<name>)``, built once (a dict, not functools.cache, which dynamo warns about)."""
    k = _KERNELS.get((module, name))
    if k is None:
        from torch_neuronx.nki_hop import wrap_nki

        from deepinfer.neuron_ops.nki import rmsnorm, swiglu

        k = _KERNELS[(module, name)] = wrap_nki(getattr({"rmsnorm": rmsnorm, "swiglu": swiglu}[module], name))
    return k


def _const(value: float, like: torch.Tensor) -> torch.Tensor:
    return torch.full((1, 1), value, dtype=torch.float32, device=like.device)


# ---------------------------------------------------------------------------------------------------------------
# RMSNorm


def norm_supported(x: torch.Tensor, extra_sf_out: torch.Tensor | None) -> bool:
    return NORM_ENABLED and x.device.type == "neuron" and x.dtype in (torch.bfloat16, torch.float32) and extra_sf_out is None


def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float, out_scale: float = 1.0) -> torch.Tensor:
    """``(x.float() * rsqrt(mean(x.float()^2) + eps) * (weight.float() * out_scale)).to(x.dtype)``, any leading
    shape."""
    dim = x.shape[-1]
    w = weight.float()
    if out_scale != 1.0:
        w = w * out_scale
    out = _kernel("rmsnorm", "rms_norm")(x.reshape(-1, dim), w.reshape(1, dim), _const(eps, x))
    return out.view(x.shape)


def ue8m0_exponent(v: torch.Tensor) -> torch.Tensor:
    """Biased exponent of ``2 ** ceil(log2(v))`` for positive normal fp32 ``v``, as int32 -- the shim's
    ``((bits(v) - 1) >> 23) + 1`` without reading the bits (a dtype bitcast in a region graph crashes neuronx-cc).
    ``exp2`` of an integer is exact on the device; ``log2`` is only trusted to within one."""
    e = torch.floor(torch.log2(v))
    e = torch.where(torch.exp2(e) > v, e - 1, e)
    e = torch.where(torch.exp2(e + 1) <= v, e + 1, e)
    e = torch.where(torch.exp2(e) == v, e, e + 1)
    return e.to(torch.int32) + 127


def pack_ue8m0_int16(e: torch.Tensor) -> torch.Tensor:
    """int32 exponent bytes ``[..., g]`` -> int16 words ``[..., ceil(g / 2)]``, low byte first (arithmetic only).

    The word is assembled in fp32 (exact: |v| < 2**17) and converted fp32 -> int16 once. An int32 -> int16 narrowing
    of the strided (``0::2`` / ``1::2``) int32 result -- with or without a select -- makes neuronx-cc emit a
    same-shape reinterpret between 4- and 2-byte dtypes and die in codegen (NCC_IBCG902 "Cannot reinterpret to
    dtype with different size", issue #13) when the exponents come from the NKI RMSNorm output."""
    if e.shape[-1] % 2:
        e = torch.nn.functional.pad(e, (0, 1))
    f = e.to(torch.float32)
    lo, hi = f[..., 0::2], f[..., 1::2]
    return (lo + hi * 256.0 - torch.where(hi >= 128.0, 65536.0, 0.0)).to(torch.int16)


def fp8_cast(
    result: torch.Tensor, num_per_channels: int, round_sf: bool, use_packed_ue8m0: bool, sf_dtype: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]:
    """The norm's per-``num_per_channels`` e4m3 cast of ``result`` ``[..., n]`` (the shim's numerics). Returns the
    e4m3 tensor and the scales: int16-packed ue8m0 words ``[..., n / npc / 2]`` with ``use_packed_ue8m0``, fp32
    ``[..., n / npc]`` otherwise."""
    from deepinfer.model_executor.layers.tilelang_ops.common import FP8_AMAX_MARGIN
    from deepinfer.neuron_ops.fp8 import FP8_MAX

    q_in = result.float()
    lead, n = q_in.shape[:-1], q_in.shape[-1]
    groups = q_in.view(*lead, n // num_per_channels, num_per_channels)
    clamped = torch.clamp_min(groups.abs().amax(dim=-1), FP8_AMAX_MARGIN)
    if round_sf:
        exponent = ue8m0_exponent(clamped * (1.0 / FP8_MAX))  # fp32 product, as the shim
        inverse = torch.exp2((127 - exponent).to(torch.float32))
        if use_packed_ue8m0:
            assert sf_dtype == torch.int16, f"{sf_dtype=}"
            sf = pack_ue8m0_int16(exponent)
        else:
            sf = torch.exp2((exponent - 127).to(torch.float32))
    else:
        inverse = torch.full_like(clamped, FP8_MAX) / clamped
        sf = clamped / FP8_MAX
    q = torch.clamp(groups * inverse.unsqueeze(-1), -FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
    return q.view(q_in.shape), sf


def norm(
    x: torch.Tensor,
    weight: torch.Tensor,
    residual: torch.Tensor | None,
    eps: float,
    out: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None,
    fmt: str,
    num_per_channels: int,
    round_sf: bool,
    use_packed_ue8m0: bool,
    return_bf16: bool,
    out_scale: float,
) -> Any:
    """``tilelang_ops.norm.norm`` on the device (same arguments and result packing; no ``extra_sf_out``).

    The scale layout is Neuron's: row-major ``[..., packs]`` int16 for packed ue8m0 whether or not the column-major
    layout was asked for (``_alloc_fp8_sf``'s Neuron branch), fp32 ``[..., groups]`` otherwise."""
    do_fp8_cast = fmt != ""
    sf_out = None
    if isinstance(out, tuple):
        out, sf_out = out
    if residual is not None:
        summed = x + residual  # in x.dtype, as the shim
        residual_out = summed
    else:
        summed, residual_out = x, None
    result = rms_norm(summed, weight, eps, out_scale)

    if do_fp8_cast:
        from deepinfer import PLATFORM

        q, sf = fp8_cast(result, num_per_channels, round_sf, use_packed_ue8m0, PLATFORM.fp8_sf_type)
        if out is not None:
            out.copy_(q.view(out.shape))
            q = out
        if sf_out is not None:
            sf_out.copy_(sf.view(sf_out.shape))
            sf = sf_out
        main: Any = (q, sf)
    else:
        if out is not None:
            out.copy_(result.view(out.shape))
            result = out
        main = result
    if return_bf16:
        main = (main, result.to(torch.bfloat16))
    if residual is not None:
        main = (main, residual_out)
    return main


# ---------------------------------------------------------------------------------------------------------------
# SwiGLU


def swiglu_supported(x: torch.Tensor, fmt: str, clamped_count: torch.Tensor | None, psum: torch.Tensor | None) -> bool:
    """bf16 in, bf16 (or fp8 via the shim's cast) out; the clamp counters and the expert prefix-sum layout stay on
    the shim. fp32 output keeps the shim's fp32 activation."""
    return (
        SWIGLU_ENABLED
        and x.device.type == "neuron"
        and x.dtype == torch.bfloat16
        and fmt in ("bf16", "e4m3")
        and clamped_count is None
        and psum is None
        and x.dim() == 2
        and x.shape[1] % 2 == 0
    )


def swiglu(
    x: torch.Tensor,
    topk_weights: torch.Tensor | None,
    routed_scaling_factor: float | None,
    clamp_value: float | None,
) -> torch.Tensor:
    """``silu(min(g, c)) * clamp(u, -c, c) [* w * routed_scaling_factor]`` for ``x = [g | u]``, bf16 ``[rows, h]``."""
    if topk_weights is None:
        w = _const(1.0, x)
    else:
        w = topk_weights.to(torch.float32).reshape(-1, 1)
        if routed_scaling_factor is not None:
            w = w * routed_scaling_factor
    limit = _const(_NO_CLAMP if clamp_value is None else float(clamp_value), x)
    return _kernel("swiglu", "swiglu_fused")(x.contiguous(), w, limit)