Escrevendo kernels

Este guia explica como escrever kernels que vão além de uma substituição forward sem estado. Ele aborda duas capacidades suportadas pela API estendida KernelConfig:

1. Transformação de parâmetros: o kernel espera pesos em um layout diferente do modelo original (por exemplo, parâmetros renomeados ou mesclados).
2. Fusão de módulos: o kernel substitui múltiplos módulos adjacentes por uma única implementação fundida.

Para kernels básicos (substituições de forward sem estado, sem alterações de parâmetros), consulte a documentação da biblioteca kernels.

Padrão de duas classes

Qualquer kernel que carrega seus próprios parâmetros segue um padrão de duas classes.

• KernelName: contém apenas a passagem forward. A biblioteca kernels usa essa classe para kernelizar o modelo, pois não permite classes de kernel com estado.
• KernelNameLayout: um nn.Module que armazena os parâmetros e aplica monkey-patching no módulo original antes que o checkpoint seja carregado. Em tempo de execução, kernelize substitui seu forward pelo forward de KernelName. Você não precisa definir forward. O Transformers injeta um automaticamente com a mesma assinatura de KernelName.forward.

A convenção de nomenclatura é rigorosa. A classe de layout deve ser nomeada {KernelName}Layout e definida no mesmo módulo de KernelName.

Transformação de parâmetros

Use este padrão quando o kernel espera pesos com nomes diferentes ou em uma forma diferente da do checkpoint original do modelo.

A classe KernelNameLayout possui a mesma assinatura de __init__ que o módulo que substitui e declara um atributo de classe conversion_mapping que informa ao Transformers como remapear as chaves do checkpoint para os novos nomes de parâmetros (veja Carregamento dinâmico de pesos para mais detalhes).

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class CustomRMSNormLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the new parameter names

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps


class CustomRMSNorm(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.scale * hidden_states.to(input_dtype)

# __init__.py
from . import layers

A biblioteca kernels descobre as classes de kernel por meio de um módulo layers (seja layers.py ou layers/__init__.py). Defina as classes de kernel e de layout diretamente em layers.py e exponha o módulo layers a partir do __init__.py do pacote. O Transformers encontra CustomRMSNormLayout no mesmo módulo de CustomRMSNorm.

Proteja as importações específicas do Transformers (WeightConverter, Concatenate, WeightRenaming) em layers.py para que o kernel também possa ser importado em ambientes em que o Transformers não está instalado.

Carregue este kernel passando o repositório e o nome da classe para KernelConfig. A chave é o nome da classe do módulo original do modelo. O valor aponta para a classe KernelName (não para o Layout) no repositório.

O valor é uma string "owner/repo:ClassName" ou uma tupla (repo_string, metadata). No formato de tupla, passe version ou revision (pelo menos um) e, opcionalmente, trust_remote_code. Use version para uma versão principal do kernel ou revision para uma tag, branch ou commit. Defina trust_remote_code=True para carregar de fora da organização confiável kernels-community. Uma string simples tem como padrão version=1 e trust_remode_code=False.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig({"RMSNorm": ("owner/my-kernel:CustomRMSNorm", {"trust_remote_code": True, "version": 1})})
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Quando o modelo é carregado, o Transformers:

1. Carrega CustomRMSNorm do repositório e procura por CustomRMSNormLayout no mesmo módulo.
2. Aplica monkey-patch em cada RMSNorm do modelo com CustomRMSNormLayout.
3. Remapeia os pesos do checkpoint usando conversion_mapping para que sejam carregados nos novos nomes de parâmetros.
4. Chama kernelize, que substitui CustomRMSNormLayout.forward por CustomRMSNorm.forward.

Fusão de módulos

Use este padrão quando um kernel substitui vários módulos adjacentes por uma única implementação fundida. Como o módulo fundido combina parâmetros de vários módulos originais, o KernelNameLayout.__init__ recebe os módulos filhos instanciados, em vez dos argumentos de seus construtores.

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class RMSNormMLPLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the fused parameter names

    def __init__(self, norm, mlp):
        super().__init__()
        self.variance_epsilon = norm.variance_epsilon
        self.scale = nn.Parameter(torch.empty_like(norm.weight))
        self.gate_up_proj = nn.Linear(
            mlp.gate_proj.in_features,
            mlp.gate_proj.out_features + mlp.up_proj.out_features,
            bias=mlp.gate_proj.bias is not None,
            device=mlp.gate_proj.weight.device,
            dtype=mlp.gate_proj.weight.dtype,
        )
        self.down_proj = nn.Linear(
            mlp.down_proj.in_features,
            mlp.down_proj.out_features,
            bias=mlp.down_proj.bias is not None,
            device=mlp.down_proj.weight.device,
            dtype=mlp.down_proj.weight.dtype,
        )
        self.act_fn = mlp.act_fn


class RMSNormMLP(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.scale * hidden_states.to(input_dtype)
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.down_proj(self.act_fn(gate) * up)

# __init__.py
from . import layers

Para fundir módulos, passe uma tupla de pares (class_name, path_pattern) como chave em KernelConfig em vez de uma string simples. Todos os padrões devem compartilhar o mesmo módulo pai (o Transformers funde os filhos nesse pai). O curinga * 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/my-kernel:RMSNormMLP", {"trust_remote_code": True, "version": 1}),
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Quando o modelo é carregado, o Transformers:

1. Carrega RMSNormMLP do repositório e localiza RMSNormMLPLayout no mesmo módulo.
2. Corresponde a cada camada do decoder em model.layers.* e constrói uma classe pai fundida cujo __init__ chama RMSNormMLPLayout(post_attention_layernorm, mlp).
3. Substitui o filho restante (mlp) por nn.Identity() para preservar a interface do módulo pai.
4. Remapeia os pesos do checkpoint usando conversion_mapping.
5. Chama kernelize, que substitui RMSNormMLPLayout.forward por RMSNormMLP.forward.

A ordem dos pares na tupla de fusão determina a ordem dos argumentos passados para KernelNameLayout.__init__.