# 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](https://github.com/huggingface/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`.

> [!IMPORTANT]
> 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](../weightconverter) para mais detalhes).

```python
# 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)
```

```python
# __init__.py
from . import layers
```

> [!NOTE]
> 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](/docs/transformers/main/pt/main_classes/kernels#transformers.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`.

```python
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.

```python
# 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)
```

```python
# __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.

```python
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`.

> [!TIP]
> A ordem dos pares na tupla de fusão determina a ordem dos argumentos passados para `KernelNameLayout.__init__`.