Update modeling_moonshot.py

modeling_moonshot.py

@@ -428,9 +428,15 @@ class Attention(nn.Module):
     ):
         # Standard scaled dot-product attention
         batch_size, q_length, num_heads, head_dim = query_states.shape
+        _, kv_length, num_kv_heads, _ = key_states.shape
         
-        # Prepare the query, key, value for attention computation
-        # (batch_size, num_heads, seq_length, head_dim)
+        # Handle grouped-query attention by repeating k/v heads if necessary
+        if num_kv_heads != num_heads:
+            # Each query head uses the corresponding key-value head (num_heads // num_kv_heads) times
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+        
+        # Prepare for attention computation (batch_size, num_heads, seq_length, head_dim)
         query_states = query_states.transpose(1, 2)
         key_states = key_states.transpose(1, 2)
         value_states = value_states.transpose(1, 2)