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