PipeInfer: Accelerating LLM Inference using Asynchronous Pipelined Speculation

Large Language Models (LLMs) have become wildly popular in recent times, especially due to their versatility in tasks like language understanding and generation. Among the various architectures, decoder-only Transformer models like OpenAI’s GPT series [1] and Meta’s Llama family [2] have gained prominence. These models are comprised of a series of decoder layers, each of which is architecturally identical, sandwiched between an input embedding layer and an output layer [2, 3]. Unlike encoder-decoder models, decoder-only models focus solely on generating output sequences based on the input they receive. During inference, each layer feeds its outputs directly to the inputs of the successive layer. Their autoregressive nature means that for generating each output token, all layers need to be evaluated iteratively. This design introduces a significant challenge, as the model size exceeds the cache size of the target processor, and a memory bandwidth bottleneck becomes evident [4, 5]. This bottleneck impacts the model’s scalability and processing speed, posing challenges in real-time applications.

Despite these challenges, the autoregressive approach is preferred for certain applications due to its ability to maintain high levels of accuracy and contextual understanding in sequence generation. Current research is actively exploring solutions to mitigate these bottlenecks, including advancements in hardware design [6], more efficient architectures [7], and novel parallel processing techniques for exploiting batched inference [8]. Batched inference is less susceptible to bandwidth limitations because of increased temporal and spatial locality, improving processor cache hit rate at the expense of increased latency.

Recently, innovative techniques aimed at mitigating the memory bandwidth bottleneck in LLMs have emerged, drawing inspiration from the concept of CPU speculative execution, including SpecInfer [8] and Staged Speculative Decoding [9]. These techniques use a smaller secondary model to generate a tree of speculative sequences, that are then batched together and run through the target model. The key to this process is that inferencing on a tree of speculations yields probability distributions for all tokens within the tree, allowing inference to jump ahead several tokens at a time. The speedup shown by these techniques is due to the greater efficiency that batched processing yields. Building on this concept, other techniques like Medusa [10] and Lookahead Decoding [11] have emerged, which adapt the speculative decoding approach by modifying the generation and verification of speculations.

A significant challenge with these speculative techniques is the latency escalation, particularly evident in scenarios of low speculation accuracy. The crux of the issue lies in the balance between computational time and effectiveness. Although each inference run may take longer, the efficiency is substantially compromised if the speculation accuracy is low, leading to fewer correct and verified tokens. This imbalance results in the computational and time costs outweighing the benefits gained from correct speculations. Additionally, in heterogeneous systems such as smartphones with CPUs, GPUs, and NPUs, synchronization costs are too high to make full utilization of all processing elements. Speculation techniques can partially resolve this by running speculation on one processing element and verification on another, but current methods only run one phase at a time and thus still show limited utilization.

In this paper, we rectify these issues by modifying the speculative inference algorithm to run multiple verification runs simultaneously, exploiting a pipeline-parallel architecture to achieve high system utilization while maintaining low communication overhead. Our design achieves remarkable resilience to interconnect and computation latency, enabling high-speed inference on low-cost clusters of disparate commodity hardware. We believe future work can build on our design to achieve higher utilization in heterogeneous systems despite interconnect latency and large throughput differences. Finally, our design is not bound to the speculative inference algorithms. We believe that this approach can be extended to other acceleration techniques, such as Lookahead Decoding or Medusa speculation heads.

Compared to pipeline-parallel speculative inference, we observe roughly a 1.5-2.15$\times$ improvement in overall generation speed in our test cases while achieving near-zero slowdown for poor speculation accuracy. Testing with Gigabit Ethernet as the interconnect revealed a tolerance to latency and throughput limitations, increasing its improvement over speculative inference in such scenarios. For well-aligned models, we observed up to 1.7$\times$ faster generation speed than pipeline-parallel speculation, and for poorly aligned models, we observed up to a 2.15$\times$ improvement.

