runtime/packed_nvfp4_linear.py

"""Inference-only resident NVFP4 linear prototype.

This module intentionally uses a narrow ctypes boundary. It proves packed
residency and Mage shape correctness; it is not yet a torch.compile/CUDA-graph
shipping operator.
"""

from __future__ import annotations

import atexit
import ctypes
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Final

import torch
from torch import nn


ABI_VERSION: Final = 1
FP4_BLOCK_ELEMENTS: Final = 16
SCALE_TILE_OUTER: Final = 128
SCALE_TILE_INNER: Final = 4
RELEASE_ROOT: Final = Path(__file__).resolve().parents[1]
DEFAULT_LIBRARY_PATH: Final = RELEASE_ROOT / "runtime" / "libmage_nvfp4_linear.so"


def round_up(value: int, multiple: int) -> int:
    if value <= 0 or multiple <= 0:
        raise ValueError("value and multiple must be positive")
    return ((value + multiple - 1) // multiple) * multiple


@dataclass(frozen=True)
class ScaleLayout:
    inner_dim: int
    outer_tiles: int
    num_bytes: int


def scale_layout(k: int, outer_columns: int) -> ScaleLayout:
    if k <= 0 or k % FP4_BLOCK_ELEMENTS:
        raise ValueError("K must be positive and divisible by 16")
    if outer_columns <= 0:
        raise ValueError("outer column count must be positive")
    inner_dim = round_up(k // FP4_BLOCK_ELEMENTS, SCALE_TILE_INNER)
    outer_tiles = (outer_columns + SCALE_TILE_OUTER - 1) // SCALE_TILE_OUTER
    return ScaleLayout(
        inner_dim=inner_dim,
        outer_tiles=outer_tiles,
        num_bytes=outer_tiles * inner_dim * SCALE_TILE_OUTER,
    )


def packed_weight_num_bytes(out_features: int, in_features: int) -> int:
    if out_features <= 0 or out_features % 8:
        raise ValueError("out_features must be positive and divisible by 8")
    if in_features <= 0 or in_features % 32:
        raise ValueError("in_features must be positive and divisible by 32")
    return out_features * in_features // 2


def padded_output_shape(input_shape: tuple[int, ...], out_features: int) -> tuple[int, int]:
    if not input_shape:
        raise ValueError("input must have at least one dimension")
    logical_m = 1
    for dimension in input_shape[:-1]:
        if dimension <= 0:
            raise ValueError("empty or negative leading dimensions are unsupported")
        logical_m *= dimension
    return round_up(logical_m, 8), out_features


def logical_output_shape(input_shape: tuple[int, ...], out_features: int) -> tuple[int, ...]:
    if not input_shape:
        raise ValueError("input must have at least one dimension")
    return (*input_shape[:-1], out_features)


class NativeNvfp4Library:
    """Typed ctypes access to the project-local native library."""

    def __init__(self, path: str | Path = DEFAULT_LIBRARY_PATH):
        self.path = Path(path).resolve()
        if not self.path.is_file():
            raise FileNotFoundError(
                f"resident NVFP4 library is not built: {self.path}"
            )
        self._library = ctypes.CDLL(str(self.path))
        self._bind()
        version = int(self._library.mage_nvfp4_abi_version())
        if version != ABI_VERSION:
            raise RuntimeError(
                f"resident NVFP4 ABI mismatch: Python={ABI_VERSION}, native={version}"
            )

    def _bind(self) -> None:
        library = self._library
        library.mage_nvfp4_abi_version.argtypes = []
        library.mage_nvfp4_abi_version.restype = ctypes.c_int
        library.mage_nvfp4_last_error.argtypes = []
        library.mage_nvfp4_last_error.restype = ctypes.c_char_p
        library.mage_nvfp4_packed_weight_bytes.argtypes = [
            ctypes.c_int,
            ctypes.c_int,
        ]
        library.mage_nvfp4_packed_weight_bytes.restype = ctypes.c_size_t
        library.mage_nvfp4_weight_scale_bytes.argtypes = [
            ctypes.c_int,
            ctypes.c_int,
        ]
        library.mage_nvfp4_weight_scale_bytes.restype = ctypes.c_size_t
        library.mage_nvfp4_pack_weight_bf16.argtypes = [
            ctypes.c_void_p,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_void_p,
            ctypes.c_size_t,
            ctypes.c_void_p,
            ctypes.c_size_t,
            ctypes.POINTER(ctypes.c_float),
        ]
        library.mage_nvfp4_pack_weight_bf16.restype = ctypes.c_int
        library.mage_nvfp4_create_context.argtypes = [
            ctypes.c_int,
            ctypes.POINTER(ctypes.c_void_p),
        ]
        library.mage_nvfp4_create_context.restype = ctypes.c_int
        library.mage_nvfp4_destroy_context.argtypes = [ctypes.c_void_p]
        library.mage_nvfp4_destroy_context.restype = ctypes.c_int
        library.mage_nvfp4_context_reserved_bytes.argtypes = [ctypes.c_void_p]
        library.mage_nvfp4_context_reserved_bytes.restype = ctypes.c_size_t
        library.mage_nvfp4_linear_forward.argtypes = [
            ctypes.c_void_p,
            ctypes.c_void_p,
            ctypes.c_void_p,
            ctypes.c_size_t,
            ctypes.c_void_p,
            ctypes.c_size_t,
            ctypes.c_void_p,
            ctypes.c_void_p,
            ctypes.c_void_p,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_int,
            ctypes.c_size_t,
        ]
        library.mage_nvfp4_linear_forward.restype = ctypes.c_int

    def _error(self, operation: str) -> RuntimeError:
        raw = self._library.mage_nvfp4_last_error()
        message = raw.decode("utf-8", errors="replace") if raw else "unknown error"
        return RuntimeError(f"{operation}: {message}")

    def pack_weight(
        self, weight: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if weight.device.type != "cpu":
            raise ValueError("native weight packer requires a CPU tensor")
        if weight.dtype != torch.bfloat16:
            raise TypeError("native weight packer requires BF16")
        if weight.ndim != 2 or not weight.is_contiguous():
            raise ValueError("weight must be contiguous [N,K]")
        n, k = map(int, weight.shape)
        packed_bytes = packed_weight_num_bytes(n, k)
        scale_bytes = scale_layout(k, n).num_bytes
        native_packed = int(
            self._library.mage_nvfp4_packed_weight_bytes(n, k)
        )
        native_scales = int(
            self._library.mage_nvfp4_weight_scale_bytes(n, k)
        )
        if (native_packed, native_scales) != (packed_bytes, scale_bytes):
            raise RuntimeError(
                "Python/native packed metadata disagreement: "
                f"Python={(packed_bytes, scale_bytes)}, "
                f"native={(native_packed, native_scales)}"
            )
        packed = torch.empty(packed_bytes, dtype=torch.uint8, device="cpu")
        scales = torch.empty(scale_bytes, dtype=torch.uint8, device="cpu")
        tensor_scale = ctypes.c_float()
        status = self._library.mage_nvfp4_pack_weight_bf16(
            ctypes.c_void_p(weight.data_ptr()),
            n,
            k,
            ctypes.c_void_p(packed.data_ptr()),
            packed.numel(),
            ctypes.c_void_p(scales.data_ptr()),
            scales.numel(),
            ctypes.byref(tensor_scale),
        )
        if status:
            raise self._error("packing BF16 weight")
        scale_tensor = torch.tensor(tensor_scale.value, dtype=torch.float32)
        return packed, scales, scale_tensor


_LIBRARY_LOCK = threading.Lock()
_LIBRARY: NativeNvfp4Library | None = None


def native_library() -> NativeNvfp4Library:
    global _LIBRARY
    with _LIBRARY_LOCK:
        if _LIBRARY is None:
            _LIBRARY = NativeNvfp4Library()
        return _LIBRARY


class ResidentContext:
    def __init__(self, library: NativeNvfp4Library, device_index: int, stream: int):
        self.library = library
        self.device_index = int(device_index)
        self.stream = int(stream)
        self._pointer = ctypes.c_void_p()
        self._lock = threading.Lock()
        status = self.library._library.mage_nvfp4_create_context(
            self.device_index, ctypes.byref(self._pointer)
        )
        if status:
            raise self.library._error("creating resident context")
        self._closed = False

    @property
    def pointer(self) -> ctypes.c_void_p:
        if self._closed:
            raise RuntimeError("resident NVFP4 context is closed")
        return self._pointer

    @property
    def reserved_bytes(self) -> int:
        return int(
            self.library._library.mage_nvfp4_context_reserved_bytes(self.pointer)
        )

    def close(self) -> None:
        with self._lock:
            if self._closed:
                return
            status = self.library._library.mage_nvfp4_destroy_context(
                self._pointer
            )
            if status:
                raise self.library._error("destroying resident context")
            self._closed = True
            self._pointer = ctypes.c_void_p()

    def forward(
        self,
        x: torch.Tensor,
        packed_weight: torch.Tensor,
        weight_scales: torch.Tensor,
        weight_scale: torch.Tensor,
        bias: torch.Tensor | None,
        output: torch.Tensor,
        logical_m: int,
        in_features: int,
        out_features: int,
    ) -> None:
        bias_pointer = (
            ctypes.c_void_p(bias.data_ptr()) if bias is not None else None
        )
        with self._lock:
            status = self.library._library.mage_nvfp4_linear_forward(
                self.pointer,
                ctypes.c_void_p(x.data_ptr()),
                ctypes.c_void_p(packed_weight.data_ptr()),
                packed_weight.numel(),
                ctypes.c_void_p(weight_scales.data_ptr()),
                weight_scales.numel(),
                ctypes.c_void_p(weight_scale.data_ptr()),
                bias_pointer,
                ctypes.c_void_p(output.data_ptr()),
                logical_m,
                in_features,
                out_features,
                self.stream,
            )
        if status:
            raise self.library._error("resident NVFP4 linear forward")


_CONTEXTS_LOCK = threading.Lock()
_CONTEXTS: dict[tuple[int, int], ResidentContext] = {}


def resident_context(device: torch.device) -> ResidentContext:
    if device.type != "cuda":
        raise ValueError("resident NVFP4 context requires CUDA")
    device_index = (
        torch.cuda.current_device() if device.index is None else int(device.index)
    )
    stream = int(torch.cuda.current_stream(device_index).cuda_stream)
    key = (device_index, stream)
    with _CONTEXTS_LOCK:
        context = _CONTEXTS.get(key)
        if context is None:
            context = ResidentContext(native_library(), device_index, stream)
            _CONTEXTS[key] = context
        return context


def close_all_contexts() -> None:
    with _CONTEXTS_LOCK:
        contexts = list(_CONTEXTS.values())
        _CONTEXTS.clear()
    errors: list[Exception] = []
    for context in contexts:
        try:
            context.close()
        except Exception as error:  # pragma: no cover - shutdown diagnostic
            errors.append(error)
    if errors:
        raise RuntimeError(
            "one or more resident NVFP4 contexts failed to close: "
            + "; ".join(map(str, errors))
        )


def _quiet_atexit_close() -> None:
    try:
        close_all_contexts()
    except Exception:
        # CUDA may already be shutting down. Explicit close_all_contexts() is
        # the auditable path; atexit is only a best-effort fallback.
        pass


atexit.register(_quiet_atexit_close)


class PackedNvfp4Linear(nn.Module):
    """Packed inference replacement for one BF16 ``nn.Linear``."""

    def __init__(
        self,
        in_features: int,
        out_features: int,
        packed_weight: torch.Tensor,
        weight_scales: torch.Tensor,
        weight_scale: torch.Tensor,
        bias: torch.Tensor | None,
    ):
        super().__init__()
        expected_weight = packed_weight_num_bytes(out_features, in_features)
        expected_scales = scale_layout(in_features, out_features).num_bytes
        if (
            packed_weight.dtype != torch.uint8
            or packed_weight.ndim != 1
            or not packed_weight.is_contiguous()
            or packed_weight.numel() != expected_weight
        ):
            raise ValueError("packed_weight has invalid dtype, shape, or size")
        if (
            weight_scales.dtype != torch.uint8
            or weight_scales.ndim != 1
            or not weight_scales.is_contiguous()
            or weight_scales.numel() != expected_scales
        ):
            raise ValueError("weight_scales has invalid dtype, shape, or size")
        if (
            weight_scale.dtype != torch.float32
            or weight_scale.numel() != 1
            or not weight_scale.is_contiguous()
        ):
            raise ValueError("weight_scale must be one contiguous FP32 value")
        if bias is not None and (
            bias.dtype != torch.bfloat16
            or bias.shape != (out_features,)
            or not bias.is_contiguous()
        ):
            raise ValueError("bias must be contiguous BF16 [out_features]")
        devices = {
            tensor.device
            for tensor in (packed_weight, weight_scales, weight_scale, bias)
            if tensor is not None
        }
        if len(devices) != 1:
            raise ValueError("all resident buffers must be on one device")

        self.in_features = int(in_features)
        self.out_features = int(out_features)
        self.register_buffer("packed_weight", packed_weight)
        self.register_buffer("weight_scales", weight_scales)
        self.register_buffer("weight_scale", weight_scale.reshape(()))
        self.register_buffer("bias", bias)

    @classmethod
    def from_linear(
        cls,
        linear: nn.Linear,
        device: torch.device | str,
        *,
        library: NativeNvfp4Library | None = None,
    ) -> "PackedNvfp4Linear":
        if not isinstance(linear, nn.Linear):
            raise TypeError("from_linear requires torch.nn.Linear")
        library = native_library() if library is None else library
        destination = torch.device(device)
        if destination.type != "cuda":
            raise ValueError("resident packed buffers must target CUDA")
        weight_cpu = (
            linear.weight.detach()
            .to(device="cpu", dtype=torch.bfloat16)
            .contiguous()
        )
        packed, scales, tensor_scale = library.pack_weight(weight_cpu)
        bias_cpu = (
            None
            if linear.bias is None
            else linear.bias.detach()
            .to(device="cpu", dtype=torch.bfloat16)
            .contiguous()
        )
        return cls(
            linear.in_features,
            linear.out_features,
            packed.to(destination),
            scales.to(destination),
            tensor_scale.to(destination),
            None if bias_cpu is None else bias_cpu.to(destination),
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if input.device.type != "cuda":
            raise ValueError("PackedNvfp4Linear requires a CUDA input")
        if input.device != self.packed_weight.device:
            raise ValueError("input and packed buffers are on different devices")
        if input.dtype != torch.bfloat16:
            raise TypeError("PackedNvfp4Linear requires BF16 input")
        if input.ndim < 1 or input.shape[-1] != self.in_features:
            raise ValueError(
                f"expected last dimension {self.in_features}, got "
                f"{tuple(input.shape)}"
            )
        if input.requires_grad:
            raise RuntimeError("resident NVFP4 prototype is inference-only")

        contiguous = input.reshape(-1, self.in_features).contiguous()
        logical_m = int(contiguous.shape[0])
        if logical_m <= 0:
            raise ValueError("empty inputs are unsupported")
        padded_m = round_up(logical_m, 8)
        padded_output = torch.empty(
            (padded_m, self.out_features),
            dtype=torch.bfloat16,
            device=input.device,
        )
        current_stream = torch.cuda.current_stream(input.device)
        # ctypes launches are invisible to the caching allocator. Explicitly
        # record every CUDA allocation read or written by the native call so a
        # tensor produced on another stream cannot be recycled while the
        # resident kernel is still using it.
        for tensor in (
            contiguous,
            self.packed_weight,
            self.weight_scales,
            self.weight_scale,
            self.bias,
            padded_output,
        ):
            if tensor is not None:
                tensor.record_stream(current_stream)
        context = resident_context(input.device)
        context.forward(
            contiguous,
            self.packed_weight,
            self.weight_scales,
            self.weight_scale,
            self.bias,
            padded_output,
            logical_m,
            self.in_features,
            self.out_features,
        )
        logical = padded_output[:logical_m]
        return logical.view(*input.shape[:-1], self.out_features)

    def extra_repr(self) -> str:
        return (
            f"in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"bias={self.bias is not None}, inference_only=True"
        )