fix llama FP8 perf issue, kvcache.update should be used since FP8 pat…

…ches KVCache Signed-off-by: Wang, Yi A <[email protected]>
huggingface · Feb 7, 2025 · 38e4777 · 38e4777
1 parent 27d1495
commit 38e4777
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/optimum/habana/transformers/models/llama/modeling_llama.py b/optimum/habana/transformers/models/llama/modeling_llama.py
@@ -668,18 +668,21 @@ def pre_attn_forward(
                         else key_states.dtype,
                         device=key_states.device,
                     )
-                    past_key.copy_(key_states)
-                    past_value.copy_(value_states)
                     # Return list instead of tuple
                     past_key_value = [past_key, past_value]
+                    key_states = self.k_cache.update(past_key_value[0], key_states, 2, token_idx, key_states.shape[-2])
+                    value_states = self.v_cache.update(
+                        past_key_value[1], value_states, 2, token_idx, value_states.shape[-2]
+                    )
+
                 elif (
                     token_idx is not None
                     and num_virtual_tokens is not None
                     and num_virtual_tokens == past_key_value[0].shape[-2]
                 ):
                     # prefix tuning case. attach past_key_value to generate first token.
-                    key_states = torch.cat((past_key_value[0], key_states), -2)
-                    value_states = torch.cat((past_key_value[1], value_states), -2)
+                    key_states = self.k_cache.update(past_key_value[0], key_states, 2, None, -1)
+                    value_states = self.v_cache.update(past_key_value[1], value_states, 2, None, -1)
                     past_key_value = (key_states, value_states)
                 else:
                     key_states = self.k_cache.update(past_key_value[0], key_states, 2, token_idx, self.inp_seq_len)

diff --git a/tests/baselines/llama_7b.json b/tests/baselines/llama_7b.json
@@ -376,7 +376,6 @@
                         "--report_to none",
                         "--max_steps 100",
                         "--peft_type prompt_tuning",
-                        "--max_seq_length 64",
                         "--lr_scheduler_type cosine",
                         "--warmup_steps 0",
                         "--weight_decay 0.05",
@@ -402,7 +401,6 @@
                         "--report_to none",
                         "--max_steps 100",
                         "--peft_type prefix_tuning",
-                        "--max_seq_length 64",
                         "--lr_scheduler_type cosine",
                         "--warmup_steps 0",
                         "--weight_decay 0.05",
@@ -428,7 +426,6 @@
                         "--report_to none",
                         "--max_steps 100",
                         "--peft_type p_tuning",
-                        "--max_seq_length 64",
                         "--lr_scheduler_type cosine",
                         "--warmup_steps 0",
                         "--weight_decay 0.05",