[cache] Allow sliding window layers to be roll-backed for speculative decoding#47447
Conversation
|
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. |
|
The failing |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
There was a problem hiding this comment.
Super clean, just one question tbh as I'm not 100% sure there
| METADATA_FIELDS = ("_from_model_config", "_commit_hash", "_original_object_hash", "transformers_version") | ||
| STATIC_CACHE_IMPLEMENTATIONS = ("static", "offloaded_static") | ||
| DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "dynamic_full", "offloaded", "quantized") | ||
| DYNAMIC_CACHE_IMPLEMENTATIONS = ("dynamic", "offloaded", "quantized") |
There was a problem hiding this comment.
Do we need som deprecation notice?
There was a problem hiding this comment.
It was only used internally for assisted decoding, that's why I did not add one - was not intended for external usage when it was added
| super().crop(max_length) | ||
| if max_length < 0: | ||
| max_length = self.get_seq_length() - abs(max_length) | ||
| max_length = self.idx_keys.shape[-2] - abs(max_length) |
There was a problem hiding this comment.
Hmm, why was this change needed?
There was a problem hiding this comment.
Because I now crop negatively instead of positively in assisted decoding, and this uncovered this bug... Basically the length is not the same as the seq length, so we would not crop correctly before
There was a problem hiding this comment.
Basically because the seq_len is cropped before in the call to super... So then using it to crop the idx is not correct anymore - let me fix by switching order instead
| tokens. | ||
| """ | ||
| # If we are beyond the sliding window, we need to be more careful | ||
| if self.get_seq_length() >= self.sliding_window: |
There was a problem hiding this comment.
I guess this can only happen on the initial prefill?
There was a problem hiding this comment.
Not necessarily, can be during decode as well, or during another forward with more than 1 token later on
|
[For maintainers] Suggested jobs to run (before merge) run-slow: diffusion_gemma, minimax_m3_vl |
* upstream/main: (39 commits) Remove deprecated training args and `is_fast` property (huggingface#46917) Consistent output shape from `get_image_features` (huggingface#46405) Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423) fix failed test cases for qwen3_omni_moe model (huggingface#47449) Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233) Move `value` padding into the attention interfaces that need it (huggingface#47451) Simplify function dispatch for linear attention (huggingface#47450) [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447) Fix double-shifted training loss in GitForCausalLM (huggingface#47395) Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895) Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379) Update bug report list (huggingface#46607) Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952) Optimize flash attention max seqlen computation in vision attention (huggingface#47170) fix: remove unreachable return in special token builder (huggingface#47420) Add Harry to slow CI (huggingface#47454) BLT: vectorize patch length processing (huggingface#47385) Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935) [Kimi] add integration tests (huggingface#47383) Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974) ...
What does this PR do?
As per the title. This follows my work in #47347. This allows to use DynamicSlidingWindowLayer even during speculative decoding (i.e. mtp) instead of forcing a full cache that would waste a lot of memory.
This has been a huge restriction of assisted decoding until now, which is now lifted