Seesaw: High-throughput LLM Inference via Model Re-sharding

Seesaw: High-throughput LLM Inference via Model Re-sharding

Speakers: Qidong Su, University of Toronto and Santa Mel; Wei Zhao; Xin Li; Chenhao Jiang; Gennady Pekhimenko

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The talk "Seesaw: High-throughput LLM Inference via Model Re-sharding" introduces a novel framework designed to significantly accelerate throughput-oriented offline large language model (LLM) text generation in distributed GPU environments. Presented by Qidong Su from the University of Toronto and Santa Mel, this work, a collaboration with researchers from the University of Toronto, the Vector Institute, Santa Mel, and Stanford University, addresses a fundamental inefficiency in current LLM inference systems. The core insight driving Seesaw is that the two primary phases of LLM text generation—pre-fill (processing input prompts) and decode (generating tokens one by one)—exhibit fundamentally different computational characteristics and thus benefit from distinct parallelization strategies.

Traditional LLM inference engines often employ a single, fixed parallelization strategy across all stages, leading to suboptimal performance for at least one of the phases. Seesaw challenges this paradigm by proposing a dynamic approach: it re-shards the model across GPUs and switches between pipeline parallelism (PP) for the pre-fill phase and tensor parallelism (TP) for the decode phase. This adaptive strategy, aptly named "Seesaw" due to its back-and-forth nature, is complemented by an innovative tiered KV cache buffering mechanism that leverages both CPU and GPU memory to mitigate the overhead of frequent re-sharding.

The significance of Seesaw lies in its ability to unlock substantial performance gains for applications heavily reliant on batch processing of LLM requests, such as large-scale data synthesis, benchmarking, and complex codebase analysis. These offline use cases prioritize overall throughput and end-to-end runtime over individual token latency, making Seesaw's optimizations particularly impactful. The framework demonstrates a remarkable 36% throughput increase compared to state-of-the-art systems like VLM, marking a significant advancement in efficient LLM deployment and operation in resource-intensive, distributed settings.

Background

▶ Watch: Introduction: Seesaw for high-throughput LLM inference (0:00)

The landscape of large language models has evolved rapidly, with model sizes continuously expanding to billions and even trillions of parameters. This exponential growth means that many contemporary LLMs can no longer fit onto a single GPU, necessitating distributed deployment across multiple accelerators. Efficiently partitioning and orchestrating these massive models across a cluster of GPUs is a critical challenge, giving rise to various parallelization strategies.

Two prevalent methods for distributing models are tensor parallelism (TP) and pipeline parallelism (PP). Tensor parallelism involves splitting individual layers of the neural network across different GPUs. Each GPU computes a portion of the layer's output, and these partial results are then aggregated, typically requiring significant communication between GPUs within a layer. Pipeline parallelism, on the other hand, slices the model vertically into disjoint stages, with each GPU or group of GPUs responsible for a sequential segment of the model layers. Execution is then pipelined, where different stages process different batches (or microbatches) concurrently.

Beyond model partitioning, understanding the distinct phases of LLM text generation is crucial. Text generation typically involves two main stages:

1. Pre-fill: This phase involves processing the input prompt, which can range from a few tokens to several thousands. During pre-fill, the model computes the initial KV (Key-Value) cache for the entire prompt. Because the input sequence can be long and often processed in large batches, this phase is typically compute-bound, meaning its performance is limited by the computational capacity of the GPUs.
2. Decode: Following pre-fill, the model enters the decode phase, where it generates output tokens one by one. Each new token is generated by attending to the previously generated tokens and the input prompt, reusing the KV cache from prior steps. Since only one new token is produced per step, and the primary operation involves fetching and processing cached data, this phase is largely memory-bound, with performance often dictated by memory bandwidth and the efficiency of KV cache access. To maximize throughput during decode, it's essential to batch multiple sequences together, allowing for more efficient memory access patterns.

The application context also significantly influences the choice of optimization targets. Online inference, exemplified by chatbots, prioritizes low latency and fast time-to-first-token (TTFT) and time-between-tokens (TBT), as users expect instant responses. In contrast, offline inference deals with large sets of requests that can be processed in parallel. Major LLM API providers often offer cheaper batch APIs for such tasks. Applications like benchmark execution, large-scale codebase analysis, and dataset synthesis (where LLMs generate vast amounts of data) fall under this category. For offline inference, the primary performance metric is throughput (tokens per second or requests per hour) or the end-to-end runtime for processing the entire dataset. It is for these throughput-oriented offline scenarios that Seesaw is specifically designed.

