kv cache

Author: lucylq

examples/models/llama2/runner/generation.py

@@ -1,9 +1,10 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 
 import argparse
-from typing import List, Optional, Tuple, TypedDict
 
 import json
+from typing import List, Optional, Tuple, TypedDict
+
 import torch
 import torch.nn.functional as F
 from executorch.examples.models.llama2.llama_transformer import ModelArgs
@@ -87,11 +88,17 @@ def generate(
             token_logprobs = torch.zeros_like(tokens, dtype=torch.float)
 
         prev_pos = 0
+        if self.params.use_kv_cache:
+            min_prompt_len = 1
 
         eos_reached = torch.tensor([False] * bsz, device="cpu")
         input_text_mask = tokens != pad_id
+        pos = torch.tensor([prev_pos], dtype=torch.int64)
         if min_prompt_len == total_len:
-            inputs = (tokens,)
+            if self.params.use_kv_cache:
+                inputs = (tokens, pos)
+            else:
+                inputs = (tokens,)
             logits = self.model.forward(inputs)  # updated forward call.
             logits = logits[0]
             token_logprobs = -F.cross_entropy(
@@ -104,7 +111,11 @@ def generate(
         stop_tokens = torch.tensor(list(self.tokenizer.stop_tokens))
 
         for cur_pos in range(min_prompt_len, total_len):
-            inputs = (tokens[:, :cur_pos],)
+            pos = torch.tensor([prev_pos], dtype=torch.int64)
+            if self.params.use_kv_cache:
+                inputs = (tokens[:, prev_pos:cur_pos], pos)
+            else:
+                inputs = (tokens[:, :cur_pos],)
             logits = self.model.forward(inputs)  # updated forward call.
             logits = logits[0]
             if temperature > 0:
@@ -116,9 +127,10 @@ def generate(
             next_token = next_token.reshape(-1)
 
             # only replace token if prompt has already been generated
-            next_token = torch.where(
-                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
-            )
+            if not self.params.use_kv_cache or cur_pos < len(prompt_tokens[0]):
+                next_token = torch.where(
+                    input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
+                )
 
             tokens[:, cur_pos] = next_token
             if logprobs:
@@ -316,6 +328,13 @@ def build_args_parser() -> argparse.ArgumentParser:
         action="store_true",
     )
 
+    parser.add_argument(
+        "--max_gen_len",
+        type=int,
+        default=10,
+        help="Maximum length of the generated response sequence.",
+    )
+
     return parser
 
 
@@ -335,7 +354,9 @@ def main() -> None:
         model_path=args.pte, tokenizer_path=args.tokenizer, model_args=model_args
     )
     result = runner.text_completion(
-        prompts=[args.prompt], max_gen_len=10, temperature=args.temperature
+        prompts=[args.prompt],
+        max_gen_len=args.max_gen_len,
+        temperature=args.temperature,
     )
     print(f"Result: {result}")