@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3bb0a03f35

[Model Runner V2] Acceptance estimation for adaptive verification (#52228)

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

Giancarlo Delfin committed Sep 14, 2026 at 15:58 UTC 3bb0a03f35269185b6358753e64932bb946cc531
9 files changed +900 -28
tests/v1/spec_decode/test_acceptance_estimator.py new
+294
@@ -0,0 +1,294 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""``predict`` writes one probability into two index spaces at once.
4 +
5 +Adaptive verification reads the estimate in this step's batch order, while the
6 +next step's refit reads it in persistent request-state slots. A mix-up between
7 +the two is silent -- every value stays a plausible probability, it just belongs
8 +to another request -- so it is asserted here rather than left to acceptance
9 +metrics to hint at.
10 +"""
11 +
12 +import pytest
13 +import torch
14 +
15 +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
16 + OnlineAcceptanceEstimator,
17 +)
18 +
19 +pytest.importorskip("triton")
20 +if not torch.cuda.is_available():
21 + pytest.skip("CUDA required for acceptance estimator tests", allow_module_level=True)
22 +
23 +MAX_NUM_REQS = 16
24 +NUM_STEPS = 4
25 +VOCAB_SIZE = 2048
26 +
27 +
28 +def _make_estimator(device: torch.device) -> OnlineAcceptanceEstimator:
29 + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device)
30 + # Distinct per-position intercepts, so reading the wrong one shows up.
31 + estimator.slope.fill_(0.7)
32 + estimator.intercepts[:] = torch.tensor([1.5, 1.0, 0.5, 0.0], device=device)
33 + return estimator
34 +
35 +
36 +def _expected_probs(
37 + estimator: OnlineAcceptanceEstimator,
38 + logits: torch.Tensor,
39 + steps: torch.Tensor,
40 + temperature: torch.Tensor,
41 +) -> torch.Tensor:
42 + """logit(q_max) by a different route than the kernel takes.
43 +
44 + The kernel shifts by the max and drops the unit term; here the argmax is
45 + masked out and the complement is reduced with logsumexp, so agreement is
46 + evidence about the value rather than about a shared formula.
47 + """
48 + temp = torch.where(temperature > 0, temperature, torch.ones_like(temperature))
49 + scaled = logits.double() / temp[:, None].double()
50 + top = scaled.max(dim=-1)
51 + rest = scaled.scatter(1, top.indices[:, None], float("-inf"))
52 + feature = (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0)
53 + weight, bias = estimator.slope, estimator.intercepts
54 + return torch.sigmoid(weight.double() * feature + bias[steps].double()).float()
55 +
56 +
57 +@pytest.mark.parametrize("tokens_per_req", [1, NUM_STEPS])
58 +def test_predict_agrees_across_slot_and_batch_order(tokens_per_req: int):
59 + """One row per request (per-step drafting) and a whole block at once.
60 +
61 + Which one it is is inferred from whether the draft step is per-row, so the
62 + two cases differ only in how ``draft_step`` is shaped.
63 + """
64 + device = torch.device("cuda")
65 + torch.manual_seed(0)
66 + estimator = _make_estimator(device)
67 + num_reqs = 5
68 +
69 + # Batch position i holds a shuffled persistent slot, as a served batch would.
70 + slots = torch.randperm(MAX_NUM_REQS, device=device)[:num_reqs].to(torch.int32)
71 + if tokens_per_req == 1:
72 + idx_mapping = slots
73 + draft_step = torch.tensor(2, device=device, dtype=torch.int32)
74 + steps = draft_step.expand(num_reqs)
75 + else:
76 + idx_mapping = slots.repeat_interleave(NUM_STEPS)
77 + draft_step = torch.arange(NUM_STEPS, device=device, dtype=torch.int32).repeat(
78 + num_reqs
79 + )
80 + steps = draft_step
81 + batch_rows = torch.arange(num_reqs, device=device).repeat_interleave(tokens_per_req)
82 + logits = torch.randn(idx_mapping.shape[0], VOCAB_SIZE, device=device) * 4.0
83 +
84 + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
85 + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5
86 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
87 +
88 + expected = _expected_probs(
89 + estimator, logits, steps.long(), temperature[idx_mapping.long()]
90 + )
91 + torch.testing.assert_close(
92 + estimator.predictions[idx_mapping.long(), steps.long()],
93 + expected,
94 + atol=2e-3,
95 + rtol=2e-3,
96 + )
97 + torch.testing.assert_close(
98 + confidence_probs[batch_rows, steps.long()], expected, atol=2e-3, rtol=2e-3
99 + )
100 +
101 +
102 +def test_predict_skips_cudagraph_padded_rows():
103 + """Padded rows carry slot -1 and must not land on a live request."""
104 + device = torch.device("cuda")
105 + torch.manual_seed(0)
106 + estimator = _make_estimator(device)
107 + num_reqs, num_padded = 3, 6
108 +
109 + slots = torch.tensor([7, 2, 11], device=device, dtype=torch.int32)
110 + idx_mapping = torch.cat(
111 + [slots, torch.full((num_padded,), -1, device=device, dtype=torch.int32)]
112 + )
113 + draft_step = torch.tensor(1, device=device, dtype=torch.int32)
114 + logits = torch.randn(num_reqs + num_padded, VOCAB_SIZE, device=device) * 4.0
115 +
116 + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
117 + temperature = torch.rand(MAX_NUM_REQS, device=device) + 0.5
118 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
119 +
120 + live = (slots.long(), torch.full((num_reqs,), 1, device=device))
121 + expected = _expected_probs(
122 + estimator, logits[:num_reqs], live[1], temperature[slots.long()]
123 + )
124 + torch.testing.assert_close(
125 + estimator.predictions[live], expected, atol=2e-3, rtol=2e-3
126 + )
127 + # Nothing outside the live (request, position) cells was touched.
128 + untouched = torch.ones(MAX_NUM_REQS, NUM_STEPS, dtype=torch.bool, device=device)
129 + untouched[live] = False
130 + assert not estimator.predictions[untouched].any()
131 + assert not confidence_probs[num_reqs:].any()
132 +
133 +
134 +# A refit takes one undamped Newton step, so it needs a realistic number of
135 +# observations per round to be well conditioned -- a handful of requests can
136 +# send the step far past the optimum. This is the batch a served step sees.
137 +REFIT_NUM_REQS = 512
138 +
139 +
140 +def _drive_round(estimator, features, num_accepted, slots):
141 + """Feed one step's graded drafts through predict's buffers into `step`."""
142 + estimator.features[slots.long()] = features
143 + estimator.predictions[slots.long()] = torch.sigmoid(
144 + estimator.slope * features + estimator.intercepts
145 + )
146 + estimator.step(
147 + slots,
148 + (num_accepted + 1).to(torch.int32),
149 + (features.shape[1] - num_accepted).to(torch.int32),
150 + )
151 +
152 +
153 +def test_refit_recovers_a_shared_slope_and_per_position_intercepts():
154 + """The fit must find the logistic that generated the observations.
155 +
156 + One slope is shared across draft positions and each keeps its own intercept,
157 + so the arrowhead solve is only correct if it recovers both from data whose
158 + positions differ in level but not in slope.
159 + """
160 + device = torch.device("cuda")
161 + torch.manual_seed(0)
162 + true_slope = 0.6
163 + true_bias = torch.tensor([1.4, 0.9, 0.4, -0.1], device=device)
164 +
165 + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
166 + estimator.REFIT_INTERVAL = 5
167 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
168 + for _ in range(60):
169 + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0
170 + accepted = torch.rand_like(features) < torch.sigmoid(
171 + true_slope * features + true_bias
172 + )
173 + # Verification stops at the first rejection, so a position's label is
174 + # conditional on every shallower one having been accepted.
175 + num_accepted = (~accepted).float().argmax(dim=1)
176 + num_accepted[accepted.all(dim=1)] = NUM_STEPS
177 + _drive_round(estimator, features, num_accepted, slots)
178 +
179 + slope, bias = estimator.slope, estimator.intercepts
180 + torch.testing.assert_close(
181 + slope[0], torch.tensor(true_slope, device=device), atol=0.1, rtol=0
182 + )
183 + torch.testing.assert_close(bias, true_bias, atol=0.3, rtol=0)
184 +
185 +
186 +def test_refit_survives_a_position_with_no_observations():
187 + """A draft position nobody reached must not stall the shared slope.
188 +
189 + Its statistics stay zero while the others accumulate, which is the degenerate
190 + case the arrowhead denominator has to tolerate.
191 + """
192 + device = torch.device("cuda")
193 + torch.manual_seed(0)
194 + estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
195 + estimator.REFIT_INTERVAL = 5
196 + slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
197 + before_slope = estimator.slope.clone()
198 + before_intercepts = estimator.intercepts.clone()
199 +
200 + # Every request is rejected at position 1, so positions 2+ are never graded.
201 + num_accepted = torch.ones(REFIT_NUM_REQS, device=device)
202 + for _ in range(20):
203 + features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device) * 8.0
204 + _drive_round(estimator, features, num_accepted, slots)
205 +
206 + assert estimator.slope.isfinite().all()
207 + assert estimator.intercepts.isfinite().all()
208 + assert not torch.equal(estimator.slope, before_slope), "slope never moved"
209 + # Positions 2 and 3 were never observed, so their intercepts must be untouched.
210 + torch.testing.assert_close(
211 + estimator.intercepts[2:], before_intercepts[2:], rtol=0, atol=0
212 + )
213 +
214 +
215 +def test_step_damping_scales_with_the_evidence():
216 + """A round moves a coefficient in proportion to the samples behind it.
217 +
218 + The Newton step itself barely changes with how much data a round carries --
219 + curvature and score both scale with it -- so scaling a round's statistics
220 + isolates the damping factor n / (n + DAMPING_OBSERVATIONS). The slope/
221 + intercept coupling is zeroed so the two parameters can be checked apart.
222 + """
223 + device = torch.device("cuda")
224 + damping = OnlineAcceptanceEstimator.DAMPING_OBSERVATIONS
225 + per_position_n = 20.0
226 +
227 + def refit_with(scale):
228 + estimator = OnlineAcceptanceEstimator(MAX_NUM_REQS, NUM_STEPS, device)
229 + estimator.info[:, 1] = 0.0 # sum(w*margin): no slope/intercept coupling
230 + estimator.info[:, 2] = 1.0 * scale # sum(w)
231 + estimator.grad[:, 1] = 0.5 * scale # sum(y - p)
232 + # The slope's information and score are accumulated as per-position
233 + # partials that the refit sums, so spread each total over the positions.
234 + estimator.info[:, 0] = 30.0 * scale / NUM_STEPS # sum(w*margin^2)
235 + estimator.grad[:, 0] = 2.0 * scale / NUM_STEPS # sum((y - p)*margin)
236 + estimator.counts[:] = per_position_n * scale
237 + estimator._steps_since_refit = estimator.REFIT_INTERVAL - 1
238 + before_slope = estimator.slope.clone()
239 + before_intercepts = estimator.intercepts.clone()
240 + empty = torch.zeros(0, dtype=torch.int32, device=device)
241 + estimator.step(empty, empty, empty)
242 + return (
243 + (estimator.slope - before_slope)[0].item(),
244 + (estimator.intercepts - before_intercepts)[0].item(),
245 + )
246 +
247 + for scale in (1.0, 10.0, 100.0):
248 + total_n = per_position_n * scale * NUM_STEPS
249 + slope_delta, bias_delta = refit_with(scale)
250 + assert slope_delta == pytest.approx(
251 + (2.0 / 30.0) * total_n / (total_n + damping), rel=1e-3
252 + )
253 + assert bias_delta == pytest.approx(
254 + 0.5 * (per_position_n * scale) / (per_position_n * scale + damping),
255 + rel=1e-3,
256 + )
257 +
258 + # A round with no observations at all must leave the coefficients alone.
259 + assert refit_with(0.0) == (0.0, 0.0)
260 +
261 +
262 +def test_feature_ignores_which_token_was_drawn():
263 + """The score must depend on the draft distribution, not on the draw.
264 +
265 + This is what keeps rejection sampling lossless: if the score read the drawn
266 + token, trimming a draft would condition on the token being trimmed and tilt
267 + the emitted distribution toward the drafter's confident modes. ``predict``
268 + takes no token argument, and this pins that -- two rows with identical
269 + logits must score identically no matter what was sampled from them.
270 + """
271 + device = torch.device("cuda")
272 + torch.manual_seed(0)
273 + estimator = _make_estimator(device)
274 +
275 + row = torch.randn(VOCAB_SIZE, device=device) * 4.0
276 + logits = row.expand(2, VOCAB_SIZE).contiguous()
277 + idx_mapping = torch.tensor([0, 1], device=device, dtype=torch.int32)
278 + draft_step = torch.tensor(0, device=device, dtype=torch.int32)
279 + temperature = torch.ones(MAX_NUM_REQS, device=device)
280 + confidence_probs = torch.zeros(MAX_NUM_REQS, NUM_STEPS, device=device)
281 +
282 + estimator.predict(logits, idx_mapping, draft_step, confidence_probs, temperature)
283 + assert estimator.features[0, 0] == estimator.features[1, 0]
284 +
285 + # And the value is the top token's log-odds, not the top-2 margin.
286 + scaled = row.double()
287 + top = scaled.max(dim=-1)
288 + rest = scaled.scatter(0, top.indices[None], float("-inf"))
289 + torch.testing.assert_close(
290 + estimator.features[0, 0].double(),
291 + (top.values - torch.logsumexp(rest, dim=-1)).clamp(-40.0, 40.0),
292 + atol=2e-3,
293 + rtol=2e-3,
294 + )
tests/v1/worker/test_gpu_autoregressive_speculator.py
+5
@@ -123,6 +123,7 @@ def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected)
123 draft_model_config=draft_model_config,
124 use_local_argmax_reduction=False,
125 draft_sample_method="greedy",
126 + enable_adaptive_verification=False,
127 )
128 vllm_config = SimpleNamespace(
129 speculative_config=speculative_config,
@@ -161,6 +162,7 @@ def test_mm_support_configured_after_model_load(monkeypatch):
162 speculator.dtype = torch.float32
163 speculator.draft_model_config = draft_model_config
164 speculator.supports_mm_inputs = False
165 + speculator.use_acceptance_estimator = False
166
167 checked_configs = []
168
@@ -195,6 +197,7 @@ def test_load_model_keeps_mm_support_for_capable_drafter(monkeypatch):
197 speculator = object.__new__(_TestSpeculator)
198 speculator.supports_mm_inputs = False
199 speculator.inputs_embeds = None
200 + speculator.use_acceptance_estimator = False
201 speculator.vllm_config = SimpleNamespace(model_config=object())
202 speculator.max_num_tokens = 4
203 speculator.hidden_size = 3
@@ -219,6 +222,7 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch):
222 speculator = object.__new__(_TestSpeculator)
223 speculator.supports_mm_inputs = False
224 speculator.inputs_embeds = None
225 + speculator.use_acceptance_estimator = False
226 speculator.vllm_config = SimpleNamespace(model_config=object())
227 draft_model = _TextOnlyDraftModel()
228 speculator.test_draft_model = draft_model
@@ -251,6 +255,7 @@ def test_multi_module_mm_support_configured_after_model_load(monkeypatch):
255 speculator = object.__new__(MultiModuleMTPSpeculator)
256 speculator.supports_mm_inputs = False
257 speculator.inputs_embeds = None
258 + speculator.use_acceptance_estimator = False
259 speculator.cached_draft_input_embeds = None
260 speculator.vllm_config = SimpleNamespace(model_config=object())
261 speculator.max_num_tokens = 4
tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
+2
@@ -60,6 +60,7 @@ def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch):
60 layer_name = "cache_only_layers.2"
61 speculator = object.__new__(ExtractHiddenStatesSpeculator)
62 speculator.vllm_config = cast(Any, SimpleNamespace())
63 + speculator.acceptance_estimator = None
64 speculator.num_hidden_states = 2
65 speculator.hidden_states = torch.zeros(4, 2, 3)
66 speculator.draft_attn_layer_names = {layer_name}
@@ -111,6 +112,7 @@ def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch):
112
113 def test_propose_requires_aux_hidden_states():
114 speculator = object.__new__(ExtractHiddenStatesSpeculator)
115 + speculator.acceptance_estimator = None
116 speculator.num_hidden_states = 2
117 input_batch = cast(
118 Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32))
vllm/config/speculative.py
+10 -2
@@ -1561,8 +1561,16 @@ class SpeculativeConfig:
1561 self.index_share_for_mtp_iteration
1562 )
1563
1564 - if self.method != "dspark" and self.enable_adaptive_verification:
1565 - raise ValueError("Adaptive verification only supported with DSpark")
1564 + if (
1565 + self.method != "dspark"
1566 + and self.enable_adaptive_verification
1567 + and self.use_local_argmax_reduction
1568 + ):
1569 + raise ValueError(
1570 + "Adaptive verification estimates per-position acceptance from "
1571 + "the draft logits, which use_local_argmax_reduction never "
1572 + "materializes. Disable one of them."
1573 + )
1574
1575 return self
1576
vllm/v1/worker/gpu/model_runner.py
+4
@@ -2120,6 +2120,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2120
2121 if self.speculator is not None:
2122 assert self.sampler is not None
2123 + if isinstance(self.speculator, DraftModelSpeculator):
2124 + self.speculator.observe_verification(
2125 + input_batch.idx_mapping, num_sampled, num_rejected
2126 + )
2127 # Let the target override the hidden state fed to the drafter
2128 # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
2129 # target returns a persistent buffer sized at max_num_batched_tokens;
vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py new
+501
@@ -0,0 +1,501 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.triton_utils import tl, triton
7 +from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
8 + _compute_max_and_sumexp,
9 +)
10 +
11 +# The feature is a log-odds, log(q / (1 - q)), which a row carrying all of its
12 +# mass on one token sends to infinity. Clamp it symmetrically.
13 +_MAX_LOG_ODDS = 40.0
14 +
15 +# Cold-start coefficients: the median slope and intercept fitted online across
16 +# DeepSeek-V4-Flash, Muse-Glimmer-30B, MiMo-V2.5-Pro, Inkling and Kimi-K2.5.
17 +# Log-loss is flat near the optimum, so one starting point serves all of them
18 +# and a real fit replaces it within a refit or two.
19 +_INIT_SLOPE = 0.5
20 +_INIT_BIAS = -0.2
21 +
22 +
23 +@triton.jit
24 +def _accumulate_kernel(
25 + info_ptr,
26 + grad_ptr,
27 + idx_mapping_ptr,
28 + num_sampled_ptr,
29 + num_rejected_ptr,
30 + features_ptr,
31 + features_stride,
32 + pred_ptr,
33 + pred_stride,
34 + counts_ptr,
35 + num_reqs,
36 + BLOCK_R: tl.constexpr,
37 +):
38 + step = tl.program_id(0).to(tl.int64)
39 + req_block = tl.arange(0, BLOCK_R)
40 + req_mask = req_block < num_reqs
41 + req_state_block = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0).to(
42 + tl.int64
43 + )
44 +
45 + num_sampled = tl.load(num_sampled_ptr + req_block, mask=req_mask, other=0).to(
46 + tl.int64
47 + )
48 + # Subtract bonus token.
49 + num_accepted = tl.maximum(num_sampled - 1, 0)
50 + num_rejected = tl.load(num_rejected_ptr + req_block, mask=req_mask, other=0).to(
51 + tl.int64
52 + )
53 + num_admitted = num_accepted + num_rejected
54 +
55 + feature = tl.load(
56 + features_ptr + req_state_block * features_stride + step,
57 + mask=req_mask,
58 + other=0.0,
59 + ).to(tl.float32)
60 + pred = tl.load(
61 + pred_ptr + req_state_block * pred_stride + step,
62 + mask=req_mask,
63 + other=0.0,
64 + ).to(tl.float32)
65 +
66 + # The accepted tokens and the first rejected draft token (if one exists) are
67 + # "observed" and given labels.
68 + observed = req_mask & (step <= num_accepted) & (step < num_admitted)
69 + label = tl.where(step < num_accepted, 1.0, 0.0)
70 + w = tl.where(observed, pred * (1.0 - pred), 0.0)
71 + resid = tl.where(observed, label - pred, 0.0)
72 +
73 + # This block's contribution to the normal equations for design row
74 + # x = [feature, eₖ]: the XᵀWX entries aₖ, bₖ, cₖ and the Xᵀr entries g₀ₖ, g₁ₖ.
75 + # aₖ and g₀ₖ are partials of scalars shared by every position, which
76 + # _refit_kernel sums.
77 + a_k = tl.sum(w * feature * feature, axis=0)
78 + b_k = tl.sum(w * feature, axis=0)
79 + c_k = tl.sum(w, axis=0)
80 + g0_k = tl.sum(resid * feature, axis=0)
81 + g1_k = tl.sum(resid, axis=0)
82 + count = tl.sum(tl.where(observed, 1.0, 0.0), axis=0)
83 + tl.store(info_ptr + step * 3 + 0, tl.load(info_ptr + step * 3 + 0) + a_k)
84 + tl.store(info_ptr + step * 3 + 1, tl.load(info_ptr + step * 3 + 1) + b_k)
85 + tl.store(info_ptr + step * 3 + 2, tl.load(info_ptr + step * 3 + 2) + c_k)
86 + tl.store(grad_ptr + step * 2 + 0, tl.load(grad_ptr + step * 2 + 0) + g0_k)
87 + tl.store(grad_ptr + step * 2 + 1, tl.load(grad_ptr + step * 2 + 1) + g1_k)
88 + tl.store(counts_ptr + step, tl.load(counts_ptr + step) + count)
89 +
90 +
91 +@triton.jit
92 +def _refit_kernel(
93 + slope_ptr,
94 + intercepts_ptr,
95 + info_ptr,
96 + info_row_stride,
97 + grad_ptr,
98 + grad_row_stride,
99 + counts_ptr,
100 + NUM_SPECULATIVE_STEPS: tl.constexpr,
101 + L2: tl.constexpr,
102 + DAMPING: tl.constexpr,
103 + BLOCK: tl.constexpr,
104 +):
105 + k = tl.arange(0, BLOCK)
106 + mask = k < NUM_SPECULATIVE_STEPS
107 +
108 + # Load the stats accumulated over this round to solve (XᵀWX + λI) Δθ = Xᵀr.
109 + a_k = tl.load(info_ptr + k * info_row_stride + 0, mask=mask, other=0.0)
110 + b_k = tl.load(info_ptr + k * info_row_stride + 1, mask=mask, other=0.0)
111 + c_k = tl.load(info_ptr + k * info_row_stride + 2, mask=mask, other=0.0) + L2
112 + a = tl.sum(a_k, axis=0) + L2
113 + g0_k = tl.load(grad_ptr + k * grad_row_stride + 0, mask=mask, other=0.0)
114 + g1_k = tl.load(grad_ptr + k * grad_row_stride + 1, mask=mask, other=0.0)
115 + g0 = tl.sum(g0_k, axis=0)
116 + n = tl.load(counts_ptr + k, mask=mask, other=0.0)
117 +
118 + # Load the current shared slope and per-position intercepts.
119 + w = tl.load(slope_ptr)
120 + bias = tl.load(intercepts_ptr + k, mask=mask, other=0.0)
121 +
122 + # Solve for the shared slope delta, Δω, in the following arrowhead system:
123 + # ⎡ a+λ b₀ b₁ … bₙ₋₁ ⎤ ⎡ Δω ⎤ ⎡ g₀ ⎤
124 + # ⎢ b₀ c₀+λ ⎥ ⎢ Δβ₀ ⎥ ⎢ g₁₀ ⎥
125 + # ⎢ b₁ c₁+λ ⎥ ⎢ Δβ₁ ⎥ = ⎢ g₁₁ ⎥
126 + # ⎢ ⋮ ⋱ ⎥ ⎢ ⋮ ⎥ ⎢ ⋮ ⎥
127 + # ⎣ bₙ₋₁ cₙ₋₁+λ ⎦ ⎣ Δβₙ₋₁ ⎦ ⎣ g₁ₙ₋₁ ⎦
128 + # XᵀWX + λI Δθ Xᵀ(Y-P)
129 + shrink = b_k * b_k / c_k
130 + coupling = b_k * g1_k / c_k
131 + step_w = (g0 - tl.sum(coupling, axis=0)) / (a - tl.sum(shrink, axis=0))
132 +
133 + # The slope learns from every position, so it is damped by the round's
134 + # total sample count.
135 + total_n = tl.sum(n, axis=0)
136 + step_w *= total_n / (total_n + DAMPING)
137 + # Mask out NaNs and steps that would drive the slope negative.
138 + step_w = tl.where((step_w == step_w) & (w + step_w >= 0.0), step_w, 0.0)
139 + # Update the shared slope.
140 + new_w = w + step_w
141 +
142 + # Solve for the per-position intercept deltas, Δβₖ, in the same system.
143 + step_bias = (g1_k - b_k * step_w) / c_k
144 + # Damp by the position's sample count.
145 + step_bias *= n / (n + DAMPING)
146 + # Mask out NaNs.
147 + step_bias = tl.where(step_bias == step_bias, step_bias, 0.0)
148 + # Update the per-position intercepts.
149 + new_bias = bias + step_bias
150 +
151 + tl.store(slope_ptr, new_w)
152 + tl.store(intercepts_ptr + k, new_bias, mask=mask)
153 +
154 + # Start the round clean by clearing the per-round statistics.
155 + tl.store(info_ptr + k * info_row_stride + 0, 0.0, mask=mask)
156 + tl.store(info_ptr + k * info_row_stride + 1, 0.0, mask=mask)
157 + tl.store(info_ptr + k * info_row_stride + 2, 0.0, mask=mask)
158 + tl.store(grad_ptr + k * grad_row_stride + 0, 0.0, mask=mask)
159 + tl.store(grad_ptr + k * grad_row_stride + 1, 0.0, mask=mask)
160 + tl.store(counts_ptr + k, 0.0, mask=mask)
161 +
162 +
163 +@triton.jit
164 +def _local_max_sumexp_kernel(
165 + local_max_ptr,
166 + local_max_stride,
167 + local_sumexp_ptr,
168 + local_sumexp_stride,
169 + logits_ptr,
170 + logits_stride,
171 + idx_mapping_ptr,
172 + idx_mapping_stride,
173 + temperature_ptr,
174 + num_tokens,
175 + vocab_size,
176 + BLOCK_SIZE: tl.constexpr,
177 +):
178 + token_idx = tl.program_id(0).to(tl.int64)
179 + if token_idx >= num_tokens:
180 + return
181 +
182 + block_idx = tl.program_id(1)
183 + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to(
184 + tl.int64
185 + )
186 + if req_state_idx < 0:
187 + # Cudagraph-padded requests carry -1. Write an empty partial so the
188 + # buffer holds nothing uninitialized; _predict_kernel skips these rows.
189 + tl.store(
190 + local_max_ptr + token_idx * local_max_stride + block_idx, float("-inf")
191 + )
192 + tl.store(local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, 0.0)
193 + return
194 +
195 + # Draft logits are pre-temperature, and acceptance is decided on the scaled
196 + # distribution the drafter actually sampled from, so scale first. A request
197 + # with temp == 0 drafts greedily, so leave its logits alone rather than dividing
198 + # by zero, matching what the sampler does.
199 + temp = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
200 + temp = tl.where(temp > 0.0, temp, 1.0)
201 +
202 + offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
203 + mask = offsets < vocab_size
204 + logits = (
205 + tl.load(
206 + logits_ptr + token_idx * logits_stride + offsets,
207 + mask=mask,
208 + other=float("-inf"),
209 + ).to(tl.float32)
210 + / temp
211 + )
212 + block_max, block_sumexp = _compute_max_and_sumexp(logits)
213 + tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, block_max)
214 + tl.store(
215 + local_sumexp_ptr + token_idx * local_sumexp_stride + block_idx, block_sumexp
216 + )
217 +
218 +
219 +@triton.jit
220 +def _predict_kernel(
221 + features_ptr,
222 + features_stride,
223 + pred_ptr,
224 + pred_stride,
225 + conf_ptr,
226 + conf_stride,
227 + slope_ptr,
228 + intercepts_ptr,
229 + local_max_ptr,
230 + local_max_stride,
231 + local_sumexp_ptr,
232 + local_sumexp_stride,
233 + idx_mapping_ptr,
234 + idx_mapping_stride,
235 + step_ptr,
236 + num_tokens,
237 + vocab_num_blocks,
238 + per_token_step: tl.constexpr,
239 + NUM_SPECULATIVE_STEPS: tl.constexpr,
240 + MAX_LOG_ODDS: tl.constexpr,
241 + PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
242 +):
243 + token_idx = tl.program_id(0).to(tl.int64)
244 + if token_idx >= num_tokens:
245 + return
246 +
247 + req_state_idx = tl.load(idx_mapping_ptr + token_idx * idx_mapping_stride).to(
248 + tl.int64
249 + )
250 + if req_state_idx < 0:
251 + # Cudagraph-padded requests carry -1. Skip them so that they don't
252 + # scatter garbage over a live request's features.
253 + return
254 +
255 + # A per-row step means the speculator sampled the whole block in one pass, so
256 + # each request owns NUM_SPECULATIVE_STEPS consecutive rows; a step shared by
257 + # every row means one row per request. Batch position, unlike req_state_idx,
258 + # is only meaningful for the step that computed it, and is what adaptive
259 + # verification indexes by.
260 + if per_token_step:
261 + step = tl.load(step_ptr + token_idx).to(tl.int64)
262 + batch_idx = token_idx // NUM_SPECULATIVE_STEPS
263 + else:
264 + step = tl.load(step_ptr).to(tl.int64)
265 + batch_idx = token_idx
266 +
267 + blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
268 + blocks_mask = blocks < vocab_num_blocks
269 + maxes = tl.load(
270 + local_max_ptr + token_idx * local_max_stride + blocks,
271 + mask=blocks_mask,
272 + other=float("-inf"),
273 + )
274 + sumexps = tl.load(
275 + local_sumexp_ptr + token_idx * local_sumexp_stride + blocks,
276 + mask=blocks_mask,
277 + other=0.0,
278 + )
279 +
280 + # The feature, logit(q(m*)) for the row's max draft logit, m*, is:
281 + # q(x) = e^(x - m*) / Σ_z e^(z - m*)
282 + # logit(q(m*)) = log(q(m*) / (1 - q(m*)))
283 + # = m* - log(Σ_{z≠m*} e^(z - m*)).
284 + # It scores the draft distribution, not the token drawn from it, so it is
285 + # a function of the prefix alone. Trimming a draft therefore never
286 + # conditions on the token being trimmed, and rejection sampling stays
287 + # lossless. Scoring the drawn token instead would select on the proposal
288 + # and bias the emitted distribution toward the drafter's confident modes.
289 +
290 + # Rescale each block's partial to the global max and sum. The most likely
291 + # token contributes exactly e^(m - m) = 1, so removing it leaves the sum
292 + # over every other token.
293 + global_max = tl.max(maxes, axis=0)
294 + total_sumexp = tl.sum(
295 + tl.where(blocks_mask, sumexps * tl.exp(maxes - global_max), 0.0), axis=0
296 + )
297 + complement_sumexp = tl.maximum(total_sumexp - 1.0, 0.0)
298 +
299 + # Shifted by m*, the numerator is e^(m* - m*) = 1, so the feature is just the
300 + # negated complement log-sum-exponential. Clamp to the min/max log-odds.
301 + feature = -tl.log(complement_sumexp)
302 + feature = tl.minimum(tl.maximum(feature, -MAX_LOG_ODDS), MAX_LOG_ODDS)
303 + tl.store(features_ptr + req_state_idx * features_stride + step, feature)
304 +
305 + # Predict the acceptance probability from the feature and current coefficients.
306 + weight = tl.load(slope_ptr)
307 + bias = tl.load(intercepts_ptr + step)
308 + prob = tl.sigmoid(weight * feature + bias)
309 + tl.store(pred_ptr + req_state_idx * pred_stride + step, prob)
310 + tl.store(conf_ptr + batch_idx * conf_stride + step, prob)
311 +
312 +
313 +class OnlineAcceptanceEstimator:
314 + """Predicts per-position acceptance, and calibrates itself while serving.
315 +
316 + Lifecycle per step, driven by ``DraftModelSpeculator``:
317 +
318 + 1. ``step`` folds the previous step's drafts, now graded by the target,
319 + into the IRLS accumulators, and periodically solves for new coefficients,
320 + independently but identically on every rank.
321 + 2. ``predict`` runs inside the captured draft graph, turning this step's
322 + draft logits into acceptance probabilities for adaptive verification.
323 +
324 + Trimming starts immediately: survival is a running product of sigmoids, so it
325 + decreases with draft position whatever the coefficients are, and an unfitted
326 + estimator degrades to uniform-depth truncation rather than to anything harmful.
327 + """
328 +
329 + # Refit every this many steps, accumulating samples in between.
330 + REFIT_INTERVAL = 100
331 + # Newton steps are damped by n / (n + DAMPING_OBSERVATIONS), so a round
332 + # carrying this many observations moves a parameter half of a full step. It
333 + # replaces a hard minimum-sample gate: a data-poor position keeps learning,
334 + # just slowly, instead of freezing until it crosses a threshold. Lower values
335 + # track a drifting workload faster, higher ones are steadier on thin data.
336 + DAMPING_OBSERVATIONS = 50.0
337 + # Ridge on the Newton solve.
338 + L2 = 1e-3
339 +
340 + def __init__(
341 + self,
342 + max_num_reqs: int,
343 + num_speculative_steps: int,
344 + device: torch.device,
345 + ):
346 + self.num_speculative_steps = num_speculative_steps
347 + self.device = device
348 + self._steps_since_refit = 0
349 + self._refits = 0
350 +
351 + # The learned coefficients used for predicting acceptance.
352 + self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device)
353 + self.intercepts = torch.full(
354 + (num_speculative_steps,), _INIT_BIAS, dtype=torch.float32, device=device
355 + )
356 + # Holds logit(max q), which is used as the feature for the logistic.
357 + # Stored in stable slots keyed by persistent request-state index.
358 + self.features = torch.zeros(
359 + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
360 + )
361 + # Holds the acceptance predictions made at draft time, computed as:
362 + # p = σ(ω * f + βₖ)
363 + # where f is feature (logit(max q)) at draft position k, ω is the learned
364 + # shared slope, and βₖ is the learned per-position intercept. Stored in
365 + # stable slots keyed by the persistent request-state index.
366 + self.predictions = torch.zeros(
367 + max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
368 + )
369 +
370 + # Per-round Newton-IRLS statistics, cleared after each refit. For observation
371 + # i the design row is xᵢ = [fᵢ, eₖ], the IRLS weight is wᵢ = pᵢ(1-pᵢ), and the
372 + # residual is yᵢ - pᵢ, where yᵢ is the label (1 for accepted, 0 for rejected).
373 + # XᵀWX is then the negative Hessian of the Bernoulli log-likelihood loss and
374 + # Xᵀ(Y-P) is its gradient, so the Newton step for θ = [ω, β₀ … βₙ₋₁] is the
375 + # the solution of (XᵀWX + λI) Δθ = Xᵀ(Y-P).
376 + #
377 + # Columns [aₖ, bₖ, cₖ] are each position's contribution to the XᵀWX arrowhead:
378 + # ⎡ a b₀ b₁ … bₙ₋₁ ⎤
379 + # ⎢ b₀ c₀ ⎥
380 + # ⎢ b₁ c₁ ⎥
381 + # ⎢ ⋮ ⋱ ⎥
382 + # ⎣ bₙ₋₁ cₙ₋₁ ⎦
383 + # bₖ and cₖ are per-position, but a = Σₖ aₖ is one scalar shared by all of
384 + # them.
385 + self.info = torch.zeros(
386 + num_speculative_steps, 3, dtype=torch.float32, device=device
387 + )
388 + # Columns [g₀ₖ, g₁ₖ] are each position's contribution to the Xᵀ(Y-P) vector:
389 + # ⎡ g₀ ⎤
390 + # ⎢ g₁₀ ⎥
391 + # ⎢ g₁₁ ⎥
392 + # ⎢ ⋮ ⎥
393 + # ⎣ g₁ₙ₋₁ ⎦
394 + # As above, g₁ₖ is per-position and g₀ = Σₖ g₀ₖ is shared.
395 + self.grad = torch.zeros(
396 + num_speculative_steps, 2, dtype=torch.float32, device=device
397 + )
398 + # Number of graded drafts per-position in the current round, zeroed after
399 + # each refit. _refit_kernel uses this to damp coefficient updates by:
400 + # n / (n + DAMPING_OBSERVATIONS)
401 + self.counts = torch.zeros(
402 + num_speculative_steps, dtype=torch.float32, device=device
403 + )
404 +
405 + def step(
406 + self,
407 + idx_mapping: torch.Tensor,
408 + num_sampled: torch.Tensor,
409 + num_rejected: torch.Tensor,
410 + ) -> None:
411 + # Accumulate the previous step's graded drafts into the IRLS statistics.
412 + num_reqs = idx_mapping.shape[0]
413 + _accumulate_kernel[(self.num_speculative_steps,)](
414 + self.info,
415 + self.grad,
416 + idx_mapping,
417 + num_sampled,
418 + num_rejected,
419 + self.features,
420 + self.features.stride(0),
421 + self.predictions,
422 + self.predictions.stride(0),
423 + self.counts,
424 + num_reqs,
425 + BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)),
426 + )
427 +
428 + self._steps_since_refit += 1
429 + if self._steps_since_refit < self.REFIT_INTERVAL:
430 + return
431 + self._steps_since_refit = 0
432 +
433 + # Fit the coefficients to the accumulated statistics gathered over the
434 + # course of the last REFIT_INTERVAL steps.
435 + _refit_kernel[(1,)](
436 + self.slope,
437 + self.intercepts,
438 + self.info,
439 + self.info.stride(0),
440 + self.grad,
441 + self.grad.stride(0),
442 + self.counts,
443 + NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
444 + L2=self.L2,
445 + DAMPING=self.DAMPING_OBSERVATIONS,
446 + BLOCK=triton.next_power_of_2(self.num_speculative_steps),
447 + )
448 + self._refits += 1
449 +
450 + def predict(
451 + self,
452 + logits: torch.Tensor,
453 + idx_mapping: torch.Tensor,
454 + draft_step: torch.Tensor,
455 + confidence_probs: torch.Tensor,
456 + temperature: torch.Tensor,
457 + ) -> None:
458 + num_tokens, vocab_size = logits.shape
459 + VOCAB_BLOCK_SIZE = 4096
460 + num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
461 + local_max = torch.empty(
462 + num_tokens, num_blocks, dtype=torch.float32, device=self.device
463 + )
464 + local_sumexp = torch.empty_like(local_max)
465 + _local_max_sumexp_kernel[(num_tokens, num_blocks)](
466 + local_max,
467 + local_max.stride(0),
468 + local_sumexp,
469 + local_sumexp.stride(0),
470 + logits,
471 + logits.stride(0),
472 + idx_mapping,
473 + idx_mapping.stride(0),
474 + temperature,
475 + num_tokens,
476 + vocab_size,
477 + BLOCK_SIZE=VOCAB_BLOCK_SIZE,
478 + )
479 + _predict_kernel[(num_tokens,)](
480 + self.features,
481 + self.features.stride(0),
482 + self.predictions,
483 + self.predictions.stride(0),
484 + confidence_probs,
485 + confidence_probs.stride(0),
486 + self.slope,
487 + self.intercepts,
488 + local_max,
489 + local_max.stride(0),
490 + local_sumexp,
491 + local_sumexp.stride(0),
492 + idx_mapping,
493 + idx_mapping.stride(0),
494 + draft_step,
495 + num_tokens,
496 + num_blocks,
497 + per_token_step=draft_step.dim() > 0,
498 + NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
499 + MAX_LOG_ODDS=_MAX_LOG_ODDS,
500 + PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks),
501 + )
vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py
+6
@@ -213,5 +213,11 @@ class DFlash2Speculator(DFlashSpeculator):
213 anchor_token_ids,
214 )
215 self._sample_path(candidate_ids, scores, num_reqs)
216 + if self.enable_adaptive_verification:
217 + self._maybe_predict_acceptance(
218 + self._selector_scores[:num_reqs].flatten(0, 1),
219 + self.sample_idx_mapping[:num_sample],
220 + self.sample_col[:num_sample],
221 + )
222 if self.draft_logits is not None:
223 self._cache_draft_logits(candidate_ids, num_sample)
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+16 -18
@@ -79,12 +79,7 @@ class DSparkSpeculator(DFlashSpeculator):
79 self.draft_model_config.hf_config, "dspark_draft_topk", None
80 )
81
82 - self.draft_token_confidence_probs = torch.empty_like(
83 - self.draft_tokens, dtype=torch.float32
84 - )
85 - self.enable_adaptive_verification = (
86 - self.speculative_config.enable_adaptive_verification
87 - )
82 + self.use_confidence_head: bool = False
83
84 def load_draft_model(
85 self,
@@ -108,13 +103,14 @@ class DSparkSpeculator(DFlashSpeculator):
103 dtype=self.draft_logits.dtype,
104 device=self.device,
105 )
111 - if self.enable_adaptive_verification and model.model.confidence_head is None:
112 - raise ValueError(
113 - "Adaptive verification needs a DSpark checkpoint with a confidence "
114 - "head, and this one has none. Pass "
115 - "enable_adaptive_verification=false in the speculative config to verify"
116 - " a fixed number of drafts instead."
117 - )
106 + self.use_confidence_head = (
107 + self.enable_adaptive_verification
108 + and model.model.confidence_head is not None
109 + )
110 + if self.use_confidence_head:
111 + # The acceptance estimator is not needed when a trained confidence head
112 + # is available.
113 + self.use_acceptance_estimator = False
114 return model
115
116 def _sample_logits(
@@ -124,8 +120,10 @@ class DSparkSpeculator(DFlashSpeculator):
120 sample_pos: torch.Tensor,
121 step: int,
122 ) -> torch.Tensor:
123 + self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step])
124 if self.draft_logits is None:
128 - return self.model.map_draft_to_target(logits.argmax(dim=-1))
125 + draft_ids = logits.argmax(dim=-1)
126 + return self.model.map_draft_to_target(draft_ids)
127
128 # Probabilistic sampling and rejection operate in target-vocabulary
129 # space. A reduced draft vocabulary is scattered into its target rows.
@@ -184,7 +182,7 @@ class DSparkSpeculator(DFlashSpeculator):
182 for i in range(n_spec):
183 # Sequential stage: Markov bias from the previously sampled token.
184 markov_embed = self.model.markov_embed(prev)
187 - if self.enable_adaptive_verification:
185 + if self.use_confidence_head:
186 confidence_markov_embeds.append(markov_embed)
187 bias = self.model.markov_bias(markov_embed)
188 logits_i = base_logits[:, i] + bias
@@ -194,7 +192,7 @@ class DSparkSpeculator(DFlashSpeculator):
192 self.draft_tokens[:num_reqs, i] = draft_sampled_i
193 prev = draft_sampled_i
194
197 - if self.enable_adaptive_verification:
195 + if self.use_confidence_head:
196 confidence = self.model.compute_confidence(
197 sample_hidden,
198 torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),
@@ -229,7 +227,7 @@ class DSparkSpeculator(DFlashSpeculator):
227
228 for i in range(n_spec):
229 markov_embed = self.model.markov_embed(prev)
232 - if self.enable_adaptive_verification:
230 + if self.use_confidence_head:
231 confidence_markov_embeds.append(markov_embed)
232 logits_i = self.model.apply_markov_bias_gathered(
233 markov_embed,
@@ -243,7 +241,7 @@ class DSparkSpeculator(DFlashSpeculator):
241 self.draft_tokens[:num_reqs, i] = draft_sampled_i
242 prev = draft_sampled_i
243
246 - if self.enable_adaptive_verification:
244 + if self.use_confidence_head:
245 confidence = self.model.compute_confidence(
246 sample_hidden,
247 torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1),
vllm/v1/worker/gpu/spec_decode/speculator.py
+62 -8
@@ -33,6 +33,9 @@ from vllm.v1.worker.gpu.dp_utils import DPSyncState
33 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
34 from vllm.v1.worker.gpu.model_states.interface import ModelState
35 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
36 +from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
37 + OnlineAcceptanceEstimator,
38 +)
39 from vllm.v1.worker.utils import AttentionGroup
40
41 if TYPE_CHECKING:
@@ -158,6 +161,15 @@ class DraftModelSpeculator(BaseSpeculator):
161 self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32)
162 self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool)
163
164 + self.draft_token_confidence_probs = torch.empty_like(
165 + self.draft_tokens, dtype=torch.float32
166 + )
167 + self.enable_adaptive_verification = (
168 + self.speculative_config.enable_adaptive_verification
169 + )
170 + self.use_acceptance_estimator = self.enable_adaptive_verification
171 + self.acceptance_estimator: OnlineAcceptanceEstimator | None = None
172 +
173 self.draft_logits: torch.Tensor | None = None
174 if self.speculative_config.draft_sample_method == "probabilistic":
175 # Pre-temperature logits, cached from the previous decode step.
@@ -227,6 +239,20 @@ class DraftModelSpeculator(BaseSpeculator):
239 type(self.model).__name__,
240 )
241
242 + if self.use_acceptance_estimator:
243 + if self.use_local_argmax_reduction:
244 + raise ValueError(
245 + "Adaptive verification without a confidence head estimates "
246 + "per-position acceptance from the draft logits, which "
247 + "use_local_argmax_reduction never materializes. Disable one "
248 + "of them."
249 + )
250 + self.acceptance_estimator = OnlineAcceptanceEstimator(
251 + self.max_num_reqs,
252 + self.num_speculative_steps,
253 + self.device,
254 + )
255 +
256 def set_eplb_state(self, eplb_state: EplbState) -> None:
257 """Inject EPLB state after construction."""
258 self.eplb_state = eplb_state
@@ -371,12 +397,6 @@ class DraftModelSpeculator(BaseSpeculator):
397 "(communication: O(2*tp_size) vs O(vocab_size))."
398 )
399
374 - def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor:
375 - if self.use_local_argmax_reduction:
376 - return self.model.get_top_tokens(hidden_states)
377 - logits = self.model.compute_logits(hidden_states)
378 - return logits.argmax(dim=-1)
379 -
400 def sample_draft(
401 self,
402 hidden_states: torch.Tensor,
@@ -408,8 +428,42 @@ class DraftModelSpeculator(BaseSpeculator):
428 idx_mapping,
429 temperature,
430 )
411 - return sampled
412 - return self._greedy_sample_draft(hidden_states)
431 + elif self.use_local_argmax_reduction:
432 + return self.model.get_top_tokens(hidden_states)
433 + else:
434 + logits = self.model.compute_logits(hidden_states)
435 + sampled = logits.argmax(dim=-1)
436 + self._maybe_predict_acceptance(logits, idx_mapping, draft_step)
437 + return sampled
438 +
439 + def _maybe_predict_acceptance(
440 + self,
441 + logits: torch.Tensor,
442 + idx_mapping: torch.Tensor,
443 + draft_step: torch.Tensor,
444 + ) -> None:
445 + if self.acceptance_estimator is not None:
446 + self.acceptance_estimator.predict(
447 + logits,
448 + idx_mapping,
449 + draft_step,
450 + self.draft_token_confidence_probs,
451 + self.temperature,
452 + )
453 +
454 + def observe_verification(
455 + self,
456 + idx_mapping: torch.Tensor,
457 + num_sampled: torch.Tensor,
458 + num_rejected: torch.Tensor,
459 + ) -> None:
460 + """Fold the target's verdict on the last drafts into the estimator.
461 +
462 + Must run before the next `propose`, which overwrites the per-slot
463 + features the verdict grades.
464 + """
465 + if self.acceptance_estimator is not None:
466 + self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected)
467
468 def prepare_watermarking(
469 self, contexts: torch.Tensor, watermarking: torch.Tensor