Support explicit head_dim on MistralBackbone (enables Magistral) (#2692)

* Support explicit head_dim on MistralBackbone (enables Magistral)

Mistral and its derivatives always had `head_dim = hidden_dim //
num_query_heads` hard-coded in `CachedMistralAttention.build`. That
assumption breaks for Magistral (mistralai/Magistral-Small-2506):

    hidden_size=5120, num_attention_heads=32  -> 5120/32 = 160
    head_dim in HF config                    -> 128

Add an optional `head_dim` argument on `MistralBackbone`,
`MistralTransformerDecoder`, and `CachedMistralAttention`. When unset,
behavior is unchanged. When provided, it overrides the hidden/heads
fallback so the Q/K/V and output einsum projections build with the
correct per-head size.

Also:
- Pass `head_dim` through the HuggingFace -> keras-hub converter and
  fall back to `.get("sliding_window")` so the `null` setting on
  Magistral no longer KeyErrors.
- Add a unit test on `MistralBackbone` that exercises the explicit
  head_dim path with sliding_window=None.
- Add a numerical-parity test that builds a small HF Mistral with
  `head_dim=12` (!= hidden/num_heads), converts it through the preset
  loader, and asserts the keras-hub forward pass matches the HF
  reference to within fp16 precision.

After this PR lands, users can load Magistral directly:

    keras_hub.models.MistralBackbone.from_preset(
        "hf://mistralai/Magistral-Small-2506"
    )

Addresses the "[Contributions Welcome] Add Magistral" item in the
KerasHub roadmap (#1836) and issue #2314.

* Address review feedback on MistralBackbone head_dim

- Rename the constructor-argument attribute `_configured_head_dim` to
  `_head_dim` on `CachedMistralAttention`, matching the style-guide rule
  that layer attributes should share the name of the `__init__` argument
  (keras-hub style guide line 542).
- Simplify `build()` to only compute `_head_dim` from `_hidden_dim`
  when the caller did not supply one, instead of maintaining a second
  attribute for the configured value.
- Update `get_config` to return the renamed attribute.
- Tweak the `MistralBackbone.head_dim` docstring: "not a multiple of"
  was inaccurate — any value that differs from
  `hidden_dim // num_query_heads` needs to be set explicitly, and `//`
  matches the code path.

* Plumb head_dim through Mistral checkpoint conversion (adds Magistral preset)

Update tools/checkpoint_conversion/convert_mistral_checkpoints.py to:
  * Add `magistral_small_2506_en` to PRESET_MAP.
  * Extract `head_dim` once from the HF config (with fallback to
    `hidden_size // num_attention_heads`) and use it in the Q/K/V/O
    reshape calls instead of the hard-coded ratio. Standard Mistral
    checkpoints take the fallback path and behavior is unchanged.
  * Pass `head_dim` through to `MistralBackbone` via `backbone_kwargs`.

Also handle the transformers 5.x rope_parameters drift in both this
script and `keras_hub/src/utils/transformers/convert_mistral.py`:
`rope_theta` no longer lives at the top level of the HF config and is
nested inside `rope_parameters`. Without this fallback, both the
script-path conversion and `from_preset("hf://...")` raise on any
recent HF Mistral/Magistral checkpoint.

Numerics verified on synthetic standard and Magistral-shaped configs
(head_dim=12 != hidden//heads, sliding_window=None) within atol=1e-2,
matching the precision bound of the fp16 weight hops.

* Use ops.convert_to_numpy in test_explicit_head_dim_matches_hf

`np.asarray()` cannot directly convert a CUDA torch tensor returned by
the keras backbone under the PyTorch backend, breaking PyTorch-GPU CI:

    TypeError: can't convert cuda:0 device type tensor to numpy.
    Use Tensor.cpu() to copy the tensor to host memory first.

`keras.ops.convert_to_numpy()` is backend-agnostic and handles the
`.cpu().numpy()` hop internally for PyTorch.

* Cover explicit head_dim path with run_backbone_test

Route test_explicit_head_dim through run_backbone_test so the Magistral
style config (explicit head_dim, sliding_window=None) also gets
serialization and mixed precision coverage, then keep the head_dim
assertion on the built attention layer.

* Add magistral_small_2507_en to Mistral PRESET_MAP

* Support Tekken tokenizer conversion for Mistral

Magistral checkpoints ship a Tekken (byte-level BPE) tokenizer as
tekken.json instead of a SentencePiece tokenizer.model, so the previous
convert_tokenizer that hardcoded get_file(preset, 'tokenizer.model')
failed on them.

Add convert_tekken.py to reconstruct BPE merge rules from the Tekken
rank-ordered byte vocabulary and re-encode it with the GPT-2
byte-to-unicode mapping. MistralTokenizer now accepts vocabulary/merges
in addition to a proto: when built from Tekken data it delegates to an
internal BytePairTokenizer configured with the Tekken split pattern,
with a tf.py_function bridge for the tf.data graph path. Existing
SentencePiece presets are unchanged.

* Parse Tekken vocab with mistral_common

Use Mistral's own mistral_common library to parse tekken.json during
conversion instead of parsing the raw JSON by hand. This is the same
backend Hugging Face delegates to for Tekken, so the byte vocabulary,
special tokens, split pattern and vocab size come from a maintained
source rather than reimplemented parsing.

mistral_common is imported lazily inside convert_tokenizer and is only
needed at conversion time; the runtime tokenizer still uses
BytePairTokenizer and does not depend on it. Added as an optional dep.

* Fix Tekken tokenizer init order on torch backend

On the torch backend a keras.Layer is also a torch.nn.Module, which
rejects submodule assignment before its own __init__ runs. Initialize
the base tokenizer before assigning the internal BytePairTokenizer so
the Tekken path works on all backends.

* Make Tekken tokenizer a standalone MistralTekkenTokenizer class

Address review feedback:
- Add MistralTekkenTokenizer(BytePairTokenizer) as a public class instead
  of turning MistralTokenizer into a dual-personality class. MistralTokenizer
  is left untouched (SentencePiece only). convert_tokenizer returns a
  MistralTekkenTokenizer directly when tekken.json is detected.
- Inline the Tekken parsing helpers into convert_mistral.py and remove the
  separate convert_tekken.py, matching how convert_gemma4.py inlines its
  tokenizer helpers.
- Drop mistral-common from requirements-common.txt; it is only needed at
  conversion time and is already lazily imported with an ImportError guard.

* Use explicit head_dim when building the generation cache

MistralCausalLM._build_cache always derived head_dim as
hidden_dim // num_query_heads, ignoring an explicitly configured head_dim.
For Magistral (hidden_dim=5120, num_query_heads=32, head_dim=128) this built
the cache with head_dim=160 and generate() failed in dot_product_attention.
Fall back to the derived value only when head_dim is not set, matching
MistralAttention. Add a regression test with an explicit head_dim.

* Align Mistral conversion script with recent converter patterns

- Load the HF reference in float32 on CPU and precompute its logits,
  then free it so only one full-precision model is in memory at a time
  (qwen3_5 pattern). This also fixes the bfloat16 .numpy() crash.
- Load the KerasHub model through from_preset("hf://...") so the script
  exercises the built-in transformers converter (which now handles
  head_dim) instead of duplicating the weight porting logic.
- Verify parameter counts, token IDs, and logits against the HF
  reference in float32 with a hard assert (atol=1e-3).
- Reload in bfloat16 and save the preset in bfloat16 (gemma4 pattern).

Author: bzantium

keras_hub/api/models/__init__.py

@@ -511,6 +511,9 @@
 from keras_hub.src.models.mistral.mistral_causal_lm_preprocessor import (
     MistralCausalLMPreprocessor as MistralCausalLMPreprocessor,
 )
+from keras_hub.src.models.mistral.mistral_tokenizer import (
+    MistralTekkenTokenizer as MistralTekkenTokenizer,
+)
 from keras_hub.src.models.mistral.mistral_tokenizer import (
     MistralTokenizer as MistralTokenizer,
 )

keras_hub/api/tokenizers/__init__.py

@@ -71,6 +71,9 @@
 from keras_hub.src.models.metaclip_2.metaclip_2_tokenizer import (
     MetaCLIP2Tokenizer as MetaCLIP2Tokenizer,
 )
+from keras_hub.src.models.mistral.mistral_tokenizer import (
+    MistralTekkenTokenizer as MistralTekkenTokenizer,
+)
 from keras_hub.src.models.mistral.mistral_tokenizer import (
     MistralTokenizer as MistralTokenizer,
 )

keras_hub/src/models/mistral/mistral_attention.py

@@ -27,13 +27,15 @@ def __init__(
         kernel_initializer="glorot_uniform",
         sliding_window=512,
         dropout=0,
+        head_dim=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self._num_query_heads = num_query_heads
         self._num_key_value_heads = num_key_value_heads
         self._sliding_window = sliding_window
         self._dropout = dropout
+        self._head_dim = head_dim
 
         self._num_key_value_groups = num_query_heads // num_key_value_heads
         self._rope_max_wavelength = rope_max_wavelength
@@ -54,7 +56,8 @@ def build(self, inputs_shape):
         # v = num key/value heads
         # h = head dim
         self._hidden_dim = inputs_shape[-1]
-        self._head_dim = self._hidden_dim // self._num_query_heads
+        if self._head_dim is None:
+            self._head_dim = self._hidden_dim // self._num_query_heads
         self._inv_norm_factor = 1.0 / math.sqrt(self._head_dim)
 
         self._query_dense = keras.layers.EinsumDense(
@@ -239,6 +242,7 @@ def get_config(self):
                 ),
                 "sliding_window": self._sliding_window,
                 "dropout": self._dropout,
+                "head_dim": self._head_dim,
             }
         )
         return config

keras_hub/src/models/mistral/mistral_backbone.py

@@ -52,7 +52,13 @@ class MistralBackbone(Backbone):
             attention layers. This controls the maximum cache size for the
             attention layers in each transformer decoder. Only `sliding_window`
             number of tokens are saved in the cache and used to generate the
-            next token. Defaults to `512`.
+            next token. Defaults to `512`. Pass `None` to disable sliding
+            window attention entirely (e.g. Magistral).
+        head_dim (int, optional): The size of each attention head. When
+            `None` (the default), falls back to `hidden_dim // num_query_heads`.
+            Set explicitly when the model's head size is not equal to
+            `hidden_dim // num_query_heads` — e.g. Magistral uses
+            `head_dim=128` with `hidden_dim=5120` and `num_query_heads=32`.
         dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
             for model computations and weights. Note that some computations,
             such as softmax and layer normalization, will always be done at
@@ -98,6 +104,7 @@ def __init__(
         rope_scaling_factor=1.0,
         layer_norm_epsilon=1e-6,
         sliding_window=512,
+        head_dim=None,
         dropout=0,
         dtype=None,
         **kwargs,
@@ -123,6 +130,7 @@ def __init__(
                 activation=ops.silu,
                 kernel_initializer=_mistral_kernel_initializer(stddev=0.02),
                 sliding_window=sliding_window,
+                head_dim=head_dim,
                 dropout=dropout,
                 dtype=dtype,
                 name=f"transformer_layer_{i}",
@@ -165,6 +173,7 @@ def __init__(
         self.num_key_value_heads = num_key_value_heads
         self.rope_scaling_factor = rope_scaling_factor
         self.sliding_window = sliding_window
+        self.head_dim = head_dim
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
 
@@ -181,6 +190,7 @@ def get_config(self):
                 "rope_scaling_factor": self.rope_scaling_factor,
                 "num_key_value_heads": self.num_key_value_heads,
                 "sliding_window": self.sliding_window,
+                "head_dim": self.head_dim,
                 "layer_norm_epsilon": self.layer_norm_epsilon,
                 "dropout": self.dropout,
             }

keras_hub/src/models/mistral/mistral_backbone_test.py

@@ -42,6 +42,26 @@ def test_num_parameters(self):
         # Reference value calculated using the PyTorch model
         self.assertEqual(model.count_params(), 2704)
 
+    def test_explicit_head_dim(self):
+        # Magistral-style config: `head_dim` is set explicitly and does not
+        # equal `hidden_dim // num_query_heads`. `sliding_window=None` is
+        # also exercised here. Run the full backbone test so the new path
+        # gets serialization and precision coverage.
+        init_kwargs = {
+            **self.init_kwargs,
+            "sliding_window": None,
+            "head_dim": 4,
+        }
+        self.run_backbone_test(
+            cls=MistralBackbone,
+            init_kwargs=init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(2, 5, 16),
+        )
+        model = MistralBackbone(**init_kwargs)
+        attention = model.transformer_layers[0]._self_attention_layer
+        self.assertEqual(attention._head_dim, 4)
+
     @pytest.mark.extra_large
     def test_smallest_preset(self):
         self.run_preset_test(

keras_hub/src/models/mistral/mistral_causal_lm.py

@@ -100,7 +100,9 @@ def _build_cache(self, token_ids):
         max_length = ops.shape(token_ids)[1]
         num_layers = self.backbone.num_layers
         num_key_value_heads = self.backbone.num_key_value_heads
-        head_dim = self.backbone.hidden_dim // self.backbone.num_query_heads
+        head_dim = self.backbone.head_dim or (
+            self.backbone.hidden_dim // self.backbone.num_query_heads
+        )
         shape = [
             batch_size,
             num_layers,

keras_hub/src/models/mistral/mistral_causal_lm_test.py

@@ -47,6 +47,30 @@ def test_causal_lm_basics(self):
             expected_output_shape=(2, 8, 10),
         )
 
+    def test_generate_with_explicit_head_dim(self):
+        # Magistral-style config where `head_dim` is set explicitly and does
+        # not equal `hidden_dim // num_query_heads`. The generation cache must
+        # be built with the explicit `head_dim`, not the derived value.
+        backbone = MistralBackbone(
+            vocabulary_size=self.preprocessor.tokenizer.vocabulary_size(),
+            num_layers=2,
+            num_query_heads=4,
+            num_key_value_heads=2,
+            hidden_dim=8,
+            intermediate_dim=16,
+            head_dim=6,
+            sliding_window=None,
+        )
+        self.assertNotEqual(
+            backbone.head_dim, backbone.hidden_dim // backbone.num_query_heads
+        )
+        causal_lm = MistralCausalLM(
+            preprocessor=self.preprocessor, backbone=backbone
+        )
+        prompt = "the quick brown fox"
+        output = causal_lm.generate(prompt)
+        self.assertTrue(prompt in output)
+
     def test_generate(self):
         causal_lm = MistralCausalLM(**self.init_kwargs)
         # String input.

keras_hub/src/models/mistral/mistral_tokenizer.py

@@ -1,8 +1,23 @@
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.mistral.mistral_backbone import MistralBackbone
+from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+try:
+    import tokenizers as hf_tokenizers
+    from tokenizers import decoders
+    from tokenizers import models as hf_models
+    from tokenizers import pre_tokenizers
+except ImportError:
+    hf_tokenizers = None
 
 
 @keras_hub_export(
@@ -20,6 +35,10 @@ class MistralTokenizer(SentencePieceTokenizer):
     Mistral models and provides a `from_preset()` method to automatically
     download a matching vocabulary for a Mistral preset.
 
+    This tokenizer is used by SentencePiece-based Mistral presets. Presets that
+    ship a Tekken (byte-level BPE) vocabulary, such as Magistral, use
+    `keras_hub.tokenizers.MistralTekkenTokenizer` instead.
+
     If input is a batch of strings (rank > 0), the layer will output a
     `tf.RaggedTensor` where the last dimension of the output is ragged.
 
@@ -55,3 +74,159 @@ def __init__(self, proto, **kwargs):
         self._add_special_token("</s>", "end_token")
         self.pad_token_id = 0
         super().__init__(proto=proto, **kwargs)
+
+
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.MistralTekkenTokenizer",
+        "keras_hub.models.MistralTekkenTokenizer",
+    ]
+)
+class MistralTekkenTokenizer(BytePairTokenizer):
+    """Mistral Tekken tokenizer layer based on byte-level BPE.
+
+    This tokenizer class handles Mistral's Tekken (tiktoken-style byte-level
+    BPE) vocabulary, used by presets such as Magistral. It is based on
+    `keras_hub.tokenizers.BytePairTokenizer`, but uses the Tekken
+    pre-tokenization regex instead of the GPT-2/Llama3 pattern hardcoded in the
+    base class, and checks for the special tokens needed by Mistral models.
+
+    The `vocabulary` and `merges` are usually produced from a `tekken.json`
+    file by the Hugging Face conversion path; see
+    `keras_hub.src.utils.transformers.convert_mistral`.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: A dict mapping token strings to integer ids, or a path to a
+            vocabulary JSON file.
+        merges: A list of BPE merge rules, or a path to a merges file.
+        split_pattern: str. The Tekken pre-tokenization regex.
+
+    Examples:
+    ```python
+    tokenizer = keras_hub.models.MistralTekkenTokenizer.from_preset(
+        "hf://mistralai/Magistral-Small-2506",
+    )
+    tokenizer("The quick brown fox jumped.")
+    tokenizer.detokenize(tokenizer("The quick brown fox jumped."))
+    ```
+    """
+
+    backbone_cls = MistralBackbone
+
+    def __init__(
+        self, vocabulary=None, merges=None, split_pattern=None, **kwargs
+    ):
+        self.split_pattern = split_pattern
+        self._add_special_token("<s>", "start_token")
+        self._add_special_token("</s>", "end_token")
+        self.pad_token_id = 0
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            unsplittable_tokens=[self.start_token, self.end_token],
+            **kwargs,
+        )
+
+    def _set_vocabulary_and_merges_tokenizers(self, vocabulary, merges):
+        self.vocabulary = vocabulary.copy()
+        self.merges = list(merges)
+        _merges = []
+        for merge in self.merges:
+            if "#version:" in merge.lstrip():
+                continue
+            a, b = str(merge).split(" ")
+            _merges.append((a, b))
+        self._tokenizer = hf_tokenizers.Tokenizer(
+            hf_models.BPE(vocab=vocabulary, merges=_merges, fuse_unk=False)
+        )
+        if self.unsplittable_tokens:
+            self._tokenizer.add_special_tokens(self.unsplittable_tokens)
+        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
+            [
+                pre_tokenizers.Split(
+                    hf_tokenizers.Regex(self.split_pattern),
+                    behavior="isolated",
+                ),
+                pre_tokenizers.ByteLevel(
+                    add_prefix_space=self.add_prefix_space, use_regex=False
+                ),
+            ]
+        )
+        self._tokenizer.decoder = decoders.ByteLevel()
+
+        # Dummy attrs for serialization compatibility with the base class.
+        if not hasattr(self, "cache"):
+            self.byte2unicode = None
+            self.unicode2byte = None
+            self.cache = None
+            self.id_to_token_map = None
+            self.token_to_id_map = None
+            self.merge_ranks_lookup_default = None
+            self.merge_ranks = None
+
+    def _set_vocabulary_and_merges_tf(self, vocabulary, merges):
+        # The base class hardcodes the GPT-2 split regex in its `tf.data`
+        # path, which does not match Tekken. We instead bridge to the
+        # `tokenizers` backend from within the graph (see `_tokenize_tf`), so
+        # there is nothing to build here.
+        self.vocabulary = vocabulary.copy()
+        self.merges = list(merges)
+
+    def _maybe_initialized_tokenizers(self):
+        if getattr(self, "_tokenizer", None) is None:
+            self._set_vocabulary_and_merges_tokenizers(
+                self.vocabulary, self.merges
+            )
+
+    @preprocessing_function
+    def _tokenize_tf(self, inputs):
+        self._maybe_initialized_tokenizers()
+
+        def _encode(string_tensor):
+            values = string_tensor.numpy()
+            strings = [v.decode("utf-8") for v in values.tolist()]
+            encodings = self._tokenizer.encode_batch(
+                strings, add_special_tokens=False
+            )
+            return tf.ragged.constant(
+                [e.ids for e in encodings], dtype=self.compute_dtype
+            )
+
+        inputs = tf.convert_to_tensor(inputs)
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
+            inputs = tf.expand_dims(inputs, 0)
+        tokens = tf.py_function(
+            _encode,
+            [inputs],
+            Tout=tf.RaggedTensorSpec(
+                shape=[None, None],
+                dtype=self.compute_dtype,
+                ragged_rank=1,
+            ),
+        )
+
+        if self.sequence_length:
+            output_shape = tokens.shape.as_list()
+            output_shape[-1] = self.sequence_length
+            tokens = tokens.to_tensor(
+                shape=output_shape,
+                default_value=getattr(self, "pad_token_id", 0),
+            )
+        if unbatched:
+            tokens = tokens[0]
+        return tokens
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"split_pattern": self.split_pattern})
+        # `unsplittable_tokens` is derived from the special tokens in the
+        # constructor, so it is not a separate config argument.
+        del config["unsplittable_tokens"]
+        return config