Deploying LLMs with vLLM and Ray

Serving Large Language Models (LLMs) in production is challenging. Models like Llama-3-70B require massive amounts of memory and compute. If you serve them naively using standard Hugging Face pipeline code, your throughput will be abysmal, and you'll suffer from out-of-memory (OOM) errors during concurrent requests.

To solve this, we combine two powerful tools: vLLM (for optimized inference) and Ray Serve (for distributed scaling).

Why vLLM? PagedAttention

vLLM's secret weapon is PagedAttention.

In traditional LLM serving, memory for the KV cache (the memory of past tokens in a sequence) is allocated contiguously. Because sequence lengths are unpredictable, memory becomes fragmented, leading to wasted VRAM and severely limiting batch size.

PagedAttention operates like an operating system's virtual memory. It divides the KV cache into blocks and allocates them dynamically. This nearly eliminates memory fragmentation, allowing vLLM to batch significantly more requests together, often achieving 2x-4x higher throughput than Hugging Face Text Generation Inference (TGI).

Why Ray Serve?

While vLLM handles single-node inference beautifully, what happens when you have hundreds of concurrent users? You need to distribute the load across multiple GPUs or multiple machines.

Ray Serve acts as a dynamic router and orchestrator. It allows you to wrap your vLLM engine in a scalable deployment, handling auto-scaling, request batching, and health checks.

The Deployment Code

Here is a complete example of deploying a Llama-3 model using Ray Serve and vLLM.

Prerequisites

You'll need a machine with at least one powerful GPU (e.g., A100 or H100) and the necessary libraries: pip install ray[serve] vllm

The Ray Serve Script (serve_llm.py)

python
import os
from fastapi import FastAPI
from ray import serve
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.sampling_params import SamplingParams
from vllm.utils import random_uuid

# 1. Initialize FastAPI for custom routing if needed
app = FastAPI()

# 2. Define the Ray Serve Deployment
@serve.deployment(
    autoscaling_config={
        "min_replicas": 1,
        "max_replicas": 4,
        "target_num_ongoing_requests_per_replica": 10,
    },
    ray_actor_options={"num_gpus": 1} # Adjust based on model size (e.g., tensor parallel size)
)
@serve.ingress(app)
class VLLMDeployment:
    def __init__(self):
        # Configure vLLM Engine
        engine_args = AsyncEngineArgs(
            model="meta-llama/Meta-Llama-3-8B-Instruct",
            tensor_parallel_size=1, # Set to >1 if splitting across multiple GPUs
            gpu_memory_utilization=0.90, # Leave 10% for PyTorch overhead
            max_num_batched_tokens=4096,
        )
        self.engine = AsyncLLMEngine.from_engine_args(engine_args)

    @app.post("/generate")
    async def generate(self, request: dict):
        prompt = request.get("prompt", "")
        max_tokens = request.get("max_tokens", 512)
        temperature = request.get("temperature", 0.7)

        sampling_params = SamplingParams(
            temperature=temperature,
            max_tokens=max_tokens,
        )

        request_id = random_uuid()
        
        # Async generation stream
        results_generator = self.engine.generate(prompt, sampling_params, request_id)
        
        final_output = None
        async for request_output in results_generator:
            final_output = request_output
            
        text = final_output.outputs[0].text
        return {"text": text}

# 3. Build the application
app_builder = VLLMDeployment.bind()

Running the Cluster

To start the Ray cluster and deploy your application, run:

bash
ray start --head
serve run serve_llm:app_builder

You can now send POST requests to http://localhost:8000/generate and watch vLLM and Ray efficiently handle the concurrent load!