vllm-project/router — L2 Distillation #

§1 TL;DR #

Rust-core (axum+tokio) request router for vLLM with 6 LB policies — headlined by a concurrent radix-tree cache-affinity algorithm — plus prefill-decode disaggregation over three KV connectors (NIXL/Mooncake/MoRI-IO), circuit breakers, K8s service discovery, and Python CLI via PyO3. Forked from SGLang Model Gateway.

§2 Project Identity & Motivation #

Motivation. Large-scale vLLM deployments need a router between clients and a pool of inference workers that is: (a) aware of KV cache locality to avoid redundant prefill, (b) able to route prefill and decode phases to separate worker pools (PD disaggregation), and (c) production-grade with circuit breakers, retries, observability. Generic HTTP load balancers (Nginx, HAProxy) satisfy none of these. A Python-only router would bottleneck under high QPS. The project solves this with a Rust data plane exposed to the vLLM Python ecosystem via PyO3.

§3 Architecture & Module Map #

Module responsibilities:

• Entry layer: main.rs uses clap to parse ~50 CLI flags into ServerConfig; the Python path (launch_router.py) parses the same flags via argparse and calls the PyO3-bound Router.start().
• Server layer: server.rs wires up an Axum router tree with middleware (API key validation against external URLs, token-bucket rate limiting, request-ID propagation from 4 default headers).
• Router layer: Four router implementations behind RouterTrait. The split is by deployment topology: Regular (homogeneous DP workers), PD (generic prefill/decode split), vLLM PD (vLLM-specific KV connectors), and OpenAI (backend passthrough).
• Policy layer: Six policies behind LoadBalancingPolicy. Stateless (Random, RoundRobin) vs. load-aware (PowerOfTwo) vs. affinity (ConsistentHash, RendezvousHash, CacheAware). Policies are per-model via PolicyRegistry.
• Core layer: Worker trait abstracts health, load tracking, circuit breaker, and DP-rank metadata. Tree provides the concurrent radix tree for cache-aware routing.
• Infrastructure: K8s service discovery via kube-rs pod watcher, ZMQ-based discovery for NCCL connectors, tokenizer support (HuggingFace + tiktoken) for policies that need request text, Prometheus + OpenTelemetry for observability.

§4 Entry Points & API Surface #

Public API (Python) #

from vllm_router import Router, PolicyType

router = Router(
    worker_urls=["http://w1:8000", "http://w2:8000"],
    policy=PolicyType.CacheAware,
    # ... ~50 config fields
)
router.start()

The Python package (vllm-router on PyPI) is a thin wrapper. PolicyType enum and Router class are the only public surface — everything else is internal Rust.

CLI entry points #

Top 10 configuration flags #

Extension points #

• Custom policies: implement LoadBalancingPolicy trait, register in PolicyFactory
• Custom routers: implement RouterTrait + WorkerManagement, register in RouterFactory
• Custom workers: implement Worker trait, use WorkerFactory
• Custom tokenizers: implement Tokenizer trait in tokenizer/traits.rs

§5 Core Data Structures #

Tree — concurrent multi-tenant radix tree (src/tree.rs) #

pub struct Tree {
    root: NodeRef,                              // Arc<Node>
    pub tenant_char_count: DashMap<TenantId, usize>,  // per-worker prefix volume
}

• Lifecycle: Created once per model by CacheAwarePolicy::init_workers. Lives for the router lifetime. Evicted periodically via evict_tenant_by_size(max_tree_size).
• Thread safety: Fully concurrent. Nodes use DashMap for children (32 shards at root, 8 for inner nodes) and RwLock for text. Tenant IDs are interned Arc.
• Mutation: Insert is lock-free on the happy path (DashMap entry); splits require write-locking the parent node's text. Epoch-based LRU with probabilistic timestamp updates (1-in-8) to reduce contention on hot nodes.
• Memory: CompactString for node text avoids heap allocation for short strings. Custom CharHasher for DashMap optimizes single-character key lookups at tree children.

Worker implementations (src/core/worker.rs) #

• Lifecycle: Created by WorkerFactory during router init or service-discovery update. Stored in WorkerRegistry (a DashMap>). Removed on health-check failure or K8s pod deletion.
• Thread safety: All mutable fields are atomic. CircuitBreaker uses internal RwLock for state transitions.
• DP addressing: DPAwareWorker uses url@rank composite keys. When forwarding, it injects X-data-parallel-rank header so the backend vLLM process selects the correct DP shard.

CircuitBreaker (src/core/circuit_breaker.rs) #

Three-state machine: Closed → Open → HalfOpen → Closed. Transitions gated by configurable failure_threshold, success_threshold, timeout_duration_secs, and window_duration_secs. State stored in RwLock.

RouterConfig (src/config/types.rs) #

~30-field struct driving all runtime behavior. Three RoutingMode variants (Regular, OpenAI, VllmPrefillDecode) control which router implementation is instantiated. Six PolicyConfig variants parameterize the selected policy. ConfigValidator runs structural checks before server startup.

§6 Critical Path Analysis & Implementation Highlights #

Hot path: HTTP request → routing decision → proxy #

Client HTTP request
  → axum handler (routes/pool_route.rs or prefill_decode_route.rs)
    → body parsing: serde_json::from_slice → ChatCompletionRequest
    → policy dispatch: LoadBalancingPolicy::select_worker_with_headers()
      → [CacheAware] Tree::prefix_match_with_counts() ← HOT INNER LOOP
      → compare match_rate vs cache_threshold
      → if hit: route to tenant worker; if miss: shortest-queue fallback
    → proxy: hyper HTTP client → selected worker URL
    → response streaming back to client (SSE for chat)

The cache-aware decision — the hardest-to-replicate insight #

CacheAwarePolicy::select_worker_with_headers (file: src/policies/cache_aware.rs) implements a two-regime algorithm:

1. Load-balanced regime (triggered when $\text{max\_load} - \text{min\_load} > \text{abs\_threshold}$ AND $\text{max\_load} > \text{min\_load} \times \text{rel\_threshold}$): Falls back to shortest-queue routing, sacrificing cache affinity to prevent worker starvation.
1. Cache-affinity regime (otherwise): Performs a character-level prefix match on the request text against a per-model radix tree. If $\text{match\_rate} = \frac{\text{matched\_chars}}{\text{input\_chars}} > \text{cache\_threshold}$, routes to the owning worker. Otherwise routes to the least-loaded worker and inserts the new prefix into the tree for future matches.

This is novel because:

• It operates on raw text characters, not token IDs, deliberately avoiding tokenization latency in the routing hot path. The approximation is acceptable because prefix sharing is character-aligned in practice (system prompts, few-shot examples).
• The radix tree is multi-tenant (keyed by worker URL), allowing a single tree per model to track which worker "owns" which prefix region.
• The two-regime fallback prevents the classic stickiness problem where cache-affinity routing creates load hotspots.

Prefix match inner loop (src/tree.rs) #

Tree::prefix_match_with_counts walks the trie character-by-character:

1. Look up first_char in current node's DashMap children — O(1) via CharHasher
2. Compare shared prefix via shared_prefix_count(remaining, node_text) — character scan
3. If full match: advance remaining pointer, descend; if partial: stop
4. Return PrefixMatchResult { tenant, matched_char_count, input_char_count }

The fast-path tenant lookup uses a cached last_tenant field (avoiding DashMap iteration) that is correct when the tree is read-heavy (which it is in production — inserts are infrequent relative to lookups).

Prefill-decode disaggregation #

Three KV connector modes affect routing semantics:

Connector	Mode	Routing behavior
NIXL	Pull-based (default)	Router selects prefill worker, prefill returns result; decode worker pulls KV
Mooncake	Push-based	Router injects transfer_id + remote_engine_id into request body; prefill pushes KV to decode
MoRI-IO	—	Separate I/O path for KV transfer

The VllmPrefillDecodeRouter (src/routers/http/vllm_pd_router.rs) coordinates this by: (a) selecting a prefill-decode worker pair via select_worker_pair, (b) modifying the request body with connector-specific metadata, (c) forwarding to prefill, then (d) for streaming, merging decode responses via logprobs_merge.rs.

Key implementation details #

1. Probabilistic LRU (1-in-8): The radix tree's epoch-based eviction only updates a node's last-access timestamp with probability 1/8 on reads. This dramatically reduces write contention on AtomicU64 timestamps under high-concurrency reads, at the cost of slightly stale eviction ordering — an acceptable trade-off for a routing cache.
1. DP rank injection: When intra_node_data_parallel_size > 1, DPAwareWorker transparently expands each physical worker URL into N virtual workers (url@0, url@1, ...). The router treats them as independent targets for load balancing, but dp_utils.rs injects X-data-parallel-rank into the forwarded request so the backend vLLM selects the right NCCL rank.

---

§7 Concurrency & Memory #

Concurrency model #

Fully async via tokio (multi-threaded runtime). The axum server handles requests concurrently on the tokio thread pool. No blocking I/O on the hot path — all worker health checks and HTTP proxying are async.

Lock hierarchy #

Lock	Type	Scope	Contention
Tree node children	DashMap (sharded)	Per-node	Low (32 shards at root)
Tree node text	RwLock	Per-node	Low (writes only on split)
Tree tenant_char_count	DashMap	Per-tree	Low
CircuitBreaker state	RwLock	Per-worker	Very low
Worker health/load	AtomicBool / AtomicUsize	Per-worker	Lock-free
WorkerRegistry	DashMap	Global	Low (writes only on add/remove)

No known ordering dependencies between locks — the design avoids nested locking by using atomics for high-frequency mutations (load counters, health flags) and DashMap sharding for concurrent reads.

Memory management #

• Rust ownership model: no GC. Arc for shared ownership of workers, tree nodes, and policy instances. DashMap for concurrent hash maps.
• Interned tenant IDs: Arc avoids per-lookup string allocation.
• CompactString: inline storage for short node texts (≤24 bytes on stack).
• Bounded radix tree: evict_tenant_by_size(max_tree_size) prevents unbounded memory growth; periodic eviction runs on a background tokio task with configurable eviction_interval_secs.

---

§8 Performance Characteristics #

Benchmark suite #

The repo includes three Criterion benchmarks in benches/:

Benchmark	What it measures
request_processing.rs	End-to-end request routing throughput (policy selection, no network)
tokenizer_benchmark.rs	Tokenizer encode/decode speed (HF vs tiktoken)
otel_disabled_path.rs	Overhead of OTel instrumentation when tracing is disabled

No published benchmark numbers in README or docs. The benchmarks measure router-internal latency, not end-to-end serving performance (which is dominated by GPU inference time).

Scaling properties #

• Worker count: All policies are O(N) or better in worker count. CacheAware is O(L) in request text length for the radix tree walk, plus O(N) for the load comparison fallback.
• Concurrent requests: Bounded by max_concurrent_requests config and tokio thread pool. Lock-free atomics for load tracking avoid serialization.
• Multi-model: PolicyRegistry maintains independent policy instances per model, so cross-model interference is zero.
• Radix tree size: Bounded by max_tree_size per model. Eviction is O(tenants) to find the largest tenant, then O(tree depth) to remove its entries.

Expected bottlenecks #

1. Tokenization (if enabled): HuggingFace tokenizer initialization is slow (~100ms per model); runtime encode is ~1ms/request but adds to P99 latency.
2. HTTP proxying: The actual forwarding via hyper is the dominant latency component — routing decisions are sub-microsecond for stateless policies, low-microsecond for CacheAware.
3. Tree contention under write-heavy workloads: If many unique prompts arrive simultaneously (cold cache), the tree sees many inserts, each requiring DashMap write locks. The 32-shard root mitigates this.

---

§9 Tech Debt & Code Quality #

Test coverage #

Three-tier test suite:

• Unit tests (py_test/unit/): arg parsing, config validation
• Integration tests (py_test/integration/): all 6 LB policies, circuit breaker state transitions, retry logic, PD routing
• E2E tests (py_test/e2e/): full router with mock workers (regular, PD, embeddings)

No pytest --cov badge or coverage numbers published. Rust-side unit tests likely exist in #[cfg(test)] modules but are not surfaced in the repo's CI config.

CI #

.buildkite/ pipeline config present. No GitHub Actions visible. Buildkite likely runs cargo test, cargo clippy, Python tests, and Docker builds.

Linter discipline #

• Rust: Cargo.toml workspace with standard edition settings; clippy likely enforced in CI
• Python: pyproject.toml present; no visible ruff or mypy config

Known warts #

1. No CHANGELOG.md — version history only from git tags
2. handler.rs and server.rs were inaccessible during L1 fetch (timeout), suggesting large files that may benefit from splitting
3. mini_lb.py is a "debug only" pure-Python mini load balancer — tech debt artifact
4. Rendezvous hashing policy exists but is not documented in the README's policy table

Dependency health #

Key Rust dependencies: axum (web framework), tokio (async runtime), hyper (HTTP), dashmap (concurrent map), pyo3 (Python bindings), serde/serde_json, clap (CLI), kube (K8s client), zeromq (ZMQ), prometheus (metrics). All are actively maintained, mainstream crates.

Key Python dependencies: maturin/setuptools-rust for PyO3 build. The Python surface is thin enough to have minimal dependency risk.

---

§10 Community Health #

Metric	Observation
Stars	233 (growing — project is young, v0.1.x)
License	Apache-2.0 (permissive)
Primary author	Byron Hsu
Org	vllm-project (corporate-backed by the vLLM team)
Bus factor	Low — appears primarily single-maintainer with the vLLM org as backstop
Governance	Part of the vllm-project GitHub org; effectively BDFL + org review
Fork origin	SGLang Model Gateway; has diverged significantly

The project is early-stage (v0.1.14) and tightly coupled to vLLM's PD disaggregation roadmap. Being under the vllm-project org provides credibility and long-term maintenance assurance, but the bus factor for the router-specific code is currently low.

---

§11 Comparison with Alternatives #

Dimension	vllm-project/router	SGLang Model Gateway	Nginx/HAProxy	Custom Python router
Language	Rust + Python (PyO3)	Python	C	Python
Cache-aware routing	Yes (radix tree)	Yes (radix tree, original impl)	No	Possible but slow
PD disaggregation	Yes (NIXL/Mooncake/MoRI)	Partial	No	Manual
LB policies	6 (including P2C, HRW)	Fewer	Many (generic)	Custom
Circuit breaker	Built-in	No	External	Manual
K8s service discovery	Native (kube-rs)	No	Ingress controller	Manual
Latency overhead	Sub-µs (Rust hot path)	~ms (Python)	Sub-µs	~ms
Observability	Prometheus + OTel	Basic	Extensive	Custom
OpenAI API compat	Full (chat, completion, embeddings, rerank, responses)	Partial	Passthrough	Custom
Ecosystem fit	vLLM-native	SGLang-native	Generic	Any
Maturity	v0.1.x (early)	Integrated in SGLang	Battle-tested	Varies

Key differentiator: The router is the only option that combines sub-microsecond Rust routing decisions with vLLM-native PD disaggregation support across three KV connector backends. The SGLang gateway has the cache-aware radix tree (it's the original source), but lacks the Rust performance and vLLM PD integration.

---

§12 Verdict & Recommendations #

When to adopt (the "yes" regime) #

• You run ≥2 vLLM workers behind a single endpoint and need cache-aware or PD-disaggregated routing
• You need vLLM-specific PD disaggregation (NIXL/Mooncake/MoRI-IO connectors)
• You want Kubernetes-native service discovery for dynamically scaling vLLM pods
• You need production-grade resilience (circuit breakers, retries) with minimal latency overhead
• Your deployment uses data parallelism with intra_node_data_parallel_size > 1 and needs transparent rank injection

When NOT to adopt (the "no" regime) #

• You use SGLang, TensorRT-LLM, or other non-vLLM backends — the PD disaggregation is vLLM-specific
• You only have a single worker — no routing needed
• You need a general-purpose API gateway with auth, rate-limiting, transformation — use a dedicated gateway (Kong, Envoy) with this router behind it
• You need stable, battle-tested production software — v0.1.x with low bus factor carries risk for mission-critical deployments
• You require custom routing logic beyond the 6 built-in policies and don't want to write Rust

Usage gotchas #

1. Cache-aware threshold tuning: The cache_threshold parameter critically affects routing quality. Too low → false cache hits cause imbalanced load; too high → no affinity benefit. Start with 0.5 and tune based on your prompt distribution.
2. DP size must match backend: intra_node_data_parallel_size in the router MUST equal the DP size configured on the vLLM workers. Mismatch causes silent routing errors (requests to wrong ranks).
3. ZMQ discovery is NCCL-only: The --vllm-discovery-address flag only works with the NCCL connector, not NIXL or Mooncake.
4. No graceful worker drain: Removing a worker (via K8s scaling or health failure) drops in-flight requests to that worker. The retry mechanism catches some, but streaming responses are lost.
5. Token bucket is per-router, not per-user: The rate_limit_tokens_per_second applies globally, not per API key. Multi-tenant fairness requires an external gateway.

Suggested contributions #

• Add integration benchmarks with published numbers (router overhead per policy under load)
• Document the Rendezvous hashing policy in the README
• Implement graceful worker drain (finish in-flight, stop new) before removal
• Add per-user rate limiting for multi-tenant deployments

---

§13 Implementation Cross-Reference #

Concept	File	Key function/struct
Cache-aware routing decision	src/policies/cache_aware.rs	CacheAwarePolicy::select_worker_with_headers
Radix tree prefix match	src/tree.rs	Tree::prefix_match_with_counts
Radix tree insert/eviction	src/tree.rs	Tree::insert, Tree::evict_tenant_by_size
Worker abstraction	src/core/worker.rs	Worker trait, BasicWorker, DPAwareWorker
Circuit breaker FSM	src/core/circuit_breaker.rs	CircuitBreaker
Retry with backoff	src/core/retry.rs	exponential backoff + jitter
PD routing (vLLM)	src/routers/http/vllm_pd_router.rs	VllmPrefillDecodeRouter
PD routing (generic)	src/routers/http/pd_router.rs	generic PD router
Policy factory	src/policies/factory.rs	PolicyFactory
Router factory	src/routers/factory.rs	RouterFactory
Config types	src/config/types.rs	RouterConfig, RoutingMode, PolicyConfig
K8s service discovery	src/service_discovery.rs	K8s pod watcher
ZMQ service discovery	src/routers/http/vllm_service_discovery.rs	ZMQ worker registration
PyO3 binding	src/lib.rs	Router, PolicyType, vllm_router_rs module
Python CLI	py_src/vllm_router/launch_router.py	main(), launch_router()
DP rank injection	src/routers/http/dp_utils.rs	X-data-parallel-rank header
Logprobs merging	src/routers/http/logprobs_merge.rs	streaming logprobs merge
Prometheus metrics	src/metrics.rs	RouterMetrics
OpenTelemetry	src/otel_trace.rs	tracing setup

核心技术壁垒: The concurrent multi-tenant radix tree (src/tree.rs) with character-level prefix matching, probabilistic LRU timestamps (1-in-8 update), custom CharHasher, interned tenant IDs, and the two-regime load-balance/cache-affinity switch in CacheAwarePolicy. Replicating this requires deep understanding of both concurrent data structure design and LLM serving cache dynamics.