vLLM x AMD: Efficient LLM Inference on AMD Instinct™ MI300X GPUs

Overview

In the previous blog, we learned that prefix caching significantly enhances the efficiency and throughput of language models, particularly for smaller sequence lengths. With APC (Advanced Prefix Caching) enabled, the number of input and output tokens per second increases notably, especially within the sequence length range of 128 to 8192 for input and 128 to 4096 for output.

This blog will explore how chunked prefill significantly boosts the performance of language models on AMD GPUs. By breaking down input sequences into smaller chunks and pre-computing outputs, chunked prefill enhances both efficiency and coherence in text generation tasks.

What is Chunked Prefill?

Chunked prefill is a technique often used in language modeling and text generation to improve the efficiency and coherence of generated outputs.

Specifically, Chunking means that the input text is divided into smaller, manageable segments or "chunks". Instead of processing an entire long sequence of text simultaneously, the model handles these chunks individually.

Prefill refers to the model’s ability to pre-compute or cache outputs for these chunks, enabling faster generation when similar or overlapping input sequences are encountered.

When to Apply Chunked Prefill

Chunked prefill is particularly useful in scenarios where there are multiple concurrent requests happen, for example:

Interactive Applications: Chunked prefill can be applied in chatbots or virtual assistants where conversations can be segmented into manageable sections.

Content Generation: In applications like automated writing or summarization, chunked prefill can indirectly benefit the user experience by providing faster and more consistent responses, which can be crucial in real-time applications. By segmenting the prefill into smaller, manageable chunks, each segment can generate with the right context, thereby improving the overall quality of the output.

Benefits of Chunked Prefill

Higher Efficiency: By processing smaller chunks, the model can utilize cached outputs for previously computed segments, reducing computation time for similar prompts.

Better Memory Management: Handling smaller segments can help manage memory more effectively in environments with limited resources or when dealing with long contexts.

Improved Coherence: Chunked processing can help maintain coherence in generated text by allowing the model to focus on relevant segments without getting overwhelmed by lengthy inputs.

Boosting Throughput with Chunked Prefill

The chart below demonstrates the performance of the Llama3.1-70B model using chunked prefill at different sequence lengths. The main finding is that the average request throughput speedup is between 1.05X and 1.4X. This indicates significant efficiency improvements in text generation, highlighting the benefits of chunked prefill in enhancing model performance, particularly for high throughput and efficient memory management.

Effects of max_num_batched_tokens

In vLLM, max_num_batched_tokens refers to the maximum number of tokens that can be processed in a single batch. By default, this value is set as 512. Adjusting this parameter can significantly impact the performance of the model in various ways:

1. Inter-Token Latency (ITL): Smaller values for max_num_batched_tokens generally result in better ITL because there are fewer prefills interrupting the decoding process.
2. Throughput: Larger values for max_num_batched_tokens can improve throughput by allowing more tokens to be processed in each batch.
3. Time to First Token (TTFT): Higher values for max_num_batched_tokens can also improve TTFT by batching more prefill requests together.
4. Memory Utilization: Adjusting max_num_batched_tokens affects how GPU memory is utilized. Higher values may require more memory, so it’s important to balance this with the available resources to avoid preemptions and recomputations.
5. Value Adjustment:
   • Higher Values: Increase throughput by allowing more tokens to be processed in parallel but may increase latency if memory is insufficient.
   • Lower Values: Improve ITL (Inter-Token Latency) by reducing interruptions during decoding but may decrease throughput.

In summary, tuning max_num_batched_tokens parameter allows the optimization of the model performance to meet specific requirements, whether it involves reducing latency, increasing throughput, or balancing memory usage.

When chunked prefill is enabled, the scheduling policy prioritizes decode requests. It batches all pending decode requests before scheduling any prefill. If a prefill request cannot fit into the max_num_batched_tokens, it is chunked.

Performance Test with Variable max_num_batched_tokens

According to the official vLLM user guide, users can tune the performance by adjusting the max_num_batched_tokens. As mentioned above, it is set to be 512 by default, but its suitability for MI300X GPU is uncertain, so we decided to run the experimental test to verify. The test has been conducted on AMD Instinct™ MI300X accelerator, with ROCm^TM software 6.2, AMD EPYC 9654 96-Core CPU processor, and Ubuntu 22.04.4 LTS.

