Upload model

config.json

@@ -22,6 +22,7 @@
   "dlm_paradigm": "bidirectional",
   "dlm_type": "llada",
   "dp_varying_mask_ratio": false,
+  "enable_self_spec": false,
   "enforce_mask": false,
   "eos_token_id": 2,
   "global_loss_avg": false,

configuration_ministral_dlm.py

@@ -112,6 +112,9 @@ class MinistralDLMConfig(PretrainedConfig):
             Adaptive permutation ratio for each block.
         ada_perm_ratio_global (`float`, *optional*):
             Adaptive permutation ratio for global.
+        enable_self_spec (`bool`, *optional*, defaults to `False`):
+            Force MinistralFlexAttention for all paradigms (including bidirectional/autoregressive).
+            Required for self speculative generation; leave False for standard eval to use faster SDPA kernels.
     """
 
     model_type = "ministral_dlm"
@@ -181,6 +184,7 @@ class MinistralDLMConfig(PretrainedConfig):
         ada_perm_ratio_per_block=None,
         ada_perm_ratio_global=None,
         ada_dlm_loss_ratio=None,
+        enable_self_spec=False,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -234,6 +238,7 @@ class MinistralDLMConfig(PretrainedConfig):
         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
+        self.enable_self_spec = enable_self_spec
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,

modeling_ministral_dlm.py

@@ -90,7 +90,8 @@ def _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block
 class MinistralFlexAttention(Ministral3Attention):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
- 
+
+        self.max_seq_length = getattr(self.config, 'max_seq_length', 4096)
         self.block_size_orig = self.config.block_size
 
         if self.config.dlm_paradigm == 'bidirectional':
@@ -151,31 +152,15 @@ class MinistralFlexAttention(Ministral3Attention):
         self.mode = mode
         self.block_size = block_size
 
-    def compute_block_mask(self, mode, q_len, block_size=None):
+    def compute_block_mask(self, mode, q_len=None, 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)
 
@@ -238,15 +223,23 @@ class MinistralFlexAttention(Ministral3Attention):
             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)
+            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
         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)
+            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
         else:
             raise ValueError(f"Unknown attention mode: {mode}")
 
+        if q_len is not None:
+            Q_LEN = q_len
+        else:
+            if mode in ['block_diff', 'sbd_block_diff']:
+                Q_LEN = self.max_seq_length * 2
+            else:
+                Q_LEN = self.max_seq_length
+
         block_mask = create_block_mask(
-            attn_mask, B=None, H=None, Q_LEN=q_len, KV_LEN=q_len
+            attn_mask, B=None, H=None, Q_LEN=Q_LEN, KV_LEN=Q_LEN
         )
 
         return block_mask
@@ -298,9 +291,9 @@ class MinistralFlexAttention(Ministral3Attention):
             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":
+        self_spec_inference_mode = getattr(self.config, "self_spec_inference_mode", None)
+        if self_spec_inference_mode is not None:
+            if self_spec_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.")
@@ -360,7 +353,7 @@ class MinistralFlexAttention(Ministral3Attention):
                 attn_output = self.o_proj(attn_output)
                 return attn_output, None
 
-            elif tidar_inference_mode == "default":
+            elif self_spec_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.")
@@ -449,11 +442,12 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         diffusion_config = copy.deepcopy(config)
         diffusion_config.diffusion_lm = True
 
+        use_flex = getattr(config, 'enable_self_spec', False)
+
         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
-
+            diffusion_config.attn_class = MinistralFlexAttention if use_flex else Ministral3Attention
             if config.dlm_paradigm == 'autoregressive':
                 diffusion_config.diffusion_lm = False
         else:
@@ -907,7 +901,6 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         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
@@ -918,7 +911,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             if use_cache and past_key_values is None:
                 past_key_values = DynamicCache()
 
-            enc_config.tidar_inference_mode = "default"
+            enc_config.self_spec_inference_mode = "default"
             input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
             outputs = self.encoder(
                 input_ids=input_ids,
@@ -937,7 +930,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
 
             return logits, past_key_values
         else:
-            enc_config.tidar_inference_mode = "quadratic"
+            enc_config.self_spec_inference_mode = "quadratic"
 
             draft_len = block_length * (block_length + 1)
             draft_input_ids = torch.cat(
@@ -994,23 +987,22 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             return logits, past_key_values
 
     @torch.no_grad()
-    def tidar_generate(
+    def self_spec_generate(
         self,
         prompt_ids: torch.Tensor,
         max_new_tokens: int = 128,
         steps: int = 128,
         block_length: int = 16,
-        threshold: Optional[float] = None,
+        ar_mix_weight: 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")
+            raise ValueError("Self speculation 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:
@@ -1064,12 +1056,12 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             )
 
             useful_token_logits = logits.view(1, block_length, block_length + 1, -1)
-            if threshold is None:
+            if ar_mix_weight 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)
+                if not (0.0 <= ar_mix_weight <= 1.0):
+                    raise ValueError("ar_mix_weight must be between 0 and 1")
+                mix_logits = useful_token_logits[:, :, 0] * ar_mix_weight + useful_token_logits[:, :, 1] * (1 - ar_mix_weight)
                 useful_token_logits[:, :, 0] = mix_logits
                 useful_token_logits[:, :, 1] = mix_logits