Implementation:Hpcaitech ColossalAI Base Extension

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Kernel Extensions, Build System, Infrastructure
Last Updated	2026-02-09 00:00 GMT

Overview

_Extension is the abstract base class for all ColossalAI kernel extensions, defining the interface for ahead-of-time (AOT) and just-in-time (JIT) compilation of C++/CUDA custom operators.

Description

The _Extension class provides the foundational abstraction for building and loading custom kernel extensions in ColossalAI. It defines abstract methods for checking hardware availability, asserting compatibility, building extensions (both AOT and JIT), and loading compiled operators. It also provides a static utility method that computes a deterministic cache directory path for JIT-compiled kernels based on the PyTorch version, device type, device version, and a hash of the ColossalAI installation path.

Usage

Use _Extension as the base class when implementing new custom kernel extensions for ColossalAI. Subclasses (such as _CppExtension and _CudaExtension) must implement all abstract methods to support both ahead-of-time and just-in-time compilation workflows.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: extensions/base_extension.py
• Lines: 1-82

Signature

class _Extension(ABC):
    def __init__(self, name: str, support_aot: bool, support_jit: bool, priority: int = 1):
        ...

    @property
    def name(self):
        ...

    @property
    def support_aot(self):
        ...

    @property
    def support_jit(self):
        ...

    @staticmethod
    def get_jit_extension_folder_path():
        ...

    @abstractmethod
    def is_available(self) -> bool:
        ...

    @abstractmethod
    def assert_compatible(self) -> None:
        ...

    @abstractmethod
    def build_aot(self) -> Union["CppExtension", "CUDAExtension"]:
        ...

    @abstractmethod
    def build_jit(self) -> Callable:
        ...

    @abstractmethod
    def load(self) -> Callable:
        ...

Import

from extensions.base_extension import _Extension

I/O Contract

Inputs

Name	Type	Required	Description
name	str	Yes	The name identifier for the extension
support_aot	bool	Yes	Whether this extension supports ahead-of-time compilation
support_jit	bool	Yes	Whether this extension supports just-in-time compilation
priority	int	No	Priority level for extension selection (default: 1)

Outputs

Name	Type	Description
is_available return	bool	Whether the required hardware is available for this extension
build_aot return	Union[CppExtension, CUDAExtension]	A PyTorch extension object for ahead-of-time compilation
build_jit return	Callable	The loaded JIT-compiled kernel operator
load return	Callable	The loaded kernel operator (from AOT or JIT)
get_jit_extension_folder_path return	str	The cache directory path for JIT-compiled extensions

Cache Directory Format

JIT-compiled kernels are cached at: ~/.cache/colossalai/torch_extensions/torch{major}.{minor}_{device_name}-{device_version}-{hash}

Where hash is the SHA-256 hash of the ColossalAI installation file path.

Usage Examples

from extensions.base_extension import _Extension

# _Extension is abstract; use a concrete subclass
# Example: Check the JIT cache path
cache_path = _Extension.get_jit_extension_folder_path()
print(cache_path)
# ~/.cache/colossalai/torch_extensions/torch2.1_cuda-12.1-<hash>

# Concrete subclasses implement all abstract methods
class MyCustomExtension(_Extension):
    def __init__(self):
        super().__init__(name="my_ext", support_aot=True, support_jit=True)

    def is_available(self) -> bool:
        return True

    def assert_compatible(self) -> None:
        pass

    def build_aot(self):
        ...

    def build_jit(self):
        ...

    def load(self):
        ...

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment