@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9464529612

[ModelRunner v2] Enable decoder token-wise pooling (#50931)

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

Taneem Ibrahim committed Aug 7, 2026 at 00:21 UTC 946452961265514622b95ea170839b214f39c2a0

8 files changed +86 -58

tests/model_executor/layers/test_pooler_methods.py

+30 -6

@@ -256,6 +256,7 @@ class TestGetSeqPoolingMethod:

256 @dataclass

257 class _FakeSchedulerConfig:

258 enable_chunked_prefill: bool = False

259 + async_scheduling: bool = False

260

261

262 @dataclass

@@ -265,10 +266,13 @@ class _FakeVllmConfig:

266

267 class TestAllPool:

268 @staticmethod

268 - def _make_all_pool(*, chunked: bool = False) -> AllPool:

269 + def _make_all_pool(

270 + *, chunked: bool = False, async_scheduling: bool = False

271 + ) -> AllPool:

272 fake_config = _FakeVllmConfig(

273 scheduler_config=_FakeSchedulerConfig(

274 enable_chunked_prefill=chunked,

275 + async_scheduling=async_scheduling,

276 ),

277 )

278 with patch(

@@ -319,6 +323,8 @@ class TestAllPool:

323 )

324 out1 = pooler(chunk1, metadata1)

325 assert out1[0] is None

326 + expected_chunk1 = chunk1.clone()

327 + chunk1.zero_()

328

329 chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

330 metadata2 = _make_metadata(

@@ -329,19 +335,37 @@ class TestAllPool:

335 metadata2.pooling_states = metadata1.pooling_states

336 out2 = pooler(chunk2, metadata2)

337 assert out2[0] is not None

332 - expected = torch.cat([chunk1, chunk2], dim=0)

338 + expected = torch.cat([expected_chunk1, chunk2], dim=0)

339 assert torch.equal(out2[0], expected)

340

335 - def test_chunked_prefill_single_shot_matches_non_chunked(self):

336 - pooler = self._make_all_pool(chunked=True)

341 + @pytest.mark.parametrize("async_scheduling", [False, True])

342 + def test_chunked_prefill_single_shot_matches_non_chunked(

343 + self, async_scheduling: bool

344 + ):

345 + pooler = self._make_all_pool(chunked=True, async_scheduling=async_scheduling)

346 hidden = torch.tensor(

347 [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]

348 )

349 + expected = [hidden[:2].clone(), hidden[2:].clone()]

350 metadata = _make_metadata([2, 3])

351 out = pooler(hidden, metadata)

352 + if async_scheduling:

353 + # Finished outputs must own their storage under async scheduling.

354 + hidden.zero_()

355 assert len(out) == 2

343 - assert torch.equal(out[0], hidden[:2])

344 - assert torch.equal(out[1], hidden[2:])

356 + assert torch.equal(out[0], expected[0])

357 + assert torch.equal(out[1], expected[1])

358 +

359 + def test_non_chunked_owns_output_under_async_scheduling(self):

360 + pooler = self._make_all_pool(async_scheduling=True)

361 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

362 + expected = [hidden[:2].clone(), hidden[2:].clone()]

363 + metadata = _make_metadata([2, 1])

364 + out = pooler(hidden, metadata)

365 + hidden.zero_()

366 + assert len(out) == 2

367 + assert torch.equal(out[0], expected[0])

368 + assert torch.equal(out[1], expected[1])

369

370 def test_chunked_prefill_mixed_finished_unfinished(self):

371 pooler = self._make_all_pool(chunked=True)

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py

+5 -3

@@ -15,9 +15,10 @@ from vllm.config import PoolerConfig

15 ["Qwen/Qwen3-Embedding-0.6B"],

16 )

17 @torch.inference_mode

18 -def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):

19 - # Keep token_embed on MRV1 when sequence pooling becomes MRV2 by default.

20 - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")

18 +def test_decoder_token_embed_model_runner_v2(

19 + hf_runner, vllm_runner, monkeypatch, model: str

20 +):

21 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")

