feat: Unnormalize waiting requests + predictive load updates for Python router (mirroring Rust) + softmax sampling to reduce thrashing (#1638)

Author: PeaBrane

examples/llm/components/kv_router.py

@@ -20,9 +20,10 @@
 from argparse import Namespace
 from typing import AsyncIterator, Tuple
 
+import numpy as np  # Add numpy import
 from components.worker import VllmWorker
 from utils.check_worker import check_required_workers
-from utils.protocol import Tokens
+from utils.protocol import LocalBlockHashes
 from utils.vllm import RouterType
 
 from dynamo.llm import AggregatedMetrics, KvIndexer, KvMetricsAggregator, OverlapScores
@@ -35,20 +36,49 @@
 logger = logging.getLogger(__name__)
 
 
+def softmax_sample_from_logits(
+    logits: dict[str, float], temperature: float = 1.0, lower_is_better: bool = True
+) -> str:
+    if not logits:
+        raise ValueError("Empty logits dictionary")
+
+    keys = list(logits.keys())
+    values = np.array(list(logits.values()))
+
+    min_val = np.min(values)
+    max_val = np.max(values)
+
+    if min_val == max_val:
+        # All values are the same, uniform probability
+        probabilities = np.ones(len(keys)) / len(keys)
+    else:
+        normalized = values / (max_val - min_val)
+        if lower_is_better:
+            normalized = -1 * normalized
+
+        scaled = normalized / temperature
+
+        exp_values = np.exp(scaled - np.max(scaled))
+        probabilities = exp_values / np.sum(exp_values)
+
+    # Sample from the probability distribution
+    return np.random.choice(keys, p=probabilities)
+
+
 def parse_args(service_name, prefix) -> Namespace:
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--min-workers",
-        type=int,
-        default=1,
-        help="Minimum number of workers required before proceeding",
-    )
     parser.add_argument(
         "--model",
         type=str,
         default="deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
         help="Model that is being served",
     )
+    parser.add_argument(
+        "--min-workers",
+        type=int,
+        default=1,
+        help="Minimum number of workers required before proceeding",
+    )
     # TODO: Read block size
     parser.add_argument(
         "--block-size",
@@ -68,6 +98,12 @@ def parse_args(service_name, prefix) -> Namespace:
         default="kv",
         help="The router type",
     )
+    parser.add_argument(
+        "--softmax-sample",
+        type=bool,
+        default=False,
+        help="Whether to do softmax sampling based on worker logits (default is to pick smallest)",
+    )
     config = ServiceConfig.get_instance()
     config_args = config.as_args(service_name, prefix=prefix)
     args = parser.parse_args(config_args)
@@ -93,8 +129,10 @@ def __init__(self):
         self.args = parse_args(self.__class__.__name__, "")
 
         self.default_metrics = {
-            "gpu_cache_usage_perc": 0.0,
+            "kv_active_blocks": 0,
+            "kv_total_blocks": 1,
             "num_requests_waiting": 0.0,
+            "gpu_cache_usage_perc": 0.0,
             "gpu_prefix_cache_hit_rate": 0.0,
         }
 
@@ -117,8 +155,36 @@ async def async_init(self):
         if self.router_type == RouterType.KV:
             self.indexer = KvIndexer(kv_listener, self.args.block_size)
         self.metrics_aggregator = KvMetricsAggregator(kv_listener)
+
+        self.active_blocks_dict = {}
+        worker_ids = self.workers_client.instance_ids()
+        for worker_id in worker_ids:
+            # [old_value, predictive_value]
+            self.active_blocks_dict[worker_id] = [0, 0]
+
         logger.info("KV Router initialized")
 
+    def _update_and_get_active_blocks(self, worker_id: str, polled_value: int) -> int:
+        """Helper routine to update waiting dict and return the desired waiting value.
+
+        This method implements a predictive mechanism for tracking waiting requests:
+        - If a new polled value is detected (different from the stored old value),
+          it updates both the old and predictive values to this new measurement and returns it
+        - If no change is detected (polled value equals old value), it returns the
+          predictive value which has been incremented based on previous routing decisions
+
+        This allows the router to account for requests that have been dispatched but
+        not yet reflected in the polled metrics.
+        """
+        old_value, predictive_value = self.active_blocks_dict[worker_id]
+
+        # Check if polled value is different from old value
+        if polled_value != old_value:
+            self.active_blocks_dict[worker_id] = [polled_value, polled_value]
+            return polled_value
+        else:
+            return predictive_value
+
     def _cost_function(
         self,
         scores: OverlapScores | None,
@@ -142,66 +208,79 @@ def _cost_function(
             (str, float): The best worker id and the corresponding score.
         """
 
-        worker_scores = {}
+        # Get all worker IDs from the client. This is needed because scores / metrics may not have values for all workers
+        # and we want all workers to be considered in the logit calculation
+        worker_ids = self.workers_client.instance_ids()
+        request_blocks = (
+            token_length + self.args.block_size - 1
+        ) // self.args.block_size
+
+        overlap_blocks_dict = {worker_id: 0 for worker_id in worker_ids}
+        new_blocks_dict = {worker_id: request_blocks for worker_id in worker_ids}
+
         if scores:
             for worker_id, score in scores.scores.items():
                 # score is number of matching blocks we multiply by block_size to get tokens
                 # and compare to token_length. The larger the cache hit the better
-                worker_scores[worker_id] = (
-                    score * self.indexer.block_size() / token_length
-                )
+                overlap_blocks_dict[worker_id] = score
+                new_blocks_dict[worker_id] = request_blocks - score
         else:
             logger.warning("Cannot get KV scores")
 
         worker_metrics = {}
-        max_waiting = 0.0
         if metrics:
             for endpoint in metrics.endpoints:
                 worker_id = endpoint.worker_id
                 worker_metrics[worker_id] = {
                     key: getattr(endpoint, key, self.default_metrics[key])
                     for key in self.default_metrics.keys()
                 }
-                max_waiting = max(
-                    max_waiting, worker_metrics[worker_id]["num_requests_waiting"]
+
+                # Update waiting value using helper routine
+                polled_active_blocks = int(
+                    worker_metrics[worker_id]["kv_active_blocks"]
                 )
+                worker_metrics[worker_id][
+                    "kv_active_blocks"
+                ] = self._update_and_get_active_blocks(worker_id, polled_active_blocks)
         else:
             logger.warning("Cannot get metrics")
 
-        # Get all worker IDs from the client. This is needed because scores / metrics may not have values for all workers
-        # and we want all workers to be considered in the logit calculation
-        worker_ids = self.workers_client.instance_ids()
-
         worker_logits = {}
         for worker_id in worker_ids:
             # Use default values if worker not in scores or metrics
-            score = worker_scores.get(worker_id, 0.0)
             metrics_dict = worker_metrics.get(worker_id, self.default_metrics)
-            gpu_cache_usage = metrics_dict["gpu_cache_usage_perc"]
+            kv_total_blocks = metrics_dict["kv_total_blocks"]
 
-            normalized_waiting = (
-                metrics_dict["num_requests_waiting"] / max_waiting
-                if max_waiting > 0
-                else 0.0
-            )
+            new_blocks = new_blocks_dict[worker_id]
+            normalized_new_blocks = new_blocks / kv_total_blocks
+            gpu_cache_usage = metrics_dict["kv_active_blocks"] / kv_total_blocks
+
+            # Use raw waiting value without normalization
+            num_requests_waiting = metrics_dict["num_requests_waiting"]
 
             # Have 1 metric that weights towards cache hit
             # 2 metrics that penalize overloaded worker and queuing
-            worker_logits[worker_id] = 2 * score - gpu_cache_usage - normalized_waiting
+            worker_logits[worker_id] = (
+                normalized_new_blocks + gpu_cache_usage + num_requests_waiting
+            )
             logger.info(
-                f"Formula for {worker_id}: {worker_logits[worker_id]:.3f} = 2.0 * {score:.3f} - {gpu_cache_usage:.3f} - {normalized_waiting:.3f}"
+                f"Formula for {worker_id}: {worker_logits[worker_id]:.3f} = {normalized_new_blocks:.3f} + {gpu_cache_usage:.3f} + {num_requests_waiting:.3f}"
             )
 
         if not worker_logits or not any(worker_logits.values()):
             logger.warning(f"All worker logits are zero. {fallback_msg}.")
             return "", 0.0
 
         # Select the worker with the highest logit
-        max_logit = max(worker_logits.values())
-        best_workers = [
-            wid for wid, logit in worker_logits.items() if logit == max_logit
-        ]
-        best_worker_id = random.choice(best_workers)
+        if self.args.softmax_sample:
+            best_worker_id = int(softmax_sample_from_logits(worker_logits))
+        else:
+            min_logit = min(worker_logits.values())
+            best_workers = [
+                wid for wid, logit in worker_logits.items() if logit == min_logit
+            ]
+            best_worker_id = random.choice(best_workers)
 
         # Log the metrics for the selected worker
         if best_worker_id:
@@ -212,15 +291,23 @@ def _cost_function(
                 f"Selected worker: {best_worker_id}, logit: {worker_logits[best_worker_id]:.3f}",
                 f"Score: {scores.scores.get(best_worker_id, 0.0) if scores else 0.0:.3f}",
                 f"GPU Cache Hit Rate: {metrics_dict['gpu_prefix_cache_hit_rate']:.3f}",
-                f"GPU Cache Usage: {metrics_dict['gpu_cache_usage_perc']:.3f}",
+                f"GPU Cache Usage: {metrics_dict['kv_active_blocks'] / metrics_dict['kv_total_blocks']:.3f}",
                 f"Requests Waiting: {metrics_dict['num_requests_waiting']}",
             ]
 
             # Log to vllm_logger
             for message in log_messages:
                 logger.info(message)
 
-        return best_worker_id, worker_scores.get(best_worker_id, 0.0)
+            # Increment predictive waiting for the selected worker before returning
+            self.active_blocks_dict[best_worker_id][1] += new_blocks_dict[
+                best_worker_id
+            ]
+
+        return (
+            best_worker_id,
+            overlap_blocks_dict[best_worker_id] * self.args.block_size / token_length,
+        )
 
     def _get_underloaded_worker(self, metrics: AggregatedMetrics | None):
         if not metrics:
@@ -248,7 +335,9 @@ def _get_underloaded_worker(self, metrics: AggregatedMetrics | None):
         return best_worker_id, kv_load[best_worker_id]
 
     @endpoint()
-    async def generate(self, request: Tokens) -> AsyncIterator[Tuple[WorkerId, float]]:
+    async def generate(
+        self, request: LocalBlockHashes
+    ) -> AsyncIterator[Tuple[WorkerId, float]]:
         metrics = await self.metrics_aggregator.get_metrics()
 
         # Quick return for KV_LOAD mode
@@ -263,19 +352,16 @@ async def generate(self, request: Tokens) -> AsyncIterator[Tuple[WorkerId, float
             return
 
         # Existing KV routing logic
-        lora_id = 0
         try:
-            scores = await self.indexer.find_matches_for_request(
-                request.tokens, lora_id
-            )
+            scores = await self.indexer.find_matches(request.hashes)
         except Exception as e:
             scores = {}
             logger.exception(f"Error finding matches: {e}. {fallback_msg}")
             yield "", 0.0
             return
 
         worker_id, prefix_hit_rate = self._cost_function(
-            scores, metrics, len(request.tokens)
+            scores, metrics, request.num_tokens
         )
 
         if worker_id:

examples/llm/components/processor.py

@@ -24,14 +24,14 @@
 from transformers import AutoTokenizer
 from utils.chat_processor import ChatProcessor, CompletionsProcessor, ProcessMixIn
 from utils.check_worker import check_required_workers
-from utils.protocol import MyRequestOutput, Tokens, vLLMGenerateRequest
+from utils.protocol import LocalBlockHashes, MyRequestOutput, vLLMGenerateRequest
 from utils.vllm import RouterType, parse_vllm_args
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest, CompletionRequest
 from vllm.outputs import RequestOutput
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 
-from dynamo.llm import KvMetricsAggregator
+from dynamo.llm import KvMetricsAggregator, compute_block_hash_for_seq_py
 from dynamo.runtime import EtcdKvCache
 from dynamo.sdk import async_on_start, depends, dynamo_context, endpoint, service
 
@@ -242,9 +242,13 @@ async def process_and_stream():
 
                 prefix_hit_rate = 0.0  # Default value
                 if self.use_router:
+                    token_ids = engine_prompt["prompt_token_ids"]
                     router_generator = await self.router_client.generate(
-                        Tokens(
-                            tokens=engine_prompt["prompt_token_ids"]
+                        LocalBlockHashes(
+                            hashes=compute_block_hash_for_seq_py(
+                                token_ids, self.engine_args.block_size
+                            ),
+                            num_tokens=len(token_ids),
                         ).model_dump_json()
                     )
                     decision = await router_generator.__anext__()

examples/llm/configs/agg_router.yaml

@@ -29,6 +29,7 @@ Processor:
 
 Router:
   min-workers: 1
+  softmax_sample: true
   common-configs: [model, block-size, router]
 
 VllmWorker:

examples/llm/utils/protocol.py

@@ -36,6 +36,11 @@ class Tokens(BaseModel):
     tokens: list[int]
 
 
+class LocalBlockHashes(BaseModel):
+    hashes: list[int]
+    num_tokens: int
+
+
 class PrefillRequest(Request):
     request_id: str
 

lib/bindings/python/rust/llm/kv.rs

@@ -481,6 +481,24 @@ impl KvIndexer {
         self.inner.block_size()
     }
 
+    fn find_matches<'p>(&self, py: Python<'p>, sequence: Vec<u64>) -> PyResult<Bound<'p, PyAny>> {
+        let indexer = self.inner.clone();
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            let local_block_hashes: Vec<llm_rs::kv_router::protocols::LocalBlockHash> = sequence
+                .into_iter()
+                .map(llm_rs::kv_router::protocols::LocalBlockHash)
+                .collect();
+
+            let rs_overlap_scores = indexer
+                .find_matches(local_block_hashes)
+                .await
+                .map_err(to_pyerr)?;
+            Ok(OverlapScores {
+                inner: rs_overlap_scores,
+            })
+        })
+    }
+
     fn find_matches_for_request<'p>(
         &self,
         py: Python<'p>,

lib/bindings/python/src/dynamo/_core.pyi

@@ -527,6 +527,18 @@ class KvIndexer:
         Create a `KvIndexer` object
         """
 
+    def find_matches(self, sequence: List[int]) -> OverlapScores:
+        """
+        Find prefix matches for the given sequence of block hashes.
+
+        Args:
+            sequence: List of block hashes to find matches for
+
+        Returns:
+            OverlapScores containing worker matching scores and frequencies
+        """
+        ...
+
     def find_matches_for_request(
         self, token_ids: List[int], lora_id: int
     ) -> OverlapScores:

lib/llm/src/kv_router.rs

@@ -73,7 +73,7 @@ pub struct KvRouterConfig {
 impl Default for KvRouterConfig {
     fn default() -> Self {
         Self {
-            overlap_score_weight: 2.0,
+            overlap_score_weight: 1.0,
             gpu_cache_usage_weight: 1.0,
             waiting_requests_weight: 1.0,
         }