What does this PR do?

Adds support for static ensemble verification in speculative decoding (assisted generation), a training-free method that increases draft token acceptance rates by relaxing the verification distribution.

Fixes #45865

Motivation

Standard speculative decoding rejects many plausible tokens because it strictly enforces that the output distribution matches the target model. Static ensemble verification blends the target and draft distributions during verification:

v(x) = w * p_target(x) + (1 - w) * q_draft(x)

This provably achieves the Pareto-optimal tradeoff between acceptance rate and distributional bias (Proposition 1 in the paper). The acceptance probability increases from 1 - TV(q,p) to 1 - w*TV(q,p).

Changes

• Adds assistant_ensemble_weight parameter to GenerationConfig (float in (0, 1], default None = lossless)
• Modifies _speculative_sampling to use ensemble acceptance ratio: v(x)/q(x) = 1 - w + w*(p(x)/q(x))
• Supports greedy decoding (do_sample=False) by comparing against argmax(v) instead of argmax(p)
• Raises clear error when used with candidate generators that do not return logits (e.g., prompt lookup)
• Adds numerical stability guard for the fallback distribution
• Adds 8 fast synthetic unit tests
• Adds documentation section to assisted_decoding.md

Usage

# Sampling mode
outputs = model.generate(
    **inputs,
    assistant_model=assistant_model,
    do_sample=True,
    assistant_ensemble_weight=0.7,
)

# Greedy mode
outputs = model.generate(
    **inputs,
    assistant_model=assistant_model,
    do_sample=False,
    assistant_ensemble_weight=0.7,
)

Design decisions

1. Fallback distribution unchanged: Since [v-q]+ = w*[p-q]+ normalizes to the same distribution as [p-q]+, we compute the standard fallback directly for numerical stability.
2. Acceptance ratio form: Uses 1 - w + w*(p_i/q_i) which is algebraically equivalent to v_i/q_i but avoids constructing the full ensemble distribution.
3. Bonus token: When all candidates are accepted, the bonus token uses p (target) rather than v, since q at position N+1 is unavailable. This is conservative and documented.
4. Greedy ensemble: For do_sample=False, compares draft tokens against argmax(v) — "greedy under the ensemble verifier."

Tests

8 fast synthetic tests (all pass, ~0.15s total):

• w=None/w=1.0 equivalence (backward compatibility)
• Deterministic acceptance ratio (w=0.7 accepts where w=1.0 rejects)
• Fallback distribution finite and normalized
• Numerical stability with near-zero residual
• Greedy ensemble: argmax(v) accepts where argmax(p) rejects
• ValueError for missing candidate logits
• GenerationConfig round-trip serialization
• Default value is None

Results (from the paper)

On CNN/DailyMail with Llama-3.1-8B-Instruct (target) + Llama-3.2-1B-Instruct (draft), temperature=0:

• Standard SD (w=1.0): ~65% acceptance rate
• Static ensemble (w=0.7): ~78% acceptance rate, ROUGE-L within 0.5pt of target

References

• Paper: DIVERSED: Relaxed Speculative Decoding via Dynamic Ensemble Verification (AISTATS 2026)
• Code: https://github.com/comeusr/diversed
• Section 3.1 describes the static ensemble method

cc @Cyrilvallez