Upload model

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

chat_utils.py

@@ -0,0 +1,267 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def add_gumbel_noise(logits, temperature):
+    '''
+    The Gumbel max is a method for sampling categorical distributions.
+    According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality.
+    Thus, we use float64.
+    '''
+    if temperature == 0:
+        return logits
+    logits = logits.to(torch.float64)
+    noise = torch.rand_like(logits, dtype=torch.float64)
+    gumbel_noise = (- torch.log(noise)) ** temperature
+    return logits.exp() / gumbel_noise
+
+
+def get_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, neg_entropy=False):
+    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+    x0 = torch.argmax(logits_with_noise, dim=-1)
+
+    if remasking == 'low_confidence':
+        # p = F.softmax(logits.to(torch.float64), dim=-1)
+        p = F.softmax(logits, dim=-1)
+        x0_p = torch.squeeze(
+            torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
+    elif remasking == 'top_p_margin':
+        # Compute probabilities
+        p = F.softmax(logits, dim=-1)                       # (B, L, V)
+        # Top-2 per position
+        top2 = torch.topk(p, k=2, dim=-1).values            # (B, L, 2)
+        margin = top2[..., 0] - top2[..., 1]                # (B, L)
+
+        # Normalize margin to [0,1] over MASKED positions per row
+        plus_inf  = torch.full_like(margin, float('inf'))
+        minus_inf = torch.full_like(margin, float('-inf'))
+        masked_for_min = torch.where(mask_index, margin, plus_inf)
+        masked_for_max = torch.where(mask_index, margin, minus_inf)
+        row_min = masked_for_min.amin(dim=1, keepdim=True)  # (B, 1)
+        row_max = masked_for_max.amax(dim=1, keepdim=True)  # (B, 1)
+        denom = (row_max - row_min)
+
+        # If denom==0 (all equal), set normalized=1 on masked; 0 elsewhere by default
+        normalized = torch.zeros_like(margin)
+        nonzero = denom > 0
+        normalized = torch.where(
+            mask_index & nonzero,
+            (margin - row_min) / (denom + 1e-12),
+            normalized
+        )
+        normalized = torch.where(
+            mask_index & (~nonzero),
+            torch.ones_like(normalized),
+            normalized
+        )
+        x0_p = normalized  # ∈ [0,1] on masked positions
+    elif remasking == 'random':
+        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+    else:
+        raise NotImplementedError(remasking)
+    
+    # Calculate negative entropy if requested
+    if neg_entropy:
+        # p = F.softmax(logits.to(torch.float64), dim=-1)
+        p = F.softmax(logits, dim=-1)
+        epsilon = 1e-10
+        log_probs = torch.log(p + epsilon)
+        confidence_scores = torch.sum(p * log_probs, dim=-1)  # negative entropy per position
+    else:
+        confidence_scores = x0_p
+    
+    x0 = torch.where(mask_index, x0, x)
+    confidence = torch.where(mask_index, confidence_scores, -np.inf)
+
+    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+    if threshold is not None:
+        num_transfer_tokens = mask_index.sum(dim=1, keepdim=True)
+    # print(f'confidence: {confidence}')
+    for j in range(confidence.shape[0]):
+        _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j])
+        transfer_index[j, select_index] = True
+        if threshold is not None:
+            for k in range(1, num_transfer_tokens[j]):
+                if confidence[j, select_index[k]] < threshold:
+                    transfer_index[j, select_index[k]] = False
+    return x0, transfer_index
+
+
+def get_num_transfer_tokens(mask_index, steps: int):
+    mask_num = mask_index.sum(dim=1, keepdim=True)
+    base = mask_num // steps
+    remainder = mask_num % steps
+    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
+    for i in range(mask_num.size(0)):
+        num_transfer_tokens[i, : int(remainder[i])] += 1
+    return num_transfer_tokens
+
+
+@torch.no_grad()
+def generate_with_prefix_cache_block_diff(
+    model,
+    prompt,
+    steps=128,
+    gen_length=128,
+    block_length=128,
+    temperature=0.,
+    remasking='low_confidence',
+    mask_id=126336,
+    threshold=None,
+    factor=None,
+    shift_logits=False,
+    neg_entropy=False,
+    causal_context=False,
+    eos_token_id=None,
+):
+    dream_style=shift_logits
+    # Initialize the accumulator
+    x_accum = prompt.clone()
+
+    assert gen_length % block_length == 0
+    num_blocks = gen_length // block_length
+
+    assert steps % num_blocks == 0
+    steps_per_block = steps // num_blocks
+
+    nfe = 0
+
+    if causal_context:
+        model_module = model.module if hasattr(model, "module") else model
+        for layer in model_module.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm=False
+
+    # Compute KV cache for the prompt initially
+    output = model(prompt, use_cache=True, use_causal_mask=causal_context)
+    past_key_values = output.past_key_values
+
+    if causal_context:
+        for layer in model_module.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm=True
+
+    # For dream_style: store the "next token logit" of the context
+    next_logits_context = None
+    if dream_style:
+        next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)
+
+    for num_block in range(num_blocks):
+        # Create a new block with mask tokens (no seeding)
+        mask_block = torch.ones(
+            (prompt.shape[0], block_length),
+            dtype=prompt.dtype,
+            device=prompt.device
+        ) * mask_id
+
+        # Append the block of masks
+        x_accum = torch.cat([x_accum, mask_block], dim=1)
+        current_block_start = prompt.size(1) + num_block * block_length
+        block_slice = slice(current_block_start, current_block_start + block_length)
+
+        # Build the initial mask for this block
+        mask_block_idx0 = (x_accum[:, block_slice] == mask_id)  # (B, Lb)
+
+        # Precompute the transfer schedule for this block
+        if dream_style:
+            # still denoise *all* positions (0..Lb-1), since none are seeded
+            schedule_mask = mask_block_idx0
+        else:
+            schedule_mask = mask_block_idx0
+
+        num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block)  # (B, steps)
+
+        # Denoise the current block
+        for i in range(steps_per_block):
+            mask_block_idx = (x_accum[:, block_slice] == mask_id)  # (B, Lb)
+            if mask_block_idx.sum() == 0:
+                break
+
+            nfe += 1
+
+            # Forward only the current noisy block using cached context
+            logits_block = model(
+                x_accum[:, block_slice],
+                past_key_values=past_key_values,
+                use_cache=False
+            ).logits
+
+            if dream_style:
+                # Align logits so that each masked position has a predictor:
+                # prepend context-next logit, then use logits_block[:-1]
+                if block_length == 1:
+                    logits_use = next_logits_context              # (B, 1, V)
+                else:
+                    logits_use = torch.cat(
+                        [next_logits_context, logits_block[:, :-1, :]],
+                        dim=1
+                    )  # (B, Lb, V)
+
+                mask_use = mask_block_idx                        # (B, Lb)
+                x_use   = x_accum[:, block_slice]                # (B, Lb)
+
+                x0, transfer_idx = get_transfer_index(
+                    logits_use, temperature, remasking, mask_use, x_use,
+                    num_transfer_tokens=num_transfer_tokens[:, i],
+                    threshold=threshold, neg_entropy=neg_entropy
+                )
+                cur = x_accum[:, block_slice].clone()
+                cur[transfer_idx] = x0[transfer_idx]
+                x_accum[:, block_slice] = cur
+
+            else:
+                # non-AR (same-position) case
+                x0, transfer_idx = get_transfer_index(
+                    logits_block, temperature, remasking, mask_block_idx,
+                    x_accum[:, block_slice],
+                    num_transfer_tokens=num_transfer_tokens[:, i],
+                    threshold=threshold, neg_entropy=neg_entropy
+                )
+                cur = x_accum[:, block_slice].clone()
+                cur[transfer_idx] = x0[transfer_idx]
+                x_accum[:, block_slice] = cur
+
+            if eos_token_id is not None:
+                block_tokens = x_accum[:, block_slice]              # (B, Lb)
+                eos_mask = (block_tokens == eos_token_id)           # (B, Lb)
+                any_eos = eos_mask.any(dim=1)                       # (B,)
+                if any_eos.any():
+                    after_eos = eos_mask.cumsum(dim=1).bool()       # (B, Lb)
+                    mask_before = (block_tokens == mask_id) & ~after_eos
+                    if (any_eos & ~mask_before.any(dim=1)).any():
+                        break
+
+        if causal_context:
+            for layer in model_module.encoder.layers:
+                if hasattr(layer.self_attn, 'diffusion_lm'):
+                    layer.self_attn.diffusion_lm=False
+
+        # after block is fully denoised, update KV cache
+        output = model(
+            x_accum[:, block_slice],
+            past_key_values=past_key_values,
+            use_cache=True,
+            use_causal_mask=causal_context
+        )
+        past_key_values = output.past_key_values
+
+        if causal_context:
+            for layer in model_module.encoder.layers:
+                if hasattr(layer.self_attn, 'diffusion_lm'):
+                    layer.self_attn.diffusion_lm=True
+
+        if dream_style and num_block < num_blocks - 1:
+            # refresh context-next logit for the next block
+            next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)
+
+        if eos_token_id is not None:
+            gen_so_far = x_accum[:, prompt.size(1):]                    # (B, gen_len_so_far)
+            is_eos = (gen_so_far == eos_token_id)                       # (B, gen_len_so_far)
+            has_eos = is_eos.any(dim=1)                                 # (B,)
+            if has_eos.all():
+                first_eos_pos = is_eos.to(torch.int64).argmax(dim=1)    # (B,)
+                max_eos = first_eos_pos.max().item()
+                return x_accum[:, : prompt.size(1) + max_eos + 1], nfe
+
+    return x_accum, nfe

