Stop hand-tuning kernels: How Neuron Agentic Development accelerates AWS Trainium optimizations

停止手动调整内核：Neuron Agentic Development 如何加速 AWS Trainium 优化 | 人工智能

停止手动调整内核：Neuron Agentic Development 如何加速 AWS Trainium 优化

随着前沿人工智能模型在规模和复杂性上的增长，开发者在每个硬件平台上都面临着一个共同的挑战：如何从模型运行的硅芯片中提取最大性能和效率。无论是为世界模型提供实时体验，还是在智能代理工作流程中支持更深层次的推理，或是在大规模上减少推理成本，硬件理论上能够提供的性能与大多数团队实际实现的性能之间仍存在显著差距。历史上，自定义内核开发是弥补这一差距的途径，但它需要深厚的架构专业知识、手动分析工作流程以及迭代优化周期，而这些是大多数团队难以承担的。

这并不一定非要如此。如果每个机器学习（ML）工程师都能像性能工程师一样工作，编写硬件感知的内核、诊断瓶颈并交付优化模型，而无需多年芯片级经验，那会怎样？如果已经在一个架构上熟练的开发者能够在几天内而不是几个月内掌握另一个架构，那又会怎样？

今天，我们宣布推出 Neuron Agentic Development 功能：这是一组人工智能代理和技能，使在 AWS Trainium 和 AWS Inferentia 上构建的开发者能够实现这些目标。首批功能使 Kiro 和 Claude 中的编码代理能够编写、调试和分析 Neuron Kernel Interface（NKI）内核，将机器学习性能工程扩展到团队中的每个开发者。来自其他架构的内核开发者可以快速扩展到 Trainium，团队可以缩短从想法到硬件优化实现的时间，而曾经限制内核开发的深厚架构知识现在通过智能工具链变得触手可及，这些工具链在每个步骤中都会引导开发者。

在本文中，我们将解释 Neuron Agentic Development 功能如何加速内核开发工作流程。

Neuron Agentic Development 技能

Neuron Agentic Development 工具包提供了五种专门的技能，这些技能遵循自然的内核开发流程：编写 → 调试 → 分析性能 → 分析。您可以单独调用这些技能以执行特定任务，或者通过 neuron-nki-agent 将它们串联在一起，该代理会根据您的请求自动选择合适的流程。要使用这些技能，请将它们添加到您的智能 IDE 的技能目录中。例如，在任何 IDE（如 VS Code、Cursor 或 Kiro）中，将技能添加到 .kiro/skills 或 .claude/skills 目录，并使它们对您的代理可用。这些技能必须在基于 Trainium 的 Amazon Elastic Compute Cloud（Amazon EC2）实例上运行。

内核编写

neuron-nki-writing 技能是您创建 NKI 内核的起点。它可以将 PyTorch、NumPy 或自然语言描述转换为正确的 NKI 代码。例如，它涵盖了尊重硬件约束的分块策略（如 128 的分区维度和 512/4096 的 PSUM 自由维度）、内存访问模式、带有显式 dst 参数的计算操作，以及针对 DMA 尺寸和 SBUF 重用的效率指南。该技能会根据任务的复杂性对任务进行分类，并仅加载所需的参考资料。

调试

neuron-nki-debugging 技能提供了一种系统的工作流程，用于解决在 Trainium 和 Inferentia 硬件上进行 NKI 编译和执行时出现的错误。例如，它涵盖了使用正确的 --target 标志进行环境设置、通过所有 28 个 NCC 错误代码的分类索引进行编译器错误解决，以及与 CPU 计算的参考值进行数值验证。

性能分析和剖析

neuron-nki-profiling 技能在硬件上捕获执行剖析信息。它配置运行时检查环境变量，运行内核，识别正确的 Neuron Execution File Format (NEFF)，并使用 neuron-explorer 捕获跟踪信息，包括 DGE（DMA 图形引擎）通知，以提供 DMA 级别的详细信息。它提取 JSON 指标并生成 neuron-nki-profile-querying 使用的 NEFF 文件。

neuron-nki-profile-querying 技能读取 NEFF 和 NTFF 文件并运行 SQL 查询，以计算性能边界、识别瓶颈引擎，并将低效问题定位到特定的 NKI 源代码行。它支持三种分析方法：通过 neuron-explorer API 服务器、直接在 parquet 上使用 DuckDB，或使用 pandas 进行自定义计算。

文档

neuron-nki-docs 技能在整个开发过程中使用。在编写过程中，它提供 API 签名和教程。在调试过程中，它解释错误代码。在剖析过程中，它澄清硬件架构的细节。可以查询特定的 nisa.* 或 nl.* API，查找错误代码，找到教程，或浏览 Trainium 1、2 和 3 的架构指南。

代理

虽然技能为单个任务提供了构建模块，但代理将多个技能组合成自主的工作流程。每个代理是一个专门的人员，能够端到端地处理多步骤的开发场景。

