iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Core] Fix internal LB load-balancing (#49204)

Browse Source 
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Nick Hill
2026-07-27 14:00:46 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent dbccc5ae32
commit 59a6b0411d
7 changed files with 188 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/engine-core-client/src/client/state.rs
+4 -8
View File
@@ -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
+82 -11
View File
@@ -8,6 +8,7 @@ import os
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.platforms import current_platform
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
+4
View File
@@ -234,6 +234,10 @@ class SchedulerInterface(ABC):
"""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
+4
View File
@@ -2154,6 +2154,10 @@ class Scheduler(SchedulerInterface):
"""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
+40 -29
View File
@@ -139,7 +139,8 @@ class DPCoordinator:
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 @@ class DPCoordinatorProc:
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 @@ class DPCoordinatorProc:
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 @@ class DPCoordinatorProc:
# 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 @@ class DPCoordinatorProc:
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
+23 -3
View File
@@ -1057,6 +1057,7 @@ class EngineCoreProc(EngineCore):
# 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 @@ class EngineCoreProc(EngineCore):
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 @@ class DPEngineCoreProc(EngineCoreProc):
# 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 @@ class DPEngineCoreProc(EngineCoreProc):
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)))
vllm/v1/engine/core_client.py
+31 -8
View File
@@ -7,7 +7,7 @@ import sys
import uuid
import weakref
from abc import ABC, abstractmethod
from collections import defaultdict, deque
from collections import Counter, defaultdict, deque
from collections.abc import Awaitable, Callable, Sequence
from concurrent.futures import Future
from dataclasses import dataclass
@@ -1271,9 +1271,11 @@ class DPAsyncMPClient(AsyncMPClient):
client_index,
)
# List of [waiting, running] pair per engine.
# List of [waiting, running, kv_cache_usage] per engine.
# Used only by DPLBAsyncMPClient subclass.
self.lb_engines: list[list[int]] = [[0, 0] for _ in self.core_engines]
self.lb_engines: list[list[int | float]] = [
[0, 0, 0.0] for _ in self.core_engines
]
self.eep_scaling_cache: ElasticScalingCache | None = None
@@ -1351,7 +1353,7 @@ class DPAsyncMPClient(AsyncMPClient):
)
if len(self.lb_engines) < new_engine_count:
self.lb_engines = self.lb_engines + [
[0, 0]
[0, 0, 0.0]
for _ in range(
new_engine_count - len(self.lb_engines)
)
@@ -1445,6 +1447,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
# To route aborts to the correct engine.
self.reqs_in_flight: dict[str, EngineIdentity] = {}
# Exact per-engine count of this client's unfinished requests.
self.engine_inflight: Counter[EngineIdentity] = Counter()
super().__init__(
vllm_config,
executor_class,
@@ -1470,14 +1475,30 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
current_counts = self.lb_engines
# TODO use P2C alg for larger DP sizes
num_engines = len(current_counts)
min_score = sys.maxsize
min_score: float = sys.maxsize
eng_index = 0
for i in range(num_engines):
# Start from client_index to help with balancing when engines
# are empty.
idx = (self.eng_start_index + i) % num_engines
waiting, running = current_counts[idx]
score = waiting * 4 + running
waiting, running, kv_cache_usage = current_counts[idx]
# Estimate engine load as the greater of the coordinator's
# latest (waiting + running) snapshot and this client's own
# in-flight count (scaled by the number of clients). The
# in-flight floor is exact and can't be erased by a snapshot
# rebind, so a burst spreads round-robin even when snapshots
# race with routing decisions; the snapshot raises the score
# when other clients or stale requests load the engine.
inflight = self.engine_inflight[self.core_engines[idx]]
score: float = max(self.client_count * inflight, waiting + running)
if waiting:
# Waiting requests are penalized in proportion to KV cache
# pressure: a queue on a KV-bound engine drains slowly, so
# new requests should strongly prefer other engines. With
# low KV usage the queue is transient (e.g. mid-burst) and
# the penalty stays off, preserving exact round-robin.
# Ramps from 0 at <=50% usage to 3x waiting at 100%.
score += waiting * 6.0 * max(0.0, kv_cache_usage - 0.5)
if score < min_score:
min_score = score
eng_index = idx
@@ -1494,6 +1515,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
chosen_engine = self.core_engines[eng_index]
# Record which engine is chosen for this request, to handle aborts.
self.reqs_in_flight[request.request_id] = chosen_engine
self.engine_inflight[chosen_engine] += 1
return chosen_engine
async def call_utility_async(self, method: str, *args) -> Any:
@@ -1513,7 +1535,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
):
if outputs.finished_requests and self.reqs_in_flight:
for req_id in outputs.finished_requests:
self.reqs_in_flight.pop(req_id, None)
if (engine := self.reqs_in_flight.pop(req_id, None)) is not None:
self.engine_inflight[engine] -= 1
@staticmethod
async def eep_process_engine_core_notification(