class OnlineAcceptanceEstimator:
"""Predicts per-position acceptance, and calibrates itself while serving.
Lifecycle per step, driven by ``DraftModelSpeculator``:
1. ``step`` folds the previous step's drafts, now graded by the target,
into the IRLS accumulators, and periodically solves for new coefficients,
independently but identically on every rank.
2. ``predict`` runs inside the captured draft graph, turning this step's
draft logits into acceptance probabilities for adaptive verification.
Trimming starts immediately: survival is a running product of sigmoids, so it
decreases with draft position whatever the coefficients are, and an unfitted
estimator degrades to uniform-depth truncation rather than to anything harmful.
"""
# Refit every this many steps, accumulating samples in between.
REFIT_INTERVAL = 100
# Newton steps are damped by n / (n + DAMPING_OBSERVATIONS), so a round
# carrying this many observations moves a parameter half of a full step. It
# replaces a hard minimum-sample gate: a data-poor position keeps learning,
# just slowly, instead of freezing until it crosses a threshold. Lower values
# track a drifting workload faster, higher ones are steadier on thin data.
DAMPING_OBSERVATIONS = 50.0
# Ridge on the Newton solve.
L2 = 1e-3
def __init__(
self,
max_num_reqs: int,
num_speculative_steps: int,
device: torch.device,
):
self.num_speculative_steps = num_speculative_steps
self.device = device
self._steps_since_refit = 0
self._refits = 0
# The learned coefficients used for predicting acceptance.
self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device)
self.intercepts = torch.full(
(num_speculative_steps,), _INIT_BIAS, dtype=torch.float32, device=device
)
# Holds logit(max q), which is used as the feature for the logistic.
# Stored in stable slots keyed by persistent request-state index.
self.features = torch.zeros(
max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
)
# Holds the acceptance predictions made at draft time, computed as:
# p = σ(ω * f + βₖ)
# where f is feature (logit(max q)) at draft position k, ω is the learned
# shared slope, and βₖ is the learned per-position intercept. Stored in
# stable slots keyed by the persistent request-state index.
self.predictions = torch.zeros(
max_num_reqs, num_speculative_steps, dtype=torch.float32, device=device
)
# Per-round Newton-IRLS statistics, cleared after each refit. For observation
# i the design row is xᵢ = [fᵢ, eₖ], the IRLS weight is wᵢ = pᵢ(1-pᵢ), and the
# residual is yᵢ - pᵢ, where yᵢ is the label (1 for accepted, 0 for rejected).
# XᵀWX is then the negative Hessian of the Bernoulli log-likelihood loss and
# Xᵀ(Y-P) is its gradient, so the Newton step for θ = [ω, β₀ … βₙ₋₁] is the
# the solution of (XᵀWX + λI) Δθ = Xᵀ(Y-P).
#
# Columns [aₖ, bₖ, cₖ] are each position's contribution to the XᵀWX arrowhead:
# ⎡ a b₀ b₁ … bₙ₋₁ ⎤
# ⎢ b₀ c₀ ⎥
# ⎢ b₁ c₁ ⎥
# ⎢ ⋮ ⋱ ⎥
# ⎣ bₙ₋₁ cₙ₋₁ ⎦
# bₖ and cₖ are per-position, but a = Σₖ aₖ is one scalar shared by all of
# them.
self.info = torch.zeros(
num_speculative_steps, 3, dtype=torch.float32, device=device
)
# Columns [g₀ₖ, g₁ₖ] are each position's contribution to the Xᵀ(Y-P) vector:
# ⎡ g₀ ⎤
# ⎢ g₁₀ ⎥
# ⎢ g₁₁ ⎥
# ⎢ ⋮ ⎥
# ⎣ g₁ₙ₋₁ ⎦
# As above, g₁ₖ is per-position and g₀ = Σₖ g₀ₖ is shared.
self.grad = torch.zeros(
num_speculative_steps, 2, dtype=torch.float32, device=device
)
# Number of graded drafts per-position in the current round, zeroed after
# each refit. _refit_kernel uses this to damp coefficient updates by:
# n / (n + DAMPING_OBSERVATIONS)
self.counts = torch.zeros(
num_speculative_steps, dtype=torch.float32, device=device
)
def step(
self,
idx_mapping: torch.Tensor,
num_sampled: torch.Tensor,
num_rejected: torch.Tensor,
) -> None:
# Accumulate the previous step's graded drafts into the IRLS statistics.
num_reqs = idx_mapping.shape[0]
_accumulate_kernel[(self.num_speculative_steps,)](
self.info,
self.grad,
idx_mapping,
num_sampled,
num_rejected,
self.features,
self.features.stride(0),
self.predictions,
self.predictions.stride(0),
self.counts,
num_reqs,
BLOCK_R=triton.next_power_of_2(max(num_reqs, 1)),
)
self._steps_since_refit += 1
if self._steps_since_refit < self.REFIT_INTERVAL:
return
self._steps_since_refit = 0
# Fit the coefficients to the accumulated statistics gathered over the
# course of the last REFIT_INTERVAL steps.
_refit_kernel[(1,)](
self.slope,
self.intercepts,
self.info,
self.info.stride(0),
self.grad,
self.grad.stride(0),
self.counts,
NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
L2=self.L2,
DAMPING=self.DAMPING_OBSERVATIONS,
BLOCK=triton.next_power_of_2(self.num_speculative_steps),
)
self._refits += 1
def predict(
self,
logits: torch.Tensor,
idx_mapping: torch.Tensor,
draft_step: torch.Tensor,
confidence_probs: torch.Tensor,
temperature: torch.Tensor,
) -> None:
num_tokens, vocab_size = logits.shape
VOCAB_BLOCK_SIZE = 4096
num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
local_max = torch.empty(
num_tokens, num_blocks, dtype=torch.float32, device=self.device
)
local_sumexp = torch.empty_like(local_max)
_local_max_sumexp_kernel[(num_tokens, num_blocks)](
local_max,
local_max.stride(0),
local_sumexp,
local_sumexp.stride(0),
logits,
logits.stride(0),
idx_mapping,
idx_mapping.stride(0),
temperature,
num_tokens,
vocab_size,
BLOCK_SIZE=VOCAB_BLOCK_SIZE,
)
_predict_kernel[(num_tokens,)](
self.features,
self.features.stride(0),
self.predictions,
self.predictions.stride(0),
confidence_probs,
confidence_probs.stride(0),
self.slope,
self.intercepts,
local_max,
local_max.stride(0),
local_sumexp,
local_sumexp.stride(0),
idx_mapping,
idx_mapping.stride(0),
draft_step,
num_tokens,
num_blocks,
per_token_step=draft_step.dim() > 0,
NUM_SPECULATIVE_STEPS=self.num_speculative_steps,
MAX_LOG_ODDS=_MAX_LOG_ODDS,
PADDED_VOCAB_NUM_BLOCKS=triton.next_power_of_2(num_blocks),
)