Codestin Search App

zheliuyu · 2025-12-11T09:20:50Z

What does this PR do?

Add local kernel loading support to transformers.KernelConfig().

Relaed issue: #42739

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? 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, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Test

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

Rocketknight1 · 2025-12-11T14:16:46Z

cc @MekkCyber

MekkCyber · 2025-12-11T14:32:19Z

Hi @zheliuyu ! Thank you for this nice PR, let me know when it's ready

zheliuyu · 2025-12-12T02:38:19Z

After submitting the draft yesterday, some unknown CI errors appeared, but they seem to have been resolved now. So, this PR is ready for review. @MekkCyber @Rocketknight1

Let me offer some context for the review:

To avoid complex validation in from_pretrained(), no new flag (e.g., use_local_kernels) was added, the existing use_kernels flag suffices. The decision to load locally should be handled when parsing the kernel_mapping field.
Specify the input format for locally loading kernels as RMSNorm":"/abs-path:package_name:layer_name. This format can leverage the existing string validation logic from store_registered_layer_names().
Add a check_kernel_from_local flag for differentiating remote and local loading.
Create a new add_to_mapping_local to parse the repo_path, package_name, and layer_name fields from kernel_mapping for kernels.LocalLayerRepository.
Param inherit_mapping should be False to avoid still loading kernel from remote when enable check_kernel_from_local=True.
Func update_kernel() is currently not in use. Remove it?

MekkCyber

Thanks @zheliuyu ! I think we could simplify a bit how we specify the kernel path, wdyt

MekkCyber · 2025-12-15T13:40:53Z

        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:package_name:layer_name'.


let's just simplify to /abs-path:layer_name

Thanks for your reply.

For LocalLayerRepository, the package_name parameter appears to be required. https://github.com/huggingface/kernels/blob/main/src/kernels/layer/layer.py#L119

but we should be able to get it from the path directly no ?

You are right.

stevhliu · 2025-12-15T23:37:25Z

Cool! Commenting to add some docs for this once #42277 is merged

zheliuyu · 2025-12-16T06:50:09Z

@MekkCyber Simplified kernel path.

Test passed

kernel_mapping = {
    "RMSNorm":
        "/root/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
)

MekkCyber

Thanks @zheliuyu ! I think it would be better to just let the user specify if they want local or hub kernels, wdyt

MekkCyber · 2025-12-16T09:38:13Z

+            if kernel is not None and kernel[0] == "/":
+                self.check_kernel_from_local = True


shouldn't be done here, let's ask the user to do this instead, in the __init__ we add an other argument use_local_kernels

It's a good idea to let users decide on local kernel loading. 🍻

Co-authored-by: Mohamed Mekkouri <[email protected]>

MekkCyber

Thanks for iterating! only few nits and good to go

MekkCyber · 2025-12-16T12:28:50Z

+                inherit_mapping = not self.use_local_kernel
+                with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):


Suggested change

inherit_mapping = not self.use_local_kernel

with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):

inherit_mapping = not kernel_config.use_local_kernel

with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):

MekkCyber · 2025-12-16T12:30:48Z

    def __init__(self, kernel_mapping={}):
        self.kernel_mapping = kernel_mapping
        self.registered_layer_names = {}
+        self.use_local_kernel = False


we need to set self.use_local_kernel to use_local_kernel passed in the constructor, and defaults to False

zheliuyu · 2025-12-16T12:41:39Z

Test passed

kernel_mapping = {
    "RMSNorm":
        "/root/liger_kernels:LigerRMSNorm",
}

+ kernel_config = KernelConfig(kernel_mapping, use_local_kernel=True)

# Load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    use_kernels=True,
    kernel_config=kernel_config
)

MekkCyber

Thank you lgtm ! will add some docs about this once the doc pr is merged cc @stevhliu

HuggingFaceDocBuilderDev · 2025-12-16T13:01:16Z

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.

* add add_to_mapping_local for KernelConfig * refactor kernel_mapping format * lint code * specify the kernel path * fix `abs/path` Co-authored-by: Mohamed Mekkouri <[email protected]> * rename check_kernel_from_local to use_local_kernel --------- Co-authored-by: Mohamed Mekkouri <[email protected]>

zheliuyu added 3 commits December 10, 2025 17:29

add add_to_mapping_local for KernelConfig

4427cfa

refactor kernel_mapping format

6777f05

Merge branch 'huggingface:main' into main

b17edb5

zheliuyu changed the title ~~Add local kernel loading support to transformers.KernelConfig().~~ Add local kernel loading support to KernelConfig(). Dec 11, 2025

lint code

db38adc

Merge branch 'main' into main

b71c5d7

zheliuyu marked this pull request as ready for review December 12, 2025 01:57

github-actions Bot requested review from ArthurZucker and Rocketknight1 December 12, 2025 01:57

MekkCyber reviewed Dec 15, 2025

View reviewed changes

zheliuyu added 2 commits December 16, 2025 10:40

Merge branch 'huggingface:main' into main

b87a9fd

specify the kernel path

5eebc2b

MekkCyber reviewed Dec 16, 2025

View reviewed changes

fix abs/path

e41024a

Co-authored-by: Mohamed Mekkouri <[email protected]>

MekkCyber reviewed Dec 16, 2025

View reviewed changes

rename check_kernel_from_local to use_local_kernel

20317d8

MekkCyber approved these changes Dec 16, 2025

View reviewed changes

Merge branch 'main' into main

d03305d

MekkCyber merged commit 2427512 into huggingface:main Dec 16, 2025
25 checks passed

zheliuyu mentioned this pull request Dec 17, 2025

How about adding local kernel loading to transformers.KernelConfig() #42739

Closed

zheliuyu mentioned this pull request Jan 6, 2026

Add huggingface/kernels part Ascend/docs#94

Merged

stevhliu mentioned this pull request Jan 14, 2026

[docs] kernels #43269

Merged

		if kernel is not None and kernel[0] == "/":
		self.check_kernel_from_local = True

		inherit_mapping = not self.use_local_kernel
		with use_kernel_mapping(kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):

Uh oh!

Conversation

zheliuyu commented Dec 11, 2025

What does this PR do?

Before submitting

Test

Uh oh!

Rocketknight1 commented Dec 11, 2025

Uh oh!

MekkCyber commented Dec 11, 2025

Uh oh!

zheliuyu commented Dec 12, 2025

Uh oh!

MekkCyber 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!

stevhliu commented Dec 15, 2025

Uh oh!

zheliuyu commented Dec 16, 2025

Test passed

Uh oh!

MekkCyber left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MekkCyber 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!

zheliuyu commented Dec 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test passed

Uh oh!

MekkCyber left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Dec 16, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zheliuyu commented Dec 16, 2025 •

edited

Loading

MekkCyber left a comment •

edited

Loading