How token-aware load balancing improves LLM throughput

Token-aware load balancing routes LLM requests based on estimated token consumption rather than simple request counts. Traditional round-robin or least-connections algorithms treat every request as equal, but a 4,000-token completion and a 50-token classification task impose vastly different loads on the same model. This guide walks through implementing token-aware routing, the metrics you need, and the tradeoffs that appear at scale.

Why request-count balancing fails for LLMs

Standard load balancers assume uniform work per request. With LLMs, that assumption breaks immediately. A single chat completion with a 32k context window and 4k output tokens can consume 100x the compute of a sentiment classification request. When you balance by request count, one heavy request monopolizes a replica while others sit idle.

The symptom: p99 latency spikes while average utilization looks healthy. You see this pattern when a few long-generation requests queue behind short ones, or when a replica handling multiple large contexts hits KV cache pressure while peers serve tiny prompts.

Token-aware load balancing solves this by estimating the token budget of each request upfront and routing to the replica with the most available token capacity.

Estimating token cost before routing

You cannot route by tokens if you don’t know the token count. The estimator runs at the gateway layer, before the request hits any model replica.

Prompt token estimation

For the prompt, you have the full text. Use the model’s tokenizer directly — this is exact, not estimated.

# gateway/token_estimator.py
from transformers import AutoTokenizer
from functools import lru_cache

@lru_cache(maxsize=8)
def get_tokenizer(model_id: str) -> AutoTokenizer:
    # Map your internal model IDs to HF tokenizer names
    tokenizer_map = {
        "gpt-4o": "cl100k_base",
        "llama-3-70b": "meta-llama/Meta-Llama-3-70B",
        "mixtral-8x7b": "mistralai/Mixtral-8x7B-v0.1",
    }
    return AutoTokenizer.from_pretrained(tokenizer_map[model_id])

def estimate_prompt_tokens(model_id: str, messages: list[dict]) -> int:
    tokenizer = get_tokenizer(model_id)
    # Apply chat template if needed
    text = tokenizer.apply_chat_template(messages, tokenize=False)
    return len(tokenizer.encode(text))

Cache tokenizers aggressively. Loading a 70B tokenizer takes ~200ms cold; warm calls are microseconds.

Output token estimation

Output tokens are unknown until generation completes. You have three strategies, ranked by accuracy:

1. Client-declared max_tokens — Use the request’s max_tokens parameter as an upper bound. Safe but often overestimates by 3-10x.
2. Historical percentiles — Track actual output tokens per (model, task_type) pair. Route using p90 or p99 of recent history.
3. Heuristic by task type — Classify the request (chat, code, summary, classification) and apply a fixed multiplier.

# gateway/output_estimator.py
from dataclasses import dataclass
from collections import defaultdict
import statistics

@dataclass
class OutputEstimate:
    estimated_tokens: int
    confidence: float  # 0.0-1.0

class OutputTokenEstimator:
    def __init__(self):
        self.history: dict[tuple[str, str], list[int]] = defaultdict(list)
        self.task_heuristics = {
            "chat": 500,
            "code": 800,
            "summary": 300,
            "classification": 20,
            "extraction": 100,
        }

    def estimate(self, model_id: str, task_type: str, max_tokens: int | None) -> OutputEstimate:
        key = (model_id, task_type)
        if len(self.history[key]) >= 100:
            p90 = statistics.quantiles(self.history[key], n=10)[8]
            return OutputEstimate(min(p90, max_tokens or p90), confidence=0.8)
        
        heuristic = self.task_heuristics.get(task_type, 500)
        return OutputEstimate(min(heuristic, max_tokens or heuristic), confidence=0.3)

    def record_actual(self, model_id: str, task_type: str, actual_tokens: int):
        self.history[(model_id, task_type)].append(actual_tokens)
        # Keep rolling window
        if len(self.history[(model_id, task_type)]) > 10000:
            self.history[(model_id, task_type)] = self.history[(model_id, task_type)][-5000:]

The estimator feeds a total_estimated_tokens = prompt_tokens + estimated_output_tokens value into the router.

Router implementation: least-token-load

The core algorithm: route each request to the replica with the lowest projected token load after accepting the request.

# gateway/router.py
import heapq
from dataclasses import dataclass, field
from threading import Lock
from typing import Optional

@dataclass(order=True)
class ReplicaState:
    projected_tokens: int
    replica_id: str = field(compare=False)
    active_requests: int = field(compare=False, default=0)
    max_context_tokens: int = field(compare=False, default=128000)

class TokenAwareRouter:
    def __init__(self, replicas: list[dict]):
        """
        replicas: [{"id": "gpu-3", "max_context_tokens": 128000}, ...]
        """
        self.heap: list[ReplicaState] = [
            ReplicaState(0, r["id"], 0, r["max_context_tokens"]) for r in replicas
        ]
        heapq.heapify(self.heap)
        self.lock = Lock()
        self.replica_map = {r.replica_id: r for r in self.heap}

    def select_replica(self, estimated_tokens: int) -> Optional[str]:
        with self.lock:
            # Pop replicas that would exceed context window
            candidates = []
            while self.heap:
                replica = heapq.heappop(self.heap)
                if replica.projected_tokens + estimated_tokens <= replica.max_context_tokens:
                    candidates.append(replica)
                    break
                candidates.append(replica)  # Still track for re-insert
            
            if not candidates:
                return None  # All replicas at capacity
            
            selected = candidates[0]
            selected.projected_tokens += estimated_tokens
            selected.active_requests += 1
            
            # Re-insert all popped replicas
            for r in candidates:
                heapq.heappush(self.heap, r)
            
            return selected.replica_id

    def release(self, replica_id: str, actual_tokens: int):
        with self.lock:
            replica = self.replica_map[replica_id]
            # Remove from heap (mark stale, rebuild periodically)
            replica.projected_tokens = max(0, replica.projected_tokens - actual_tokens)
            replica.active_requests = max(0, replica.active_requests - 1)
            # Rebuild heap every N releases to handle staleness
            if self._should_rebuild():
                self._rebuild_heap()

    def _should_rebuild(self) -> bool:
        # Simple heuristic: rebuild every 100 releases
        return hasattr(self, '_release_count') and (self._release_count := self._release_count + 1) % 100 == 0

    def _rebuild_heap(self):
        self.heap = list(self.replica_map.values())
        heapq.heapify(self.heap)

This is a minimal implementation. Production systems need:

• Staleness handling (replicas report actual usage periodically)
• Health checks (remove unhealthy replicas from heap)
• Priority lanes (VIP customers, latency-sensitive tasks)
• Cross-region awareness

Integrating with the request path

The estimator and router sit in your gateway’s request flow:

# gateway/main.py
from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel
from typing import Optional

app = FastAPI()

prompt_estimator = PromptTokenEstimator()
output_estimator = OutputTokenEstimator()
router = TokenAwareRouter(replicas=load_replica_config())

class ChatRequest(BaseModel):
    model: str
    messages: list[dict]
    max_tokens: Optional[int] = None
    task_type: Optional[str] = "chat"  # Client hint, optional

@app.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest, http_request: Request):
    # 1. Estimate prompt tokens (exact)
    prompt_tokens = prompt_estimator.estimate(request.model, request.messages)
    
    # 2. Estimate output tokens
    output_estimate = output_estimator.estimate(
        request.model, 
        request.task_type, 
        request.max_tokens
    )
    
    total_estimated = prompt_tokens + output_estimate.estimated_tokens
    
    # 3. Route
    replica_id = router.select_replica(total_estimated)
    if not replica_id:
        raise HTTPException(503, "All replicas at token capacity")
    
    # 4. Forward to replica, track actual usage
    try:
        response = await forward_to_replica(replica_id, request)
        actual_output = response.usage.completion_tokens
        output_estimator.record_actual(request.model, request.task_type, actual_output)
        router.release(replica_id, prompt_tokens + actual_output)
        return response
    except Exception as e:
        router.release(replica_id, total_estimated)  # Release estimate on error
        raise

Handling heterogeneous model fleets

Real fleets mix model sizes (7B, 70B, 400B+), quantization levels (FP16, INT4, FP8), and hardware (H100, A100, L4). Token-aware routing must account for token throughput per second, not just token capacity.

Normalize by throughput

# gateway/capacity.py
REPLICA_SPECS = {
    "gpu-3": {"model": "llama-3-70b", "quant": "FP8", "gpu": "H100", "tokens_per_sec": 4500},
    "gpu-7": {"model": "llama-3-70b", "quant": "INT4", "gpu": "A100", "tokens_per_sec": 2800},
    "gpu-12": {"model": "mixtral-8x7b", "quant": "FP16", "gpu": "H100", "tokens_per_sec": 3200},
}

def normalize_tokens(tokens: int, replica_id: str) -> float:
    """Convert raw tokens to 'H100-FP8-equivalent tokens' for fair comparison."""
    spec = REPLICA_SPECS[replica_id]
    baseline = 4500  # H100 FP8 70B
    return tokens * (baseline / spec["tokens_per_sec"])

Use normalized tokens in the heap ordering. A 70B INT4 replica on A100 shows higher projected load for the same token count, so the router naturally prefers faster replicas for large jobs.

Context window awareness

Different models have different context limits. The router must reject or truncate requests that exceed a replica’s window before routing.

def can_fit(model_id: str, prompt_tokens: int, estimated_output: int, replica_id: str) -> bool:
    spec = REPLICA_SPECS[replica_id]
    max_context = MODEL_CONTEXT_WINDOWS[spec["model"]]
    return prompt_tokens + estimated_output <= max_context

Metrics to expose

You cannot tune what you cannot see. Export these from the gateway:

Alert on estimation_error_ratio p99 > 2.0 — your output estimator is drifting.

Common pitfalls

Over-relying on max_tokens

Clients set max_tokens=4096 but generate 50 tokens. If you route by the declared max, you waste capacity. Always blend with historical actuals. The OutputTokenEstimator above does this, but you must seed history before trusting it.

Ignoring KV cache pressure

Token count correlates with KV cache usage, but not perfectly. Multi-turn conversations reuse prefix cache; parallel requests with different prompts fragment it. Track active_sequences alongside token load, and add a penalty factor for replicas with high sequence counts.

# In ReplicaState, add:
kv_cache_pressure: float = field(compare=False, default=0.0)  # 0.0-1.0

# In select_replica, adjust projected_tokens:
effective_load = replica.projected_tokens * (1 + replica.kv_cache_pressure * 0.5)

Stale projected loads

The heap accumulates error as actual tokens diverge from estimates. Rebuild the heap from ground truth every 50-100 releases, or have replicas push actual usage every 10 seconds via a side channel.

Cold-start tokenizer latency

First request for a model triggers tokenizer download/load. Pre-warm tokenizers at gateway startup, or use a lightweight approximation (character count / 3.5 for English) for the first N requests while the real tokenizer loads in background.

Tradeoffs

Token-aware routing adds ~1-2ms latency at the gateway (tokenizer + heap ops). At high QPS, batch estimate calls or use a faster approximate tokenizer (tiktoken Rust bindings, or a character-based heuristic for routing only).

When to add priority lanes

Not all tokens are equal. A 500-token legal summary for a paying customer should preempt a 2000-token hobbyist chat. Extend the router with priority bands:

@dataclass(order=True)
class PriorityRequest:
    priority: int  # Lower = higher priority
    estimated_tokens: int
    request_id: str = field(compare=False)
    replica_id: str = field(compare=False)

Maintain separate heaps per priority band, or use a single heap with (priority, projected_tokens) as the sort key. Drain higher-priority bands first. This prevents token-aware balancing from starving latency-sensitive traffic.

Testing the router

Unit test the estimator and router in isolation. Integration test with a fake replica that simulates variable generation times.

# tests/test_router.py
import pytest
from gateway.router import TokenAwareRouter, ReplicaState

def test_routes_to_least_loaded():
    router = TokenAwareRouter([
        {"id": "gpu-1", "max_context_tokens": 1000},
        {"id": "gpu-2", "max_context_tokens": 1000},
    ])
    
    # First request goes to gpu-1 (tie-break by heap order)
    assert router.select_replica(100) == "gpu-1"
    
    # Second request goes to gpu-2 (lower projected load)
    assert router.select_replica(100) == "gpu-2"
    
    # Third request goes to gpu-1 (both at 100, tie-break)
    assert router.select_replica(100) == "gpu-1"

def test_rejects_when_full():
    router = TokenAwareRouter([{"id": "gpu-1", "max_context_tokens": 500}])
    assert router.select_replica(300) == "gpu-1"
    assert router.select_replica(300) is None  # Would exceed 500

def test_release_updates_load():
    router = TokenAwareRouter([{"id": "gpu-1", "max_context_tokens": 1000}])
    router.select_replica(200)
    router.release("gpu-1", 150)  # Actual less than estimated
    # Next request should see 50 projected
    assert router.select_replica(600) == "gpu-1"  # 50 + 600 <= 1000

Load test with a mix of request sizes. Verify p99 latency improves vs round-robin at same replica count.

Scaling the gateway

The router is a single point of coordination. At ~10k QPS, the lock becomes a bottleneck. Options:

1. Shard by model — One router instance per model, stateless, horizontally scalable.
2. Lock-free heap — Use a concurrent priority queue (e.g., folly::ConcurrentPriorityQueue in C++, or partition the heap by replica ID modulo N).
3. Push to replicas — Replicas advertise capacity via etcd/Consul; clients or a thin proxy route directly. Removes central router but adds eventual consistency.

For most teams, sharding by model with a stateless router per shard handles 50k+ QPS on a single CPU core.

Summary

Token-aware load balancing replaces request-count equality with token-budget awareness. The implementation requires:

1. Exact prompt tokenization at the gateway
2. Output token estimation via historical percentiles blended with client hints
3. A least-projected-load router using a min-heap, normalized by replica throughput
4. Continuous feedback: record actuals, rebuild heap, alert on estimation drift

Start with the estimator and router in a single gateway process. Measure estimation_error_ratio and replica_token_utilization variance. When variance drops and p99 latency improves, you’ve validated the approach. Then consider priority lanes, cross-region routing, and replica-autoscaling hooks — all of which become tractable once token load is a first-class routing signal.