config.json

@@ -0,0 +1,79 @@
+{
+  "ada_dlm_loss_ratio": null,
+  "ada_perm_ratio_global": null,
+  "ada_perm_ratio_per_block": null,
+  "adaptive_mask_rate": false,
+  "ar_loss_weight": 1.0,
+  "architectures": [
+    "MinistralDiffEncoderModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_implementation": "sdpa",
+  "auto_map": {
+    "AutoConfig": "configuration_ministral_dlm.MinistralDLMConfig",
+    "AutoModel": "modeling_ministral_dlm.MinistralDiffEncoderModel"
+  },
+  "block_size": 32,
+  "bos_token_id": 1,
+  "diff_loss_weight": 1,
+  "dlm_arch": "encoder",
+  "dlm_loss_weight": null,
+  "dlm_paradigm": "bidirectional",
+  "dlm_type": "llada",
+  "dp_varying_mask_ratio": false,
+  "enforce_mask": false,
+  "eos_token_id": 2,
+  "global_loss_avg": false,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 3072,
+  "initializer_range": 0.02,
+  "intermediate_size": 9216,
+  "mask_token_id": 100,
+  "max_position_embeddings": 262144,
+  "mlp_bias": false,
+  "model_type": "ministral_dlm",
+  "multi_sampling": null,
+  "num_ar_layers": 0,
+  "num_attention_heads": 32,
+  "num_diffusion_layers": 0,
+  "num_hidden_layers": 26,
+  "num_key_value_heads": 8,
+  "num_skip_loss_tokens": 0,
+  "prefix_ratio": 0.8,
+  "random_length_prob": 0,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "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_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": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "tok_mask_half_life_ratio": null,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.55.4",
+  "use_cache": false,
+  "vocab_size": 131072
+}

configuration_ministral_dlm.py

@@ -0,0 +1,247 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Ministral DLM model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MinistralDLMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Ministral3Model`] for diffusion language models.
+    It is used to instantiate a Ministral model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 131072):
+            Vocabulary size of the Ministral model.
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 14336):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 34):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            Number of key_value heads for Grouped Query Attention.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function.
+        max_position_embeddings (`int`, *optional*, defaults to 262144):
+            The maximum sequence length.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        rope_parameters (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings.
+            Default uses YaRN scaling with factor=16, original_max_position_embeddings=16384.
+        attention_bias (`bool`, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in up_proj, down_proj and gate_proj layers.
+        sliding_window (`int`, *optional*, defaults to None):
+            Sliding window attention size.
+        mask_token_id (`int`, *optional*, defaults to -1):
+            Token ID for masking in diffusion.
+        dlm_type (`str`, *optional*, defaults to 'llada'):
+            Type of diffusion language model ('llada', 'dream').
+        random_length_prob (`float`, *optional*):
+            Probability of using random lengths during training.
+        num_ar_layers (`int`, *optional*, defaults to 0):
+            Number of autoregressive layers.
+        num_diffusion_layers (`int`, *optional*, defaults to 0):
+            Number of diffusion layers.
+        diff_loss_weight (`float`, *optional*, defaults to 1):
+            Weight for diffusion loss.
+        enforce_mask (`bool`, *optional*, defaults to False):
+            Whether to enforce masking.
+        prefix_ratio (`float`, *optional*, defaults to 0.8):
+            Ratio for prefix in prefix_bidirectional mode.
+        dlm_paradigm (`str`, *optional*, defaults to 'bidirectional'):
+            Paradigm for diffusion ('bidirectional', 'autoregressive', 'prefix_bidirectional', 'efficient_block_diff', 'block_diff', 'sbd_block_diff').
+        dlm_arch (`str`, *optional*, defaults to 'encoder'):
+            Architecture type ('encoder', 'encoder_decoder').
+        block_size (`int`, *optional*, defaults to 32):
+            Block size for block diffusion paradigms.
+        tok_mask_half_life_ratio (`float`, *optional*):
+            Half-life ratio for token masking.
+        adaptive_mask_rate (`bool`, *optional*, defaults to False):
+            Whether to use adaptive mask rate.
+        multi_sampling (`int`, *optional*):
+            Number of samples for multi-sampling.
+        num_skip_loss_tokens (`int`, *optional*, defaults to 0):
+            Number of tokens to skip in loss calculation.
+        dlm_loss_weight (`float`, *optional*):
+            Weight for diffusion LM loss.
+        ar_loss_weight (`float`, *optional*, defaults to 1.0):
+            Weight for autoregressive loss in sbd_block_diff paradigm. Use 10000 to only use AR loss.
+        global_loss_avg (`bool`, *optional*, defaults to False):
+            Whether to use global loss average.
+        dp_varying_mask_ratio (`bool`, *optional*, defaults to False):
+            Whether to use varying mask ratio for each DP rank during sampling.
+        ada_perm_ratio_per_block (`float`, *optional*):
+            Adaptive permutation ratio for each block.
+        ada_perm_ratio_global (`float`, *optional*):
+            Adaptive permutation ratio for global.
+    """
+
+    model_type = "ministral_dlm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Ministral`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        hidden_size=4096,
+        intermediate_size=14336,
+        num_hidden_layers=34,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=262144,
+        initializer_range=0.02,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        rope_parameters=None,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        sliding_window=None,
+        attn_implementation="sdpa",
+        mask_token_id=-1,
+        dlm_type='llada',
+        random_length_prob=None,
+        num_ar_layers=0,
+        num_diffusion_layers=0,
+        diff_loss_weight=1,
+        enforce_mask=False,
+        prefix_ratio=0.8,
+        dlm_paradigm='bidirectional',
+        dlm_arch='encoder',
+        block_size=32,
+        tok_mask_half_life_ratio=None,
+        adaptive_mask_rate=False,
+        multi_sampling=None,
+        num_skip_loss_tokens=0,
+        dlm_loss_weight=None,
+        ar_loss_weight=1.0,
+        global_loss_avg=False,
+        dp_varying_mask_ratio=False,
+        ada_perm_ratio_per_block=None,
+        ada_perm_ratio_global=None,
+        ada_dlm_loss_ratio=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        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
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.sliding_window = sliding_window
+        
+        rope_config_validation(self)
+        
+        self.attn_implementation = attn_implementation
+        
+        self.mask_token_id = mask_token_id
+        self.dlm_type = dlm_type
+        self.random_length_prob = random_length_prob
+        self.num_ar_layers = num_ar_layers
+        self.num_diffusion_layers = num_diffusion_layers
+        self.diff_loss_weight = diff_loss_weight
+        self.enforce_mask = enforce_mask
+        self.prefix_ratio = prefix_ratio
+        self.dlm_paradigm = dlm_paradigm
+        self.dlm_arch = dlm_arch
+        self.block_size = block_size
+        self.tok_mask_half_life_ratio = tok_mask_half_life_ratio
+        self.adaptive_mask_rate = adaptive_mask_rate
+        self.multi_sampling = multi_sampling
+        self.num_skip_loss_tokens = num_skip_loss_tokens
+        self.dlm_loss_weight = dlm_loss_weight
+        self.ar_loss_weight = ar_loss_weight
+        self.global_loss_avg = global_loss_avg
+        self.dp_varying_mask_ratio = dp_varying_mask_ratio
+        self.ada_perm_ratio_per_block = ada_perm_ratio_per_block
+        self.ada_perm_ratio_global = ada_perm_ratio_global
+        self.ada_dlm_loss_ratio = ada_dlm_loss_ratio
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["MinistralDLMConfig"]
+

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.55.4",
+  "use_cache": false
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6510fd1ca747463dc8077e1b144f9bc3a1d64c45331d9fff945a00d914cf7cfc
+size 7663347000

modeling_ministral.py

@@ -0,0 +1,551 @@
+from collections.abc import Callable
+from typing import Optional, Union
+
+import torch
+from torch import nn
+
+from transformers.utils.generic import check_model_inputs
+
+from transformers.activations import ACT2FN
+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.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
+# from transformers.utils.generic import maybe_autocast
+from .configuration_ministral_dlm import MinistralDLMConfig
+
+#ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask_older_torch
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+# @use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def _get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
+    scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
+    return scaling.unsqueeze(-1)
+
+
+# @use_kernelized_func(apply_rotary_pos_emb)
+class Ministral3Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+
+        self.diffusion_lm = config.diffusion_lm
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        query_states = query_states * _get_llama_4_attn_scale(
+            cache_position,
+            self.config.rope_parameters.get("llama_4_scaling_beta"),
+            self.config.rope_parameters.get("original_max_position_embeddings"),
+        ).to(query_states.dtype)
+
+        if past_key_values is not None:
+            if use_cache:
+                # sin and cos are specific to RoPE models; cache_position needed for the static cache
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            else:  ## if use_cache == False, do not update cache
+                old_k, old_v = past_key_values.layers[self.layer_idx].keys, past_key_values.layers[self.layer_idx].values
+                key_states   = torch.cat([old_k, key_states], dim=-2)
+                value_states = torch.cat([old_v, value_states], dim=-2)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        if self.diffusion_lm:
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                None,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                is_causal=False,
+                **kwargs,
+            )
+
+        else:
+            attn_output, attn_weights = attention_interface(
+                self,
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=getattr(self.config, "sliding_window", None),  # main diff with Llama
+                **kwargs,
+            )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Ministral3MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Ministral3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Ministral3RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Ministral3DecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+    
+        if hasattr(config, 'attn_class'):
+            attn_class = config.attn_class
+        else:
+            attn_class = Ministral3Attention
+
+        self.self_attn = attn_class(config=config, layer_idx=layer_idx)
+        self.mlp = Ministral3MLP(config)
+        self.input_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+class Ministral3PreTrainedModel(PreTrainedModel):
+    config: MinistralDLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Ministral3DecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": Ministral3DecoderLayer,
+        "attentions": Ministral3Attention,
+    }
+
+
+class Ministral3RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: MinistralDLMConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = inv_freq
+
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: Optional[MinistralDLMConfig] = None,
+        device: Optional["torch.device"] = None,
+        seq_len: Optional[int] = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        # device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        # with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+
+        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        cos = emb.cos() * self.attention_scaling
+        sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class Ministral3Model(Ministral3PreTrainedModel):
+    def __init__(self, config: MinistralDLMConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Ministral3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Ministral3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            # past_key_values = DynamicCache(config=self.config)
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if kwargs.get("use_causal_mask", False):
+            mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
+            causal_mask = mask_function(
+                config=self.config,
+                input_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+            )
+            
+        else:
+            causal_mask = None
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+@auto_docstring
+class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Ministral3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Ministral3ForCausalLM
+
+        >>> model = Ministral3ForCausalLM.from_pretrained("meta-ministral3/Ministral3-2-7b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-ministral3/Ministral3-2-7b-hf")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class Ministral3ForTokenClassification(GenericForTokenClassification, Ministral3PreTrainedModel):
+    pass
+
+
+class Ministral3ForSequenceClassification(GenericForSequenceClassification, Ministral3PreTrainedModel):
+    pass
+
+
+class Ministral3ForQuestionAnswering(GenericForQuestionAnswering, Ministral3PreTrainedModel):
+    pass
+
+
+__all__ = [
+    "Ministral3ForCausalLM",
+    "Ministral3ForQuestionAnswering",
+    "Ministral3Model",
+    "Ministral3PreTrainedModel",
+    "Ministral3ForSequenceClassification",
+    "Ministral3ForTokenClassification",
+]

modeling_ministral_dlm.py

@@ -0,0 +1,1112 @@
+import copy
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple, Union
+import random
+import os
+import sys
+import json
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput
+from transformers.utils import ModelOutput
+
+from torch.nn.attention.flex_attention import BlockMask, flex_attention, create_block_mask, or_masks
+
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+
+from transformers.processing_utils import Unpack
+
+from transformers.cache_utils import Cache, DynamicCache
+
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.generation import GenerationMixin
+
+import math
+
+from .chat_utils import generate_with_prefix_cache_block_diff
+from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
+from .configuration_ministral_dlm import MinistralDLMConfig
+
+
+@dataclass
+class MinistralDiffOutputWithPast(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    causal_logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+# @torch.compile(dynamic=True, mode="reduce-overhead")
+# @torch.compile(mode="default")
+# @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
+@torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False)
+def fused_flex_attention(q, k, v, block_mask=None):
+    return flex_attention(q, k, v, block_mask=block_mask)
+
+
+def _crop_dynamic_cache(past_key_values: DynamicCache, max_length: int):
+    """Crop a DynamicCache to max_length, compatible with both old and new transformers."""
+    if hasattr(past_key_values, 'crop'):
+        past_key_values.crop(max_length)
+    else:
+        for layer_idx in range(len(past_key_values)):
+            past_key_values.key_cache[layer_idx] = past_key_values.key_cache[layer_idx][:, :, :max_length]
+            past_key_values.value_cache[layer_idx] = past_key_values.value_cache[layer_idx][:, :, :max_length]
+        past_key_values._seen_tokens = max_length
+
+
+def _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block_length: int):
+    """After quadratic decoding, extract only draft tokens (first of each block) from cache."""
+    for layer_idx in range(len(past_key_values)):
+        if hasattr(past_key_values, 'layers'):
+            layer_cache = past_key_values.layers[layer_idx]
+            k, v = layer_cache.keys, layer_cache.values
+        else:
+            k = past_key_values.key_cache[layer_idx]
+            v = past_key_values.value_cache[layer_idx]
+
+        clean_k, draft_k = k[:, :, :clean_len], k[:, :, clean_len::block_length + 1]
+        clean_v, draft_v = v[:, :, :clean_len], v[:, :, clean_len::block_length + 1]
+        new_k = torch.cat([clean_k, draft_k], dim=2)
+        new_v = torch.cat([clean_v, draft_v], dim=2)
+
+        if hasattr(past_key_values, 'layers'):
+            layer_cache.keys = new_k
+            layer_cache.values = new_v
+        else:
+            past_key_values.key_cache[layer_idx] = new_k
+            past_key_values.value_cache[layer_idx] = new_v
+
+    past_key_values._seen_tokens = clean_len + block_length
+
+
+# with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb
+class MinistralFlexAttention(Ministral3Attention):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+ 
+        self.block_size_orig = self.config.block_size
+
+        if self.config.dlm_paradigm == 'bidirectional':
+            self.bidirectional_mask = self.compute_block_mask(mode='bidirectional')
+        elif self.config.dlm_paradigm == 'autoregressive':
+            self.autoregressive_mask = self.compute_block_mask(mode='autoregressive')
+        elif self.config.dlm_paradigm == 'block_diff':
+            self.block_diff_mask = None
+        elif self.config.dlm_paradigm == 'sbd_block_diff':
+            self.sbd_block_diff_mask = None
+        else:
+            raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}")
+
+        self.block_size = self.block_size_orig
+        self.mode = self.config.dlm_paradigm
+        self._quadratic_block_mask = {}
+
+        import torch._dynamo.config as dcfg
+        dcfg.cache_size_limit = 512
+
+
+    def _get_sbd_inference_quadratic_decoding_block_mask(self, block_length: int):
+        if block_length not in self._quadratic_block_mask:
+            draft_len = block_length * (block_length + 1)
+
+            def quadratic(b, h, q_idx, kv_idx):
+                first_clean = torch.logical_and(
+                    kv_idx % (block_length + 1) == 0,
+                    kv_idx < draft_len,
+                )
+                first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
+                block_q = q_idx // (block_length + 1)
+                block_kv = kv_idx // (block_length + 1)
+                same_block = torch.logical_and(block_q == block_kv, q_idx < draft_len)
+                same_block_except_first = torch.logical_and(
+                    same_block,
+                    q_idx % (block_length + 1) != 0,
+                )
+                draft_part = torch.logical_or(first_clean, same_block_except_first)
+                clean_part = kv_idx >= draft_len
+                return torch.logical_or(draft_part, clean_part)
+
+            block_mask = create_block_mask(
+                quadratic,
+                B=None,
+                H=None,
+                Q_LEN=draft_len,
+                KV_LEN=draft_len + self.config.max_position_embeddings,
+                device="cuda",
+            )
+
+            self._quadratic_block_mask[block_length] = block_mask
+
+        return self._quadratic_block_mask[block_length]
+
+
+    def set_attention_mode(self, mode, block_size=None):
+        self.mode = mode
+        self.block_size = block_size
+
+    def compute_block_mask(self, mode, q_len, block_size=None):
+
+        def bidirectional_mask(b, h, q, kv): 
+            return (q >= kv) | (q < kv)
+        
+        def autoregressive_mask(b, h, q, kv):
+            return (q >= kv)
+        
+        def block_diff_mask(block_size, b, h, q_idx, kv_idx, n):
+            """
+            Constructs the specialized block diffusion attention mask for training
+            composed of three masks:
+            - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
+            - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
+            - **Block Causal Mask (M_BC)**: Attention to update x0
+            Args:
+                b, h: Batch and head indices (ignored for mask logic).
+                q_idx, kv_idx: Query and Key indices.
+                seq_len: Total sequence length.
+                block_size: Defines the block structure.
+            Returns:
+                A boolean attention mask.
+            """
+
+            # Indicate whether token belongs to xt or x0
+            x0_flag_q = (q_idx >= n)
+            x0_flag_kv = (kv_idx >= n)
+
+            # Compute block indices
+            block_q = torch.where(x0_flag_q == 1,
+                                    (q_idx - n) // block_size,
+                                    q_idx // block_size)
+            block_kv = torch.where(x0_flag_kv == 1,
+                                    (kv_idx - n) // block_size,
+                                    kv_idx // block_size)
+
+            # **1. Block Diagonal Mask (M_BD) **
+            block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv)
+
+            # **2. Offset Block-Causal Mask (M_OBC) **
+            offset_block_causal = (
+                (block_q > block_kv)
+                & (x0_flag_kv == 1)
+                & (x0_flag_q == 0)
+            )
+
+            # **3. Block-Causal Mask (M_BC) **
+            block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)
+
+            # **4. Combine Masks **
+            return block_diagonal | offset_block_causal | block_causal
+        
+
+        def sbd_block_diff_mask(block_size, b, h, q_idx, kv_idx, n):
+            x0_flag_q = (q_idx >= n)
+            x0_flag_kv = (kv_idx >= n)
+
+            # Compute block indices
+            block_q = torch.where(x0_flag_q == 1,
+                                    (q_idx - n) // block_size,
+                                    q_idx // block_size)
+            block_kv = torch.where(x0_flag_kv == 1,
+                                    (kv_idx - n) // block_size,
+                                    kv_idx // block_size)
+
+            # **1. Block Diagonal Mask (M_BD) **
+            block_diagonal = (block_q == block_kv) & (x0_flag_kv == 0) & (x0_flag_q == 0)
+
+            # **2. Offset Block-Causal Mask (M_OBC) **
+            offset_block_causal = (
+                (block_q > block_kv)
+                & (x0_flag_kv == 1)
+                & (x0_flag_q == 0)
+            )
+
+            # **3. Fully Causal Mask (M_BC) **
+            fully_causal = (q_idx >= kv_idx) & (x0_flag_kv == 1) & (x0_flag_q == 1)
+
+            # **4. Combine Masks **
+            return block_diagonal | offset_block_causal | fully_causal
+
+        if mode == 'bidirectional':
+            attn_mask = bidirectional_mask
+        elif mode == 'autoregressive':
+            attn_mask = autoregressive_mask
+        elif mode == 'block_diff':
+            assert block_size is not None
+            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, q_len//2)
+        elif mode == 'sbd_block_diff':
+            assert block_size is not None
+            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, q_len//2)
+        else:
+            raise ValueError(f"Unknown attention mode: {mode}")
+
+        block_mask = create_block_mask(
+            attn_mask, B=None, H=None, Q_LEN=q_len, KV_LEN=q_len
+        )
+
+        return block_mask
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        is_training: bool = True,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+
+        if self.mode in ['block_diff', 'sbd_block_diff'] and is_training:
+            # Split query and key states in half along sequence length dimension
+            q1, q2 = query_states.chunk(2, dim=2)
+            k1, k2 = key_states.chunk(2, dim=2)
+            
+            # Apply RoPE independently to each half
+            q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin)
+            q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin)
+            
+            # Recombine the halves
+            query_states = torch.cat([q1, q2], dim=2)
+            key_states = torch.cat([k1, k2], dim=2)
+        else:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        query_states = query_states * _get_llama_4_attn_scale(
+            cache_position,
+            self.config.rope_parameters.get("llama_4_scaling_beta"),
+            self.config.rope_parameters.get("original_max_position_embeddings"),
+        ).to(query_states.dtype)
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        tidar_inference_mode = getattr(self.config, "tidar_inference_mode", None)
+        if tidar_inference_mode is not None:
+            if tidar_inference_mode == "quadratic":
+                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
+                if block_length is None:
+                    raise ValueError("SBD quadratic decoding requires block_length in config.")
+                if past_key_values is not None:
+                    seq_len = key_states.shape[2]
+                    draft_len = block_length * (block_length + 1)
+
+                    clean_keys = key_states[:, :, :-draft_len]
+                    draft_keys = key_states[:, :, -draft_len:]
+                    clean_values = value_states[:, :, :-draft_len]
+                    draft_values = value_states[:, :, -draft_len:]
+                    key_states = torch.cat([draft_keys, clean_keys], dim=2)
+                    value_states = torch.cat([draft_values, clean_values], dim=2)
+
+                    block_mask: BlockMask = self._get_sbd_inference_quadratic_decoding_block_mask(
+                        block_length=block_length
+                    )
+                    block_mask.seq_lengths = (draft_len, seq_len)
+                else:
+                    seq_len = query_states.shape[2]
+                    draft_len = block_length * (block_length + 1)
+                    clean_len = seq_len - draft_len
+
+                    def _causal_mask(b, h, q_idx, kv_idx):
+                        return torch.logical_and(q_idx >= kv_idx, q_idx < clean_len)
+
+                    def _draft2clean_mask(b, h, q_idx, kv_idx):
+                        full_clean = torch.logical_and(q_idx >= clean_len, kv_idx <= clean_len)
+                        first_clean = torch.logical_and(
+                            q_idx >= clean_len, (kv_idx - clean_len) % (block_length + 1) == 0
+                        )
+                        first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
+                        return torch.logical_or(full_clean, first_clean)
+
+                    def _draft_mask(b, h, q_idx, kv_idx):
+                        block_q = (q_idx - clean_len) // (block_length + 1)
+                        block_kv = (kv_idx - clean_len) // (block_length + 1)
+                        quadrant = torch.logical_and(q_idx >= clean_len, kv_idx >= clean_len)
+                        same_block = torch.logical_and(block_q == block_kv, quadrant)
+                        same_block_except_first = torch.logical_and(
+                            same_block,
+                            (q_idx - clean_len) % (block_length + 1) != 0,
+                        )
+                        return torch.logical_and(block_q == block_kv, same_block_except_first)
+
+                    mask = or_masks(_causal_mask, _draft2clean_mask)
+                    mask = or_masks(mask, _draft_mask)
+
+                    block_mask = create_block_mask(
+                        mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len,
+                    )
+
+                key_states = repeat_kv(key_states, self.num_key_value_groups)
+                value_states = repeat_kv(value_states, self.num_key_value_groups)
+                attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+                attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+                attn_output = self.o_proj(attn_output)
+                return attn_output, None
+
+            elif tidar_inference_mode == "default":
+                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
+                if block_length is None:
+                    raise ValueError("SBD default decoding requires block_length in config.")
+                seq_len = query_states.shape[2]
+                prefix_len = seq_len - block_length
+
+                def _clean_q_mask(b, h, q_idx, kv_idx):
+                    return torch.logical_and(q_idx >= kv_idx, q_idx < prefix_len)
+
+                def _noisy_q_mask(b, h, q_idx, kv_idx):
+                    return q_idx >= prefix_len
+
+                block_mask = create_block_mask(
+                    or_masks(_clean_q_mask, _noisy_q_mask),
+                    B=None,
+                    H=None,
+                    Q_LEN=seq_len,
+                    KV_LEN=seq_len,
+                )
+
+                key_states = repeat_kv(key_states, self.num_key_value_groups)
+                value_states = repeat_kv(value_states, self.num_key_value_groups)
+                attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+                attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+                attn_output = self.o_proj(attn_output)
+                return attn_output, None
+        
+        else:
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+            if self.mode == 'bidirectional':
+                if self.bidirectional_mask is None or q_len != self.bidirectional_mask.shape[-2]:
+                    block_mask = self.compute_block_mask(mode='bidirectional', q_len=q_len)
+                else:
+                    block_mask = self.bidirectional_mask
+
+            elif self.mode == 'autoregressive':
+                if self.autoregressive_mask is None or q_len != self.autoregressive_mask.shape[-2]:
+                    block_mask = self.compute_block_mask(mode='autoregressive', q_len=q_len)
+                else:
+                    block_mask = self.autoregressive_mask
+
+            elif self.mode == 'block_diff':
+                if self.block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.block_diff_mask.shape[-2]:
+                    block_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size, q_len=q_len)
+                else:
+                    block_mask = self.block_diff_mask
+            elif self.mode == 'sbd_block_diff':
+                if self.sbd_block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]:
+                    block_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size, q_len=q_len)
+                else:
+                    block_mask = self.sbd_block_diff_mask
+            else:
+                raise ValueError(f"Unknown attention mode: {self.mode}")
+
+            attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+            attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+
+            attn_output = self.o_proj(attn_output)
+
+            return attn_output, None
+
+
+def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor:
+    """Return a Bool mask of length len(log_w) with exactly k True."""
+    g = -torch.log(-torch.log(torch.rand_like(log_w) + 1e-9) + 1e-9)
+    topk = torch.topk(log_w + g, k).indices
+    mask = torch.zeros_like(log_w, dtype=torch.bool)
+    mask[topk] = True
+    return mask
+
+
+class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
+    """
+    A single model with:
+      - a bidirectional encoder + diffusion‐LM head over A
+      - a causal decoder + LM head over B, conditioned on F_A
+    """
+
+    def __init__(self, config: MinistralDLMConfig):
+        super().__init__(config)
+
+        self.mask_token_id = config.mask_token_id
+
+        diffusion_config = copy.deepcopy(config)
+        diffusion_config.diffusion_lm = True
+
+        if config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            diffusion_config.attn_class = MinistralFlexAttention
+        elif config.dlm_paradigm in ['bidirectional', 'autoregressive']:
+            diffusion_config.attn_class = Ministral3Attention
+
+            if config.dlm_paradigm == 'autoregressive':
+                diffusion_config.diffusion_lm = False
+        else:
+            raise ValueError(f"Unsupported DLM paradigm: {config.dlm_paradigm}")
+        
+        self.encoder = Ministral3Model(diffusion_config)
+        self.diffusion_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.vocab_size = config.vocab_size
+
+        self.current_iter_ratio = None
+
+        self.post_init()
+
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.encoder.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.diffusion_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.diffusion_head = new_embeddings
+
+
+    def forward_process(self, input_ids, eps=1e-3, block_size=None, loss_mask=None):
+        b, l = input_ids.shape
+        device = input_ids.device
+
+        if self.config.dp_varying_mask_ratio:
+            # Enable different random seeds for each DP rank during sampling
+            import torch.distributed as dist
+            dp_rank = 0
+            if dist.is_initialized():
+                try:
+                    dp_rank = dist.get_rank()
+                except Exception:
+                    dp_rank = 0
+            # Use a local generator to avoid affecting global RNG state
+            generator = torch.Generator(device=device)
+            generator.manual_seed(torch.seed() + dp_rank)
+        else:
+            generator = None
+            
+        if self.config.adaptive_mask_rate:
+            assert block_size is not None
+
+            # --- simple linear window mapping ---
+            bs_min = getattr(self.config, "t_bs_min", 16)
+            bs_max = getattr(self.config, "t_bs_max", 128)
+            w = getattr(self.config, "t_window_width", 0.6)  # fixed width
+
+            # fraction in [0,1] (unclamped first)
+            frac = (float(block_size) - float(bs_min)) / max(1.0, float(bs_max - bs_min))
+            # upper bound decreases linearly from 1.0 -> 0.5
+            u_max = 1.0 - w * frac
+            # clamp to [0.6, 1.0] to handle bs outside [bs_min, bs_max]
+            u_max = max(0.6, min(1.0, u_max))
+            u_min = u_max - w  # ensures width = w
+
+            # sample t ~ Uniform(u_min, u_max)
+            t = u_min + (u_max - u_min) * torch.rand(b, device=device, generator=generator)
+        else:
+            t = torch.rand(b, device=device, generator=generator)
+        
+        p_mask = (1 - eps) * t + eps  # shape: (b,)
+        p_mask = p_mask[:, None].expand(-1, l)  # shape: (b, l)
+
+        masked_indices = torch.rand((b, l), device=device) < p_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)        
+
+        return noisy_batch, masked_indices, p_mask
+
+
+    def forward_process_exp(
+        self,
+        input_ids: torch.Tensor,
+        eps: float = 1e-3,
+        block_size: int | None = None,
+        half_life_ratio: float = 0.25, # λ = ln 2 / (half_life_ratio·L)
+        loss_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Two-stage corruption with optional per-block sampling.
+        • Stage 1:  m ~ U(eps, 1)   →   k = round(m · len)  (exact budget).
+        • Stage 2:  sample exactly k positions with weights
+                    w_i(m) = exp[ λ · (1−m) · i ]   (late-heavy when m→0,
+                                                     uniform when m→1).
+          If `block_size` is given, the procedure is run *independently*
+          inside each contiguous block of that length (last block may be shorter).
+          When block_size is provided, m is sampled per-block and p_mask is per-block.
+        Args
+        ----
+        input_ids : (B, L)  LongTensor
+        eps       : minimum corruption ratio
+        block_size: if not None, operate block-wise with per-block m sampling
+        half_life_ratio : controls steepness when m→0
+        """
+        B, L = input_ids.shape
+        device = input_ids.device
+        dtype  = torch.float32
+
+        masked_indices = torch.zeros((B, L), dtype=torch.bool, device=device)
+        p_mask = torch.zeros((B, L), dtype=dtype, device=device)
+
+        # ---------- Stage 1 & 2: whole-sentence or block-wise -------------------
+        for b in range(B):
+            if block_size is None:
+                # ---------- Per-batch sampling (original behavior) ----------
+                m = eps + (1.0 - eps) * torch.rand(1, device=device).item()   # scalar
+                k_tot = int(round(m * L))
+                k_tot = max(1, min(k_tot, L))  # clamp to [1, L]
+                
+                # Fill p_mask for this batch
+                p_mask[b, :] = m
+                
+                slope = 1.0 - m          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+                
+                # ------- single pool over the whole sentence -------------
+                lam_base = math.log(2.0) / (half_life_ratio * L) # base decay rate (λ when slope=1)
+
+                pos   = torch.arange(L, device=device, dtype=dtype)
+                log_w = (lam_base * slope * pos).clone()
+
+                masked_indices[b] = gumbel_topk(log_w, k_tot)
+
+            else:
+                # ---------- Per-block sampling ----------
+                num_blocks = math.ceil(L / block_size)
+                lam_base = math.log(2.0) / (half_life_ratio * block_size) # base decay rate (λ when slope=1)
+
+                for blk in range(num_blocks):
+                    start = blk * block_size
+                    end   = min((blk + 1) * block_size, L)
+                    blk_len = end - start
+
+                    # Sample m per block
+                    m_blk = eps + (1.0 - eps) * torch.rand(1, device=device).item()
+                    
+                    # Fill p_mask for this block
+                    p_mask[b, start:end] = m_blk
+                    
+                    # per-block budget
+                    k_blk = int(round(m_blk * blk_len))
+                    k_blk = max(0, min(k_blk, blk_len))
+                    if k_blk == 0:
+                        continue
+
+                    slope = 1.0 - m_blk          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+
+                    pos   = torch.arange(blk_len, device=device, dtype=dtype)
+                    log_w = lam_base * slope * pos
+
+                    blk_mask = gumbel_topk(log_w, k_blk)
+                    masked_indices[b, start:end] = blk_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)
+        return noisy_batch, masked_indices, p_mask
+    
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor]   = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor]       = None,
+        split_len: Optional[int]                 = None,
+        past_key_values: Optional[Cache]         = None,
+        block_size: Optional[int]                = None,
+        block_diff_ppl: bool                     = False,
+        eps: float                               = 1e-3,
+        is_teacher: bool                        = False,
+        masked_indices: Optional[torch.Tensor]   = None,
+        p_mask: Optional[torch.Tensor]           = None,
+        teacher_logits: Optional[torch.Tensor]   = None,
+        masked_indices_teacher: Optional[torch.Tensor] = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        ce_loss_weight: float = 1.0,
+        output_last_hidden_states_only: bool = False,
+        skip_loss: bool = False,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+
+        batch_size, seq_len = input_ids.shape
+
+        if self.config.dlm_paradigm == 'bidirectional' or self.config.dlm_paradigm == 'autoregressive':
+            if labels is not None and torch.rand(1) < self.config.random_length_prob:
+                random_length = torch.randint(2, input_ids.shape[1] + 1, (1,))
+                input_ids = input_ids[:, :random_length]
+                labels = labels[:, :random_length]
+                
+                if attention_mask is not None:
+                    attention_mask = attention_mask[:, :random_length]
+                if position_ids is not None:
+                    position_ids = position_ids[:, :random_length]
+                if loss_mask is not None:
+                    loss_mask = loss_mask[:, :random_length]
+
+        elif self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if labels is not None and block_size is None:
+                if torch.rand(1) < self.config.random_length_prob:
+                    block_size = torch.randint(1, 8, (1,)).item() * 4  ## [4, 32] divisible by 4
+                else:
+                    block_size = self.config.block_size
+
+        else:
+            raise ValueError(f"Unknown dLM paradigm: {self.config.dlm_paradigm}")
+
+        if labels is not None and self.config.dlm_paradigm != 'autoregressive':
+            if masked_indices is not None:
+                # assert p_mask is not None
+
+                if loss_mask is not None:
+                    masked_indices[loss_mask == 0] = 0
+
+                noisy_inputs = torch.where(masked_indices, self.mask_token_id, input_ids)
+
+            else:
+                if self.config.tok_mask_half_life_ratio is not None:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process_exp(input_ids, eps=eps, block_size=block_size, half_life_ratio=self.config.tok_mask_half_life_ratio, loss_mask=loss_mask)
+                else:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask)
+
+        else:
+            noisy_inputs = input_ids
+            masked_indices = None
+            p_mask = None
+
+        if self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            for layer in self.encoder.layers:
+                if hasattr(layer.self_attn, 'set_attention_mode'):
+                    layer.self_attn.set_attention_mode(self.config.dlm_paradigm, block_size=block_size)
+
+        input_ids_len = noisy_inputs.shape[1]
+        if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len, device=noisy_inputs.device).unsqueeze(0)
+            noisy_inputs = torch.cat([noisy_inputs, input_ids], dim=1)
+
+        if block_diff_ppl:
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len // 2, device=noisy_inputs.device).unsqueeze(0)
+
+        enc_out  = self.encoder(
+            past_key_values=past_key_values,
+            input_ids=noisy_inputs,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_training=(labels is not None) or (block_diff_ppl),
+            **kwargs,
+        )
+
+        if output_last_hidden_states_only:
+            return BaseModelOutput(last_hidden_state=enc_out.last_hidden_state)
+
+        logits = self.diffusion_head(enc_out.last_hidden_state)  # (batch, len_B, vocab)
+        causal_logits = None
+
+        if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if self.config.dlm_paradigm == 'sbd_block_diff':
+                causal_logits = logits[:, input_ids_len:]
+            else:
+                causal_logits = None
+
+            logits = logits[:, :input_ids_len]
+
+        loss = None
+        if labels is not None and not skip_loss:
+            if self.config.dlm_paradigm == 'autoregressive':
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                
+                if loss_mask is None:
+                    loss_fct = CrossEntropyLoss()
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    loss = loss_fct(shift_logits, shift_labels)
+
+                else:
+                    loss_mask = loss_mask[..., 1:].contiguous()
+
+                    loss_fct = CrossEntropyLoss(reduction='none')
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    shift_labels = shift_labels.to(shift_logits.device)
+                    
+                    token_losses = loss_fct(shift_logits, shift_labels)
+                                    
+                    flat_loss_mask = loss_mask.reshape(-1)
+                    loss = token_losses[flat_loss_mask == 1].sum() / flat_loss_mask.sum()
+
+            else:
+                # Handle DREAM vs LLADA style losses
+                if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
+                    logits = logits[..., :-1, :].contiguous()
+                    labels = labels[..., 1:].contiguous()
+                    masked_indices = masked_indices[:, 1:]
+                    p_mask = p_mask[:, 1:]
+
+                if self.config.ada_perm_ratio_per_block is not None:
+                    # Only compute loss for the top ada_perm_ratio_per_block tokens by confidence within each block
+                    block_size = self.config.block_size
+                    batch_size, seq_len = masked_indices.shape
+                    num_blocks = seq_len // block_size
+                    
+                    # Get the max logit (confidence) for each position
+                    confidence = logits.max(dim=-1).values.detach()  # (batch_size, seq_len)
+                    
+                    # Create a mask for tokens to include in loss
+                    selected_mask = torch.zeros_like(masked_indices, dtype=torch.bool)
+                    
+                    for blk in range(num_blocks):
+                        start = blk * block_size
+                        end = min((blk + 1) * block_size, seq_len)
+                        
+                        # Get masked indices within this block
+                        block_masked = masked_indices[:, start:end]  # (batch_size, block_len)
+                        block_confidence = confidence[:, start:end]  # (batch_size, block_len)
+                        
+                        for b in range(batch_size):
+                            # Get positions that are masked in this block for this batch
+                            masked_positions = torch.where(block_masked[b])[0]
+                            num_masked = len(masked_positions)
+                            
+                            if num_masked > 0:
+                                # Number of tokens to keep (top by confidence)
+                                k = min(max(1, int(block_size * self.config.ada_perm_ratio_per_block)), num_masked)
+                                
+                                # Get confidence values for masked positions
+                                masked_confidence = block_confidence[b, masked_positions]
+                                
+                                # Get indices of top-k confident tokens
+                                _, topk_indices = torch.topk(masked_confidence, k)
+                                selected_positions = masked_positions[topk_indices]
+                                
+                                # Mark these positions in the selected mask
+                                selected_mask[b, start + selected_positions] = True
+                    
+                    # Calculate loss only for selected positions
+                    token_loss = torch.nn.functional.cross_entropy(
+                        logits[selected_mask],
+                        labels[selected_mask],
+                        reduction='none'
+                    ) / p_mask[selected_mask]
+
+                    num_mask_tokens = selected_mask.sum()
+
+                else:
+                    # Calculate token-wise cross entropy loss for masked positions in B
+                    token_loss = torch.nn.functional.cross_entropy(
+                        logits[masked_indices], 
+                        labels[masked_indices], 
+                        reduction='none'
+                    ) / p_mask[masked_indices]
+
+                    num_mask_tokens = masked_indices.sum()
+
+                if self.config.global_loss_avg:
+                    loss = token_loss.sum()
+                else:
+                    loss = token_loss.sum() / num_mask_tokens
+                
+                if self.config.ada_dlm_loss_ratio is not None:
+                    assert self.current_iter_ratio is not None
+                    assert self.config.dlm_loss_weight is not None
+
+                    dlm_loss_weight = min(self.config.dlm_loss_weight, self.current_iter_ratio / self.config.ada_dlm_loss_ratio * self.config.dlm_loss_weight)
+                    loss = dlm_loss_weight * loss
+
+                elif self.config.dlm_loss_weight is not None:
+                    loss = self.config.dlm_loss_weight * loss
+
+                if self.config.dlm_paradigm == 'sbd_block_diff':
+                    causal_logits = causal_logits[..., :-1, :].contiguous()
+                    causal_logits = causal_logits.view(-1, causal_logits.size(-1))
+
+                    if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
+                        causal_labels = labels.view(-1)
+                    else: 
+                        causal_labels = labels[..., 1:].contiguous().view(-1)
+                    
+                    if self.config.global_loss_avg:
+                        loss_fct = CrossEntropyLoss(reduction='sum')
+                        ar_loss = loss_fct(causal_logits, causal_labels)
+
+                        self.loss_diffusion = loss.detach().item() / num_mask_tokens
+                        self.loss_ar = ar_loss.detach().item() / seq_len
+
+                        loss = loss + self.config.ar_loss_weight * ar_loss
+                    else:
+                        loss_fct = CrossEntropyLoss()
+                        ar_loss = loss_fct(causal_logits, causal_labels)
+
+                        self.loss_diffusion = loss.detach().item()
+                        self.loss_ar = ar_loss.detach().item()
+
+                        loss = loss + self.config.ar_loss_weight * ar_loss
+                
+                if self.config.global_loss_avg:
+                    if self.config.dlm_paradigm == 'sbd_block_diff':
+                        loss = (loss, num_mask_tokens + int(self.config.ar_loss_weight * seq_len))
+                    else: 
+                        loss = (loss, num_mask_tokens)
+
+        return MinistralDiffOutputWithPast(
+            loss=loss if not is_teacher else logits,
+            logits=logits,
+            causal_logits=causal_logits,
+            past_key_values=enc_out.past_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+    def generate(self, prompt_ids, max_new_tokens, steps, block_length, shift_logits, threshold, causal_context=True, temperature=0, eos_token_id=None):
+        out_ids, nfe = generate_with_prefix_cache_block_diff(
+                        model=self,
+                        prompt=prompt_ids,
+                        gen_length=max_new_tokens,
+                        steps=steps,
+                        block_length=block_length,
+                        remasking="low_confidence",
+                        temperature=temperature,
+                        mask_id=self.mask_token_id,
+                        threshold=threshold,
+                        shift_logits=shift_logits,
+                        neg_entropy=False,
+                        causal_context=causal_context,
+                        eos_token_id=eos_token_id,
+                    )
+
+        return out_ids, nfe
+
+
+    @torch.no_grad()
+    def sbd_inference_diffusion_quadratic(
+        self,
+        clean_input_ids: Optional[torch.Tensor],
+        draft_input_ids: torch.Tensor,
+        block_length: int,
+        draft_only: bool = False,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+    ):
+        """SBD quadratic inference (injected by build_hf_tidar_repo)."""
+        enc_config = self.encoder.config
+        enc_config.use_sbd_objective = True
+        enc_config.block_length = block_length
+
+        if draft_only:
+            assert clean_input_ids is not None
+
+            if use_cache and past_key_values is None:
+                past_key_values = DynamicCache()
+
+            enc_config.tidar_inference_mode = "default"
+            input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
+            outputs = self.encoder(
+                input_ids=input_ids,
+                position_ids=None,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                is_training=False,
+            )
+
+            hidden_states = outputs.last_hidden_state
+            logits = self.diffusion_head(hidden_states)
+
+            past_key_values = getattr(outputs, "past_key_values", None)
+            if use_cache and past_key_values is not None:
+                _crop_dynamic_cache(past_key_values, clean_input_ids.shape[1])
+
+            return logits, past_key_values
+        else:
+            enc_config.tidar_inference_mode = "quadratic"
+
+            draft_len = block_length * (block_length + 1)
+            draft_input_ids = torch.cat(
+                [
+                    draft_input_ids.view(-1, block_length, 1),
+                    torch.full(
+                        (draft_input_ids.shape[0], block_length, block_length),
+                        fill_value=self.config.mask_token_id,
+                        device=draft_input_ids.device,
+                    ),
+                ],
+                dim=-1,
+            ).view(-1, draft_len)
+
+            if use_cache:
+                assert past_key_values is not None, (
+                    "Past key values should be provided when using cache, e.g. run draft_only=True first."
+                )
+                assert clean_input_ids is None, (
+                    "Clean input ids should already be in cache, thus none should be provided."
+                )
+                clean_len = past_key_values.get_seq_length()
+                input_ids = draft_input_ids
+            else:
+                clean_len = clean_input_ids.shape[1]
+                input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
+
+            per_block_position_ids = torch.arange(
+                clean_len, clean_len + block_length + 1, device=draft_input_ids.device
+            )[None,].repeat(block_length, 1)
+            per_block_position_ids += torch.arange(block_length, device=draft_input_ids.device).view(-1, 1)
+
+            if use_cache:
+                position_ids = per_block_position_ids.view(-1)[None,]
+            else:
+                clean_position_ids = torch.arange(clean_len, device=draft_input_ids.device)
+                position_ids = torch.cat([clean_position_ids, per_block_position_ids.view(-1)], dim=-1)[None,]
+
+            outputs = self.encoder(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                is_training=False,
+            )
+
+            hidden_states = outputs.last_hidden_state
+            logits = self.diffusion_head(hidden_states)
+            past_key_values = getattr(outputs, "past_key_values", None)
+
+            if use_cache and past_key_values is not None:
+                _extract_draft_kv_cache(past_key_values, clean_len, block_length)
+
+            return logits, past_key_values
+
+    @torch.no_grad()
+    def tidar_generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        steps: int = 128,
+        block_length: int = 16,
+        threshold: Optional[float] = None,
+        temperature: float = 0.0,
+        mask_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+    ):
+        """TiDAR quadratic speculative decoding (injected by build_hf_tidar_repo)."""
+        self.config.use_sbd_objective = True
+        self.config.dlm_paradigm = "sbd"
+
+        if prompt_ids.shape[0] != 1:
+            raise ValueError("TiDAR quadratic decoding currently requires batch_size == 1")
+
+        token_mask_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, "eos_token_id", None)
+
+        x = torch.full(
+            (1, prompt_ids.shape[1] + max_new_tokens + block_length * 2),
+            token_mask_id,
+            dtype=torch.long,
+            device=prompt_ids.device,
+        )
+        x[:, : prompt_ids.shape[1]] = prompt_ids.clone()
+
+        if max_new_tokens % block_length != 0:
+            raise ValueError("max_new_tokens must be divisible by block_length")
+        num_blocks = max_new_tokens // block_length
+        if steps % num_blocks != 0:
+            raise ValueError("steps must be divisible by (max_new_tokens // block_length)")
+
+        prompt_len = prompt_ids.shape[1]
+        nfe = 0
+        nfe += 1
+        logits, past_key_values = self.sbd_inference_diffusion_quadratic(
+            clean_input_ids=x[:, :prompt_len],
+            draft_input_ids=x[:, prompt_len : prompt_len + block_length],
+            block_length=block_length,
+            draft_only=True,
+            use_cache=True,
+        )
+
+        logits_proposal = logits[:, prompt_len - 1 : prompt_len + block_length]
+        logits_proposal[:, 1] = logits_proposal[:, 0]
+        logits_proposal = logits_proposal[:, 1:]
+        x0_proposal = torch.argmax(logits_proposal, dim=-1)
+        x[:, prompt_len : prompt_len + block_length] = x0_proposal
+
+        total_accept_token = 0
+        while True:
+            nfe += 1
+            block_start = prompt_len + total_accept_token
+            block_end = block_start + block_length
+            draft_input_ids = x[:, block_start:block_end]
+
+            logits, past_key_values = self.sbd_inference_diffusion_quadratic(
+                clean_input_ids=None,
+                draft_input_ids=draft_input_ids,
+                block_length=block_length,
+                draft_only=False,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+
+            useful_token_logits = logits.view(1, block_length, block_length + 1, -1)
+            if threshold is None:
+                useful_token_logits[:, :, 1] = useful_token_logits[:, :, 0]
+            else:
+                if not (0.0 <= threshold <= 1.0):
+                    raise ValueError("threshold must be between 0 and 1")
+                mix_logits = useful_token_logits[:, :, 0] * threshold + useful_token_logits[:, :, 1] * (1 - threshold)
+                useful_token_logits[:, :, 0] = mix_logits
+                useful_token_logits[:, :, 1] = mix_logits
+
+            if temperature > 0:
+                useful_token_logits = useful_token_logits / temperature
+
+            useful_token_pred = torch.argmax(useful_token_logits, dim=-1)
+            new_draft_input_ids = useful_token_pred[:, 0, 1:]
+            accept_cnt = 1
+
+            while accept_cnt < block_length:
+                if useful_token_pred[:, accept_cnt - 1, 0].item() != draft_input_ids[:, accept_cnt].item():
+                    break
+                new_draft_input_ids = useful_token_pred[:, accept_cnt, 1:]
+                accept_cnt += 1
+
+            x[:, block_start : block_start + accept_cnt] = draft_input_ids[:, :accept_cnt]
+
+            # EoS early stopping: all accepted tokens are finalized left-to-right,
+            # so if any is EoS we can truncate and return immediately.
+            if eos_token_id is not None:
+                accepted = x[0, block_start : block_start + accept_cnt]
+                eos_positions = (accepted == eos_token_id).nonzero(as_tuple=True)[0]
+                if len(eos_positions) > 0:
+                    first_eos_rel = eos_positions[0].item()
+                    total_accept_token += first_eos_rel + 1
+                    output_end = prompt_len + total_accept_token
+                    return x[:, :output_end], nfe
+
+            x[:, block_start + accept_cnt : block_start + accept_cnt + block_length] = new_draft_input_ids
+            past_key_values.crop(block_start + accept_cnt)
+            total_accept_token += accept_cnt
+
+            if total_accept_token >= max_new_tokens:
+                break
+
+        return x[:, : -(block_length * 2)], nfe
+
+
+__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]