quantized_linear.py

"""Portable Mage-VL FP8 and NVFP4 linear modules.

This file is loaded as Hugging Face remote code.  The checkpoint config
selects one format before the state dictionary is materialized, so the
original BF16 language projection weights are never allocated or requested.
"""

from __future__ import annotations

import hashlib
import os
import types
from functools import lru_cache
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from torch import nn


LANGUAGE_PROJECTION_ROLES = {
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_proj",
    "up_proj",
    "down_proj",
}
_SMALLM_SOURCE_ROOT: Path | None = None


def _quantize_nvfp4_activation(value: torch.Tensor) -> Any:
    """Apply the checkpoint's exact dynamic NVFP4 activation policy."""

    from comfy_kitchen.tensor import QuantizedTensor

    return QuantizedTensor.from_float(value, "TensorCoreNVFP4Layout")


def _configure_smallm_source(model_name_or_path: str) -> None:
    """Resolve native sources from a local repo or Hugging Face snapshot."""

    global _SMALLM_SOURCE_ROOT
    candidate = Path(model_name_or_path).expanduser()
    local_root = candidate / "native" / "smallm_gemv"
    required = ("smallm_gemv.cpp", "smallm_gemv.cu", "smallm_gemv.h")
    if all((local_root / name).is_file() for name in required):
        _SMALLM_SOURCE_ROOT = local_root.resolve()
        return
    if not model_name_or_path:
        raise RuntimeError("Mage-VL small-M source repository is unspecified")

    from transformers.utils.hub import cached_file

    resolved = [
        Path(
            cached_file(
                model_name_or_path,
                f"native/smallm_gemv/{name}",
            )
        )
        for name in required
    ]
    parents = {path.parent.resolve() for path in resolved}
    if len(parents) != 1:
        raise RuntimeError(
            "small-M native sources resolved to different directories: "
            f"{sorted(str(value) for value in parents)}"
        )
    _SMALLM_SOURCE_ROOT = parents.pop()


def _smallm_source_root() -> Path:
    if _SMALLM_SOURCE_ROOT is None:
        raise RuntimeError(
            "small-M native sources were not configured during model setup"
        )
    return _SMALLM_SOURCE_ROOT


@lru_cache(maxsize=1)
def _load_smallm_extension() -> Any:
    from torch.utils.cpp_extension import load

    source_root = _smallm_source_root()
    configured_build = os.environ.get("MAGE_VL_SMALLM_BUILD_DIR")
    build_root = (
        Path(configured_build).expanduser().resolve()
        if configured_build
        else Path(__file__).resolve().parent / ".native_build"
    )
    build_root.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
    os.environ.setdefault("MAX_JOBS", "4")
    source_hash = hashlib.sha256(
        b"".join(
            (source_root / name).read_bytes()
            for name in (
                "smallm_gemv.cpp",
                "smallm_gemv.cu",
                "smallm_gemv.h",
            )
        )
    ).hexdigest()[:12]
    return load(
        name=f"mage_vl_smallm_gemv_{source_hash}",
        sources=[
            str(source_root / "smallm_gemv.cpp"),
            str(source_root / "smallm_gemv.cu"),
        ],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
        extra_include_paths=[str(source_root)],
        build_directory=str(build_root),
        with_cuda=True,
        verbose=False,
        is_python_module=True,
    )


def _smallm_nvfp4_linear(
    value: torch.Tensor,
    *,
    qdata: torch.Tensor,
    weight_block_scale: torch.Tensor,
    weight_scale: torch.Tensor,
    bias: torch.Tensor | None,
) -> torch.Tensor:
    return _load_smallm_extension().linear(
        value.contiguous(),
        qdata,
        weight_block_scale,
        weight_scale,
        bias,
    )


def _resolve_parent(root: nn.Module, module_name: str) -> tuple[nn.Module, str]:
    parent_name, separator, leaf = module_name.rpartition(".")
    if not separator:
        return root, module_name
    return root.get_submodule(parent_name), leaf