The core problem arises from the fact that existing LLM inference systems, like VLM, typically commit to a single, fixed parallelization strategy (e.g., pure tensor parallelism or pure pipeline parallelism) throughout both the pre-fill and decode phases. However, as the research shows, the compute-bound nature of pre-fill and the memory-bound nature of decode inherently favor different strategies, leading to a compromise in overall efficiency.

Key Findings

▶ Watch: Model parallelism (tensor, pipeline) and generation stages (pre-fill, decode) (2:00)

The central discovery of the Seesaw project is that the pre-fill and decode stages of large language model inference exhibit distinct computational bottlenecks, making them optimally suited for different distributed parallelization strategies. This fundamental observation forms the bedrock of Seesaw's design and its significant performance improvements.

Specifically, the key findings are:

• Divergent Parallelism Preferences:
• Pre-fill is primarily compute-bound due to processing potentially long input sequences in large batches. It benefits significantly from pipeline parallelism (PP) because PP effectively reduces communication overhead, which becomes a bottleneck when compute is abundant.
• Decode is primarily memory-bound as it generates tokens one by one, heavily relying on efficient access to the KV cache. It benefits from tensor parallelism (TP) due to TP's ability to accelerate memory access (weight loading) by distributing memory bandwidth requirements across multiple GPUs.
• Dynamic Re-sharding as an Optimal Strategy: Based on these divergent preferences, Seesaw proposes that the most effective way to maximize throughput for LLM inference is to dynamically switch between these parallelization strategies. It employs pipeline parallelism during the pre-fill phase and then re-shards the model to tensor parallelism for the decode phase. This dynamic adaptation ensures that each stage operates under its most favorable parallelization scheme.
• Re-sharding Overhead Mitigation: A critical challenge in dynamic re-sharding is the significant overhead associated with moving tens or hundreds of gigabytes of model weights between GPUs during transitions. The widely used continuous batching scheduling policy, which frequently interleaves pre-fill and decode operations, exacerbates this problem, making re-sharding a bottleneck.
• Tiered KV Cache Buffering: To address the re-sharding overhead, Seesaw introduces an innovative tiered KV cache buffering strategy. This mechanism leverages both CPU memory and GPU memory for KV cache storage. By intelligently offloading KV cache from GPU to CPU during pre-fill and reloading it during decode (overlapping these swaps with communication), Seesaw significantly reduces the frequency of costly re-sharding events while maintaining high GPU memory utilization. This strategy ensures that the GPU remains heavily utilized for decode operations, even when new requests are not immediately ready for pre-fill, thus optimizing overall throughput.
• Superior Throughput Performance: Through extensive evaluation, Seesaw demonstrates a substantial improvement in end-to-end throughput. It achieves up to a 36% throughput increase compared to VLM, a state-of-the-art LLM inference engine that relies on a fixed parallelization strategy. On L4 GPUs with an archive summarization dataset, Seesaw specifically showed a 30% throughput increase. This significant gain validates the efficacy of dynamic re-sharding combined with intelligent KV cache management for throughput-oriented offline LLM inference.

Technical Deep Dive

▶ Watch: Key insight: Pre-fill and decode favor different parallelisms (3:20)

Seesaw's innovation stems from a deep understanding of the computational characteristics of LLM inference phases and the intricate trade-offs between different parallelization strategies.

Analysis of Parallelization Strategies

The talk provides a detailed breakdown of why tensor parallelism (TP) and pipeline parallelism (PP) behave differently in pre-fill and decode:

• Tensor Parallelism (TP):
• Memory Access & Computation: When scaling from one to two GPUs using TP, both memory bandwidth and compute FLOPS are effectively doubled. This accelerates memory access (loading model weights) and computation by a factor of two.
• Communication Overhead: The major drawback of TP is its heavy communication overhead. Since each layer's computation is split, partial results must be aggregated across GPUs after each sub-computation, leading to frequent and substantial data transfers. This can become a bottleneck, especially for compute-bound tasks where GPUs might spend more time waiting for data than computing.
• Decode Advantage: In the memory-bound decode phase, the ability of TP to accelerate weight loading (memory access) is highly beneficial, as it alleviates the memory bandwidth bottleneck.

