单机 Trainium2 跑 320B GLM 模型：NKI 巨核推理引擎开源
WildCard727 · reddit · 2026-11-05
RunAnywhere 团队为 AWS Trainium2 写了一个开源推理引擎 trnrt，在单台 trn2.48xlarge 上服务 320B MoE、1M 上下文的 GLM-5.3-Flash，并公开了完整 profile 数据。
核心发现：decode 时间大头不是模型数学
- 早期版本中，batch-1 下某 KDA（线性注意力）层的 recurrence 仅占 attention 部分的 2%；其余是 mHC 预投影（36%）、XLA 胶水（33%）、转置（10%）、stream mix（10%）等小算子
- batch-8 decode 一步 24.2ms，对比 6.1ms 的内存下限；瓶颈是每条约 153ns 固定开销的小向量指令链、每个 kernel 边界都写回 HBM、每步 125 次跨芯片通信
三条 kernel 重写规则
- 工作集留在片上：每个 NeuronCore 有 28MiB SBUF，只让层输入输出和流式权重过 HBM
- 用满逻辑核的两个物理核：权重流需分到两核，合计 672–708 GB/s
- 把阻塞下一步的 collective 移进 kernel 内部，让独立工作与其重叠
效果与踩坑
- batch-8 decode 从 32.2ms 降到 24.2ms；融合稀疏注意力层省 3.2ms，KDA+MoE kernel 再省 3.2ms
- 失败教训：融合编译器已重叠的接缝（预估省 2-4ms 实际只省 0.85ms）、孤立 benchmark 表现好的 kernel 反而让整步变慢、simulator 通过但真机出错（PSUM 切片累加、kernel 内 collective 缺 fence、共享 HBM 读写竞争）
- 每个 kernel 逐字节校验或对照高精度参考，误差上限设为上一版 1.05 倍；稀疏注意力 prefill 保留一处近似，影响 1.2% 的选择
端到端在 SemiAnalysis AgentX 基准上自测达单机 28.9 agent 请求/分钟，16 agents 时中位解码 35.9 tok/s。无 GPU 对比基线：vLLM Neuron 插件不注册 GLM，NxD Inference 已处维护模式。代码开源于 GitHub，docs/RESULTS.md 含复现命令。
所属事件:RunAnywhere 开源 Trainium 推理引擎 TrnRT，单机跑 320B GLM 模型(2 条相关)→
「Infra」频道最新
- AI 算力生意暗藏 WeWork 式风险:长债短收的结构性错配 — bookwormengr · 2026-11-05
- SK 海力士:存储供不应求将持续到 2030 年 — Beth_Kindig · 2026-11-05
- Macrocosmos 发布首个公开模型 IOTA Decide 4B：全球分布式 GPU 训练、Apache-2.0 开源 — markjeffrey · 2026-11-05
- 「焦耳换 GDP」框架:最高价值的 AI 任务未必最省能 — arena · 2026-11-05
- 开发者在 Mac Studio 上把 Meta Muse Glimmer 30B 提速最高 21% — TimDarcet · 2026-11-05
- Mach 2 Medium 量化 Qwen 3.8：仅损 1.5% 省 10.9GB，超 Unsloth 同档 — pbaylies · 2026-11-05