22 chunk_size = 10

23 n_prompt_tokens = [55, 56, 57]

24 token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]

@@ -33,6 +34,7 @@ def test_embed_models(hf_runner, vllm_runner, monkeypatch, model: str):

34 enable_chunked_prefill=True,

35 enable_prefix_caching=True,

36 ) as vllm_model:

37 + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner

38 vllm_outputs = vllm_model.token_embed(

39 [TokensPrompt(prompt_token_ids=t) for t in token_prompts],

40 )

tests/models/language/pooling/test_jina_reranker_v3.py

+3 -1

@@ -45,8 +45,10 @@ INSTRUCTION = (

45 )

48 -def test_offline(vllm_runner):

48 +def test_offline(vllm_runner, monkeypatch):

49 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")

50 with vllm_runner(model_name, runner="pooling") as llm_runner:

51 + assert llm_runner.llm.llm_engine.vllm_config.use_v2_model_runner

52 llm = llm_runner.get_llm()

53 _test_offline_1_v_1(llm)

54 _test_offline_1_v_n(llm)

tests/models/language/pooling/test_reward.py

+10

@@ -98,6 +98,7 @@ def test_prm_models(

98 hf_runner,

99 vllm_runner,

100 math_step_prompts,

101 + monkeypatch,

102 model: str,

103 dtype: str,

104 ) -> None:

@@ -107,12 +108,14 @@ def test_prm_models(

108 if current_platform.is_cpu():

109 pytest.skip("CPU only supports V1")

110

111 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")

112 with vllm_runner(

113 model,

114 max_model_len=1024,

115 dtype=dtype,

116 kv_cache_memory_bytes=64 * MiB_bytes,

117 ) as vllm_model:

118 + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner

119 vllm_outputs = vllm_model.token_classify(math_step_prompts)

120

121 with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:

@@ -145,18 +148,25 @@ def test_prm_models(

148 def test_prm_models_with_golden_outputs(

149 vllm_runner,

150 math_step_prompts,

151 + monkeypatch,

152 model: str,

153 dtype: str,

154 ) -> None:

155 if not FIXTURE_REWARD_RESULT.get(model):

156 pytest.skip(f"No available golden outputs for {model}.")

157

158 + use_v2_model_runner = not current_platform.is_cpu()

159 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2_model_runner else "0")

160 with vllm_runner(

161 model,

162 max_model_len=1024,

163 dtype=dtype,

164 kv_cache_memory_bytes=64 * MiB_bytes,

165 ) as vllm_model:

166 + assert (

167 + vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner

168 + is use_v2_model_runner

169 + )

170 vllm_outputs = vllm_model.token_classify(math_step_prompts)

171

172 golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])

tests/models/language/pooling/test_splade_sparse_pooler.py

+8 -15

@@ -231,7 +231,7 @@ def test_pooling_runner_supports_encoder_token_embedding() -> None:

231 assert runner.supported_tasks == {"token_embed"}

232

233

234 -def test_pooling_runner_rejects_decoder_token_classification() -> None:

234 +def test_pooling_runner_supports_decoder_token_classification() -> None:

235 model = MagicMock()

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

237 vllm_config = MagicMock()

@@ -239,13 +239,12 @@ def test_pooling_runner_rejects_decoder_token_classification() -> None:

239 vllm_config.model_config.attn_type = "decoder"

240 vllm_config.model_config.get_pooling_task.return_value = "token_classify"

241

242 - with pytest.raises(ValueError, match="selects 'token_classify'") as exc_info:

243 - PoolingRunner(model, vllm_config)

242 + runner = PoolingRunner(model, vllm_config)

243

245 - assert "Set an explicitly supported task" not in str(exc_info.value)

244 + assert runner.supported_tasks == {"token_classify"}

245

246

248 -def test_pooling_runner_filters_decoder_token_classification() -> None:

247 +def test_pooling_runner_keeps_decoder_token_classification() -> None:

248 model = MagicMock()

249 model.pooler.get_supported_tasks.return_value = {"embed", "token_classify"}

