主要收获
- PyTorch 和 vLLM 对 AI 生态系统都至关重要,并且越来越多地被一起用于前沿生成式 AI 应用,包括大规模推理、后训练和代理系统。
- 随着 PyTorch 基金会转型为一个伞式基金会,我们很高兴看到项目被从超大规模企业到初创公司等各类客户使用和支持。
- vLLM 正在利用更广泛的 PyTorch 生态系统加速创新,受益于 torch.compile、TorchAO、FlexAttention 等项目,并合作支持异构硬件和复杂的并行化。
- 各团队(及其他)正在合作,为大规模推理和后训练构建 PyTorch 原生支持和集成。
甚至在 vLLM 加入 PyTorch 基金会之前,我们已经看到 PyTorch 和 vLLM 共同被世界上部署大规模 LLM 的一些顶级公司广泛采用。有趣的是,这两个项目有许多共同之处,包括:强大的富有远见的领导者;由多个实体(包括工业界和学术界)的提交者组成的广泛、自然形成的多边治理结构;以及对开发者体验的压倒性关注。
此外,在过去一年多的时间里,我们看到这两个项目支撑了许多最受欢迎的开源 LLM,包括各种 Llama 和 DeepSeek 模型。鉴于这些相似之处以及项目之间的互补性,看到所有不同的集成点确实令人兴奋。
PyTorch → vLLM 集成
在各个点进行集成的总体目标是释放性能并为用户带来新功能。这包括对 Llama 模型以及更广泛的开放模型的优化和支持。
torch.compile:torch.compile 是一个编译器,可以优化 PyTorch 代码,以最少的用户工作量提供快速性能。虽然手动调整模型的性能可能需要几天、几周甚至几个月,但这种方法对于大量模型来说是不切实际的。相反,torch.compile 提供了一个方便的解决方案来优化模型性能。vLLM 默认使用 torch.compile 为其大部分模型生成优化的内核。最近的基准测试显示,对于 Llama4、Qwen3 和 Gemma3 等流行模型,torch.compile 在 CUDA 上实现了显著加速,范围从 1.05 倍到 1.9 倍。
TorchAO:我们很高兴地宣布,TorchAO 现已正式支持作为 vLLM 中的量化解决方案。此次集成带来了使用 Int4、Int8 和 FP8 数据类型的高性能推理能力,并将支持专为 B200 GPU 设计的 MXFP8、MXFP4 和 NVFP4 优化。此外,我们正在努力计划对 AMD GPU 提供 FP8 推理支持,以扩展高性能量化推理的硬件兼容性。
TorchAO 的量化 API 由强大的高性能内核集合提供支持,包括来自 PyTorch Core、FBGEMM 和 gemlite 的内核。TorchAO 技术旨在与 torch.compile 结合使用。这意味着更简单的实现和从 PT2 获得的自动性能提升。编写更少的代码,获得更好的性能。
此次集成最令人兴奋的方面之一是它实现了无缝工作流程:vLLM 用户现在可以使用 TorchTitan 进行 float8 训练,使用 TorchTune 进行量化感知训练 (QAT),然后直接通过 vLLM 加载和部署其量化模型以进行生产推理。这种端到端管道显著简化了从模型训练和微调到部署的路径,使高级量化技术更容易为开发者所用。
FlexAttention:vLLM 现在包含了 FlexAttention——一个为灵活性设计的新注意力后端。FlexAttention 提供了一个可编程的注意力框架,允许开发者定义自定义注意力模式,从而更容易支持新颖的模型设计而无需进行大量后端修改。
这个由 torch.compile 启用的后端生成 JIT 融合内核。这使得在保持非标准注意力模式性能的同时保持灵活性。FlexAttention 目前在 vLLM 中处于早期开发阶段,尚未准备好用于生产。我们将继续投资于此集成,并计划使其成为 vLLM 建模工具包的强大组成部分。目标是简化对新兴注意力模式和模型架构的支持,从而更容易弥合研究创新与可部署推理之间的差距。
异构硬件:PyTorch 团队与不同的硬件供应商合作,为不同类型的硬件后端提供了坚实的支持,包括 NVIDIA GPU、AMD GPU、Intel GPU、Google TPU 等。vLLM 推理引擎利用 PyTorch 作为代理与不同的硬件对话,这显著简化了对异构硬件的支持。
此外,PyTorch 工程师与其他 vLLM 贡献者密切合作,以支持下一代 NVIDIA GPU。例如,我们已经彻底测试了vLLM 在 Blackwell 上对 FlashInfer 的支持,进行了性能比较,并调试了精度问题。
PyTorch 团队还与 AMD 合作,增强了对 vLLM + Llama4 的支持,例如AMD 上的 Llama 4 首日支持以及MI300x 上的 Llama4 性能优化。
并行化:在 Meta,我们在生产中利用不同类型的并行化及其组合。流水线并行化 (PP) 是其中重要的一种。vLLM 中的原始 PP 硬性依赖于 Ray。然而,并非所有用户都利用 Ray 来管理他们的服务和协调不同的主机。PyTorch 团队开发了支持纯 torchrun 的 PP,并进一步优化了其方法以重叠微批次之间的计算。此外,PyTorch 工程师还开发了视觉编码器的数据并行化,这对于多模态模型的性能至关重要。
持续集成 (CI):鉴于 vLLM 对 PyTorch 生态系统至关重要,我们正在合作确保项目之间的 CI 具有良好的测试覆盖率、资金充足,并且整个社区可以依赖所有这些集成。仅仅集成 API 是不够的;重要的是要建立 CI,以确保随着 vLLM 和 PyTorch 都发布新版本和新功能,随着时间的推移不会出现任何问题。更具体地说,我们正在测试 vLLM 主分支和 PyTorch nightly 版本的组合,我们相信这将为我们和社区提供监控两个项目之间集成状态所需的信号。在 Meta,我们开始将一些开发工作转移到 vLLM 主分支上,以对 vLLM 的各种正确性和性能方面进行压力测试。此外,vLLM v1 的性能仪表板现已在hud.pytorch.org上提供。
下一步是……
这仅仅是个开始。我们正在共同努力构建以下高级功能:
1. 大规模模型推理:主要目标是确保 vLLM 在云产品上高效大规模运行,展示关键能力(预填充-解码解耦、多节点并行、高性能内核和通信、上下文感知路由和容错),以扩展到数千个节点,并成为企业构建的稳定基础。在第二季度,Meta 工程师已经在 vLLM 引擎和 KV 连接器 API 的基础上原型化了解耦集成。团队正在与社区合作尝试新策略,并计划将最成功的策略上游,以进一步推动 vLLM 的可能性。
硬件:H100 GPU,96 GB HBM2e,AMD Genoa CPU,CUDA 12.4
2. 强化学习后训练:推理时间计算正迅速成为 LLM 和代理系统的关键。我们正在研究端到端原生后训练,该训练将大规模 RL 与 vLLM 作为系统的推理骨干相结合。
干杯!
- PyTorch 团队(Meta)和 vLLM 团队