"""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", ]