250 vllm_config = MagicMock()

@@ -255,10 +254,10 @@ def test_pooling_runner_filters_decoder_token_classification() -> None:

254

255 runner = PoolingRunner(model, vllm_config)

256

258 - assert runner.supported_tasks == {"embed"}

257 + assert runner.supported_tasks == {"embed", "token_classify"}

258

259

261 -def test_pooling_runner_filters_decoder_token_embedding() -> None:

260 +def test_pooling_runner_keeps_decoder_token_embedding() -> None:

261 model = MagicMock()

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

263 vllm_config = MagicMock()

@@ -268,12 +267,10 @@ def test_pooling_runner_filters_decoder_token_embedding() -> None:

267

268 runner = PoolingRunner(model, vllm_config)

269

271 - assert runner.supported_tasks == {"embed"}

270 + assert runner.supported_tasks == {"embed", "token_embed"}

271

272

274 -def test_pooling_runner_gates_embed_token_classification_to_encoder() -> None:

275 - # BGE-M3's combined task emits per-token weights alongside the embedding,

276 - # so it is enabled only where prefill is unchunked.

273 +def test_pooling_runner_supports_embed_token_classification() -> None:

274 model = MagicMock()

275 model.pooler.get_supported_tasks.return_value = {"embed", "embed&token_classify"}

276 vllm_config = MagicMock()

@@ -285,7 +282,3 @@ def test_pooling_runner_gates_embed_token_classification_to_encoder() -> None:

282 "embed",

283 "embed&token_classify",

284 }

288 -

289 - vllm_config.model_config.attn_type = "decoder"

290 - vllm_config.model_config.get_pooling_task.return_value = "embed"

291 - assert PoolingRunner(model, vllm_config).supported_tasks == {"embed"}

vllm/model_executor/layers/pooler/tokwise/methods.py

+21 -4

@@ -40,9 +40,15 @@ class AllPool(TokenPoolingMethod):

40 scheduler_config = vllm_config.scheduler_config

42 self.enable_chunked_prefill = scheduler_config.enable_chunked_prefill

43 + # With async scheduling, the next step may overwrite the hidden states

44 + # buffer before outputs are copied to CPU, so they must own their storage.

45 + self.clone_finished = bool(scheduler_config.async_scheduling)

47 def extra_repr(self) -> str:

45 - return f"enable_chunked_prefill={self.enable_chunked_prefill}"

48 + return (

49 + f"enable_chunked_prefill={self.enable_chunked_prefill}, "

50 + f"clone_finished={self.clone_finished}"

51 + )

53 def forward(

54 self,

@@ -58,18 +64,29 @@ class AllPool(TokenPoolingMethod):

64 )

66 if not self.enable_chunked_prefill:

67 + if self.clone_finished:

68 + return [hs.clone() for hs in hidden_states_lst]

69 return hidden_states_lst

71 pooling_states = pooling_metadata.pooling_states

72 + finished_mask = pooling_cursor.is_finished().tolist()

74 # If chunked_prefill is enabled

75 # 1. first store the chunked hidden_states in pooling_states.hidden_states_cache

67 - for p, hs_chunk in zip(pooling_states, hidden_states_lst):

68 - p.hidden_states_cache.append(hs_chunk)

76 + for p, hs_chunk, finished in zip(

77 + pooling_states, hidden_states_lst, finished_mask

78 + ):

79 + # Own data retained across steps or returned without concatenation.

80 + needs_owned_storage = not finished or (

81 + self.clone_finished and not p.hidden_states_cache

82 + )

83 + p.hidden_states_cache.append(

84 + hs_chunk.clone() if needs_owned_storage else hs_chunk

85 + )

87 # 2. Once prefill is finished, send hidden_states_cache to PoolerHead

88 output_list = list[TokenPoolingMethodOutputItem]()

72 - for p, finished in zip(pooling_states, pooling_cursor.is_finished()):

89 + for p, finished in zip(pooling_states, finished_mask):

90 if finished:

91 hidden_states_cache = p.hidden_states_cache

