iterative-kernel-optimization-loop

# 迭代式 Kernel 优化循环 ## 概览 这个技能是整个系统的编排层，它把其他技能串成自动化、可迭代的优化循环。 它把 `awesome_kernel_skills` 从静态知识集合，变成一个类似 [AutoKernel](https://github.com/RightNow-AI/autokernel) `program.md` 的智能体驱动优化引擎。 智能体（Cursor / Claude Code / Codex）就是执行引擎， 而这个文件定义了它该做什么、先后顺序是什么，以及什么时候停止。 ## 优化闭环 ```text ┌──────────────────────────────┐ │ 输入 │ │ - kernel 源码 (.py) │ │ - 目标指标与阈值 │ │ - 目标 GPU（nvidia/amd） │ └──────────────┬───────────────┘ ▼ ┌─── 阶段 0：BASELINE ────────────────┐ │ 运行 test_*.py，记录基线性能 │ │ 如果没有测试，先创建测试 │ └──────────────┬──────────────────────┘ ▼ ┌────────────────────────────────────────────┐ │ 阶段 1：PROFILE │ ┌───▶│ 技能：profiling/SKILL.md │ │ │ 动作：运行 NCU 或 rocprof │ │ │ 输出：metrics.json │ │ └──────────────┬─────────────────────────────┘ │ ▼ │ ┌────────────────────────────────────────────┐ │ │ 阶段 2：DIAGNOSE │ │ │ 技能：bottleneck-diagnosis/SKILL.md │ │ │ 动作：分类 memory/compute/latency │ │ │ 输出：瓶颈类型 + tier 优先级 │ │ └──────────────┬─────────────────────────────┘ │ ▼ │ ┌────────────────────────────────────────────┐ │ │ 阶段 3：OPTIMIZE │ │ │ 技能：按优先级使用 tier1~5 + kernel skill │ │ │ 动作：一次只应用一个优化点 │ │ │ 规则：改动小、可回退 │ │ └──────────────┬─────────────────────────────┘ │ ▼ │ ┌────────────────────────────────────────────┐ │ │ 阶段 4：VERIFY │ │ │ 技能：verification/SKILL.md │ │ │ 动作：运行测试，检查正确性 │ │ │ 若失败：回退改动，尝试下一个 tier │ │ └──────────────┬─────────────────────────────┘ │ ▼ │ ┌────────────────────────────────────────────┐ │ │ 阶段 5：BENCHMARK │ │ │ 技能：benchmark/SKILL.md │ │ │ 动作：测量 TFLOPS 或 GBps │ │ │ 并与上一轮结果对比 │ │ └──────────────┬─────────────────────────────┘ │ ▼ │ ┌────────────────────────────────────────────┐ │ │ 阶段 6：DECIDE │ │ │ - 达标？ → 结束并记录 │ │ │ - 有提升？ → 记录并继续 │ │ │ - 有回退？ → 回滚并试下一项 │ │ │ - tier 用尽？ → 结束并保留最佳结果 │ │ │ - 到达最大轮数？ → 结束并保留最佳结果 │ │ └──────────────┬─────────────────────────────┘ │ │ 未结束 └───────────────────┘ ``` ## 配置项 在进入循环前，先确定这些参数： ```yaml kernel_path: skills/kernels/softmax/triton_template.py test_path: skills/kernels/softmax/test_softmax.py target_metric: gbps # "tflops" or "gbps" target_value: 800 # 目标 GBps（或 TFLOPS） target_pct_of_peak: 0.70 # 另一种写法：达到峰值的 70% max_iterations: 5 # 最多进行 N 轮优化 gpu: auto # "nvidia"、"amd" 或 "auto"（自动检测） ``` ## 分阶段说明 ### 阶段 0：基线 1. 通过 `scripts/detect_gpu.py` 检测 GPU 2. 运行 `test_*.py --json`，获取基线正确性和性能 3. 将结果记录进优化日志 ```markdown ## Iteration 0 (Baseline) - GPU: NVIDIA RTX 4090 - Correctness: PASS - GBps: 1747.0 - Target: 2000 GBps (70% of peak) - Gap: 253 GBps (12.6%) ``` ### 阶段 1：性能分析 阅读 `profiling/SKILL.md` 并执行： - NVIDIA: `ncu --set full -o iter{N}_profile python test_*.py` - AMD: `rocprof --stats python test_*.py` 把关键指标提取成结构化结果。 ### 阶段 2：瓶颈诊断 阅读 `bottleneck-diagnosis/SKILL.md` 并按流程图判断： ```text IF dram_bw_pct high AND compute_pct low → memory-bound → Tier 1-2, then 3 ELIF compute_pct high → compute-bound → Tier 3, then 1, 5 ELIF both low → latency-bound → Tier 4, then 3 ``` 输出一个 tier 优先级列表，例如 `[tier2, tier1, tier3]`。 ### 阶段 3：优化 针对最高优先级且尚未尝试过的 tier： 1. 阅读对应的 `tier{N}/SKILL.md` 2. 阅读对应算子的 `SKILL.md` 3. 只应用一个聚焦的改动 4. 常见改动如下： | Tier | 常见改动 | |------|---------------| | 1 | 调整 `BLOCK_SIZE`、`num_warps`、`GROUP_SIZE_M` | | 2 | 改善 coalescing、增加向量化加载、修复 bank conflict | | 3 | 融合 epilogue、切换 fast math、修正 accumulator dtype | | 4 | 改成 persistent kernel，或加入 Split-K | | 5 | 使用架构特定能力（`TMA`、`WGMMA`、MFMA 友好形状） | **规则**：修改前先保留副本，确保可以回退。 ### 阶段 4：验证 运行 `test_*.py`，只做正确性检查： - **PASS**：继续做 benchmark - **FAIL**：回退改动，把该 tier+change 标记为已尝试且失败，再选下一项 ### 阶段 5：基准测试 运行 `test_*.py --json` 做完整 benchmark： - 记录 TFLOPS 或 GBps - 和上一轮对比 ### 阶段 6：决策 ```python if current_metric >= target_value: status = "TARGET_MET" elif current_metric > best_metric: best_metric = current_metric best_kernel = save_snapshot() status = "IMPROVED" elif current_metric <= prev_metric: revert_to_previous() mark_tier_change_as_tried() status = "REGRESSED" if iteration >= max_iterations or all_tiers_exhausted(): restore_best_kernel() status = "MAX_ITERATIONS" ``` ## 优化日志格式 每轮结束后，把结果追加到 `experiments/README.md`： ```markdown ## Optimization: {kernel_name} on {gpu} | Iter | Action | Tier | Correctness | Metric | Delta | Status | |------|--------|------|-------------|--------|-------|--------| | 0 | baseline | - | PASS | 1747 GBps | - | - | | 1 | increase BLOCK_SIZE 4096→8192 | T1 | PASS | 1820 GBps | +4.2% | improved | | 2 | vectorized loads (float4) | T2 | PASS | 1955 GBps | +7.4% | improved | | 3 | multi-row=8 | T1 | FAIL | - | - | reverted | | 4 | fuse with layernorm | T3 | PASS | 2010 GBps | +2.8% | TARGET_MET | Final: 2010 GBps (70.2% of peak), 4 iterations, target met. ``` ## 约束规则 1. **每轮只改一个点**：不要同时做两个优化，不然无法归因 2. **先验证再测性能**：错误的 kernel 不准做 benchmark 3. **性能回退就回滚**：除非某个改动是为更大的下一步铺路，否则不要保留变慢的版本，并且要写清楚原因 4. **尊重平台差异**：使用架构特性前先看 `tier5-arch-nvidia` 或 `tier5-arch-amd-rocm` 5. **默认最多 5 轮**：如果 5 轮之后还卡住，问题很可能已经不是调参，而是算法层面的限制 6. **完整记录**：每一轮都要记录，避免以后重复试错 ## 当无法使用 Profiling 工具时 如果 NCU / rocprof 无法运行，就退化为： 1. 只用 `test_*.py --json` 的 benchmark 数据做判断 2. 从该 kernel 自己的 `SKILL.md` 推断它的瓶颈类型 3. 按对应的 tier 优先级执行 4. 优先从 Tier 1 的 autotune 配置开始，因为这是最安全的入口 ## Agent 指令 当收到“迭代优化一个 kernel”的任务时： 1. 先读这个 skill 2. 根据用户要求或 GPU 峰值规格确定目标 3. 先跑阶段 0 获取基线 4. 按 Profile → Diagnose → Optimize → Verify → Benchmark → Decide 的闭环执行 5. 每轮结束都更新优化日志 6. 达标、达到最大轮数、或可行 tier 全部用尽时停止 7. 最后用总结表报告结果 ## 与其他技能的关系 ```text optimize-loop（本技能）─── 负责编排 ───▶ 其他所有技能 │ ├─ Phase 1 → profiling/SKILL.md ├─ Phase 2 → bottleneck-diagnosis/SKILL.md ├─ Phase 3 → tier1~5/SKILL.md + kernels/*/SKILL.md ├─ Phase 4 → verification/SKILL.md └─ Phase 5 → benchmark/SKILL.md ``` ## 参考资料 - [AutoKernel program.md](https://github.com/RightNow-AI/autokernel) —— 原始的智能体驱动 kernel 迭代优化循环 - [KernelSkill dual-level memory](https://arxiv.org/abs/2603.10085) —— 可复用的优化技能与短期回溯记忆机制