[Feat][Router] Add TTFT Routing

[chickeyton]

FIX #583

An answer to the feature request TTFT Routing

Dependancies

This PR depends on the to be merged feature FullLookUp in LMCache. Therefore, this PR can not be merged until FullLookUp be merged and be released in an offiical version of LMCache

Design

Motivation

The current PrefixAwareRouter and KvawareRouter (in production-stack) route user request purely base on the length the cached prefix. However, it may not be optimal due to the selected instance may have a long pending request queue. A more efficient way is estimating TTFT(or prefill workload) of the incoming request for each instance, select the one with the least estimated TTFT(or prefill workload).

TTFT & Prefill Workload Estimation

TTFT esitmation require 1) prefill computation workload estimation & 2) computation speed profiling, but since there are different parallelisms going on and there are lots of factors, it often give unreliable numbers about the speed, the solution is simply assume the computation speed of all instances are the same, then comparing workload is equivalent to comparing TTFT, the experiment result supports that point.

Prefill Workload Estimation:

def workload(num_prefix_tokens, num_cached_tokens):
    """A comparable number to represent the size of """
    """computation amount (Q•K dot products or alike operations) """
    """in prefill by trapezoid area formula"""
    top = num_cached_tokens + 1
    bottom = num_prefix_tokens
    height = num_prefix_tokens - num_cached_tokens
    return (top  + bottom) * height / 2

It is not the actual no. of float operations required, since each model architecture introduce differences, it is not practical to have a tailor made formula for each architecture, so we use a comparable number to represent the size of computation amount (no. of Q•K dot products or alike operations).

For each instance, we forcast the pending workload with the addtion workload from the new request:

def estimate_workload(instance, new_req):
    queue_workload = sum([workload(req.num_prefill_tokens, req.num_cache_tokens) for req in instance.request_queue])
    new_workload = workload(new_req.num_prefill_tokens, new_req.num_cache_tokens)
    return queue_workload + new_workload 

num_cache_tokens: No. of cached tokens in prefix, queried from the FullLookup service of LMCache, if share cache is enabled, the max no. of cached tokens among all vllm instance will be used.

Actual Code Changes:

1. Added TtftRouter, like KvawareRouter, it host a centralized LMCache registry service, the vllm instances using the LMCache connector and register/unregister prefix caches from that registry service. On request routing, TtftRouter do a FullLookup (to be merged feature in LMCache) from the registry service to retrieve

2. Added optional commandline argument --tokenizer <model> for specifing the tokenizer path or HF model name, KvawareRouter and TtftRouter will load that tokenizer immediately after starting the LMCache registry service. If not specified they will load the tokenizer with the model specified in the first vllm endpoint on routing the first user request.

3. Added optional commandline argument --lmcache-instances <comman separated list of LMCache instance id> used with --service-discovery static only, if two or more vllm endpoints share a same IP then --lmcache-instances must be specified, otherwise the QueryInstMsg to LMCache queries in KvawareRouter and TtftRouter will give wrong results

4. Added optional commandline argument --enable-share-cache, specify that if p2p transfer is enabled in LMCache.

5. Added prefill_todo_workload measures to RequestStats referring to queue_workload

6. Added RequestStatsCacheInfo that returns by TtftRouter.route_request() and kept track by RequestStatsMonitor

Usage:

Assuming two vllm endpoints, vllm router, and Redis service are hosted on the same machine

1. prepare 2 LMCache yaml config files: ~/lmcache_config_0.yaml and ~/lmcache_config_1.yaml:

example lmcache_config_0.yaml

chunk_size: 256
controller_url: localhost:65432
distributed_url: localhost:8700
enable_controller: true
enable_p2p: false
lmcache_instance_id: instance_0
lmcache_worker_port: 8710
local_cpu: true
local_disk: /tmp/lmcache_0
lookup_url: localhost:8600
max_local_cpu_size: 1.0
max_local_disk_size: 10.0

2. Start Redis server: redis-server --port 8600

3. Start the first vllm endpoint:

