diff --git a/examples/kvc_aware_router/plot_metrics.py b/examples/kvc_aware_router/plot_metrics.py index 7575ee50866..e955c33cb7e 100644 --- a/examples/kvc_aware_router/plot_metrics.py +++ b/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).""" diff --git a/tests/workers/rollout/router/strategies/test_route.py b/tests/workers/rollout/router/strategies/test_route.py index feb64635232..30ae92bfcec 100644 --- a/tests/workers/rollout/router/strategies/test_route.py +++ b/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") diff --git a/verl/trainer/config/rollout/router/kvcaware.yaml b/verl/trainer/config/rollout/router/kvcaware.yaml index b99fbdd32aa..a80bba40e0a 100644 --- a/verl/trainer/config/rollout/router/kvcaware.yaml +++ b/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 diff --git a/verl/workers/rollout/router/kvcaware/balancer.py b/verl/workers/rollout/router/kvcaware/balancer.py index 971d9eb4cef..fcce8362cfd 100644 --- a/verl/workers/rollout/router/kvcaware/balancer.py +++ b/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,9 +228,18 @@ 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, @@ -236,6 +247,7 @@ def acquire_server(self, request_id: str, prompt_ids: list[int] | None = None) - self._store, replicas, request_id, + gpu_hash_strs, ) dt_ms = (time.perf_counter() - t0) * 1000 self._route_time_total_s += dt_ms / 1000.0 @@ -243,6 +255,15 @@ def acquire_server(self, request_id: str, prompt_ids: list[int] | None = None) - 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)}, " diff --git a/verl/workers/rollout/router/kvcaware/collectors/collector.py b/verl/workers/rollout/router/kvcaware/collectors/collector.py index 4a4d010887b..4b82c88dd1a 100644 --- a/verl/workers/rollout/router/kvcaware/collectors/collector.py +++ b/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( diff --git a/verl/workers/rollout/router/kvcaware/store/data_store.py b/verl/workers/rollout/router/kvcaware/store/data_store.py index e557a4ed4e9..cedb06c9980 100644 --- a/verl/workers/rollout/router/kvcaware/store/data_store.py +++ b/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) ─── diff --git a/verl/workers/rollout/router/kvcaware/store/per_replica_store.py b/verl/workers/rollout/router/kvcaware/store/per_replica_store.py index 90e5ae87602..f18e996e0f1 100644 --- a/verl/workers/rollout/router/kvcaware/store/per_replica_store.py +++ b/verl/workers/rollout/router/kvcaware/store/per_replica_store.py @@ -63,7 +63,7 @@ def get(self, node_id: str, key: str | None = None) -> Any | dict[str, Any]: return node[key] return METRIC_SPECS[key]["default"] - def incr(self, node_id: str, key: str, delta: int | float = 1) -> int | float: + def incr(self, node_id: str, key: str, delta: int | float = 1, floor: int | float | None = None) -> int | float: """Apply a numeric delta to one key for one node (inflight ±1). Unlike ``refresh`` (batch merge overwrite), this is an incremental @@ -75,6 +75,11 @@ def incr(self, node_id: str, key: str, delta: int | float = 1) -> int | float: node_id: Target node. key: Canonical metric key (must be in ``METRIC_SPECS``). delta: Signed delta to add (default +1). + floor: If given, clamp the post-delta value to ``max(new, floor)``. + Needed for gauges fed by two independent, unsynchronized + sources (e.g. acquire-time estimate vs. eviction-driven KV + events) where a burst of decrements could otherwise drive the + value negative. Returns: The new value of ``key`` after applying ``delta``. @@ -84,9 +89,15 @@ def incr(self, node_id: str, key: str, delta: int | float = 1) -> int | float: """ if key not in METRIC_SPECS: raise KeyError(f"Unknown metric key '{key}'. Valid keys: {sorted(METRIC_SPECS.keys())}") - return self._apply(node_id, {key: delta})[key] - - def incr_many(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, int | float]: + floors = {key: floor} if floor is not None else None + return self._apply(node_id, {key: delta}, floor=floors)[key] + + def incr_many( + 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. Same incremental semantics as :meth:`incr` (reads current, adds delta, @@ -98,6 +109,8 @@ def incr_many(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, i Args: node_id: Target node. deltas: ``{canonical_key: signed_delta}`` (all keys in ``METRIC_SPECS``). + floor: Optional ``{canonical_key: floor_value}`` for keys whose + post-delta value should be clamped to ``max(new, floor_value)``. Returns: ``{canonical_key: new_value}`` for every key in ``deltas``. @@ -110,15 +123,28 @@ def incr_many(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, i bad = [k for k in deltas if k not in METRIC_SPECS] if bad: raise KeyError(f"Unknown metric keys: {sorted(bad)}. Valid keys: {sorted(METRIC_SPECS.keys())}") - return self._apply(node_id, deltas) - - def _apply(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, int | float]: + return self._apply(node_id, deltas, floor=floor) + + def _apply( + self, + node_id: str, + deltas: dict[str, int | float], + floor: dict[str, int | float] | None = None, + ) -> dict[str, int | float]: """Apply signed deltas to one node under one lock (caller validates keys). Each key falls back to its ``METRIC_SPECS`` default before adding the delta, so the writer stays stateless — it only emits the +/-delta. Shared by :meth:`incr` (single key) and :meth:`incr_many` (batch). + Args: + floor: Optional ``{canonical_key: floor_value}`` — keys present here + are clamped to ``max(current + delta, floor_value)`` after the + delta is applied. Needed when a gauge is decremented by a source + independent of whatever incremented it (e.g. KV-event-driven + eviction vs. acquire-time estimate), where accumulated drift + could otherwise push the value below a meaningful floor (usually 0). + Returns: ``{canonical_key: new_value}`` after applying each delta — computed from the loop-local ``current + delta`` under the same lock, so there @@ -129,6 +155,8 @@ def _apply(self, node_id: str, deltas: dict[str, int | float]) -> dict[str, int new_values: dict[str, int | float] = {} for key, delta in deltas.items(): new = node.get(key, METRIC_SPECS[key]["default"]) + delta + if floor is not None and key in floor: + new = max(new, floor[key]) node[key] = new new_values[key] = new return new_values diff --git a/verl/workers/rollout/router/kvcaware/strategies/kvc_aware.py b/verl/workers/rollout/router/kvcaware/strategies/kvc_aware.py index 48b6ebf966e..c714f5c3b6c 100644 --- a/verl/workers/rollout/router/kvcaware/strategies/kvc_aware.py +++ b/verl/workers/rollout/router/kvcaware/strategies/kvc_aware.py @@ -234,12 +234,17 @@ def score( store: DataStore, replicas: list[ReplicaInfo], request_id: str | None = None, + gpu_hash_strs: list[str] | None = None, ) -> list[float]: """Score each replica. Larger is better. After the sticky short-circuit misses, the ``slow_cut`` mode selects the fallback scoring: ``prefix-load-aware`` → ``S = α·S_cache + (1-α)·S_load``; ``least-inflight`` → ``-INFLIGHT_COUNT`` (verl GlobalRequestLoadBalancer-style). + + ``gpu_hash_strs`` is the caller-computed prefix-hash chain (shared + across replicas, computed once in ``route()``). When ``None``, resolve + here — backward-compat for direct callers that don't pre-resolve. """ if not isinstance(replicas, list): raise StrategyError(f"replicas must be a list, got {type(replicas).__name__}") @@ -256,7 +261,8 @@ def score( if self.slow_cut == SlowCut.LEAST_INFLIGHT: return [-store.get_metric(r.replica_id, MetricKey.INFLIGHT_COUNT) for r in replicas] # Hash-resolving slow_cuts share one resolution across all replicas. - gpu_hash_strs = resolve_prefix_hashes(prompt_ids or [], request_id, store) + if gpu_hash_strs is None: + gpu_hash_strs = resolve_prefix_hashes(prompt_ids or [], request_id, store) if self.slow_cut == SlowCut.PREFIX_LOAD_AWARE: return self._prefix_load_aware(store, replicas, gpu_hash_strs) if self.slow_cut == SlowCut.CAPACITY_TOKEN_AWARE: @@ -357,29 +363,31 @@ def _capacity_token_scores( prompt_ids: list[int], gpu_hash_strs: list[str], ) -> list[float]: - """Capacity-gated token routing (discrete: winner=STICKY_TOP_SCORE, rest 0). + """Capacity-gated token routing (discrete: top replicas get STICKY_TOP_SCORE). For each replica ``i``:: - avail[i] = cap × (1 - kv_cache_usage_perc[i]) # free tokens (no cache) - need[i] = len(prompt_ids) × (1 - gpu_hit[i]) # prefill this req adds - remaining[i] = avail[i] - need[i] # free tokens after assign + need[i] = len(prompt_ids) × (1 - gpu_hit[i]) # prefill this req adds + remaining[i] = avail[i] - need[i] # free after assign eligible[i] = avail[i] >= cap × (1 - load_threshold) # pure capacity gate - pick ``argmin(inflight_tokens)`` (least in-flight tokens wins) to keep - the first wave from collapsing onto ``pool[0]``. - Otherwise pick ``argmax(eligible, remaining)``. + Per-replica switch — whichever replica gets polled first switches + first. Ranking: replicas above the capacity gate rank before those + below, each group ordered by ``remaining`` descending. Replicas tied + at the top ``remaining`` all get ``STICKY_TOP_SCORE`` — ``route()`` + breaks the tie at random (anti pool[0]-collapse). """ n = len(replicas) cap = self._total_token_capacity(store) plen = len(prompt_ids) if prompt_ids else 0 rows: list[dict] = [] for replica in replicas: - kv_perc = store.get_metric(replica.replica_id, MetricKey.KV_CACHE_USAGE_PERC) or 0.0 + kv_perc = store.get_metric(replica.replica_id, MetricKey.KV_CACHE_USAGE_PERC) or 0 inflight = store.get_metric(replica.replica_id, MetricKey.INFLIGHT_COUNT) or 0 inflight_tokens = store.get_metric(replica.replica_id, MetricKey.INFLIGHT_TOKENS) or 0 s_cache, gpu_hit = self._cache_score(store, replica, gpu_hash_strs) - avail = cap * (1.0 - kv_perc) + + avail = cap - inflight_tokens need = plen * (1.0 - gpu_hit) remaining = avail - need if emitter.enabled(): @@ -399,37 +407,52 @@ def _capacity_token_scores( ) thresh = cap * (1.0 - self.load_threshold) - cold_start = store.get_sticky_binding(request_id) is None - if cold_start: - top = min(range(n), key=lambda i: rows[i]["inflight_tokens"]) - logger.info("score(): CAPACITY_TOKEN_AWARE cold start → min inflight_tokens") - else: - eligible = [i for i in range(n) if rows[i]["avail"] >= thresh] - if not eligible: - top = max(range(n), key=lambda i: rows[i]["remaining"]) - logger.info("score(): CAPACITY_TOKEN_AWARE no eligible → max remaining") - else: - top = max(eligible, key=lambda i: rows[i]["remaining"]) + # Unified ranking: eligible (avail >= thresh) first, then by remaining + # desc, then by need asc. Among replicas with equal remaining, the one + # with the smallest ``need`` wins — need = plen×(1-gpu_hit), so smaller + # need means more prefix-cache overlap (locality-first tie-break). Only + # replicas tying on (eligible, remaining, need) get STICKY_TOP_SCORE + # together; route() then breaks that residual tie at random. + # ``eligible`` is a bool (1 = above gate) so it sorts before 0. + order = sorted( + range(n), + key=lambda i: (1 if rows[i]["avail"] >= thresh else 0, rows[i]["remaining"], -rows[i]["need"]), + reverse=True, + ) + best_remaining = rows[order[0]]["remaining"] + best_need = rows[order[0]]["need"] + top_idx = [ + i + for i in order + if rows[i]["remaining"] == best_remaining and rows[i]["need"] == best_need and rows[i]["avail"] >= thresh + ] + # When no replica clears the gate, ``order``'s head is the largest + # ``remaining`` (then smallest ``need``) among the all-overloaded set; + # tie those too. + if not top_idx: + top_idx = [order[0]] for i, row in enumerate(rows): - tag = " ← WINNER" if i == top else "" + tag = " ← WINNER" if i in top_idx else "" logger.info( f"score(): replica={row['replica'].replica_id} kv_perc={row['kv_perc']:.3f} " f"gpu_hit={row['gpu_hit']:.3f} inflight={row['inflight']} " f"avail={row['avail']:.0f} need={row['need']:.0f} " f"max_num_batched_tokens={self._max_num_batched_tokens} inflight_tokens={row['inflight_tokens']:} " - f"remaining={row['remaining']:.0f}{tag}" + f"remaining={row['remaining']:.0f} {tag}" ) - winner = rows[top]["replica"].replica_id + winners = [rows[i]["replica"].replica_id for i in top_idx] logger.info( - f"score(): CAPACITY_TOKEN_AWARE winner={winner} " - f"(kv_perc={rows[top]['kv_perc']:.3f}, remaining={rows[top]['remaining']:.0f})" + f"score(): CAPACITY_TOKEN_AWARE winners={winners} " + f"(remaining={rows[top_idx[0]]['remaining']:.0f}, " + f"kv_perc={rows[top_idx[0]]['kv_perc']:.3f})" ) # Per-replica capacity signal for the plot (mirrors route-load in prefix-load-aware). cap_loads = {row["replica"].replica_id: row["remaining"] for row in rows} logger.info(f"route-capacity remaining={cap_loads}") scores = [0.0] * n - scores[top] = STICKY_TOP_SCORE + for i in top_idx: + scores[i] = STICKY_TOP_SCORE return scores diff --git a/verl/workers/rollout/router/kvcaware/strategies/routing.py b/verl/workers/rollout/router/kvcaware/strategies/routing.py index a8dce41ae8a..79243ab6cc2 100644 --- a/verl/workers/rollout/router/kvcaware/strategies/routing.py +++ b/verl/workers/rollout/router/kvcaware/strategies/routing.py @@ -43,6 +43,7 @@ def score( store: Any, replicas: list[Any], request_id: str | None = None, + gpu_hash_strs: list[str] | None = None, ) -> list[float]: """Score each replica. Larger is better; negatives are allowed. @@ -51,6 +52,13 @@ def score( return a pre-built score list that places it first when it is not overloaded (see ``KVCacheAwareStrategy``). Strategies that ignore stickiness accept ``request_id`` and proceed with their own scoring. + + ``gpu_hash_strs`` is the caller-computed prefix-hash chain (shared + across all replicas, computed once in ``route()`` from + ``prompt_ids`` + ``request_id``). Strategies that need it + (``prefix-load-aware``, ``capacity-token-aware``) consume it directly + instead of re-resolving; ``None`` means the caller did not pre-resolve + and the strategy may compute it itself. """ ... @@ -66,6 +74,7 @@ def route( store: Any, replicas: list[Any], request_id: str | None = None, + gpu_hash_strs: list[str] | None = None, ) -> list[str]: """Return replica ids ranked best-first. @@ -79,6 +88,11 @@ def route( store: ``DataStore`` for metric + sticky-session queries. replicas: ``[ReplicaInfo, ...]`` — candidate replicas. request_id: session id for sticky-session routing (may be ``None``). + gpu_hash_strs: caller-computed prefix-hash chain for ``prompt_ids`` + (shared across replicas). Computed once by the caller + (``acquire_server``) and threaded through here so strategies and the + post-route dispatch-recorder share one resolution. ``None`` → + resolve inside the strategy (backward-compat). Returns: Replica ids sorted by total score, best first. Falls back to random @@ -100,6 +114,7 @@ def route( store, replicas, request_id, + gpu_hash_strs, ) if len(scores) != n: raise ValueError(f"{name}.score() returned {len(scores)} scores, expected {n}") @@ -113,7 +128,18 @@ def route( for idx in range(n): final[idx] += weight * scores[idx] + # Rank best-first. Among replicas tied at the top score, pick one at random + # (instead of the stable sort's always-first bias) so a true cold start or + # a same-prompt rollout wave spreads across the pool rather than collapsing + # onto pool[0]. Lower ranks keep their stable order — only the winner set + # is randomized. ranking = sorted(range(n), key=lambda idx: _rank_key(final[idx]), reverse=True) + top_score = _rank_key(final[ranking[0]]) + top_ties = [idx for idx in ranking if _rank_key(final[idx]) == top_score] + if len(top_ties) > 1: + chosen = random.choice(top_ties) + # Move the random pick to the head; keep the rest in stable order. + ranking = [chosen] + [idx for idx in ranking if idx != chosen] scores_str = ", ".join(f"{replicas[idx].replica_id}={final[idx]:.4f}" for idx in ranking) logger.info(f"route(): replicas={n} ranking=[{scores_str}]") return [replicas[idx].replica_id for idx in ranking]