AI API Gateway Token-Aware Load Balancing - Distribute Work, Not Request Counts

Why request counts mislead#

Request-count balancing works when requests are roughly equal. AI requests are not: one might contain a couple of hundred tokens and return a short answer, while another carries a very long context, images, tools and an extended stream. Sending one request to each route looks balanced on a dashboard while a single provider does most of the computation.

The fix is to balance the work itself.

Define and estimate a work unit#

Input tokens, expected output tokens, images, audio duration, tool rounds, cache reads and model-specific compute can all matter. Pick a normalized cost estimate that reflects the bottleneck you are balancing, and be precise about naming it.

Estimate before dispatch: use the model tokenizer where available, count system prompts and tool schemas, inspect media dimensions or duration, and derive an output estimate from the request limit and historical behavior. The estimate does not need to be perfect, but it must be bounded, versioned and conservative enough that obviously large requests cannot look cheap.

Then track committed work per route: estimated input, reserved output, observed generation rate, queued work, retries and long-lived streams. One slot running an enormous generation can be busier than ten short requests elsewhere.

Capacity dimensions and model throughput#

A provider can have spare request concurrency but no tokens-per-minute headroom, or free text capacity while its vision workers are saturated. Keep request, token, media, connection and tool limits distinct when they fail independently, and apply a single scalar score only after hard constraints pass.

The same token count costs different amounts across model families, context lengths, quantization and hardware. Learn service rates per immutable model revision, route, region, modality and request shape, and reset or discount learned rates after a rollout instead of inheriting confidence from a different model.

A practical route score combines queued and committed work, expected service rate, network latency, policy, price and deadline, choosing the route most likely to finish safely. Keep tie-breaking deterministic so identical workers do not cause unstable routing.

Long jobs, caching and streams#

Large requests can block interactive traffic, while always favoring short work starves long jobs. Size classes, weighted fair queues or reserved capacity help, and the contract should state maximum queue delay, deadline behavior and whether large work may move to batch capacity.

Prompt caching complicates the estimate. A large prompt with a valid provider-side cache may cost less than a smaller uncached one, but cache locality creates route affinity. Include cache key, expected hit, freshness and route ownership, and if the hit is uncertain, reserve for the miss and correct the accounting once usage is reported.

Streams are ongoing work. Input cost arrives early but output continues until completion or cancellation, so update committed load using observed token rate and remaining output, and do not free capacity when the first token is sent.

Reconciliation, abuse and policy#

Compare predicted input, output, media, tool and duration costs with provider-reported and gateway-observed values, update estimators gradually, and keep the original prediction in request lineage. Large systematic errors may point to a tokenizer or adapter bug.

Guard against abuse. Clients can set huge output limits to hold reservations, or omit metadata that reveals expensive media. Apply policy caps, authenticated quotas, request-shape validation and reservation expiry, and do not trust client-supplied token counts unless they are verifiable.

Load balancing picks a route; concurrency control decides whether that route can take more work. Do not send traffic to a low-score route whose controller has just contracted capacity. Residency, authorization, model allowlists, quality gates and tenant contracts define the eligible set first, and scoring happens only inside it.

Failover, billing, fairness and measurement#

When a provider or region fails, queued and uncertain attempts may still hold capacity. Move only work whose state is known, reserve conservatively on the fallback, avoid piling every large request onto one surviving route, and apply retry budgets so failure does not multiply token load.

Record predicted work, selected route, pricing version, actual usage, retries and cache outcomes so customers can see why a request cost what it did. A multi-provider gateway such as Crazyrouter depends on this record to keep routing decisions explainable. Enforce tenant token budgets, weighted shares and burst limits based on reconciled work, not admission estimates alone.

Test tiny prompts, maximum context, long output, early stop, media, tools, cache hits and misses, tokenizer mismatch, cancellation and adversarial output limits. Compare against request-count balancing on production-like traffic, because a sophisticated score is only worth having if outcomes improve.