LLM Inference: GPU vs Trainium2

Reference Architecture — vLLM on Amazon EKS, GPU & Trainium side-by-side

One EKS cluster serves Qwen3-8B on both accelerators. Each pod owns its accelerator and runs vLLM; the Neuron device plugin exposes NeuronCores, and node scaling is handled by Karpenter / EKS Auto Mode (provisions the right instance per pod), not pod-level HPA.

☸️ Amazon EKS Cluster

GPU + Trainium node pools, one control plane

┌────────────────────────────── Amazon EKS control plane (managed) ──────────────────────────────┐
│                                                                                                  │
│   Ingress / ALB  ──▶  vLLM Service  ──▶  routes by model / node pool                              │
│                                                                                                  │
│   ┌─────────────────── GPU node pool ───────────────────┐   ┌──────────── Trainium node pool ────────────┐ │
│   │  g6e.4xlarge  (1× NVIDIA L40S, 48GB)             │   │  trn2.3xlarge (1× Trainium2 chip, 4 cores) │ │
│   │  ┌───────────────────────────────────┐            │   │  ┌────────────────────────────────────────┐ │ │
│   │  │ vLLM pod                           │            │   │  │ vLLM pod                                  │ │ │
│   │  │  Qwen3-8B · CUDA · TP=1            │            │   │  │  Qwen3-8B · Neuron · TP=4  →  4 cores    │ │ │
│   │  │  requests: nvidia.com/gpu: 1       │            │   │  │  requests: aws.amazon.com/neuron: 1       │ │ │
│   │  └───────────────────────────────────┘            │   │  └────────────────────────────────────────┘ │ │
│   │   NVIDIA device plugin                            │   │   Neuron device plugin (DaemonSet)         │ │
│   └────────────────────────────────────────────────────┘   └──────────────────────────────────────────────┘ │
│                                                                                                  │
│   Karpenter / EKS Auto Mode — provisions g6e or trn2 nodes on demand to match pending pods         │
│   Observability: neuron-monitor ─▶ CloudWatch   ·   Neuron Explorer (offline profile viewer)      │
│                                                                                                  │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

GPU node pool

• g6e.4xlarge — 1× NVIDIA L40S (48 GB)
• vLLM, tensor-parallel = 1 (single GPU)
• Pod requests nvidia.com/gpu: 1
• NVIDIA device plugin advertises the GPU
• $3.004/hr on-demand (us-east-2)

Trainium node pool

• trn2.3xlarge — 1× Trainium2 chip = 4 NeuronCores (96 GB)
• vLLM, tensor-parallel = 4 → one core per TP rank
• Pod requests aws.amazon.com/neuron: 1 (whole device)
• Neuron device plugin advertises cores/devices
• $2.235/hr Capacity Block (sa-east-1)

Scaling — Karpenter / EKS Auto Mode (not HPA): pods declare their accelerator request; Karpenter watches pending pods and launches the right node type (g6e or trn2) to fit them, then consolidates when idle. Scale-out adds nodes/accelerators sized to the workload — a better fit for one-accelerator-per-pod LLM serving than replica-count HPA on a fixed node set.

🟢 NVIDIA L40S (GPU)

g6e.4xlarge | $3.004/hr

Ready

TTFT (ms)

tok/s

in tok

out tok

cost

🟣 AWS Trainium2

trn2.3xlarge (1 chip, 4 NeuronCores) | TP=4 | ~$2.24/hr

Ready

TTFT (ms)

tok/s

in tok

out tok

cost

🟢 NVIDIA L40S

g6e.4xlarge | $3.00/hr

Ready

tok/s (aggregate)

$/1M tokens

GPU Utilization (approx from throughput)

Tokens: --

Time: --s

Reqs: --

🟣 AWS Trainium2

trn2.3xlarge (1 chip, 4 NeuronCores) | TP=4 | ~$2.24/hr

Ready

tok/s (aggregate)

$/1M tokens

NeuronCore Utilization (live neuron-monitor)

Tokens: --

Time: --s

Reqs: --

Metric	GPU p50	p95	p99	TRN p50	p95	p99
TTFT (ms)	–	–	–	–	–	–
End-to-end latency (ms)	–	–	–	–	–	–
Throughput (tok/s)	–	–

GPU Serving

Standard CUDA path

# Instance: g6e.4xlarge (NVIDIA L40S, 46GB)
# Ubuntu 24.04 | CUDA 13.2 | vLLM 0.24

# Install
pip install vllm

# Serve
vllm serve Qwen/Qwen3-8B \
    --dtype bfloat16 \
    --max-model-len 2048 \
    --host 0.0.0.0 --port 8000

# Query (OpenAI-compatible API)
curl http://localhost:8000/v1/chat/completions \
  -d '{
    "model": "Qwen/Qwen3-8B",
    "messages": [{"role":"user","content":"Hello"}],
    "max_tokens": 256
  }'

Trainium Serving

Same API, different backend

# Instance: trn2.3xlarge (4 NeuronCores, 96GB)
# Ubuntu 24.04 | Neuron SDK 2.25 | vLLM 0.16 + NxDI

# Activate (pre-installed on Neuron DLAMI)
source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate

# Serve
export VLLM_NEURON_FRAMEWORK=neuronx-distributed-inference

vllm serve Qwen/Qwen3-8B \
    --tensor-parallel-size 4 \
    --max-model-len 2048 \
    --host 0.0.0.0 --port 8000

# Query — IDENTICAL API! Same curl command works.
curl http://localhost:8000/v1/chat/completions \
  -d '{
    "model": "Qwen/Qwen3-8B",
    "messages": [{"role":"user","content":"Hello"}],
    "max_tokens": 256
  }'

PyTorch Native on Neuron — Zero Friction Porting

With PyTorch Native, moving from GPU to Trainium is as simple as changing device='cuda' → device='neuron' and using torch.compile(backend="neuron"). No custom compilation steps, no model rewrites — same PyTorch code, different hardware.

GPU Code (PyTorch + CUDA)

Standard PyTorch workflow

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    torch_dtype=torch.bfloat16,
)

# Move to device
model = model.to("cuda")
model.eval()

# Compile for performance
model.forward = torch.compile(
    model.forward,
    # uses default inductor backend
)

# Tokenize
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
inputs = tokenizer("Hello", return_tensors="pt").to("cuda")

# Generate
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=128)

print(tokenizer.decode(output[0]))

Trainium Code (PyTorch Native)

3 lines changed

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    torch_dtype=torch.bfloat16,
)

# Move to device
model = model.to("neuron")           # ← just change device
model.eval()

# Compile for Neuron
model.forward = torch.compile(
    model.forward,
    backend="neuron",        # ← add backend
)

# Tokenize
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
inputs = tokenizer("Hello", return_tensors="pt").to("neuron")  # ← device

# Generate — IDENTICAL call
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=128)

print(tokenizer.decode(output[0]))

Watch it end-to-end: our GPU training script → Trainium

TorchTitan Qwen3-8B example ↗

Take a standard TorchTitan Qwen3-8B training run and move it to Trainium2 by applying a 4-line diff — no custom kernels, no model rewrite. Same torchrun … torchtitan.train command, different device.

Commands and the TorchTitan.diff shown are verbatim from the AWS Neuron repo. Training step metrics in this clip are illustrative for pacing, not a measured benchmark — see the Throughput tab for real, server-side numbers.

Neuron Explorer — Qwen3-8B Inference Profile

Perfetto-style timeline: Tensor Engine, Vector Engine, DMA, Memory — trn2.3xlarge (4 NeuronCores)

Open Full Screen ↗

☸️ vLLM on EKS — Live Cluster

Qwen3-8B served from a vLLM pod on a Trainium2 node (sa-east-1)

$ kubectl get nodes
NAME                                          STATUS   ROLES    VERSION
ip-172-31-17-145.sa-east-1.compute.internal   Ready    <none>   v1.32.0

$ kubectl get node -o jsonpath='{...allocatable}'
aws.amazon.com/neuron: 1   aws.amazon.com/neuroncore: 8

$ kubectl get pods -l app=vllm-qwen3 -o wide
NAME                          READY   STATUS    NODE
vllm-qwen3-5bb9dbd7bf-scw59   1/1     Running   ip-172-31-17-145

$ curl trn2-node:8001/v1/models
served model: Qwen/Qwen3-8B   (vLLM 0.16 + PyTorch Native on Neuron)
                

☸️ EKS Deployment Architecture

vLLM on Trainium with Kubernetes-native scaling

┌─────────────────────────────────────────────────────────┐
│                    Amazon EKS Cluster                    │
├─────────────────────────────────────────────────────────┤
│                                                          │
│  ┌──────────────────┐    ┌──────────────────────────┐   │
│  │  Ingress / ALB    │───▶│   vLLM Service (ClusterIP)│   │
│  └──────────────────┘    └──────────────────────────┘   │
│                                    │                     │
│           ┌────────────────────────┼────────────┐       │
│           ▼                        ▼            ▼       │
│  ┌─────────────────┐  ┌─────────────────┐  ┌────────┐  │
│  │  vLLM Pod (TP=4) │  │  vLLM Pod (TP=4) │  │  ...   │  │
│  │  1 Trainium chip │  │  1 Trainium chip │  │        │  │
│  └─────────────────┘  └─────────────────┘  └────────┘  │
│        trn2 node            trn2 node                   │
│                                                          │
│  ┌──────────────────────────────────────────────────┐   │
│  │  Neuron Device Plugin (DaemonSet)                 │   │
│  │  advertises aws.amazon.com/neuron to the scheduler│   │
│  └──────────────────────────────────────────────────┘   │
│                                                          │
│  ┌──────────────────────────────────────────────────┐   │
│  │  HPA (scale replicas on latency / queue depth)   │   │
│  └──────────────────────────────────────────────────┘   │
│                                                          │
├─────────────── Observability Stack ─────────────────────┤
│                                                          │
│  ┌────────────┐  ┌─────────────┐  ┌────────────────┐   │
│  │ CloudWatch  │  │  OTel Agent  │  │ Neuron Monitor │   │
│  │ Container   │  │  (metrics +  │  │  (NeuronCore   │   │
│  │ Insights    │  │   traces)    │  │   utilization) │   │
│  └────────────┘  └─────────────┘  └────────────────┘   │
│                                                          │
└─────────────────────────────────────────────────────────┘

Key Components (as deployed):
• Neuron Device Plugin — DaemonSet that discovers NeuronCores and advertises aws.amazon.com/neuron as a schedulable resource
• vLLM Deployment — pod requests aws.amazon.com/neuron: 1; scheduler places it on the trn2 node
• Amazon VPC CNI + kube-proxy — pod networking and Service routing
• neuron-monitor → CloudWatch — per-core utilization, memory, execution stats (dashboard →)

Why it matters:
• Same vLLM + PyTorch Native stack as standalone, now Kubernetes-scheduled ✓
• Trainium exposed as a native K8s resource ✓
• Scale by adding replicas / trn2 nodes ✓

vLLM Deployment

Neuron device plugin — the manifest we deployed

# Neuron device plugin (DaemonSet) advertises NeuronCores
# kubectl apply -f k8s-neuron-device-plugin.yml

---
apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-qwen3
spec:
  replicas: 1
  template:
    spec:
      nodeSelector:
        aws.amazon.com/neuron.present: "true"
      containers:
      - name: vllm
        image: public.ecr.aws/neuron/pytorch-inference-\
                 vllm-neuronx:0.16.0-sdk2.29.0
        args: ["python", "-m", "vllm.entrypoints.openai.api_server",
               "--model", "/models/Qwen3-8B",
               "--served-model-name", "Qwen/Qwen3-8B",
               "--tensor-parallel-size", "4",
               "--max-model-len", "512",
               "--dtype", "bfloat16"]
        resources:
          limits:
            aws.amazon.com/neuron: 1
        volumeMounts:
        - name: model
          mountPath: /models/Qwen3-8B
      volumes:
      - name: model
        hostPath: { path: /home/ubuntu/local-models/Qwen/Qwen3-8B }

Live CloudWatch Dashboard

neuron-monitor → CloudWatch (auto-refresh every 10s)

LLM Inference on Trainium2 with vLLM & PyTorch Native

Reference Architecture — vLLM on Amazon EKS, GPU & Trainium side-by-side

☸️ Amazon EKS Cluster

GPU node pool

Trainium node pool

🟢 NVIDIA L40S (GPU)

🟣 AWS Trainium2

🟢 NVIDIA L40S

🟣 AWS Trainium2

GPU Serving

Trainium Serving

PyTorch Native on Neuron — Zero Friction Porting

GPU Code (PyTorch + CUDA)

Trainium Code (PyTorch Native)

Watch it end-to-end: our GPU training script → Trainium

Neuron Explorer — Qwen3-8B Inference Profile

☸️ vLLM on EKS — Live Cluster

☸️ EKS Deployment Architecture

vLLM Deployment

Live CloudWatch Dashboard