As shown in Figure 5:

• Throughput increases from 2373 to 3774 tokens/s as max_num_batched_tokens varies from 256 to 4096.
• Throughput stabilizes and does not significantly increase beyond 4096 max_num_batched_tokens.
• Optimal throughput is achieved when max_num_batched_tokens is set to greater than 2048.
• Users need to determine the proper value for max_num_batched_tokens based on their application scenario to avoid performance degradation.

It can be concluded that for each application scenario, users need to determine the proper value of this feature when applied, because both smaller and larger number would result in performance degradation. It is recommended to set it as > 2048 for optimal throughput.

Latency Analysis

In this analysis, we compared the performance variations by enabling and disabling chunked prefill under different input and output sequence length conditions. For latency, the lower, the better!

With Chunked Prefill
The P99 Latency increases significantly as the in_out_seq_length increases.

Without Chunked Prefill
The P99 Latency also increases with the in_out_seq_length; however, the values are significantly lower compared to the chunked prefill scenario.

Overall, the P99 Latency increases with the in_out_seq_length in both scenarios, but the increase is much more pronounced with using chunked prefill. Without chunked prefill, the latency values are notably lower and increase at a more gradual rate. In summary, the chunked prefill feature may lead to increased latency under specific conditions, particularly when the system configuration and workload are not optimally balanced. Here are some scenarios where latency might increase:

Insufficient GPU Memory
If the GPU memory is not sufficient to handle the larger batch sizes created by chunked prefill, it can lead to frequent preemptions and recomputations. This overhead can increase the end-to-end latency.

High Preemption Rate
When the system frequently encounters preemptions due to insufficient KV cache space, the recomputation of pre-empted requests can adversely affect latency. This is more likely to happen if the gpu_memory_utilization is set too low.

Suboptimal max_num_batched_tokens
Setting the max_num_batched_tokens too high can lead to inefficient GPU utilization and increased latency.

Large Model Sizes
For very large models, the overhead of managing multiple smaller chunks can become significant, leading to increased latency.

To mitigate these issues, it's crucial to tune the parameters such as gpu_memory_utilization, max_num_batched_tokens, and tensor_parallel_size based on your specific hardware and workload. Proper tuning can help achieve a balance between throughput and latency, ensuring optimal performance.

Conclusion

In part 2 of the AMD vLLM blog series, we delved into the performance impacts of using vLLM chunked prefill for LLM inference on AMD GPUs. Our findings indicated that while chunked prefill can lead to significant latency increases, especially under conditions of high preemption rates or insufficient GPU memory, careful tuning of system parameters can help mitigate these issues. By optimizing settings such as gpu_memory_utilization, max_num_batched_tokens, and tensor_parallel_size, we can strike a better balance between throughput and latency, ensuring more efficient performance.

As we continue our exploration of advanced vLLM techniques for enhanced LLM performance on AMD GPU, part 3 will introduce the feature of speculative decoding.

Overview

Part 1 and Part 2 of the AMD vLLM blog series discussed vLLM features such as prefix caching, chunked prefill and tuning system parameters to optimize performance. We now introduce speculative decoding. This innovative technique further enhances language model performance on AMD GPUs by significantly speeding up the text generation process.

What is Speculative Decoding?

Speculative decoding is a technique used in natural language processing and text generation, designed to improve the efficiency and quality of generated text through predictions of future tokens. This technique is commonly utilized in scenarios such as:

Prediction of Future Tokens: speculative decoding involves generating a sequence of tokens in a way that allows the model to explore multiple potential continuations of a given input. Instead of producing one token at a time, the model speculates on what might come next.

Parallel Processing: By generating multiple possible continuations in parallel, the model can quickly assess several options before finalizing the output.

When to Apply Speculative Decoding?

You can apply speculative decoding to various application use cases. Let’s take the following as examples:

• Chatbots and Conversational Agents: Speculative decoding can enhance the responsiveness and relevance of responses by generating multiple potential replies and selecting the best one.
• Content Creation: This technique is useful in automated writing tools where generating coherent and engaging text is crucial.

Significant Latency Reduction and Enhanced Performance with Speculative Decoding

