From 2b91bc983d1e7fe3ef659ef518ba84e12fc2c11b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 21 Mar 2025 16:53:44 -0700
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 364c0444bd..4dfe214cf0 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -422,6 +422,7 @@ def LlamaAttention_fast_forward(
         V = torch.cat([past_key_value[1], V], dim = 2)
     pass
     past_key_value = (K, V) if use_cache else None
+    print(bsz, q_len, past_key_value[0].shape, past_key_value[1].shape)
 
     # Attention module
     if (not HAS_FLASH_ATTENTION and HAS_XFORMERS and attention_mask is None):