Implementation:Hiyouga LLaMA Factory V1 Kernel Registry

Knowledge Sources	Hiyouga_LLaMA_Factory
Domains	Machine Learning, Software Architecture
Last Updated	2026-02-06 19:00 GMT

Overview

Registry is a class that manages the registration and lookup of hardware-specific kernel implementations, exporting the register_kernel decorator for auto-registration.

Description

The Registry class maintains a class-level dictionary (_kernels) mapping kernel ID strings to BaseKernel subclasses. The register class method acts as a decorator that validates the kernel class inherits from BaseKernel, checks that the kernel's target device matches the current accelerator (skipping registration if not), and ensures no duplicate kernel IDs. The get method retrieves a registered kernel by ID, and get_registered_kernels returns the full registry dictionary. A module-level alias register_kernel is exported for convenient decorator usage.

Usage

Use the @register_kernel decorator on BaseKernel subclasses to register them automatically at import time. Use Registry.get(kernel_id) to retrieve a kernel class by its ID for application. Only kernels matching the current hardware accelerator will be registered.

Code Reference

Source Location

Repository: Hiyouga_LLaMA_Factory
File: src/llamafactory/v1/plugins/model_plugins/kernels/registry.py
Lines: 1-96

Signature

class Registry:
    _kernels: dict[str, type[BaseKernel]] = {}

    @classmethod
    def register(cls, kernel_cls: type[BaseKernel]) -> type[BaseKernel] | None: ...

    @classmethod
    def get(cls, kernel_id: str) -> type[BaseKernel] | None: ...

    @classmethod
    def get_registered_kernels(cls) -> dict[str, type[BaseKernel]]: ...

register_kernel = Registry.register

Import

from llamafactory.v1.plugins.model_plugins.kernels.registry import Registry, register_kernel

I/O Contract

Inputs

Name	Type	Required	Description
kernel_cls (register)	type[BaseKernel]	Yes	The kernel class to register; must inherit from BaseKernel
kernel_id (get)	str	Yes	The unique identifier of the kernel to retrieve

Outputs

Name	Type	Description
register	type[BaseKernel] or None	The registered kernel class if device matches, None otherwise
get	type[BaseKernel] or None	The kernel class if found, None otherwise
get_registered_kernels	dict[str, type[BaseKernel]]	Dictionary mapping all registered kernel IDs to their classes

Usage Examples

from llamafactory.v1.plugins.model_plugins.kernels.registry import Registry, register_kernel
from llamafactory.v1.plugins.model_plugins.kernels.base import BaseKernel
from llamafactory.v1.accelerator.helper import DeviceType

# Register a kernel using the decorator
@register_kernel
class MyKernel(BaseKernel):
    _kernel_id = "my_custom_kernel"
    _device = DeviceType.CUDA

    @classmethod
    def apply(cls, **kwargs):
        model = kwargs.get("model")
        return model

# Retrieve a registered kernel
kernel_cls = Registry.get("my_custom_kernel")
if kernel_cls is not None:
    model = kernel_cls.apply(model=model)

# List all registered kernels
all_kernels = Registry.get_registered_kernels()
for kid, kcls in all_kernels.items():
    print(f"Kernel: {kid} -> {kcls.__name__}")

Related Pages

Hiyouga_LLaMA_Factory_V1_Kernel_Base - Abstract base class that registered kernels must extend
Hiyouga_LLaMA_Factory_V1_NPU_RMSNorm - Example kernel that uses @register_kernel

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment