[Paper Review] 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 Systems4 citations

TL;DR

Teola proposes a fine-grained, primitive-level dataflow graph orchestration framework for LLM-based applications, enabling end-to-end optimization across both LLM and non-LLM components. By exposing low-level task primitives and their dependencies, Teola achieves up to 2.09x speedup through topology-aware batching and cross-module parallelization and pipelining, outperforming existing module-based frameworks.

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.

Motivation & Objective

• To address the suboptimal end-to-end performance in LLM-based applications caused by coarse-grained, module-level orchestration.
• To expose a larger optimization space by modeling workflows at the primitive operation level instead of high-level modules.
• To enable joint optimization of parallelization, pipelining, and scheduling across LLM and non-LLM components.
• To improve application-level latency by decoupling frontend orchestration from backend execution and enabling application-aware scheduling.
• To build a system that supports holistic, end-to-end performance optimization beyond isolated LLM inference improvements.

Proposed method

• Represent each LLM-based application workflow as a primitive-level dataflow graph, where each primitive is a symbolic operation with metadata.
• Use a Graph Optimizer to parse user queries into application-specific primitive graphs and apply targeted optimization passes for efficient execution.
• Implement a two-tier runtime scheduler: an upper-tier scheduler for query-level execution graphs and a lower-tier scheduler per engine for batching and processing related primitives.
• Leverage topological depth and request correlations in the primitive graph to enable topology-aware batching and improve resource utilization.
• Integrate with existing execution engines (e.g., vLLM) via Ray for distributed scheduling and execution, maintaining compatibility with frameworks like LlamaIndex and LangChain.
• Expose engine-specific attributes and relationships between primitives to guide scheduling decisions and enable cross-module optimizations.

Experimental results

Research questions

• RQ1 Can fine-grained orchestration at the primitive level significantly reduce end-to-end latency in LLM-based applications?
• RQ2 How does primitive-level dataflow modeling improve scheduling and batching compared to module-level chaining?
• RQ3 To what extent can cross-module parallelization and pipelining be exploited for performance gains in heterogeneous LLM workflows?
• RQ4 Can application-aware scheduling based on primitive dependencies and topological depth improve system throughput and latency?
• RQ5 How does Teola’s end-to-end optimization compare to existing frameworks in real-world LLM application workloads?

Key findings

• Teola achieves up to 2.09x speedup over existing systems across diverse LLM applications, including search-empowered generation and document question answering.
• Non-LLM components contribute to over 50% of end-to-end latency in some RAG-based workflows, highlighting the need for holistic optimization.
• Primitive-level dataflow graphs expose a significantly larger optimization space than module-based chains, enabling cross-module parallelism and pipelining.
• Topology-aware batching based on primitive graph structure improves request fusion and resource utilization in the runtime scheduler.
• The framework’s two-tier scheduling architecture effectively decouples orchestration logic from engine execution, enabling efficient, application-aware request scheduling.
• Teola’s design is compatible with existing LLM inference optimizations (e.g., vLLM, KV cache reuse), allowing synergistic performance gains.