[Core] Fix internal LB load-balancing

rust/src/engine-core-client/src/client/state.rs

@@ -74,17 +74,13 @@ impl EngineRoutingState {
     ///
     /// Scheduler stats can raise the load estimate above the frontend-local
     /// view, but they should not lower it below requests this frontend has
-    /// already admitted. Waiting requests still get the same extra penalty
-    /// as the original `waiting * 4 + running` score.
+    /// already admitted.
     fn routing_score(&self) -> usize {
-        const WAITING_WEIGHT: usize = 4;
-
         let Some(stats) = self.last_scheduler_stats else {
             return self.inflight;
         };
 
-        let scheduler_total = stats.running + stats.waiting;
-        self.inflight.max(scheduler_total) + stats.waiting * (WAITING_WEIGHT - 1)
+        self.inflight.max(stats.running + stats.waiting)
     }
 
     /// Replace the local routing view with a fresh real scheduler snapshot.
@@ -750,7 +746,7 @@ mod tests {
     }
 
     #[test]
-    fn routing_score_keeps_extra_waiting_penalty() {
+    fn routing_score_counts_waiting_without_extra_penalty() {
         let state = EngineRoutingState {
             inflight: 1,
             last_scheduler_stats: Some(EngineLoadSnapshot {
@@ -759,7 +755,7 @@ mod tests {
             }),
         };
 
-        assert_eq!(state.routing_score(), 14);
+        assert_eq!(state.routing_score(), 5);
     }
 
     #[test]

tests/v1/engine/test_engine_core_client.py

@@ -8,6 +8,7 @@
 import signal
 import time
 import uuid
+from collections import Counter
 from concurrent.futures import Future
 from dataclasses import dataclass
 from threading import Thread
@@ -27,7 +28,11 @@
 from vllm.pooling_params import LateInteractionParams, PoolingParams
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils.torch_utils import set_default_torch_num_threads
-from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest
+from vllm.v1.engine import (
+    EngineCoreOutputs,
+    EngineCoreReadyResponse,
+    EngineCoreRequest,
+)
 from vllm.v1.engine.core import EngineCore
 from vllm.v1.engine.core_client import (
     AsyncMPClient,
@@ -198,13 +203,19 @@ def _make_pooling_request(
     )
 
 
-def test_dplb_late_interaction_sticky_routing():
+def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncMPClient:
     client = object.__new__(DPLBAsyncMPClient)
-    client.client_count = 1
+    client.client_count = client_count
     client.reqs_in_flight = {}
-    client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
-    client.lb_engines = [[0, 0], [0, 0], [0, 0]]
+    client.engine_inflight = Counter()
+    client.core_engines = [bytes([i, 0]) for i in range(num_engines)]
+    client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]
     client.eng_start_index = 0
+    return client
+
+
+def test_dplb_late_interaction_sticky_routing():
+    client = _make_dplb_client()
 
     query_key = "rerank-abc-query-0"
     query_request = _make_pooling_request(
@@ -223,12 +234,8 @@ def test_dplb_late_interaction_sticky_routing():
 
 
 def test_dplb_non_late_interaction_still_uses_lb():
-    client = object.__new__(DPLBAsyncMPClient)
-    client.client_count = 1
-    client.reqs_in_flight = {}
-    client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
-    client.lb_engines = [[2, 1], [0, 0], [1, 0]]
-    client.eng_start_index = 0
+    client = _make_dplb_client()
+    client.lb_engines = [[2, 1, 0.0], [0, 0, 0.0], [1, 0, 0.0]]
 
     request = make_request(SamplingParams(max_tokens=1))
     chosen_engine = client.get_core_engine_for_request(request)
@@ -237,6 +244,70 @@ def test_dplb_non_late_interaction_still_uses_lb():
     assert client.lb_engines[1][0] == 1
 
 
+def test_dplb_burst_round_robins_despite_snapshot_rebinds():
+    """A stats snapshot rebind wipes the optimistic lb_engines increments;
+    the exact in-flight floor must keep a burst spreading round-robin."""
+    client = _make_dplb_client(num_engines=4)
+
+    for _ in range(4):
+        client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))
+    # Coordinator snapshot arrives, not yet reflecting the 4 routed requests.
+    client.lb_engines = [[0, 0, 0.0] for _ in range(4)]
+    for _ in range(4):
+        client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))
+
+    assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2]
+
+
+def test_dplb_snapshot_backpressure_overrides_inflight():
+    """An engine reported heavily loaded by the coordinator is avoided even
+    when this client has routed nothing to it."""
+    client = _make_dplb_client(num_engines=2)
+    client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]]
+
+    chosen = client.get_core_engine_for_request(
+        make_request(SamplingParams(max_tokens=1))
+    )
+
+    assert chosen == client.core_engines[1]
+
+
+def test_dplb_kv_pressure_amplifies_waiting_penalty():
+    """A waiting queue on a KV-bound engine (slow drain) is penalized, while
+    the same queue with low KV usage is not (e.g. transient burst)."""
+    client = _make_dplb_client(num_engines=2)
+    # Engine 0 has a smaller total but is KV-bound with a queue.
+    client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]]
+
+    chosen = client.get_core_engine_for_request(
+        make_request(SamplingParams(max_tokens=1))
+    )
+    assert chosen == client.core_engines[1]
+
+    # Same counts without KV pressure: the smaller total wins.
+    client = _make_dplb_client(num_engines=2)
+    client.lb_engines = [[5, 10, 0.2], [0, 20, 0.2]]
+
+    chosen = client.get_core_engine_for_request(
+        make_request(SamplingParams(max_tokens=1))
+    )
+    assert chosen == client.core_engines[0]
+
+
+def test_dplb_finished_requests_release_inflight():
+    client = _make_dplb_client(num_engines=2)
+
+    req = make_request(SamplingParams(max_tokens=1))
+    engine = client.get_core_engine_for_request(req)
+    assert client.engine_inflight[engine] == 1
+
+    outputs = EngineCoreOutputs(finished_requests={req.request_id})
+    asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs))
+
+    assert client.engine_inflight[engine] == 0
+    assert req.request_id not in client.reqs_in_flight
+
+
 def test_apply_ready_response_syncs_block_size():
     import msgspec
 

vllm/v1/core/sched/interface.py

@@ -234,6 +234,10 @@ def get_request_counts(self) -> tuple[int, int]:
         """Returns (num_running_reqs, num_waiting_reqs)."""
         raise NotImplementedError
 
+    def get_kv_cache_usage(self) -> float:
+        """Returns the fraction of the KV cache currently in use (0.0-1.0)."""
+        return 0.0
+
     @abstractmethod
     def make_stats(self) -> "SchedulerStats | None":
         """Make a SchedulerStats object for logging.

vllm/v1/core/sched/scheduler.py

@@ -2110,6 +2110,10 @@ def get_request_counts(self) -> tuple[int, int]:
         """Returns (num_running_reqs, num_waiting_reqs)."""
         return len(self.running), len(self.waiting) + len(self.skipped_waiting)
 
+    def get_kv_cache_usage(self) -> float:
+        """Returns the fraction of the KV cache currently in use (0.0-1.0)."""
+        return self.kv_cache_manager.usage
+
     def add_request(self, request: Request) -> None:
         existing = self.requests.get(request.request_id)
         if existing is not None:

vllm/v1/engine/coordinator.py

@@ -139,7 +139,8 @@ def shutdown(self, timeout: float | None = None) -> None:
 
 class EngineState:
     def __init__(self):
-        self.request_counts = [0, 0]  # [waiting, running]
+        # [waiting, running, kv_cache_usage]
+        self.request_counts: list[int | float] = [0, 0, 0.0]
 
 
 class DPCoordinatorProc:
@@ -202,7 +203,7 @@ def process_input_socket(
         stats_changed = False
         last_stats_step = -1
         last_stats_wave = -1
-        last_step_counts: list[list[int]] | None = None
+        last_step_counts: list[list[int | float]] | None = None
 
         with (
             make_zmq_socket(
@@ -260,8 +261,12 @@ def process_input_socket(
                 wait_for = self.stats_update_interval_ms if stats_changed else 5000
 
                 # Wait at least 50ms to ensure we've received all stats for
-                # the current step.
-                min_timeout = 50 if last_step_counts is None else 0
+                # the current step. Only applicable to lockstep (MoE) DP;
+                # non-lockstep engines have no synchronized step boundaries.
+                if self.enable_wave_coordination and last_step_counts is None:
+                    min_timeout = 50
+                else:
+                    min_timeout = 0
 
                 events = poller.poll(timeout=max(min_timeout, wait_for - elapsed))
                 if not events:
@@ -374,32 +379,38 @@ def process_input_socket(
                         # 1. Updated request load stats - update our local
                         # state with these.
                         stats = self.engines[eng_index].request_counts
-                        stats_step = scheduler_stats.step_counter
-                        stats_wave = scheduler_stats.current_wave
-                        if (
-                            stats_wave > last_stats_wave
-                            or stats_wave == last_stats_wave
-                            and stats_step > last_stats_step
-                        ):
-                            if stats_changed:
-                                last_step_counts = self._get_engine_counts(do_copy=True)
-                            last_stats_step = stats_step
-                            last_stats_wave = stats_wave
-                        elif stats_wave != last_stats_wave or (
-                            stats_step != last_stats_step
-                        ):
-                            logger.warning(
-                                "Received stats for out-of-order "
-                                "step (%d, %d) from engine %d (expected "
-                                "> (%d, %d))",
-                                stats_wave,
-                                stats_step,
-                                eng_index,
-                                last_stats_wave,
-                                last_stats_step,
-                            )
+                        if self.enable_wave_coordination:
+                            # Steps are synchronized across lockstep (MoE) DP
+                            # ranks; snapshot counts at step boundaries.
+                            stats_step = scheduler_stats.step_counter
+                            stats_wave = scheduler_stats.current_wave
+                            if (
+                                stats_wave > last_stats_wave
+                                or stats_wave == last_stats_wave
+                                and stats_step > last_stats_step
+                            ):
+                                if stats_changed:
+                                    last_step_counts = self._get_engine_counts(
+                                        do_copy=True
+                                    )
+                                last_stats_step = stats_step
+                                last_stats_wave = stats_wave
+                            elif stats_wave != last_stats_wave or (
+                                stats_step != last_stats_step
+                            ):
+                                logger.warning(
+                                    "Received stats for out-of-order "
+                                    "step (%d, %d) from engine %d (expected "
+                                    "> (%d, %d))",
+                                    stats_wave,
+                                    stats_step,
+                                    eng_index,
+                                    last_stats_wave,
+                                    last_stats_step,
+                                )
                         stats[0] = scheduler_stats.num_waiting_reqs
                         stats[1] = scheduler_stats.num_running_reqs
+                        stats[2] = scheduler_stats.kv_cache_usage
                         stats_changed = True
 
                     # Wave coordination: handle wave completion and start notifications
@@ -452,7 +463,7 @@ def _send_start_wave(
         wave_encoded = msgspec.msgpack.encode((wave, exclude_engine_index))
         socket.send_multipart((EngineCoreRequestType.START_DP_WAVE.value, wave_encoded))
 
-    def _get_engine_counts(self, do_copy=False) -> list[list[int]]:
+    def _get_engine_counts(self, do_copy=False) -> list[list[int | float]]:
         """Return list of [waiting, running] count lists for each engine."""
         if do_copy:
             return [copy.copy(e.request_counts) for e in self.engines]

vllm/v1/engine/core.py

@@ -1057,6 +1057,7 @@ def __init__(
             # Only publish request queue stats to coordinator for "internal"
             # and "hybrid" LB modes.
             self.publish_dp_lb_stats = internal_dp_balancing
+            self.last_counts = (0, 0)
 
             self.addresses = addresses
             self.process_input_queue_block = True
@@ -1376,11 +1377,27 @@ def run_busy_loop(self):
         while self._handle_shutdown():
             # 1) Poll the input queue until there is work to do.
             self._process_input_queue()
+            # Publish request counts before and after GPU step to ensure freshness.
+            self._maybe_publish_request_counts()
             # 2) Step the engine core and return the outputs.
             self._process_engine_step()
+            self._maybe_publish_request_counts()
 
         raise SystemExit
 
+    def _maybe_publish_request_counts(self):
+        if not self.publish_dp_lb_stats:
+            return
+
+        # Publish our request counts (if they've changed).
+        counts = self.scheduler.get_request_counts()
+        if counts != self.last_counts:
+            self.last_counts = counts
+            stats = SchedulerStats(
+                *counts, kv_cache_usage=self.scheduler.get_kv_cache_usage()
+            )
+            self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats)))
+
     def _process_input_queue(self):
         """Exits when an engine step needs to be performed."""
 
@@ -1890,7 +1907,6 @@ def __init__(
         # finished with DP peers every N steps.
         self.step_counter = 0
         self.current_wave = 0
-        self.last_counts = (0, 0)
 
         # Two-phase pause protocol state. When pending_pause is True, the
         # engine keeps stepping (dummy batches) while waiting for all DP
@@ -2027,12 +2043,16 @@ def _maybe_publish_request_counts(self):
         if not self.publish_dp_lb_stats:
             return
 
-        # Publish our request counts (if they've changed).
+        # Publish our request counts (if they've changed), stamped with the
+        # lockstep-synchronized step counter and wave number.
         counts = self.scheduler.get_request_counts()
         if counts != self.last_counts:
             self.last_counts = counts
             stats = SchedulerStats(
-                *counts, step_counter=self.step_counter, current_wave=self.current_wave
+                *counts,
+                kv_cache_usage=self.scheduler.get_kv_cache_usage(),
+                step_counter=self.step_counter,
+                current_wave=self.current_wave,
             )
             self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats)))