Tải các kernel

Một kernel hoạt động như một giải pháp thay thế trực tiếp cho các thao tác PyTorch tiêu chuẩn. Nó thay thế phương thức forward bằng cài đặt kernel tối ưu mà không làm hỏng mã mô hình.

Hướng dẫn này cho thấy cách nạp các kernel để tăng tốc suy luận.

Cài đặt Transformers với phiên bản được hỗ trợ của gói kernels.

pip install -U "transformers[kernels]"

Đặt use_kernels=True trong from_pretrained() để tải các kernel hiệu năng cao nhất có sẵn trên Hub cho thiết bị của bạn. Thao tác này thay thế các phép toán PyTorch được hỗ trợ bằng cài đặt kernel.

from transformers import AutoModelForCausalLM

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

Bản ánh xạ mặc định phụ thuộc vào mô hình, thiết bị và chế độ thực thi. Bảng dưới đây liệt kê các kho Hub phổ biến cho các thiết bị được hiển thị. Khi không có kernel nào được đăng ký cho một tổ hợp cụ thể, thao tác sẽ quay về PyTorch tiêu chuẩn.

Bảng này không phải là toàn diện. Các mô hình có thể đăng ký thêm các lớp và hàm, và một số ánh xạ chỉ có sẵn cho các chế độ thực thi cụ thể.

GPU AMD báo cáo loại thiết bị của chúng là cuda trong PyTorch. Transformers phát hiện ROCm khi chạy và định tuyến các thao tác được hỗ trợ đến các kernel AMD ở trên, bao gồm các bản build AITER như kernels-community/aiter-rope. Bạn không cần tự đặt loại thiết bị.

Duyệt các kernel có sẵn trong tổ chức kernels-community.

Các kernel chú ý

Tải các kernel chú ý từ Hub bằng đối số attn_implementation.

from transformers import AutoModelForCausalLM

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

Lưu ý rằng đối với các attention kernel, bất kỳ thứ gì không thuộc về kho lưu trữ kernels-community (được tin cậy - chúng tôi có thể thêm nhiều kho lưu trữ được tin cậy khác trong tương lai) sẽ yêu cầu sử dụng thêm tham số allow_all_kernels=True (tương tự như tham số trust_remote_code=True cho các model không thuộc HF). Điều này là do việc nạp một kernel có thể dẫn đến việc thực thi mã tùy ý trên máy chủ, và chúng tôi không thể xác minh mọi kho lưu trữ, vì vậy bạn cần cho phép điều này một cách rõ ràng.

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

Các kernel cụ thể, chẳng hạn như attention, chấp nhận nhiều định dạng khác nhau.

• @v2.1.0 cố định vào một thẻ hoặc nhánh cụ thể.
• @>=2.0,<3.0 đặt các ràng buộc về phiên bản ngữ nghĩa.

from transformers import AutoModelForCausalLM

# pin to a specific version
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    attn_implementation="kernels-community/flash-attn2@v2.1.0",
    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"
)

Dự phòng cho FlashAttention

Việc yêu cầu attn_implementation="flash_attention_2", "flash_attention_3" hoặc "flash_attention_4" sẽ tự động chuyển sang kernel tương ứng trên Hub khi gói flash-attn đã biên dịch chưa được cài đặt hoặc thiết bị của bạn không phải là CUDA.

Hỗ trợ FlashAttention-4 đang ở giai đoạn beta. API và hành vi có thể thay đổi.

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

Nhận thức chế độ

Kernels tự động thích ứng với chế độ huấn luyện và suy luận dựa trên trạng thái model.training của 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()

Bật chế độ huấn luyện và suy luận một cách tường minh bằng đối số mode trong hàm kernelize(). Chế độ huấn luyện cũng hỗ trợ một chế độ torch.compile bổ sung.

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 tùy chỉnh các kernel được sử dụng trong một mô hình.

Các khóa kernel_mapping là tên được đăng ký bởi mô hình. Chúng có thể tham chiếu đến một lớp, chẳng hạn như "RMSNorm", hoặc một hàm đã đăng ký, chẳng hạn như "rotary_pos_emb". Ký tự phân tách : đặt tên cho một mục kernel cụ thể bên trong kho lưu trữ và ánh xạ nó đến lớp hoặc hàm đó.

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

Chỉ định các triển khai kernel khác nhau cho từng loại thiết bị.

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

Metadata của kernel

Thêm một dict metadata để kiểm soát việc nạp build kernel nào.

from transformers import KernelConfig

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

Các nhánh cũ hơn của kho kernel có thể không có bản dựng cho các phiên bản PyTorch và CUDA của bạn. Nếu một kernel không tải được, hãy thử một version mới hơn trước khi kết luận phần cứng của bạn không được hỗ trợ.

Các ánh xạ kế thừa

Một KernelConfig kế thừa ánh xạ kernel mặc định của Transformers cho các kernel trên Hub, và các mục trong kernel_mapping ghi đè lên mặc định cho các lớp và hàm tương ứng. Đặt inherit_mapping=False để chỉ sử dụng các mục trong kernel_mapping. Mọi thứ bạn bỏ qua sẽ quay về PyTorch tiêu chuẩn. Điều này hữu ích khi benchmark một kernel cụ thể hoặc kiểm tra một cài đặt tùy chỉnh mà không áp dụng các mặc định khác.

Cấu hình bên dưới ánh xạ hàm rotary_pos_emb sang một kernel RoPE và giữ mọi thao tác khác trên PyTorch. Tên hàm phải được đăng ký bởi model.

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

Nối module

Nối các module liền kề thành một kernel duy nhất bằng cách truyền một tuple gồm các cặp (class_name, path_pattern) làm key trong KernelConfig. Tất cả các pattern phải cùng chia sẻ một module cha. * khớp với bất kỳ đoạn path đơn nào.

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 yêu cầu kho kernel phải cung cấp một lớp KernelNameLayout đồng hành bên cạnh lớp KernelName. Xem hướng dẫn Viết kernel để biết cách triển khai một lớp như vậy.

Local kernels

Tải kernel từ đường dẫn tệp cục bộ bằng cách đặt use_local_kernel=True trong KernelConfig. Thao tác này sẽ tải từ đường dẫn hệ thống tệp cục bộ thay vì từ kho Hub.

Các kernel cục bộ sử dụng /abs/path:layer_name thay vì định dạng 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
)

Tắt các kernel trên toàn cục

Đặt biến môi trường USE_HUB_KERNELS để vô hiệu hóa Hub kernels ở mọi nơi mà không cần thay đổi mã của bạn.

export USE_HUB_KERNELS=0  # or OFF or NO

Xử lý sự cố

Việc tích hợp kernel phụ thuộc vào sự phối hợp hoạt động giữa phần cứng, trình điều khiển và các phiên bản gói. Các phần sau đây đề cập đến các lỗi thường gặp.

Sự cố cài đặt

Lỗi import có nghĩa là gói kernels bị thiếu hoặc phiên bản của nó nằm ngoài phạm vi mà Transformers hỗ trợ. Cài đặt lại thông qua extra để có được phiên bản tương thích.

pip install -U "transformers[kernels]"

Lỗi tải Kernel

Nếu các kernel cụ thể không tải được, hãy thử các bước sau.

• Kiểm tra sự tương thích phần cứng của bạn với các yêu cầu về kernel.
• Kiểm tra xem trình điều khiển CUDA/ROCm/Metal của bạn đã được cập nhật lên phiên bản mới nhất chưa.
• Tham khảo tài liệu của kho kernel để biết các vấn đề đã biết.

Tương thích thiết bị

Không phải tất cả các kernel đều hỗ trợ mọi thiết bị. Thư viện sẽ quay lại các thao tác PyTorch tiêu chuẩn nếu một kernel không khả dụng cho phần cứng của bạn. Kiểm tra tài liệu của kho kernel để biết về hỗ trợ cụ thể cho từng thiết bị.

Tài nguyên

• Kernels
• bài đăng blog Nâng cấp Mô hình của Bạn trong 5 Phút với Hugging Face Kernel Hub
• Khám phá các kernel trong tổ chức kernels-community