从 KernelAgent 到 triton-agent：Roofline 分析与 Beam Search 的工程导入方案

May 3, 2026

从 KernelAgent 到 triton-agent：Roofline 分析与 Beam Sea

背景

在研究 meta-pytorch/KernelAgent——Meta PyTorch 团队的 GPU Kernel 自主生成与优化项目——的过程中，我发现它和我的 triton-agent（面向华为 Ascend NPU 的 Triton 算子优化 CLI）有很多概念上的重叠。KernelAgent 的两项核心能力——Roofline 驱动的自动终止和Beam Search 并行探索——是 triton-agent 当前缺失的。这篇文章详细分析这两项能力的工程实现，以及如何把它们引入到 triton-agent 的 skill 系统中。

两个项目的架构对比

先说清楚两者的差异，才能理解"导入"的适配成本。

维度KernelAgenttriton-agent
硬件目标NVIDIA CUDA + Intel XPU华为 Ascend NPU（昇腾）
Profiling 工具NCU (NVIDIA Compute Utility)msprof (Ascend Toolkit)
优化引擎直接调 GPT-5/Claude APICode Agent 后端（Codex/OpenCode/pi/Claude）
优化策略单路径 + Beam Search 可选单路径顺序迭代 (opt-round-1 → 2 → ...)
终止条件Roofline ≥95% SOL 自动停人工判断 + min_rounds 门槛
知识库GPU specs 配置文件 + RAG21 个 pattern card + 4 个 symptom card

关键差异：triton-agent 的优化循环完全依赖 code agent 的自主决策，CLI 只管编排和 gate。KernelAgent 则是 CLI/Orchestrator 深度参与每轮决策——这正是 Roofline 和 Beam Search 能工作的前提。

P0：Roofline 分析 + 自动终止

KernelAgent 的实现

KernelAgent 的 roofline 模块（kernel_perf_agent/kernel_opt/roofline/ncu_roofline.py，~254 行）是一个独立、可配置的分析器。核心架构：

# 简化后的核心逻辑（完整代码见文末仓库链接）
@dataclass
class RooflineConfig:
    threshold_pct: float = 95.0       # SOL% 达到此值认为触顶
    early_stop: bool = True           # 触顶时自动终止
    convergence_rounds: int = 5       # 连续 N 轮无改进触发停止
    min_improvement_pct: float = 0.1  # 最小改进阈值
    underutilized_threshold: float = 60.0  # 双低 = underutilized

class RooflineAnalyzer:
    def analyze(self, ncu_metrics: dict) -> RooflineResult:
        compute_sol = ncu_metrics["sm__throughput.avg.pct_of_peak_sustained_elapsed"]
        memory_sol = ncu_metrics["gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed"]
        efficiency = max(compute_sol, memory_sol)

        bottleneck = self._classify_bottleneck(compute_sol, memory_sol)
        # 逻辑：双低(<60%) → underutilized
        #       memory_sol >= compute_sol → memory-bound
        #       否则 → compute-bound

        return RooflineResult(
            compute_sol_pct=compute_sol,
            memory_sol_pct=memory_sol,
            efficiency_pct=efficiency,
            at_roofline=efficiency >= self.config.threshold_pct,
            bottleneck=bottleneck,
        )

    def should_stop(self, result: RooflineResult) -> tuple[bool, str]:
        self._efficiency_history.append(result.efficiency_pct)
        # 条件1：达到 roofline 阈值且 early_stop 开启
        if self.config.early_stop and result.at_roofline:
            return True, f"At roofline ({result.efficiency_pct}% SOL)"
        # 条件2：连续 N 轮效率无改进
        if len(self._efficiency_history) >= self.config.convergence_rounds:
            recent = self._efficiency_history[-self.config.convergence_rounds:]
            if max(recent) - min(recent) < self.config.min_improvement_pct:
                return True, f"Converged (improvement < {self.config.min_improvement_pct}%)"
        return False, ""

在 OptimizationOrchestrator.optimize_kernel() 中，每轮优化后都运行 roofline 检查，达到阈值直接 break 出循环并输出 early_stop_reason。

triton-agent 当前的状态

triton-agent 的 optimize 已有 profiling 能力：

  • triton-npu-profile-operator skill：调用 msprof 采集算子级性能数据
  • triton-npu-analyze-round-performance skill：基于 profiling 数据做深层诊断
  • opt-note.md 和 perf-analysis.md：记录每轮性能变化

但缺失的是：没有把 profiling 数据量化成单一效率指标，没有形式化的终止条件。当前优化循环的收敛判断完全依赖 code agent 在 opt-note.md 中的主观判断，以及用户设置的 --min-rounds 硬门槛。

从 contract.json 可以看到，当前 round state 的 required fields 包含 perf_artifact、comparison_target、next_recommendation，但没有 efficiency_pct 或 bottleneck 分类字段。

引入方案：新增 triton-npu-roofline 模块

分三步实现。每步都是独立的增量，降低风险。

Step 1：硬件 Spec 知识卡

在 skills/triton-npu-optimize-knowledge/references/ 下新增 hardware_specs.md，记录 Ascend 各芯片的理论峰值：

# Ascend NPU Hardware Specifications

## A5 (Ascend 950)
- FP16 compute peak: 512 TFLOPS
- Memory bandwidth: 2048 GB/s (HBM2e)
- UB size: 256 KB per core
- L1 cache: 512 KB per core
- L2 cache: 96 MB (shared)
- SM count: 80

## A3 (Ascend 910B)
- FP16 compute peak: 320 TFLOPS
- Memory bandwidth: 1536 GB/s
- ...

用 --target-chip A5（已有参数）选择对应配置。

Step 2：在 optimize skill 的 layered analysis 中加入 roofline 层

修改 skills/triton-npu-optimize/SKILL.md 的 Stage 2（Layered Analysis），在 profiling diagnosis 和 IR attribution 之间插入 roofline 层级：

### roofline analysis（新增）

- After each round's benchmark succeeds, calculate SOL efficiency.
- Read `../triton-npu-optimize-knowledge/references/hardware_specs.md` for peak values.
- Compute: efficiency_pct = max(achieved_compute / peak_compute, achieved_bw / peak_bw) × 100
- Classify bottleneck: memory-bound / compute-bound / underutilized.
- Record in `opt-round-N/roofline.json`:
  {
    "efficiency_pct": 72.3,
    "compute_sol_pct": 55.1,
    "memory_sol_pct": 72.3,
    "bottleneck": "memory",
    "at_roofline": false
  }
- Auto-termination: if efficiency_pct >= 95% OR last 5 rounds show improvement < 0.1%,
  record in opt-note.md: `## Early Termination: roofline converged`.

这只是 skill 指令层面的改动——告诉 code agent 在每轮 benchmark 后多做一件事。零 CLI 代码变更。

Step 3：CLI 层面支持自动终止

在 src/triton_agent/optimize/checks.py 中新增 roofline check 函数，hook 进 run_loop：

def check_roofline_termination(workdir: Path, config: RooflineConfig) -> tuple[bool, str]:
    """扫描所有 opt-round-N/roofline.json，判断是否应该终止"""
    rounds = sorted(workdir.glob("opt-round-*/roofline.json"))
    if not rounds:
        return False, ""

    efficiency_history = []
    for r in rounds:
        data = json.loads(r.read_text())
        efficiency_history.append(data["efficiency_pct"])

    latest = efficiency_history[-1]

    # 条件1：达到 roofline 阈值（默认 95%）
    if config.early_stop and latest >= config.threshold_pct:
        return True, f"At roofline ({latest:.1f}% >= {config.threshold_pct}%)"

    # 条件2：连续收敛
    if len(efficiency_history) >= config.convergence_rounds:
        recent = efficiency_history[-config.convergence_rounds:]
        if max(recent) - min(recent) < config.min_improvement_pct:
            return True, f"Converged (Δ{max(recent)-min(recent):.2f}% < {config.min_improvement_pct}%)"

    return False, ""

在 SupervisedOptimizeAdapter.run_supervisor() 中，GateDecision.PASS_STOP 的触发条件从纯人工判断扩展为：人工判断 OR roofline check。

改动量估算：新增 ~80 行 Python + ~50 行 skill 指令 + 1 个 hardware_specs.md，不改变 optimize 的 round 目录结构契约。

P1：Beam Search 优化

KernelAgent 的实现

KernelAgent 的 beam search（triton_kernel_agent/opt_worker_component/searching/strategy/beam_search.py，~169 行）是一个干净的策略实现：

class BeamSearchStrategy(SearchStrategy):
    """每轮同时探索 N 个 top kernel × M 个瓶颈方向 = N×M 个 worker"""

    def __init__(self, num_top_kernels: int = 2, num_bottlenecks: int = 2):
        self.top_kernels: list[ProgramEntry] = []  # beam 中的 top-N kernel

    def initialize(self, initial_program: ProgramEntry):
        self.top_kernels = [initial_program] * self.num_top_kernels

    def select_candidates(self, round_num: int) -> list[dict]:
        """为每个 (kernel, bottleneck) 对生成一个候选任务"""
        candidates = []
        for rank, kernel in enumerate(self.top_kernels):
            for bottleneck_id in range(1, self.num_bottlenecks + 1):
                candidates.append({
                    "parent": kernel,
                    "bottleneck_id": bottleneck_id,
                    "kernel_rank": rank,
                })
        return candidates  # 共 N × M 个

    def update_with_results(self, results: list[dict], round_num: int):
        """合并新结果 → 按性能排序 → 保留 top-N"""
        new_entries = [r["program_entry"] for r in results if r.get("success")]
        all_candidates = self.top_kernels + new_entries
        all_candidates.sort(key=lambda x: x.metrics.time_ms)
        self.top_kernels = all_candidates[:self.num_top_kernels]

关键设计：bottleneck_id 驱动探索多样性。每个 worker 的 OptimizationOrchestrator 收到不同的 bottleneck_id 参数，LLM 瓶颈分析返回的是一组瓶颈排序列表，bottleneck_id=1 选第一个方向（最可能的瓶颈），bottleneck_id=2 选第二个方向（次可能），从而实现从不同角度同时攻击同一 kernel。

triton-agent 当前的状态

triton-agent 的 optimize 是严格的单路径：每个 round 只有一个优化方向。从 run_loop.py 可以看到，supervised 模式下每个 round 就是 worker→supervisor→gate decision 的单线流程。虽然有 optimize-batch 支持多 workspace 并行，但那是多个不同算子的独立优化，而非同一个算子的多方向探索。

不过基础设施已有雏形：

  • TRITON_AGENT_BATCH_NPU_DEVICES 已经支持多设备分配（如 0,3-5,8-9）
  • npu_affinity.py 实现了设备亲和性绑定
  • OptimizeRunLoop 已经有 supervised loop 的架构（worker + supervisor + gate）

引入方案：新增 BeamSearchOptimizeRunLoop

架构改变

并行化发生在 round 级别而非 loop 级别——每轮的 N×M 个 worker 并发执行，supervisor 汇总后 gate 决定继续或停止。

# 当前架构（单路径）
for round in 1..max_rounds:
    worker_result = agent.optimize(direction)
    gate = supervisor.check(round_dir)
    if gate == PASS_STOP: break

# Beam Search 架构（并行探索）
beam = [baseline_kernel] * K
for round in 1..max_rounds:
    candidates = []
    for kernel in beam:
        for bottleneck_id in 1..M:
            workers.append(spawn(agent.optimize(kernel, bottleneck_id)))
    results = await gather(workers)
    beam = merge_and_top_k(beam, results, K)
    if roofline_check(beam[0]): break

具体实现步骤

Step 1：新增 --strategy beam_search 参数和 BeamSearchConfig

在 src/triton_agent/optimize/models.py 中添加：

@dataclass(frozen=True)
class BeamSearchConfig:
    num_top_kernels: int = 2    # beam size K
    num_bottlenecks: int = 2    # 每个 kernel 探索 M 个方向
    max_beam_rounds: int = 5    # beam search 最大轮次

Step 2：新增 BeamSearchOptimizeAdapter in execution.py

复用现有 SupervisedOptimizeAdapter 作为单 worker 适配器，在外层包装 beam search 逻辑：

class BeamSearchOptimizeAdapter:
    def __init__(self, runner, artifacts_manager, config: BeamSearchConfig):
        self._single_adapter = SupervisedOptimizeAdapter(runner, ...)
        self._beam: list[tuple[Path, float]] = []  # (round_dir, time_ms)

    def run_round(self, request: AgentRequest) -> GateResult:
        """并行启动 N×M 个 worker，汇总后 gate"""
        candidates = self._select_candidates()
        results = parallel_map(self._run_single_worker, candidates,
                               device_affinity=self._assign_devices)
        self._update_beam(results)
        return self._gate_check()

    def _run_single_worker(self, candidate: BeamCandidate) -> WorkerResult:
        """每个 worker 在自己的 opt-round-N-worker-M/ 目录下运行"""
        worker_request = replace(request,
            workdir=request.workdir / f"opt-round-{candidate.round}-worker-{candidate.worker_id}",
            optimize_role="worker",
            bottleneck_hint=candidate.bottleneck_id,
        )
        return self._single_adapter.run_worker(worker_request)

Step 3：设备亲和性调度

利用现有 npu_affinity.py 和 TRITON_AGENT_BATCH_NPU_DEVICES：

  • 2×2 beam search = 4 个并发 worker，需要 4 个 NPU 设备
  • 每个 worker 绑定独立设备，避免 profiling 互相干扰
  • 如果设备不足 N×M，降级为轮询分批执行

Step 4：瓶颈多样性（bottleneck_id 映射）

triton-agent 的瓶颈分析是通过 skill 指令让 code agent 自主完成的，没有 KernelAgent 的 BottleneckAnalyzer 返回排序列表。需要调整 triton-npu-optimize skill，让 code agent 在分析阶段输出多个候选方向：

## bottleneck analysis（扩展）

After profiling, list the top M candidate bottleneck directions in priority order.
Write them to `opt-round-N/bottleneck_candidates.json`:
[
  {"rank": 1, "category": "memory", "root_cause": "...", "recommended_fix": "..."},
  {"rank": 2, "category": "compute", "root_cause": "...", "recommended_fix": "..."}
]

CLI 侧 beam search adapter 解析这个 JSON，把 rank=1 分给 worker-1，rank=2 分给 worker-2。

Step 5：Merge & Gate

每轮结束后，扫描所有 worker 的 benchmark 结果，按性能排序，保留 top-K 进入下一轮。Gate 逻辑增加 beam search 特有的条件：

  • top-1 kernel 达到 roofline → PASS_STOP
  • beam 中所有 kernel 连续两轮无改进 → PASS_STOP
  • 否则 → PASS_CONTINUE

改动量估算：新增 ~200 行 Python（BeamSearchOptimizeAdapter + BeamSearchConfig + merge/gate 逻辑）+ 并行调度 ~50 行 + skill 指令扩展 ~40 行 + bottleneck_candidates.json 契约。

P2：LLM Trace 完整持久化

问题

triton-agent 的 optimize 循环依赖 code agent（Codex/OpenCode）执行实际优化。当前 attempts.md 只记录高层决策（hypothesis、performance delta、pass/fail），不保存 code agent 的完整 prompt 和回复。这意味着：

  • 无法复现某一轮的具体优化过程
  • 无法分析 code agent 为什么做了某个优化决策
  • 无法从成功的优化中提炼出可复用的 pattern

KernelAgent 的做法

