KernelEvolve

KernelEvolve#

KernelEvolve is a self-evolving kernel rewriting application API tool based on graph agent, used to rewrite PyTorch kernel code into Triton implementation through a self-evolving mode, and the module has already integrated a GPU evaluator.

1. Conditions and Restrictions#

  • The configured model must support Function calling; otherwise, this API cannot be used.

  • Evolution results are related to model capabilities. If unable to generate, you can try regenerating multiple times.

  • The Kernel format specification must conform to the following format, start with Model, implement the kernel part in forward, and include get_inputs and get_init_inputs for initializing and testing parameters; otherwise, the generated kernel may not be properly evaluated:

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Compute C = diag(A) * B + D
    A: (N,)
    B: (N, M)
    D: (N, M)
    C: (N, M)
    """

    def __init__(self, BLOCK_M=128):
        super(Model, self).__init__()
        self.BLOCK_M = BLOCK_M

    def forward(self, A, B, D):
        return torch.diag(A) @ B + D


def get_inputs():
    N, M = 4096, 4096
    A = torch.randn(N, dtype=torch.float32)
    B = torch.randn(N, M, dtype=torch.float32)
    D = torch.randn(N, M, dtype=torch.float32)
    return [A, B, D]


def get_init_inputs():
    return []

2. User Guide#

Using kernel evolve to complete kernel evolution rewriting only requires four steps: 1. First, please prepare the kernel code to be rewritten, for example:

torch.diag(A) @ B + D

Therefore, please rewrite in Model format and write the initialization and test parameter acquisition functions:

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Compute C = diag(A) * B + D
    A: (N,)
    B: (N, M)
    D: (N, M)
    C: (N, M)
    """

    def __init__(self, BLOCK_M=128):
        super(Model, self).__init__()
        self.BLOCK_M = BLOCK_M

    def forward(self, A, B, D):
        return torch.diag(A) @ B + D


def get_inputs():
    N, M = 4096, 4096
    A = torch.randn(N, dtype=torch.float32)
    B = torch.randn(N, M, dtype=torch.float32)
    D = torch.randn(N, M, dtype=torch.float32)
    return [A, B, D]


def get_init_inputs():
    return []

  1. After the code is ready, please import the dependency LLMConfig and configure the model parameters.

LLMConfig(
    model_class="PanguClient",
    model_name='Pangu_38b_5.0.3.1',
    api_key="xxxx",
    base_url="xxxx",
    default_headers={"csb-token": "xxxx"}
)

  1. You can directly use evofabric.app.kernel_evolve.GPUEvaluator as the evaluator for evolutionary evaluation of GPU operators. Additionally, you can inherit BaseEvaluator to implement the evaluation method as a reference for evolution according to your needs:

class GPUKernelEvaluator(BaseEvaluator):
    def evaluate(self, initial_code, evolve_code) -> Metrics:
        logger.info(f"Initial code: {initial_code}")
        logger.info(f"Evolve code: {evolve_code}")
        metrics = {
            "speedup": 1.5,
            "original_time": 380,
            "optimized_time": 190,
        }
        return Metrics(**metrics)

  1. Use the previous configuration as the constructor for the evolver, initialize KernelEvolve, start the evolutionary evaluation, and obtain the rewritten results:

kernel_evolve = KernelEvolve(
    initial_code=original_code,
    llm_config=config,
    evaluator=evaluator)
success_flag, result = kernel_evolve.evolve()

  • If the rewrite is successful, return success_flag as True;

  • If the rewrite fails, an error message will be returned.

Contents