From 1fa07c58fb0ad19592eb27bed082b44ca76c64f4 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Thu, 7 Dec 2023 00:13:29 +1100
Subject: [PATCH] Fix generation

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

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index bfc0e26c0..0a89f6bf3 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -74,6 +74,91 @@ def original_apply_o(self, X):
 pass
 
 
+def LlamaAttention_fast_forward_inference(
+    self,
+    hidden_states:  torch.Tensor,
+    past_key_value: Optional[Tuple[torch.Tensor]],
+    position_ids,
+):
+    """
+        https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L406
+        Fast inference using KV cache.
+        QK^T can be computed in 4 chunks
+
+        [Q, q] @ [K, k].T where q, k are the new tokens.
+        [QK^T, Qk^T]
+        [qK^T, qk^T]
+
+        Since the attention mask wipes Qk^T, we just get
+        [QK^T,    0]
+        [qK^T, qk^T]
+
+        Since softmax is row-wise, we get
+        softmax([QK^T,    0])
+        softmax([qK^T, qk^T])
+
+        We then multiply by   [V]
+                              [v]
+        softmax([QK^T,    0]) [softmax(QK^T)V] *
+        softmax([qK^T, qk^T]) [softmax([qK^T, qk^T]) @ [V, v]]
+
+        But notice * [softmax(QK^T)V] is just the last attention.
+        We just need to compute the last final row.
+
+        This means we can pass in a row of Q, but we need to
+        remember K and V, which are called the KV cache.
+    """
+    Xn = hidden_states
+    bsz, _, _ = hidden_states.size()
+    K1, V1 = past_key_value
+
+    Wq = self.q_proj.weight
+    Wk = self.k_proj.weight
+    Wv = self.v_proj.weight
+    Wo = self.o_proj.weight
+
+    n_heads    = self.num_heads
+    n_groups   = self.num_key_value_groups
+    n_kv_heads = self.num_key_value_heads
+    head_dim   = self.head_dim
+    assert(n_kv_heads * n_groups == n_heads)
+
+    Qn, Kn, Vn = original_apply_qkv(self, Xn)
+    Qn = Qn.view(bsz, 1, n_heads,    head_dim).transpose(1, 2)
+    Kn = Kn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
+    Vn = Vn.view(bsz, 1, n_kv_heads, head_dim).transpose(1, 2)
+
+    kv_seq_len = K1.shape[-2] + 1
+    cos, sin = self.rotary_emb(Vn, seq_len = kv_seq_len)
+    Qn, Kn = inplace_rope_embedding(Qn, Kn, cos, sin, position_ids)
+    
+    # New KV cache
+    Kn = torch.cat([K1, Kn], dim = 2)
+    Vn = torch.cat([V1, Vn], dim = 2)
+
+    # Grouped query attention
+    # K = repeat_kv(K, n_groups)
+    # V = repeat_kv(V, n_groups)
+    if n_groups != 1:
+        _, _, cached_len, _ = Kn.shape
+        Kn = Kn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
+        Vn = Vn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
+        Kn = Kn.reshape(bsz, n_heads, cached_len, head_dim)
+        Vn = Vn.reshape(bsz, n_heads, cached_len, head_dim)
+    pass
+
+    # Attention
+    A = torch.matmul(Qn, Kn.transpose(2, 3))
+    A *= 1.0 / (self.head_dim**0.5)
+    A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(A.dtype)
+    A = torch.matmul(A, Vn)
+    A = A.transpose(1, 2)
+    A = A.reshape(bsz, 1, self.hidden_size)
+    A = original_apply_o(self, A)
+    return A, (Kn, Vn)
+pass
+
+
 # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L320
 def LlamaAttention_fast_forward(
     self,
@@ -91,6 +176,17 @@ def LlamaAttention_fast_forward(
     bsz, q_len, _ = hidden_states.size()
     Q, K, V = self.apply_qkv(self, hidden_states)
 
+    # Check for inference
+    if use_cache and past_key_value is not None and q_len == 1:
+        A, past_key_value = LlamaAttention_fast_forward_inference(
+            self,
+            hidden_states,
+            past_key_value,
+            position_ids,
+        )
+        return A, None, past_key_value
+    pass
+
     n_heads    = self.num_heads
     n_groups   = self.num_key_value_groups
     n_kv_heads = self.num_key_value_heads