[MRV2][Spec Decode] Avoid rejection sampler OOM by chunking (#48630)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Michael Goin committed
Jul 23, 2026 at 06:13 UTC
ac36a7a1e7eb8f03f9ec2b6bf643f1002b205794
12 files changed
+390
-67
tests/v1/spec_decode/test_rejection_sampler_utils.py
+62
@@ -412,3 +412,65 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int)
412
f"Block verification mean accepted length {mean_block:.4f} is worse "
413
f"than standard {mean_standard:.4f}."
414
)
415
+
416
+
417
+@pytest.mark.parametrize("has_draft_logits", [True, False])
418
+def test_chunked_requests_match_full_batch(has_draft_logits: bool):
419
+ torch.manual_seed(7)
420
+ device = "cuda"
421
+ num_reqs = 5
422
+ num_speculative_steps = 3
423
+ vocab_size = 257
424
+
425
+ target_logits = torch.randn(vocab_size, device=device)
426
+ draft_logits = torch.randn(vocab_size, device=device)
427
+ inputs = _build_rejection_sample_inputs(
428
+ target_logits,
429
+ draft_logits,
430
+ num_speculative_steps,
431
+ temperature=0.6,
432
+ num_trials=num_reqs,
433
+ )
434
+ padded_target_logits = torch.empty(
435
+ inputs["target_logits"].shape[0], vocab_size + 3, device=device
436
+ )
437
+ padded_target_logits[:, :vocab_size].copy_(inputs["target_logits"])
438
+ inputs["target_logits"] = padded_target_logits[:, :vocab_size]
439
+ assert inputs["target_logits"].stride(-1) == 1
440
+ assert not inputs["target_logits"].is_contiguous()
441
+ if not has_draft_logits:
442
+ inputs["draft_logits"] = None
443
+
444
+ sampled, num_sampled = rejection_sample(
445
+ **inputs, num_speculative_steps=num_speculative_steps
446
+ )
447
+
448
+ sampled_chunks = []
449
+ num_sampled_chunks = []
450
+ for start, end in ((0, 2), (2, 5)):
451
+ lo = start * (num_speculative_steps + 1)
452
+ hi = end * (num_speculative_steps + 1)
453
+ chunk_inputs = dict(inputs)
454
+ for name in (
455
+ "target_logits",
456
+ "draft_sampled",
457
+ "pos",
458
+ "expanded_idx_mapping",
459
+ "expanded_local_pos",
460
+ ):
461
+ chunk_inputs[name] = inputs[name][lo:hi]
462
+ chunk_inputs["cu_num_logits"] = inputs["cu_num_logits"][start : end + 1] - lo
463
+ chunk_inputs["idx_mapping"] = inputs["idx_mapping"][start:end]
464
+
465
+ chunk_sampled, chunk_num_sampled = rejection_sample(
466
+ **chunk_inputs, num_speculative_steps=num_speculative_steps
467
+ )
468
+ sampled_chunks.append(chunk_sampled)
469
+ num_sampled_chunks.append(chunk_num_sampled)
470
+
471
+ chunked_sampled = torch.cat(sampled_chunks)
472
+ chunked_num_sampled = torch.cat(num_sampled_chunks)
473
+ assert torch.equal(chunked_num_sampled, num_sampled)
474
+ steps = torch.arange(num_speculative_steps + 1, device=device)
475
+ valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1)
476
+ assert torch.equal(chunked_sampled[valid], sampled[valid])
tests/v1/test_outputs.py
+26
-1
@@ -2,7 +2,32 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from unittest import TestCase
4
5
-from vllm.v1.outputs import LogprobsLists
5
+import torch
6
+
7
+from vllm.v1.outputs import LogprobsLists, LogprobsTensors
8
+
9
+
10
+def test_logprobs_tensors_cat():
11
+ first = LogprobsTensors(
12
+ torch.tensor([[1, 2]]),
13
+ torch.tensor([[0.1, 0.2]]),
14
+ torch.tensor([1]),
15
+ )
16
+ second = LogprobsTensors(
17
+ torch.tensor([[3, 4]]),
18
+ torch.tensor([[0.3, 0.4]]),
19
+ torch.tensor([2]),
20
+ )
21
+
22
+ result = LogprobsTensors.cat([first, second], [0, 1, 2])
23
+
24
+ assert result.logprob_token_ids.tolist() == [[1, 2], [3, 4]]
25
+ assert result.logprobs.tolist() == (
26
+ first.logprobs.tolist() + second.logprobs.tolist()
27
+ )
28
+ assert result.selected_token_ranks.tolist() == [1, 2]
29
+ assert result.cu_num_generated_tokens == [0, 1, 2]
30
+ assert LogprobsTensors.cat([first]) is first
31
32
33
class TestLogprobsLists(TestCase):
tests/v1/worker/test_gpu_rejection_sampler_chunking.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import MethodType, SimpleNamespace
5
+from typing import get_args
6
+
7
+import numpy as np
8
+import pytest
9
+import torch
10
+
11
+from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
12
+from vllm.platforms import current_platform
13
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler import (
14
+ RejectionSampler,
15
+ _iter_request_chunks,
16
+)
17
+
18
+
19
+def test_iter_request_chunks_preserves_request_boundaries():
20
+ cu_num_logits = np.array([0, 3, 4, 11, 13], dtype=np.int32)
21
+
22
+ assert list(_iter_request_chunks(cu_num_logits, max_chunk_logits=5)) == [
23
+ (0, 2),
24
+ (2, 3),
25
+ (3, 4),
26
+ ]
27
+
28
+
29
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
30
+@pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode))
31
+def test_chunked_scores_match_full_batch(logprobs_mode: str):
32
+ device = torch.device("cuda")
33
+ cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32)
34
+ num_logits_per_req = np.diff(cu_num_logits_np)
35
+ idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32)
36
+ input_batch = SimpleNamespace(
37
+ num_reqs=4,
38
+ cu_num_logits_np=cu_num_logits_np,
39
+ cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device),
40
+ idx_mapping_np=idx_mapping_np,
41
+ idx_mapping=torch.from_numpy(idx_mapping_np).to(device),
42
+ expanded_idx_mapping=torch.from_numpy(
43
+ np.repeat(idx_mapping_np, num_logits_per_req)
44
+ ).to(device),
45
+ expanded_local_pos=torch.from_numpy(
46
+ np.concatenate(
47
+ [np.arange(count, dtype=np.int32) for count in num_logits_per_req]
48
+ )
49
+ ).to(device),
50
+ )
51
+ rejection_sampler = object.__new__(RejectionSampler)
52
+ rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode)
53
+ rejection_sampler.num_speculative_steps = 3
54
+
55
+ def fake_verify(
56
+ self,
57
+ logits,
58
+ _draft_logits,
59
+ _draft_sampled,
60
+ _pos,
61
+ cu_num_logits,
62
+ idx_mapping,
63
+ *_mappings,
64
+ ):
65
+ num_sampled = torch.diff(cu_num_logits).to(torch.int32)
66
+ sampled = (
67
+ idx_mapping.to(torch.int64).unsqueeze(1) + torch.arange(4, device=device)
68
+ ) % logits.shape[1]
69
+ return logits.float() + 1, sampled, num_sampled
70
+
71
+ rejection_sampler._verify = MethodType(fake_verify, rejection_sampler)
72
+ logits = torch.arange(170, dtype=torch.float32, device=device).view(10, 17)
73
+
74
+ sampled, num_sampled, chunked_logprobs = rejection_sampler._verify_in_chunks(
75
+ logits,
76
+ input_batch,
77
+ draft_logits=None,
78
+ draft_sampled=torch.arange(10, device=device),
79
+ pos=torch.arange(10, device=device),
80
+ max_chunk_logits=5,
81
+ max_num_logprobs=2,
82
+ )
83
+ score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits
84
+ full_logprobs = rejection_sampler._get_logprobs_tensors(
85
+ sampled,
86
+ num_sampled,
87
+ score_logits,
88
+ input_batch.cu_num_logits,
89
+ input_batch.cu_num_logits_np,
90
+ max_num_logprobs=2,
91
+ )
92
+
93
+ assert sampled[:, 0].tolist() == idx_mapping_np.tolist()
94
+ assert num_sampled.tolist() == num_logits_per_req.tolist()
95
+ assert chunked_logprobs is not None
96
+ assert full_logprobs is not None
97
+ assert torch.equal(
98
+ chunked_logprobs.logprob_token_ids,
99
+ full_logprobs.logprob_token_ids,
100
+ )
101
+ assert torch.equal(chunked_logprobs.logprobs, full_logprobs.logprobs)
102
+ assert torch.equal(
103
+ chunked_logprobs.selected_token_ranks,
104
+ full_logprobs.selected_token_ranks,
105
+ )
106
+ assert (
107
+ chunked_logprobs.cu_num_generated_tokens
108
+ == full_logprobs.cu_num_generated_tokens
109
+ )
vllm/config/model.py
+4
@@ -90,6 +90,10 @@ ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
90
LogprobsMode = Literal[
91
"raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs"
92
]
93
+PROCESSED_LOGPROBS_MODES: tuple[LogprobsMode, ...] = (
94
+ "processed_logits",
95
+ "processed_logprobs",
96
+)
97
HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig]
98
ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"]
99
LayerBlockType = Literal["attention", "linear_attention", "mamba"]
vllm/v1/outputs.py
+27
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
+from collections.abc import Sequence
6
from copy import copy
7
from dataclasses import dataclass, field
8
from typing import TYPE_CHECKING, NamedTuple, TypeAlias
@@ -90,6 +91,32 @@ class LogprobsTensors(NamedTuple):
91
self.selected_token_ranks[mask],
92
)
93
94
+ @staticmethod
95
+ def cat(
96
+ tensors: Sequence["LogprobsTensors"],
97
+ cu_num_generated_tokens: list[int] | None = None,
98
+ ) -> "LogprobsTensors":
99
+ """Concatenate flattened logprob tensors."""
100
+ assert tensors
101
+ assert cu_num_generated_tokens is not None or all(
102
+ tensor.cu_num_generated_tokens is None for tensor in tensors
103
+ )
104
+ if len(tensors) == 1:
105
+ tensor = tensors[0]
106
+ if cu_num_generated_tokens is None:
107
+ return tensor
108
+ return tensor._replace(cu_num_generated_tokens=cu_num_generated_tokens)
109
+ return LogprobsTensors(
110
+ logprob_token_ids=torch.cat(
111
+ [tensor.logprob_token_ids for tensor in tensors]
112
+ ),
113
+ logprobs=torch.cat([tensor.logprobs for tensor in tensors]),
114
+ selected_token_ranks=torch.cat(
115
+ [tensor.selected_token_ranks for tensor in tensors]
116
+ ),
117
+ cu_num_generated_tokens=cu_num_generated_tokens,
118
+ )
119
+
120
@staticmethod
121
def empty_cpu(
122
num_positions: int, num_tokens_per_position: int
vllm/v1/sample/ops/topk_topp_sampler.py
+8
-12
@@ -7,7 +7,7 @@ import torch.nn as nn
7
8
from vllm import envs
9
from vllm._aiter_ops import rocm_aiter_ops
10
-from vllm.config.model import LogprobsMode
10
+from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
11
from vllm.logger import init_logger
12
from vllm.platforms import CpuArchEnum, current_platform
13
from vllm.triton_utils import HAS_TRITON
@@ -87,7 +87,7 @@ class TopKTopPSampler(nn.Module):
87
# FlashInfer doesn't expose post-top-k/top-p logits/logprobs,
88
# so it can't be used when the configured mode requires them.
89
can_use_flashinfer = (
90
- logprobs_mode not in ("processed_logits", "processed_logprobs")
90
+ logprobs_mode not in PROCESSED_LOGPROBS_MODES
91
and flashinfer_sampler_supported()
92
)
93
self.forward = (
@@ -108,7 +108,7 @@ class TopKTopPSampler(nn.Module):
108
else:
109
self.forward = self.forward_native
110
elif (
111
- logprobs_mode not in ("processed_logits", "processed_logprobs")
111
+ logprobs_mode not in PROCESSED_LOGPROBS_MODES
112
and rocm_aiter_ops.is_enabled()
113
):
114
self.aiter_ops = None
@@ -165,7 +165,7 @@ class TopKTopPSampler(nn.Module):
165
return self.forward_native(logits, generators, k, p)
166
if self.use_fp64_gumbel:
167
return self.forward_native(logits, generators, k, p)
168
- assert self.logprobs_mode not in ("processed_logits", "processed_logprobs"), (
168
+ assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, (
169
"FlashInfer does not support returning logits/logprobs"
170
)
171
# flashinfer sampling functions expect contiguous logits.
@@ -236,10 +236,9 @@ class TopKTopPSampler(nn.Module):
236
return self.forward_native(logits, generators, k, p)
237
if self.use_fp64_gumbel:
238
return self.forward_native(logits, generators, k, p)
239
- assert self.logprobs_mode not in (
240
- "processed_logits",
241
- "processed_logprobs",
242
- ), "aiter sampler does not support returning logits/logprobs."
239
+ assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, (
240
+ "aiter sampler does not support returning logits/logprobs."
241
+ )
242
if self.aiter_ops is None and not self._init_aiter_ops():
243
return self.forward_native(logits, generators, k, p)
244
return self.aiter_sample(logits, k, p, generators), None
@@ -300,10 +299,7 @@ class TopKTopPSampler(nn.Module):
299
logits.shape[0], dtype=torch.int64, device=logits.device
300
)
301
logits_to_return = None
303
- if (
304
- self.logprobs_mode == "processed_logits"
305
- or self.logprobs_mode == "processed_logprobs"
306
- ):
302
+ if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
303
logits_to_return = torch.empty_like(logits)
304
305
assert len(generators) != logits.shape[0], (
vllm/v1/sample/rejection_sampler.py
+2
-4
@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING
10
import torch
11
import torch.nn as nn
12
13
+from vllm.config.model import PROCESSED_LOGPROBS_MODES
14
from vllm.logger import init_logger
15
from vllm.triton_utils import tl, triton
16
from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
@@ -67,10 +68,7 @@ class RejectionSampler(nn.Module):
68
self.sampler = sampler
69
self.use_fp64_gumbel = getattr(sampler, "use_fp64_gumbel", False)
70
logprobs_mode = self.sampler.logprobs_mode
70
- self.is_processed_logprobs_mode = logprobs_mode in (
71
- "processed_logprobs",
72
- "processed_logits",
73
- )
71
+ self.is_processed_logprobs_mode = logprobs_mode in PROCESSED_LOGPROBS_MODES
72
self.is_logits_logprobs_mode = logprobs_mode in (
73
"raw_logits",
74
"processed_logits",
vllm/v1/worker/gpu/sample/prompt_logprob.py
+1
-9
@@ -132,15 +132,7 @@ class PromptLogprobsWorker:
132
133
if prompt_logprobs_list:
134
# Merge the in-progress logprobs.
135
- logprobs = LogprobsTensors(
136
- logprob_token_ids=torch.cat(
137
- [x.logprob_token_ids for x in prompt_logprobs_list]
138
- ),
139
- logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]),
140
- selected_token_ranks=torch.cat(
141
- [x.selected_token_ranks for x in prompt_logprobs_list]
142
- ),
143
- )
135
+ logprobs = LogprobsTensors.cat(prompt_logprobs_list)
136
prompt_logprobs_list.clear()
137
138
if logprobs is None:
vllm/v1/worker/gpu/sample/sampler.py
+3
-6
@@ -5,7 +5,7 @@ import numpy as np
5
import torch
6
7
import vllm.envs as envs
8
-from vllm.config.model import LogprobsMode
8
+from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode
9
from vllm.sampling_params import SamplingParams
10
from vllm.v1.sample.ops.topk_topp_sampler import (
11
apply_top_k_top_p,
@@ -100,7 +100,7 @@ class Sampler:
100
)
101
102
if return_logprobs:
103
- if self.logprobs_mode in ("processed_logprobs", "processed_logits"):
103
+ if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
104
logits = processed_logits
105
expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
106
cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
@@ -221,10 +221,7 @@ class Sampler:
221
# any greedy requests or per-request seeds, or if post-processed
222
# logprobs need to be returned for any requests.
223
(top_k is None and top_p is None)
224
- or (
225
- return_logprobs
226
- and self.logprobs_mode in ("processed_logprobs", "processed_logits")
227
- )
224
+ or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES)
225
or self.sampling_states.any_greedy(idx_mapping_np)
226
or self.sampling_states.any_explicit_seed(idx_mapping_np)
227
)
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+142
-31
@@ -1,8 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterator
4
+
5
+import numpy as np
6
import torch
7
8
from vllm.config import SpeculativeConfig
9
+from vllm.config.model import PROCESSED_LOGPROBS_MODES
10
from vllm.triton_utils import tl, triton
11
from vllm.v1.outputs import LogprobsTensors
12
from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
@@ -19,6 +23,29 @@ from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
23
rejection_sample,
24
)
25
26
+# Cap on the FP32 target-logits buffer materialized by apply_sampling_params.
27
+# TODO(mgoin): Chunking is a workaround. The rejection kernels already upcast
28
+# per vocab block on load and apply ops like temperature and gumbel, so folding
29
+# sampling-param application into those kernels would remove this buffer and
30
+# its traffic entirely.
31
+MAX_CHUNK_BYTES = 2**30 # 1GB
32
+_FP32_BYTES = 4
33
+
34
+
35
+def _iter_request_chunks(
36
+ cu_num_logits: np.ndarray, max_chunk_logits: int
37
+) -> Iterator[tuple[int, int]]:
38
+ """Yield maximally packed request ranges without splitting requests."""
39
+ assert max_chunk_logits > 0
40
+ num_reqs = cu_num_logits.size - 1
41
+ start = 0
42
+ while start < num_reqs:
43
+ max_logit = int(cu_num_logits[start]) + max_chunk_logits
44
+ end = int(np.searchsorted(cu_num_logits, max_logit, side="right") - 1)
45
+ end = min(num_reqs, max(start + 1, end))
46
+ yield start, end
47
+ start = end
48
+
49
50
@triton.jit
51
def _flatten_sampled_kernel(
@@ -66,18 +93,17 @@ class RejectionSampler:
93
94
def _get_logprobs_tensors(
95
self,
69
- input_batch: InputBatch,
96
sampled: torch.Tensor,
97
num_sampled: torch.Tensor,
98
logits: torch.Tensor,
99
+ cu_num_logits: torch.Tensor,
100
+ cu_num_logits_np: np.ndarray,
101
+ max_num_logprobs: int,
102
) -> LogprobsTensors | None:
74
- max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
75
- input_batch.idx_mapping_np
76
- )
103
if max_num_logprobs == NO_LOGPROBS:
104
return None
105
80
- num_reqs = input_batch.cu_num_logits.shape[0] - 1
106
+ num_reqs = cu_num_logits.shape[0] - 1
107
num_logits = logits.shape[0]
108
flat_sampled = torch.zeros(
109
num_logits, dtype=sampled.dtype, device=sampled.device
@@ -87,48 +113,48 @@ class RejectionSampler:
113
sampled,
114
sampled.stride(0),
115
num_sampled,
90
- input_batch.cu_num_logits,
116
+ cu_num_logits,
117
num_warps=1,
118
)
93
- expanded_logits = num_logits != input_batch.idx_mapping.shape[0]
119
+ expanded_logits = num_logits != num_reqs
120
return compute_topk_scores(
121
logits,
122
max_num_logprobs,
123
flat_sampled,
98
- input_batch.cu_num_logits_np.tolist() if expanded_logits else None,
124
+ cu_num_logits_np.tolist() if expanded_logits else None,
125
logits_mode=self.sampler.logprobs_mode
126
in ("raw_logits", "processed_logits"),
127
)
128
103
- def __call__(
129
+ def _verify(
130
self,
131
logits: torch.Tensor,
106
- input_batch: InputBatch,
107
- draft_logits: torch.Tensor | None = None,
108
- ) -> SamplerOutput:
109
- # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
110
- # that num_nans is computed before applying penalties and temperature.
111
- num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
112
-
113
- draft_sampled = input_batch.input_ids[input_batch.logits_indices]
114
- pos = input_batch.positions[input_batch.logits_indices]
132
+ draft_logits: torch.Tensor | None,
133
+ draft_sampled: torch.Tensor,
134
+ pos: torch.Tensor,
135
+ cu_num_logits: torch.Tensor,
136
+ idx_mapping: torch.Tensor,
137
+ idx_mapping_np: np.ndarray,
138
+ expanded_idx_mapping: torch.Tensor,
139
+ expanded_local_pos: torch.Tensor,
140
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
141
processed_logits = self.sampler.apply_sampling_params(
142
logits,
117
- input_batch.expanded_idx_mapping,
118
- input_batch.idx_mapping_np,
143
+ expanded_idx_mapping,
144
+ idx_mapping_np,
145
pos,
146
draft_sampled,
121
- input_batch.expanded_local_pos,
147
+ expanded_local_pos,
148
)
149
sampled, num_sampled = rejection_sample(
150
processed_logits,
151
draft_logits,
152
draft_sampled,
127
- input_batch.cu_num_logits,
153
+ cu_num_logits,
154
pos,
129
- input_batch.idx_mapping,
130
- input_batch.expanded_idx_mapping,
131
- input_batch.expanded_local_pos,
155
+ idx_mapping,
156
+ expanded_idx_mapping,
157
+ expanded_local_pos,
158
self.sampler.sampling_states.temperature.gpu,
159
self.sampler.sampling_states.seeds.gpu,
160
self.num_speculative_steps,
@@ -136,13 +162,98 @@ class RejectionSampler:
162
use_fp64=self.sampler.use_fp64_gumbel,
163
use_block_verification=self.use_block_verification,
164
)
139
- logprobs_tensors = self._get_logprobs_tensors(
165
+ return processed_logits, sampled, num_sampled
166
+
167
+ def _verify_in_chunks(
168
+ self,
169
+ logits: torch.Tensor,
170
+ input_batch: InputBatch,
171
+ draft_logits: torch.Tensor | None,
172
+ draft_sampled: torch.Tensor,
173
+ pos: torch.Tensor,
174
+ max_chunk_logits: int,
175
+ max_num_logprobs: int,
176
+ ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:
177
+ cu_num_logits_np = input_batch.cu_num_logits_np
178
+ use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES
179
+ sampled_chunks: list[torch.Tensor] = []
180
+ num_sampled_chunks: list[torch.Tensor] = []
181
+ logprobs_chunks: list[LogprobsTensors] = []
182
+
183
+ for start, end in _iter_request_chunks(cu_num_logits_np, max_chunk_logits):
184
+ lo = int(cu_num_logits_np[start])
185
+ hi = int(cu_num_logits_np[end])
186
+ chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo
187
+ chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo
188
+ # draft_logits uses persistent request-state indices and stays global.
189
+ processed_logits, sampled, num_sampled = self._verify(
190
+ logits[lo:hi],
191
+ draft_logits,
192
+ draft_sampled[lo:hi],
193
+ pos[lo:hi],
194
+ chunk_cu_num_logits,
195
+ input_batch.idx_mapping[start:end],
196
+ input_batch.idx_mapping_np[start:end],
197
+ input_batch.expanded_idx_mapping[lo:hi],
198
+ input_batch.expanded_local_pos[lo:hi],
199
+ )
200
+ chunk_logprobs = self._get_logprobs_tensors(
201
+ sampled,
202
+ num_sampled,
203
+ processed_logits if use_processed_logits else logits[lo:hi],
204
+ chunk_cu_num_logits,
205
+ chunk_cu_num_logits_np,
206
+ max_num_logprobs,
207
+ )
208
+ if chunk_logprobs is not None:
209
+ logprobs_chunks.append(chunk_logprobs)
210
+ del processed_logits
211
+ sampled_chunks.append(sampled)
212
+ num_sampled_chunks.append(num_sampled)
213
+
214
+ if len(sampled_chunks) == 1:
215
+ logprobs_tensors = logprobs_chunks[0] if logprobs_chunks else None
216
+ return sampled_chunks[0], num_sampled_chunks[0], logprobs_tensors
217
+
218
+ logprobs_tensors = None
219
+ if logprobs_chunks:
220
+ expanded_logits = logits.shape[0] != input_batch.num_reqs
221
+ logprobs_tensors = LogprobsTensors.cat(
222
+ logprobs_chunks,
223
+ cu_num_generated_tokens=(
224
+ cu_num_logits_np.tolist() if expanded_logits else None
225
+ ),
226
+ )
227
+
228
+ sampled = torch.cat(sampled_chunks)
229
+ num_sampled = torch.cat(num_sampled_chunks)
230
+ return sampled, num_sampled, logprobs_tensors
231
+
232
+ def __call__(
233
+ self,
234
+ logits: torch.Tensor,
235
+ input_batch: InputBatch,
236
+ draft_logits: torch.Tensor | None = None,
237
+ ) -> SamplerOutput:
238
+ # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
239
+ # that num_nans is computed before applying penalties and temperature.
240
+ num_nans = get_num_nans(logits) if self.sampler.compute_nans else None
241
+
242
+ draft_sampled = input_batch.input_ids[input_batch.logits_indices]
243
+ pos = input_batch.positions[input_batch.logits_indices]
244
+
245
+ max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
246
+ input_batch.idx_mapping_np
247
+ )
248
+ max_chunk_logits = max(1, MAX_CHUNK_BYTES // (logits.shape[1] * _FP32_BYTES))
249
+ sampled, num_sampled, logprobs_tensors = self._verify_in_chunks(
250
+ logits,
251
input_batch,
141
- sampled,
142
- num_sampled,
143
- processed_logits
144
- if self.sampler.logprobs_mode in ("processed_logprobs", "processed_logits")
145
- else logits,
252
+ draft_logits,
253
+ draft_sampled,
254
+ pos,
255
+ max_chunk_logits,
256
+ max_num_logprobs,
257
)
258
259
num_sampled, num_rejected = get_num_sampled_and_rejected(
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+4
@@ -888,6 +888,10 @@ def rejection_sample(
888
use_fp64: bool = False,
889
use_block_verification: bool = False,
890
) -> tuple[torch.Tensor, torch.Tensor]:
891
+ assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
892
+ assert draft_logits is None or (
893
+ draft_logits.ndim == 3 and draft_logits.stride(-1) == 1
894
+ )
895
num_reqs = cu_num_logits.shape[0] - 1
896
num_logits, vocab_size = target_logits.shape
897
draft_logits_stride_0 = 0
vllm/v1/worker/gpu_model_runner.py
+2
-4
@@ -37,6 +37,7 @@ from vllm.config import (
37
update_config,
38
)
39
from vllm.config.cache import CacheConfig
40
+from vllm.config.model import PROCESSED_LOGPROBS_MODES
41
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
42
from vllm.distributed.eplb.eplb_state import EplbState
43
from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
@@ -6215,10 +6216,7 @@ class GPUModelRunner(
6216
# memory during profile_run.
6217
# No .clone() of logits: warmup output is discarded, so any in-place
6218
# mutation by forward_native does not affect correctness.
6218
- if self.sampler.logprobs_mode not in (
6219
- "processed_logits",
6220
- "processed_logprobs",
6221
- ):
6219
+ if self.sampler.logprobs_mode not in PROCESSED_LOGPROBS_MODES:
6220
self.sampler(
6221
logits=logits,
6222
sampling_metadata=replace(