dispatch_kernel_agent.py 是 KernelAgent 系统中的调度组件，负责将 subgraph_extractor.py 生成的子图（JSON 格式）转换为具体的 Triton 内核生成任务，并调度 TritonKernelAgent 来生成和验证这些内核。
Dispatcher架构图如下，其功能概括是：读subgraphs.json，把每个子图转成含 reference code 的精确 Triton 生成spec，交给独立的 TritonKernelAgent 实例并发生成，产出kernel.py + summary.json。
dispatch_kernel_agent.py 在 KernelAgent 系统中扮演着桥梁的角色，它将高层的子图分解结果转化为具体的 Triton 内核生成任务。其核心价值在于：
dispatch_kernel_agent.py 在 流水线 中的位置如下：
原始模型 → orchestrator.py（fuse）→ subgraph_extractor.py（extract）→ dispatch_kernel_agent.py → compose_end_to_end.py
run 函数会把子图发给 KernelAgent，来并行生成Triton 内核
def run(
subgraphs_path: Path,
out_dir: Path,
agent_model: str | None = None,
jobs: int = 1,
target_platform: str = "cuda",
max_iters: int = 10,
) -> Path:
"""Dispatch subgraphs to KernelAgent with optional parallelism.
jobs controls the number of concurrent subgraph generations. Default=1
preserves previous behavior and avoids GPU/LLM contention.
"""
# Submit tasks with bounded concurrency
jobs = max(1, int(jobs or 1))
ordered_inputs: list[tuple[int, dict[str, Any]]] = list(enumerate(items, start=1))
results: dict[int, dict[str, Any]] = {}
if jobs == 1: # 串行处理
for pair in ordered_inputs:
i, res = _handle_one(pair)
results[i] = res
else: # 并发处理
with _futures.ThreadPoolExecutor(max_workers=jobs) as ex:
future_map = {
ex.submit(_handle_one, pair): pair[0] for pair in ordered_inputs
}
for fut in _futures.as_completed(future_map):
i, res = fut.result()
results[i] = res
_handle_one 函数会调用 KernelAgent 生成算子。
def _handle_one(idx_item: tuple[int, dict[str, Any]]) -> tuple[int, dict[str, Any]]:
idx, item = idx_item
sid = str(item.get("id", f"subgraph_{idx}"))
pdesc = _synthesize_problem_description(item, target_platform=platform)
sg_dir = out_dir / sid
sg_dir.mkdir(parents=True, exist_ok=True)
(sg_dir / "problem.txt").write_text(pdesc, encoding="utf-8")
# Pin KernelAgent concurrency defaults: 4 workers, max_iters rounds
# 为每个子图创建独立的 TritonKernelAgent 实例
local_agent = TritonKernelAgent(
num_workers=4,
max_rounds=max_iters,
model_name=agent_model,
target_platform=platform,
)
# 生成算子
try:
result = local_agent.generate_kernel(
problem_description=pdesc, test_code=None
)
_handle_one 函数会调用 _synthesize_problem_description 来生成 问题描述
生成包含子图信息、形状、操作序列的问题描述
def _synthesize_problem_description(
item: dict[str, Any], target_platform: PlatformConfig
) -> str:
id_ = str(item.get("id", "unknown"))
type_ = str(item.get("type", ""))
layout = item.get("data_layout") or "NCHW"
dtype = item.get("dtype") or "float32"
input_shape = item.get("input_shape")
output_shape = item.get("output_shape")
inputs_multi = item.get("inputs")
weights_fused = item.get("weights_fused")
weights_orig = item.get("weights_original")
source = item.get("source") or {}
ref_code, _ = _build_reference_code(item)
# Get device string for the platform
header = textwrap.dedent(
f"""
Implement a Triton kernel that computes the following subgraph end-to-end.
Subgraph ID: {id_}
Type: {type_}
Data layout: {layout}
DType: {dtype}
Target Platform: {target_platform.name}
Device String: {target_platform.device_string}
Shapes:
- input: {_fmt_shape(inputs_multi[0]) if isinstance(inputs_multi, list) else _fmt_shape(input_shape)}
{("- input2: " + _fmt_shape(inputs_multi[1])) if isinstance(inputs_multi, list) and len(inputs_multi) > 1 else ""}
- output: {_fmt_shape(output_shape)}
Weights (fused): {json.dumps(weights_fused, indent=2) if isinstance(weights_fused, dict) else "null"}
Weights (original): {json.dumps(weights_orig, indent=2) if isinstance(weights_orig, dict) else "null"}
Operations in order (with parameters):
{json.dumps(item.get("ops", []), indent=2)}
Requirements:
- Return a complete Python file with a @triton.jit kernel and a wrapper function named kernel_function(...).
- kernel_function must accept input tensor(s) and any required weights/bias parameters (match shapes above).
- Implement the exact semantics of the listed ops in the given order for the provided shapes.
- Use {layout} layout and {dtype} dtype semantics.
- Allocate inputs, weights, intermediates, and outputs on device='{target_platform.device_string}' and keep them there throughout forward/verification.
- CPU is acceptable only for metadata, scalars, and export serialization—avoid `.cpu()` or `.to('cpu')` on compute tensors.
- The test will import kernel_function and compare to the reference implementation below.
Test tolerance policy (enforced in generated tests):
- Default tolerances: rtol=1e-3, atol=1e-3.
- Absolute cap: NEVER exceed rtol=1e-2 or atol=1e-2 in torch.allclose.
- For float16/bfloat16 inputs: use rtol=1e-2, atol=1e-2 at most (do not go higher).
- Include a one-line comment if you relax from default; never exceed the cap.
Reference PyTorch implementation (exact semantics to match):
"""
).strip()
src_code_block = "" # optional original snippet for context
if isinstance(source, dict) and source.get("code"):
mod = source.get("module", "Model")
code = str(source.get("code"))
src_code_block = f"\nOriginal source snippet ({mod}):\n```python\n{code}\n```\n"
problem = header + "\n\n```python\n" + ref_code + "```\n" + src_code_block
return problem
其中，_build_reference_code 生成参考实现，即根据操作类型生成对应的 PyTorch 代码
def _build_reference_code(item: dict[str, Any]) -> tuple[str, list[str]]:
"""Return (reference_code_str, param_names) implementing the subgraph.
param_names are additional parameters to reference() beyond the first input(s).
"""
ops: list[dict[str, Any]] = [
op for op in (item.get("ops") or []) if isinstance(op, dict)
]
lines: list[str] = ["import torch", "import torch.nn.functional as F", ""]
params: list[str] = []
# 省略其他代码
dispatch_kernel_agent.py 与 subgraph_extractor.py 的交互
输入：接收 subgraphs.json 文件
处理：解析子图结构，提取操作和形状信息
依赖：依赖于子图提取阶段的输出
dispatch_kernel_agent.py 与 TritonKernelAgent 的交互
调用：为每个子图实例化 TritonKernelAgent
传递：传递合成的问题描述和平台配置
接收：接收生成的内核代码和验证结果
dispatch_kernel_agent.py 与 compose_end_to_end.py 的交互
输出：生成 summary.json，记录每个子图的内核生成结果
用途：为合成阶段提供已验证的 Triton 内核
输出目录结构如下。
out_dir/
├─ <subgraph_id_1>/
│ ├─ problem.txt # 合成的问题描述
│ └─ kernel.py # 生成的 Triton 内核
├─ <subgraph_id_2>/
│ ├─ problem.txt
│ └─ kernel.py
└─ summary.json # 所有子图的生成结果汇总
摘要文件格式如下
[
{
"id": "subgraph_1",
"success": true,
"worker_id": "worker_1",
"rounds": 3,
"session_dir": "/path/to/session",
"kernel_path": "/path/to/kernel.py"
},
{
"id": "subgraph_2",
"success": false,
"message": "generation failed...",
"session_dir": "/path/to/session"
}
]
triton_kernel_agent/agent.py 实现了 TritonKernelAgent 类，这是 Triton 内核生成系统的主要代理类，负责协调整个内核生成过程。
TritonKernelAgent (agent.py)
↓
WorkerManager (manager.py)
↓
VerificationWorker (worker.py)
↓
Kernel Generation & Refinement Loop
TritonKernelAgent 是 Triton 内核生成系统的核心协调者，负责：
它是连接问题描述和实际 Triton 内核实现的关键组件，通过协调多个子组件来实现高效、可靠的内核生成。
该方法明确三大核心组件的强依赖逻辑：_generate_kernel_seeds 生成的多版本初始内核种子，必须基于 generated_test_code （标准化测试代码）进行开发适配；run_verification 则以 _generate_kernel_seeds 的内核种子为验证对象，以 generated_test_code 为验证标准，完成多版本内核的并行有效性检测，三者形成「测试代码标准化→内核种子生成→并行验证筛选」的严格执行链路。
_generate_test 生成标准化测试代码（参考代码仅作为适配依据），确保后续内核种子生成、验证环节使用统一的测试基准，避免因测试代码格式不统一导致的验证失效。_generate_kernel_seeds 生成批量初始内核种子，为并行验证提供多版本候选，相比单版本生成大幅提升「筛选出可通过测试内核」的概率。manager.run_verification 对多版本内核种子做并行验证，利用多工作器同时检测内核是否通过标准化测试，大幅缩短验证耗时，适配批量内核的快速筛选需求。
def generate_kernel(
self, problem_description: str, test_code: str | None = None
) -> dict[str, Any]:
"""
Generate an optimized Triton kernel for the given problem.
Args:
problem_description: Description of the kernel to generate
test_code: Optional test code (generated if not provided)
The test code should:
1. Import the kernel function: from kernel import kernel_function
2. Test the kernel and return True/False
3. Exit with code 0 on success, 1 on failure
Returns:
Dictionary with results including successful kernel
"""
# Always generate test code using LLM (even if test is provided as reference)
generated_test_code = self._generate_test(problem_description, test_code)
# Use the generated test code in standardized format
test_code = generated_test_code
# Log inputs
import time
# Add microseconds to ensure unique directory names
timestamp = (
datetime.now().strftime("%Y%m%d_%H%M%S")
+ f"_{int(time.time() * 1000000) % 1000000}"
)
session_dir = self.log_dir / f"session_{timestamp}"
session_dir.mkdir(exist_ok=True)
with open(session_dir / "problem.txt", "w") as f:
f.write(problem_description)
with open(session_dir / "test.py", "w") as f:
f.write(test_code)
# Generate kernel seeds
kernel_seeds = self._generate_kernel_seeds(problem_description, test_code)
# Save seeds
for i, kernel in enumerate(kernel_seeds):
with open(session_dir / f"seed_{i}.py", "w") as f:
f.write(kernel)
# Run parallel verification with session directory for worker logs
result = self.manager.run_verification(
kernel_seeds=kernel_seeds,
test_code=test_code,
problem_description=problem_description,
session_log_dir=session_dir,
)
# Process results
if result and result["success"]:
# Save successful kernel
with open(session_dir / "final_kernel.py", "w") as f:
f.write(result["kernel_code"])
# Save full result
with open(session_dir / "result.json", "w") as f:
json.dump(result, f, indent=2)
return {
"success": True,
"kernel_code": result["kernel_code"],
"worker_id": result["worker_id"],
"rounds": result["rounds"],
"session_dir": str(session_dir),
}
else:
return {
"success": False,
"message": "Failed to generate working kernel",
"session_dir": str(session_dir),
}
该方法是基于 LLM 生成 Triton/CUDA 内核代码配套测试代码的核心功能模块，专为 PyTorch KernelAgent 设计，核心目标是为待实现的 GPU 内核（最终写入 kernel.py）自动生成可直接运行的标准化测试代码，支撑内核代码的语法校验、真机运行验证、功能正确性检测，是 LLM 生成 GPU 内核流水线中「验证环节」的关键组成部分。
kernel 模块导入内核函数（因内核最终写入工作目录的 kernel.py），确保测试代码与内核代码的调用路径一致，无运行路径错误。kernel.py 内部），贴合 Triton/CUDA 内核的测试习惯。
def _generate_test(
self, problem_description: str, provided_test_code: str | None = None
) -> str:
"""
Generate test code for the problem using OpenAI API.
The test must import from 'kernel' module since each worker writes
the kernel to 'kernel.py' in their working directory.
Args:
problem_description: Description of the problem
provided_test_code: Optional reference test code provided by user
Returns:
Generated test code in standardized format
"""
# Use LLM provider if available; no mock fallback allowed
if not self.provider:
raise RuntimeError(
"Unable to generate test code: no LLM provider available and mock fallback disabled"
)
# Use LLM provider if available
if self.provider:
try:
self.logger.info(f"Generating test code using {self.model_name}")
# Create prompt for test generation using template
prompt = self.prompt_manager.render_test_generation_prompt(
problem_description=problem_description,
provided_test_code=provided_test_code,
)
# Call LLM API
messages = [{"role": "user", "content": prompt}]
response_text = self._call_llm(messages, max_tokens=24000)
self.logger.info("Raw test generation response:\n%s", response_text)
# Extract test code from response
test_code = self._extract_code_from_response(response_text)
if test_code:
self.logger.info(
f"Successfully generated test code using {self.model_name}"
)
return test_code
else:
self.logger.error("Failed to extract valid code from LLM response")
raise ValueError("No valid code found in LLM response")
except Exception as e:
self.logger.error(f"Error generating test with LLM API: {e}")
raise
# Mock test generation (fallback)
self.logger.info("Generating test code (mock implementation)")
# If provided test code exists, create a basic wrapper
if provided_test_code:
test_code = '''"""
Test for kernel implementation (adapted from provided test).
"""
import torch
def test_kernel():
"""Test the kernel implementation."""
from kernel import kernel_function
# Adapted from provided test code
try:
# Create test data (standardized format)
test_input = torch.randn(1024, device='cuda')
# Call kernel_function as a normal Python function
result = kernel_function(test_input)
# Basic validation
if result is not None:
print("Test passed!")
return True
else:
print("Test failed: No result returned")
return False
except Exception as e:
print(f"Test failed: {e}")
return False
if __name__ == "__main__":
import sys
success = test_kernel()
sys.exit(0 if success else 1)
'''
else:
test_code = '''"""
Test for kernel implementation.
"""
import torch
def test_kernel():
"""Test the kernel implementation."""
from kernel import kernel_function
# Mock test - replace with actual test logic
try:
# Create test data
test_input = torch.randn(1024, device='cuda')
# Call kernel_function as a normal Python function
# (kernel launch logic is handled inside kernel.py)
result = kernel_function(test_input)
print("Test passed!")
return True
except Exception as e:
print(f"Test failed: {e}")
return False
if __name__ == "__main__":
import sys
success = test_kernel()
sys.exit(0 if success else 1)
'''
return test_code
该方法是基于 LLM 批量生成 Triton 内核初始实现代码（Kernel Seeds）的核心模块，为 PyTorch KernelAgent 提供多版本的初始内核候选代码，所有生成代码需适配指定测试代码并遵循统一封装规范，是 LLM 生成 Triton 内核流水线中「初始代码生成环节」的核心，为后续内核筛选、调优提供多版本基础素材。
num_seeds），未指定时默认匹配工作器数量（self.num_workers），生成多版本初始内核，为后续筛选可用内核提供样本基础。temperature），提升多版本内核的多样性。kernel_function 作为内核启动包装函数，该函数统一处理参数接收、Triton 内核启动逻辑，与测试代码的调用方式完全匹配，无调用规范冲突。
def _generate_kernel_seeds(
self, problem_description: str, test_code: str, num_seeds: int | None = None
) -> list[str]:
"""
Generate initial kernel implementations using OpenAI API.
Args:
problem_description: Description of the kernel to generate
test_code: Test code that the kernel must pass
num_seeds: Number of kernel variations to generate
Returns:
List of kernel implementation strings
"""
if num_seeds is None:
num_seeds = self.num_workers
# Use LLM provider if available
if self.provider:
try:
self.logger.info(
f"Generating {num_seeds} kernel seeds using {self.model_name}"
)
# Create prompt with Triton guidelines using template
prompt = self.prompt_manager.render_kernel_generation_prompt(
problem_description=problem_description, test_code=test_code
)
kernels = []
messages = [{"role": "user", "content": prompt}]
# Use provider's multiple response capability
max_completion_tokens = 20000
if self.provider.supports_multiple_completions():
# Provider supports native multiple completions
responses = self.provider.get_multiple_responses(
self.model_name,
messages,
n=num_seeds,
temperature=0.8,
max_tokens=max_completion_tokens,
high_reasoning_effort=self.high_reasoning_effort,
)
for i, response in enumerate(responses):
kernel_code = self._extract_code_from_response(response.content)
if kernel_code:
kernels.append(kernel_code)
else:
self.logger.warning(
f"Failed to extract code from kernel seed {i}"
)
else:
# Provider doesn't support multiple completions, make individual calls
for i in range(num_seeds):
response_text = self._call_llm(
messages,
max_tokens=max_completion_tokens,
temperature=0.8 + (i * 0.1),
)
kernel_code = self._extract_code_from_response(response_text)
if kernel_code:
kernels.append(kernel_code)
else:
self.logger.warning(
f"Failed to extract code from kernel seed {i}"
)
if kernels:
self.logger.info(
f"Successfully generated {len(kernels)} kernel seeds"
)
return kernels
else:
self.logger.error(
"Failed to extract any valid kernels from LLM responses"
)
raise ValueError("No valid kernel code found in any LLM response")
except Exception as e:
self.logger.error(f"Error generating kernels with LLM API: {e}")
# Fall back to mock implementation
# Mock kernel generation (fallback)
self.logger.info(f"Generating {num_seeds} kernel seeds (mock implementation)")
kernels = []
for i in range(num_seeds):
# Simpler mock that still demonstrates the wrapper pattern
if i == 2: # Third kernel will pass
kernel = '''"""
Kernel implementation - working version.
"""
def kernel_function(*args, **kwargs):
"""Wrapper function that handles kernel launch."""
# Mock implementation that passes tests
# In real kernels, this would launch a Triton kernel
return True
'''
else:
kernel = f'''"""
Kernel implementation {i + 1}.
"""
def kernel_function(*args, **kwargs):
"""Wrapper function that handles kernel launch."""
# Mock implementation that fails
raise NotImplementedError('Mock kernel not implemented')
'''
kernels.append(kernel)
return kernels
本模块与 PromptManager 在如下时机进行交互。
初始化时创建
# Initialize prompt manager with resolved config
self.prompt_manager = PromptManager(target_platform=self._platform_config)
生成测试时
# Create prompt for test generation using template
prompt = self.prompt_manager.render_test_generation_prompt(
problem_description=problem_description,
provided_test_code=provided_test_code,
)
生成内核时
# Create prompt with Triton guidelines using template
prompt = self.prompt_manager.render_kernel_generation_prompt(
problem_description=problem_description, test_code=test_code
)
优化时
# Create refinement prompt using template
prompt = self.prompt_manager.render_kernel_refinement_prompt(
problem_description=problem_description,
test_code=test_code,
kernel_code=kernel_code,
error_info=error_info,
history_context=history_context,
)
会在生成内核时进行验证
# Initialize worker manager
self.manager = WorkerManager(
num_workers=self.num_workers,
max_rounds=self.max_rounds,
log_dir=self.log_dir,
openai_api_key=os.getenv("OPENAI_API_KEY"),
openai_model=self.model_name,
high_reasoning_effort=self.high_reasoning_effort,
target_platform=self._platform_config.name,
)
# Run parallel verification with session directory for worker logs
result = self.manager.run_verification(
kernel_seeds=kernel_seeds,
test_code=test_code,
problem_description=problem_description,
session_log_dir=session_dir,
)
render_kernel_generation_prompt 函数生成了用于生成kernel的prompt。
def render_kernel_generation_prompt(
self,
problem_description: str,
test_code: str,
triton_guidelines: str | None = None,
) -> str:
"""
Render the kernel generation prompt.
Args:
problem_description: Description of the kernel to generate
test_code: Test code that the kernel must pass
triton_guidelines: Optional guidelines (if None, loads from template)
Returns:
Rendered prompt string
"""
template = self.templates["kernel_generation"]
# Load triton guidelines if not provided
if triton_guidelines is None:
triton_guidelines = self.render_triton_guidelines()
return template.render(
problem_description=problem_description,
test_code=test_code,
triton_guidelines=triton_guidelines,
kernel_guidance=self.target_platform.kernel_guidance,
)
render_triton_guidelines 函数会获取生成 Triton的指导准则。
def render_triton_guidelines(self) -> str:
"""
Render the Triton guidelines.
Returns:
Rendered guidelines string
"""
template = self.templates["triton_guidelines"]
return template.render()
template_files 如下：
# Define template mappings
template_files = {
"test_generation": "test_generation.j2",
"kernel_generation": "kernel_generation.j2",
"kernel_refinement": "kernel_refinement.j2",
"triton_guidelines": "triton_guidelines.j2",
}
四个模板如下：
该文件 定义Triton内核编程的基本规范和最佳实践。
结构规范：
内存访问模式：
优化技术：
运行时限制：
该文件生成全新的 Triton 内核实现
该文件会基于已有内核和错误信息进行修复和改进
triton_guidelines.j2：
kernel_generation.j2：
kernel_refinement.j2：
triton_guidelines.j2：
kernel_generation.j2：
kernel_refinement.j2：
triton_guidelines.j2：
kernel_generation.j2：
kernel_refinement.j2：
triton_guidelines.j2：
kernel_generation.j2：
kernel_refinement.j2：
manager.py 是 KernelAgent 系统中的 WorkerManager 类实现文件，负责管理多个内核验证工作者进程的并行执行。在系统架构中的位置如下。
TritonKernelAgent
↓
WorkerManager (manager.py)
↓
VerificationWorker (worker.py)
↓
Kernel Generation & Testing
manager.py 在 KernelAgent 架构中扮演着并行计算调度器的角色，通过多进程并行执行多个内核验证任务，显著提升了内核生成的成功率和效率。其设计考虑了容错性、性能和资源管理等多个方面，是整个系统能够高效运行的关键组件之一。
使用场景
WorkerManager 类负责多进程协调，其工作负载分配原则如下：
class WorkerManager:
def __init__(
self,
num_workers: int = 4,
max_rounds: int = 10,
history_size: int = 8,
log_dir: str | None = None,
openai_api_key: str | None = None,
openai_model: str = "gpt-5",
high_reasoning_effort: bool = True,
target_platform: str = "cuda",
):
# 初始化多个工作者进程
self.num_workers = num_workers
self.workers: list[mp.Process] = []
会为 为工作者创建临时工作目录。
@contextmanager
def temp_workdirs(self) -> list[Path]:
"""Create temporary working directories for workers."""
workdirs = []
try:
for i in range(self.num_workers):
workdir = Path(tempfile.mkdtemp(prefix=f"worker_{i}_"))
workdirs.append(workdir)
self.logger.info(f"Created workdir for worker {i}: {workdir}")
yield workdirs
finally:
# Cleanup
for workdir in workdirs:
if workdir.exists():
shutil.rmtree(workdir)
self.logger.info(f"Cleaned up workdir: {workdir}")
启动：使用 multiprocessing.Process 启动工作者进程
监控：监控工作者进程状态
清理：在完成或异常情况下终止进程
成功信号
self.success_event = mp.Event() # 跨进程成功信号
当任一工作者找到解决方案时，设置事件
其他工作者检测到信号后停止工作
结果收集
self.result_queue = mp.Queue() # 结果队列
收集各工作者的结果
确保最成功的结果被返回
关键设计特点如下：
容错机制
超时处理：进程超时后强制终止
异常恢复：捕获和处理各种异常情况
资源清理：确保临时文件和进程被正确清理
性能优化
并发执行：利用多核并行验证
早期退出：一旦找出成功解，立即停止其他进程
资源限制：控制并发数量，避免资源耗尽
配置灵活性
验证执行流程如下。
def run_verification(
self,
kernel_seeds: list[str],
test_code: str,
problem_description: str,
session_log_dir: Path | None = None,
) -> dict[str, Any | None]:
"""
Run parallel verification on multiple kernel seeds.
Args:
kernel_seeds: List of initial kernel implementations
test_code: Test code to verify kernel correctness
problem_description: Description of the problem
session_log_dir: Optional session directory for worker logs
Returns:
Dictionary with successful kernel and metadata, or None
"""
self.logger.info(f"Starting verification with {len(kernel_seeds)} seeds")
# Reset cross-worker success signal for a fresh run
try:
self.success_event.clear()
except Exception:
pass
# Reset workers list to avoid holding stale processes between runs
self.workers = []
# Determine where to put worker logs
if session_log_dir:
workers_parent_dir = Path(session_log_dir) / "workers"
workers_parent_dir.mkdir(exist_ok=True)
else:
workers_parent_dir = self.workers_dir
# 1. 准备工作目录
with self.temp_workdirs() as workdirs:
# Start workers 启动工作者进程
for i, (kernel, workdir) in enumerate(zip(kernel_seeds, workdirs)):
worker_log_dir = workers_parent_dir / f"worker_{i}"
worker_log_dir.mkdir(exist_ok=True)
args = (
i,
kernel,
test_code,
problem_description,
workdir,
worker_log_dir,
self.max_rounds,
self.history_size,
self.success_event,
self.result_queue,
self.openai_api_key,
self.openai_model,
self.high_reasoning_effort,
self.target_platform,
)
process = mp.Process(target=worker_process, args=args)
process.start()
self.workers.append(process)
self.logger.info(f"Started worker {i}")
# Wait for any worker to succeed or all to finish
successful_result = None
# 等待结果或超时
while any(w.is_alive() for w in self.workers):
try:
# Check for results with timeout
result = self.result_queue.get(timeout=1.0)
if result["success"]:
successful_result = result
self.logger.info(f"Worker {result['worker_id']} succeeded!")
# Signal all workers to stop
# 信号其他进程停止
self.success_event.set()
break
except queue.Empty:
continue
# Wait for all workers to finish
# 清理进程
for worker in self.workers:
worker.join(timeout=5.0)
if worker.is_alive():
self.logger.warning(f"Terminating worker {worker.pid}")
worker.terminate()
# Collect any remaining results
while not self.result_queue.empty():
try:
result = self.result_queue.get_nowait()
if result["success"] and successful_result is None:
successful_result = result
except queue.Empty:
break
return successful_result
工作者进程入口
def worker_process(
worker_id: int,
kernel_code: str,
test_code: str,
problem_description: str,
workdir: Path,
log_dir: Path,
max_rounds: int,
history_size: int,
success_event: mp.Event,
result_queue: mp.Queue,
openai_api_key: str | None,
openai_model: str,
high_reasoning_effort: bool,
target_platform: str,
):
"""
Worker process for kernel verification and refinement.
This is run in a separate process.
"""
# Import here to avoid issues with multiprocessing
from .worker import VerificationWorker
worker = VerificationWorker(
worker_id=worker_id,
workdir=workdir,
log_dir=log_dir,
max_rounds=max_rounds,
history_size=history_size,
openai_api_key=openai_api_key,
openai_model=openai_model,
high_reasoning_effort=high_reasoning_effort,
target_platform=target_platform,
)
result = worker.run(
kernel_code=kernel_code,
test_code=test_code,
problem_description=problem_description,
success_event=success_event,
)
result_queue.put(result)
与 TritonKernelAgent 的交互
与 VerificationWorker 的交互
与系统资源的交互
VerificationWorker 是专门用于验证和优化单个 Triton 内核实现的工作流程。
TritonKernelAgent
↓
WorkerManager
↓
VerificationWorker
↓
Kernel Generation & Refinement Loop
VerificationWorker 是 TritonKernelAgent 系统中负责单个内核验证和优化的核心组件，它实现了闭环的验证-优化循环，通过 LLM 驱动的方式自动改进内核实现，同时确保符合 Triton 编程规范和系统约束。
违规检测机制
LLM 交互机制
并发安全
错误处理
日志记录
与 WorkerManager 的交互
与 PromptManager 的交互
与 LLM 提供商的交互
_refine_kernel函数的作用：根据错误信息使用 LLM 来修复和改进当前的Triton内核实现。
这个函数是整个 Triton 内核生成过程中的关键组件，它实现了自动化错误修复和代码优化，使系统能够自主解决遇到的问题。
反馈循环
约束维护
错误处理机制
输入参数
输出结果
_refine_kernel 调用 render_kernel_refinement_prompt。
提示构建：使用PromptManager生成专门的修复提示，其中包含：
VerificationWorker使用subprocess而非exec来执行生成的代码，其原因如下：
不过项目也提供了_run_test_multiprocess作为替代方案（通过multiprocessing.Process+exec）。
VerificationWorker中的DISALLOWED_TORCH_PATTERNS列表给出3种被禁止的模式。该模式用于检测生成的 Triton 内核代码中是否非法使用了 PyTorch 计算，确保内核是真正用 Triton 实现的，而非偷偷调用 PyTorch。
另外，这也是一种安全防线，防止LLM生成的内核代码通过帧内省“作弊"：
LLM有时会“发现“这些捷径，因为它被训练来通过测试一这些规则，确保通过测试的唯一方式是真正实现正确的Triton内核。
被禁止的模式举例如下：
_getframe() → 禁止帧内省（防止从调用者窃取测试数据）检测时会先用_strip_comments_and_strings()去掉注释和文档字符串，避免误报。
Verificationworker的history使用了deque(maxlen=history_size)（默认 8）。为什么要限制 history 窗口大小而不是保留全部历史？其原因如下：
测试执行
def _run_test(self) -> tuple[bool, str, str]:
"""
Run the test script and capture results.
Returns:
Tuple of (success, stdout, stderr)
"""
cmd = [sys.executable, str(self.test_file)]
try:
result = subprocess.run(
cmd,
cwd=str(self.workdir),
capture_output=True,
text=True,
timeout=30, # 30 second timeout
)
success = result.returncode == 0
if success:
self.logger.info("Test passed")
else:
self.logger.error(
"Test failed. Exit code: %s, stderr: %s",
result.returncode,
result.stderr[:500],
)
return success, result.stdout, result.stderr
except subprocess.TimeoutExpired:
self.logger.error("Test timed out")
return False, "", "Test execution timed out after 30 seconds"
except Exception as e:
self.logger.error(f"Test execution error: {e}")
return False, "", str(e)
PyTorch 计算违规检测
DISALLOWED_TORCH_PATTERNS = [
(
re.compile(r"\bimport\s+torch\.nn(\b|\s+as\b)"),
"importing torch.nn modules is not allowed",
),
(
re.compile(r"\bfrom\s+torch\s+import\s+nn\b"),
"importing torch.nn modules is not allowed",
),
(
re.compile(r"\bimport\s+torch\.nn\.functional\s+as\s+F\b"),
"aliasing torch.nn.functional as F is not allowed",
),
(re.compile(r"\btorch\.nn\."), "torch.nn module usage is not allowed"),
(
re.compile(r"\btorch\.nn\.functional\b"),
"torch.nn.functional usage is not allowed",
),
(
re.compile(r"\bF\.[A-Za-z_]+\("),
"torch.nn.functional alias calls (F.*) are not allowed",
),
(re.compile(r"\btorch\.conv"), "torch convolution helpers are not allowed"),
(
re.compile(
r"\btorch\.(relu|sigmoid|tanh|softmax|gelu|mish|hardtanh|max_pool|avg_pool)[A-Za-z0-9_]*\("
),
"PyTorch activation/pooling helpers are not allowed",
),
(
re.compile(r"\bclass\s+\w+\s*\(\s*nn\.Module"),
"Subclassing torch.nn.Module is not allowed",
),
(
re.compile(r"\.forward\("),
"Calling .forward() indicates torch.nn module usage and is not allowed",
),
(
re.compile(r"\btorch\.ops\.aten\b"),
"Low-level torch.ops.aten.* calls are not allowed; implement these ops directly in Triton kernels instead of relying on PyTorch compute",
),
# Generic tensor-tensor math that must be implemented in Triton kernels
(
re.compile(r"\btorch\.(matmul|mm|bmm)\s*\("),
"PyTorch matmul/mm/bmm tensor-tensor ops are not allowed; implement these in Triton kernels",
),
(
re.compile(r"\.(matmul|mm|bmm)\s*\("),
"Tensor.matmul/mm/bmm methods are not allowed; implement these in Triton kernels",
),
(
re.compile(r"\btorch\.einsum\s*\("),
"torch.einsum is not allowed; implement this contraction with Triton primitives",
),
(
re.compile(r"\.einsum\s*\("),
"Tensor.einsum is not allowed; implement this contraction with Triton primitives",
),
# Introspection / frame inspection that can be used to steal test locals
(
re.compile(r"\bimport\s+inspect\b"),
"inspect-based reflection is not allowed inside kernel files",
),
(
re.compile(r"\binspect\.(stack|currentframe|getouterframes)\s*\("),
"inspect stack/frame introspection is not allowed in kernels",
),
(
re.compile(r"\bsys\._getframe\s*\("),
"sys._getframe is not allowed in kernels; do not access caller frames",
),
(
re.compile(r"\.f_locals\b|\.f_globals\b"),
"Accessing frame locals/globals (f_locals/f_globals) from kernels is not allowed",
),
(
re.compile(r"\bglobals\s*\("),
"globals() is not allowed in kernels; avoid depending on ambient test state",
),
(
re.compile(r"\blocals\s*\("),
"locals() is not allowed in kernels; avoid depending on caller scopes",
),
]
LLM 驱动优化
def _refine_kernel(
self,
kernel_code: str,
error_info: dict[str, str],
problem_description: str,
test_code: str,
) -> str:
"""
Refine kernel based on error information using OpenAI API.
Uses multi-turn dialogue by incorporating history of previous attempts.
"""
if self.provider:
try:
self.logger.info(f"Refining kernel using {self.openai_model}")
# Build context from history
# 构建历史上下文
history_context = ""
if self.history:
history_context = "\n\nPREVIOUS ATTEMPTS:\n"
for i, round_data in enumerate(self.history):
# 添加历史记录
history_context += f"\nAttempt {i + 1}:\n"
history_context += f"Kernel code:\n```python\n{round_data['kernel_code'][:500]}...\n```\n"
if round_data.get("stderr"):
history_context += f"Error: {round_data['stderr'][:200]}\n"
if round_data.get("stdout"):
history_context += f"Output: {round_data['stdout'][:200]}\n"
# Create refinement prompt using template
# 使用模板创建优化提示
prompt = self.prompt_manager.render_kernel_refinement_prompt(
problem_description=problem_description,
test_code=test_code,
kernel_code=kernel_code,
error_info=error_info,
history_context=history_context,
)
# Call LLM API
messages = [{"role": "user", "content": prompt}]
response_text = self._call_llm(messages, max_tokens=20000)
# Extract refined kernel from response
# 从响应中提取优化后的内核
refined_kernel = self._extract_code_from_response(response_text)
if refined_kernel:
self.logger.info(
f"Successfully refined kernel using {self.openai_model}"
)
return refined_kernel
else:
self.logger.error("Failed to extract valid code from LLM response")
# Return original kernel if extraction fails
return kernel_code
except Exception as e:
self.logger.error(f"Error refining kernel with LLM API: {e}")
# Fall back to mock refinement
# Mock refinement (fallback)
self.logger.info("Refining kernel (mock implementation)")
# For testing, make a simple modification
if "error" in error_info.get("stderr", "").lower():
# Add a comment to show refinement happened
return f"# Refinement attempt {len(self.history) + 1}\n{kernel_code}"
return kernel_code
def run(
self,
kernel_code: str,
test_code: str,
problem_description: str,
success_event: mp.Event,
) -> dict[str, Any]:
"""
Run verification and refinement loop.
Args:
kernel_code: Initial kernel implementation
test_code: Test code to verify kernel
problem_description: Problem description for context
success_event: Shared event to check if another worker succeeded
Returns:
Dictionary with results
"""
self.logger.info(f"Starting verification for worker {self.worker_id}")
# 运行验证和优化循环
current_kernel = kernel_code
for round_num in range(self.max_rounds):
# Check if another worker has succeeded
# 检查是否有其他工作进程已经成功
if success_event.is_set():
self.logger.info("Another worker succeeded, stopping")
return {
"worker_id": self.worker_id,
"success": False,
"stopped_early": True,
"rounds": round_num,
}
self.logger.info(f"Round {round_num + 1}/{self.max_rounds}")
# 更新内核文件
# Write files - test only on first round, kernel every round
if round_num == 0:
# First round: write both kernel and test
self._write_files(current_kernel, test_code)
else:
# Subsequent rounds: only update kernel, test remains unchanged
self._write_kernel(current_kernel)
# 检测违规
violation = self._detect_pytorch_compute(current_kernel)
if violation:
message = f"Disallowed PyTorch usage detected: {violation}"
self.logger.error(message)
self._log_round(round_num + 1, False, current_kernel, "", message)
error_info = {
"stdout": "",
"stderr": message,
"history": list(self.history),
}
# 违规处理
current_kernel = self._refine_kernel(
current_kernel, error_info, problem_description, test_code
)
continue
# Run test 执行测试
success, stdout, stderr = (
self._run_test()
if os.getenv("KA_PROCESS_USE_SYS_EXECUTABLE", "1") == "1"
else _run_test_multiprocess(self.logger, self.workdir, self.test_file)
)
# Log round 记录轮次
self._log_round(round_num + 1, success, current_kernel, stdout, stderr)
if success:
self.logger.info(
f"Success! Kernel passed test in round {round_num + 1}"
)
return {
"worker_id": self.worker_id,
"success": True,
"kernel_code": current_kernel,
"rounds": round_num + 1,
"history": list(self.history),
}
# Refine kernel for next round
# 为下一轮优化内核
error_info = {
"stdout": stdout,
"stderr": stderr,
"history": list(self.history),
}
current_kernel = self._refine_kernel(
current_kernel, error_info, problem_description, test_code
)
# Max rounds reached without success
self.logger.warning(f"Max rounds ({self.max_rounds}) reached without success")
return {
"worker_id": self.worker_id,
"success": False,
"max_rounds_reached": True,
"rounds": self.max_rounds,
"history": list(self.history),
}
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阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。