@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 39052dbca8
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

Support token_type_ids in V1 with less code changes (#21985)

Signed-off-by: Max de Bayser <mbayser@br.ibm.com>

Maximilien de Bayser committed Aug 11, 2025 at 02:54 UTC 39052dbca87616a549ab152713f1a3020b2f4eb8
10 files changed +235 -130
tests/entrypoints/openai/test_rerank.py
+3 -1
@@ -126,7 +126,9 @@ def test_invocations(server: RemoteOpenAIServer):
126 invocation_output["results"]):
127 assert rerank_result.keys() == invocations_result.keys()
128 assert rerank_result["relevance_score"] == pytest.approx(
129 - invocations_result["relevance_score"], rel=0.01)
129 + invocations_result["relevance_score"], rel=0.05)
130 + # TODO: reset this tolerance to 0.01 once we find
131 + # an alternative to flash_attn with bfloat16
132
133
134 @pytest.mark.asyncio
tests/entrypoints/openai/test_score.py
+3 -1
@@ -220,7 +220,9 @@ class TestModel:
220 invocation_output["data"]):
221 assert score_data.keys() == invocation_data.keys()
222 assert score_data["score"] == pytest.approx(
223 - invocation_data["score"], rel=0.01)
223 + invocation_data["score"], rel=0.05)
224 + # TODO: reset this tolerance to 0.01 once we find
225 + # an alternative to flash_attn with bfloat16
226
227 def test_activation(self, server: RemoteOpenAIServer, model: dict[str,
228 Any]):
tests/models/language/pooling/test_scoring.py
+9
@@ -23,6 +23,15 @@ TEXTS_2 = [
23 "The capital of Germany is Berlin.",
24 ]
25
26 +
27 +@pytest.fixture(autouse=True)
28 +def v1(run_with_both_engines):
29 + # Simple autouse wrapper to run both engines for each test
30 + # This can be promoted up to conftest.py to run for every
31 + # test in a package
32 + pass
33 +
34 +
35 DTYPE = "half"
36
37
vllm/entrypoints/llm.py
+26 -28
@@ -28,11 +28,15 @@ from vllm.entrypoints.chat_utils import (ChatCompletionMessageParam,
28 apply_mistral_chat_template,
29 parse_chat_messages,
30 resolve_chat_template_content_format)
31 +# yapf conflicts with isort for this block
32 +# yapf: disable
33 from vllm.entrypoints.score_utils import (ScoreContentPartParam,
34 ScoreMultiModalParam,
35 _cosine_similarity,
36 _validate_score_input_lens,
37 + compress_token_type_ids,
38 get_score_prompt)
39 +# yapf: enable
40 from vllm.entrypoints.utils import (_validate_truncation_size,
41 log_non_default_args)
42 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
@@ -1329,6 +1333,7 @@ class LLM:
1333
1334 model_config = self.llm_engine.model_config
1335 pooling_params.verify("score", model_config)
1336 + pooling_params_list = list[PoolingParams]()
1337
1338 tokenization_kwargs: dict[str, Any] = {}
1339
@@ -1339,38 +1344,31 @@ class LLM:
1344
1345 input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
1346
1342 - if model_config.is_multimodal_model:
1343 - for q, d in input_pairs:
1344 - _, engine_prompt = get_score_prompt(
1345 - model_config=model_config,
1346 - data_1=q,
1347 - data_2=d,
1348 - tokenizer=tokenizer,
1349 - tokenization_kwargs=tokenization_kwargs,
1350 - )
1347 + model_config = self.llm_engine.model_config
1348
1352 - parsed_prompts.append(engine_prompt)
1353 - else:
1354 - for q, t in input_pairs:
1355 - if model_config.use_pad_token:
1356 - # cross_encoder models defaults to using pad_token.
1357 - prompt_inputs = tokenizer(
1358 - text=q, # type: ignore[arg-type]
1359 - text_pair=t, # type: ignore[arg-type]
1360 - **tokenization_kwargs)
1361 - else:
1362 - # `llm as reranker` models defaults to not using pad_token.
1363 - prompt_inputs = tokenizer(
1364 - text=q + t, # type: ignore[operator]
1365 - **tokenization_kwargs)
1366 - engine_prompt = TokensPrompt(
1367 - prompt_token_ids=prompt_inputs["input_ids"],
1368 - token_type_ids=prompt_inputs.get("token_type_ids"))
1369 - parsed_prompts.append(engine_prompt)
1349 + for q, d in input_pairs:
1350 + _, engine_prompt = get_score_prompt(
1351 + model_config=model_config,
1352 + data_1=q,
1353 + data_2=d,
1354 + tokenizer=tokenizer,
1355 + tokenization_kwargs=tokenization_kwargs,
1356 + )
1357 +
1358 + if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
1359 + "token_type_ids", None)):
1360 + params = pooling_params.clone()
1361 + compressed = compress_token_type_ids(token_type_ids)
1362 + params.extra_kwargs = {"compressed_token_type_ids": compressed}
1363 + pooling_params_list.append(params)
1364 + else:
1365 + pooling_params_list.append(pooling_params)
1366 +
1367 + parsed_prompts.append(engine_prompt)
1368
1369 self._validate_and_add_requests(
1370 prompts=parsed_prompts,
1373 - params=pooling_params,
1371 + params=pooling_params_list,
1372 use_tqdm=use_tqdm,
1373 lora_request=lora_request,
1374 )
vllm/entrypoints/openai/serving_score.py
+31 -51
@@ -7,6 +7,7 @@ from typing import Any, Optional, Union
7
8 from fastapi import Request
9
10 +from vllm import envs
11 from vllm.config import ModelConfig
12 from vllm.engine.protocol import EngineClient
13 from vllm.entrypoints.logger import RequestLogger
@@ -17,11 +18,15 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse, RerankDocument,
18 ScoreResponseData, UsageInfo)
19 from vllm.entrypoints.openai.serving_engine import OpenAIServing
20 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
21 +# yapf conflicts with isort for this block
22 +# yapf: disable
23 from vllm.entrypoints.score_utils import (ScoreContentPartParam,
24 ScoreMultiModalParam,
25 _cosine_similarity,
26 _validate_score_input_lens,
27 + compress_token_type_ids,
28 get_score_prompt)
29 +# yapf: enable
30 from vllm.entrypoints.utils import _validate_truncation_size
31 from vllm.inputs.data import TokensPrompt
32 from vllm.logger import init_logger
@@ -158,6 +163,8 @@ class ServingScores(OpenAIServing):
163 tokenizer=tokenizer,
164 tokenization_kwargs=tokenization_kwargs,
165 )
166 + self._validate_input(request, engine_prompt["prompt_token_ids"],
167 + full_prompt)
168 if request.mm_processor_kwargs is not None:
169 engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs
170
@@ -188,64 +195,27 @@ class ServingScores(OpenAIServing):
195
196 input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]
197
191 - if self.model_config.is_multimodal_model:
198 + preprocess_async = make_async(self._preprocess_score,
199 + executor=self._tokenizer_executor)
200
193 - preprocess_async = make_async(self._preprocess_score,
194 - executor=self._tokenizer_executor)
201 + preprocessed_prompts = await asyncio.gather(
202 + *(preprocess_async(request=request,
203 + tokenizer=tokenizer,
204 + tokenization_kwargs=tokenization_kwargs,
205 + data_1=t1,
206 + data_2=t2) for t1, t2 in input_pairs))
207
196 - preprocessed_prompts = await asyncio.gather(
197 - *(preprocess_async(request=request,
198 - tokenizer=tokenizer,
199 - tokenization_kwargs=tokenization_kwargs,
200 - data_1=t1,
201 - data_2=t2) for t1, t2 in input_pairs))
202 -
203 - for full_prompt, engine_prompt in preprocessed_prompts:
204 - request_prompts.append(full_prompt)
205 - engine_prompts.append(engine_prompt)
206 -
207 - else:
208 - tokenize_async = make_async(tokenizer.__call__,
209 - executor=self._tokenizer_executor)
210 - use_pad_token = self.model_config.use_pad_token
211 -
212 - if use_pad_token:
213 - # cross_encoder models defaults to using pad_token.
214 - tokenized_prompts = await asyncio.gather(*(
215 - tokenize_async(
216 - text=t1, # type: ignore[arg-type]
217 - text_pair=t2, # type: ignore[arg-type]
218 - **tokenization_kwargs) for t1, t2 in input_pairs))
219 - else:
220 - # `llm as reranker` models defaults to not using pad_token.
221 - tokenized_prompts = await asyncio.gather(*(
222 - tokenize_async(
223 - text=t1 + # type: ignore[operator]
224 - t2,
225 - **tokenization_kwargs) for t1, t2 in input_pairs))
226 -
227 - for prompt_inputs, (t1, t2) in zip(tokenized_prompts, input_pairs):
228 - sep_token = tokenizer.sep_token if (tokenizer.sep_token
229 - and use_pad_token) else ''
230 - request_prompt = f"{t1}{sep_token}{t2}"
231 -
232 - input_ids = prompt_inputs["input_ids"]
233 - text_token_prompt = \
234 - self._validate_input(request, input_ids, request_prompt)
235 - engine_prompt = TokensPrompt(
236 - prompt_token_ids=text_token_prompt["prompt_token_ids"],
237 - token_type_ids=prompt_inputs.get("token_type_ids"))
238 -
239 - request_prompts.append(request_prompt)
240 - engine_prompts.append(engine_prompt)
208 + for full_prompt, engine_prompt in preprocessed_prompts:
209 + request_prompts.append(full_prompt)
210 + engine_prompts.append(engine_prompt)
211
212 # Schedule the request and get the result generator.
213 generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
214
245 - pooling_params = request.to_pooling_params()
215 + default_pooling_params = request.to_pooling_params()
216
217 try:
248 - pooling_params.verify("score", self.model_config)
218 + default_pooling_params.verify("score", self.model_config)
219 except ValueError as e:
220 return self.create_error_response(str(e))
221
@@ -254,9 +224,19 @@ class ServingScores(OpenAIServing):
224
225 self._log_inputs(request_id_item,
226 request_prompts[i],
257 - params=pooling_params,
227 + params=default_pooling_params,
228 lora_request=lora_request)
229
230 + if envs.VLLM_USE_V1 and (token_type_ids := engine_prompt.pop(
231 + "token_type_ids", None)):
232 + pooling_params = default_pooling_params.clone()
233 + compressed = compress_token_type_ids(token_type_ids)
234 + pooling_params.extra_kwargs = {
235 + "compressed_token_type_ids": compressed
236 + }
237 + else:
238 + pooling_params = (default_pooling_params)
239 +
240 generator = self.engine_client.encode(
241 engine_prompt,
242 pooling_params,
vllm/entrypoints/score_utils.py
+37 -3
@@ -184,15 +184,49 @@ def get_score_prompt(
184 model_config,
185 tokenizer,
186 )
187 + from vllm.model_executor.model_loader import get_model_cls
188
188 - full_prompt = apply_score_template(model_config, prompt_1, prompt_2)
189 -
190 - prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
189 + model = get_model_cls(model_config)
190 + if supports_score_template(model):
191 + full_prompt = apply_score_template(model_config, prompt_1, prompt_2)
192 + prompt_inputs = tokenizer(full_prompt, **tokenization_kwargs)
193 + elif model_config.use_pad_token:
194 + # cross_encoder models defaults to using pad_token.
195 + prompt_inputs = tokenizer(text=prompt_1,
196 + text_pair=prompt_2,
197 + **tokenization_kwargs)
198 + full_prompt = tokenizer.decode(prompt_inputs["input_ids"])
199 + else:
200 + # `llm as reranker` models defaults to not using pad_token.
201 + full_prompt = prompt_1 + prompt_2
202 + prompt_inputs = tokenizer(text=full_prompt, **tokenization_kwargs)
203
204 engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["input_ids"])
205
206 + if (token_type_ids := prompt_inputs.get("token_type_ids")) is not None:
207 + engine_prompt["token_type_ids"] = token_type_ids
208 +
209 post_process_tokens(model_config, engine_prompt)
210
211 if mm_data is not None:
212 engine_prompt["multi_modal_data"] = mm_data
213 return full_prompt, engine_prompt
214 +
215 +
216 +def compress_token_type_ids(token_type_ids: list[int]) -> int:
217 + """
218 + Return position of the first 1 or the length of the list
219 + if not found.
220 + """
221 + first_one = len(token_type_ids)
222 + err_msg = "Token type ids are expected to be a sequence"\
223 + " of zeros followed by a sequence of ones"
224 + for i, type_id in enumerate(token_type_ids):
225 + if type_id == 0 and first_one < i:
226 + raise ValueError(err_msg)
227 + elif type_id == 1 and first_one > i:
228 + first_one = i
229 + elif type_id > 1:
230 + raise ValueError(err_msg)
231 +
232 + return first_one
vllm/model_executor/models/bert.py
+63 -25
@@ -28,7 +28,7 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata
28 from vllm.sequence import IntermediateTensors
29 from vllm.tasks import PoolingTask
30
31 -from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only
31 +from .interfaces import SupportsCrossEncoding, SupportsQuant
32 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
33
34
@@ -60,21 +60,13 @@ class BertEmbedding(nn.Module):
60 self,
61 input_ids: torch.Tensor,
62 position_ids: torch.Tensor,
63 - token_type_ids: Optional[torch.Tensor] = None,
63 ) -> torch.Tensor:
65 - input_shape = input_ids.size()
64
67 - # Input embeddings.
68 - inputs_embeds = self.word_embeddings(input_ids)
65 + token_type_ids = _decode_token_type_ids(input_ids)
66
70 - # Position embeddings.
67 + inputs_embeds = self.word_embeddings(input_ids)
68 position_embeddings = self.position_embeddings(position_ids)
69
73 - if token_type_ids is None:
74 - token_type_ids = torch.zeros(input_shape,
75 - dtype=torch.long,
76 - device=inputs_embeds.device)
77 -
70 token_type_embeddings = self.token_type_embeddings(token_type_ids)
71
72 embeddings = inputs_embeds + token_type_embeddings + position_embeddings
@@ -350,25 +342,23 @@ class BertModel(nn.Module, SupportsQuant):
342 ) -> None:
343 super().__init__()
344
353 - config = vllm_config.model_config.hf_config
354 - self.embeddings = embedding_class(config)
345 + self.config = vllm_config.model_config.hf_config
346 + self.embeddings = embedding_class(self.config)
347 self.encoder = BertEncoder(vllm_config=vllm_config,
348 prefix=f"{prefix}.encoder")
349
350 def forward(
351 self,
352 input_ids: torch.Tensor,
361 - position_ids: torch.Tensor,
353 + positions: torch.Tensor,
354 intermediate_tensors: Optional[IntermediateTensors] = None,
355 inputs_embeds: Optional[torch.Tensor] = None,
364 - token_type_ids: Optional[torch.Tensor] = None,
356 ) -> torch.Tensor:
357 if inputs_embeds is not None:
358 hidden_states = inputs_embeds
359 else:
360 hidden_states = self.embeddings(input_ids=input_ids,
370 - position_ids=position_ids,
371 - token_type_ids=token_type_ids)
361 + position_ids=positions)
362 return self.encoder(hidden_states)
363
364 def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
@@ -468,13 +458,11 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
458 self,
459 input_ids: torch.Tensor,
460 positions: torch.Tensor,
471 - token_type_ids: Optional[torch.Tensor] = None,
461 intermediate_tensors: Optional[IntermediateTensors] = None,
462 inputs_embeds: Optional[torch.Tensor] = None,
463 ) -> torch.Tensor:
464 return self.model(input_ids=input_ids,
476 - position_ids=positions,
477 - token_type_ids=token_type_ids,
465 + positions=positions,
466 inputs_embeds=inputs_embeds,
467 intermediate_tensors=intermediate_tensors)
468
@@ -508,8 +496,53 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
496 })
497
498
511 -class BertForSequenceClassification(nn.Module, SupportsV0Only,
512 - SupportsCrossEncoding, SupportsQuant):
499 +# Here we encode the token type ids together with the input ids.
500 +# Since we use int 32 for the input IDs and the vocabulary size
501 +# is way lower than 2**31, there is room to encode additional
502 +# bits. At the same time, for cross-encoder use cases, the
503 +# token type ids are only 0 or 1, requiring only 1 bit.
504 +# This means that we can store the token type ids in the 31st
505 +# bit. We void the 32nd bit because that would produce a negative
506 +# number, which could be used to signal other things.
507 +#
508 +# The reason for all of this is that all the tensors that are
509 +# passed as input to the forward function of a module marked
510 +# with @support_torch_compile have to be persistent. So to
511 +# avoid adding more persistent tensors in the model runner, we
512 +# encode more information in the same persistent tensor.
513 +#
514 +# Since the *ForClassification module is outside of the BertModel
515 +# which is compiled, we can do the encoding here and then separate
516 +# the information again in the Embedding layer. Since with bit masks
517 +# we can do this entirely with torch operations and without branching,
518 +# it works with torch compile.
519 +
520 +TOKEN_TYPE_SHIFT = 30
521 +
522 +
523 +def _encode_token_type_ids(input_ids: torch.Tensor,
524 + token_type_ids: torch.Tensor) -> None:
525 + # input_ids can be padded to the right
526 + input_ids[:token_type_ids.shape[0]].bitwise_or_(
527 + token_type_ids << TOKEN_TYPE_SHIFT)
528 +
529 +
530 +def _decode_token_type_ids(input_ids: torch.Tensor) -> torch.Tensor:
531 +
532 + ids_mask = torch.ones(input_ids.shape,
533 + dtype=torch.int32,
534 + device=input_ids.device) << TOKEN_TYPE_SHIFT
535 + tokens_mask = ids_mask.bitwise_not()
536 +
537 + token_type_ids = input_ids.bitwise_and(ids_mask) >> TOKEN_TYPE_SHIFT
538 +
539 + input_ids.bitwise_and_(tokens_mask)
540 +
541 + return token_type_ids
542 +
543 +
544 +class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
545 + SupportsQuant):
546 """A model that uses Bert to provide embedding functionalities.
547
548 This class encapsulates the BertModel and provides an interface for
@@ -567,8 +600,13 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
600 inputs_embeds: Optional[torch.Tensor] = None,
601 token_type_ids: Optional[torch.Tensor] = None,
602 ) -> torch.Tensor:
603 +
604 + if token_type_ids is not None:
605 + assert self.bert.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)
606 + assert input_ids is not None
607 + _encode_token_type_ids(input_ids, token_type_ids)
608 +
609 return self.bert(input_ids=input_ids,
571 - position_ids=positions,
610 + positions=positions,
611 inputs_embeds=inputs_embeds,
573 - intermediate_tensors=intermediate_tensors,
574 - token_type_ids=token_type_ids)
612 + intermediate_tensors=intermediate_tensors)
vllm/model_executor/models/roberta.py
+17 -19
@@ -14,13 +14,16 @@ from vllm.model_executor.layers.pooler import (ClassifierPooler, CLSPool,
14 DispatchPooler, Pooler)
15 from vllm.model_executor.layers.vocab_parallel_embedding import (
16 VocabParallelEmbedding)
17 -from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
17 +from vllm.model_executor.models.bert import (TOKEN_TYPE_SHIFT,
18 + BertEmbeddingModel, BertModel,
19 + _decode_token_type_ids,
20 + _encode_token_type_ids)
21 from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
22 maybe_prefix)
23 from vllm.sequence import IntermediateTensors
24
25 from .bert_with_rope import BertWithRope, JinaRobertaModel
23 -from .interfaces import SupportsCrossEncoding, SupportsV0Only
26 +from .interfaces import SupportsCrossEncoding
27
28
29 class RobertaEmbedding(nn.Module):
@@ -53,17 +56,12 @@ class RobertaEmbedding(nn.Module):
56 self,
57 input_ids: torch.Tensor,
58 position_ids: torch.Tensor,
56 - token_type_ids: Optional[torch.Tensor] = None,
59 ) -> torch.Tensor:
58 - input_shape = input_ids.size()
59 - inputs_embeds = self.word_embeddings(input_ids)
60
61 - # Position embeddings.
61 + token_type_ids = _decode_token_type_ids(input_ids)
62 +
63 + inputs_embeds = self.word_embeddings(input_ids)
64 position_embeddings = self.position_embeddings(position_ids)
63 - if token_type_ids is None:
64 - token_type_ids = torch.zeros(input_shape,
65 - dtype=torch.long,
66 - device=inputs_embeds.device)
65
66 token_type_embeddings = self.token_type_embeddings(token_type_ids)
67 embeddings = inputs_embeds + token_type_embeddings + position_embeddings
@@ -107,7 +105,6 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
105 self,
106 input_ids: torch.Tensor,
107 positions: torch.Tensor,
110 - token_type_ids: Optional[torch.Tensor] = None,
108 intermediate_tensors: Optional[IntermediateTensors] = None,
109 inputs_embeds: Optional[torch.Tensor] = None,
110 ) -> torch.Tensor:
@@ -119,9 +116,8 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
116 position_ids=positions,
117 padding_idx=self.padding_idx)
118
122 - return self.model(input_ids,
123 - positions,
124 - token_type_ids=token_type_ids,
119 + return self.model(input_ids=input_ids,
120 + positions=positions,
121 inputs_embeds=inputs_embeds,
122 intermediate_tensors=intermediate_tensors)
123
@@ -153,8 +149,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
149 return loader.load_weights(weights_list, mapper=mapper)
150
151
156 -class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
157 - SupportsV0Only):
152 +class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
153 """A model that uses Roberta to provide embedding functionalities.
154
155 This class encapsulates the BertModel and provides an interface for
@@ -226,11 +221,14 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
221 replace_roberta_positions(input_ids=input_ids,
222 position_ids=positions,
223 padding_idx=self.padding_idx)
224 + if token_type_ids is not None:
225 + assert self.roberta.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)
226 + assert input_ids is not None
227 + _encode_token_type_ids(input_ids, token_type_ids)
228 return self.roberta(input_ids=input_ids,
230 - position_ids=positions,
229 + positions=positions,
230 inputs_embeds=inputs_embeds,
232 - intermediate_tensors=intermediate_tensors,
233 - token_type_ids=token_type_ids)
231 + intermediate_tensors=intermediate_tensors)
232
233
234 # Adapted from transformers
vllm/pooling_params.py
+6 -2
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from copy import deepcopy
5 -from typing import TYPE_CHECKING, Optional
5 +from typing import TYPE_CHECKING, Any, Optional
6
7 import msgspec
8
@@ -46,6 +46,9 @@ class PoolingParams(
46 requires_token_ids: bool = False
47 """Internal use only."""
48
49 + extra_kwargs: Optional[dict[str, Any]] = None
50 + """Internal use only."""
51 +
52 output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
53
54 @property
@@ -167,7 +170,8 @@ class PoolingParams(
170 f"softmax={self.softmax}, "
171 f"step_tag_id={self.step_tag_id}, "
172 f"returned_token_ids={self.returned_token_ids}, "
170 - f"requires_token_ids={self.requires_token_ids})")
173 + f"requires_token_ids={self.requires_token_ids}, "
174 + f"extra_kwargs={self.extra_kwargs})")
175
176 def __post_init__(self) -> None:
177 assert self.output_kind == RequestOutputKind.FINAL_ONLY,\
vllm/v1/worker/gpu_model_runner.py
+40
@@ -336,6 +336,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
336
337 self.reorder_batch_threshold: Optional[int] = None
338
339 + def _init_model_kwargs(self, num_tokens: int):
340 + model_kwargs = dict[str, Any]()
341 + num_reqs = self.input_batch.num_reqs
342 +
343 + pooling_params = self.input_batch.pooling_metadata.pooling_params
344 +
345 + num_pooling_reqs = len(pooling_params)
346 +
347 + if num_pooling_reqs == 0:
348 + return model_kwargs
349 +
350 + assert num_pooling_reqs == num_reqs
351 +
352 + token_type_id_requests = dict[int, Any]()
353 + for i, param in enumerate(pooling_params):
354 + if param.extra_kwargs is not None and \
355 + (token_types := param.extra_kwargs.get(
356 + "compressed_token_type_ids")) is not None:
357 + token_type_id_requests[i] = token_types
358 +
359 + if len(token_type_id_requests) == 0:
360 + return model_kwargs
361 +
362 + seq_lens = self.seq_lens[:num_reqs]
363 + token_type_ids = []
364 +
365 + for i in range(num_reqs):
366 + pos = token_type_id_requests.get(i, seq_lens[i])
367 + ids = (torch.arange(seq_lens[i]) >= pos).int()
368 + token_type_ids.append(ids)
369 +
370 + model_kwargs["token_type_ids"] = torch.concat(token_type_ids).to(
371 + device=self.device)
372 + return model_kwargs
373 +
374 def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
375 """
376 Update the order of requests in the batch based on the attention
@@ -1504,12 +1539,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1539 input_ids = None
1540 inputs_embeds = self.inputs_embeds[:num_input_tokens]
1541 model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
1542 + model_kwargs = self._init_model_kwargs(num_scheduled_tokens)
1543 else:
1544 # For text-only models, we use token ids as input.
1545 # While it is possible to use embeddings as input just like the
1546 # multimodal models, it is not desirable for performance since
1547 # then the embedding layer is not included in the CUDA graph.
1548 input_ids = self.input_ids[:num_input_tokens]
1549 + model_kwargs = self._init_model_kwargs(num_input_tokens)
1550 inputs_embeds = None
1551 model_mm_kwargs = {}
1552 if self.uses_mrope:
@@ -1548,6 +1585,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1585 model_mm_kwargs,
1586 device=self.device,
1587 ),
1588 + **model_kwargs,
1589 )
1590
1591 if self.use_aux_hidden_state_outputs:
@@ -2211,6 +2249,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2249
2250 with self.maybe_dummy_run_with_lora(self.lora_config,
2251 num_scheduled_tokens):
2252 + model_kwargs = self._init_model_kwargs(num_tokens)
2253 if self.supports_mm_inputs:
2254 input_ids = None
2255 inputs_embeds = self.inputs_embeds[:num_tokens]
@@ -2252,6 +2291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2291 model_mm_kwargs,
2292 device=self.device,
2293 ),
2294 + **model_kwargs,
2295 )
2296
2297 if self.use_aux_hidden_state_outputs: