huggingface--transformers/tests/kernels/test_kernels.py

# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Run the test: CUDA_VISIBLE_DEVICES=0 RUN_SLOW=1 pytest -sv tests/kernels/test_kernels.py


import copy
import os
import tempfile
import types
from unittest.mock import MagicMock, patch

import torch
from huggingface_hub import snapshot_download

from transformers import AutoModelForCausalLM, AutoTokenizer, KernelConfig
from transformers.integrations.hub_kernels import (
    _HUB_KERNEL_MAPPING,
    _KERNEL_MODULE_MAPPING,
    is_kernel,
    lazy_load_kernel,
    load_and_register_attn_kernel,
)
from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.monkey_patching import clear_patch_mapping, get_patch_mapping, register_patch_mapping
from transformers.testing_utils import (
    TestCasePlus,
    cleanup,
    require_kernels,
    require_rocm,
    require_torch_accelerator,
    slow,
    torch_device,
)
from transformers.utils.import_utils import is_kernels_available
from transformers.utils.kernel_config import add_to_mapping_local


if is_kernels_available():
    from kernels import Device, LocalLayerRepository, Mode, kernelize

    import transformers.integrations.hub_kernels as hub_kernels_pkg


@require_kernels
@slow
class TestHubKernels(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        cls.model_id = "unsloth/Llama-3.2-1B-Instruct"
        cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_id)
        cls.model_kernelized = AutoModelForCausalLM.from_pretrained(
            cls.model_id, use_kernels=True, device_map=torch_device
        )
        cls.model_not_kernelized = AutoModelForCausalLM.from_pretrained(
            cls.model_id, use_kernels=False, device_map=torch_device
        )
        cls.input = "Hello"

    @classmethod
    def tearDownClass(cls):
        for attr in [
            "model_kernelized",
            "model_not_kernelized",
            "tokenizer",
        ]:
            if hasattr(cls, attr):
                try:
                    delattr(cls, attr)
                except Exception as e:
                    print(f"Could not delete attribute {attr}: {e}")

        # Clear any temporary kernel module cache entries populated by tests
        try:
            keys_to_remove = [
                k for k, v in list(_KERNEL_MODULE_MAPPING.items()) if v is None or isinstance(v, types.ModuleType)
            ]
            for k in keys_to_remove:
                _KERNEL_MODULE_MAPPING.pop(k, None)
        except Exception as e:
            print(f"Could not clear kernel module cache: {e}")

    def setUp(self):
        self._pre_test_patch_mapping = get_patch_mapping()

    def tearDown(self):
        # Restore monkey patch state to avoid leaking kernel patches across tests.
        clear_patch_mapping()
        if self._pre_test_patch_mapping:
            register_patch_mapping(self._pre_test_patch_mapping)
        # Free accelerator memory/cache and trigger GC
        cleanup(torch_device, gc_collect=True)

    @require_torch_accelerator
    def test_forward(self):
        tokenized_input = self.tokenizer(self.input, return_tensors="pt").input_ids.to(self.model_kernelized.device)
        output_ = self.model_kernelized.generate(tokenized_input, max_new_tokens=10, do_sample=False)
        output = self.tokenizer.decode(output_[0], skip_special_tokens=True)

        self.EXPECTED_OUTPUT = set()
        self.EXPECTED_OUTPUT.add("Hello, I'm looking for a reliable and trustworthy online")
        self.EXPECTED_OUTPUT.add("Hello! I'm excited to be a part of this")

        self.assertTrue(output in self.EXPECTED_OUTPUT)

    @require_rocm
    def test_rocm_rotary_kernel_forward_matches_baseline(self):
        """
        Regression test for the ROCm `rotary_pos_emb` function kernel (`kernels-community/aiter-rope`).

        On ROCm, `use_kernels=True` dispatches `apply_rotary_pos_emb` to the `aiter-rope` shim. A stale shim
        (e.g. the dropped `position_ids` signature mismatch fixed in
        https://github.com/huggingface/transformers/pull/46810) only blows up at runtime here, so comparing the
        kernelized forward against the non-kernelized baseline catches such breakages.
        """
        tokenized_input = self.tokenizer(self.input, return_tensors="pt").input_ids.to(torch_device)
        with torch.no_grad():
            kernelized_out = self.model_kernelized(tokenized_input).logits
            baseline_out = self.model_not_kernelized(tokenized_input).logits

        torch.testing.assert_close(baseline_out, kernelized_out, atol=1e-3, rtol=1e-3)

    def test_getter_use_kernels(self):
        self.assertTrue(self.model_kernelized.use_kernels)
        self.assertFalse(self.model_not_kernelized.use_kernels)

    def assert_kernelized_forward_is_different(self, kernelized_model, not_kernelized_model):
        """
        Iterate over modules and check if the forward method is different between
        the kernelized and not kernelized models. Break on first difference, else continue.
        Finally, assert that at least one forward is different.
        """
        found_difference = False
        for (name1, module1), (name2, module2) in zip(
            kernelized_model.named_modules(), not_kernelized_model.named_modules()
        ):
            # Only compare modules with the same name
            if name1 != name2:
                continue
            # Check if both modules have a 'forward' attribute
            if hasattr(module1, "forward") and hasattr(module2, "forward"):
                # Compare the code objects of the forward methods
                code1 = getattr(module1.forward, "__code__", None)
                code2 = getattr(module2.forward, "__code__", None)
                if code1 is not None and code2 is not None:
                    if code1 is not code2:
                        found_difference = True
                        break
        self.assertTrue(
            found_difference,
            "No module's forward method was different between kernelized and not kernelized models.",
        )

    def assert_kernelized_forward_is_the_same(self, model_1, model_2):
        """
        Iterate over modules and check if the forward method is the same between
        the kernelized and not kernelized models. Break on first difference, else continue.
        Finally, assert that at least one forward is the same.
        """
        no_difference = True
        for (name1, module1), (name2, module2) in zip(model_1.named_modules(), model_2.named_modules()):
            # Only compare modules with the same name
            if name1 != name2:
                continue
            # Check if both modules have a 'forward' attribute
            if hasattr(module1, "forward") and hasattr(module2, "forward"):
                # Compare the code objects of the forward methods
                code1 = getattr(module1.forward, "__code__", None)
                code2 = getattr(module2.forward, "__code__", None)
                if code1 is not None and code2 is not None:
                    if code1 != code2:
                        no_difference = False
                        break
        self.assertTrue(
            no_difference,
            "All module's forward methods were the same between the two models",
        )

    def test_kernelize(self):
        model = copy.deepcopy(self.model_not_kernelized)
        kernelize(model, mode=Mode.INFERENCE, device=Device(type=model.device.type))  # type: ignore[arg-type]
        self.assert_kernelized_forward_is_different(model, self.model_not_kernelized)
        self.assert_kernelized_forward_is_the_same(model, self.model_kernelized)
        del model

    def test_setter_use_kernels(self):
        model = copy.deepcopy(self.model_not_kernelized)
        model.use_kernels = True
        self.assertTrue(model.use_kernels)
        self.assert_kernelized_forward_is_different(model, self.model_not_kernelized)
        self.assert_kernelized_forward_is_the_same(model, self.model_kernelized)
        del model

    def test_unkernelize(self):
        model = copy.deepcopy(self.model_kernelized)

        with self.assertLogs("transformers.modeling_utils", level="WARNING") as cm:
            model.use_kernels = False

        self.assertTrue(
            any(
                "Disabling kernels at runtime is a no-op as there is no 'unkernelize' routine; keeping current kernels active."
                in msg
                for msg in cm.output
            )
        )

        self.assertFalse(model.use_kernels)
        del model

    def test_kernels_mapping(self):
        kernel_config = KernelConfig(kernel_mapping={"RMSNorm": "kernels-community/layer-norm:LlamaRMSNorm"})
        model = AutoModelForCausalLM.from_pretrained(
            "unsloth/Llama-3.2-1B-Instruct", use_kernels=True, device_map=torch_device, kernel_config=kernel_config
        )

        EXPECTED_OUTPUT = set()
        EXPECTED_OUTPUT.add("Hello, I'm looking for a reliable and trustworthy online")

        tokenized_input = self.tokenizer(self.input, return_tensors="pt").input_ids.to(model.device)
        output = model.generate(tokenized_input, max_new_tokens=10, do_sample=False)
        output = self.tokenizer.decode(output[0], skip_special_tokens=True)
        self.assertTrue(output in EXPECTED_OUTPUT)

        del model

    def test_kernels_mapping_explicit_version(self):
        kernel_config = KernelConfig(
            kernel_mapping={"RMSNorm": ("kernels-community/layer-norm:LlamaRMSNorm", {"version": 1})}
        )
        model = AutoModelForCausalLM.from_pretrained(
            "unsloth/Llama-3.2-1B-Instruct", use_kernels=True, device_map=torch_device, kernel_config=kernel_config
        )

        EXPECTED_OUTPUT = set()
        EXPECTED_OUTPUT.add("Hello, I'm looking for a reliable and trustworthy online")

        tokenized_input = self.tokenizer(self.input, return_tensors="pt").input_ids.to(model.device)
        output = model.generate(tokenized_input, max_new_tokens=10, do_sample=False)
        output = self.tokenizer.decode(output[0], skip_special_tokens=True)
        self.assertTrue(output in EXPECTED_OUTPUT)

        del model

    @require_torch_accelerator
    def test_kernel_fusion(self):
        model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random"
        kernel_config = KernelConfig(
            {
                (
                    ("RMSNorm", "model.layers.*.post_attention_layernorm"),
                    ("MLP", "model.layers.*.mlp"),
                ): (
                    "AntonV/dummy-rmsnorm-mlp-with-transformations-and-init:RMSNormMLP",
                    {"revision": "d582b66bea8e567dd06e683eca611648cfe53a7b", "trust_remote_code": True},
                ),
            }
        )

        tokenizer = AutoTokenizer.from_pretrained(model_id)
        inputs = tokenizer("Hello, how are you?", return_tensors="pt")

        baseline = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, device_map=torch_device)
        baseline.eval()
        inputs = {k: v.to(torch_device) for k, v in inputs.items()}
        with torch.no_grad():
            baseline_out = baseline(**inputs).logits
        del baseline

        fused = AutoModelForCausalLM.from_pretrained(
            model_id, use_kernels=True, kernel_config=kernel_config, device_map=torch_device
        )
        fused.eval()
        with torch.no_grad():
            fused_out = fused(**inputs).logits

        torch.testing.assert_close(baseline_out, fused_out, atol=1e-4, rtol=1e-4)

        decoder_layers = [
            (name, m)
            for name, m in fused.named_modules()
            if hasattr(m, "post_attention_layernorm") and hasattr(m, "mlp")
        ]
        self.assertTrue(len(decoder_layers) > 0, "No decoder layers found")
        for name, layer in decoder_layers:
            self.assertIsInstance(
                layer.mlp,
                torch.nn.Identity,
                f"{name}.mlp should be nn.Identity after fusion",
            )
            self.assertTrue(
                hasattr(layer.post_attention_layernorm, "kernel_layer_name")
                or hasattr(type(layer.post_attention_layernorm), "kernel_layer_name"),
                f"{name}.post_attention_layernorm should carry kernel_layer_name after fusion",
            )

        del fused

    @require_torch_accelerator
    def test_kernel_replacement_with_layout(self):
        model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random"
        kernel_config = KernelConfig(
            {
                "RMSNorm": (
                    "AntonV/dummy-rmsnorm-kernel-with-init:CustomRMSNorm",
                    {"revision": "e5dcd4fe743b81fa2b065964ef9a108107496c4f", "trust_remote_code": True},
                )
            }
        )

        tokenizer = AutoTokenizer.from_pretrained(model_id)
        inputs = tokenizer("Hello, how are you?", return_tensors="pt")

        baseline = AutoModelForCausalLM.from_pretrained(model_id, use_kernels=True, device_map=torch_device)
        baseline.eval()
        inputs = {k: v.to(torch_device) for k, v in inputs.items()}
        original_rmsnorm_cls = type(next(m for m in baseline.modules() if "RMSNorm" in type(m).__name__))
        with torch.no_grad():
            baseline_out = baseline(**inputs).logits
        del baseline

        model = AutoModelForCausalLM.from_pretrained(
            model_id, use_kernels=True, kernel_config=kernel_config, device_map=torch_device
        )
        model.eval()
        with torch.no_grad():
            model_out = model(**inputs).logits

        torch.testing.assert_close(baseline_out, model_out, atol=1e-4, rtol=1e-4)

        replaced = [m for m in model.modules() if hasattr(type(m), "kernel_layer_name")]
        self.assertTrue(len(replaced) > 0, "No replaced kernel layout modules found")
        for m in replaced:
            self.assertNotIsInstance(m, original_rmsnorm_cls)

        del model

    def test_faulty_fusion_incomplete_pattern(self):
        model_id = "michaelbenayoun/qwen3-tiny-4kv-heads-4layers-random"
        # "layers.*.post_attention_layernorm" is missing the leading "model." segment.
        # re.fullmatch("layers.\w+", "model.layers.0") returns None, so no module
        # is ever matched and the function raises ValueError.
        kernel_config = KernelConfig(
            {
                (
                    ("RMSNorm", "layers.*.post_attention_layernorm"),
                    ("MLP", "layers.*.mlp"),
                ): (
                    "AntonV/dummy-rmsnorm-mlp-with-transformations-and-init:RMSNormMLP",
                    {"revision": "d582b66bea8e567dd06e683eca611648cfe53a7b", "trust_remote_code": True},
                ),
            }
        )
        with self.assertRaises(ValueError):
            _ = AutoModelForCausalLM.from_pretrained(
                model_id,
                use_kernels=True,
                kernel_config=kernel_config,
                device_map=torch_device,
            )

    def test_faulty_kernel_mapping_layer_name(self):
        kernel_config = KernelConfig(kernel_mapping={"RMSNorm1": "kernels-community/layer-norm:LlamaRMSNorm"})
        with self.assertRaises(ValueError):
            _ = AutoModelForCausalLM.from_pretrained(
                "unsloth/Llama-3.2-1B-Instruct", use_kernels=True, device_map=torch_device, kernel_config=kernel_config
            )

    def test_faulty_kernel_mapping_type(self):
        kernel_config = KernelConfig(kernel_mapping={"RMSNorm": 1})
        with self.assertRaises(ValueError):
            _ = AutoModelForCausalLM.from_pretrained(
                "unsloth/Llama-3.2-1B-Instruct", use_kernels=True, device_map=torch_device, kernel_config=kernel_config
            )