• Pipeline Parallelism (PP):
• Communication Overhead: PP generally has lower communication overhead compared to TP. Communication primarily occurs as peer-to-peer transfers between adjacent stages in the pipeline, typically sending intermediate activations from one GPU to the next.
• Microbatching Problem: A significant challenge with PP is the microbatching requirement. To fully utilize all GPUs in a pipeline, the overall batch of requests must be partitioned into smaller microbatches. Without microbatching, GPUs in later stages would often sit idle, waiting for the output from earlier stages (the "pipeline bubble"). For example, with two GPUs and four requests, the first GPU processes the first half of the model for the first microbatch, then sends its output to the second GPU. During this transfer, the first GPU cannot immediately start processing the next token for the same request. It must wait for the first token to complete or start processing a different microbatch. This means that in any given time step, only 1/PP fraction of the total requests can be actively processed by a specific GPU within its stage. Consequently, even though each GPU holds only 1/PP fraction of the model weights, the effective "loading weights" time for the entire batch is not proportionally accelerated, as it takes PP times more steps to process the full batch compared to TP for certain operations.
• Pre-fill Advantage: In the compute-bound pre-fill phase, where communication is a primary bottleneck, the reduced communication overhead of PP makes it a more favorable strategy.

Seesaw's Dynamic Re-sharding Strategy

The core of Seesaw's approach is to exploit these differences:

1. Pre-fill Phase: Seesaw configures the distributed system to use pipeline parallelism. This minimizes communication overhead, which is critical for the compute-intensive processing of large input prompts.
2. Decode Phase: Once pre-fill is complete, Seesaw re-shards the model weights across the GPUs to switch to tensor parallelism. This maximizes memory access efficiency, which is crucial for the memory-bound, token-by-token generation.

The entire process resembles a "seesaw" as the system dynamically transitions between these two optimal configurations.

Addressing Re-sharding Overhead: Tiered KV Cache Buffering

While dynamic switching is theoretically optimal, the practical challenge lies in the cost of re-sharding. Moving tens to hundreds of gigabytes of model weights between GPUs is an expensive operation that can quickly negate the benefits if done too frequently.

• Problem with Continuous Batching: The prevalent continuous batching scheduling policy, while efficient for individual requests, leads to a high re-sharding frequency. In continuous batching, the system tries to keep GPUs busy by constantly admitting new pre-fill requests as soon as memory is available (e.g., when some decode requests finish and free up KV cache space). This constant interleaving of pre-fill and decode phases means that even the completion of a single request can trigger a pre-fill-to-decode or decode-to-pre-fill transition, making re-sharding the bottleneck.
• Suboptimal Decode-Prioritizing Scheduling: An alternative, "decode prioritizing" scheduling, could reduce re-sharding frequency by processing an entire batch of decode requests before switching back to pre-fill. However, this often leads to significant GPU memory underutilization during decode, as memory freed by completed requests isn't immediately filled, resulting in suboptimal throughput.

Seesaw's solution is tiered KV cache buffering:

• Dual-Tier Storage: Seesaw utilizes both CPU memory (host memory) and GPU memory (device memory) for storing the KV cache.
• Pre-fill Offloading: During the pre-fill phase, if GPU memory becomes constrained, Seesaw intelligently offloads some of the KV cache from the GPU to the CPU.
• Decode Reloading: During the decode phase, if the system needs to refill GPU memory (e.g., to accommodate more active sequences or to prepare for an upcoming pre-fill), it reloads KV cache directly from the CPU back to the GPU.
• Overlapped Swaps: Crucially, these CPU-GPU KV cache swaps are overlapped with communication, ensuring minimal additional overhead.
• Optimized Transition Logic: From a memory utilization perspective, Seesaw switches from pre-fill to decode only when both CPU and GPU memory are fully utilized. During decode, if GPU memory needs refilling, it prioritizes reloading KV cache from the CPU. Only once all KV cache from CPU is loaded back to GPU, and more capacity is needed, does the system consider switching back to a new pre-fill phase. This strategy ensures:
• High GPU Memory Utilization: The GPU memory remains highly utilized for decoding, which is essential for maximizing decode throughput.
• Low Re-sharding Frequency: By buffering KV cache on the CPU, Seesaw significantly delays the need to switch back to a new pre-fill phase, thereby minimizing the costly re-sharding operations.

This sophisticated interplay between dynamic parallelization and intelligent memory management allows Seesaw to achieve high throughput while effectively managing the complexities of distributed LLM inference.

Experimental Setup & Results

▶ Watch: Seesaw's dynamic re-sharding strategy for optimal performance (6:00)

To validate the effectiveness of Seesaw, the researchers conducted a series of comprehensive experiments, comparing its performance against a state-of-the-art LLM inference engine, VLM.

Experimental Setup

