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>