[Generation] Add static ensemble verification for lossy speculative decoding (#45979)
* feat(generation): Add static ensemble verification for lossy speculative decoding
Add assistant_ensemble_weight parameter to GenerationConfig that enables
static ensemble verification, a training-free method that increases draft
token acceptance rates by relaxing the verification distribution.
The verification distribution becomes v(x) = w * p(x) + (1-w) * q(x),
where w is the ensemble weight. This provably achieves the Pareto-optimal
tradeoff between acceptance rate and distributional bias.
Changes:
- Add assistant_ensemble_weight to GenerationConfig
- Modify _speculative_sampling to use ensemble acceptance ratio
- Support greedy decoding with argmax(v) in Case 2
- Add error for incompatible candidate generators (no logits)
- Add numerical stability guard for fallback distribution
- Add 8 fast synthetic unit tests
- Add documentation section to assisted_decoding.md
Reference: Wang & Kasa et al., "DIVERSED: Relaxed Speculative Decoding
via Dynamic Ensemble Verification", AISTATS 2026.
https://arxiv.org/abs/2604.07622
Fixes #45865
* Address Cyril's review comments
- Enforce 0.0 < assistant_ensemble_weight < 1.0 strictly in GenerationConfig.validate()
(boundaries and out-of-range values now raise at config time).
- Simplify runtime guards in _speculative_sampling and the greedy verification path
to rely on the new validate() check instead of repeating the < 1.0 condition inline.
- Tighten the prompt-lookup guard: any non-None weight now raises (previously only
raised for w < 1.0).
- Use modern type annotation 'float | None = None' on _speculative_sampling.
- Remove '1.0 keeps decoding lossless' from the docstring; only None is lossless now.
- Move tests from tests/generation/test_static_ensemble.py into existing files:
- 4 sampling-behaviour tests into UtilsFunctionsTest in tests/generation/test_utils.py
- 2 config tests into tests/generation/test_configuration_utils.py
- Delete tests/generation/test_static_ensemble.py.
Addresses review at https://github.com/huggingface/transformers/pull/45979#pullrequestreview-3490039503
* style: ruff format new test methods
* Address Cyril's follow-up nits
- Drop intermediate 'w = assistant_ensemble_weight' aliases in both
_speculative_sampling and the greedy verification path; inline the
parameter name directly for readability.
- Scope the numerical-stability guard on 'p_prime_sum' to the ensemble
path only (when assistant_ensemble_weight is not None). Standard
(lossless) SD keeps its original fallback behavior.
---------
Co-authored-by: Siva Rajesh Kasa <kasasiva@amazon.com>
Co-authored-by: Siva Rajesh Kasa <kasakh@users.noreply.github.com>