The benefits of speculative decoding lie in its ability to substantially reduce the latency and enhance performance of text generation models. This section highlights some of the key benefits that this technique brings to natural language processing tasks.

• Faster Generation: This method can significantly speed up text generation by evaluating multiple possibilities simultaneously, thereby reducing the time needed for serial token generation.
• Enhanced Coherence and Quality: By taking various continuations into account, the model is able to select the most contextually appropriate or relevant output, resulting in more coherent and higher-quality text.
• Exploration of Creativity: In creative writing applications, speculative decoding facilitates the exploration of diverse narrative paths or responses, thereby enriching the richness of creative output.

As shown in Figure 7, for the  amd/Llama-3.1-70B-Instruct-FP8-KV model selected, by using the amd/Llama-3.1-8B-Instruct-FP8-KV as its corresponding speculative model, the average request latency declines from 3.3 to 2 seconds, achieving about 1.65X speedup when speculative decoding is enabled.

As shown in Figure 8, for the amd/Llama-3.1-405B-Instruct-FP8-KV model, by using the amd/Llama-3.1-8B-Instruct-FP8-KV as its corresponding speculative model, the average request latency declines from 4 to 3 seconds, achieving about 1.3X speedup when speculative decoding is enabled.

Conclusion

Speculative decoding is a powerful technique that enhances the efficiency and quality of text generation by predicting multiple potential continuations for a given input, resulting in faster and more coherent outputs. For users interested in trying it out, a detailed step-by-step guide will be provided in part 4 of this blog series on how to run these features on AMD Instinct™ MI300X accelerators. We will walk you through each process, ensuring your replication of the optimization steps effectively. Furthermore, the final AMD vLLM blog will include a roadmap detailing the upcoming support for vLLM on AMD platforms, outlining new features and enhancements that can be expected in the near future.

Overview

This is part 4 of a four-part AMD vLLM blog series. We previously introduced the advanced vLLM techniques of prefix caching, chunked prefill, and speculative decoding, and have showcased their significant improvements on large model inference performance on AMD GPUs. This final blog will transition from theory to practice. Below you will find a comprehensive hands-on guide for those eager to reproduce these optimizations and witness firsthand the enhancements they offer. This guide will provide step-by-step instructions designed to ensure you can effectively implement these techniques on the AMD Instinct™ MI300X accelerator.

Implementation Guide

Prerequisites

The implementation has been conducted on AMD Instinct MI300X accelerator, with ROCm™ 6.2, AMD EPYC 9654 96-Core CPU processor, and Ubuntu 22.04.4 LTS.

To get the optimized performance, it is recommended to disable automatic NUMA balancing. Failure to do so may result in the GPU being unresponsive until the periodic balancing is finalized. For further details, please refer to the AMD Instinct MI300X system optimization guide.

# disable automatic NUMA balancing
sh -c 'echo 0 > /proc/sys/kernel/numa_balancing'
# check if NUMA balancing is disabled (returns 0 if disabled)
cat /proc/sys/kernel/numa_balancing
0
sudo rocm-smi --setperfdeterminism 1900

For MI300X (gfx942) users, to achieve optimal performance, please refer to MI300x tuning guide for performance optimization and tuning tips on system and workflow level.

The NCCL_MIN_NCHANNELS argument controls the minimum number of communication channels that NCCL will use for collective operations. The default value for this is generally 1, but it can be increased for scaling out communication across multiple GPUs, especially when using multiple nodes or large-scale setups.

export NCCL_MIN_NCHANNELS=112

Step 1: Pull the Docker image from Docker Hub.

docker pull rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4

Inside the docker, you are ready to run the following commands. Note that for other supported GPUs beyond AMD Instinct MI300X, vLLM can be installed using the instructions at docs.vllm.ai.

export VLLM_USE_TRITON_FLASH_ATTN=0

To use CK flash-attention or PyTorch naive attention, please use this flag to turn off triton flash attention.

Step 2: Set up the Path 

cd /app/vllm/benchmarks

Please modify the path as your own path, Eg,

MODEL PATH

/workspace/home/models/Meta-Llama-3.1-70B-Instruct

DATASET PATH

/app/vllm/benchmarks/dataset/ShareGPT_V3_unfiltered_cleaned_split.json

Additional Optimizations

Optimizing CPU Overheads in LLM Serving

In the realm of LLMs, the CPU plays a crucial role in ensuring smooth and efficient operations. In addition to above hands-on guide, this blog also examines the key tasks that are managed by the CPUs and explain why optimizing these tasks is essential for improved performance.

Key Tasks of CPU in LLM Serving

The CPU undertakes several critical responsibilities to maintain the efficiency and responsiveness of LLM services. There are 4 key tasks for the CPU in LLM Serving:

• Handling HTTP Requests and Streaming Output Tokens:  The CPU is responsible for managing incoming HTTP requests, which are essential for interacting with the LLM. Additionally, it streams the output tokens generated by the model, ensuring that responses are delivered promptly.
• Scheduling Requests at Every Step: Efficient scheduling of requests is vital to maintain seamless operation flow. The CPU schedules these requests at each step, coordinating the tasks to ensure that the model runs smoothly.
• Preparing Input Tensors and Gathering Output Tensors: Prior to the data processing, the CPU prepares the input tensors. After processing, it gathers the output tensors, which are subsequently transformed into meaningful responses.
• Converting Output Tokens into Strings: The final step involves converting the output tokens generated by the model into strings that can be understood and used by end-users.

To illustrate the impact of CPU overheads, consider the example of serving the Llama-3-8B model (Figure 9). In this scenario, the GPU experienced idleness for up to 60% of the time. This idleness was primarily due to the time spent by the CPU on scheduling requests, preparing inputs, and gathering outputs. This example highlights the importance of optimizing CPU tasks to minimize idle time and improve overall efficiency.

By identifying and addressing these CPU overheads, it is possible to enhance the performance of LLM serving, ensuring faster and more efficient processing of requests. This not only improves the user experience but also maximizes the utilization of available resources.

Recommended Strategies to Minimize CPU Overheads

To address the inefficiencies caused by CPU overheads and enhance the performance of large language model serving, consider the following strategies to optimize CPU overheads:

1. Separating API server and engine processes using multiprocessing and ZeroMQ to remove API server overheads from the critical path, thereby achieving great performance under high QPS (queries per second).
2. Implementing multi-step scheduling to schedule multiple steps ahead, thereby amortizing the cost of scheduling and input preparation.
3. Adopting asynchronous output processing to overlap de-tokenization and stop string matching with model execution, which adds one extra step per request but significantly increases GPU utilization outweighs the cost.

Scalable System Architecture

To achieve a more scalable system architecture, several key strategies are being implemented. 

Asynchronous scheduling is being generalized to enhance both output processing and multi-step scheduling, while incremental input preparation aims to reuse input and output tensors from previous steps, thereby improving efficiency.

In the context of vLLM and torch.compile, the optimization process involves four steps: first, optimizing certain operators using torch.compile; second, capturing the full model with TorchDynamo; third, compiling the full model with TorchInductor; and finally, adding custom optimizations on top of torch.compile.

Additionally, the development of multi-modal models is being advanced with contributions from model vendors, introducing new modalities such as audio with models like Ultravox. These models also feature support for dynamic image sizes, multi-image interfaces, and image embeddings as input, along with new optimizations like optimized kernels and tensor parallelism support for vision encoders.

vLLM Roadmap

The future vLLM roadmap will mainly focus on re-architecture, with deeper integration into the PyTorch ecosystem using tools like torch.compile and Triton to optimize performance and compatibility. Enhancements will include better support for Vision-Language Models (VLMs) and the introduction of new modalities such as speech and video. Additionally, efficient support for long context scenarios is being developed, including disaggregating serving and key-value (KV) cache transfer processes, as well as implementing sequence parallelism to ensure smoother and faster processing.

Conclusion

The strategic partnership between AMD and open-source innovators like vLLM aims to utilize AMD Instinct GPUs, ROCm software, and features like prefix caching, chunked prefill, and speculative decoding for optimal performance. By integrating advanced memory optimization techniques, efficient execution algorithms, and further optimizations on the vLLM roadmap with an emphasis on minimizing CPU overhead in LLM serving, this partnership intends to improve the capability and efficiency of large model inference, making it more accessible and effective for real-world applications. Stay tuned for more updates as we continue to innovate and enhance vLLM features on AMD. For more inquiries, contact us at amd_ai_mkt@amd.com.