92 if len(hidden_states_cache) == 1:

vllm/v1/worker/gpu/model_runner.py

+1 -3

@@ -313,9 +313,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

313 # Do not rely on pooling_runner here, since this information is needed

314 # on the first PP rank, while pooling_runner is only initialized

315 # on the last PP rank.

316 - tasks.extend(

317 - PoolingRunner.get_supported_tasks(self.model, self.model_config)

318 - )

316 + tasks.extend(PoolingRunner.get_supported_tasks(self.model))

317 return tuple(tasks)

318

319 def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None:

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

+8 -26

@@ -6,7 +6,7 @@ import numpy as np

6 import torch

7 import torch.nn as nn

9 -from vllm.config import ModelConfig, VllmConfig

9 +from vllm.config import VllmConfig

10 from vllm.model_executor.models import VllmModelForPooling, is_pooling_model

11 from vllm.pooling_params import PoolingParams

12 from vllm.tasks import PoolingTask

@@ -20,11 +20,6 @@ from vllm.v1.worker.gpu.states import RequestState

20 _SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(

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

22 )

23 -# `embed&token_classify` concatenates a fixed-size embedding with per-token

24 -# weights, so its output length scales with the prompt like the others here.

25 -_TOKEN_TASKS: frozenset[PoolingTask] = frozenset(

26 - {"token_embed", "token_classify", "embed&token_classify"}

27 -)

25 class PoolingRunner:

@@ -33,26 +28,23 @@ class PoolingRunner:

28 self.model_config = vllm_config.model_config

29 self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs

30 model_tasks = tuple(sorted(self.model.pooler.get_supported_tasks()))

36 - enabled_tasks = self._get_enabled_tasks(self.model_config)

31 selected_task = self.model_config.get_pooling_task(model_tasks)

38 - if selected_task not in enabled_tasks:

32 + if selected_task not in _SUPPORTED_TASKS:

33 hint = (

34 "Set an explicitly supported task or VLLM_USE_V2_MODEL_RUNNER=0."

41 - if enabled_tasks.intersection(model_tasks)

35 + if _SUPPORTED_TASKS.intersection(model_tasks)

36 else "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."

37 )

38 raise ValueError(

39 "Model Runner V2 supports pooling tasks "

46 - f"{sorted(enabled_tasks)}, but this model selects "

40 + f"{sorted(_SUPPORTED_TASKS)}, but this model selects "

41 f"{selected_task!r} from {list(model_tasks)}. {hint}"

42 )

49 - self.supported_tasks = frozenset(

50 - self.get_supported_tasks(model, self.model_config)

51 - )

43 + self.supported_tasks = frozenset(self.get_supported_tasks(model))

44 if not self.supported_tasks:

45 raise ValueError(

46 "Model Runner V2 supports pooling tasks "

55 - f"{sorted(enabled_tasks)}, but this model supports "

47 + f"{sorted(_SUPPORTED_TASKS)}, but this model supports "

48 f"{list(model_tasks)}. "

49 "Set VLLM_USE_V2_MODEL_RUNNER=0 to use this model."

50 )

@@ -62,20 +54,10 @@ class PoolingRunner:

54 self.late_interaction_runner = LateInteractionRunner()

56 @staticmethod

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

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

67 - if model_config.attn_type == "encoder_only":

68 - return _SUPPORTED_TASKS

69 - return _SUPPORTED_TASKS - _TOKEN_TASKS

70 -

71 - @staticmethod

72 - def get_supported_tasks(

73 - model: nn.Module, model_config: ModelConfig

74 - ) -> list[PoolingTask]:

57 + def get_supported_tasks(model: nn.Module) -> list[PoolingTask]:

58 if not is_pooling_model(model):

59 return []

77 - enabled_tasks = PoolingRunner._get_enabled_tasks(model_config)

78 - return sorted(model.pooler.get_supported_tasks() & enabled_tasks)

60 + return sorted(model.pooler.get_supported_tasks() & _SUPPORTED_TASKS)

62 def add_request(

63 self,