runtime/torch_ops_native.py

"""Loader and fake/meta shim for the compiled resident NVFP4 torch op.

The compiled bridge lives in ``native/torch_op`` and, when built, moves CUDA
execution out of Python+ctypes and into a C++ dispatcher implementation that
calls the existing resident C ABI directly.

This shim keeps CPU/meta tests lightweight:

- if the compiled bridge exists, load it first;
- otherwise define a temporary Python schema fallback for shape-only tests;
- always register fake/meta implementations in Python.
"""

from __future__ import annotations

import atexit
from pathlib import Path

import torch

from packed_nvfp4_linear import PackedNvfp4Linear, scale_layout


RELEASE_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_EXTENSION_PATH = RELEASE_ROOT / "runtime" / "libmage_nvfp4_torch_op.so"

OP_NAMESPACE = "mage_nvfp4"
OP_NAME = "sm120_linear_native"
OP_QUALNAME = f"{OP_NAMESPACE}::{OP_NAME}"

_SCHEMA_FALLBACK_LIB: torch.library.Library | None = None
_META_LIB: torch.library.Library | None = None
_EXTENSION_LOADED = False
_FAKE_REGISTERED = False
_META_REGISTERED = False


def _expected_weight_bytes(out_features: int, in_features: int) -> int:
    if out_features <= 0 or out_features % 8:
        raise ValueError("resident NVFP4 requires out_features divisible by 8")
    if in_features <= 0 or in_features % 32:
        raise ValueError("resident NVFP4 requires in_features divisible by 32")
    return (out_features * in_features) // 2


def _check_common_shapes(
    input: torch.Tensor,
    packed_weight: torch.Tensor,
    weight_scales: torch.Tensor,
    weight_scale: torch.Tensor,
    bias: torch.Tensor | None,
    in_features: int,
    out_features: int,
) -> None:
    if input.ndim < 1:
        raise ValueError("resident NVFP4 input must have at least one dimension")
    if input.dtype != torch.bfloat16:
        raise TypeError("resident NVFP4 input must be bfloat16")
    if input.requires_grad:
        raise RuntimeError("resident NVFP4 torch op is inference-only")
    if int(input.shape[-1]) != int(in_features):
        raise ValueError(
            f"expected input last dimension {in_features}, got {tuple(input.shape)}"
        )
    if packed_weight.dtype != torch.uint8 or packed_weight.ndim != 1:
        raise ValueError("packed_weight must be a 1D uint8 tensor")
    if weight_scales.dtype != torch.uint8 or weight_scales.ndim != 1:
        raise ValueError("weight_scales must be a 1D uint8 tensor")
    if weight_scale.dtype != torch.float32 or weight_scale.numel() != 1:
        raise ValueError("weight_scale must be a scalar or length-1 float32 tensor")
    if bias is not None and (
        bias.dtype != torch.bfloat16 or bias.ndim != 1 or int(bias.numel()) != int(out_features)
    ):
        raise ValueError("bias must be a 1D bfloat16 tensor with length out_features")
    for tensor in (packed_weight, weight_scales, weight_scale, bias):
        if tensor is not None and tensor.device != input.device:
            raise ValueError("input and resident buffers must be on the same device")
        if tensor is not None and not tensor.is_contiguous():
            raise ValueError("resident packed buffers must be contiguous")

    expected_weight_bytes = _expected_weight_bytes(int(out_features), int(in_features))
    expected_scale_bytes = scale_layout(int(in_features), int(out_features)).num_bytes
    if int(packed_weight.numel()) != expected_weight_bytes:
        raise ValueError(
            f"packed_weight size mismatch: expected {expected_weight_bytes}, got {packed_weight.numel()}"
        )
    if int(weight_scales.numel()) != expected_scale_bytes:
        raise ValueError(
            f"weight_scales size mismatch: expected {expected_scale_bytes}, got {weight_scales.numel()}"
        )


def _logical_output(input: torch.Tensor, out_features: int) -> torch.Tensor:
    return input.new_empty((*input.shape[:-1], int(out_features)), dtype=torch.bfloat16)


def ensure_native_sm120_schema(
    *,
    extension_path: str | Path = DEFAULT_EXTENSION_PATH,
    allow_python_schema_fallback: bool = True,
) -> bool:
    global _EXTENSION_LOADED, _SCHEMA_FALLBACK_LIB
    extension_path = Path(extension_path)
    if not _EXTENSION_LOADED and extension_path.is_file():
        torch.ops.load_library(str(extension_path.resolve()))
        _EXTENSION_LOADED = True
        return True
    if _EXTENSION_LOADED:
        return True
    if not allow_python_schema_fallback:
        return False
    if _SCHEMA_FALLBACK_LIB is None:
        lib = torch.library.Library(OP_NAMESPACE, "FRAGMENT")
        lib.define(
            "sm120_linear_native(Tensor input, Tensor packed_weight, Tensor weight_scales, Tensor weight_scale, Tensor? bias, int in_features, int out_features) -> Tensor"
        )
        lib.define("clear_native_contexts() -> ()")
        _SCHEMA_FALLBACK_LIB = lib
    return False


def _register_meta_impl() -> None:
    global _META_LIB, _META_REGISTERED
    if _META_REGISTERED:
        return
    _META_LIB = torch.library.Library(OP_NAMESPACE, "IMPL", "Meta")
    _META_LIB.impl(
        OP_NAME,
        lambda input, packed_weight, weight_scales, weight_scale, bias, in_features, out_features: (
            _check_common_shapes(
                input,
                packed_weight,
                weight_scales,
                weight_scale,
                bias,
                in_features,
                out_features,
            ),
            _logical_output(input, int(out_features)),
        )[1],
    )
    _META_REGISTERED = True


def _register_fake_impl() -> None:
    global _FAKE_REGISTERED
    if _FAKE_REGISTERED:
        return

    @torch.library.register_fake(OP_QUALNAME)
    def _fake(
        input: torch.Tensor,
        packed_weight: torch.Tensor,
        weight_scales: torch.Tensor,
        weight_scale: torch.Tensor,
        bias: torch.Tensor | None,
        in_features: int,
        out_features: int,
    ) -> torch.Tensor:
        _check_common_shapes(
            input,
            packed_weight,
            weight_scales,
            weight_scale,
            bias,
            in_features,
            out_features,
        )
        return _logical_output(input, int(out_features))

    _FAKE_REGISTERED = True


def initialize_native_sm120_op(
    *,
    extension_path: str | Path = DEFAULT_EXTENSION_PATH,
    allow_python_schema_fallback: bool = True,
) -> bool:
    loaded = ensure_native_sm120_schema(
        extension_path=extension_path,
        allow_python_schema_fallback=allow_python_schema_fallback,
    )
    _register_meta_impl()
    _register_fake_impl()
    return loaded


def close_native_contexts() -> None:
    """Synchronize and release native contexts when the compiled bridge is loaded."""
    if _EXTENSION_LOADED:
        torch.ops.mage_nvfp4.clear_native_contexts()


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


class PackedNvfp4LinearNativeOp(PackedNvfp4Linear):
    """Thin wrapper over resident packed buffers using the native torch op."""

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if input.device.type not in {"cuda", "meta"}:
            raise ValueError("PackedNvfp4LinearNativeOp requires a CUDA or meta input")
        return torch.ops.mage_nvfp4.sm120_linear_native(
            input,
            self.packed_weight,
            self.weight_scales,
            self.weight_scale,
            self.bias,
            self.in_features,
            self.out_features,
        )


initialize_native_sm120_op()
atexit.register(_quiet_atexit_close)


__all__ = [
    "DEFAULT_EXTENSION_PATH",
    "OP_NAME",
    "OP_NAMESPACE",
    "OP_QUALNAME",
    "PackedNvfp4LinearNativeOp",
    "close_native_contexts",
    "ensure_native_sm120_schema",
    "initialize_native_sm120_op",
]