Load Balancing for LLMs

Standard load balancing assumes requests are roughly equal in cost and duration. LLM requests violate both assumptions: a 50-token query takes 200ms while a 100K-token context takes 30 seconds. Round-robin sends the 100K request to the same instance as three quick ones, creating a 30-second head-of-line block. LLM load balancing must account for variable request duration, KV cache state, and prefix sharing.

---

1. Why Standard Strategies Fail

Strategy	Assumption	LLM reality	Problem
Round-robin	Equal request cost	Costs vary 100×	Unbalanced load
Least-connections	Equal connection weight	One long stream ≠ one short	Still unbalanced
Random	Equal probability works	Same issues as round-robin	Same problems
IP hash	Sticky sessions help	No per-user state (typically)	Wastes affinity

All these strategies ignore the two things that matter most: how busy each instance actually is (queue depth in tokens, not requests) and what KV cache state each instance has (prefix sharing).

---

2. Queue-Depth-Aware Routing

Route to the instance with the fewest pending tokens, not the fewest connections:

class QueueDepthRouter:
    """Route to instance with lowest token queue depth."""
    
    def __init__(self, backends: list[Backend]):
        self.backends = backends
    
    async def select(self, request: ChatRequest) -> Backend:
        # Estimate tokens for this request
        input_tokens = count_tokens(request.messages)
        estimated_output = request.max_tokens or 4096
        request_tokens = input_tokens + estimated_output
        
        # Get queue state from each backend
        backend_loads = []
        for backend in self.backends:
            status = await backend.get_status()  # Returns pending tokens, active requests
            backend_loads.append({
                "backend": backend,
                "pending_tokens": status.pending_tokens,
                "active_requests": status.active_requests,
                "estimated_completion": status.pending_tokens / status.throughput_tok_s,
            })
        
        # Select backend with lowest pending tokens
        best = min(backend_loads, key=lambda b: b["pending_tokens"])
        return best["backend"]

Impact: 20-40% improvement in P99 latency compared to round-robin, because long requests don’t block short ones on the same instance.

---

3. Prefix-Cache-Affinity Routing

When requests share common prefixes (system prompts, RAG context), routing them to the same instance enables KV cache reuse:

class PrefixAffinityRouter:
    """Route requests with shared prefixes to the same instance for KV cache reuse."""
    
    def __init__(self, backends: list[Backend], affinity_weight=0.6):
        self.backends = backends
        self.affinity_weight = affinity_weight  # Balance affinity vs load
    
    async def select(self, request: ChatRequest) -> Backend:
        # Extract prefix (system prompt + first N tokens)
        prefix = self._extract_prefix(request.messages)
        prefix_hash = hashlib.md5(prefix.encode()).hexdigest()
        
        # Score each backend
        scores = []
        for backend in self.backends:
            status = await backend.get_status()
            
            # Affinity score: does this backend have the prefix cached?
            has_prefix = prefix_hash in status.cached_prefixes
            affinity_score = 1.0 if has_prefix else 0.0
            
            # Load score: inverse of queue depth (lower is better)
            max_tokens = max(b.pending_tokens for b in self.backends) or 1
            load_score = 1.0 - (status.pending_tokens / max_tokens)
            
            # Combined score
            combined = (self.affinity_weight * affinity_score + 
                       (1 - self.affinity_weight) * load_score)
            scores.append((backend, combined))
        
        # Select highest-scoring backend
        return max(scores, key=lambda x: x[1])[0]
    
    def _extract_prefix(self, messages: list[dict]) -> str:
        """Extract cacheable prefix — typically system prompt."""
        system = next((m["content"] for m in messages if m["role"] == "system"), "")
        return system[:2000]  # First 2000 chars as prefix key

Impact: Prefix cache hit rate 40-80% (vs 0% with round-robin), saving significant prefill computation.

For how prefix caching works inside the engine, see Prefix Caching.

---

4. Token-Budget-Aware Routing

Each instance has a maximum token throughput. Route based on remaining capacity:

class TokenBudgetRouter:
    """Route based on remaining token-processing capacity per instance."""
    
    def __init__(self, backends: list[Backend]):
        self.backends = backends
    
    async def select(self, request: ChatRequest) -> Backend:
        estimated_tokens = count_tokens(request.messages) + (request.max_tokens or 4096)
        
        candidates = []
        for backend in self.backends:
            status = await backend.get_status()
            remaining_capacity = status.max_tokens_per_sec - status.current_tokens_per_sec
            
            # Can this backend handle the request without exceeding capacity?
            if remaining_capacity > estimated_tokens / 10:  # rough: need 10s budget
                candidates.append((backend, remaining_capacity))
        
        if not candidates:
            # All backends near capacity — pick least loaded
            return min(self.backends, key=lambda b: b.pending_tokens)
        
        # Pick backend with most remaining capacity
        return max(candidates, key=lambda x: x[1])[0]

---

5. Combined Strategy

Production systems combine multiple signals:

class CombinedLLMRouter:
    """Multi-signal routing for LLM serving."""
    
    WEIGHTS = {
        "queue_depth":     0.35,  # Avoid overloaded instances
        "prefix_affinity": 0.35,  # Maximize cache hits
        "token_budget":    0.20,  # Stay within throughput limits
        "health":          0.10,  # Prefer healthy instances
    }
    
    async def select(self, request: ChatRequest) -> Backend:
        prefix_hash = self._prefix_hash(request)
        est_tokens = self._estimate_tokens(request)
        
        scores = []
        for backend in self.backends:
            s = await backend.get_status()
            
            score = (
                self.WEIGHTS["queue_depth"] * self._queue_score(s) +
                self.WEIGHTS["prefix_affinity"] * self._affinity_score(s, prefix_hash) +
                self.WEIGHTS["token_budget"] * self._budget_score(s, est_tokens) +
                self.WEIGHTS["health"] * self._health_score(s)
            )
            scores.append((backend, score))
        
        return max(scores, key=lambda x: x[1])[0]
    
    def _queue_score(self, status):
        """0.0 = fully loaded, 1.0 = empty queue."""
        return max(0, 1 - status.pending_tokens / status.max_pending)
    
    def _affinity_score(self, status, prefix_hash):
        """1.0 if prefix cached, 0.0 if not."""
        return 1.0 if prefix_hash in status.cached_prefixes else 0.0
    
    def _budget_score(self, status, est_tokens):
        """1.0 = plenty of capacity, 0.0 = at limit."""
        remaining = status.max_tok_s - status.current_tok_s
        return max(0, min(1, remaining / (est_tokens / 5)))
    
    def _health_score(self, status):
        """1.0 = healthy, 0.0 = unhealthy."""
        if status.error_rate > 0.05:
            return 0.0
        return 1.0 - status.error_rate * 10

---

6. Handling Streaming Connections

LLM responses stream for seconds to minutes. The load balancer must handle long-lived connections:

# Nginx configuration for LLM streaming
upstream llm_backends {
    least_conn;  # Basic, but better than round-robin for LLMs
    server vllm-1:8000;
    server vllm-2:8000;
    server vllm-3:8000;
    keepalive 64;
}

server {
    listen 443 ssl;

    location /v1/chat/completions {
        proxy_pass http://llm_backends;
        
        # Critical for SSE streaming:
        proxy_buffering off;           # Don't buffer response
        proxy_cache off;               # Don't cache
        proxy_http_version 1.1;        # Required for chunked transfer
        proxy_set_header Connection ""; # Keep-alive
        
        # Long timeouts for generation:
        proxy_read_timeout 300s;       # 5 min for long generations
        proxy_send_timeout 300s;
        
        # Headers for SSE:
        add_header X-Accel-Buffering no;
    }
}

Connection Draining

When taking an instance offline (for updates or scaling down), drain gracefully:

async def drain_instance(backend: Backend, timeout=120):
    """Stop sending new requests, wait for active ones to complete."""
    backend.accepting_new = False  # Router stops sending new requests
    
    start = time.time()
    while backend.active_requests > 0:
        if time.time() - start > timeout:
            logger.warning(f"Drain timeout: {backend.active_requests} requests still active")
            break
        await asyncio.sleep(1)
    
    await backend.shutdown()

---

7. Health Checking

LLM backends can fail in subtle ways — running but producing garbage, or accepting requests but timing out:

class LLMHealthChecker:
    """Deep health check for LLM inference instances."""
    
    async def check(self, backend: Backend) -> HealthStatus:
        # Level 1: Is the process alive?
        try:
            status = await backend.get_status(timeout=2)
        except Exception:
            return HealthStatus.DOWN
        
        # Level 2: Is it accepting requests?
        if status.queue_depth > status.max_queue:
            return HealthStatus.OVERLOADED
        
        # Level 3: Is it producing valid output?
        # (Periodic canary request — every 30s)
        if time.time() - self.last_canary > 30:
            try:
                response = await backend.generate(
                    messages=[{"role": "user", "content": "Say hello"}],
                    max_tokens=5,
                    timeout=10,
                )
                if "hello" not in response.lower():
                    return HealthStatus.DEGRADED  # Running but output is wrong
                self.last_canary = time.time()
            except Exception:
                return HealthStatus.DEGRADED
        
        return HealthStatus.HEALTHY

---

Self-Check

• Explain why round-robin and least-connections fail for LLM serving
• Implement queue-depth-aware routing that accounts for token count
• Design prefix-cache-affinity routing for workloads with shared system prompts
• Configure Nginx for SSE streaming with proper timeouts and buffer settings
• Build a deep health checker that validates model output quality

---

References

• vLLM: Efficient Memory Management for LLM Serving
• Splitwise: Efficient Generative LLM Inference with Disaggregation
• For prefix caching internals, see Prefix Caching
• For KV-cache-aware routing, see llm-d