Carregando kernels

Um kernel funciona como uma substituição direta para as operações padrão do PyTorch. Ele substitui o método forward pela implementação otimizada do kernel sem quebrar o código do modelo.

Este guia mostra como carregar kernels para acelerar a inferência.

Instale o Transformers com a versão suportada do pacote kernels.

pip install -U "transformers[kernels]"

Defina use_kernels=True em from_pretrained() para carregar os kernels mais performáticos disponíveis no Hub para o seu dispositivo. Isso substitui as operações suportadas do PyTorch pela implementação do kernel.

from transformers import AutoModelForCausalLM

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

O mapeamento padrão depende do modelo, do dispositivo e do modo de execução. A tabela abaixo lista os repositórios comuns do Hub para os dispositivos mostrados. Quando nenhum kernel está registrado para uma combinação específica, a operação recorre ao PyTorch padrão.

A tabela não é exaustiva. Os modelos podem registrar camadas e funções adicionais, e alguns mapeamentos estão disponíveis apenas para modos de execução específicos.

GPUs da AMD relatam seu tipo de dispositivo como cuda no PyTorch. O Transformers detecta o ROCm em tempo de execução e direciona as operações suportadas para os kernels da AMD acima, incluindo builds do AITER, como kernels-community/aiter-rope. Você não precisa definir o tipo de dispositivo manualmente.

Navegue pelos kernels disponíveis na organização kernels-community.

Kernels de atenção

Carregue kernels de atenção do Hub com o argumento attn_implementation.

from transformers import AutoModelForCausalLM

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

Observe que, para kernels de atenção, qualquer coisa que não faça parte do repositório kernels-community (que é confiável – podemos adicionar mais repositórios confiáveis no futuro) exigirá o uso de um kwarg adicional allow_all_kernels=True (similar ao kwarg trust_remote_code=True para modelos não-HF). Isso ocorre porque o carregamento de um kernel pode levar à execução de código arbitrário na máquina host, e não podemos verificar todos os repositórios, portanto, você precisa permitir isso explicitamente.

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, aceitam vários formatos.

• @v2.1.0 fixa uma tag ou branch específica.
• @>=2.0,<3.0 define restrições de versionamento semântico.

from transformers import AutoModelForCausalLM

# pin to a specific version
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/[email protected]",
    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 do FlashAttention

Solicitar attn_implementation="flash_attention_2", "flash_attention_3" ou "flash_attention_4" recorre ao kernel correspondente do Hub quando o pacote flash-attn compilado não está instalado ou quando seu dispositivo não é CUDA.

O suporte ao FlashAttention-4 está em beta. As APIs e o comportamento podem mudar.

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

Consciência do modo

Os kernels se adaptam automaticamente aos modos de treinamento e inferência com base no estado model.training do 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()

Ative explicitamente os modos de treinamento e inferência com o argumento mode na função kernelize(). O modo de treinamento também suporta um modo adicional do 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 quais kernels são usados em um modelo.

As chaves de kernel_mapping são nomes registrados pelo modelo. Elas podem se referir a uma camada, como "RMSNorm", ou a uma função registrada, como "rotary_pos_emb". O separador : nomeia uma entrada de kernel específica dentro do repositório e a mapeia para aquela camada ou função.

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 implementações 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"
        }
    }
)

Metadados do kernel

Adicione um dicionário de metadados para controlar qual build do kernel é carregada.

from transformers import KernelConfig

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

Ramificações mais antigas de um repositório de kernel podem não ter builds para as suas versões de PyTorch e CUDA. Se um kernel falhar ao carregar, tente uma version mais recente antes de concluir que o seu hardware não é suportado.

Mapeamentos herdados

Um KernelConfig herda o mapeamento de kernels padrão do Transformers para kernels do Hub, e as entradas em kernel_mapping substituem o padrão para as camadas e funções correspondentes. Defina inherit_mapping=False para usar apenas as entradas em kernel_mapping. Tudo o que você deixar de fora volta ao PyTorch padrão. Isso é útil ao fazer benchmark de um kernel específico ou ao testar uma implementação personalizada sem aplicar os outros padrões.

A configuração abaixo mapeia a função rotary_pos_emb para um kernel RoPE e mantém todas as outras operações no PyTorch. O nome da função deve ser registrado pelo 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",
)

Fusão de módulos

Fusione módulos adjacentes em um único kernel passando uma tupla de pares (class_name, path_pattern) como chave em KernelConfig. Todos os padrões devem compartilhar o mesmo módulo pai. * corresponde a qualquer segmento único de caminho.

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

O Fusion exige que o repositório do kernel forneça uma classe KernelNameLayout complementar junto com a classe KernelName. Consulte o guia Writing kernels para saber como implementar uma.

Kernels locais

Carregue kernels a partir de caminhos de arquivo locais com use_local_kernel=True em KernelConfig. Isso carrega de um caminho do sistema de arquivos local em vez de um repositório do Hub.

Kernels locais usam /abs/path:layer_name em vez do formato do 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
)

Desativar kernels globalmente

Defina a variável de ambiente USE_HUB_KERNELS para desativar os kernels do Hub em todos os lugares sem alterar seu código.

export USE_HUB_KERNELS=0  # or OFF or NO

Solução de problemas

A integração do kernel depende do funcionamento conjunto do hardware, dos drivers e das versões dos pacotes. As seções a seguir abordam falhas comuns.

Problemas de instalação

Erros de importação indicam que o pacote de kernels está ausente ou que sua versão está fora da faixa suportada pelo Transformers. Reinstale por meio do extra para obter uma versão compatível.

pip install -U "transformers[kernels]"

Falhas no carregamento de kernels

Se kernels específicos falharem ao carregar, tente o seguinte.

• Verifique a compatibilidade do seu hardware com os requisitos do kernel.
• Verifique se seus drivers CUDA/ROCm/Metal estão atualizados.
• Consulte a documentação do repositório do kernel para problemas conhecidos.

Compatibilidade de dispositivos

Nem todos os kernels suportam todos os dispositivos. A biblioteca recorre às operações padrão do PyTorch quando um kernel não está disponível para o seu hardware. Consulte a documentação do repositório de kernels para verificar o suporte específico a cada dispositivo.

Recursos

• repositório Kernels
• post no blog Reforce seus modelos em 5 minutos com o Hugging Face Kernel Hub
• Descubra kernels na organização kernels-community