@require_kernels
class TestKernelsEnv(TestCasePlus):
    def test_disable_hub_kernels(self):
        import importlib

        original_state = hub_kernels_pkg.__dict__.copy()

        try:
            with patch.dict(os.environ, {"USE_HUB_KERNELS": "OFF"}):
                importlib.reload(hub_kernels_pkg)
                self.assertFalse(hub_kernels_pkg._kernels_enabled)
        finally:
            hub_kernels_pkg.__dict__.clear()
            hub_kernels_pkg.__dict__.update(original_state)

    def test_enable_hub_kernels(self):
        import importlib

        original_state = hub_kernels_pkg.__dict__.copy()

        try:
            with patch.dict(os.environ, {"USE_HUB_KERNELS": "ON"}):
                importlib.reload(hub_kernels_pkg)
                self.assertTrue(hub_kernels_pkg._kernels_enabled)
        finally:
            hub_kernels_pkg.__dict__.clear()
            hub_kernels_pkg.__dict__.update(original_state)


@require_kernels
class TestKernelUtilities(TestCasePlus):
    def test_is_kernel_regex(self):
        valid = [
            "org/model",
            "org/model@main",
            "org/model:my_func",
            "org/model@v1.2.3:my_func",
            "flash|org/model@rev:fn",
        ]
        invalid = [
            "org//model",
            "org/model:too:many",
            "org/model@rev:fn:extra",
            "/org/model",
            "org:model",
        ]
        for s in valid:
            self.assertTrue(is_kernel(s.split("|")[-1]))
        for s in invalid:
            self.assertFalse(is_kernel(s))

    def test_lazy_load_kernel_success_and_cache(self):
        sentinel = types.ModuleType("sentinel_kernel_module")

        def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False):
            self.assertIn(repo_id, {"kernels-community/causal-conv1d"})
            self.assertFalse(allow_all_kernels)
            return sentinel

        patched_module_mapping = copy.copy(_KERNEL_MODULE_MAPPING)
        patched_module_mapping.pop("causal-conv1d", None)

        with patch.dict(
            lazy_load_kernel.__globals__,
            {
                "_KERNEL_MODULE_MAPPING": patched_module_mapping,
                "get_kernel": fake_get_kernel,
                "ALLOW_ALL_KERNELS": False,
            },
        ):
            mod1 = lazy_load_kernel("causal-conv1d", mapping=patched_module_mapping)
            self.assertIs(mod1, sentinel)

            mod2 = lazy_load_kernel("causal-conv1d", mapping=patched_module_mapping)
            self.assertIs(mod2, sentinel)

    def test_lazy_load_kernel_unknown(self):
        name = "unknown-kernel-name"
        _KERNEL_MODULE_MAPPING.pop(name, None)
        mod = lazy_load_kernel(name)
        self.assertIsNone(mod)
        self.assertIn(name, _KERNEL_MODULE_MAPPING)
        # Cleanup cache entry to avoid growth across tests
        _KERNEL_MODULE_MAPPING.pop(name, None)

    def test_lazy_load_kernel_version(self):
        name = "causal-conv1d"
        version_spec = ">=0.0.4,<0.1.0"

        sentinel_mod = types.ModuleType("sentinel_kernel_module")
        call_count = {"n": 0}

        def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False):
            call_count["n"] += 1
            self.assertEqual(repo_id, "kernels-community/causal-conv1d")
            self.assertIsNone(revision)
            self.assertEqual(version, version_spec)
            self.assertFalse(allow_all_kernels)
            return sentinel_mod

        patched_hub_mapping = copy.deepcopy(_HUB_KERNEL_MAPPING)
        patched_hub_mapping[name] = {
            "repo_id": "kernels-community/causal-conv1d",
            "version": version_spec,
        }

        patched_module_mapping = copy.copy(_KERNEL_MODULE_MAPPING)
        patched_module_mapping.pop(name, None)

        with patch.dict(
            lazy_load_kernel.__globals__,
            {
                "_HUB_KERNEL_MAPPING": patched_hub_mapping,
                "_KERNEL_MODULE_MAPPING": patched_module_mapping,
                "get_kernel": fake_get_kernel,
                "ALLOW_ALL_KERNELS": False,
            },
        ):
            mod1 = lazy_load_kernel(name, mapping=patched_module_mapping)
            mod2 = lazy_load_kernel(name, mapping=patched_module_mapping)

            self.assertIs(mod1, sentinel_mod)
            self.assertIs(mod2, sentinel_mod)
            self.assertEqual(call_count["n"], 1)


@require_kernels
class TestAttentionKernelRegistration(TestCasePlus):
    def test_trust_remote_code_for_attention_kernels(self):
        """
        Test that using an untrusted kernel (any repo outside `kernels-community`) as attention requires
        passing an expplicit `allow_all_kernels=True`
        """
        from transformers import LlamaConfig, LlamaModel

        config = LlamaConfig(num_hidden_layers=2, hidden_size=32, intermediate_size=64, vocab_size=100)
        model = LlamaModel(copy.deepcopy(config))
        untrusted_kernel = "untrusted/flash_attention_2"
        trusted_kernel = "kernels-community/flash-attn2"

        with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(tmpdirname)

            # Test that an untrusted kernel will raise an error without the flag
            with self.assertRaisesRegex(
                ValueError,
                "Kernel repository 'untrusted/flash_attention_2' could not verify publisher trust status. Set trust_remote_code=True to allow loading kernels from untrusted sources.",
            ):
                _ = LlamaModel.from_pretrained(tmpdirname, attn_implementation=untrusted_kernel)

            def dummy_lazy_import(*args, **kwargs):
                pass

            # Test that it works with the flag - though the repo does not exist, so patch the dispatch
            with patch("transformers.modeling_utils.lazy_import_flash_attention", dummy_lazy_import):
                model = LlamaModel.from_pretrained(
                    tmpdirname, attn_implementation=untrusted_kernel, allow_all_kernels=True
                )
                self.assertEqual(model.config._attn_implementation, untrusted_kernel)

            # Test that a trusted kernel does not need trust_remote_code
            model = LlamaModel.from_pretrained(tmpdirname, attn_implementation=trusted_kernel)
            self.assertEqual(model.config._attn_implementation, trusted_kernel)

    def test_load_and_register_flash_attn_like_kernel(self):
        kernel_obj = types.SimpleNamespace(flash_attn_varlen_func=lambda *a, **k: None)

        with (
            patch("transformers.integrations.hub_kernels.get_kernel", return_value=kernel_obj),
            patch("transformers.modeling_flash_attention_utils.lazy_import_flash_attention", return_value=None),
        ):
            attn_impl = "org/model"
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_MASK_ATTENTION_FUNCTIONS`: {e}")

    def test_load_and_register_named_function_kernel(self):
        def my_attention(*args, **kwargs):
            return None

        kernel_obj = types.SimpleNamespace(my_func=my_attention)
        with patch("transformers.integrations.hub_kernels.get_kernel", return_value=kernel_obj):
            attn_impl = "org/model:my_func"
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_MASK_ATTENTION_FUNCTIONS`: {e}")

    def test_add_to_mapping_local(self):
        repo_path = "/abs/path/kernel"
        compatible_mapping = {}
        add_to_mapping_local("RMSNorm", "cuda", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping)

        repo = compatible_mapping["RMSNorm"]["cuda"][Mode.INFERENCE]
        self.assertIsInstance(repo, LocalLayerRepository)
        self.assertEqual(repo.layer_name, "LlamaRMSNorm")

        with self.assertRaisesRegex(ValueError, "Only cuda, rocm, xpu, npu, neuron and tpu devices supported"):
            add_to_mapping_local("RMSNorm", "cpu", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping)

    @slow
    @require_torch_accelerator
    def test_add_to_mapping_local_then_load(self):
        repo_path = snapshot_download("kernels-community/layer-norm")
        compatible_mapping = {}
        add_to_mapping_local("RMSNorm", "cuda", f"{repo_path}:LlamaRMSNorm", Mode.INFERENCE, compatible_mapping)

        repo = compatible_mapping["RMSNorm"]["cuda"][Mode.INFERENCE]
        self.assertIsInstance(repo, LocalLayerRepository)
        self.assertEqual(repo.layer_name, "LlamaRMSNorm")

        layer_cls = repo.load()
        self.assertTrue(issubclass(layer_cls, torch.nn.Module))


