커널 작성

이 가이드는 상태 없는 forward 대체를 넘어선 커널 작성 방법을 설명합니다. 확장된 KernelConfig API가 지원하는 두 가지 기능을 다룹니다:

Parameter transformation: 커널은 원래 모델과 다른 레이아웃의 가중치를 기대합니다(예: 이름이 변경되었거나 병합된 파라미터).
Module fusion: 커널이 여러 인접한 모듈을 단일 융합된 구현으로 대체합니다.

기본 커널(파라미터 변경이 없는 무상태 forward 대체)의 경우 kernels 라이브러리 문서를 참조하세요.

두 클래스 패턴

kernel에 자체 파라미터를 포함하는 모든 kernel은 두 개의 클래스 패턴을 따릅니다.

KernelName: forward 패스만 포함합니다. kernels 라이브러리는 상태가 있는(stateful) 커널 클래스를 허용하지 않으므로 이 클래스를 사용하여 모델을 커널화합니다.
KernelNameLayout: 체크포인트가 로드되기 전에 매개변수를 보유하고 원본 모듈을 monkey-patch하는 nn.Module입니다. 런타임에서 kernelize는 forward를 KernelName의 forward로 대체합니다. forward를 정의할 필요가 없습니다. Transformers가 KernelName.forward와 동일한 시그니처를 가진 forward를 자동으로 주입합니다.

네이밍 규칙은 엄격합니다. 레이아웃 클래스는 {KernelName}Layout으로 명명되어야 하며, KernelName과 동일한 모듈에서 정의되어야 합니다.

파라미터 변환

kernel이 원래 모델 체크포인트와 다른 이름이나 다른 형식의 가중치를 기대하는 경우 이 패턴을 사용하세요.

KernelNameLayout 클래스는 대체하는 모듈과 동일한 __init__ 시그니처를 가지며, Transformers가 체크포인트 키를 새로운 매개변수 이름으로 어떻게 재매핑할지 알려주는 conversion_mapping 클래스 속성을 선언합니다(자세한 내용은 동적 가중치 로딩을 참조).

# 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

kernels 라이브러리는 layers 모듈(layers.py 또는 layers/__init__.py)을 통해 커널 클래스를 검색합니다. 커널 클래스와 레이아웃 클래스를 모두 layers.py에 직접 정의하고, 패키지 __init__.py에서 layers 모듈을 노출하세요. Transformers는 CustomRMSNorm과 동일한 모듈에서 CustomRMSNormLayout을 찾습니다.

layers.py에서 Transformers 전용 import(WeightConverter, Concatenate, WeightRenaming)를 가드하여, Transformers가 설치되지 않은 환경에서도 커널을 import할 수 있도록 합니다.

KernelConfig에 repo와 클래스 이름을 전달하여 이 커널을 로드합니다. 키는 모델의 원래 모듈 클래스 이름입니다. 값은 repo의 KernelName 클래스(Layout가 아님)를 가리킵니다.

값은 "owner/repo:ClassName" 문자열이거나 (repo_string, metadata) 튜플입니다. 튜플 형식에서는 version 또는 revision(최소 하나)을 전달하고, 선택적으로 trust_remote_code를 지정할 수 있습니다. 주요 커널 버전에는 version을, 태그, 브랜치 또는 커밋에는 revision을 사용하세요. 신뢰할 수 있는 kernels-community 조직 외부에서 로드하려면 trust_remote_code=True로 설정하세요. 문자열만 단독으로 사용할 경우 기본값은 version=1 및 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",
)

모델이 로드되면 Transformers는:

CustomRMSNorm을 저장소에서 로드하고 같은 모듈에서 CustomRMSNormLayout을 찾습니다.
모델의 모든 RMSNorm을 CustomRMSNormLayout으로 monkey-patch합니다.
conversion_mapping를 사용하여 체크포인트 가중치를 재매핑하여 새로운 파라미터 이름으로 로드되도록 합니다.
kernelize를 호출하여 CustomRMSNormLayout.forward를 CustomRMSNorm.forward로 대체합니다.

모듈 융합

커널이 여러 인접 모듈을 단일 융합 구현으로 대체하는 경우 이 패턴을 사용하세요. 융합 모듈은 여러 원래 모듈의 매개변수를 결합하므로, KernelNameLayout.__init__는 생성자 인수가 아닌 인스턴스화된 자식 모듈을 수신합니다.

# 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

모듈을 융합하려면, KernelConfig의 키로 일반 문자열 대신 (class_name, path_pattern) 쌍의 튜플을 전달합니다. 모든 패턴은 동일한 부모 모듈을 공유해야 합니다(Transformers는 해당 부모의 자식들을 융합합니다). * 와일드카드는 단일 경로 세그먼트를 임의로 일치시킵니다.

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

모델이 로드되면 Transformers는:

RMSNormMLP를 저장소에서 로드하고 같은 모듈에서 RMSNormMLPLayout을 찾습니다.
model.layers.*의 모든 decoder layer에 매칭하여, __init__에서 RMSNormMLPLayout(post_attention_layernorm, mlp)를 호출하는 fused parent class를 생성합니다.
남은 자식 모듈(mlp)을 nn.Identity()로 교체하여 부모 모듈의 인터페이스를 유지합니다.
conversion_mapping를 사용하여 체크포인트 가중치를 다시 매핑합니다.
kernelize를 호출하여 RMSNormMLPLayout.forward를 RMSNormMLP.forward로 대체합니다.

fusion 튜플에서 쌍의 순서는 KernelNameLayout.__init__에 전달되는 인자의 순서를 결정합니다.

Update on GitHub