🚨 Kernels=True ?!

conftest.py

@@ -30,6 +30,7 @@
 import _pytest
 import pytest
 
+from transformers import PreTrainedModel
 from transformers.testing_utils import (
     HfDoctestModule,
     HfDocTestParser,
@@ -39,7 +40,7 @@
     patch_testing_methods_to_collect_info,
     patch_torch_compile_force_graph,
 )
-from transformers.utils import enable_tf32
+from transformers.utils import ENV_VARS_TRUE_VALUES, enable_tf32
 from transformers.utils.network_logging import register_network_debug_plugin
 
 
@@ -213,6 +214,10 @@ def wrapper(*args, **kwargs):
 # they become normal warnings - i.e. the tests still need to test the current functionality
 warnings.simplefilter(action="ignore", category=FutureWarning)
 
+# we patch `from_pretrained` temporarily to be controlled via env variables
+_from_pretrained_kernels_patch = None
+_original_from_pretrained = PreTrainedModel.from_pretrained.__func__
+
 
 def pytest_configure(config):
     # Shared directory for the read-only cache fallback events. The controller creates it and
@@ -298,6 +303,32 @@ def pytest_configure(config):
     os.environ["DISABLE_SAFETENSORS_CONVERSION"] = "true"
     register_network_debug_plugin(config)
 
+    # Kernels related patching
+    global _from_pretrained_kernels_patch, _original_from_pretrained
+    os.environ["USE_HUB_KERNELS"] = "0"  # Disable kernels for tests by default
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        # Manipulate default as per env variables
+        kwargs.setdefault("use_kernels", os.environ.get("USE_HUB_KERNELS", "NO").upper() in ENV_VARS_TRUE_VALUES)
+        return _original_from_pretrained(cls, *args, **kwargs)
+
+    _from_pretrained_kernels_patch = mock.patch.object(
+        PreTrainedModel,
+        "from_pretrained",
+        from_pretrained,
+    )
+    _from_pretrained_kernels_patch.start()
+
+
+def pytest_unconfigure(config):
+    # Disable default kwargs per env for kernels activation
+    global _from_pretrained_kernels_patch
+
+    if _from_pretrained_kernels_patch is not None:
+        _from_pretrained_kernels_patch.stop()
+        _from_pretrained_kernels_patch = None
+
 
 def pytest_collection_modifyitems(items):
     for item in items:

docker/transformers-all-latest-gpu/Dockerfile

@@ -143,9 +143,6 @@ RUN python3 -m pip install --no-cache-dir einops
 # For `VibeVoice` (added in PR #40546)
 RUN python3 -m pip install --no-cache-dir diffusers
 
-# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs
-RUN python3 -m pip uninstall -y kernels
-
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop

docker/transformers-pytorch-amd-gpu/Dockerfile

@@ -42,9 +42,6 @@ RUN cd transformers && python3 setup.py develop
 # Remove nvml and nvidia-ml-py as it is not compatible with ROCm. apex is not tested on NVIDIA either.
 RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y
 
-# `kernels` may causes many failing tests
-RUN python3 -m pip uninstall -y kernels
-
 # On ROCm, torchcodec is required to decode audio files. 0.10 matches the
 # torch 2.10 C++ ABI in the base image above;
 RUN python3 -m pip install --no-cache-dir "torchcodec==0.10"

docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile

@@ -48,6 +48,3 @@ RUN python3 -c "from deepspeed.launcher.runner import main"
 
 # Remove nvml as it is not compatible with ROCm
 RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y
-
-# `kernels` may causes many failing tests
-RUN python3 -m pip uninstall -y kernels

docker/transformers-pytorch-deepspeed-latest-gpu/Dockerfile

@@ -50,9 +50,6 @@ RUN python3 -m pip uninstall -y deepspeed
 # TODO: Find out why test fail.
 RUN DS_BUILD_CPU_ADAM=1 DS_BUILD_FUSED_ADAM=1 python3 -m pip install deepspeed --no-build-isolation --config-settings="--build-option=build_ext" --config-settings="--build-option=-j8" --no-cache -v --disable-pip-version-check 2>&1
 
-# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs
-RUN python3 -m pip uninstall -y kernels
-
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop

docker/transformers-pytorch-deepspeed-nightly-gpu/Dockerfile

@@ -62,9 +62,6 @@ RUN python3 -m pip uninstall -y deepspeed
 #RUN git clone https://github.com/pytorch/TensorRT.git
 #RUN cd TensorRT/py && python3 setup.py install --fx-only
 
-# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs
-RUN python3 -m pip uninstall -y kernels
-
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop

docker/transformers-pytorch-gpu/Dockerfile

@@ -28,9 +28,6 @@ RUN [ ${#TORCH_AUDIO} -gt 0 ] && VERSION='torchaudio=='TORCH_AUDIO'.*' ||  VERSI
 RUN python3 -m pip install --no-cache-dir git+https://github.com/facebookresearch/detectron2.git pytesseract
 RUN python3 -m pip install -U "itsdangerous<2.1.0"
 
-# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs
-RUN python3 -m pip uninstall -y kernels
-
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop

docker/transformers-quantization-latest-gpu/Dockerfile

@@ -85,9 +85,6 @@ RUN python3 -m pip install --no-cache-dir torchao
 # Add transformers in editable mode
 RUN python3 -m pip install --no-cache-dir -e ./transformers[dev]
 
-# `kernels` may give different outputs (within 1e-5 range) even with the same model (weights) and the same inputs
-RUN python3 -m pip uninstall -y kernels
-
 # Uninstall flash-attn installed by autoawq, it causes issues here : https://github.com/huggingface/transformers/actions/runs/15915442841/job/44892146131
 RUN python3 -m pip uninstall -y flash-attn
 

setup.py

@@ -175,7 +175,7 @@ def deps_list(*pkgs):
 
 extras = {}
 
-extras["torch"] = deps_list("torch", "accelerate")
+extras["torch"] = deps_list("torch", "accelerate", "kernels")
 extras["vision"] = deps_list("torchvision", "Pillow")
 extras["audio"] = deps_list("torchaudio", "librosa", "pyctcdecode", "phonemizer")
 if PYTHON_MINOR_VERSION < 13:
@@ -186,7 +186,6 @@ def deps_list(*pkgs):
     "datasets", "ruff", "GitPython", "urllib3", "libcst", "rich", "ty", "tomli", "transformers-mlinter"
 )
 extras["docs"] = deps_list("hf-doc-builder")
-extras["kernels"] = deps_list("kernels")
 extras["sentencepiece"] = deps_list("sentencepiece", "protobuf")
 extras["tiktoken"] = deps_list("tiktoken", "blobfile")
 extras["mistral-common"] = deps_list("mistral-common[image]")
@@ -249,7 +248,6 @@ def deps_list(*pkgs):
     + extras["vision"]
     + extras["audio"]
     + extras["video"]
-    + extras["kernels"]
     + extras["timm"]
     + extras["sentencepiece"]
     + extras["tiktoken"]

src/transformers/integrations/accelerate.py

@@ -1021,3 +1021,27 @@ def wrapped(self, *args, **kwargs):
         return wrapped
 
     return decorator
+
+
+def device_map_uses_accelerator(device_map: dict | str | int | None) -> bool:
+    """Detect whether the device map contains any non-cpu device"""
+    if device_map is None:
+        return False
+
+    # Single device
+    if not isinstance(device_map, dict):
+        if isinstance(device_map, int):
+            return True
+        device = torch.device(device_map)
+        return device.type not in {"cpu", "meta"}
+
+    # Multi device
+    for target in device_map.values():
+        if target == "disk":
+            continue
+        if isinstance(target, int):
+            return True
+        if torch.device(target).type not in {"cpu", "meta"}:
+            return True
+
+    return False

src/transformers/integrations/hub_kernels.py

@@ -21,12 +21,12 @@
 from contextlib import contextmanager
 from pathlib import Path
 from types import ModuleType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from ..conversion_mapping import get_checkpoint_conversion_mapping, register_checkpoint_conversion_mapping
+from ..modeling_flash_attention_utils import FLASH_ATTN_KERNEL_VERSIONS
 from ..monkey_patching import register_patch_mapping
 from ..utils import ENV_VARS_TRUE_VALUES, logging
-from ..utils.generic import is_flash_attention_requested
 from ..utils.import_utils import (
     KERNELS_MAX_VERSION,
     KERNELS_MIN_VERSION,
@@ -58,8 +58,17 @@
 )
 
 
-_TRANSFORMERS_USE_HUB_KERNELS = os.environ.get("USE_HUB_KERNELS", "YES").upper()
-_kernels_enabled = _TRANSFORMERS_USE_HUB_KERNELS in ENV_VARS_TRUE_VALUES
+def _kernels_enabled() -> bool:
+    return os.environ.get("USE_HUB_KERNELS", "YES").upper() in ENV_VARS_TRUE_VALUES
+
+
+def _ensure_kernels_enabled() -> None:
+    if not _kernels_enabled():
+        value = os.environ.get("USE_HUB_KERNELS", "YES")
+        raise ValueError(
+            "Hub kernels are disabled through the environment variable "
+            f"`USE_HUB_KERNELS={value}`. Set `USE_HUB_KERNELS=YES` to enable them."
+        )
 
 
 # Maps from func name to the internal module path
@@ -98,37 +107,21 @@
     )
     from kernels import use_kernelized_func as _kernels_use_kernelized_func
 
+    # We change the strategy to prepare the metadata like kernels but the exchange
+    # at runtime can be determined by the env variable `USE_HUB_KERNELS`
     def use_kernel_forward_from_hub(layer_name: str):
-        if _kernels_enabled:
-            return _kernels_use_kernel_forward_from_hub(layer_name)
-        else:
-            logger.warning_once(
-                f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"
-            )
-            return lambda cls: cls
+        return _kernels_use_kernel_forward_from_hub(layer_name)
 
     def use_kernelized_func(module_names: list[Callable] | Callable):
-        if _kernels_enabled:
-            if isinstance(module_names, Callable):
-                module_names = [module_names]
-            return _kernels_use_kernelized_func(*module_names)
-        else:
-            logger.warning_once(
-                f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"
-            )
-            return lambda cls: cls
+        if isinstance(module_names, Callable):
+            module_names = [module_names]
+        return _kernels_use_kernelized_func(*module_names)
 
     def use_kernel_func_from_hub(layer_name: str):
-        if _kernels_enabled:
-            logger.warning_once(
-                "`use_kernel_func_from_hub` is deprecated in transformers v5.16 and will be removed in the future. Please use `use_kernel_forward_from_hub` instead."
-            )
-            return _kernels_use_kernel_forward_from_hub(layer_name)
-        else:
-            logger.warning_once(
-                f"kernels hub usage is disabled through the environment USE_HUB_KERNELS={_TRANSFORMERS_USE_HUB_KERNELS}"
-            )
-            return lambda cls: cls
+        logger.warning_once(
+            "`use_kernel_func_from_hub` is deprecated in transformers v5.16 and will be removed in the future. Please use `use_kernel_forward_from_hub` instead."
+        )
+        return _kernels_use_kernel_forward_from_hub(layer_name)
 
     # The default kernel mapping is built lazily (see `get_kernel_mapping_transformers`) so that simply
     # importing transformers (or `transformers.pipeline`) does not instantiate any `LayerRepository` /
@@ -556,14 +549,13 @@ def get_kernel_mapping_transformers() -> dict:
         return _KERNEL_MAPPING_CACHE
 
     def register_kernel_mapping_transformers(mapping=None):
+        _ensure_kernels_enabled()
         if mapping is None:
             mapping = get_kernel_mapping_transformers()
         register_kernel_mapping(mapping)
 
 else:
-    _kernels_enabled = False
-
-    # Stub to make decorators in transformers work when `kernels`
+    # Stub to make decorators int transformers work when `kernels`
     # is not installed.
     def use_kernel_forward_from_hub(*args, **kwargs):
         def decorator(cls):
@@ -624,16 +616,6 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1},
 }
 
-# Flash attention version -> major version of its hub kernel repo. Flash attention flavors that are not
-# listed here, and all other attention kernels, use `_DEFAULT_ATTN_KERNEL_VERSION`.
-_FLASH_ATTN_KERNEL_VERSION_MAPPING: dict[int, int] = {
-    # v3 is the first version shipping the Torch stable ABI (CUDA/ROCm) and Torch 2.13 builds (incl. XPU)
-    2: 3,
-    # FA4 is still in beta -> only v0 has been released
-    4: 0,
-}
-_DEFAULT_ATTN_KERNEL_VERSION = 1
-
 _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {}
 
 
@@ -647,10 +629,7 @@ def is_kernel(attn_implementation: str | None) -> bool:
 
 def get_attn_kernel_version(repo_id: str) -> int:
     """Return the major version of the hub kernel repo `repo_id` to load, e.g. `3` for `kernels-community/flash-attn2`."""
-    for flash_attn_version, kernel_version in _FLASH_ATTN_KERNEL_VERSION_MAPPING.items():
-        if is_flash_attention_requested(requested_attention_implementation=repo_id, version=flash_attn_version):
-            return kernel_version
-    return _DEFAULT_ATTN_KERNEL_VERSION
+    return FLASH_ATTN_KERNEL_VERSIONS.get(repo_id, 1)
 
 
 def load_and_register_attn_kernel(
@@ -692,9 +671,7 @@ def load_and_register_attn_kernel(
 
     # create revision xor version
     rev = rev.strip() if rev else None
-    version = None
-    if rev is None:
-        version = get_attn_kernel_version(repo_id)
+    version = get_attn_kernel_version(repo_id) if rev is None else None
 
     # Load the kernel from hub
     try:
@@ -736,7 +713,7 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _
         logger.warning_once(f"Kernel {kernel_name} not found in _HUB_KERNEL_MAPPING")
         mapping[kernel_name] = None
         return None
-    if is_kernels_available() and _kernels_enabled:
+    if is_kernels_available() and _kernels_enabled():
         try:
             repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"]
             revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None)
@@ -779,10 +756,11 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _
     return mapping[kernel_name]
 
 
-def kernelize(model: "PreTrainedModel", mode: "Mode | None" = None):
+def kernelize(model: "PreTrainedModel", mode: "Mode | None" = None, kernel_config: Optional["KernelConfig"] = None):
     """Temporarily register hidden kernel wrappers so `kernelize` can discover and replace them."""
     if not is_kernels_available():
         raise ImportError(_MISSING_KERNELS_MESSAGE)
+    _ensure_kernels_enabled()
 
     def get_device(device_type):
         if device_type == "cuda" and is_rocm_platform():
@@ -792,9 +770,10 @@ def get_device(device_type):
     mode = Mode.INFERENCE if not model.training else Mode.TRAINING if mode is None else mode
     device = get_device(model.device.type)
 
-    if model.kernel_config is not None:
-        inherit_mapping = not model.kernel_config.use_local_kernel
-        with use_kernel_mapping(model.kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
+    used_kernel_config = kernel_config or model.kernel_config
+    if used_kernel_config is not None:
+        inherit_mapping = not used_kernel_config.use_local_kernel and used_kernel_config.inherit_mapping
+        with use_kernel_mapping(used_kernel_config.kernel_mapping, inherit_mapping=inherit_mapping):
             _kernels_kernelize(model, device=device, mode=mode)
     else:
         _kernels_kernelize(model, device=device, mode=mode)

src/transformers/modeling_flash_attention_utils.py

@@ -70,6 +70,14 @@ def is_flash_attn_available():
     "flash_attention_4": "kernels-community/flash-attn4",
 }
 
+FLASH_ATTN_KERNEL_VERSIONS = {
+    "kernels-community/flash-attn2": 3,
+    "kernels-community/flash-attn3": 1,
+    "kernels-community/vllm-flash-attn3": 1,
+    "kernels-community/aiter-flash-attn": 1,
+    "kernels-community/flash-attn4": 0,
+    "kernels-community/metal-flash-sdpa": 1,
+}
 
 # Meta information on each mainline FA compatibility:
 #   1. The import structure and availability