• Hardware: The evaluation was performed on different GPU hardware, with specific results reported for L4 GPUs, a common choice for LLM inference due to their balance of performance and cost efficiency. The setup involved multiple GPUs in a distributed configuration, though the talk primarily focuses on single-node distributed scenarios. The CPU memory used for KV cache buffering was as large as the summation of the GPU memory in the node, which the authors deem a "reasonable setup" given the comparative cost of CPU vs. GPU memory.
• Models: Seesaw was tested with various LLM sizes, including a 15 billion parameter model, which was used for detailed runtime breakdown analysis.
• Baselines: The primary baseline was VLM, a widely recognized and high-performing LLM engine. For VLM, the researchers selected its optimal fixed parallelization strategy (either pure tensor parallelism or pure pipeline parallelism) for each specific workload and model size to ensure a fair comparison against the best possible static configuration.
• Seesaw Configuration: Seesaw was configured with chunk pre-fill enabled, and the optimal chunk size was carefully tuned for the specific experimental conditions.
• Workloads and Datasets: Two distinct types of workloads were used to represent common LLM inference patterns:

1. Summarization Dataset: Characterized by inputs that are significantly longer than the generated outputs (e.g., the Archive summarization dataset). This type of workload heavily exercises the pre-fill phase.
2. ShareGPT Dataset: Features inputs and outputs of roughly similar lengths. This dataset provides a balanced test for both pre-fill and decode phases.

Headline Results

The evaluation demonstrated significant performance gains for Seesaw:

• Overall Throughput Increase: Seesaw achieved an impressive 36% throughput increase compared to VLM across various test scenarios.
• L4 GPU Performance: Specifically, on L4 GPUs with the Archive summarization dataset, Seesaw delivered a 30% throughput increase over VLM. This highlights Seesaw's effectiveness in scenarios with long inputs, where pre-fill optimization is critical.

Runtime Breakdown Analysis

To provide deeper insight into the source of these speedups, the researchers presented a runtime breakdown for a 15 billion parameter model:

• The breakdown clearly illustrated that during the pre-fill phase, pipeline parallelism (PP) consistently outperformed other strategies, confirming the analytical prediction that PP's reduced communication overhead is advantageous for compute-bound operations.
• Conversely, during the decode phase, tensor parallelism (TP) was shown to be the fastest, validating the hypothesis that TP's efficient memory access and weight loading capabilities are superior for memory-bound, token-by-token generation.
• Seesaw, by dynamically combining these optimal strategies, achieved the fastest execution in both phases, leading to its superior overall end-to-end throughput.

Applicability and Limitations Discussed

The talk also touched upon the applicability of Seesaw and its current limitations:

• Latency-Oriented Tasks: The speaker explicitly stated that Seesaw is designed for throughput-oriented offline tasks and is not the optimal solution for latency-sensitive online applications. While the underlying principles of analyzing memory access, compute, and communication overhead still apply, a different system design would be needed for online scenarios.
• Prompt Length Sensitivity: The talk acknowledged that the optimal strategy might depend on the prompt length and input/output ratio. A sensitivity study, detailed in the paper, confirms that Seesaw's technique remains applicable in most common cases.
• Multi-Node Scenarios: The current implementation of Seesaw is primarily for single-node distributed scenarios. Extending it to multi-node environments presents additional challenges, such as designing faster, more centralized storage for efficient weight and KV cache exchange during re-sharding. The speaker noted that tensor parallelism is particularly hard to scale beyond a single node, suggesting that pipeline parallelism or expert parallelism would become "a must-have" in such larger deployments, and further research would be needed to see how Seesaw's dynamic approach extends.

These results unequivocally demonstrate Seesaw's capability to significantly enhance the efficiency of LLM inference in throughput-critical applications by intelligently adapting its parallelization strategy and managing KV cache across memory tiers.

Practical Implications

▶ Watch: Addressing re-sharding overhead and scheduling challenges (6:50)

Seesaw's dynamic re-sharding and tiered KV cache buffering offer significant practical implications for various stakeholders involved in the deployment and operation of large language models.

For Practitioners and Model Deployers

