Escribir kernels

Esta guía explica cómo escribir kernels que van más allá de una sustitución de forward sin estado. Aborda dos capacidades que la API extendida de KernelConfig admite:

1. Transformación de parámetros: el kernel espera los pesos en un formato diferente al del modelo original (por ejemplo, parámetros renombrados o fusionados).
2. Fusión de módulos: el kernel reemplaza varios módulos adyacentes con una única implementación fusionada.

Para los kernels básicos (sustituciones de forward sin estado y sin cambios de parámetros), consulte la documentación de la biblioteca kernels.

Patrón de dos clases

Cualquier kernel que lleva sus propios parámetros sigue un patrón de dos clases.

• KernelName: contiene solo el paso forward. La biblioteca kernels utiliza esta clase para kernelizar el modelo porque no permite clases de kernel con estado.
• KernelNameLayout: un nn.Module que mantiene los parámetros y aplica un monkey-patch al módulo original antes de que se cargue el checkpoint. En tiempo de ejecución, kernelize reemplaza su forward con el forward de KernelName. No es necesario que definas forward. Transformers inyecta uno automáticamente con la misma firma que KernelName.forward.

La convención de nomenclatura es estricta. La clase de distribución debe llamarse {KernelName}Layout y definirse en el mismo módulo que KernelName.

Transformación de parámetros

Use este patrón cuando el kernel espera los pesos con nombres diferentes o en una forma distinta a la del punto de control original del modelo.

La clase KernelNameLayout tiene la misma firma de __init__ que el módulo que reemplaza y declara un atributo de clase conversion_mapping que indica a Transformers cómo remapear las claves del checkpoint a los nuevos nombres de parámetros (consulte Carga dinámica de pesos para más detalles).

# 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

La biblioteca kernels descubre las clases de kernel a través de un módulo layers (ya sea layers.py o layers/__init__.py). Defina tanto las clases de kernel como las de layout directamente en layers.py y exponga el módulo layers desde el __init__.py del paquete. Transformers encuentra CustomRMSNormLayout en el mismo módulo que CustomRMSNorm.

Guarde las importaciones específicas de Transformers (WeightConverter, Concatenate, WeightRenaming) en layers.py para que el kernel también pueda importarse en entornos donde Transformers no está instalado.

Cargue este kernel pasando el nombre del repositorio y la clase a KernelConfig. La clave es el nombre de la clase del módulo original del modelo. El valor apunta a la clase KernelName (no a Layout) en el repositorio.

El valor es una cadena "owner/repo:ClassName" o una tupla (repo_string, metadata). En la forma de tupla, pase version o revision (al menos uno) y, opcionalmente, trust_remote_code. Use version para una versión principal del kernel o revision para una etiqueta, rama o commit. Establezca trust_remote_code=True para cargar desde fuera de la organización kernels-community de confianza. Una cadena simple tiene por defecto version=1 y 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",
)

Cuando el modelo se carga, Transformers:

1. Carga CustomRMSNorm desde el repositorio y busca CustomRMSNormLayout en el mismo módulo.
2. Aplica parches de forma dinámica a cada RMSNorm del modelo con CustomRMSNormLayout.
3. Reasigna los pesos del checkpoint usando conversion_mapping para que se carguen con los nuevos nombres de parámetros.
4. Llama a kernelize, que reemplaza CustomRMSNormLayout.forward con CustomRMSNorm.forward.

Fusión de módulos

Use este patrón cuando un kernel reemplace varios módulos adyacentes por una única implementación fusionada. Dado que el módulo fusionado combina parámetros de varios módulos originales, el KernelNameLayout.__init__ recibe los módulos hijos instanciados en lugar de sus argumentos de constructor.

# 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 fusionar módulos, pase una tupla de pares (class_name, path_pattern) como clave en KernelConfig en lugar de una cadena simple. Todos los patrones deben compartir el mismo módulo padre (Transformers fusiona los hijos de ese padre). El comodín * 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/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",
)

Cuando el modelo se carga, Transformers:

1. Carga RMSNormMLP del repositorio y encuentra RMSNormMLPLayout en el mismo módulo.
2. Coincide con cada capa decodificadora en model.layers.* y construye una clase padre fusionada cuyo __init__ llama a RMSNormMLPLayout(post_attention_layernorm, mlp).
3. Reemplaza el hijo restante (mlp) con nn.Identity() para preservar la interfaz del módulo padre.
4. Remapea los pesos del checkpoint usando conversion_mapping.
5. Llama a kernelize, que reemplaza RMSNormMLPLayout.forward con RMSNormMLP.forward.

El orden de las parejas en la tupla de fusión determina el orden de los argumentos pasados a KernelNameLayout.__init__.