transformers

ffddd251 - Fix `StaticCache` for Mllama and enable `torch.compile` (#48141)

Fix `StaticCache` for Mllama and enable `torch.compile` (#48141) * Fix StaticCache allocation for Mllama cross-attention layers `generate(cache_implementation="static")` crashed on Mllama, both eagerly (`index_copy_(): index out of bounds`) and under `torch.compile` (`expand: attempting to expand a dimension of length 6404 -> 47!`). `_prepare_static_cache` sizes every layer of the `StaticCache` from `max_cache_len`, which counts text tokens. Mllama interleaves cross-attention layers into the same layer list, and those cache the vision tower states instead, whose length is `num_images * max_num_tiles * num_patches` (6404 for the 11B checkpoint, against a `max_cache_len` of 47 in the repro). Re-allocate the cross-attention layers with that length, mirroring the `is_encoder_decoder` branch of `_prepare_static_cache`, which likewise sizes the cross-attention cache from the encoder output rather than from `max_cache_len`. Mllama cannot reuse `EncoderDecoderCache` directly because it has a single interleaved stack instead of two separate ones. The resulting cache is fully static, so `Cache.is_compileable` is `True` and `generate` would auto-compile. Mllama appends a row to `cross_attention_mask` at every decoding step, so `forward` sees a new input shape each step and dynamo recompiles the whole graph, measured ~17x slower than eager on the 11B checkpoint. `_can_compile_fullgraph = False` was meant to prevent exactly this but is never read by `generate`, so honour it in `_valid_auto_compile_criteria`. Benchmarked on Llama-3.2-11B-Vision-Instruct (A100, bf16, sdpa, 64 tokens, median of 5): dynamic 2.014s/31.78 tok/s (unchanged from main, identical tokens), static 2.491s/25.69 tok/s (crashed on main). Static and dynamic differ slightly in output, as expected from the different KV layout. * Scope the auto-compile override to `MllamaForConditionalGeneration` `MllamaPreTrainedModel` is reused by `blt` through modular, so an override placed there leaks into `BltPreTrainedModel` and breaks `utils/check_modular_conversion.py`. Only `MllamaForConditionalGeneration` builds the cross-attention static cache, so put the override there. Also derive the cross-attention cache length from `self.model.vision_model` (`max_num_tiles`, `num_patches`) instead of recomputing it from the config, so it stays in sync with the length the vision tower actually produces. * Enable torch.compile for Mllama with a static cache Now that the cross-attention layers are sized correctly, the static cache is compileable. Three things kept the traced graph from being static: - `MllamaTextCrossAttention.forward` and `MllamaTextModel.forward` branched on `past_key_values.get_seq_length(...)`, which is a 0-dim tensor for compileable layers. Both only need to know whether the cross-attention layer is populated, so use `layer.is_initialized` (a python bool) through a small helper. - `cross_attention_mask` grows by one row per decoded token and was passed to `forward` in full, so dynamo recompiled at every step. Slice it in `prepare_inputs_for_generation`, like the generic implementation does for `attention_mask`, including the `clone` that keeps the stride consistent. - Drop the now-incorrect `_can_compile_fullgraph = False` from `MllamaPreTrainedModel` and set it on `MllamaForConditionalGeneration`. Llama-3.2-11B-Vision-Instruct, bf16, sdpa, 64 new tokens: 42.70 tok/s compiled against 32.13 tok/s with the dynamic cache, with one graph, no graph breaks and no recompiles. Also update `_check_attentions_for_generate` to account for static caches, and add a test asserting the compiled output matches eager with a single graph. * Shorten comments * Shorten comment * Drop PR reference from comment * Don't touch blt, and skip the compile test without an accelerator `MllamaForCausalLM` already sets `_can_compile_fullgraph = True` explicitly, so the `False` on `MllamaPreTrainedModel` was only ever read by `MllamaForConditionalGeneration`. Overriding it there instead of removing it leaves `blt`, which inherits from `MllamaPreTrainedModel` through modular, completely untouched. `generate` only auto-compiles on an accelerator, so `test_generate_compile_matches_eager` counted 0 graphs on the CPU runners. * Move the redundant `_can_compile_fullgraph` down to blt `MllamaForCausalLM` and `MllamaForConditionalGeneration` both declare the flag themselves, so the `False` on `MllamaPreTrainedModel` was unreachable, and its comment ("static cache cannot have different shapes for each layer") is exactly what this PR disproves. `blt` inherits that base class through modular, so the line moves into `modular_blt.py` verbatim: every `Blt*` class still resolves to `False` (`BltForCausalLM` sets it explicitly, the rest now inherit the same default from `PreTrainedModel`), and `tests/models/blt/` gives an identical result before and after. * Revert "Move the redundant `_can_compile_fullgraph` down to blt" This reverts commit 3b6fa6fd45012718fe390969f95368afc9501dd9. * Reapply "Move the redundant `_can_compile_fullgraph` down to blt" This reverts commit 386a4dd54c269b3981e13ec4364ce8e7ecd7e016. * Shorten the `_can_compile_fullgraph` comment * Address review: reuse next_sequence_length, handle early-initialized caches, drop dead cross-attention branch * Read the cross-attention cache length from the config instead of the vision module