Motivation

Prob based LB can disturb cache aware when the load is actually balanced. Switching to threshold based LB can help improve the cache hit rate and also makes the behavior more deterministic.

Modifications

1. bench_serving: for generated shared prefix dataset, Remove argument-based data caching and use auto caching with the key as the dataset params, so users don't have to manually configure the arguments
2. Shortest queue load balancing is triggered when the load is imbalanced. The imbalanced is defined as (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. By default it only uses abs_threshold, and rel_threshold was set to a very small value.

Benchmark

Benchmark Results

TLDR: the perf is on par with v1 on perfectly divided and low hit rate cases. In one long sys prompt case, it can outperform original RR DP, while it is slower than v1 because one worker will have "abs_threshold" more requests.

Generated Shared Prefix Dataset

python bench_serving.py --host 127.0.0.1 --port 30000 --dataset-name generated-shared-prefix \
    --generated-input-path ~/.cache/gen.json --generated-input-save-path ~/.cache/gen.json

SharedGPT Dataset

python bench_serving.py --host 127.0.0.1 --port 30000

Multi Turn Dataset

python long_prompt_multi_turn.py --port 30000 --tokenizer "/shared/public/elr-models/meta-llama/Meta-Llama-3.1-8B-Instruct/07eb05b21d191a58c577b4a45982fe0c049d0693/" | tee client.log

Generated Shared Prefix Dataset but only has one system prompt

python bench_serving.py --host 127.0.0.1 --port 30000 --dataset-name generated-shared-prefix --gen-num-groups 8 --gen-num-groups 1 --gen-prompts-per-group 1024

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.