Add after-load fusion for static quantized MLPs

docs/source/en/fusion_mapping.md

@@ -58,9 +58,11 @@ This lets a fusion use a different runtime module structure while still loading
 
 Note: With the current monkey-patching mechanism, fusion registration is class-level: one compatible module class maps to one fused replacement class.
 
+Some fusions require loaded quantized weights and scales. Those fusions still use `fusion_config` for validation, but are applied after checkpoint loading by the relevant quantizer instead of using pre-initialization monkey patching.
+
 ## Current fusion families
 
-Currently, `fusion_config` supports one fusion family:
+Currently, `fusion_config` supports these fusion families:
 
 - `patch_embeddings`
   Enable with:
@@ -71,6 +73,33 @@ Currently, `fusion_config` supports one fusion family:
 
   Effect:
   Replaces compatible `nn.Conv3d` patch embedding projections with equivalent flattened `nn.Linear` projections at runtime.
+- `static_quantized_mlp`
+  Configure the Hub kernels per MLP pattern:
+
+  ```python
+  fusion_config = {
+      "static_quantized_mlp": {
+          "input_quant": {
+              "repo_id": "owner/static-fp8-input-quant",
+              "version": 1,
+          },
+          "gated_mlp": {
+              "repo_id": "owner/static-fp8-gated-mlp",
+              "version": 1,
+          },
+          "dense_gelu_mlp": {
+              "repo_id": "owner/static-fp8-dense-gelu-mlp",
+              "version": 1,
+          },
+      }
+  }
+  ```
+
+  Effect:
+  After static FP8 weights and activation scales are loaded, replaces eligible static per-tensor FineGrainedFP8
+  `gate_proj/up_proj/down_proj` MLPs and `fc1/GELU/fc2` MLPs with fused Hub kernels. Unsupported quantization
+  layouts or activations keep the existing module path. The configured repositories must expose the functions required
+  by the selected pattern. Set `trust_remote_code=True` only for kernel repositories you trust.
 
 ## Extending fusion mapping
 

src/transformers/fusion_mapping.py

@@ -67,6 +67,36 @@ def make_transforms(self, config: "PretrainedConfig") -> list[WeightTransform]:
         """Build the weight transforms needed to load and save the fused runtime layout."""
         raise NotImplementedError
 
+    def validate_options(self, options: bool | Mapping[str, Any]) -> None:
+        """Validate user options for this fusion family."""
+        return
+
+
+class AfterLoadFusionSpec:
+    """Base recipe for fusions that require loaded weights and buffers."""
+
+    def validate_options(self, options: bool | Mapping[str, Any]) -> None:
+        """Validate user options for this fusion family."""
+        raise NotImplementedError
+
+    def apply(self, model: "PreTrainedModel", options: bool | Mapping[str, Any]) -> "PreTrainedModel":
+        """Apply the fusion after the model weights have been loaded."""
+        raise NotImplementedError
+
+
+class StaticQuantizedMLPAfterLoadFusionSpec(AfterLoadFusionSpec):
+    """Fuse eligible static quantized MLP modules after quantized weights and scales are loaded."""
+
+    def validate_options(self, options: bool | Mapping[str, Any]) -> None:
+        from .integrations.static_quantized_mlp import get_static_quantized_mlp_fusion_spec
+
+        get_static_quantized_mlp_fusion_spec(options)
+
+    def apply(self, model: "PreTrainedModel", options: bool | Mapping[str, Any]) -> "PreTrainedModel":
+        from .integrations.static_quantized_mlp import replace_with_static_quantized_mlp
+
+        return replace_with_static_quantized_mlp(model, options)
+
 
 class _FusedPatchEmbeddingMixin:
     def __init__(self, *args, **kwargs):
@@ -232,7 +262,10 @@ def _register_module_fusion(
     register_checkpoint_conversion_mapping(model_type, converters, overwrite=True)
 
 
-_FUSION_REGISTRY: dict[str, ModuleFusionSpec] = {"patch_embeddings": PatchEmbeddingsFusionSpec()}
+_FUSION_REGISTRY: dict[str, ModuleFusionSpec | AfterLoadFusionSpec] = {
+    "patch_embeddings": PatchEmbeddingsFusionSpec(),
+    "static_quantized_mlp": StaticQuantizedMLPAfterLoadFusionSpec(),
+}
 
 
 def _iter_enabled_fusions(fusion_config: Mapping[str, bool | Mapping[str, Any]]) -> list[str]:
@@ -248,6 +281,7 @@ def _iter_enabled_fusions(fusion_config: Mapping[str, bool | Mapping[str, Any]])
             raise ValueError(
                 f"Invalid fusion config for {fusion_name}: expected `True`, `False`, or a mapping of options."
             )
+        _FUSION_REGISTRY[fusion_name].validate_options(fusion_options)
         enabled_fusions.append(fusion_name)
     return enabled_fusions
 
@@ -267,4 +301,22 @@ def register_fusion_patches(
         return
 
     for fusion_name in _iter_enabled_fusions(fusion_config):
-        _register_module_fusion(cls, config, fusion_name, _FUSION_REGISTRY[fusion_name])
+        spec = _FUSION_REGISTRY[fusion_name]
+        if isinstance(spec, ModuleFusionSpec):
+            _register_module_fusion(cls, config, fusion_name, spec)
+
+
+def apply_after_load_fusions(
+    model: "PreTrainedModel", fusion_config: Mapping[str, bool | Mapping[str, Any]] | None = None
+) -> "PreTrainedModel":
+    """Apply fusions that require loaded weights and buffers."""
+
+    if not fusion_config:
+        return model
+
+    for fusion_name in _iter_enabled_fusions(fusion_config):
+        spec = _FUSION_REGISTRY[fusion_name]
+        if isinstance(spec, AfterLoadFusionSpec):
+            model = spec.apply(model, fusion_config[fusion_name])
+
+    return model

src/transformers/integrations/__init__.py

@@ -153,6 +153,12 @@
     "quanto": ["replace_with_quanto_layers"],
     "sinq": ["SinqDeserialize", "SinqQuantize"],
     "spqr": ["replace_with_spqr_linear"],
+    "static_quantized_mlp": [
+        "StaticFP8DenseGELUMLP",
+        "StaticFP8GatedMLP",
+        "get_static_quantized_mlp_fusion_spec",
+        "replace_with_static_quantized_mlp",
+    ],
     "vptq": ["replace_with_vptq_linear"],
 }
 
@@ -309,6 +315,12 @@
     from .quanto import replace_with_quanto_layers
     from .sinq import SinqDeserialize, SinqQuantize
     from .spqr import replace_with_spqr_linear
+    from .static_quantized_mlp import (
+        StaticFP8DenseGELUMLP,
+        StaticFP8GatedMLP,
+        get_static_quantized_mlp_fusion_spec,
+        replace_with_static_quantized_mlp,
+    )
     from .vptq import replace_with_vptq_linear
 
     try: