NVIDIA Blackwell 上的 GPT-OSS 性能优化:突破帕累托最优边界

阅读时间 8 分钟
vLLM 与 NVIDIA 团队

TL;DR: vLLM 与 NVIDIA 团队携手开源社区,在 NVIDIA Blackwell GPU 上运行 gpt-oss-120b 模型时取得了重大的性能里程碑。通过与 FlashInfer 的深度集成、利用 torch.compile 实现的新型算子融合,以及一系列推理运行时特性,我们刷新了该模型的性能帕累托前沿(Pareto frontier)记录——在最大吞吐量提升 38% 的同时,实现了交互性提升 13%。

本文详细介绍了这一工程历程、技术突破及重现结果的操作指南。持续的基准测试结果也可在 SemiAnalysis Inference MAXvLLM Recipes 上查看。

目录


简介

仅针对单一指标(如最大吞吐量或单批次延迟)进行优化往往不足以满足实际部署需求。不同的应用场景对延迟限制和请求并发量的要求各异。因此,真正的挑战在于优化帕累托前沿:即在每 GPU 每秒令牌数 (TPS)(总拥有成本,TCO)与每用户 TPS(交互性)之间实现最佳权衡的曲线。推动该曲线向右上方移动,意味着在为单个用户提供更快生成速度的同时,允许更多用户共享硬件资源。SemiAnalysis InferenceMAX 已经明确了针对现代 GPU 上的 LLM 推理工作负载进行性能数据测量、报告和改进这一关键需求。

主要应用场景之一是服务 OpenAI 的 gpt-oss-120b 模型,这是一款原生 4 位量化 (MXFP4) 的混合专家 (MoE) LLM。它在保持同等规模模型 SOTA 准确性的同时,具备强大的智能体能力。在近期的 SemiAnalysis InferenceMAX 展示中,vLLM 证明了其在 NVIDIA 最新的 Blackwell (B200/GB200) 架构上高效处理该工作负载的能力。

优化的核心在于硬件与软件的协同设计。NVIDIA B200/GB200 GPU 引入了原生 FP4 TensorCore 和每 GPU 192GB HBM 等强大特性,这对服务 gpt-oss 等大型 MoE 模型至关重要。为了充分利用这些硬件,vLLM 和 NVIDIA 团队集成了 FlashInfer,并采取了严谨的优化策略,聚焦于算子融合、通信开销缩减以及主机-设备端异步重叠。

FlashInfer 集成与基于 torch.compile 的融合优化

为了最大化 Blackwell Tensor Core 的利用率,vLLM 将 FlashInfer 作为其处理 Attention、MoE 以及其他计算密集型和融合算子的主要后端。

1. 关键计算算子集成:

  • MoE 后端: 我们启用了 trtllm-gen (PR23819)cutlass (PR23696) 两种后端来配合 FlashInfer 进行 MoE 操作。这使得 vLLM 能够为专家路由和计算选择性能最佳的算子。FlashInfer 不仅为 LLM 提供了最高性能的算子,还包含了即时编译 (JIT)、自动调优和算子缓存功能,极大改善了有高性能算子开发需求用户的体验。
  • FP8 KV 缓存: 以 FP8 精度存储 KV 缓存允许引擎在相同的 KV 缓存容量预算下服务更多并发请求。此外,在 FP8 精度下执行部分 Attention 操作也降低了 Attention 计算的复杂度和内存压力。为实现该用例的最佳性能,vLLM 已集成 FlashInfer 的优化 Attention 算子(PR25674)

2. 基于 torch.compile 的图融合: 我们优化的很大一部分工作聚焦于通过算子融合来减少内存访问和算子启动开销。vLLM 没有采用硬编码的融合优化,而是基于 torch.compile 构建了一套广泛的基础架构来自动执行算子融合。这种方法不仅提升了性能,还显著降低了启用、泛化和维护此类改进的工作量。

  • AR + Norm 融合: 我们实现了 AllReduce (AR) 和 RMSNorm 操作的融合。这对于张量并行 (TP) 部署尤为重要,因为通信开销往往会成为瓶颈,详细信息请参见 PR20691
  • Pad + Quant & Finalize + Slice: 我们正在积极推进融合通道(PR30647),用于 Padding/Quantization 和 Finalize/Slice 操作,以进一步简化 MoE 执行路径,预计将带来 6% 的性能提升。

随着我们识别并开发新的融合操作,团队将继续通过这一基础设施带来自动的性能提升。

运行时改进

在像 Blackwell 这样的下一代硬件上,GPU 速度极快,以至于 CPU(主机)往往成为瓶颈,难以足够快地分发算子以保持 GPU 繁忙。此外,prepare_batch、请求调度和采样逻辑也需要繁重的 CPU 侧计算。这种“主机开销”表现为算子执行之间的间隙,从而降低性能和 GPU 的整体利用率。

为了解决这个问题,我们在 vLLM 中实现了异步调度 (Async Scheduling)流间隔 (Stream Interval),有效地消除了主机侧的开销。

异步调度:

  • 机制: 该调度器将 CPU 的请求调度与 GPU 的执行进行了解耦。通过允许 CPU 在 GPU 仍处理当前批次的同时准备下一批次请求,我们有效地隐藏了主机端开销。
  • 影响: 这种优化对于 gpt-oss 模型至关重要,特别是在高吞吐量和最小延迟场景中。在更强大的 GPU(H200、B200、GB200)上,预计会有约 10% 的性能提升。
  • 配置: 最近的 vLLM 版本中已默认开启此功能。

流间隔:

  • 机制: 该功能通过在发送生成令牌到客户端前进行缓冲,降低了网络响应的粒度。引擎不是为每个单一令牌触发网络调用,而是等待直到达到指定的缓冲区大小(“间隔”)。关键在于,该实现通过确保第一个令牌始终被立即发送(保持首字延迟 TTFT 较低)来保留响应性,而后续令牌则进行批量发送。
  • 影响: 通过减少 HTTP/gRPC 响应分发的频率,这显著降低了与网络 I/O 和序列化相关的 CPU 开销。在高并发基准测试中(例如,gpt-oss-20b 加上 1024 个并发请求),此优化缓解了输出队列瓶颈,实现了 57% 的端到端性能提升并改善了输出令牌时间 (TPOT)。
  • 配置: 用户可以使用 --stream-interval <num_tokens> 参数配置此行为。默认值为 1(标准流式处理),但在高吞吐量部署中增加此值(例如设为 10)对于降低主机开销非常有效。

部署方案

大部分优化已在最新 vLLM 版本中默认应用。此外,为了在 Blackwell GPU (B200/GB200) 上重现 gpt-oss 的优化性能,我们建议在您的 vLLM 部署方案中使用以下配置。它们也可以在 vLLM Recipes 页面上找到。

推荐配置标志

  • 图捕获
    • --cuda-graph-capture-size 2048
  • 调度
    • --api-server-count 20--stream-interval 20:这有助于将 HTTP API 服务器开销与推理引擎解耦,从而在高并发下稳定性能。
  • MoE 后端
    • 显式启用针对 FP8/FP4 MoE 优化的 Cutlass 后端以确保最大吞吐量:VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1

结果

InferenceMax 发布以来,这些优化措施的综合效果带来了显著的性能提升。值得注意的是,最大吞吐量下性能提高了 38%,最小延迟下性能提高了 13%。


这些改进不仅仅适用于单一用例,而是覆盖整个帕累托曲线,从而使整个 vLLM 社区受益

后续步骤

我们针对 gpt-oss 的工作仍在继续。以下是推动帕累托前沿进一步发展的积极工程方向。该列表也可在 Issue 30758 中找到。

存算分离

通过将预填充 (Prefill) 阶段和解码 (Decode) 阶段分离到不同的 GPU 上,我们有望实现更高的单 GPU 吞吐量。我们目前正在实验此设置,并寻找能够实现更佳性能的正确配置。

数据+专家并行性能

我们的推算表明,使用 DEP2(在 2 个 GPU 上进行 Attention DP + MoE EP)可能比在相同延迟(TPS/用户)下的 TP1 和 TP2 实现更高的每 GPU 吞吐量。然而,目前由于 MoE 算子选择问题,DEP2 的性能暂时劣于 TP1/TP2。我们正在积极努力解决此问题。

最小延迟性能

我们针对最小延迟场景(更具体地说是 TP8 并发 8)识别出了一些性能优化机会:

  • RoPE+Q+Cache 融合:FlashInfer 中已有该算子,vLLM 的集成工作正在进行中。
  • 路由器 gemm 和 fc_qkv/fc_o_proj gemm:我们可以使用具有更好性能和 PDL 支持的专用微型 gemm 算子。

致谢

我们要感谢 vLLM 社区中许多有才华的人,他们为这一努力共同做出了贡献。

  • Red Hat:Michael Goin, Alexander Matveev, Lucas Wilkinson, Luka Govedič, Wentao Ye, Ilia Markov, Matt Bonanni, Varun Sundar Rabindranath, Bill Nell, Tyler Michael Smith, Robert Shaw
  • NVIDIA:Po-Han Huang, Pavani Majety, Shu Wang, Elvis Chen, Zihao Ye, Duncan Moss, Kaixi Hou, Siyuan Fu, Benjamin Chislett, Xin Li, Vadim Gimpelson, Minseok Lee, Amir Samani, Elfie Guo, Lee Nau, Kushan Ahmadian, Grace Ho, Pen Chun Li
  • vLLM:Chen Zhang, Yongye Zhu, Bowen Wang, Kaichao You, Simon Mo, Woosuk Kwon, Zhuohan Li
  • Meta:Yang Chen, Xiaozhu Meng, Boyuan Feng, Lu Fang