• Cost Reduction and Throughput Maximization: For applications that require processing large volumes of LLM requests (e.g., data synthesis, large-scale content generation, batch sentiment analysis, or code generation for entire repositories), Seesaw directly translates to higher throughput. This means more work can be done in less time, or the same amount of work can be completed using fewer GPU resources, leading to substantial cost savings on cloud infrastructure or more efficient utilization of on-premise clusters.
• Enhanced Capability for Offline Tasks: Seesaw makes it more feasible and economical to leverage LLMs for tasks traditionally performed offline, where latency is less critical than overall processing speed. This could unlock new use cases and expand the reach of LLM-powered automation.
• Optimized Resource Utilization: By dynamically adapting parallelization strategies and intelligently managing KV cache across CPU and GPU memory, Seesaw ensures that expensive GPU resources are utilized more effectively, minimizing idle time and maximizing computational efficiency.
• Trade-offs: Practitioners must understand that Seesaw is specifically optimized for throughput-oriented offline inference. It is not designed for, nor would it be optimal for, low-latency online interactive applications like chatbots, where time-to-first-token is paramount. Deployers need to choose the right inference engine based on their specific workload requirements.

For Infrastructure Teams

• Advanced System Design: Seesaw provides a blueprint for building more sophisticated and adaptive distributed inference systems. It highlights the importance of not only choosing the right parallelization strategy but also dynamically adjusting it based on the workload phase.
• Intelligent Memory Management: The tiered KV cache buffering mechanism underscores the value of heterogeneous memory architectures (CPU + GPU) for LLM inference. Infra teams can design their systems to leverage host memory as an extension of GPU memory, optimizing KV cache management to reduce costly data transfers and re-sharding operations. This approach can inform future hardware and software co-design efforts.
• Complexity Management: Implementing dynamic re-sharding and tiered memory management adds a layer of complexity to the inference stack. Infrastructure teams will need robust orchestration, monitoring, and debugging tools to manage these dynamic transitions and ensure system stability. The benefits in throughput, however, often justify this increased complexity for high-volume workloads.
• Future-Proofing: As LLMs continue to grow, dynamic adaptation will become increasingly important. Seesaw offers insights into how to build systems that can scale and perform efficiently even with ever-larger models and diverse workloads.

For Model Builders and Researchers

• Performance Bottleneck Awareness: The work reinforces the idea that different parts of an LLM inference pipeline have different performance bottlenecks. Model builders can consider these insights when designing new architectures or training paradigms, potentially optimizing models for easier dynamic re-sharding or more efficient KV cache management.
• Beyond Static Parallelism: Seesaw pushes the boundaries of distributed ML, moving beyond static parallelization choices to a more adaptive, intelligent approach. This opens up new avenues for research into dynamic resource allocation, workload-aware scheduling, and heterogeneous computing in LLM inference.
• Hardware-Software Co-design: The interplay between CPU and GPU memory for KV cache highlights the importance of hardware-software co-design. Researchers can explore how future hardware architectures could better support such tiered memory systems or accelerate re-sharding operations.

In essence, Seesaw demonstrates that significant performance gains are achievable in LLM inference by moving away from "one-size-fits-all" parallelization strategies towards a more intelligent, adaptive system that understands and responds to the varying demands of different computational phases. While introducing complexity, the throughput benefits for offline, batch-oriented LLM applications are substantial, making it a critical advancement for the field.

Key Takeaways

• Phase-Specific Optimization: The pre-fill and decode phases of LLM text generation have distinct computational characteristics (compute-bound vs. memory-bound) and thus favor different parallelization strategies.
• Dynamic Parallelism is Key: Pipeline Parallelism (PP) is optimal for the compute-bound pre-fill phase due to its lower communication overhead, while Tensor Parallelism (TP) is optimal for the memory-bound decode phase due to its efficient memory access and weight loading capabilities.
• Seesaw's Core Strategy: Seesaw dynamically switches between PP for pre-fill and TP for decode through model re-sharding, leading to superior performance in both stages.
• Tiered KV Cache for Efficiency: To mitigate the high overhead of frequent re-sharding, Seesaw employs tiered KV cache buffering, leveraging both CPU and GPU memory. This strategy reduces transition frequency while maintaining high GPU utilization.
• Significant Throughput Gains: Seesaw achieves up to a 36% throughput increase compared to state-of-the-art fixed-strategy systems like VLM, demonstrating its effectiveness for throughput-oriented offline LLM inference.
• Targeted Application: This innovation is specifically designed for and highly beneficial to offline, batch-oriented LLM applications (e.g., data synthesis, benchmarking) where overall throughput is the primary metric, rather than low-latency online inference.

About the Speaker(s)

The talk was presented by Qidong Su, representing the University of Toronto and Santa Mel. He is one of the key researchers behind the Seesaw project, a collaborative effort involving experts from the University of Toronto, the Vector Institute, Santa Mel, and Stanford University. His work focuses on optimizing large model inference in distributed computing environments, with a particular emphasis on throughput-oriented scenarios.