Integração e uso do Kernels Hub

A biblioteca kernels permite que kernels de computação otimizados sejam carregados diretamente do Hub.
Você pode encontrar kernels em organizações dedicadas ou pesquisando pela kernel tag no Hub.

Kernels são trechos de código otimizados que auxiliam no desenvolvimento, treinamento e inferência de modelos. Aqui, focaremos na sua integração com o TRL, mas consulte os recursos acima para saber mais sobre eles.

Instalação

Para usar kernels com o TRL, você precisará instalar a biblioteca no seu ambiente Python:

pip install kernels

Usando Kernels do Hub no TRL

Os kernels podem substituir diretamente as implementações de atenção, eliminando a necessidade de compilar manualmente backends de atenção como Flash Attention e aumentando a velocidade de treinamento apenas ao obter o kernel de atenção respectivo do Hub.

Você pode especificar um kernel ao carregar um modelo:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2"  # other options: kernels-community/vllm-flash-attn3, kernels-community/paged-attention
)

Ou ao executar um script de treinamento do TRL:

python sft.py ... --attn_implementation kernels-community/flash-attn2

Ou usando o TRL CLI:

trl sft ... --attn_implementation kernels-community/flash-attn2

Agora você pode aproveitar backends de atenção mais rápidos com um kernel pré-otimizado para a configuração de hardware do seu Hub, acelerando tanto o desenvolvimento quanto o treinamento.

Escolhendo e fixando uma versão de kernel

Os repositórios de kernels no Hub são versionados como branches (v1, v2, v3, …), e o Transformers seleciona uma versão padrão para cada repositório. Essa versão padrão pode mudar de uma release do Transformers para a próxima, de modo que o mesmo valor de attn_implementation nem sempre resolve para o mesmo build.

Para controlar qual build é carregado, acrescente uma revisão ao id do repositório, seja um branch de versão ou um SHA de commit:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "your-model-name",
    attn_implementation="kernels-community/flash-attn2@v2",  # or a commit SHA to pin an exact build
)

O pinning é útil para tornar uma execução de treinamento reproduzível e para permanecer em uma build conhecida e estável quando uma versão mais recente apresenta regressões. Tenha em mente que cada versão distribui builds apenas para uma faixa de versões do Torch e do CUDA, portanto, uma versão fixada pode não ter uma variante para o seu ambiente.

attn_implementation="flash_attention_2" recorre ao kernel do Hub kernels-community/flash-attn2 quando o pacote flash-attn não está instalado, portanto, você pode estar usando um kernel do Hub sem ter solicitado um explicitamente.

Os builds v3 para CUDA 12.8 atualmente falham na passagem reversa para modelos que usam grouped-query attention. Se o seu ambiente usa essa versão do CUDA, fixe @v2 até que huggingface/kernels-community#1085 seja fechado.

Comparando Implementações de Atenção

Avaliamos diversas implementações de atenção disponíveis em transformers, juntamente com diferentes backends de kernel, utilizando TRL e SFT.
Os experimentos foram executados em uma única H100 GPU com CUDA 12.9, aproveitando Qwen3-8B com batch size de 8, acumulação de gradiente de 1 e precisão bfloat16.
Lembre-se de que os resultados apresentados aqui são específicos para esta configuração e podem variar com diferentes configurações de treinamento.

A figura a seguir ilustra tanto a latência (tempo por etapa de treinamento) quanto a memória alocada de pico para as diferentes implementações de atenção e backends de kernel.
As implementações baseadas em kernel apresentam desempenho equivalente ao da atenção instalada de forma personalizada, e aumentar o max_length do modelo melhora ainda mais o desempenho. O consumo de memória é semelhante em todas as implementações, sem diferenças significativas. Obtemos o mesmo desempenho, porém com menos atrito, conforme descrito em a seção a seguir.

Flash Attention vs. Kernels do Hub

Compilar o Flash Attention a partir do código-fonte pode ser demorado, frequentemente levando de alguns minutos a horas, dependendo do seu hardware, da configuração de CUDA/PyTorch e da disponibilidade de wheels pré-compilados.

Em contraste, o Hugging Face Kernels oferece um fluxo de trabalho muito mais rápido e confiável. Os desenvolvedores não precisam se preocupar com configurações complexas — tudo é tratado automaticamente. Em nossos benchmarks, os kernels estavam prontos para uso em cerca de 2,5 segundos, sem necessidade de compilação. Isso permite que você comece o treinamento quase instantaneamente, acelerando significativamente o desenvolvimento. Basta especificar a versão desejada e o kernels cuida do resto.