• neuron-nki-agent 是 NKI 开发的统一入口点。它根据您的请求（编写、调试、剖析或文档查询）自动选择合适的工作流程，并协调适当的技能。这是默认的起点。

• neuron-nki-writing-agent 专注于内核编写。它将 PyTorch、NumPy 或自然语言描述转换为 NKI 代码，并处理现有内核的修改。

• neuron-nki-debugging-agent 自主解决编译器错误，通过分析错误、在文档中搜索修复方法并应用更正。它跟踪迭代（最多 10 次），并在遇到困难时逐步简化。

• neuron-nki-docs-agent 是一个轻量级的文档导航器，用于 API 签名、错误代码解释、教程和架构细节。

• neuron-nki-profile-analysis-agent 运行两个独立的技能以识别性能瓶颈。它使用 neuron-nki-profile 技能在硬件上捕获执行剖析信息：设置环境变量、运行内核、识别 NEFF，并运行 neuron-explorer 捕获以生成剖析 parquet 文件。然后，它使用 neuron-nki-profile-querying 技能对这些 parquet 文件运行 SQL 查询，以计算性能边界、识别瓶颈引擎，并将低效问题定位到特定的 NKI 源代码行。

实践应用：优化自定义 softmax 内核

以下的逐步操作展示了这些代理能力如何在实际中协同工作。您将探索两个内核：一个 softmax 内核（步骤 1 和 2）和一个 SwiGLU 内核（步骤 3 和 4），这将展示在实际工作负载上的剖析过程。

假设你有一个 PyTorch 的 softmax 操作，它是推理流水线中的瓶颈，你希望编写一个自定义的 NKI 内核，将其与前面的 scale 操作融合。

步骤 0：设置你的实例和环境

要开始使用：

• 通过 AWS MLCBs 使用 AWS Neuron 深度学习 AMI（DLAMI）启动一个 trn2.3xlarge 实例。此处以 São Paulo（sa-east-1）和 Melbourne（ap-southeast-4）作为示例 AWS 区域。有关其他支持区域的完整 Trainium 可用性列表，请参阅相关文档。

• 使用 SSH 连接到该实例。

• 安装 Kiro：curl -fsSL https://cli.kiro.dev/install | bash

• 按照 neuron-agentic-development 仓库中的说明安装 Neuron Agentic 开发技能。

注意：运行 trn2.3xlarge 实例时每小时会产生费用。完成本教程后，请记得终止实例以避免持续费用。

有关更详细的实例设置和配置说明，请参阅 Neuron DLAMI 设置指南。

从远程终端验证 neuron 设备是否可见：

# 确认 neuron 设备是否可见
neuron-ls

# 确认 neuron-explorer 是否可用
which neuron-explorer && neuron-explorer --version

DLAMI 中预装了一个虚拟环境，路径如下：

~opt/aws_neuronx_venv_pytorch_2_9

使用以下命令激活它：

source ~opt/aws_neuronx_venv_pytorch_2_9/bin/activate

环境设置完成后，可以通过以下命令开始开发内核：

kiro-cli --agent neuron-nki-agent

步骤 1：编写内核

在交互式 Kiro CLI 会话中，输入以下提示：“编写一个 NKI 内核，计算缩放后的 softmax：在最后一个维度上计算 softmax(x * scale)，输入形状为 [batch, seq_len, hidden_dim]，使用 bfloat16。”

代理使用 nisa.activation(np.exp, ...) 来进行硬件加速的 exp 操作，使用 float32 累加以确保数值稳定性，并在自由维度上进行适当的分块，生成一个完整的三阶段内核（行最大值、exp 总和、归一化）。它解释了其设计决策：每个行一个程序实例，P_MAX=128（与 128 分区的硬件限制匹配），F_MAX=2048（与 Trainium 上 2048 元素的自由维度限制匹配），以及 bfloat16 输出转换。

图 1：NKI 代理编写内核。

步骤 2：在硬件上进行调试

请代理运行内核并验证与 PyTorch 参考的数值一致性。

代理立即遇到一个问题：nisa.tensor_tensor 不会自动广播归约结果，因此每行的最大值和总和值无法直接应用于整个隐藏维度。代理查阅了 NKI 参考模式，识别出正确的广播机制（通过 .ap() 的 stride-0 访问视图），并据此重写内核。

在将修复后的内核同步到实例并在设备上运行后：

PASS: shape=(2, 128, 512), max_diff=0.000008
PASS: shape=(4, 256, 1024), max_diff=0.000004
PASS: shape=(1, 1, 64), max_diff=0.000061
PASS: shape=(2, 300, 768), max_diff=0.000007

所有测试通过。

所有四个案例都通过了，最大误差远在 bfloat16 容忍范围内，确认内核在真实的 Trainium 硬件上数值正确。

