Files
llmgw/llmgw-plan.md
T
ilkerandClaude Opus 5 28029eb43b
CI / test (push) Has been cancelled
Switch the project language to English
Comments, doc comments, test names, error and log messages, CLI help, README,
Makefile, CI step names and the design document are now all in English.
No behaviour changes; the full suite passes under the race detector.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-16 01:15:21 +02:00

53 KiB
Raw Blame History

llmgw — A Single-Binary LLM Gateway for Multi-Fabric InfiniBand and Heterogeneous GPUs

Status: Design / Plan v0.2 What changed in v0.2: existing LLM gateway and router projects were surveyed (§2.5) and their critical designs folded in — above all separating admission from scheduling (§7.2), the plugin-based filter → score → pick pipeline (§7.3), and the three-tier model of prefix/KV cache awareness (§7.5). Goal: one Go binary (running under different roles) that serves the LLMs on our own GPU servers behind a single OpenAI-compatible endpoint. Behind it sit several independent InfiniBand fabrics, nodes without IB, and different GPU types.


1. Scope

1.1 In scope

  • OpenAI-compatible API (/v1/chat/completions, /v1/completions, /v1/embeddings, /v1/models), including SSE streaming.
  • An agent on every GPU node: GPU and IB discovery; starting, supervising and stopping inference engine processes (vLLM / SGLang).
  • Control plane: node inventory, model catalog, replica placement, state tracking.
  • Router, in two layers: admission (priority bands, tenant fairness, queueing under saturation) and scheduling (plugin-based filter/score/pick with queue, KV-cache and prefix awareness), plus fallback, retry, cooldown and circuit breaking.
  • Fabric awareness: a multi-node replica never spans two fabrics; a TP group never mixes GPU types.
  • API keys, rate limits, token accounting, Prometheus metrics.

1.2 Out of scope (at least for v1)

  • Inference itself. Go only orchestrates and proxies; vLLM/SGLang run the model.
  • A Kubernetes operator (bare metal / systemd is the target; K8s can come later).
  • Prefill/decode disaggregation (a v2 goal, see §12).
  • External providers (OpenAI, Anthropic) — the architecture stays open to them, but not in v1.

2. Terminology

Term Meaning
Node A physical server running the agent.
Fabric / Island An independent IB fabric. Nodes without IB belong to the none island and each counts as its own single-node island.
GPU Class A group of GPUs with the same model and VRAM (e.g. h100-80g, a100-40g, l40s-48g).
Pool A (fabric, gpu class) pair. The unit of placement.
Model A logical model in the catalog (e.g. llama-3.3-70b). This is the name users see.
Profile How a model runs on a particular GPU class (TP, PP, quantization, max length, engine arguments).
Replica One running copy of a model. May span 1..N nodes. This is the unit the router sees.
Shard A replica's part on a single node (one engine process).
Routing Profile The scorer set and weights the router uses for a class of request (chat / batch / long-context). Not to be confused with the model Profile above.
Band A priority class (platinum/gold/silver/bronze). The unit of admission decisions.
Block The unit of KV cache accounting; a fixed number of tokens (16 in vLLM, 64 by default in the GIE indexer).

2.5 Reference Implementations and What We Took

Rather than design from scratch, we inherited the problems already solved by mature projects in this space.

Project What it is What llmgw takes from it
Gateway API Inference Extension (EPP) The K8s standard; an Endpoint Picker running over Envoy ext_proc Admission/scheduling separation, the plugin-based filter → score → pick pipeline, the model server metric contract (§7.6)
llm-d "Well-lit paths" built on GIE KV-event-based precise prefix index, priority bands with tenant fairness, the saturation detector
AIBrix (a vLLM project) An Envoy ext_proc gateway plus a heterogeneous GPU optimizer The heterogeneous GPU SLO/cost model (§6.2.1), queueing at the gateway (slo-least-load-pulling), disabling prefix affinity under load imbalance
NVIDIA Dynamo Radix-tree KV router plus SLA planner The shape of the cost function (§7.4) and router_temperature softmax sampling against herding
SGLang router Cache-aware load balancing over an approximate radix tree v1's prefix answer: approximate tree plus threshold logic (§7.5 tier 1)
vLLM production-stack / vllm-project/router vLLM's own routers Session affinity (measured 96% vs 93% cache hit), consistent hashing
LiteLLM The most widely deployed multi-provider proxy Virtual key → team/project → budget hierarchy, cooldowns, keeping spend writes off the request path, budget-triggered fallback
Envoy AI Gateway An Envoy-based AI gateway (v1.0, June 2026) Token-based (not merely request-based) rate limiting

What we deliberately did not take: a Kubernetes dependency (bare metal is the target), the Envoy ext_proc architecture (single binary is the goal — we write the router directly in Go), semantic routing / LLM-based model selection (out of scope), and external provider management (not in v1).


3. Architecture

                    ┌───────────────────────────────────────────┐
  Clients ────────► │  llmgw server                             │
  (OpenAI API)      │  ┌────────────┐  ┌──────────────────────┐ │
                    │  │  Router    │  │  Control Plane       │ │
                    │  │ (data)     │◄─┤  - Inventory         │ │
                    │  │            │  │  - Catalog           │ │
                    │  └─────┬──────┘  │  - Scheduler         │ │
                    │        │         │  - Store (SQLite)    │ │
                    │        │         └─────────▲────────────┘ │
                    └────────┼───────────────────┼──────────────┘
                 HTTP (eth)  │                   │ gRPC stream (mTLS, agent dials out)
        ┌────────────────────┼───────────────────┼─────────────────────┐
        ▼                    ▼                   │                     ▼
  ┌───────────┐        ┌───────────┐       ┌───────────┐         ┌───────────┐
  │ llmgw     │  IB-A  │ llmgw     │       │ llmgw     │  IB-B   │ llmgw     │
  │ agent     │◄──────►│ agent     │       │ agent     │◄───────►│ agent     │  ...
  │ vLLM rank0│  NCCL  │ vLLM rank1│       │ SGLang    │         │ (no IB)   │
  └───────────┘        └───────────┘       └───────────┘         └───────────┘
      Fabric A island                        Fabric B island       none island

3.1 One binary, several roles

llmgw server   --config server.yaml     # router + control plane (same process in v1)
llmgw agent    --config agent.yaml      # on every GPU node
llmgw ctl ...                           # CLI: list nodes, deploy models, drain, ...
llmgw sim-engine                        # fake OpenAI-compatible engine for testing

router and control must eventually be able to run as separate processes (llmgw server --roles=router), so the dependency between them goes through an interface (see §6.4).

3.2 Build note

go-nvml requires cgo. Put the agent's NVML code behind a //go:build nvml tag and use a fake / nvidia-smi-based discovery path in untagged builds. That keeps the server side cgo-free and lets development machines without GPUs run the tests.


4. Project Layout

llmgw/
├── cmd/llmgw/main.go              # cobra root; subcommands
├── api/
│   └── proto/llmgw/v1/*.proto     # agent <-> control plane
├── internal/
│   ├── config/                    # YAML + env via koanf
│   ├── agent/
│   │   ├── agent.go               # main loop, control plane connection
│   │   ├── discovery/
│   │   │   ├── gpu_nvml.go        # //go:build nvml
│   │   │   ├── gpu_fallback.go
│   │   │   ├── ib.go              # reads /sys/class/infiniband
│   │   │   └── topology.go        # NVLink, PCIe, NUMA, GPU<->HCA affinity
│   │   ├── supervisor/            # engine process management (exec, restart, logs)
│   │   ├── engines/
│   │   │   ├── engine.go          # Engine interface
│   │   │   ├── vllm.go
│   │   │   └── sglang.go
│   │   └── health/                # polls engine /health and /metrics
│   ├── control/
│   │   ├── inventory/             # node and GPU state
│   │   ├── catalog/               # model and profile definitions
│   │   ├── scheduler/             # placement algorithm
│   │   ├── reconciler/            # desired vs actual loop
│   │   └── agentsvc/              # gRPC server
│   ├── router/
│   │   ├── openai/                # HTTP handlers, request parsing
│   │   ├── proxy/                 # streaming proxy to upstream
│   │   ├── admission/             # saturation detector, band queues (§7.2)
│   │   ├── scheduler/             # Filter/Scorer/Picker pipeline (§7.3)
│   │   ├── prefix/                # tier 0 session hash, tier 1 approx tree, tier 2 kv index
│   │   ├── endpoint/              # live endpoint state, inflight, staleness
│   │   ├── breaker/               # circuit breaker + cooldown
│   │   └── scrape/                # reads engine /metrics (the §7.6 contract)
│   ├── auth/                      # API keys, org/team/project, budgets
│   ├── ratelimit/                 # token bucket (requests + tokens)
│   ├── usage/                     # token counting and recording
│   ├── store/                     # SQLite (modernc.org/sqlite), migrations
│   ├── obs/                       # slog, prometheus, otel
│   └── sim/                       # fake engine
├── deploy/
│   ├── systemd/
│   └── examples/
├── test/e2e/
└── docs/

4.1 Libraries

Need Choice
CLI github.com/spf13/cobra
Config github.com/knadh/koanf/v2
RPC connectrpc.com/connect (gRPC compatible, HTTP/2, easy to debug)
GPU github.com/NVIDIA/go-nvml
DB modernc.org/sqlite (cgo-free) + pressly/goose migrations
Metrics github.com/prometheus/client_golang, parsing via github.com/prometheus/common/expfmt
Tracing go.opentelemetry.io/otel
Logging log/slog (stdlib)
Rate limiting golang.org/x/time/rate (plus our own bucket for token limits)
Hashing github.com/cespare/xxhash/v2
HA (v2) github.com/hashicorp/raft or etcd
Testing stdlib testing, github.com/stretchr/testify (optional)

5. Agent

5.1 GPU discovery

For every GPU:

  • index, UUID, name, VRAM, compute capability (for FP8 and similar support), driver/CUDA version
  • NUMA node, PCIe bus id
  • NVLink peers (NVML GetNvLinkRemotePciInfo) → which GPUs pair efficiently for TP within a node
  • health: XID errors, ECC, temperature, memory in use (are there non-engine processes?)

GPUClass is derived automatically: normalize(name) + "-" + round(vram_gb) + "g" → h100-80g. It can be overridden in configuration.

5.2 IB discovery and fabric identity

From /sys/class/infiniband/<hca>/ports/<n>/:

  • state (ACTIVE?), phys_state, rate, link_layer (InfiniBand / Ethernet → RoCE)
  • gids/0 → subnet prefix
  • sm_lid, lid
  • the HCA's PCIe address and NUMA node → GPU affinity

How is fabric identity determined? Automatic detection is not reliable, because OpenSM's default subnet prefix (fe80::) is the same on every fabric. In order of preference:

  1. An explicit label in the agent config: fabrics: [{name: ib-a, hcas: [mlx5_0, mlx5_1]}] (recommended)
  2. If every OpenSM has a unique subnet_prefix, automatic mapping from the GID prefix (a prefix → fabric table in the server config)
  3. Otherwise the node is marked fabric=unknown and excluded from multi-node placement.

Bonus verification (v1.1): for two nodes claiming the same fabric, the agents verify reachability with a short ibv_rc_pingpong-style test or ibping. A mismatch raises an alert.

A node may have HCAs attached to more than one fabric, and the model must allow it (Node.Fabrics []FabricAttachment). A replica, however, picks exactly one fabric and uses only that fabric's HCAs.

5.3 Engine interface

type Engine interface {
    Name() string                                   // "vllm", "sglang"
    Validate(p Profile, gpus []GPU) error           // compute capability, quantization fit, ...
    BuildCommand(spec ShardSpec) (*exec.Cmd, error) // builds arguments and environment
    HealthURL(spec ShardSpec) string
    MetricsURL(spec ShardSpec) string
    ParseMetrics(r io.Reader) (EngineStats, error)  // engine-specific metric names
}

type ShardSpec struct {
    ReplicaID   string
    Rank        int      // 0 = leader (the one serving HTTP)
    NNodes      int
    LeaderAddr  string   // dist init address (ethernet IP)
    LeaderPort  int
    GPUIndices  []int
    HCAs        []string // HCAs of the chosen fabric
    HTTPPort    int
    Profile     Profile
    ModelPath   string
}

type EngineStats struct {
    // The mandatory contract from §7.6 — without these the router cannot route.
    Running, Waiting int
    KVCacheUsage     float64 // 0..1   (vllm:kv_cache_usage_perc)
    // Optional.
    BlockSize, NumGPUBlocks int
    PrefixHitRate          float64
    TTFTp50, TPOTp50       time.Duration // from histograms
    DrainRate              float64       // observed completions/sec; for the `load` scorer
    LoRA                   LoRAState     // v2
    ScrapedAt              time.Time     // for staleness detection — the router uses this
}

Supporting a new engine is defined as: being able to fill in the mandatory three from §7.6.

Multi-node launch:

  • SGLang supports it natively: --nnodes N --node-rank R --dist-init-addr host:port --tp T.
  • vLLM: depending on the version, either Ray (ray start --head on the leader, ray start --address on the workers) or the newer multi-node multiprocessing flags. Check the documentation for the vLLM version you deploy; vllm.go should offer both paths behind a launchMode.

Environment generation (critical):

CUDA_VISIBLE_DEVICES=<GPUIndices>
NCCL_IB_HCA=<HCAs, comma separated>    # the chosen fabric only
NCCL_SOCKET_IFNAME=<mgmt interface>     # bootstrap
NCCL_IB_DISABLE=1                       # single node, or no IB
NCCL_CROSS_NIC / NCCL_NET_GDR_LEVEL     # optional, from the profile
GLOO_SOCKET_IFNAME=<mgmt interface>

5.4 Supervisor

  • Start processes in their own process group (SysProcAttr.Setpgid) and kill the whole group when stopping (SIGTERM → timeout → SIGKILL).
  • stdout/stderr → line-oriented slog plus a ring buffer (the last N lines, for llmgw ctl logs).
  • State machine: Pending → Starting → Loading → Ready → Draining → Stopped | Failed.
  • Ready only once /health returns 200 and /v1/models lists the model.
  • Model loading takes a long time; startupTimeout comes from the profile (e.g. 20 min).
  • On a crash: report to the control plane and let the control plane decide about restarting. Restarting a single shard of a multi-node replica accomplishes nothing; the whole replica must be rebuilt.
  • If the agent restarts, try to adopt still-running engine processes (pid file plus cmdline check) and clean up whatever cannot be adopted.

5.5 Model weights

  • v1: a shared filesystem (NFS/Lustre/Weka) or a path pre-populated on the node. The agent only checks existence and size.
  • v1.1: the agent downloads from HF/S3 into a local cache, reports progress, and enforces a disk quota.

5.6 Agent ↔ control plane protocol

The agent dials out (NAT- and firewall-friendly) over a bidirectional stream:

service AgentService {
  rpc Connect(stream AgentMessage) returns (stream ControlMessage);
}

message AgentMessage {
  oneof msg {
    Register     register  = 1; // node info, GPUs, fabrics, versions
    Heartbeat    heartbeat = 2; // every 5s: GPU health, shard states
    ShardStatus  shard     = 3; // sent on every state change
    EngineStats  stats     = 4; // every 2s (the router may also scrape directly)
  }
}

message ControlMessage {
  oneof msg {
    StartShard   start = 1;
    StopShard    stop  = 2;
    DrainShard   drain = 3;
    Ack          ack   = 4;
  }
}
  • Commands are idempotent: a repeated StartShard with the same shard_id is a no-op.
  • mTLS, with node identity taken from the certificate CN. A simple internal CA in v1; llmgw ctl cert issue.
  • Three missed heartbeats mark the node Unreachable; after 60s it becomes Lost and its replicas are rescheduled.

6. Control Plane

6.1 Data model (abridged)

type Node struct {
    ID, Hostname, MgmtAddr string
    Labels     map[string]string
    GPUs       []GPU
    Fabrics    []FabricAttachment // empty means no IB
    State      NodeState          // Ready, Cordoned, Draining, Unreachable, Lost
    LastSeen   time.Time
}

type FabricAttachment struct {
    Fabric string   // "ib-a"
    HCAs   []string
    Rate   int      // Gbps
}

type Model struct {
    Name      string            // what users see
    Aliases   []string
    Source    string            // path or HF id
    Profiles  []Profile
    Scaling   Scaling           // min/max replicas, target queue, ...
    Fallbacks []string          // other model names
}

type Profile struct {
    Name        string
    GPUClass    string          // "h100-80g"
    Engine      string          // "vllm"
    TP, PP      int
    GPUsPerNode int             // usually TP
    Nodes       int             // PP*... ; a fabric is mandatory when > 1
    Args        []string        // --quantization fp8 --max-model-len 32768 ...
    Env         map[string]string
    Weight      int             // capacity weight for the router (from benchmarks)
    Priority    int             // placement preference
    StartupTimeout time.Duration
}

type Replica struct {
    ID, Model, Profile string
    Fabric   string             // set for multi-node replicas
    Shards   []Shard
    State    ReplicaState
    Endpoint string             // rank0's HTTP address
}

6.2 Scheduler (placement)

Input: the desired replica count, the profile list, and the inventory.

Hard constraints:

  1. All shards share one GPUClass.
  2. When Nodes > 1, every node must attach to the same Fabric and have ACTIVE HCAs on it.
  3. GPUsPerNode free GPUs, preferably in the same NVLink domain and on the same NUMA node as the HCA.
  4. The node is Ready and not Cordoned.
  5. The engine's Validate must pass (e.g. an FP8 profile never lands on a card with compute capability < 8.9).

Soft preferences (scored):

  • the profile's Priority
  • bin packing: fill partially used nodes first, keeping empty nodes free for large models
  • spread a model's replicas across fabrics and nodes (failure domain distribution)
  • GPU-to-HCA NUMA affinity

Algorithm (v1, simple and deterministic):

for each profile in model.Profiles sorted by Priority desc:
    pools := groupNodes(by GPUClass == profile.GPUClass, by Fabric if profile.Nodes > 1)
    for each pool sorted by score:
        nodes := pickNodes(pool, profile.Nodes, profile.GPUsPerNode)  // greedy + NVLink group selection
        if nodes != nil:
            return Placement{profile, pool.Fabric, nodes, gpuAssignments}
return ErrUnschedulable (with reasons — essential for debugging)

groupNodes: nodes without IB join the pool for Nodes == 1 profiles, and never for Nodes > 1.

llmgw ctl explain-placement <model> must show why each pool was rejected.

6.2.1 Heterogeneous GPUs: the SLO/capacity model

Comparing GPU classes through a single scalar "weight" is too crude; the model taken from AIBrix's heterogeneous GPU optimizer is more accurate. For each (model, profile) pair, ctl bench (§10) produces a capacity profile:

type CapacityProfile struct {
    Model, Profile string
    GPUClass       string
    // Maximum concurrency that still meets the SLO; the real capacity measure.
    MaxConcurrencyAtSLO int
    ThroughputTokPerSec float64
    TTFTp99, TPOTp99    time.Duration
    RelativeCost        float64   // relative GPU-hour cost, from configuration
    MeasuredAt          time.Time
    EngineVersion       string    // the profile expires when the version changes
}

An example SLO: TPOT p99 ≤ 50ms && TTFT p99 ≤ 2s.

The profile is used in three places:

  1. Router weight: Profile.Weight is no longer hand-written; it derives from MaxConcurrencyAtSLO.
  2. Admission thresholds: TargetWaitingPerReplica (§7.2.1) is suggested from the measurement.
  3. Placement and autoscaling objective: "minimize total cost while meeting the SLO". For a predicted load, the replica count needed per GPU class is ceil(load / MaxConcurrencyAtSLO), and the cheapest combination wins. In v1 this only produces advice via ctl plan <model> --target-rps N; automatic application arrives with the v1.1 autoscaler.

When EngineVersion or the profile arguments change the measurement goes stale. llmgw must surface that (llmgw_capacity_profile_stale{model,profile}) rather than silently reusing an obsolete number.

6.3 Reconciler

  • Every 5s, or on an event: compare desired (model scaling) with actual (replica states).
  • Missing replica → schedule → send StartShard for each shard (workers first, then rank0, or whatever order the engine requires).
  • Any shard in a replica Failed → stop the whole replica and reschedule with backoff.
  • Excess replicas → Drain first (the router stops sending new requests, existing ones finish), then Stop.
  • Node Lost → treat its replicas as Failed.

Autoscaling (v1.1): a per-model targetWaitingPerReplica, averaged from router metrics, with cooldown-guarded scaling. Shrinking lower-priority models when no GPUs are free (preemption) is v2.

6.4 Interface to the router

type EndpointSource interface {
    // model name -> Ready replicas (+ pool, weight, fabric, gpuclass, capacity profile)
    Endpoints(model string) []Endpoint
    Subscribe() <-chan struct{} // change notification
}

In-memory in v1; once the router and control plane split apart, the same interface is backed by a watch stream.

The router always caches this list: even if the control plane goes down it keeps routing to the replicas it knows about, which mitigates the single-point-of-failure risk in §14. When the cache ages past endpoint_cache_max_age (say 5 min) it only raises a warning metric; it never stops traffic.

6.5 Store

  • SQLite in WAL mode. Tables: nodes, models, profiles, replicas, shards, api_keys, usage_rollup, events.
  • Inventory and replica state can be rebuilt from the agents; what genuinely needs to persist is the catalog, the keys and the usage records.
  • HA (v2): Raft-replicated state, or an external Postgres/etcd.

7. Router

7.1 Request lifecycle

Request
 ├─ 1. Auth       : key → project/team, permitted models, priority band, routing profile
 ├─ 2. Quota      : RPM/TPM and budget (in-memory counters; no DB on the request path)
 ├─ 3. Parse      : model, stream, messages/prompt, max_tokens, user.
 │                  The rest of the body passes through as untouched JSON.
 ├─ 4. Resolve    : model/alias → candidate pools (fabric × gpu_class)
 ├─ 5. ADMISSION  : "now?"      §7.2   ← waits here while the pool is saturated
 ├─ 6. SCHEDULING : "where?"    §7.3   filter → score → pick
 ├─ 7. Proxy      : stream to upstream, propagate cancellation           §7.8
 └─ 8. After the request (async): usage, spend, metrics, logs

The critical separation (taken from GIE / llm-d). When a request proceeds and where it goes are two distinct decisions and must stay distinct. Kept apart, the router can hold requests in its own central queue and apply priority and tenant fairness before they enter an engine queue.

Why this matters: an engine's own queue is FIFO and has no notion of tenants or priority. Once a request is in there, no prioritization is possible. When a batch job fills the vLLM queue with 500 requests, the TTFT of an interactive chat request arriving behind them collapses and there is nothing the router can do. v0.1 had no such layer; for a shared cluster this was the most serious omission.

7.2 Admission control (flow control)

7.2.1 Saturation detector

One saturation signal is produced per pool, combining two metrics (llm-d's approach — queue depth alone is not enough, KV pressure is a genuine exhaustion signal too):

// per replica, then aggregated across the pool as p90 or mean
func saturation(e EngineStats, t Thresholds) float64 {
    q  := float64(e.Waiting) / t.TargetWaitingPerReplica   // e.g. target 4
    kv := e.KVCacheUsage / t.KVCeiling                      // e.g. ceiling 0.90
    return math.Max(q, kv)
}

The goal is to keep a small but healthy buffer on the engine: enough requests for continuous batching to stay efficient, but not enough for the queue to balloon. Set TargetWaitingPerReplica too low and the GPU starves; too high and queueing delay returns. Both thresholds must be tunable per model/profile, and ctl bench (§10) should propose them.

7.2.2 Priority bands

Every API key maps to a band, and the band sets the saturation ceiling:

bands:
  - { name: platinum, priority: 400, saturation_ceiling: 1.00 }  # never held back
  - { name: gold,     priority: 300, saturation_ceiling: 0.85 }
  - { name: silver,   priority: 200, saturation_ceiling: 0.60 }
  - { name: bronze,   priority: 100, saturation_ceiling: 0.30 }  # batch jobs

The dispatch loop scans bands from highest to lowest. Once pool saturation exceeds a band's ceiling, that band's requests wait in the queue. The result: as the cluster fills, batch traffic slows down first and interactive traffic is protected. (In llm-d's measurements platinum, gold and silver stayed under a second while bronze batch fell far behind — exactly the desired behaviour.)

7.2.3 Fairness and queue management

  • Fairness within a band: when several tenants share a band, dispatch turns are distributed round-robin. One tenant's traffic spike cannot consume the whole band.
  • Capacity: each band's queue is bounded. When it is full, new requests get an immediate 503 — far better than accumulating silently.
  • TTL: a request that waits longer than queue_ttl is dropped (504). There is no point sending the GPU work the client has already given up on.
  • Client disconnect: if the connection drops while queued, the request is discarded quietly (context cancellation).
  • Displacement (v1.1): to make room for a higher band, drop the newest requests from the lowest band — the oldest ones are already near their TTL and have invested the most waiting.
type Admission interface {
    // Blocks until the request may proceed; errors on cancellation or TTL.
    Acquire(ctx context.Context, req *RequestInfo) (Lease, error)
}

Lease.Release() is called once the response completes; it also maintains the router's inflight counter.

7.3 Scheduler: filter → score → pick

The list of "strategies" from v0.1 is replaced by a plugin pipeline (the GIE EPP model). No strategy is a separate monolith; scorers are combined by weight.

type Filter interface { // drops candidates (hard constraints)
    Name() string
    Filter(ctx *SchedCtx, eps []Endpoint) []Endpoint
}

type Scorer interface { // 0..1 score; 1 is best
    Name() string
    Score(ctx *SchedCtx, eps []Endpoint) map[EndpointID]float64
}

type Picker interface { // chooses from the scores
    Pick(ctx *SchedCtx, scores map[EndpointID]float64) Endpoint
}

type RoutingProfile struct {
    Name    string
    Filters []Filter
    Scorers []WeightedScorer   // {Scorer, Weight}
    Picker  Picker
}

Total score: score(ep) = Σ wᵢ · scoreᵢ(ep), with each scorer normalizing to 0..1 internally.

A two-stage decision. Because the topology is heterogeneous, the pool is chosen first and the replica second:

  1. Pool selection: the pool with the highest Priority; if every replica there is saturated (§7.2.1), spill over to the next pool. This prevents falling back to the A100 pool before the H100 pool is actually full.
  2. Replica selection: filter/score/pick within the chosen pool.

Routing profiles, from configuration (a per-model default, overridable per request via the x-llmgw-profile header):

Profile Weights Use
chat prefix 3.0, queue 1.0, kv 1.0 Multi-turn chat and agents — high cache reuse
batch queue 3.0, kv 1.0, prefix 0.5 Bulk work — throughput matters, TTFT does not
long-context kv 3.0, prefix 2.0, queue 1.0 Long prompts — KV pressure is the real bottleneck
debug round-robin picker Troubleshooting

7.4 Scorers and their formulas

Scorer Formula Notes
queue 1 − min(1, waiting / target_waiting) Queue depth
kv 1 − kv_cache_usage KV occupancy; a hard penalty above the ceiling
prefix matched_blocks / prompt_blocks §7.5
load 1 − norm(running / drain_rate) AIBrix: drain_rate is the observed completion rate. running alone cannot spot a slow replica.
inflight 1 − norm(router_inflight) The router's own count; compensates for scrape lag. At a 2s scrape interval, herding is unavoidable without it.
pool-priority constant, from the profile's Priority Pool preference
session session hash match §7.5 tier 0
lora (v2) is the adapter loaded vllm:lora_requests_info

An alternative cost formulation (Dynamo). Modelling prefill and decode cost directly in blocks may be more accurate than normalized scores:

cost(ep) = w_overlap · (prompt_blocks − cached_blocks(ep)) + decode_blocks(ep)
choice   = argmin cost

A high w_overlap favours cache reuse (better TTFT); a low one evens out load (better ITL). Both can be expressed as a Scorer, so the pipeline survives either way: weighted scores in v1, with this cost form added as a scorer in v1.1 and compared using ctl bench.

Pickers:

  • max-score: the highest score. Deterministic but prone to herding.
  • p2c (power of two choices): two random candidates, the better one. The default once there are many replicas.
  • softmax(temperature): Dynamo's router_temperature. Treats scores as logits and samples. temperature=0 degenerates to argmax. A small temperature breaks the oscillation that synchronized decisions cause among equally scored replicas.

7.5 Prefix / KV cache awareness — three tiers

This is the single largest win the router can deliver, but the cost gap between the tiers is enormous. They stay separate so they can be enabled per model.

Tier 0 — session affinity (v1, cheap, high return)

The client-supplied x-session-id header (falling back to OpenAI's user field) maps to a replica via bounded-load consistent hashing. If that replica is overloaded — beyond load_factor times the average — the next one on the ring is used. Roughly 100 lines of code; production-stack measured cache hits going from about 93% to 96% versus round robin. For multi-turn chat and agent workloads this is the best cost/benefit ratio available.

Tier 1 — approximate radix tree (v1.1, moderate cost)

The router maintains a tree that estimates each worker's cache from the request prefixes it has seen (the SGLang router's approach). It does not know the true cache state; it updates lazily and costs little.

The thresholds — the part that matters in real implementations:

  • cache_threshold: below this match ratio, treat it as "no cache" and let load balancing decide.
  • balance_abs_threshold and balance_rel_threshold: if max_load − min_load > abs and max_load / min_load > rel, ignore the prefix entirely and send to the least loaded replica. (AIBrix adds the same guard.) Without it, prefix affinity drowns one replica: the cache hit rate rises while p99 latency collapses.
  • eviction_interval plus max_tree_size: keep the tree from growing without bound.

Which branch a decision took must be recorded as a metric: cache_hit | cache_miss_min_load | load_balance_override | stale | no_tree | no_healthy_worker. Without that, tuning prefix routing is guesswork.

Tier 2 — precise KV index (v2, expensive, optional)

vLLM publishes block events over ZMQ with --kv-events-config (BlockStored, BlockRemoved). The router subscribes and maintains a block_hash → {replica} map, giving it a true global view.

The trap — a cost the plan must acknowledge: block hashes must be computed identically to vLLM's token-id-based scheme. That forces the router to tokenize requests with the same tokenizer as the engine. On the Go side that means loading a tokenizer per model (HF tokenizers via cgo, or a pure-Go implementation), tracking version compatibility, and paying the tokenization CPU cost per request. This is why AIBrix requires a "remote tokenizer". Tiers 0 and 1 need no tokenizer — hence they are v1 and this is v2.

Event loss: if a gap appears in the event sequence numbers, request the missing range from the replay endpoint; if that fails, drop the index for that replica and fall back to tier 1 for it. Returning to approximate information beats silently routing on wrong data. If a replica restarts, its index is reset.

7.6 The engine metric contract

This is the formal definition of "adding a new engine". If an engine exposes these three metrics in Prometheus format, llmgw can route to it (GIE's model server protocol):

Contract field vLLM metric Mandatory
Waiting vllm:num_requests_waiting yes
Running vllm:num_requests_running yes
KVCacheUsage vllm:kv_cache_usage_perc yes
BlockSize, NumGPUBlocks vllm:cache_config_info no
LoRA state vllm:lora_requests_info (max_lora, running_lora_adapters, waiting_lora_adapters) no

Correction: v0.1 said vllm:gpu_cache_usage_perc; the correct name is vllm:kv_cache_usage_perc. Metric names drift between versions, so the ParseMetrics mapping must be overridable from configuration, and the agent must report the engine version when it registers.

Scrape rules:

  • Every Ready endpoint every 1–2s, one goroutine per endpoint, with jitter.
  • A failed scrape marks the endpoint stale; its score drops but it is not removed.
  • Three consecutive failures trip the breaker (§7.7).
  • So that decisions are never made purely on stale data, the inflight scorer (§7.4) stays enabled at all times.

7.7 Reliability

  • Retry: only before the first byte has been written upstream, only on connection errors, 503 or 429, always against a different endpoint, at most N times. No retries once the stream has started.
  • Cooldown (the LiteLLM lesson): an endpoint that errors leaves the candidate pool for cooldown_seconds. A long fallback list plus retries without cooldown equals a retry storm: one endpoint's rate limit spreads across the whole cluster.
  • Circuit breaker: a closed/open/half-open state machine per endpoint, with exponential backoff.
  • Fallback: when a model has no usable endpoints, fall back to a model in its Fallbacks list (if the key permits it). The response carries an x-llmgw-served-model header. Budget fallback: when a key exceeds its budget, optionally downgrade to a cheaper model rather than rejecting the request (§7.9).
  • Timeouts: separate limits for connection, first token (TTFT) and total, plus an inter-token idle timeout while streaming.
  • Cancellation: when the client disconnects, cancel the upstream request too (context propagation) so the GPU stops working. This applies to queued requests as well.
  • Draining: a Draining replica receives no new requests; active ones run to completion or time out.
  • Control plane loss: the router keeps serving from its last known endpoint list (§14).

7.8 Streaming proxy details

  • http.Transport with a high MaxIdleConnsPerHost and HTTP/1.1 keep-alive (engines rarely speak h2c).
  • Read lines with bufio.Reader, pass data: lines through, and call http.Flusher.Flush() after every event.
  • If the client did not send stream_options.include_usage, the router may add it (configurable) to capture usage from the final chunk, and then withhold that chunk from a client that did not ask for it.
  • Normalize upstream error JSON into OpenAI's format.
  • Response headers: x-llmgw-replica, x-llmgw-pool, x-llmgw-served-model, x-llmgw-queue-ms (time spent in admission) and x-llmgw-cache-decision (§7.5).

7.9 Identity, quota and usage (tenancy)

LiteLLM's production-proven model:

Organization ─► Team ─► Project ─► API Key

Each level carries a budget (in tokens or notional cost units), RPM/TPM limits and a permitted model list. A key additionally carries its priority band (§7.2.2) and its default routing profile (§7.3).

  • API keys look like sk-llmgw-...; only the SHA-256 hash is stored. A leaked key is rotated without touching upstream credentials.
  • Rate limiting: RPM via x/time/rate plus TPM. The token limit is reserved at request start as max_tokens + estimated prompt and reconciled against the real figure when the response completes. (The prompt estimate is bytes/4 in v1; it becomes exact once the tier 2 tokenizer arrives.)
  • Budgets: periodic (daily/monthly) spend. On exhaustion: 429, or a budget fallback. Zero-cost models are exempt from budget checks, so a key whose budget is spent can still be given a "cheap model" fallback.
  • Spend writes stay off the request path: usage records go to a channel and are batched into SQLite, with hourly rollups. Nothing writes to the database on the request path.
  • Replica pinning (debug): the x-llmgw-replica header, scoped to the key. If the pinned replica is in cooldown or gone, the request falls through to normal routing.
  • v1 is a single instance with in-memory limiters. With multiple instances this becomes Redis or consistent-hash stickiness (v2).

8. Observability

Prometheus metrics (examples):

llmgw_requests_total{model,pool,status}
llmgw_request_duration_seconds{model,pool}
llmgw_ttft_seconds{model,pool}
llmgw_tokens_total{model,direction="prompt|completion"}
llmgw_endpoint_inflight{replica}
llmgw_endpoint_score{replica}
llmgw_replicas{model,state}
llmgw_gpu_free{node,gpu_class,fabric}
llmgw_schedule_failures_total{model,reason}
llmgw_agent_connected{node}

# admission / flow control (§7.2)
llmgw_admission_queue_depth{band,model}
llmgw_admission_wait_seconds{band,model}          # histogram
llmgw_admission_rejected_total{band,reason="full|ttl|client_gone"}
llmgw_pool_saturation{pool}                       # 0..1, the detector's output

# routing decisions (§7.5) — prefix routing cannot be tuned without these
llmgw_cache_decision_total{model,decision="hit|miss_min_load|balance_override|stale|no_tree"}
llmgw_prefix_match_ratio{model}                   # histogram
llmgw_scorer_value{scorer,replica}
llmgw_endpoint_metrics_stale{replica}
llmgw_capacity_profile_stale{model,profile}
  • OpenTelemetry traces: router span → upstream span (the traceparent header is forwarded to the engine).
  • Structured logs (slog JSON) with request_id everywhere.
  • llmgw ctl top: live queue depth and occupancy per pool.

9. Configuration Examples

9.1 server.yaml

listen: ":8080"
admin_listen: ":9090"          # metrics, admin API, pprof
agent_listen: ":7443"
tls:
  ca: /etc/llmgw/ca.pem
  cert: /etc/llmgw/server.pem
  key: /etc/llmgw/server-key.pem
store:
  path: /var/lib/llmgw/llmgw.db

fabrics:
  - name: ib-a
    gid_prefix: "fec0:0000:0000:000a"   # optional automatic mapping
  - name: ib-b
    gid_prefix: "fec0:0000:0000:000b"

router:
  scrape_interval: 2s
  inject_usage: true
  retries: 2
  cooldown_seconds: 15            # how long a failing endpoint leaves the candidate pool
  endpoint_cache_max_age: 5m      # keep serving this long if the control plane goes down

  # §7.2 admission — the saturation detector
  admission:
    target_waiting_per_replica: 4   # proposed by ctl bench
    kv_ceiling: 0.90
    queue_ttl: 30s
    queue_capacity_per_band: 2000
    bands:
      - { name: platinum, priority: 400, saturation_ceiling: 1.00 }
      - { name: gold,     priority: 300, saturation_ceiling: 0.85 }
      - { name: silver,   priority: 200, saturation_ceiling: 0.60 }
      - { name: bronze,   priority: 100, saturation_ceiling: 0.30 }

  # §7.3 routing profiles
  default_profile: chat
  profiles:
    chat:
      scorers: { prefix: 3.0, queue: 1.0, kv: 1.0, inflight: 1.0 }
      picker:  { type: p2c }
    batch:
      scorers: { queue: 3.0, kv: 1.0, load: 1.0, prefix: 0.5 }
      picker:  { type: max-score }
    long-context:
      scorers: { kv: 3.0, prefix: 2.0, queue: 1.0 }
      picker:  { type: softmax, temperature: 0.2 }

  # §7.5 prefix awareness
  prefix:
    session_affinity: { enabled: true, header: x-session-id, load_factor: 1.25 }  # tier 0
    approx_tree:                                                                  # tier 1
      enabled: false          # v1.1
      cache_threshold: 0.30
      balance_abs_threshold: 16
      balance_rel_threshold: 1.5
      max_tree_size: 2000000
      eviction_interval: 60s
    kv_events: { enabled: false }   # tier 2 — needs a tokenizer, v2

models:
  - name: llama-3.3-70b
    source: /models/Llama-3.3-70B-Instruct
    scaling: { min: 2, max: 6 }
    fallbacks: [qwen2.5-32b]
    routing_profile: chat
    slo: { ttft_p99: 2s, tpot_p99: 50ms }   # ctl bench measures capacity against this SLO
    profiles:
      - name: h100-tp4-fp8
        gpu_class: h100-80g
        engine: vllm
        tp: 4
        nodes: 1
        priority: 100
        weight: 100
        args: ["--quantization", "fp8", "--max-model-len", "32768"]
      - name: a100-tp8
        gpu_class: a100-80g
        engine: vllm
        tp: 8
        nodes: 1
        priority: 50
        weight: 45

  - name: deepseek-large
    source: /models/DeepSeek-V3
    scaling: { min: 1, max: 2 }
    profiles:
      - name: h100-2node
        gpu_class: h100-80g
        engine: sglang
        tp: 16
        nodes: 2           # → must live on a single fabric
        gpus_per_node: 8
        priority: 100
        weight: 100
        startup_timeout: 40m

9.2 agent.yaml

node_id: gpu-a-01
server: "llmgw-server.internal:7443"
mgmt_interface: eth0
tls: { ca: ..., cert: ..., key: ... }
labels: { rack: r12 }
fabrics:
  - name: ib-a
    hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3]
port_range: [20000, 20999]        # engine HTTP and dist-init ports
engines:
  vllm:   { bin: /opt/venvs/vllm/bin/vllm, launch_mode: ray }
  sglang: { python: /opt/venvs/sglang/bin/python }
work_dir: /var/lib/llmgw-agent

10. Benchmarking and Weighting

llmgw ctl bench <model> --profile <p>:

  • Drives a target replica at increasing concurrency with fixed prompt and output lengths.
  • Measures TTFT p50/p95, TPOT, throughput (tok/s) and the maximum concurrency that still meets the SLO.
  • The output is the CapacityProfile from §6.2.1: MaxConcurrencyAtSLO, throughput, TTFT/TPOT p99.
  • --apply updates three things at once: the profile's weight, target_waiting_per_replica, and the autoscaling target.
  • This is the only sound way to compare GPU classes; never use an estimated weight.
  • --sweep-prefix mode measures behaviour at 0/50/90% prefix cache hit rates, so routing profile weights (§7.3) can be tuned from measurements instead of guesses.
  • A profile goes stale when the engine version changes, so every ctl bench result is tagged with the engine version.

11. Roadmap

Phase 0 — Skeleton (1 week)

  • go mod init, cobra commands, koanf config, slog, Makefile, golangci-lint
  • sim-engine: OpenAI-compatible fake engine (tunable latency, token rate, queue, vLLM-format /metrics)
  • CI: go test ./..., race detector

Phase 1 — Static router (2 weeks)

  • OpenAI handlers, streaming proxy, cancellation propagation
  • Static endpoint list from configuration
  • The Filter/Scorer/Picker pipeline (§7.3) — a plugin pipeline from the start, not one "strategy"
  • queue, kv, inflight scorers; max-score and p2c pickers
  • Session affinity (§7.5 tier 0) — cheap, an early win
  • Metrics scraping (the §7.6 contract), staleness detection, retry, cooldown, breaker
  • Prometheus metrics plus decision metrics (§8)
  • Exit criterion: 1000 concurrent streams behind 3 sim-engines; upstream cancellation on client disconnect verified; router overhead p99 < 2 ms.

Phase 1.5 — Admission control (1 week) ⭐

  • Saturation detector (§7.2.1)
  • Priority bands, round-robin tenant fairness within a band, queue capacity and TTL
  • Admission interface and lease, cancellation while queued
  • llmgw_admission_* metrics
  • Exit criterion: while 500 concurrent batch requests flow through the bronze band, the TTFT p99 of interactive requests in the platinum band stays under 1.5× its idle value. Do not move on before writing this test — it is the entire reason phase 1.5 exists.

Phase 2 — Agent: discovery and single-node engines (2 weeks)

  • GPU discovery (NVML plus fallback), IB discovery, topology
  • Engine interface, single-node vLLM and SGLang
  • Supervisor (process groups, log buffer, state machine, adoption)
  • connect-go stream, register/heartbeat, mTLS
  • Exit criterion: start and stop a model on a real node via ctl.

Phase 3 — Control plane (2–3 weeks)

  • SQLite store and migrations
  • Catalog, inventory, scheduler (single-node profiles), reconciler
  • The router's EndpointSource becomes dynamic
  • ctl: nodes, models, replicas, cordon, drain, explain-placement, logs
  • Exit criterion: shutting down a node moves its replicas elsewhere while the router keeps serving.

Phase 4 — Multi-node and fabrics (2 weeks)

  • Fabric identity (config plus GID prefix), verification
  • Multi-node placement (the fabric constraint), port allocation, rank-ordered startup
  • NCCL environment generation, HCA selection
  • Replica-level failure handling (one dead shard kills the replica)
  • Exit criterion: a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it to the other one.

Phase 5 — Production features (2–3 weeks)

  • Tenancy: org/team/project/key hierarchy, budgets, RPM/TPM, usage records (§7.9)
  • Budget fallback, model allowlists, key→band mapping
  • Routing profiles from configuration plus the x-llmgw-profile header
  • Approximate radix tree (§7.5 tier 1), with its thresholds and decision metrics
  • Fallback models and pool spillover
  • ctl bench → CapacityProfile (§6.2.1), ctl plan
  • OTel tracing
  • systemd units, documentation

Phase 6 — v2 ideas

  • Precise KV index (§7.5 tier 2) — vLLM KV events plus tokenizer integration
  • Dynamo-style block cost function, tuning the softmax picker
  • Admission displacement (dropping requests from lower bands)
  • Autoscaling and priority-based preemption
  • Control plane HA (Raft)
  • Weight downloading and cache management
  • Prefill/decode disaggregation (§12)
  • LoRA adapter routing (affinity to a replica that already has the adapter)
  • External provider support (as a fallback)
  • Admin web UI

12. Prefill/Decode Disaggregation (v2 notes)

  • Prefill and decode replicas must live on the same fabric (KV transfer goes over RDMA).
  • A two-stage flow in the router: send to a prefill replica → KV handle → decode replica.
  • Engine support varies sharply by version (vLLM KV connectors, SGLang PD mode, NVIDIA Dynamo/NIXL). Wrapping the engine's own built-in PD router as a single "endpoint" is less risky than implementing this ourselves.
  • What llm-d and Dynamo show: PD's benefit lies mainly in long context. For short-prompt chat traffic the KV transfer cost eats the gain. Measure your workload's ISL/OSL distribution with the benchmark in §10 first, and enable PD only for long-context models.
  • The scheduler gains a Role: prefill|decode field and a "paired placement on the same fabric" constraint.

13. Test Strategy

  • Unit: the scheduler (table-driven: heterogeneous inventory plus expected placement or failure reason), balancer scores, SSE parsing, metric parsing (fixtures taken from real vLLM/SGLang /metrics output).
  • Property-based: the invariants "a multi-node replica lives on one fabric" and "one GPU class per replica" always hold in scheduler output (testing/quick or rapid).
  • Integration: server plus N fake agents (fake GPU/IB inventory) plus sim-engines, in one process.
  • Chaos: kill agents, break the network (toxiproxy), crash engines, slow engines, and kill the control plane to verify the router keeps serving from its cache.
  • Priority isolation: phase 1.5's exit criterion as a permanent regression test — high-priority TTFT is preserved under low-priority load. This test must stay in CI.
  • Routing quality: give sim-engine a realistic prefix cache model (block-level LRU) and measure the hit rate against synthetic multi-turn traffic. The comparison between round robin, session affinity and the approximate tree must be numeric.
  • Load: k6 or our own Go load generator; measure the router's own overhead (target: p99 under 2 ms of added latency).
  • Real hardware: at the end of every phase, a smoke test on at least one node with IB and one without.

14. Risks and Open Questions

Topic Note
Engine version differences CLI flags and metric names change often → make engine adapters version-aware and overridable from configuration.
vLLM multi-node launch Ray or built-in multiprocessing? Must be tested on the version actually deployed.
Fabric detection Can a unique subnet prefix be set on each OpenSM? If not, manual labelling is mandatory.
Multiple fabrics per node Does this occur? If so, HCA-to-GPU NUMA matching matters.
Model storage Is there a shared filesystem, and what are the load times? Startup timeouts follow from this.
Mixed driver/CUDA versions Should engine virtualenvs differ per GPU class? (The engines config must be overridable per profile.)
Single control plane A SPOF in v1. The router caches the endpoint list (§6.4) and keeps working through a brief control plane outage.
Security Engine ports must be reachable only by the router (firewall or bind address).
Tokenizer consistency The precise KV index (§7.5 tier 2) forces the router to use the same tokenizer as the engine. In Go that means cgo or a separate tokenizer service, which is why it is deferred to v2; v1's prefix approach needs no tokenizer. The same limitation makes the TPM prompt estimate approximate in v1 (bytes/4).
Prefix affinity backfiring Without the load-imbalance guard (the §7.5 tier 1 thresholds), prefix routing raises the cache hit rate while drowning a single replica and wrecking p99. Prefix affinity must not be enabled before those thresholds are.
Tuning the admission thresholds Too low a target_waiting_per_replica starves the GPU; too high brings queueing delay back. This single number governs the system's behaviour, so measure it with ctl bench rather than guessing.
KV event loss ZMQ is fire-and-forget; a lost event makes the router silently misroute. Sequence-gap detection and falling back to tier 1 are mandatory (§7.5).
Metric staleness At 1000 requests/sec with a 2s scrape interval, decisions rest entirely on stale data. That is why the inflight scorer is not optional.

15. First Steps

  1. mkdir llmgw && cd llmgw && go mod init <module-path>/llmgw
  2. cmd/llmgw/main.go: an empty cobra skeleton with the server, agent, ctl and sim-engine subcommands.
  3. internal/sim: the fake engine — make this the first real code, because everything else is tested against it.
  4. internal/router/proxy: streaming proxy plus a cancellation test.
  5. internal/router/scheduler: the Filter/Scorer/Picker interfaces plus the queue, kv and inflight scorers. Do not write a single "least-load" function — build the pipeline from the start; coming back to it later is expensive.
  6. internal/router/admission: saturation detector and band queues. Phase 1.5.
  7. Write the end-to-end tests for the phase 1 and phase 1.5 exit criteria, then move to phase 2.

Use internal/router/scheduler/ rather than internal/router/balancer/ (the tree in §4 reflects this): "balancer" understates the job — this component does more than balance load.