Update modeling_quiet.py

modeling_quiet.py

@@ -194,20 +194,13 @@ class QuietRotaryEmbedding(nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
     def forward(self, x, seq_len=None):
-        if seq_len is None:
-            seq_len = x.shape[-2]
-        if seq_len > self.max_position_embeddings:
-            # If the sequence length is greater than the maximum position embeddings,
-            # extend the rotary embeddings by repeating them
-            num_repeats = (seq_len - 1) // self.max_position_embeddings + 1
-            cos_cached = self.cos_cached.repeat(num_repeats, 1)[:seq_len]
-            sin_cached = self.sin_cached.repeat(num_repeats, 1)[:seq_len]
-        else:
-            cos_cached = self.cos_cached[:seq_len]
-            sin_cached = self.sin_cached[:seq_len]
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
         return (
-            cos_cached.to(dtype=x.dtype),
-            sin_cached.to(dtype=x.dtype),
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
         )
 
 
@@ -355,7 +348,6 @@ class QuietAttention(nn.Module):
                     "with a layer index."
                 )
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-        
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)