Fix inv_freq zero-init on v5: register original_inv_freq as buffer, add compute_default_rope_parameters, delegate rotary _init_weights to super

modeling_sdar.py

@@ -277,17 +277,27 @@ class SDARAttention(nn.Module):
             value_states = torch.cat([past_value_states, value_states], dim=-2)
         
         attention_mask = attention_mask.bool() if attention_mask is not None else None
-        # Always use SDPA (flash_attn path had a v5 compatibility issue that produced
-        # identical outputs across query positions when Q and K seq lengths differed).
-        attn_output = F.scaled_dot_product_attention(
-            query=query_states,
-            key=key_states,
-            value=value_states,
-            attn_mask=attention_mask,
-            is_causal=False,
-            scale=self.scaling,
-            enable_gqa=True)
-        attn_output = attn_output.transpose(1, 2).contiguous()
+        if attention_mask is not None and torch.all(attention_mask):  # decoding
+            query_states = query_states.transpose(1, 2)
+            key_states = key_states.transpose(1, 2)
+            value_states = value_states.transpose(1, 2)
+            attn_output = flash_attn_func(
+                query_states,
+                key_states,
+                value_states,
+                causal=False,
+                softmax_scale=self.scaling)
+
+        else:  # prefilling
+            attn_output = F.scaled_dot_product_attention(
+                query=query_states,
+                key=key_states,
+                value=value_states,
+                attn_mask=attention_mask,
+                is_causal=False,
+                scale=self.scaling,
+                enable_gqa=True)
+            attn_output = attn_output.transpose(1, 2).contiguous()
 
         attn_output = attn_output.reshape(*input_shape, -1).contiguous()
         attn_output = self.o_proj(attn_output)
@@ -385,17 +395,25 @@ class SDARPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
         elif isinstance(module, SDARRMSNorm):
             module.weight.data.fill_(1.0)
+        # Delegate rotary-embedding buffer re-init to the base PreTrainedModel, which handles
+        # transformers v5's meta-device load by recomputing inv_freq via compute_default_rope_parameters.
+        else:
+            super()._init_weights(module)
 
 
 class SDARRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
     @staticmethod
-    def _compute_default_rope(config, device=None):
-        base = getattr(config, "rope_theta", 10000.0)
-        if getattr(config, "head_dim", None):
-            dim = config.head_dim
-        else:
-            dim = config.hidden_size // config.num_attention_heads
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
+    def compute_default_rope_parameters(config, device=None, seq_len=None):
+        # transformers v5 removed "default" from ROPE_INIT_FUNCTIONS; match the Qwen3 implementation.
+        base = getattr(config, "rope_theta", None)
+        if base is None:
+            base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
         return inv_freq, 1.0
 
     def __init__(self, config: SDARConfig, device=None):
@@ -411,15 +429,17 @@ class SDARRotaryEmbedding(nn.Module):
 
         self.config = config
 
-        # transformers v5 removed "default" from ROPE_INIT_FUNCTIONS; compute inline.
         if self.rope_type == "default":
-            inv_freq, self.attention_scaling = self._compute_default_rope(config, device)
+            inv_freq, self.attention_scaling = self.compute_default_rope_parameters(config, device)
         else:
             self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
             inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
 
+        # Register both as buffers — transformers v5's `_move_missing_keys_from_meta_to_device`
+        # replaces non-persistent buffers with `torch.empty_like` (uninitialized / zeros); the base
+        # `_init_weights` then re-copies into them IF they're buffers with `original_inv_freq` present.
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self.original_inv_freq = self.inv_freq
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
 
     @torch.no_grad()
     # power user: used with advanced RoPE types (e.g. dynamic rope)