Repository: vllm-project/production-stack
Pull request #432: [Feat] Prefix aware routing implementation based on hash trie
Author: login: KuntaiDu
Created: 2025-05-15T05:49:00Z | Merged: 2025-05-16T16:04:53Z

This PR implements prefix-cache-aware routing based on hash trie.

**BEFORE SUBMITTING, PLEASE READ THE CHECKLIST BELOW AND FILL IN THE DESCRIPTION ABOVE**

---

- [x] Make sure the code changes pass the [pre-commit](../CONTRIBUTING.md) checks.
- [x] Sign-off your commit by using <code>-s</code> when doing <code>git commit</code>
- [x] Try to classify PRs for easy understanding of the type of changes, such as `[Bugfix]`, `[Feat]`, and `[CI]`.

<details>
<!-- inside this <details> section, markdown rendering does not work, so we use raw html here. -->
<summary><b> Detailed Checklist (Click to Expand) </b></summary>

<p>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.</p>

<h3>PR Title and Classification</h3>
<p>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:</p>
<ul>
    <li><code>[Bugfix]</code> for bug fixes.</li>
    <li><code>[CI/Build]</code> for build or continuous integration improvements.</li>
    <li><code>[Doc]</code> for documentation fixes and improvements.</li>
    <li><code>[Feat]</code> for new features in the cluster (e.g., autoscaling, disaggregated prefill, etc.).</li>
    <li><code>[Router]</code> for changes to the <code>vllm_router</code> (e.g., routing algorithm, router observability, etc.).</li>
    <li><code>[Misc]</code> for PRs that do not fit the above categories. Please use this sparingly.</li>
</ul>
<p><strong>Note:</strong> If the PR spans more than one category, please include all relevant prefixes.</p>

<h3>Code Quality</h3>

<p>The PR need to meet the following code quality standards:</p>

<ul>
    <li>Pass all linter checks. Please use <code>pre-commit</code> to format your code. See <code>README.md</code> for installation.</li>
    <li>The code need to be well-documented to ensure future contributors can easily understand the code.</li>
    <li> Please include sufficient tests to ensure the change is stay correct and robust. This includes both unit tests and integration tests.</li>
</ul>

<h3>DCO and Signed-off-by</h3>
<p>When contributing changes to this project, you must agree to the <a href="https://github.com/vllm-project/vllm/blob/main/DCO">DCO</a>. Commits must include a <code>Signed-off-by:</code> header which certifies agreement with the terms of the DCO.</p>
<p>Using <code>-s</code> with <code>git commit</code> will automatically add this header.</p>

<h3>What to Expect for the Reviews</h3>

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.



Diff:
diff --git a/pyproject.toml b/pyproject.toml
index cf9e4d70e..17565a6bc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: Apache Software License",
 ]
-# Should be the same as src/router/requirements.txt
+# Should be the same as src/vllm_router/requirements.txt
 dependencies = [
     "aiofiles==24.1.0",
     "fastapi==0.115.8",
@@ -22,6 +22,7 @@ dependencies = [
     "sentry-sdk[fastapi,httpx]==2.27.0",
     "uhashring==2.3",
     "uvicorn==0.34.0",
+    "xxhash==3.5.0",
 ]
 
 [project.scripts]
diff --git a/src/vllm_router/parsers/parser.py b/src/vllm_router/parsers/parser.py
index 2d53d64d1..3cd1cd75b 100644
--- a/src/vllm_router/parsers/parser.py
+++ b/src/vllm_router/parsers/parser.py
@@ -144,7 +144,7 @@ def parse_args():
     parser.add_argument(
         "--routing-logic",
         type=str,
-        choices=["roundrobin", "session", "kvaware"],
+        choices=["roundrobin", "session", "kvaware", "prefixaware"],
         help="The routing logic to use",
     )
     parser.add_argument(
diff --git a/src/vllm_router/prefix/hashtrie.py b/src/vllm_router/prefix/hashtrie.py
new file mode 100644
index 000000000..f7763bf19
--- /dev/null
+++ b/src/vllm_router/prefix/hashtrie.py
@@ -0,0 +1,104 @@
+# Copyright 2024-2025 The vLLM Production Stack Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+import random
+from typing import Any, Dict, Generator, Set, Tuple
+
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.endpoints = set()
+
+        # assign a lock for each trie node.
+        # this assures that each node will only be accessed by one co-routine
+        # at a time.
+        self.lock = asyncio.Lock()
+
+
+class HashTrie:
+    def __init__(self, chunk_size: int = 128):
+        """
+        Initialize the HashTrie.
+        Args:
+            chunk_size (int): the string chunk size (in terms of # characters)
+        """
+        self.root = TrieNode()
+        self.chunk_size = chunk_size
+
+    def _chunk_and_hash(self, request: str) -> Generator[int, None, None]:
+        """
+        Chunk and hash the request.
+        Args:
+            request (str): The request to chunk and hash.
+        Returns:
+            Generator[int, None, None]: A generator that yields a hash for each
+            chunk.
+        """
+
+        for i in range(0, len(request), self.chunk_size):
+            yield xxhash.xxh64(request[i : i + self.chunk_size]).intdigest()
+
+    async def insert(self, request: str, endpoint: str) -> None:
+        """
+        Insert the request and endpoint into the trie.
+        Args:
+            request (str): The request to insert.
+            endpoint (str): The endpoint to insert.
+        """
+        node = self.root
+        node.endpoints.add(endpoint)
+        for chunk_hash in self._chunk_and_hash(request):
+            async with node.lock:
+                if chunk_hash not in node.children:
+                    node.children[chunk_hash] = TrieNode()
+                node = node.children[chunk_hash]
+            node.endpoints.add(endpoint)
+
+    async def longest_prefix_match(
+        self, request: str, available_endpoints: Set[str] = set()
+    ) -> Tuple[int, Set[str]]:
+        """
+        Find the longest matching prefix using hashed chunks.
+        Args:
+            request (str): The request to find the longest matching prefix.
+            available_endpoints (Set[str]): The endpoints that are available.
+        """
+        node = self.root
+        match_length = 0
+        chunk_hashes = self._chunk_and_hash(request)
+        selected_endpoints = available_endpoints
+
+        for i, chunk_hash in enumerate(chunk_hashes):
+            async with node.lock:
+                if chunk_hash in node.children:
+
+                    node = node.children[chunk_hash]
+
+                    # reached longest prefix match in currently-available endpoints.
+                    if not node.endpoints.intersection(selected_endpoints):
+                        break
+
+                    match_length += self.chunk_size
+                    selected_endpoints = node.endpoints.intersection(selected_endpoints)
+                else:
+                    break
+
+        return match_length, selected_endpoints
diff --git a/src/vllm_router/requirements.txt b/src/vllm_router/requirements.txt
index fc9e00a21..7f7dd85e5 100644
--- a/src/vllm_router/requirements.txt
+++ b/src/vllm_router/requirements.txt
@@ -8,3 +8,4 @@ python-multipart==0.0.20
 sentry-sdk[fastapi,httpx]==2.27.0
 uhashring==2.3
 uvicorn==0.34.0
+xxhash==3.5.0
diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py
index 35308814d..2ee608232 100644
--- a/src/vllm_router/routers/routing_logic.py
+++ b/src/vllm_router/routers/routing_logic.py
@@ -15,6 +15,7 @@
 import abc
 import asyncio
 import enum
+import random
 import socket
 import threading
 from typing import Dict, List
@@ -46,6 +47,7 @@ class RoutingLogic(str, enum.Enum):
     ROUND_ROBIN = "roundrobin"
     SESSION_BASED = "session"
     KVAWARE = "kvaware"
+    PREFIXAWARE = "prefixaware"
 
 
 class RoutingInterface(metaclass=SingletonABCMeta):
@@ -285,6 +287,59 @@ async def route_request(
             return self.instance_id_to_ip[instance_id.best_instance_id]
 
 
+class PrefixAwareRouter(RoutingInterface):
+    """
+    Route the request to the appropriate engine URL by where the longest
+    prefix match is found.
+
+    In this class, we assume that there is no eviction of prefix cache.
+    """
+
+    def __init__(self: int):
+        if hasattr(self, "_initialized"):
+            return
+        from vllm_router.prefix.hashtrie import HashTrie
+
+        self.hashtrie = HashTrie()
+        self._initialized = True
+
+    async def route_request(
+        self,
+        endpoints: List[EndpointInfo],
+        engine_stats: Dict[str, EngineStats],
+        request_stats: Dict[str, RequestStats],
+        request: Request,
+        request_json: Dict,
+    ) -> str:
+        """
+        Route the request to the appropriate engine URL by where the longest
+        prefix match is found.
+
+        In this routing logic, we do not consider the eviction of prefix cache.
+
+        Args:
+            endpoints (List[EndpointInfo]): The list of engine URLs
+            engine_stats (Dict[str, EngineStats]): The engine stats indicating
+               the 'physical' load of each engine
+            request_stats (Dict[str, RequestStats]): The request stats
+               indicating the request-level performance of each engine
+            request (Request): The incoming request
+            request_json (Dict): The request body (needed for finding the
+            longest prefix match)
+        """
+
+        available_endpoints = set(endpoint.url for endpoint in endpoints)
+        _, matched_endpoint = await self.hashtrie.longest_prefix_match(
+            request_json["prompt"], available_endpoints
+        )
+
+        selected_endpoint = random.choice(list(matched_endpoint))
+
+        await self.hashtrie.insert(request_json["prompt"], selected_endpoint)
+
+        return selected_endpoint
+
+
 # Instead of managing a global _global_router, we can define the initialization functions as:
 def initialize_routing_logic(
     routing_logic: RoutingLogic, *args, **kwargs
@@ -300,6 +355,9 @@ def initialize_routing_logic(
         router = KvawareRouter(kwargs.get("lmcache_controller_port"))
         router.start_kv_manager()
         return router
+    elif routing_logic == RoutingLogic.PREFIXAWARE:
+        logger.info("Initializing prefix-aware routing logic")
+        return PrefixAwareRouter()
     else:
         raise ValueError(f"Invalid routing logic {routing_logic}")
 
@@ -316,7 +374,12 @@ def reconfigure_routing_logic(
 
 def get_routing_logic() -> RoutingInterface:
     # Look up in our singleton registry which router (if any) has been created.
-    for cls in (SessionRouter, RoundRobinRouter, KvawareRouter):
+    for cls in (
+        SessionRouter,
+        RoundRobinRouter,
+        KvawareRouter,
+        PrefixAwareRouter,
+    ):
         if cls in SingletonABCMeta._instances:
             return cls()
     raise ValueError("The global router has not been initialized")
diff --git a/src/vllm_router/services/request_service/request.py b/src/vllm_router/services/request_service/request.py
index c80ae7dc6..09149a349 100644
--- a/src/vllm_router/services/request_service/request.py
+++ b/src/vllm_router/services/request_service/request.py
@@ -22,7 +22,7 @@
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from vllm_router.log import init_logger
-from vllm_router.routers.routing_logic import KvawareRouter
+from vllm_router.routers.routing_logic import KvawareRouter, PrefixAwareRouter
 from vllm_router.service_discovery import get_service_discovery
 from vllm_router.services.request_service.rewriter import (
     get_request_rewriter,
@@ -216,7 +216,9 @@ async def route_general_request(
         )
 
     logger.debug(f"Routing request {request_id} for model: {requested_model}")
-    if isinstance(request.app.state.router, KvawareRouter):
+    if isinstance(request.app.state.router, KvawareRouter) or isinstance(
+        request.app.state.router, PrefixAwareRouter
+    ):
         server_url = await request.app.state.router.route_request(
             endpoints, engine_stats, request_stats, request, request_json
         )


Reviews:
user: login: Shaoting-Feng
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-05-15T23:00:03Z
user: login: KuntaiDu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-16T07:11:35Z
user: login: KuntaiDu
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-16T07:11:39Z
user: login: YuhanLiu11
body: Thanks for fixing the comments! 
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-16T15:59:07Z
user: login: Shaoting-Feng
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-16T16:00:43Z

Review comments:
diff_hunk: @@ -0,0 +1,98 @@
+# Copyright 2024-2025 The vLLM Production Stack Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+from collections import defaultdict
+from typing import Any, Dict, Generator, Set, Tuple
+
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end = False
+        self.endpoints = set()
+
+
+class HashTrie:
+    def __init__(self, chunk_size: int = 128):
+        """
+        Initialize the HashTrie.
+        Args:
+            chunk_size (int): the string chunk size (in terms of # characters)
+        """
+        self.root = TrieNode()
+        self.chunk_size = chunk_size
+
+    def _chunk_and_hash(self, request: str) -> Generator[int, None, None]:
+        """
+        Chunk and hash the request.
+        Args:
+            request (str): The request to chunk and hash.
+        Returns:
+            Generator[int, None, None]: A generator that yields hashes of 32-character chunks.
path: src/vllm_router/prefix/hashtrie.py
user: login: Shaoting-Feng
body: hardcoded 32-character in comment
created_at: 2025-05-15T22:45:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,98 @@
+# Copyright 2024-2025 The vLLM Production Stack Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+from collections import defaultdict
+from typing import Any, Dict, Generator, Set, Tuple
+
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end = False
+        self.endpoints = set()
+
+
+class HashTrie:
+    def __init__(self, chunk_size: int = 128):
path: src/vllm_router/prefix/hashtrie.py
user: login: Shaoting-Feng
body: Do we need to protect `insert` and `longest_prefix_match` with a mutex?
created_at: 2025-05-15T22:59:51Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,98 @@
+# Copyright 2024-2025 The vLLM Production Stack Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+from collections import defaultdict
+from typing import Any, Dict, Generator, Set, Tuple
+
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end = False
+        self.endpoints = set()
+
+
+class HashTrie:
+    def __init__(self, chunk_size: int = 128):
path: src/vllm_router/prefix/hashtrie.py
user: login: KuntaiDu
body: fixed.
created_at: 2025-05-16T07:11:35Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,98 @@
+# Copyright 2024-2025 The vLLM Production Stack Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import random
+from collections import defaultdict
+from typing import Any, Dict, Generator, Set, Tuple
+
+import xxhash
+
+logger = logging.getLogger(__name__)
+
+
+class TrieNode:
+    def __init__(self):
+        self.children = {}
+        self.is_end = False
+        self.endpoints = set()
+
+
+class HashTrie:
+    def __init__(self, chunk_size: int = 128):
+        """
+        Initialize the HashTrie.
+        Args:
+            chunk_size (int): the string chunk size (in terms of # characters)
+        """
+        self.root = TrieNode()
+        self.chunk_size = chunk_size
+
+    def _chunk_and_hash(self, request: str) -> Generator[int, None, None]:
+        """
+        Chunk and hash the request.
+        Args:
+            request (str): The request to chunk and hash.
+        Returns:
+            Generator[int, None, None]: A generator that yields hashes of 32-character chunks.
path: src/vllm_router/prefix/hashtrie.py
user: login: KuntaiDu
body: fixed.
created_at: 2025-05-16T07:11:39Z
author_association: COLLABORATOR

PR comments:
user: login: KuntaiDu
created_at: 2025-05-16T07:18:32Z
body: @Shaoting-Feng @YuhanLiu11 Fixed the comments. PTAL.
author_association: COLLABORATOR
user: login: KuntaiDu
created_at: 2025-05-16T07:18:57Z
body: I also did a small test and made sure that when all requests share the same prefix they will only be routed to one instance.
author_association: COLLABORATOR
user: login: YuhanLiu11
created_at: 2025-05-16T15:58:16Z
body: > I also did a small test and made sure that when all requests share the same prefix they will only be routed to one instance.

That's great! 
author_association: COLLABORATOR