AREX Feed Article
RadixArk 与 AWS Annapurna Labs 发布首个公开版 SGLang-Neuron v0.1.0：SGLang 推理栈以 PyTorch 原生插件落地 Trainium
11 月 12 日，RadixArk 的 SGLang 团队与 AWS Annapurna Labs 联合发布并开源 ，这是 SGLang 推理服务栈面向 AWS Trainium 芯片的首个公开版本，以 public preview 形态提供。它是一个树外（out-of-tree）插件，安装后经 SGLang 的插件入口注册 neuron 平台，启动时以 SGLANG_PLATFORM=neuron 选中；预览期经由 AWS Deep Learning Container 的 SGLang-Neuron 镜像提供支持。 独立记录了同一发布：LLM 跑在 Trn3 实例的同一条全异步运行时上，多模态扩散模型经 SGLang Diffusion 跑在 Trn2 实例上。
这次发布里最大的部署是 Kimi-K3：按官方博客的说法，它跨三个 Trainium3 节点做 PD 分离（prefill-decode disaggregation，把预填充与解码分置在不同服务器）——1 个预填充节点加 2 个解码节点，KV 缓存经 NIXL 在 EFA 网络上跨节点传输，官方称在混合 4K～1M token 提示、满槽负载下达到每用户 67 token/秒；GLM-5.2（MXFP4）与 Kimi-K3 也都被标注支持 1M token 上下文。
图：SGLang-Neuron 插件在 SGLang 技术栈中的位置（来源：LMSYS 博客 Figure 1）
十月封闭 beta 之后，首个公开版本多了什么
10 月 22 日，SGLang 团队曾宣布 ：SGLang 的 Neuron 后端在 TorchNeuron 还处私测阶段时就已开始在其上开发，当时的 beta 只向一小批团队开放，覆盖 Qwen3 稠密与 MoE（mixture-of-experts，专家混合）、DeepSeek-V2-Lite、GLM-4、GLM-5.2/5.3 的推理服务和 RL（强化学习）rollout。v0.1.0 博客这样区分两个版本：beta 证明了 SGLang 的模型组件可以经 TorchNeuron 基本原样地搬到 Trainium，这个版本要解决的是怎么把它们做快——大部分工程花在三处：步骤怎么发射、一步计算能收进多大的单个内核、KV 缓存不在显存时放在哪里。
支持的模型清单也扩大了。release notes 的亮点列表里有 GPT-OSS-120B、Qwen3.5-397B-A17B、GLM-5.2（MXFP4）与 Kimi-K3；博客的完整支持表里既有 Qwen3-0.6B/8B/30B-A3B、GLM-4-9B、DeepSeek-V2-Lite，也有 Kimi-K3.1 和 GLM-5.3，release notes 并注明本轮 LLM 均为纯文本服务。致谢部分还点名了 Yotta Labs 对 SGLang on Trainium 的贡献，包括更早的 Mini-SGLang-Neuron 工作。
插件架构：SGLang 核心不动，Neuron 平台按需接入
职责切分很严格：SGLang 核心保留 API server、调度器、radix cache 和分离式服务协议；插件提供平台要求的部件——注册在模型目录里的 Neuron 模型实现（服务器启动前先校验 checkpoint、并行布局与标志位，所有模型共用一条启动路径）、覆盖注意力/MoE/投影与整步 megakernel（一步收进一次内核调用）的 NKI（Neuron Kernel Interface）内核、张量/专家/数据并行加流水线并行与解码上下文并行（DCP）、带编译管理器和分桶器的图执行器，以及 KV 缓存池：MHA（multi-head attention，多头注意力）模型用分页全量 KV 池，MLA（multi-head latent attention，多头潜在注意力）与 DSA（DeepSeek sparse attention，DeepSeek 稀疏注意力）注意力走独立的隐空间池。投机解码和跑在设备上的采样器也由插件提供。
底座是 TorchNeuron——AWS 的 Trainium 原生 PyTorch 后端。它把 Trainium 注册成 PyTorch 设备并充当 TorchDynamo 后端，torch.compile 的图经 Neuron 编译器落地，NKI 内核插进同一张编译图；每个 SGLang 调度 rank 绑定自己的 NeuronCore，在其上构建模型和 KV 池。Trainium 执行的是编译好的程序（NEFF），图执行器预先编译一组固定 shape 桶，每一步派发到匹配的桶上。
全异步运行时与 megakernel：两步之间不让 NeuronCore 空转
全异步运行时针对的目标很直接：设备在两步之间不空闲。单流执行时，两张图之间的一次主机拷贝要先等上一张图跑完，又挡住下一张图，NeuronCore 就闲了下来。device_io 层用双流分工解决：默认流跑图，旁流跑所有主机拷贝，两者互不排队；下一步的输入在图发射前就排上旁流，与上一张图在设备上的执行重叠；回读只在对应图的发射事件之后入队，采样出的 token 很小、可以直接读。SGLANG_NEURON_MAX_INFLIGHT_LAUNCHES 默认为 2，调度器因此能在第 n 步跑完前发射第 n+1 步。博客称主机开销由此藏进启动队列——对应 GPU 上把主机开销藏进设备计算的同一条思路——而且对插件上所有模型生效，不是某一条调优过的路径。
megakernel 改变一步计算的形态。默认编译路径把一个解码步拆成许多小内核，内核边界由编译器决定，中间结果在各内核之间往返 HBM；megakernel 把一个解码步的每一层收进一次 NKI 内核调用，权重按前导层轴打包、KV 池原地写入，内存搬运端到端可控。各阶段在平台上仍是 Neuron 后端的 torch.compile 图，只是每个阶段只调用一次内核；调度、KV 缓存、PD 分离与结果处理照旧归 SGLang。一份 MegakernelSpec 契约管住整步的规划与执行，新模型接入不再需要新的运行时子类。
走这条路径的模型：博客称 GPT-OSS-120B 是第一个，解码经 release notes 称作 giga kernel 的路径，其上叠流水线并行与 DCP；GLM-5.2 的预填充和解码都用 megakernel；Kimi-K3 为预填充配了专用 NKI 内核；Qwen3.5-397B-A17B 按 release notes 跑全模型 megakernel。
PD 分离怎么搭：预填充写 KV、解码稳 TPOT
博客认为在 Trainium 上，预填充与解码最好跑成不同的服务器，理由有三：其一，聚合部署在两种阶段间切换时要付程序切换栅栏（program-switch fence），分置后每个服务器只编译、预热自己阶段的图；其二，解码步之间不再插进预填充前向，每输出 token 的耗时（TPOT）稳定在中位数；其三，两个角色还能各拿一套 TP/EP/DP（张量/专家/数据并行）布局、独立扩容。机制上，预填充服务器算出提示的 KV 缓存和首个 token，再用 NIXL WRITE 经 EFA 直接写进解码服务器的缓存池；预填充侧做成流水线，计算跑在 KV 传输前面，首 token 不必等交接完成。一个部署可以跑任意数量的预填充和解码服务器（xPyD），前面挂镜像自带的负载均衡器。模型可以声明命名解码点（如 128K 与 1M），预填充和解码两端用同一点位编译，编译出的形状才能对上。
图：Kimi-K3 在三个 Trainium3 节点上的 PD 分离拓扑（来源：LMSYS 博客 Figure 2）
Kimi-K3 的运行门槛写在 release notes 里：需要 Trn3 UltraServer——至少两台经交换机互连的 trn3.48xlarge。博客给出的拓扑是：1 个预填充节点跑 8 级流水线（PP8×TP16×DCP16），配分块键控前缀缓存与 HiCache；2 个解码节点（DP8×TP32×DCP32）组成解码端，其上有 8 个 DP-attention 组、专家并行到 EP256，运行 DSPARK 投机解码。模型用 FP8 的 MLA 注意力加 MXFP4 专家，KV 存在 MLA 隐空间池里。
GLM-5.2（MXFP4）是 MoE 加 DSA 注意力，release notes 写明仅以分离方式服务（1 个预填充服务器加 1 个解码服务器），配 MTP（multi-token prediction，多 token 预测）投机解码，并设 128K 与 1M 两个命名解码点；它跑在公开的 Neuron SDK 2.33.0 上。1M 解码点瞄准的是会话上下文逐轮增长的长上下文 agent 编码类负载。
HiCache 与远程 DRAM：被逐出的前缀不必重算
Agent 类负载会反复重发越来越长的前缀；设备缓存池满了，radix cache 就要逐出页面，而被逐出的前缀下一次命中通常意味着整段重算。HiCache 给这些页面一个去处：被逐出的 radix cache 页面保留在主机内存，之后命中时直接装回。radix 树、主机分配器与传输引擎用的是 SGLang 上游代码，插件只提供为 Trainium 内存搬运方式定制的主机池——整页、层优先排布的锁页（pinned）缓冲；页面搬运走 nrt_tensor_read_batch 与 nrt_tensor_write_batch 的批量 DMA（direct memory access，直接内存访问），每个方向一条专用线程；聚合部署时，等待装载的批次被挂起，正在跑的批次继续解码。
效果数字同样出自发布方：GPT-OSS-120B 在 TP8、32K 上下文下，HiCache L2 带来 72%～75% 的输出吞吐提升。v0.1.0 还在 L2 之下加了一层远程 DRAM，经 NIXL LIBFABRIC 访问：会话超出本机主机池时，页面溢出到另一节点的内存而不是被丢弃；主机预算在该节点各 rank 与 --dp-size 副本之间共享，避免超卖同一份 DRAM。
扩散模型、RL rollout 与公开预览的边界
扩散侧，SGLang Diffusion 新增了 Neuron 平台，装了 torch-neuronx 就会自动选中；SGLang 仍负责启动器、调度器、图像与视频 API 及实时 WebSocket 协议，插件提供 pipeline、编译图和 NKI 内核。五款模型都以 BF16 权重跑在一台 trn2.48xlarge 上，每个模型的 recipe（模型配方）里写了启动命令、请求格式与实测延迟，并设了对照参考的精度门：FLUX.1-dev 出 1024×1024 与 2048×2048 的文生图；Wan2.2 A14B 出 81 帧、480p/720p/1080p 的文生视频与图生视频；MiniMax-H3 经 t2va（text-to-video-and-audio，文生视频加音频）或首尾帧条件生成 124 帧 768p 视频加配乐；FLUX.3-Action 是机器人动作策略模型，从三路相机画面、机器人状态与指令返回接下来的 32 个机器人动作，三个 DROID checkpoint（base、gd、sd）都经 OpenPI WebSocket/HTTP 动作 API 服务；LingBot-World-V2 从单张图片沿相机路径实时流式生成 480p 世界，DiT 用分页 KV 缓存维护因果历史，与文本编码器、流式 VAE（variational autoencoder，变分自编码器）分用不同内核组，客户端走与 GPU 一致的实时视频 WebSocket 协议。
芯片布局上还有一个具体约束：trn2.48xlarge 的 16 颗芯片组成 4×4 torus（环面互连），只有构成环的芯片组才能路由集合通信，而 SGLang 按固定 rank 顺序生成的 CP（context parallelism，上下文并行）组和 Wan2.2 的 CFG（classifier-free guidance，无分类器引导）组会被 fabric 拒绝；插件在 rank 生成器上加了一个小钩子重命名 rank，让每个 TP 组落在同一行、每个 CP 组排成一列。引导尺度、时间步与采样系数都作为设备张量进图，请求之间改参数不触发重编译。
推理只是 RL 后训练循环的一半：RadixArk 的开源 RL 框架 Miles 用 SGLang 生成 rollout、经可插拔后端训练，在 Trainium 上走 TorchTitan 后端跑在 TorchNeuron 上，双方在 Qwen3-30B-A3B 上端到端验证了整条链路。
想跑起来的门槛写在 release notes 里：LLM 需要 trn3.48xlarge 实例，扩散模型需要 trn2.48xlarge，主机装 Neuron SDK 2.33.0 对应的驱动（aws-neuronx-dkms 2.31.7.0），再从 public.ecr.aws 拉取 v0.1.0 镜像启动；的 examples 目录为除 Qwen3-MoE 外的每个 LLM 提供带校验标志的启动脚本和 GSM8K 精度检查。
release notes 写明这是 public preview，功能、API 与支持的模型都可能在后续版本变化。博客列出的下一步——面向循环状态模型的统一框架、带印记与声明参数的内核包契约、Kimi-K3.1 投机解码、把五个运行时子类收敛成一个 host pattern——目前都还是计划；两份一手记录给出的吞吐与延迟数字都只有发布方口径，记录里没有附独立第三方的对照基准。