Expected Speedup of Speculative Decoding

With a draft model that proposes γ tokens per step and an acceptance rate α per token (assumed independent), the expected number of tokens produced per target-model pass is

E[\text{tokens}] = \frac{1 - \alpha^{\gamma + 1}}{1 - \alpha}

(it's γ + 1 when α = 1). If one draft forward pass costs c times a target pass, each step costs γ·c + 1 target-pass equivalents, so the speedup over plain decoding is E[tokens] / (γ·c + 1) (Leviathan et al., 2023).

Implement:

expected_tokens(alpha, gamma)
speedup(alpha, gamma, c)
best_gamma(alpha, c, max_gamma=20): the γ in 1..max_gamma with the highest speedup (ties: the smaller γ)
estimate_alpha(accepted_counts, gamma): estimate α from observed runs, where each count is how many of γ drafts were accepted before the first rejection (γ means all accepted). Use the MLE for a truncated geometric: α = total_accepted / (total_accepted + number_of_rejections).