[Router] Add prefix_min_match_length threshold to PrefixAwareRouter

Add a configurable --prefix-min-match-length option for the
prefixaware routing logic. When the longest prefix match is shorter
than this threshold, the request falls back to QPS-based routing
instead of using the matched endpoint, mitigating load hotspots
caused by long shared prefixes (e.g. common system prompts).

Defaults to 0, which preserves the original behavior.

Partially addresses #957.

Signed-off-by: Max Li <hitliqiwei@gmail.com>

Author: ywc668

src/tests/test_prefixaware_router.py

@@ -0,0 +1,137 @@
+from typing import Any, Dict
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_router.routers.routing_logic import (
+    PrefixAwareRouter,
+    cleanup_routing_logic,
+)
+
+
+@pytest.fixture(autouse=True)
+def cleanup_router():
+    cleanup_routing_logic()
+    yield
+    cleanup_routing_logic()
+
+
+class EndpointInfo:
+    def __init__(self, url: str):
+        self.url = url
+
+
+class RequestStats:
+    def __init__(self, qps: float):
+        self.qps = qps
+
+
+class Request:
+    def __init__(self, headers: Dict[str, str], body: Dict[str, Any] = None):
+        self.headers = headers
+        self.body = body
+
+
+@pytest.mark.asyncio
+async def test_route_falls_back_to_qps_when_match_below_threshold():
+    """
+    When the longest prefix match is shorter than prefix_min_match_length,
+    the request should NOT use the matched endpoint. It should fall back to
+    QPS-based routing and pick the engine with the lowest QPS.
+    """
+
+    endpoints = [
+        EndpointInfo(url="http://engine1.com"),
+        EndpointInfo(url="http://engine2.com"),
+    ]
+    request_stats = {
+        "http://engine1.com": RequestStats(qps=10),
+        "http://engine2.com": RequestStats(qps=5),
+    }
+    request = Request(headers={})
+    request_json = {"prompt": "some prompt text"}
+
+    router = PrefixAwareRouter(prefix_min_match_length=4096)
+
+    fake_hashtrie = AsyncMock()
+    fake_hashtrie.longest_prefix_match.return_value = (
+        100,
+        {"http://engine1.com"},
+    )
+    router.hashtrie = fake_hashtrie
+
+    url = await router.route_request(
+        endpoints, None, request_stats, request, request_json
+    )
+
+    assert url == "http://engine2.com"
+
+    fake_hashtrie.insert.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_route_uses_matched_endpoint_when_match_above_threshold():
+    """
+    When the longest prefix match is no shorter than prefix_min_match_length,
+    the request should use the matched endpoint.
+    """
+
+    endpoints = [
+        EndpointInfo(url="http://engine1.com"),
+        EndpointInfo(url="http://engine2.com"),
+    ]
+    request_stats = {}
+    request = Request(headers={})
+    request_json = {"prompt": "some prompt text"}
+
+    router = PrefixAwareRouter(prefix_min_match_length=50)
+
+    fake_hashtrie = AsyncMock()
+    fake_hashtrie.longest_prefix_match.return_value = (
+        5000,
+        {"http://engine1.com"},
+    )
+    router.hashtrie = fake_hashtrie
+
+    url = await router.route_request(
+        endpoints, None, request_stats, request, request_json
+    )
+
+    assert url == "http://engine1.com"
+
+    fake_hashtrie.insert.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_default_threshold_zero_preserves_original_behavior():
+    """
+    When --prefix-min-match-length is not provided, prefix_min_match_length
+    defaults to 0. Even when there is no prefix match at all (match_length 0),
+    the request still uses the matched endpoint instead of falling back,
+    preserving the original behavior before this change.
+    """
+
+    endpoints = [
+        EndpointInfo(url="http://engine1.com"),
+        EndpointInfo(url="http://engine2.com"),
+    ]
+    request_stats = {}
+    request = Request(headers={})
+    request_json = {"prompt": "some prompt text"}
+
+    router = PrefixAwareRouter()
+
+    fake_hashtrie = AsyncMock()
+    fake_hashtrie.longest_prefix_match.return_value = (
+        0,
+        {"http://engine1.com"},
+    )
+    router.hashtrie = fake_hashtrie
+
+    url = await router.route_request(
+        endpoints, None, request_stats, request, request_json
+    )
+
+    assert url == "http://engine1.com"
+
+    fake_hashtrie.insert.assert_awaited_once()

src/vllm_router/app.py

@@ -279,6 +279,7 @@ def initialize_all(app: FastAPI, args):
         prefill_model_labels=args.prefill_model_labels,
         decode_model_labels=args.decode_model_labels,
         kv_aware_threshold=args.kv_aware_threshold,
+        prefix_min_match_length=args.prefix_min_match_length,
         max_instance_failover_reroute_attempts=args.max_instance_failover_reroute_attempts,
         lmcache_health_check_interval=args.lmcache_health_check_interval,
         lmcache_worker_timeout=args.lmcache_worker_timeout,

src/vllm_router/parsers/parser.py

@@ -447,6 +447,16 @@ def parse_args():
         help="The threshold for kv-aware routing.",
     )
 
+    parser.add_argument(
+        "--prefix-min-match-length",
+        type=int,
+        default=0,
+        help="The minimum prefix match length, in characters, required for "
+        "prefixaware routing to reuse a matched endpoint. If the longest "
+        "prefix match is shorter than this value, the request falls back "
+        "to QPS-based routing. Defaults to 0, which disables the threshold.",
+    )
+
     parser.add_argument(
         "--max-instance-failover-reroute-attempts",
         type=int,

src/vllm_router/routers/routing_logic.py

@@ -436,12 +436,16 @@ class PrefixAwareRouter(RoutingInterface):
     In this class, we assume that there is no eviction of prefix cache.
     """
 
-    def __init__(self: int):
+    def __init__(
+        self,
+        prefix_min_match_length: int = 0,
+    ):
         if hasattr(self, "_initialized"):
             return
         from vllm_router.prefix.hashtrie import HashTrie
 
         self.hashtrie = HashTrie()
+        self.prefix_min_match_length = prefix_min_match_length
         self._initialized = True
 
     async def route_request(
@@ -496,10 +500,13 @@ async def route_request(
             prompt = request_json["prompt"]
 
         available_endpoints = set(endpoint.url for endpoint in endpoints)
-        _, matched_endpoint = await self.hashtrie.longest_prefix_match(
+        match_length, matched_endpoint = await self.hashtrie.longest_prefix_match(
             prompt, available_endpoints
         )
 
+        if match_length < self.prefix_min_match_length:
+            return self._qps_routing(endpoints, request_stats)
+
         selected_endpoint = random.choice(list(matched_endpoint))
 
         await self.hashtrie.insert(prompt, selected_endpoint)
@@ -684,7 +691,9 @@ def initialize_routing_logic(
         router.start_kv_manager()
     elif routing_logic == RoutingLogic.PREFIXAWARE:
         logger.info("Initializing prefix-aware routing logic")
-        router = PrefixAwareRouter()
+        router = PrefixAwareRouter(
+            prefix_min_match_length=kwargs.get("prefix_min_match_length", 0),
+        )
     elif routing_logic == RoutingLogic.DISAGGREGATED_PREFILL:
         logger.info("Initializing disaggregated prefill routing logic")
         router = DisaggregatedPrefillRouter(