KernelAgent 的 OptimizationOrchestrator 每轮保存：

  • round{N:03d}_opt_prompt.txt — 完整的优化 prompt（包含 kernel 代码、性能数据、瓶颈分析结果）
  • round{N:03d}_opt_reply.txt — LLM 原始响应
  • round{N:03d}_strategy.json — 瓶颈分析的结构化结果
  • round{N:03d}_reflexion.txt — LLM 反思响应

这些文件全部存在 round artifact 目录下，可独立重放。

引入方案

这是优先级最低但实现最简单的改进——纯增量，不改变任何契约：

Step 1：在 opt-round-N/ 下新增 trace/ 子目录

不改变现有 contract.json 的 required fields，trace 是可选附加品：

opt-round-N/
  state.json          # 保持现有 contract
  benchmark_result.json
  perf-analysis.md
  trace/              # 新增：可选 trace 目录
    prompt.txt        # code agent 收到的完整 prompt
    response.md       # code agent 的响应（含 tool calls）
    session_id.txt    # agent session ID（用于关联到会话历史）

Step 2：在 optimize execution 中记录 trace

triton-agent 的 code agent 调用经过 RecoveryRunnerAdapter.run() 和 SupervisedOptimizeAdapter.run_worker()。在 AgentRequest 中新增一个可选的 trace_dir 参数，AgentRunner 如果支持 trace 模式，就在执行前保存 prompt、执行后保存 response。

@dataclass(frozen=True)
class AgentRequest:
    # ... existing fields ...
    trace_dir: Path | None = None  # 如果设置，agent backend 应保存 trace

Step 3：在 run_loop 中启用 trace

如果 --trace flag 被设置，每个 round 的 worker request 携带 trace_dir 指向当前 round 目录的 trace/ 子目录：

# 在 orchestration.py 中
if request.enable_trace:
    trace_dir = request.workdir / f"opt-round-{round_num}" / "trace"
    trace_dir.mkdir(parents=True, exist_ok=True)
    worker_request = replace(request, trace_dir=trace_dir)

Step 4（可选）：提炼为 pattern

当 trace 积累足够多后，可以定期扫描成功的 trace，让 code agent 提炼出新的 pattern card，自动补充到 triton-npu-optimize-knowledge 知识库。这是未来的增强方向。

改动量：~60 行 Python + AgentRequest 新增 1 个 optional 字段 + --trace CLI 参数

P3：AutoRouter — AST 级算子分类与自动路由

问题

triton-agent 当前要求用户手动选择命令：

  • convert：PyTorch → Triton NPU 转换
  • optimize：已有 Triton 算子的性能优化
  • gen-test：生成测试

用户需要自行判断该走哪条路径。对于新手使用者，这是障碍。

KernelAgent 的做法

Fuser/auto_agent.py 实现了三层路由决策：

第一层：AST 静态分析（~200 行，零 LLM 调用）

class Complexity:
    has_control_flow: bool
    has_attention_like: bool
    has_conv_transpose: bool
    has_group_norm: bool
    has_conv: bool
    chain_len_estimate: int

    def route_to_fuser(self) -> bool:
        # 触发 fuser pipeline 的条件（任一满足）：
        # 1. 包含 attention-like 模式（softmax + QK、multihead attention、einsum+bmm）
        #    BUT：如果是短链（≤3 步）且无 gn/transpose/control，走轻量路径
        # 2. conv_transpose2d 存在
        # 3. 显式控制流（if/for/while）
        # 4. group_norm + conv + 长链 ≥4 步
        # 否则走 KernelAgent 直接生成

第二层：LLM 辅助决策（可选）

当静态分析无法确定时，调用一个小型 LLM（浅层 reasoning），输出 JSON：{"route_strategy": "kernelagent", "confidence": 0.85}。支持四种策略：kernelagent、fuser、kernel_then_fuser（先轻量、失败则回退）、fuser_then_kernel。

第三层：缓存 + 自动回退

决策结果按文件 SHA256 缓存到 .fuse/router_cache.json。如果首选路径失败，自动尝试备选路径。

引入方案：新增 triton-npu-autoroute skill

Step 1：AST 分析器

复用 triton-agent 已有的 triton-npu-optimize/scripts/extract_code_facts.py（已经在 optimize 的 pattern triage 阶段用于代码结构分析），扩展它以输出算子分类：

# 扩展 extract_code_facts.py 输出中添加 classification
python3 scripts/extract_code_facts.py a.py --format json

# 新增输出字段：
{
  "code_facts": {...},        # 现有字段
  "classification": {          # 新增
    "operator_type": "matmul",  # elementwise | reduction | matmul | attention | custom
    "has_control_flow": false,
    "op_chain_depth": 3,
    "estimated_complexity": "medium",
    "suggested_action": "convert",   # convert | optimize | convert_then_optimize
    "relevant_patterns": ["classic-matmul", "tiling", "cache_use"]
  }
}

Step 2：新增 triton-npu-autoroute skill

# SKILL.md 核心流程
## Auto Route Decision

1. Run `python3 scripts/extract_code_facts.py  --format json` 
   on the input operator file.

2. Read `classification.suggested_action`:
   - `convert`: the operator hasn't been converted to Triton NPU yet → use triton-npu-convert-pytorch-operator
   - `optimize`: the operator is already Triton → use triton-npu-optimize  
   - `convert_then_optimize`: complex multi-op module → convert first, then optimize

3. If `classification.relevant_patterns` is non-empty, pre-read those 
   pattern cards before starting the actual conversion/optimization.

4. If static analysis confidence is low, escalate to LLM-based routing:
   - Send the operator code + static analysis results to the user's configured model
   - Ask for JSON: {"action": "convert|optimize", "confidence": 0.0-1.0, "reasoning": "..."}
   - Apply the decision

Step 3：CLI 集成

新增 triton-agent autoroute --input a.py 命令，自动选择最优路径并执行：

uv run triton-agent autoroute --input a.py --verify

等价于自动判定后执行 convert 或 optimize。

改动量：extend extract_code_facts.py ~40 行 + 新 skill SKILL.md ~80 行 + CLI 命令 ~30 行

P4：Subgraph 拆解 + 并行转换

问题

triton-agent 的 convert 命令一次只能处理一个算子文件。当面对一个包含多个子模块的 PyTorch module（例如 encoder-decoder 结构），用户必须手动拆解后分别调用 convert，再手动组合。这是最复杂但也最有挑战性的能力导入。

KernelAgent 的做法

KernelAgent 的 Fuser pipeline 是一个完整的子图拆解→并行生成→端到端组合流程：

Stage 1: Fuser Orchestrator（Fuser/orchestrator.py）

把 PyTorch module 的 __init__ 和 forward 重写为可融合模块，打包成 tar.gz 存档。

Stage 2: Subgraph Extractor（Fuser/subgraph_extractor.py，~414 行）

LLM 分析 fused code，输出 subgraphs.json：

[
  {
    "id": "sg_0",
    "description": "Conv2d 3→64, kernel=3, stride=2",
    "ops": ["conv2d", "relu"],
    "inputs": [{"name": "x", "shape": [1, 3, 224, 224], "dtype": "float32"}],
    "outputs": [{"shape": [1, 64, 112, 112], "dtype": "float32"}],
    "weights": {"weight": [64, 3, 3, 3], "bias": [64]}
  },
  {
    "id": "sg_1",
    "description": "BatchNorm2d 64 + ReLU",
    "ops": ["batch_norm", "relu"],
    "inputs": [{"shape": [1, 64, 112, 112], "dtype": "float32"}],
    ...
  }
]

关键设计：按 shape signature 去重。如果两个子图的输入/输出/权重 shape 完全一致，只生成一次 kernel，减少重复工作。

Stage 3: Dispatcher（Fuser/dispatch_kernel_agent.py，~511 行）

把每个 unique subgraph 转换为具体的 Triton kernel 生成任务，使用 concurrent.futures 并行调用 TritonKernelAgent：

with concurrent.futures.ProcessPoolExecutor(max_workers=jobs) as pool:
    futures = {pool.submit(_run_one, item, ...): item["id"] for item in items}
    for fut in concurrent.futures.as_completed(futures):
        results.append(fut.result())

Stage 4: Composer（Fuser/compose_end_to_end.py，~537 行）

LLM 把验证通过的 kernel 拼回一个完整的端到端 Triton 程序：

  • 输入：原始 PyTorch problem + subgraphs.json + 各 subgraph 的 kernel + summary.json
  • 输出：单个 composed_kernel.py，包含所有 @triton.jit kernel + kernel_function(...) wrapper + 自测试
  • 可选 --verify：生成后立即执行自测试验证数值正确性

引入方案：分阶段实施

P4 改动量大，建议分两个 milestone 实施：

Milestone 1：Subgraph 拆解 + 并行 convert（最小可行方案）

Step 1：新增 triton-npu-convert-module skill

## Module Decomposition

1. Read the multi-op PyTorch module file.
2. Analyze the `forward()` method and `__init__()` to identify distinct 
   computational subgraphs.
3. For each subgraph, record:
   - Operations and their order
   - Input/output shapes and dtypes
   - Parameter tensors (from __init__)
4. Output `subgraphs.json`:
   [
     {"id": "sg_0", "ops": [...], "inputs": [...], "outputs": [...], "weights": {...}},
     ...
   ]
5. Deduplicate by shape signature: if two subgraphs have identical 
   (input_shapes, output_shapes, weight_shapes), keep only one.
6. For each unique subgraph, spawn a parallel convert worker using 
   the existing `triton-npu-convert-pytorch-operator` skill.
7. Workers run concurrently, each in its own `subgraphs/sg_N/` directory.
8. Compose verified kernels into `composed_kernel.py`.

Step 2：CLI 并行调度

复用 convert-batch 的并发架构和 TRITON_AGENT_BATCH_NPU_DEVICES 设备分配：

uv run triton-agent convert-module --input model.py --max-concurrency 4

Step 3：Compose 阶段

生成 composed_kernel.py：把 subgraph 描述 + 各 kernel 代码喂给 code agent，让它生成一个统一的 Triton module，包含所有 @triton.jit kernel 和调用逻辑。

Milestone 2：Compose 验证 + 自动优化（后续）

Compose 完成后自动运行 correctness test 和 benchmark——这部分可以完全复用现有的 triton-npu-gen-test + triton-npu-run-eval skill。

改动量：新 skill ~150 行 + CLI 协调器 ~100 行 + 并行调度复用现有批量架构 ~30 行。Compose 阶段是最不稳定的部分，依赖 code agent 对多 kernel 拼接的正确理解。

风险警告

  • LLM 的子图拆解质量：KernelAgent 用 GPT-5/o4-mini 直接调 API，可控性高。triton-agent 依赖 code agent 做拆解，code agent 本身的决策质量会影响拆解结果。
  • Compose 正确性：把多个 Triton kernel 拼成一个正确的端到端程序是 LLM 最容易出错的步骤——shape 对齐、内存布局、调用顺序都可能出错。
  • 建议：先从简单场景开始（sequential operators, known shapes），逐步扩展到动态 shape 和复杂拓扑。

能力导入总览

优先级能力改动量新文件破坏现有契约？核心价值
P0Roofline 自动终止~150 行hardware_specs.md + roofline.json 契约否让优化循环有客观终止条件
P1Beam Search~300 行BeamSearchOptimizeAdapter + bottleneck_candidates.json否（复用 supervised loop）多方向并行探索，大幅提升优化效率
P2LLM Trace 持久化~60 行trace/ 子目录否（增量可选）可复现、可调试、可提炼 pattern
P3AutoRouter AST 路由~150 行triton-npu-autoroute skill + 扩展 extract_code_facts.py否降低使用门槛，自动选择最优路径
P4Subgraph 拆解~300 行triton-npu-convert-module skill + subgraphs.json 契约否（新命令）解决 multi-op module 的真实痛点

相关仓库