Chargement des noyaux

Un noyau fonctionne comme un substitut direct des opérations standard de PyTorch. Il remplace la méthode forward par l'implémentation optimisée du noyau sans casser le code du modèle.

Ce guide montre comment charger des noyaux pour accélérer l'inférence.

Installez Transformers avec la version prise en charge du package kernels.

pip install -U "transformers[kernels]"

Définissez use_kernels=True dans from_pretrained() pour charger les noyaux les plus performants disponibles sur le Hub pour votre appareil. Cela remplace les opérations PyTorch prises en charge par l’implémentation du noyau.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

La correspondance par défaut dépend du modèle, de l'appareil et du mode d'exécution. Le tableau ci-dessous répertorie les dépôts Hub courants pour les appareils présentés. Lorsqu'aucun noyau n'est enregistré pour une combinaison particulière, l'opération revient au PyTorch standard.

Le tableau n'est pas exhaustif. Les modèles peuvent enregistrer des couches et des fonctions supplémentaires, et certaines correspondances ne sont disponibles que pour des modes d'exécution spécifiques.

Les GPU AMD signalent leur type de périphérique comme cuda dans PyTorch. Transformers détecte ROCm à l’exécution et oriente les opérations prises en charge vers les noyaux AMD ci-dessus, y compris les builds AITER tels que kernels-community/aiter-rope. Vous n’avez pas besoin de définir le type de périphérique vous-même.

Parcourez les noyaux disponibles dans l'organisation kernels-community.

Noyaux d'attention

Chargez les noyaux d'attention depuis le Hub avec l'argument attn_implementation.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    device_map="auto"
)

Notez que pour les noyaux d’attention, tout ce qui ne fait pas partie du dépôt kernels-community (qui est de confiance – nous pourrons ajouter d’autres dépôts de confiance à l’avenir) nécessitera l’utilisation d’un argument supplémentaire allow_all_kernels=True (de manière similaire à l’argument trust_remote_code=True pour les modèles non-HF). Cela est dû au fait que le chargement d’un noyau peut entraîner l’exécution de code arbitraire sur la machine hôte, et que nous ne pouvons pas vérifier chaque dépôt, il est donc nécessaire de l’autoriser explicitement.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="random-repo/random-attention",
    allow_all_kernels=True,
    device_map="auto"
)

Certains noyaux spécifiques, comme l'attention, acceptent plusieurs formats.

• @v2.1.0 est épinglé à une balise ou une branche spécifique.
• @>=2.0,<3.0 définit des contraintes de versionnage sémantique.

from transformers import AutoModelForCausalLM

# pin to a specific version
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@v2.1.0",
    device_map="auto"
)
# use semantic versioning constraints
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@>=2.0,<3.0",
    device_map="auto"
)

Repli sur FlashAttention

Demander attn_implementation="flash_attention_2", "flash_attention_3" ou "flash_attention_4" bascule sur le noyau Hub correspondant si le package flash-attn compilé n’est pas installé ou si votre appareil n’est pas CUDA.

La prise en charge de FlashAttention-4 est en version bêta. Les API et le comportement peuvent changer.

from transformers import AutoModelForCausalLM

# uses the compiled flash-attn package if present, otherwise the kernels-community/flash-attn2 Hub kernel
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="flash_attention_2",
    device_map="auto",
)

Prise en compte du mode

Les noyaux s'adaptent automatiquement aux modes entraînement et inférence en fonction de l'état model.training de PyTorch.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

# Switch to inference mode - uses inference-optimized kernels
model.eval()
with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=50)

# Switch to training mode - uses training-optimized kernels with gradient support
model.train()
loss = model(input_ids, labels=labels).loss
loss.backward()

Activez explicitement les modes d’entraînement et d’inférence à l’aide de l’argument mode dans la fonction kernelize(). Le mode d’entraînement prend également en charge un mode torch.compile supplémentaire.

from kernels import Mode
from transformers import kernelize

# inference optimized kernels
kernelize(model, mode=Mode.INFERENCE)

# training optimized kernels
kernelize(model, mode=Mode.TRAINING)

# training and torch-compile friendly kernels
kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)

KernelConfig

KernelConfig personnalise les noyaux utilisés dans un modèle.

Les clés de kernel_mapping sont des noms enregistrés par le modèle. Elles peuvent faire référence à une couche, comme "RMSNorm", ou à une fonction enregistrée, comme "rotary_pos_emb". Le séparateur : désigne une entrée de noyau spécifique dans le dépôt et l’associe à cette couche ou à cette fonction.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": "kernels-community/liger-kernels:LigerRMSNorm",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto"
)

