Optimizing Multi-Low-Rank Adaptation for Mixture of Experts Models in vLLM

This heading encapsulates the main focus of the content, highlighting both the technical aspect of Multi-Low-Rank Adaptation and its application within Mixture of Experts models using the vLLM framework.

Leveraging Multi-LoRA for Efficient AI Model Inference

In today’s landscape of artificial intelligence, organizations and individuals deploying multiple custom AI models, particularly those in the Mixture of Experts (MoE) family, are likely wrestling with the challenge of underutilized GPU capacity. When traffic to individual models fluctuates, dedicated compute resources can remain idle, leading to wasted investment. To tackle this issue head-on, we have joined forces with the vLLM community and developed a cutting-edge solution for Multi-Low-Rank Adaptation (Multi-LoRA) serving of popular open-source MoE models like GPT-OSS and Qwen.

What is Multi-LoRA?

Multi-LoRA is an innovative approach that allows organizations to fine-tune models efficiently. Unlike methods that require retraining entire model weights, multi-LoRA maintains the original weights and introduces small, adaptable adapters into the model’s layers. This means that at inference time, multiple custom models can share the same GPU by simply swapping these adapters in and out per request.

For instance, imagine five customers each using only 10% of a dedicated GPU. With Multi-LoRA, instead of needing five separate GPUs, a single GPU can now serve all five customers efficiently, transforming a scenario of underutilized resources into a streamlined operation.

Implementing Multi-LoRA Inference for MoE Models

Our journey into implementing Multi-LoRA inference for MoE models in vLLM began with an understanding of how MoE models function. These models comprise numerous specialized neural networks known as experts. A router intelligently directs each input token to the most relevant experts, thereby optimizing computational resources by activating only a fraction of the total model’s parameters per token.

The "expand-then-compress" pattern within MoE models permits rich transformations while ensuring consistent output sizes. However, challenges arise when integrating Multi-LoRA due to the need for efficient resource management across disparate users and tasks.

Key to this integration was the development of the fused_moe_lora kernel, which combines LoRA operations with the fused_moe kernel for seamless incorporation of Multi-LoRA adapters.

Performance Enhancements

After the initial implementation, we identified bottlenecks using NVIDIA Nsight Systems and NVIDIA Nsight Compute. Here’s how we enhanced performance:

Execution Optimizations

Initially, the multi-LoRA Time to First Token (TTFT) was significantly higher than that of the base model. We discovered that every context length was treated as a compile-time constant, leading to unnecessary recompilation. We resolved this by adding compiler hints for variable reuse, drastically improving latency.

By introducing early exit logic for layers without LoRA adapters and implementing Programmatic Dependent Launch (PDL) for overlapping kernel execution, we reduced idle GPU time dramatically.

Kernel Optimizations

Further performance issues stemmed from the execution of skinny matrices, which required inventive solutions like Split-K for improved load balancing and CTA swizzling for enhanced memory usage. We also optimized the identification of unnecessary masking and dot product operations, leading to an impressive reduction in overhead.

Tuning for Amazon Environments

Specific tuning parameters for Triton kernels tailored for MoE LoRA serving allowed us to unlock additional performance improvements. By enabling customized configurations through easily accessible paths, we ensured that Amazon SageMaker AI and Amazon Bedrock users could leverage these advancements.

Measurable Results

Through collaborative efforts within the vLLM community, we’ve not only implemented Multi-LoRA serving but also optimized it effectively. For example, we’ve achieved up to 454% output token per second (OTPS) improvements and an 87% reduction in TTFT for GPT-OSS 20B. Such enhancements allow for practical, real-world deployment of models without the burden of underutilized resources.

Conclusion

The advent of Multi-LoRA has revolutionized how we think about serving multiple AI models, especially within the MoE framework. Organizations can now harness the power of AI without the associated costs of underutilized GPU capacity.

With our latest enhancements available in vLLM 0.15.0 and beyond, users can look forward to efficient deployment on platforms like Amazon SageMaker AI and Amazon Bedrock, yielding significant latency improvements and increased overall performance.

Acknowledgments

We extend our gratitude to the vLLM community for their contributions and collaboration. A special mention to our dedicated team members who made this initiative possible.

By integrating these optimizations, not only are we able to deliver significant performance benefits, but we are also paving the way for a more efficient, cost-effective future in AI deployment. Get started today—delve into the new features, optimize your models, and witness firsthand the transformative power of Multi-LoRA.