add kv cache to eval

Author: lucylq

examples/models/llama2/eval_llama_lib.py

@@ -41,13 +41,15 @@ def __init__(
         model: nn.Module,
         tokenizer: Union[SentencePieceTokenizer, Tiktoken],
         max_seq_length: Optional[int] = None,
+        use_kv_cache: bool = False,
     ):
         device = "cuda" if torch.cuda.is_available() else "cpu"
         super().__init__(device=device)
         self._model = model
         self._tokenizer = tokenizer
         self._device = torch.device(device)
         self._max_seq_length = 2048 if max_seq_length is None else max_seq_length
+        self._use_kv_cache = use_kv_cache
 
     @property
     def eot_token_id(self):
@@ -83,7 +85,15 @@ def tok_decode(self, tokens):
         return decoded
 
     def _model_call(self, inps):
-        return self._model(inps)
+        if self._use_kv_cache:
+            result_logits = []
+            for pos in range(self._max_seq_length):
+                pos_tensor = torch.tensor([pos], dtype=torch.int64)
+                logits = self._model(inps[:, pos : pos + 1], pos_tensor)
+                result_logits.append(logits)
+            return torch.cat(result_logits, dim=1)
+        else:
+            return self._model(inps)
 
     def _model_generate(self, context, max_length, eos_token_id):
         raise Exception("unimplemented")
@@ -100,9 +110,11 @@ def __init__(
         model: str,
         tokenizer: Union[SentencePieceTokenizer, Tiktoken],
         max_seq_length: Optional[int] = None,
+        use_kv_cache: bool = False,
     ):
         super().__init__(None, tokenizer, max_seq_length)
         self._model = model  # Expects model to be path to a .pte file
+        self._use_kv_cache = use_kv_cache
 
         from executorch.extension.pybindings.portable_lib import _load_for_executorch
 
@@ -111,9 +123,17 @@ def __init__(
     def _model_call(self, inps):
         # Given inps (tokens), return the logits from a single forward call
         # inps: Tensor of shape (1, max_seq_len - 1)
-        # logits: Tensor of shape (1, max_seq_len - 1, 32000)
-        result = self._et_model.forward((inps,))
-        return result[0]
+        # logits: Tensor of shape (1, max_seq_len - 1, vocab_size)
+        if self._use_kv_cache:
+            result_logits = []
+            for pos in range(self._max_seq_length):
+                pos_tensor = torch.tensor([pos], dtype=torch.int64)
+                logits = self._et_model.forward((inps[:, pos : pos + 1], pos_tensor))
+                result_logits.append(logits[0])
+            return torch.cat(result_logits, dim=1)
+        else:
+            result = self._et_model.forward((inps,))
+            return result[0]
 
 
 class ETRunnerEvalWrapper(GPTFastEvalWrapper):
@@ -139,7 +159,7 @@ def _model_call(self, inps):
 
         # Example:
         # inps: Tensor of shape (1, N)
-        # logits: Tensor of shape (1, N, 32000)
+        # logits: Tensor of shape (1, N, vocab_size)
         pass
 
 
@@ -212,6 +232,7 @@ def gen_eval_wrapper(
             # Exported model takes at most (max_seq_length - 1) tokens.
             # Note that the eager model takes at most max_seq_length tokens.
             max_seq_length=args.max_seq_length - 1,
+            use_kv_cache=args.use_kv_cache,
         )
 
     # GPTFastEvalWrapper: Create a wrapper around a pre-exported model
@@ -225,6 +246,7 @@ def gen_eval_wrapper(
         model=model,
         tokenizer=tokenizer,
         max_seq_length=args.max_seq_length,
+        use_kv_cache=args.use_kv_cache,
     )