Carga de kernels

Un kernel funciona como un reemplazo directo para las operaciones estándar de PyTorch. Sustituye el método forward por la implementación optimizada del kernel sin romper el código del modelo.

Esta guía muestra cómo cargar kernels para acelerar la inferencia.

Instale Transformers con la versión compatible del paquete kernels.

pip install -U "transformers[kernels]"

Establezca use_kernels=True en from_pretrained() para cargar los kernels más eficientes disponibles en el Hub para su dispositivo. Esto reemplaza las operaciones de PyTorch compatibles con la implementación del kernel.

from transformers import AutoModelForCausalLM

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

El mapeo predeterminado depende del modelo, el dispositivo y el modo de ejecución. La tabla siguiente enumera los repositorios comunes de Hub para los dispositivos mostrados. Cuando no hay ningún kernel registrado para una combinación en particular, la operación recurre a PyTorch estándar.

La tabla no es exhaustiva. Los modelos pueden registrar capas y funciones adicionales, y algunas correspondencias solo están disponibles para modos de ejecución específicos.

Las GPU de AMD informan su tipo de dispositivo como cuda en PyTorch. Transformers detecta ROCm en tiempo de ejecución y enruta las operaciones compatibles a los kernels de AMD anteriores, incluidas las compilaciones de AITER como kernels-community/aiter-rope. No es necesario configurar el tipo de dispositivo manualmente.

Explora los kernels disponibles en la organización kernels-community.

Kernels de atención

Cargue los núcleos de atención desde el Hub con el argumento attn_implementation.

from transformers import AutoModelForCausalLM

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

Tenga en cuenta que, en el caso de los núcleos de atención, todo lo que no forme parte del repositorio kernels-community (que es de confianza; es posible que añadamos más repositorios de confianza en el futuro) requerirá el uso de un argumento adicional allow_all_kernels=True (similar al argumento trust_remote_code=True para modelos no de HF). Esto se debe a que la carga de un núcleo puede provocar la ejecución de código arbitrario en la máquina anfitriona, y no podemos verificar cada repositorio, por lo que debe permitirlo explícitamente.

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"
)

Kernels específicos, como attention, aceptan varios formatos.

• @v2.1.0 fija una etiqueta o rama específica.
• @>=2.0,<3.0 establece restricciones de versionado semántico.

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"
)

Fallback de FlashAttention

Solicitar attn_implementation="flash_attention_2", "flash_attention_3" o "flash_attention_4" recurre al kernel correspondiente de Hub cuando el paquete flash-attn compilado no está instalado o su dispositivo no es CUDA.

El soporte para FlashAttention-4 está en beta. Las API y el comportamiento pueden cambiar.

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",
)

Conciencia del modo

Los kernels se adaptan automáticamente a los modos de entrenamiento e inferencia en función del estado 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()

Habilite explícitamente los modos de entrenamiento e inferencia con el argumento mode en la función kernelize(). El modo de entrenamiento también admite un modo adicional de torch.compile.

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 personaliza qué kernels se utilizan en un modelo.

Las claves de kernel_mapping son nombres registrados por el modelo. Pueden referirse a una capa, como "RMSNorm", o a una función registrada, como "rotary_pos_emb". El separador : nombra una entrada de kernel específica dentro del repositorio y la asigna a esa capa o función.

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"
)

Especifique diferentes implementaciones de kernel para cada tipo de dispositivo.

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"
        }
    }
)

Metadatos del kernel

Añade un diccionario de metadatos para controlar qué compilación del kernel se carga.

from transformers import KernelConfig

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

Las ramas más antiguas de un repositorio de kernel pueden no tener compilaciones para tus versiones de PyTorch y CUDA. Si un kernel no se carga, prueba una version más reciente antes de concluir que tu hardware no es compatible.

Mapeos heredados

Un KernelConfig hereda el mapeo de kernels de Transformers por defecto para los kernels de Hub, y las entradas en kernel_mapping sobrescriben el valor por defecto para las capas y funciones correspondientes. Establece inherit_mapping=False para usar solo las entradas de kernel_mapping. Todo lo que omitas volverá a PyTorch estándar. Esto es útil al realizar pruebas de rendimiento de un kernel específico o al probar una implementación personalizada sin aplicar los demás valores por defecto.

La configuración siguiente asigna la función rotary_pos_emb a un kernel de RoPE y deja todas las demás operaciones en PyTorch. El nombre de la función debe estar registrado por el modelo.

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",
)

Fusión de módulos

Fusiona módulos adyacentes en un solo kernel pasando una tupla de pares (class_name, path_pattern) como clave en KernelConfig. Todos los patrones deben compartir el mismo módulo padre. * coincide con cualquier segmento de ruta individual.

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 requiere que el repositorio del kernel proporcione una clase KernelNameLayout de acompañamiento junto con la clase KernelName. Consulte la guía Writing kernels para saber cómo implementar una.

Kernels locales

Cargue kernels desde rutas de archivos locales con use_local_kernel=True en KernelConfig. Esto carga desde una ruta del sistema de archivos local en lugar de un repositorio de Hub.

Los kernels locales usan /abs/path:layer_name en lugar del formato del 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
)

Desactivar kernels globalmente

Establezca la variable de entorno USE_HUB_KERNELS para deshabilitar los kernels de Hub en todas partes sin cambiar su código.

export USE_HUB_KERNELS=0  # or OFF or NO

Solución de problemas

La integración del kernel depende de que el hardware, los controladores y las versiones de los paquetes funcionen de manera conjunta. Las siguientes secciones abordan los fallos más comunes.

Problemas de instalación

Los errores de importación indican que el paquete de kernels no está presente o que su versión está fuera del rango que Transformers admite. Vuelva a instalarlo a través del extra para obtener una versión compatible.

pip install -U "transformers[kernels]"

Fallos en la carga de kernels

Si los kernels específicos no se cargan, prueba lo siguiente.

• Comprueba la compatibilidad de tu hardware con los requisitos del kernel.
• Verifique que sus controladores de CUDA/ROCm/Metal estén actualizados.
• Consulte la documentación del repositorio del kernel para conocer los problemas conocidos.

Compatibilidad de dispositivos

No todos los kernels admiten todos los dispositivos. La biblioteca recurre a operaciones estándar de PyTorch si un kernel no está disponible para tu hardware. Consulta la documentación del repositorio de kernels para conocer el soporte específico de cada dispositivo.

Recursos

• repositorio Kernels
• artículo de blog Mejora tus modelos en 5 minutos con el Hugging Face Kernel Hub
• Descubre kernels en la organización kernels-community