[论文解读] Teola: Towards End-to-End Optimization of LLM-based Applications

Xin Tan, Yimin Jiang+2 more|arXiv (Cornell University)|Jun 29, 2024

Distributed and Parallel Computing Systems被引用 4

一句话总结

Teola 提出了一种面向基于大语言模型（LLM）应用的细粒度、原语级数据流图编排框架，实现了 LLM 与非 LLM 组件之间的端到端优化。通过暴露底层任务原语及其依赖关系，Teola 利用拓扑感知的批处理以及跨模块并行化和流水线处理，实现了最高 2.09 倍的性能提升，优于现有的基于模块的框架。

ABSTRACT

Large language model (LLM)-based applications consist of both LLM and non-LLM components, each contributing to the end-to-end latency. Despite great efforts to optimize LLM inference, end-to-end workflow optimization has been overlooked. Existing frameworks employ coarse-grained orchestration with task modules, which confines optimizations to within each module and yields suboptimal scheduling decisions. We propose fine-grained end-to-end orchestration, which utilizes task primitives as the basic units and represents each query's workflow as a primitive-level dataflow graph. This explicitly exposes a much larger design space, enables optimizations in parallelization and pipelining across primitives of different modules, and enhances scheduling to improve application-level performance. We build Teola, a novel orchestration framework for LLM-based applications that implements this scheme. Comprehensive experiments show that Teola can achieve up to 2.09x speedup over existing systems across various popular LLM applications. The code is available at https://github.com/NetX-lab/Ayo.

研究动机与目标

• 解决由于粗粒度、基于模块的编排方式导致的基于大语言模型（LLM）应用端到端性能不佳的问题。
• 通过在原语操作层面建模工作流，而非高层模块，暴露更大的优化空间。
• 实现 LLM 与非 LLM 组件之间并行化、流水线处理和调度的联合优化。
• 通过将前端编排与后端执行解耦，并支持应用感知的调度，降低应用级延迟。
• 构建一个超越孤立 LLM 推理优化的系统，实现全面的、端到端的性能优化。

提出的方法

• 将每个基于大语言模型（LLM）的应用工作流表示为原语级数据流图，其中每个原语为带有元数据的符号化操作。
• 使用图优化器将用户查询解析为特定应用的原语图，并应用针对性的优化步骤以实现高效执行。
• 实现两级运行时调度器：上层调度器负责查询级执行图，下层调度器针对每个引擎负责批处理和相关原语的处理。
• 利用原语图中的拓扑深度和请求相关性，实现拓扑感知的批处理，提升资源利用率。
• 通过 Ray 与现有执行引擎（如 vLLM）集成，实现分布式调度与执行，保持与 LlamaIndex 和 LangChain 等框架的兼容性。
• 暴露引擎特定的属性以及原语之间的关系，以指导调度决策，并支持跨模块优化。

实验结果

研究问题

• RQ1 在原语级别进行细粒度编排是否能显著降低基于大语言模型（LLM）应用的端到端延迟？
• RQ2 与基于模块的链式结构相比，原语级别的数据流建模在调度与批处理方面有何改进？
• RQ3 在异构 LLM 工作流中，跨模块并行化与流水线处理能带来多大程度的性能增益？
• RQ4 基于原语依赖关系和拓扑深度的应用感知调度，是否能提升系统吞吐量与延迟表现？
• RQ5 在真实世界的大语言模型（LLM）应用工作负载中，Teola 的端到端优化与现有框架相比表现如何？

主要发现

• Teola 在多种 LLM 应用中实现了最高 2.09 倍的性能提升，涵盖增强搜索的生成任务与文档问答任务。
• 在某些基于检索增强生成（RAG）的工作流中，非 LLM 组件贡献了超过 50% 的端到端延迟，凸显了整体优化的必要性。
• 原语级数据流图相比基于模块的链式结构，暴露了显著更大的优化空间，支持跨模块并行化与流水线处理。
• 基于原语图结构的拓扑感知批处理可提升运行时调度器中的请求融合与资源利用率。
• 该框架的两级调度架构有效实现了编排逻辑与引擎执行的解耦，支持高效、应用感知的请求调度。
• Teola 的设计与现有 LLM 推理优化技术（如 vLLM、KV 缓存复用）兼容，可实现协同的性能增益。