Add runtime/fp4_bridge_runtime.py

runtime/fp4_bridge_runtime.py

@@ -0,0 +1,1190 @@
+"""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",
+]