LMCACHE_USE_EXPERIMENTAL=True \
LMCACHE_CONFIG_FILE=~/lmcache_config_0.yaml \
CUDA_VISIBLE_DEVICES=0 \
vllm serve Qwen/Qwen3-0.6B --host 0.0.0.0 --port 9081 \
--gpu-memory-utilization 0.8 \
--kv-transfer-config '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}'

4. Start the second vllm endpoint like the first one with proper namings, port and CUDA_VISIBLE_DEVICES, LMCACHE_CONFIG_FILE values

5. Start vllm-router with TtftRouter:

vllm-router --port 9091 \
    --service-discovery static \
    --static-backends "http://localhost:9081,http://localhost:9082" \
    --static-models "Qwen/Qwen2-7B,Qwen/Qwen2-7B" \
    --engine-stats-interval 10 \
    --log-stats \
    --routing-logic ttft \
    --lmcache-controller-port 65432 \
    --lmcache-instances "instance_0,instance_1" \
    --session-key SESSION \
    --tokenizer "Qwen/Qwen2-7B"

--routing-logic set to ttft
--lmcache-controller-port must be same as the controller_url's port specified in lmcache_config_*.yaml
--lmcache-instances must be listing the lmcache_instance_id value specified in lmcache_config_*.yaml as two vllm endpoints share the same IP

6. Good to send request to the vllm-router now:

curl http://localhost:9091/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen2-7B",
    "prompt": "Hello, my name is",
    "max_tokens": 50
  }'

Experiment 1

Setup

Model: Qwen/Qwen2-7B x 2 vllm instances
GPU: NVIDIA GeForce RTX 3090 (24GB) X 2
Routers:

TTFT

vllm-router --port 9091 \
    --service-discovery static \
    --static-backends "http://localhost:9081,http://localhost:9082" \
    --static-models "Qwen/Qwen2-7B,Qwen/Qwen2-7B" \
    --engine-stats-interval 10 \
    --log-stats \
    --routing-logic ttft \
    --lmcache-controller-port 65432 \
    --lmcache-instances "instance_0,instance_1" \
    --session-key SESSION \
    --tokenizer "Qwen/Qwen2-7B"

Kvaware

vllm-router --port 9091 \
    --service-discovery static \
    --static-backends "http://localhost:9081,http://localhost:9082" \
    --static-models "Qwen/Qwen2-7B,Qwen/Qwen2-7B" \
    --engine-stats-interval 10 \
    --log-stats \
    --routing-logic kvaware \
    --lmcache-controller-port 65432 \
    --lmcache-instances "instance_0,instance_1" \
    --session-key SESSION \
    --kv-aware-threshold 999999999 \
    --tokenizer "Qwen/Qwen2-7B"

QPS (fallback of Kvaware by setting '--kv-aware-threshold 0')

vllm-router --port 9091 \
    --service-discovery static \
    --static-backends "http://localhost:9081,http://localhost:9082" \
    --static-models "Qwen/Qwen2-7B,Qwen/Qwen2-7B" \
    --engine-stats-interval 10 \
    --log-stats \
    --routing-logic kvaware \
    --lmcache-controller-port 65432 \
    --lmcache-instances "instance_0,instance_1" \
    --session-key SESSION \
    --kv-aware-threshold 0 \
    --tokenizer "Qwen/Qwen2-7B"

Benchmark Script (from production-stack/benchmarks/multi-round-qa):

python ./multi-round-qa.py --num-users 5 --num-rounds 10 --qps 100 --max-unfinished-queries 20 --shared-system-prompt 1 --user-history-prompt 200 --answer-len 100 --model "Qwen/Qwen2-7B" --sharegpt --init-user-id 0 --base-url http://localhost:9091/v1 --time 180

Results

TTFT

==================== Performance summary ======================
  QPS: 100.0000 reqs/s

  Processing speed: 11.4516 reqs/s

  Requests on-the-fly: 0

  Input tokens per second: 8549.3601 tokens/s

  Output tokens per second: 746.7229 tokens/s

  Average generation throughput (per request): 42.3025 tokens/req/s

  Average TTFT: 0.3968s

Time range: 1758281191.9479365 - 1758281374.5431235 (182.60s)
===============================================================

Kvaware

