fix: RotaryEmbedding lazy cache build (evita garbage da meta-device init)

modeling_quark.py

@@ -28,14 +28,18 @@ class RotaryEmbedding(nn.Module):
         super().__init__()
         inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self._build_cache(max_seq_len)
-
-    def _build_cache(self, seq_len):
-        t     = torch.arange(seq_len, device=self.inv_freq.device).float()
-        freqs = torch.outer(t, self.inv_freq)
+        self.max_seq_len = max_seq_len
+        self._max = 0   # forza build al primo forward, mai nell'__init__
+        self.cos_cache = None
+        self.sin_cache = None
+
+    def _build_cache(self, seq_len, device, dtype):
+        # Ricalcola sempre da inv_freq corrente (mai cache stantia da meta-device)
+        t     = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.outer(t, self.inv_freq.to(device=device, dtype=torch.float32))
         emb   = torch.cat([freqs, freqs], dim=-1)
-        self.register_buffer("cos_cache", emb.cos()[None, None], persistent=False)
-        self.register_buffer("sin_cache", emb.sin()[None, None], persistent=False)
+        self.cos_cache = emb.cos()[None, None].to(dtype)
+        self.sin_cache = emb.sin()[None, None].to(dtype)
         self._max = seq_len
 
     @staticmethod
@@ -45,11 +49,17 @@ class RotaryEmbedding(nn.Module):
 
     def forward(self, q, k):
         T = q.size(2)
-        if T > self._max:
-            self._build_cache(T)
+        # Ricostruisce la cache se: mai costruita, troppo corta, o device/dtype cambiati
+        needs_rebuild = (
+            self.cos_cache is None
+            or T > self._max
+            or self.cos_cache.device != q.device
+            or self.cos_cache.dtype  != q.dtype
+        )
+        if needs_rebuild:
+            self._build_cache(max(T, self.max_seq_len), q.device, q.dtype)
         cos = self.cos_cache[:, :, :T, :]
         sin = self.sin_cache[:, :, :T, :]
-        # Identico a train.py — nessun cast, broadcast naturale
         q = q * cos + self._rotate_half(q) * sin
         k = k * cos + self._rotate_half(k) * sin
         return q, k