[feature] add inflight to s_load makesure select best replica in star…

examples/kvc_aware_router/plot_metrics.py

@@ -348,6 +348,50 @@ def draw(self, ax, data, colors, order) -> None:
         if self.ylim_top is not None:
             ax.set_ylim(top=self.ylim_top)
         ax.grid(True, alpha=0.3)
+        self._annotate_stats(ax, data, order)
+
+    # -- statistics annotation --
+    def stats(self, data: dict, order: list) -> str:
+        """Panel-wide aggregate mean/std over all replicas' points.
+
+        Accumulating panels (cumulative totals — evict, cumul. completed)
+        override to report the last value instead, since mean/std of a
+        monotonically rising series is not informative.
+        """
+        vals = []
+        for rep in order:
+            for _t, v in self.points_for(data, rep):
+                if v == v:  # skip NaN (matplotlib skips them; stats should too)
+                    vals.append(float(v))
+        return self._mean_std_str(vals)
+
+    def _annotate_stats(self, ax, data, order) -> None:
+        """Draw the panel-wide μ/σ (or last) annotation in the top-right."""
+        txt = self.stats(data, order)
+        if not txt:
+            return
+        ax.text(
+            0.99,
+            0.95,
+            txt,
+            transform=ax.transAxes,
+            ha="right",
+            va="top",
+            fontsize=7,
+            color="#333333",
+            bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="#bbbbbb", alpha=0.75),
+        )
+
+    @staticmethod
+    def _mean_std_str(vals: list) -> str:
+        if not vals:
+            return ""
+        import statistics
+
+        m = statistics.mean(vals)
+        # population stdev — panels are full samples, not a drawn subset.
+        s = statistics.pstdev(vals) if len(vals) > 1 else 0.0
+        return f"μ={m:.3g} σ={s:.3g}"
 
     def summarize(self, data: dict, rep: str) -> str:
         return self._summary(self.points_for(data, rep)) if self._summary else ""
@@ -504,6 +548,22 @@ class CumulativePanel(Panel):
     :meth:`accumulate` for other cumulative semantics.
     """
 
+    def stats(self, data: dict, order: list) -> str:
+        """Cumulative panels report the final aggregate value, not μ/σ.
+
+        A monotonically rising series' mean/std (≈ mid/quarter-span) is not
+        informative; the last value is the realized total — what matters here.
+        Sums the last value across replicas (one evictions/completed total).
+        """
+        last_vals = []
+        for rep in order:
+            pts = self.points_for(data, rep)
+            if pts and pts[-1][1] == pts[-1][1]:
+                last_vals.append(float(pts[-1][1]))
+        if not last_vals:
+            return ""
+        return f"Σlast={sum(last_vals):.3g}"
+
     def derive(self, retained: dict):
         out = {}
         for rep, hist in retained.items():
@@ -754,6 +814,12 @@ def summarize(self, data, rep):
         pts = data.get(self.key, {}).get(self._GLOBAL, [])
         return self._summary(pts) if self._summary else ""
 
+    def stats(self, data, order) -> str:
+        # Global single-line series: aggregate over _GLOBAL, not per-replica.
+        pts = data.get(self.key, {}).get(self._GLOBAL, [])
+        vals = [float(v) for _t, v in pts if v == v]
+        return self._mean_std_str(vals)
+
 
 class StickyOverloadPanel(Panel):
     """Load of the sticky-bound replica evaluated in is_overloaded (one replica per dispatch)."""

tests/workers/rollout/router/strategies/test_route.py

@@ -100,7 +100,7 @@ class _DummyConfig:
             pass
 
         class _DummyStrategy:
-            def score(self, prompt_ids, provider, replicas, request_id=None, sticky_table=None):
+            def score(self, prompt_ids, provider, replicas, request_id=None, gpu_hash_strs=None):
                 return [0.0] * len(replicas)
 
         StrategyRegistry.register(_DummyConfig, _DummyStrategy)
@@ -149,17 +149,17 @@ class ConstantStrategy:
     def __init__(self, scores):
         self._scores = scores
 
-    def score(self, prompt_ids, provider, replicas, request_id=None, sticky_table=None):
+    def score(self, prompt_ids, provider, replicas, request_id=None, gpu_hash_strs=None):
         return list(self._scores)
 
 
 class BadLengthStrategy:
-    def score(self, prompt_ids, provider, replicas, request_id=None, sticky_table=None):
+    def score(self, prompt_ids, provider, replicas, request_id=None, gpu_hash_strs=None):
         return [1.0]
 
 
 class RaisingStrategy:
-    def score(self, prompt_ids, provider, replicas, request_id=None, sticky_table=None):
+    def score(self, prompt_ids, provider, replicas, request_id=None, gpu_hash_strs=None):
         raise KeyError("boom")
 
 

verl/trainer/config/rollout/router/kvcaware.yaml

@@ -32,7 +32,7 @@ strategies:
 # Collector connection-type tuning (HTTP polling + ZMQ long-connection).
 collector:
   http_polling:
-    polling_interval: 5
+    polling_interval: 0.05
     http_timeout: 10
   long_connection:
     base_retry_delay: 1.0

verl/workers/rollout/router/kvcaware/balancer.py

@@ -38,6 +38,8 @@
     StrategyRegistry,
     route,
 )
+from .types import Layer
+from .utils.prefix_cache import resolve_prefix_hashes
 
 logger = get_router_logger("balancer")
 
@@ -226,23 +228,42 @@ def acquire_server(self, request_id: str, prompt_ids: list[int] | None = None) -
         Raises ``RuntimeError`` if no replica is available. ``request_id`` is
         forwarded so strategies can short-circuit to a sticky-bound replica;
         ``on_acquire`` then refreshes the binding.
+
+        The prompt's prefix-hash chain is resolved once here and threaded
+        through ``route()`` → ``score()`` (so strategies don't re-resolve), then
+        used to record the dispatch in the GPU reverse index — populating the
+        locality signal **immediately**, before vLLM's KV events arrive
+        (6-9s batched). KV-event ``removed``/``clear`` still evict later.
         """
         replicas = [ReplicaInfo(replica_id=sid) for sid in self._servers]
         self._route_calls += 1
+        # Resolve the prefix-hash chain once: shared by score() (locality) and
+        # the post-route dispatch record (populate reverse index immediately).
+        gpu_hash_strs = resolve_prefix_hashes(prompt_ids or [], request_id, self._store) if prompt_ids else None
         t0 = time.perf_counter()
         ranking = route(
             self._strategies,
             prompt_ids,
             self._store,
             replicas,
             request_id,
+            gpu_hash_strs,
         )
         dt_ms = (time.perf_counter() - t0) * 1000
         self._route_time_total_s += dt_ms / 1000.0
         self._route_time_max_ms = max(self._route_time_max_ms, dt_ms)
         if not ranking:
             raise RuntimeError("no available replica to route to")
         server_id = ranking[0]
+
+        if request_id is not None and prompt_ids:
+            gpu_hit = 0.0
+            if gpu_hash_strs:
+                gpu_hit = self._store.get_layer_prefix_hit_rate(server_id, gpu_hash_strs, Layer.GPU) or 0.0
+            miss_tokens = int(len(prompt_ids) * (1.0 - gpu_hit))
+            self._store.set_per_request(request_id, "miss_tokens", miss_tokens)
+        if gpu_hash_strs:
+            self._store.record_dispatched_prefix(server_id, gpu_hash_strs)
         self._fire("on_acquire", request_id, server_id, prompt_ids)
         logger.info(
             f"request={request_id} routed to server={server_id} (ranking={ranking}, pool={list(self._servers)}, "

verl/workers/rollout/router/kvcaware/collectors/collector.py

@@ -26,7 +26,7 @@
 from ..insight import WriteEvent, WriteKind, emitter
 from ..logging import get_router_logger
 from ..store.data_store import DataStore
-from ..types import EmitKey, MetricKey
+from ..types import EmitKey, Layer, MetricKey
 from .decoder import Decoder, KVCacheUpdate, MetricsUpdate, StickyUpdate
 from .transport.base import Transport
 
@@ -149,7 +149,16 @@ def handler(raw_data: bytes | str, node_id: str) -> None:
                 tmp_loop.close()
 
     def _write_kv_update(self, update: KVCacheUpdate) -> None:
-        """Write KVCacheUpdate via DataStore, then emit a periodic kv-events tally."""
+        """Write KVCacheUpdate via DataStore, then emit a periodic kv-events tally.
+
+        A real ``BlockRemoved`` (GPU layer) is the only thing allowed to
+        decrement ``INFLIGHT_TOKENS`` (see ``_write_metrics_update`` — release
+        no longer does). This makes the gauge "cumulative acquire-time
+        estimate − cumulative real eviction", tracking actual GPU KV pressure
+        instead of request lifetime. Uses the decoder-learned ``block_size``
+        (BlockRemoved events don't carry one — only BlockStored does), so no
+        decrement happens until at least one BlockStored has been seen.
+        """
         if update.block_size is not None:
             self._data_store.set_block_size(update.block_size)
         if update.clear_all:
@@ -161,6 +170,16 @@ def _write_kv_update(self, update: KVCacheUpdate) -> None:
             if hashes:
                 self._data_store.add_kv_blocks(update.node_id, hashes, layer=layer)
 
+        n_removed_gpu = len(update.remove_blocks.get(Layer.GPU, []))
+        block_size = self._data_store.get_block_size()
+        if n_removed_gpu and block_size:
+            self._data_store.incr_metric(
+                update.node_id,
+                MetricKey.INFLIGHT_TOKENS,
+                -(n_removed_gpu * block_size),
+                floor=0,
+            )
+
         # Tally for periodic summary — observe BlockStored/BlockRemoved flow.
         n_added = sum(len(v) for v in update.add_blocks.values())
         n_removed = sum(len(v) for v in update.remove_blocks.values())
@@ -222,7 +241,29 @@ def _write_metrics_update(self, update: MetricsUpdate) -> None:
                     deltas[MetricKey.INFLIGHT_TURN_SUM] = -self._data_store.get_per_request(
                         update.request_id, "turn", 0
                     )
-            new_values = self._data_store.incr_metrics(update.node_id, deltas)
+            # INFLIGHT_TOKENS approximates KV occupancy: acquire adds the
+            # acquire-time miss_tokens estimate (plen × (1 - gpu_hit)) stashed
+            # in PerRequestStore by acquire_server. It is deliberately NOT
+            # subtracted on release — request completion does not free KV
+            # blocks; vLLM keeps them in the free-but-cached pool until real
+            # LRU eviction. The decrement instead comes from BlockRemoved
+            # events (see ``_write_kv_update``), so the gauge tracks
+            # "cumulative estimated occupancy − cumulative real eviction",
+            # closer to vLLM's own kv_cache_usage_perc than a request-lifetime
+            # counter. Falls back to prompt_len (the old symmetric semantics)
+            # only on acquire when miss_tokens was never set — e.g. requests
+            # routed without a strategy that resolves gpu_hit; release always
+            # drops the INFLIGHT_TOKENS delta the decoder emitted and frees
+            # the per-request entry once its purpose (acquire lookup) is done.
+            if MetricKey.INFLIGHT_TOKENS in deltas and update.request_id is not None:
+                if is_acquire:
+                    miss_tokens = self._data_store.get_per_request(update.request_id, "miss_tokens", None)
+                    if miss_tokens is not None:
+                        deltas[MetricKey.INFLIGHT_TOKENS] = int(miss_tokens)
+                elif is_release:
+                    deltas.pop(MetricKey.INFLIGHT_TOKENS, None)
+                    self._data_store.del_per_request(update.request_id, "miss_tokens")
+            new_values = self._data_store.incr_metrics(update.node_id, deltas, floor={MetricKey.INFLIGHT_TOKENS: 0})
             if emitter.enabled() and (is_acquire or is_release):
                 emitter.on_write(
                     WriteEvent(

verl/workers/rollout/router/kvcaware/store/data_store.py

@@ -123,6 +123,29 @@ def add_kv_blocks(self, node_id: str, block_hashes: list[str], layer: Layer = La
         """
         self._kv.add_blocks(node_id, block_hashes, layer=layer)
 
+    def record_dispatched_prefix(self, replica_id: str, hash_strs: list[str]) -> None:
+        """Mark ``hash_strs`` as dispatched to ``replica_id`` in the GPU reverse index.
+
+        Called at acquire time (see ``KVCAwareBalancer.acquire_server``) so the
+        prefix→replica reverse index that ``get_layer_prefix_hit_rate`` walks is
+        populated **immediately**, without waiting for vLLM's KV events (which
+        arrive in 6-9s batches). The same hashes later arrive via KV-event
+        ``BlockStored`` → ``add_kv_blocks``; that path is idempotent
+        (``add_blocks`` dedups via set), so the backfill is a no-op.
+
+        KV-event ``BlockRemoved``/``AllBlocksCleared`` still evict entries from
+        the reverse index — so once vLLM actually LRU-evicts a prefix, the
+        locality signal decays correctly (with the KV-event batch delay).
+
+        Args:
+            replica_id: The replica the prompt was routed to.
+            hash_strs: The prompt's full-block chained prefix hashes (the same
+                ``gpu_hash_strs`` ``score()`` consumes — computed once in
+                ``acquire_server`` and threaded through ``route()``).
+        """
+        if hash_strs:
+            self._kv.add_blocks(replica_id, hash_strs, layer=Layer.GPU)
+
     def remove_kv_blocks(self, node_id: str, block_hashes: list[str], layer: Layer = Layer.GPU) -> None:
         """Remove KV cache blocks from a node.
 
@@ -206,29 +229,44 @@ def per_replica_block_counts(self) -> dict[str, int]:
 
     # ── PerReplicaStore incremental write ──────────────────────────────────
 
-    def incr_metric(self, node_id: str, key: str, delta: int | float = 1) -> int | float:
+    def incr_metric(
+        self, node_id: str, key: str, delta: int | float = 1, floor: int | float | None = None
+    ) -> int | float:
         """Apply a signed delta to one metric for one node (inflight ±1).
 
         Routes to ``PerReplicaStore.incr`` (not ``refresh``) so a stateless delta
         emitter (``InflightDecoder``) can move a running counter without
         tracking the absolute value itself.
 
+        Args:
+            floor: If given, clamp the post-delta value to ``max(new, floor)``.
+
         Returns:
             The new value of ``key``.
         """
-        return self._metrics.incr(node_id, key, delta)
+        return self._metrics.incr(node_id, key, delta, floor=floor)
 
-    def incr_metrics(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, int | float]:
+    def incr_metrics(
+        self,
+        node_id: str,
+        deltas: dict[str, int | float],
+        floor: dict[str, int | float] | None = None,
+    ) -> dict[str, int | float]:
         """Apply multiple signed deltas to one node under a single lock.
 
         Batched variant of :meth:`incr_metric` for the ``on_acquire`` decoder,
         which emits several deltas per dispatch (INFLIGHT / DISPATCHED /
         PROMPT_LEN_SUM) — batching avoids one lock cycle per key on the hot path.
 
+        Args:
+            floor: Optional ``{canonical_key: floor_value}`` passed through to
+                ``PerReplicaStore.incr_many`` — clamps keys decremented by a
+                source independent of whatever incremented them.
+
         Returns:
             ``{canonical_key: new_value}`` for every key in ``deltas``.
         """
-        return self._metrics.incr_many(node_id, deltas)
+        return self._metrics.incr_many(node_id, deltas, floor=floor)
 
     # ── Sticky bindings (a per-request value stored under _STICKY_KEY) ───