Repository navigation
[generate] Slice cached decoder inputs of encoder-decoder models without decoder mask - #50013
tarekziade wants to merge 1 commit into
Conversation
…out decoder mask When `generate` restarts from an existing cache, `_prefill` only slices the already-cached inputs when it knows the mask length. Encoder-decoder models usually have no `decoder_attention_mask` (Whisper even pops it), and `_mask_length` is only recorded for decoder-only models, so the FULL `decoder_input_ids` were processed again on top of the cache. This hits encoder-decoder assistants in assisted decoding: the assistant cache grows by the full sequence at every drafting round, until Whisper's decoder indexes past its 448 learned positions (device-side assert on CUDA, which then poisons every following test of the worker). Fixes #50011 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
run-slow: whisper, bart, t5, mbart, marian, pegasus, m2m_100, bigbird_pegasus, t5gemma, seamless_m4t, moonshine, encoder_decoder, speech_encoder_decoder, vision_encoder_decoder |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Cyrilvallez
left a comment
Cyrilvallez
left a comment
There was a problem hiding this comment.
The fix is in general wrong, see my comment! For any small first generation loop (small prefix + small new number of tokens, it's very likely to have a subsequent prefill with more tokens than the existing Cache!
I opened #50027 to supersede this with a better fix!
| elif self.config.is_encoder_decoder: | ||
| # The decoder mask is usually absent (e.g. assistant models): the decoder inputs are the FULL sequence | ||
| # if they are longer than the cache | ||
| mask_length = input_ids.shape[1] if input_ids.shape[1] > past_length else -1 |
There was a problem hiding this comment.
This is wrong, they may be longer than the Cache even when restarting, e.g. any small first prefil + generation with a longer subsequent prefill
|
Closing this in the meantime in favor of the linked PR! (Also note that most of the tests that are linked here are in fact passing on |
Fixes #50011
test_validate_assistant(Whisper + distil-whisper assistant) triggers a device-side assert that takes down every later test in the worker.Cause: when
generaterestarts from a cache,_prefillonly slices off the cached inputs if it knows the mask length. Encoder-decoder models have no decoder mask, and_mask_length(#48814) is only set for decoder-only. So an encoder-decoder assistant re-prefills the fulldecoder_input_idsevery drafting round, its cache doubles, and Whisper indexes position 448 (max is 447).Fix: for encoder-decoder without a decoder mask, treat
decoder_input_idsas the full sequence when longer than the cache. Added a tiny-BART test (fails on main, passes here); the slow test passes locally.🤖 Generated with Claude Code