@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 408195ec59
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 9, 20:45 UTC

[Model Runner V2] Refactor Prompt Logprobs (#32811)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Jan 21, 2026 at 15:12 UTC 408195ec59e5631c834c8ce5e8ecf5f8f1c7b47a
4 files changed +230 -142
vllm/v1/worker/gpu/model_runner.py
+17 -102
@@ -22,7 +22,6 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22 from vllm.v1.kv_cache_interface import KVCacheConfig
23 from vllm.v1.outputs import (
24 EMPTY_MODEL_RUNNER_OUTPUT,
25 - LogprobsTensors,
25 ModelRunnerOutput,
26 )
27 from vllm.v1.worker.gpu.async_utils import AsyncOutput
@@ -51,8 +50,8 @@ from vllm.v1.worker.gpu.input_batch import (
50 )
51 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
52 from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
54 -from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
53 from vllm.v1.worker.gpu.sample.output import SamplerOutput
54 +from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
55 from vllm.v1.worker.gpu.sample.sampler import Sampler
56 from vllm.v1.worker.gpu.spec_decode import init_speculator
57 from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
@@ -156,6 +155,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
155 device=self.device,
156 logprobs_mode=self.model_config.logprobs_mode,
157 )
158 + self.prompt_logprobs_worker = PromptLogprobsWorker(self.max_num_reqs)
159
160 # CUDA graphs.
161 self.cudagraph_manager = CudaGraphManager(
@@ -416,10 +416,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
416 self.req_states.remove_request(req_id)
417 if self.supports_mm_inputs:
418 self.encoder_runner.remove_request(req_id)
419 + self.prompt_logprobs_worker.remove_request(req_id)
420 for req_id in scheduler_output.finished_req_ids:
421 self.req_states.remove_request(req_id)
422 if self.supports_mm_inputs:
423 self.encoder_runner.remove_request(req_id)
424 + self.prompt_logprobs_worker.remove_request(req_id)
425
426 def free_states(self, scheduler_output: SchedulerOutput) -> None:
427 if self.supports_mm_inputs:
@@ -438,7 +440,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
440 prompt_len=prompt_len,
441 prefill_token_ids=new_req_data.prefill_token_ids,
442 num_computed_tokens=new_req_data.num_computed_tokens,
441 - sampling_params=new_req_data.sampling_params,
443 lora_request=new_req_data.lora_request,
444 )
445 req_index = self.req_states.req_id_to_index[req_id]
@@ -461,6 +462,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
462 self.sampler.add_request(
463 req_index, prompt_len, new_req_data.sampling_params
464 )
465 + self.prompt_logprobs_worker.add_request(
466 + req_id, req_index, new_req_data.sampling_params
467 + )
468
469 if scheduler_output.scheduled_new_reqs:
470 self.req_states.apply_staged_writes()
@@ -729,104 +733,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
733 )
734 return sampler_output, num_sampled, num_rejected
735
732 - def compute_prompt_logprobs(
733 - self,
734 - hidden_states: torch.Tensor,
735 - input_batch: InputBatch,
736 - ) -> dict[str, LogprobsTensors]:
737 - idx_mapping_np = input_batch.idx_mapping_np
738 - needs_prompt_logprobs = self.req_states.needs_prompt_logprobs[idx_mapping_np]
739 - if not np.any(needs_prompt_logprobs):
740 - # No request asks for prompt logprobs.
741 - return {}
742 -
743 - prompt_lens = self.req_states.prompt_len[idx_mapping_np]
744 - # NOTE(woosuk): -1 because the last prompt token's hidden state is not
745 - # needed for prompt logprobs.
746 - computed_prefill = self.req_states.num_computed_prefill_tokens[idx_mapping_np]
747 - includes_prompt = computed_prefill < prompt_lens - 1
748 - # NOTE(woosuk): If the request was resumed after preemption, its prompt
749 - # logprobs must have been computed before preemption. Skip.
750 - resumed_after_prompt = (
751 - prompt_lens < self.req_states.prefill_len.np[idx_mapping_np]
752 - )
753 - needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
754 - if not np.any(needs_prompt_logprobs):
755 - return {}
756 -
757 - # Just to be safe, clone the input ids.
758 - n = input_batch.num_tokens
759 - # Shift the input ids by one.
760 - token_ids = torch.empty_like(input_batch.input_ids[:n])
761 - token_ids[: n - 1] = input_batch.input_ids[1:n]
762 - # To avoid out-of-bound access, set the last token id to 0.
763 - token_ids[n - 1] = 0
764 -
765 - # Handle chunked prompts.
766 - pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
767 - is_prompt_chunked = pos_after_step < prompt_lens
768 - prefill_token_ids = self.req_states.prefill_token_ids.gpu
769 - query_start_loc_np = input_batch.query_start_loc_np
770 - for i, req_id in enumerate(input_batch.req_ids):
771 - if not needs_prompt_logprobs[i]:
772 - continue
773 - if not is_prompt_chunked[i]:
774 - continue
775 - # The prompt is chunked. Get the next prompt token.
776 - req_idx = input_batch.idx_mapping_np[i]
777 - idx = int(query_start_loc_np[i + 1] - 1)
778 - # NOTE(woosuk): This triggers two GPU operations.
779 - next_prompt_token = prefill_token_ids[req_idx, pos_after_step[i]]
780 - token_ids[idx] = next_prompt_token
781 -
782 - # NOTE(woosuk): We mask out logprobs for negative tokens.
783 - prompt_logprobs, prompt_ranks = compute_prompt_logprobs(
784 - token_ids,
785 - hidden_states[:n],
786 - self.model.compute_logits,
787 - )
788 -
789 - prompt_token_ids = token_ids.unsqueeze(-1)
790 - prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
791 - for i, req_id in enumerate(input_batch.req_ids):
792 - if not needs_prompt_logprobs[i]:
793 - continue
794 -
795 - start_idx = query_start_loc_np[i]
796 - end_idx = query_start_loc_np[i + 1]
797 - assert start_idx < end_idx, (
798 - f"start_idx ({start_idx}) >= end_idx ({end_idx})"
799 - )
800 - logprobs = LogprobsTensors(
801 - logprob_token_ids=prompt_token_ids[start_idx:end_idx],
802 - logprobs=prompt_logprobs[start_idx:end_idx],
803 - selected_token_ranks=prompt_ranks[start_idx:end_idx],
804 - )
805 -
806 - req_extra_data = self.req_states.extra_data[req_id]
807 - prompt_logprobs_list = req_extra_data.in_progress_prompt_logprobs
808 - if is_prompt_chunked[i]:
809 - # Prompt is chunked. Do not return the logprobs yet.
810 - prompt_logprobs_list.append(logprobs)
811 - continue
812 -
813 - if prompt_logprobs_list:
814 - # Merge the in-progress logprobs.
815 - prompt_logprobs_list.append(logprobs)
816 - logprobs = LogprobsTensors(
817 - logprob_token_ids=torch.cat(
818 - [x.logprob_token_ids for x in prompt_logprobs_list]
819 - ),
820 - logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]),
821 - selected_token_ranks=torch.cat(
822 - [x.selected_token_ranks for x in prompt_logprobs_list]
823 - ),
824 - )
825 - prompt_logprobs_list.clear()
826 -
827 - prompt_logprobs_dict[req_id] = logprobs
828 - return prompt_logprobs_dict
829 -
736 def postprocess(
737 self,
738 input_batch: InputBatch,
@@ -1002,7 +908,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
908 sampler_output, num_sampled, num_rejected = self.sample(
909 hidden_states, input_batch, grammar_output
910 )
1005 - prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
911 + prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
912 + self.model.compute_logits,
913 + hidden_states,
914 + input_batch,
915 + self.req_states.prefill_token_ids.gpu,
916 + self.req_states.num_computed_tokens.gpu,
917 + self.req_states.prompt_len,
918 + self.req_states.prefill_len.np,
919 + self.req_states.num_computed_prefill_tokens,
920 + )
921
922 # Prepare the model runner output.
923 model_runner_output = ModelRunnerOutput(
vllm/v1/worker/gpu/sample/logprob.py
-29
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Callable
3
4 import torch
5
@@ -137,31 +136,3 @@ def compute_topk_logprobs(
136 logprobs=logprobs,
137 selected_token_ranks=token_ranks,
138 )
140 -
141 -
142 -def compute_prompt_logprobs(
143 - prompt_token_ids: torch.Tensor,
144 - prompt_hidden_states: torch.Tensor,
145 - logits_fn: Callable[[torch.Tensor], torch.Tensor],
146 -) -> tuple[torch.Tensor, torch.Tensor]:
147 - # Since materializing the full prompt logits can take too much memory,
148 - # we compute it in chunks.
149 - CHUNK_SIZE = 1024
150 - logprobs = []
151 - ranks = []
152 - prompt_token_ids = prompt_token_ids.to(torch.int64)
153 - for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
154 - end_idx = start_idx + CHUNK_SIZE
155 - # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
156 - prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
157 - prompt_logprobs = compute_topk_logprobs(
158 - prompt_logits,
159 - 0, # num_logprobs
160 - prompt_token_ids[start_idx:end_idx],
161 - )
162 - logprobs.append(prompt_logprobs.logprobs)
163 - ranks.append(prompt_logprobs.selected_token_ranks)
164 -
165 - logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
166 - ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
167 - return logprobs, ranks
vllm/v1/worker/gpu/sample/prompt_logprob.py new
+212
@@ -0,0 +1,212 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 +
5 +import numpy as np
6 +import torch
7 +
8 +from vllm.sampling_params import SamplingParams
9 +from vllm.triton_utils import tl, triton
10 +from vllm.v1.outputs import LogprobsTensors
11 +from vllm.v1.worker.gpu.input_batch import InputBatch
12 +from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
13 +
14 +
15 +class PromptLogprobsWorker:
16 + def __init__(self, max_num_reqs: int):
17 + self.max_num_reqs = max_num_reqs
18 +
19 + self.uses_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
20 + # req_idx -> list of in-progress LogprobsTensors
21 + self.in_progress_prompt_logprobs: dict[str, list[LogprobsTensors]] = {}
22 +
23 + def add_request(self, req_id: str, req_idx: int, sampling_params: SamplingParams):
24 + # For now, only support prompt logprobs for the prompt tokens (not top-k).
25 + uses_prompt_logprobs = sampling_params.prompt_logprobs is not None
26 + if uses_prompt_logprobs:
27 + self.uses_prompt_logprobs[req_idx] = True
28 + self.in_progress_prompt_logprobs[req_id] = []
29 + else:
30 + self.uses_prompt_logprobs[req_idx] = False
31 +
32 + def remove_request(self, req_id: str) -> None:
33 + self.in_progress_prompt_logprobs.pop(req_id, None)
34 +
35 + def compute_prompt_logprobs(
36 + self,
37 + logits_fn: Callable[[torch.Tensor], torch.Tensor],
38 + hidden_states: torch.Tensor,
39 + input_batch: InputBatch,
40 + # [max_num_reqs, max_model_len]
41 + prefill_token_ids: torch.Tensor,
42 + # [max_num_reqs]
43 + num_computed_tokens: torch.Tensor,
44 + # [max_num_reqs]
45 + prompt_lens: np.ndarray,
46 + # [max_num_reqs]
47 + prefill_lens: np.ndarray,
48 + # [max_num_reqs]
49 + num_computed_prefill_tokens: np.ndarray,
50 + ) -> dict[str, LogprobsTensors]:
51 + idx_mapping_np = input_batch.idx_mapping_np
52 + needs_prompt_logprobs = self.uses_prompt_logprobs[idx_mapping_np]
53 + if not np.any(needs_prompt_logprobs):
54 + # Common case: No request asks for prompt logprobs.
55 + return {}
56 +
57 + prompt_lens = prompt_lens[idx_mapping_np]
58 + # NOTE(woosuk): -1 because the last prompt token's hidden state is not
59 + # needed for prompt logprobs.
60 + computed_prefill = num_computed_prefill_tokens[idx_mapping_np]
61 + includes_prompt = computed_prefill < prompt_lens - 1
62 + # NOTE(woosuk): If the request was resumed after preemption, its prompt
63 + # logprobs must have been computed before preemption. Skip.
64 + resumed_after_prompt = prompt_lens < prefill_lens[idx_mapping_np]
65 + needs_prompt_logprobs &= includes_prompt & ~resumed_after_prompt
66 + if not np.any(needs_prompt_logprobs):
67 + return {}
68 +
69 + # Get the prompt logprobs token_ids.
70 + prompt_logprobs_token_ids = get_prompt_logprobs_token_ids(
71 + input_batch.num_tokens,
72 + input_batch.query_start_loc,
73 + input_batch.idx_mapping,
74 + num_computed_tokens,
75 + prefill_token_ids,
76 + )
77 + # Compute the prompt logprobs.
78 + prompt_logprobs, prompt_ranks = compute_prompt_logprobs_with_chunking(
79 + prompt_logprobs_token_ids,
80 + hidden_states[: input_batch.num_tokens],
81 + logits_fn,
82 + )
83 +
84 + pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
85 + is_prompt_chunked = pos_after_step < prompt_lens
86 +
87 + query_start_loc_np = input_batch.query_start_loc_np
88 + prompt_token_ids = prompt_logprobs_token_ids.unsqueeze(-1)
89 + prompt_logprobs_dict: dict[str, LogprobsTensors] = {}
90 + for i, req_id in enumerate(input_batch.req_ids):
91 + if not needs_prompt_logprobs[i]:
92 + continue
93 +
94 + start_idx = query_start_loc_np[i]
95 + end_idx = query_start_loc_np[i + 1]
96 + assert start_idx < end_idx, (
97 + f"start_idx ({start_idx}) >= end_idx ({end_idx})"
98 + )
99 + if not is_prompt_chunked[i]:
100 + end_idx -= 1
101 + logprobs = LogprobsTensors(
102 + logprob_token_ids=prompt_token_ids[start_idx:end_idx],
103 + logprobs=prompt_logprobs[start_idx:end_idx],
104 + selected_token_ranks=prompt_ranks[start_idx:end_idx],
105 + )
106 +
107 + prompt_logprobs_list = self.in_progress_prompt_logprobs[req_id]
108 + if is_prompt_chunked[i]:
109 + # Prompt is chunked. Do not return the logprobs yet.
110 + prompt_logprobs_list.append(logprobs)
111 + continue
112 +
113 + if prompt_logprobs_list:
114 + # Merge the in-progress logprobs.
115 + prompt_logprobs_list.append(logprobs)
116 + logprobs = LogprobsTensors(
117 + logprob_token_ids=torch.cat(
118 + [x.logprob_token_ids for x in prompt_logprobs_list]
119 + ),
120 + logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]),
121 + selected_token_ranks=torch.cat(
122 + [x.selected_token_ranks for x in prompt_logprobs_list]
123 + ),
124 + )
125 + prompt_logprobs_list.clear()
126 +
127 + prompt_logprobs_dict[req_id] = logprobs
128 + return prompt_logprobs_dict
129 +
130 +
131 +@triton.jit
132 +def _prompt_logprobs_token_ids_kernel(
133 + prompt_logprobs_token_ids_ptr,
134 + query_start_loc_ptr,
135 + idx_mapping_ptr,
136 + num_computed_tokens_ptr,
137 + prefill_token_ids_ptr,
138 + prefill_token_ids_stride,
139 + BLOCK_SIZE: tl.constexpr,
140 +):
141 + batch_idx = tl.program_id(0)
142 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
143 +
144 + query_start = tl.load(query_start_loc_ptr + batch_idx)
145 + query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
146 + query_len = query_end - query_start
147 +
148 + num_computed_tokens = tl.load(num_computed_tokens_ptr + req_state_idx)
149 + for i in range(0, query_len, BLOCK_SIZE):
150 + block = i + tl.arange(0, BLOCK_SIZE)
151 + mask = block < query_len
152 + # NOTE(woosuk): We should shift the pos by one
153 + # because the logprob is computed for the next token.
154 + target_pos = num_computed_tokens + 1 + block
155 + token_ids = tl.load(
156 + prefill_token_ids_ptr
157 + + req_state_idx * prefill_token_ids_stride
158 + + target_pos,
159 + mask=mask,
160 + )
161 + tl.store(
162 + prompt_logprobs_token_ids_ptr + query_start + block, token_ids, mask=mask
163 + )
164 +
165 +
166 +def get_prompt_logprobs_token_ids(
167 + num_tokens: int,
168 + query_start_loc: torch.Tensor,
169 + idx_mapping: torch.Tensor,
170 + num_computed_tokens: torch.Tensor,
171 + prefill_token_ids: torch.Tensor,
172 +) -> torch.Tensor:
173 + token_ids = torch.empty(num_tokens, dtype=torch.int64, device=idx_mapping.device)
174 + num_reqs = idx_mapping.shape[0]
175 + _prompt_logprobs_token_ids_kernel[(num_reqs,)](
176 + token_ids,
177 + query_start_loc,
178 + idx_mapping,
179 + num_computed_tokens,
180 + prefill_token_ids,
181 + prefill_token_ids.stride(0),
182 + BLOCK_SIZE=1024,
183 + )
184 + return token_ids
185 +
186 +
187 +def compute_prompt_logprobs_with_chunking(
188 + prompt_token_ids: torch.Tensor,
189 + prompt_hidden_states: torch.Tensor,
190 + logits_fn: Callable[[torch.Tensor], torch.Tensor],
191 +) -> tuple[torch.Tensor, torch.Tensor]:
192 + # Since materializing the full prompt logits can take too much memory,
193 + # we compute it in chunks.
194 + CHUNK_SIZE = 1024
195 + logprobs = []
196 + ranks = []
197 + prompt_token_ids = prompt_token_ids.to(torch.int64)
198 + for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
199 + end_idx = start_idx + CHUNK_SIZE
200 + # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
201 + prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
202 + prompt_logprobs = compute_topk_logprobs(
203 + prompt_logits,
204 + 0, # num_logprobs
205 + prompt_token_ids[start_idx:end_idx],
206 + )
207 + logprobs.append(prompt_logprobs.logprobs)
208 + ranks.append(prompt_logprobs.selected_token_ranks)
209 +
210 + logprobs = torch.cat(logprobs, dim=0) if len(logprobs) > 1 else logprobs[0]
211 + ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
212 + return logprobs, ranks
vllm/v1/worker/gpu/states.py
+1 -11
@@ -1,13 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from dataclasses import dataclass, field
3 +from dataclasses import dataclass
4
5 import numpy as np
6 import torch
7
8 from vllm.lora.request import LoRARequest
9 -from vllm.sampling_params import SamplingParams
10 -from vllm.v1.outputs import LogprobsTensors
9 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
10
11 NO_LORA_ID = 0
@@ -76,8 +74,6 @@ class RequestState:
74 self.lora_ids = np.zeros(self.max_num_reqs, dtype=np.int32)
75 self.lora_ids.fill(NO_LORA_ID)
76
79 - self.needs_prompt_logprobs = np.zeros(self.max_num_reqs, dtype=bool)
80 -
77 @property
78 def num_reqs(self) -> int:
79 return len(self.req_id_to_index)
@@ -88,7 +84,6 @@ class RequestState:
84 prompt_len: int,
85 prefill_token_ids: list[int],
86 num_computed_tokens: int,
91 - sampling_params: SamplingParams,
87 lora_request: LoRARequest | None,
88 ) -> None:
89 assert len(self.free_indices) > 0, "No free indices"
@@ -112,10 +107,6 @@ class RequestState:
107 else:
108 self.lora_ids[req_idx] = NO_LORA_ID
109
115 - # For now, only support prompt logprobs for the prompt tokens.
116 - needs_prompt_logprobs = sampling_params.prompt_logprobs is not None
117 - self.needs_prompt_logprobs[req_idx] = needs_prompt_logprobs
118 -
110 def apply_staged_writes(self) -> None:
111 self.prefill_len.copy_to_uva()
112 self.prefill_token_ids.apply_write()
@@ -151,4 +142,3 @@ class RequestState:
142 @dataclass
143 class ExtraData:
144 lora_request: LoRARequest | None
154 - in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)