[Whisper] Fix speculative decoding: UnboundLocalError, cache corruption, and speed regression (#48000)
* [Whisper] Fix cache type and layer count for speculative decoding with WhisperForCausalLM
Fix 1 — is_updated UnboundLocalError (WhisperAttention.forward)
PR #39956 wrapped the `is_updated` assignment in an isinstance guard but left it
consumed unconditionally below. The bug was dormant while Whisper always received
EncoderDecoderCache. PR #43679 changed generate() to pre-initialize a plain DynamicCache
for is_encoder_decoder=False models, activating the latent bug. Fix: initialize
is_updated=False before the guard, identical to PR #40517 which fixed 55 other models.
Fix 2a — Wrong cache type: DynamicCache instead of EncoderDecoderCache (WhisperDecoder.forward)
generate() uses config.is_encoder_decoder to pick the cache type. WhisperForCausalLM
sets is_encoder_decoder=False, so PR #43679 made generate() pre-initialize a plain
DynamicCache. But WhisperForCausalLM uses cross-attention: self- and cross-attention
share the same layer_idx slots in DynamicCache and overwrite each other. Fix: when
generate() passes an empty DynamicCache and encoder_hidden_states are present, convert
it to EncoderDecoderCache before the first forward pass.
Fix 2b — Wrong layer count: 32 slots pre-allocated for a 2-layer decoder (WhisperDecoder.forward)
DynamicCache(config=self.config) calls config.get_text_config(decoder=True), which for
WhisperConfig returns self unchanged (flat config, no nested decoder sub-config). It then
reads config.num_hidden_layers, which WhisperConfig maps via attribute_map to
encoder_layers=32. For distil-whisper-large-v2, encoder_layers=32 but decoder_layers=2,
so 32 cache slots are pre-allocated for a 2-layer decoder. Slots 2-31 are uninitialized
and crash on crop(-n). Fix: deepcopy the config and override num_hidden_layers with
decoder_layers before passing to DynamicCache.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [Whisper] Restore absolute cache crop in _update_past_and_masks for speculative decoding
Fix 3 — Cache size accumulation (candidate_generator.py)
Commit 5ec7b73615 ("crop with negative always") changed _update_past_and_masks to always
remove a fixed number of tokens (crop(-1)) instead of cropping to an absolute target.
When multiple speculative candidates are rejected in a round, the cache grows by several
tokens but only 1 is trimmed. Over successive rounds, cache.get_seq_length() exceeds
input_ids.shape[1]. _prefill then computes next_sequence_length = input_ids.shape[1] -
cache.get_seq_length() which goes negative, causing input_ids[:, -n:] to produce a [1, 0]
empty tensor that crashes in WhisperAttention.forward() at the reshape.
Verified by bisect: at the parent of 5ec7b73615 (68075f8104), fixes 1+2a+2b alone are
sufficient (speed assertion only). At 5ec7b73615 with only fixes 1+2a+2b the reshape
crash reproduces. Adding this fix restores the test to speed-assertion-only.
Fix: compute the absolute target size from input_ids and only crop if the cache exceeds
it, using a negative delta to comply with the current Cache.crop(negative) API.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [Generate] Sync adaptive confidence threshold to generation_config for ConfidenceCriteria
PR #42702 introduced `self.assistant_generation_config` (a deepcopy with defaults applied)
to hold the assistant model's config, but the adaptive threshold update
(`update_candidate_strategy`) was only syncing the new best threshold to
`self.assistant_generation_config.assistant_confidence_threshold`, not to
`self.generation_config.assistant_confidence_threshold`.
`ConfidenceCriteria` is constructed each round from `self.generation_config`, so it was
always using the initial 0.4 threshold regardless of what the ROC-curve adaptation computed.
This caused the assistant to always generate the maximum number of speculative tokens per
round, making assisted generation slower than non-assisted on the distil-whisper test.
Fix: also update `self.generation_config.assistant_confidence_threshold` whenever the
adaptive threshold is updated, so each new round's ConfidenceCriteria reflects the current
best threshold.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [Whisper] Format: wrap long EncoderDecoderCache line
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [AudioFlamingo3] Fix is_updated UnboundLocalError in cross-attention (sync with modular)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---------
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>