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