orchestrator 是 KernelAgent 系统中的一个核心组件，负责协调和管理多个工作进程(worker)，实现并行执行任务并从中选择最优结果。Fuser/orchestrator.py 文件实现了 Orchestrator 类，用于多进程协调任务执行。其功能用一句话概括：fork N 个Worker竞赛，首个 PASS胜出，其余终止，产物打包返回。
orchestrator 的精简架构如下：
前文提到，KernelAgent 的完整流水线（pipeline）分为四个阶段？
其实，实际代码和官方文档有出入，实际代码中，Orchestrator 是 Extract 阶段的子组件，负责"重写"这一步，而 subgraph_extractor 是 Extract 阶段的完整入口。或者说，Orchestrator = 多 worker 竞赛生成融合代码，是 Pipeline 第一步 Extract 阶段的执行引擎。
即，Extract 阶段包含两步，Orchestrator 负责第一步：
Extract 阶段（subgraph_extractor.py 统一入口）
|
|-- Step 1a: Orchestrator 重写代码
| 输入：原始 PyTorch 问题文件
| 过程：多 Worker 竞赛，LLM 将代码重构为可融合子模块
| 输出：重构后的 code.py（仍是 PyTorch，但已拆分为 nn.Module 子模块）
|
|-- Step 1b: LLM 分析子图
输入：原始问题 + 重构后的 code.py
过程：单次 LLM 调用，提取 shapes/ops/dtypes
输出：subgraphs.json（JSON 数组，描述每个子图的精确形状签名）
在 subgraph_extractor.py 的 extract_subgraphs_to_json() 函数中可以清晰看到：
Orchestrator (Fuser/orchestrator.py) 的功能是 将 PyTorch 模块重写为可融合的子模块代码。Fuser 阶段的核心目的是在保持程序语义的前提下，重新组织模型结构以促进更有效的操作融合。通过将相关的连续操作打包成独立的模块，同时保留控制流结构，该阶段为后续的子图提取、Triton 内核生成和最终合成提供了优化的基础。这种结构化的方法既提高了性能潜力，又保持了代码的可读性和功能完整性。
具体来说，Orchestrator 的核心职责是：
Orchestrator 的关键机制如下：
输出：一个重构后的 Python 文件（仍然是 PyTorch 代码，但已拆分为 nn.Module 子模块），这个文件随后被 subgraph_extractor 分析产出 subgraphs.json 。
Fuse 阶段是将原始模型中的多个连续操作融合成更少的子模块，以便后续生成更高效的 Triton 内核。以下面代码为例，其具体变化分析，是从线性结构到模块化结构。
输入：任意复杂度的原始 PyTorch 模型（包含多个单独的 nn 层），即单一 forward 方法，顺序执行多个操作
x → conv → bn → tanh → max_pool2d → norm
class Model(nn.Module):
def forward(self, x):
if x.sum() > 0:
x = self.conv(x)
x = self.bn(x)
x = torch.tanh(x)
x = F.max_pool2d(x, 2)
return self.norm(x)
x → branch（融合了conv+bn+tanh+max_pool2d) → norm
class FusedModel(nn.Module):
def __init__(self, channels: int):
super().__init__()
self.branch = ConvBnTanhMaxPool(channels=channels) # 融合模块
self.norm = ChannelwiseNorm(channels=channels) # 独立模块
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.sum() > 0: # 控制流 intact（保持不变）
x = self.branch(x) # 单个调用执行多个操作
return self.norm(x)
Fuse 阶段的具体操作如下：
为何这样做可以性能优化？
减少内核启动开销
提高内存访问效率
Triton 内核生成优化
控制流保留的重要性
工作进程入口函数：_worker_process_main 函数
总结报告生成 ResultSummary 对象
持久化存储
KernelAgent 采用多 Worker 并行执行的架构设计，这是其核心的性能优化策略之一。每个 Worker 作为独立的进程运行相同的问题解决任务，但可能使用不同的参数或策略。
LLM 固有的不确定性
问题复杂性
参数多样性
随机性和探索性
时间效率
资源利用
KernelAgent 的主要目标是快速获得可用的融合内核，而不是进行详尽的统计分析。而且，在大多数情况下，第一个成功解决方案已经满足需求。
早期停止机制
时间成本：GPU 计算资源昂贵，让所有 Worker 运行到结束会浪费大量计算时间
能源成本：不必要的计算会消耗更多电力
硬件资源：释放 GPU 资源供其他任务使用
机会成本
解决方案等价性
统一输出
选择最佳方案
确定性结果
KernelAgent 当前采用竞争性并行策略，即多个 Worker 同时运行，但只要有一个成功就立即停止其他进程。
竞争性队列
验证标准
优先级顺序：
选择逻辑位于 worker.py 中：
if rr.passed:
self.logger.info("PASS at iter %d via %s", k, rr.validator_used)
try:
self.winner_queue.put({
"worker_id": self.cfg.worker_id,
"iter": k,
"validator": rr.validator_used,
"runs_dir": str(run_root),
"artifacts_dir": str(self.dirs["artifacts"]),
}, timeout=0.1)
except queue.Full:
pass
选择后的行为
立即终止其他进程
p.terminate() 和 p.kill() 强制终止进程结果包装和记录
将胜出者的成果打包到压缩文件中，包含代码文件和运行目录内容
记录胜出者信息到 summary.json 文件
在专门的目录中保存胜出者信息
在 KernelAgent 系统中，prompt 是指导 LLM（大语言模型）生成特定解决方案的关键输入，特别是在 orchestrator 的 worker 过程中用于生成融合子图模块。
render_prompt 函数定义了 prompt。
构建 Prompt 的硬性要求列表如下：
python 代码块包围完整解决方案重发
动态错误上下文（迭代优化）
@dataclass(frozen=True)
class RenderedPrompt:
system: str
user: str
extras: dict[str, Any]
def render_prompt(
problem_path: Path,
variant_index: int,
attempt_index: int,
error_context: str | None,
enable_reasoning_extras: bool,
seed: int | None = None,
model_name: str | None = None,
) -> RenderedPrompt:
"""Render system+user prompts and extras for the Responses API (deterministic)."""
content = problem_path.read_text(encoding="utf-8")
user = build_user_prompt(
attempt_index=attempt_index,
problem_file_content=content,
error_context=error_context,
variant_index=variant_index,
)
extras: dict[str, Any] = {}
if seed is not None:
extras["seed"] = seed
if enable_reasoning_extras:
# Use high reasoning effort for GPT-5 per policy
extras["reasoning"] = {"effort": "high"}
# Align with Responses API text options for clearer outputs
text_options: dict[str, Any] = {"format": {"type": "text"}}
if model_name:
if model_name.startswith("gpt-5"):
text_options["verbosity"] = "high"
elif model_name.startswith("o4-mini"):
text_options["verbosity"] = "medium"
extras["text"] = text_options
return RenderedPrompt(system=SYSTEM_PROMPT, user=user, extras=extras)
简洁明了
SYSTEM_PROMPT 为 "Return a single runnable Python file only."BASE_DEVELOPER_PROMPT 是 User Prompt。
此处做了角色设定：
BASE_DEVELOPER_PROMPT = (
"You are an expert PyTorch engineer focused on inference-only graph fusion.\n\n"
"Hard requirements:\n"
"- Return ONE runnable Python file, fenced as a single ```python block.\n"
"- Each fused subgraph must be represented by its own nn.Module class with a clearly documented forward; do not leave raw nn.* ops inline in the top-level Model.\n"
"- Include a function run_tests() that validates numerical equivalence to the original using helpers in the problem file. "
"On success, run_tests() must print 'PASS' and exit(0).\n"
"- If you cannot implement run_tests(), then at minimum print the exact sentinel ALL_TESTS_PASSED and exit(0) when tests succeed.\n"
"- No network or file I/O outside the current directory. Avoid extra dependencies.\n"
"- Deterministic: set seeds where relevant.\n\n"
"Fusion guidance:\n"
"- Detect scaled dot-product attention patterns and aggressively fuse the entire block (QKV linears, splits/reshapes, scaled QK^T, causal masking, ReLU or gating, applying V, and head merge) into a single attention subgraph whenever feasible.\n"
"- Only decompose attention into smaller subgraphs when you are certain fusion is impossible.\n\n"
"Iteration contract:\n"
"- On each attempt, re-emit the entire single-file solution.\n"
"- When ERROR_CONTEXT is provided, carefully analyze and fix issues, then re-emit the whole file.\n"
)
具体构建代码如下。
def build_user_prompt(
attempt_index: int,
problem_file_content: str,
error_context: str | None,
variant_index: int,
) -> str:
parts: list[str] = []
parts.append(_variant_line(variant_index))
parts.append("")
parts.append(BASE_DEVELOPER_PROMPT)
parts.append("")
parts.append(f"ATTEMPT: {attempt_index}")
if error_context:
parts.append("")
parts.append("ERROR_CONTEXT:")
parts.append(error_context.strip())
parts.append("")
parts.append("PROBLEM_FILE_CONTENT:")
parts.append(problem_file_content)
return "\n".join(parts)
其中 _variant_line 代码如下。
def _variant_line(idx: int) -> str:
i = idx % len(VARIANT_WORDINGS)
return VARIANT_WORDINGS[i]
VARIANT_WORDINGS 是在 prompting.py 文件中定义的一个元组，包含四个不同的提示词变体。
VARIANT_WORDINGS: tuple[str, str, str, str] = (
"Rewrite the provided model into fusable subgraph modules with explicit input/output shapes.",
"Refactor the given model into fusion-friendly submodules, specifying exact tensor shapes.",
"Decompose the model into subgraphs suitable for fusion; document all input/output shapes.",
"Split the model into fusable modules and clearly state the shape contracts for each.",
"Every fused subgraph must be packaged as its own nn.Module (no inline nn.* ops at top level)",
)
# 注意这里有5个元素！
这些变体用于在多 Worker 并行执行时，让不同的 Worker 使用略微不同的提示词来引导 LLM 生成多样化的解决方案。
设计意图如下：
多样性探索
避免重复
这种设计巧妙地利用了提示词工程中的微妙差异来增加解决方案的多样性，同时保持了任务目标的一致性。
各变体详细分析如下。
第一个变体为："Rewrite the provided model into fusable subgraph modules with explicit input/output shapes."
关键词汇
侧重点
第二个变体如下："Refactor the given model into fusion-friendly submodules, specifying exact tensor shapes."
关键词汇
侧重点
"Decompose the model into subgraphs suitable for fusion; document all input/output shapes."
关键词汇
侧重点
"Split the model into fusable modules and clearly state the shape contracts for each."
关键词汇
侧重点
动词选择差异
术语差异
语义细微差别
先说结论，Worker 进程数与 VARIANT_WORDINGS 数目不一定一致。
这种设计提供了灵活性，允许用户根据资源情况调整并发数量，而不受限于提示词变体的数量。
根据代码分析，Worker 进程数由 OrchestratorConfig 中的 workers 参数决定，这个参数是在运行时传入的。
VARIANT_WORDINGS 固定包含 4 个不同的提示词变体：
取模运算分配
在 _make_worker_cfg 方法中，每个 worker 的 variant_index 是这样确定的：
variant_index = idx % 4 # 使用硬编码的 4
循环分配机制
% 4 操作符意味着最多只会循环前 4 个变体，第 5 个变体实际上不会被使用实际情况：不一定相等
--workers 4）分配策略
# 示例：如果有 6 个 workers
worker_0: variant_index = 0 % 4 = 0 # 使用第 1 个变体
worker_1: variant_index = 1 % 4 = 1 # 使用第 2 个变体
worker_2: variant_index = 2 % 4 = 2 # 使用第 3 个变体
worker_3: variant_index = 3 % 4 = 3 # 使用第 4 个变体
worker_4: variant_index = 4 % 4 = 0 # 重新使用第 1 个变体
worker_5: variant_index = 5 % 4 = 1 # 重新使用第 2 个变体
设计考虑
灵活性
多样性
关于迭代优化。我们可以用 KernelBench 的研究来看。KernelBench框架使模型能够在迭代优化过程中接收并利用反馈。这些真实信号包括NVCC编译器错误信息、执行统计数据（例如正确性检查和挂钟时间），以及PyTorch分析器（操作时间分解）。
他们在多轮过程中为模型提供每次生成的反馈：在初始生成后，向模型提供其之前的生成结果G，以及当前生成对应的编译器/执行反馈E和/或分析器输出P。
然后将每次生成及其后续反馈定义为一轮（turn），并在N轮内运行这一迭代优化过程。利用执行反馈有助于减少错误，并随时间提升整体加速效果。
研究人员发现迭代优化在不同模型和KernelBench的各个级别上均持续提升了性能。
此外，通过分析迭代优化轨迹，他们发现模型在执行反馈E的帮助下能更有效地自我纠正，尤其是在修复与执行错误相关的问题上。DeepSeek-R1在Level 1和Level 2上，经过10轮优化后，能在超过90%的任务中生成功能正确的内核。然而，剩余的错误内核几乎总是由于功能不正确而失败，这可能是因为正确性反馈的颗粒度不如执行失败信息细致。
worker.py 是 KernelAgent 系统中执行层面的核心组件，是单个工作进程的实现，负责执行特定的融合任务。
worker.py 实现了从问题描述到有效解决方案的完整迭代流程。它通过与 LLM 的交互、代码执行验证和错误反馈机制，不断改进解决方案，并参与多进程竞争以成为最终的获胜者。
worker.py 的核心职责如下：
WorkerState 数据类如下。
@dataclass
class WorkerState:
worker_id: str
iter_index: int
last_response_id: str | None
last_error: str | None
passed: bool
字段说明
Worker 定义如下。每个工作进程在自己的工作目录中运行，避免文件冲突，保留完整的执行历史。
class Worker:
def __init__(
self,
cfg: WorkerConfig,
problem_path: Path,
winner_queue: Any,
cancel_event: Any,
on_delta: Callable[[str, None]] | None = None,
) -> None:
self.cfg = cfg
self.problem_path = problem_path
self.winner_queue = winner_queue
self.cancel_event = cancel_event
self.on_delta = on_delta
self.logger = setup_file_logger(
cfg.workspace_dir / "logs" / "worker.log", name=f"worker-{cfg.worker_id}"
)
self.dirs = _ensure_dirs(cfg.workspace_dir)
Worker 的业务逻辑在 run 函数中实现。
代码提取
extracted = extract_single_python_file(result.get("output_text", ""))
代码验证
首先会做重复检测， signature = ops 列表 + input shapes + output shapes + weight shapes + layout + dtype，JSON 序列化后比较。
sha = sha256_of_code(extracted.code)
status, owner = register_digest(
self.cfg.shared_digests_dir, sha, self.cfg.worker_id, k
)
此处子图去重（dedup）是基于shape signature而非代码文本。其理由如下：
代码执行
其次会调用 run_candidate 函数来执行：
run_candidate 在 runner.py 中定义，其中验证标准如下：
_PASS_REGEX.search(out_text)（查找 "PASS"）_SENTINEL in out_text（查找 "ALL_TESTS_PASSED"）获胜者竞争机制
获胜者队列为 winner_queue。
self.winner_queue.put({
"worker_id": self.cfg.worker_id,
"iter": k,
"validator": rr.validator_used,
"runs_dir": str(run_root),
"artifacts_dir": str(self.dirs["artifacts"]),
}, timeout=0.1)
WorkerManager如何实现“个worker 成功后立刻停止其他worker“？具体如下：
使用multiprocessing.Event（self.success_event）作为跨进程信号：
错误处理机制
错误上下文在 state.last_error 中传递。
state.last_error = f"RUN_FAIL:{rr.reason}\nSTDOUT_TAIL:\n{out_tail}\nSTDERR_TAIL:\n{err_tail}"
错误分类如下：
迭代恢复
Orchestrator 的交互
通过 winner_queue 向协调器报告成功
响应 cancel_event 信号
与 Runner 的交互
调用 run_candidate 执行代码
接收 RunResult 验证结果
与 Prompting 系统的交互
使用 render_prompt 生成提示
利用 VARIANT_WORDINGS 的不同变体
def run(self) -> None:
state = WorkerState(
worker_id=self.cfg.worker_id,
iter_index=0,
last_response_id=None,
last_error=None,
passed=False,
)
_write_json(self.cfg.workspace_dir / "state.json", asdict(state))
for k in range(1, self.cfg.max_iters + 1):
if self.cancel_event.is_set():
self.logger.info("cancel seen; exiting")
return
state.iter_index = k
_write_json(self.cfg.workspace_dir / "state.json", asdict(state))
# Render prompt
rp = render_prompt(
problem_path=self.problem_path,
variant_index=self.cfg.variant_index,
attempt_index=k,
error_context=state.last_error,
enable_reasoning_extras=self.cfg.enable_reasoning_extras,
model_name=self.cfg.model,
)
prompt_path = self.dirs["prompts"] / f"iteration_{k}.txt"
prompt_path.write_text(rp.user, encoding="utf-8")
"""
Temporary MUX to support Relay while we migrate to OpenAI Responses
API.
Uses EventAdapter for OpenAI otherwise Provider inferface
"""
provider = get_model_provider(self.cfg.model)
if provider.name != "openai":
# Call LLM directly using provider
messages: list[dict[str, str]] = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": rp.user},
]
try:
response = provider.get_response(
self.cfg.model, messages, max_tokens=16000, **rp.extras
)
result = {
"output_text": response.content or "",
"response_id": response.response_id or None,
"error": None,
}
except Exception as e:
error = f"stream_error: {e.__class__.__name__}: {e}"
result = {
"output_text": "",
"response_id": None,
"error": error,
}
else:
# Stream via EventAdapter
jsonl_path = self.dirs["responses"] / f"iteration_{k}.stream.jsonl"
adapter = EventAdapter(
model=self.cfg.model,
store_responses=self.cfg.store_responses,
timeout_s=self.cfg.llm_timeout_s,
jsonl_path=jsonl_path,
stop_event=self.cancel_event,
on_delta=self.on_delta,
)
result = adapter.stream(
system_prompt=SYSTEM_PROMPT, user_prompt=rp.user, extras=rp.extras
)
state.last_response_id = result.get("response_id")
_write_json(
self.dirs["responses"] / f"iteration_{k}.final.json",
result,
)
if self.cancel_event.is_set():
self.logger.info("cancel after streaming; exiting")
return
# Extract code
try:
extracted = extract_single_python_file(result.get("output_text", ""))
except Exception as e:
state.last_error = f"EXTRACT_FAIL: {e}"
self.logger.warning("iteration %d: extract failed: %s", k, e)
continue
iter_art_dir = self.dirs["artifacts"] / f"iteration_{k}"
latest_dir = self.dirs["artifacts"] / "latest"
iter_art_dir.mkdir(parents=True, exist_ok=True)
(iter_art_dir / "code.py").write_text(extracted.code, encoding="utf-8")
(latest_dir / "code.py").write_text(extracted.code, encoding="utf-8")
# Dedup registration
sha = sha256_of_code(extracted.code)
status, owner = register_digest(
self.cfg.shared_digests_dir, sha, self.cfg.worker_id, k
)
if status == "duplicate_cross_worker":
self.logger.info("duplicate across workers (owner=%s); exiting", owner)
return
if status == "duplicate_same_worker":
self.logger.info("duplicate in same worker; continuing")
continue
# Execute
run_root = self.dirs["runs"] / f"iteration_{k}"
run_root.mkdir(parents=True, exist_ok=True)
rr = run_candidate(
artifacts_code_path=latest_dir / "code.py",
run_root=run_root,
timeout_s=self.cfg.run_timeout_s,
isolated=self.cfg.isolated,
deny_network=self.cfg.deny_network,
cancel_event=self.cancel_event,
)
if rr.passed:
self.logger.info("PASS at iter %d via %s", k, rr.validator_used)
try:
self.winner_queue.put(
{
"worker_id": self.cfg.worker_id,
"iter": k,
"validator": rr.validator_used,
"runs_dir": str(run_root),
"artifacts_dir": str(self.dirs["artifacts"]),
},
timeout=0.1,
)
except queue.Full:
pass
state.passed = True
_write_json(self.cfg.workspace_dir / "state.json", asdict(state))
return
# Build ERROR_CONTEXT and continue
out_tail = _tail_text(rr.stdout_path)
err_tail = _tail_text(rr.stderr_path)
state.last_error = f"RUN_FAIL: {rr.reason}\nSTDOUT_TAIL:\n{out_tail}\nSTDERR_TAIL:\n{err_tail}"
_write_json(self.cfg.workspace_dir / "state.json", asdict(state))
# Done all iterations
self.logger.info("exhausted max_iters without PASS")
runner.py 是 KernelAgent 系统中安全执行候选程序的模块，负责在隔离环境中运行生成的代码并验证其正确
runner.py 是 KernelAgent 系统中至关重要的验证组件，它提供了安全、隔离的执行环境来测试生成的代码。通过多层次的验证机制和严格的资源限制，确保只有真正正确的解决方案才能被认为是成功的，从而保证了整个系统的可靠性和安全性。
runner.py 的职责如下：
创建安全的执行环境
运行生成的 Python 代码
捕获和分析执行结果
根据预定义规则判断是否通过验证
RunResult 如下。
@dataclass (frozen=True)
class RunResult:
rc: int # 返回码
passed: bool # 是否通过验证
validator_used: str # 使用的验证器类型
reason: str # 原因说明
t_started: float # 开始时间
t_finished: float # 结束时间
stdout_path: Path # 标准输出文件路径
stderr_path: Path # 标准错误文件路径
验证常量如下：
_SENTINEL = "ALL_TESTS_PASSED" # 通用哨兵字符串
_PASS_REGEX = re.compile (r"\bPASS\b") # PASS 正则表达式
run_candidate 函数为 runner.py 的主要逻辑。
输入参数如下：
artifacts_code_path: 候选代码文件路径
run_root: 执行根目录
timeout_s: 超时时间（秒）
isolated: 是否使用隔离环境
deny_network: 是否禁用网络
cancel_event：取消事件
进程管理
在子进程中使用 procss group
p = subprocess.Popen(
argv,
cwd=str(run_dir),
stdin=subprocess.DEVNULL,
stdout=f_out,
stderr=f_err,
start_new_session=True,
env=env,
)
创建执行环境
创建唯一标识的执行目录
避免命名冲突
run_dir = (
run_root
/ f"attempt_{int(time.time() * 1000)}_{os.getpid()}_{random.randint(0, 9999):04d}"
)
run_dir.mkdir(parents=True, exist_ok=False)
准备执行文件
复制候选代码到执行目录
使用 candidate_main.p 避免与 Python 标准库的 code.py 冲突
exec_filename = "candidate_main.py"
code_dst = run_dir / exec_filename
st = artifacts_code_path.lstat()
if not stat.S_ISREG(st.st_mode) or stat.S_ISLNK(st.st_mode):
raise ValueError("artifacts_code_path must be a regular file (no symlink)")
shutil.copy2(artifacts_code_path, code_dst)
网络限制
如果启用网络限制，在执行目录中创建 sitecustomize.py 来阻止网络连接
if deny_network:
_write_sitecustomize_block_network(run_dir)
环境准备
构建执行命令行参数
应用环境变量白名单
argv = [sys.executable, "-u"]
if isolated and not deny_network:
argv.append("-I")
argv.append(exec_filename)
env = _allowlist_env()
验证逻辑
在 run_candidate 函数中实现优先级验证顺序
if rc == 0:
# Prefer explicit run_tests PASS if present in stdout
if _PASS_REGEX.search(out_text): # 最高优先级
passed = True
validator = "run_tests"
reason = "run_tests printed PASS and exited 0"
elif _SENTINEL in out_text: # 次优先级
passed = True
validator = "sentinel"
reason = "sentinel ALL_TESTS_PASSED found and exited 0"
else: # 基础条件
passed = False
validator = "unknown"
if scan_truncated:
reason = (
"rc==0 but neither PASS nor sentinel found (scan_truncated=true)"
)
else:
reason = "rc==0 but neither PASS nor sentinel found"
else:
passed = False
reason = f"nonzero exit code: {rc}"
_allowlist_env() 完成了环境隔离
def _allowlist_env() -> dict[str, str]:
allow: dict[str, str] = {}
for k, v in os.environ.items():
if k == "PATH":
allow[k] = v
elif k == "PYTHONPATH":
# sanitize: keep only absolute, existing dirs
parts = [p for p in v.split(os.pathsep) if p]
keep: list[str] = []
for p in parts:
try:
pp = os.path.abspath(p)
if os.path.isabs(pp) and os.path.isdir(pp):
keep.append(pp)
except Exception:
continue
if keep:
allow["PYTHONPATH"] = os.pathsep.join(keep)
elif k.startswith("LANG") or k.startswith("LC_"):
allow[k] = v
# Determinism and small resource caps
allow["PYTHONHASHSEED"] = "0"
allow.setdefault("OMP_NUM_THREADS", "1")
allow.setdefault("MKL_NUM_THREADS", "1")
allow.setdefault("OPENBLAS_NUM_THREADS", "1")
return allow
_write_sitecustomize_block_network 完成了网络限制
def _write_sitecustomize_block_network(dst_dir: Path) -> None:
code = (
"import socket\n"
"def _block(*a, **k):\n raise RuntimeError('network disabled')\n"
"class _Blocked(socket.socket):\n def connect(self, *a, **k):\n _block()\n def connect_ex(self, *a, **k):\n _block()\n"
"socket.socket = _Blocked\n"
"socket.create_connection = _block\n"
)
(dst_dir / "sitecustomize.py").write_text(code, encoding="utf-8")
时间限制
使用 subprocess.TimeoutExpired 处理超时
支持外部取消事件
与 Worker 的交互
worker.py 调用 run_candidate 来验证生成的代码
根据 RunResult 决定是否成功
与 Prompting 系统的交互
验证结果影响后续提示的错误上下文
未通过的执行结果会被用作改进提示的依据
此处一个特点是：验证通过的判定逻辑是rc== 0 &&（PASS ∈ stdout || sentinel ∈ stdout）。为什么仅rc==0不够？其原因如下：
误判风险：某些Python 程序即使内部测试失败也可能以rc=0退出（如try-except吞掉了异常）
明确的成功信号：要求主动打印PASS或ALL_TESTS_PASSED是一种“肯定声明"（positiveassertion），证明测试逻辑确实执行了且判定为通过。
防止空程序通过：一个空文件或只有pass的文件也会rc=0，但不会输出这些标识。
渐进容错：两种 validator（run_tests识别PASS，sentinel识别ALL_TESTS_PASSED）提供了灵活性，适应不同格式的测试脚本。
此设计实质上是双重确认：进程退出码+输出内容共同决定验证结果。
def run_candidate(
artifacts_code_path: Path,
run_root: Path,
timeout_s: int,
isolated: bool,
deny_network: bool,
cancel_event: "threading.Event" | None = None,
) -> RunResult:
"""
Execute a candidate program in a fresh run directory under run_root.
- Copies artifacts_code_path to run_dir/candidate_main.py
- Runs [sys.executable, '-u', 'candidate_main.py'] with optional -I (isolated)
- If deny_network, injects sitecustomize.py to block sockets and do NOT use -I
- Captures stdout/stderr to files; kills on timeout or cancel_event
- Classifies pass/fail according to design precedence
"""
run_dir = (
run_root
/ f"attempt_{int(time.time() * 1000)}_{os.getpid()}_{random.randint(0, 9999):04d}"
)
run_dir.mkdir(parents=True, exist_ok=False)
# Prepare working files. We intentionally avoid the name "code.py" here because
# Python's stdlib exposes a module with that name, and PyTorch's import stack
# (via pdb -> code) would accidentally load the candidate file instead of the
# stdlib module, leading to partially initialised torch packages.
exec_filename = "candidate_main.py"
code_dst = run_dir / exec_filename
st = artifacts_code_path.lstat()
if not stat.S_ISREG(st.st_mode) or stat.S_ISLNK(st.st_mode):
raise ValueError("artifacts_code_path must be a regular file (no symlink)")
shutil.copy2(artifacts_code_path, code_dst)
if deny_network:
_write_sitecustomize_block_network(run_dir)
stdout_path = run_dir / "stdout.txt"
stderr_path = run_dir / "stderr.txt"
argv = [sys.executable, "-u"]
if isolated and not deny_network:
argv.append("-I")
argv.append(exec_filename)
env = _allowlist_env()
t_started = time.time()
(run_dir / "EXEC_STARTED").write_text(str(t_started), encoding="utf-8")
# Run the candidate (via subprocess or multiprocess)
rc, t_finished = (
_run_candidate(
run_dir,
argv,
env,
stdout_path,
stderr_path,
t_started,
timeout_s,
cancel_event,
)
if os.getenv("FUSER_COMPOSE_USE_SYS_EXECUTABLE", "1") == "1"
else _run_candidate_multiprocess(
exec_filename,
run_dir,
argv,
env,
stdout_path,
stderr_path,
t_started,
timeout_s,
cancel_event,
)
)
# Read bounded scan for classification
out_text, scan_truncated = _read_all_text_bounded(stdout_path, MAX_SCAN_BYTES)
# Classification
passed = False
validator = "unknown"
reason = ""
if rc == 0:
# Prefer explicit run_tests PASS if present in stdout
if _PASS_REGEX.search(out_text):
passed = True
validator = "run_tests"
reason = "run_tests printed PASS and exited 0"
elif _SENTINEL in out_text:
passed = True
validator = "sentinel"
reason = "sentinel ALL_TESTS_PASSED found and exited 0"
else:
passed = False
validator = "unknown"
if scan_truncated:
reason = (
"rc==0 but neither PASS nor sentinel found (scan_truncated=true)"
)
else:
reason = "rc==0 but neither PASS nor sentinel found"
else:
passed = False
reason = f"nonzero exit code: {rc}"
return RunResult(
rc=rc,
passed=passed,
validator_used=validator,
reason=reason,
t_started=t_started,
t_finished=t_finished,
stdout_path=stdout_path,
stderr_path=stderr_path,
)
KernelFalcon: Autonomous GPU Kernel Generation via Deep Agents
Automating GPU Kernel Generation with DeepSeek-R1 and Inference Time Scaling
DeepSeek-R1自写CUDA内核跑分屠榜！斯坦福学霸狂飙GPU编程自动化挑战人类
大模型能否为不同硬件平台生成高性能内核？南大、浙大提出跨平台内核生成评测框架MultiKernelBench
AKG kernel Agent：利用multi-agent进行kernel的生成和迁移
AKG KERNEL AGENT: A MULTI-AGENT FRAMEWORK FOR CROSS-PLATFORM KERNEL SYNTHESIS
RL 猛刷 CUDA 核：CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning
MultiKernelBench: A Multi-Platform Benchmark for Kernel Generation
Ouyang A, Guo S, Arora S, et al. Kernelbench: Can llms write efficient gpu kernels?[J]. arXiv preprint arXiv:2502.10517, 2025.
Baronio, Carlo, et al. "Kevin: Multi-turn rl for generating cuda kernels."arXiv preprint arXiv:2507.11948(2025).
Li, Shangzhan, et al. "Autotriton: Automatic triton programming with reinforcement learning in llms."arXiv preprint arXiv:2507.05687(2025).
Li, Jianling, et al. "Tritonbench: Benchmarking large language model capabilities for generating triton operators."Findings of the Association for Computational Linguistics: ACL 2025. 2025.
Tjarko Lange, Robert, et al. "Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization."arXiv e-prints(2025): arXiv-2509.
Chen, Wentao, et al. "CUDA-LLM: LLMs Can Write Efficient CUDA Kernels."arXiv preprint arXiv:2506.09092(2025).
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。