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 7, 14:00 UTC
[Model Runner V2] Support pooling models (#35120)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 27, 2026 at 18:03 UTC
86ac7bcf8483d87951a876cd2ed28341f60c95e0
6 files changed
+209
-14
vllm/v1/worker/gpu/async_utils.py
+36
@@ -70,6 +70,42 @@ class AsyncOutput(AsyncModelRunnerOutput):
70
return self.model_runner_output
71
72
73
+class AsyncPoolingOutput(AsyncModelRunnerOutput):
74
+ def __init__(
75
+ self,
76
+ model_runner_output: ModelRunnerOutput,
77
+ pooler_output: torch.Tensor,
78
+ is_valid: torch.Tensor | None,
79
+ main_stream: torch.cuda.Stream,
80
+ copy_stream: torch.cuda.Stream,
81
+ copy_event: torch.cuda.Event,
82
+ ):
83
+ self.model_runner_output = model_runner_output
84
+ self.pooler_output = pooler_output
85
+ self.is_valid = is_valid
86
+ self.copy_event = copy_event
87
+
88
+ with stream(copy_stream, main_stream):
89
+ copy_stream.wait_stream(main_stream)
90
+ self.pooler_output_cpu = self.pooler_output.to("cpu", non_blocking=True)
91
+ if self.is_valid is not None:
92
+ self.is_valid_cpu = self.is_valid.to("cpu", non_blocking=True)
93
+ else:
94
+ self.is_valid_cpu = None
95
+ self.copy_event.record(copy_stream)
96
+
97
+ def get_output(self) -> ModelRunnerOutput:
98
+ self.copy_event.synchronize()
99
+ pooler_output = self.pooler_output_cpu.unbind(dim=0)
100
+ if self.is_valid_cpu is not None:
101
+ is_valid_cpu = self.is_valid_cpu.tolist()
102
+ for i, is_valid in enumerate(is_valid_cpu):
103
+ if not is_valid:
104
+ pooler_output[i] = None
105
+ self.model_runner_output.pooler_output = pooler_output
106
+ return self.model_runner_output
107
+
108
+
109
def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
110
return x.to("cpu", non_blocking=True).numpy()
111
vllm/v1/worker/gpu/input_batch.py
+32
@@ -499,6 +499,38 @@ def post_update(
499
)
500
501
502
+@triton.jit
503
+def _post_update_pool_kernel(
504
+ idx_mapping_ptr,
505
+ num_computed_tokens_ptr,
506
+ query_start_loc_ptr,
507
+):
508
+ batch_id = tl.program_id(0)
509
+ query_start = tl.load(query_start_loc_ptr + batch_id)
510
+ query_end = tl.load(query_start_loc_ptr + batch_id + 1)
511
+ query_len = query_end - query_start
512
+
513
+ req_state_idx = tl.load(idx_mapping_ptr + batch_id)
514
+ num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
515
+ tl.store(num_computed_tokens_ptr + req_state_idx, num_computed + query_len)
516
+
517
+
518
+def post_update_pool(
519
+ # [num_reqs]
520
+ idx_mapping: torch.Tensor,
521
+ # [max_num_reqs]
522
+ num_computed_tokens: torch.Tensor,
523
+ # [num_reqs + 1]
524
+ query_start_loc: torch.Tensor,
525
+) -> None:
526
+ num_reqs = idx_mapping.shape[0]
527
+ _post_update_pool_kernel[(num_reqs,)](
528
+ idx_mapping,
529
+ num_computed_tokens,
530
+ query_start_loc,
531
+ )
532
+
533
+
534
@triton.jit
535
def _expand_idx_mapping_kernel(
536
idx_mapping_ptr,
vllm/v1/worker/gpu/model_runner.py
+90
-14
@@ -38,13 +38,14 @@ from vllm.logger import init_logger
38
from vllm.model_executor.model_loader import get_model_loader
39
from vllm.multimodal import MULTIMODAL_REGISTRY
40
from vllm.sequence import IntermediateTensors
41
+from vllm.tasks import SupportedTask
42
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
43
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
44
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
45
from vllm.v1.kv_cache_interface import KVCacheConfig
46
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
47
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
47
-from vllm.v1.worker.gpu.async_utils import AsyncOutput
48
+from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
49
from vllm.v1.worker.gpu.attn_utils import (
50
build_slot_mappings_by_layer,
51
get_kv_cache_spec,
@@ -66,6 +67,7 @@ from vllm.v1.worker.gpu.input_batch import (
67
expand_idx_mapping,
68
get_num_sampled_and_rejected,
69
post_update,
70
+ post_update_pool,
71
prepare_pos_seq_lens,
72
prepare_prefill_inputs,
73
)
@@ -77,6 +79,7 @@ from vllm.v1.worker.gpu.kv_connector import (
79
from vllm.v1.worker.gpu.lora_utils import LoraState
80
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
81
from vllm.v1.worker.gpu.model_states import ModelState
82
+from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
83
from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
84
from vllm.v1.worker.gpu.sample.output import SamplerOutput
85
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
@@ -119,7 +122,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
122
self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
123
self.cache_config.cache_dtype
124
]
122
- self.is_pooling_model = False
125
126
self.vocab_size = self.model_config.get_vocab_size()
127
self.max_model_len = self.model_config.max_model_len
@@ -217,6 +219,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
219
# KV Connector if configured.
220
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
221
222
+ # Pooling models.
223
+ self.is_pooling_model = self.model_config.runner_type == "pooling"
224
+ self.pooling_runner: PoolingRunner | None = None
225
+
226
# For transferring state from execute_model to subsequent sample_tokens call.
227
self.execute_model_state: tuple | None = None
228
@@ -224,9 +230,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
230
self.max_model_len = max_model_len
231
self.req_states.max_model_len = max_model_len
232
227
- @staticmethod
228
- def get_supported_tasks() -> tuple[str]:
229
- return ("generate",)
233
+ def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
234
+ tasks: list[SupportedTask] = []
235
+ if self.model_config.runner_type == "generate":
236
+ tasks.append("generate")
237
+ if self.pooling_runner is not None:
238
+ tasks.extend(self.pooling_runner.get_supported_pooling_tasks())
239
+ return tuple(tasks)
240
241
def load_model(self, *args, **kwargs) -> None:
242
time_before_load = time.perf_counter()
@@ -263,6 +273,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
273
274
# Initialize the components that require the model.
275
self.model_state = ModelState(self.vllm_config, self.model, self.device)
276
+ if self.is_pooling_model:
277
+ self.pooling_runner = PoolingRunner(self.model)
278
279
def get_model(self) -> nn.Module:
280
return self.model
@@ -388,16 +400,24 @@ class GPUModelRunner(LoRAModelRunnerMixin):
400
expanded_local_pos,
401
)
402
403
+ @torch.inference_mode()
404
+ def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
405
+ assert self.pooling_runner is not None
406
+ self.pooling_runner.dummy_pooler_run(hidden_states)
407
+
408
@torch.inference_mode()
409
def profile_run(self) -> None:
410
hidden_states, sample_hidden_states = self._dummy_run(
411
self.max_num_tokens, skip_attn=True
412
)
413
397
- # Only run sampler on last PP rank (non-last ranks return None).
414
+ # Only run sampler/pooler on last PP rank (non-last ranks return None).
415
if self.is_last_pp_rank:
416
assert sample_hidden_states is not None
400
- self._dummy_sampler_run(sample_hidden_states)
417
+ if self.pooling_runner is None:
418
+ self._dummy_sampler_run(sample_hidden_states)
419
+ else:
420
+ self._dummy_pooler_run(hidden_states)
421
422
if self.speculator is not None:
423
num_tokens_across_dp = make_num_tokens_across_dp(
@@ -505,7 +525,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
525
for new_req_data in scheduler_output.scheduled_new_reqs:
526
assert new_req_data.prompt_token_ids is not None
527
assert new_req_data.prefill_token_ids is not None
508
- assert new_req_data.sampling_params is not None
528
req_id = new_req_data.req_id
529
prompt_len = len(new_req_data.prompt_token_ids)
530
self.req_states.add_request(
@@ -523,14 +542,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
542
self.block_tables.append_block_ids(
543
req_index, new_req_data.block_ids, overwrite=True
544
)
526
- self.sampler.add_request(
527
- req_index, prompt_len, new_req_data.sampling_params
528
- )
529
- self.prompt_logprobs_worker.add_request(
530
- req_id, req_index, new_req_data.sampling_params
531
- )
545
self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)
546
547
+ if new_req_data.sampling_params is not None:
548
+ self.sampler.add_request(
549
+ req_index, prompt_len, new_req_data.sampling_params
550
+ )
551
+ self.prompt_logprobs_worker.add_request(
552
+ req_id, req_index, new_req_data.sampling_params
553
+ )
554
+
555
if scheduler_output.scheduled_new_reqs:
556
self.req_states.apply_staged_writes()
557
self.sampler.apply_staged_writes()
@@ -1083,3 +1104,58 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1104
1105
def take_draft_token_ids(self) -> DraftTokenIds | None:
1106
return self.draft_tokens_handler.get_draft_tokens()
1107
+
1108
+ @torch.inference_mode()
1109
+ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None:
1110
+ if self.execute_model_state is None:
1111
+ # The prior execute_model call must have failed.
1112
+ return None
1113
+
1114
+ input_batch, _, _, _, hidden_states, _, kv_connector_output = (
1115
+ self.execute_model_state
1116
+ )
1117
+ self.execute_model_state = None
1118
+
1119
+ if not self.is_last_pp_rank:
1120
+ self.postprocess_pool(input_batch)
1121
+ return None
1122
+
1123
+ assert self.pooling_runner is not None
1124
+ pooler_output, is_valid = self.pooling_runner.pool(
1125
+ hidden_states, input_batch, self.req_states
1126
+ )
1127
+ self.postprocess_pool(input_batch)
1128
+
1129
+ # Build the model runner output.
1130
+ model_runner_output = ModelRunnerOutput(
1131
+ req_ids=input_batch.req_ids,
1132
+ req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
1133
+ kv_connector_output=kv_connector_output,
1134
+ )
1135
+ async_output = AsyncPoolingOutput(
1136
+ model_runner_output=model_runner_output,
1137
+ pooler_output=pooler_output,
1138
+ is_valid=is_valid,
1139
+ main_stream=self.main_stream,
1140
+ copy_stream=self.output_copy_stream,
1141
+ copy_event=self.output_copy_event,
1142
+ )
1143
+ if self.use_async_scheduling:
1144
+ return async_output
1145
+ return async_output.get_output()
1146
+
1147
+ def postprocess_pool(self, input_batch: InputBatch) -> None:
1148
+ # Update the number of computed tokens.
1149
+ post_update_pool(
1150
+ input_batch.idx_mapping,
1151
+ self.req_states.num_computed_tokens.gpu,
1152
+ input_batch.query_start_loc,
1153
+ )
1154
+
1155
+ # Update the number of computed prefill tokens.
1156
+ idx_mapping_np = input_batch.idx_mapping_np
1157
+ computed_prefill = self.req_states.num_computed_prefill_tokens
1158
+ computed_prefill[idx_mapping_np] += input_batch.num_scheduled_tokens
1159
+ np.minimum(
1160
+ computed_prefill, self.req_states.prefill_len.np, out=computed_prefill
1161
+ )
vllm/v1/worker/gpu/pool/__init__.py
vllm/v1/worker/gpu/pool/pooling_runner.py
new
+45
@@ -0,0 +1,45 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
+
5
+import torch
6
+import torch.nn as nn
7
+import torch.nn.functional as F
8
+
9
+from vllm.model_executor.models import VllmModelForPooling, is_pooling_model
10
+from vllm.tasks import PoolingTask
11
+from vllm.v1.worker.gpu.input_batch import InputBatch
12
+from vllm.v1.worker.gpu.states import RequestState
13
+
14
+
15
+# NOTE(woosuk): Currently, this class only supports the "LAST" pooling task
16
+# on decoder-only models. How to support other pooling tasks and models
17
+# is to be determined.
18
+class PoolingRunner:
19
+ def __init__(self, model: nn.Module):
20
+ self.model = cast(VllmModelForPooling, model)
21
+
22
+ def get_supported_pooling_tasks(self) -> list[PoolingTask]:
23
+ if not is_pooling_model(self.model):
24
+ return []
25
+ assert "embed" in self.model.pooler.get_supported_tasks()
26
+ return ["embed"]
27
+
28
+ def pool(
29
+ self,
30
+ hidden_states: torch.Tensor,
31
+ input_batch: InputBatch,
32
+ req_states: RequestState,
33
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
34
+ # TODO(woosuk): Support different types of pooling tasks.
35
+ last_hidden_states = hidden_states[input_batch.logits_indices]
36
+ # TODO(woosuk): Make normalization optional.
37
+ last_hidden_states = F.normalize(last_hidden_states, p=2, dim=-1)
38
+
39
+ prompt_len = req_states.prompt_len.gpu[input_batch.idx_mapping]
40
+ is_valid = input_batch.seq_lens == prompt_len
41
+ return last_hidden_states, is_valid
42
+
43
+ def dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
44
+ F.normalize(hidden_states, p=2, dim=-1)
45
+ return
vllm/v1/worker/gpu_worker.py
+6
@@ -700,6 +700,12 @@ class Worker(WorkerBase):
700
output = self.model_runner.execute_model(
701
scheduler_output, intermediate_tensors
702
)
703
+ if (
704
+ self.use_v2_model_runner
705
+ and self.model_runner.is_pooling_model
706
+ and output is None
707
+ ):
708
+ output = self.model_runner.pool() # type: ignore
709
if isinstance(
710
output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType
711
):