From e5aff89937ea1b987c99b8f972c4223816ddf3d2 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 06:50:06 +0000 Subject: [PATCH 01/13] Fix static-cache recompilation with chunked prefill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With `prefill_chunk_size` set, `generate()` compiles the prefill (via `get_compiled_call()` per chunk). It hits a condition in cache.update that leads to a graph specialization between the first call to generate and the second (one has the cache init, the second doesn't), hence a recompilation. The fix is to force a call to early_initialization outside of the compile region when chunked prefill. Also corrects the `lazy_initialization` docstring, which claimed `early_initialization` was unnecessary for `compile` because the prefill is never compiled internally — not true once `prefill_chunk_size` is set. Fixes #46421 Co-Authored-By: Claude Opus 4.8 --- src/transformers/cache_utils.py | 8 ++------ src/transformers/generation/utils.py | 28 ++++++++++++++++++++++++++-- tests/utils/test_cache_utils.py | 27 +++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 8 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 1f4f4b9c957d..fdfba6a088bb 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -300,12 +300,8 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens devices, dtypes etc later on for each `update` (which could break the static dynamo addresses as well). If this is unwanted, one can call `early_initialization(...)` on the Cache directly, which will call this - function ahead-of-time (this is required for `torch.export` for example). Note that for `compile`, as we - internally don't compile the prefill, this is guaranteed to have been called already when compiling. - If compiling the prefill as well, e.g. calling `model.compile(...)` before `generate` with a static cache, - it is still supported in general, but without guarantees depending on the compilation options (e.g. cuda graphs, - i.e. `mode="reduce-overhead"` is known to fail). But it will in general work correctly, and prefill should - not be compiled anyway for performances! + function ahead-of-time (this is required for `torch.export` or for example). It is also required whenever the + prefill itself ends up in a compiled region (with chunked prefill for instance). """ self.dtype, self.device = key_states.dtype, key_states.device self.max_batch_size, self.num_heads = key_states.shape[:2] diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index a4e3376ff224..14835dd75f65 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1918,12 +1918,36 @@ def _prepare_cache_for_generation( f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, " "and the layer structure will be inferred automatically." ) - model_kwargs[cache_name] = self._prepare_static_cache( + cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size + cache = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation, - batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size, + batch_size=cache_batch_size, max_cache_len=max_cache_length, model_kwargs=model_kwargs, ) + # With chunked prefill, the prefill itself runs inside a compiled region. + # Initializing the cache eagerly here to avoid a recompilation on the second generate() call. See #46421. + if ( + generation_config.prefill_chunk_size is not None + and isinstance(cache, StaticCache) + and not cache.is_initialized + ): + text_config = self.config.get_text_config(decoder=True) + head_dim = ( + getattr(text_config, "head_dim", None) + or text_config.hidden_size // text_config.num_attention_heads + ) + num_key_value_heads = ( + getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads + ) + cache.early_initialization( + batch_size=cache_batch_size, + num_heads=num_key_value_heads, + head_dim=head_dim, + dtype=self.dtype, + device=self.device, + ) + model_kwargs[cache_name] = cache elif generation_config.cache_implementation == "quantized": if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache(): raise ValueError( diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index 769639e5f612..1378fbe3b8f3 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -14,6 +14,7 @@ import copy import unittest +from unittest.mock import patch import pytest from packaging import version @@ -488,6 +489,32 @@ def test_static_cache_no_cuda_graph_skips(self): model.generate(**inputs, max_new_tokens=2, cache_implementation="static") self.assertNotIn("cuda", cap.err.lower()) + def test_chunked_prefill_initializes_static_cache_eagerly(self): + """ + With chunked prefill the prefill runs inside a compiled region, where the static cache's lazy initialization + cannot tag its tensors via `torch._dynamo.mark_static_address`, so a later call recompiles. `generate` must + therefore initialize the static cache eagerly, and only when it compiles the prefill (i.e. when + `prefill_chunk_size` is set). Regression test for #46421. + """ + model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM" + model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device) + tokenizer = AutoTokenizer.from_pretrained(model_repo) + inputs = tokenizer(["The quick brown fox jumps over the lazy dog"], return_tensors="pt").to(torch_device) + generation_kwargs = {"max_new_tokens": 3, "do_sample": False, "cache_implementation": "static"} + + # Without chunked prefill, the static cache is left to lazily initialize on the (eager) prefill. + with patch.object(Cache, "early_initialization", autospec=True) as eager_init: + model.generate(**inputs, **generation_kwargs) + eager_init.assert_not_called() + + # Drop the (now initialized) cache so the chunked-prefill call re-allocates and initializes a fresh one. + model._cache = None + # With chunked prefill, `generate` must initialize the static cache itself, before the prefill runs. + prefill_chunk_size = max(inputs.input_ids.shape[-1] // 2, 1) + with patch.object(Cache, "early_initialization", autospec=True, wraps=Cache.early_initialization) as init: + model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size) + init.assert_called_once() + @require_torch_multi_accelerator @slow def test_static_cache_multi_accelerator(self): From 90b0b61939d7f028c153d3f665953608c51e4c0c Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 06:50:35 +0000 Subject: [PATCH 02/13] Honor cache_config max_cache_len on the static cache path `GenerationConfig.cache_config` is documented as carrying arguments for the key-value cache, but `_prepare_cache_for_generation` read it only on the quantized branch. The static branch hardcoded `max_cache_len = max_length - 1`, so `cache_config={"max_cache_len": N}` with `cache_implementation="static"` was silently ignored. Fixes #46424 Co-Authored-By: Claude Opus 4.8 --- src/transformers/generation/utils.py | 7 +++++++ tests/utils/test_cache_utils.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 14835dd75f65..961e4d8cc4dc 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1918,6 +1918,13 @@ def _prepare_cache_for_generation( f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, " "and the layer structure will be inferred automatically." ) + # A user-provided `cache_config["max_cache_len"]` lets the static cache be sized for the worst case + # across calls, so that later calls with a longer prompt or a larger `max_new_tokens` (up to that + # ceiling) reuse the same cache instead of triggering a reallocation (and a `torch.compile` + # recompilation). Without this, the cache is sized to the current call's `max_length` only. See #46424. + cache_config = generation_config.cache_config or {} + if cache_config.get("max_cache_len") is not None: + max_cache_length = max(max_cache_length, cache_config["max_cache_len"]) cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size cache = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation, diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index 1378fbe3b8f3..2f5d3b2c22d0 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -489,6 +489,34 @@ def test_static_cache_no_cuda_graph_skips(self): model.generate(**inputs, max_new_tokens=2, cache_implementation="static") self.assertNotIn("cuda", cap.err.lower()) + def test_static_cache_honors_cache_config_max_cache_len(self): + """ + `cache_config["max_cache_len"]` must size the auto-allocated `StaticCache` on the static path, so a cache + pinned to a worst-case length is reused across calls instead of reallocating (and recompiling). Without this, + the cache is sized to the current call's `max_length` only. Regression test for #46424. + """ + # Llama has only full-attention layers, so the cache length is not capped by a sliding window. + model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM" + model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device) + tokenizer = AutoTokenizer.from_pretrained(model_repo) + inputs = tokenizer(["The quick brown fox"], return_tensors="pt").to(torch_device) + + max_new_tokens = 5 + # Deliberately ask for a cache much larger than the natural `max_length - 1` for this call. + natural_max_cache_len = inputs.input_ids.shape[-1] + max_new_tokens - 1 + requested_max_cache_len = natural_max_cache_len + 64 + + out = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + cache_implementation="static", + cache_config={"max_cache_len": requested_max_cache_len}, + return_dict_in_generate=True, + ) + self.assertIsInstance(out.past_key_values, StaticCache) + self.assertEqual(out.past_key_values.max_cache_len, requested_max_cache_len) + def test_chunked_prefill_initializes_static_cache_eagerly(self): """ With chunked prefill the prefill runs inside a compiled region, where the static cache's lazy initialization From 994e1d9407c5bf6342ae504a7137a5a3def8a094 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 07:55:43 +0000 Subject: [PATCH 03/13] Make StaticCache.early_initialization robust to linear attention layers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `early_initialization` initialized every layer via the `(num_heads, head_dim)` key/value layout, but linear attention layers (mamba/conv/linear_attention) don't use that layout — they track their own statically-shaped conv/recurrent states, initialized lazily on update. Initializing them through the key/value path is wrong and crashes on the first call to prefill. Fixes #46439 Co-Authored-By: Claude Opus 4.8 --- src/transformers/cache_utils.py | 17 ++++++++++++++++- tests/utils/test_cache_utils.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index fdfba6a088bb..3ff49657576c 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -1061,6 +1061,17 @@ def early_initialization( ) for layer, layer_num_heads, layer_head_dim in zip(self.layers, num_heads, head_dim): + # Only attention layers use the `(num_heads, head_dim)` key/value layout and an `is_initialized` flag, so + # they are the only ones we can early-initialize here. Linear attention layers (mamba/conv/...) track their + # own conv/recurrent states, whose shapes (conv kernel size, SSM state size, ...) are not derivable from + # `(num_heads, head_dim)`, so there is no correct way to early-init them through this path: they are + # skipped and lazily initialize from the real states on their first update. (Consequence: under chunked + # prefill they may still trigger a recompile, since that lazy init then happens inside the compiled + # region.) Already-initialized layers are skipped too, so this is a no-op on a reused cache, which + # preserves the static tensor addresses across `compile` calls. + is_initialized = getattr(layer, "is_initialized", None) + if is_initialized is None or is_initialized: + continue # Note that the initialization needs all dimensions (except -2), as well as device and dtype, so we use # this fake tensor approach. It has size 0 on the -2 dimension, so it does not allocate any data (it only # creates an empty tensor with correct shape, dtype and device), which is very efficient and practical @@ -1206,7 +1217,11 @@ def is_compileable(self) -> bool: @property def is_initialized(self) -> bool: """Return whether the cache data is initialized""" - return len(self.layers) > 0 and all(layer.is_initialized for layer in self.layers) + # Layers without an `is_initialized` flag (e.g. linear attention layers, which track their own conv/recurrent + # states) don't use this lazy-init mechanism, so skip them; a hybrid cache then reflects its attention layers. + flags = [getattr(layer, "is_initialized", None) for layer in self.layers] + flags = [flag for flag in flags if flag is not None] + return len(flags) > 0 and all(flags) @property def is_sliding(self) -> list[bool]: diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index 2f5d3b2c22d0..478c3d2a2ce7 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -56,6 +56,7 @@ convert_and_export_with_cache, pipeline, ) + from transformers.cache_utils import LinearAttentionLayer, StaticLayer from transformers.integrations.executorch import export_with_dynamic_cache @@ -109,6 +110,37 @@ def _random_kvs(config): self.assertTrue(cached_keys.shape == (1, 1, 10, 128)) self.assertTrue(cached_values.shape == (1, 1, 10, 128)) + def test_early_initialization_does_not_corrupt_linear_attention_layers(self): + """ + Regression test: `early_initialization` initialized *every* layer through the attention `(num_heads, head_dim)` + key/value layout, but linear attention layers track their own statically-shaped conv/recurrent states. So it + pre-allocated their `conv_states` with the wrong shape (a spurious 0-length axis) and flagged them + initialized; the first real `update_conv_state` then hit `conv_states.copy_(...)` with mismatched shapes and + raised a `RuntimeError` (the customer-facing failure, observed e.g. on Neuron). Linear attention layers must + be left untouched so they lazily take the correct shape on their first update. + """ + config = LlamaConfig(num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, hidden_size=32) + config.layer_types = ["full_attention", "linear_attention"] + cache = StaticCache(config=config, max_cache_len=8) + self.assertIsInstance(cache.layers[0], StaticLayer) + linear_layer = cache.layers[1] + self.assertIsInstance(linear_layer, LinearAttentionLayer) + + cache.early_initialization(batch_size=1, num_heads=2, head_dim=8, dtype=torch.float32, device=torch_device) + # The attention layer is initialized via the key/value layout, as expected. + self.assertTrue(cache.layers[0].is_initialized) + + # Simulate the first forward updating the linear layer's conv state with its real (conv) shape. Before the + # fix, `early_initialization` had pre-allocated `conv_states` with the wrong key/value shape, so this raised + # `RuntimeError: ... must match ...`; with the fix the layer lazily takes the correct shape here instead. + conv_states = torch.zeros((1, 8, 4), dtype=torch.float32, device=torch_device) # (batch, channels, kernel) + updated_conv_states = linear_layer.update_conv_state(conv_states) + self.assertEqual(updated_conv_states.shape, conv_states.shape) + + # The cache-level `is_initialized` must also tolerate linear attention layers (which don't expose the flag); + # before the fix this raised `AttributeError`. It reflects the attention layer's state. + self.assertTrue(cache.is_initialized) + def _skip_on_failed_cache_prerequisites(test, cache_implementation): """Function to skip tests on failed cache prerequisites, given a cache implementation""" From 7b2aa4adddd142f616b5f460e8cf904258180b32 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 09:44:27 +0000 Subject: [PATCH 04/13] Let create_masks_for_generate defer for models with custom layer masks `generate()` pre-builds attention masks for compilable caches via helpers registered in `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING`. No helpers exist for hybrid linear-attention models (Qwen3.5, Qwen3-Next, ...), since they build their own masks in the forward, so this raised a cryptic `KeyError: 'linear_attention'` and made `cache_implementation="static"` unusable with these models. This skips the early creation of the mask for these models. Co-Authored-By: Claude Opus 4.8 --- src/transformers/masking_utils.py | 8 +++++++- tests/utils/test_masking_utils.py | 24 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/transformers/masking_utils.py b/src/transformers/masking_utils.py index 2e755bdaee60..fe550b31cb0d 100644 --- a/src/transformers/masking_utils.py +++ b/src/transformers/masking_utils.py @@ -1504,8 +1504,14 @@ def create_masks_for_generate( # If the attribute exist, we need several masks if hasattr(effective_config, "layer_types"): + layer_types = effective_config.layer_types + # Hybrid models with a layer type that has no registered attention-mask function (e.g. `linear_attention` + # in Qwen3.5 / Qwen3-Next) build their own per-layer masks in the forward, so their masks cannot be + # pre-built here. Defer to the model by returning the raw attention mask, rather than raising a `KeyError`. + if any(layer_type not in LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING for layer_type in layer_types): + return attention_mask causal_masks = {} - for layer_pattern in set(effective_config.layer_types): + for layer_pattern in set(layer_types): causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs) return causal_masks # In this case, all layers are sliding diff --git a/tests/utils/test_masking_utils.py b/tests/utils/test_masking_utils.py index b780574f30f3..8b91593a38b5 100644 --- a/tests/utils/test_masking_utils.py +++ b/tests/utils/test_masking_utils.py @@ -32,6 +32,7 @@ create_bidirectional_mask, create_causal_mask, create_chunked_causal_mask, + create_masks_for_generate, find_packed_sequence_indices, ) @@ -373,3 +374,26 @@ def test_bidirectional_mask_skip_eager(self): self.assertTrue(padded_mask[0] is not None) self.assertTrue(padded_mask[1] is not None) + + def test_create_masks_for_generate_defers_for_unmapped_layer_types(self): + """ + `create_masks_for_generate` pre-builds attention masks for compilable caches by mapping each + `config.layer_types` entry through `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING`. Hybrid models with a layer type + that has no mask function (e.g. `linear_attention` in Qwen3.5 / Qwen3-Next, which build their own per-layer + masks in the forward) must not raise `KeyError` here; instead the raw attention mask is returned so the model + can build its own masks. + """ + config = LlamaConfig(num_hidden_layers=2) + config.layer_types = ["full_attention", "linear_attention"] + attention_mask = torch.ones((1, 5), dtype=torch.long, device=torch_device) + + # Must not raise (previously `KeyError: 'linear_attention'`), and defers the raw mask to the model. + out = create_masks_for_generate( + config, inputs_embeds=None, attention_mask=attention_mask, past_key_values=None + ) + self.assertIs(out, attention_mask) + + # `None` (no padding) is deferred too. + self.assertIsNone( + create_masks_for_generate(config, inputs_embeds=None, attention_mask=None, past_key_values=None) + ) From 533d3d6240cea06fec3597c91339afdbb340b95f Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 11:52:49 +0000 Subject: [PATCH 05/13] Make StaticCache.max_cache_len robust to linear attention layers The `max_cache_len` property iterated every layer, but linear attention layers have no `max_cache_len`, so the auto-allocated static-cache reuse check in `_prepare_static_cache` raised `AttributeError` on a hybrid model. Skip them. Co-Authored-By: Claude Opus 4.8 --- src/transformers/cache_utils.py | 5 +++-- tests/utils/test_cache_utils.py | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 3ff49657576c..607bffa8d3ba 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -1203,8 +1203,9 @@ def max_batch_size(self) -> int: @property def max_cache_len(self) -> int: """Return the maximum cache length of the cache""" - values = [layer.max_cache_len for layer in self.layers] - return max(values) + # Linear attention layers have no `max_cache_len`; skip them so a hybrid cache reports its attention layers'. + values = [layer.max_cache_len for layer in self.layers if hasattr(layer, "max_cache_len")] + return max(values) if values else 0 @property def is_compileable(self) -> bool: diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index 478c3d2a2ce7..ccc3dec2cd4f 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -141,6 +141,14 @@ def test_early_initialization_does_not_corrupt_linear_attention_layers(self): # before the fix this raised `AttributeError`. It reflects the attention layer's state. self.assertTrue(cache.is_initialized) + def test_max_cache_len_ignores_linear_attention_layers(self): + """`max_cache_len` must skip linear attention layers (which have no such attribute), else the static-cache + reuse check in `_prepare_static_cache` raises `AttributeError` on a hybrid model.""" + config = LlamaConfig(num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2, hidden_size=32) + config.layer_types = ["full_attention", "linear_attention"] + cache = StaticCache(config=config, max_cache_len=8) + self.assertEqual(cache.max_cache_len, 8) + def _skip_on_failed_cache_prerequisites(test, cache_implementation): """Function to skip tests on failed cache prerequisites, given a cache implementation""" From 40ddb52968953b04c5f49cc44c91a281c0acc47e Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 15:46:10 +0000 Subject: [PATCH 06/13] review(generate): skip/shard eager static-cache init under TP and multi-device device_map --- src/transformers/generation/utils.py | 14 +++++++++----- tests/utils/test_cache_utils.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 961e4d8cc4dc..4fd7c4da27c9 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1934,22 +1934,26 @@ def _prepare_cache_for_generation( ) # With chunked prefill, the prefill itself runs inside a compiled region. # Initializing the cache eagerly here to avoid a recompilation on the second generate() call. See #46421. + # It uses a single device and `num_heads // tp_size`, so only when the model is not device-mapped across + # devices and the heads shard evenly; otherwise fall back to lazy init. + text_config = self.config.get_text_config(decoder=True) + tp_size = getattr(self, "_tp_size", None) or 1 + num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads + device_mapped = hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1 + compatible = not device_mapped and num_key_value_heads % tp_size == 0 if ( generation_config.prefill_chunk_size is not None and isinstance(cache, StaticCache) and not cache.is_initialized + and compatible ): - text_config = self.config.get_text_config(decoder=True) head_dim = ( getattr(text_config, "head_dim", None) or text_config.hidden_size // text_config.num_attention_heads ) - num_key_value_heads = ( - getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads - ) cache.early_initialization( batch_size=cache_batch_size, - num_heads=num_key_value_heads, + num_heads=num_key_value_heads // tp_size, head_dim=head_dim, dtype=self.dtype, device=self.device, diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index ccc3dec2cd4f..efb16365b42f 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -583,6 +583,25 @@ def test_chunked_prefill_initializes_static_cache_eagerly(self): model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size) init.assert_called_once() + # Under TP, eager init is still used but the head count is sharded by `_tp_size` (one device per rank). + tc = model.config.get_text_config(decoder=True) + num_kv_heads = getattr(tc, "num_key_value_heads", None) or tc.num_attention_heads + model._cache = None + model._tp_size = num_kv_heads # divides evenly -> 1 head per rank + with patch.object(Cache, "early_initialization", autospec=True) as tp_init: + model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size) + tp_init.assert_called_once() + self.assertEqual(tp_init.call_args.kwargs["num_heads"], 1) + model._tp_size = None + + # A multi-device `device_map` (single `model.device` can't cover all layers) skips eager init -> lazy. + model._cache = None + model.hf_device_map = {"a": 0, "b": 1} + with patch.object(Cache, "early_initialization", autospec=True) as md_init: + model.generate(**inputs, **generation_kwargs, prefill_chunk_size=prefill_chunk_size) + md_init.assert_not_called() + del model.hf_device_map + @require_torch_multi_accelerator @slow def test_static_cache_multi_accelerator(self): From aa28be8af7a45f5b0e54cd691545a53eea8cfe09 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 5 Jun 2026 18:48:14 +0200 Subject: [PATCH 07/13] fix: typo --- src/transformers/cache_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 607bffa8d3ba..c28293297693 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -300,7 +300,7 @@ def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tens devices, dtypes etc later on for each `update` (which could break the static dynamo addresses as well). If this is unwanted, one can call `early_initialization(...)` on the Cache directly, which will call this - function ahead-of-time (this is required for `torch.export` or for example). It is also required whenever the + function ahead-of-time (this is required for `torch.export` for example). It is also required whenever the prefill itself ends up in a compiled region (with chunked prefill for instance). """ self.dtype, self.device = key_states.dtype, key_states.device From f1593a989ae4c05fe1a337363905df9b70b001b7 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Tue, 9 Jun 2026 13:33:34 +0000 Subject: [PATCH 08/13] review(generate): move eager static-cache init into _prepare_static_cache --- src/transformers/generation/utils.py | 58 ++++++++++++++-------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 4fd7c4da27c9..7e6d7d6e0101 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1772,7 +1772,12 @@ def _prepare_generation_config( return generation_config, model_kwargs def _prepare_static_cache( - self: "GenerativePreTrainedModel", cache_implementation: str, batch_size: int, max_cache_len: int, model_kwargs + self: "GenerativePreTrainedModel", + cache_implementation: str, + batch_size: int, + max_cache_len: int, + prefill_chunk_size: int | None, + model_kwargs, ) -> Cache: """ Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a @@ -1818,6 +1823,27 @@ def _prepare_static_cache( "offloading": offload_cache, } self._cache = EncoderDecoderCache(self._cache, StaticCache(**cross_attention_cache_kwargs)) + elif prefill_chunk_size is not None: + # Chunked prefill compiles the prefill, so eagerly init the fresh cache to avoid a recompile next call + # (#46421) -- only when single-device and the heads shard evenly across TP ranks (else lazy init). + text_config = self.config.get_text_config(decoder=True) + tp_size = getattr(self, "_tp_size", None) or 1 + num_key_value_heads = ( + getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads + ) + device_mapped = hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1 + if not device_mapped and num_key_value_heads % tp_size == 0: + head_dim = ( + getattr(text_config, "head_dim", None) + or text_config.hidden_size // text_config.num_attention_heads + ) + self._cache.early_initialization( + batch_size=batch_size, + num_heads=num_key_value_heads // tp_size, + head_dim=head_dim, + dtype=self.dtype, + device=self.device, + ) else: self._cache.reset() return self._cache @@ -1926,39 +1952,13 @@ def _prepare_cache_for_generation( if cache_config.get("max_cache_len") is not None: max_cache_length = max(max_cache_length, cache_config["max_cache_len"]) cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size - cache = self._prepare_static_cache( + model_kwargs[cache_name] = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation, batch_size=cache_batch_size, max_cache_len=max_cache_length, + prefill_chunk_size=generation_config.prefill_chunk_size, model_kwargs=model_kwargs, ) - # With chunked prefill, the prefill itself runs inside a compiled region. - # Initializing the cache eagerly here to avoid a recompilation on the second generate() call. See #46421. - # It uses a single device and `num_heads // tp_size`, so only when the model is not device-mapped across - # devices and the heads shard evenly; otherwise fall back to lazy init. - text_config = self.config.get_text_config(decoder=True) - tp_size = getattr(self, "_tp_size", None) or 1 - num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads - device_mapped = hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1 - compatible = not device_mapped and num_key_value_heads % tp_size == 0 - if ( - generation_config.prefill_chunk_size is not None - and isinstance(cache, StaticCache) - and not cache.is_initialized - and compatible - ): - head_dim = ( - getattr(text_config, "head_dim", None) - or text_config.hidden_size // text_config.num_attention_heads - ) - cache.early_initialization( - batch_size=cache_batch_size, - num_heads=num_key_value_heads // tp_size, - head_dim=head_dim, - dtype=self.dtype, - device=self.device, - ) - model_kwargs[cache_name] = cache elif generation_config.cache_implementation == "quantized": if self.config.is_encoder_decoder or not self._supports_default_dynamic_cache(): raise ValueError( From 835e90f793b66c071f450dd568dd3514f3483d8f Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Tue, 9 Jun 2026 13:34:14 +0000 Subject: [PATCH 09/13] review(generate): isolate eager static-cache init args in a helper --- src/transformers/generation/utils.py | 35 +++++++++++++++++----------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 7e6d7d6e0101..d5a6bad10706 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1771,6 +1771,23 @@ def _prepare_generation_config( return generation_config, model_kwargs + def _get_static_cache_init_shape(self: "GenerativePreTrainedModel") -> tuple[int, int] | None: + """ + Returns the per-rank `(num_heads, head_dim)` to eagerly initialize a `StaticCache`, with the head count sharded + for tensor parallelism. Returns `None` when the cache cannot be initialized on a single device -- the model is + device-mapped across devices, or the head count does not shard evenly across TP ranks -- in which case it + should be lazily initialized instead. + """ + if hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1: + return None + text_config = self.config.get_text_config(decoder=True) + tp_size = getattr(self, "_tp_size", None) or 1 + num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads + if num_key_value_heads % tp_size != 0: + return None + head_dim = getattr(text_config, "head_dim", None) or text_config.hidden_size // text_config.num_attention_heads + return num_key_value_heads // tp_size, head_dim + def _prepare_static_cache( self: "GenerativePreTrainedModel", cache_implementation: str, @@ -1825,21 +1842,13 @@ def _prepare_static_cache( self._cache = EncoderDecoderCache(self._cache, StaticCache(**cross_attention_cache_kwargs)) elif prefill_chunk_size is not None: # Chunked prefill compiles the prefill, so eagerly init the fresh cache to avoid a recompile next call - # (#46421) -- only when single-device and the heads shard evenly across TP ranks (else lazy init). - text_config = self.config.get_text_config(decoder=True) - tp_size = getattr(self, "_tp_size", None) or 1 - num_key_value_heads = ( - getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads - ) - device_mapped = hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1 - if not device_mapped and num_key_value_heads % tp_size == 0: - head_dim = ( - getattr(text_config, "head_dim", None) - or text_config.hidden_size // text_config.num_attention_heads - ) + # (#46421). Skipped (-> lazy init) when it can't be initialized on a single device. + init_shape = self._get_static_cache_init_shape() + if init_shape is not None: + num_heads, head_dim = init_shape self._cache.early_initialization( batch_size=batch_size, - num_heads=num_key_value_heads // tp_size, + num_heads=num_heads, head_dim=head_dim, dtype=self.dtype, device=self.device, From 6ec33423aecd96bbd2b0e7f7bd2354767a9576a2 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Tue, 9 Jun 2026 14:20:37 +0000 Subject: [PATCH 10/13] review(cache): gate early-init via a supports_early_init layer flag instead of getattr --- src/transformers/cache_utils.py | 21 ++++++--------------- tests/utils/test_cache_utils.py | 3 +++ 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index c28293297693..7b8b9ef7e186 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -38,6 +38,7 @@ class CacheLayerMixin(ABC): """Base, abstract class for a single layer's cache.""" is_compileable = False + supports_early_init = True # Subclasses can set ``layer_type`` to auto-register themselves in # ``LAYER_TYPE_CACHE_MAPPING`` at import time (used by ``DynamicCache`` to dispatch # per-layer cache classes from ``config.layer_types``). @@ -700,6 +701,8 @@ class LinearAttentionCacheLayerMixin(ABC): # All shapes are static by essence in a LinearAttention layer, so it is compileable is_compileable = True + # Linear attention layers track their own conv/recurrent states; they don't use the key/value early-init path. + supports_early_init = False def __init__(self): self.conv_states: torch.Tensor | None = None @@ -1061,16 +1064,7 @@ def early_initialization( ) for layer, layer_num_heads, layer_head_dim in zip(self.layers, num_heads, head_dim): - # Only attention layers use the `(num_heads, head_dim)` key/value layout and an `is_initialized` flag, so - # they are the only ones we can early-initialize here. Linear attention layers (mamba/conv/...) track their - # own conv/recurrent states, whose shapes (conv kernel size, SSM state size, ...) are not derivable from - # `(num_heads, head_dim)`, so there is no correct way to early-init them through this path: they are - # skipped and lazily initialize from the real states on their first update. (Consequence: under chunked - # prefill they may still trigger a recompile, since that lazy init then happens inside the compiled - # region.) Already-initialized layers are skipped too, so this is a no-op on a reused cache, which - # preserves the static tensor addresses across `compile` calls. - is_initialized = getattr(layer, "is_initialized", None) - if is_initialized is None or is_initialized: + if not layer.supports_early_init or layer.is_initialized: continue # Note that the initialization needs all dimensions (except -2), as well as device and dtype, so we use # this fake tensor approach. It has size 0 on the -2 dimension, so it does not allocate any data (it only @@ -1218,11 +1212,8 @@ def is_compileable(self) -> bool: @property def is_initialized(self) -> bool: """Return whether the cache data is initialized""" - # Layers without an `is_initialized` flag (e.g. linear attention layers, which track their own conv/recurrent - # states) don't use this lazy-init mechanism, so skip them; a hybrid cache then reflects its attention layers. - flags = [getattr(layer, "is_initialized", None) for layer in self.layers] - flags = [flag for flag in flags if flag is not None] - return len(flags) > 0 and all(flags) + layers = [layer for layer in self.layers if layer.supports_early_init] + return len(layers) > 0 and all(layer.is_initialized for layer in layers) @property def is_sliding(self) -> list[bool]: diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index efb16365b42f..227c95911a29 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -125,6 +125,9 @@ def test_early_initialization_does_not_corrupt_linear_attention_layers(self): self.assertIsInstance(cache.layers[0], StaticLayer) linear_layer = cache.layers[1] self.assertIsInstance(linear_layer, LinearAttentionLayer) + # Linear attention layers opt out of the key/value early-init path; attention layers opt in. + self.assertTrue(cache.layers[0].supports_early_init) + self.assertFalse(linear_layer.supports_early_init) cache.early_initialization(batch_size=1, num_heads=2, head_dim=8, dtype=torch.float32, device=torch_device) # The attention layer is initialized via the key/value layout, as expected. From bf237fd289ed263d82079615c873a20b6cb2a8ca Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Mon, 15 Jun 2026 11:25:14 +0000 Subject: [PATCH 11/13] Revert "Honor cache_config max_cache_len on the static cache path" This reverts commit 90b0b61939d7f028c153d3f665953608c51e4c0c. --- src/transformers/generation/utils.py | 7 ------- tests/utils/test_cache_utils.py | 28 ---------------------------- 2 files changed, 35 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index d5a6bad10706..67641e942e81 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1953,13 +1953,6 @@ def _prepare_cache_for_generation( f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, " "and the layer structure will be inferred automatically." ) - # A user-provided `cache_config["max_cache_len"]` lets the static cache be sized for the worst case - # across calls, so that later calls with a longer prompt or a larger `max_new_tokens` (up to that - # ceiling) reuse the same cache instead of triggering a reallocation (and a `torch.compile` - # recompilation). Without this, the cache is sized to the current call's `max_length` only. See #46424. - cache_config = generation_config.cache_config or {} - if cache_config.get("max_cache_len") is not None: - max_cache_length = max(max_cache_length, cache_config["max_cache_len"]) cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size model_kwargs[cache_name] = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation, diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index 227c95911a29..ac012f55e8d9 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -532,34 +532,6 @@ def test_static_cache_no_cuda_graph_skips(self): model.generate(**inputs, max_new_tokens=2, cache_implementation="static") self.assertNotIn("cuda", cap.err.lower()) - def test_static_cache_honors_cache_config_max_cache_len(self): - """ - `cache_config["max_cache_len"]` must size the auto-allocated `StaticCache` on the static path, so a cache - pinned to a worst-case length is reused across calls instead of reallocating (and recompiling). Without this, - the cache is sized to the current call's `max_length` only. Regression test for #46424. - """ - # Llama has only full-attention layers, so the cache length is not capped by a sliding window. - model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM" - model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device) - tokenizer = AutoTokenizer.from_pretrained(model_repo) - inputs = tokenizer(["The quick brown fox"], return_tensors="pt").to(torch_device) - - max_new_tokens = 5 - # Deliberately ask for a cache much larger than the natural `max_length - 1` for this call. - natural_max_cache_len = inputs.input_ids.shape[-1] + max_new_tokens - 1 - requested_max_cache_len = natural_max_cache_len + 64 - - out = model.generate( - **inputs, - max_new_tokens=max_new_tokens, - do_sample=False, - cache_implementation="static", - cache_config={"max_cache_len": requested_max_cache_len}, - return_dict_in_generate=True, - ) - self.assertIsInstance(out.past_key_values, StaticCache) - self.assertEqual(out.past_key_values.max_cache_len, requested_max_cache_len) - def test_chunked_prefill_initializes_static_cache_eagerly(self): """ With chunked prefill the prefill runs inside a compiled region, where the static cache's lazy initialization From 4ab7a915f4d92a66fa3716c1dc76fd686efb6871 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Mon, 15 Jun 2026 11:29:49 +0000 Subject: [PATCH 12/13] feat(generate): add max_cache_len generation flag for static caches Expose a documented top-level `GenerationConfig.max_cache_len` to pre-size the static cache for the worst case across calls, so later calls with a longer prompt or larger `max_new_tokens` (up to that ceiling) reuse the same cache instead of reallocating and recompiling. Replaces the previously silently-ignored `cache_config["max_cache_len"]` path with an explicit, validated flag. Fixes #46424 Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generation/configuration_utils.py | 12 ++++++++ src/transformers/generation/utils.py | 6 ++++ tests/utils/test_cache_utils.py | 28 +++++++++++++++++++ 3 files changed, 46 insertions(+) diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index 3742d782a78b..ad440aac7422 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -173,6 +173,12 @@ class GenerationConfig(PushToHubMixin): our [cache documentation](https://huggingface.co/docs/transformers/en/kv_cache) for further information. cache_config (`dict`, *optional*, default to `None`): Arguments used in the key-value cache class can be passed in `cache_config`. + max_cache_len (`int`, *optional*): + Only used with static caches (`cache_implementation` set to `"static"` or `"offloaded_static"`). + Pre-sizes the cache to this length instead of the current call's `max_length`. Set it once to the + largest call you expect so that repeated `generate()` calls with a longer prompt or a larger + `max_new_tokens` (up to this ceiling) reuse the same cache instead of triggering a reallocation and a + `torch.compile` recompilation. > Parameters for manipulation of the model output logits @@ -388,6 +394,7 @@ def __init__(self, **kwargs): self.use_cache = kwargs.pop("use_cache", None) self.cache_implementation = kwargs.pop("cache_implementation", None) self.cache_config = kwargs.pop("cache_config", None) + self.max_cache_len = kwargs.pop("max_cache_len", None) # Parameters for manipulation of the model output logits self.temperature = kwargs.pop("temperature", None) @@ -654,6 +661,11 @@ def validate(self, strict=False, user_set_attributes: set[str] | None = None): f"Invalid `cache_implementation` ({self.cache_implementation}). Choose one of: " f"{valid_cache_implementations}" ) + if self.max_cache_len is not None and self.cache_implementation not in ALL_STATIC_CACHE_IMPLEMENTATIONS: + logger.warning_once( + f"`max_cache_len` is only used with static caches ({STATIC_CACHE_IMPLEMENTATIONS}); it will be " + f"ignored with `cache_implementation={self.cache_implementation!r}`." + ) # 1.3. Performance attributes if self.compile_config is not None and not isinstance(self.compile_config, CompileConfig): raise ValueError( diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index 67641e942e81..8e6c1a5f4c64 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1953,6 +1953,12 @@ def _prepare_cache_for_generation( f"and will be removed in v5.13. Please only use one of {STATIC_CACHE_IMPLEMENTATIONS}, " "and the layer structure will be inferred automatically." ) + # `max_cache_len` lets the static cache be sized for the worst case across calls, so that later calls + # with a longer prompt or a larger `max_new_tokens` (up to that ceiling) reuse the same cache instead of + # triggering a reallocation (and a `torch.compile` recompilation). Without it, the cache is sized to the + # current call's `max_length` only. See #46424. + if generation_config.max_cache_len is not None: + max_cache_length = max(max_cache_length, generation_config.max_cache_len) cache_batch_size = max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size model_kwargs[cache_name] = self._prepare_static_cache( cache_implementation=generation_config.cache_implementation, diff --git a/tests/utils/test_cache_utils.py b/tests/utils/test_cache_utils.py index ac012f55e8d9..36a4fe958416 100644 --- a/tests/utils/test_cache_utils.py +++ b/tests/utils/test_cache_utils.py @@ -532,6 +532,34 @@ def test_static_cache_no_cuda_graph_skips(self): model.generate(**inputs, max_new_tokens=2, cache_implementation="static") self.assertNotIn("cuda", cap.err.lower()) + def test_static_cache_honors_max_cache_len(self): + """ + `generation_config.max_cache_len` must size the auto-allocated `StaticCache` on the static path, so a cache + pinned to a worst-case length is reused across calls instead of reallocating (and recompiling). Without this, + the cache is sized to the current call's `max_length` only. Regression test for #46424. + """ + # Llama has only full-attention layers, so the cache length is not capped by a sliding window. + model_repo = "hf-internal-testing/tiny-random-LlamaForCausalLM" + model = AutoModelForCausalLM.from_pretrained(model_repo).to(torch_device) + tokenizer = AutoTokenizer.from_pretrained(model_repo) + inputs = tokenizer(["The quick brown fox"], return_tensors="pt").to(torch_device) + + max_new_tokens = 5 + # Deliberately ask for a cache much larger than the natural `max_length - 1` for this call. + natural_max_cache_len = inputs.input_ids.shape[-1] + max_new_tokens - 1 + requested_max_cache_len = natural_max_cache_len + 64 + + out = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=False, + cache_implementation="static", + max_cache_len=requested_max_cache_len, + return_dict_in_generate=True, + ) + self.assertIsInstance(out.past_key_values, StaticCache) + self.assertEqual(out.past_key_values.max_cache_len, requested_max_cache_len) + def test_chunked_prefill_initializes_static_cache_eagerly(self): """ With chunked prefill the prefill runs inside a compiled region, where the static cache's lazy initialization From 736803ad9cad32f53e6ad2da2845efb8ba179834 Mon Sep 17 00:00:00 2001 From: David Corvoysier Date: Fri, 19 Jun 2026 08:12:07 +0000 Subject: [PATCH 13/13] review: MLA models cannot be lazy initialized --- src/transformers/generation/utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/transformers/generation/utils.py b/src/transformers/generation/utils.py index f6361dd27b6f..937011e8a351 100644 --- a/src/transformers/generation/utils.py +++ b/src/transformers/generation/utils.py @@ -1787,16 +1787,19 @@ def _prepare_generation_config( def _get_static_cache_init_shape(self: "GenerativePreTrainedModel") -> tuple[int, int] | None: """ Returns the per-rank `(num_heads, head_dim)` to eagerly initialize a `StaticCache`, with the head count sharded - for tensor parallelism. Returns `None` when the cache cannot be initialized on a single device -- the model is - device-mapped across devices, or the head count does not shard evenly across TP ranks -- in which case it - should be lazily initialized instead. + for tensor parallelism. Returns `None` when the cache cannot be early initialized. """ if hasattr(self, "hf_device_map") and len(set(self.hf_device_map.values())) > 1: + # The model layers are on different devices return None text_config = self.config.get_text_config(decoder=True) tp_size = getattr(self, "_tp_size", None) or 1 num_key_value_heads = getattr(text_config, "num_key_value_heads", None) or text_config.num_attention_heads if num_key_value_heads % tp_size != 0: + # The model cannot be evenly sharded by head + return None + if getattr(text_config, "qk_head_dim", None) is not None: + # MLA models have distinct key (`qk_head_dim`) and value (`v_head_dim`) sizes. return None head_dim = getattr(text_config, "head_dim", None) or text_config.hidden_size // text_config.num_attention_heads return num_key_value_heads // tp_size, head_dim