扩展框架

如何使用新平台和方法扩展 KernelGenBench。

添加新芯片后端

步骤 1：添加设备检测

Edit src/runtime/__init__.py:

def _detect_device_name() -> str:
    """Detect current hardware platform."""
    # Add detection logic for your platform
    if is_my_platform():
        return "my_platform"
    # ... existing detection logic ...

    return "unknown"

步骤 2：创建平台配置

Add platform constraints to the DEVICE_CONSTRAINTS dict in src/runtime/__init__.py:

class MyPlatformConfig:
    """Configuration for my_platform."""

    name = "my_platform"

    # Numerical tolerance
    rtol = 1e-4
    atol = 1e-4

    # Constraints
    max_threads_per_block = 1024

    # Import statements for generated code
    triton_import = """
import triton
import triton.language as tl
"""

步骤 3：添加提示模板

创建 agent_bench/templates/my_platform/：

# Template for my_platform

{operator_schema}

Write a Triton kernel for this operator.
Use the following imports:

```python
{platform_imports}
```

步骤 4：更新依赖

创建 requirements/requirements_my_platform.txt：

torch>=2.0
triton>=3.0
# Platform-specific packages

添加新智能体方法

按照以下步骤将新的智能体评估方法集成到 KernelGenBench。

步骤 1：创建方法目录

mkdir -p agent_bench/methods/my_method/templates

步骤 2：创建方法配置

创建 agent_bench/methods/my_method/config.yaml：

name: my_method
description: My custom agent method
max_iterations: 10
timeout: 600

步骤 3：创建指令模板

创建 agent_bench/methods/my_method/templates/instructions.md：

# Task

Generate a Triton kernel for the following operator: {operator_name}

## Schema

{operator_schema}

## Requirements

1. Implement all functionality
2. Handle edge cases
3. Optimize for performance

步骤 4：注册方法

编辑 agent_bench/methods/__init__.py：

METHODS = {
    # ... existing methods ...
    "my_method": MyMethodRunner,
}

步骤 5：创建测试脚本

创建 agent_bench/test_my_method.sh：

#!/bin/bash
# Test script for my_method
python run.py --method my_method "$@"

添加新评估指标

Add custom analysis to scripts/analyze/analyze.py:

def compute_my_metric(results):
    """Compute custom metric."""
    # Implementation
    return value

测试扩展

# Test new platform
python -c "from runtime import get_device_type; print(get_device_type())"

# Test new method
cd agent_bench && bash test_my_method.sh add --device-count 1