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