a

Dockerfile

@@ -66,6 +66,8 @@ COPY --chown=user . $HOME/app
 RUN wget -c -nv https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin
 RUN wget -c -nv https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin
 RUN wget -c -nv https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin
+RUN wget -c -nv https://huggingface.co/kirp/TinyLlama-1.1B-Chat-v0.2-bin/resolve/main/tok_tl-chat.bin
+RUN wget -c -nv https://huggingface.co/kirp/TinyLlama-1.1B-Chat-v0.2-bin/resolve/main/tl-chat.bin
 
 # CMD ["mojo", "llama2.mojo"]
 CMD ["python3", "gradio_app.py"]

gradio_app.py

@@ -1,36 +1,91 @@
 import gradio as gr
 import subprocess
 import sys
-import os
+from pathlib import Path
 
-
-async def generate(prompt):
-    # os.environ["PROMPT"] = prompt
+async def generate(prompt, model_name, seed=0, temperature=0.5, num_tokens=256):
     # stream stout
+    base = ""#"../model/"
+    tokenizer_name = "tokenizer.bin"
+    if model_name == "tl-chat.bin":
+        tokenizer_name = 'tok_tl-chat.bin'
     process = subprocess.Popen(
-        ["mojo", "llama2.mojo"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        [
+            "mojo",
+            "llama2.mojo",
+            Path(base + model_name),
+            "-s",
+            str(seed),
+            "-n",
+            str(num_tokens),
+            "-t",
+            str(temperature),
+            "-i",
+            prompt,
+            "-z",
+            Path(base + tokenizer_name)
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
     )
     text = ""
     for char in iter(lambda: process.stdout.read(1), b""):
-        char_decoded = char.decode()
-        sys.stdout.write(char_decoded)
+        char_decoded = char.decode("utf-8", errors="ignore")
         text += char_decoded
         yield text
 
 
-output_text = gr.Textbox(label="Generated Text")
-
-demo = gr.Interface(
-    fn=generate,
-    inputs=None,
-    outputs=output_text,
-    description="""
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
 # llama2.🔥
 ## [Mojo](https://docs.modular.com/mojo/) implementation of [llama2.c](https://github.com/karpathy/llama2.c) by [@tairov](https://github.com/tairov)
 Source: https://github.com/tairov/llama2.mojo
-    """,
-    allow_flagging="never",
-)
+    """
+    )
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", placeholder="Add your prompt here...")
+            seed = gr.Slider(
+                minimum=0,
+                maximum=2**53,
+                value=0,
+                step=1,
+                label="Seed",
+                randomize=True,
+            )
+            temperature = gr.Slider(
+                minimum=0.0, maximum=2.0, step=0.01, value=0.0, label="Temperature"
+            )
+            num_tokens = gr.Slider(
+                minimum=1, maximum=256, value=256, label="Number of tokens"
+            )
+            model_name = gr.Dropdown(
+                ["stories15M.bin", "stories42M.bin", "stories110M.bin", "tl-chat.bin"],
+                value="stories15M.bin",
+                label="Model Size",
+            )
+            with gr.Row():
+                stop = gr.Button("Stop")
+                run = gr.Button("Run")
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(label="Generated Text")
+
+    # update maximum number of tokens based on model size
+    model_name.change(
+        lambda x: gr.update(maximum=1024)
+        if x == "stories110M.bin" or x == "stories42M.bin" or x == "tl-chat.bin"
+        else gr.update(maximum=256),
+        model_name,
+        num_tokens,
+        queue=False,
+    )
+    click_event = run.click(
+        fn=generate,
+        inputs=[prompt, model_name, seed, temperature, num_tokens],
+        outputs=output_text,
+    )
+    stop.click(fn=None, inputs=None, outputs=None, cancels=[click_event])
 
 demo.queue()
 demo.launch(server_name="0.0.0.0")

llama2.mojo

@@ -1,194 +1,220 @@
+from algorithm import sum
+from algorithm import vectorize, parallelize
+from builtin import string
 from math import round
-import math
-
 from memory import memset_zero, memcpy
+from memory.buffer import Buffer
 from memory.unsafe import DTypePointer
+from python import Python
 from random import rand
-from sys.info import simdwidthof
-from builtin import string
-import time
-import random
-import os
-
-from runtime.llcl import num_cores
-
 from read import BufReader, File
-from memory.buffer import Buffer
-
-from python import Python
+from runtime.llcl import num_cores, Runtime
+from sys import argv
+from tensor import Tensor, TensorShape, TensorSpec
 
 # The SIMD vector width.
-from algorithm import vectorize, parallelize
-from algorithm import sum
+from sys.info import simdwidthof
+import math
+import os
+import random
+import time
 
-alias nelts = (2 * simdwidthof[DType.float32]())
+var workers = 0
+
+alias nelts = (2*simdwidthof[DType.float32]())
 
 alias PointerString = Pointer[UInt8]
 alias BufferPtrType = DTypePointer[DType.uint8]
 alias BufferPtrFloat32 = DTypePointer[DType.float32]
 alias PointerStrings = Pointer[PointerString]
+alias TensorF32 = Tensor[DType.float32]
+
+
+struct TensorSlice:
+    # Provides a view into a tensor representing a 1D slice on its first or first 2 dimensions.
+    # Same function signatures as Tensor but without owning the data.
+    var _data: BufferPtrFloat32
+    var _shape: TensorShape
+
+    fn __init__(inout self, t: TensorF32, layer: Int) raises:
+        let elements_per_layer = t.num_elements() // t.dim(0)
+        self._data = t.data().offset(layer * elements_per_layer)
+        if t.rank() == 2:
+            self._shape = TensorShape(t.dim(1))
+        elif t.rank() == 3:
+            self._shape = TensorShape(t.dim(1), t.dim(2))
+        else:
+            # Compiler complains if _shape not defined
+            self._shape = TensorShape(1)
+            raise Error("TensorSlice: rank greater than 3 not implemented.")
+
+    fn __init__(inout self, t: TensorF32, layer: Int, row: Int) raises:
+        let elements_per_layer = t.num_elements() // t.dim(0)
+        let elements_per_row = elements_per_layer // t.dim(1)
+        self._data = t.data().offset(
+            layer * elements_per_layer + row * elements_per_row
+        )
+        if t.rank() == 3:
+            self._shape = TensorShape(t.dim(2))
+        elif t.rank() == 1:
+            # Compiler complains if _shape not defined
+            self._shape = TensorShape(1)
+            raise Error(
+                "Trying to slice a 1D Tensor by layer and row.  This requires a 3D"
+                " Tensor."
+            )
+        else:
+            # Compiler complains if _shape not defined
+            self._shape = TensorShape(1)
+            raise Error("TensorSlice: rank greater than 3 not implemented.")
 
+    fn data(self) -> BufferPtrFloat32:
+        return self._data
 
-struct Matrix3:
-    var data: BufferPtrFloat32
-    var rows: Int
-    var cols: Int
-    var layers: Int
-    var allocated: Int
-
-    fn __init__(inout self, layers: Int, rows: Int, cols: Int):
-        self.data = BufferPtrFloat32.alloc(0)
-        self.rows = rows
-        self.cols = cols
-        self.layers = layers
-        self.allocated = 0
+    fn shape(self) -> TensorShape:
+        return self._shape
 
-    @always_inline
-    fn alloc(inout self, fill: Int = 0):
-        self.data = BufferPtrFloat32.alloc(self.size())
-        self.allocated = 1
-        if fill == 1:
-            self.zero()
+    fn num_elements(self) -> Int:
+        return self._shape.num_elements()
 
-    @always_inline
-    fn alloc_zero(inout self):
-        self.alloc(1)
+    fn dim(self, idx: Int) -> Int:
+        return self._shape[idx]
 
-    @always_inline
-    fn set_buf_ptr(inout self, ptr: BufferPtrFloat32):
-        self.data = ptr
+    fn rank(self) -> Int:
+        return self._shape.rank()
 
-    fn __del__(owned self):
-        if self.allocated == 1:
-            self.data.free()
+    fn simd_load[nelts: Int](self, idx: Int) -> SIMD[DType.float32, nelts]:
+        return self._data.simd_load[nelts](idx)
 
-    @always_inline
-    fn zero(inout self):
-        memset_zero(self.data, self.layers * self.rows * self.cols)
+    fn simd_load[nelts: Int](self, *indices: Int) -> SIMD[DType.float32, nelts]:
+        if len(VariadicList(indices)) > 2:
+            print(
+                "Warning: TensorSlice only supports 1D and 2D indexing.  Results are"
+                " unlikely to be correct."
+            )
+        return self.simd_load[nelts](indices[0] * self._shape[1] + indices[1])
 
-    @always_inline
-    fn size(inout self) -> Int:
-        return self.layers * self.cols * self.rows
+    fn simd_load[
+        nelts: Int
+    ](self, indices: StaticIntTuple[2]) -> SIMD[DType.float32, nelts]:
+        return self._data.simd_load[nelts](indices[0] * self._shape[1] + indices[1])
 
-    @always_inline
-    fn __getitem__(self, z: Int, y: Int, x: Int) -> Float32:
-        return self.load[1](z, y, x)
+    fn __getitem__(self, idx: Int) -> SIMD[DType.float32, 1]:
+        return self._data.simd_load[1](idx)
 
-    @always_inline
-    fn load[nelts: Int](self, z: Int, y: Int, x: Int) -> SIMD[DType.float32, nelts]:
-        return self.data.simd_load[nelts](z * self.layers + y * self.cols + x)
+    fn simd_store[nelts: Int](self, idx: Int, val: SIMD[DType.float32, nelts]):
+        return self._data.simd_store[nelts](idx, val)
 
-    @always_inline
-    fn __setitem__(self, z: Int, y: Int, x: Int, val: Float32):
-        return self.store[1](z, y, x, val)
+    fn __setitem__(self, idx: Int, val: SIMD[DType.float32, 1]):
+        return self.simd_store[1](idx, val)
 
-    @always_inline
-    fn store[nelts: Int](self, z: Int, y: Int, x: Int, val: SIMD[DType.float32, nelts]):
-        self.data.simd_store[nelts](z * self.layers + y * self.cols + x, val)
 
+fn read_val_int(inout buf: FileBuf) raises -> Int:
+    # DTypePointer[DType.ui8](buf.data).bitcast[DType.ui8]()
+    let data = buf.data.offset(buf.get_offset()).bitcast[DType.int32]()
+    let result = data.load(0)
+    buf.move_offset(4)
+    return result.to_int()
 
-struct Matrix:
-    var data: BufferPtrFloat32
-    var rows: Int
-    var cols: Int
-    var allocated: Int
 
-    fn __init__(inout self, rows: Int, cols: Int):
-        self.data = BufferPtrFloat32.alloc(0)
-        self.rows = rows
-        self.cols = cols
-        self.allocated = 0
+fn read_val_float32(inout buf: FileBuf) raises -> Float32:
+    # DTypePointer[DType.ui8](buf.data).bitcast[DType.ui8]()
+    let val = buf.data.offset(buf.get_offset()).bitcast[DType.float32]().load(0)
+    buf.move_offset(4)
+    return val
 
-    fn __init__(inout self, cols: Int):
-        self.data = BufferPtrFloat32.alloc(0)
-        self.rows = 1
-        self.cols = cols
-        self.allocated = 0
 
-    fn __del__(owned self):
-        if self.allocated == 1:
-            self.data.free()
-
-    fn alloc(inout self, fill: Int = 0):
-        self.data = BufferPtrFloat32.alloc(self.size())
-        self.allocated = 1
-        if fill == 1:
-            self.zero()
+fn read_val_str(inout buf: FileBuf, slen: Int) raises -> PointerString:
+    let str = PointerString.alloc(slen + 1)
+    for i in range(slen):
+        str.store(i, buf.data.load(buf.get_offset()))
+        buf.move_offset(1)
+    str.store(slen, 0)
 
-    fn alloc_zero(inout self):
-        self.alloc(1)
+    return str
 
-    fn zero(inout self):
-        memset_zero(self.data, self.rows * self.cols)
 
-    fn set_buf_ptr(inout self, ptr: BufferPtrFloat32):
-        self.data = ptr
+fn str_len(s: PointerString) -> Int:
+    var len = 0
+    while s[len] != 0:
+        len += 1
+    return len
 
-    # set buf ptr with redefined rows, colss
-    fn set_buf_ptr(inout self, ptr: BufferPtrFloat32, rows: Int, cols: Int):
-        self.data = ptr
-        self.rows = rows
-        self.cols = cols
 
-    fn size(inout self) -> Int:
-        return self.cols * self.rows
+# not optimal concat
+fn str_concat(s1: PointerString, s2: PointerString) -> PointerString:
+    let l1 = str_len(s1)
+    let l2 = str_len(s2)
+    let str = PointerString.alloc(l1 + l2 + 1)
+    memcpy[UInt8](str, s1, l1)
+    memcpy[UInt8](str.offset(l1), s2, l2)
+    str.store(l1 + l2, 0)
+    return str
 
-    @always_inline
-    fn __getitem__(self, y: Int, x: Int) -> Float32:
-        return self.load[1](y, x)
 
-    @always_inline
-    fn __getitem__(self, x: Int) -> Float32:
-        return self.load[1](0, x)
+fn str_to_ptr(s: String) -> PointerString:
+    let ret = PointerString.alloc(len(s) + 1)
+    for i in range(len(s)):
+        ret.store(i, ord(s[i]))
+    ret.store(len(s), 0)
+    return ret
 
-    @always_inline
-    fn load[nelts: Int](self, y: Int, x: Int) -> SIMD[DType.float32, nelts]:
-        return self.data.simd_load[nelts](y * self.cols + x)
 
-    @always_inline
-    fn __setitem__(self, y: Int, x: Int, val: Float32):
-        return self.store[1](y, x, val)
+fn string_compare(a: PointerString, b: PointerString) -> Int:
+    var index = 0
+    while a[index] != 0 and b[index] != 0:
+        if a[index] < b[index]:
+            return -1
+        if a[index] > b[index]:
+            return 1
 
-    @always_inline
-    fn __setitem__(self, x: Int, val: Float32):
-        return self.store[1](0, x, val)
+        index += 1
 
-    @always_inline
-    fn store[nelts: Int](self, y: Int, x: Int, val: SIMD[DType.float32, nelts]):
-        self.data.simd_store[nelts](y * self.cols + x, val)
+    if a[index] != 0 and b[index] == 0:
+        return 1
 
-    @always_inline
-    fn load[nelts: Int](self, x: Int) -> SIMD[DType.float32, nelts]:
-        return self.data.simd_load[nelts](x)
+    if a[index] == 0 and b[index] != 0:
+        return -1
 
-    @always_inline
-    fn store[nelts: Int](self, x: Int, val: SIMD[DType.float32, nelts]):
-        self.data.simd_store[nelts](x, val)
+    return 0
 
 
-fn read_val_int(inout buf: FileBuf) -> Int:
-    # DTypePointer[DType.ui8](buf.data).bitcast[DType.ui8]()
-    let data = buf.data.offset(buf.offset).bitcast[DType.uint32]()
-    let result = data.simd_load[1](0)
-    buf.offset += 4
-    return result.to_int()
+# Quicksort helper function to find the partition position
+fn partition(
+    inout array: PointerStrings, inout indices: DynamicVector[Int], low: Int, high: Int
+) -> Int:
+    let pivot = array[high]
+    var ii = low - 1
+    for jj in range(low, high):
+        if string_compare(pivot, array[jj]) == 1:
+            # If element smaller than pivot, swap
+            ii = ii + 1
 
+            let tmp = array[ii]
+            let tmp_idx = indices[ii]
+            array.store(ii, array[jj])
+            indices[ii] = indices[jj]
+            array.store(jj, tmp)
+            indices[jj] = tmp_idx
 
-fn read_val_float32(inout buf: FileBuf) -> Float32:
-    # DTypePointer[DType.ui8](buf.data).bitcast[DType.ui8]()
-    let val = buf.data.offset(buf.offset).bitcast[DType.float32]().simd_load[1](0)
-    buf.offset += 4
-    return val
+    # Swap the pivot element
+    let tmp = array[ii + 1]
+    let tmp_idx = indices[ii + 1]
+    array.store(ii + 1, array[high])
+    indices[ii + 1] = indices[high]
+    array.store(high, tmp)
+    indices[high] = tmp_idx
 
+    return ii + 1
 
-fn read_val_str(inout buf: FileBuf, slen: Int) -> PointerString:
-    let str = PointerString.alloc(slen + 1)
-    for i in range(slen):
-        str.store(i, buf.data.simd_load[1](buf.offset))
-        buf.offset += 1
-    str.store(slen, 0)
 
-    return str
+fn quicksort(
+    inout array: PointerStrings, inout indices: DynamicVector[Int], low: Int, high: Int
+):
+    if low < high:
+        let pi = partition(array, indices, low, high)
+        quicksort(array, indices, low, pi - 1)
+        quicksort(array, indices, pi + 1, high)
 
 
 struct FileBuf:
@@ -201,36 +227,111 @@ struct FileBuf:
         self.offset = 0
         self.size = 0
 
-    fn move_offset(inout self, size: Int):
-        self.offset += size
+    fn __del__(owned self):
+        self.data.free()
+
+    fn move_offset(inout self, size: Int) raises:
+        let new_offset = self.offset + size
+        if new_offset > self.size:
+            raise Error("Resulting offset will be past the end of the FileBuf")
+        if new_offset < 0:
+            raise Error("Resulting offset will be before the beginning of the FileBuf")
+        self.offset = new_offset
 
-    fn bitcast_offset_float32(inout self, size: Int) -> BufferPtrFloat32:
+    fn bitcast_offset_f32(inout self, size: Int) raises -> BufferPtrFloat32:
         let ret = self.data.offset(self.offset).bitcast[DType.float32]()
-        self.offset += size * sizeof[DType.float32]()
+        self.move_offset(size * sizeof[DType.float32]())
         return ret
 
+    fn get_offset(self) raises -> Int:
+        if self.offset > self.size:
+            raise Error("Offset is past the end of the FileBuf")
+        if self.offset < 0:
+            raise Error("Offset is before the beginning of the FileBuf")
+        return self.offset
+
+
+fn wrap(token: PointerString) -> PointerString:
+    if string_compare(token, str_to_ptr("\\n")) == 0:
+        return str_to_ptr("<0x0A>")
+    if string_compare(token, str_to_ptr("\\t")) == 0:
+        return str_to_ptr("<0x09>")
+    if string_compare(token, str_to_ptr("'")) == 0:
+        return str_to_ptr("<0x27>")
+    elif string_compare(token, str_to_ptr('"')) == 0:
+        return str_to_ptr("<0x22>")
+    return token
+
 
 struct Tokenizer:
     var vocab: PointerStrings
     var vocab_scores: BufferPtrFloat32
     var max_token_length: Int
     var vocab_size: Int
+    var sorted_vocab: PointerStrings
+    var sorted_indices: DynamicVector[Int]
 
-    fn __init__(inout self, vocab_size: Int):
+    fn __init__(inout self, vocab_size: Int, inout buf: FileBuf) raises -> None:
         self.vocab_size = vocab_size
-        self.vocab = PointerStrings.alloc(vocab_size)
-        self.vocab_scores = BufferPtrFloat32.alloc(vocab_size)
-        self.max_token_length = 0
+        self.max_token_length = read_val_int(buf)
+        self.vocab_scores = BufferPtrFloat32.alloc(self.vocab_size)
+        self.vocab = PointerStrings.alloc(self.vocab_size)
+        # lazy load sorted vocab
+        self.sorted_vocab = PointerStrings.alloc(0)
+        self.sorted_indices = DynamicVector[Int](0)
+
+        # read vocab_scores & vocab values (tokens)
+        for i in range(0, self.vocab_size):
+            self.vocab_scores.store(i, read_val_float32(buf))
+            let slen = read_val_int(buf)
+            self.vocab.store(i, read_val_str(buf, slen))
+
+        return None
+
+    # sort vocab by string_compare
+    fn sort(inout self) -> None:
+        if len(self.sorted_indices) < self.vocab_size:
+            self.sorted_indices = DynamicVector[Int](self.vocab_size)
+            self.sorted_vocab = PointerStrings.alloc(self.vocab_size)
+            for ii in range(self.vocab_size):
+                self.sorted_vocab.store(ii, self.vocab[ii])
+                self.sorted_indices.push_back(ii)
+
+        let n = self.vocab_size
+        quicksort(self.sorted_vocab, self.sorted_indices, 0, n - 1)
+        return None
+
+    # Binary search that returns -1 if string is not found
+    fn find(inout self, token_o: PointerString) -> Int:
+        let token = wrap(token_o)
+        let n = self.vocab_size
+        if len(self.sorted_indices) < n:
+            self.sort()
+        var left = 0
+        var right = n - 1
+        while left <= right:
+            let mid = left + (right - left) // 2
+            let comparison = string_compare(self.sorted_vocab[mid], token)
+            if comparison == 0:
+                return self.sorted_indices[mid]
+            if comparison < 0:
+                left = mid + 1
+            else:
+                right = mid - 1
+        return -1
 
 
 struct Config:
     var dim: Int
+    var kv_dim: Int
     var hidden_dim: Int
     var n_layers: Int
     var n_heads: Int
     var n_kv_heads: Int
+    var kv_mul: Int
     var vocab_size: Int
     var seq_len: Int
+    var head_size: Int
 
     fn __init__(inout self):
         self.dim = 0
@@ -240,109 +341,93 @@ struct Config:
         self.n_kv_heads = 0
         self.vocab_size = 0
         self.seq_len = 0
+        self.kv_dim = 0
+        self.kv_mul = 0
+        self.head_size = 0
 
 
 struct RunState:
-    var x: Matrix  # activation at current time stamp (dim,)
-    var xb: Matrix  # same, but inside a residual branch (dim,)
-    var xb2: Matrix  # an additional buffer just for convenience (dim,)
-    var hb: Matrix  # buffer for hidden dimension in the ffn (hidden_dim,)
-    var hb2: Matrix  # buffer for hidden dimension in the ffn (hidden_dim,)
-    var q: Matrix  # query (dim,)
-    var k: Matrix  # key (dim,)
-    var v: Matrix  # value (dim,)
-    var att: Matrix  # buffer for scores/attention values (n_heads, seq_len)
-    var logits: Matrix  # output logits
-    var key_cache: Matrix3  # (layer, seq_len, dim)
-    var value_cache: Matrix3  # (layer, seq_len, dim)
-
-    fn __init__(inout self, config: Config):
-        self.x = Matrix(config.dim)
-        self.x.alloc_zero()
-        self.xb = Matrix(config.dim)
-        self.xb.alloc_zero()
-        self.xb2 = Matrix(config.dim)
-        self.xb2.alloc_zero()
-        self.hb = Matrix(config.hidden_dim)
-        self.hb.alloc_zero()
-        self.hb2 = Matrix(config.hidden_dim)
-        self.hb2.alloc_zero()
-        self.q = Matrix(config.dim)
-        self.q.alloc_zero()
-        self.k = Matrix(config.dim)
-        self.k.alloc_zero()
-        self.v = Matrix(config.dim)
-        self.v.alloc_zero()
-        self.att = Matrix(config.n_heads, config.seq_len)
-        self.att.alloc_zero()
-        self.logits = Matrix(config.vocab_size)
-        self.logits.alloc_zero()
-        self.key_cache = Matrix3(config.n_layers, config.seq_len, config.dim)
-        self.key_cache.alloc_zero()
-        self.value_cache = Matrix3(config.n_layers, config.seq_len, config.dim)
-        self.value_cache.alloc_zero()
+    var x: TensorF32  # activation at current time stamp (dim,)
+    var xb: TensorF32  # same, but inside a residual branch (dim,)
+    var xb2: TensorF32  # an additional buffer just for convenience (dim,)
+    var hb: TensorF32  # buffer for hidden dimension in the ffn (hidden_dim,)
+    var hb2: TensorF32  # buffer for hidden dimension in the ffn (hidden_dim,)
+    var q: TensorF32  # query (dim,)
+    var k: TensorSlice  # key (kv_dim,)
+    var v: TensorSlice  # value (kv_dim,)
+    var att: TensorF32  # buffer for scores/attention values (n_heads, seq_len)
+    var logits: TensorF32  # output logits
+    var key_cache: TensorF32  # (layer, seq_len, dim)
+    var value_cache: TensorF32  # (layer, seq_len, dim)
+
+
+    fn __init__(inout self, config: Config) raises:
+        self.x = TensorF32(config.dim)
+        self.xb = TensorF32(config.dim)
+        self.xb2 = TensorF32(config.dim)
+        self.hb = TensorF32(config.hidden_dim)
+        self.hb2 = TensorF32(config.hidden_dim)
+        self.q = TensorF32(config.dim)
+        self.att = TensorF32(config.n_heads, config.seq_len)
+        self.logits = TensorF32(config.vocab_size)
+        self.key_cache = TensorF32(config.n_layers, config.seq_len, config.kv_dim)
+        self.value_cache = TensorF32(config.n_layers, config.seq_len, config.kv_dim)
+        # So their updates flow to the caches, k and v are slices with shared memory.
+        # Initialize with placeholders. The real tensors reference layer and position during forward pass.
+        self.k = TensorSlice(TensorF32(TensorShape(1, config.kv_dim)), 1)
+        self.v = TensorSlice(TensorF32(TensorShape(1, config.kv_dim)), 1)
+
 
 
 struct TransformerWeights:
-    var token_embedding_table: Matrix
-    var freq_cis_real: Matrix
-    var freq_cis_imag: Matrix
-    var rms_att_weight: Matrix
-    var wq: Matrix3
-    var wk: Matrix3
-    var wv: Matrix3
-    var wo: Matrix3
-    var rms_ffn_weight: Matrix
-    var w1: Matrix3
-    var w3: Matrix3
-    var w2: Matrix3
-    var rms_final_weight: Matrix
-    var wcls: Matrix
-
-    fn __init__(inout self, config: Config, shared_weights: Int, inout buf: FileBuf):
-        self.token_embedding_table = Matrix(config.vocab_size, config.dim)
-        # set buf ptr to buf data from file
-        self.token_embedding_table.set_buf_ptr(
-            buf.bitcast_offset_float32(self.token_embedding_table.size())
-        )
-        self.rms_att_weight = Matrix(config.n_layers, config.dim)
-        self.rms_att_weight.set_buf_ptr(
-            buf.bitcast_offset_float32(self.rms_att_weight.size())
-        )
-        self.wq = Matrix3(config.n_layers, config.dim, config.dim)
-        self.wq.set_buf_ptr(buf.bitcast_offset_float32(self.wq.size()))
-        self.wk = Matrix3(config.n_layers, config.dim, config.dim)
-        self.wk.set_buf_ptr(buf.bitcast_offset_float32(self.wk.size()))
-        self.wv = Matrix3(config.n_layers, config.dim, config.dim)
-        self.wv.set_buf_ptr(buf.bitcast_offset_float32(self.wv.size()))
-        self.wo = Matrix3(config.n_layers, config.dim, config.dim)
-        self.wo.set_buf_ptr(buf.bitcast_offset_float32(self.wo.size()))
-        self.rms_ffn_weight = Matrix(config.n_layers, config.dim)
-        self.rms_ffn_weight.set_buf_ptr(
-            buf.bitcast_offset_float32(self.rms_ffn_weight.size())
-        )
-        self.w1 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
-        self.w1.set_buf_ptr(buf.bitcast_offset_float32(self.w1.size()))
-        self.w2 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
-        self.w2.set_buf_ptr(buf.bitcast_offset_float32(self.w2.size()))
-        self.w3 = Matrix3(config.n_layers, config.dim, config.hidden_dim)
-        self.w3.set_buf_ptr(buf.bitcast_offset_float32(self.w3.size()))
-        self.rms_final_weight = Matrix(config.dim)
-        self.rms_final_weight.set_buf_ptr(
-            buf.bitcast_offset_float32(self.rms_final_weight.size())
-        )
-        self.freq_cis_real = Matrix(config.seq_len, (config.dim // config.n_heads) // 2)
-        self.freq_cis_real.set_buf_ptr(
-            buf.bitcast_offset_float32(self.freq_cis_real.size())
-        )
-        self.freq_cis_imag = Matrix(config.seq_len, (config.dim // config.n_heads) // 2)
-        self.freq_cis_imag.set_buf_ptr(
-            buf.bitcast_offset_float32(self.freq_cis_imag.size())
-        )
-        self.wcls = Matrix(
-            config.vocab_size, config.dim
-        )  # if shared_weights else rest_floats
-        self.wcls.set_buf_ptr(self.token_embedding_table.data)
+    var token_embedding_table: TensorF32
+    var freq_cis_real: TensorF32
+    var freq_cis_imag: TensorF32
+    var rms_att_weight: TensorF32
+    var wq: TensorF32
+    var wk: TensorF32
+    var wv: TensorF32
+    var wo: TensorF32
+    var rms_ffn_weight: TensorF32
+    var w1: TensorF32
+    var w3: TensorF32
+    var w2: TensorF32
+    var rms_final_weight: TensorF32
+    var wcls: TensorF32
+
+    fn __init__(
+        inout self, config: Config, shared_weights: Int, inout buf: FileBuf
+    ) raises:
+        fn load_weights(inout buf: FileBuf, *dims: Int) raises -> TensorF32:
+            # Ensure returned Tensor doesn't share a pointer with FileBuf
+            let shape = TensorShape(dims)
+            let result_data = BufferPtrFloat32.alloc(shape.num_elements())
+            memcpy(
+                result_data,
+                buf.bitcast_offset_f32(shape.num_elements()),
+                shape.num_elements(),
+            )
+            return TensorF32(result_data, shape)
+
+        self.token_embedding_table = load_weights(buf, config.vocab_size, config.dim)
+        self.rms_att_weight = load_weights(buf, config.n_layers, config.dim)
+        self.wq = load_weights(buf, config.n_layers, config.dim, config.dim)
+        self.wk = load_weights(buf, config.n_layers, config.kv_dim, config.dim)
+        self.wv = load_weights(buf, config.n_layers, config.kv_dim, config.dim)
+        self.wo = load_weights(buf, config.n_layers, config.dim, config.dim)
+        self.rms_ffn_weight = load_weights(buf, config.n_layers, config.dim)
+        self.w1 = load_weights(buf, config.n_layers, config.hidden_dim, config.dim)
+        self.w2 = load_weights(buf, config.n_layers, config.dim, config.hidden_dim)
+        self.w3 = load_weights(buf, config.n_layers, config.hidden_dim, config.dim)
+        self.rms_final_weight = load_weights(buf, config.dim)
+        # maybe need modifying for different model
+        # config.head_size // 2 for stories and tinyllama-1.1
+        self.freq_cis_real = load_weights(buf, config.seq_len, config.head_size // 2)
+        self.freq_cis_imag = load_weights(buf, config.seq_len, config.head_size // 2)
+        if shared_weights:
+            self.wcls = self.token_embedding_table
+        else:
+            self.wcls = load_weights(buf, config.vocab_size, config.dim)
 
 
 fn read_file(file_name: String, inout buf: FileBuf) raises:
@@ -375,270 +460,323 @@ fn config_init(inout config: Config, inout buf: FileBuf) raises:
     config.n_kv_heads = read_val_int(buf)
     config.vocab_size = read_val_int(buf)
     config.seq_len = read_val_int(buf)
+    config.head_size = config.dim // config.n_heads
+    config.kv_dim = (config.n_kv_heads * config.dim) // config.n_heads
+    config.kv_mul = config.n_heads // config.n_kv_heads
     return None
 
 
-fn tokenizer_init(inout tok: Tokenizer, inout buf: FileBuf) -> None:
-    tok.max_token_length = read_val_int(buf)
-    tok.vocab_scores = BufferPtrFloat32.alloc(tok.vocab_size)
-    tok.vocab = PointerStrings.alloc(tok.vocab_size)
-
-    # read vocab_scores & vocab values (tokens)
-    for i in range(0, tok.vocab_size):
-        tok.vocab_scores.simd_store[1](i, read_val_float32(buf))
-        let slen = read_val_int(buf)
-        tok.vocab.store(i, read_val_str(buf, slen))
-
-    tok.vocab_scores = buf.data.offset(buf.offset).bitcast[DType.float32]()
-    buf.offset += tok.vocab_size * 4
-    return None
+@always_inline
+fn accum(inout a: TensorF32, b: TensorF32) -> None:
+    let size = a.dim(0)
 
+    @parameter
+    fn _acc[_nelts: Int](j: Int):
+        a.simd_store[_nelts](j, a.simd_load[_nelts](j) + b.simd_load[_nelts](j))
 
-fn accum(inout a: BufferPtrFloat32, b: BufferPtrFloat32, size: Int) -> None:
-    for i in range(size):
-        let val = a.offset(i).simd_load[1](0) + b.offset(i).simd_load[1](0)
-        a.offset(i).simd_store[1](0, val)
+    vectorize[nelts, _acc](size)
 
 
+@always_inline
 fn rmsnorm(
     inout o: BufferPtrFloat32, x: BufferPtrFloat32, weight: BufferPtrFloat32, size: Int
 ) -> None:
     # Calculate sum of squares
-    var ss: Float32 = 0.0
-    for i in range(size):
-        let xx = x.offset(i).simd_load[1](0) ** 2
-        ss += xx
+    var tmp = SIMD[DType.float32, nelts](0)
+
+    @parameter
+    fn _sum2[_nelts: Int](j: Int):
+        if _nelts < nelts:
+            tmp[0] += (x.offset(j).simd_load[_nelts](0) ** 2).reduce_add()
+        else:
+            tmp += x.offset(j).simd_load[nelts](0) ** 2
+
+    vectorize[nelts, _sum2](size)
+
+    var ss: Float32 = tmp.reduce_add()
     ss = ss / size + 1e-5
     ss = 1.0 / math.sqrt(ss)
+
     # Normalize and scale
-    for j in range(size):
-        let val = weight.offset(j).simd_load[1](0) * (ss * x.offset(j).simd_load[1](0))
-        o.offset(j).simd_store[1](0, val)
-
-
-fn softmax(inout x: BufferPtrFloat32, size: Int) -> None:
-    # Find max value (for numerical stability)
-    var max_val: Float32 = x.offset(0).simd_load[1](0)
-    for i in range(size):
-        let xi = x.offset(i).simd_load[1](0)
-        if xi > max_val:
-            max_val = xi
-    # Exp and sum
+    @parameter
+    fn _norm[_nelts: Int](j: Int):
+        let val = weight.simd_load[_nelts](j) * ss * x.simd_load[_nelts](j)
+        o.offset(j).simd_store[_nelts](0, val)
+
+    vectorize[nelts, _norm](size)
+
+
+@always_inline
+fn rmsnorm(inout o: TensorF32, x: TensorF32, weight: TensorF32):
+    rmsnorm(o._ptr, x.data(), weight.data(), weight.dim(weight.rank() - 1))
+
+
+@always_inline
+fn rmsnorm(inout o: TensorF32, x: TensorF32, weight: TensorSlice):
+    rmsnorm(o._ptr, x.data(), weight.data(), weight.dim(weight.rank() - 1))
+
+
+@always_inline
+fn softmax(inout x: TensorF32) -> None:
+    softmax(x, 0, x.dim(0))
+
+
+@always_inline
+fn softmax(inout x: TensorF32, start: Int, end: Int):
+    var max_val: Float32 = -1e9
+
+    @parameter
+    fn _max[_nelts: Int](ii: Int):
+        let val = x.simd_load[_nelts](start + ii).reduce_max()
+        if val > max_val:
+            max_val = val
+
+    vectorize[nelts, _max](end - start)
+
     var ssum: Float32 = 0.0
-    for i in range(size):
-        let xi = x.offset(i).simd_load[1](0)
-        x.offset(i).simd_store[1](0, math.exp(xi - max_val))
-        ssum += x.offset(i).simd_load[1](0)
-    # Normalize
-    for i in range(size):
-        let xi = x.offset(i).simd_load[1](0)
-        x.offset(i).simd_store[1](0, xi / ssum)
-
-
-fn matmul_naive(C: Matrix, x: Matrix, w: Matrix) -> None:
-    # W(d,n) @ X(n,) -> C (d,)
-    # By far the most amount of time is spent inside this little function
-    for i in range(w.rows):
-        C[i] = 0.0
-        for j in range(w.cols):
-            C[i] += x[j] * w[i, j]
-
-
-fn matmul_vectorized(C: Matrix, A: Matrix, B: Matrix):
-    for i in range(0, B.rows):
+
+    @parameter
+    fn _exp[_nelts: Int](ii: Int):
+        x.simd_store[_nelts](
+            start + ii, math.exp(x.simd_load[_nelts](start + ii) - max_val)
+        )
+        ssum += x.simd_load[_nelts](start + ii).reduce_add()
+
+    vectorize[nelts, _exp](end - start)
+
+    @parameter
+    fn _norm[_nelts: Int](ii: Int):
+        x.simd_store[_nelts](start + ii, x.simd_load[_nelts](start + ii) / ssum)
+
+    vectorize[nelts, _norm](end - start)
+
+
+@always_inline
+fn matmul_parallelized(C: BufferPtrFloat32,A: BufferPtrFloat32,B: BufferPtrFloat32,rows: Int,cols: Int,):
+    @parameter
+    fn compute_row(i: Int):
         var tmp = SIMD[DType.float32, nelts](0)
 
         @parameter
         fn dot[_nelts: Int](j: Int):
-            if _nelts < nelts: # take care of tail array elements with length <  nelts
-                tmp[0] += (A.load[_nelts](j) * B.load[_nelts](i, j)).reduce_add()
+            if _nelts < nelts:  # take care of tail array elements with length <  nelts
+                tmp[0] += (
+                A.simd_load[_nelts](j) * B.simd_load[_nelts](i * cols + j)
+                ).reduce_add()
             else:
-                tmp += A.load[nelts](j) * B.load[nelts](i, j)
+                tmp += A.simd_load[nelts](j) * B.simd_load[nelts](i * cols + j)
 
-        vectorize[nelts, dot](B.cols)
-        C[i] = tmp.reduce_add()
+        vectorize[nelts, dot](cols)
+        C.store(i, tmp.reduce_add())
+    
 
-fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix):
-    @parameter
-    fn calc_row(i: Int):
-        var T = BufferPtrFloat32.alloc(nelts)
-        var Tbuf = Buffer[nelts, DType.float32](T)
-        memset_zero(T, nelts)
-        @parameter
-        fn dot[nelts: Int](j: Int):
-            T.simd_store[nelts](
-                0, T.simd_load[nelts](0) + A.load[nelts](j) * B.load[nelts](i, j)
-            )
+    parallelize[compute_row](rows, workers)
 
-        vectorize[nelts, dot](B.cols)
-        C[i] = sum[nelts, DType.float32](Tbuf)
+    
 
-    parallelize[calc_row](B.rows)
+    
+    
+@always_inline
+fn matmul(C: TensorF32, A: TensorF32, B: TensorF32) raises:
+    # B (d,n) @ A (n,) -> C (d,)
+    matmul_dimension_checks(A.shape(), B.shape())
+    matmul_parallelized(C.data(), A.data(), B.data(), B.dim(0), B.dim(1))
 
 
-fn matmul(inout C: Matrix, A: Matrix, B: Matrix) -> None:
+@always_inline
+fn matmul(C: TensorF32, A: TensorF32, B: TensorSlice) raises:
     # B (d,n) @ A (n,) -> C (d,)
-    matmul_vectorized(C, A, B)
-    # matmul_parallelized(C, A, B)
+    matmul_dimension_checks(A.shape(), B.shape())
+    matmul_parallelized(C.data(), A.data(), B.data(), B.dim(0), B.dim(1))
+
+
+@always_inline
+fn matmul(C: TensorSlice, A: TensorF32, B: TensorSlice) raises:
+    # B (d,n) @ A (n,) -> C (d,)
+    matmul_dimension_checks(A.shape(), B.shape())
+    matmul_parallelized(C.data(), A.data(), B.data(), B.dim(0), B.dim(1))
+
+
+fn matmul_dimension_checks(a: TensorShape, b: TensorShape) raises:
+    if a[0] != b[1]:
+        raise Error(
+            "matmul dimension mismatch. A rows (dim 0) not equal to B columns (dim 1)"
+        )
+    if b.rank() != 2:
+        raise Error("matmul expects B to be a 2D matrix")
 
 
+# Apply RoPE rotation to the q and k vectors for each head
+# rotate odd and even dim
+@always_inline
+fn rope_rotation_llama(
+    inout state: RunState,
+    freq_cis_real_row: TensorSlice,
+    freq_cis_imag_row: TensorSlice,
+    config: Config,
+) -> None:
+    # stories model, llama2
+    let head_size = config.head_size
+    @parameter
+    fn head_loop(i:Int):
+        # Simple vectorization with (head_size // 2) steps gave junk transformer output.
+        # Maybe because the nelt ranges end up overlapping between the steps.
+        for j in range(0, config.head_size, 2):
+            let fcr = freq_cis_real_row[j // 2]
+            let fci = freq_cis_imag_row[j // 2]
+            let q0 = state.q[i * head_size + j]
+            let q1 = state.q[i * head_size + j + 1]
+            state.q[i * head_size + j] = q0 * fcr - q1 * fci
+            state.q[i * head_size + j + 1] = q0 * fci + q1 * fcr
+            if i < config.n_kv_heads:
+                let k0 = state.k[i * head_size + j]
+                let k1 = state.k[i * head_size + j + 1]
+                state.k[i * head_size + j] = k0 * fcr - k1 * fci
+                state.k[i * head_size + j + 1] = k0 * fci + k1 * fcr
+    parallelize[head_loop](config.n_heads, workers)
+
+
+
+@always_inline
 fn transformer(
     token: Int,
     pos: Int,
     config: Config,
     inout state: RunState,
     weights: TransformerWeights,
-) -> None:
+) raises -> None:
     # A few convenience variables
-    var x = state.x.data
     let dim = config.dim
     let hidden_dim = config.hidden_dim
-    let head_size = dim // config.n_heads
-
-    # tmp matrix for matmul operations
-    var tmpw = Matrix(0, 0)
+    let head_size = config.head_size
+    let kv_dim = config.kv_dim
+    let kv_mul = config.kv_mul
 
     # Copy the token embedding into x
-    let content_row = weights.token_embedding_table.data.offset(token * dim)
-    memcpy[DType.float32](x, content_row, config.dim)
+    let content_row = weights.token_embedding_table.data().offset(token * dim)
+    memcpy[DType.float32](state.x.data(), content_row, dim)
 
     # Pluck out the "pos" row of freq_cis_real and freq_cis_imag
-    let freq_cis_real_row = weights.freq_cis_real.data.offset(pos * head_size // 2)
-    let freq_cis_imag_row = weights.freq_cis_imag.data.offset(pos * head_size // 2)
+    let freq_cis_real_row = TensorSlice(weights.freq_cis_real, pos)
+    let freq_cis_imag_row = TensorSlice(weights.freq_cis_imag, pos)
 
     # Forward all the layers
     for l in range(config.n_layers):
         # Attention rmsnorm
-        rmsnorm(state.xb.data, x, weights.rms_att_weight.data.offset(l * dim), dim)
-
+        rmsnorm(state.xb, state.x, TensorSlice(weights.rms_att_weight, l))
         # QKV matmuls for this position
-        tmpw.set_buf_ptr(weights.wq.data.offset(l * dim * dim), dim, dim)
-        matmul(state.q, state.xb, tmpw)
+        matmul(state.q, state.xb, TensorSlice(weights.wq, l))
 
-        tmpw.set_buf_ptr(weights.wk.data.offset(l * dim * dim), dim, dim)
-        matmul(state.k, state.xb, tmpw)
+        let loff = l * config.seq_len * config.kv_dim
+        state.k = TensorSlice(state.key_cache, l, pos)
+        matmul(state.k, state.xb, TensorSlice(weights.wk, l))
 
-        tmpw.set_buf_ptr(weights.wv.data.offset(l * dim * dim), dim, dim)
-        matmul(state.v, state.xb, tmpw)
+        state.v = TensorSlice(state.value_cache, l, pos)
+        matmul(state.v, state.xb, TensorSlice(weights.wv, l))
 
         # Apply RoPE rotation to the q and k vectors for each head
-        for h in range(config.n_heads):
-            # Get the q and k vectors for this head
-            let q = state.q.data.offset(h * head_size)
-            let k = state.k.data.offset(h * head_size)
-
-            # Rotate q and k by the freq_cis_real and freq_cis_imag
-            for i in range(0, head_size, 2):
-                let q0 = q.offset(i).simd_load[1](0)
-                let q1 = q.offset(i + 1).simd_load[1](0)
-                let k0 = k.offset(i).simd_load[1](0)
-                let k1 = k.offset(i + 1).simd_load[1](0)
-                let fcr = freq_cis_real_row.offset(i // 2).simd_load[1](0)
-                let fci = freq_cis_imag_row.offset(i // 2).simd_load[1](0)
-                q.offset(i).simd_store[1](0, q0 * fcr - q1 * fci)
-                q.offset(i + 1).simd_store[1](0, q0 * fci + q1 * fcr)
-                k.offset(i).simd_store[1](0, k0 * fcr - k1 * fci)
-                k.offset(i + 1).simd_store[1](0, k0 * fci + k1 * fcr)
-
-        # Save key,value at this time step (pos) to our kv cache
-        let loff = l * config.seq_len * dim  # kv cache layer offset for convenience
-        let key_cache_row = state.key_cache.data.offset(loff + pos * dim)
-        let value_cache_row = state.value_cache.data.offset(loff + pos * dim)
-        memcpy[DType.float32](key_cache_row, state.k.data, config.dim)
-        memcpy[DType.float32](value_cache_row, state.v.data, config.dim)
-
-        # Multihead attention. Iterate over all heads
-        for h in range(config.n_heads):
+        rope_rotation_llama(state, freq_cis_real_row, freq_cis_imag_row, config)
+
+        memset_zero(state.xb.data(), state.xb.num_elements())
+
+        # Multihead attention. Iterate over all heads in parallel.
+        @parameter
+        fn loop_over_heads(h:Int):
             # Get the query vector for this head
-            let q = state.q.data.offset(h * head_size)
+            let q_offset = h * head_size
 
-            # Attention scores for this head
-            var att = state.att.data.offset(h * config.seq_len)
+            # Index of attention scores for this head
+            let att_offset = h * config.seq_len
 
             # Iterate over all timesteps, including the current one
             for t in range(pos + 1):
-                # Get the key vector for this head and at this timestep
-                let k = state.key_cache.data.offset(loff + t * dim + h * head_size)
+                # Starting index of the key vector for this head and at this timestep
+                let k_offset = loff + t * kv_dim + (h // kv_mul) * head_size
                 # Calculate the attention score as the dot product of q and k
                 var score: Float32 = 0.0
-                for i in range(head_size):
-                    score += q.offset(i).simd_load[1](0) * k.offset(i).simd_load[1](0)
+
+                @parameter
+                fn score_fn[_nelts: Int](i: Int):
+                    score += (
+                        state.q.simd_load[_nelts](q_offset + i)
+                        * state.key_cache.simd_load[_nelts](k_offset + i)
+                    ).reduce_add()
+
+                vectorize[nelts, score_fn](head_size)
                 score /= math.sqrt[DType.float32, 1](head_size)
 
                 # Save the score to the attention buffer
-                att.offset(t).simd_store[1](0, score)
+                state.att[att_offset + t] = score
 
             # Softmax the scores to get attention weights, from 0..pos inclusively
-            softmax(att, pos + 1)
-
+            softmax(state.att, att_offset, att_offset + pos + 1)
             # Weighted sum of the values, store back into xb
-            let xb = state.xb.data.offset(h * head_size)
-            memset_zero(xb, head_size)
+            let xb_offset = h * head_size
             for t in range(pos + 1):
-                # Get the value vector for this head and at this timestep
-                let v = state.value_cache.data.offset(loff + t * dim + h * head_size)
+                # Starting index of the value vector for this head and at this timestep
+                let v_offset = loff + t * kv_dim + (h // kv_mul) * head_size
+
                 # Get the attention weight for this timestep
-                let a = att.offset(t).simd_load[1](0)
+                let a = state.att[att_offset + t]
                 # Accumulate the weighted value into xb
-                for i in range(head_size):
-                    let xbi = xb.offset(i).simd_load[1](0) + a * v.offset(i).simd_load[
-                        1
-                    ](0)
-                    xb.offset(i).simd_store[1](0, xbi)
-        # Final matrix multiplication to get the output of the attention
-        tmpw.set_buf_ptr(weights.wo.data.offset(l * dim * dim), dim, dim)
-        matmul(state.xb2, state.xb, tmpw)
 
-        # Residual connection back into x
-        accum(x, state.xb2.data, dim)
+                @parameter
+                fn xb_accumulate[_nelts: Int](i: Int):
+                    let xbi = state.xb.simd_load[_nelts](
+                        xb_offset + i
+                    ) + a * state.value_cache.simd_load[_nelts](v_offset + i)
+                    state.xb.simd_store[_nelts](xb_offset + i, xbi)
 
+                vectorize[nelts, xb_accumulate](head_size)
+
+        parallelize[loop_over_heads](config.n_heads, workers)
+        # Final matrix multiplication to get the output of the attention
+        matmul(state.xb2, state.xb, TensorSlice(weights.wo, l))
+        # Residual connection back into x
+        accum(state.x, state.xb2)
         # FFN rmsnorm
-        rmsnorm(state.xb.data, x, weights.rms_ffn_weight.data.offset(l * dim), dim)
+        rmsnorm(state.xb, state.x, TensorSlice(weights.rms_ffn_weight, l))
 
         # Calculate self.w1(x) and self.w3(x) for FFN
-        tmpw.set_buf_ptr(weights.w1.data.offset(l * dim * hidden_dim), hidden_dim, dim)
-        matmul(state.hb, state.xb, tmpw)
-
-        tmpw.set_buf_ptr(weights.w3.data.offset(l * dim * hidden_dim), hidden_dim, dim)
-        matmul(state.hb2, state.xb, tmpw)
-
-        # Apply SiLU activation function (silu(x) = x * sigmoid(x))
-        for i in range(hidden_dim):
-            let hbi = state.hb[i]
-            state.hb[i] = hbi * (1.0 / (1.0 + math.exp(-hbi)))
+        matmul(state.hb, state.xb, TensorSlice(weights.w1, l))
 
-        # Elementwise multiply with w3(x)
-        for i in range(hidden_dim):
-            state.hb[i] = state.hb[i] * state.hb2[i]
+        matmul(state.hb2, state.xb, TensorSlice(weights.w3, l))
 
+        @parameter
+        fn silu[_nelts: Int](i: Int):
+            let initial_hb = state.hb.simd_load[_nelts](i)
+            # Apply SiLU activation function (silu(x) = x * sigmoid(x))
+            let hbi = initial_hb * (1.0 / (1.0 + math.exp(-initial_hb)))
+            # Elementwise multiply with w3(x)
+            state.hb.simd_store[_nelts](i, hbi * state.hb2.simd_load[_nelts](i))
+
+        vectorize[nelts, silu](hidden_dim)
         # Final matrix multiplication to get the output of the FFN
-        tmpw.set_buf_ptr(weights.w2.data.offset(l * dim * hidden_dim), dim, hidden_dim)
-        matmul(state.xb, state.hb, tmpw)
+        matmul(state.xb, state.hb, TensorSlice(weights.w2, l))
 
         # Residual connection
-        accum(x, state.xb.data, dim)
+        accum(state.x, state.xb)
 
     # Final rmsnorm
-    rmsnorm(x, x, weights.rms_final_weight.data, dim)
+    rmsnorm(state.x, state.x, weights.rms_final_weight)
 
     # Classifier into logits
-    tmpw.set_buf_ptr(weights.wcls.data, config.vocab_size, dim)
-    matmul(state.logits, state.x, tmpw)
+    matmul(state.logits, state.x, weights.wcls)
 
 
-fn argmax(v: Matrix) -> Int:
+fn argmax(v: TensorF32) -> Int:
     # return argmax of v
     var max_i: Int = 0
     var max_p: Float32 = v[0]
-    for i in range(v.cols):
+    for i in range(v.dim(0)):
         if v[i] > max_p:
             max_i = i
             max_p = v[i]
     return max_i
 
 
-fn sample(probabilities: Matrix) -> Int:
-    let n = probabilities.cols
+fn sample(probabilities: TensorF32) -> Int:
+    let n = probabilities.dim(0)
     # Sample index from probabilities, they must sum to 1
     # get random value within (min, max) float32 range
     let r = DTypePointer[DType.float32].alloc(1)
@@ -646,12 +784,64 @@ fn sample(probabilities: Matrix) -> Int:
     var cdf: Float32 = 0.0
     for i in range(n):
         cdf += probabilities[i]
-        if r.simd_load[1](0) < cdf:
+        if r.load(0) < cdf:
             return i
     return n - 1  # In case of rounding errors
 
 
+fn bpe_encode(inout tokens: DynamicVector[Int], text: String, inout tok: Tokenizer):
+    for pos in range(len(text)):
+        let char = str_to_ptr(text[pos])
+        let id = tok.find(char)
+
+        if id == -1:
+            print("Not a good prompt token at pos ", pos)
+            return
+        tokens.push_back(id)
+
+    while True:
+        var best_score = Float32(-1e10)
+        var best_id = -1
+        var best_idx = -1
+
+        for i in range(len(tokens) - 1):
+            # Check if we can merge the pair (tokens[i], tokens[i+1])
+            let str = str_concat(tok.vocab[tokens[i]], tok.vocab[tokens[i + 1]])
+            let id = tok.find(str)
+            if id != -1 and tok.vocab_scores.load(id) > best_score:
+                best_score = tok.vocab_scores.load(id)
+                best_id = id
+                best_idx = i
+
+        if best_idx == -1:
+            # We couldn't find any more pairs to merge, so we're done
+            break
+
+        # Merge the consecutive pair (best_idx, best_idx+1) into new token best_id
+        tokens[best_idx] = best_id
+        # Delete token at position best_idx+1, shift the entire sequence back 1
+        var _tokens = DynamicVector[Int]()
+        for i in range(0, best_idx + 1):
+            _tokens.push_back(tokens[i])
+        for i in range(best_idx + 2, len(tokens)):
+            _tokens.push_back(tokens[i])
+        tokens = _tokens
+
+
+fn str2num(d: Int) -> Int:
+    # covert Hex to decimal
+    if d >= ord("A"):
+        return d - ord("A") + 10
+    return d - ord("0")
+
+
 fn print_str(s: PointerString):
+    # print raw byte like <0x0A>
+    if (s[1].to_int() == ord("0")) and (s[2].to_int() == ord("x")):
+        let d1: Int = s[3].to_int()
+        let d2: Int = s[4].to_int()
+        print_no_newline(chr(str2num(d1) * 16 + str2num(d2)))
+        return
     # print all chars till null character
     var p: Int = 0
     while s[p].to_int() != 0:
@@ -664,22 +854,76 @@ fn time_in_ms() -> Int:
     return time.now() // 1_000_000
 
 
+fn print_usage():
+    print("Usage: mojo llama2.mojo <checkpoint> [options]")
+    print(
+        'Example: mojo llama2.mojo stories15M.bin -s 99 -n 256 -t 0.5 -i "Llama is an'
+        ' animal"'
+    )
+    print("Options:")
+    print("  -s <int>    random seed, default time.now()")
+    print("  -t <float>  temperature in [0,1.0], default 1.0")
+    print("  -n <int>    number of steps to run for, default 256. 0 = max_seq_len")
+    print("  -i <string> input prompt")
+    print("  -z          tokenizer path")
+    print("  -j          number of workers to use, default num_cores()")
+
+
 fn main() raises:
-    print("num hardware threads: ", num_cores(), " SIMD vector width: ", nelts)
-    let checkpoint = "stories15M.bin"
-    # let checkpoint = "stories110M.bin"
-    let tokenizer = "tokenizer.bin"
-    let temperature = 0.0
+    workers = num_cores()
+    var tokenizer = StringRef("tokenizer.bin")
+    var checkpoint = StringRef("stories15M.bin")
+    var temperature = 0.9
     var steps = 256
-    let prompt = ""
-    let rng_seed: Int = time.now()
+    var prompt = String("")
+    var rng_seed: Int = time.now()
+
+    @parameter
+    fn argparse() raises -> Int:
+        let args = argv()
+        if len(args) < 2:
+            return 0
+        checkpoint = args[1]
+        for i in range(2, len(args), 2):
+            if args[i] == "-p":
+                print("Option not supported: ", args[i])
+            if args[i] == "-n":
+                steps = atol(args[i + 1])
+            if args[i] == "-z":
+                tokenizer = args[i + 1]
+            if args[i] == "-s":
+                rng_seed = atol(args[i + 1])
+            if args[i] == "-i":
+                prompt = args[i + 1]
+            if args[i] == "-j":
+                workers = atol(args[i + 1])
+            if args[i] == "-t":
+                let val = args[i + 1]
+                temperature = 0.0
+                # hacky parse float, keep only 1 digit
+                for c in range(0, len(val)):
+                    if val[c] == ".":
+                        temperature += atol(val[c + 1]) * 0.1
+                        break
+                    else:
+                        temperature = atol(val[c])
+                if temperature < -1e9 or temperature > (1 + 1e9):
+                    print("Wrong temperature value", temperature)
+                    return 0
+        return 1
+
+    let res = argparse()
+    if res == 0:
+        print_usage()
+        return
+
+    print("num parallel workers:", workers, " SIMD width:", nelts)
     random.seed(rng_seed)
     var fbuf: FileBuf = FileBuf()
     var tbuf: FileBuf = FileBuf()
     var config: Config = Config()
 
     read_file(checkpoint, fbuf)
-    print("checkpoint size: ", fbuf.size)
     config_init(config, fbuf)
 
     # negative vocab size is hacky way of signaling unshared weights. bit yikes.
@@ -690,51 +934,63 @@ fn main() raises:
 
     let weights: TransformerWeights = TransformerWeights(config, shared_weights, fbuf)
 
-    var tok: Tokenizer = Tokenizer(config.vocab_size)
-
     if steps <= 0 or steps > config.seq_len:
         steps = config.seq_len
 
     # Read in the tokenizer.bin file
     read_file(tokenizer, tbuf)
-    tokenizer_init(tok, tbuf)
+    var tok = Tokenizer(config.vocab_size, tbuf)
+
+    # print the layers number and vocab size
+    print("checkpoint size: ", fbuf.size, "[", fbuf.size // 1024 // 1024, "MB ]", 
+        "| n layers:", config.n_layers, "| vocab size:", tok.vocab_size)
 
     # Create and initialize the application RunState
     var state = RunState(config)
 
+    # Process the prompt, if any
+    var prompt_tokens = DynamicVector[Int]()
+
+    if prompt:
+        bpe_encode(prompt_tokens, prompt, tok)
+
     # Start the main loop
     var start = 0  # Used to time our code, only initialized after the first iteration
     var next_token = 0  # Will store the next token in the sequence
     # Initialize with token 1 (=BOS), as done in Llama-2 sentencepiece tokenizer
     var token = 1
-    var pos = 0  # Position in the sequence
-    # Explicitly print the initial BOS token for stylistic symmetry reasons
-
-    print("<s>")
 
+    # Position in the sequence
+    var pos = 0
     while pos < steps:
         # Forward the transformer to get logits for the next token
         transformer(token, pos, config, state, weights)
 
-        # Sample the next token
-        if temperature == 0.0:
-            # Greedy argmax sampling: take the token with the highest probability
-            next_token = argmax(state.logits)
+        if pos < len(prompt_tokens):
+            next_token = prompt_tokens[pos]
         else:
-            # Apply the temperature to the logits
-            for q in range(config.vocab_size):
-                state.logits[q] = state.logits[q] / temperature
-            # Apply softmax to the logits to get the probabilities for the next token
-            softmax(state.logits.data, config.vocab_size)
-            # Sample from this distribution to get the next token
-            next_token = sample(state.logits)
-
+            # Sample the next token
+            if temperature == 0.0:
+                # Greedy argmax sampling: take the token with the highest probability
+                next_token = argmax(state.logits)
+            else:
+                # Apply the temperature to the logits
+                for q in range(config.vocab_size):
+                    state.logits[q] = state.logits[q] / temperature
+
+                # Apply softmax to the logits to get the probabilities for the next token
+                softmax(state.logits)
+                # Sample from this distribution to get the next token
+                next_token = sample(state.logits)
+
+            # Finish generating when EOS, BOS appear
+            if next_token == 1 or next_token == 2:
+                break
         var token_str: PointerString = tok.vocab[next_token]
         if token == 1 and token_str[0] == ord(" "):
             token_str = token_str.offset(1)
 
         print_str(token_str)
-        # flush?
 
         # Advance forward
         token = next_token
@@ -744,4 +1000,4 @@ fn main() raises:
             start = time_in_ms()
 
     let end = time_in_ms()
-    print("\nachieved tok/s: ", (steps - 1) / (end - start) * 1000)
+    print("\nachieved tok/s: ", (pos - 1) / (end - start) * 1000)