runtime/fp4_bridge_runtime.py

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