DeepInfer/deepinfer/neuron_ops/fp8.py

"""fp8 on trn2: keep every ``float8_e4m3fn`` value inside the range of the hardware's e4m3 (issue #24).

trn2 has no e4m3fn. With ``--experimental-unsafe-fp8e4m3fn-as-fp8e4m3`` the compiler accepts ``float8_e4m3fn``
tensors and reads their bits as IEEE-style e4m3 (bias 7, max 240, exponent field 15 = inf/NaN). Both formats decode
identically whenever the exponent field is below 15, i.e. for |x| < 256. So on Neuron:

* activations / KV records quantize to ``FP8_MAX`` (224, as vllm-neuron#71) instead of 448. With power-of-two
  (ue8m0) or ``amax / max`` scales the scale just doubles, so dequantized values are unchanged except in the lowest
  subnormal binade;
* checkpoint fp8 weights (full ±448 range) are rescaled once at load: values halved, scales doubled
  (``rescale_fp8_weights``), exact except the lowest subnormal bit.

``DEEPINFER_NEURON_FP8_MAX=448`` restores the e4m3fn range (host-mode debugging only: the device cannot decode it).
"""

from __future__ import annotations

import functools
import os
import shlex

import numpy as np
import torch

FP8_MAX = float(os.environ.get("DEEPINFER_NEURON_FP8_MAX", "224"))
_HW_E4M3_FLAG = "--experimental-unsafe-fp8e4m3fn-as-fp8e4m3"


def configure_compiler() -> None:
    """Let neuronx-cc (eager and torch.compile) take float8_e4m3fn tensors as the hardware e4m3."""
    if FP8_MAX > 240.0:
        return
    os.environ.setdefault("UNSAFE_FP8FNCAST", "1")
    flags = os.environ.get("NEURON_CC_FLAGS", "")
    if _HW_E4M3_FLAG not in flags:
        # The option belongs to hlo2tensorizer; neuronx-cc rejects it as a top-level flag (NCC_EARG002). All
        # hlo2tensorizer options must share one --internal-hlo2tensorizer-options (a second one replaces the first),
        # so append to an existing one if the user set it.
        key = "--internal-hlo2tensorizer-options="
        if key in flags:
            head, _, rest = flags.partition(key)
            value, sep, tail = rest.partition(" --")
            value = value.strip("'\"")
            flags = f"{head}{key}'{value} {_HW_E4M3_FLAG}'" + (f" --{tail}" if sep else "")
        else:
            flags = f"{flags} {key}{_HW_E4M3_FLAG}".strip()
        os.environ["NEURON_CC_FLAGS"] = flags


def halve_e4m3fn_bytes(b: torch.Tensor) -> torch.Tensor:
    """``e4m3fn(x / 2)`` from the bits of ``x`` (uint8), round-to-nearest-even; integer ops only, any device."""
    v = b.to(torch.int16)
    sign = v & 0x80
    mag = v & 0x7F
    exp = mag >> 3
    man = mag & 7
    # exponent >= 2: decrement the exponent. exponent 1 -> subnormal (8 + m) / 2; subnormal m -> m / 2 (RNE).
    sub_in = torch.where(exp == 1, man + 8, man)
    half = sub_in >> 1
    rne = half + ((sub_in & 1) & (half & 1))
    out = torch.where(exp >= 2, mag - 8, torch.where(exp <= 1, rne, mag))
    out = torch.where(mag == 0x7F, mag, out)  # NaN stays NaN
    return (out | sign).to(torch.uint8)


_CHUNK_BYTES = 1 << 26


def _halve_inplace(data: torch.Tensor) -> None:
    """In place, on the host in fixed-size chunks (bounded int16 temporaries: the Engram tables are ~200 GB).
    Device weights make one host round trip: eager device ops recompile per shape / view offset, which made this
    step take over half an hour for the full model."""
    if data.device.type != "cpu":
        host = data.cpu()
        _halve_inplace(host)
        data.copy_(host.to(data.device))
        return
    flat = data.view(-1).view(torch.uint8)
    table = _halve_table().numpy()
    for start in range(0, flat.numel(), _CHUNK_BYTES):
        chunk = flat[start : start + _CHUNK_BYTES]
        # A clone per chunk: .numpy() on the parameter itself would pin its storage as non-resizable, and the
        # runtime adaptation later frees merged weights with untyped_storage().resize_(0).
        arr = chunk.clone().numpy()
        np.take(table, arr, out=arr)  # one 256-entry lookup per byte
        chunk.copy_(torch.from_numpy(arr))


@functools.cache
def _halve_table() -> torch.Tensor:
    # Explicitly on the host: model loading runs under a default-device context that maps to the NeuronCore.
    return halve_e4m3fn_bytes(torch.arange(256, dtype=torch.int32, device="cpu").to(torch.uint8))


def _double_scale(sf: torch.Tensor) -> None:
    if sf.device.type != "cpu":  # small: on the host, no per-shape eager compiles on the device
        host = sf.cpu()
        _double_scale(host)
        sf.copy_(host.to(sf.device))
        return
    if sf.dtype in (torch.float32, torch.bfloat16, torch.float16):
        sf.mul_(2)
        return
    # ue8m0 exponents (float8_e8m0fnu, or bytes packed into int words): exponent + 1.
    sf_bytes = sf.view(torch.uint8)
    assert not bool((sf_bytes == 255).any()), "ue8m0 scale overflow"
    sf_bytes.add_(1)


def rescale_fp8_weights(model: torch.nn.Module) -> int:
    """Halve every fp8 weight (and double its scale) so it fits the hardware e4m3. Returns the count rescaled.

    Run once, on the checkpoint parameters right after loading (before runtime adaptation merges or repacks them).
    Parameters sharing storage (tied weights) are rescaled once.
    """
    if FP8_MAX > 240.0:
        return 0
    by_storage: dict[int, torch.nn.Parameter] = {}
    for p in model.parameters():
        if p.dtype != torch.float8_e4m3fn or p.numel() == 0:
            continue
        ptr = p.untyped_storage().data_ptr()
        if getattr(p, "sf", None) is not None or ptr not in by_storage:
            by_storage[ptr] = p  # prefer the copy that carries the scale (tied aliases may not)
    count = 0
    for p in by_storage.values():
        sf = getattr(p, "sf", None)
        assert sf is not None, "fp8 weight without a scale cannot be rescaled"
        assert p.is_contiguous()
        _halve_inplace(p.data.view(torch.uint8))
        _double_scale(sf.data if isinstance(sf, torch.nn.Parameter) else sf)
        count += 1
    return count


def describe() -> str:
    return f"fp8 max {FP8_MAX:g}, NEURON_CC_FLAGS={shlex.quote(os.environ.get('NEURON_CC_FLAGS', ''))}"