@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8d524ce79f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 22:35 UTC

[BugFix] Improve internal DP load balancing (#21617)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Aug 2, 2025 at 03:45 UTC 8d524ce79ffd0571d6a576cb9a5c21feab187246

7 files changed +122 -59

vllm/entrypoints/openai/api_server.py

@@ -199,6 +199,8 @@ async def build_async_engine_client_from_engine_args(

199

200 from vllm.v1.engine.async_llm import AsyncLLM

201 async_llm: Optional[AsyncLLM] = None

202 + client_count = client_config.pop(

203 + "client_count") if client_config else 1

204 client_index = client_config.pop(

205 "client_index") if client_config else 0

206 try:

@@ -208,6 +210,7 @@ async def build_async_engine_client_from_engine_args(

210 enable_log_requests=engine_args.enable_log_requests,

211 disable_log_stats=engine_args.disable_log_stats,

212 client_addresses=client_config,

213 + client_count=client_count,

214 client_index=client_index)

215

216 # Don't keep the dummy data in memory

vllm/v1/engine/async_llm.py

@@ -57,6 +57,7 @@ class AsyncLLM(EngineClient):

57 start_engine_loop: bool = True,

58 stat_loggers: Optional[list[StatLoggerFactory]] = None,

59 client_addresses: Optional[dict[str, str]] = None,

60 + client_count: int = 1,

61 client_index: int = 0,

62 ) -> None:

63 """

@@ -120,6 +121,7 @@ class AsyncLLM(EngineClient):

121 executor_class=executor_class,

122 log_stats=self.log_stats,

123 client_addresses=client_addresses,

124 + client_count=client_count,

125 client_index=client_index,

126 )

127

@@ -156,6 +158,7 @@ class AsyncLLM(EngineClient):

158 enable_log_requests: bool = False,

159 disable_log_stats: bool = False,

160 client_addresses: Optional[dict[str, str]] = None,

161 + client_count: int = 1,

162 client_index: int = 0,

163 disable_log_requests: bool = True, # Deprecated, will be removed

164 ) -> "AsyncLLM":

@@ -176,6 +179,7 @@ class AsyncLLM(EngineClient):

179 log_stats=not disable_log_stats,

180 usage_context=usage_context,

181 client_addresses=client_addresses,

182 + client_count=client_count,

183 client_index=client_index,

184 )

185

vllm/v1/engine/coordinator.py

+73 -37

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import copy

4 import multiprocessing

5 import time

6 import weakref

@@ -65,18 +66,14 @@ class DPCoordinator:

67 # Assume coordinator is colocated with front-end procs when not in

68 # either external or hybrid DP LB mode.

69 + local_only = not (external_lb or hybrid_lb)

70 front_publish_address = get_engine_client_zmq_addr(

69 - local_only=not external_lb and not hybrid_lb, host=host)

71 + local_only=local_only, host=host)

73 local_only_eng = dp_size == parallel_config.data_parallel_size_local

74 back_publish_address = get_engine_client_zmq_addr(local_only_eng, host)

75 back_output_address = get_engine_client_zmq_addr(local_only_eng, host)

75 - # When in external LB mode, load stats aren't published, only changes

76 - # to request wave / running state, so we don't need to rate-limit the

77 - # updates to the front-end proc(s).

78 - min_stats_update_interval_ms = 0 if external_lb else 100

79 -

77 context = get_mp_context()

78 self.proc: multiprocessing.Process = context.Process(

79 target=DPCoordinatorProc.run_coordinator,

@@ -86,7 +83,6 @@ class DPCoordinator:

83 "front_publish_address": front_publish_address,

84 "back_output_address": back_output_address,

85 "back_publish_address": back_publish_address,

89 - "min_stats_update_interval_ms": min_stats_update_interval_ms,

86 },

87 daemon=True)

88 self.proc.start()

@@ -125,10 +121,6 @@ class DPCoordinatorProc:

121

122 self.stats_update_interval_ms = min_stats_update_interval_ms

123

128 - self.current_wave = 0

129 - self.engines_running = False

130 - self.stats_changed = False

131 -

124 @staticmethod

125 def run_coordinator(

126 engine_count: int,

@@ -155,6 +147,16 @@ class DPCoordinatorProc:

147

148 decoder = MsgpackDecoder(EngineCoreOutputs)

149

150 + # For tracking request wave progression.

151 + current_wave = 0

152 + engines_running = False

153 +

154 + # For tracking request counts for internal load-balancing.

155 + stats_changed = False

156 + last_stats_step = -1

157 + last_stats_wave = -1

158 + last_step_counts: Optional[list[list[int]]] = None

159 +

160 with make_zmq_socket(

161 path=front_publish_address, # IPC

162 ctx=self.ctx,

@@ -191,21 +193,33 @@ class DPCoordinatorProc:

193 while True:

194 elapsed = int(time.time() * 1000) - last_publish_time

195 # Send at stats_update_interval_ms interval if the stats have

194 - # changed, or otherwise every 4 seconds.

196 + # changed, or otherwise every 5 seconds.

197 wait_for = (self.stats_update_interval_ms

196 - if self.stats_changed else 4000)

197 - events = poller.poll(timeout=max(0, wait_for - elapsed))

198 + if stats_changed else 5000)

199 +

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

201 + # the current step.

202 + min_timeout = 50 if last_step_counts is None else 0

203 +

204 + events = poller.poll(timeout=max(min_timeout, wait_for -

205 + elapsed))

206 if not events:

207 # Poller timeout - publish current stats to front-ends.

200 - engine_req_counts_list = self._get_engine_counts()

201 - to_publish = (engine_req_counts_list, self.current_wave,

202 - self.engines_running)

208 + if last_step_counts is not None:

209 + engine_req_counts_list = last_step_counts

210 + last_step_counts = None

211 + else:

212 + engine_req_counts_list = self._get_engine_counts()

213 + stats_changed = False

214 +

215 + to_publish = (engine_req_counts_list, current_wave,

216 + engines_running)

217 publish_front.send(msgspec.msgpack.encode(to_publish))

218 last_publish_time = int(time.time() * 1000)

205 - self.stats_changed = False

219 continue

220

221 events = dict(events)

222 + wave_state_changed = False

223

224 if publish_front in events:

225 buffer = publish_front.recv()

@@ -232,7 +246,7 @@ class DPCoordinatorProc:

246 # current_wave

247 # we note that 0 is the wave number for the new

248 # engine

235 - self.engines_running = False

249 + engines_running = False

250 logger.info(

251 "DPCoordinator scaled up from %s to %s "

252 "engines", current_count, new_engine_count)

@@ -248,15 +262,15 @@ class DPCoordinatorProc:

262 # engines are paused, so that we can wake the other

263 # engines.

264 engine_to_exclude, wave = decoded

251 - if not self.engines_running:

252 - if wave < self.current_wave:

265 + if not engines_running:

266 + if wave < current_wave:

267 # If the wave number is stale, ensure the message

268 # is handled by all the engines.

269 engine_to_exclude = None

270

257 - self.engines_running = True

258 - self.stats_changed = True

259 - self._send_start_wave(publish_back, self.current_wave,

271 + engines_running = True

272 + wave_state_changed = True

273 + self._send_start_wave(publish_back, current_wave,

274 engine_to_exclude)

275

276 if output_back in events:

@@ -274,36 +288,56 @@ class DPCoordinatorProc:

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

289 # state with these.

290 stats = self.engines[eng_index].request_counts

291 + stats_step = scheduler_stats.step_counter

292 + stats_wave = scheduler_stats.current_wave

293 + if (stats_wave > last_stats_wave

294 + or stats_wave == last_stats_wave

295 + and stats_step > last_stats_step):

296 + if stats_changed:

297 + last_step_counts = self._get_engine_counts(

298 + do_copy=True)

299 + last_stats_step = stats_step

300 + last_stats_wave = stats_wave

301 + elif stats_wave != last_stats_wave or (

302 + stats_step != last_stats_step):

303 + logger.warning(

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

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

306 + "> (%d, %d))", stats_wave, stats_step,

307 + eng_index, last_stats_wave, last_stats_step)

308 stats[0] = scheduler_stats.num_waiting_reqs

309 stats[1] = scheduler_stats.num_running_reqs

279 - self.stats_changed = True

310 + stats_changed = True

311

312 if (wave := outputs.wave_complete) is not None:

313 # 2. Notification from rank 0 engine that we've

314 # moved into the global paused state

315 # (engines_running==False).

285 - if self.current_wave <= wave:

316 + if current_wave <= wave:

317 new_wave = wave + 1

318 logger.debug("Moving DP wave from %d to %d.",

288 - self.current_wave, new_wave)

289 - self.current_wave = new_wave

290 - self.engines_running = False

291 - self.stats_changed = True

319 + current_wave, new_wave)

320 + current_wave = new_wave

321 + engines_running = False

322 + wave_state_changed = True

323 elif (wave := outputs.start_wave) is not None and (

293 - wave > self.current_wave or

294 - (wave == self.current_wave

295 - and not self.engines_running)):

324 + wave > current_wave or

325 + (wave == current_wave and not engines_running)):

326 # 3. The engine received request for a non-current wave

327 # so we must ensure that other engines progress to the

328 # next wave (race condition handling).

329 logger.debug(

330 "Starting wave %d after notification of "

331 "stale wave request from engine.", wave)

302 - self.current_wave = wave

303 - self.engines_running = True

304 - self.stats_changed = True

332 + current_wave = wave

333 + engines_running = True

334 + wave_state_changed = True

335 self._send_start_wave(publish_back, wave, eng_index)

336

337 + if wave_state_changed:

338 + message = (None, current_wave, engines_running)

339 + publish_front.send(msgspec.msgpack.encode(message))

340 +

341 @staticmethod

342 def _send_start_wave(socket: zmq.Socket, wave: int,

343 exclude_engine_index: Optional[int]):

@@ -316,6 +350,8 @@ class DPCoordinatorProc:

350 socket.send_multipart(

351 (EngineCoreRequestType.START_DP_WAVE.value, wave_encoded))

352

319 - def _get_engine_counts(self) -> list[list[int]]:

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

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

355 + if do_copy:

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

357 return [e.request_counts for e in self.engines]

vllm/v1/engine/core.py

+8 -5

@@ -928,7 +928,7 @@ class DPEngineCoreProc(EngineCoreProc):

928 ):

929 # Counts forward-passes of the model so that we can synchronize

930 # finished with DP peers every N steps.

931 - self.counter = 0

931 + self.step_counter = 0

932 self.current_wave = 0

933 self.last_counts = (0, 0)

934

@@ -999,7 +999,9 @@ class DPEngineCoreProc(EngineCoreProc):

999 counts = self.scheduler.get_request_counts()

1000 if counts != self.last_counts:

1001 self.last_counts = counts

1002 - stats = SchedulerStats(*counts)

1002 + stats = SchedulerStats(*counts,

1003 + step_counter=self.step_counter,

1004 + current_wave=self.current_wave)

1005 self.output_queue.put_nowait(

1006 (-1, EngineCoreOutputs(scheduler_stats=stats)))

1007

@@ -1041,15 +1043,16 @@ class DPEngineCoreProc(EngineCoreProc):

1043 self.output_queue.put_nowait(

1044 (client_index,

1045 EngineCoreOutputs(wave_complete=self.current_wave)))

1046 + # Increment wave count and reset step counter.

1047 self.current_wave += 1

1048 + self.step_counter = 0

1049

1050 def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:

1051

1052 # Optimization - only perform finish-sync all-reduce every 32 steps.

1049 - self.counter += 1

1050 - if self.counter != 32:

1053 + self.step_counter += 1

1054 + if self.step_counter % 32 != 0:

1055 return True

1052 - self.counter = 0

1056

1057 return ParallelConfig.has_unfinished_dp(self.dp_group,

1058 local_unfinished)

vllm/v1/engine/core_client.py

+29 -17

@@ -86,11 +86,12 @@ class EngineCoreClient(ABC):

86 executor_class: type[Executor],

87 log_stats: bool,

88 client_addresses: Optional[dict[str, str]] = None,

89 + client_count: int = 1,

90 client_index: int = 0,

91 ) -> "MPClient":

92 parallel_config = vllm_config.parallel_config

93 client_args = (vllm_config, executor_class, log_stats,

93 - client_addresses, client_index)

94 + client_addresses, client_count, client_index)

95 if parallel_config.data_parallel_size > 1:

96 if parallel_config.data_parallel_external_lb:

97 # External load balancer - client per DP rank.

@@ -727,6 +728,7 @@ class AsyncMPClient(MPClient):

728 executor_class: type[Executor],

729 log_stats: bool,

730 client_addresses: Optional[dict[str, str]] = None,

731 + client_count: int = 1,

732 client_index: int = 0):

733 super().__init__(

734 asyncio_mode=True,

@@ -929,11 +931,12 @@ class DPAsyncMPClient(AsyncMPClient):

931 executor_class: type[Executor],

932 log_stats: bool,

933 client_addresses: Optional[dict[str, str]] = None,

934 + client_count: int = 1,

935 client_index: int = 0):

936 self.current_wave = 0

937

938 super().__init__(vllm_config, executor_class, log_stats,

936 - client_addresses, client_index)

939 + client_addresses, client_count, client_index)

940

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

942 # Used only by DPLBAsyncMPClient subclass.

@@ -1029,7 +1032,11 @@ class DPAsyncMPClient(AsyncMPClient):

1032 counts, wave, running = msgspec.msgpack.decode(buf)

1033 self.current_wave = wave

1034 self.engines_running = running

1032 - self.lb_engines = counts[count_slice]

1035 + if counts is not None:

1036 + sliced_counts = counts[count_slice]

1037 + self.lb_engines = sliced_counts

1038 + logger.debug("Received counts: %s (%s)", sliced_counts,

1039 + count_slice)

1040

1041 resources.stats_update_task = asyncio.create_task(

1042 run_engine_stats_update_task())

@@ -1065,40 +1072,45 @@ class DPLBAsyncMPClient(DPAsyncMPClient):

1072 executor_class: type[Executor],

1073 log_stats: bool,

1074 client_addresses: Optional[dict[str, str]] = None,

1075 + client_count: int = 1,

1076 client_index: int = 0):

1077

1078 + self.client_count = client_count

1079 +

1080 # To route aborts to the correct engine.

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

1082

1083 super().__init__(vllm_config, executor_class, log_stats,

1074 - client_addresses, client_index)

1084 + client_addresses, client_count, client_index)

1085

1086 assert len(self.core_engines) > 1

1087

1088 + self.eng_start_index = (len(self.core_engines) *

1089 + self.client_index) // client_count

1090 +

1091 def get_core_engine_for_request(

1092 self, request: EngineCoreRequest) -> EngineIdentity:

1093 # Engines are in rank order.

1094 + current_counts = self.lb_engines

1095 if (eng_index := request.data_parallel_rank) is None:

1082 - if not self.lb_engines:

1096 + if not current_counts:

1097 return self.core_engine

1098 # TODO use P2C alg for larger DP sizes

1085 - num_engines = len(self.lb_engines)

1086 - min_counts = [sys.maxsize, sys.maxsize]

1099 + num_engines = len(current_counts)

1100 + min_score = sys.maxsize

1101 eng_index = 0

1102 for i in range(num_engines):

1103 # Start from client_index to help with balancing when engines

1104 # are empty.

1091 - idx = (self.client_index + i) % num_engines

1092 - counts = self.lb_engines[idx]

1093 - if counts < min_counts:

1094 - min_counts = counts

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

1106 + waiting, running = current_counts[idx]

1107 + score = waiting * 4 + running

1108 + if score < min_score:

1109 + min_score = score

1110 eng_index = idx

1096 - # Adjust local counts for better balancing between stats updates

1097 - # from the coordinator (which happen every 100ms).

1098 - if min_counts[0]:

1099 - min_counts[0] += 1

1100 - else:

1101 - min_counts[1] += 1

1111 + # Increment local waiting count for better balancing between stats

1112 + # updates from the coordinator (which happen every 100ms).

1113 + current_counts[eng_index][0] += self.client_count

1114

1115 chosen_engine = self.core_engines[eng_index]

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

vllm/v1/metrics/stats.py

@@ -33,6 +33,10 @@ class SchedulerStats:

33 num_running_reqs: int = 0

34 num_waiting_reqs: int = 0

36 + # These are used for internal DP load-balancing.

37 + step_counter: int = 0

38 + current_wave: int = 0

39 +

40 kv_cache_usage: float = 0.0

42 prefix_cache_stats: PrefixCacheStats = field(

vllm/v1/utils.py

@@ -154,6 +154,7 @@ class APIServerProcessManager:

154 client_config = {

155 "input_address": in_addr,

156 "output_address": out_addr,

157 + "client_count": num_servers,

158 "client_index": i

159 }

160 if stats_update_address is not None: