NKI-Agent：特定领域的 微调 和用于神经元内核生成的代理工具

最近基于 LLM 的内核生成的代理方法在 CUDA 上取得了令人印象深刻的结果。对于 AWS Trainium 和 Inferentia 等新兴人工智能加速器来说，自动化内核生成和优化在很大程度上仍未得到解决。通过Neuron Kernel Interface (NKI) 为这些芯片编写内核尤其具有挑战性：开发人员必须驾驭多引擎架构、基于图块的编程以及跨多级内存层次结构的显式数据移动。此外，该领域不存在公开可用的训练数据、基准或工具增强代理。我们推出了 NKI-Agent，这是第一个将特定领域的监督 微调 (SFT) 与用于 NKI 内核生成的编译-验证-修复代理循环相结合的系统。我们将现有的 CUDA-Agent 框架适应 Neuron 硬件，策划 6,000 个 NKI 内核生成任务进行训练，并构建 NKIBench，这是一个跨三个难度级别的 250 个任务基准测试。在真实的 Trn1 硬件上进行评估，配备 Claude Opus 4.8 和rank感知系统提示的 NKI-Agent 在 150 个任务的 NKIBench 上实现了 77.3% 的通过率。我们表明工具的使用至关重要：Opus 4.8 在没有代理工具的单次模式下得分为 6%。在 60 个任务子集中，我们表明经过 SFT 训练的 Qwen3-Coder-30B-A3B 以 1/100 的成本实现了 25.0% 的通过率，优于 Claude Sonnet 4 (15.0%)。我们还报告说，具有二元编译成功奖励的组相对策略优化 (GRPO) 未能比 SFT 有所改进，这为基于 RL 的内核生成的奖励设计提供了指导。