@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 59a6b0411d

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

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nick Hill committed Jul 27, 2026 at 07:00 UTC 59a6b0411d1817c712728a17fa8c2e0cdb4cf1f9

7 files changed +188 -59

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

+4 -8

@@ -74,17 +74,13 @@ impl EngineRoutingState {

74 ///

75 /// Scheduler stats can raise the load estimate above the frontend-local

76 /// view, but they should not lower it below requests this frontend has

77 - /// already admitted. Waiting requests still get the same extra penalty

78 - /// as the original `waiting * 4 + running` score.

77 + /// already admitted.

78 fn routing_score(&self) -> usize {

80 - const WAITING_WEIGHT: usize = 4;

81 -

79 let Some(stats) = self.last_scheduler_stats else {

80 return self.inflight;

81 };

86 - let scheduler_total = stats.running + stats.waiting;

87 - self.inflight.max(scheduler_total) + stats.waiting * (WAITING_WEIGHT - 1)

83 + self.inflight.max(stats.running + stats.waiting)

84 }

86 /// Replace the local routing view with a fresh real scheduler snapshot.

@@ -750,7 +746,7 @@ mod tests {

746 }

747

748 #[test]

753 - fn routing_score_keeps_extra_waiting_penalty() {

749 + fn routing_score_counts_waiting_without_extra_penalty() {

750 let state = EngineRoutingState {

751 inflight: 1,

752 last_scheduler_stats: Some(EngineLoadSnapshot {

@@ -759,7 +755,7 @@ mod tests {

755 }),

756 };

757

762 - assert_eq!(state.routing_score(), 14);

758 + assert_eq!(state.routing_score(), 5);

759 }

760

761 #[test]

tests/v1/engine/test_engine_core_client.py

+82 -11

@@ -8,6 +8,7 @@ import os

8 import signal

9 import time

10 import uuid

11 +from collections import Counter

12 from concurrent.futures import Future

13 from dataclasses import dataclass

14 from threading import Thread

@@ -27,7 +28,11 @@ from vllm.platforms import current_platform

28 from vllm.pooling_params import LateInteractionParams, PoolingParams

29 from vllm.usage.usage_lib import UsageContext

30 from vllm.utils.torch_utils import set_default_torch_num_threads

30 -from vllm.v1.engine import EngineCoreReadyResponse, EngineCoreRequest

31 +from vllm.v1.engine import (

32 + EngineCoreOutputs,

33 + EngineCoreReadyResponse,

34 + EngineCoreRequest,

35 +)

36 from vllm.v1.engine.core import EngineCore

37 from vllm.v1.engine.core_client import (

38 AsyncMPClient,

@@ -198,13 +203,19 @@ def _make_pooling_request(

203 )

204

205

201 -def test_dplb_late_interaction_sticky_routing():

206 +def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncMPClient:

207 client = object.__new__(DPLBAsyncMPClient)

203 - client.client_count = 1

208 + client.client_count = client_count

209 client.reqs_in_flight = {}

205 - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]

206 - client.lb_engines = [[0, 0], [0, 0], [0, 0]]

210 + client.engine_inflight = Counter()

211 + client.core_engines = [bytes([i, 0]) for i in range(num_engines)]

212 + client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]

213 client.eng_start_index = 0

214 + return client

215 +

216 +

217 +def test_dplb_late_interaction_sticky_routing():

218 + client = _make_dplb_client()

219

220 query_key = "rerank-abc-query-0"

221 query_request = _make_pooling_request(

@@ -223,12 +234,8 @@ def test_dplb_late_interaction_sticky_routing():

234

235

236 def test_dplb_non_late_interaction_still_uses_lb():

226 - client = object.__new__(DPLBAsyncMPClient)

227 - client.client_count = 1

228 - client.reqs_in_flight = {}

229 - client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]

230 - client.lb_engines = [[2, 1], [0, 0], [1, 0]]

231 - client.eng_start_index = 0

237 + client = _make_dplb_client()

238 + client.lb_engines = [[2, 1, 0.0], [0, 0, 0.0], [1, 0, 0.0]]

239

240 request = make_request(SamplingParams(max_tokens=1))

241 chosen_engine = client.get_core_engine_for_request(request)

@@ -237,6 +244,70 @@ def test_dplb_non_late_interaction_still_uses_lb():

244 assert client.lb_engines[1][0] == 1

245

246

247 +def test_dplb_burst_round_robins_despite_snapshot_rebinds():

248 + """A stats snapshot rebind wipes the optimistic lb_engines increments;

249 + the exact in-flight floor must keep a burst spreading round-robin."""

250 + client = _make_dplb_client(num_engines=4)

251 +

252 + for _ in range(4):

253 + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))

254 + # Coordinator snapshot arrives, not yet reflecting the 4 routed requests.

255 + client.lb_engines = [[0, 0, 0.0] for _ in range(4)]

256 + for _ in range(4):

257 + client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))

258 +

259 + assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2]

260 +

261 +

262 +def test_dplb_snapshot_backpressure_overrides_inflight():

263 + """An engine reported heavily loaded by the coordinator is avoided even

264 + when this client has routed nothing to it."""

265 + client = _make_dplb_client(num_engines=2)

266 + client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]]

267 +

268 + chosen = client.get_core_engine_for_request(

269 + make_request(SamplingParams(max_tokens=1))

270 + )

271 +

272 + assert chosen == client.core_engines[1]

273 +

274 +

275 +def test_dplb_kv_pressure_amplifies_waiting_penalty():

276 + """A waiting queue on a KV-bound engine (slow drain) is penalized, while

277 + the same queue with low KV usage is not (e.g. transient burst)."""

278 + client = _make_dplb_client(num_engines=2)

279 + # Engine 0 has a smaller total but is KV-bound with a queue.

280 + client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]]

281 +

282 + chosen = client.get_core_engine_for_request(

283 + make_request(SamplingParams(max_tokens=1))

284 + )

285 + assert chosen == client.core_engines[1]

286 +

287 + # Same counts without KV pressure: the smaller total wins.

288 + client = _make_dplb_client(num_engines=2)

289 + client.lb_engines = [[5, 10, 0.2], [0, 20, 0.2]]

290 +

291 + chosen = client.get_core_engine_for_request(

292 + make_request(SamplingParams(max_tokens=1))

293 + )

294 + assert chosen == client.core_engines[0]

295 +

296 +

297 +def test_dplb_finished_requests_release_inflight():

298 + client = _make_dplb_client(num_engines=2)

299 +

300 + req = make_request(SamplingParams(max_tokens=1))

301 + engine = client.get_core_engine_for_request(req)

302 + assert client.engine_inflight[engine] == 1

303 +

304 + outputs = EngineCoreOutputs(finished_requests={req.request_id})

305 + asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs))

306 +

307 + assert client.engine_inflight[engine] == 0

308 + assert req.request_id not in client.reqs_in_flight

309 +

310 +

311 def test_apply_ready_response_syncs_block_size():

312 import msgspec

313

vllm/v1/core/sched/interface.py

@@ -234,6 +234,10 @@ class SchedulerInterface(ABC):

234 """Returns (num_running_reqs, num_waiting_reqs)."""

235 raise NotImplementedError

236

237 + def get_kv_cache_usage(self) -> float:

238 + """Returns the fraction of the KV cache currently in use (0.0-1.0)."""

239 + return 0.0

240 +

241 @abstractmethod

242 def make_stats(self) -> "SchedulerStats | None":

