Add runtime/torch_ops_native.py

runtime/torch_ops_native.py

@@ -0,0 +1,233 @@
+"""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",
+]