EvoMem：让 LLM 进化式代码优化不再”失忆”

一句话总结

EvoMem 为 LLM 驱动的进化式代码搜索引入持久化记忆，将每次成功的优化变异提炼为结构化建议供后续运行复用，在 GPU kernel 优化等任务中减少重复探索、加速收敛。

为什么需要这个？

进化式代码优化的工作方式

“让 LLM 帮我优化这段 CUDA kernel”——这个想法已经有了成熟的实现框架。典型流程是：

给 LLM 当前代码，让它生成变异版本
实际编译运行，测量性能（时间、带宽利用率等）
保留性能更好的变体，淘汰差的
重复迭代直到收敛

本质上是遗传算法，只是”变异算子”换成了 LLM。这个框架能发现人工难以想到的优化——比如特定 warp 大小与 tile size 的组合、隐式的 L2 cache 利用方式。

核心问题：每次都从零开始

但有一个根本性的缺陷：这些知识是一次性的。

假设第 1 次运行发现了”对 M×N 矩阵乘法，tile_size=64 比 tile_size=32 在 A100 上快 18%”，这个发现在运行结束后就消失了。第 2 次运行——哪怕是相邻的相似任务——又得重新探索一遍。

量化这个问题：GPU kernel 优化中，每次 LLM 调用耗时 2-5 秒，一次完整进化搜索需要数百次迭代。重复探索导致：

40-60% 的 LLM 调用预算浪费在”已知”方向
共性优化（向量化加载、shared memory tiling、__ldg 只读缓存）被反复”发明”
跨任务的迁移知识完全丢失

EvoMem 的解法直接：把每次成功的变异提炼成可复用建议，持久存储，下次检索引导。

核心原理

直觉：给进化算法装上”经验库”

类比 CUDA 工程师的成长过程。新手每遇到内存带宽瓶颈，都要重新分析：是 global memory 访问不连续？还是 cache miss 率太高？经验丰富的工程师会直接问：”stride 是多少？用 float4 加载试过吗？shared memory tile 对齐了吗？”

EvoMem 就是把这种”经验”结构化存储，让 LLM 在每次变异时能”翻阅”过去的成功案例。

两个阶段

阶段一：提炼与存储（Post-run Extraction）

每次进化运行结束后，对成功的变异事件执行：

让 LLM 分析”原始代码 → 优化代码”的差异，提炼出一句话建议
记录任务上下文（kernel 类型、矩阵规模、GPU 架构）
带 provenance（来源信息）存入记忆库，明确这条建议在哪个场景被验证过

阶段二：检索与引导（During-evolution Retrieval）

新任务开始时：

基于当前任务描述和代码，检索 top-K 条相关建议
将建议注入 LLM 的变异提示词
LLM 在有历史经验参考下生成变异，避免完全随机探索

关键设计：K 不能太大（通常 3-5 条），避免 LLM 被大量历史信息淹没而失去探索性。引导方式是”参考性”而非”强制性”。

记忆的结构

EvoMem 存储的不是代码本身，而是任务感知的结构化建议：

任务类型: 矩阵乘法 CUDA kernel
优化建议: 将 tile_size 从 32 增大到 64，利用 A100 更大的 L1 cache
验证条件: 矩阵规模 ≥ 2048×2048，SM 占用率不受 shared memory 限制
历史提升: +23% (A100, CUDA 12.0)
标签: [shared_memory, tile_size, L1_cache]

这种结构既可跨运行复用，又带有清晰的适用边界。

代码实现

下面实现一个简化版 EvoMem，专注于 CUDA kernel 优化场景。

Baseline：无记忆的进化优化器

from dataclasses import dataclass
from typing import List

@dataclass
class OptResult:
    code: str
    runtime_ms: float
    improvement: float  # 相对 baseline 的提升比例

def evaluate_kernel(cuda_code: str, benchmark_args: dict) -> float:
    """编译运行 CUDA kernel，返回执行时间 (ms)"""
    # nvcc 编译 + 运行 benchmark，多次取均值
    # ... (完整实现省略)
    return 0.0

