LLM Inference Serving: Architecture, Routing & Auto-Scaling › Load Balancing
Language AI HandbookAll chapters LLM Inference Serving: Architecture, Routing & Auto-Scaling › Load BalancingOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Load balancing distributes requests across multiple model instances to maximize throughput and maintain consistent latency. While load balancing is a well-studied problem in distributed systems, LLM serving presents unique challenges that traditional load balancing algorithms don't address well. The variable cost of requests, the long-running nature of generation, and the importance of memory locality all complicate the problem.
Think of load balancing as managing a team of workers who do jobs of wildly varying duration. If you assign jobs by turns (round-robin), a worker who gets a ten-hour job will still be busy long after everyone else has cycled through dozens of short jobs, creating severe imbalance. If you assign by fewest current jobs, you might send the next job to a worker already doing a ten-hour task just because they technically have fewer total jobs than someone with three five-minute tasks. The right approach considers how much work each worker has remaining, not just how many items are on their list.
The basic challenge for LLM load balancing is that "load" is not a single number. An endpoint serving ten short requests that will each complete in 200ms has very different remaining work than an endpoint serving one request that will generate 2000 tokens over 40 seconds. Connection count captures neither the diversity in prompt length nor the diversity in requested output length. Token-aware approaches exist precisely because connection-count metrics mislead the balancer into poor decisions.
Load balancing also interacts with higher-level concerns like session affinity and KV cache efficiency. The best endpoint for a new request from the standpoint of pure load might not be the best endpoint when you factor in that routing to a specific instance would allow the server to reuse cached KV states from a prior request in the same conversation. These multi-factor decisions are a microcosm of the broader production serving challenge: each individual optimization creates interference with others, and achieving good overall performance requires reasoning about the whole system simultaneously.
Traditional Algorithms
Standard load balancing approaches provide a foundation for understanding the problem, even if they don't fully address LLM-specific challenges.
Round-robin distributes requests sequentially across endpoints in a fixed order. The first request goes to endpoint 0, the second to endpoint 1, and so on, cycling back to 0 after reaching the last endpoint. This predictable approach is easy to implement and gives each endpoint the same request count. However, it ignores actual load, which matters when requests have varying computational costs. A long generation request and a short one count equally, even though they impose vastly different burdens.
Least connections routes each new request to the endpoint with the fewest active requests. This approach adapts better to variable-duration requests because endpoints processing long requests naturally accumulate fewer connections over time. However, it still treats all requests as equal cost, which fails to capture the difference between generating 10 tokens and generating 1000 tokens.
Weighted distribution assigns requests proportionally to endpoint capacity, which is useful when endpoints have different hardware. An H100 GPU might receive twice as many requests as an A100 because it can process them twice as fast. Weights can be configured manually based on hardware specs or adjusted dynamically based on observed performance.
import random
from collections import defaultdict
from typing import List
class LoadBalancer:
"""Implements various load balancing strategies."""
def __init__(self, endpoints: List[ModelEndpoint]):
self.endpoints = endpoints
self.round_robin_index = 0
self.request_counts = defaultdict(int)
def round_robin(self) -> Optional[ModelEndpoint]:
"""Simple round-robin selection."""
healthy = [e for e in self.endpoints if e.is_healthy]
if not healthy:
return None
selected = healthy[self.round_robin_index % len(healthy)]
self.round_robin_index += 1
return selected
def least_connections(self) -> Optional[ModelEndpoint]:
"""Select endpoint with lowest current load."""
healthy = [e for e in self.endpoints if e.is_healthy]
if not healthy:
return None
return min(healthy, key=lambda e: e.current_load)
def weighted_random(self) -> Optional[ModelEndpoint]:
"""Random selection weighted by available capacity."""
healthy = [e for e in self.endpoints if e.is_healthy]
if not healthy:
return None
total_capacity = sum(e.available_capacity for e in healthy)
if total_capacity == 0:
return None
weights = [e.available_capacity / total_capacity for e in healthy]
return random.choices(healthy, weights=weights)[0]# Create endpoints with varying capacities and loads
endpoints = [
ModelEndpoint("gpu-0", "llama", capacity=100, current_load=30),
ModelEndpoint("gpu-1", "llama", capacity=100, current_load=70),
ModelEndpoint("gpu-2", "llama", capacity=150, current_load=50),
]
balancer = LoadBalancer(endpoints)
# Simulate 1000 routing decisions with each strategy
strategies = {
"Round Robin": balancer.round_robin,
"Least Connections": balancer.least_connections,
"Weighted Random": balancer.weighted_random,
}
results = {}
for name, strategy in strategies.items():
# Reset state
balancer.round_robin_index = 0
counts = defaultdict(int)
for _ in range(1000):
selected = strategy()
if selected:
counts[selected.endpoint_id] += 1
results[name] = dict(counts)Request Distribution (1000 requests): -------------------------------------------------- Round Robin: gpu-0: 334 (33.4%) gpu-1: 333 (33.3%) gpu-2: 333 (33.3%) Least Connections: gpu-0: 1000 (100.0%) gpu-1: 0 (0.0%) gpu-2: 0 (0.0%) Weighted Random: gpu-0: 346 (34.6%) gpu-1: 165 (16.5%) gpu-2: 489 (48.9%)
The results illustrate the behavioral differences between strategies. Round-robin distributes evenly regardless of load, giving each endpoint roughly 33% of requests. Least connections concentrates traffic on the least-loaded endpoint (gpu-0), which may improve its latency but could overload it if requests are long-running. Weighted random accounts for available capacity, directing more requests to endpoints with headroom while maintaining some randomness to avoid perfect synchronization effects.
Token-Aware Load Balancing
Traditional metrics like connection count fail to capture LLM workload characteristics because they treat all requests as having equal cost. In reality, a request generating 1000 tokens consumes far more GPU resources and time than one generating 10 tokens, even though both count as single connections. Similarly, a request with a 4000-token prompt requires substantially more prefill computation than one with a 100-token prompt.
The key insight is that LLM load should be measured in tokens, not requests. When we measure load in tokens, we can compute a much more accurate picture of how much remaining work each endpoint has. An endpoint with three requests each generating 20 tokens has roughly the same remaining work as one request generating 60 tokens. An endpoint-count metric would show the first endpoint as having 3x the load when in fact they are equivalent.
Token-aware load balancing accounts for the actual computational cost of requests by tracking token counts rather than just request counts. We also need to distinguish between tokens in the prefill queue (waiting to be processed) and tokens in the decode phase (currently being generated), because these phases have very different costs. The decode phase is the throughput bottleneck in most production systems because memory bandwidth limits how quickly individual decode steps execute, so active decode tokens should be weighted more heavily than queued prefill tokens.
from dataclasses import dataclass
@dataclass
class TokenAwareEndpoint:
"""Endpoint that tracks load in tokens, not just requests."""
endpoint_id: str
max_tokens_per_second: int # Throughput capacity
current_prompt_tokens: int = 0 # Tokens in prefill queue
current_decode_tokens: int = 0 # Tokens being decoded
active_requests: int = 0
def estimated_load(self) -> float:
"""
Estimate current load as fraction of capacity.
Weights decode tokens more heavily since decode is the bottleneck.
"""
# Prefill is compute-bound, decode is memory-bound
# Decode throughput is typically 10-50x lower than prefill
prefill_cost = self.current_prompt_tokens * 0.1
decode_cost = self.current_decode_tokens * 1.0
total_cost = prefill_cost + decode_cost
return total_cost / self.max_tokens_per_second
class TokenAwareBalancer:
"""Load balancer that considers token counts."""
def __init__(self, endpoints: List[TokenAwareEndpoint]):
self.endpoints = endpoints
def select_endpoint(
self, prompt_tokens: int, estimated_output_tokens: int
) -> Optional[TokenAwareEndpoint]:
"""Select endpoint considering the incoming request's token cost."""
healthy = [e for e in self.endpoints if e.estimated_load() < 0.95]
if not healthy:
return None
# Find endpoint that will have lowest load after adding this request
def projected_load(endpoint: TokenAwareEndpoint) -> float:
current = endpoint.estimated_load()
added_cost = prompt_tokens * 0.1 + estimated_output_tokens
added_load = added_cost / endpoint.max_tokens_per_second
return current + added_load
return min(healthy, key=projected_load)# Demonstrate token-aware balancing
token_endpoints = [
TokenAwareEndpoint(
"gpu-0",
max_tokens_per_second=100,
current_prompt_tokens=500,
current_decode_tokens=20,
),
TokenAwareEndpoint(
"gpu-1",
max_tokens_per_second=100,
current_prompt_tokens=100,
current_decode_tokens=60,
),
]
token_balancer = TokenAwareBalancer(token_endpoints)
# Route requests with different characteristics
requests = [
(100, 50), # Short prompt, short output
(2000, 100), # Long prompt, medium output
(100, 500), # Short prompt, long output
]
routing_decisions = []
for prompt_len, output_len in requests:
selected = token_balancer.select_endpoint(prompt_len, output_len)
routing_decisions.append((prompt_len, output_len, selected))Endpoint Token Loads:
gpu-0: 70.0% utilized
Prefill queue: 500 tokens
Active decode: 20 tokens
gpu-1: 70.0% utilized
Prefill queue: 100 tokens
Active decode: 60 tokens
Routing Decisions:
100 prompt + 50 output -> gpu-0
2000 prompt + 100 output -> gpu-0
100 prompt + 500 output -> gpu-0Token-aware balancing makes smarter routing decisions by considering the actual work each endpoint is performing. In this example, gpu-0 has more prompt tokens queued for prefill, but gpu-1 has more active decode tokens. Because decode is the bottleneck phase (each token requires reading the entire KV cache), the long-output request routes to gpu-0, which has less active decoding. This prevents one endpoint from becoming decode-saturated while another sits idle during its compute-bound prefill phase.
Session Affinity Considerations
Some applications benefit from routing subsequent requests to the same endpoint to enable optimization opportunities. If a conversation spans multiple API calls with the same system prompt or shared context, routing to the same instance enables KV cache reuse through prefix caching, as discussed earlier in this part. Rather than recomputing attention states for the repeated prefix, the server can reuse cached values, significantly accelerating prefill.
However, session affinity conflicts with optimal load balancing because it constrains routing decisions. If a particular session's preferred endpoint becomes overloaded, the system must choose between maintaining affinity (and accepting higher latency) or breaking affinity (and losing cache benefits). A simple compromise routes requests with known session IDs to consistent endpoints using hash-based assignment, while distributing new sessions based on current load.
The trade-off between affinity and balance is most pronounced for long-running conversations with consistent system prompts. For a chatbot with a fixed system prompt that consumes 500 tokens, routing all turns of the same conversation to the same instance could save 500 tokens of prefill computation per turn after the first, a potentially large saving for high-volume deployments. Whether this benefit outweighs the load balancing cost depends on your traffic patterns: if conversations are short and system prompts are short, the benefit is minimal; if conversations are long and system prompts are long, affinity provides meaningful efficiency gains.
import hashlib
class AffinityBalancer:
"""Balancer with optional session affinity."""
def __init__(self, endpoints: List[ModelEndpoint]):
self.endpoints = endpoints
def select_endpoint(
self, session_id: Optional[str] = None
) -> Optional[ModelEndpoint]:
healthy = [e for e in self.endpoints if e.is_healthy]
if not healthy:
return None
if session_id:
# Hash session to consistently select same endpoint
hash_val = int(hashlib.md5(session_id.encode()).hexdigest(), 16)
index = hash_val % len(healthy)
return healthy[index]
else:
# No session: use least connections
return min(healthy, key=lambda e: e.current_load)affinity_balancer = AffinityBalancer(endpoints)
# Same session always routes to same endpoint
session = "user-123-conv-456"
session_routes = []
for i in range(3):
selected = affinity_balancer.select_endpoint(session_id=session)
session_routes.append((i + 1, selected))
# Different sessions may route differently
users = ["user-100", "user-200", "user-300"]
user_routes = []
for user in users:
selected = affinity_balancer.select_endpoint(session_id=user)
user_routes.append((user, selected))Session 'user-123-conv-456' routes to: Request 1: gpu-0 Request 2: gpu-0 Request 3: gpu-0 Different sessions: user-100: gpu-0 user-200: gpu-1 user-300: gpu-0
Consistent hashing ensures the same session always reaches the same endpoint, assuming the endpoint pool remains stable. This deterministic mapping enables KV cache hits for repeated prompts within a conversation. When endpoints are added or removed, consistent hashing minimizes disruption by changing the mapping for only a fraction of sessions rather than all of them. More sophisticated approaches can combine affinity preferences with load awareness, preferring the affinity endpoint when it has capacity but falling back to other endpoints when necessary.