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

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 10, 20:15 UTC

[BugFix] Fix async scheduling for pooling models (#31584)

Signed-off-by: njhill <nickhill123@gmail.com>

Nick Hill committed Dec 31, 2025 at 14:48 UTC 6c2cfb62ffb81b1ec3c39d2501d1e51ed6e76a4e

7 files changed +131 -92

vllm/v1/executor/ray_utils.py

+5 -2

@@ -104,8 +104,11 @@ try:

104 scheduler_output, intermediate_tensors

105 )

106 if isinstance(output, IntermediateTensors):

107 - output = scheduler_output, grammar_output, output

108 - elif not get_pp_group().is_last_rank:

107 + return scheduler_output, grammar_output, output

108 +

109 + if isinstance(output, AsyncModelRunnerOutput):

110 + output = output.get_output()

111 + if not get_pp_group().is_last_rank:

112 # Case where there are no scheduled requests

113 # but may still be finished requests.

114 assert not output or not output.req_ids

vllm/v1/outputs.py

+7 -18

@@ -151,21 +151,23 @@ class ModelRunnerOutput:

151 # num_generated_tokens is the number of tokens

152 # generated in the current step. It can be different for

153 # each request due to speculative/jump decoding.

154 - sampled_token_ids: list[list[int]]

154 + sampled_token_ids: list[list[int]] = field(default_factory=list)

155

156 # [num_reqs, max_num_logprobs + 1]

157 # [num_reqs, max_num_logprobs + 1]

158 # [num_reqs]

159 - logprobs: LogprobsLists | None

159 + logprobs: LogprobsLists | None = None

160

161 # req_id -> (token_ids, logprobs, ranks)

162 # [prompt_len, num_prompt_logprobs]

163 # [prompt_len, num_prompt_logprobs]

164 # [prompt_len]

165 - prompt_logprobs_dict: dict[str, LogprobsTensors | None]

165 + prompt_logprobs_dict: dict[str, LogprobsTensors | None] = field(

166 + default_factory=dict

167 + )

168

169 # [num_reqs, hidden_size]

168 - pooler_output: list[torch.Tensor | None]

170 + pooler_output: list[torch.Tensor | None] | None = None

171

172 kv_connector_output: KVConnectorOutput | None = None

173

@@ -225,21 +227,8 @@ def make_empty_encoder_model_runner_output(

227 req_ids=req_ids,

228 req_id_to_index=req_id_to_index,

229 sampled_token_ids=sampled_token_ids,

228 - logprobs=None,

229 - prompt_logprobs_dict={},

230 pooler_output=pooler_output,

231 - kv_connector_output=None,

232 - ec_connector_output=None,

233 - num_nans_in_logits=None,

231 )

232

233

237 -EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(

238 - req_ids=[],

239 - req_id_to_index={},

240 - sampled_token_ids=[],

241 - logprobs=None,

242 - prompt_logprobs_dict={},

243 - pooler_output=[],

244 - num_nans_in_logits=None,

245 -)

234 +EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[], req_id_to_index={})

vllm/v1/pool/metadata.py

+20 -28

@@ -2,6 +2,7 @@

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

3 from dataclasses import dataclass

5 +import numpy as np

6 import torch

8 from vllm.pooling_params import PoolingParams

@@ -91,36 +92,27 @@ class PoolingMetadata:

93 def build_pooling_cursor(

94 self,

94 - num_scheduled_tokens: list[int],

95 + num_scheduled_tokens_np: np.ndarray,

96 seq_lens_cpu: torch.Tensor,

97 device: torch.device,

98 ):

98 - self.pooling_cursor = build_pooling_cursor(

99 - num_scheduled_tokens, seq_lens_cpu, self.prompt_lens, device

100 - )

99 + n_seq = len(num_scheduled_tokens_np)

100 + prompt_lens = self.prompt_lens

101

102 + assert len(prompt_lens) == n_seq

103

103 -def build_pooling_cursor(

104 - num_scheduled_tokens: list[int],

105 - seq_lens_cpu: torch.Tensor,

106 - prompt_lens: torch.Tensor,

107 - device: torch.device,

108 -):

109 - assert len(prompt_lens) == len(num_scheduled_tokens)

110 -

111 - n_seq = len(num_scheduled_tokens)

112 - index = list(range(n_seq))

113 - num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, device="cpu")

114 - cumsum = torch.zeros(

115 - n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"

116 - )

117 - torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])

118 - cumsum = cumsum.to(device, non_blocking=True)

119 - return PoolingCursor(

120 - index=index,

121 - first_token_indices_gpu=cumsum[:n_seq],

122 - last_token_indices_gpu=cumsum[1:] - 1,

123 - prompt_lens_cpu=prompt_lens,

124 - seq_lens_cpu=seq_lens_cpu,

125 - num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,

126 - )

104 + index = list(range(n_seq))

105 + num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)

106 + cumsum = torch.zeros(

107 + n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"

108 + )

109 + torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])

110 + cumsum = cumsum.to(device, non_blocking=True)

111 + self.pooling_cursor = PoolingCursor(

112 + index=index,

113 + first_token_indices_gpu=cumsum[:n_seq],

114 + last_token_indices_gpu=cumsum[1:] - 1,

115 + prompt_lens_cpu=prompt_lens,

116 + seq_lens_cpu=seq_lens_cpu,

117 + num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,

118 + )

vllm/v1/worker/gpu/model_runner.py

+1 -5

@@ -968,11 +968,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

968 # Only for compatibility with the existing model runner and scheduler.

969 req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},

970 sampled_token_ids=None, # type: ignore

971 - logprobs=None,

972 - prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore

973 - pooler_output=[],

974 - kv_connector_output=None,

975 - num_nans_in_logits=None,

971 + prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]

972 )

973 async_output = AsyncOutput(

974 model_runner_output=model_runner_output,

vllm/v1/worker/gpu_model_runner.py

+92 -35

@@ -254,6 +254,50 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

254 return output

255

256

257 +class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):

258 + def __init__(

259 + self,

260 + model_runner_output: ModelRunnerOutput,

261 + raw_pooler_output: PoolerOutput,

262 + finished_mask: list[bool],

263 + async_output_copy_stream: torch.cuda.Stream,

264 + ):

265 + self._model_runner_output = model_runner_output

266 + self._finished_mask = finished_mask

267 +

268 + # Event on the copy stream so we can synchronize the non-blocking copy.

269 + self.async_copy_ready_event = torch.Event()

270 +

271 + # Keep a reference to the device tensors to avoid them being

272 + # deallocated until we finish copying it to the host.

273 + self._raw_pooler_output = raw_pooler_output

274 +

275 + # Initiate the copy on a separate stream, but do not synchronize it.

276 + default_stream = torch.cuda.current_stream()

277 + with torch.cuda.stream(async_output_copy_stream):

278 + async_output_copy_stream.wait_stream(default_stream)

279 + self._raw_pooler_output_cpu = json_map_leaves(

280 + lambda x: None if x is None else x.to("cpu", non_blocking=True),

281 + self._raw_pooler_output,

282 + )

283 + self.async_copy_ready_event.record()

284 +

285 + def get_output(self) -> ModelRunnerOutput:

286 + """Copy the device tensors to the host and return a ModelRunnerOutput.

287 + This function blocks until the copy is finished.

288 + """

289 + self.async_copy_ready_event.synchronize()

290 +

291 + # Release the device tensors once the copy has completed.

292 + del self._raw_pooler_output

293 +

294 + self._model_runner_output.pooler_output = [

295 + out if include else None

296 + for out, include in zip(self._raw_pooler_output_cpu, self._finished_mask)

297 + ]

298 + return self._model_runner_output

299 +

300 +

301 class ExecuteModelState(NamedTuple):

302 """Ephemeral cached state transferred between execute_model() and

303 sample_tokens(), after execute_model() returns None."""

@@ -2476,17 +2520,19 @@ class GPUModelRunner(

2520 hidden_states: torch.Tensor,

2521 num_scheduled_tokens: int,

2522 num_scheduled_tokens_np: np.ndarray,

2479 - ) -> ModelRunnerOutput:

2480 - assert self.input_batch.num_reqs == len(self.input_batch.pooling_params), (

2523 + kv_connector_output: KVConnectorOutput | None,

2524 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput:

2525 + num_reqs = self.input_batch.num_reqs

2526 + assert num_reqs == len(self.input_batch.pooling_params), (

2527 "Either all or none of the requests in a batch must be pooling request"

2528 )

2529

2530 hidden_states = hidden_states[:num_scheduled_tokens]

2485 - seq_lens_cpu = self.seq_lens.cpu[: self.input_batch.num_reqs]

2531 + seq_lens_cpu = self.seq_lens.cpu[:num_reqs]

2532

2533 pooling_metadata = self.input_batch.get_pooling_metadata()

2534 pooling_metadata.build_pooling_cursor(

2489 - num_scheduled_tokens_np.tolist(), seq_lens_cpu, device=hidden_states.device

2535 + num_scheduled_tokens_np, seq_lens_cpu, device=hidden_states.device

2536 )

2537

2538 model = cast(VllmModelForPooling, self.model)

@@ -2494,27 +2540,41 @@ class GPUModelRunner(

2540 hidden_states=hidden_states,

2541 pooling_metadata=pooling_metadata,

2542 )

2543 +

2544 + finished_mask = [

2545 + seq_len == prompt_len

2546 + for seq_len, prompt_len in zip(seq_lens_cpu, pooling_metadata.prompt_lens)

2547 + ]

2548 +

2549 + model_runner_output = ModelRunnerOutput(

2550 + req_ids=self.input_batch.req_ids.copy(),

2551 + req_id_to_index=self.input_batch.req_id_to_index.copy(),

2552 + kv_connector_output=kv_connector_output,

2553 + )

2554 +

2555 + if raw_pooler_output is None or not any(finished_mask):

2556 + model_runner_output.pooler_output = [None] * num_reqs

2557 + return model_runner_output

2558 +

2559 + if self.use_async_scheduling:

2560 + return AsyncGPUPoolingModelRunnerOutput(

2561 + model_runner_output=model_runner_output,

2562 + raw_pooler_output=raw_pooler_output,

2563 + finished_mask=finished_mask,

2564 + async_output_copy_stream=self.async_output_copy_stream,

2565 + )

2566 +

2567 raw_pooler_output = json_map_leaves(

2498 - lambda x: x.to("cpu", non_blocking=True) if x is not None else x,

2568 + lambda x: None if x is None else x.to("cpu", non_blocking=True),

2569 raw_pooler_output,

2570 )

2571 self._sync_device()

2572

2503 - pooler_output: list[torch.Tensor | None] = []

2504 - for raw_output, seq_len, prompt_len in zip(

2505 - raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens

2506 - ):

2507 - output = raw_output if seq_len == prompt_len else None

2508 - pooler_output.append(output)

2509 -

2510 - return ModelRunnerOutput(

2511 - req_ids=self.input_batch.req_ids,

2512 - req_id_to_index=self.input_batch.req_id_to_index,

2513 - sampled_token_ids=[],

2514 - logprobs=None,

2515 - prompt_logprobs_dict={},

2516 - pooler_output=pooler_output,

2517 - )

2573 + model_runner_output.pooler_output = [

2574 + out if include else None

2575 + for out, include in zip(raw_pooler_output, finished_mask)

2576 + ]

2577 + return model_runner_output

2578

2579 def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int:

2580 # Pad tokens to multiple of tensor_parallel_size when

@@ -3036,7 +3096,7 @@ class GPUModelRunner(

3096 self,

3097 scheduler_output: "SchedulerOutput",

3098 intermediate_tensors: IntermediateTensors | None = None,

3039 - ) -> ModelRunnerOutput | IntermediateTensors | None:

3099 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None:

3100 if self.execute_model_state is not None:

3101 raise RuntimeError(

3102 "State error: sample_tokens() must be called "

@@ -3244,11 +3304,12 @@ class GPUModelRunner(

3304

3305 if self.is_pooling_model:

3306 # Return the pooling output.

3247 - output = self._pool(

3248 - hidden_states, num_scheduled_tokens, num_scheduled_tokens_np

3307 + return self._pool(

3308 + hidden_states,

3309 + num_scheduled_tokens,

3310 + num_scheduled_tokens_np,

3311 + kv_connector_output,

3312 )

3250 - output.kv_connector_output = kv_connector_output

3251 - return output

3313

3314 sample_hidden_states = hidden_states[logits_indices]

3315 logits = self.model.compute_logits(sample_hidden_states)

@@ -3437,7 +3498,6 @@ class GPUModelRunner(

3498 sampled_token_ids=valid_sampled_token_ids,

3499 logprobs=logprobs_lists,

3500 prompt_logprobs_dict=prompt_logprobs_dict,

3440 - pooler_output=[],

3501 kv_connector_output=kv_connector_output,

3502 ec_connector_output=ec_connector_output

3503 if self.supports_mm_inputs

@@ -4508,17 +4568,14 @@ class GPUModelRunner(

4568 max_num_reqs = self.scheduler_config.max_num_seqs

4569 num_reqs = min(num_tokens, max_num_reqs)

4570 min_tokens_per_req = num_tokens // num_reqs

4511 - num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs

4512 - num_scheduled_tokens_list[-1] += num_tokens % num_reqs

4513 - assert sum(num_scheduled_tokens_list) == num_tokens

4514 - assert len(num_scheduled_tokens_list) == num_reqs

4571 + num_scheduled_tokens_np = np.full(num_reqs, min_tokens_per_req)

4572 + num_scheduled_tokens_np[-1] += num_tokens % num_reqs

4573 + assert np.sum(num_scheduled_tokens_np) == num_tokens

4574 + assert len(num_scheduled_tokens_np) == num_reqs

4575

4576 req_num_tokens = num_tokens // num_reqs

4577

4518 - dummy_prompt_lens = torch.tensor(

4519 - num_scheduled_tokens_list,

4520 - device="cpu",

4521 - )

4578 + dummy_prompt_lens = torch.from_numpy(num_scheduled_tokens_np)

4579 dummy_token_ids = torch.zeros(

4580 (num_reqs, req_num_tokens), dtype=torch.int32, device=self.device

4581 )

@@ -4537,7 +4594,7 @@ class GPUModelRunner(

4594 )

4595

4596 dummy_metadata.build_pooling_cursor(

4540 - num_scheduled_tokens_list,

4597 + num_scheduled_tokens_np,

4598 seq_lens_cpu=dummy_prompt_lens,

4599 device=hidden_states.device,

4600 )

vllm/v1/worker/gpu_worker.py

+4 -2

@@ -575,7 +575,7 @@ class Worker(WorkerBase):

575 @torch.inference_mode()

576 def execute_model(

577 self, scheduler_output: "SchedulerOutput"

578 - ) -> ModelRunnerOutput | None:

578 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

579 intermediate_tensors = None

580 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

581 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -624,7 +624,9 @@ class Worker(WorkerBase):

624 output = self.model_runner.execute_model(

625 scheduler_output, intermediate_tensors

626 )

627 - if isinstance(output, ModelRunnerOutput | NoneType):

627 + if isinstance(

628 + output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType

629 + ):

630 return output

631

632 assert isinstance(output, IntermediateTensors)

vllm/v1/worker/worker_base.py

+2 -2

@@ -124,7 +124,7 @@ class WorkerBase:

124

125 def execute_model(

126 self, scheduler_output: SchedulerOutput

127 - ) -> ModelRunnerOutput | None:

127 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

128 """If this method returns None, sample_tokens should be called immediately after

129 to obtain the ModelRunnerOutput.

130

@@ -362,7 +362,7 @@ class WorkerWrapperBase:

362 scheduler_output: SchedulerOutput,

363 *args,

364 **kwargs,

365 - ) -> ModelRunnerOutput | None:

365 + ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:

366 self._apply_mm_cache(scheduler_output)

367

368 assert self.worker is not None