Add runtime/packed_nvfp4_linear.py

runtime/packed_nvfp4_linear.py

@@ -0,0 +1,485 @@
+"""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"
+        )