We introduce several key contributions to the field of speculative decoding in large language models:

• •

  Asynchronous Speculation and Inference: We enhanced speculative inference by integrating physically separate compute pipelines to run single-token inference or tree verification concurrently with the generation of the speculation tree. This modification enables simultaneous processing, significantly improving computational efficiency and reducing latency. Time-to-first-token latency reached near-parity with non-speculative iterative inference, while system utilization doubled.

• •

  Continuous Speculation: By leveraging asynchronous speculation, we devised a method of continuously generating speculations in small micro-batches rather than large single batches, improving the end-to-end latency and reducing the penalty for low speculation acceptance rates. With continuous speculation, we observed that the latency reduction was directly proportional to the reduction in batch size. We also observed continuous speculation allowed PipeInfer to adapt to low-bandwidth scenarios. Continuous speculation improved PipeInfer’s generation speed up to 1.5$\times.$

• •

  Pipelined KV Cache Multibuffering: To preserve the causality relationship of the generated tokens, we segmented the KV cache sequences into private sections for each speculative run. Cache operations are pipelined to maintain coherence during inference, allowing speculative runs to avoid computation of tokens shared by previous runs, even before they have been completed. By exploiting this ability, we improved computation throughput by a factor proportional to the alignment of the speculative model.

• •

  Early Inference Cancellation: We devised a method of flushing invalidated runs from the pipeline by back-propagating an asynchronous cancellation signal, reducing the performance impact of continuous speculation with poorly aligned speculation models. Somewhat counter-intuitively, we observed greater speedups up to 2.15$\times$ for poorly aligned models thanks to early inference cancellation.

Large Language Models based on the Transformer decoder architecture, including Llama 2 [2] and GPT [1], are built up from a series of decoder layers. Each decoder layer contains an attention module and a multi-layer perceptron [3]. The attention module calculates the scores for all tokens in the sequence, but the key and value vectors are oftentimes cached to prevent needless recomputation.

During inference, each decoder layer is evaluated in sequence, requiring the weights for the attention module and MLP to be loaded, as well as the relevant entries from the KV cache. The sheer size of most modern models necessitates multiple gigabytes of memory transfers during a single inference run, resulting in a memory bandwidth bound for small batch sizes [9]. This bandwidth limitation mirrors similar constraints in CPU pipelines, causing significant stalls as the processing elements wait for values from memory [4]. Inspired by CPU speculative execution, many similar designs for speculative inference were created to alleviate this bottleneck [9, 8, 10].

Speculative decoding has previously been explored through SpecInfer [8] and Staged Speculative Decoding [9].

SpecInfer utilized multiple speculative models executed in parallel with each other, but did not run the speculative and target models in parallel, resulting in increased end-to-end latency. Staged Speculative Decoding took a slightly different approach, where the speculative models were themselves speculatively inferenced, improving overall speed but incurring even greater latency penalties, as the staged speculations were also not run in parallel with their respective targets.

Medusa [10] takes a slightly different approach to speculative decoding, adding new sampling heads to the target model that produce the speculations without a secondary model. Medusa does not incur significant latency overhead but requires training new sampling heads for the target model.

Lookahead Decoding [11] takes a different approach entirely, opting to use Jacobi iteration to generate multiple tokens simultaneously by generating N-grams based on the trajectory of the current tokens. The generated n-grams are cached and later verified in a separate stage. Lookahead Decoding exhibits high utilization and low latency on single-node systems, but does not account for multi-node systems or systems with slow interconnects.

A recent work named SPEED [13] uses similar speculative techniques to enhance decoding efficiency, but instead targets inference on a single GPU. Additionally, SPEED generates the speculations from the hidden states of previous layers, eliminating the need for a separate speculative model. This approach contrasts with PipeInfer, which utilizes a secondary model for speculation generation. While this difference impacts the computational overhead, it also influences the flexibility and adaptability of the system to different modeling scenarios. Moreover, invalidation of speculations results in a pause for the entire system in SPEED, while PipeInfer continues inferencing all other valid runs at the same time invalidation and flushing operations are executed. Finally, SPEED targets specialized models implementing parameter sharing, while PipeInfer works across a variety of off-the-shelf models.

