update

Author: Xinyu

ape/ape_gemma.py

@@ -39,7 +39,6 @@ def gemma_attention_prefill_prefix(
     query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
     key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
     value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
-    print(query_states.shape, key_states.shape, value_states.shape)
 
     cos, sin = self.rotary_emb(value_states, position_ids)
     query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

demo_ape.py

@@ -58,24 +58,24 @@ def enable_attention_prefill_prefix(model_name, model):
 
 def enable_attention_prefill_context(model_name, model):
     if "llama" in args.model:
-        from src.ape_llama import enable_llama_attention_prefill_context
+        from ape.ape_llama import enable_llama_attention_prefill_context
         enable_llama_attention_prefill_context(model)
     elif "mistral" in model_name:
-        from src.ape_mistral import enable_mistral_attention_prefill_context
+        from ape.ape_mistral import enable_mistral_attention_prefill_context
         enable_mistral_attention_prefill_context(model)
     elif "gemma" in model_name:
-        from src.ape_gemma import enable_gemma_attention_prefill_context
+        from ape.ape_gemma import enable_gemma_attention_prefill_context
         enable_gemma_attention_prefill_context(model)
 
 def enable_attention_prefill_query(model_name, model, temperature, scale):
     if "llama" in args.model:
-        from src.ape_llama import enable_llama_attention_prefill_query
+        from ape.ape_llama import enable_llama_attention_prefill_query
         enable_llama_attention_prefill_query(model, temperature, scale)
     elif "mistral" in model_name:
-        from src.ape_mistral import enable_mistral_attention_prefill_query
+        from ape.ape_mistral import enable_mistral_attention_prefill_query
         enable_mistral_attention_prefill_query(model, temperature, scale)
     elif "gemma" in model_name:
-        from src.ape_gemma import enable_gemma_attention_prefill_query
+        from ape.ape_gemma import enable_gemma_attention_prefill_query
         enable_gemma_attention_prefill_query(model, temperature, scale)
 
 def seed_everything(seed):