@require_kernels
class TestUseKernelsLifecycle(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        cls.model_id = "unsloth/Llama-3.2-1B-Instruct"
        cls.model = AutoModelForCausalLM.from_pretrained(cls.model_id, use_kernels=False, device_map=torch_device)

    @classmethod
    def tearDownClass(cls):
        # Delete large objects to drop references early
        if hasattr(cls, "model"):
            try:
                del cls.model
            except Exception as e:
                print(f"Could not delete model: {e}")

    def tearDown(self):
        # Free accelerator memory/cache and trigger GC
        cleanup(torch_device, gc_collect=True)

    def test_setting_use_kernels_twice_does_not_rekernelize(self):
        with (
            patch.object(hub_kernels_pkg, "register_kernel_mapping_transformers") as mock_register,
            patch.object(hub_kernels_pkg, "_kernels_kernelize") as mock_kernelize,
        ):
            self.model.use_kernels = True

            self.assertTrue(self.model.use_kernels)
            # Check that both registraton and the underlying kernelize call happened
            mock_register.assert_called_once_with()
            self.assertEqual(mock_kernelize.call_count, 1)

            self.model.use_kernels = True

            mock_register.assert_called_once_with()
            self.assertEqual(mock_kernelize.call_count, 1)

    def test_train_eval_calls_kernelize_with_correct_mode(self):
        last_modes = []

        def spy_kernelize(model, device=None, mode=None):
            last_modes.append(mode)

        with patch.object(hub_kernels_pkg, "_kernels_kernelize", side_effect=spy_kernelize):
            self.model.use_kernels = True
            self.model.train(True)
            self.assertTrue(any(m == Mode.TRAINING for m in last_modes))
            self.model.eval()
            self.assertTrue(any(m == Mode.INFERENCE for m in last_modes))


@require_kernels
class TestKernelMappingDeviceFiltering(TestCasePlus):
    """Test that kernel mappings correctly filter by current device."""

    def test_multi_device_mapping_filters_correctly(self):
        """
        Test that when a kernel_mapping contains multiple devices (cuda, rocm),
        only the current device's kernel is registered.
        Regression test for issue where ROCm overwrote CUDA mapping.
        """
        kernel_mapping = {
            "RMSNorm": {
                "cuda": "kernels-community/layer-norm:LlamaRMSNorm",
                "rocm": "kernels-community/layer-norm:LlamaRMSNorm",
            }
        }

        kernel_config = KernelConfig(kernel_mapping)

        # Create a mock model on CUDA device
        mock_model = MagicMock()
        mock_model.training = False

        # Mock parameter with CUDA device
        mock_param = MagicMock()
        mock_param.device.type = "cuda"
        mock_model.parameters.return_value = iter([mock_param])

        # Mock named_modules with RMSNorm layer
        mock_layer = MagicMock()
        mock_layer.kernel_layer_name = "RMSNorm"
        mock_model.named_modules.return_value = [("layers.0", mock_layer)]

        # Trigger the mapping creation
        kernel_config.create_compatible_mapping(mock_model)

        # Verify results
        result_mapping = kernel_config.kernel_mapping

        self.assertIn("RMSNorm", result_mapping, "RMSNorm should be in mapping")
        backends = list(result_mapping["RMSNorm"].keys())

        # Assert only CUDA is present, not ROCm
        self.assertIn("cuda", backends, "CUDA backend should be registered")
        self.assertNotIn("rocm", backends, "ROCm backend should NOT be registered on CUDA device")

    def test_single_device_mapping_still_works(self):
        """
        Test that single-device mappings continue to work as expected.
        """
        kernel_mapping = {"RMSNorm": "kernels-community/layer-norm:LlamaRMSNorm"}

        kernel_config = KernelConfig(kernel_mapping)

        # Create a mock model
        mock_model = MagicMock()
        mock_model.training = False

        mock_param = MagicMock()
        mock_param.device.type = "cuda"
        mock_model.parameters.return_value = iter([mock_param])

        mock_layer = MagicMock()
        mock_layer.kernel_layer_name = "RMSNorm"
        mock_model.named_modules.return_value = [("layers.0", mock_layer)]
        kernel_config.create_compatible_mapping(mock_model)

        result_mapping = kernel_config.kernel_mapping
        self.assertIn("RMSNorm", result_mapping, "RMSNorm should be in mapping")