Load balancing with Kong AI Gateway

AI Gateway provides load balancing capabilities to distribute requests across multiple LLM models. You can use these features to improve fault tolerance, optimize resource utilization, and balance traffic across your AI systems.

In AI Gateway, load balancing is configured on the AI Model entity through config.balancer and targets.

Load balancing algorithms

AI Gateway supports multiple load balancing strategies for distributing traffic across AI models. Each algorithm addresses different goals: balancing load, improving cache-hit ratios, reducing latency, or providing failover reliability.

The following table describes the available algorithms for AI Model entities and considerations for selecting one.

For examples of each algorithm, see Algorithm examples in the AI Model entity reference.

Request routing rules

Routing rules allow clients to send an alias instead of the actual model name in the request. This decouples the external model identifier from the internal provider model, enabling flexible routing without changing client code.

An AI Model can have an optional config.route.model field, matching on the request body, a header, or the path. For example, when config.route.model is set to { body_param: model, values: [alias-value] } and a client sends "model": "alias-value" in the request body, AI Gateway routes to the matching AI Model. This feature works independently of load balancing algorithms. The routing rule determines which AI Model (or target set) handles the request, and the configured load balancing algorithm selects the final backend within that set.

Retry and fallback

The load balancer includes built-in support for retries and fallbacks. When a request fails, the balancer can automatically retry the same target or redirect the request to a different target model.

How retry and fallback works

1. Client sends a request.
2. The load balancer selects a target based on the configured algorithm (round-robin, lowest-latency, etc.).

3. If the target fails (based on defined failover_criteria), the balancer:

   • Retries the same or another target.
   • Fallbacks to another available target.
4. If retries are exhausted without success, the load balancer returns a failure to the client.

 
flowchart LR
    Client(((Application))) --> LBLB
    subgraph AIGateway
        LBLB[/Load Balancer/]
    end
    LBLB -->|Request| AIProvider1(AI Model Provider 1)
    AIProvider1 --> Decision1{Is Success?}
    Decision1 -->|Yes| Client
    Decision1 -->|No| AIProvider2(AI Model Provider 2)
    subgraph Retry
        AIProvider2 --> Decision2{Is Success?}
    end
    Decision2 ------>|Yes| Client
  

Figure 1: A simplified diagram of fallback and retry processing in AI Gateway’s load balancer.

Retry and fallback configuration

The AI Gateway load balancer supports fine-grained control over failover behavior on the Model entity. Use failover_criteria to define when a request should retry on the next target model. By default, retries occur on error and timeout. An error means a failure occurred while connecting to the server, forwarding the request, or reading the response header. A timeout indicates that any of those stages exceeded the allowed time.

You can add more criteria to adjust retry behavior as needed:

Retry and fallback scenarios

You can customize the AI Gateway load balancer to fit different application needs, such as minimizing latency, enabling sticky sessions, or optimizing for cost. The table below maps common scenarios to key configuration options that control load balancing behavior:

Health check and circuit breaker

For AI Model entity, circuit breaker behavior is controlled through the balancer configuration on the AI Model. Use these settings to fail fast when a target model is unhealthy and to retry or fall back to another target instead of waiting for repeated slow responses.

The load balancer supports health checks and circuit breakers to improve reliability. If the number of unsuccessful attempts to a target reaches config.balancer.max_fails, the load balancer stops sending requests to that target until it reconsiders the target after the period defined by config.balancer.fail_timeout. The diagram below illustrates this behavior:

Consider an example where config.balancer.max_fails is 3 and config.balancer.fail_timeout is 10 seconds. When failed requests for a target reach 3, the target is marked unhealthy and the load balancer stops sending requests to it. After 10 seconds, the target is reconsidered. If the request to this target still fails, the target remains unhealthy and the load balancer continues to exclude it. If the request succeeds, the target is marked healthy again and recovers from the circuit breaker.

The failure counter tracks total failures, not consecutive failures. If a target receives 2 failed requests, then 1 successful request within the timeout window, the counter remains at 2. The counter resets only when a successful request occurs after config.balancer.fail_timeout has elapsed since the last failed request.

If all targets become unhealthy simultaneously, requests fail with HTTP 500.