Загрузка ядер

Ядро работает как замена стандартных операций PyTorch. Оно подменяет метод forward оптимизированной реализацией ядра, не нарушая код модели.

В этом руководстве показано, как загружать ядра для ускорения вывода.

Установите Transformers с поддерживаемой версией пакета kernels.

pip install -U "transformers[kernels]"

Установите use_kernels=True в from_pretrained(), чтобы загрузить наиболее производительные ядра, доступные на Hub для вашего устройства. Это заменяет поддерживаемые операции PyTorch реализацией на ядрах.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

По умолчанию отображение зависит от модели, устройства и режима выполнения. В таблице ниже перечислены общие репозитории Hub для показанных устройств. Если для определённой комбинации ядро не зарегистрировано, операция выполняется с использованием стандартного PyTorch.

Таблица не является исчерпывающей. Модели могут регистрировать дополнительные слои и функции, а некоторые сопоставления доступны только для определённых режимов выполнения.

Видеокарты AMD сообщают о своём типе устройства как cuda в PyTorch. Transformers определяет ROCm во время выполнения и направляет поддерживаемые операции в ядра AMD, описанные выше, включая сборки AITER, такие как kernels-community/aiter-rope. Вам не нужно самостоятельно задавать тип устройства.

Просмотрите доступные ядра в организации kernels-community.

Ядра внимания

Загружайте ядра внимания из Hub с помощью аргумента attn_implementation.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    device_map="auto"
)

Обратите внимание, что для attention-ядер всё, что не входит в репозиторий kernels-community (которому доверяем — в будущем мы можем добавить ещё доверенные репозитории), потребует использования дополнительного аргумента allow_all_kernels=True (аналогично аргументу trust_remote_code=True для моделей, не относящихся к HF). Это связано с тем, что загрузка ядра может привести к выполнению произвольного кода на хост-машине, и мы не можем проверить каждый репозиторий, поэтому вы должны явно разрешить это.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="random-repo/random-attention",
    allow_all_kernels=True,
    device_map="auto"
)

Определённые ядра, такие как attention, поддерживают несколько форматов.

• @v2.1.0 фиксирует конкретный тег или ветку.
• @>=2.0,<3.0 задаёт ограничения на семантическую версию.

from transformers import AutoModelForCausalLM

# pin to a specific version
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/[email protected]",
    device_map="auto"
)
# use semantic versioning constraints
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@>=2.0,<3.0",
    device_map="auto"
)

Резервный вариант FlashAttention

Запрос attn_implementation="flash_attention_2", "flash_attention_3" или "flash_attention_4" приводит к использованию соответствующего ядра из Hub, если скомпилированный пакет flash-attn не установлен или ваше устройство не является CUDA.

Поддержка FlashAttention-4 находится в бета-версии. API и поведение могут измениться.

from transformers import AutoModelForCausalLM

# uses the compiled flash-attn package if present, otherwise the kernels-community/flash-attn2 Hub kernel
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="flash_attention_2",
    device_map="auto",
)

Осведомлённость о режиме

Ядра автоматически адаптируются к режимам обучения и инференса на основе состояния model.training в PyTorch.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    device_map="auto"
)

# Switch to inference mode - uses inference-optimized kernels
model.eval()
with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=50)

# Switch to training mode - uses training-optimized kernels with gradient support
model.train()
loss = model(input_ids, labels=labels).loss
loss.backward()

Явно включайте режимы обучения и вывода с помощью аргумента mode в функции kernelize(). Режим обучения также поддерживает дополнительный режим torch.compile.

from kernels import Mode
from transformers import kernelize

# inference optimized kernels
kernelize(model, mode=Mode.INFERENCE)

# training optimized kernels
kernelize(model, mode=Mode.TRAINING)

# training and torch-compile friendly kernels
kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE)

KernelConfig

KernelConfig настраивает, какие ядра используются в модели.

Ключи kernel_mapping — это имена, зарегистрированные моделью. Они могут ссылаться на слой, например "RMSNorm", или на зарегистрированную функцию, например "rotary_pos_emb". Разделитель : указывает на конкретный элемент ядра внутри репозитория и сопоставляет его с этим слоем или функцией.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": "kernels-community/liger-kernels:LigerRMSNorm",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto"
)

Укажите различные реализации ядра для каждого типа устройства.

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": {
            "cuda": "kernels-community/liger-kernels:LigerRMSNorm",
            "rocm": "kernels-community/liger-kernels:LigerRMSNorm",
            "xpu": "kernels-community/rmsnorm:RMSNorm"
        }
    }
)

Метаданные ядра

Добавьте словарь метаданных для управления тем, какая сборка ядра загружается.

from transformers import KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "RMSNorm": ("kernels-community/liger-kernels:LigerRMSNorm", {"version": 3}),
    }
)

Старые ветки репозитория ядра могут не содержать сборок для ваших версий PyTorch и CUDA. Если ядро не удаётся загрузить, попробуйте более новую version, прежде чем делать вывод о том, что ваше оборудование не поддерживается.

Наследуемые сопоставления

KernelConfig наследует стандартное сопоставление ядер Transformers для ядер Hub, а записи в kernel_mapping переопределяют стандартные значения для соответствующих слоев и функций. Установите inherit_mapping=False, чтобы использовать только записи из kernel_mapping. Все, что вы не укажете, будет выполняться с помощью стандартного PyTorch. Это полезно при бенчмаркинге конкретного ядра или тестировании собственной реализации без применения остальных значений по умолчанию.

Конфигурация ниже сопоставляет функцию rotary_pos_emb с ядром RoPE и оставляет все остальные операции на PyTorch. Имя функции должно быть зарегистрировано моделью.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    kernel_mapping={
        "rotary_pos_emb": (
            "kernels-community/rotary:apply_rotary_transformers",
            {"version": 2},
        ),
    },
    inherit_mapping=False,
)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Слияние модулей

Соедините соседние модули в одно ядро, передав кортеж пар (class_name, path_pattern) в качестве ключа в KernelConfig. Все паттерны должны иметь один и тот же родительский модуль. * соответствует любому одиночному сегменту пути.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    {
        (
            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
            ("MLP",     "model.layers.*.mlp"),
        ): "owner/fused-rmsnorm-mlp:RMSNormMLP",
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

Для работы Fusion репозиторий ядра должен предоставлять класс KernelNameLayout наряду с классом KernelName. О том, как его реализовать, см. в руководстве Writing kernels.

Локальные ядра

Загружайте ядра из локальных путей к файлам с помощью use_local_kernel=True в KernelConfig. Это позволяет загружать ядра из локальной файловой системы, а не из репозитория Hub.

Локальные ядра используют /abs/path:layer_name вместо формата Hub org/repo:layer_name.

from transformers import KernelConfig, AutoModelForCausalLM

kernel_mapping = {
    "RMSNorm": "/path/to/liger-kernels:LigerRMSNorm",
}
kernel_config = KernelConfig(kernel_mapping, use_local_kernel=True)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype="auto",
    device_map="auto",
    use_kernels=True,
    kernel_config=kernel_config
)

Отключить ядра глобально

Установите переменную окружения USE_HUB_KERNELS, чтобы отключить ядра Hub везде, не изменяя код.

export USE_HUB_KERNELS=0  # or OFF or NO

Устранение неполадок

Интеграция ядра зависит от совместной работы аппаратного обеспечения, драйверов и версий пакетов. В следующих разделах описаны типичные сбои.

Проблемы с установкой

Ошибки импорта означают, что пакет kernels отсутствует или его версия выходит за пределы диапазона, поддерживаемого Transformers. Переустановите его через extra, чтобы получить совместимую версию.

pip install -U "transformers[kernels]"

Сбои при загрузке ядра

Если загрузка конкретных ядер не выполняется, попробуйте следующее.

• Проверьте совместимость вашего оборудования с требованиями к ядру.
• Проверьте, что драйверы CUDA/ROCm/Metal обновлены до последней версии.
• Ознакомьтесь с документацией репозитория ядра, чтобы узнать о известных проблемах.

Совместимость с устройствами

Не все ядра поддерживают все устройства. Если ядро недоступно для вашего оборудования, библиотека переключается на стандартные операции PyTorch. Проверьте документацию репозитория ядер для информации о поддержке конкретных устройств.

Ресурсы

• репозиторий Kernels
• статья в блоге «Улучшите свои модели за 5 минут с помощью Hugging Face Kernel Hub»
• Находите ядра в организации kernels-community