GLM-5.3-Flash 在 AWS Trainium2 上的 NKI 层级 Megakernels 优化揭示了解码步骤的时间去向

作者为 GLM-5.3-Flash（320B MoE，1M-token 上下文）在 AWS Trainium2 上构建了推理引擎并开源代码，实测显示单个 decode 步骤耗时 24.2 ms，其中模型实际数学运算仅占极小部分，大部分时间消耗在内存交互、XLA 胶水代码及跨芯片通信上。该工作将每层注意力机制重写为单一 NKI 内核，遵循三个核心规则：保持工作集在片上以利用 28 MiB SBUF 减少 HBM 访问，同时激活物理核心的双核以分担权重流，并将阻塞后续操作的集合操作移入内核内部实现重叠执行。

性能分析表明，从 10 月 17 日到 11 月 1 日，batch-8 的 decode 步骤从 32.2 ms 降至 24.2 ms，主要归功于稀疏注意力层的融合（节省 3.2 ms）以及 KDA 和 MoE 内核的联合发布（再省 3.2 ms）。剩余时间中，6.1 ms 为内存下限，10.0 ms 为向量引擎指令链，4.6 ms 为内核内等待，2.8 ms 为暴露的集合操作。测试发现编译器已自动重叠的融合操作反而收益甚微，且孤立基准测试具有误导性，例如某 Dense FFN 内核单独运行快但放入完整步骤后反而拖慢整体。此外，NKI 模拟器存在准确性偏差，三个通过模拟的内核在设备上产生错误答案，最终精度误差控制在上一构建版本的 1.05 倍以内。

在 SemiAnalysis AgentX 基准测试中，单台设备达到每分钟 28.9 个代理请求，16 个并发时中位解码速度为 35.9 tokens/s。事实层面确认了 NKI 内核设计对特定硬件架构的显著优化效果；推断层面指出在缺乏 GPU 对比基线的情况下，该方案展示了针对非 GPU 硬件进行极致优化的可行性路径；猜测层面认为若其他厂商能复现此类细粒度内核优化策略，可能改变当前大模型推理成本结构，但目前尚无公开证据支持通用性结论。