Update modeling_Llamoe.py

modeling_Llamoe.py

@@ -575,9 +575,9 @@ class LlamoeSdpaAttention(LlamoeAttention):
 
         cos, sin = self.rotary_emb(value_states, position_ids, seq_len=None)
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, None)
-
+        
         past_key_value = getattr(self, "past_key_value", past_key_value)
-
+        
         if past_key_value is not None:
             # sin and cos are specific to RoPE models; cache_position needed for the static cache
             cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
@@ -585,6 +585,8 @@ class LlamoeSdpaAttention(LlamoeAttention):
 
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
+        print("after_rb_key_states:",key_states)
+        print("after_rb_value_states:",value_states)
 
         causal_mask = attention_mask
         if attention_mask is not None and cache_position is not None: