[kernels] When multiple device exist in `kernel_mapping`, calling `add_to_mapping()` will repeatedly overwrite the `repo_name`.

[zheliuyu]

System Info

device: A100-PCIE-40GB
torch: 2.7.1
kernels: main
transformers: main

Who can help?

@MekkCyber @drbh

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

1. First, the _KERNEL_MAPPING in hub_kernels.py must be cleared to ensure kernels are specified entirely through kernel_mapping.

2. Run script

from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig
import torch

model_id = "meta-llama/Llama-3.2-1B"

kernel_mapping = {
    "RMSNorm": {
        "cuda":
            "kernels-community/layer_norm:LlamaRMSNorm"
        "rocm":
            "kernels-community/layer_norm:LlamaRMSNorm"
    }
}

kernel_config = KernelConfig(kernel_mapping)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    use_kernels=True,
    kernel_config=kernel_config
)

When multiple device exist in kernel_mapping, calling add_to_mapping() will repeatedly overwrite the repo_name.

Since ROCm is configured as the second device while I'm currently working in an A100 env, the RMSNorm kernels remain inactive.

Expected behavior

Does the repeated invocation of add_to_mapping() causing overwrites in kernel_mapping represent an issue that needs to be addressed?

# In transformers/src/transformers/utils/kernel_config.py
elif isinstance(kernel, dict):
    for device, repo_name in kernel.items():
        add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping)

[i3hz]

Alright so I've been trying to fix this issue for a bit and I've run into some issues . My idea was to just check if the device from
for device, repo_name is kernel.items(): was same as the current device(via infer_device) and if it is then we call add_to_mapping.

elif isinstance(kernel, dict):
                for device, repo_name in kernel.items():
                    if device!=current_device:
                        continue
                    add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping)

Would this fix the issue ?
I'm having some trouble with the reproduction script but my kernel mapping does change to {'RMSNorm': {'cuda': {<Mode.INFERENCE: 4>: <kernels.layer.LayerRepository object at 0x72d3c59f5490>}}} after we load the model

[Aaraviitkgp]

@i3hz @zheliuyu I have made a pr on this issue you can check if there is any issue.

[i3hz]

Well I was planning to make the PR but I wasn't sure if that's the intended fix .
In any case your PR does seem consistent with my logic .

[zheliuyu]

Alright so I've been trying to fix this issue for a bit and I've run into some issues . My idea was to just check if the device from for device, repo_name is kernel.items(): was same as the current device(via infer_device) and if it is then we call add_to_mapping.

elif isinstance(kernel, dict):
for device, repo_name in kernel.items():
if device!=current_device:
continue
add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping)
Would this fix the issue ? I'm having some trouble with the reproduction script but my kernel mapping does change to {'RMSNorm': {'cuda': {<Mode.INFERENCE: 4>: <kernels.layer.LayerRepository object at 0x72d3c59f5490>}}} after we load the model

This is indeed one approach to resolve this issue.

Let's review this PR together. #42466

[ada-ggf25]

Hi! Didn't saw that you were already working on this Issue, I also started working on it once it came out. I also just did a Pull Request (#42486) to try to fix this issue, if you want to take a look at my solution. I hope it helps!