def llm_mutate(code: str, prompt_hint: str = "") -> str:
    """调用 LLM 生成代码变异"""
    system = "你是 CUDA 优化专家。对给定 kernel 生成一个优化变体，保持接口不变。"
    user = f"{prompt_hint}\n\n当前代码:\n{code}"
    # ... LLM API 调用省略
    return code

def baseline_search(initial_code: str, n_iter: int = 50) -> OptResult:
    """基础进化搜索，无记忆引导"""
    best_code, best_time = initial_code, evaluate_kernel(initial_code, {})
    baseline_time = best_time

    for i in range(n_iter):
        candidate = llm_mutate(best_code)  # 无引导，纯随机变异
        t = evaluate_kernel(candidate, {})
        if t < best_time:
            best_code, best_time = candidate, t
            print(f"Iter {i}: 提升 {(1 - t/baseline_time)*100:.1f}%")

    return OptResult(best_code, best_time, 1 - best_time/baseline_time)

性能分析：对 2048×2048 矩阵乘法 kernel，50 次迭代中约 30-35 次是无效变异（语法等价或性能持平），有效探索不足 40%。

EvoMem：持久化记忆存储与检索

from dataclasses import dataclass

@dataclass
class OptResult:
    code: str
    runtime_ms: float
    improvement: float  # 相对 baseline 的提升比例

def evaluate_kernel(cuda_code: str) -> float:
    # ... (编译运行，多次取均值)
    return 0.0

def llm_mutate(code: str) -> str:
    # ... (LLM API 调用)
    return code

def baseline_search(initial_code: str, n_iter: int = 50) -> OptResult:
    """基础进化搜索，无记忆引导"""
    best_code, best_time = initial_code, evaluate_kernel(initial_code)
    baseline_time = best_time

    for i in range(n_iter):
        candidate = llm_mutate(best_code)  # 无引导，纯随机变异
        t = evaluate_kernel(candidate)
        if t < best_time:
            best_code, best_time = candidate, t

    return OptResult(best_code, best_time, 1 - best_time / baseline_time)

建议提炼模块

@dataclass
class MemoryEntry:
    task_desc: str
    advice: str
    improvement: float
    tags: List[str]
    gpu_arch: str

class EvoMemStore:
    def __init__(self, path: str = "evomem.jsonl"):
        self.path = Path(path)
        self.entries: List[MemoryEntry] = self._load()

    def _load(self) -> List[MemoryEntry]:
        # ... (文件读取与反序列化省略)
        return [MemoryEntry(**json.loads(line)) for line in self.path.read_text().splitlines() if line.strip()]

    def store(self, entry: MemoryEntry):
        self.entries.append(entry)
        # ... (追加写入 JSONL 省略)

    def retrieve(self, task_desc: str, gpu_arch: str, top_k: int = 3) -> List[MemoryEntry]:
        """词袋重叠打分 + 同架构加权（生产环境用向量检索）"""
        task_words = set(task_desc.lower().split())
        scored = [
            (len(task_words & set(e.task_desc.lower().split())) + (0.3 if e.gpu_arch == gpu_arch else 0.0), e)
            for e in self.entries
        ]
        return [e for _, e in sorted(scored, key=lambda x: -x[0])[:top_k]]

完整的 EvoMem 搜索循环

