@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 652ba59229

[Model Runner V2] Enable encoder token embedding (#50574)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

Taneem Ibrahim committed Aug 1, 2026 at 06:55 UTC 652ba59229499eb65fc4115b7feadeddf9bcb75d

4 files changed +106 -7

tests/models/language/pooling/test_colbert.py

+13 -2

@@ -379,8 +379,16 @@ def test_colbert_embed_not_supported(

379 vllm_model.embed([TEXTS_1[0]])

380

381

382 -@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys()))

383 -def test_colbert_hf_comparison(vllm_runner, backend):

382 +@pytest.mark.parametrize(

383 + ("backend", "use_v2"),

384 + [

385 + pytest.param("bert", True, id="bert-v2"),

386 + pytest.param("modernbert", True, id="modernbert-v2"),

387 + pytest.param("jina", False, id="jina-v1"),

388 + pytest.param("lfm2", False, id="lfm2-v1"),

389 + ],

390 +)

391 +def test_colbert_hf_comparison(vllm_runner, monkeypatch, backend, use_v2):

392 """Test that vLLM ColBERT embeddings match HuggingFace for each backend."""

393 from transformers import AutoTokenizer

394

@@ -393,6 +401,8 @@ def test_colbert_hf_comparison(vllm_runner, backend):

401 assert isinstance(extra_kwargs, dict)

402 test_texts = [TEXTS_1[0], TEXTS_2[0]]

403

404 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0")

405 +

406 with vllm_runner(

407 model_name,

408 runner="pooling",

@@ -401,6 +411,7 @@ def test_colbert_hf_comparison(vllm_runner, backend):

411 enforce_eager=True,

412 **extra_kwargs,

413 ) as vllm_model:

414 + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner == use_v2

415 vllm_outputs = vllm_model.token_embed(test_texts)

416

417 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tests/models/language/pooling/test_splade_sparse_pooler.py

+67 -1

