vLLM on Neuron — Llama 3.1 8B (trn1)¶

This tutorial walks you through running the vLLM framework on AWS Trainium to serve large language models at high speed.

Overview¶

vLLM is a high-performance LLM inference engine based on PagedAttention. It integrates with the AWS Neuron SDK to run natively on Trainium.

Key Features¶

Continuous Batching: Maximize throughput with dynamic batching
PagedAttention: Memory-efficient KV cache management
Tensor Parallelism: Distribute models across multiple NeuronCores
OpenAI-compatible API: Drop-in replacement for existing applications

Prerequisites¶

AWS account with access to trn1.32xlarge or larger instances
An AWS account with access to Trainium instances
Docker installed

Step 1: Environment Setup¶

# Launch EC2 instance with Neuron Deep Learning AMI
aws ec2 run-instances \
  --image-id ami-0xxxxxxxxxxxxx \
  --instance-type trn1.32xlarge \
  --key-name my-key \
  --security-group-ids sg-xxxxxxxx

AMI Selection

Use the latest Deep Learning AMI Neuron (Ubuntu 22.04). Neuron drivers and runtime are pre-installed.

Step 2: Install vLLM¶

# Create virtual environment
python -m venv vllm-env
source vllm-env/bin/activate

# Install vLLM (Neuron build)
pip install vllm[neuron]

Step 3: Compile Model¶

Models must be compiled to Neuron format for execution on Neuron devices:

# Compile Llama 3.1 8B model
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3.1-8B-Instruct \
  --device neuron \
  --tensor-parallel-size 2 \
  --max-model-len 4096 \
  --block-size 4096

Compilation Time

First-time compilation may take 15-30 minutes. Compiled models are cached for faster subsequent loads.

Step 4: Run Inference Server¶

# Launch OpenAI-compatible API server
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3.1-8B-Instruct \
  --device neuron \
  --tensor-parallel-size 2 \
  --port 8000

Step 5: Test¶

# Send inference request with curl
curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "prompt": "The benefits of AWS Trainium are",
    "max_tokens": 100,
    "temperature": 0.7
  }'

Performance Benchmarks¶

Model	Instance	TP	Throughput	Latency (P50)
Llama 3.1 8B	trn1.2xlarge	2	~800 tok/s	~35ms
Llama 3.1 70B	trn1.32xlarge	32	~400 tok/s	~80ms

Clean Up¶

# Terminate EC2 instance
aws ec2 terminate-instances --instance-ids i-xxxxxxxxxxxxxxxxx

Cost Prevention

Always terminate unused trn1 instances. trn1.32xlarge costs approximately $21.50/hour (On-Demand).