跳到正文
LMSYS Blog· SGLang, NVIDIA and AMD teams·· 2 天前精选AI 评分60

SGLang 发布 Low-Latency Mode,通过插件框架在 NVIDIA 和 AMD GPU 上降低推理延迟

SGLang Low-Latency Mode: Kernels and Runtime

AI 导读

SGLang 团队推出 Low-Latency Mode,通过插件框架引入专用内核,在 NVIDIA GB300 和 AMD MI355X GPU 上显著降低推理延迟。

推荐理由

原文给出 SGLang 低延迟模式的具体内核优化路径和 NVIDIA、AMD 两套实测数字,可帮助读者评估低并发推理场景的收益。

正文 · AI 翻译

低延迟推理为针对特定模型、硬件和工作负载的优化提供了许多机会。这些优化通常需要大量的内核融合以及跨实现复制代码,使其难以维护。为支持这种特化,我们推出SGLang 低延迟模式。它通过插件框架将专用内核引入 SGLang,并在 NVIDIA 和 AMD GPU 上大幅降低延迟。在四块 NVIDIA GB300 GPU 上的 AgentX 轨迹回放中,单并发会话下 P90 每用户 token 生成速度从671.2 提升至 928.4 tokens/s(1.38×)。在四块 AMD MI355X GPU 上,四并发会话下从331.6 提升至 516.0 tokens/s(1.56×)。每项结果均为同一硬件上低延迟模式与 SGLang 基线的对比。

AgentX performance on NVIDIA GB300 and AMD MI355X at 1, 4 and 8 concurrent sessions.

AgentX · DeepSeek-V4.1-Flash with DSpark · NVIDIA GB300 与 AMD MI355X。

目录

  1. Kernels for Low-Latency Inference
  2. 插件系统:设计与集成
  3. 未来计划

面向低延迟推理的内核

NVIDIA 低延迟内核

在 NVIDIA GPU 上,我们通过重叠相互依赖的内核与通信,并针对小批量特化注意力、MoE 和草稿采样来降低延迟。这些内核位于 sgl-nvidia-mega 中,在 NVIDIA GB300 GPU 上为 DeepSeek-V4.1-Flash with DSpark 加速目标验证和草稿执行。

时间花在哪里

在单请求下,DSpark 每步验证六行 token,每个内核的工作量很少,而依赖关系和跨 GPU 通信会增加延迟。在预热后的 CUDA-graph 性能分析中,该插件将验证图从 1,205 个内核减少到 1,189 个,但将解码器层延迟中位数从 96.8 µs 降至 68.4 µs。收益主要来自减少等待和重叠工作,其中注意力区域的降幅最大。下表列出了这些变化,随后我们将逐一介绍主要优化。

单个解码器层的一个区域,单请求(中位数)SGLang使用插件
注意力,从上一 MoE 输出到输出投影33.4 µs18.3 µs
注意力输出投影与张量并行归约10.9 µs9.4 µs
MoE,从路由到归约后的输出44.4 µs36.5 µs
整层,从一次路由步骤到下一次96.8 µs68.4 µs

四块 GPU 上每组 780–800 个实例的中位数。各区域中位数之和不等于整层中位数。性能分析会影响计时;这些测量用于解释机制,而非报告服务性能。

低延迟专用优化

1. 在其输入存在之前启动下一个内核。 程序化依赖启动(PDL)允许内核在前驱内核仍在运行时启动。它可以在等待所需输入之前准备屏障、设置张量内存并加载权重。SGLang 已使用 PDL。该插件将依赖释放提前,使更多此类准备工作与正在进行的计算重叠。例如,MoE 归约在开始轮询对等结果时,就释放下一层的第一个投影。

注意力路径还使用了针对小批量专门优化的投影 GEMM,并在加载 query 时应用 RoPE,从而省去一个单独的内核。在依赖关系允许的情况下,注意力会在等待 query 之前就开始收集 KV 数据。在单个请求下,该插件还会在步骤开始时将独立的 Engram 查找和投影移到侧流上。它们的 all-reduce 仍留在主流上,以保持跨 GPU 的通信顺序。

2. 从 GEMM epilogue 发布投影输出。 在注意力输出投影之后,每个 GPU 持有一个部分结果,该结果必须跨 GPU 进行归约。SGLang 的标准路径使用两个内核:GEMM 写出其输出,然后一个融合内核将其读回,发送给对等节点,并执行归约、mHC 更新和归一化。

该插件保留两个内核,但将发送移入 GEMM 的 epilogue。每个 CTA 在其他 tile 仍在计算时,直接通过 NVLink 多播存储将已完成的 tile 从寄存器发布到对等 GPU。消费者通过 PDL 启动,并在归约之前等待已发布的结果。这更早地启动通信,并避免仅为发送而重新读取部分输出。消费者的计算保持不变。

Attention output handoff in SGLang and the NVIDIA plugin for DeepSeek-V4.1-Flash on GB300

该插件将跨 GPU 发送移入 GEMM epilogue。两条路径都使用两个内核;方框大小不代表持续时间。

3. 在路由完成之前启动专家 GEMM。 在低并发下,路由可能占据 MoE 路径的相当大一部分。在最多八行的情况下,每个专家选中的 token 可容纳在一个八行 tile 中。该插件使用位图和种群计数来放置这些 tile,无需扫描或排序,同时保留 SGLang 的 top-k 选择。它提前启动第一个专家 GEMM,使设置与路由重叠,并跳过填充权重中的全零 tile。标准内核已经使用八行 tile;收益来自更简单的路由、更早的执行和更少的浪费工作。在性能分析运行中,MoE 窗口中位数从 44.4 µs 降至 36.5 µs,而第二个 GEMM 或归约几乎没有变化。

Measured MoE timeline at one request with and without the NVIDIA plugin

来自同一 GPU、层和重放的一对匹配跟踪。GEMM1 的条形包含其对路由的等待,因此更早的开始可能产生更长的条形,同时缩短整体 MoE 窗口。

4. 将顺序草稿采样保留在一个内核中。 DSpark 的 Markov head 顺序提出 token。该插件将全部五个提议合并为每个 GPU 一次内核启动,并在跨 GPU 协调 token 选择时将 head 权重保留在寄存器中。这将性能分析中的提议跨度从 85.1 µs 降至 34.2 µs,而不改变 SGLang 的推测 worker 或接受逻辑。

服务结果

在 8K 输入 / 1K 输出的工作负载上,NVIDIA 插件在四块 GB300 GPU 上、1–8 个并发请求时将每用户解码速度提升 12–38%。两次运行使用相同的设置和模拟平均接受长度 3.51,仅更改插件标志。

Per-user decode speed versus aggregate output throughput per GPU for DeepSeek-V4.1-Flash on four GB300 GPUs, with and without the NVIDIA plugin

DeepSeek-V4.1-Flash 搭配 DSpark · 4 × GB300 · 8K 输入 / 1K 输出。

AMD 低延迟内核

在 AMD GPU 上,我们通过将解码器层的大部分计算和通信合并到少数持久内核中来降低延迟。这些内核位于 sgl-amd-mono 中,在 AMD Instinct MI355X 和 MI350X GPU 上为 DeepSeek-V4.1-Flash 搭配 DSpark 加速目标验证和草稿执行。

在低并发下,每个解码器层会运行数十个小内核和两次张量并行 all-reduce。三项成本占主导:

  • 启动开销:HIP graph 消除了 CPU 的启动工作,但每个小内核在 GPU 上仍需 2-4 µs。
  • 空闲边界:每个内核必须完成后下一个才能开始,因此当一个内核排空、下一个内核启动时,计算单元处于空闲状态,权重加载也会暂停。
  • 通信停顿:每一层都要等待两次跨 GPU 的 all-reduce,期间没有计算在运行。

Schematic timeline of one DeepSeek-V4.1 layer at batch size 1. The standard path runs dozens of small kernels with an idle gap at every launch boundary and two all-reduces that stop computation. With the AMD plugin, the layer runs as two FlyDSL launches, K1 and K2, where weights that do not depend on routing load early and partial sums go to peer GPUs while computation continues.

该插件用两次持久化启动替代数十个小内核,提前加载权重,并将跨 GPU 求和与计算重叠。方框大小不代表时长。

每层两次持久化启动

持久化 FlyDSL 内核将每层的大部分工作合并为两次启动。K1 处理 mHC 更新、query 和 KV 投影以及 KV-cache 写入。K2 处理 attention、输出投影、剩余的 mHC 更新、MoE 以及两次 all-reduce。索引层在两者之间加入压缩、K2i 打分和 top-k。启动边界使 K1 的 KV 写入在 K2 读取之前对整个 GPU 可见。

启动内部各阶段的重叠

每个阶段产生的值都会与当前解码步的标签一起存储,因此下一个阶段(无论在同一 GPU 还是另一 GPU 上)只要标签匹配就会立即开始处理该值。步与步之间无需单独的标志或重置。不依赖路由的权重会在其输入到达之前加载:K1 的投影权重在 norm 完成前加载,K2 中 mHC 混合权重和共享专家的权重在它们要相乘的数据之前加载。路由专家权重在路由器选定专家后加载。

将 all-reduce 引入内核

每个输出投影 tile 计算完成后,K2 将其推送到每个 GPU 的内存,每个 GPU 按 rank 顺序对各部分求和,同时其他工作组仍在生成 tile(推送与归约阶段)。MoE 输出以相同方式归约。每次跨 GPU 等待都有低于 SGLang 看门狗超时的截止时间;错过截止时间会在每个 GPU 上停止启动,而不是让服务器挂起。

AMD 实现与 decode graph 兼容,并直接集成到 SGLang 的 graph runner 中。内核在初始化期间编译,缓冲区在初始化期间准备。随后 SGLang 为每种支持的形状捕获专用执行路径,并通过现有的服务流程重放。

服务结果

在 AgentX trace 回放中,AMD 插件在四块 MI355X GPU 上、1、4 和 8 个并发会话时,将 P90 每用户 token 生成速度提升了 45–56%,如概览图所示。两次运行均使用相同的 AgentX 配方,配备 DSpark 推测解码和 TP4。

插件系统:设计与集成

专用内核可能需要大量代码。硬件、模型架构和并行配置的每种组合都可能需要自己的实现。不同的工作负载还会带来进一步的专门化。更多组合意味着更多需要审查、测试和维护的代码。将每个实现都引入 SGLang 会使核心代码库越来越难以维护,并给审查者和维护者带来越来越多的工作量。

External kernel libraries register model specifications with SGLang, which validates configurations, builds the plugin and dispatches supported steps.

外部内核库通过注册、验证、构建和运行时调度连接到 SGLang。该图展示了共享设计;AMD 和 NVIDIA 接口仍在演进中。

插件系统通过声明其需求和支持配置的模型规范,将这些外部库连接到 SGLang。内核团队可以独立开发其实现,同时复用 SGLang 的调度和运行时资源。

该集成将启动准备与执行分离:

  • 注册:插件为其支持的每种模型架构注册一个规范。该规范描述实现的需求,并提供构建它的入口点。这为 SGLang 提供了一种发现外部实现的方式,同时将模型支持的细节保留在拥有内核的库中。
  • 守卫:专用内核会对硬件、精度、并行性和服务特性做出假设。SGLang 在加载权重之前检查声明的需求,以便不兼容的配置尽早失败。在构建插件之前,对已加载模型运行额外检查。这些检查区分不支持的服务器配置与受支持的服务器遇到插件覆盖范围之外的步骤。
  • 构建:权重准备和执行资源在初始化期间设置,以便请求可以复用它们。该集成适应每种实现的需求:AMD 可以在权重转换之前保留检查点张量,而 NVIDIA 接收一个绑定,通过它可以在运行时状态可用时访问该状态。SGLang 在确定 KV 缓存大小之前构建插件,以便其分配被计入,然后在模型的前向路径中安装调度。
  • 运行时:专用执行和通用执行共存,因为插件可能只覆盖某些步骤类型和 token 数量。安装的调度器使用插件的支持条件来选择其前向或 SGLang 的原始前向。张量并行 rank 必须在覆盖范围上达成一致,以避免进入不兼容的集合通信。此选择适合 ModelRunner 的现有执行流程。在图捕获期间,为相应形状记录所选路径;服务随后可以重放该图,而无需重复 Python 调度。

这种设计让内核团队能够交付专注的优化,而无需将每个专用实现添加到 SGLang 的核心代码库中。它将集成工作集中在显式接口上,并允许每个库逐步扩展其覆盖范围。专用内核可以服务它们优化的场景,而 SGLang 的通用前向处理受支持配置内的剩余步骤。

未来计划

更广泛的模型支持

Megakernel/Monokernel 的开发正在针对 GLM-5.3、K3、MiniMax M3 以及其他模型在 NVIDIA 和 AMD GPU 上进行。我们计划在这些实现准备就绪后将其集成到插件框架中。这些低延迟内核显示出进一步提升性能的潜力。

CI 与性能验证

外部内核实现位于 SGLang 常规代码审查流程之外。因此,插件接口处的强大 CI 覆盖对于验证这些实现满足 SGLang 的正确性和运行时要求至关重要。我们计划测试兼容性检查、权重绑定、前向输出和状态更新,包括 KV 缓存写入。随着 SGLang 和外部库的发展,端到端测试还将覆盖图执行、回退行为和性能回归。

演进中的共享插件接口

尽管插件框架现已就位,但仍处于早期阶段。我们预计将通过长期迭代来完善其接口和设计,同时集成更多内核,并从生产环境的使用中学习经验。

来源:LMSYS Blog · lmsys.org