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 os
import shlex

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:
    flat = data.view(-1)
    for start in range(0, flat.numel(), _CHUNK_BYTES):  # bounded int16 temporaries (Engram tables are ~200 GB)
        chunk = flat[start : start + _CHUNK_BYTES]
        chunk.copy_(halve_e4m3fn_bytes(chunk))


def _double_scale(sf: torch.Tensor) -> None:
    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', ''))}"