@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6c117cff7d

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, 05:30 UTC

[Frontend] Pass API server count to each process (#23717)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Sep 20, 2025 at 01:15 UTC 6c117cff7d0110c74f97f68cda009595a2fdae5e

12 files changed +221 -51

benchmarks/kernels/benchmark_w8a8_block_fp8.py

+1 -1

@@ -11,13 +11,13 @@ from datetime import datetime

11 from typing import Any

13 import torch

14 -import triton

14 from tqdm import tqdm

16 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

17 _w8a8_block_fp8_matmul,

18 )

19 from vllm.platforms import current_platform

20 +from vllm.triton_utils import triton

21 from vllm.utils import FlexibleArgumentParser

23 mp.set_start_method("spawn", force=True)

examples/others/tensorize_vllm_model.py

+1 -8

@@ -1,8 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 -import argparse

5 -import dataclasses

4 import json

5 import logging

6 import os

@@ -327,12 +325,7 @@ def main():

325

326

327 if args.command == "serialize":

330 - eng_args_dict = {f.name: getattr(args, f.name) for f in

331 - dataclasses.fields(EngineArgs)}

332 -

333 - engine_args = EngineArgs.from_cli_args(

334 - argparse.Namespace(**eng_args_dict)

335 - )

328 + engine_args = EngineArgs.from_cli_args(args)

329

330 input_dir = tensorizer_dir.rstrip('/')

331 suffix = args.suffix if args.suffix else uuid.uuid4().hex

tests/entrypoints/test_api_server_process_manager.py

+1 -1

@@ -60,7 +60,7 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update):

60 global WORKER_RUNTIME_SECONDS

61 WORKER_RUNTIME_SECONDS = 0.5

63 - # Copy the args to avoid mutating the

63 + # Copy the args to avoid mutating them

64 args = api_server_args.copy()

66 if not with_stats_update:

tests/v1/test_external_lb_dp.py

+48 -4

@@ -9,6 +9,7 @@ from contextlib import AsyncExitStack

9 import openai # use the official client for correctness check

10 import pytest

11 import pytest_asyncio

12 +import requests

14 from tests.utils import RemoteOpenAIServer

15 from vllm.platforms import current_platform

@@ -70,6 +71,8 @@ class ExternalLBServerManager:

71 sargs,

72 auto_port=False,

73 env_dict={

74 + "VLLM_SERVER_DEV_MODE":

75 + "1",

76 current_platform.device_control_env_var:

77 ",".join(

78 str(

@@ -127,11 +130,19 @@ def default_server_args():

130

131

132 @pytest.fixture(scope="module", params=[1, 4])

130 -def servers(request, default_server_args):

133 +def server_manager(request, default_server_args):

134 api_server_count = request.param

132 - with ExternalLBServerManager(MODEL_NAME, DP_SIZE, api_server_count,

133 - default_server_args) as server_list:

134 - yield server_list

135 + server_manager = ExternalLBServerManager(MODEL_NAME, DP_SIZE,

136 + api_server_count,

137 + default_server_args)

138 +

139 + with server_manager:

140 + yield server_manager

141 +

142 +

143 +@pytest.fixture

144 +def servers(server_manager):

145 + return server_manager.servers

146

147

148 @pytest_asyncio.fixture

@@ -144,6 +155,39 @@ async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):

155 ]

156

157

158 +def _get_parallel_config(server: RemoteOpenAIServer):

159 + response = requests.get(server.url_for("server_info?config_format=json"))

160 + response.raise_for_status()

161 +

162 + vllm_config = response.json()["vllm_config"]

163 + return vllm_config["parallel_config"]

164 +

165 +

166 +def test_external_lb_server_info(server_manager):

167 + servers = server_manager.servers

168 + api_server_count = server_manager.api_server_count

169 +

170 + for i, (server, _) in enumerate(servers):

171 + print(f"Testing {i=}")

172 +

173 + # Each request will hit one of the API servers

174 + # `n_reqs` is set so that there is a good chance each server

175 + # receives at least one request

176 + n_reqs = 2 * api_server_count * api_server_count

177 + parallel_configs = [

178 + _get_parallel_config(server) for _ in range(n_reqs)

179 + ]

180 + api_process_counts = [

181 + c["_api_process_count"] for c in parallel_configs

182 + ]

183 + api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]

