topk_topp_sampler.py 14.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6

import torch
import torch.nn as nn
7
from packaging import version
8
9

from vllm import envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config.model import LogprobsMode
12
from vllm.logger import init_logger
13
from vllm.platforms import CpuArchEnum, current_platform
14
15
16
17
18

logger = init_logger(__name__)


class TopKTopPSampler(nn.Module):
19
20
21
22
23
24
    """
    Module that performs optional top-k and top-p filtering followed by
    weighted random sampling of logits.

    Implementations may update the logits tensor in-place.
    """
25

26
    def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs") -> None:
27
        super().__init__()
28
29
30
        self.logprobs_mode = logprobs_mode
        # flashinfer optimization does not apply if intermediate
        # logprobs/logits after top_k/top_p need to be returned
31
32
33
34
        if (
            logprobs_mode not in ("processed_logits", "processed_logprobs")
            and current_platform.is_cuda()
        ):
35
            if envs.VLLM_USE_FLASHINFER_SAMPLER:
36
37
38
39
40
41
42
43
44
45
                from vllm.v1.attention.backends.flashinfer import FlashInferBackend

                capability = current_platform.get_device_capability()
                assert capability is not None
                if not FlashInferBackend.supports_compute_capability(capability):
                    capability_str = capability.as_version_str()
                    raise RuntimeError(
                        "FlashInfer does not support compute capability "
                        f"{capability_str}, unset VLLM_USE_FLASHINFER_SAMPLER=1."
                    )
46
                # Users must opt in explicitly via VLLM_USE_FLASHINFER_SAMPLER=1.
47
48
49
50
                logger.info_once(
                    "Using FlashInfer for top-p & top-k sampling.",
                    scope="global",
                )
51
                self.forward = self.forward_cuda
52
            else:
53
54
55
56
                logger.debug_once(
                    "FlashInfer top-p/top-k sampling is available but disabled "
                    "by default. Set VLLM_USE_FLASHINFER_SAMPLER=1 to opt in "
                    "after verifying accuracy for your workloads."
57
                )
58
                self.forward = self.forward_native
59

60
        elif current_platform.is_cpu():
61
62
63
64
65
            arch = current_platform.get_cpu_architecture()
            # Fall back to native implementation for POWERPC and RISCV.
            # On PowerPC argmax produces incorrect output with torch.compile.
            # PR: https://github.com/vllm-project/vllm/pull/26987
            if arch in (CpuArchEnum.RISCV, CpuArchEnum.POWERPC):
66
67
68
                self.forward = self.forward_native
            else:
                self.forward = self.forward_cpu
69
70
71
72
        elif (
            logprobs_mode not in ("processed_logits", "processed_logprobs")
            and rocm_aiter_ops.is_enabled()
        ):
73
74
            try:
                import aiter.ops.sampling  # noqa: F401
75

76
77
78
79
80
81
82
83
84
85
86
                self.aiter_ops = torch.ops.aiter
                logger.info_once(
                    "Using aiter sampler on ROCm (lazy import, sampling-only)."
                )
                self.forward = self.forward_hip
            except ImportError:
                logger.warning_once(
                    "aiter.ops.sampling is not available on ROCm. "
                    "Falling back to forward_native implementation."
                )
                self.forward = self.forward_native
87
88
        else:
            self.forward = self.forward_native
89
90

        self.apply_top_k_top_p = apply_top_k_top_p
91
92
93
94

    def forward_native(
        self,
        logits: torch.Tensor,
95
        generators: dict[int, torch.Generator],
96
97
98
        k: torch.Tensor | None,
        p: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
99
100
101
102
103
        """
        PyTorch-native implementation of top-k and top-p sampling.

        The logits tensor may be updated in-place.
        """
104
105
        logits = self.apply_top_k_top_p(logits, k, p)
        logits_to_return = None
106
        if self.logprobs_mode == "processed_logits":
107
            logits_to_return = logits
108
        elif self.logprobs_mode == "processed_logprobs":
109
            logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
110
        probs = logits.softmax(dim=-1, dtype=torch.float32)
111
        return random_sample(probs, generators), logits_to_return
112
113
114
115

    def forward_cuda(
        self,
        logits: torch.Tensor,
116
        generators: dict[int, torch.Generator],
117
118
119
        k: torch.Tensor | None,
        p: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
120
        """More optimized implementation for top-k and top-p sampling."""
121
122
123
124
125
        # We prefer `random_sample` over `flashinfer_sample` when sorting is
        # not needed. This is because `random_sample` does not require
        # CPU-GPU synchronization while `flashinfer_sample` does.
        if (k is None and p is None) or generators:
            if generators:
126
127
128
129
130
                logger.debug_once(
                    "FlashInfer 0.2.3+ does not support "
                    "per-request generators. Falling back to "
                    "PyTorch-native implementation."
                )
131
            return self.forward_native(logits, generators, k, p)
132
133
134
        assert self.logprobs_mode not in ("processed_logits", "processed_logprobs"), (
            "FlashInfer does not support returning logits/logprobs"
        )
135
136
137
        # flashinfer sampling functions expect contiguous logits.
        # In flex_attn/triton_attn fp32 inference, logits can be non-contiguous
        # because of slicing operation in logits_processor.
138
        return flashinfer_sample(logits.contiguous(), k, p, generators), None
139

140
141
142
143
    def forward_cpu(
        self,
        logits: torch.Tensor,
        generators: dict[int, torch.Generator],
144
145
146
        k: torch.Tensor | None,
        p: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
        """
        PyTorch-native implementation of top-k and top-p sampling for CPU.

        The logits tensor may be updated in-place.
        """
        logits = self.apply_top_k_top_p(logits, k, p)
        logits_to_return = None
        if self.logprobs_mode == "processed_logits":
            logits_to_return = logits
        elif self.logprobs_mode == "processed_logprobs":
            logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)

        if len(generators) != logits.shape[0]:
            return compiled_random_sample(logits), logits_to_return
        else:
            probs = logits.softmax(dim=-1, dtype=torch.float32)
            q = torch.empty_like(probs)
            q.exponential_()
            for i, generator in generators.items():
                q[i].exponential_(generator=generator)

            return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return

170
171
172
173
174
175
176
    def forward_hip(
        self,
        logits: torch.Tensor,
        generators: dict[int, torch.Generator],
        k: torch.Tensor | None,
        p: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
177
178
        # FIXME: Fix aiter_sampler's accuracy issue and remove this flag
        DISABLE_AITER_SAMPLER = True
179
180
181
182
183
184
185
186
187
188
189
190
        """Optimized ROCm/aiter path (same structure as forward_cuda)."""
        if (k is None and p is None) or generators:
            if generators:
                logger.warning_once(
                    "aiter sampler does not support per-request generators; "
                    "falling back to PyTorch-native."
                )
            return self.forward_native(logits, generators, k, p)
        assert self.logprobs_mode not in (
            "processed_logits",
            "processed_logprobs",
        ), "aiter sampler does not support returning logits/logprobs."
191
192
        if DISABLE_AITER_SAMPLER:
            return self.forward_native(logits, generators, k, p)
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
        return self.aiter_sample(logits, k, p, generators), None

    def aiter_sample(
        self,
        logits: torch.Tensor,
        k: torch.Tensor | None,
        p: torch.Tensor | None,
        generators: dict[int, torch.Generator],
    ) -> torch.Tensor:
        """Sample from logits using aiter ops."""
        use_top_k = k is not None
        use_top_p = p is not None
        # Joint k+p path
        if use_top_p and use_top_k:
            probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
            next_token_ids = self.aiter_ops.top_k_top_p_sampling_from_probs(
                probs,
                None,
                *_to_tensor_scalar_tuple(k),
                *_to_tensor_scalar_tuple(p),
                deterministic=True,
            )
            return next_token_ids.view(-1)
        # Top-p only path
        elif use_top_p:
            probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
            next_token_ids = self.aiter_ops.top_p_sampling_from_probs(
                probs, None, *_to_tensor_scalar_tuple(p), deterministic=True
            )
            return next_token_ids.view(-1)
        # Top-k only path
        elif use_top_k:
            probs = logits.softmax(dim=-1, dtype=torch.float32).contiguous()
            renorm_probs = self.aiter_ops.top_k_renorm_probs(
                probs, *_to_tensor_scalar_tuple(k)
            )
            return torch.multinomial(renorm_probs, num_samples=1).view(-1)
        raise RuntimeError("aiter_sample was called with no active top-k or top-p.")

232

233
234
235
236
237
238
239
240
241
242
# Note: this is a workaround for
# https://github.com/pytorch/pytorch/pull/151218
@torch.compile(dynamic=True)
def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:
    probs = logits.softmax(dim=-1, dtype=torch.float32)
    q = torch.empty_like(probs)
    q.exponential_()
    return probs.div(q).argmax(dim=-1).view(-1)


243
244
def apply_top_k_top_p(
    logits: torch.Tensor,
245
246
    k: torch.Tensor | None,
    p: torch.Tensor | None,
247
248
249
) -> torch.Tensor:
    """Apply top-k and top-p masks to the logits.

250
251
252
253
    If a top-p is used, this function will sort the logits tensor,
    which can be slow for large batches.

    The logits tensor may be updated in-place.
254
    """
255
256
257
258
259
260
261
    if p is None:
        if k is None:
            return logits

        # Avoid sorting vocab for top-k only case.
        return apply_top_k_only(logits, k)

262
263
    logits_sort, logits_idx = logits.sort(dim=-1, descending=False)

264
    if k is not None:
265
        # Apply top-k.
266
        top_k_mask = logits_sort.size(1) - k.to(torch.long)  # shape: B
267
268
269
270
271
        # Get all the top_k values.
        top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))
        top_k_mask = logits_sort < top_k_mask
        logits_sort.masked_fill_(top_k_mask, -float("inf"))

272
    if p is not None:
273
274
        # Apply top-p.
        probs_sort = logits_sort.softmax(dim=-1)
275
        probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
276
277
278
279
280
281
282
283
284
285
        top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
        # at least one
        top_p_mask[:, -1] = False
        logits_sort.masked_fill_(top_p_mask, -float("inf"))

    # Re-sort the probabilities.
    logits = logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort)
    return logits


286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
def apply_top_k_only(
    logits: torch.Tensor,
    k: torch.Tensor,
) -> torch.Tensor:
    """
    Apply top-k mask to the logits.

    This implementation doesn't involve sorting the entire vocab.

    The logits tensor may be updated in-place.
    """
    no_top_k_mask = k == logits.shape[1]
    # Set non-top-k rows to 1 so that we can gather.
    k = k.masked_fill(no_top_k_mask, 1)
    max_top_k = k.max()
    # topk.values tensor has shape [batch_size, max_top_k].
    # Convert top k to 0-based index in range [0, max_top_k).
303
    k_index = k.sub_(1).unsqueeze(1)
304
    top_k_mask = logits.topk(max_top_k, dim=1).values.gather(1, k_index.long())
305
306
307
308
309
310
    # Handle non-topk rows.
    top_k_mask.masked_fill_(no_top_k_mask.unsqueeze(1), -float("inf"))
    logits.masked_fill_(logits < top_k_mask, -float("inf"))
    return logits


311
312
def random_sample(
    probs: torch.Tensor,
313
    generators: dict[int, torch.Generator],
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
) -> torch.Tensor:
    """Randomly sample from the probabilities.

    We use this function instead of torch.multinomial because torch.multinomial
    causes CPU-GPU synchronization.
    """
    q = torch.empty_like(probs)
    # NOTE(woosuk): To batch-process the requests without their own seeds,
    # which is the common case, we first assume that every request does
    # not have its own seed. Then, we overwrite the values for the requests
    # that have their own seeds.
    if len(generators) != probs.shape[0]:
        q.exponential_()
    if generators:
        # TODO(woosuk): This can be slow because we handle each request
        # one by one. Optimize this.
        for i, generator in generators.items():
            q[i].exponential_(generator=generator)
    return probs.div_(q).argmax(dim=-1).view(-1)


def flashinfer_sample(
336
    logits: torch.Tensor,
337
338
    k: torch.Tensor | None,
    p: torch.Tensor | None,
339
    generators: dict[int, torch.Generator],
340
) -> torch.Tensor:
341
    """Sample from the logits using FlashInfer.
342
343
344
345

    Statistically, this function is equivalent to the `random_sample` function.
    However, this function is faster because it avoids sorting the logits tensor
    via rejection sampling.
346

347
348
349
350
351
352
353
354
    NOTE: The outputs of this function do not necessarily match the outputs of
    the `random_sample` function. It only guarantees that the outputs are
    statistically equivalent.

    NOTE: This function includes CPU-GPU synchronization, while `random_sample`
    does not. Call this function at the end of the forward pass to minimize
    the synchronization overhead.
    """
355
356
357
358
359
360
361
    import flashinfer

    if version.parse(flashinfer.__version__) < version.parse("0.2.3"):
        raise ImportError(
            "FlashInfer version >= 0.2.3 required for top-k and top-p sampling. "
        )

362
363
    assert not (k is None and p is None)
    if k is None:
364
        # Top-p only.
365
        probs = logits.softmax(dim=-1, dtype=torch.float32)
366
        next_token_ids = flashinfer.sampling.top_p_sampling_from_probs(
367
368
            probs, p, deterministic=True
        )
369
    elif p is None:
370
        # Top-k only.
371
        probs = logits.softmax(dim=-1, dtype=torch.float32)
372
        next_token_ids = flashinfer.sampling.top_k_sampling_from_probs(
373
374
            probs, k, deterministic=True
        )
375
376
    else:
        # Both top-k and top-p.
377
        next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(
378
379
            logits, k, p, deterministic=True
        )
380

381
    return next_token_ids.view(-1)
382
383
384
385
386
387
388


def _to_tensor_scalar_tuple(x):
    if isinstance(x, torch.Tensor):
        return (x, 0)
    else:
        return (None, x)