[rust] cache-aware DP - approx tree

[ByronHsu]

1. Algorithm Overview

#1732

We propose a method to approximate worker-side radix trees on the router side using request and response information flow. The router maintains approximated trees ("approx trees") that mirror the cache state of each worker's radix tree ("worker tree").

Core Algorithm

Given N workers, the algorithm operates as follows:

1. Initialize N approx trees on the router, one corresponding to each worker
2. For each incoming request:
   • Match request's input_ids against all approx trees
   • Select the worker with highest cache match
   • Speculation Phase: Before forwarding the request, insert input_ids into the selected approx tree
   • Correction Phase: Upon receiving the worker's response, remove non-cached tokens from the approx tree

The speculation phase anticipates the worker tree's future state, while the correction phase aligns the approx tree with the worker's actual cache state. This forward-backward mechanism enables continuous self-adjustment of the approximation.

Load Balancing Strategies

1. Cache Threshold with Shortest Queue

1. Maintain N approx trees on the router
2. For each request:
   • Calculate cache match rates across all workers
   • If highest match rate > threshold: Select highest-matching worker
   • Else: Select worker with shortest request queue
   • Apply speculation and correction phases

2. Variance-Based Load Balancing

1. Maintain N approx trees on the router
2. For each request:
   • Calculate mean (μ) and standard deviation (σ) of queue lengths
   • Filter out workers with queue length > μ + Kσ
   • Among remaining workers, select highest cache match
   • Apply speculation and correction phases

Parameters

• K: variance threshold parameter
• μ: mean queue length
• σ: standard deviation of queue lengths

See details of google doc version here

2. Changes

1. Implemented the core algorithm and "Cache Threshold with Shortest Queue" LB strategy in rust
2. Rust python binding for approx tree
3. Rust cli for approx tree
4. A benchmark file which simulates the case of long system prompt + relatively shorter QA.
5. Demo python binding usage under py_src

• main.py: launch a minimal router server with existing worker
• dp_demo.py: mimic the case like the current --dp. Launch --dp workers and a router. This part of code can be moved to sglang core afterwards and let sglang depends on sglang-router. Users can install sglang-router as an optional to sglang.

Benchmark

# original
python -m sglang.launch_server --model-path <model path> --host 127.0.0.1 --port 30000 --dp 8

# approx tree
python dp_demo.py --model-path <model path> --local-tokenizer-path <local tokenizer.json path> --dp 8

# benchmark client
python long_prompt_multi_turn.py --port 30000  --tokenizer <tokenizer path>

We can see clear improvement from the original method due to the high cache hitting.

Reference: google sheet

Follow-up

There are still many follow-ups. Notably

1. Python binding CI publishing setup (sglang-router pypi)
2. Various optimization in rust (marked as TODO in the code`)
3. Implementation of the correction phase and Variance-Based Load Balancing

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[jayzhan211]

Dashmap may helps here

[ByronHsu]

Actually i want to lock on the node of radixtree instead of the whole hashmap. Do you have any suggestion for that?

[jayzhan211]

I think we can start from the simple one (lock whole map). Concurrent map is quite complex. And may try out this comment later on if you are interested
https://users.rust-lang.org/t/locking-only-one-entry-in-hashmap/84764/4

[jayzhan211]

If worker_urls is mostly for read only, instead of owned string, can use Arc<str> for thread-safe. Vec<Arc<str>> is probably suitable in this case

[ByronHsu]

For now Vec is read-only, but once we add dynamic scaling, it will be read+write and mostly read. Vec<Arc<str>> looks great