@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7df331c66b

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 12, 16:00 UTC

[BugFix] Fix chunked prompt logprobs + preemption (#29071)

Nick Hill committed Nov 22, 2025 at 13:07 UTC 7df331c66b242a3109a95fb434a1badd196966d7

6 files changed +127 -31

tests/conftest.py

+23 -4

@@ -853,6 +853,7 @@ class VllmRunner:

853 @staticmethod

854 def _final_steps_generate_w_logprobs(

855 req_outputs: list[RequestOutput],

856 + include_prompt_token_ids: bool = False,

857 ) -> list[TokensTextLogprobsPromptLogprobs]:

858 outputs: list[TokensTextLogprobsPromptLogprobs] = []

859 for req_output in req_outputs:

@@ -861,9 +862,26 @@ class VllmRunner:

862 output_str = sample.text

863 output_ids = list(sample.token_ids)

864 output_logprobs = sample.logprobs

864 - outputs.append(

865 - (output_ids, output_str, output_logprobs, req_output.prompt_logprobs)

866 - )

865 + if include_prompt_token_ids:

866 + outputs.append(

867 + ( # type: ignore[arg-type]

868 + output_ids,

869 + output_str,

870 + output_logprobs,

871 + req_output.prompt_token_ids,

872 + req_output.prompt_logprobs,

873 + )

874 + )

875 + else:

876 + outputs.append(

877 + (

878 + output_ids,

879 + output_str,

880 + output_logprobs,

881 + req_output.prompt_logprobs,

882 + )

883 + )

884 +

885 return outputs

886

887 def generate_w_logprobs(

@@ -873,6 +891,7 @@ class VllmRunner:

891 images: PromptImageInput | None = None,

892 audios: PromptAudioInput | None = None,

893 videos: PromptVideoInput | None = None,

894 + include_prompt_token_ids: bool = False,

895 **kwargs: Any,

896 ) -> list[TokensTextLogprobs] | list[TokensTextLogprobsPromptLogprobs]:

897 inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)

@@ -882,7 +901,7 @@ class VllmRunner:

901 )

902

903 toks_str_logsprobs_prompt_logprobs = self._final_steps_generate_w_logprobs(

885 - req_outputs

904 + req_outputs, include_prompt_token_ids

905 )

906 # Omit prompt logprobs if not required by sampling params

907 return (

tests/v1/sample/test_logprobs.py

+76

@@ -605,3 +605,79 @@ def test_spec_decode_logprobs(

605 )

606 assert ref_logprob.rank == spec_logprob.rank

607 assert ref_logprob.decoded_token == spec_logprob.decoded_token

608 +

609 +

610 +def test_prompt_logprobs_with_chunking_and_preemption():

611 + """Test that prompt logprobs are correctly returned when using

612 + both chunked prefill and preemption.

613 +

614 + This test ensures that the num_prompt_logprobs tracking persists

615 + across preemptions and prefill chunks.

616 + """

617 +

618 + # Create prompts that will trigger chunking and preemption

619 + prompts = [

620 + "The following numbers of the sequence "

621 + + ", ".join(str(i) for i in range(10))

622 + + " are:",

623 + "In one word, the capital of France is ",

624 + ] + [f"Tell me about the number {i}: " for i in range(32)]

625 +

626 + sampling_params = SamplingParams(

627 + temperature=0.0,

628 + max_tokens=40,

629 + min_tokens=20,

630 + prompt_logprobs=2, # Request prompt logprobs

631 + )

632 +

633 + with VllmRunner(

634 + "Qwen/Qwen3-0.6B",

635 + max_model_len=512,

636 + enable_chunked_prefill=True,

637 + max_num_batched_tokens=48, # Force prefill chunking

638 + num_gpu_blocks_override=32, # Force preemptions

639 + disable_log_stats=False,

640 + gpu_memory_utilization=0.25,

641 + ) as vllm_model:

642 + metrics_before = vllm_model.llm.get_metrics()

643 +

644 + # Generate with prompt logprobs using generate_w_logprobs which

645 + # returns (output_ids, output_str, output_logprobs, prompt_logprobs)

646 + outputs = vllm_model.generate_w_logprobs(

647 + prompts, sampling_params=sampling_params, include_prompt_token_ids=True

648 + )

649 +

650 + # Verify that all outputs have prompt logprobs

651 + for i, output in enumerate(outputs):

652 + _, _, _, prompt_token_ids, prompt_logprobs = output

653 + assert prompt_logprobs is not None and len(prompt_logprobs) > 0, (

654 + f"Output {i} missing prompt logprobs"

655 + )

656 + assert len(prompt_logprobs) == len(prompt_token_ids), (

657 + "Unexpected number of prompt logprob positions"

658 + )

659 +

660 + # Each position should have the requested number of logprobs

661 + for pos, logprobs_dict in enumerate(prompt_logprobs):

662 + if logprobs_dict is not None: # First token may be None

663 + assert (

664 + sampling_params.prompt_logprobs

665 + <= len(logprobs_dict)

666 + <= sampling_params.prompt_logprobs + 1

667 + ), (

668 + f"Output {i} position {pos} has {len(logprobs_dict)} "

669 + f"logprobs, expected {sampling_params.prompt_logprobs}"

670 + )

671 +

672 + # Check that we actually had preemptions

673 + metrics_after = vllm_model.llm.get_metrics()

674 + preemptions_before = next(

675 + (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0

676 + )

677 + preemptions_after = next(

678 + (m.value for m in metrics_after if m.name == "vllm:num_preemptions"), 0

679 + )

680 + preemptions = preemptions_after - preemptions_before

681 + assert preemptions > 0, "Test did not trigger any preemptions"

682 +

683 + print(f"Test passed with {preemptions} preemptions")

vllm/v1/worker/gpu_input_batch.py

-14

@@ -219,9 +219,6 @@ class InputBatch:

219 self.generators: dict[int, torch.Generator] = {}

220

221 self.num_logprobs: dict[str, int] = {}

222 - # NOTE(rob): num_prompt_logprobs only includes reqs

223 - # that are currently in the prefill phase.

224 - self.num_prompt_logprobs: dict[str, int] = {}

222

223 # To accumulate prompt logprobs tensor chunks across prefill steps.

224 self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

@@ -385,12 +382,6 @@ class InputBatch:

382 if sampling_params.logprobs == -1

383 else sampling_params.logprobs

384 )

388 - if sampling_params.prompt_logprobs is not None:

389 - self.num_prompt_logprobs[req_id] = (

390 - self.vocab_size

391 - if sampling_params.prompt_logprobs == -1

392 - else sampling_params.prompt_logprobs

393 - )

385

386 if sampling_params.allowed_token_ids:

387 self.has_allowed_token_ids.add(req_id)

@@ -488,7 +479,6 @@ class InputBatch:

479 self.repetition_penalties_reqs.discard(req_id)

480 self.generators.pop(req_index, None)

481 self.num_logprobs.pop(req_id, None)

491 - self.num_prompt_logprobs.pop(req_id, None)

482 self.in_progress_prompt_logprobs_cpu.pop(req_id, None)

483

484 self.has_allowed_token_ids.discard(req_id)

@@ -972,10 +962,6 @@ class InputBatch:

962 def max_num_logprobs(self) -> int | None:

963 return max(self.num_logprobs.values()) if self.num_logprobs else None

964

975 - @property

976 - def no_prompt_logprob(self) -> bool:

977 - return not self.num_prompt_logprobs

978 -

965 @property

966 def no_allowed_token_ids(self) -> bool:

967 return len(self.has_allowed_token_ids) == 0

vllm/v1/worker/gpu_model_runner.py

+17 -3

@@ -393,6 +393,9 @@ class GPUModelRunner(

393

394 # Request states.

395 self.requests: dict[str, CachedRequestState] = {}

396 + # NOTE(rob): num_prompt_logprobs only includes reqs

397 + # that are currently in the prefill phase.

398 + self.num_prompt_logprobs: dict[str, int] = {}

399 self.comm_stream = torch.cuda.Stream()

400

401 # Input Batch

@@ -687,6 +690,7 @@ class GPUModelRunner(

690 # Remove finished requests from the cached states.

691 for req_id in scheduler_output.finished_req_ids:

692 self.requests.pop(req_id, None)

693 + self.num_prompt_logprobs.pop(req_id, None)

694 # Remove the finished requests from the persistent batch.

695 # NOTE(woosuk): There could be an edge case where finished_req_ids and

696 # scheduled_req_ids overlap. This happens when a request is aborted and

@@ -755,6 +759,13 @@ class GPUModelRunner(

759 )

760 self.requests[req_id] = req_state

761

762 + if sampling_params and sampling_params.prompt_logprobs is not None:

763 + self.num_prompt_logprobs[req_id] = (

764 + self.input_batch.vocab_size

765 + if sampling_params.prompt_logprobs == -1

766 + else sampling_params.prompt_logprobs

767 + )

768 +

769 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

770 if self.uses_mrope:

771 self._init_mrope_positions(req_state)

@@ -2671,7 +2682,7 @@ class GPUModelRunner(

2682 scheduler_output, self.vllm_config

2683 )

2684 if self.cache_config.kv_sharing_fast_prefill:

2674 - assert not self.input_batch.num_prompt_logprobs, (

2685 + assert not self.num_prompt_logprobs, (

2686 "--kv-sharing-fast-prefill produces incorrect "

2687 "logprobs for prompt tokens, tokens, please disable "

2688 "it when the requests need prompt logprobs"

@@ -3436,7 +3447,7 @@ class GPUModelRunner(

3447 hidden_states: torch.Tensor,

3448 num_scheduled_tokens: dict[str, int],

3449 ) -> dict[str, LogprobsTensors | None]:

3439 - num_prompt_logprobs_dict = self.input_batch.num_prompt_logprobs

3450 + num_prompt_logprobs_dict = self.num_prompt_logprobs

3451 if not num_prompt_logprobs_dict:

3452 return {}

3453

@@ -3447,7 +3458,10 @@ class GPUModelRunner(

3458 # maintainable loop over optimal performance.

3459 completed_prefill_reqs = []

3460 for req_id, num_prompt_logprobs in num_prompt_logprobs_dict.items():

3450 - num_tokens = num_scheduled_tokens[req_id]

3461 + num_tokens = num_scheduled_tokens.get(req_id)

3462 + if num_tokens is None:

3463 + # This can happen if the request was preempted in prefill stage.

3464 + continue

3465

3466 # Get metadata for this request.

3467 request = self.requests[req_id]

vllm/v1/worker/tpu_input_batch.py

-10

@@ -149,9 +149,6 @@ class InputBatch:

149 self.generators: dict[int, torch.Generator] = {}

150

151 self.num_logprobs: dict[str, int] = {}

152 - # NOTE(rob): num_prompt_logprobs only includes reqs

153 - # that are currently in the prefill phase.

154 - self.num_prompt_logprobs: dict[str, int] = {}

152

153 # To accumulate prompt logprobs tensor chunks across prefill steps.

154 self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

@@ -256,8 +253,6 @@ class InputBatch:

253

254 if sampling_params.logprobs is not None:

255 self.num_logprobs[req_id] = sampling_params.logprobs

259 - if sampling_params.prompt_logprobs is not None:

260 - self.num_prompt_logprobs[req_id] = sampling_params.prompt_logprobs

256 if sampling_params.logit_bias is not None:

257 self.logit_bias[req_index] = sampling_params.logit_bias

258

@@ -317,7 +312,6 @@ class InputBatch:

312 self.repetition_penalties_reqs.discard(req_id)

313 self.generators.pop(req_index, None)

314 self.num_logprobs.pop(req_id, None)

320 - self.num_prompt_logprobs.pop(req_id, None)

315 self.in_progress_prompt_logprobs_cpu.pop(req_id, None)

316

317 # LoRA

@@ -584,10 +578,6 @@ class InputBatch:

578 def max_num_logprobs(self) -> int | None:

579 return max(self.num_logprobs.values()) if self.num_logprobs else None

580

587 - @property

588 - def no_prompt_logprob(self) -> bool:

589 - return not self.num_prompt_logprobs

590 -

581 @property

582 def no_allowed_token_ids(self) -> bool:

583 return len(self.has_allowed_token_ids) == 0

vllm/v1/worker/tpu_model_runner.py

+11

@@ -247,6 +247,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

247

248 # Request states.

249 self.requests: dict[str, CachedRequestState] = {}

250 + # NOTE(rob): num_prompt_logprobs only includes reqs

251 + # that are currently in the prefill phase.

252 + self.num_prompt_logprobs: dict[str, int] = {}

253

254 # Initialize input batch early to avoid AttributeError in _update_states

255 self.input_batch = InputBatch(

@@ -420,6 +423,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

423 # Remove finished requests from the cached states.

424 for req_id in scheduler_output.finished_req_ids:

425 self.requests.pop(req_id, None)

426 + self.num_prompt_logprobs.pop(req_id, None)

427

428 # Remove the finished requests from the persistent batch.

429 # NOTE(woosuk): There could be an edge case where finished_req_ids and

@@ -477,6 +481,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

481 lora_request=new_req_data.lora_request,

482 )

483

484 + if sampling_params and sampling_params.prompt_logprobs is not None:

485 + self.num_prompt_logprobs[req_id] = (

486 + self.input_batch.vocab_size

487 + if sampling_params.prompt_logprobs == -1

488 + else sampling_params.prompt_logprobs

489 + )

490 +

491 req_ids_to_add.append(req_id)

492

493 # Update the states of the running/resumed requests.