Deploying LLMs on Multi-node EKS: A Ray LLM DLC Approach

Deploying large language models (LLMs) across multiple nodes has become a whole lot easier with Ray LLM Deep Learning Containers (DLC) on Amazon EKS. This kind of multi-node setup is common for tensor, pipeline, and data parallelism, as well as prefill/decode disaggregation inference workloads, which means we get better performance by running different inference phases on separate GPU nodes. That's a big deal.

So, What's New with Ray LLM DLC?

The Ray LLM DLC extends AWS Deep Learning Containers to distributed LLM serving. Basically, the DLC packs all the necessary components—framework, dependencies, and GPU stack—into a pre-built Docker image. You can deploy it without the usual hassle, which makes managing the GPU environment simpler and ensures everything plays nice together. That's usually a headache in multi-node deployments, so this is a nice fix.

Ray Serve deployments can span multiple nodes. Customers use multi-node Ray Serve for several reasons. Some want to run foundation or open-weight models that are too large for a single GPU. Others need more throughput than one replica can deliver. The DLC includes NVIDIA Inference Transfer Library (NIXL) and Elastic Fabric Adapter (EFA) to efficiently manage KV-cache transfers between nodes. What does this mean for us? It does wonders for applications like Retrieval Augmented Generation (RAG), document summarization, and tasks that need a bit more intelligence.

The Importance of Prefill/Decode Disaggregation

LLM inference involves two stages: prefill and decode. Prefill is the heavy lifter, chewing through input prompts in parallel and building a key-value (KV) cache. Decode is all about generating responses, and it's really hungry for memory bandwidth. When these phases share a single GPU, they fight over resources, which is bad for performance. By splitting them onto separate nodes, we can tweak hardware and serving configurations to suit each phase, boosting performance.

The prefill phase processes the input prompt in parallel and builds the KV cache. Then, during decode, this cache is read from and extended to generate responses one token at a time. Prefill is compute-heavy, while decode chugs along with the need for speed in memory bandwidth. By keeping them apart, we avoid resource battles, especially useful for workloads with long prompts or high concurrency. The catch? You need fast KV-cache transfer between nodes to make the performance gains stick. The Ray LLM DLC uses EFA to ensure we get the high-bandwidth communication needed.

Steps to Set Up Ray LLM DLC on Amazon EKS

Here's how you can deploy a model using Ray LLM DLC on Amazon EKS:

1. Set Up the EKS Cluster: Use the provided script to get your EKS cluster running. This includes a CPU system node for system workloads and the Ray head. Make sure nodes live in private subnets and have outbound access via a NAT Gateway. You'll need to provision the VPC, OIDC provider, and core add-ons.
2. Add GPU Worker Nodes: You'll want to provision two GPU nodes with EFA enabled. These have to be in the same Availability Zone because EFA can't hop across zones. These GPU nodes will handle the prefill and decode phases, leveraging the high-speed EFA interconnect for KV-cache transfers.
3. Get the KubeRay Operator Running: Use Helm to deploy the KubeRay operator, which will manage the Ray cluster and automate connections between the head and worker nodes. This involves setting up the operator to keep tabs on the RayService and integrating head and worker nodes smoothly.
4. Deploy the RayService: Set up your manifest to create the RayService, which will run the prefill and decode phases on separate GPU nodes. The manifest should cover GPU allocation, memory use, and how the KV-cache transfer will happen.

After you're all set up, you can call the model using the OpenAI-compatible API the DLC exposes. No need for a custom schema here. The pre-built application within the serveConfigV2 block of the manifest handles inference requests using build_pd_openai_app to configure the prefill and decode phases.

Potential Areas of Concern

From what I've seen, splitting inference phases can really amp up performance, but it needs careful resource and KV-cache transfer speed monitoring. For example, making sure the NIXL uses the correct backend (LIBFABRIC) is crucial because it defaults to slower TCP if misconfigured. Checking that EFA is set up properly is also key to keeping the high-speed communication this architecture depends on.

In a fast-paced environment like at India Today Group, where quick responses and high throughput matter, this kind of architecture can keep AI-assisted content workflows from lagging and improve the efficiency of smart search systems. Keeping tabs on performance metrics for both the prefill and decode nodes is important to make sure they're hitting their targets.

Limitations and Things to Keep in Mind

While the Ray LLM DLC setup can boost performance for multi-node inference tasks, there are trade-offs. You need a speedy interconnect for KV-cache transfers, making network configuration crucial. Also, the setup demands enough AWS service quotas, especially for GPU instances, which might be a hurdle for some. Proper infrastructure configuration is essential to meet the demands of high-throughput inference tasks.

Don't forget the costs that go along with such setups. High-performance GPU instances and data transfers aren't cheap. Balancing performance needs with budget constraints is important. You'll want to consider scalability too, as you might need to tweak the deployment based on workload changes.

Ray LLM DLC on Amazon EKS lays a strong foundation for large-scale inference workloads, but it requires careful planning and resource management to make the most of it. Regular audits of resource usage and cost analysis will help keep things efficient and cost-effective.

FAQs

• How does Ray LLM DLC boost performance?
  By moving the prefill and decode phases to different nodes, Ray LLM DLC cuts down on resource contention, allowing each phase to run optimally according to its needs. This separation helps us tailor hardware configurations that match the unique demands of each phase.
• What applications gain the most from this setup?
  It's great for applications that need high throughput and low latency, like document summarization and agentic tasks. Handling large-scale input prompts and generating responses efficiently makes it ideal for scenarios with high concurrency.
• Is EFA necessary?
  Yes, EFA is essential for maintaining high-speed node-to-node communication, which is crucial for the performance gains from this setup. Without EFA, KV-cache transfer could bottleneck the system, undoing the benefits of splitting the phases.
• What do I need to deploy Ray LLM DLC?
  You need an AWS account with billing enabled, enough service quotas, and several CLI tools like AWS CLI, kubectl, eksctl, helm, and envsubst. Ensuring these prerequisites are met is vital for a successful deployment.