| 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 | ) |