Extended & simplified n-to-1 kernel fusion via KernelConfig

[michaelbenayoun]

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).

• Parameter transformation: handle cases where a kernel expects weights in a different layout than the original modeling (e.g. fused linears).

Compared to previous PR, this approach is more explicit and way simpler, putting much of the burden to the kernel authors.

How it works

The kernel author needs to define two classes:

• KernelName: defines the forward pass, used by the kernels library to kernelize the model
• KernelNameLayout: defines the conversion_mapping as well as an __init__ method. This is used to monkey-patch the model

Having two classes because the kernels library prevents us from having stateful kernel classes.
While it might not be as pleasing as having one big class, it separates concerns.

Script for the examples

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", use_kernels=True)
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_repo_id = "michaelbenayoun/dummy-rmsnorm-mlp-with-transformations-and-init:RMSNormMLP"
kernel_repo_id = "michaelbenayoun/dummy-rmsnorm-kernel-with-init:CustomRMSNorm"
kernel_config = KernelConfig(
    {
        # (
        #     ("RMSNorm", "model.layers.*.post_attention_layernorm"),
        #     ("MLP",     "model.layers.*.mlp"),
        # ): kernel_repo_id,
        "RMSNorm": 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())

Example 1: Parameter transformation, no fusion

In this case, the KernelNameLayout class's __init__ method has the same signature as the module being replaced.

import torch
import torch.nn as nn

from transformers.conversion_mapping import WeightRenaming

class CustomRMSNormLayout(nn.Module):
    conversion_mapping = [
        WeightRenaming(
            source_patterns=r"(.*(?:input_layernorm|post_attention_layernorm|norm)\.)weight",
            target_patterns=r"\1scale",
        ),
    ]

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        pass  # replaced at runtime by kernelize


class CustomRMSNorm(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        print("This dummy kernel is used")
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.scale * hidden_states.to(input_dtype)


class layers:
    CustomRMSNorm = CustomRMSNorm

Example 2: Fusion and parameter transformation

Compared to the first example, here we will fuse two modules in the original model into one module.
Because of this, the __init__ method does not have the same signature, but rather take the instantiated modules it's fusing.

import torch
import torch.nn as nn

from transformers import Concatenate, WeightConverter
from transformers.conversion_mapping import WeightRenaming


class RMSNormMLPLayout(nn.Module):
    conversion_mapping = [
        # norm.weight → scale (placed at post_attention_layernorm.scale)
        WeightRenaming(
            source_patterns=r"(.*post_attention_layernorm\.)weight",
            target_patterns=r"\1scale",
        ),
        # mlp.gate_proj + mlp.up_proj → post_attention_layernorm.gate_up_proj (concat)
        WeightConverter(
            ["mlp.gate_proj", "mlp.up_proj"],
            "post_attention_layernorm.gate_up_proj",
            [Concatenate(dim=0)],
        ),
        # mlp.down_proj.* → post_attention_layernorm.down_proj.*
        WeightRenaming(
            source_patterns=r"(.*\.)mlp\.(down_proj\..*)",
            target_patterns=r"\1post_attention_layernorm.\2",
        ),
    ]

    def __init__(self, norm, mlp):
        super().__init__()
        self.variance_epsilon = norm.variance_epsilon
        self.scale = nn.Parameter(torch.empty_like(norm.weight))
        self.gate_up_proj = nn.Linear(
            mlp.gate_proj.in_features,
            mlp.gate_proj.out_features + mlp.up_proj.out_features,
            bias=False,
            device=mlp.gate_proj.weight.device,
            dtype=mlp.gate_proj.weight.dtype,
        )
        self.down_proj = mlp.down_proj
        self.act_fn = mlp.act_fn

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        pass


class RMSNormMLP(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.scale * hidden_states.to(input_dtype)
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.down_proj(self.act_fn(gate) * up)


class layers:
    RMSNormMLP = RMSNormMLP

[HuggingFaceDocBuilderDev]

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.

[ArthurZucker]

Much better!

[ArthurZucker]

Suggested change

	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)
	kernelize(self, device=Device(type=self.device.type), mode=mode, self.kernel_config)

let's reduce surface as much as possible

[michaelbenayoun]

kernelize is defined in kernels. I can make a PR there, but for now it cannot be changed here.

[ArthurZucker]

okay! we can also just create def kernelize to put in kernels utils!

[ArthurZucker]

I don't think we need to iterate over all the modules!
We could register like we do for the tp plan with explicit path, we like explicitness in general!

{ "layers.*.self_attn.q_proj" : XXXX} 

[ArthurZucker]

MOST important comment IMO if the contract is more like this we have a lot of simplifications no?

[michaelbenayoun]

We already have this contract.

kernel_config = KernelConfig(
    {
        (
            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
            ("MLP",     "model.layers.*.mlp"),
        ): kernel_repo_id,
    },
)

I will update this loop

[ArthurZucker]

let's do both in a single func!

[michaelbenayoun]

Done

[ArthurZucker]

to remove

[michaelbenayoun]

Done

[ArthurZucker]

this is super important needs to be documented well:

• we replace the fused cls by identity
• thus we have to patch some inits, etc etc c
  also do we even have to patch inits when the proper class replaces the one that holds them?

[michaelbenayoun]

Done

[ArthurZucker]

Much much better! Its just missing a piece of doc / update the doc for monkey patching, maybe some bench if you have but that's fine for another PR !

Ty for iterating its quite nice now!

[ArthurZucker]

Nice

[ArthurZucker]

nice

[ArthurZucker]

Suggested change

	meta_model = cls(config)
	meta__modules = cls(config).named_modules()

we only need these

[ArthurZucker]

ah maybe it gets updated but that's good, you can'tupdate twice so its even better in a way no? (to not re-compute the named modules)

[ArthurZucker]

okay! we can also just create def kernelize to put in kernels utils!

[michaelbenayoun]

For the kernelize refactor, it is done here: #46520.