Extended & simplified n-to-1 kernel fusion via KernelConfig
Author: michaelbenayoun
Pull request: yes

# 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

```python
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.

```python
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.

```python
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
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46339). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

For the `kernelize` refactor, it is done here: https://github.com/huggingface/transformers/pull/46520.

The doc PR is available here: #46754