0x00 概述

dispatch_kernel_agent.py 是 KernelAgent 系统中的调度组件，负责将 subgraph_extractor.py 生成的子图（JSON 格式）转换为具体的 Triton 内核生成任务，并调度 TritonKernelAgent 来生成和验证这些内核。

Dispatcher架构图如下，其功能概括是：读subgraphs.json，把每个子图转成含 reference code 的精确 Triton 生成spec，交给独立的 TritonKernelAgent 实例并发生成，产出kernel.py + summary.json。

0x01 Dispatch Kernel Agent 的作用

dispatch_kernel_agent.py 在 KernelAgent 系统中扮演着桥梁的角色，它将高层的子图分解结果转化为具体的 Triton 内核生成任务。其核心价值在于：

• 自动化任务分配：将复杂的融合模型分解为独立的子图任务
• 标准化问题描述：为每个子图生成适合 Triton 内核生成的描述
• 并行处理能力：支持并发生成多个子图的 Triton 内核
• 结果整合：收集和整理所有子图的内核生成结果，为后续的合成阶段做准备

dispatch_kernel_agent.py 在 流水线 中的位置如下：

原始模型 → orchestrator.py（fuse）→ subgraph_extractor.py（extract）→ dispatch_kernel_agent.py → compose_end_to_end.py

1.1 整体功能

并发处理机制

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] = []

    # 省略其他代码

1.2 与系统其他组件的交互

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 内核

1.3 生成结果管理

输出目录结构如下。

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"
  }
]

0x02 TritonKernelAgent

triton_kernel_agent/agent.py 实现了 TritonKernelAgent 类，这是 Triton 内核生成系统的主要代理类，负责协调整个内核生成过程。

TritonKernelAgent (agent.py) 
    ↓
WorkerManager (manager.py)
    ↓
VerificationWorker (worker.py)
    ↓
Kernel Generation & Refinement Loop

2.1 核心功能

TritonKernelAgent 是 Triton 内核生成系统的核心协调者，负责：

• 配置管理：处理环境变量和默认配置
• 资源初始化：初始化 LLM 提供商、日志记录和子组件
• 测试生成：使用 LLM 生成适当的测试代码
• 内核种子生成：生成多个初始内核实现变体
• 验证协调：协调 WorkerManager 运行并行验证
• 结果处理：处理生成结果并返回适当的响应

它是连接问题描述和实际 Triton 内核实现的关键组件，通过协调多个子组件来实现高效、可靠的内核生成。

主生成方法

该方法明确三大核心组件的强依赖逻辑：_generate_kernel_seeds 生成的多版本初始内核种子，必须基于 generated_test_code （标准化测试代码）进行开发适配；run_verification 则以 _generate_kernel_seeds 的内核种子为验证对象，以 generated_test_code 为验证标准，完成多版本内核的并行有效性检测，三者形成「测试代码标准化→内核种子生成→并行验证筛选」的严格执行链路。

核心特色

1. 测试代码强制标准化，统一验证基准：无论用户是否提供参考测试代码，均通过 _generate_test 生成标准化测试代码（参考代码仅作为适配依据），确保后续内核种子生成、验证环节使用统一的测试基准，避免因测试代码格式不统一导致的验证失效。
2. 全流程会话化归档，可追溯可复现：为每次内核生成任务创建唯一时间戳会话目录，归档问题描述、标准化测试代码、所有内核种子、最终有效内核及验证结果，实现全流程可追溯，便于问题排查与结果复现。
3. 多版本内核种子生成，提升有效率：调用 _generate_kernel_seeds 生成批量初始内核种子，为并行验证提供多版本候选，相比单版本生成大幅提升「筛选出可通过测试内核」的概率。
4. 并行验证筛选，提升效率：通过 manager.run_verification 对多版本内核种子做并行验证，利用多工作器同时检测内核是否通过标准化测试，大幅缩短验证耗时，适配批量内核的快速筛选需求。
5. 标准化结果返回，贴合工程使用：成功时返回有效内核代码、工作器 ID、验证轮次、会话目录等核心信息；失败时明确返回失败状态与原因，结果格式统一，便于上层模块调用与后续处理。

逻辑关系图

代码

    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 内核流水线中「验证环节」的关键组成部分。

核心特色

1. LLM 主导生成，支持参考代码适配：优先调用配置的 LLM 服务商（如 OpenAI），通过 Prompt 模板渲染生成贴合问题描述的测试代码；若用户提供参考测试代码，会基于参考代码适配生成，无参考时则生成通用标准化测试，兼顾灵活性与贴合性。
2. 强约束的标准化输出：强制要求生成的测试代码从 kernel 模块导入内核函数（因内核最终写入工作目录的 kernel.py），确保测试代码与内核代码的调用路径一致，无运行路径错误。
3. 完整的代码提取与异常处理：调用 LLM 后会从返回结果中提取有效代码，无有效代码则直接抛出异常；全流程记录日志（生成开始、原始响应、成功 / 失败状态），并捕获 LLM 调用、代码提取中的所有异常，便于问题排查。
4. 无 Mock 兜底限制，保证生成有效性：仅当未配置 LLM 服务商时才触发 Mock 兜底，且兜底逻辑禁用「Mock 回退开关」，避免无实际能力的空生成，确保测试代码要么由 LLM 专业生成，要么由兜底逻辑生成基础可用代码。
5. 适配 GPU 内核测试特性：兜底测试代码默认基于 PyTorch 实现，针对 CUDA 设备设计测试数据，内核函数以普通 Python 函数方式调用（内核启动逻辑封装在 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 内核流水线中「初始代码生成环节」的核心，为后续内核筛选、调优提供多版本基础素材。

核心特色

1. 批量生成多版本候选内核，支持数量灵活配置：可指定生成内核数量（num_seeds），未指定时默认匹配工作器数量（self.num_workers），生成多版本初始内核，为后续筛选可用内核提供样本基础。
2. LLM 生成强绑定测试代码，确保适配性：生成 Prompt 中融入用户提供的测试代码，要求 LLM 生成的内核必须能对接该测试代码，从源头保证内核与测试的兼容性，避免后续测试环节的基础适配问题。
3. 原生多响应 + 兜底循环调用，适配不同 LLM 服务商能力：智能适配 LLM 服务商能力 —— 支持原生多响应的服PyTorch KernelAgent 源码解读 ---（5）--- Dispatcher规范，贴合工程落地：强制生成的内核遵循固定封装模式，需实现 kernel_function 作为内核启动包装函数，该函数统一处理参数接收、Triton 内核启动逻辑，与测试代码的调用方式完全匹配，无调用规范冲突。
4. 完整的代码提取与容错机制：对每个 LLM 响应单独提取有效内核代码，单版本提取失败仅记录警告不中断整体流程，全量提取失败则抛出明确异常；全流程捕获 LLM 调用异常，异常后自动触发 Mock 兜底，保证功能不中断。

逻辑关系图

代码