| 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(do_not_specialize=["num_reqs"]) |
| 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(do_not_specialize=["num_tokens"]) |
| 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(do_not_specialize=["num_tokens"]) |
| 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=max(256, triton.next_power_of_2(num_reqs)), |
| 426 | num_warps=1 if num_reqs <= 256 else 4, |
| 427 | ) |
| 428 | |
| 429 | self._steps_since_refit += 1 |
| 430 | if self._steps_since_refit < self.REFIT_INTERVAL: |
| 431 | return |
| 432 | self._steps_since_refit = 0 |
| 433 | |
| 434 | # Fit the coefficients to the accumulated statistics gathered over the |
| 435 | # course of the last REFIT_INTERVAL steps. |
| 436 | _refit_kernel[(1,)]( |
| 437 | self.slope, |
| 438 | self.intercepts, |
| 439 | self.info, |
| 440 | self.info.stride(0), |
| 441 | self.grad, |
| 442 | self.grad.stride(0), |
| 443 | self.counts, |
| 444 | NUM_SPECULATIVE_STEPS=self.num_speculative_steps, |
| 445 | L2=self.L2, |
| 446 | DAMPING=self.DAMPING_OBSERVATIONS, |
| 447 | BLOCK=triton.next_power_of_2(self.num_speculative_steps), |
| 448 | ) |
| 449 | self._refits += 1 |
| 450 | |
| 451 | def predict( |
| 452 | self, |
| 453 | logits: torch.Tensor, |
| 454 | idx_mapping: torch.Tensor, |
| 455 | draft_step: torch.Tensor, |
| 456 | confidence_probs: torch.Tensor, |
| 457 | temperature: torch.Tensor, |
| 458 | ) -> None: |
| 459 | num_tokens, vocab_size = logits.shape |
| 460 | VOCAB_BLOCK_SIZE = 4096 |
| 461 | num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE) |
| 462 | local_max = torch.empty( |
| 463 | num_tokens, num_blocks, dtype=torch.float32, device=self.device |
| 464 | ) |
| 465 | local_sumexp = torch.empty_like(local_max) |
| 466 | _local_max_sumexp_kernel[(num_tokens, num_blocks)]( |
| 467 | local_max, |
| 468 | local_max.stride(0), |
| 469 | local_sumexp, |
| 470 | local_sumexp.stride(0), |
| 471 | logits, |
| 472 | logits.stride(0), |
| 473 | idx_mapping, |
| 474 | idx_mapping.stride(0), |
| 475 | temperature, |
| 476 | num_tokens, |
| 477 | vocab_size, |
| 478 | BLOCK_SIZE=VOCAB_BLOCK_SIZE, |
| 479 | ) |
| 480 | _predict_kernel[(num_tokens,)]( |
| 481 | self.features, |
| 482 | self.features.stride(0), |
| 483 | self.predictions, |
| 484 | self.predictions.stride(0), |
| 485 | confidence_probs, |
| 486 | confidence_probs.stride(0), |
| 487 | self.slope, |
| 488 | self.intercepts, |
| 489 | local_max, |
| 490 | local_max.stride(0), |
| 491 | local_sumexp, |
| 492 | local_sumexp.stride(0), |
| 493 | idx_mapping, |
| 494 | idx_mapping.stride(0), |
| 495 | draft_step, |
| 496 | num_tokens, |
| 497 | num_blocks, |
| 498 | per_token_step=draft_step.dim() > 0, |
| 499 | NUM_SPECULATIVE_STEPS=self.num_speculative_steps, |
| 500 | MAX_LOG_ODDS=_MAX_LOG_ODDS, |
| 501 | PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks), |
| 502 | ) |