Clean up rope params; ensure transformers 4.55/5.0 compatibility

- Remove duplicate top-level rope_scaling block and stray rope_theta from config.json
- Remove duplicate 'type' key from rope_parameters
- For 3B-Base/8B-Base: set max_position_embeddings=4096 and factor=0.25 to match training
- Mirror rope_theta and rope_scaling from rope_parameters in MinistralDLMConfig for v4.55 yarn
- Drop unused sdpa_mask_older_torch import (removed in transformers v5.0)
- Bump transformers_version to 5.0.0
- In linear_spec_generate_mp, guard direct past_kv.key_cache / value_cache access behind a hasattr(past_kv, 'layers') check so v5.0's DynamicCache API works too

config.json

@@ -54,27 +54,13 @@
     "mscale_all_dim": 1.0,
     "original_max_position_embeddings": 16384,
     "rope_theta": 1000000000.0,
-    "rope_type": "yarn",
-    "type": "yarn"
+    "rope_type": "yarn"
   },
-  "rope_scaling": {
-    "beta_fast": 32.0,
-    "beta_slow": 1.0,
-    "factor": 16.0,
-    "llama_4_scaling_beta": 0.1,
-    "mscale": 1.0,
-    "mscale_all_dim": 1.0,
-    "original_max_position_embeddings": 16384,
-    "rope_theta": 1000000.0,
-    "rope_type": "yarn",
-    "type": "yarn"
-  },
-  "rope_theta": 1000000000.0,
   "sliding_window": null,
   "tie_word_embeddings": false,
   "tok_mask_half_life_ratio": null,
   "torch_dtype": "bfloat16",
-  "transformers_version": "4.55.4",
+  "transformers_version": "5.0.0",
   "use_cache": false,
   "vocab_size": 131072
 }

configuration_ministral_dlm.py

@@ -156,7 +156,6 @@ class MinistralDLMConfig(PretrainedConfig):
         tie_word_embeddings=False,
         rope_theta=1000000.0,
         rope_parameters=None,
-        rope_scaling=None,
         attention_bias=False,
         attention_dropout=0.0,
         mlp_bias=False,
@@ -204,9 +203,11 @@ class MinistralDLMConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.rms_norm_eps = rms_norm_eps
         self.use_cache = use_cache
-        self.rope_theta = rope_theta
         self.rope_parameters = rope_parameters
-        self.rope_scaling = rope_scaling
+        # `rope_theta` is read at the top level by transformers v4.55's yarn impl; mirror from rope_parameters when present.
+        self.rope_theta = (rope_parameters or {}).get("rope_theta", rope_theta)
+        # v4.55 reads rope params from `rope_scaling`; in v5.0 `rope_scaling` is a property alias for rope_parameters.
+        self.rope_scaling = rope_parameters
         self.attention_bias = attention_bias
         self.attention_dropout = attention_dropout
         self.mlp_bias = mlp_bias

modeling_ministral.py

@@ -11,7 +11,7 @@ from transformers.cache_utils import Cache, DynamicCache
 from transformers.generation import GenerationMixin
 # from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask, ALL_MASK_ATTENTION_FUNCTIONS, sdpa_mask_older_torch
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask, ALL_MASK_ATTENTION_FUNCTIONS
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import (
     GenericForQuestionAnswering,

modeling_ministral_dlm.py

@@ -1489,9 +1489,15 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
                     layer.self_attn.diffusion_lm = val
 
         def _crop_cache(kv, length):
+            # transformers 4.55 exposes .key_cache/.value_cache lists; 5.0 moved them under .layers[i].keys/.values.
             for li in range(len(kv)):
-                kv.key_cache[li] = kv.key_cache[li][:, :, :length]
-                kv.value_cache[li] = kv.value_cache[li][:, :, :length]
+                if hasattr(kv, 'layers'):
+                    layer = kv.layers[li]
+                    layer.keys = layer.keys[:, :, :length]
+                    layer.values = layer.values[:, :, :length]
+                else:
+                    kv.key_cache[li] = kv.key_cache[li][:, :, :length]
+                    kv.value_cache[li] = kv.value_cache[li][:, :, :length]
             kv._seen_tokens = length
 
         # ----- tree verify helpers (inlined) -----
@@ -1546,9 +1552,14 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
                     candidate_blocks[pi, p] = combo[ci]
 
             # Expand KV cache batch dimension (shared, no copy)
-            for li in range(len(past_kv.key_cache)):
-                past_kv.key_cache[li] = past_kv.key_cache[li].expand(num_paths, -1, -1, -1)
-                past_kv.value_cache[li] = past_kv.value_cache[li].expand(num_paths, -1, -1, -1)
+            for li in range(len(past_kv)):
+                if hasattr(past_kv, 'layers'):
+                    layer = past_kv.layers[li]
+                    layer.keys = layer.keys.expand(num_paths, -1, -1, -1)
+                    layer.values = layer.values.expand(num_paths, -1, -1, -1)
+                else:
+                    past_kv.key_cache[li] = past_kv.key_cache[li].expand(num_paths, -1, -1, -1)
+                    past_kv.value_cache[li] = past_kv.value_cache[li].expand(num_paths, -1, -1, -1)
 
             # Batched causal verify — uses flash attention + GQA
             _set_dlm(False)
@@ -1583,9 +1594,14 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             accepted_toks = ar_tokens[best_pidx:best_pidx+1, :best_acc]
 
             # Extract winning path's KV cache slice
-            for li in range(len(past_kv.key_cache)):
-                past_kv.key_cache[li] = past_kv.key_cache[li][best_pidx:best_pidx+1].contiguous()
-                past_kv.value_cache[li] = past_kv.value_cache[li][best_pidx:best_pidx+1].contiguous()
+            for li in range(len(past_kv)):
+                if hasattr(past_kv, 'layers'):
+                    layer = past_kv.layers[li]
+                    layer.keys = layer.keys[best_pidx:best_pidx+1].contiguous()
+                    layer.values = layer.values[best_pidx:best_pidx+1].contiguous()
+                else:
+                    past_kv.key_cache[li] = past_kv.key_cache[li][best_pidx:best_pidx+1].contiguous()
+                    past_kv.value_cache[li] = past_kv.value_cache[li][best_pidx:best_pidx+1].contiguous()
             _crop_cache(past_kv, cache_len + best_acc)
 
             return accepted_toks, best_acc, past_kv, accepted_toks[:, -1:]