Exploiting pipeline parallelism to improve inference latency while simultaneously running speculation has been explored previously [14]. However, this work performed speculation via similar methods to SPEED, predicting a single next token from the hidden states of intermediate layers, requiring additional sampling heads to be trained. Additionally, this method only speculates one token in advance at a time, limiting the possible speedup compared to other methods utilizing separate speculative models or heads.

Work has also been done using tensor parallelism, distributing the sub-layer operations across the nodes instead of whole layers [15]. However, such methods have been found to suffer from extreme interconnect bandwidth bottlenecks, even for extremely slow compute nodes like Raspberry Pi 4s. Even as few as eight nodes incurred higher synchronization time than the time spent running the inference itself.

Non-speculative techniques have also been explored in the area of inference acceleration. A prominent technique is based around early exit inference, in which some later model layers are skipped if the output from an earlier layer can be used instead. CALM [16], for example, trains a classifier model to detect when an earlier layer is sufficiently confident in its output. Similarly, Depth-Adaptive Transformers [17] utilize classifiers to predict the depth at which inference can halt, either per-token or per-sequence. Yet another early-exit style strategy called LITE [18] accelerates inference while maintaining output accuracy, a problem that plagues other similar systems. Early exit decoding shows promise but requires training classifiers for the target model.

Other non-speculative acceleration techniques focus on modifying the precision at which inference is conducted. Such techniques include quantization strategies like AWQ [19] and QuIP [20] or pruning strategies like SparseGPT [21] and SparseML [22]. Still other techniques such as SqueezeLLM [23] combine sparsity and quantization. Quantization and pruning approaches generally require an offline conversion step to compress the original model weights but can be substantially less intensive than a full pre-training run.

PipeInfer does not require any pre-training or conversion steps, but also does not conflict with quantization or pruning techniques.

Our proposed method, PipeInfer, enhances speculative inference and is composed of four major components: Asynchronous Speculation, Continuous Speculation, Pipelined KV Cache Multibuffering, and Early Inference Cancellation. Our reference implementation is built on top of llama.cpp [12], with inter-node communication accomplished with MPI [24]. An overall system diagram is shown in Figure 1.

PipeInfer enhances the efficiency and processing speed of LLMs by modifying the speculative inference algorithm to support multiple verification runs within a pipeline-parallel architecture. We achieved not only optimized system utilization and minimized communication overhead but also demonstrated exceptional resilience to latency and computational delays, especially in cost-effective, heterogeneous hardware environments. The results show a remarkable 2.15$\times$ improvement in generation speed, maintaining high performance even with low speculation accuracy. Additionally, our approach exhibits robustness in latency and throughput-constrained environments, achieving high CPU utilization.

PipeInfer exhibits far-improved performance in multiple scenarios, and future work may extend it to other inference acceleration strategies, further improving its performance. We believe that Lookahead decoding [11] and Medusa [10] would benefit greatly from PipeInfer augmentation. We also believe self-speculation techniques like SPEED [13] or PPD [14] could complement PipeInfer’s external speculative model approach.

PipeInfer may also be extended to support hybrid parallelization via multi-GPU nodes, applying tensor parallelism at the local node level and maintaining pipeline parallelism across the cluster. Alternatively, bandwidth bottlenecks resulting from the PCIe bus could be alleviated through the application of PipeInfer on a single node.

Bottlenecks in a pipeline can also be mitigated by adding new nodes in parallel with the slowest nodes, acting as load-balancers. When the primary node is busy, the secondary nodes can take over the computation of the designated layers. PipeInfer’s use of many independent and simultaneous inference runs allows such load-balancing without affecting the end results, so long as the ordering remains consistent.