Repository: huggingface/transformers
Pull request: Add local kernel loading support to KernelConfig().

# What does this PR do?

Add local kernel loading support to transformers.KernelConfig().

Relaed issue: https://github.com/huggingface/transformers/issues/42739

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Test
```python
import time
import logging
from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig


# Set the level to `DEBUG` to see which kernels are being called.
logging.basicConfig(level=logging.DEBUG)

model_name = "/root/Qwen3"

kernel_mapping = {
    "RMSNorm":
        "/root/liger_kernels:liger_kernels:LigerRMSNorm",
}

kernel_config = KernelConfig(kernel_mapping)

# Load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    kernel_config=kernel_config
)

# Prepare the model input
prompt = "Output the first 20 digits of pi."
messages = [{"role": "user", "content": prompt}]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=False,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# Warm_up
for _ in range(2):
    generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()

# Print Runtime
for _ in range(5):
    start_time = time.time()
    generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
    print("runtime: ", time.time() - start_time)
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
    print("content:", content)
```


**Output:**
```
use_kernel_func_from_hub is not available in the installed kernels version. Please upgrade kernels to use this feature.
A kernel_config was provided but use_kernels is False; setting use_kernels=True automatically. To suppress this warning, explicitly set use_kernels to True.
Loading weights: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 665.31it/s, Materializing param=model.norm.weight]
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
INFO:root:Using layer `LigerRMSNorm` from repo `/root/liger_kernels` (package: liger_kernels), layer `LigerRMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.INFERENCE
INFO:root:Using layer `LigerRMSNorm` from repo `/root/liger_kernels` (package: liger_kernels), layer `LigerRMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.INFERENCE
/root/miniconda3/lib/python3.12/site-packages/kernels/layer.py:918: UserWarning: 
No kernel mapping found for layer `SiLU`. Check if the layer name matches one of the kernels in the mapping or add the kernel you want to use to the mapping. Defaulting to original forward implementation.
  warnings.warn(
INFO:root:Using layer `LigerRMSNorm` from repo `/root/liger_kernels` (package: liger_kernels), layer `LigerRMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.INFERENCE

...

INFO:root:Using layer `LigerRMSNorm` from repo `/root/liger_kernels` (package: liger_kernels), layer `LigerRMSNorm`
DEBUG:root:kernelize mode: Mode.INFERENCE, repo mode: Mode.INFERENCE
runtime:  1.0340442657470703
content: The first 20 digits of π (pi) are:

**3.14159265358979323846**.
runtime:  0.9599497318267822
content: The first 20 digits of π are:

**3.14159265358979323846**.

Let me know if you'd like the next 20 digits!
runtime:  0.8574888706207275
content: The first 20 digits of π are:

**3.14159265358979323846**

Let me know if you'd like the next 10 digits or more!
runtime:  0.7771825790405273
content: The first 20 digits of π are:

**3.14159265358979323846**

Let me know if you'd like more digits!
runtime:  0.6569499969482422
content: The first 20 digits of π are:  
**3.14159265358979323846**.
```


Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -3533,7 +3533,9 @@ def set_use_kernels(self, use_kernels, kernel_config):
 
                 # This is a context manager to override the default kernel mapping
                 # We are calling kernelize inside this context manager using the use_kernels setter
-                with use_kernel_mapping(kernel_config.kernel_mapping):
+                # Param inherit_mapping should be False to avoid still loading kernel from remote
+                inherit_mapping = not kernel_config.use_local_kernel
+                with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
                     self.use_kernels = True
             # We use the default kernel mapping in .integrations.hub_kernels
             else:

filename: src/transformers/utils/kernel_config.py
patch: @@ -71,14 +71,36 @@ def add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping):
     }
 
 
+def add_to_mapping_local(layer_name, device, repo_name, mode, compatible_mapping):
+    from pathlib import Path
+
+    from kernels import LocalLayerRepository
+
+    if device not in ["cuda", "rocm", "xpu", "npu"]:
+        raise ValueError(f"Only cuda, rocm, xpu and npu devices supported, got: {device}")
+    repo_layer_name = repo_name.split(":")[1]
+    repo_path = repo_name.split(":")[0]
+    repo_package_name = repo_path.split("/")[-1]
+    compatible_mapping[layer_name] = {
+        device: {
+            mode: LocalLayerRepository(
+                repo_path=Path(repo_path),
+                package_name=repo_package_name,
+                layer_name=repo_layer_name,
+            )
+        }
+    }
+
+
 class KernelConfig(PushToHubMixin):
     """
     Kernel configuration class. This class is used to configure the kernel mapping for a model.
     """
 
-    def __init__(self, kernel_mapping={}):
+    def __init__(self, kernel_mapping={}, use_local_kernel=False):
         self.kernel_mapping = kernel_mapping
         self.registered_layer_names = {}
+        self.use_local_kernel = use_local_kernel
 
     def update_kernel(self, repo_id, registered_name, layer_name, device, mode, revision=None):
         from kernels import LayerRepository
@@ -105,6 +127,7 @@ def sanitize_kernel_mapping(self, model):
         2. Each kernel value is either a string of the form 'org/repo:layer_name' or a dict mapping device types ("cuda", "rocm", "xpu", "npu") to such strings.
         3. Each device key in a dict is one of "cuda", "rocm", "xpu", or "npu".
         4. Each repo_name is a valid repository and layer name in the format 'org/repo:layer_name' (i.e., a string containing both a slash and a colon).
+        5. If a local path is detected, it should be in the format '/abs/path:layer_name'. The absolute path must include the `package_name`, like "/home/user/layer_norm".
 
         Args:
             model: The model instance whose modules are checked for registered kernel_layer_name attributes.
@@ -114,14 +137,13 @@ def sanitize_kernel_mapping(self, model):
                         or if a repo_name is not a valid 'org/repo:layer_name' string.
         """
         MAPPING_FORMAT = """
+        For single device form remote
         {
             "RMSNorm":
                 "kernels-community/layer_norm:LlamaRMSNorm",
             ...
         },
-
-        or
-
+        For multiple devices form remote
         {
             "RMSNorm": {
                 "cuda":
@@ -132,6 +154,23 @@ def sanitize_kernel_mapping(self, model):
             },
             ...
         }
+        For single device form local
+        {
+            "RMSNorm":
+                "/abs/path:LlamaRMSNorm",
+            ...
+        },
+        For multiple devices form local
+        {
+            "RMSNorm": {
+                "cuda":
+                    "/abs/path:LlamaRMSNorm",
+                "rocm":
+                    "/abs/path:LlamaRMSNorm",
+                ...
+            },
+            ...
+        }
         """
         self.store_registered_layer_names(model)
         # Validate that the kernel mapping is a dict
@@ -149,7 +188,7 @@ def sanitize_kernel_mapping(self, model):
             if isinstance(kernel, str):
                 if "/" not in kernel or ":" not in kernel:
                     raise ValueError(
-                        f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name'), got: {kernel}"
+                        f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name' or '/abs/path:layer_name'), got: {kernel}"
                     )
 
             elif isinstance(kernel, dict):
@@ -159,9 +198,8 @@ def sanitize_kernel_mapping(self, model):
 
                     if not isinstance(repo_name, str) or "/" not in repo_name or ":" not in repo_name:
                         raise ValueError(
-                            f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name'), got: {repo_name}"
+                            f"Kernel mapping for '{layer_name}' must be a valid repo name with a layer name (e.g., 'org/repo:layer_name' or '/abs/path:layer_name'), got: {repo_name}"
                         )
-
             else:
                 raise ValueError(f"Kernel mapping must follow the format: {MAPPING_FORMAT}, got: {kernel}")
 
@@ -174,18 +212,13 @@ def create_compatible_mapping(self, model, compile=False):
                 ...
             },
 
-            or
+            or for local path:
 
             {
-                "RMSNorm": {
-                    "cuda":
-                        "kernels-community/layer_norm:LlamaRMSNorm",
-                    "rocm":
-                        "kernels-community/layer_norm:LlamaRMSNorm",
-                    ...
-                },
+                "RMSNorm":
+                    "/home/user/liger_kernels:LigerRMSNorm",
                 ...
-            }
+            },
 
         into a nested mapping:
 
@@ -200,6 +233,20 @@ def create_compatible_mapping(self, model, compile=False):
                 }
             }
 
+            or for local path:
+
+            {
+                "RMSNorm": {
+                    "cuda": {
+                        Mode.INFERENCE: LocalLayerRepository(
+                            repo_path=Path("/home/user/liger_kernels"),
+                            package_name="liger_kernels",
+                            layer_name="LigerRMSNorm",
+                        )
+                    }
+                }
+            }
+
         that's compatible with the kernels library.
 
         The device is inferred from the model's parameters if not provided.
@@ -217,11 +264,17 @@ def create_compatible_mapping(self, model, compile=False):
 
             if isinstance(kernel, str):
                 repo_name = kernel
-                add_to_mapping(layer_name, current_device, repo_name, mode, compatible_mapping)
+                if not self.use_local_kernel:
+                    add_to_mapping(layer_name, current_device, repo_name, mode, compatible_mapping)
+                else:
+                    add_to_mapping_local(layer_name, current_device, repo_name, mode, compatible_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)
+                    if not self.use_local_kernel:
+                        add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping)
+                    else:
+                        add_to_mapping_local(layer_name, device, repo_name, mode, compatible_mapping)
 
         self.kernel_mapping = compatible_mapping