def evomem_search(
    initial_code: str,
    task_desc: str,
    memory: EvoMemStore,
    n_iter: int = 50,
    gpu_arch: str = "A100"
) -> OptResult:
    best_code, best_time = initial_code, evaluate_kernel(initial_code, {})
    baseline_time = best_time
    successful_mutations = []  # 本次成功变异，运行后批量提炼

    for i in range(n_iter):
        # 检索相关历史建议
        relevant = memory.retrieve(task_desc, gpu_arch, top_k=3)
        hint = ""
        if relevant:
            advice_list = "\n".join([
                f"- {e.advice}（历史提升: {e.improvement*100:.1f}%, 标签: {', '.join(e.tags)}）"
                for e in relevant
            ])
            hint = f"参考以下历史优化经验（可选择性采纳）:\n{advice_list}"

        candidate = llm_mutate(best_code, prompt_hint=hint)
        t = evaluate_kernel(candidate, {})

        if t < best_time:
            improvement = 1 - t / best_time
            successful_mutations.append((best_code, candidate, improvement))
            best_code, best_time = candidate, t

    # Post-run: 只提炼提升显著的变异（>5%）存入记忆库
    for orig, opt, imp in successful_mutations:
        if imp > 0.05:
            entry = extract_advice(llm, orig, opt, imp, task_desc)
            memory.store(entry)

    return OptResult(best_code, best_time, 1 - best_time/baseline_time)

常见错误：过度信任记忆

def evomem_search(initial_code, task_desc, memory, n_iter=50, gpu_arch="A100"):
    best_code, best_time = initial_code, evaluate_kernel(initial_code, {})
    successful_mutations = []

    for i in range(n_iter):
        # 检索相关历史建议，构造提示
        relevant = memory.retrieve(task_desc, gpu_arch, top_k=3)
        # ... (hint 构造省略)

        candidate = llm_mutate(best_code, prompt_hint=hint)
        t = evaluate_kernel(candidate, {})

        if t < best_time:
            successful_mutations.append((best_code, candidate, 1 - t / best_time))
            best_code, best_time = candidate, t

    # Post-run: 只提炼提升显著的变异（>5%）存入记忆库
    for orig, opt, imp in successful_mutations:
        if imp > 0.05:
            memory.store(extract_advice(llm, orig, opt, imp, task_desc))

    return OptResult(best_code, best_time, 1 - best_time / evaluate_kernel(initial_code, {}))

引导强度需要平衡：太强会使 LLM 陷入局部最优，丧失探索性；太弱则记忆形同虚设。

性能实测

以下数据来自论文实验（A100 80GB, CUDA 12.0，多次运行取均值）：

场景	无记忆	EvoMem	备注
首次任务	基线	≈ 相同	记忆库为空时无收益
相邻同类 kernel 复用	基线	搜索加速 ~20-30%	记忆迁移效果显著
跨域迁移（不同 kernel 类型）	基线	收益不稳定	视任务相似度而定
LLM 调用次数节省	基线	减少 ~15-25%	减少无效探索

诚实说明：论文明确指出”variability across tasks”——EvoMem 不是银弹。在任务差异较大时，历史记忆可能引入偏见。具体收益依赖于任务的相似性和记忆库的积累量。

什么时候用 / 不用？

适用场景	不适用场景
同类 kernel 的多次优化任务	任务差异极大，记忆难以迁移
有固定优化目标和评估指标	单次运行，无法积累记忆
团队共享优化经验库	记忆库 GPU 架构与当前不匹配
资源受限，需减少 LLM 调用	追求极限性能（记忆可能限制探索范围）

调试技巧

检查记忆库质量：记忆质量比数量更重要。

def audit_memory(memory: EvoMemStore):
    for entry in memory.entries:
        if entry.improvement < 0.03:
            print(f"低效建议（提升<3%，考虑清理）: {entry.advice}")
        if len(entry.advice) > 80:
            print(f"建议过长（精炼不足）: {entry.advice[:40]}...")
        if not entry.tags:
            print(f"缺标签（影响检索精度）: {entry.advice}")

验证 LLM 真正改变了 GPU 行为：用 Nsight Compute 对比变异前后的 l1tex__t_bytes_pipe_lsu_mem_global_op_ld.sum（global memory 加载字节数）和 sm__warps_active.avg.pct_of_peak_sustained_active（warp 占用率）。语法变体不会改变这些指标，真正的优化会。

记忆冷启动策略：首次使用时记忆库为空，建议预置一批人工整理的 CUDA 优化经验（shared memory tiling、向量化加载、bank conflict 避免等），作为初始种子。

EvoMem：让 LLM 进化式代码优化不再"失忆"