1. 项目概述
KernelEvolve是Meta公司为解决大规模推荐系统中kernel开发与优化问题而提出的自动化系统。在Meta的广告推荐系统中,每天需要运行海量推理任务,涉及上千个模型、严格的延迟要求(sub-second级别)以及多种硬件平台(包括NVIDIA GPU、AMD GPU和Meta自研的MTIA v3加速器)。传统依赖人工编写和优化kernel的方式已无法满足需求,KernelEvolve旨在通过AI代理自动完成这一过程。
1.1 核心问题解析
Meta面临的kernel开发挑战主要体现在三个维度的异构性:
-
硬件异构:不同平台具有不同的内存层次结构、编程抽象和代际特性。例如:
- NVIDIA的CUDA、Tensor Core、TMA等
- AMD的ROCm/HIP、Infinity Cache等
- MTIA的自定义C++ kernel DSL和特定硬件特性
-
模型异构:推荐系统包含多种模型架构,如:
- 检索模型
- 早期/晚期排序模型
- 基于Transformer的序列模型
- 嵌入密集型模型
-
kernel异构:不仅包含GEMM运算,还有大量:
- 数据预处理算子
- 特征转换
- 稀疏归一化
- 哈希/分桶/截断操作
- 算子融合场景
这三重异构性导致了组合爆炸问题,使得传统手工优化方式难以应对。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 系统设计与架构
2.1 整体架构
KernelEvolve的核心架构如图5所示,包含以下关键组件:
- 树搜索/状态机:将kernel优化建模为图搜索问题
- LLM合成器:负责生成和修改kernel代码
- 上下文记忆子代理:管理运行时上下文信息
- 深度搜索子代理:执行更深入的优化探索
- 持久知识库:存储各平台约束、优化指南和代码示例
- 评估工具链:包括TritonBench、性能分析器和硬件解释器
- 元数据存储:记录kernel候选及其性能指标
2.2 树搜索机制
系统将kernel优化视为图搜索问题,其中:
- 每个节点代表一个kernel实现候选
- 每条边代表一次代码变换
每轮搜索执行三个步
