fixes

Author: andrej

iron/applications/llama_3.2_1b/llama_cpu.py

@@ -291,8 +291,8 @@ def llama_forward_pass(config, state):
 
 
 def main():
-    prompt = "The capital of France is "
     args = harness.parse_args()
+    prompt = harness.get_prompt(args.prompt_len)
     config, state = harness.init(args.weights_path, args.tokenizer_path, prompt=prompt)
     print(prompt, end="", flush=True)
     harness.generate(config, state, llama_forward_pass, num_tokens=args.num_tokens)

iron/applications/llama_3.2_1b/llama_inference_harness.py

@@ -178,6 +178,12 @@ def parse_args():
     parser.add_argument(
         "tokenizer_path", type=str, help="Path to the tokenizer model (tiktoken file)"
     )
+    parser.add_argument(
+        "--prompt-len",
+        type=int,
+        default=2048,
+        help="Length of the input prompt in tokens (default: 2048)",
+    )
     parser.add_argument(
         "--num-tokens",
         type=int,
@@ -187,6 +193,13 @@ def parse_args():
     return parser.parse_args()
 
 
+def get_prompt(prompt_len):
+    with open("prompt.txt", "r") as f:
+        prompt = f.read()
+    prompt = prompt[:prompt_len]
+    return prompt
+
+
 def init(
     weights_path,
     tokenizer_path,
@@ -249,12 +262,13 @@ def generate(config, state, forward_pass, num_tokens=100, use_kv_cache=True):
     t_decode = t_decode_end - t_decode_start
     sys.stderr.write("\n\n=== Performance Statistics ===\n")
     sys.stderr.write(f"[Prefill] Time to first token:   {t_prefill:7.3f} s\n")
-    sys.stderr.write(
-        f"[Decode]  Time per token (mean): {t_decode / (n_tokens_generated - 1):7.3f} s\n"
-    )
-    sys.stderr.write(
-        f"[Decode]  Tokens per second:     {(n_tokens_generated - 1) / t_decode:7.3f}\n"
-    )
+    if n_tokens_generated > 1:
+        sys.stderr.write(
+            f"[Decode]  Time per token (mean): {t_decode / (n_tokens_generated - 1):7.3f} s\n"
+        )
+        sys.stderr.write(
+            f"[Decode]  Tokens per second:     {(n_tokens_generated - 1) / t_decode:7.3f}\n"
+        )
     sys.stderr.write(
         f"[Total]   Time per token (mean): {(t_prefill + t_decode) / n_tokens_generated:7.3f} s\n"
     )

iron/applications/llama_3.2_1b/llama_npu.py

@@ -1326,14 +1326,15 @@ def llama_forward_pass(config, state):
 
 
 def main():
-    global aie_ops, aie_buffers
-    prompt = "The capital of France is "
-    # with open('prompt.txt', 'r') as f:
-    #    prompt = f.read()
-    # prompt = prompt[:max_seq_len]
-
+    global aie_ops, aie_buffers, max_seq_len
     args = harness.parse_args()
 
+    assert (
+        max_seq_len >= args.prompt_len + args.num_tokens
+    ), "max_seq_len must be at least prompt_len + num_tokens"
+
+    prompt = harness.get_prompt(args.prompt_len)
+
     config, state = harness.init(args.weights_path, args.tokenizer_path, prompt=prompt)
 
     aie_ops = AIELlamaOperators(config, max_seq_len)

iron/applications/llama_3.2_1b/test.py

@@ -5,13 +5,14 @@
 import subprocess
 import pytest
 from pathlib import Path
+import os
 
 test_dir = Path(__file__).parent
-weights_dir = Path("/srv")
+weights_dir = Path(os.environ.get("IRON_EXAMPLE_WEIGHTS_DIR", "/srv"))
 
 
 def generate_test_params():
-    prompt_lengths = [2048, 13]
+    prompt_lengths = [1024, 13]
     num_tokens_list = [40, 1]
 
     params = []
@@ -32,7 +33,7 @@ def generate_test_params():
 )
 @pytest.mark.parametrize("prompt_len,num_tokens", params, ids=names)
 def test_llama_3_2_1b(prompt_len, num_tokens):
-    command = f"python3 {test_dir}/llama_npu.py {weights_dir}/llama3.2-1b/model.safetensors {weights_dir}/llama3.2-1b/tokenizer.model --num-tokens {num_tokens}"
+    command = f"python3 {test_dir}/llama_npu.py {weights_dir}/llama3.2-1b/model.safetensors {weights_dir}/llama3.2-1b/tokenizer.model --num-tokens {num_tokens} --prompt-len {prompt_len}"
 
     result = subprocess.run(
         command,

iron/common/fusion.py

@@ -240,6 +240,7 @@ def load_elf(op):
 def patch_elf(elf_data, patches):
     for i, patch in patches.items():
         val, mask = patch
+        val = np.uint64(val)
         mask = np.uint64(mask)  # avoid numpy overflow errors
         elf_data[i] = np.uint32((elf_data[i] & ~mask) | (val & mask))
     return elf_data