trnRT 用单个 AWS Trainium2 实例运行 GLM-5.3-Flash，解码提速至 24.2ms

回到博客trnRT · 2026年11月4日 · 34分钟阅读 一台AWS trn2.48xlarge在SemiAnalysis的AgentX基准测试上以每分钟最高28.9个agent请求的速度服务GLM-5.3-Flash，每分钟处理357万个token，并将百万token的提示在约三分钟内产出首个token。NKI层megakernel将其解码步从32.2毫秒缩短至24.2毫秒。 GLM-5.3-Flash是一个3200亿参数、百万token的MoE模型。我们在一台AWS trn2.48xlarge上对其进行服务，并在SemiAnalysis的AgentX基准测试上测量，该基准重现了真实编码agent会话的形态：冗长且大部分已缓存的提示、子agent，以及等待每个答案的用户。我们没有找到此前在Trainium上发布的AgentX结果。为了实现这一目标，我们将模型的解码重写为每半个层一个NKI kernel，对于稀疏注意力层则是每整个层一个kernel，并构建了一个前缀缓存，将每个agent的上下文保留在设备和主机内存中。以下是这样做带来的收益、在AWS上的成本、时间仍然花在哪里，以及哪些尝试没有奏效。 整篇博文浓缩在一张图里：随着agent增加，一台机器在AgentX上能服务多少。最多16个agent时，请求中位数在约一秒内开始；到32个时，请求需要排队等待解码槽位。我们尚未在相同的trace上运行GPU，因此这里没有P5线（见局限性）。 关键结果（一台AWS trn2.48xlarge，16颗Trainium2芯片，两个服务实例；AgentX运行于2026年11月2日） 伴随本文的引擎代码位于github.com/RunanywhereAI/trnrt（Apache-2.0）。 一个3200亿参数、百万token上下文的模型首先需要内存：每个MoE层中288个专家的FP8权重、每个会话的线性注意力层循环状态，以及每个agent上下文的稀疏注意力层缓存。在AWS自己的价目表上，一台trn2.48xlarge每TB加速器内存的租金不到八卡H200或B200机器的一半。 美国东部（俄亥俄）的EC2 ML容量块：价格来自AWS的容量块定价，加速器内存来自其实例页面（Trn2、P5、P6），读取于2026年11月3日。这些是租赁价格，不是服务成本。 每GB更便宜的机器只是起点：每个请求的成本取决于引擎。在AgentX上16个agent时，一台trn2.48xlarge按容量块价格计算每个agent请求成本为0.031美元。我们尚未在相同的trace上于这些GPU运行GLM-5.3-Flash，因此我们不对Trainium2与GPU的每请求表现做出任何断言；该对比在我们的待办清单上。我们能说的是还有多少余量：如今一个解码步耗时是其内存下限的四倍（24.2毫秒对6.1毫秒），因此随着引擎改进，同一台机器每小时可以服务更多agent。本文其余部分讲述这些时间花在哪里，以及我们做了什么。 GLM-5.3-Flash有45层、3200亿参数，每个token激活约180亿参数。有三点使其难以服务： 一个解码步：34个KDA层和11个DSA层，每一层都包裹在mHC中。 一台trn2.48xlarge有16颗Trainium2芯片，每颗配备96 GiB HBM、带宽2.9 TB/s，以及八个物理NeuronCore。每个NeuronCore将四个计算引擎与一个大型软件管理的暂存区SBUF结合在一起。这些引擎分别是矩阵引擎、向量引擎、标量引擎和通用GpSimd引擎。Trainium2的默认逻辑核心设置LNC = 2将物理核心两两配对，因此整机有64个逻辑核心。一个kernel将每一对视为两个并排运行的程序。我们每台机器运行两个服务实例，每个实例在8颗芯片上跨32个逻辑核心进行张量并行；在实例内部，稀疏注意力层还将长上下文分散到八组rank上（上下文并行）。这些实例位于一个路由器之后，该路由器将每个agent会话保持在同一个实例上，以便其前缀保持缓存。 每台机器两个服务实例，每个实例持有完整的模型副本。 这些kernel用NKI（Neuron Kernel Interface）编写：这是AWS为Trainium提供的、嵌入Python的kernel语言。一个NKI kernel在tile层面工作。它决定每条指令由哪个引擎执行，并显式地在HBM与SBUF之间移动数据。在GPU上，最接近的类比是用Triton或CUDA编写，而不是依赖编译器的融合。 当我们在十月初开始时，没有任何AWS支持或主流开源引擎在Trainium上服务GLM-5.3-Flash。截至11月3日： 唯一在Trainium2上服务GLM-5.3-Flash的引擎是Kiln，一个独立的Trainium开源推理引擎，于2026年10月5日首次发布。在我们的机器上，它重现了自己发布的基准成绩：在8K token提示、256 token回答、128并发请求下，每秒输出229.2个token，而v0.2.0发布的数据为229.5。 Kiln也是我们的基线，我们是在自己的机器上测量它，而不是依赖其发布的数据。10月9日，我们在同一台trn2.48xlarge上运行了Kiln v0.2.0（commit 607ac2e，即tag之后的一个README commit），采用我们为其找到的最佳配置：其发布的Trainium2配置，解码批大小为1、2和4，作为两个实例置于其路由器之后。该配置是为Kiln的基准形态（8K token提示）构建的；对于百万token提示，v0.2.0有一个单独的、需显式启用的长上下文引擎，我们没有运行它，因此那一行引用的是Kiln发布的数据。Agent流量则不同。提示很长且大部分已缓存，每个用户都在等待解码。这设定了我们的两个目标：agent负载下的百万token上下文，以及每个实例约八个并发用户时的快速解码。与我们的机器上的Kiln v0.2.0相比：

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）