fix: cast float32 prima di RoPE per evitare overflow

modeling_quark.py

@@ -71,16 +71,15 @@ class GroupedQueryAttention(nn.Module):
 
     def forward(self, x, **kwargs):
         B, T, _ = x.shape
-        q = self.q_proj(x).view(B, T, self.n_heads,    self.head_dim).transpose(1, 2)
-        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        orig_dtype = x.dtype
+        # Cast a float32 prima di tutto per evitare overflow in RoPE e SDPA
+        q = self.q_proj(x).view(B, T, self.n_heads,    self.head_dim).transpose(1, 2).float()
+        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2).float()
+        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2).float()
         q, k = self.rope(q, k)
         if self.n_groups > 1:
             k = k.repeat_interleave(self.n_groups, dim=1)
             v = v.repeat_interleave(self.n_groups, dim=1)
-        # SDPA in float32 per stabilità numerica, poi riporta al dtype originale
-        orig_dtype = v.dtype
-        q, k, v = q.float(), k.float(), v.float()
         out = F.scaled_dot_product_attention(
             q, k, v, is_causal=True,
             dropout_p=self.drop if self.training else 0.0,