def _empty_like_source(
    source: nn.Linear,
    shape: tuple[int, ...],
    dtype: torch.dtype,
) -> torch.Tensor:
    return torch.empty(shape, dtype=dtype, device=source.weight.device)


class MageVLScaledFP8Linear(nn.Module):
    """W8A8 prefill with optional resident-weight W8A16 small-M decode."""

    def __init__(
        self,
        source: nn.Linear,
        *,
        role: str,
        smallm_backend: str,
        smallm_threshold: int,
        smallm_roles: set[str],
    ) -> None:
        super().__init__()
        if smallm_backend not in {"off", "w8a16_gemv"}:
            raise ValueError(f"unsupported FP8 small-M backend: {smallm_backend}")
        if smallm_threshold <= 0:
            raise ValueError("FP8 small-M threshold must be positive")
        self.in_features = int(source.in_features)
        self.out_features = int(source.out_features)
        self.role = role
        self.smallm_backend = (
            smallm_backend if role in smallm_roles else "off"
        )
        self.smallm_threshold = int(smallm_threshold)
        self.register_buffer(
            "qdata",
            _empty_like_source(
                source,
                (self.out_features, self.in_features),
                torch.float8_e4m3fn,
            ),
            persistent=True,
        )
        self.register_buffer(
            "weight_scale",
            _empty_like_source(source, (), torch.float32),
            persistent=True,
        )
        if source.bias is None:
            self.bias_bf16 = None
        else:
            self.register_buffer(
                "bias_bf16",
                _empty_like_source(
                    source,
                    (self.out_features,),
                    torch.bfloat16,
                ),
                persistent=True,
            )

    def _weight_quantized_tensor(self) -> Any:
        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreFP8Layout

        params = TensorCoreFP8Layout.Params(
            scale=self.weight_scale,
            orig_dtype=torch.bfloat16,
            orig_shape=(self.out_features, self.in_features),
        )
        return QuantizedTensor(
            self.qdata,
            "TensorCoreFP8Layout",
            params,
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        from comfy_kitchen.tensor import QuantizedTensor

        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if (
            self.smallm_backend == "w8a16_gemv"
            and flattened.shape[0] <= self.smallm_threshold
        ):
            from .fp8_decode_runtime import smallm_fp8_linear

            output = smallm_fp8_linear(
                flattened,
                qdata=self.qdata,
                weight_scale=self.weight_scale,
                bias=self.bias_bf16,
            )
            return output.reshape(*input_shape[:-1], self.out_features)
        quantized_input = QuantizedTensor.from_float(
            flattened,
            "TensorCoreFP8Layout",
        )
        output = F.linear(
            quantized_input,
            self._weight_quantized_tensor(),
            None,
        )
        if self.bias_bf16 is not None:
            output = output + self.bias_bf16
        return output.reshape(*input_shape[:-1], self.out_features)


class MageVLNVFP4Linear(nn.Module):
    """Native W4A4 prefill with optional packed-weight W4A16 small-M decode."""

    def __init__(
        self,
        source: nn.Linear,
        *,
        role: str,
        smallm_backend: str,
        smallm_threshold: int,
        smallm_roles: set[str],
    ) -> None:
        super().__init__()
        self.in_features = int(source.in_features)
        self.out_features = int(source.out_features)
        self.role = role
        self.smallm_backend = (
            smallm_backend if role in smallm_roles else "off"
        )
        self.smallm_threshold = int(smallm_threshold)
        self.register_buffer(
            "qdata",
            _empty_like_source(
                source,
                (self.out_features, self.in_features // 2),
                torch.uint8,
            ),
            persistent=True,
        )
        self.register_buffer(
            "weight_scale",
            _empty_like_source(source, (), torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "weight_block_scale",
            _empty_like_source(
                source,
                (self.out_features, self.in_features // 16),
                torch.float8_e4m3fn,
            ),
            persistent=True,
        )
        if source.bias is None:
            self.bias_bf16 = None
        else:
            self.register_buffer(
                "bias_bf16",
                _empty_like_source(
                    source,
                    (self.out_features,),
                    torch.bfloat16,
                ),
                persistent=True,
            )

    def _weight_quantized_tensor(self) -> Any:
        from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout

        params = TensorCoreNVFP4Layout.Params(
            scale=self.weight_scale,
            orig_dtype=torch.bfloat16,
            orig_shape=(self.out_features, self.in_features),
            block_scale=self.weight_block_scale,
        )
        return QuantizedTensor(
            self.qdata,
            "TensorCoreNVFP4Layout",
            params,
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if (
            self.smallm_backend == "w4a16_gemv"
            and flattened.shape[0] <= self.smallm_threshold
        ):
            output = _smallm_nvfp4_linear(
                flattened,
                qdata=self.qdata,
                weight_block_scale=self.weight_block_scale,
                weight_scale=self.weight_scale,
                bias=self.bias_bf16,
            )
            return output.reshape(*input_shape[:-1], self.out_features)

        quantized_input = _quantize_nvfp4_activation(flattened)
        output = F.linear(
            quantized_input,
            self._weight_quantized_tensor(),
            None,
        )
        if self.bias_bf16 is not None:
            output = output + self.bias_bf16
        return output.reshape(*input_shape[:-1], self.out_features)


def _smallm_policy(
    quantization: dict[str, Any],
    *,
    format_name: str,
) -> tuple[str, int, set[str]]:
    backend = os.environ.get(
        "MAGE_VL_SMALLM_BACKEND",
        str(quantization.get("smallm_backend", "off")),
    )
    supported_backend = (
        "w8a16_gemv"
        if format_name == "scaled_fp8_w8a8"
        else "w4a16_gemv"
    )
    if backend not in {"off", supported_backend}:
        raise ValueError(f"unsupported MAGE_VL_SMALLM_BACKEND: {backend}")
    threshold = int(
        os.environ.get(
            "MAGE_VL_SMALLM_THRESHOLD",
            str(quantization.get("smallm_threshold", 1)),
        )
    )
    if threshold <= 0:
        raise ValueError("MAGE_VL_SMALLM_THRESHOLD must be positive")
    configured_roles = quantization.get(
        "smallm_roles",
        sorted(LANGUAGE_PROJECTION_ROLES),
    )
    role_text = os.environ.get(
        "MAGE_VL_SMALLM_ROLES",
        ",".join(str(value) for value in configured_roles),
    )
    roles = {value.strip() for value in role_text.split(",") if value.strip()}
    if not roles <= LANGUAGE_PROJECTION_ROLES:
        raise ValueError(
            f"invalid small-M roles: {sorted(roles - LANGUAGE_PROJECTION_ROLES)}"
        )
    return backend, threshold, roles


def _environment_flag(name: str, default: bool) -> bool:
    value = os.environ.get(name)
    if value is None:
        return bool(default)
    normalized = value.strip().lower()
    if normalized in {"1", "true", "yes", "on"}:
        return True
    if normalized in {"0", "false", "no", "off"}:
        return False
    raise ValueError(f"{name} must be one of 1/0, true/false, yes/no, or on/off")


def _runtime_policy(
    quantization: dict[str, Any],
    *,
    format_name: str,
) -> tuple[str, dict[str, Any]]:
    """Resolve a named runtime profile without changing stored weights."""

    profiles = quantization.get("runtime_profiles")
    if format_name != "native_nvfp4_w4a4" or profiles is None:
        return "legacy", dict(quantization)
    if not isinstance(profiles, dict) or not profiles:
        raise TypeError("runtime_profiles must be a non-empty dictionary")
    default_profile = str(quantization.get("default_profile", "hybrid_fast"))
    profile_name = os.environ.get("MAGE_VL_RUNTIME_PROFILE", default_profile)
    if profile_name not in profiles:
        raise ValueError(
            f"unsupported MAGE_VL_RUNTIME_PROFILE: {profile_name}; "
            f"available={sorted(profiles)}"
        )
    profile = profiles[profile_name]
    if not isinstance(profile, dict):
        raise TypeError(f"runtime profile {profile_name} must be a dictionary")
    effective = dict(quantization)
    effective.update(profile)
    return profile_name, effective


def _pure_w4a4_projection(
    module: MageVLNVFP4Linear,
    quantized_input: Any,
) -> torch.Tensor:
    output = F.linear(
        quantized_input,
        module._weight_quantized_tensor(),
        None,
    )
    if module.bias_bf16 is not None:
        output = output + module.bias_bf16
    return output


def install_nvfp4_shared_gate_up_activation(
    model: nn.Module,
    *,
    threshold: int = 1,
) -> dict[str, Any]:
    """Reuse one exact dynamic NVFP4 activation for gate and up."""

    if threshold <= 0:
        raise ValueError("shared gate/up activation threshold must be positive")
    installed = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.mlp"
        mlp = model.get_submodule(name)
        if hasattr(mlp, "_nvfp4_shared_gate_up_original_forward"):
            raise RuntimeError(f"{name}: shared gate/up activation is installed")
        for role in ("gate_proj", "up_proj"):
            module = getattr(mlp, role, None)
            if not isinstance(module, MageVLNVFP4Linear):
                raise TypeError(
                    f"{name}.{role}: expected MageVLNVFP4Linear, "
                    f"got {type(module).__name__}"
                )
            if module.smallm_backend != "off":
                raise RuntimeError(
                    f"{name}.{role}: shared activation requires pure W4A4"
                )
        original_forward = mlp.forward
        object.__setattr__(
            mlp,
            "_nvfp4_shared_gate_up_original_forward",
            original_forward,
        )
        object.__setattr__(
            mlp,
            "_nvfp4_shared_gate_up_threshold",
            int(threshold),
        )

        def shared_forward(self, value: torch.Tensor):
            input_shape = tuple(value.shape)
            flattened = value.reshape(-1, input_shape[-1]).contiguous()
            if flattened.shape[0] > self._nvfp4_shared_gate_up_threshold:
                return self._nvfp4_shared_gate_up_original_forward(value)
            quantized_input = _quantize_nvfp4_activation(flattened)
            gate = _pure_w4a4_projection(self.gate_proj, quantized_input)
            up = _pure_w4a4_projection(self.up_proj, quantized_input)
            intermediate = (F.silu(gate) * up).reshape(
                *input_shape[:-1],
                self.gate_proj.out_features,
            )
            return self.down_proj(intermediate)

        object.__setattr__(mlp, "forward", types.MethodType(shared_forward, mlp))
        installed.append(name)
    return {
        "feature": "shared_exact_dynamic_nvfp4_gate_up_activation",
        "threshold": int(threshold),
        "installed_module_count": len(installed),
        "fallback": "original_qwen_mlp_forward",
        "stored_weight_payload": "unchanged",
    }


def restore_nvfp4_shared_gate_up_activation(model: nn.Module) -> dict[str, Any]:
    """Restore MLP callables patched for exact NVFP4 activation reuse."""

    restored = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.mlp"
        mlp = model.get_submodule(name)
        original = getattr(
            mlp,
            "_nvfp4_shared_gate_up_original_forward",
            None,
        )
        if original is None:
            continue
        object.__setattr__(mlp, "forward", original)
        object.__delattr__(mlp, "_nvfp4_shared_gate_up_original_forward")
        object.__delattr__(mlp, "_nvfp4_shared_gate_up_threshold")
        restored.append(name)
    return {
        "feature": "shared_exact_dynamic_nvfp4_gate_up_activation",
        "restored_module_count": len(restored),
    }


class _NVFP4SharedQKVActivationCoordinator:
    """Reuse one exact dynamic NVFP4 activation across Q/K/V."""

    def __init__(
        self,
        q_proj: MageVLNVFP4Linear,
        k_proj: MageVLNVFP4Linear,
        v_proj: MageVLNVFP4Linear,
        *,
        threshold: int,
    ) -> None:
        self.projections = {"q": q_proj, "k": k_proj, "v": v_proj}
        self.original_forwards = {
            role: module.forward for role, module in self.projections.items()
        }
        self.threshold = int(threshold)
        self.pending: dict[str, Any] | None = None

    @staticmethod
    def _signature(value: torch.Tensor) -> tuple[Any, ...]:
        return (
            value.data_ptr(),
            tuple(value.shape),
            tuple(value.stride()),
            value.storage_offset(),
            value.device,
            value.dtype,
        )

    def forward(self, role: str, value: torch.Tensor) -> torch.Tensor:
        input_shape = tuple(value.shape)
        flattened = value.reshape(-1, input_shape[-1]).contiguous()
        if flattened.shape[0] > self.threshold:
            self.pending = None
            return self.original_forwards[role](value)
        signature = self._signature(value)
        if role == "q":
            quantized_input = _quantize_nvfp4_activation(flattened)
            outputs = {
                projection_role: _pure_w4a4_projection(
                    projection,
                    quantized_input,
                ).reshape(*input_shape[:-1], projection.out_features)
                for projection_role, projection in self.projections.items()
            }
            self.pending = {
                "signature": signature,
                "k": outputs["k"],
                "v": outputs["v"],
            }
            return outputs["q"]
        if self.pending is None or self.pending["signature"] != signature:
            return self.original_forwards[role](value)
        output = self.pending[role]
        if role == "v":
            self.pending = None
        return output


def install_nvfp4_shared_qkv_activation(
    model: nn.Module,
    *,
    threshold: int = 1,
) -> dict[str, Any]:
    """Patch pure-W4A4 Q/K/V to share one exact dynamic activation."""

    if threshold <= 0:
        raise ValueError("shared QKV activation threshold must be positive")
    installed = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.self_attn"
        attention = model.get_submodule(name)
        if hasattr(attention, "_nvfp4_shared_qkv_coordinator"):
            raise RuntimeError(f"{name}: shared QKV activation is installed")
        projections = {}
        for role in ("q", "k", "v"):
            module = getattr(attention, f"{role}_proj", None)
            if not isinstance(module, MageVLNVFP4Linear):
                raise TypeError(
                    f"{name}.{role}_proj: expected MageVLNVFP4Linear, "
                    f"got {type(module).__name__}"
                )
            if module.smallm_backend != "off":
                raise RuntimeError(
                    f"{name}.{role}_proj: shared activation requires pure W4A4"
                )
            projections[role] = module
        coordinator = _NVFP4SharedQKVActivationCoordinator(
            projections["q"],
            projections["k"],
            projections["v"],
            threshold=threshold,
        )
        object.__setattr__(attention, "_nvfp4_shared_qkv_coordinator", coordinator)
        for role, module in projections.items():
            def shared_forward(
                self,
                value: torch.Tensor,
                *,
                _role: str = role,
                _coordinator: _NVFP4SharedQKVActivationCoordinator = coordinator,
            ):
                return _coordinator.forward(_role, value)

            object.__setattr__(
                module,
                "forward",
                types.MethodType(shared_forward, module),
            )
        installed.append(name)
    return {
        "feature": "shared_exact_dynamic_nvfp4_qkv_activation",
        "threshold": int(threshold),
        "installed_module_count": len(installed),
        "fallback": "separate_q_k_v_projection_forwards",
        "stored_weight_payload": "unchanged",
    }


def restore_nvfp4_shared_qkv_activation(model: nn.Module) -> dict[str, Any]:
    """Restore Q/K/V callables patched for exact NVFP4 activation reuse."""

    restored = []
    for layer in range(36):
        name = f"language_model.layers.{layer}.self_attn"
        attention = model.get_submodule(name)
        coordinator = getattr(attention, "_nvfp4_shared_qkv_coordinator", None)
        if coordinator is None:
            continue
        for role in ("q", "k", "v"):
            module = getattr(attention, f"{role}_proj")
            object.__setattr__(module, "forward", coordinator.original_forwards[role])
        coordinator.pending = None
        object.__delattr__(attention, "_nvfp4_shared_qkv_coordinator")
        restored.append(name)
    return {
        "feature": "shared_exact_dynamic_nvfp4_qkv_activation",
        "restored_module_count": len(restored),
    }


def apply_mage_vl_quantization(
    model: nn.Module,
    config: Any,
) -> None:
    """Replace all 252 Qwen language projections before checkpoint loading."""

    quantization = getattr(config, "mage_vl_quantization", None)
    if not quantization:
        return
    if not isinstance(quantization, dict):
        raise TypeError("mage_vl_quantization must be a dictionary")
    format_name = quantization.get("format")
    if format_name not in {"scaled_fp8_w8a8", "native_nvfp4_w4a4"}:
        raise ValueError(f"unsupported Mage-VL quantization: {format_name}")

    runtime_profile, effective_quantization = _runtime_policy(
        quantization,
        format_name=format_name,
    )
    backend, threshold, smallm_roles = _smallm_policy(
        effective_quantization,
        format_name=format_name,
    )
    fused_gate_up = False
    fused_qkv = False
    fused_gate_up_threshold = 1
    fused_qkv_threshold = 1
    shared_gate_up = False
    shared_qkv = False
    shared_gate_up_threshold = 1
    shared_qkv_threshold = 1
    if format_name == "scaled_fp8_w8a8":
        fused_gate_up = _environment_flag(
            "MAGE_VL_FP8_FUSED_GATE_UP",
            bool(quantization.get("fused_gate_up", False)),
        )
        fused_qkv = _environment_flag(
            "MAGE_VL_FP8_FUSED_QKV",
            bool(quantization.get("fused_qkv", False)),
        )
        fused_gate_up_threshold = int(
            os.environ.get(
                "MAGE_VL_FP8_FUSED_GATE_UP_THRESHOLD",
                str(quantization.get("fused_gate_up_threshold", 1)),
            )
        )
        fused_qkv_threshold = int(
            os.environ.get(
                "MAGE_VL_FP8_FUSED_QKV_THRESHOLD",
                str(quantization.get("fused_qkv_threshold", 1)),
            )
        )
        if fused_gate_up_threshold <= 0 or fused_qkv_threshold <= 0:
            raise ValueError("FP8 fusion thresholds must be positive")
        if backend != "off" or fused_gate_up or fused_qkv:
            from .fp8_decode_runtime import configure_fp8_decode_sources

            configure_fp8_decode_sources(
                str(getattr(config, "_name_or_path", ""))
            )
    if format_name == "native_nvfp4_w4a4" and backend != "off":
        _configure_smallm_source(str(getattr(config, "_name_or_path", "")))
    if format_name == "native_nvfp4_w4a4":
        shared_gate_up = _environment_flag(
            "MAGE_VL_W4A4_SHARED_GATE_UP_ACTIVATION",
            bool(effective_quantization.get("shared_gate_up_activation", False)),
        )
        shared_qkv = _environment_flag(
            "MAGE_VL_W4A4_SHARED_QKV_ACTIVATION",
            bool(effective_quantization.get("shared_qkv_activation", False)),
        )
        shared_gate_up_threshold = int(
            os.environ.get(
                "MAGE_VL_W4A4_SHARED_GATE_UP_THRESHOLD",
                str(effective_quantization.get("shared_gate_up_threshold", 1)),
            )
        )
        shared_qkv_threshold = int(
            os.environ.get(
                "MAGE_VL_W4A4_SHARED_QKV_THRESHOLD",
                str(effective_quantization.get("shared_qkv_threshold", 1)),
            )
        )
        if shared_gate_up_threshold <= 0 or shared_qkv_threshold <= 0:
            raise ValueError("NVFP4 shared-activation thresholds must be positive")
        if backend != "off" and (shared_gate_up or shared_qkv):
            raise ValueError(
                "NVFP4 shared-activation features require MAGE_VL_SMALLM_BACKEND=off"
            )
    installed = []
    for layer in range(36):
        for branch, roles in (
            ("self_attn", ("q_proj", "k_proj", "v_proj", "o_proj")),
            ("mlp", ("gate_proj", "up_proj", "down_proj")),
        ):
            for role in roles:
                name = f"language_model.layers.{layer}.{branch}.{role}"
                parent, leaf = _resolve_parent(model, name)
                source = getattr(parent, leaf)
                if not isinstance(source, nn.Linear):
                    raise TypeError(
                        f"{name}: expected nn.Linear, got "
                        f"{type(source).__name__}"
                    )
                if format_name == "scaled_fp8_w8a8":
                    replacement = MageVLScaledFP8Linear(
                        source,
                        role=role,
                        smallm_backend=backend,
                        smallm_threshold=threshold,
                        smallm_roles=smallm_roles,
                    )
                else:
                    replacement = MageVLNVFP4Linear(
                        source,
                        role=role,
                        smallm_backend=backend,
                        smallm_threshold=threshold,
                        smallm_roles=smallm_roles,
                    )
                setattr(parent, leaf, replacement)
                installed.append(name)
    if len(installed) != 252:
        raise RuntimeError(
            f"expected 252 quantized language projections, got {len(installed)}"
        )
    runtime_manifest = {
        "runtime_version": int(quantization.get("runtime_version", 1)),
        "runtime_profile": runtime_profile,
        "format": format_name,
        "smallm_backend": backend,
        "smallm_threshold": threshold,
        "smallm_roles": sorted(smallm_roles),
        "fused_gate_up": fused_gate_up,
        "fused_gate_up_threshold": fused_gate_up_threshold,
        "fused_qkv": fused_qkv,
        "fused_qkv_threshold": fused_qkv_threshold,
        "shared_gate_up_activation": shared_gate_up,
        "shared_gate_up_threshold": shared_gate_up_threshold,
        "shared_qkv_activation": shared_qkv,
        "shared_qkv_threshold": shared_qkv_threshold,
    }
    if format_name == "scaled_fp8_w8a8":
        from .fp8_decode_runtime import (
            install_fp8_fused_gate_up,
            install_fp8_fused_qkv,
        )

        if fused_gate_up:
            runtime_manifest["gate_up_install"] = install_fp8_fused_gate_up(
                model,
                threshold=fused_gate_up_threshold,
            )
        if fused_qkv:
            runtime_manifest["qkv_install"] = install_fp8_fused_qkv(
                model,
                threshold=fused_qkv_threshold,
            )
    if format_name == "native_nvfp4_w4a4":
        if shared_gate_up:
            runtime_manifest["shared_gate_up_install"] = (
                install_nvfp4_shared_gate_up_activation(
                    model,
                    threshold=shared_gate_up_threshold,
                )
            )
        if shared_qkv:
            runtime_manifest["shared_qkv_install"] = (
                install_nvfp4_shared_qkv_activation(
                    model,
                    threshold=shared_qkv_threshold,
                )
            )
    object.__setattr__(model, "_mage_vl_runtime_manifest", runtime_manifest)


__all__ = [
    "MageVLNVFP4Linear",
    "MageVLScaledFP8Linear",
    "apply_mage_vl_quantization",
    "install_nvfp4_shared_gate_up_activation",
    "install_nvfp4_shared_qkv_activation",
    "restore_nvfp4_shared_gate_up_activation",
    "restore_nvfp4_shared_qkv_activation",
]