==================== Performance summary ======================
  QPS: 100.0000 reqs/s

  Processing speed: 11.2663 reqs/s

  Requests on-the-fly: 0

  Input tokens per second: 8403.1305 tokens/s

  Output tokens per second: 734.8160 tokens/s

  Average generation throughput (per request): 44.4633 tokens/req/s

  Average TTFT: 0.4788s

Time range: 1758282338.741013 - 1758282520.3442867 (181.60s)
===============================================================

QPS

==================== Performance summary ======================
  QPS: 100.0000 reqs/s

  Processing speed: 11.0099 reqs/s

  Requests on-the-fly: 0

  Input tokens per second: 8202.4903 tokens/s

  Output tokens per second: 717.3891 tokens/s

  Average generation throughput (per request): 45.0055 tokens/req/s

  Average TTFT: 0.5119s

Time range: 1758281833.2226286 - 1758282015.6940422 (182.47s)
===============================================================

Avg. TTFT:
TTFT: 0.3968s (Kvaware -17.1%)
Kvaware: 0.4788s
QPS: 0.5119s

Experiment 2

Setup

Same as experiment 1

Benchmark script is a talior made one simulates a scenario that the first same 256 tokens are always be cached then followed by 20k random tokens

import os
import numpy as np
import time
import requests
from multiprocessing import Process, Value, Array, Lock

base_url = "http://localhost:9091/v1"
model = "Qwen/Qwen2-7B"

num_requests = 100
num_workers  = 10
num_max_active_requests = 5


word_pool = ["hi", "hellow", "yes", "no", "cat", "dog", "pig", "game", "coffee", "cake", "noodles", "burger", "football", "tennis", "ship", "car", "ship", "boat"]

active_requests = Value('i', 0)
finish_times = Array('d', [0] * num_requests)
lock = Lock()


def gen_prompt(num_words):
    rand_nums = np.random.randint(0, len(word_pool), num_words)
    return " ".join([word_pool[n] for n in rand_nums])


def http_request(prompt):
    json_obj = {}
    json_obj["model"] = model
    json_obj["prompt"] = prompt
    json_obj["max_tokens"] = 1

    url = base_url + "/completions"

    response = requests.post(url, json=json_obj)
    if(response.status_code != 200):
        raise ValueError(f"status_code:{response.status_code} is not 200")


fixed_prefix = gen_prompt(260)


def request_proc(worker_id, num_requests, active_requests, finish_times, lock):
    np.random.seed(worker_id)
    for i in range(num_requests):
        while True:
            time.sleep(0.01)
            with lock:
                if active_requests.value >= num_max_active_requests:
                    continue
                active_requests.value = active_requests.value + 1
                break

        prompt2 = gen_prompt(20000)
        start_time = time.time()
        http_request(fixed_prefix + ' ' + prompt2)
        elapsed_time = time.time() - start_time
        with lock:
            finish_times[worker_id + i]=elapsed_time
            active_requests.value = active_requests.value - 1


worker_requests = num_requests // num_workers
processes = []
for w in range(num_workers):
    p = Process(target=request_proc, args=(w, worker_requests, active_requests, finish_times, lock))
    p.start()
    processes.append(p)

for p in processes:
    p.join()


print(f"avg finish time: {np.mean(finish_times[:])}s")

Avg. TTFT:
TTFT: 2.4981s (QPS -18.5%)
Kvaware: 5.3141s
QPS: 3.0669s

In this experiement, Kvaware always route the requests to the same vllm instance, so it perform poorly

Conculsion

TTFT routing shows significant reduction in TTFT compare to Kvaware & QPS routing by estimating the prefill workload of the vllm instances. And it is more robust than Kvaware routing, consider that, A client keep sending requests with a fixed system prompt of a few hundreds tokens, the Kvaware routing will always route the request to the same vllm instance while TTFT routing provides a good load balancing in that use case.

---

• Make sure the code changes pass the pre-commit checks.
• Sign-off your commit by using -s when doing git commit
• Try to classify PRs for easy understanding of the type of changes, such as [Bugfix], [Feat], and [CI].

Detailed Checklist (Click to Expand)

Thank you for your contribution to production-stack! Before submitting the pull request, please ensure the PR meets the following criteria. This helps us maintain the code quality and improve the efficiency of the review process.

PR Title and Classification

Please try to classify PRs for easy understanding of the type of changes. The PR title is prefixed appropriately to indicate the type of change. Please use one of the following:

• [Bugfix] for bug fixes.
• [CI/Build] for build or continuous integration improvements.
• [Doc] for documentation fixes and improvements.
• [Feat] for new features in the cluster (e.g., autoscaling, disaggregated prefill, etc.).
• [Router] for changes to the vllm_router (e.g., routing algorithm, router observability, etc.).
• [Misc] for PRs that do not fit the above categories. Please use this sparingly.

Note: If the PR spans more than one category, please include all relevant prefixes.

Code Quality

The PR need to meet the following code quality standards:

• Pass all linter checks. Please use pre-commit to format your code. See README.md for installation.
• The code need to be well-documented to ensure future contributors can easily understand the code.
• Please include sufficient tests to ensure the change is stay correct and robust. This includes both unit tests and integration tests.

DCO and Signed-off-by

When contributing changes to this project, you must agree to the DCO. Commits must include a Signed-off-by: header which certifies agreement with the terms of the DCO.

Using -s with git commit will automatically add this header.

What to Expect for the Reviews

We aim to address all PRs in a timely manner. If no one reviews your PR within 5 days, please @-mention one of YuhanLiu11
, Shaoting-Feng or ApostaC.

[gemini-code-assist]

Summary of Changes

Hello @chickeyton, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates a novel Time To First Token (TTFT) routing mechanism into the vLLM router. This enhancement aims to optimize request distribution across available backends by predicting the fastest path to the first generated token, thereby improving overall system responsiveness and efficiency. The changes also include necessary updates to the request monitoring and benchmarking tools to support and evaluate this new routing approach.

Highlights

• New TTFT Routing Logic: Introduced a new Time To First Token (TTFT) routing strategy that intelligently selects the optimal backend by estimating the time it will take to generate the first token, considering KV cache hits and current engine load.
• Enhanced Prompt Handling: Implemented a robust utility to extract prompts from various request formats (completion, chat, multimodal) for accurate tokenization, which is essential for effective KV cache lookups in the new routing logic.
• Improved Request Statistics: Extended the request statistics monitoring to track critical metrics such as 'engine_prefill_tps' (tokens per second during prefill) and 'uncomputed_prefix_tokens', providing the necessary data for the TTFT routing algorithm's estimations.
• Benchmarking Tool Updates: Modified the multi-round QA benchmark to include a 'max_unfinished_queries' parameter, allowing for better control over concurrent requests and more accurate simulation of load during performance testing of the new routing strategy.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new TTFT (Time To First Token) routing strategy, which is a significant feature. The implementation spans routing logic, statistics collection, and request processing, and includes updates to the benchmarking tool. The overall approach is sound, but I've identified a critical bug in the new TtftRouter where an exception is created but not raised. I've also left several comments to improve code clarity, maintainability, and robustness, such as renaming misleading variables, removing magic numbers, and improving an algorithm's efficiency. Addressing these points will strengthen the new routing logic.

[chickeyton]

This PR depends on to be merged FullLookup feature of LMCache, so cannot be merged yet, but please take a look of the design if u have time, thanks!

@KuntaiDu @ApostaC

[chickeyton]

/gemini review

[chickeyton]

/gemini summary

[gemini-code-assist]

Code Review

This pull request introduces a new TTFT (Time-To-First-Token) based routing logic, which is a significant feature enhancement. The implementation includes the new TtftRouter, supporting CLI arguments, and necessary updates to statistics collection. The code is well-structured, but I've found a few critical issues and areas for improvement. Specifically, there are critical bugs in the TtftRouter's logic for finding the best-matched cache and in the initialization of a data class for request stats. There are also opportunities to improve the robustness of argument parsing and the correctness of an interval-merging utility class. Addressing these points will be crucial for the stability and correctness of this new feature.

[gemini-code-assist]

Summary of Changes

This pull request introduces a new routing logic, TtftRouter, to the vllm-router component. The primary goal is to optimize request distribution across vLLM instances by estimating the Time-To-First-Token (TTFT) for incoming requests. This estimation considers factors like pending request queue time and KV cache transfer time, leading to more intelligent routing decisions than purely prefix-based methods. The implementation is a partial step towards a more comprehensive "Ultimate Design" for TTFT estimation, with future work planned for real-time measurement of prefill computation and transfer speeds within the vLLM engine and LMCache.

Highlights

• New TTFT Routing Logic: Introduced a new routing logic, TtftRouter, which routes requests based on an estimated Time-To-First-Token (TTFT). This aims to improve load balancing by considering factors like pending request queue time and KV cache transfer time, leading to more efficient request distribution.
• Enhanced LMCache Integration: The new TtftRouter leverages LMCache's FullLookUp feature to retrieve detailed KV cache information, which is crucial for accurate TTFT estimation. This integration requires a pending LMCache pull request to be merged.
• Flexible Router Configuration: Added new command-line arguments --tokenizer to specify the tokenizer model and --lmcache-instances to map LMCache instance IDs to backend URLs. These additions provide greater flexibility in configuring the router setup.
• Improved Request Statistics: The RequestStats now includes engine_prefill_comp_speed and uncomputed_prefix_tokens, offering more granular metrics for informed routing decisions. A TimePeriods utility was also introduced to accurately calculate active prefill periods.
• Benchmarking Tool Enhancement: The multi-round-qa.py benchmarking script has been updated with a --max-unfinished-queries argument, allowing users to better control and limit the number of concurrent unfinished queries during stress tests.
• Code Refinements and Bug Fixes: Several issues identified during code review were addressed, including a critical constructor typo in RequestStatsCacheInfo, a misleading parameter name in KvawareRouter (renamed to instance_id_to_url), and a more robust prompt extraction utility.

Changelog

• benchmarks/multi-round-qa/multi-round-qa.py
  • Added max_unfinished_queries to WorkloadConfig and as a command-line argument.
  • Implemented logic to pause new query creation if max_unfinished_queries limit is reached.
• src/vllm_router/app.py
  • Introduced create_instance_id_to_url helper function to map LMCache instance IDs to backend URLs.
  • Updated initialize_all to pass tokenizer name and instance ID to URL mapping to router initialization.
• src/vllm_router/parsers/parser.py
  • Added ttft as a choice for --routing-logic.
  • Introduced new command-line arguments: --lmcache-instances and --tokenizer.
• src/vllm_router/routers/routing_logic.py
  • Implemented the TtftRouter class, including its __init__, start_kv_manager, route_request, _find_best_matched, _find_best_ttft, _estimate_ttft, _get_instance_url, _calc_transfer_time, and _fallback_routing methods.
  • Added TTFT to the RoutingLogic enum.
  • Created a new utility function extract_prompt to handle prompt extraction from various request formats.
  • Refactored KvawareRouter and PrefixAwareRouter to use the new extract_prompt function.
  • Renamed instance_id_to_ip to instance_id_to_url in KvawareRouter for clarity.
  • Added tokenizer_name and instance_id_to_url parameters to KvawareRouter and TtftRouter constructors.
  • Defined CACHE_LOC_TO_TRANS_TIME and DEFAULT_CACHE_TRANS_TIME constants in TtftRouter.
  • Updated initialize_routing_logic to correctly instantiate and start TtftRouter.
• src/vllm_router/services/request_service/request.py
  • Modified process_request to accept an optional cache_info argument.
  • Updated route_general_request to handle the new tuple return type from TtftRouter (containing server URL and cache_info).
  • Extended isinstance checks to include TtftRouter for routing logic.
• src/vllm_router/stats/request_stats.py
  • Added engine_prefill_comp_speed and uncomputed_prefix_tokens fields to RequestStats.
  • Introduced TimePeriods class for managing and computing lengths of overlapping time intervals.
  • Added RequestStatsCacheInfo dataclass to store num_prefix_tokens and num_cached_tokens.
  • Updated RequestStatsMonitor to track cache_infos for requests.
  • Implemented _calc_engine_prefill_comp_speed and _get_uncomputed_prefix_tokens methods to compute new request statistics.
  • Corrected the constructor name for RequestStatsCacheInfo from __int__ to __init__.

Activity

• Author chickeyton requested review, noting a dependency on a pending LMCache PR.
• A critical bug was identified by gemini-code-assist[bot] where a ValueError was created but not raised in TtftRouter, which was subsequently fixed by chickeyton.
• The parameter instance_id_to_ip in KvawareRouter was suggested to be renamed to instance_id_to_url for accuracy by gemini-code-assist[bot], and chickeyton implemented this change.
• A typo in lmcache_contorller_port was pointed out by gemini-code-assist[bot] and fixed by chickeyton.
• Hardcoded transfer times were recommended to be replaced with named constants in TtftRouter by gemini-code-assist[bot], leading to the introduction of CACHE_LOC_TO_TRANS_TIME and DEFAULT_CACHE_TRANS_TIME.
• The route_request method was advised by gemini-code-assist[bot] to explicitly return uncached_prefix_tokens (now RequestStatsCacheInfo) along with the URL, a change chickeyton adopted.
• Debugging print statements were removed from the code as suggested by gemini-code-assist[bot].
• A discussion occurred regarding the efficiency of the TimePeriods.union method, with gemini-code-assist[bot] suggesting a more standard interval merging algorithm, though chickeyton responded asserting the current implementation's O(N) complexity.
• A critical typo (__int__ instead of __init__) in RequestStatsCacheInfo's constructor was found by gemini-code-assist[bot].
• A more robust implementation for parsing lmcache_instances to handle whitespace and empty strings was suggested by gemini-code-assist[bot].
• The type hint for TtftRouter.route_request was suggested to be corrected by gemini-code-assist[bot] to accurately reflect its tuple return type.

[chickeyton]

/gemini review

[chickeyton]

/gemini summary

[gemini-code-assist]

Code Review

This pull request introduces a new TtftRouter routing logic that aims to reduce Time To First Token (TTFT) by selecting instances with the lowest estimated prefill workload. The changes are extensive, touching routing logic, request processing, and statistics monitoring. Overall, this is a valuable feature addition. However, I've identified a few critical issues, including a potential runtime error due to an incomplete refactoring, a missing return statement that could cause crashes, and a memory leak in the statistics monitor. Addressing these issues is crucial for the stability and correctness of the new functionality.

[gemini-code-assist]

Summary of Changes

This pull request introduces a significant enhancement to the vllm-router by adding a new TTFT (Time-To-First-Token) routing logic. This new router intelligently distributes incoming requests across vLLM instances by estimating the prefill workload, which includes considering pending requests and the efficiency of KV cache utilization. The goal is to minimize the time it takes to receive the first token of a response, thereby improving overall system responsiveness and load balancing. The implementation integrates deeply with LMCache's advanced lookup capabilities and provides new configuration options for greater operational flexibility.

Highlights

• New TTFT Routing Logic: Introduced a new TtftRouter that routes requests based on an estimated Time-To-First-Token (TTFT), aiming to optimize load balancing by considering prefill workload and KV cache utilization across instances.
• Enhanced LMCache Integration: The TtftRouter leverages LMCache's FullLookup feature to retrieve detailed KV cache information, which is critical for accurate TTFT estimation. This feature depends on a pending LMCache pull request.
• Flexible Router Configuration: Added new command-line arguments: --tokenizer for specifying the tokenizer model, --lmcache-instances for mapping LMCache instance IDs to backend URLs, and --enable-share-cache to indicate if p2p transfer is enabled in LMCache.
• Improved Request Statistics: The RequestStats now includes prefill_todo_workload to track unfinished prefill computation, and a new RequestStatsCacheInfo class is used to return cache-related details from the TtftRouter.
• Benchmarking Tool Enhancement: The multi-round-qa.py benchmarking script has been updated with a --max-unfinished-queries argument, allowing better control over concurrent unfinished queries during stress tests.

Changelog

• benchmarks/multi-round-qa/multi-round-qa.py
  • Added max_unfinished_queries to WorkloadConfig and as a command-line argument.
  • Implemented logic to pause new query creation if max_unfinished_queries limit is reached.
• src/vllm_router/app.py
  • Introduced create_instance_id_to_url helper function to map LMCache instance IDs to backend URLs.
  • Updated initialize_all to pass tokenizer name and instance ID to URL mapping to router initialization.
• src/vllm_router/parsers/parser.py
  • Added ttft as a choice for the --routing-logic argument.
  • Introduced new command-line arguments: --lmcache-instances, --tokenizer, and --enable-shared-cache.
• src/vllm_router/routers/routing_logic.py
  • Implemented the TtftRouter class, including its initialization, KV manager startup, request routing, workload estimation, and instance URL resolution methods.
  • Added TTFT to the RoutingLogic enum.
  • Created a new utility function extract_prompt to handle prompt extraction from various request formats.
  • Refactored KvawareRouter and PrefixAwareRouter to use the new extract_prompt function.
  • Renamed instance_id_to_ip to instance_id_to_url in KvawareRouter for clarity.
  • Added tokenizer_name and instance_id_to_url parameters to KvawareRouter and TtftRouter constructors.
  • Updated initialize_routing_logic to correctly instantiate and start TtftRouter.
  • Modified RoutingInterface.route_request return type to support RequestStatsCacheInfo.
• src/vllm_router/services/request_service/request.py
  • Modified process_request to accept an optional cache_info argument.
  • Updated route_general_request to handle the new tuple return type from TtftRouter (containing server URL and cache_info).
  • Extended isinstance checks to include TtftRouter for routing logic.
  • Passed cache_info to request.app.state.request_stats_monitor.on_new_request.
• src/vllm_router/stats/request_stats.py
  • Added prefill_workload function for calculating prefill computation workload.
  • Added prefill_todo_workload field to RequestStats.
  • Introduced RequestStatsCacheInfo class to store num_prefix_tokens and num_cached_tokens.
  • Updated RequestStatsMonitor to track cache_infos for requests.
  • Modified on_new_request to accept an optional cache_info parameter.
  • Added _get_prefill_todo_workload method to calculate pending prefill workload.
  • Updated get_request_stats to include an optional urls parameter and the prefill_todo_workload in the returned statistics.

Activity

• The author, chickeyton, initially noted that this PR depends on a pending LMCache FullLookup feature and cannot be merged until that dependency is resolved.
• A critical bug was identified where a ValueError was created but not raised in TtftRouter; this was subsequently fixed by chickeyton.
• A typo in the RequestStatsCacheInfo constructor (__int__ instead of __init__) was found and corrected by chickeyton.
• The parameter instance_id_to_ip in KvawareRouter was renamed to instance_id_to_url for improved clarity and consistency, a change implemented by chickeyton.
• A typo in lmcache_contorller_port was pointed out and fixed by chickeyton.
• Hardcoded transfer times in TtftRouter were replaced with named constants (CACHE_LOC_TO_TRANS_TIME, DEFAULT_CACHE_TRANS_TIME) for better readability and configurability, as suggested and implemented by chickeyton.
• The TtftRouter.route_request method was modified to explicitly return RequestStatsCacheInfo along with the server URL, making the data flow more explicit.
• Debugging print statements were removed from the code.
• The parsing logic for lmcache_instances was made more robust to handle extra whitespace and empty strings, ensuring correct instance ID extraction.
• The type hint for TtftRouter.route_request was corrected to accurately reflect its tuple return type of (str, RequestStatsCacheInfo).
• A discussion occurred between chickeyton and the bot regarding the time complexity of the TimePeriods.union method, with chickeyton asserting the current O(N) complexity while the bot suggested an O(N log N) sort-and-merge approach for improved robustness and readability, acknowledging potential practical performance similarities for small N.

[chickeyton]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new TTFT (Time To First Token) based routing logic for the vLLM router. The goal is to improve request scheduling by estimating the prefill workload on each instance and routing to the one with the lowest expected TTFT. The changes are comprehensive, touching the routing logic, statistics monitoring, command-line argument parsing, and benchmarking tools. The overall design is sound and the implementation is mostly solid. However, I've identified a few critical bugs that could lead to runtime errors or incorrect behavior, along with several suggestions for improving code quality and robustness. Please address the critical issues before merging.

[gemini-code-assist]

This line calls self.query_manager, but this method was removed from KvawareRouter in this same pull request. This will cause an AttributeError at runtime. You should call self.kv_manager.handle_orchestration_message directly, as you've done in other parts of the code.

                    endpoint_instance_id = await self.kv_manager.handle_orchestration_message(query_message)

[gemini-code-assist]

This logic is flawed. If best_matched_info is True, the function will fall through the if block and implicitly return None because there is no return statement. This is a critical bug. Additionally, if best_matched_info is False, num_cached_tokens is discarded, leading to incorrect stats. The logic should be simplified to always set num_cached_tokens on cache_info and return the result.

            cache_info.num_cached_tokens = num_cached_tokens
            return best_inst_url, cache_info

[gemini-code-assist]

The stats object can be None if request_stats.get(url, None) returns None (e.g., for a new or idle endpoint). Accessing stats.prefill_todo_workload will raise an AttributeError. You should handle the case where stats is None, for example by assuming a pending workload of 0.

        workload = ((stats.prefill_todo_workload if stats else 0) +
                    prefill_workload(num_prefix_tokens, num_cache_tokens))

[gemini-code-assist]

This initialization can be made more concise by using the or operator.

        self.instance_id_to_url = instance_id_to_url or {}

[gemini-code-assist]

This initialization can be made more concise by using the or operator. This same pattern is used in KvawareRouter and could be updated there as well.

        self.instance_id_to_url = instance_id_to_url or {}

[gemini-code-assist]

The attribute self.cached_prefix_tokens is initialized here but it's never used anywhere in the class. It appears to be dead code and can be removed to improve clarity.

[gemini-code-assist]

This logic to extract the IP address from the URL is complex and can be fragile. Using urllib.parse.urlparse would be a more robust and readable way to get the hostname. You'll need to add from urllib.parse import urlparse at the top of the file. This same improvement can be applied to KvawareRouter as well.

                ip=urlparse(endpoint.url).hostname

[gemini-code-assist]

This class can be simplified by using the @dataclass decorator, which automatically generates methods like __init__ and __repr__. It would make the code more concise and idiomatic. The dataclasses module is already imported in this file.

Suggested change

	class RequestStatsCacheInfo:
	"""
	Cache information.
	"""
	def __init__(self):
	self.num_prefix_tokens : int = 0
	self.num_cached_tokens : int = 0
	@dataclass
	class RequestStatsCacheInfo:
	"""
	Cache information.
	"""
	num_prefix_tokens: int = 0
	num_cached_tokens: int = 0

[chickeyton]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new TTFT-based routing logic, which is a significant feature enhancement. The goal is to improve request routing by estimating prefill workload and selecting the instance with the lowest load, which should lead to better TTFT. The implementation is comprehensive, with changes across routing, statistics monitoring, and application configuration. The code is well-structured, and the inclusion of benchmark results is very helpful. I've identified a couple of areas for minor improvements related to code clarity and performance optimization.

[gemini-code-assist]

This check for best_matched_info being None appears to be unreachable. If matched_infos is empty, the function returns None at the beginning. If it's not empty, best_matched_info is guaranteed to be assigned a value in the first iteration of the loop. Removing this check would improve code clarity.

[gemini-code-assist]

The loop iterates through all endpoints to find the URL for a given instance_id, even after a match is found. This can be inefficient, especially with a large number of endpoints. Consider adding a break statement after finding the matching instance to avoid unnecessary queries.

Suggested change

	for endpoint in endpoints:
	msg = QueryInstMsg(
	event_id="",
	ip=urlparse(endpoint.url).hostname
	)
	ret_msg = await self.kv_manager.handle_orchestration_message(msg)
	self.instance_id_to_url[ret_msg.instance_id] = endpoint.url
	if ret_msg.instance_id == instance_id:
	url = endpoint.url
	for endpoint in endpoints:
	msg = QueryInstMsg(
	event_id="",
	ip=urlparse(endpoint.url).hostname
	)
	ret_msg = await self.kv_manager.handle_orchestration_message(msg)
	self.instance_id_to_url[ret_msg.instance_id] = endpoint.url
	if ret_msg.instance_id == instance_id:
	url = endpoint.url
	break