Add runtime/fused_gelu_up_runtime.py

runtime/fused_gelu_up_runtime.py

@@ -0,0 +1,236 @@
+"""Experimental Python bridge for the side-by-side fused NVFP4 GELU-up ABI."""
+
+from __future__ import annotations
+
+import ctypes
+from pathlib import Path
+from typing import Any
+
+
+class FusedGeluUpLibrary:
+    def __init__(self, library_path: Path, torch: Any):
+        self.torch = torch
+        self._enabled = True
+        self._closed = False
+        self._installed_modules: tuple[str, ...] = ()
+        self.library_path = library_path.resolve()
+        self.library = ctypes.CDLL(
+            str(self.library_path),
+            mode=ctypes.RTLD_LOCAL,
+        )
+        self.library.mage_nvfp4_last_error.argtypes = []
+        self.library.mage_nvfp4_last_error.restype = ctypes.c_char_p
+        self.library.mage_nvfp4_create_context.argtypes = [
+            ctypes.c_int,
+            ctypes.POINTER(ctypes.c_void_p),
+        ]
+        self.library.mage_nvfp4_create_context.restype = ctypes.c_int
+        self.library.mage_nvfp4_destroy_context.argtypes = [ctypes.c_void_p]
+        self.library.mage_nvfp4_destroy_context.restype = ctypes.c_int
+        self.forward_function = self.library.mage_nvfp4_gelu_up_forward
+        self.forward_function.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,
+        ]
+        self.forward_function.restype = ctypes.c_int
+        context = ctypes.c_void_p()
+        status = self.library.mage_nvfp4_create_context(
+            torch.cuda.current_device(),
+            ctypes.byref(context),
+        )
+        if status:
+            raise RuntimeError(self.last_error("creating fused GELU-up context"))
+        self.context = context
+
+    @property
+    def enabled(self) -> bool:
+        return self._enabled and not self._closed and bool(self.context)
+
+    @property
+    def installed_modules(self) -> tuple[str, ...]:
+        return self._installed_modules
+
+    def set_enabled(self, enabled: bool) -> None:
+        if enabled and self._closed:
+            raise RuntimeError("fused GELU-up runtime is closed")
+        self._enabled = bool(enabled)
+
+    def set_installed_modules(self, modules: list[str]) -> None:
+        self._installed_modules = tuple(modules)
+
+    def last_error(self, operation: str) -> str:
+        raw = self.library.mage_nvfp4_last_error()
+        message = raw.decode("utf-8", errors="replace") if raw else "unknown error"
+        return f"{operation}: {message}"
+
+    def forward(self, value: Any, projection: Any) -> Any:
+        if not self.enabled:
+            raise RuntimeError("fused GELU-up runtime is unavailable")
+        torch = self.torch
+        flattened = value.reshape(-1, projection.in_features).contiguous()
+        logical_m = int(flattened.shape[0])
+        padded_m = ((logical_m + 7) // 8) * 8
+        output = torch.empty(
+            (padded_m, projection.out_features),
+            dtype=torch.bfloat16,
+            device=flattened.device,
+        )
+        status = self.forward_function(
+            self.context,
+            ctypes.c_void_p(flattened.data_ptr()),
+            ctypes.c_void_p(projection.packed_weight.data_ptr()),
+            projection.packed_weight.numel(),
+            ctypes.c_void_p(projection.weight_scales.data_ptr()),
+            projection.weight_scales.numel(),
+            ctypes.c_void_p(projection.weight_scale.data_ptr()),
+            ctypes.c_void_p(projection.bias.data_ptr()),
+            ctypes.c_void_p(output.data_ptr()),
+            logical_m,
+            projection.in_features,
+            projection.out_features,
+            int(torch.cuda.current_stream().cuda_stream),
+        )
+        if status:
+            raise RuntimeError(self.last_error("running fused GELU-up"))
+        return output.narrow(0, 0, logical_m).view(
+            (*value.shape[:-1], projection.out_features)
+        )
+
+    def close(self) -> None:
+        self._enabled = False
+        self._closed = True
+        if self.context:
+            status = self.library.mage_nvfp4_destroy_context(self.context)
+            self.context = ctypes.c_void_p()
+            if status:
+                raise RuntimeError(
+                    self.last_error("destroying fused GELU-up context")
+                )
+
+
+def install_fused_gelu_up(
+    transformer: Any,
+    *,
+    library_path: Path,
+    torch: Any,
+    stream_names: tuple[str, ...] = ("img_mlp", "txt_mlp"),
+) -> tuple[FusedGeluUpLibrary, dict[str, Any]]:
+    """Replace all Mage GELU-up modules while preserving packed parameters."""
+
+    allowed_streams = ("img_mlp", "txt_mlp")
+    requested_streams = tuple(dict.fromkeys(stream_names))
+    if not requested_streams or any(
+        name not in allowed_streams for name in requested_streams
+    ):
+        raise RuntimeError(
+            "fused GELU-up streams must be a non-empty subset of "
+            "('img_mlp', 'txt_mlp')"
+        )
+    runtime = FusedGeluUpLibrary(library_path, torch)
+
+    class FusedGeluUp(torch.nn.Module):
+        is_mage_fused_gelu_up_wrapper = True
+        _xpo3_fused_gelu_up = True
+
+        def __init__(self, activation: Any):
+            super().__init__()
+            self.original_activation = activation
+
+        @property
+        def proj(self) -> Any:
+            return self.original_activation.proj
+
+        @property
+        def projection(self) -> Any:
+            return self.original_activation.proj
+
+        def forward(self, hidden_states: Any) -> Any:
+            if runtime.enabled:
+                return runtime.forward(hidden_states, self.projection)
+            return self.original_activation(hidden_states)
+
+    installed: list[str] = []
+    skipped_bf16: list[str] = []
+    skipped_replaced: list[str] = []
+    try:
+        for block_index, block in enumerate(transformer.transformer_blocks):
+            for stream_name in requested_streams:
+                feed_forward = getattr(block, stream_name)
+                if not hasattr(feed_forward, "net"):
+                    if stream_name == "img_mlp":
+                        skipped_replaced.append(
+                            f"transformer_blocks.{block_index}.{stream_name}"
+                        )
+                        continue
+                    raise RuntimeError(
+                        f"missing GELU net at block {block_index} {stream_name}"
+                    )
+                activation = feed_forward.net[0]
+                if getattr(activation, "approximate", None) != "tanh":
+                    raise RuntimeError(
+                        f"expected tanh GELU at block {block_index} {stream_name}"
+                    )
+                projection = getattr(activation, "proj", None)
+                required = (
+                    "packed_weight",
+                    "weight_scales",
+                    "weight_scale",
+                    "bias",
+                    "in_features",
+                    "out_features",
+                )
+                if projection is None:
+                    raise RuntimeError(
+                        f"missing GELU projection at block {block_index} "
+                        f"{stream_name}"
+                    )
+                if any(
+                    not hasattr(projection, name) for name in required
+                ):
+                    skipped_bf16.append(
+                        f"transformer_blocks.{block_index}."
+                        f"{stream_name}.net.0"
+                    )
+                    continue
+                if (
+                    int(projection.in_features) != 3072
+                    or int(projection.out_features) != 12288
+                ):
+                    raise RuntimeError(
+                        "unexpected Mage GELU-up dimensions at "
+                        f"block {block_index} {stream_name}"
+                    )
+                module_name = (
+                    f"transformer_blocks.{block_index}.{stream_name}.net.0"
+                )
+                feed_forward.net[0] = FusedGeluUp(activation)
+                installed.append(module_name)
+    except Exception:
+        runtime.close()
+        raise
+
+    runtime.set_installed_modules(installed)
+    return runtime, {
+        "mode": "nvfp4_gelu_up",
+        "library_path": str(library_path.resolve()),
+        "enabled": runtime.enabled,
+        "stream_names": list(requested_streams),
+        "module_count": len(installed),
+        "modules": installed,
+        "installed_modules": installed,
+        "skipped_bf16_count": len(skipped_bf16),
+        "skipped_bf16_modules": skipped_bf16,
+        "skipped_replaced_count": len(skipped_replaced),
+        "skipped_replaced_modules": skipped_replaced,
+    }