@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f8b4ce0c0

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 6, 20:30 UTC

[Model Runner V2] Do not initialize sampler for non-last PP ranks (#36824)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Mar 11, 2026 at 20:55 UTC 2f8b4ce0c0ece2dfc3a5c61c51631606ee879ec0

3 files changed +75 -50

vllm/v1/worker/gpu/input_batch.py

+11 -8

@@ -438,17 +438,20 @@ def _post_update_kernel(

438

439 for i in range(num_sampled):

440 token_id = tl.load(sampled_tokens_ptr + req_id * sampled_tokens_stride + i)

441 - token_ptr = (

442 - output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + token_id

443 - )

444 - count = tl.load(token_ptr)

445 - count += 1

446 - tl.store(token_ptr, count)

441 tl.store(

442 all_token_ids_ptr + req_state_idx * all_token_ids_stride + total_len + i,

443 token_id,

444 )

445

446 + if output_bin_counts_ptr is not None:

447 + token_ptr = (

448 + output_bin_counts_ptr

449 + + req_state_idx * output_bin_counts_stride

450 + + token_id

451 + )

452 + count = tl.load(token_ptr)

453 + tl.store(token_ptr, count + 1)

454 +

455 query_start = tl.load(query_start_loc_ptr + req_id)

456 query_end = tl.load(query_start_loc_ptr + req_id + 1)

457 query_len = query_end - query_start

@@ -467,7 +470,7 @@ def post_update(

470 # [max_num_reqs]

471 last_sampled_tokens: torch.Tensor,

472 # [max_num_reqs, vocab_size]

470 - output_bin_counts: torch.Tensor,

473 + output_bin_counts: torch.Tensor | None,

474 # [num_reqs, num_speculative_steps + 1]

475 sampled_tokens: torch.Tensor,

476 # [num_reqs]

@@ -487,7 +490,7 @@ def post_update(

490 num_computed_tokens,

491 last_sampled_tokens,

492 output_bin_counts,

490 - output_bin_counts.stride(0),

493 + output_bin_counts.stride(0) if output_bin_counts is not None else 0,

494 sampled_tokens,

495 sampled_tokens.stride(0),

496 num_sampled,

vllm/v1/worker/gpu/model_runner.py

+60 -39

@@ -183,6 +183,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

183 # Draft tokens propagation - for spec-dec + struct outputs.

184 self.draft_tokens_handler = DraftTokensHandler(self.device)

185

186 + # Pooling models.

187 + self.is_pooling_model = self.model_config.runner_type == "pooling"

188 + self.pooling_runner: PoolingRunner | None = None

189 +

190 # General request states.

191 self.req_states = RequestState(

192 max_num_reqs=self.max_num_reqs,

@@ -199,20 +203,34 @@ class GPUModelRunner(LoRAModelRunnerMixin):

203 max_num_tokens=self.max_num_tokens,

204 device=self.device,

205 )

202 - self.sampler = Sampler(

203 - max_num_reqs=self.max_num_reqs,

204 - vocab_size=self.vocab_size,

205 - device=self.device,

206 - req_states=self.req_states,

207 - logprobs_mode=self.model_config.logprobs_mode,

208 - num_speculative_tokens=self.num_speculative_steps + 1,

209 - )

210 - self.rejection_sampler = RejectionSampler(

211 - self.sampler,

212 - num_speculative_steps=self.num_speculative_steps,

213 - use_strict_rejection_sampling=use_strict_rejection_sampling,

214 - )

215 - self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)

206 +

207 + self.sampler: Sampler | None = None

208 + self.rejection_sampler: RejectionSampler | None = None

209 + self.prompt_logprobs_worker: PromptLogprobsWorker | None = None

210 + self.structured_outputs_worker: StructuredOutputsWorker | None = None

211 + if self.is_last_pp_rank and not self.is_pooling_model:

212 + # Initialize sampling-related workers.

213 + # These components are only set up on the last PP rank and

214 + # for generative (non-pooling) models.

215 + self.sampler = Sampler(

216 + max_num_reqs=self.max_num_reqs,

217 + vocab_size=self.vocab_size,

218 + device=self.device,

219 + req_states=self.req_states,

220 + logprobs_mode=self.model_config.logprobs_mode,

221 + num_speculative_tokens=self.num_speculative_steps + 1,

222 + )

223 + self.rejection_sampler = RejectionSampler(

224 + self.sampler,

225 + num_speculative_steps=self.num_speculative_steps,

226 + use_strict_rejection_sampling=use_strict_rejection_sampling,

227 + )

228 + self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)

229 + self.structured_outputs_worker = StructuredOutputsWorker(

230 + max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),

231 + vocab_size=self.vocab_size,

232 + device=self.device,

233 + )

234

235 # CUDA graphs.

236 self.decode_query_len = self.num_speculative_steps + 1

@@ -222,21 +240,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

240 self.compilation_config.cudagraph_mode,

241 decode_query_len=self.decode_query_len,

242 )

225 - # Structured outputs worker.

226 - self.structured_outputs_worker = StructuredOutputsWorker(

227 - max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),

228 - vocab_size=self.vocab_size,

229 - device=self.device,

230 - )

243 # LoRA-related workers.

244 self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)

245 # KV Connector if configured.

246 self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

247

236 - # Pooling models.

237 - self.is_pooling_model = self.model_config.runner_type == "pooling"

238 - self.pooling_runner: PoolingRunner | None = None

239 -

248 # For transferring state from execute_model to subsequent sample_tokens call.

249 self.execute_model_state: ExecuteModelState | None = None

250

@@ -248,8 +256,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

256 tasks: list[SupportedTask] = []

257 if self.model_config.runner_type == "generate":

258 tasks.extend(self.model_state.get_supported_generation_tasks())

251 - if self.pooling_runner is not None:

252 - tasks.extend(self.pooling_runner.get_supported_pooling_tasks())

259 + if self.is_pooling_model:

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

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

262 + # on the last PP rank.

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

264 return tuple(tasks)

265

266 def load_model(self, *args, **kwargs) -> None:

@@ -289,7 +300,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

300 self.model_state = init_model_state(

301 self.vllm_config, self.model, self.encoder_cache, self.device

302 )

292 - if self.is_pooling_model:

303 + if self.is_pooling_model and self.is_last_pp_rank:

304 self.pooling_runner = PoolingRunner(self.model)

305

306 def get_model(self) -> nn.Module:

@@ -420,6 +431,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

431

432 # dummy run the eagle speculator's propose to ensure DP/EP sync.

433 if self.speculator is not None:

434 + assert self.sampler is not None

435 self.speculator.propose(

436 input_batch=input_batch,

437 attn_metadata=attn_metadata,

@@ -457,10 +469,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

469 # NOTE(woosuk): During the initial memory profiling, the sampler may skip

470 # top_k, top_p, and logprobs, using less GPU memory than what is possible

471 # during actual execution.

460 - self.sampler(

461 - logits,

462 - dummy_input_batch,

463 - )

472 + assert self.sampler is not None

473 + self.sampler(logits, dummy_input_batch)

474

475 @torch.inference_mode()

476 def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:

@@ -558,7 +568,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

568 self.req_states.remove_request(req_id)

569 if self.encoder_cache is not None:

570 self.encoder_cache.remove_request(req_id)

561 - self.prompt_logprobs_worker.remove_request(req_id)

571 + if self.prompt_logprobs_worker is not None:

572 + self.prompt_logprobs_worker.remove_request(req_id)

573 self.lora_state.remove_request(req_id)

574

575 def free_states(self, scheduler_output: SchedulerOutput) -> None:

@@ -589,18 +600,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):

600 )

601 self.lora_state.add_request(req_id, req_index, new_req_data.lora_request)

602

592 - if new_req_data.sampling_params is not None:

603 + if self.is_last_pp_rank and new_req_data.sampling_params is not None:

604 + assert self.sampler is not None

605 self.sampler.add_request(

606 req_index, prompt_len, new_req_data.sampling_params

607 )

608 + assert self.prompt_logprobs_worker is not None

609 self.prompt_logprobs_worker.add_request(

610 req_id, req_index, new_req_data.sampling_params

611 )

612

613 if scheduler_output.scheduled_new_reqs:

614 self.req_states.apply_staged_writes()

602 - self.sampler.apply_staged_writes()

615 self.model_state.apply_staged_writes()

616 + if self.sampler is not None:

617 + self.sampler.apply_staged_writes()

618

619 def update_requests(self, scheduler_output: SchedulerOutput) -> None:

620 # Add new blocks for the existing requests.

@@ -788,6 +802,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

802 logits = self.model.compute_logits(sample_hidden_states)

803 if grammar_output is not None:

804 # Apply grammar bitmask to the logits in-place.

805 + assert self.structured_outputs_worker is not None

806 self.structured_outputs_worker.apply_grammar_bitmask(

807 logits,

808 input_batch,

@@ -797,12 +812,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

812

813 if input_batch.num_draft_tokens == 0:

814 # No draft tokens (common case).

800 - sampler_output = self.sampler(

801 - logits,

802 - input_batch,

803 - )

815 + assert self.sampler is not None

816 + sampler_output = self.sampler(logits, input_batch)

817 else:

818 # Rejection sampling for spec decoding.

819 + assert self.rejection_sampler is not None

820 sampler_output = self.rejection_sampler(

821 logits,

822 input_batch,

@@ -831,11 +845,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

845 num_rejected: torch.Tensor,

846 ) -> None:

847 # Update the number of computed tokens.

848 + if self.is_last_pp_rank:

849 + assert self.sampler is not None

850 + output_bin_counts = self.sampler.penalties_state.output_bin_counts

851 + else:

852 + output_bin_counts = None

853 post_update(

854 input_batch.idx_mapping,

855 self.req_states.num_computed_tokens.gpu,

856 self.req_states.last_sampled_tokens,

838 - self.sampler.penalties_state.output_bin_counts,

857 + output_bin_counts,

858 sampled_tokens,

859 num_sampled,

860 num_rejected,

@@ -1076,6 +1095,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1095 # Broadcast to non-last PP ranks (handles spec decode multi-token).

1096 pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)

1097

1098 + assert self.prompt_logprobs_worker is not None

1099 prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(

1100 self.model.compute_logits,

1101 hidden_states,

@@ -1115,6 +1135,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1135 input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected

1136 )

1137 if self.speculator is not None:

1138 + assert self.sampler is not None

1139 draft_tokens = self.speculator.propose(

1140 input_batch,

1141 attn_metadata,

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

+4 -3

@@ -19,10 +19,11 @@ class PoolingRunner:

19 def __init__(self, model: nn.Module):

20 self.model = cast(VllmModelForPooling, model)

22 - def get_supported_pooling_tasks(self) -> list[PoolingTask]:

23 - if not is_pooling_model(self.model):

22 + @staticmethod

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

24 + if not is_pooling_model(model):

25 return []

25 - assert "embed" in self.model.pooler.get_supported_tasks()

26 + assert "embed" in model.pooler.get_supported_tasks()

27 return ["embed"]

29 def pool(