"""Portable selected-block FP4 image-MLP bridge runtime for the ComfyUI plugin.""" from __future__ import annotations import ctypes import math from dataclasses import dataclass, field from pathlib import Path from typing import Any, Iterable, Mapping RUNTIME_ROOT = Path(__file__).resolve().parent ABI_VERSION = 1 UP_IN_FEATURES = 3072 UP_OUT_FEATURES = 12288 DOWN_OUT_FEATURES = 3072 TRANSFORMER_BLOCK_COUNT = 12 BLOCK0_IMG_MLP_MODULE = "transformer_blocks.0.img_mlp" UP_MODULE = f"{BLOCK0_IMG_MLP_MODULE}.net.0.proj" DOWN_MODULE = f"{BLOCK0_IMG_MLP_MODULE}.net.2" def _tensor_bytes(tensor: Any) -> int: return int(tensor.numel() * tensor.element_size()) def resolve_existing_runtime_file( path: str | Path, label: str, *, base_dir: str | Path | None = None, ) -> Path: root = RUNTIME_ROOT if base_dir is None else Path(base_dir).expanduser().resolve() candidate = Path(path).expanduser() resolved = (root / candidate).resolve() if not candidate.is_absolute() else candidate.resolve() if not resolved.is_file(): raise RuntimeError(f"{label} must resolve to an existing file: {resolved}") return resolved def _require_positive_finite(value: float, label: str) -> float: if not math.isfinite(value) or value <= 0.0: raise RuntimeError(f"{label} must be finite and positive; got {value!r}") return float(value) def _require_block_index(block_index: int) -> int: value = int(block_index) if value < 0 or value >= TRANSFORMER_BLOCK_COUNT: raise RuntimeError( f"bridge block index must be in [0, {TRANSFORMER_BLOCK_COUNT - 1}], got {value}" ) return value def _module_names_for_block(block_index: int) -> tuple[str, str, str]: base = f"transformer_blocks.{block_index}.img_mlp" return ( base, f"{base}.net.0.proj", f"{base}.net.2", ) def normalize_block_tensor_scales( block_tensor_scales: str | Mapping[int | str, float | str], ) -> dict[int, float]: if isinstance(block_tensor_scales, str): text = block_tensor_scales.strip() if not text: raise RuntimeError("bridge block scale map must not be empty") parsed: dict[int | str, float | str] = {} for entry in text.split(","): item = entry.strip() if not item: raise RuntimeError("bridge block scale map contains an empty entry") if "=" not in item: raise RuntimeError( "bridge block scale map entries must look like block=scale" ) block_text, scale_text = item.split("=", 1) block_key = block_text.strip() if block_key in parsed: raise RuntimeError(f"duplicate bridge block index {block_key}") parsed[block_key] = scale_text.strip() block_tensor_scales = parsed if not block_tensor_scales: raise RuntimeError("bridge candidate requires at least one selected block") normalized: dict[int, float] = {} for raw_block_index, raw_scale in block_tensor_scales.items(): try: block_index = _require_block_index(int(raw_block_index)) except (TypeError, ValueError) as error: raise RuntimeError( f"invalid bridge block index {raw_block_index!r}" ) from error if block_index in normalized: raise RuntimeError(f"duplicate bridge block index {block_index}") try: scale_value = float(raw_scale) except (TypeError, ValueError) as error: raise RuntimeError( f"invalid bridge fixed tensor scale {raw_scale!r} for block {block_index}" ) from error normalized[block_index] = _require_positive_finite( scale_value, f"bridge fixed tensor scale for block {block_index}", ) return dict(sorted(normalized.items())) def normalize_selected_blocks( selected_blocks: None | int | str | Iterable[int], *, allowed_blocks: Iterable[int] | None = None, ) -> set[int]: if selected_blocks is None: return set() if isinstance(selected_blocks, int): items: list[int | str] = [selected_blocks] elif isinstance(selected_blocks, str): text = selected_blocks.strip() if not text: raise RuntimeError("selected_blocks must not be empty") items = [item.strip() for item in text.split(",")] else: items = list(selected_blocks) normalized = {_require_block_index(int(item)) for item in items} if allowed_blocks is not None: allowed = {_require_block_index(int(item)) for item in allowed_blocks} unknown = sorted(normalized - allowed) if unknown: raise RuntimeError(f"selected blocks were not installed: {unknown}") return normalized def _require_packed_projection( projection: Any, *, label: str, in_features: int, out_features: int, torch: Any, ) -> None: required = ( "packed_weight", "weight_scales", "weight_scale", "bias", "in_features", "out_features", ) if projection is None or any(not hasattr(projection, name) for name in required): raise RuntimeError(f"{label} is not a packed NVFP4 projection") if int(projection.in_features) != in_features: raise RuntimeError( f"{label} in_features changed: {int(projection.in_features)} != {in_features}" ) if int(projection.out_features) != out_features: raise RuntimeError( f"{label} out_features changed: {int(projection.out_features)} != {out_features}" ) if projection.bias is None: raise RuntimeError(f"{label} requires a resident BF16 bias") tensors = ( ("packed_weight", projection.packed_weight, torch.uint8), ("weight_scales", projection.weight_scales, torch.uint8), ("weight_scale", projection.weight_scale, torch.float32), ("bias", projection.bias, torch.bfloat16), ) for tensor_label, tensor, dtype in tensors: if tensor.device.type != "cuda": raise RuntimeError(f"{label}.{tensor_label} must be CUDA-resident") if tensor.dtype != dtype: raise RuntimeError( f"{label}.{tensor_label} dtype changed: {tensor.dtype} != {dtype}" ) if not tensor.is_contiguous(): raise RuntimeError(f"{label}.{tensor_label} must be contiguous") def _require_exact_numel( tensor: Any, *, expected: int, label: str, ) -> None: if int(tensor.numel()) != int(expected): raise RuntimeError( f"{label} size changed: {int(tensor.numel())} != {int(expected)}" ) def _record_stream_for_tensors(current_stream: Any, *tensors: Any) -> None: for tensor in tensors: if tensor is not None: tensor.record_stream(current_stream) def _resolve_up_projection(activation: Any, module_name: str) -> tuple[Any, str]: projection = getattr(activation, "proj", None) if projection is None: projection = getattr(activation, "projection", None) if projection is None: raise RuntimeError(f"{module_name} is missing a .proj/.projection projection") if ( getattr(activation, "is_mage_fused_gelu_up_wrapper", False) or getattr(activation, "_xpo3_fused_gelu_up", False) ): return projection, "fused_gelu_up_wrapper" if getattr(activation, "approximate", None) == "tanh": return projection, "tanh_gelu" raise RuntimeError( f"{module_name} activation changed; expected tanh GELU or _xpo3_fused_gelu_up wrapper" ) class BridgeUpLibrary: def __init__(self, path: Path, device: int): self.path = path self.library = ctypes.CDLL(str(path), mode=ctypes.RTLD_LOCAL) self.context = ctypes.c_void_p() self._bind() if int(self.library.mage_nvfp4_bridge_up_abi_version()) != ABI_VERSION: raise RuntimeError("bridge-up ABI mismatch") status = self.library.mage_nvfp4_bridge_up_create_context( int(device), ctypes.byref(self.context) ) if status: raise self.error("creating bridge-up context") def _bind(self) -> None: library = self.library library.mage_nvfp4_bridge_up_abi_version.argtypes = [] library.mage_nvfp4_bridge_up_abi_version.restype = ctypes.c_int library.mage_nvfp4_bridge_up_last_error.argtypes = [] library.mage_nvfp4_bridge_up_last_error.restype = ctypes.c_char_p for suffix in ( "input_bytes", "weight_bytes", "weight_scale_bytes", "output_fp4_bytes", "output_scale_bytes", ): function = getattr(library, f"mage_nvfp4_bridge_up_{suffix}") function.argtypes = [ctypes.c_int, ctypes.c_int] function.restype = ctypes.c_size_t library.mage_nvfp4_bridge_up_bias_bytes.argtypes = [ctypes.c_int] library.mage_nvfp4_bridge_up_bias_bytes.restype = ctypes.c_size_t library.mage_nvfp4_bridge_up_scalar_bytes.argtypes = [] library.mage_nvfp4_bridge_up_scalar_bytes.restype = ctypes.c_size_t library.mage_nvfp4_bridge_up_create_context.argtypes = [ ctypes.c_int, ctypes.POINTER(ctypes.c_void_p), ] library.mage_nvfp4_bridge_up_create_context.restype = ctypes.c_int library.mage_nvfp4_bridge_up_destroy_context.argtypes = [ctypes.c_void_p] library.mage_nvfp4_bridge_up_destroy_context.restype = ctypes.c_int library.mage_nvfp4_bridge_up_context_reserved_bytes.argtypes = [ ctypes.c_void_p ] library.mage_nvfp4_bridge_up_context_reserved_bytes.restype = ctypes.c_size_t forward = library.mage_nvfp4_bridge_up_forward forward.argtypes = [ 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_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_size_t, ] forward.restype = ctypes.c_int self.forward_function = forward def error(self, operation: str) -> RuntimeError: raw = self.library.mage_nvfp4_bridge_up_last_error() message = raw.decode("utf-8", errors="replace") if raw else "unknown native error" return RuntimeError(f"{operation}: {message}") def helper(self, suffix: str, *dimensions: int) -> int: function = getattr(self.library, f"mage_nvfp4_bridge_up_{suffix}") value = int(function(*dimensions)) if value == 0: raise self.error(f"querying bridge-up {suffix}") return value def scalar_bytes(self) -> int: value = int(self.library.mage_nvfp4_bridge_up_scalar_bytes()) if value == 0: raise self.error("querying bridge-up scalar bytes") return value @property def reserved_bytes(self) -> int: return int( self.library.mage_nvfp4_bridge_up_context_reserved_bytes(self.context) ) def forward( self, input_bf16: Any, packed_weight: Any, packed_weight_scales: Any, weight_tensor_scale: Any, bias_bf16: Any, bridge_norm_constant: Any, output_fp4: Any, output_scales: Any, logical_m: int, stream: int, ) -> None: status = self.forward_function( self.context, ctypes.c_void_p(input_bf16.data_ptr()), _tensor_bytes(input_bf16), ctypes.c_void_p(packed_weight.data_ptr()), _tensor_bytes(packed_weight), ctypes.c_void_p(packed_weight_scales.data_ptr()), _tensor_bytes(packed_weight_scales), ctypes.c_void_p(weight_tensor_scale.data_ptr()), _tensor_bytes(weight_tensor_scale), ctypes.c_void_p(bias_bf16.data_ptr()), _tensor_bytes(bias_bf16), ctypes.c_void_p(bridge_norm_constant.data_ptr()), _tensor_bytes(bridge_norm_constant), ctypes.c_void_p(output_fp4.data_ptr()), _tensor_bytes(output_fp4), ctypes.c_void_p(output_scales.data_ptr()), _tensor_bytes(output_scales), int(logical_m), UP_IN_FEATURES, UP_OUT_FEATURES, int(stream), ) if status: raise self.error("running bridge-up") def close(self) -> None: if not self.context: return status = self.library.mage_nvfp4_bridge_up_destroy_context(self.context) self.context = ctypes.c_void_p() if status: raise self.error("destroying bridge-up context") class PrequantizedDownLibrary: def __init__(self, path: Path, device: int): self.path = path self.library = ctypes.CDLL(str(path), mode=ctypes.RTLD_LOCAL) self.context = ctypes.c_void_p() self._bind() if ( int(self.library.mage_nvfp4_prequantized_down_abi_version()) != ABI_VERSION ): raise RuntimeError("bridge-down ABI mismatch") status = self.library.mage_nvfp4_prequantized_down_create_context( int(device), ctypes.byref(self.context) ) if status: raise self.error("creating bridge-down context") def _bind(self) -> None: library = self.library prefix = "mage_nvfp4_prequantized_down_" abi = getattr(library, f"{prefix}abi_version") abi.argtypes = [] abi.restype = ctypes.c_int last_error = getattr(library, f"{prefix}last_error") last_error.argtypes = [] last_error.restype = ctypes.c_char_p for suffix in ( "activation_bytes", "activation_scale_bytes", "weight_bytes", "weight_scale_bytes", "output_bytes", ): function = getattr(library, f"{prefix}{suffix}") function.argtypes = [ctypes.c_int, ctypes.c_int] function.restype = ctypes.c_size_t bias_bytes = getattr(library, f"{prefix}bias_bytes") bias_bytes.argtypes = [ctypes.c_int] bias_bytes.restype = ctypes.c_size_t create = getattr(library, f"{prefix}create_context") create.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_void_p)] create.restype = ctypes.c_int destroy = getattr(library, f"{prefix}destroy_context") destroy.argtypes = [ctypes.c_void_p] destroy.restype = ctypes.c_int reserved = getattr(library, f"{prefix}context_reserved_bytes") reserved.argtypes = [ctypes.c_void_p] reserved.restype = ctypes.c_size_t forward = getattr(library, f"{prefix}forward") forward.argtypes = [ 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_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_size_t, ] forward.restype = ctypes.c_int self.forward_function = forward def error(self, operation: str) -> RuntimeError: raw = self.library.mage_nvfp4_prequantized_down_last_error() message = raw.decode("utf-8", errors="replace") if raw else "unknown native error" return RuntimeError(f"{operation}: {message}") def helper(self, suffix: str, *dimensions: int) -> int: function = getattr(self.library, f"mage_nvfp4_prequantized_down_{suffix}") value = int(function(*dimensions)) if value == 0: raise self.error(f"querying bridge-down {suffix}") return value @property def reserved_bytes(self) -> int: return int( self.library.mage_nvfp4_prequantized_down_context_reserved_bytes( self.context ) ) def forward( self, activation_fp4: Any, activation_scales: Any, activation_tensor_scale: Any, packed_weight: Any, packed_weight_scales: Any, weight_tensor_scale: Any, bias_bf16: Any, output_bf16: Any, logical_m: int, stream: int, ) -> None: status = self.forward_function( self.context, ctypes.c_void_p(activation_fp4.data_ptr()), _tensor_bytes(activation_fp4), ctypes.c_void_p(activation_scales.data_ptr()), _tensor_bytes(activation_scales), ctypes.c_void_p(activation_tensor_scale.data_ptr()), _tensor_bytes(activation_tensor_scale), ctypes.c_void_p(packed_weight.data_ptr()), _tensor_bytes(packed_weight), ctypes.c_void_p(packed_weight_scales.data_ptr()), _tensor_bytes(packed_weight_scales), ctypes.c_void_p(weight_tensor_scale.data_ptr()), _tensor_bytes(weight_tensor_scale), ctypes.c_void_p(bias_bf16.data_ptr()), _tensor_bytes(bias_bf16), ctypes.c_void_p(output_bf16.data_ptr()), _tensor_bytes(output_bf16), int(logical_m), UP_OUT_FEATURES, DOWN_OUT_FEATURES, int(stream), ) if status: raise self.error("running bridge-down") def close(self) -> None: if not self.context: return status = self.library.mage_nvfp4_prequantized_down_destroy_context( self.context ) self.context = ctypes.c_void_p() if status: raise self.error("destroying bridge-down context") @dataclass class _BridgeBuffers: payload: Any scales: Any output: Any @dataclass(frozen=True) class _BridgeBlockBinding: block_index: int module: str up_module: str down_module: str up_projection: Any down_projection: Any fixed_tensor_scale: float bridge_tensor_scale: Any bridge_norm_constant: Any activation_mode: str @dataclass class _BlockTelemetry: native_calls: int = 0 fallback_calls: int = 0 fallback_reasons: dict[str, int] = field(default_factory=dict) last_route: str | None = None last_logical_m: int | None = None last_stream: int | None = None def record(self, *, native: bool, reason: str, logical_m: int | None, stream: int | None) -> None: self.last_route = reason self.last_logical_m = logical_m self.last_stream = stream if native: self.native_calls += 1 return self.fallback_calls += 1 self.fallback_reasons[reason] = self.fallback_reasons.get(reason, 0) + 1 def snapshot(self) -> dict[str, Any]: return { "native_calls": self.native_calls, "fallback_calls": self.fallback_calls, "fallback_reasons": dict(sorted(self.fallback_reasons.items())), "last_route": self.last_route, "last_logical_m": self.last_logical_m, "last_stream": self.last_stream, } class BlockImgMlpBridgeRuntime: def __init__( self, *, bridge_up_library_path: Path, bridge_down_library_path: Path, torch: Any, device_index: int | None = None, bridge_up_factory: Any = BridgeUpLibrary, bridge_down_factory: Any = PrequantizedDownLibrary, ) -> None: self.torch = torch self.bridge_up_library_path = Path(bridge_up_library_path).resolve() self.bridge_down_library_path = Path(bridge_down_library_path).resolve() self.device_index = ( int(device_index) if device_index is not None else int(torch.cuda.current_device()) ) self.bridge_up = bridge_up_factory( self.bridge_up_library_path, self.device_index, ) self.bridge_down = bridge_down_factory( self.bridge_down_library_path, self.device_index, ) scalar_bytes = self.bridge_up.scalar_bytes() if scalar_bytes != 4: raise RuntimeError( f"bridge scalar ABI changed: expected 4 bytes, got {scalar_bytes}" ) self._buffers_by_m: dict[int, _BridgeBuffers] = {} self._bindings_by_block: dict[int, _BridgeBlockBinding] = {} self._telemetry_by_block: dict[int, _BlockTelemetry] = {} self._enabled = True self._enabled_blocks: set[int] = set() self._bound_stream: int | None = None self._closed = False def _telemetry(self, block_index: int) -> _BlockTelemetry: return self._telemetry_by_block.setdefault(block_index, _BlockTelemetry()) def _logical_m_reason(self, logical_m: int) -> str | None: if logical_m <= 0: return "logical_m_non_positive" if logical_m % 8: return "logical_m_not_divisible_by_8" if logical_m > 6400: return "logical_m_above_max" return None def _require_logical_m(self, logical_m: int) -> None: reason = self._logical_m_reason(logical_m) if reason == "logical_m_non_positive": raise RuntimeError("bridge candidate requires a positive logical M") if reason == "logical_m_not_divisible_by_8": raise RuntimeError( "bridge candidate only supports logical M divisible by 8" ) if reason == "logical_m_above_max": raise RuntimeError( "bridge candidate only supports logical M up to 6400" ) def _buffers(self, logical_m: int, device: Any) -> _BridgeBuffers: self._require_logical_m(logical_m) cached = self._buffers_by_m.get(int(logical_m)) if cached is not None: return cached payload_bytes = self.bridge_up.helper("output_fp4_bytes", logical_m, UP_OUT_FEATURES) scale_bytes = self.bridge_up.helper("output_scale_bytes", logical_m, UP_OUT_FEATURES) output_bytes = self.bridge_down.helper("output_bytes", logical_m, DOWN_OUT_FEATURES) output_elements = output_bytes // self.torch.tensor( [], dtype=self.torch.bfloat16, ).element_size() expected_elements = int(logical_m) * DOWN_OUT_FEATURES if output_elements != expected_elements: raise RuntimeError( "bridge-down output byte contract changed: " f"{output_elements} elements != {expected_elements}" ) buffers = _BridgeBuffers( payload=self.torch.empty(payload_bytes, dtype=self.torch.uint8, device=device), scales=self.torch.empty(scale_bytes, dtype=self.torch.uint8, device=device), output=self.torch.empty( expected_elements, dtype=self.torch.bfloat16, device=device, ).view(int(logical_m), DOWN_OUT_FEATURES), ) self._buffers_by_m[int(logical_m)] = buffers return buffers def bind_block( self, *, block_index: int, module: str, up_module: str, down_module: str, up_projection: Any, down_projection: Any, fixed_tensor_scale: float, activation_mode: str, ) -> _BridgeBlockBinding: checked_block_index = _require_block_index(block_index) if checked_block_index in self._bindings_by_block: raise RuntimeError(f"bridge block {checked_block_index} already installed") scale_value = _require_positive_finite( fixed_tensor_scale, f"bridge fixed tensor scale for block {checked_block_index}", ) binding = _BridgeBlockBinding( block_index=checked_block_index, module=module, up_module=up_module, down_module=down_module, up_projection=up_projection, down_projection=down_projection, fixed_tensor_scale=scale_value, bridge_tensor_scale=self.torch.tensor( [scale_value], device=f"cuda:{self.device_index}", dtype=self.torch.float32, ), bridge_norm_constant=self.torch.tensor( [1.0 / scale_value], device=f"cuda:{self.device_index}", dtype=self.torch.float32, ), activation_mode=activation_mode, ) self._bindings_by_block[checked_block_index] = binding self._enabled_blocks.add(checked_block_index) self._telemetry(checked_block_index) return binding def installed_block_indices(self) -> list[int]: return sorted(self._bindings_by_block) @property def enabled(self) -> bool: return self._enabled and not self._closed @property def enabled_block_indices(self) -> list[int]: return sorted(self._enabled_blocks) def set_enabled( self, enabled: bool, selected_blocks: None | int | str | Iterable[int] = None, ) -> None: if enabled and self._closed: raise RuntimeError("FP4 bridge runtime is closed") if selected_blocks is None: self._enabled = bool(enabled) return blocks = normalize_selected_blocks( selected_blocks, allowed_blocks=self._bindings_by_block, ) if enabled: self._enabled_blocks.update(blocks) else: self._enabled_blocks.difference_update(blocks) def set_active_blocks( self, selected_blocks: None | int | str | Iterable[int], ) -> None: if self._closed: raise RuntimeError("FP4 bridge runtime is closed") self._enabled_blocks = normalize_selected_blocks( selected_blocks, allowed_blocks=self._bindings_by_block, ) def reset_telemetry(self) -> None: self._telemetry_by_block = { block_index: _BlockTelemetry() for block_index in self._bindings_by_block } def _logical_m_from_shape(self, hidden_states: Any) -> int | None: shape = getattr(hidden_states, "shape", None) if shape is None: return None if len(shape) < 1: return None return int(math.prod(int(value) for value in shape[:-1])) if len(shape) > 1 else 1 def _current_stream_handle(self, hidden_states: Any) -> int: current_stream = self.torch.cuda.current_stream(hidden_states.device) return int(current_stream.cuda_stream) def _route_reason( self, hidden_states: Any, binding: _BridgeBlockBinding, ) -> tuple[bool, str, int | None, int | None]: logical_m = self._logical_m_from_shape(hidden_states) if self._closed: return False, "runtime_closed", logical_m, None if not self._enabled: return False, "runtime_disabled", logical_m, None if binding.block_index not in self._enabled_blocks: return False, "block_disabled", logical_m, None if getattr(hidden_states, "device", None) is None: return False, "missing_device", logical_m, None if hidden_states.device.type != "cuda": return False, f"device_{hidden_states.device.type}", logical_m, None if getattr(hidden_states, "dtype", None) != self.torch.bfloat16: return False, f"dtype_{hidden_states.dtype}", logical_m, None if getattr(hidden_states, "ndim", 0) < 1: return False, "shape_rank", logical_m, None if int(hidden_states.shape[-1]) != UP_IN_FEATURES: return False, "shape_last_dim", logical_m, None if logical_m is None: return False, "logical_m_unknown", logical_m, None logical_reason = self._logical_m_reason(int(logical_m)) if logical_reason is not None: return False, logical_reason, int(logical_m), None try: stream = self._current_stream_handle(hidden_states) except Exception: return False, "stream_query_failed", int(logical_m), None if self._bound_stream is not None and self._bound_stream != stream: return False, "stream_mismatch", int(logical_m), stream return True, "native", int(logical_m), stream def _native_forward( self, hidden_states: Any, binding: _BridgeBlockBinding, *, logical_m: int, stream: int, ) -> Any: flattened = hidden_states.reshape(-1, UP_IN_FEATURES).contiguous() buffers = self._buffers(logical_m, flattened.device) current_stream = self.torch.cuda.current_stream(hidden_states.device) if self._bound_stream is None: self._bound_stream = stream self.bridge_up.forward( flattened, binding.up_projection.packed_weight, binding.up_projection.weight_scales, binding.up_projection.weight_scale, binding.up_projection.bias, binding.bridge_norm_constant, buffers.payload, buffers.scales, logical_m, stream, ) self.bridge_down.forward( buffers.payload, buffers.scales, binding.bridge_tensor_scale, binding.down_projection.packed_weight, binding.down_projection.weight_scales, binding.down_projection.weight_scale, binding.down_projection.bias, buffers.output, logical_m, stream, ) _record_stream_for_tensors( current_stream, flattened, binding.up_projection.packed_weight, binding.up_projection.weight_scales, binding.up_projection.weight_scale, binding.up_projection.bias, binding.bridge_norm_constant, buffers.payload, buffers.scales, binding.bridge_tensor_scale, binding.down_projection.packed_weight, binding.down_projection.weight_scales, binding.down_projection.weight_scale, binding.down_projection.bias, buffers.output, ) return buffers.output.view(*hidden_states.shape[:-1], DOWN_OUT_FEATURES) def forward_or_fallback( self, hidden_states: Any, binding: _BridgeBlockBinding, fallback_module: Any, ) -> Any: can_launch, reason, logical_m, stream = self._route_reason( hidden_states, binding, ) telemetry = self._telemetry(binding.block_index) if not can_launch: telemetry.record( native=False, reason=reason, logical_m=logical_m, stream=stream, ) return fallback_module(hidden_states) try: output = self._native_forward( hidden_states, binding, logical_m=int(logical_m), stream=int(stream), ) except Exception: telemetry.record( native=False, reason="native_error", logical_m=logical_m, stream=stream, ) # A native error can occur after work was enqueued. Do not mix an # ordinary fallback launch into that same stream. raise telemetry.record( native=True, reason="native", logical_m=logical_m, stream=stream, ) return output def telemetry_snapshot(self) -> dict[str, Any]: return { str(block_index): self._telemetry(block_index).snapshot() for block_index in sorted(self._bindings_by_block) } def report(self) -> dict[str, Any]: return { "bridge_up_library": str(self.bridge_up_library_path), "bridge_down_library": str(self.bridge_down_library_path), "block_indices": self.installed_block_indices(), "enabled": self.enabled, "enabled_block_indices": self.enabled_block_indices, "logical_m_contract": { "divisible_by_8": True, "max_supported": 6400, "fallback_on_unsupported": True, }, "toggle_contract": { "default_enabled": True, "global_gate": "set_enabled(bool)", "block_gate": "set_enabled(bool, selected_blocks=...)", "native_route_requires_global_and_block_enable": True, }, "telemetry_by_block": self.telemetry_snapshot(), "closed": self._closed, } def close(self) -> None: if self._closed: return self._closed = True self._enabled = False if self.torch.cuda.is_available(): self.torch.cuda.synchronize(self.device_index) close_error = None for library in (self.bridge_down, self.bridge_up): try: library.close() except Exception as error: # noqa: BLE001 if close_error is None: close_error = error self._bindings_by_block.clear() self._enabled_blocks.clear() self._buffers_by_m.clear() if close_error is not None: raise close_error def install_selected_img_mlp_bridges( transformer: Any, *, bridge_up_library_path: str | Path, bridge_down_library_path: str | Path, block_tensor_scales: str | Mapping[int | str, float | str], torch: Any, enabled: bool = True, device_index: int | None = None, bridge_up_factory: Any = BridgeUpLibrary, bridge_down_factory: Any = PrequantizedDownLibrary, ) -> tuple[BlockImgMlpBridgeRuntime, dict[str, Any]]: """Replace selected image MLP blocks with the chained FP4 bridge pair.""" import torch.nn as nn if not torch.cuda.is_available(): raise RuntimeError("bridge candidate installation requires CUDA") if transformer.training: raise RuntimeError("bridge candidate expects an eval-mode transformer") normalized_scales = normalize_block_tensor_scales(block_tensor_scales) runtime = BlockImgMlpBridgeRuntime( bridge_up_library_path=resolve_existing_runtime_file( bridge_up_library_path, "bridge-up library", ), bridge_down_library_path=resolve_existing_runtime_file( bridge_down_library_path, "bridge-down library", ), torch=torch, device_index=device_index, bridge_up_factory=bridge_up_factory, bridge_down_factory=bridge_down_factory, ) installed: list[dict[str, Any]] = [] try: for block_index in sorted(normalized_scales): module_name, up_module_name, down_module_name = _module_names_for_block( block_index ) block = transformer.transformer_blocks[block_index] feed_forward = block.img_mlp net = getattr(feed_forward, "net", None) if net is None or len(net) != 3: raise RuntimeError( f"{module_name} layout changed; expected 3 net entries" ) activation = net[0] up_projection, activation_mode = _resolve_up_projection( activation, up_module_name, ) dropout = net[1] if not isinstance(dropout, nn.Dropout): raise RuntimeError(f"{module_name}.net.1 changed; expected Dropout") down_projection = net[2] try: _require_packed_projection( up_projection, label=up_module_name, in_features=UP_IN_FEATURES, out_features=UP_OUT_FEATURES, torch=torch, ) _require_packed_projection( down_projection, label=down_module_name, in_features=UP_OUT_FEATURES, out_features=DOWN_OUT_FEATURES, torch=torch, ) except RuntimeError as error: raise RuntimeError( f"{module_name} cannot use the bridge candidate because it is not a packed NVFP4 image MLP block: {error}" ) from error if up_projection.packed_weight.device != down_projection.packed_weight.device: raise RuntimeError( f"{module_name} requires up/down projections on one device" ) _require_exact_numel( up_projection.packed_weight, expected=runtime.bridge_up.helper( "weight_bytes", UP_OUT_FEATURES, UP_IN_FEATURES, ), label=f"{up_module_name}.packed_weight", ) _require_exact_numel( up_projection.weight_scales, expected=runtime.bridge_up.helper( "weight_scale_bytes", UP_OUT_FEATURES, UP_IN_FEATURES, ), label=f"{up_module_name}.weight_scales", ) _require_exact_numel( up_projection.bias, expected=runtime.bridge_up.helper("bias_bytes", UP_OUT_FEATURES) // up_projection.bias.element_size(), label=f"{up_module_name}.bias", ) _require_exact_numel( down_projection.packed_weight, expected=runtime.bridge_down.helper( "weight_bytes", DOWN_OUT_FEATURES, UP_OUT_FEATURES, ), label=f"{down_module_name}.packed_weight", ) _require_exact_numel( down_projection.weight_scales, expected=runtime.bridge_down.helper( "weight_scale_bytes", DOWN_OUT_FEATURES, UP_OUT_FEATURES, ), label=f"{down_module_name}.weight_scales", ) _require_exact_numel( down_projection.bias, expected=runtime.bridge_down.helper("bias_bytes", DOWN_OUT_FEATURES) // down_projection.bias.element_size(), label=f"{down_module_name}.bias", ) binding = runtime.bind_block( block_index=block_index, module=module_name, up_module=up_module_name, down_module=down_module_name, up_projection=up_projection, down_projection=down_projection, fixed_tensor_scale=normalized_scales[block_index], activation_mode=activation_mode, ) class BlockImgMlpBridge(nn.Module): def __init__( self, installed_binding: _BridgeBlockBinding, original_feed_forward: Any, ) -> None: super().__init__() self.binding = installed_binding self.fallback_module = original_feed_forward self._xpo3_fp4_bridge = True self._xpo3_fp4_bridge_block_index = int( installed_binding.block_index ) def forward(self, hidden_states: Any) -> Any: return runtime.forward_or_fallback( hidden_states, self.binding, self.fallback_module, ) block.img_mlp = BlockImgMlpBridge(binding, feed_forward) installed.append( { "block_index": block_index, "module": module_name, "up_module": up_module_name, "down_module": down_module_name, "fixed_tensor_scale": float(binding.fixed_tensor_scale), "bridge_norm_constant": float(binding.bridge_norm_constant.item()), "activation_mode": activation_mode, } ) runtime.set_enabled(bool(enabled)) except Exception: runtime.close() raise metadata = runtime.report() metadata.update( { "mode": "selected_img_mlp", "block_count": len(installed), "blocks": installed, "fixed_tensor_scales_by_block": { str(entry["block_index"]): entry["fixed_tensor_scale"] for entry in installed }, "bridge_up_context_reserved_bytes": runtime.bridge_up.reserved_bytes, "bridge_down_context_reserved_bytes": runtime.bridge_down.reserved_bytes, "resident_scalar_bytes": ( sum( _tensor_bytes(binding.bridge_tensor_scale) for binding in runtime._bindings_by_block.values() ) + sum( _tensor_bytes(binding.bridge_norm_constant) for binding in runtime._bindings_by_block.values() ) ), "reuses_existing_projection_buffers": True, "validation": { "no_bf16_post_gelu_intermediate_materialized": True, "bf16_input_required": True, "bf16_output_preserved": True, "dropout_bypassed_in_eval_only": True, "single_cuda_stream_required_for_native_route": True, "unsupported_conditions_fallback_before_native_launch": True, }, } ) return runtime, metadata def install_block0_img_mlp_bridge( transformer: Any, *, bridge_up_library_path: str | Path, bridge_down_library_path: str | Path, fixed_tensor_scale: float, torch: Any, enabled: bool = True, device_index: int | None = None, bridge_up_factory: Any = BridgeUpLibrary, bridge_down_factory: Any = PrequantizedDownLibrary, ) -> tuple[BlockImgMlpBridgeRuntime, dict[str, Any]]: runtime, metadata = install_selected_img_mlp_bridges( transformer, bridge_up_library_path=bridge_up_library_path, bridge_down_library_path=bridge_down_library_path, block_tensor_scales={0: fixed_tensor_scale}, torch=torch, enabled=enabled, device_index=device_index, bridge_up_factory=bridge_up_factory, bridge_down_factory=bridge_down_factory, ) metadata["mode"] = "block0_img_mlp" return runtime, metadata __all__ = [ "BLOCK0_IMG_MLP_MODULE", "DOWN_MODULE", "TRANSFORMER_BLOCK_COUNT", "UP_MODULE", "BlockImgMlpBridgeRuntime", "install_block0_img_mlp_bridge", "install_selected_img_mlp_bridges", "normalize_block_tensor_scales", "normalize_selected_blocks", "resolve_existing_runtime_file", ]