transformers
a5f41e09 - [Whisper] Fix speculative decoding: UnboundLocalError, cache corruption, and speed regression (#48000)

Commit
2 days ago
[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>
Author
Parents
Loading