图 2：NKI 代理调试其错误。

步骤 3：分析内核执行的性能

在内核编译并产生数值正确的结果后，下一步是在硬件上对执行情况进行分析，以识别性能瓶颈并指导优化。

为了在实际工作负载上展示性能分析和分析，这一步使用了 SwiGLU MLP 内核，这是大型语言模型（LLMs）中常见的模块。

将代理指向 SwiGLU 内核，并让它分析性能数据。代理首先将内核编译为 NEFF，并通过 neuron-explorer 捕获 NTFF 跟踪记录。然后，它对内核进行两部分的调查，首先查看内核级别的统计信息和性能边界，然后通过在指令执行级别查询性能数据，深入分析具体的低效之处。

首先，代理对捕获的性能数据进行全面的边界分析，并发现多个值得进一步调查的性能缺口：

图 3：NKI 代理提取内核的摘要统计信息并计算性能边界。

它发现多个值得进一步调查的性能缺口。Tensor Engine（TE）主导了执行过程，但效率较低。此外，它还存在较大的空闲时间，这表明可能值得调查其最可能的依赖项（DMA 引擎），在该引擎中我们可以看到既冗余又低效的工作。

图 4：NKI 代理调查性能数据中的低效之处并提供分析。

这些调查帮助我们审核性能缺口并确定优先优化的方向。尽管瓶颈引擎（Tensor Engine）的低效性本应是优化的首要目标，但代理发现 NKI 的 matmul 指令已经接近其峰值效率。相比之下，我们发现 DMA 指令远低于目标大小（低效），并且我们还在重复加载所有输入八次（冗余）。我们甚至找到了导致次优数据传输的三条确切的 NKI 代码行。解决这些代码行可能会减少 TE 的空闲时间并提高内核的延迟。

需要了解的内容

在使用 Neuron Agentic Development 技能和代理时，请记住以下注意事项。

• 性能分析和调试技能需要在实际的 Trainium 或 Inferentia 实例上执行。

• 编写和文档技能可以在任何地方使用。

• 所有技能都针对当前的 NKI API（在发布时为 0.4.0）。技能支持 Trainium1（gen2）、Trainium2（gen3）和 Trainium 3（gen4），使用适当的 --target 标志。

• 技能和代理是协同设计的。顶级代理会根据需要自动调用性能分析和调试技能。

清理

为了避免持续产生费用，请终止在步骤 0 中创建的 trn2.3xlarge 实例。你可以通过 AWS 管理控制台（EC2 > 实例，选择实例，然后选择 实例状态 > 终止）来完成此操作，或者运行以下命令：

aws ec2 terminate-instances --instance-ids <your-instance-id>

在关闭控制台之前，请确认实例状态显示为“已终止”。

下一步

内核编写和性能分析技能降低了在 Trainium 上编写高性能内核的门槛，但它们只是更广泛愿景的第一部分。

目前，开发人员使用性能分析结果来指导下一轮内核编辑。这个迭代周期（分析、诊断、重构、重新分析）是耗时最多的部分。我们希望使这个循环完全自动化。例如，代理可以自主迭代内核，直到其达到性能目标，而无需开发人员解释每个分析结果并手动编写下一个修复方案。

我们也从性能开发人员那里了解到，自定义内核只是更广泛挑战的一部分。开发人员希望他们的模型能够在Trainium上运行，而无需担心模型代码和语法的移植、解决操作符的差距、应用模型级别的优化，以及在大规模上验证正确性。我们希望将同样的智能代理方法应用于这一更广泛的问题。

总而言之，我们的愿景是利用Trainium和Neuron SDK支持前沿模型的下一轮创新，并利用Neuron智能代理开发工具套件，实现从新模型架构的实验到大规模运行生产模型的各种用例的领先性价比。

随着这些能力的逐步成熟，我们将分享更多内容。要立即开始使用当前可用的功能，请访问Neuron智能代理开发的GitHub仓库。

与我们一同构建

Neuron智能代理开发功能现已开放。立即开始：克隆neuron-agentic-development仓库，并在几分钟内编写你的第一个NKI内核。

以下是开始使用的步骤：

• 从neuron-nki-agent开始。它会根据你的请求选择合适的流程，为你提供端到端的完全自主体验。

• 运行技能示例。直接调用单个技能（例如，/neuron-nki-writing）来完成特定任务，或者在你的内核产生正确结果后，链式调用/neuron-nki-profiling和/neuron-nki-profile-querying。

• 如果你遇到问题或有想法，请在GitHub上提交问题。我们正与社区一起积极开发，并将尽快回复你。

• 反馈贡献。提交PR，分享你构建的内核，并帮助我们让这些工具变得更好。

我们公开构建这些功能，因为最好的开发工具是由使用它们的开发者塑造的。与我们一同构建吧。

关于作者

'"`