@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a8481be7a9

[Rust Frontend][Perf] Use dedicated runtime for HTTP/request-processing/ZMQ (#46051)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Jun 23, 2026 at 12:03 UTC a8481be7a9e9a4faa1fbf386f3e46e5c2b4acdc2

13 files changed +328 -15

rust/src/chat/src/multimodal.rs

@@ -365,6 +365,7 @@ impl MultimodalModelInfo {

365 let processor = self.image_processor.raw;

366 let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();

367

368 + // TODO: is it still necessary given that we've already in a dedicated runtime?

369 tokio::task::spawn_blocking(move || {

370 processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}"))

371 })

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

+19 -11

@@ -15,6 +15,7 @@ use crate::protocol::handshake::EngineCoreReadyResponse;

15 use crate::protocol::lora::LoraRequest;

16 use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};

17 use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};

18 +use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};

19 use crate::transport::{self, ConnectedEngine};

21 pub(crate) mod imp;

@@ -201,6 +202,8 @@ pub struct EngineCoreClient {

202 coordinator: Option<CoordinatorHandle>,

203 abort_tx: mpsc::UnboundedSender<AbortRequest>,

204

205 + /// Runtime used to send messages to the engine and drive all background tasks.

206 + runtime: BackgroundShutdownRuntime,

207 // Background tasks

208 output_task: AbortOnDropHandle<()>,

209 dispatcher_task: AbortOnDropHandle<()>,

@@ -280,21 +283,22 @@ impl EngineCoreClient {

283 let (output_tx, output_rx) = mpsc::channel(64);

284 let (abort_tx, abort_rx) = mpsc::unbounded_channel();

285 let engines = connected.engines;

286 + let runtime = build_zmq_runtime();

287 let inner = Arc::new(ClientInner::new(

288 connected.input_send,

289 + runtime.handle().clone(),

290 config.model_name.clone(),

291 &engines,

292 ));

288 - let output_task = AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(

293 + let output_task = AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(

294 connected.output_socket,

295 output_tx,

296 )));

292 - let dispatcher_task = AbortOnDropHandle::new(tokio::spawn(run_output_dispatcher_loop(

293 - inner.clone(),

294 - output_rx,

295 - )));

297 + let dispatcher_task = AbortOnDropHandle::new(

298 + runtime.spawn(run_output_dispatcher_loop(inner.clone(), output_rx)),

299 + );

300 let abort_task =

297 - AbortOnDropHandle::new(tokio::spawn(run_abort_loop(inner.clone(), abort_rx)));

301 + AbortOnDropHandle::new(runtime.spawn(run_abort_loop(inner.clone(), abort_rx)));

302

303 // If any engine reported a dp_stats_address in its ready response, use it

304 // as the external coordinator address.

@@ -307,13 +311,13 @@ impl EngineCoreClient {

311 CoordinatorHandle::new_inproc(coordinator_transport.input_socket);

312 let (coordinator_output_tx, coordinator_output_rx) = mpsc::channel(64);

313 let coordinator_output_task =

310 - AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(

314 + AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(

315 coordinator_transport.output_socket,

316 coordinator_output_tx,

317 )));

314 - let coordinator_task = AbortOnDropHandle::new(tokio::spawn(

315 - runner.run(coordinator_output_rx, inner.clone()),

316 - ));

318 + let coordinator_task = AbortOnDropHandle::new(

319 + runtime.spawn(runner.run(coordinator_output_rx, inner.clone())),

320 + );

321 (

322 Some(handle),

323 Some(coordinator_output_task),

@@ -327,7 +331,7 @@ impl EngineCoreClient {

331 {

332 let (handle, service) = CoordinatorHandle::connect_external(address).await?;

333 let coordinator_task =

330 - AbortOnDropHandle::new(tokio::spawn(service.run(inner.clone())));

334 + AbortOnDropHandle::new(runtime.spawn(service.run(inner.clone())));

335 (Some(handle), None, Some(coordinator_task))

336 } else {

337 (None, None, None)

@@ -341,6 +345,7 @@ impl EngineCoreClient {

345 inner,

346 coordinator,

347 abort_tx,

348 + runtime,

349 output_task,

350 dispatcher_task,

351 abort_task,

@@ -737,6 +742,7 @@ impl EngineCoreClient {

742 let Self {

743 inner,

744 abort_tx,

745 + runtime,

746 output_task,

747 dispatcher_task,

748 abort_task,

@@ -757,6 +763,8 @@ impl EngineCoreClient {

763

764 tasks.iter().for_each(|t| t.abort());

765 join_all(tasks).await;

766 + drop(inner);

767 + drop(runtime);

768

769 info!("engine-core client shut down");

770 Ok(())

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

+19 -3

@@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH};

5 use arc_swap::ArcSwapOption;

6 use parking_lot::Mutex;

7 use thiserror_ext::AsReport as _;

8 +use tokio::runtime::Handle;

9 use tokio::sync::mpsc;

10 use tracing::{debug, info, trace, warn};

11 use vllm_metrics::METRICS;

@@ -26,6 +27,8 @@ use crate::{Error, Result, transport};

28 pub(crate) struct ClientInner {

29 input_send: RouterSendHalf,

30 + /// The runtime handle used for sending messages to the engine.

31 + handle: Handle,

32 model_name: String,

33 request_reg: Mutex<RequestRegistry>,

34 utility_reg: Mutex<UtilityRegistry>,

@@ -37,11 +40,13 @@ impl ClientInner {

40 /// handshake completes.

41 pub fn new(

42 input_send: RouterSendHalf,

43 + handle: Handle,

44 model_name: String,

45 engines: &[ConnectedEngine],

46 ) -> Self {

47 Self {

48 input_send,

49 + handle,

50 model_name,

51 request_reg: Mutex::new(RequestRegistry::new(engines)),

52 utility_reg: Mutex::new(UtilityRegistry::default()),

@@ -213,9 +218,19 @@ impl ClientInner {

218 // frames instead of always producing a single msgpack frame.

219 let payload = encode_msgpack(payload)?;

220 let mut input_send = self.input_send.clone();

216 - transport::send_message(&mut input_send, engine_id, request_type.to_frame(), payload)

217 - .await?;

218 - Ok(())

221 + let engine_id = engine_id.clone();

222 +

223 + self.handle

224 + .spawn(async move {

225 + transport::send_message(

226 + &mut input_send,

227 + &engine_id,

228 + request_type.to_frame(),

229 + payload,

230 + )

231 + .await

232 + })

233 + .await?

234 }

235

236 /// Handle an abort request by sending the abort message to the engine.

@@ -434,6 +449,7 @@ mod tests {

449 let (send, _) = socket.split();

450 ClientInner::new(

451 send,

452 + Handle::current(),

453 "test-model".to_string(),

454 &[ConnectedEngine {

455 engine_id: EngineId::from(b"engine-0"),

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

@@ -29,6 +29,8 @@ pub enum Error {

29 Io(#[from] std::io::Error),

30 #[error("transport error")]

31 Transport(#[from] zeromq::ZmqError),

32 + #[error("ZMQ runtime task failed")]

33 + ZmqRuntimeTask(#[from] tokio::task::JoinError),

34 #[error("engine core reported fatal failure")]

35 EngineCoreDead,

36 #[error("startup handshake timed out while waiting for {stage} after {timeout:?}")]

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

@@ -4,6 +4,7 @@ mod error;

4 mod metrics;

5 pub mod mock_engine;

6 pub mod protocol;

7 +pub mod runtime;

8 #[cfg(any(test, feature = "test-util"))]

9 pub mod test_utils;

10 mod transport;

rust/src/engine-core-client/src/runtime.rs new

+75

@@ -0,0 +1,75 @@

1 +use std::mem::ManuallyDrop;

2 +use std::ops::{Deref, DerefMut};

3 +use std::sync::OnceLock;

4 +use std::sync::atomic::{AtomicUsize, Ordering};

5 +

6 +use tokio::runtime::Runtime;

7 +

8 +/// A wrapper around [`Runtime`] that shuts down the runtime in the background when dropped.

9 +///

10 +/// This can be useful in some cases, because sometimes we want to drop the runtime without

11 +/// blocking the current thread, for example, when it's nested inside another runtime.

12 +pub struct BackgroundShutdownRuntime(ManuallyDrop<Runtime>);

13 +

14 +impl Drop for BackgroundShutdownRuntime {

15 + fn drop(&mut self) {

16 + // Safety: The runtime is only dropped once here.

17 + let runtime = unsafe { ManuallyDrop::take(&mut self.0) };

18 + runtime.shutdown_background();

19 + }

20 +}

21 +

22 +impl Deref for BackgroundShutdownRuntime {

23 + type Target = Runtime;

24 +

25 + fn deref(&self) -> &Self::Target {

26 + &self.0

27 + }

28 +}

29 +

30 +impl DerefMut for BackgroundShutdownRuntime {

31 + fn deref_mut(&mut self) -> &mut Self::Target {

32 + &mut self.0

33 + }

34 +}

35 +

36 +impl From<Runtime> for BackgroundShutdownRuntime {

37 + fn from(runtime: Runtime) -> Self {

38 + Self(ManuallyDrop::new(runtime))

39 + }

40 +}

41 +

42 +const ZMQ_WORKER_THREADS_ENV: &str = "VLLM_RS_ZMQ_WORKER_THREADS";

43 +/// The number of tasks running on the ZMQ runtime is fixed and expected to remain

44 +/// small, and multiple engines share the same ZMQ socket. Therefore, based on

45 +/// benchmarks, a default value of 4 is generally sufficient.

46 +const DEFAULT_ZMQ_WORKER_THREADS: usize = 4;

47 +

48 +static ZMQ_RUNTIME_SEQUENCE: OnceLock<AtomicUsize> = OnceLock::new();

49 +

50 +/// Build a Tokio runtime for ZMQ tasks. Multiple calls to this function will

51 +/// return multiple runtimes with distinct thread name suffixes.

52 +pub(crate) fn build_zmq_runtime() -> BackgroundShutdownRuntime {

53 + let sequence = ZMQ_RUNTIME_SEQUENCE

54 + .get_or_init(|| AtomicUsize::new(0))

55 + .fetch_add(1, Ordering::Relaxed);

56 +

57 + tokio::runtime::Builder::new_multi_thread()

58 + .worker_threads(zmq_worker_threads())

59 + .thread_name_fn(move || format!("vllm-zmq-{sequence}"))

60 + .enable_all()

61 + .build()

62 + .expect("failed to build vLLM ZMQ runtime")

63 + .into()

64 +}

65 +

66 +/// Get the number of worker threads to use for the ZMQ runtime. If env var

67 +/// `VLLM_RS_ZMQ_WORKER_THREADS` is set and a valid positive integer, it will be used.

68 +/// Otherwise, the default value of `DEFAULT_ZMQ_WORKER_THREADS` will be used.

69 +fn zmq_worker_threads() -> usize {

70 + std::env::var(ZMQ_WORKER_THREADS_ENV)

71 + .ok()

72 + .and_then(|value| value.parse::<usize>().ok())

73 + .filter(|value| *value > 0)

74 + .unwrap_or(DEFAULT_ZMQ_WORKER_THREADS)

75 +}

rust/src/server/Cargo.toml

@@ -30,6 +30,7 @@ tokio-stream.workspace = true

30 tokio-util.workspace = true

31 tonic.workspace = true

32 tonic-prost.workspace = true

33 +tower.workspace = true

34 tower-http.workspace = true

35 tracing.workspace = true

36 tracing-futures.workspace = true

rust/src/server/src/lib.rs

+5 -1

@@ -7,6 +7,7 @@ mod listener;

7 mod lora;

8 mod middleware;

9 mod routes;

10 +mod runtime;

11 mod server_info;

12 mod state;

13 mod utils;

@@ -157,6 +158,9 @@ where

158 .with_context(|| format!("failed to bind gRPC listener on {grpc_host}:{grpc_port}"))?;

159 let addr = grpc_listener.local_addr()?;

160 let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));

161 + let svc = TonicServer::builder()

162 + .layer(middleware::request_runtime_layer(state.clone()))

163 + .add_service(svc);

164 info!(%addr, "starting gRPC server");

165 Some((grpc_listener, svc))

166 } else {

@@ -238,7 +242,7 @@ where

242 shutdown.cancelled().await;

243 return Ok(());

244 };

241 - let server = TonicServer::builder().add_service(svc).serve_with_incoming_shutdown(

245 + let server = svc.serve_with_incoming_shutdown(

246 TcpListenerStream::new(grpc_listener),

247 shutdown.cancelled_owned(),

248 );

rust/src/server/src/middleware/mod.rs

@@ -2,10 +2,12 @@ mod auth;

2 mod cors;

3 mod load;

4 mod metrics;

5 +mod offload;

6 mod request_id;

8 pub use auth::authenticate_api_key;

9 pub use cors::{cors_layer, strip_cors_on_no_origin};

10 pub use load::track_server_load;

11 pub use metrics::track_http_metrics;

12 +pub(crate) use offload::request_runtime_layer;

13 pub use request_id::set_request_id_header;

rust/src/server/src/middleware/offload.rs new

+134

@@ -0,0 +1,134 @@

1 +use std::sync::Arc;

2 +use std::task::{Context, Poll};

3 +

4 +use axum::http::Request;

5 +use axum::response::{IntoResponse, Response};

6 +use futures::future::BoxFuture;

7 +use tokio_util::task::AbortOnDropHandle;

8 +use tonic::Status;

9 +use tower::Service;

10 +use tower::layer::layer_fn;

11 +use tracing::error;

12 +

13 +use crate::error::{ApiError, server_error};

14 +use crate::state::AppState;

15 +

16 +/// Request paths that are run on the request runtime.

17 +///

18 +/// These routes can perform CPU-heavy request preparation, including JSON

19 +/// extraction, validation, chat-template rendering, tokenization, request

20 +/// lowering, and engine submission. Lightweight operational routes stay on the

21 +/// HTTP runtime.

22 +const OFFLOADED_PATHS: &[&str] = &[

23 + // HTTP routes:

24 + "/v1/chat/completions",

25 + "/v1/completions",

26 + "/tokenize",

27 + "/detokenize",

28 + "/inference/v1/generate",

29 + // gRPC routes:

30 + "/vllm.Generate/Generate",

31 + "/vllm.Generate/GenerateStream",

32 +];

33 +

34 +/// Return a Tower layer that runs selected data-plane requests on the request runtime,

35 +/// so that we can offset heavy request parsing and preprocessing from the HTTP runtime.

36 +pub(crate) fn request_runtime_layer<S>(

37 + state: Arc<AppState>,

38 +) -> impl tower::Layer<S, Service = RequestRuntimeService<S>> + Clone {

39 + layer_fn(move |inner| RequestRuntimeService {

40 + inner,

41 + state: state.clone(),

42 + })

43 +}

44 +

45 +/// Service produced by [`request_runtime_layer`].

46 +#[derive(Clone)]

47 +pub(crate) struct RequestRuntimeService<S> {

48 + inner: S,

49 + state: Arc<AppState>,

50 +}

51 +

52 +impl<S, B> Service<Request<B>> for RequestRuntimeService<S>

53 +where

54 + S: Service<Request<B>> + Clone + Send + 'static,

55 + S::Future: Send + 'static,

56 + S::Response: RequestRuntimeErrorResponse + Send + 'static,

57 + S::Error: Send + 'static,

58 + B: Send + 'static,

59 +{

60 + type Response = S::Response;

61 + type Error = S::Error;

62 + type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

63 +

64 + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

65 + self.inner.poll_ready(cx)

66 + }

67 +

68 + fn call(&mut self, req: Request<B>) -> Self::Future {

69 + if !should_offload(req.uri().path()) {

70 + return Box::pin(self.inner.call(req));

71 + }

72 +

73 + // Axum extractors and route handlers execute inside the inner service,

74 + // so offloading here moves request parsing and preprocessing off the

75 + // HTTP runtime without wrapping each handler manually. For streaming

76 + // HTTP responses, the response body is still polled on the HTTP runtime.

77 + let clone = self.inner.clone();

78 + let mut inner = std::mem::replace(&mut self.inner, clone);

79 + let task = AbortOnDropHandle::new(self.state.request_runtime().spawn(inner.call(req)));

80 +

81 + Box::pin(async move {

82 + match task.await {

83 + Ok(result) => result,

84 + Err(error) => {

85 + error!(%error, "request runtime task failed");

86 + Ok(S::Response::request_runtime_error_response())

87 + }

88 + }

89 + })

90 + }

91 +}

92 +

93 +trait RequestRuntimeErrorResponse {

94 + fn request_runtime_error_response() -> Self;

95 +}

96 +

97 +impl RequestRuntimeErrorResponse for Response {

98 + fn request_runtime_error_response() -> Self {

99 + server_error!("request runtime task failed").into_response()

100 + }

101 +}

102 +

103 +impl RequestRuntimeErrorResponse for axum::http::Response<tonic::body::Body> {

104 + fn request_runtime_error_response() -> Self {

105 + Status::internal("request runtime task failed").into_http()

106 + }

107 +}

108 +

109 +fn should_offload(path: &str) -> bool {

110 + OFFLOADED_PATHS.contains(&path)

111 +}

112 +

113 +#[cfg(test)]

114 +mod tests {

115 + use super::should_offload;

116 +

117 + #[test]

118 + fn offloads_generation_and_tokenization_paths() {

119 + assert!(should_offload("/v1/chat/completions"));

120 + assert!(should_offload("/v1/completions"));

121 + assert!(should_offload("/tokenize"));

122 + assert!(should_offload("/detokenize"));

123 + assert!(should_offload("/inference/v1/generate"));

124 + assert!(should_offload("/vllm.Generate/Generate"));

125 + assert!(should_offload("/vllm.Generate/GenerateStream"));

126 + }

127 +

128 + #[test]

129 + fn passes_through_lightweight_paths() {

130 + assert!(!should_offload("/health"));

131 + assert!(!should_offload("/metrics"));

132 + assert!(!should_offload("/v1/models"));

133 + }

134 +}

rust/src/server/src/routes.rs

@@ -108,6 +108,7 @@ fn build_router_with_options(

108 let enable_api_key_auth = state.has_api_keys();

109 let mut router = router

110 .with_state(state.clone())

111 + .layer(middleware::request_runtime_layer(state.clone()))

112 .layer(from_fn_with_state(

113 state.clone(),

114 middleware::track_server_load,

rust/src/server/src/runtime.rs new

+54

@@ -0,0 +1,54 @@

1 +use tokio::runtime::Builder;

2 +use tracing::{info, warn};

3 +use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;

4 +

5 +const REQUEST_WORKER_THREADS_ENV: &str = "VLLM_RS_REQUEST_WORKER_THREADS";

6 +const DEFAULT_MAX_REQUEST_WORKER_THREADS: usize = 32;

7 +

8 +/// Build a Tokio runtime for heavyweight request paths outside the HTTP runtime.

9 +///

10 +/// The server middleware uses this runtime for inference and tokenization

11 +/// routes so CPU-heavy request preparation does not monopolize the HTTP

12 +/// runtime's worker queue. Dropping the wrapper shuts the runtime down in the

13 +/// background.

14 +pub(crate) fn build_request_runtime() -> BackgroundShutdownRuntime {

15 + Builder::new_multi_thread()

16 + .enable_all()

17 + .thread_name("vllm-request")

18 + .worker_threads(request_worker_threads())

19 + .build()

20 + .expect("failed to build request runtime")

21 + .into()

22 +}

23 +

24 +/// Get the number of worker threads to use for the request runtime.

25 +///

26 +/// If `VLLM_RS_REQUEST_WORKER_THREADS` is set to a valid positive integer, it is

27 +/// used directly. Otherwise, the runtime uses available parallelism capped by

28 +/// `DEFAULT_MAX_REQUEST_WORKER_THREADS`.

29 +fn request_worker_threads() -> usize {

30 + if let Some(value) = std::env::var_os(REQUEST_WORKER_THREADS_ENV) {

31 + match value.to_string_lossy().parse::<usize>() {

32 + Ok(worker_threads) if worker_threads > 0 => return worker_threads,

33 + _ => warn!(

34 + value = %value.to_string_lossy(),

35 + "ignoring invalid {REQUEST_WORKER_THREADS_ENV}"

36 + ),

37 + }

38 + }

39 +

40 + std::thread::available_parallelism()

41 + .map(|parallelism| {

42 + let available = parallelism.get();

43 + let worker_threads = available.min(DEFAULT_MAX_REQUEST_WORKER_THREADS);

44 + if worker_threads < available {

45 + info!(

46 + available_parallelism = available,

47 + capped_worker_threads = worker_threads,

48 + "capping request runtime worker threads, set {REQUEST_WORKER_THREADS_ENV} to override"

49 + );

50 + }

51 + worker_threads

52 + })

53 + .unwrap_or(DEFAULT_MAX_REQUEST_WORKER_THREADS)

54 +}

rust/src/server/src/state.rs

+14

@@ -1,16 +1,20 @@

1 use std::sync::Arc;

2 +use std::sync::OnceLock;

3 use std::sync::atomic::{AtomicU64, Ordering};

5 use serde_json::Value;

6 use sha2::{Digest, Sha256};

7 +use tokio::runtime::Runtime;

8 use tokio::time::{Duration, Instant, sleep_until};

9 use tracing::warn;

10 use vllm_chat::ChatLlm;

11 use vllm_engine_core_client::EngineCoreClient;

12 use vllm_engine_core_client::protocol::lora::LoraRequest;

13 +use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;

15 use crate::config::{ApiServerOptions, CorsConfig};

16 use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};

17 +use crate::runtime::build_request_runtime;

18 use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};

20 const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);

@@ -42,6 +46,8 @@ pub struct AppState {

46 lora_manager: LoraManager,

47 /// Backend model path reported as `root` for base-model cards.

48 model_path: Option<String>,

49 + /// Lazily initialized runtime for heavyweight request paths.

50 + request_runtime: OnceLock<BackgroundShutdownRuntime>,

51 }

53 impl AppState {

@@ -68,6 +74,7 @@ impl AppState {

74 server_load: AtomicU64::new(0),

75 lora_manager: LoraManager::new(),

76 model_path: None,

77 + request_runtime: OnceLock::new(),

78 }

79 }

@@ -185,6 +192,12 @@ impl AppState {

192 self.chat.engine_core_client()

193 }

194

195 + /// Runtime used by middleware to isolate heavyweight request handlers from

196 + /// the HTTP reactor.

197 + pub(crate) fn request_runtime(&self) -> &Runtime {

198 + self.request_runtime.get_or_init(build_request_runtime)

199 + }

200 +

201 /// Return the current in-flight inference request count for the `/load`

202 /// endpoint.

203 pub fn server_load(&self) -> u64 {

@@ -214,6 +227,7 @@ impl AppState {

227 match Arc::try_unwrap(self) {

228 Ok(state) => {

229 state.chat.shutdown().await?;

230 + drop(state.request_runtime); // shutdown in background

231 return Ok(());

232 }

233 Err(state) => self = state,