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 15, 08:30 UTC
Include private attributes in API documentation (#18614)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
May 23, 2025 at 15:18 UTC
d4c291976085e4bdfabe22ed7c69534263c4cc7b
3 files changed
+43
-49
mkdocs.yaml
+1
@@ -66,6 +66,7 @@ plugins:
66
options:
67
show_symbol_type_heading: true
68
show_symbol_type_toc: true
69
+ filters: []
70
summary:
71
modules: true
72
show_if_no_docstring: true
vllm/model_executor/layers/rejection_sampler.py
+18
-17
@@ -262,16 +262,16 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):
262
True, then a token can be accepted, else it should be
263
rejected.
264
265
- Given {math}`q(\hat{x}_{n+1}|x_1, \dots, x_n)`, the probability of
266
- {math}`\hat{x}_{n+1}` given context {math}`x_1, \dots, x_n` according
267
- to the target model, and {math}`p(\hat{x}_{n+1}|x_1, \dots, x_n)`, the
265
+ Given $q(\hat{x}_{n+1}|x_1, \dots, x_n)$, the probability of
266
+ $\hat{x}_{n+1}$ given context $x_1, \dots, x_n$ according
267
+ to the target model, and $p(\hat{x}_{n+1}|x_1, \dots, x_n)$, the
268
same conditional probability according to the draft model, the token
269
is accepted with probability:
270
271
- :::{math}
271
+ $$
272
\min\left(1, \frac{q(\hat{x}_{n+1}|x_1, \dots, x_n)}
273
{p(\hat{x}_{n+1}|x_1, \dots, x_n)}\right)
274
- :::
274
+ $$
275
276
This implementation does not apply causality. When using the output,
277
if a token is rejected, subsequent tokens should not be used.
@@ -314,30 +314,31 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):
314
target model is recovered (within hardware numerics).
315
316
The probability distribution used in this rejection case is constructed
317
- as follows. Given {math}`q(x|x_1, \dots, x_n)`, the probability of
318
- {math}`x` given context {math}`x_1, \dots, x_n` according to the target
319
- model and {math}`p(x|x_1, \dots, x_n)`, the same conditional probability
317
+ as follows. Given $q(x|x_1, \dots, x_n)$, the probability of
318
+ $x$ given context $x_1, \dots, x_n$ according to the target
319
+ model and $p(x|x_1, \dots, x_n)$, the same conditional probability
320
according to the draft model:
321
322
- :::{math}
322
+ $$
323
x_{n+1} \sim (q(x|x_1, \dots, x_n) - p(x|x_1, \dots, x_n))_+
324
- :::
324
+ $$
325
326
- where {math}`(f(x))_+` is defined as:
326
+ where $(f(x))_+$ is defined as:
327
328
- :::{math}
328
+ $$
329
(f(x))_+ = \frac{\max(0, f(x))}{\sum_x \max(0, f(x))}
330
- :::
330
+ $$
331
332
See https://github.com/vllm-project/vllm/pull/2336 for a visualization
333
of the draft, target, and recovered probability distributions.
334
335
Returns a tensor of shape [batch_size, k, vocab_size].
336
337
- Note: This batches operations on GPU and thus constructs the recovered
338
- distribution for all tokens, even if they are accepted. This causes
339
- division-by-zero errors, so we use self._smallest_positive_value to
340
- avoid that. This introduces some drift to the distribution.
337
+ Note:
338
+ This batches operations on GPU and thus constructs the recovered
339
+ distribution for all tokens, even if they are accepted. This causes
340
+ division-by-zero errors, so we use self._smallest_positive_value to
341
+ avoid that. This introduces some drift to the distribution.
342
"""
343
_, k, _ = draft_probs.shape
344
vllm/model_executor/layers/typical_acceptance_sampler.py
+24
-32
@@ -93,29 +93,27 @@ class TypicalAcceptanceSampler(SpecDecodeDeterministicBaseSampler):
93
Evaluates and returns a mask of accepted tokens based on the
94
posterior probabilities.
95
96
- Parameters:
97
- ----------
98
- target_probs : torch.Tensor
99
- A tensor of shape (batch_size, k, vocab_size) representing
100
- the probabilities of each token in the vocabulary for each
101
- position in the proposed sequence. This is the distribution
102
- generated by the target model.
103
- draft_token_ids : torch.Tensor
104
- A tensor of shape (batch_size, k) representing the proposed
105
- token ids.
96
+ Args:
97
+ target_probs (torch.Tensor): A tensor of shape
98
+ (batch_size, k, vocab_size) representing the probabilities of
99
+ each token in the vocabulary for each position in the proposed
100
+ sequence. This is the distribution generated by the target
101
+ model.
102
+ draft_token_ids (torch.Tensor): A tensor of shape (batch_size, k)
103
+ representing the proposed token ids.
104
105
A draft token_id x_{n+k} is accepted if it satisfies the
106
following condition
107
110
- :::{math}
108
+ $$
109
p_{\text{original}}(x_{n+k} | x_1, x_2, \dots, x_{n+k-1}) >
110
\min \left( \epsilon, \delta * \exp \left(
111
-H(p_{\text{original}}(
112
\cdot | x_1, x_2, \ldots, x_{n+k-1})) \right) \right)
115
- :::
113
+ $$
114
117
- where {math}`p_{\text{original}}` corresponds to target_probs
118
- and {math}`\epsilon` and {math}`\delta` correspond to hyperparameters
115
+ where $p_{\text{original}}$ corresponds to target_probs
116
+ and $\epsilon$ and $\delta$ correspond to hyperparameters
117
specified using self._posterior_threshold and self._posterior_alpha
118
119
This method computes the posterior probabilities for the given
@@ -126,13 +124,10 @@ class TypicalAcceptanceSampler(SpecDecodeDeterministicBaseSampler):
124
returns a boolean mask indicating which tokens can be accepted.
125
126
Returns:
129
- -------
130
- torch.Tensor
131
- A boolean tensor of shape (batch_size, k) where each element
132
- indicates whether the corresponding draft token has been accepted
133
- or rejected. True indicates acceptance and false indicates
134
- rejection.
135
-
127
+ torch.Tensor: A boolean tensor of shape (batch_size, k) where each
128
+ element indicates whether the corresponding draft token has
129
+ been accepted or rejected. True indicates acceptance and false
130
+ indicates rejection.
131
"""
132
device = target_probs.device
133
candidates_prob = torch.gather(
@@ -156,17 +151,14 @@ class TypicalAcceptanceSampler(SpecDecodeDeterministicBaseSampler):
151
The recovered token ids will fill the first unmatched token
152
by the target token.
153
159
- Parameters
160
- ----------
161
- target_probs : torch.Tensor
162
- A tensor of shape (batch_size, k, vocab_size) containing
163
- the target probability distribution
164
-
165
- Returns
166
- -------
167
- torch.Tensor
168
- A tensor of shape (batch_size, k) with the recovered token
169
- ids which are selected from target probs.
154
+ Args:
155
+ target_probs (torch.Tensor): A tensor of shape
156
+ (batch_size, k, vocab_size) containing the target probability
157
+ distribution.
158
+
159
+ Returns:
160
+ torch.Tensor: A tensor of shape (batch_size, k) with the recovered
161
+ token ids which are selected from target probs.
162
"""
163
max_indices = torch.argmax(target_probs, dim=-1)
164