Skip to content

[kernels] change import time in KernelConfig #42004

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 4, 2025
Merged

[kernels] change import time in KernelConfig #42004

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2f9ef07
change import time
MekkCyber Nov 4, 2025
1db33c9
style
MekkCyber Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions src/transformers/utils/kernel_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,12 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from ..utils import PushToHubMixin, is_kernels_available, is_torch_available
from ..utils import PushToHubMixin, is_torch_available


if is_kernels_available():
from kernels import LayerRepository, Mode

if is_torch_available():
import torch

Expand Down Expand Up @@ -58,6 +55,8 @@ def infer_device(model):


def add_to_mapping(layer_name, device, repo_name, mode, compatible_mapping):
from kernels import LayerRepository

if device not in ["cuda", "rocm", "xpu"]:
raise ValueError(f"Only cuda, rocm, and xpu devices supported, got: {device}")
repo_layer_name = repo_name.split(":")[1]
Expand All @@ -82,6 +81,8 @@ def __init__(self, kernel_mapping={}):
self.registered_layer_names = {}

def update_kernel(self, repo_id, registered_name, layer_name, device, mode, revision=None):
from kernels import LayerRepository

self.kernel_mapping[registered_name] = {
device: {
mode: LayerRepository(
Expand Down Expand Up @@ -204,6 +205,8 @@ def create_compatible_mapping(self, model, compile=False):
The device is inferred from the model's parameters if not provided.
The Mode is inferred from the model's training state.
"""
from kernels import Mode

compatible_mapping = {}
for layer_name, kernel in self.kernel_mapping.items():
# Infer Mode: use Mode.TRAINING if model is training, else use Mode.INFERENCE
Expand Down