This repository provides a replication of STARK, an LLM agentic framework for GPU kernel optimization. STARK systematically explores the kernel design space through multi-agent collaboration, grounded instruction, dynamic context management, and strategic search.
Designing and optimizing GPU kernels is an extremely challenging task due to the complex interactions between thread scheduling, memory hierarchy utilization, and hardware-specific characteristics. Existing LLM-based approaches often treat LLMs as single-shot generators or naive iterative tools, which are insufficient for navigating the irregular kernel optimization landscape.
STARK addresses these limitations with three key components:
- Collaborative multi-agent workflow: Decomposes the optimization process into specialized roles—a Plan Agent for high-level strategies, a Code Agent for precise CUDA implementation, and a Debug Agent for targeted fixes.
- Grounded instruction & dynamic context window: Bridges the gap between high-level strategies and low-level code by inserting explicit span anchors (Grounded Instruction) into the code. It also provides each agent with a tailored, dynamic view of past attempts and evaluations.
-
Strategic search with tree memory: Employs an adapted
$\epsilon$ -greedy search policy over a persistent tree memory to balance exploration of new strategies and exploitation of known good kernels, avoiding local optima.
The repository is organized to clearly reflect the components detailed in the STARK paper:
main.py: The entry point of the framework. It implements the main strategic search and optimization loop (Algorithm 1).core/agent/: Implementations of the specialized LLM agents.plan_agent.py: Proposes targeted transformations and emits grounded instructions with span anchors.code_agent.py: Consumes instructions and translates them into executable custom GPU-kernel code.debug_agent.py: Repairs promising but failing candidates using local context and compiler/runtime logs.single_agent.py: A monolithic agent implementation, primarily used for ablation studies.
core/memory/tree.py: Implements the persistent search tree memory and leaderboard to store candidates and their observations.core/search/policy.py: Implements the strategic search policies (e.g.,EpsilonGreedyPolicywith root throttling and dead-branch pruning,BestOfKPolicy).env/: Manages the execution and compilation environment.runner.py: Benchmarks and evaluates candidate GPU kernels in an isolated subprocess.compiler.py: Performs compilation environment diagnostics (checks for ninja, nvcc, and C++ compilers).benchmark.py: Loads the reference tasks to be optimized.
prompts/: Contains the dynamic prompt templates (templates.py) and response parsers (parser.py) for the agents.utils/: Utility functions, including theLLMClient(llm_client.py) and system logger (logger.py).
Install the required basic dependencies:
pip install -r requirements.txtNote: A CUDA-enabled PyTorch environment, along with nvcc and a C++ compiler (like g++ or cl.exe), is required on your system for kernel compilation.
Start the STARK optimization process using main.py:
python main.py --ref_path <path_to_reference_model.py> [options]Key Arguments:
-
--ref_path: Path to the reference PyTorch model file (Required). -
-B,--budget: Maximum number of optimization attempts in the search tree (Default: 30). -
--provider: LLM Provider to use (e.g.,openai,anthropic, or a custom provider) (Default:openai). -
--model: Model name to use for the LLMs (Default:gpt-4o). -
--epsilon: Exploration rate for the$\epsilon$ -greedy selection policy (Default: 0.3). -
--ablation: Ablation mode to test different components (nonefor full STARK,search_agentfor Single-Agent with search,ma_onlyfor Multi-Agent with Best-of-K) (Default:none).
Example:
python main.py --ref_path ./kernelBench/level1/task.py -B 30 --provider openai --model gpt-4o设计和优化 GPU 内核由于涉及线程调度、内存层次结构利用以及硬件特性的复杂交互，是一项极具挑战性的任务。现有的基于大语言模型 (LLM) 的方法通常将 LLM 视为单次生成器或简单的迭代工具，无法有效应对不规则的内核优化设计空间。
STARK 框架通过以下三个核心组件解决这些局限性：
- 协作式多智能体工作流 (Multi-Agent Collaboration)：将优化过程分解为专业角色——负责制定高层策略的 Plan Agent (规划智能体)，负责精确实现 CUDA 代码的 Code Agent (编码智能体)，以及负责修复错误的 Debug Agent (调试智能体)。
- 基础指令与动态上下文 (Grounded Instruction & Dynamic Context)：通过在代码中插入明确的锚点 (Span Anchors) 来弥合高层策略与底层代码之间的鸿沟。同时为每个智能体提供量身定制的、包含历史尝试和评估结果的动态上下文视图。
-
基于树形记忆的策略搜索 (Strategic Search with Tree Memory)：在持久化的树形记忆上采用自适应的
$\epsilon$ -贪心搜索策略，平衡对新策略的“探索 (Exploration)”与对已知优秀内核的“利用 (Exploitation)”，避免陷入局部最优。
本仓库的文件结构与 STARK 论文中描述的组件紧密对应：
main.py: 框架的入口文件，实现了核心的策略搜索和主优化循环（对应论文中的 Algorithm 1）。core/agent/: 包含各类专属 LLM 智能体的实现。plan_agent.py: 提出目标转换策略，并发出带有代码锚点标记的基础指令。code_agent.py: 接收指令，将其转化为可执行的自定义 GPU 内核代码。debug_agent.py: 结合局部上下文和编译/运行日志，修复有希望但运行失败的候选代码。single_agent.py: 单一智能体实现，主要用于消融实验对比。
core/memory/tree.py: 实现了持久化的搜索树记忆 (Tree Memory) 和排行榜 (Leaderboard)，用于存储候选方案及其运行观察结果。core/search/policy.py: 实现了策略搜索算法（例如带有根节点限流和死区剪枝的EpsilonGreedyPolicy，以及BestOfKPolicy）。env/: 管理代码执行和编译环境。runner.py: 在独立的子进程中对候选 GPU 内核进行基准测试和性能评估。compiler.py: 执行编译环境诊断（检查 ninja、nvcc 和 C++ 编译器是否就绪）。benchmark.py: 加载需要被优化的参考任务。
prompts/: 包含供智能体使用的动态 Prompt 模板 (templates.py) 和响应解析器 (parser.py)。utils/: 实用工具包，包含 LLM 客户端 (llm_client.py) 和系统日志记录器 (logger.py)。
安装基础的 Python 依赖包：
pip install -r requirements.txt注意：要进行内核编译，你的系统还需要预先配置好支持 CUDA 的 PyTorch 环境，以及 nvcc 和 C++ 编译器（如 g++ 或 cl.exe）。
使用 main.py 启动 STARK 优化流程：
python main.py --ref_path <参考模型文件路径> [选项]核心参数说明:
-
--ref_path: 待优化的 PyTorch 参考模型文件路径 (必填)。 -
-B,--budget: 搜索树中允许的最大优化尝试次数 (默认: 30)。 -
--provider: 调用的 LLM API 供应商 (如openai,anthropic或自定义) (默认:openai)。 -
--model: 使用的 LLM 模型名称 (默认:gpt-4o)。 -
--epsilon:$\epsilon$ -贪心选择策略的探索率 (默认: 0.3)。 -
--ablation: 消融实验模式，用于测试不同组件 (none代表完整的 STARK，search_agent代表带搜索的单智能体，ma_only代表使用 Best-of-K 的多智能体) (默认:none)。
运行示例:
python main.py --ref_path ./kernelBench/level1/task.py -B 30 --provider openai --model gpt-4o