AIBrix Router

AIBrix Router#

Overview#

The AIBrix Router is a pluggable, intelligent traffic management component embedded in the AIBrix LLM serving stack. It is designed as an Envoy Gateway extension via external processing hooks, serving as the single entry point for all LLM inference requests.

This gateway abstracts away the underlying complexity of managing multiple models, LoRA adapters, heterogeneous GPU backends, and diverse scaling strategies.

Core Principle#

The router maintains a high-frequency local cache of pod metrics via periodic pulls and subscriptions. This allows it to apply sophisticated multi-objective routing logic without blocking on live queries to pods, keeping latency low on the hot path and enabling scaling to thousands of QPS.

The active routing strategy can be set per-request via the routing-strategy HTTP header, or globally via the ROUTING_ALGORITHM environment variable, or through a model config profile annotation.

Detailed Sequence Flow#

Client
  │
  │  POST /v1/chat/completions
  ▼
Envoy
  │
  │  External Processing Hook
  ▼
GatewayPlugin
  │
  │  Make routing decision
  ▼
Router ──────────────────▶ Cache
  │      Query pod metrics    │
  │      & KV state           │
  │ ◀─────────────────────────┘
  │      Return latest metrics
  │
  │  Apply routing algorithm
  │
  │  Forward request to selected pod
  ▼
InferencePod
  │
  │  Return streamed tokens
  ▼
GatewayPlugin ──▶ Envoy ──▶ Client

Supported Routing Strategies#

AIBrix ships with a set of built-in algorithms, each optimized for different workload patterns.

General load balancing

random: routes to a randomly selected pod. Useful as a baseline or when pods are homogeneous and load is uniform.
least-request: routes to the pod with the fewest in-flight requests.
least-busy-time: routes to the pod with the least cumulative busy processing time.
least-latency: routes to the pod with the lowest average processing latency.
least-kv-cache: routes to the pod with the smallest current KV cache occupancy (least VRAM used).
least-gpu-cache: routes to the pod with the lowest GPU cache utilization.
least-utilization: routes to the pod with the lowest overall utilization score.
load-balance: capacity-aware weighted least-request routing. Scores each pod as (running_requests + λ·queued_requests) / EWMA(output_tokens_per_second) × (1 + α·(1 − kv_free)²) and selects the pod with the lowest score. The token rate is the gateway-observed rate at which the pod completes output tokens (so faster hardware is learned dynamically and receives proportionally more traffic), and kv_free is the pod’s free KV-cache fraction. λ defaults to 0 (the running count already includes requests queued in the engine) and α to 2. A pod with less than 10% free KV cache scores +Inf and is skipped; if every pod is below that threshold the request still routes, to the pod with the most KV headroom. The three knobs are AIBRIX_LOAD_BALANCE_QUEUED_WEIGHT, AIBRIX_LOAD_BALANCE_KV_PRESSURE_ALPHA and AIBRIX_LOAD_BALANCE_KV_CRITICAL_FREE. A pod whose engine reports no KV-cache usage, or reports NaN, is treated as fully free (no penalty and no guardrail). Ties are broken using least combined GPU+CPU KV-cache usage (falling back to a random pick if cache metrics are unavailable for the tied pods). A pod with no token-rate estimate yet is scored at the mean estimate of the others (uniform capacity when none has one). See pkg/plugins/gateway/algorithms/load_balance.md. Its load-imbalance gate, which restricts candidates to the least-loaded pods when load is severely skewed, is applied centrally by the gateway ahead of whichever strategy actually routes each request — not just when load-balance itself is selected (see pkg/plugins/gateway/ENV_VARS.md).
throughput: routes to the pod that has processed the fewest total weighted tokens, favoring underloaded pods.
power-of-two: applies the power-of-two choices algorithm — samples two pods and selects the better one.

KV-cache aware

prefix-cache: routes to a pod that already has a KV cache matching the request’s prompt prefix; the best prefix-matched pod within a stddev load threshold is selected. Purely about prefix caching — it does not itself gate on cluster-wide load imbalance (the gateway applies that gate centrally ahead of it; see load-balance above). Supports standard mode (local hash table) and KV sync mode (real-time distributed index, enabled via AIBRIX_PREFIX_CACHE_KV_EVENT_SYNC_ENABLED=true).
prefix-cache-preble: routes considering both prefix cache hits and pod load. Based on Preble: Efficient Distributed Prompt Scheduling for LLM Serving.

Fairness

vtc-basic: routes using a hybrid score that balances per-user token fairness and pod utilization. A simplified variant of the Virtual Token Counter (VTC) algorithm. See VTC-artifact for background.

SLO-aware

slo: routes with awareness of per-request service-level objectives.
slo-pack-load: SLO-aware routing that packs load onto fewer pods to improve efficiency.
slo-least-load: SLO-aware routing that spreads load to the least loaded pod.
slo-least-load-pulling: variant of slo-least-load that pulls metrics directly rather than relying on the cached snapshot.

Specialized

pd: prefill-decode disaggregation routing. Splits processing between dedicated prefill pods and decode pods for optimized end-to-end latency.
session-affinity: sticky session routing. Encodes the target pod’s address (IP:Port) as a base64 value in the x-session-id response header. Subsequent requests carrying that header are routed to the same pod. Falls back to a random ready pod and issues a new session ID if the original pod is unavailable.

Auto-blended capacity awareness

Every strategy above, except the exclusive ones (pd, slo/slo-*), an explicit standalone load-balance selection, and a bare session-affinity selection, silently gets load-balance’s capacity-aware scoring blended in behind the scenes — and least-request too, when the selected strategy doesn’t already route by request count, to keep multi-port/data-parallel pod routing working under the blend. The caller never sees this: ctx.Algorithm, response headers, and Validate() all still reflect exactly the strategy that was requested. This keeps any single strategy from steering traffic at an already-hot pod even outside the load-imbalance gate described above. Set AIBRIX_ROUTING_AUTO_BLEND_LOAD_BALANCE_WEIGHT=0 to disable it (see pkg/plugins/gateway/ENV_VARS.md). A bare prefix-cache request uses a 5:4 (1.25:1) lean toward cache affinity over load-balance instead of the flat weight-1 append, and does not receive least-request. session-affinity gets no auto-blend at all: its scoring is binary (the resolved pod vs. everything else), so a smaller load-balance weight could never change the outcome anyway — it keeps running its own Route()/ScoreAll() unblended.

How to Extend Routing Algorithms#

The routing framework is designed to be highly pluggable. All routing logic is expressed through the Router interface:

// Router defines the interface for routing logic to select target pods.
type Router interface {
    // Route selects a target pod from the provided list of pods.
    // The input pods is guaranteed to be non-empty and contain only routable pods.
    Route(ctx *RoutingContext, readyPodList PodList) (string, error)
}

Parameters

ctx *RoutingContext: Per-request context carrying routing inputs. Key fields include:
- Algorithm — the active routing strategy name.
- Model — the model name extracted from the request body.
- Message — the raw prompt text (available for token-level decisions).
- User — the optional user identifier (used by fairness-based algorithms).
- ReqHeaders — a copy of the incoming HTTP request headers.
- ConfigProfile — resolved model config profile (strategy override, RPS limit, etc.), or nil if not set.
readyPodList PodList: Pre-filtered list of pods that are healthy and eligible for routing. Guaranteed non-empty.

Return value

The IP address of the selected pod (e.g. "10.0.0.5:8080"), or a non-nil error if selection fails.

Steps to add a new algorithm

Implement the Router interface in a new *.go file under pkg/plugins/gateway/algorithms/.

Declare a typed constant for the strategy name:

const RouterMyStrategy types.RoutingAlgorithm = "my-strategy"

Register the constructor in an init() function so it is picked up at startup:

func init() {
    Register(RouterMyStrategy, NewMyStrategyRouter)
}

Specify it per-request via the routing-strategy HTTP header, or set ROUTING_ALGORITHM=my-strategy to use it as the default.

Additional router interfaces in pkg/types/router.go support optional capabilities:

QueueRouter — for routers that maintain an internal queue and expose queue length.
FallbackRouter — enables chaining by delegating to a secondary router when the primary cannot make a decision.