243 """Make a SchedulerStats object for logging.

vllm/v1/core/sched/scheduler.py

@@ -2154,6 +2154,10 @@ class Scheduler(SchedulerInterface):

2154 """Returns (num_running_reqs, num_waiting_reqs)."""

2155 return len(self.running), len(self.waiting) + len(self.skipped_waiting)

2156

2157 + def get_kv_cache_usage(self) -> float:

2158 + """Returns the fraction of the KV cache currently in use (0.0-1.0)."""

2159 + return self.kv_cache_manager.usage

2160 +

2161 def add_request(self, request: Request) -> None:

2162 existing = self.requests.get(request.request_id)

2163 if existing is not None:

vllm/v1/engine/coordinator.py

+40 -29

@@ -139,7 +139,8 @@ class DPCoordinator:

139

140 class EngineState:

141 def __init__(self):

142 - self.request_counts = [0, 0] # [waiting, running]

142 + # [waiting, running, kv_cache_usage]

143 + self.request_counts: list[int | float] = [0, 0, 0.0]

144

145

146 class DPCoordinatorProc:

@@ -202,7 +203,7 @@ class DPCoordinatorProc:

203 stats_changed = False

204 last_stats_step = -1

205 last_stats_wave = -1

205 - last_step_counts: list[list[int]] | None = None

206 + last_step_counts: list[list[int | float]] | None = None

207

208 with (

209 make_zmq_socket(

@@ -260,8 +261,12 @@ class DPCoordinatorProc:

261 wait_for = self.stats_update_interval_ms if stats_changed else 5000

262

263 # Wait at least 50ms to ensure we've received all stats for

263 - # the current step.

264 - min_timeout = 50 if last_step_counts is None else 0

264 + # the current step. Only applicable to lockstep (MoE) DP;

265 + # non-lockstep engines have no synchronized step boundaries.

266 + if self.enable_wave_coordination and last_step_counts is None:

267 + min_timeout = 50

268 + else:

269 + min_timeout = 0

270

271 events = poller.poll(timeout=max(min_timeout, wait_for - elapsed))

272 if not events:

@@ -374,32 +379,38 @@ class DPCoordinatorProc:

379 # 1. Updated request load stats - update our local

380 # state with these.

381 stats = self.engines[eng_index].request_counts

377 - stats_step = scheduler_stats.step_counter

378 - stats_wave = scheduler_stats.current_wave

379 - if (

380 - stats_wave > last_stats_wave

381 - or stats_wave == last_stats_wave

382 - and stats_step > last_stats_step

383 - ):

384 - if stats_changed:

385 - last_step_counts = self._get_engine_counts(do_copy=True)

386 - last_stats_step = stats_step

387 - last_stats_wave = stats_wave

388 - elif stats_wave != last_stats_wave or (

389 - stats_step != last_stats_step

390 - ):

391 - logger.warning(

392 - "Received stats for out-of-order "

393 - "step (%d, %d) from engine %d (expected "

394 - "> (%d, %d))",

395 - stats_wave,

396 - stats_step,

397 - eng_index,

398 - last_stats_wave,

399 - last_stats_step,

400 - )

382 + if self.enable_wave_coordination:

383 + # Steps are synchronized across lockstep (MoE) DP

384 + # ranks; snapshot counts at step boundaries.

385 + stats_step = scheduler_stats.step_counter

386 + stats_wave = scheduler_stats.current_wave

387 + if (

388 + stats_wave > last_stats_wave

389 + or stats_wave == last_stats_wave

390 + and stats_step > last_stats_step

391 + ):

392 + if stats_changed:

393 + last_step_counts = self._get_engine_counts(

394 + do_copy=True

395 + )

396 + last_stats_step = stats_step

397 + last_stats_wave = stats_wave

398 + elif stats_wave != last_stats_wave or (

399 + stats_step != last_stats_step

400 + ):

401 + logger.warning(

402 + "Received stats for out-of-order "

403 + "step (%d, %d) from engine %d (expected "

404 + "> (%d, %d))",

405 + stats_wave,

406 + stats_step,

407 + eng_index,

408 + last_stats_wave,

409 + last_stats_step,

410 + )

411 stats[0] = scheduler_stats.num_waiting_reqs

412 stats[1] = scheduler_stats.num_running_reqs

413 + stats[2] = scheduler_stats.kv_cache_usage

414 stats_changed = True

415

416 # Wave coordination: handle wave completion and start notifications

@@ -452,7 +463,7 @@ class DPCoordinatorProc:

463 wave_encoded = msgspec.msgpack.encode((wave, exclude_engine_index))

464 socket.send_multipart((EngineCoreRequestType.START_DP_WAVE.value, wave_encoded))

465

455 - def _get_engine_counts(self, do_copy=False) -> list[list[int]]:

466 + def _get_engine_counts(self, do_copy=False) -> list[list[int | float]]:

467 """Return list of [waiting, running] count lists for each engine."""

468 if do_copy:

469 return [copy.copy(e.request_counts) for e in self.engines]

vllm/v1/engine/core.py

+23 -3

@@ -1057,6 +1057,7 @@ class EngineCoreProc(EngineCore):

1057 # Only publish request queue stats to coordinator for "internal"

1058 # and "hybrid" LB modes.

1059 self.publish_dp_lb_stats = internal_dp_balancing

1060 + self.last_counts = (0, 0)

1061

1062 self.addresses = addresses

1063 self.process_input_queue_block = True

@@ -1376,11 +1377,27 @@ class EngineCoreProc(EngineCore):

1377 while self._handle_shutdown():

1378 # 1) Poll the input queue until there is work to do.

1379 self._process_input_queue()

1380 + # Publish request counts before and after GPU step to ensure freshness.

1381 + self._maybe_publish_request_counts()

1382 # 2) Step the engine core and return the outputs.

1383 self._process_engine_step()

1384 + self._maybe_publish_request_counts()

1385

1386 raise SystemExit

1387

1388 + def _maybe_publish_request_counts(self):

1389 + if not self.publish_dp_lb_stats:

1390 + return

1391 +

1392 + # Publish our request counts (if they've changed).

1393 + counts = self.scheduler.get_request_counts()

1394 + if counts != self.last_counts:

1395 + self.last_counts = counts

1396 + stats = SchedulerStats(

1397 + *counts, kv_cache_usage=self.scheduler.get_kv_cache_usage()

1398 + )

1399 + self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats)))

1400 +

1401 def _process_input_queue(self):

1402 """Exits when an engine step needs to be performed."""

1403

@@ -1890,7 +1907,6 @@ class DPEngineCoreProc(EngineCoreProc):

1907 # finished with DP peers every N steps.

1908 self.step_counter = 0

1909 self.current_wave = 0

1893 - self.last_counts = (0, 0)

1910

1911 # Two-phase pause protocol state. When pending_pause is True, the

1912 # engine keeps stepping (dummy batches) while waiting for all DP

@@ -2027,12 +2043,16 @@ class DPEngineCoreProc(EngineCoreProc):

2043 if not self.publish_dp_lb_stats:

2044 return

2045

2030 - # Publish our request counts (if they've changed).

2046 + # Publish our request counts (if they've changed), stamped with the

2047 + # lockstep-synchronized step counter and wave number.

2048 counts = self.scheduler.get_request_counts()

2049 if counts != self.last_counts:

2050 self.last_counts = counts

2051 stats = SchedulerStats(

2035 - *counts, step_counter=self.step_counter, current_wave=self.current_wave

2052 + *counts,

2053 + kv_cache_usage=self.scheduler.get_kv_cache_usage(),

2054 + step_counter=self.step_counter,

2055 + current_wave=self.current_wave,

2056 )

2057 self.output_queue.put_nowait((-1, EngineCoreOutputs(scheduler_stats=stats)))

2058

vllm/v1/engine/core_client.py

+31 -8

@@ -7,7 +7,7 @@ import sys

7 import uuid

8 import weakref

9 from abc import ABC, abstractmethod

10 -from collections import defaultdict, deque

10 +from collections import Counter, defaultdict, deque

11 from collections.abc import Awaitable, Callable, Sequence

12 from concurrent.futures import Future

13 from dataclasses import dataclass

@@ -1271,9 +1271,11 @@ class DPAsyncMPClient(AsyncMPClient):

1271 client_index,

1272 )

1273

1274 - # List of [waiting, running] pair per engine.

1274 + # List of [waiting, running, kv_cache_usage] per engine.

1275 # Used only by DPLBAsyncMPClient subclass.

1276 - self.lb_engines: list[list[int]] = [[0, 0] for _ in self.core_engines]

1276 + self.lb_engines: list[list[int | float]] = [

1277 + [0, 0, 0.0] for _ in self.core_engines

1278 + ]

1279

1280 self.eep_scaling_cache: ElasticScalingCache | None = None

1281

@@ -1351,7 +1353,7 @@ class DPAsyncMPClient(AsyncMPClient):

1353 )

1354 if len(self.lb_engines) < new_engine_count:

1355 self.lb_engines = self.lb_engines + [

1354 - [0, 0]

1356 + [0, 0, 0.0]

1357 for _ in range(

1358 new_engine_count - len(self.lb_engines)

1359 )

@@ -1445,6 +1447,9 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1447 # To route aborts to the correct engine.

1448 self.reqs_in_flight: dict[str, EngineIdentity] = {}

1449

1450 + # Exact per-engine count of this client's unfinished requests.

1451 + self.engine_inflight: Counter[EngineIdentity] = Counter()

1452 +

1453 super().__init__(

1454 vllm_config,

1455 executor_class,

@@ -1470,14 +1475,30 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1475 current_counts = self.lb_engines

1476 # TODO use P2C alg for larger DP sizes

1477 num_engines = len(current_counts)

1473 - min_score = sys.maxsize

1478 + min_score: float = sys.maxsize

1479 eng_index = 0

1480 for i in range(num_engines):

1481 # Start from client_index to help with balancing when engines

1482 # are empty.

1483 idx = (self.eng_start_index + i) % num_engines

1479 - waiting, running = current_counts[idx]

1480 - score = waiting * 4 + running

1484 + waiting, running, kv_cache_usage = current_counts[idx]

1485 + # Estimate engine load as the greater of the coordinator's

1486 + # latest (waiting + running) snapshot and this client's own

1487 + # in-flight count (scaled by the number of clients). The

1488 + # in-flight floor is exact and can't be erased by a snapshot

1489 + # rebind, so a burst spreads round-robin even when snapshots

1490 + # race with routing decisions; the snapshot raises the score

1491 + # when other clients or stale requests load the engine.

1492 + inflight = self.engine_inflight[self.core_engines[idx]]

1493 + score: float = max(self.client_count * inflight, waiting + running)

1494 + if waiting:

1495 + # Waiting requests are penalized in proportion to KV cache

1496 + # pressure: a queue on a KV-bound engine drains slowly, so

1497 + # new requests should strongly prefer other engines. With

1498 + # low KV usage the queue is transient (e.g. mid-burst) and

1499 + # the penalty stays off, preserving exact round-robin.

1500 + # Ramps from 0 at <=50% usage to 3x waiting at 100%.

1501 + score += waiting * 6.0 * max(0.0, kv_cache_usage - 0.5)

1502 if score < min_score:

1503 min_score = score

1504 eng_index = idx

@@ -1494,6 +1515,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1515 chosen_engine = self.core_engines[eng_index]

1516 # Record which engine is chosen for this request, to handle aborts.

1517 self.reqs_in_flight[request.request_id] = chosen_engine

1518 + self.engine_inflight[chosen_engine] += 1

1519 return chosen_engine

1520

1521 async def call_utility_async(self, method: str, *args) -> Any:

@@ -1513,7 +1535,8 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1535 ):

1536 if outputs.finished_requests and self.reqs_in_flight:

1537 for req_id in outputs.finished_requests:

1516 - self.reqs_in_flight.pop(req_id, None)

1538 + if (engine := self.reqs_in_flight.pop(req_id, None)) is not None:

1539 + self.engine_inflight[engine] -= 1

1540

1541 @staticmethod

1542 async def eep_process_engine_core_notification(