@@ -15,8 +15,10 @@ from vllm.model_executor.models.bert import (

15 from vllm.platforms import current_platform

16 from vllm.pooling_params import PoolingParams

17 from vllm.utils.torch_utils import PIN_MEMORY

18 +from vllm.v1.pool.late_interaction_runner import LateInteractionRunner

19 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

20 from vllm.v1.worker.gpu.input_batch import InputBatch

21 +from vllm.v1.worker.gpu.model_runner import GPUModelRunner

22 from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner

23 from vllm.v1.worker.gpu.states import RequestState

@@ -138,9 +140,11 @@ def test_pooling_runner_stores_only_required_token_ids() -> None:

140 runner.pooling_params = {}

141 runner.pooling_states = {}

142 runner.prompt_token_ids = {}

143 + runner.late_interaction_runner = MagicMock()

144

142 - runner.add_request(1, PoolingParams(task="embed"), [101, 102])

145 + runner.add_request("req-1", 1, PoolingParams(task="embed"), [101, 102])

146 runner.add_request(

147 + "req-2",

148 2,

149 PoolingParams(task="embed", requires_token_ids=True),

150 [101, 11, 102],

@@ -150,6 +154,42 @@ def test_pooling_runner_stores_only_required_token_ids() -> None:

154 torch.testing.assert_close(runner.prompt_token_ids[2], torch.tensor([101, 11, 102]))

155

156

157 +def test_pooling_runner_releases_aborted_late_interaction_doc() -> None:

158 + runner = PoolingRunner.__new__(PoolingRunner)

159 + runner.late_interaction_runner = LateInteractionRunner()

160 +

161 + query_key = "query-abort"

162 + late_interaction_runner = runner.late_interaction_runner

163 + late_interaction_runner._query_cache[query_key] = torch.ones(2, 4)

164 + late_interaction_runner._query_uses[query_key] = 1

165 + late_interaction_runner._doc_query_keys["doc-req"] = query_key

166 +

167 + runner.on_requests_finished({"doc-req"})

168 +

169 + assert not late_interaction_runner._query_cache

170 + assert not late_interaction_runner._query_uses

171 + assert not late_interaction_runner._doc_query_keys

172 +

173 +

174 +def test_encoder_cache_reset_clears_late_interaction_state() -> None:

175 + # Cached query embeddings are only invalidated by weight reloads, which

176 + # reach the runner via reset_encoder_cache. Resetting the multi-modal

177 + # cache is unrelated and must leave them intact.

178 + runner = GPUModelRunner.__new__(GPUModelRunner)

179 + runner.encoder_cache = MagicMock()

180 + runner.pooling_runner = MagicMock()

181 +

182 + runner.reset_mm_cache()

183 +

184 + runner.encoder_cache.reset_mm_cache.assert_called_once_with()

185 + runner.pooling_runner.clear.assert_not_called()

186 +

187 + runner.reset_encoder_cache()

188 +

189 + runner.encoder_cache.reset_encoder_cache.assert_called_once_with()

190 + runner.pooling_runner.clear.assert_called_once_with()

191 +

192 +

193 def test_pooling_runner_rejects_unsupported_selected_task() -> None:

194 model = MagicMock()

195 model.pooler.get_supported_tasks.return_value = {

@@ -182,6 +222,19 @@ def test_pooling_runner_supports_encoder_token_classification() -> None:

222 assert runner.supported_tasks == {"token_classify"}

223

224

225 +def test_pooling_runner_supports_encoder_token_embedding() -> None:

226 + model = MagicMock()

227 + model.pooler.get_supported_tasks.return_value = {"token_embed"}

228 + vllm_config = MagicMock()

229 + vllm_config.scheduler_config.max_num_seqs = 2

230 + vllm_config.model_config.attn_type = "encoder_only"

231 + vllm_config.model_config.get_pooling_task.return_value = "token_embed"

232 +

233 + runner = PoolingRunner(model, vllm_config)

234 +

235 + assert runner.supported_tasks == {"token_embed"}

236 +

237 +

238 def test_pooling_runner_rejects_decoder_token_classification() -> None:

239 model = MagicMock()

240 model.pooler.get_supported_tasks.return_value = {"token_classify"}

@@ -207,3 +260,16 @@ def test_pooling_runner_filters_decoder_token_classification() -> None:

260 runner = PoolingRunner(model, vllm_config)

261

262 assert runner.supported_tasks == {"embed"}

263 +

264 +

265 +def test_pooling_runner_filters_decoder_token_embedding() -> None:

266 + model = MagicMock()

267 + model.pooler.get_supported_tasks.return_value = {"embed", "token_embed"}

268 + vllm_config = MagicMock()

269 + vllm_config.scheduler_config.max_num_seqs = 2

270 + vllm_config.model_config.attn_type = "decoder"

271 + vllm_config.model_config.get_pooling_task.return_value = "embed"

272 +

273 + runner = PoolingRunner(model, vllm_config)

274 +

275 + assert runner.supported_tasks == {"embed"}

vllm/v1/worker/gpu/model_runner.py

@@ -743,6 +743,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

743 def reset_encoder_cache(self) -> None:

744 if self.encoder_cache is not None:

745 self.encoder_cache.reset_encoder_cache()

746 + if self.pooling_runner is not None:

747 + self.pooling_runner.clear()

748

749 def profile_cudagraph_memory(self) -> int:

750 # NOTE(woosuk): It is TBD whether we keep this API or not.

@@ -813,6 +815,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

815

816 def finish_requests(self, scheduler_output: SchedulerOutput) -> None:

817 finished_req_ids = scheduler_output.finished_req_ids

818 + if self.pooling_runner is not None:

819 + # Preempted docs keep their query-use reservation until rescheduled.

820 + self.pooling_runner.on_requests_finished(finished_req_ids)

821 preempted_req_ids = scheduler_output.preempted_req_ids

822 if preempted_req_ids:

823 finished_req_ids = finished_req_ids.union(preempted_req_ids)

@@ -857,6 +862,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

862 if self.pooling_runner is not None:

863 assert new_req_data.pooling_params is not None

864 self.pooling_runner.add_request(

865 + req_id,

866 req_index,

867 new_req_data.pooling_params,

868 new_req_data.prompt_token_ids,

vllm/v1/worker/gpu/pool/pooling_runner.py

+20 -4

@@ -12,13 +12,15 @@ from vllm.pooling_params import PoolingParams

12 from vllm.tasks import PoolingTask

13 from vllm.utils.torch_utils import PIN_MEMORY

14 from vllm.v1.outputs import PoolerOutput

15 +from vllm.v1.pool.late_interaction_runner import LateInteractionRunner

16 from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates

17 from vllm.v1.worker.gpu.input_batch import InputBatch

18 from vllm.v1.worker.gpu.states import RequestState

20 _SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(

20 - {"embed", "classify", "token_classify"}

21 + {"embed", "classify", "token_embed", "token_classify"}

22 )

23 +_TOKEN_TASKS: frozenset[PoolingTask] = frozenset({"token_embed", "token_classify"})

26 class PoolingRunner:

@@ -53,13 +55,14 @@ class PoolingRunner:

55 self.pooling_params: dict[int, PoolingParams] = {}

56 self.pooling_states: dict[int, PoolingStates] = {}

57 self.prompt_token_ids: dict[int, torch.Tensor] = {}

58 + self.late_interaction_runner = LateInteractionRunner()

60 @staticmethod

61 def _get_enabled_tasks(model_config: ModelConfig) -> frozenset[PoolingTask]:

59 - # Token classification is validated only for unchunked encoder-only prefill.

62 + # Token-wise pooling is validated only for unchunked encoder-only prefill.

63 if model_config.attn_type == "encoder_only":

64 return _SUPPORTED_TASKS

62 - return _SUPPORTED_TASKS - {"token_classify"}

65 + return _SUPPORTED_TASKS - _TOKEN_TASKS

67 @staticmethod

68 def get_supported_tasks(

@@ -72,6 +75,7 @@ class PoolingRunner:

76 def add_request(

77 self,

78 + req_id: str,

79 req_index: int,

80 pooling_params: PoolingParams,

81 prompt_token_ids: list[int],

@@ -83,6 +87,7 @@ class PoolingRunner:

87 f"Supported tasks: {sorted(self.supported_tasks)}"

88 )

89 self.model.pooler.get_pooling_updates(task).apply(pooling_params)

90 + self.late_interaction_runner.register_request(req_id, pooling_params)

91 self.pooling_params[req_index] = pooling_params

92 self.pooling_states[req_index] = PoolingStates()

93 if pooling_params.requires_token_ids:

@@ -96,6 +101,12 @@ class PoolingRunner:

101 state.clean()

102 self.prompt_token_ids.pop(req_index, None)

103

104 + def on_requests_finished(self, req_ids: set[str]) -> None:

105 + self.late_interaction_runner.on_requests_finished(req_ids)

106 +

107 + def clear(self) -> None:

108 + self.late_interaction_runner.clear()

109 +

110 def _get_pooling_metadata(

111 self,

112 input_batch: InputBatch,

@@ -153,8 +164,13 @@ class PoolingRunner:

164 query_start_loc_gpu=input_batch.query_start_loc[: num_reqs + 1],

165 )

166 pooler_output = self.model.pooler(hidden_states, pooling_metadata)

156 -

167 finished_mask = pooling_metadata.get_pooling_cursor().is_finished().tolist()

168 + pooler_output = self.late_interaction_runner.postprocess_pooler_output(

169 + raw_pooler_output=pooler_output,

170 + pooling_params=pooling_metadata.pooling_params,

171 + req_ids=input_batch.req_ids,

172 + finished_mask=finished_mask,

173 + )

174 return pooler_output, finished_mask

175

176 def _dummy_pooler_run_task(