Extended n-to-1 kernel fusion via `KernelConfig` by michaelbenayoun · Pull Request #45666 · huggingface/transformers

michaelbenayoun · 2026-04-27T21:37:20Z

What does this PR do?

Extends the KernelConfig API with two orthogonal capabilities:

Module fusion: specify how Transformers modules should be fused together before a custom kernel is applied (n-to-1 replacement).
Weight transformation in the live model: handle cases where a kernel expects weights in a different layout than the original modeling (e.g. fused linears).

These two capabilities are independent and can be combined.

Module fusion

Users can now specify a tuple of layers as a KernelConfig key to fuse multiple modules into a single kernel target:

  kernel_config = KernelConfig(
      {
          (
              ("RMSNorm", "model.layers.*.post_attention_layernorm"),
              ("MLP",     "model.layers.*.mlp"),
          ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP",
      },
  )
  model = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, kernel_config=kernel_config)

Internally, before model instantiation, register_kernel_fusions:

Meta-instantiates the model to discover parent classes containing all target children.
Creates a FusedModuleBase subclass representing the fused layer and a fused parent class that instantiates it.
Registers monkey patches via the patch API so the model is built with the fused structure from the start.
Resolves the tuple key to a scalar kernel_layer_name so the downstream kernelize pipeline is unchanged.

Weight transformation in the live model

Some kernels require the actual live parameters to be restructured (e.g. concatenating weights). This PR extends WeightTransform with a transform_model method that subclasses can override to restructure the live model graph before weights are loaded. The default is a no-op.

Kernels that require weight manipulation should provide a conversion_mapping class attribute; register_kernel_fusions picks it up automatically and appends it to the inferred transforms.

Related PRs

Model patching API #43917 — Original monkey patching API
[inference_fusion] convert conv3d patch embed to linear #45041 — FusedPatching support
n-to-1 kernel fusion via KernelConfig #45363 — Related open PR

HuggingFaceDocBuilderDev · 2026-04-27T21:49:14Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

michaelbenayoun · 2026-05-12T10:03:30Z

+            kernel_config = KernelConfig(
+                {
+                    (
+                        ("RMSNorm", "model.layers.*.post_attention_layernorm"),
+                        ("MLP", "model.layers.*.mlp"),
+                    ): "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations:RMSNormMLP",
+                }
+            )
+        else:
+            kernel_config = KernelConfig(
+                {
+                    (
+                        ("RMSNorm", "model.layers.*.post_attention_layernorm"),
+                        ("MLP", "model.layers.*.mlp"),
+                    ): "michaelbenayoun/dummy-rmsnorm-mlp:RMSNormMLP",
+                }
+            )


I can upload and use my test repo somewhere more official.

michaelbenayoun · 2026-05-12T10:04:04Z

+            if self.kernel_config is not None:
+                from kernels import use_kernel_mapping
+
+                inherit_mapping = not self.kernel_config.use_local_kernel
+                with use_kernel_mapping(self.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
+                    kernelize(self, device=Device(type=self.device.type), mode=mode)
+            else:
+                kernelize(self, device=Device(type=self.device.type), mode=mode)


These changes are required to make model.eval() work. Otherwise it wont apply the kernels.

Cyrilvallez

Hey @michaelbenayoun! Thanks for the PR!

In general, we do not want to add such complexity in core_model_loading.py. All model manipulations/module changes should happen before, and register proper ConversionOps so that the loader can create the correct weights. Model loading being a critical part of the lib, and having a lot of non trivial tweaks and tricks for perfs, adding such model operations at the same time will quickly become hard to understand/debug/profile.
The workflow we want to aim for here is:
initialize meta model -> identify which modules need to be fused/unfused etc -> monkey patch the modules and create WeightConversion for them -> defer to loader to create proper weights
This way we keep a clean loading flow, that is easy to understand and maintain.

cc @vasqu @IlyasMoutawwakil as well here, for kernels and monkey-patching API, if you can have an eye as well! This is non trivial and we need to make sure we have a clean integration for fused kernels!

Cyrilvallez · 2026-05-18T05:06:22Z

+    def transform_model(self, model: PreTrainedModel, layer_name: str) -> None:
+        source_names = list(self.layer_targets.get(layer_name, []))
+        if not source_names:
+            return
+        source_name = source_names[0]
+        if source_name == layer_name:
+            return
+
+        source_mod_path, _, source_attr = source_name.rpartition(".")
+        target_mod_path, _, target_attr = layer_name.rpartition(".")
+
+        # Idempotency guard: if the target already holds the parameter/buffer, the
+        # model is already in the correct shape (e.g. from monkey-patching before
+        # from_pretrained). Nothing to move.
+        try:
+            target_mod = model.get_submodule(target_mod_path) if target_mod_path else model
+            if target_attr in target_mod._parameters or target_attr in target_mod._buffers:
+                return
+        except AttributeError:
+            pass
+
+        try:
+            source_mod = model.get_submodule(source_mod_path) if source_mod_path else model
+        except AttributeError:
+            return  # already moved
+
+        if source_attr in source_mod._parameters:
+            obj = source_mod._parameters.pop(source_attr)
+            _get_submodule_or_create(
+                model, target_mod_path, original_module=source_mod, future_weight=obj
+            ).register_parameter(target_attr, obj)
+        elif source_attr in source_mod._buffers:
+            obj = source_mod._buffers.pop(source_attr)
+            _get_submodule_or_create(model, target_mod_path).register_buffer(target_attr, obj)
+
+        # Cleanup any empty modules that may be left after moving the parameter/buffer.
+        self.post_transformation_cleanup(model, source_names)


All these transforms should happen on the meta model, upstream of loading weights into it

It is already the case.
The function invoking the transform_model methods is called here: https://github.com/huggingface/transformers/pull/45666/changes#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaR4257

dacorvo · 2026-05-18T07:41:52Z

@michaelbenayoun do you have benchmark numbers comparing the full graphs obtained using torch.compile with and without these fused NKI kernels (and without any NKI kernels) ? My experience with NKI so far is that it does not produce faster graphs, although it can help working around compiler failures.

vasqu · 2026-05-26T13:51:03Z

Sorry got a bit delayed. Lmk when I should review seeing a lot of wip commits so assuming it's more of a draft for now?

michaelbenayoun · 2026-05-26T14:32:04Z

Sorry got a bit delayed. Lmk when I should review seeing a lot of wip commits so assuming it's more of a draft for now?

I just addressed @Cyrilvallez's comment. We do not change anything besides on the meta device, at instantiation-time and without transform_model methods.

michaelbenayoun · 2026-05-28T09:25:07Z

Gently pinging

ArthurZucker

My main concern with this approach is how non explicit it is, hiding a lot of stuff + it does not follow the standard contract:

model layers should be changes post meta init. (you are not running with all context manager etc)
weight conversion should not be Outside, but passed to the main conversion loading script. If you change conversion, be explicit
You need to supports distributed, quantization etc. Or not but that has to be explicit.
There should be a single point where model is modified, and where weights are modified.

Also to help review would help a fully documented example of what you put on the hub in just pure python:

conversion_mapping = [
--
WeightConverter(
["MLP.gate_proj", "MLP.up_proj"],
"MLP.gate_up_proj",
[Concatenate(dim=0)],
),
]
 
class RMSNormMLP(nn.Module):
"""
Fused RMSNorm + MLP kernel.
 
The fused module container exposes the source modules as named children:
- self.RMSNorm  — the RMSNorm layer (Qwen3RMSNorm / LlamaRMSNorm)
- self.MLP      — the MLP layer (Qwen3MLP / LlamaMLP)
 
This kernel runs them sequentially, making it a drop-in correctness
reference. A real kernel would fuse the two operations into a single
CUDA kernel.
"""
 
conversion_mapping = conversion_mapping
 
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states = self.RMSNorm(hidden_states)
gate, up = self.MLP.gate_up_proj(hidden_states).chunk(2, dim=-1)
hidden_states = self.MLP.down_proj(self.MLP.act_fn(gate) * up)
return hidden_states
 
 
class layers:
RMSNormMLP = RMSNormMLP

ArthurZucker · 2026-05-28T13:47:18Z

        if _torch_distributed_available and device_mesh is not None:  # add hooks to nn.Modules: no weights
            model = distribute_model(model, tp_plan, distributed_config, device_mesh, tp_size)

        # Prepare the full device map


why don't you call the model changes here, once everything is quantized and distributed ?

ArthurZucker · 2026-05-28T13:47:49Z

+        for module in meta_model.modules():
+            module_cls = type(module)
+            if module_cls in seen:
+                continue
+            if not all(hasattr(module, name) for name in child_names):
+                continue
+            seen.add(module_cls)
+
+            if kernel_conversion_mapping:
+                for source_name, child_name in zip(source_names, child_names):
+                    child_cls = type(getattr(module, child_name))
+                    if child_cls in seen_children:
+                        continue
+                    seen_children.add(child_cls)
+                    prefix_dot = source_name + "."
+                    relevant = [
+                        c
+                        for c in kernel_conversion_mapping
+                        if any(p.startswith(prefix_dot) for p in c._original_source_patterns)
+                    ]
+                    if relevant:
+                        stripped = [_strip_converter_prefix(c, source_name) for c in relevant]
+                        patch_mapping[child_cls.__name__] = _make_converted_child_class(child_cls, stripped)


for me it really does not make sense to have to do the conversions first before anything

ArthurZucker · 2026-05-28T13:50:37Z

+    return parent_mod._modules[mod_name]
+
+
+def _apply_weight_conversions(module: "nn.Module", conversion_mapping: list) -> None:


should never be done

ArthurZucker · 2026-05-28T13:51:04Z

+        return f"{self.__class__.__name__}(fused=({names}))"
+
+
+@functools.cache


this is the core part of your changes we need a documented example etc

ArthurZucker · 2026-05-28T13:53:29Z

+        # 1. Load the kernel class to get any conversion_mapping it declares.
+        kernel_cls = _try_load_kernel_class(repo_str, use_local=kernel_config.use_local_kernel)
+        kernel_conversion_mapping = (
+            list(kernel_cls.conversion_mapping)


This I like!

but why does the conversion mapping not include the fusing itself?

ArthurZucker · 2026-05-28T13:55:31Z

+
+        # 2. Meta-device scan — find parent classes that contain all target children.
+        with torch.device("meta"):
+            meta_model = cls(config)


for every single kernel mapped we re-init a meta model?

ArthurZucker · 2026-05-28T13:56:29Z

+        seen: set[type] = set()
+        seen_children: set[type] = set()
+        patch_mapping: dict[str, type] = {}
+        for module in meta_model.modules():


pretty sure you can avoid the double for loop or the order is just weird since you can iterate just once on all meta modules, collect into layer_groups, then fuse the inners groups.

ArthurZucker · 2026-05-28T13:57:48Z

+def make_fused_parent_class(
+    parent_cls: type,
+    child_names: list[str],
+    source_names: list[str],
+    kernel_layer_name: str,
+) -> type:


this is somewhat fine, but IDK how you are gonna setup a "contract" that the forward is always properly "kept" or something?

ArthurZucker · 2026-05-28T13:59:18Z

+    return WeightConverter(new_src, new_tgt, operations=converter.operations)
+
+
+def _apply_converter_to_module(module: "nn.Module", converter) -> None:


mmm not sure we should do that

michaelbenayoun added 13 commits April 10, 2026 14:52

feat: module fusion API for kernels

b387190

fix: improve __repr__ for fused modules

6bc9402

wip: integration to KernelConfig

62d4454

wip: add temporary example

4082fe1

wip: pattern matching in KernelConfig and actual kernel repo

ac4a699

refactor: move relevant code to hub_kernels.py

e13111f

docs: reformat docstring

d9d53f0

refactor: remove comment

e1c7f3f

Merge branch 'main' into fused_kernels

db0b7f0

Merge branch 'main' into fused_kernels

e21d06e

refactor: update example script for testing

bd640ae

wip: remove apply_fusions method

323b000

wip: add core feature for integration with the current fusing API

fe3002d

evalstate mentioned this pull request Apr 28, 2026

Cumulative defect fixes from recent Transformers PRs evalstate/transformers#41

Open

michaelbenayoun added 2 commits April 28, 2026 16:07

fix: move kernel mapping patching to kernelize

b541453

wip: update example script

b3d73a7

This was referenced Apr 29, 2026

Cumulative feature and defect updates from recent Transformers PRs evalstate/transformers#42

Open

Cumulative defect fixes from recent Transformers PRs evalstate/transformers#43

Open

michaelbenayoun added 11 commits May 1, 2026 10:50

wip: add transform_model method for WeightTransform

0845222

wip: conversion_mapping in Kernel

ecfd97d

Merge branch 'main' into extended_kernels_api

973a616

wip: remove things from __all__

0f0a64b

wip: remove imports

91177ae

fix: remove register_fusion_pattern path

2636c06

Merge branch 'main' into extended_kernels_api

573d9f0

fix: remove unused attribute

f7c15bd

wip: update experimentation script

f9d4299

refactor: add convert as abstract method

847dbd4

style: reformat hub_kernels.py

4443c9a

michaelbenayoun added 10 commits May 7, 2026 13:51

Merge branch 'main' into extended_kernels_api

d1bd06d

test: fix hub_kernels package reload

fb0a748

style: ruff

ded8b5f

Merge branch 'main' into extended_kernels_api

6d6411b

refactor: do not create dynamic classes in test

318553e

refactor: no dynamic class creation in tests

2a24760

refactor: test

8136781

Merge branch 'main' into extended_kernels_api

aa7743a

fix: TYPE_CHECKING imports were broken

17f4a9e

Merge branch 'main' into extended_kernels_api

8ec88d1

michaelbenayoun commented May 13, 2026

View reviewed changes

michaelbenayoun requested review from ArthurZucker and Cyrilvallez May 13, 2026 14:31

michaelbenayoun marked this pull request as ready for review May 13, 2026 14:33

Cyrilvallez reviewed May 18, 2026

View reviewed changes

michaelbenayoun added 7 commits May 19, 2026 11:58

wip: get rid of transform_model methods

fabadca

wip: move tests

e1f3a83

wip: make conversion happen before fused module instantiation

68c3659

refactor

041182f

wip: move conversion_mapping inside the init

519e673

wip: without any transform_model

88e0aee

wip: remove dead code

ad0e24e

michaelbenayoun added 2 commits May 26, 2026 16:01

wip: api imrpovement

3924cf3

wip: refactor

c948834

michaelbenayoun requested a review from Cyrilvallez May 27, 2026 15:44

ArthurZucker reviewed May 28, 2026

View reviewed changes

		return parent_mod._modules[mod_name]


		def _apply_weight_conversions(module: "nn.Module", conversion_mapping: list) -> None:

		return f"{self.__class__.__name__}(fused=({names}))"


		@functools.cache

		return WeightConverter(new_src, new_tgt, operations=converter.operations)


		def _apply_converter_to_module(module: "nn.Module", converter) -> None:

Uh oh!

Conversation

michaelbenayoun commented Apr 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Module fusion

Weight transformation in the live model

Related PRs

Uh oh!

HuggingFaceDocBuilderDev commented Apr 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dacorvo commented May 18, 2026

Uh oh!

vasqu commented May 26, 2026

Uh oh!

michaelbenayoun commented May 26, 2026

Uh oh!

michaelbenayoun commented May 28, 2026

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

michaelbenayoun commented Apr 27, 2026 •

edited

Loading