@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5daf62271d
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 5, 11:00 UTC

[Model Runner V2] Fuse probabilistic rejection sample kernels (#38496)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Apr 7, 2026 at 17:37 UTC 5daf62271d207e8cbf8694b881a081a3c57a0fff
5 files changed +886 -377
.buildkite/test_areas/model_runner_v2.yaml
+2
@@ -100,11 +100,13 @@ steps:
100 - vllm/v1/worker/gpu/
101 - vllm/v1/worker/gpu_worker.py
102 - tests/v1/spec_decode/test_max_len.py
103 + - tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
104 - tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
105 - tests/v1/e2e/spec_decode/test_spec_decode.py
106 commands:
107 - set -x
108 - export VLLM_USE_V2_MODEL_RUNNER=1
109 - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
110 + - pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
111 - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
112 - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py new
+215
@@ -0,0 +1,215 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
10 + probabilistic_rejection_sample,
11 +)
12 +
13 +VOCAB_SIZE = 4096
14 +
15 +# Skip if no CUDA - Triton kernel requires GPU
16 +pytest.importorskip("triton")
17 +if not torch.cuda.is_available():
18 + pytest.skip("CUDA required for rejection sampler tests", allow_module_level=True)
19 +
20 +
21 +def _build_rejection_sample_inputs(
22 + target_logits_1d: torch.Tensor,
23 + draft_logits_1d: torch.Tensor,
24 + num_speculative_steps: int,
25 + temperature: float,
26 + num_trials: int,
27 +) -> dict:
28 + device = target_logits_1d.device
29 + vocab_size = target_logits_1d.shape[0]
30 + K = num_speculative_steps
31 + num_logits = num_trials * (K + 1)
32 +
33 + target_logits = target_logits_1d.unsqueeze(0).expand(num_logits, -1).contiguous()
34 + draft_logits = (
35 + draft_logits_1d.view(1, 1, vocab_size).expand(num_trials, K, -1).contiguous()
36 + )
37 +
38 + draft_probs = torch.softmax(draft_logits_1d, dim=0)
39 + draft_tokens = torch.multinomial(
40 + draft_probs.expand(num_trials, -1), K, replacement=True
41 + )
42 + draft_sampled_2d = torch.zeros(num_trials, K + 1, dtype=torch.int64, device=device)
43 + draft_sampled_2d[:, 1:] = draft_tokens
44 + draft_sampled = draft_sampled_2d.reshape(-1)
45 +
46 + cu_num_logits = torch.arange(num_trials + 1, dtype=torch.int32, device=device) * (
47 + K + 1
48 + )
49 + pos = torch.arange(num_logits, dtype=torch.int32, device=device)
50 + idx_mapping = torch.arange(num_trials, dtype=torch.int32, device=device)
51 + expanded_idx_mapping = torch.arange(
52 + num_trials, dtype=torch.int32, device=device
53 + ).repeat_interleave(K + 1)
54 + expanded_local_pos = torch.arange(K + 1, dtype=torch.int32, device=device).repeat(
55 + num_trials
56 + )
57 + temp_tensor = torch.full(
58 + (num_trials,), temperature, dtype=torch.float32, device=device
59 + )
60 + seed = torch.arange(num_trials, dtype=torch.int64, device=device)
61 +
62 + return dict(
63 + target_logits=target_logits,
64 + draft_logits=draft_logits,
65 + draft_sampled=draft_sampled,
66 + cu_num_logits=cu_num_logits,
67 + pos=pos,
68 + idx_mapping=idx_mapping,
69 + expanded_idx_mapping=expanded_idx_mapping,
70 + expanded_local_pos=expanded_local_pos,
71 + temperature=temp_tensor,
72 + seed=seed,
73 + )
74 +
75 +
76 +def _assert_distribution_match(
77 + sampled_tokens: torch.Tensor,
78 + target_probs: torch.Tensor,
79 + device: str,
80 + label: str = "",
81 + min_expected: float = 5.0,
82 +):
83 + """
84 + Assert sampled tokens match the target distribution via a
85 + chi-squared goodness-of-fit test. This is done by computing
86 + observed vs expected token counts (target_probs * num_samples),
87 + then checking that the chi-squared statistic is below a conservative
88 + threshold. The threshold is set at df + 10*sqrt(2*df), which
89 + corresponds to ~10 sigma under the chi-squared distribution's
90 + normal approximation, effectively disallowing false positives.
91 +
92 + NOTE: Tokens with expected count < min_expected are merged into
93 + a single "other" bin to minimize chi-squared noise.
94 + """
95 + num_samples = sampled_tokens.shape[0]
96 + vocab_size = target_probs.shape[0]
97 +
98 + observed = torch.zeros(vocab_size, device=device, dtype=torch.float32)
99 + observed.scatter_add_(0, sampled_tokens, torch.ones(num_samples, device=device))
100 + expected = target_probs * num_samples
101 +
102 + sufficient = expected >= min_expected
103 + obs_main = observed[sufficient]
104 + exp_main = expected[sufficient]
105 +
106 + obs_other = observed[~sufficient].sum().unsqueeze(0)
107 + exp_other = expected[~sufficient].sum().unsqueeze(0)
108 +
109 + if exp_other.item() >= min_expected:
110 + obs_all = torch.cat([obs_main, obs_other])
111 + exp_all = torch.cat([exp_main, exp_other])
112 + else:
113 + obs_all = obs_main
114 + exp_all = exp_main
115 +
116 + chi2 = ((obs_all - exp_all) ** 2 / exp_all).sum().item()
117 + df = obs_all.shape[0] - 1
118 + if df < 1:
119 + # All samples were merged into < 2 bins, which is too
120 + # few to evaluate.
121 + return
122 +
123 + threshold = df + 10 * math.sqrt(2 * df)
124 + prefix = f"[{label}] " if label else ""
125 + assert chi2 < threshold, (
126 + f"{prefix}Chi-squared test failed: chi2={chi2:.1f}, "
127 + f"df={df}, threshold={threshold:.1f}. "
128 + f"Output distribution does not match target distribution."
129 + )
130 +
131 +
132 +@pytest.mark.parametrize(
133 + "num_speculative_steps,temperature",
134 + [
135 + (1, 0.6),
136 + (3, 0.6),
137 + (1, 1.0),
138 + (3, 1.0),
139 + ],
140 +)
141 +def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: float):
142 + """
143 + Verify that rejection sampling produces the target distribution.
144 + This is done by simulating many independent trials of speculative
145 + decoding (from a fixed target and draft distribution). We then
146 + run rejection sample on all of the trials (requests), and verify
147 + that the sampled tokens at every position follow the target
148 + distribution p(x).
149 + """
150 +
151 + torch.manual_seed(42)
152 + device = "cuda"
153 + num_trials = 10 * VOCAB_SIZE
154 +
155 + target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
156 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
157 +
158 + if temperature > 0:
159 + target_logits_1d /= temperature
160 + draft_logits_1d /= temperature
161 +
162 + inputs = _build_rejection_sample_inputs(
163 + target_logits_1d,
164 + draft_logits_1d,
165 + num_speculative_steps,
166 + temperature=temperature,
167 + num_trials=num_trials,
168 + )
169 +
170 + sampled, num_sampled = probabilistic_rejection_sample(
171 + **inputs, num_speculative_steps=num_speculative_steps
172 + )
173 +
174 + target_probs = torch.softmax(target_logits_1d, dim=0)
175 + for pos in range(num_speculative_steps + 1):
176 + accepted_mask = num_sampled >= pos + 1
177 + _assert_distribution_match(
178 + sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
179 + )
180 +
181 +
182 +@pytest.mark.parametrize("num_speculative_steps", [1, 3])
183 +def test_greedy_rejection_sample(num_speculative_steps: int):
184 + """
185 + Verify that greedy (temperature=0) always outputs the target argmax
186 + at every accepted position.
187 + """
188 +
189 + torch.manual_seed(42)
190 + device = "cuda"
191 + num_trials = 10 * VOCAB_SIZE
192 +
193 + target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
194 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
195 +
196 + inputs = _build_rejection_sample_inputs(
197 + target_logits_1d,
198 + draft_logits_1d,
199 + num_speculative_steps,
200 + temperature=0.0,
201 + num_trials=num_trials,
202 + )
203 +
204 + sampled, num_sampled = probabilistic_rejection_sample(
205 + **inputs, num_speculative_steps=num_speculative_steps
206 + )
207 +
208 + target_argmax = target_logits_1d.argmax().item()
209 +
210 + steps = torch.arange(num_speculative_steps + 1, device=device).unsqueeze(0)
211 + accepted_mask = steps < num_sampled.unsqueeze(1)
212 +
213 + assert (sampled[accepted_mask] == target_argmax).all(), (
214 + "Greedy sampling produced tokens that are not the target argmax"
215 + )
vllm/v1/worker/gpu/sample/gumbel.py
+54 -25
@@ -65,36 +65,20 @@ def tl_rand64(seed, offset, includes_zero: tl.constexpr):
65
66
67 @triton.jit
68 -def _gumbel_sample_kernel(
69 - local_argmax_ptr,
70 - local_argmax_stride,
71 - local_max_ptr,
72 - local_max_stride,
73 - processed_logits_ptr,
74 - processed_logits_stride,
75 - logits_ptr,
76 - logits_stride,
68 +def gumbel_block_argmax(
69 + logits,
70 + block,
71 + mask,
72 + token_idx,
73 expanded_idx_mapping_ptr,
74 + temp_ptr,
75 seeds_ptr,
76 pos_ptr,
80 - temp_ptr,
81 - vocab_size,
82 - BLOCK_SIZE: tl.constexpr,
77 + processed_logits_ptr,
78 + processed_logits_stride,
79 APPLY_TEMPERATURE: tl.constexpr,
80 ):
85 - token_idx = tl.program_id(0)
81 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
87 -
88 - block_idx = tl.program_id(1)
89 - block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
90 - mask = block < vocab_size
91 - logits = tl.load(
92 - logits_ptr + token_idx * logits_stride + block,
93 - mask=mask,
94 - other=float("-inf"),
95 - )
96 - logits = logits.to(tl.float32)
97 -
82 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
83 if temp != 0.0 and APPLY_TEMPERATURE:
84 # Apply temperature.
@@ -102,8 +86,8 @@ def _gumbel_sample_kernel(
86 # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.
87 logits = logits / temp
88
105 - # Store the temperature-applied logits.
89 if processed_logits_ptr is not None:
90 + # Store the temperature-applied logits.
91 tl.store(
92 processed_logits_ptr + req_state_idx * processed_logits_stride + block,
93 logits,
@@ -126,6 +110,51 @@ def _gumbel_sample_kernel(
110 logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
111
112 value, idx = tl.max(logits, axis=0, return_indices=True)
113 + return value, idx
114 +
115 +
116 +@triton.jit
117 +def _gumbel_sample_kernel(
118 + local_argmax_ptr,
119 + local_argmax_stride,
120 + local_max_ptr,
121 + local_max_stride,
122 + processed_logits_ptr,
123 + processed_logits_stride,
124 + logits_ptr,
125 + logits_stride,
126 + expanded_idx_mapping_ptr,
127 + seeds_ptr,
128 + pos_ptr,
129 + temp_ptr,
130 + vocab_size,
131 + BLOCK_SIZE: tl.constexpr,
132 + APPLY_TEMPERATURE: tl.constexpr,
133 +):
134 + token_idx = tl.program_id(0)
135 + block_idx = tl.program_id(1)
136 + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
137 + mask = block < vocab_size
138 + logits = tl.load(
139 + logits_ptr + token_idx * logits_stride + block,
140 + mask=mask,
141 + other=float("-inf"),
142 + )
143 + logits = logits.to(tl.float32)
144 +
145 + value, idx = gumbel_block_argmax(
146 + logits,
147 + block,
148 + mask,
149 + token_idx,
150 + expanded_idx_mapping_ptr,
151 + temp_ptr,
152 + seeds_ptr,
153 + pos_ptr,
154 + processed_logits_ptr,
155 + processed_logits_stride,
156 + APPLY_TEMPERATURE=APPLY_TEMPERATURE,
157 + )
158 token_id = block_idx * BLOCK_SIZE + idx
159 tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
160 tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
vllm/v1/worker/gpu/spec_decode/probabilistic_rejection_sampler_utils.py new
+612
@@ -0,0 +1,612 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.triton_utils import tl, triton
6 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand64
7 +
8 +
9 +@triton.jit
10 +def _compute_block_max_and_sumexp(logits):
11 + block_max = tl.max(logits, axis=0)
12 + block_sumexp = tl.where(
13 + block_max > float("-inf"),
14 + tl.sum(tl.exp(logits - block_max)),
15 + 0.0,
16 + )
17 + return block_max, block_sumexp
18 +
19 +
20 +@triton.jit
21 +def _compute_global_lse(
22 + local_max_ptr,
23 + local_max_stride,
24 + local_sumexp_ptr,
25 + local_sumexp_stride,
26 + logit_idx,
27 + vocab_num_blocks,
28 + PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
29 +):
30 + blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
31 + blocks_mask = blocks < vocab_num_blocks
32 + maxes = tl.load(
33 + local_max_ptr + logit_idx * local_max_stride + blocks,
34 + mask=blocks_mask,
35 + other=float("-inf"),
36 + )
37 + sumexps = tl.load(
38 + local_sumexp_ptr + logit_idx * local_sumexp_stride + blocks,
39 + mask=blocks_mask,
40 + other=0.0,
41 + )
42 + global_max = tl.max(maxes, axis=0)
43 + global_lse = global_max + tl.log(tl.sum(sumexps * tl.exp(maxes - global_max)))
44 + return global_lse
45 +
46 +
47 +@triton.jit
48 +def _compute_block_max_and_sumexp_kernel(
49 + # [num_logits, num_blocks]
50 + target_local_argmax_ptr,
51 + target_local_argmax_stride,
52 + # [num_logits, num_blocks]
53 + target_local_max_ptr,
54 + target_local_max_stride,
55 + # [num_logits, num_blocks]
56 + target_local_sumexp_ptr,
57 + target_local_sumexp_stride,
58 + # [num_logits, num_blocks]
59 + draft_local_max_ptr,
60 + draft_local_max_stride,
61 + # [num_logits, num_blocks]
62 + draft_local_sumexp_ptr,
63 + draft_local_sumexp_stride,
64 + # [num_logits, V]
65 + target_logits_ptr,
66 + target_logits_stride,
67 + # [max_num_reqs, num_speculative_steps, V]
68 + draft_logits_ptr,
69 + draft_logits_stride_0,
70 + draft_logits_stride_1,
71 + # [num_logits]
72 + expanded_idx_mapping_ptr,
73 + # [num_logits]
74 + expanded_local_pos_ptr,
75 + # [max_num_reqs]
76 + temp_ptr,
77 + vocab_size,
78 + num_speculative_steps,
79 + BLOCK_SIZE: tl.constexpr,
80 +):
81 + logit_idx = tl.program_id(0)
82 + draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
83 +
84 + if draft_step_idx >= num_speculative_steps:
85 + # Bonus token. Max/argmax and summed exponentials are not needed.
86 + return
87 +
88 + req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
89 + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
90 +
91 + block_idx = tl.program_id(1)
92 + block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
93 + mask = block_offsets < vocab_size
94 +
95 + if temp == 0.0:
96 + # Greedy sampling. Only the target max/argmax are needed.
97 + target_logits = tl.load(
98 + target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
99 + mask=mask,
100 + other=float("-inf"),
101 + ).to(tl.float32)
102 + value, idx = tl.max(target_logits, axis=0, return_indices=True)
103 + token_id = block_idx * BLOCK_SIZE + idx
104 + tl.store(
105 + target_local_argmax_ptr
106 + + logit_idx * target_local_argmax_stride
107 + + block_idx,
108 + token_id,
109 + )
110 + tl.store(
111 + target_local_max_ptr + logit_idx * target_local_max_stride + block_idx,
112 + value,
113 + )
114 + else:
115 + # Get local draft max and summed exponentials.
116 + draft_logits = tl.load(
117 + draft_logits_ptr
118 + + req_state_idx * draft_logits_stride_0
119 + + draft_step_idx * draft_logits_stride_1
120 + + block_offsets,
121 + mask=mask,
122 + other=float("-inf"),
123 + ).to(tl.float32)
124 + draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
125 + tl.store(
126 + draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
127 + draft_max,
128 + )
129 + tl.store(
130 + draft_local_sumexp_ptr + logit_idx * draft_local_sumexp_stride + block_idx,
131 + draft_sumexp,
132 + )
133 + # Get local target max and summed exponentials.
134 + target_logits = tl.load(
135 + target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
136 + mask=mask,
137 + other=float("-inf"),
138 + ).to(tl.float32)
139 + target_max, target_sumexp = _compute_block_max_and_sumexp(target_logits)
140 + tl.store(
141 + target_local_max_ptr + logit_idx * target_local_max_stride + block_idx,
142 + target_max,
143 + )
144 + tl.store(
145 + target_local_sumexp_ptr
146 + + logit_idx * target_local_sumexp_stride
147 + + block_idx,
148 + target_sumexp,
149 + )
150 +
151 +
152 +@triton.jit
153 +def _probabilistic_rejection_kernel(
154 + # [num_reqs, num_speculative_steps + 1]
155 + sampled_ptr,
156 + sampled_stride,
157 + # [num_reqs]
158 + rejected_steps_ptr,
159 + # [num_reqs]
160 + target_rejected_logsumexp_ptr,
161 + # [num_reqs]
162 + draft_rejected_logsumexp_ptr,
163 + # [num_logits, V]
164 + target_logits_ptr,
165 + target_logits_stride,
166 + # [num_logits, num_blocks]
167 + target_local_argmax_ptr,
168 + target_local_argmax_stride,
169 + # [num_logits, num_blocks]
170 + target_local_max_ptr,
171 + target_local_max_stride,
172 + # [num_logits, num_blocks]
173 + target_local_sumexp_ptr,
174 + target_local_sumexp_stride,
175 + # [num_logits]
176 + draft_sampled_ptr,
177 + # [max_num_reqs, num_speculative_steps, V]
178 + draft_logits_ptr,
179 + draft_logits_stride_0,
180 + draft_logits_stride_1,
181 + # [num_logits, num_blocks]
182 + draft_local_max_ptr,
183 + draft_local_max_stride,
184 + # [num_logits, num_blocks]
185 + draft_local_sumexp_ptr,
186 + draft_local_sumexp_stride,
187 + # [num_reqs + 1]
188 + cu_num_logits_ptr,
189 + # [num_reqs]
190 + idx_mapping_ptr,
191 + # [max_num_reqs]
192 + temp_ptr,
193 + # [max_num_reqs]
194 + seed_ptr,
195 + # [num_logits]
196 + pos_ptr,
197 + vocab_num_blocks,
198 + PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
199 +):
200 + req_idx = tl.program_id(0)
201 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
202 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
203 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
204 + num_tokens = end_idx - start_idx
205 + seed = tl.load(seed_ptr + req_state_idx)
206 + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
207 +
208 + rejected_step = 0
209 + target_lse = 0.0
210 + draft_lse = 0.0
211 + accepted = True
212 + for i in range(num_tokens - 1):
213 + if accepted:
214 + logit_idx = start_idx + i
215 + draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1)
216 + if temp == 0.0:
217 + # Greedy sampling. Accept IFF draft matches target argmax.
218 + # NOTE: Target argmax is stored directly so that resampling
219 + # can be skipped upon rejection.
220 + target_blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
221 + target_blocks_mask = target_blocks < vocab_num_blocks
222 + target_local_max = tl.load(
223 + target_local_max_ptr
224 + + logit_idx * target_local_max_stride
225 + + target_blocks,
226 + mask=target_blocks_mask,
227 + other=float("-inf"),
228 + )
229 + max_target_block_idx = tl.argmax(target_local_max, axis=0)
230 + target_argmax = tl.load(
231 + target_local_argmax_ptr
232 + + logit_idx * target_local_argmax_stride
233 + + max_target_block_idx
234 + )
235 + accepted &= target_argmax == draft_sampled
236 + tl.store(sampled_ptr + req_idx * sampled_stride + i, target_argmax)
237 + else:
238 + target_logit = tl.load(
239 + target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
240 + ).to(tl.float32)
241 + draft_logit = tl.load(
242 + draft_logits_ptr
243 + + req_state_idx * draft_logits_stride_0
244 + + i * draft_logits_stride_1
245 + + draft_sampled
246 + ).to(tl.float32)
247 + target_lse = _compute_global_lse(
248 + target_local_max_ptr,
249 + target_local_max_stride,
250 + target_local_sumexp_ptr,
251 + target_local_sumexp_stride,
252 + logit_idx,
253 + vocab_num_blocks,
254 + PADDED_VOCAB_NUM_BLOCKS,
255 + )
256 + draft_lse = _compute_global_lse(
257 + draft_local_max_ptr,
258 + draft_local_max_stride,
259 + draft_local_sumexp_ptr,
260 + draft_local_sumexp_stride,
261 + logit_idx,
262 + vocab_num_blocks,
263 + PADDED_VOCAB_NUM_BLOCKS,
264 + )
265 + target_log_prob = target_logit - target_lse
266 + draft_log_prob = draft_logit - draft_lse
267 + pos = tl.load(pos_ptr + logit_idx)
268 + u = tl_rand64(seed, pos, includes_zero=False)
269 + # Probability ratio test: p(x) > u * q(x)
270 + # Equivalent log form: log_p(x) > log(u) + log_q(x)
271 + accepted &= target_log_prob > tl.log(u) + draft_log_prob
272 + tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
273 + rejected_step += accepted
274 + tl.store(rejected_steps_ptr + req_idx, rejected_step)
275 + tl.store(target_rejected_logsumexp_ptr + req_idx, target_lse)
276 + tl.store(draft_rejected_logsumexp_ptr + req_idx, draft_lse)
277 +
278 +
279 +@triton.jit
280 +def _resample_kernel(
281 + # [num_reqs, num_blocks]
282 + resampled_local_argmax_ptr,
283 + resampled_local_argmax_stride,
284 + # [num_reqs, num_blocks]
285 + resampled_local_max_ptr,
286 + resampled_local_max_stride,
287 + # [num_logits, V]
288 + target_logits_ptr,
289 + target_logits_stride,
290 + # [num_reqs]
291 + target_rejected_logsumexp_ptr,
292 + # [max_num_reqs, num_speculative_steps, V]
293 + draft_logits_ptr,
294 + draft_logits_stride_0,
295 + draft_logits_stride_1,
296 + # [num_reqs]
297 + draft_rejected_logsumexp_ptr,
298 + # [num_reqs]
299 + rejected_step_ptr,
300 + # [num_reqs + 1]
301 + cu_num_logits_ptr,
302 + # [num_logits]
303 + expanded_idx_mapping_ptr,
304 + # [max_num_reqs]
305 + temp_ptr,
306 + # [max_num_reqs]
307 + seed_ptr,
308 + # [num_logits]
309 + pos_ptr,
310 + vocab_size,
311 + BLOCK_SIZE: tl.constexpr,
312 +):
313 + req_idx = tl.program_id(0)
314 + resample_idx = tl.load(rejected_step_ptr + req_idx)
315 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
316 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
317 + resample_token_idx = start_idx + resample_idx
318 + req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx)
319 +
320 + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
321 + is_bonus = resample_token_idx == end_idx - 1
322 + if temp == 0.0 and not is_bonus:
323 + # Greedy + non-bonus token. No resampling needed because
324 + # the target argmax is already in the sampled tensor.
325 + return
326 +
327 + block_idx = tl.program_id(1)
328 + block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
329 + mask = block < vocab_size
330 +
331 + # Compute the residual logits to resample the rejected token
332 + # from. In the case of no rejections (bonus token), we directly
333 + # use the target logits.
334 + if is_bonus:
335 + residual_logits = tl.load(
336 + target_logits_ptr + resample_token_idx * target_logits_stride + block,
337 + mask=mask,
338 + other=float("-inf"),
339 + ).to(tl.float32)
340 + else:
341 + target_logits = tl.load(
342 + target_logits_ptr + resample_token_idx * target_logits_stride + block,
343 + mask=mask,
344 + other=float("-inf"),
345 + ).to(tl.float32)
346 + draft_logits = tl.load(
347 + draft_logits_ptr
348 + + req_state_idx * draft_logits_stride_0
349 + + resample_idx * draft_logits_stride_1
350 + + block,
351 + mask=mask,
352 + other=float("-inf"),
353 + ).to(tl.float32)
354 + target_lse = tl.load(target_rejected_logsumexp_ptr + req_idx)
355 + draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx)
356 + target_log_probs = target_logits - target_lse
357 + draft_log_probs = draft_logits - draft_lse
358 + # Compute the residual: max(p(x) - q(x), 0)
359 + # Equivalent log form: log(max(exp(log_p(x)) - exp(log_q(x)), 0))
360 + # The more numerically stable form is:
361 + # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
362 + ratio = tl.exp(draft_log_probs - target_log_probs)
363 + residual_logits = tl.where(
364 + ratio < 1.0,
365 + target_log_probs + tl.log(1 - ratio),
366 + float("-inf"),
367 + ).to(tl.float32)
368 +
369 + # Resample the rejected/bonus token.
370 + value, idx = gumbel_block_argmax(
371 + residual_logits,
372 + block,
373 + mask,
374 + resample_token_idx,
375 + expanded_idx_mapping_ptr,
376 + temp_ptr,
377 + seed_ptr,
378 + pos_ptr,
379 + None,
380 + 0,
381 + APPLY_TEMPERATURE=False,
382 + )
383 + token_id = block_idx * BLOCK_SIZE + idx
384 + tl.store(
385 + resampled_local_argmax_ptr
386 + + req_idx * resampled_local_argmax_stride
387 + + block_idx,
388 + token_id,
389 + )
390 + tl.store(
391 + resampled_local_max_ptr + req_idx * resampled_local_max_stride + block_idx,
392 + value,
393 + )
394 +
395 +
396 +@triton.jit
397 +def _insert_resampled_kernel(
398 + # [num_reqs, num_speculative_steps + 1]
399 + sampled_ptr,
400 + sampled_stride,
401 + # [num_reqs]
402 + num_sampled_ptr,
403 + # [num_reqs, num_blocks]
404 + resampled_local_argmax_ptr,
405 + resampled_local_argmax_stride,
406 + # [num_reqs, num_blocks]
407 + resampled_local_max_ptr,
408 + resampled_local_max_stride,
409 + resample_num_blocks,
410 + # [num_reqs + 1]
411 + cu_num_logits_ptr,
412 + # [num_reqs]
413 + expanded_idx_mapping_ptr,
414 + # [max_num_reqs]
415 + temp_ptr,
416 + PADDED_RESAMPLE_NUM_BLOCKS: tl.constexpr,
417 +):
418 + req_idx = tl.program_id(0)
419 + num_sampled = tl.load(num_sampled_ptr + req_idx)
420 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
421 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
422 + resample_token_idx = start_idx + num_sampled
423 + req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx)
424 +
425 + # Increment the number of sampled tokens.
426 + tl.store(num_sampled_ptr + req_idx, num_sampled + 1)
427 +
428 + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
429 + is_bonus = resample_token_idx == end_idx - 1
430 + if temp == 0.0 and not is_bonus:
431 + # Greedy + non-bonus token. The target argmax is already
432 + # in the sampled tensor.
433 + return
434 +
435 + # Insert the resampled token.
436 + block = tl.arange(0, PADDED_RESAMPLE_NUM_BLOCKS)
437 + mask = block < resample_num_blocks
438 + resampled_local_max = tl.load(
439 + resampled_local_max_ptr + req_idx * resampled_local_max_stride + block,
440 + mask=mask,
441 + other=float("-inf"),
442 + )
443 + resampled_max_block_idx = tl.argmax(resampled_local_max, axis=0)
444 + resampled = tl.load(
445 + resampled_local_argmax_ptr
446 + + req_idx * resampled_local_argmax_stride
447 + + resampled_max_block_idx,
448 + )
449 + tl.store(
450 + sampled_ptr + req_idx * sampled_stride + num_sampled,
451 + resampled,
452 + )
453 +
454 +
455 +def probabilistic_rejection_sample(
456 + # [num_logits, V]
457 + target_logits: torch.Tensor,
458 + # [max_num_reqs, num_speculative_steps, V]
459 + draft_logits: torch.Tensor,
460 + # [num_logits]
461 + draft_sampled: torch.Tensor,
462 + # [num_reqs + 1]
463 + cu_num_logits: torch.Tensor,
464 + # [num_logits]
465 + pos: torch.Tensor,
466 + # [num_reqs]
467 + idx_mapping: torch.Tensor,
468 + # [num_logits]
469 + expanded_idx_mapping: torch.Tensor,
470 + # [num_logits]
471 + expanded_local_pos: torch.Tensor,
472 + # [max_num_reqs]
473 + temperature: torch.Tensor,
474 + # [max_num_reqs]
475 + seed: torch.Tensor,
476 + num_speculative_steps: int,
477 +) -> tuple[torch.Tensor, torch.Tensor]:
478 + num_reqs = cu_num_logits.shape[0] - 1
479 + num_logits, vocab_size = target_logits.shape
480 +
481 + # Gather draft logits, compute target argmax for greedy, and
482 + # compute per-block LSE and max for non-greedy requests.
483 + VOCAB_BLOCK_SIZE = 8192
484 + vocab_num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
485 + padded_vocab_num_blocks = triton.next_power_of_2(vocab_num_blocks)
486 + target_local_argmax = target_logits.new_empty(
487 + num_logits, vocab_num_blocks, dtype=torch.int64
488 + )
489 + target_local_max = target_logits.new_empty(
490 + num_logits, vocab_num_blocks, dtype=torch.float32
491 + )
492 + target_local_sumexp = target_logits.new_empty(
493 + num_logits, vocab_num_blocks, dtype=torch.float32
494 + )
495 + draft_local_max = target_logits.new_empty(
496 + num_logits, vocab_num_blocks, dtype=torch.float32
497 + )
498 + draft_local_sumexp = target_logits.new_empty(
499 + num_logits, vocab_num_blocks, dtype=torch.float32
500 + )
501 + _compute_block_max_and_sumexp_kernel[(num_logits, vocab_num_blocks)](
502 + target_local_argmax,
503 + target_local_argmax.stride(0),
504 + target_local_max,
505 + target_local_max.stride(0),
506 + target_local_sumexp,
507 + target_local_sumexp.stride(0),
508 + draft_local_max,
509 + draft_local_max.stride(0),
510 + draft_local_sumexp,
511 + draft_local_sumexp.stride(0),
512 + target_logits,
513 + target_logits.stride(0),
514 + draft_logits,
515 + draft_logits.stride(0),
516 + draft_logits.stride(1),
517 + expanded_idx_mapping,
518 + expanded_local_pos,
519 + temperature,
520 + vocab_size,
521 + num_speculative_steps,
522 + BLOCK_SIZE=VOCAB_BLOCK_SIZE,
523 + )
524 +
525 + # Sample up until the first rejected/bonus token, and store
526 + # the step.
527 + sampled = draft_sampled.new_empty(
528 + num_reqs, num_speculative_steps + 1, dtype=torch.int64
529 + )
530 + num_sampled = sampled.new_empty(num_reqs)
531 + target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
532 + draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
533 + _probabilistic_rejection_kernel[(num_reqs,)](
534 + sampled,
535 + sampled.stride(0),
536 + num_sampled,
537 + target_rejected_logsumexp,
538 + draft_rejected_logsumexp,
539 + target_logits,
540 + target_logits.stride(0),
541 + target_local_argmax,
542 + target_local_argmax.stride(0),
543 + target_local_max,
544 + target_local_max.stride(0),
545 + target_local_sumexp,
546 + target_local_sumexp.stride(0),
547 + draft_sampled,
548 + draft_logits,
549 + draft_logits.stride(0),
550 + draft_logits.stride(1),
551 + draft_local_max,
552 + draft_local_max.stride(0),
553 + draft_local_sumexp,
554 + draft_local_sumexp.stride(0),
555 + cu_num_logits,
556 + idx_mapping,
557 + temperature,
558 + seed,
559 + pos,
560 + vocab_num_blocks,
561 + PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
562 + num_warps=1,
563 + )
564 +
565 + # Resample the rejected/bonus tokens.
566 + RESAMPLE_BLOCK_SIZE = 1024
567 + resample_num_blocks = triton.cdiv(vocab_size, RESAMPLE_BLOCK_SIZE)
568 + padded_resample_num_blocks = triton.next_power_of_2(resample_num_blocks)
569 + resampled_local_argmax = target_logits.new_empty(
570 + num_reqs, resample_num_blocks, dtype=torch.int64
571 + )
572 + resampled_local_max = target_logits.new_empty(
573 + num_reqs, resample_num_blocks, dtype=torch.float64
574 + )
575 + _resample_kernel[(num_reqs, resample_num_blocks)](
576 + resampled_local_argmax,
577 + resampled_local_argmax.stride(0),
578 + resampled_local_max,
579 + resampled_local_max.stride(0),
580 + target_logits,
581 + target_logits.stride(0),
582 + target_rejected_logsumexp,
583 + draft_logits,
584 + draft_logits.stride(0),
585 + draft_logits.stride(1),
586 + draft_rejected_logsumexp,
587 + num_sampled,
588 + cu_num_logits,
589 + expanded_idx_mapping,
590 + temperature,
591 + seed,
592 + pos,
593 + vocab_size,
594 + BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
595 + )
596 +
597 + # Insert the resampled tokens into the output sampled.
598 + _insert_resampled_kernel[(num_reqs,)](
599 + sampled,
600 + sampled.stride(0),
601 + num_sampled,
602 + resampled_local_argmax,
603 + resampled_local_argmax.stride(0),
604 + resampled_local_max,
605 + resampled_local_max.stride(0),
606 + resample_num_blocks,
607 + cu_num_logits,
608 + expanded_idx_mapping,
609 + temperature,
610 + PADDED_RESAMPLE_NUM_BLOCKS=padded_resample_num_blocks,
611 + )
612 + return sampled, num_sampled
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+3 -352
@@ -7,11 +7,13 @@ from vllm.triton_utils import tl, triton
7 from vllm.v1.outputs import LogprobsTensors
8 from vllm.v1.worker.gpu.input_batch import InputBatch
9 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
10 -from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample, tl_rand64
10 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
11 from vllm.v1.worker.gpu.sample.output import SamplerOutput
12 from vllm.v1.worker.gpu.sample.sampler import Sampler
13 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
14 +from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
15 + probabilistic_rejection_sample,
16 +)
17 from vllm.v1.worker.gpu.spec_decode.synthetic_rejection_sampler_utils import (
18 compute_synthetic_rejection_sampler_params,
19 synthetic_rejection_sample,
@@ -75,357 +77,6 @@ def strict_rejection_sample(
77 return sampled, num_sampled
78
79
78 -@triton.jit
79 -def _gather_draft_logits_and_target_argmax_kernel(
80 - local_target_argmax_ptr,
81 - local_target_argmax_stride,
82 - local_target_max_ptr,
83 - local_target_max_stride,
84 - # [num_logits, V]
85 - out_draft_logits_ptr,
86 - out_draft_logits_stride,
87 - # [num_logits, V]
88 - target_logits_ptr,
89 - target_logits_stride,
90 - # [max_num_reqs, num_speculative_steps, V]
91 - draft_logits_ptr,
92 - draft_logits_stride_0,
93 - draft_logits_stride_1,
94 - # [num_logits]
95 - expanded_idx_mapping_ptr,
96 - # [num_logits]
97 - expanded_local_pos_ptr,
98 - # [max_num_reqs]
99 - temp_ptr,
100 - vocab_size,
101 - num_speculative_steps,
102 - BLOCK_SIZE: tl.constexpr,
103 -):
104 - logit_idx = tl.program_id(0)
105 - req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
106 - draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
107 -
108 - block_idx = tl.program_id(1)
109 - block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
110 - mask = block_offsets < vocab_size
111 - temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
112 -
113 - if temp == 0.0:
114 - # Greedy sampling. Get the target logits argmax.
115 - target_logits = tl.load(
116 - target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
117 - mask=mask,
118 - other=float("-inf"),
119 - ).to(tl.float32)
120 - value, idx = tl.max(target_logits, axis=0, return_indices=True)
121 - token_id = block_idx * BLOCK_SIZE + idx
122 - tl.store(
123 - local_target_argmax_ptr
124 - + logit_idx * local_target_argmax_stride
125 - + block_idx,
126 - token_id,
127 - )
128 - tl.store(
129 - local_target_max_ptr + logit_idx * local_target_max_stride + block_idx,
130 - value,
131 - )
132 - elif draft_step_idx < num_speculative_steps:
133 - draft_logits = tl.load(
134 - draft_logits_ptr
135 - + req_state_idx * draft_logits_stride_0
136 - + draft_step_idx * draft_logits_stride_1
137 - + block_offsets,
138 - mask=mask,
139 - other=float("-inf"),
140 - ).to(tl.float32)
141 - tl.store(
142 - out_draft_logits_ptr + logit_idx * out_draft_logits_stride + block_offsets,
143 - draft_logits,
144 - mask=mask,
145 - )
146 -
147 -
148 -@triton.jit
149 -def _probabilistic_rejection_kernel(
150 - # [num_reqs, num_speculative_steps + 1]
151 - sampled_ptr,
152 - sampled_stride,
153 - # [num_reqs]
154 - rejected_steps_ptr,
155 - # [num_reqs]
156 - rejected_pos_ptr,
157 - # [num_logits]
158 - draft_sampled_ptr,
159 - # [num_logits, V]
160 - target_probs_ptr,
161 - target_probs_stride,
162 - # [num_logits, V]
163 - draft_probs_ptr,
164 - draft_probs_stride,
165 - # [num_logits, num_blocks]
166 - local_target_argmax_ptr,
167 - local_target_argmax_stride,
168 - # [num_logits, num_blocks]
169 - local_target_max_ptr,
170 - local_target_max_stride,
171 - # [num_reqs + 1]
172 - cu_num_logits_ptr,
173 - # [num_logits]
174 - pos_ptr,
175 - # [num_reqs]
176 - idx_mapping_ptr,
177 - # [max_num_reqs]
178 - temp_ptr,
179 - # [max_num_reqs]
180 - seeds_ptr,
181 - NUM_BLOCKS: tl.constexpr,
182 - PADDED_NUM_BLOCKS: tl.constexpr,
183 -):
184 - req_idx = tl.program_id(0)
185 - start_idx = tl.load(cu_num_logits_ptr + req_idx)
186 - num_tokens = tl.load(cu_num_logits_ptr + req_idx + 1) - start_idx
187 - req_state_idx = tl.load(idx_mapping_ptr + req_idx)
188 - seed = tl.load(seeds_ptr + req_state_idx)
189 - temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
190 -
191 - rejected_step = 0
192 - accepted = True
193 - for i in range(num_tokens - 1):
194 - if accepted:
195 - logit_idx = start_idx + i
196 - draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1)
197 - if temp == 0.0:
198 - # Greedy sampling. Only accept the sampled draft token if
199 - # it exactly matches the target argmax.
200 - block_offsets = tl.arange(0, PADDED_NUM_BLOCKS)
201 - block_mask = block_offsets < NUM_BLOCKS
202 - local_max = tl.load(
203 - local_target_max_ptr
204 - + logit_idx * local_target_max_stride
205 - + block_offsets,
206 - mask=block_mask,
207 - other=float("-inf"),
208 - )
209 - max_block = tl.argmax(local_max, axis=0)
210 - target_argmax = tl.load(
211 - local_target_argmax_ptr
212 - + logit_idx * local_target_argmax_stride
213 - + max_block
214 - )
215 - accepted &= target_argmax == draft_sampled
216 - else:
217 - target_prob = tl.load(
218 - target_probs_ptr + logit_idx * target_probs_stride + draft_sampled
219 - ).to(tl.float64)
220 - draft_prob = tl.load(
221 - draft_probs_ptr + logit_idx * draft_probs_stride + draft_sampled
222 - ).to(tl.float64)
223 - pos = tl.load(pos_ptr + logit_idx)
224 - u = tl_rand64(seed, pos, includes_zero=False)
225 - accepted &= target_prob > u * draft_prob
226 - tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
227 - rejected_step += accepted
228 - tl.store(rejected_steps_ptr + req_idx, rejected_step)
229 - pos_val = tl.load(pos_ptr + start_idx + rejected_step)
230 - tl.store(rejected_pos_ptr + req_idx, pos_val)
231 -
232 -
233 -@triton.jit
234 -def _compute_residual_logits_kernel(
235 - # [num_reqs, V]
236 - residual_logits_ptr,
237 - residual_logits_stride,
238 - # [num_logits, V]
239 - target_probs_ptr,
240 - target_probs_stride,
241 - # [num_logits, V]
242 - draft_probs_ptr,
243 - draft_probs_stride,
244 - # [num_logits, V]
245 - target_logits_ptr,
246 - target_logits_stride,
247 - # [num_reqs]
248 - rejected_step_ptr,
249 - # [num_reqs + 1]
250 - cu_num_logits_ptr,
251 - # [num_reqs]
252 - idx_mapping_ptr,
253 - # [max_num_reqs]
254 - temp_ptr,
255 - vocab_size,
256 - BLOCK_SIZE: tl.constexpr,
257 -):
258 - req_idx = tl.program_id(0)
259 - block_idx = tl.program_id(1)
260 -
261 - req_state_idx = tl.load(idx_mapping_ptr + req_idx)
262 - start_idx = tl.load(cu_num_logits_ptr + req_idx)
263 - end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
264 - rejected_logit_idx = start_idx + tl.load(rejected_step_ptr + req_idx)
265 - temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
266 - block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
267 - mask = block_offsets < vocab_size
268 -
269 - if temp == 0.0 or (rejected_logit_idx == end_idx - 1):
270 - # Greedy sampling / bonus token. In either case, use the
271 - # target logits directly to reduce numerical error.
272 - residual_logits = tl.load(
273 - target_logits_ptr
274 - + rejected_logit_idx * target_logits_stride
275 - + block_offsets,
276 - mask=mask,
277 - other=float("-inf"),
278 - )
279 - else:
280 - target_probs = tl.load(
281 - target_probs_ptr + rejected_logit_idx * target_probs_stride + block_offsets,
282 - mask=mask,
283 - other=0.0,
284 - )
285 - draft_probs = tl.load(
286 - draft_probs_ptr + rejected_logit_idx * draft_probs_stride + block_offsets,
287 - mask=mask,
288 - other=0.0,
289 - )
290 - residual_probs = tl.maximum(target_probs - draft_probs, 0.0)
291 - residual_logits = tl.log(residual_probs)
292 -
293 - tl.store(
294 - residual_logits_ptr + req_idx * residual_logits_stride + block_offsets,
295 - residual_logits,
296 - mask=mask,
297 - )
298 -
299 -
300 -def probabilistic_rejection_sample(
301 - # [num_logits, V]
302 - target_logits: torch.Tensor,
303 - # [max_num_reqs, num_speculative_steps, V]
304 - draft_logits: torch.Tensor,
305 - # [num_logits]
306 - draft_sampled: torch.Tensor,
307 - # [num_reqs + 1]
308 - cu_num_logits: torch.Tensor,
309 - # [num_logits]
310 - pos: torch.Tensor,
311 - # [num_reqs]
312 - idx_mapping: torch.Tensor,
313 - # [num_logits]
314 - expanded_idx_mapping: torch.Tensor,
315 - # [num_logits]
316 - expanded_local_pos: torch.Tensor,
317 - # [max_num_reqs]
318 - temperature: torch.Tensor,
319 - # [max_num_reqs]
320 - seed: torch.Tensor,
321 - num_speculative_steps: int,
322 -) -> tuple[torch.Tensor, torch.Tensor]:
323 - num_reqs = cu_num_logits.shape[0] - 1
324 - num_logits, vocab_size = target_logits.shape
325 -
326 - BLOCK_SIZE = 1024
327 - num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
328 -
329 - # Gather draft logits and target argmax for greedy sampling.
330 - gathered_draft_logits = target_logits.new_empty(target_logits.shape)
331 - local_target_argmax = target_logits.new_empty(
332 - num_logits, num_blocks, dtype=torch.int64
333 - )
334 - local_target_max = target_logits.new_empty(
335 - num_logits, num_blocks, dtype=torch.float32
336 - )
337 - _gather_draft_logits_and_target_argmax_kernel[(num_logits, num_blocks)](
338 - local_target_argmax,
339 - local_target_argmax.stride(0),
340 - local_target_max,
341 - local_target_max.stride(0),
342 - gathered_draft_logits,
343 - gathered_draft_logits.stride(0),
344 - target_logits,
345 - target_logits.stride(0),
346 - draft_logits,
347 - draft_logits.stride(0),
348 - draft_logits.stride(1),
349 - expanded_idx_mapping,
350 - expanded_local_pos,
351 - temperature,
352 - vocab_size,
353 - num_speculative_steps,
354 - BLOCK_SIZE=BLOCK_SIZE,
355 - )
356 -
357 - # Compute target and draft probs.
358 - target_probs = torch.softmax(target_logits, dim=-1)
359 - draft_probs = torch.softmax(gathered_draft_logits, dim=-1)
360 -
361 - # Rejection sample.
362 - # [num_reqs, num_speculative_steps + 1]
363 - sampled = draft_sampled.new_empty(
364 - num_reqs, num_speculative_steps + 1, dtype=torch.int64
365 - )
366 - # [num_reqs]
367 - rejected_steps = sampled.new_empty(num_reqs)
368 - # [num_reqs]
369 - rejected_pos = pos.new_empty(num_reqs)
370 - _probabilistic_rejection_kernel[(num_reqs,)](
371 - sampled,
372 - sampled.stride(0),
373 - rejected_steps,
374 - rejected_pos,
375 - draft_sampled,
376 - target_probs,
377 - target_probs.stride(0),
378 - draft_probs,
379 - draft_probs.stride(0),
380 - local_target_argmax,
381 - local_target_argmax.stride(0),
382 - local_target_max,
383 - local_target_max.stride(0),
384 - cu_num_logits,
385 - pos,
386 - idx_mapping,
387 - temperature,
388 - seed,
389 - num_warps=1,
390 - NUM_BLOCKS=num_blocks,
391 - PADDED_NUM_BLOCKS=triton.next_power_of_2(num_blocks),
392 - )
393 -
394 - # Compute the logits and positions to resample the rejected/bonus
395 - # tokens from.
396 - # [num_reqs, vocab_size]
397 - residual_logits = target_logits.new_empty(num_reqs, vocab_size)
398 - _compute_residual_logits_kernel[(num_reqs, num_blocks)](
399 - residual_logits,
400 - residual_logits.stride(0),
401 - target_probs,
402 - target_probs.stride(0),
403 - draft_probs,
404 - draft_probs.stride(0),
405 - target_logits,
406 - target_logits.stride(0),
407 - rejected_steps,
408 - cu_num_logits,
409 - idx_mapping,
410 - temperature,
411 - vocab_size,
412 - BLOCK_SIZE=BLOCK_SIZE,
413 - )
414 -
415 - # Gumbel sample tokens from the residual distribution.
416 - resampled = gumbel_sample(
417 - residual_logits,
418 - idx_mapping,
419 - temperature,
420 - seed,
421 - rejected_pos,
422 - apply_temperature=False,
423 - )
424 - sampled.scatter_(1, rejected_steps.unsqueeze(1), resampled.unsqueeze(1))
425 -
426 - return sampled, rejected_steps + 1
427 -
428 -
80 @triton.jit
81 def _flatten_sampled_kernel(
82 # [num_logits]