Multi-node LLM inference scaling using AWS Trainium and vLLM for conversational AI shopping assistant

rufus · E-commerce · 2025 · source Agentic: no · promotional bias: med · source-backed · extracted with prompt 0.5.6 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: Large-scale conversational shopping assistant inference

Amazon's Rufus team needed to serve increasingly large custom LLMs for millions of shopping-assistant users after model size outgrew single-node accelerator memory.

Constraints:

Model complexity exceeded single-node memory capacity.
Distributed inference introduces model sharding, communication, and fault-tolerance complexity.
Prime Day traffic requires large-scale reliability.
Multi-node collective operations need low-latency high-bandwidth networking.

Architecture

Rufus uses multi-node inference on AWS Trainium with vLLM, ECS, leader/follower inference units, hybrid parallelism, topology-aware placement, EFA networking, Triton, NxDI, proxies, and containerized rolling deployments.

Diagram confidence: reconstructed (reconstructed from prose).

flowchart LR
  load-balancing-infrastructure["Load Balancing Infrastructure"]
  proxy-layer["Proxy Layer"]
  leader-node["Leader Node"]
  follower-nodes["Follower Nodes"]
  trainium-chips["Trainium Chips"]
  ecs-multi-node-unit["ECS Multi-Node Unit"]
  efa-and-nxdi-communication["EFA and NxDI Communication"]
  load-balancing-infrastructure --> proxy-layer
  proxy-layer --> load-balancing-infrastructure
  proxy-layer --> leader-node
  leader-node --> follower-nodes
  leader-node --> proxy-layer
  follower-nodes --> leader-node
  follower-nodes --> trainium-chips
  trainium-chips --> leader-node
  trainium-chips --> follower-nodes
  ecs-multi-node-unit --> leader-node
  ecs-multi-node-unit --> follower-nodes
  efa-and-nxdi-communication --> leader-node
  efa-and-nxdi-communication --> follower-nodes

Drag to pan · scroll or use +/− to zoom · reconstructed from the source's described component topology

Components

Component	Role	Interacts with
`Load Balancing Infrastructure`	Routes conversational AI traffic toward healthy multi-node inference units.	`Proxy Layer`
`Proxy Layer`	Monitors health and load and reports routing signals.	`Load Balancing Infrastructure`, `Leader Node`
`Leader Node`	Runs Triton and vLLM for scheduling, batching, and coordination.	`Follower Nodes`, `Proxy Layer`
`Follower Nodes`	Execute distributed model computations in parallel with the leader.	`Leader Node`, `Trainium Chips`
`Trainium Chips`	Provide accelerator capacity for custom LLM inference.	`Leader Node`, `Follower Nodes`
`ECS Multi-Node Unit`	Packages leader and follower containers as deployable distributed replicas.	`Leader Node`, `Follower Nodes`
`EFA and NxDI Communication`	Provides low-latency collective operations for distributed inference.	`Leader Node`, `Follower Nodes`

Data flow

Traffic reaches a proxy and leader node. The leader schedules and batches requests with Triton and vLLM, broadcasts inputs to followers, and the nodes execute distributed computation using Gloo for input movement and NxDI/EFA for high-bandwidth collectives.

“The leader node serves as the primary orchestration unit, running both NVIDIA Triton Inference Server and the vLLM engine to handle request scheduling, batching, and coordination.”summary, stated

“The follower nodes operate as distributed execution units that receive broadcasted model inputs and execute distributed model computations in parallel with the leader.”summary, stated

Why this architecture

The leader/follower abstraction centralizes scheduling and batching while distributing computation across Trainium nodes. Hybrid parallelism matches different inference phases, and topology-aware placement reduces communication bottlenecks.

“By centralizing scheduling and batching logic in the leader node, the system maintains the benefits of vLLM's continuous batching capabilities while distributing the computational load across multiple nodes.”summary, stated

“The team implemented sophisticated hybrid parallelism strategies through integration with AWS's Neuron SDK to maximize hardware utilization across their distributed infrastructure.”summary, stated

Models used: custom Rufus language models · vLLM · AWS Trainium

Scale

Rufus serves millions of customers.
“deploying increasingly large custom language models for their generative AI shopping assistant serving millions of customers.”summary, stated
The deployment spans tens of thousands of Trainium chips.
“successfully deploy across tens of thousands of Trainium chips”summary, stated
The system supported Prime Day traffic.
“supporting Prime Day traffic while delivering the performance and reliability required for production-scale conversational AI.”summary, stated

Alternatives considered (roads not taken)

Single-node inference

Rejected because: Larger custom models no longer fit into one accelerator instance.

“as model size and capability requirements grow, single accelerator instances become insufficient to host the entire model in memory.”summary, stated

Static parallelism for all inference phases

Rejected because: Prefill and decoding have different computational patterns, so hybrid strategies improve utilization.

“The prefill phase typically involves processing longer sequences but smaller batch sizes, making sequence-level parallelism more effective.”summary, stated

Patterns

`durable-execution` (operations)

ECS and cell-by-cell rolling deployments treat distributed replicas as operational units.

“consistent rolling deployments on a cell-by-cell basis across their large-scale fleet.”summary, stated

`observability-tracing` (operations)

Proxy-layer health monitoring and load reporting support routing and reliability.

“A proxy layer positioned between the load balancing infrastructure and multi-node inference units provides continuous health monitoring and load reporting.”summary, stated

Tradeoffs

Centralize scheduling in a leader

Chose Leader/follower architecture over Fully decentralized inference scheduling.

Rationale: Central scheduling preserves continuous batching while distributing computation.

“centralizing scheduling and batching logic in the leader node”summary, stated

Use hybrid parallelism

Chose Context parallelism for prefill and data parallelism for decoding over One fixed parallelism strategy.

Rationale: Different inference phases benefit from different partitioning.

“During the context encoding (prefill) phase, context parallelism splits inputs along the sequence dimension”summary, stated

Gotchas (what bit them)

Distributed inference adds collective-communication bottlenecks.

Root cause: Multi-node model execution requires high-bandwidth, low-latency cross-node operations.

Resolution: Use EFA networking, Neuron Distributed Inference, and topology-aware placement.

“high-bandwidth, low-latency communication required for efficient distributed matrix operations.”summary, stated

Distributed replicas are harder to deploy and monitor than single containers.

Root cause: Leader and follower nodes must stay healthy and coordinated.

Resolution: Use ECS abstractions, proxy health monitoring, and load reporting.

“treats distributed model replicas as single deployable units.”summary, stated

Evaluation

The case emphasizes production reliability, throughput, hardware utilization, health monitoring, and traffic support rather than a single offline benchmark.

Metrics: Prime Day traffic support · hardware utilization · health monitoring · load reporting

“The health monitoring capability is particularly important in distributed inference scenarios”summary, stated

Outcomes

Rufus deployed multi-node LLM inference across tens of thousands of Trainium chips.
The system supported Prime Day traffic.
The architecture enabled models larger than single-node memory capacity.

Millions of customers, tens of thousands of Trainium chips, and Prime Day production traffic.

“The case study details how Amazon's engineering team addressed the significant technical challenges of scaling custom large language model inference beyond single-node capacity”summary, stated