Run multi-node inference workloads using Ray LLM Deep Learning Containers

來源原標題

Run multi-node inference workloads using Ray LLM Deep Learning Containers

作者

Ananth Raghavendra

原始標籤

Amazon Elastic Kubernetes Service · Intermediate (200) · Technical How-to

原始連結

https://aws.amazon.com/blogs/machine-learning/run-multi-node-inference-workloads-using-ray-llm-deep-learning-containers/

摘要原文

This post shows how to deploy one large language model (LLM) across two GPU nodes on Amazon Elastic Kubernetes Service (Amazon EKS) with the multi-node Ray LLM Deep Learning Containers . It uses separate prefill and decode nodes behind one OpenAI-compatible endpoint. This kind of multi-node setup is common for tensor, pipeline, and data parallelism, as well as prefill/decode disaggregation inference workloads. In a previous post , we introduced the Ray Serve DLC and used it to stand up a small model behind an HTTP endpoint on a single GPU node. You had a working inference endpoint, and the container keeps the GPU stack compatible, including CUDA libraries, so you did not need to write boilerplate code. Ray is a distributed compute framework, so 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. Still others want to separate the two phases of inference to improve performance. In each case, customers must build the dependency chain themselves: choosing compatible libraries for the GPU stack and debugging issues when a component drifts out of alignment. This is undifferentiated work that takes time and effort away from delivering models. AWS Deep Learning Containers (DLCs) have long solved this kind of problem for training and inferencing workloads. DLCs are pre-built, performance-optimized Docker images that bundle a framework, its dependencies, and the GPU stack into a tested, patched combination you can pull and use immediately. With the launch of the Ray LLM DLC for multi-node deployments, that same approach now extends to multi-node inference. You get a container purpose-built for serving models behind an HTTP endpoint, maintained and tested by AWS, with the full inference stack already assembled. This post walks through implementing multi-node disaggregated inference using multi-node Ray LLM DLC: deploying a single large language model across two GPU-powered nodes on Amazon EKS cooperating behind a single endpoint. While one of the GPU nodes handles the prefill phase, the other GPU node handles the decode phase. Separating the phases prevents compute-heavy long prompts from stalling token generation for requests already in progress. This approach is well suited for workloads such as Retrieval Augmented Generation (RAG), document summarization, and agentic applications. We use the multi-node Ray LLM DLC, which ships Ray Serve LLM , NVIDIA Inference Transfer Library (NIXL) , and the Elastic Fabric Adapter (EFA) stack, and exposes an OpenAI-compatible API with no application code. The complete code is available in the accompanying repository . LLM inference has two distinct phases. Prefill processes the entire input prompt in parallel and builds the key-value (KV) cache. Decode then generates the response one token at a time, reading and extending that cache. Prefill is compute bound, and Decode is bound by memory bandwidth. When both phases share one GPU, they contend for the same resources, and the bursty, compute-heavy prefill work interferes with the steady, latency-sensitive decode work. Separating the prefill phase from the decode phase means you can select the hardware and serving properties for each phase, improving overall performance at scale. PD disaggregation places the two phases on separate nodes, so each phase runs without interfering with the other and can be scaled independently. The trade-off is that the KV cache built during prefill must be handed to the decode node, and that transfer should be fast enough not to erase the benefit. That is exactly why the interconnect between the nodes matters. AWS DLCs are pre-built, performance-optimized Docker images that bundle a framework, its dependencies, and the GPU stack into a tested, patched combination you can pull and use immediately. All components are validated together before each release, so there is no version drift between the CUDA runtime, the framework, and the serving layer. Security patches are applied at build time. The Ray LLM DLC extends this to distributed LLM serving. It ships vLLM as the inference engine, Ray Serve’s HTTP as the serving layer, NIXL for KV-cache transfer between nodes, and the EFA stack ( libfabric and aws-ofi-nccl ) for high-bandwidth node-to-node communication. Because the Ray Serve LLM library includes a built-in application builder, the entire serving application is a block of YAML. There’s no handler class, no model archiver, and no custom image to build. A single helper, build_pd_openai_app , takes a prefill configuration and a decode configuration and wires the KV-cache transfer between them behind one OpenAI-compatible endpoint. The following diagram shows how the deployment is laid out on Amazon EKS. A user request reaches the Ray head on the CPU system node, and the two GPU worker nodes run the prefill and decode phases of the model. Figure 1: Multi-node inference architecture on Amazon EKS, with the Ray head on a CPU system node and the prefill and decode phases on two GPU worker nodes The deployment is an EKS cluster inside a virtual private cloud (VPC). Nodes run in private subnets with outbound access through a NAT Gateway. A CPU system node (m7i.xlarge) hosts system workloads, the KubeRay operator, and the Ray head pod, which is CPU-only. Two GPU worker nodes (g6.8xlarge) run the Ray workers: one serving the prefill phase, and the other serving the decode phase. Both nodes pull the Ray LLM DLC from Amazon Elastic Container Registry (Amazon ECR) . KubeRay watches the RayService you apply and wires the head-to-worker join automatically. Ray Serve places the prefill and decode deployments on the two GPU workers. The KV cache flows from the prefill worker to the decode worker over NIXL, using EFA through NIXL’s libfabric backend. The GPU node group is pinned to a single Availability Zone (AZ) because EFA cannot cross AZs. To follow along with this post, you need: Verify that your AWS credentials are active with the following command: The application is the serveConfigV2 block in manifest/rayservice.yaml . It uses build_pd_openai_app , which takes a prefill_config and a decode_config and connects the KV-cache transfer between them: A few settings carry most of the weight. max_model_len: 4096 sizes the context to fit the measured KV cache on the GPU. dtype: bfloat16 and gpu_memory_utilization: 0.90 set precision and how much VRAM vLLM can claim. enforce_eager: true skips CUDA graph capture to reduce peak GPU memory (model- and input-dependent). The kv_transfer_config setting is what makes disaggregation work: both phases use NixlConnector with kv_role: kv_both and a distinct engine_id (prefill on one, decode on the other). The critical setting is kv_connector_extra_config: { backends: ["LIBFABRIC"] } . Without that line, NIXL defaults to UCX, which runs at TCP speed even on EFA. The rest of the manifest, in rayClusterConfig , describes the pods. The head pod is CPU-only ( num-gpus: "0" ) and pinned to the system node group. Each worker pod requests one GPU and one EFA interface ( nvidia.com/gpu: "1", vpc.amazonaws.com/efa: "1" ) and targets the GPU nodes through nodeSelector: { role: gpu-worker } . It runs privileged: true because NIXL registers CUDA VRAM with the EFA NIC. The IPC_LOCK capability alone isn’t enough. The workers set FI_PROVIDER=efa , read VLLM_HOST_IP from the pod IP, and mount a shared-memory volume. Memory limits sit well above the roughly 18 GiB checkpoint so the OS page cache can hold it. For Qwen3.5-9B model loading on g6.8xlarge with the preceding cache configuration, the model weights loading drops from around 145 seconds to about 2 seconds on restart. All scripts share a single configuration file, scripts/env.sh . The script contains defaults for CLUSTER_NAME , REGION , K8S_VERSION , and NAMESPACE .