Spécifiez différentes implémentations de noyau pour chaque type d’appareil.

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": {
            "cuda": "kernels-community/liger-kernels:LigerRMSNorm",
            "rocm": "kernels-community/liger-kernels:LigerRMSNorm",
            "xpu": "kernels-community/rmsnorm:RMSNorm"
        }
    }
)

Métadonnées du noyau

Ajouter un dictionnaire de métadonnées pour contrôler quelle compilation de noyau est chargée.

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": ("kernels-community/liger-kernels:LigerRMSNorm", {"version": 3}),
    }
)

Les anciennes branches d’un dépôt de noyaux peuvent ne pas disposer de compilations pour vos versions de PyTorch et de CUDA. Si un noyau échoue au chargement, essayez une version plus récente avant de conclure que votre matériel n’est pas pris en charge.

Mappages hérités

Un KernelConfig hérite de la correspondance de kernels Transformers par défaut pour les kernels du Hub, et les entrées de kernel_mapping remplacent la valeur par défaut pour les couches et fonctions correspondantes. Définissez inherit_mapping=False pour n’utiliser que les entrées de kernel_mapping. Tout ce que vous omettez retombe sur PyTorch standard. Cela est utile lors du benchmark d’un kernel spécifique ou du test d’une implémentation personnalisée sans appliquer les autres valeurs par défaut.

La configuration ci-dessous associe la fonction rotary_pos_emb à un noyau RoPE et laisse toutes les autres opérations sur PyTorch. Le nom de la fonction doit être enregistré par le modèle.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "rotary_pos_emb": (
            "kernels-community/rotary:apply_rotary_transformers",
            {"version": 2},
        ),
    },
    inherit_mapping=False,
)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Fusion de modules

Fusionnez les modules adjacents en un seul noyau en passant un tuple de paires (class_name, path_pattern) comme clé dans KernelConfig. Tous les motifs doivent partager le même module parent. * correspond à n'importe quel segment de chemin unique.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    {
        (
            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
            ("MLP",     "model.layers.*.mlp"),
        ): "owner/fused-rmsnorm-mlp:RMSNormMLP",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Fusion exige que le dépôt de noyaux fournisse une classe KernelNameLayout accompagnant la classe KernelName. Consultez le guide Writing kernels pour savoir comment en implémenter une.

Noirs locaux

Chargez les noyaux à partir de chemins de fichiers locaux avec use_local_kernel=True dans KernelConfig. Cela charge depuis un chemin du système de fichiers local au lieu d'un dépôt Hub.

Les noyaux locaux utilisent /abs/path:layer_name au lieu du format Hub org/repo:layer_name.

from transformers import KernelConfig, AutoModelForCausalLM

kernel_mapping = {
    "RMSNorm": "/path/to/liger-kernels:LigerRMSNorm",
}
kernel_config = KernelConfig(kernel_mapping, use_local_kernel=True)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype="auto",
    device_map="auto",
    use_kernels=True,
    kernel_config=kernel_config
)

Désactiver les noyaux globalement

Définissez la variable d'environnement USE_HUB_KERNELS pour désactiver les noyaux du Hub partout sans modifier votre code.

export USE_HUB_KERNELS=0  # or OFF or NO

Dépannage

L'intégration du noyau dépend de la compatibilité entre le matériel, les pilotes et les versions des paquets. Les sections suivantes traitent des pannes courantes.

Problèmes d'installation

Les erreurs d’importation indiquent que le package kernels est manquant ou que sa version est hors de la plage prise en charge par Transformers. Réinstallez-le via l’extra pour obtenir une version compatible.

pip install -U "transformers[kernels]"

Échecs de chargement des noyaux

Si des noyaux spécifiques échouent au chargement, essayez ce qui suit.

• Vérifiez la compatibilité de votre matériel avec les exigences du noyau.
• Vérifiez que vos pilotes CUDA/ROCm/Metal sont à jour.
• Consultez la documentation du dépôt de noyaux pour les problèmes connus.

Compatibilité des appareils

Tous les noyaux ne prennent pas en charge tous les appareils. La bibliothèque bascule sur les opérations standard de PyTorch si un noyau n'est pas disponible pour votre matériel. Consultez la documentation du dépôt de noyaux pour la prise en charge spécifique à chaque appareil.

Ressources

• répertoire Kernels
• l'article de blog Améliorez vos modèles en 5 minutes avec le Hugging Face Kernel Hub
• Découvrez les noyaux dans l'organisation kernels-community