Extended n-to-1 kernel fusion via KernelConfig

fused_qwen_example.py

@@ -0,0 +1,49 @@
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig
+
+
+model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+inputs = tokenizer("Hello, how are you?", return_tensors="pt")
+
+# --- baseline: plain model, no fusion ---
+print("=" * 60)
+print("Loading baseline model (no fusion)...")
+baseline = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda")
+baseline.eval()
+inputs = {k: v.to(baseline.device) for k, v in inputs.items()}
+
+with torch.no_grad():
+    baseline_out = baseline(**inputs).logits
+print("Baseline output shape:", baseline_out.shape)
+# del baseline
+
+# --- fused model ---
+print("=" * 60)
+print("Loading fused model...")
+
+kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP"
+#kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP"
+kernel_config = KernelConfig(
+    {
+        (
+            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
+            ("MLP",     "model.layers.*.mlp"),
+        ): kernel_repo_id,
+    },
+)
+
+fused_model = AutoModelForCausalLM.from_pretrained(
+    model_id, use_kernels=True, kernel_config=kernel_config, device_map="cuda"
+)
+fused_model.eval()
+print(fused_model)
+
+with torch.no_grad():
+    fused_out = fused_model(**inputs).logits
+print("Fused output shape:", fused_out.shape)
+
+# --- compare ---
+print("=" * 60)
+print("Max diff fused vs baseline:", (fused_out - baseline_out).abs().max().item())

src/transformers/core_model_loading.py

@@ -19,7 +19,7 @@
 import os
 import re
 import traceback
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Callable
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -80,7 +80,7 @@ def build_glob_alternation(
     return alternation, src_group_to_glob, tgt_group_to_glob
 
 
-class ConversionOps:
+class ConversionOps(ABC):
     """Base class for weight conversion operations."""
 
     def __repr__(self):

src/transformers/integrations/__init__.py

@@ -67,6 +67,7 @@
     ],
     "hqq": ["prepare_for_hqq_linear"],
     "hub_kernels": [
+        "FusedModuleBase",
         "LayerRepository",
         "lazy_load_kernel",
         "register_kernel_mapping",
@@ -221,6 +222,7 @@
     from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear
     from .hqq import prepare_for_hqq_linear
     from .hub_kernels import (
+        FusedModuleBase,
         LayerRepository,
         lazy_load_kernel,
         register_kernel_mapping,