184 +

185 + assert all(c == api_server_count

186 + for c in api_process_counts), api_process_counts

187 + assert all(0 <= r < api_server_count

188 + for r in api_process_ranks), api_process_ranks

189 +

190 +

191 @pytest.mark.asyncio

192 @pytest.mark.parametrize(

193 "model_name",

tests/v1/test_hybrid_lb_dp.py

+49 -5

@@ -9,6 +9,7 @@ from contextlib import AsyncExitStack

9 import openai # use the official client for correctness check

10 import pytest

11 import pytest_asyncio

12 +import requests

14 from tests.utils import RemoteOpenAIServer

15 from tests.v1.test_utils import check_request_balancing

@@ -92,6 +93,8 @@ class HybridLBServerManager:

93 sargs,

94 auto_port=False,

95 env_dict={

96 + "VLLM_SERVER_DEV_MODE":

97 + "1",

98 current_platform.device_control_env_var:

99 ",".join(

100 str(

@@ -150,12 +153,20 @@ def default_server_args():

153

154

155 @pytest.fixture(scope="module", params=[1, 4])

153 -def servers(request, default_server_args):

156 +def server_manager(request, default_server_args):

157 api_server_count = request.param

155 - with HybridLBServerManager(MODEL_NAME, DP_SIZE, api_server_count,

156 - default_server_args, DP_SIZE_LOCAL,

157 - TP_SIZE) as server_list:

158 - yield server_list

158 + server_manager = HybridLBServerManager(MODEL_NAME, DP_SIZE,

159 + api_server_count,

160 + default_server_args, DP_SIZE_LOCAL,

161 + TP_SIZE)

162 +

163 + with server_manager:

164 + yield server_manager

165 +

166 +

167 +@pytest.fixture

168 +def servers(server_manager):

169 + return server_manager.servers

170

171

172 @pytest_asyncio.fixture

@@ -168,6 +179,39 @@ async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):

179 ]

180

181

182 +def _get_parallel_config(server: RemoteOpenAIServer):

183 + response = requests.get(server.url_for("server_info?config_format=json"))

184 + response.raise_for_status()

185 +

186 + vllm_config = response.json()["vllm_config"]

187 + return vllm_config["parallel_config"]

188 +

189 +

190 +def test_hybrid_dp_server_info(server_manager):

191 + servers = server_manager.servers

192 + api_server_count = server_manager.api_server_count

193 +

194 + for i, (server, _) in enumerate(servers):

195 + print(f"Testing {i=}")

196 +

197 + # Each request will hit one of the API servers

198 + # `n_reqs` is set so that there is a good chance each server

199 + # receives at least one request

200 + n_reqs = 2 * api_server_count * api_server_count

201 + parallel_configs = [

202 + _get_parallel_config(server) for _ in range(n_reqs)

203 + ]

204 + api_process_counts = [

205 + c["_api_process_count"] for c in parallel_configs

206 + ]

207 + api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]

208 +

209 + assert all(c == api_server_count

210 + for c in api_process_counts), api_process_counts

211 + assert all(0 <= r < api_server_count

212 + for r in api_process_ranks), api_process_ranks

213 +

214 +

215 @pytest.mark.asyncio

216 @pytest.mark.parametrize(

217 "model_name",

tests/v1/test_internal_lb_dp.py

+49 -8

@@ -10,6 +10,7 @@ from typing import Optional, cast

10 import openai # use the official client for correctness check

11 import pytest

12 import pytest_asyncio

13 +import requests

15 from tests.utils import RemoteOpenAIServer

16 from tests.v1.test_utils import check_request_balancing

@@ -101,6 +102,8 @@ class MultinodeInternalLBServerManager:

102 sargs,

103 auto_port=False,

104 env_dict={

105 + "VLLM_SERVER_DEV_MODE":

106 + "1",

107 current_platform.device_control_env_var:

108 ",".join(

109 str(

@@ -214,7 +217,10 @@ class APIOnlyServerManager:

217 self.model_name,

218 api_server_args,

219 auto_port=False,

217 - env_dict={}) # No GPUs needed for API-only server

220 + env_dict={

221 + "VLLM_SERVER_DEV_MODE": "1",

222 + # No GPUs needed for API-only server

223 + })

224 server.__enter__()

225 print(f"API-only server started successfully with "

226 f"{self.api_server_count} API servers")

@@ -293,14 +299,21 @@ def default_server_args():

299

300

301 @pytest.fixture(scope="module", params=[1, 4])

296 -def servers(request, default_server_args):

302 +def server_manager(request, default_server_args):

303 api_server_count = request.param

298 - with MultinodeInternalLBServerManager(MODEL_NAME, DP_SIZE,

299 - api_server_count,

300 - default_server_args,

301 - DP_SIZE // NUM_NODES,

302 - TP_SIZE) as server_list:

303 - yield server_list

304 + server_manager = MultinodeInternalLBServerManager(MODEL_NAME, DP_SIZE,

305 + api_server_count,

306 + default_server_args,

307 + DP_SIZE // NUM_NODES,

308 + TP_SIZE)

309 +

310 + with server_manager:

311 + yield server_manager

312 +

313 +

314 +@pytest.fixture

315 +def servers(server_manager):

316 + return server_manager.servers

317

318

319 @pytest.fixture(scope="module", params=[1, 4])

@@ -331,6 +344,34 @@ async def api_only_client(api_only_servers: list[tuple[RemoteOpenAIServer,

344 yield client

345

346

347 +def _get_parallel_config(server: RemoteOpenAIServer):

348 + response = requests.get(server.url_for("server_info?config_format=json"))

349 + response.raise_for_status()

350 +

351 + vllm_config = response.json()["vllm_config"]

352 + return vllm_config["parallel_config"]

353 +

354 +

355 +def test_multinode_dp_server_info(server_manager):

356 + head_server = server_manager.servers[0][0]

357 + api_server_count = server_manager.api_server_count

358 +

359 + # Each request will hit one of the API servers

360 + # `n_reqs` is set so that there is a good chance each server

361 + # receives at least one request

362 + n_reqs = 2 * api_server_count * api_server_count

363 + parallel_configs = [

364 + _get_parallel_config(head_server) for _ in range(n_reqs)

365 + ]

366 + api_process_counts = [c["_api_process_count"] for c in parallel_configs]

367 + api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]

368 +

369 + assert all(c == api_server_count

370 + for c in api_process_counts), api_process_counts

371 + assert all(0 <= r < api_server_count

372 + for r in api_process_ranks), api_process_ranks

373 +

374 +

375 @pytest.mark.asyncio

376 @pytest.mark.parametrize(

377 "model_name",

vllm/config/parallel.py

+25

@@ -193,6 +193,25 @@ class ParallelConfig:

193 not change by dcp, it simply reuse the GPUs of TP group, and tp_size

194 needs to be divisible by dcp_size."""

195

196 + _api_process_count: int = 1

197 + """

198 + The number of API processes initialized.

199 +

200 + Note:

201 + This is an internal config that is only valid for and

202 + should only be set by API server scale-out.

203 + """

204 +

205 + _api_process_rank: int = 0

206 + """

207 + The rank of this API process, or `-1` for engine core processes

208 + under API server scale-out.

209 +

210 + Note:

211 + This is an internal config that is only valid for and

212 + should only be set by API server scale-out.

213 + """

214 +

215 @property

216 def world_size_across_dp(self) -> int:

217 """world_size_across_dp is TPxPPxDP, it is the size of the world

@@ -428,6 +447,12 @@ class ParallelConfig:

447 if self.distributed_executor_backend is None and self.world_size == 1:

448 self.distributed_executor_backend = "uni"

449

450 + if not -1 <= self._api_process_rank < self._api_process_count:

451 + raise ValueError(

452 + "Invalid value of `_api_process_rank`. "

453 + f"Expected to be `-1` or `[0, {self._api_process_count})`, "

454 + f"but found: {self._api_process_rank}")

455 +

456 @property

457 def use_ray(self) -> bool:

458 return self.distributed_executor_backend == "ray" or (

vllm/engine/arg_utils.py

+8 -1

@@ -333,6 +333,8 @@ class EngineArgs:

333 enable_eplb: bool = ParallelConfig.enable_eplb

334 expert_placement_strategy: ExpertPlacementStrategy = \

335 ParallelConfig.expert_placement_strategy

336 + _api_process_count: int = ParallelConfig._api_process_count

337 + _api_process_rank: int = ParallelConfig._api_process_rank

338 num_redundant_experts: int = EPLBConfig.num_redundant_experts

339 eplb_window_size: int = EPLBConfig.window_size

340 eplb_step_interval: int = EPLBConfig.step_interval

@@ -952,7 +954,10 @@ class EngineArgs:

954 # Get the list of attributes of this dataclass.

955 attrs = [attr.name for attr in dataclasses.fields(cls)]

956 # Set the attributes from the parsed arguments.

955 - engine_args = cls(**{attr: getattr(args, attr) for attr in attrs})

957 + engine_args = cls(**{

958 + attr: getattr(args, attr)

959 + for attr in attrs if hasattr(args, attr)

960 + })

961 return engine_args

962

963 def create_model_config(self) -> ModelConfig:

@@ -1366,6 +1371,8 @@ class EngineArgs:

1371 worker_cls=self.worker_cls,

1372 worker_extension_cls=self.worker_extension_cls,

1373 decode_context_parallel_size=self.decode_context_parallel_size,

1374 + _api_process_count=self._api_process_count,

1375 + _api_process_rank=self._api_process_rank,

1376 )

1377

1378 speculative_config = self.create_speculative_config(

vllm/entrypoints/cli/serve.py

+6 -12

@@ -135,23 +135,20 @@ def run_headless(args: argparse.Namespace):

135 def run_multi_api_server(args: argparse.Namespace):

136

137 assert not args.headless

138 - num_api_servers = args.api_server_count

138 + num_api_servers: int = args.api_server_count

139 assert num_api_servers > 0

140

141 - orig_mm_processor_cache_gb = args.mm_processor_cache_gb

142 -

141 if num_api_servers > 1:

142 setup_multiprocess_prometheus()

143

146 - # Not compatible with API server scale-out

147 - args.mm_processor_cache_gb = 0

148 -

144 listen_address, sock = setup_server(args)

145

146 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)

147 + engine_args._api_process_count = num_api_servers

148 + engine_args._api_process_rank = -1

149 +

150 usage_context = UsageContext.OPENAI_API_SERVER

151 vllm_config = engine_args.create_engine_config(usage_context=usage_context)

154 - model_config = vllm_config.model_config

152

153 if num_api_servers > 1:

154 if not envs.VLLM_USE_V1:

@@ -161,10 +158,6 @@ def run_multi_api_server(args: argparse.Namespace):

158 raise ValueError("VLLM_ALLOW_RUNTIME_LORA_UPDATING cannot be used "

159 "with api_server_count > 1")

160

164 - if model_config.is_multimodal_model and orig_mm_processor_cache_gb > 0:

165 - logger.warning("Multi-modal processor cache is disabled because "

166 - "it is not compatible with `api_server_count > 1`.")

167 -

161 executor_class = Executor.get_class(vllm_config)

162 log_stats = not engine_args.disable_log_stats

163

@@ -221,9 +214,10 @@ def run_api_server_worker_proc(listen_address,

214 client_config=None,

215 **uvicorn_kwargs) -> None:

216 """Entrypoint for individual API server worker processes."""

217 + client_config = client_config or {}

218 + server_index = client_config.get("client_index", 0)

219

220 # Set process title and add process-specific prefix to stdout and stderr.

226 - server_index = client_config.get("client_index", 0) if client_config else 0

221 set_process_title("APIServer", str(server_index))

222 decorate_logs()

223

vllm/entrypoints/openai/api_server.py

+29 -9

@@ -17,13 +17,14 @@ from argparse import Namespace

17 from collections.abc import AsyncGenerator, AsyncIterator, Awaitable

18 from contextlib import asynccontextmanager

19 from http import HTTPStatus

20 -from typing import Annotated, Any, Callable, Optional

20 +from typing import Annotated, Any, Callable, Literal, Optional

22 import prometheus_client

23 import pydantic

24 import regex as re

25 import uvloop

26 -from fastapi import APIRouter, Depends, FastAPI, Form, HTTPException, Request

26 +from fastapi import (APIRouter, Depends, FastAPI, Form, HTTPException, Query,

27 + Request)

28 from fastapi.exceptions import RequestValidationError

29 from fastapi.middleware.cors import CORSMiddleware

30 from fastapi.responses import JSONResponse, Response, StreamingResponse

@@ -166,6 +167,9 @@ async def build_async_engine_client(

167 # Context manager to handle engine_client lifecycle

168 # Ensures everything is shutdown and cleaned up on error/exit

169 engine_args = AsyncEngineArgs.from_cli_args(args)

170 + if client_config:

171 + engine_args._api_process_count = client_config.get("client_count", 1)

172 + engine_args._api_process_rank = client_config.get("client_index", 0)

173

174 if disable_frontend_multiprocessing is None:

175 disable_frontend_multiprocessing = bool(

@@ -209,8 +213,12 @@ async def build_async_engine_client_from_engine_args(

213

214 from vllm.v1.engine.async_llm import AsyncLLM

215 async_llm: Optional[AsyncLLM] = None

212 - client_count = client_config.pop("client_count") if client_config else 1

213 - client_index = client_config.pop("client_index") if client_config else 0

216 +

217 + # Don't mutate the input client_config

218 + client_config = dict(client_config) if client_config else {}

219 + client_count = client_config.pop("client_count", 1)

220 + client_index = client_config.pop("client_index", 0)

221 +

222 try:

223 async_llm = AsyncLLM.from_vllm_config(

224 vllm_config=vllm_config,

@@ -956,9 +964,22 @@ if envs.VLLM_SERVER_DEV_MODE:

964 logger.warning("SECURITY WARNING: Development endpoints are enabled! "

965 "This should NOT be used in production!")

966

967 + PydanticVllmConfig = pydantic.TypeAdapter(VllmConfig)

968 +

969 @router.get("/server_info")

960 - async def show_server_info(raw_request: Request):

961 - server_info = {"vllm_config": str(raw_request.app.state.vllm_config)}

970 + async def show_server_info(

971 + raw_request: Request,

972 + config_format: Annotated[Literal["text", "json"],

973 + Query()] = "text",

974 + ):

975 + vllm_config: VllmConfig = raw_request.app.state.vllm_config

976 + server_info = {

977 + "vllm_config":

978 + str(vllm_config)

979 + if config_format == "text" else PydanticVllmConfig.dump_python(

980 + vllm_config, mode="json", fallback=str)

981 + # fallback=str is needed to handle e.g. torch.dtype

982 + }

983 return JSONResponse(content=server_info)

984

985 @router.post("/reset_prefix_cache")

@@ -1856,8 +1877,6 @@ async def run_server_worker(listen_address,

1877 if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:

1878 ToolParserManager.import_tool_parser(args.tool_parser_plugin)

1879

1859 - server_index = client_config.get("client_index", 0) if client_config else 0

1860 -

1880 # Load logging config for uvicorn if specified

1881 log_config = load_log_config(args.log_config_file)

1882 if log_config is not None:

@@ -1873,7 +1892,8 @@ async def run_server_worker(listen_address,

1892 vllm_config = await engine_client.get_vllm_config()

1893 await init_app_state(engine_client, vllm_config, app.state, args)

1894

1876 - logger.info("Starting vLLM API server %d on %s", server_index,

1895 + logger.info("Starting vLLM API server %d on %s",

1896 + vllm_config.parallel_config._api_process_rank,

1897 listen_address)

1898 shutdown_task = await serve_http(

1899 app,

vllm/multimodal/cache.py

+2 -1

@@ -494,7 +494,8 @@ def _enable_processor_cache(

494

495 def _enable_ipc_cache(vllm_config: "VllmConfig") -> bool:

496 parallel_config = vllm_config.parallel_config

497 - supports_ipc_cache = (parallel_config.data_parallel_size == 1

497 + supports_ipc_cache = ((parallel_config._api_process_count == 1

498 + and parallel_config.data_parallel_size == 1)

499 or parallel_config.data_parallel_external_lb)

500

501 return supports_ipc_cache

vllm/v1/engine/core_client.py

+2 -1

@@ -437,7 +437,7 @@ class MPClient(EngineCoreClient):

437 self.engines_running = False

438

439 self.stats_update_address: Optional[str] = None

440 - if client_addresses is not None:

440 + if client_addresses:

441 # Engines are managed externally to this client.

442 input_address = client_addresses["input_address"]

443 output_address = client_addresses["output_address"]

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

774 client_addresses=client_addresses,

775 )

776

777 + self.client_count = client_count

778 self.client_index = client_index

779 self.outputs_queue = asyncio.Queue[Union[EngineCoreOutputs,

780 Exception]]()