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
Remove runtime checks based on pooling params (#24051)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
Maximilien de Bayser committed
Sep 2, 2025 at 00:54 UTC
d59c986444a701b39369453eff0a8ba324bd565f
2 files changed
+14
-23
vllm/v1/worker/gpu_input_batch.py
+6
-11
@@ -704,17 +704,12 @@ class InputBatch:
704
logitsprocs=self.logitsprocs,
705
)
706
707
- @property
708
- def pooling_metadata(self) -> PoolingMetadata:
709
- if len(self.pooling_params) == 0:
710
- pooling_params = []
711
- else:
712
- # Note, for now this assumes that all request in the batch
713
- # are either sampling or pooling requests
714
- assert len(self.req_ids) == len(self.pooling_params)
715
- pooling_params = [
716
- self.pooling_params[req_id] for req_id in self.req_ids
717
- ]
707
+ def get_pooling_params(self) -> list[PoolingParams]:
708
+ assert len(self.req_ids) == len(self.pooling_params)
709
+ return [self.pooling_params[req_id] for req_id in self.req_ids]
710
+
711
+ def get_pooling_metadata(self) -> PoolingMetadata:
712
+ pooling_params = self.get_pooling_params()
713
714
return PoolingMetadata(
715
prompt_lens=torch.from_numpy(
vllm/v1/worker/gpu_model_runner.py
+8
-12
@@ -138,7 +138,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
138
self.kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[
139
cache_config.cache_dtype]
140
141
- self.is_pooling_model = model_config.pooler_config is not None
141
+ self.is_pooling_model = (model_config.runner_type == 'pooling')
142
self.is_multimodal_raw_input_only_model = (
143
model_config.is_multimodal_raw_input_only_model)
144
@@ -332,17 +332,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
332
333
def _init_model_kwargs(self, num_tokens: int):
334
model_kwargs = dict[str, Any]()
335
- num_reqs = self.input_batch.num_reqs
336
-
337
- num_pooling_reqs = len(self.input_batch.pooling_params)
335
339
- if num_pooling_reqs == 0:
336
+ if not self.is_pooling_model:
337
return model_kwargs
338
342
- # This does nontrivial work.
343
- pooling_params = self.input_batch.pooling_metadata.pooling_params
344
-
345
- assert num_pooling_reqs == num_reqs
339
+ num_reqs = self.input_batch.num_reqs
340
+ pooling_params = self.input_batch.get_pooling_params()
341
342
token_type_id_requests = dict[int, Any]()
343
for i, param in enumerate(pooling_params):
@@ -456,7 +451,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
451
else:
452
generator = None
453
459
- if pooling_params:
454
+ if self.is_pooling_model:
455
+ assert pooling_params is not None
456
task = pooling_params.task
457
assert task is not None, "You did not set `task` in the API"
458
@@ -1437,7 +1433,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1433
" a batch must be pooling request"
1434
1435
hidden_states = hidden_states[:num_scheduled_tokens]
1440
- pooling_metadata = self.input_batch.pooling_metadata
1436
+ pooling_metadata = self.input_batch.get_pooling_metadata()
1437
pooling_metadata.build_pooling_cursor(num_scheduled_tokens_np.tolist(),
1438
device=hidden_states.device)
1439
seq_lens_cpu = self.seq_lens.cpu[:self.input_batch.num_reqs]
@@ -1609,7 +1605,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1605
all_gather_group=get_tp_group())
1606
logits = None
1607
else:
1612
- if self.input_batch.pooling_params:
1608
+ if self.is_pooling_model:
1609
return self._pool(hidden_states, num_scheduled_tokens,
1610
num_scheduled_tokens_np, kv_connector_output)
1611