transformers
12b6da1f - Fix sliding window cache when assistant model calls generate (#48280)

Commit
39 days ago
Fix sliding window cache when assistant model calls generate (#48280) * Fix over-wide KV states from recorded sliding window cache updates With assisted or speculative decoding, DynamicSlidingWindowLayer runs with past recording active so that crop can roll back rejected drafts. While recording, update stored and returned the unbounded concatenation, but get_mask_sizes kept advertising sliding_window - 1 + query_length columns for the attention mask. Drafting bursts update the layer once per token with no crop in between, so the second and later tokens of a burst handed attention more key/value columns than the mask has: eager attention fails with a size mismatch and SDPA silently attends beyond-window states, making assisted decoding crash or diverge on any model whose sequence reaches the sliding window. Return only the trailing advertised window while recording; the full buffer stays stored internally for crop to roll back, and paths without recording are unchanged. Tests: a layer-level test pins returned width to get_mask_sizes across consecutive uncropped updates, an end-to-end test checks assisted decoding stays equal to greedy search with a prompt past the window, and a multi-token draft burst under eager attention crashes before the fix. Test Plan: python -m pytest tests/utils/test_cache_utils.py -k "recording_returns_advertised" -q --no-header -p no:cacheprovider fail-before: 1 failed (returned 5 != 4), pass-after: 1 passed python -m pytest tests/generation/test_utils.py -k "multi_token_draft" ... fail-before: 1 failed, RuntimeError "The size of tensor a (13) must match the size of tensor b (6)" pass-after: 1 passed (3 consecutive runs) python -m pytest tests/utils/test_cache_utils.py::CacheCroppingTests -q --no-header -p no:cacheprovider 5 passed python -m pytest tests/generation/test_utils.py -k "assisted" -q --no-header -p no:cacheprovider 22 passed, 3 skipped * slightly change and improve comment * improve tests a bit * correctly parameterize --------- Co-authored-by: Cyril Vallez <cyril.vallez@gmail.com>
Author
Parents
Loading