[
  {
    "path": ".gitignore",
    "content": "# Compiled class file\n*.class\n\n# Log file\n*.log\n\n# BlueJ files\n*.ctxt\n\n# Mobile Tools for Java (J2ME)\n.mtj.tmp/\n\n# Package Files #\n*.jar\n*.war\n*.nar\n*.ear\n*.zip\n*.tar.gz\n*.rar\n\n# virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml\nhs_err_pid*\nreplay_pid*\n"
  },
  {
    "path": "LICENSE",
    "content": "MIT License\n\nCopyright (c) 2023 Alfonso² Peterssen\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "Llama2.java",
    "content": "///usr/bin/env jbang \"$0\" \"$@\" ; exit $?\n//JAVA 21\n//COMPILE_OPTIONS --enable-preview -source 21 --add-modules=jdk.incubator.vector\n//RUNTIME_OPTIONS --enable-preview --add-modules=jdk.incubator.vector\n//NATIVE_OPTIONS  --enable-preview --add-modules=jdk.incubator.vector --initialize-at-build-time=Llama2 -Dllama2.VectorAPI=false\n\n/* Inference for Llama-2 Transformer model in pure Java */\n\n// ----------------------------------------------------------------------------\n// Transformer model\n\nimport jdk.incubator.vector.FloatVector;\nimport jdk.incubator.vector.VectorOperators;\nimport jdk.incubator.vector.VectorSpecies;\n\nimport java.io.BufferedInputStream;\nimport java.io.IOException;\nimport java.io.InputStream;\nimport java.lang.foreign.MemorySegment;\nimport java.lang.foreign.Arena;\nimport java.nio.ByteBuffer;\nimport java.nio.ByteOrder;\nimport java.nio.FloatBuffer;\nimport java.nio.channels.FileChannel;\nimport java.nio.charset.StandardCharsets;\nimport java.nio.file.Paths;\nimport java.nio.file.StandardOpenOption;\nimport java.util.*;\nimport java.util.stream.IntStream;\n\nfinal class Config {\n    final int dim; // transformer dimension\n    final int hidden_dim; // for ffn layers\n    final int n_layers; // number of layers\n    final int n_heads; // number of query heads\n    final int n_kv_heads; // number of key/value heads (can be < query heads because of multiquery)\n    final int vocab_size; // vocabulary size, usually 256 (byte-level)\n    final int seq_len; // max sequence length\n    final boolean shared_weights;\n    final int head_size;\n\n    Config(ByteBuffer buffer) {\n        this.dim = buffer.getInt();\n        this.hidden_dim = buffer.getInt();\n        this.n_layers = buffer.getInt();\n        this.n_heads = buffer.getInt();\n        this.n_kv_heads = buffer.getInt();\n        int vocab_size = buffer.getInt();\n        this.vocab_size = Math.abs(vocab_size);\n        this.seq_len = buffer.getInt();\n        this.shared_weights = vocab_size > 0;\n        this.head_size = dim / n_heads;\n    }\n\n    @Override\n    public String toString() {\n        return \"Config{\" +\n                \"dim=\" + dim +\n                \", hidden_dim=\" + hidden_dim +\n                \", n_layers=\" + n_layers +\n                \", n_heads=\" + n_heads +\n                \", n_kv_heads=\" + n_kv_heads +\n                \", vocab_size=\" + vocab_size +\n                \", seq_len=\" + seq_len +\n                \", shared_weights=\" + shared_weights +\n                \", head_size=\" + head_size +\n                '}';\n    }\n}\n\nfinal class Weights {\n    // token embedding table\n    final FloatBuffer token_embedding_table; // (vocab_size, dim)\n    // weights for rmsnorms\n    final FloatBuffer[] rms_att_weight; // (layer, dim) rmsnorm weights\n    // weights for matmuls. note dim == n_heads * head_size\n    final FloatBuffer[] wq; // (layer, dim, n_heads * head_size)\n    final FloatBuffer[] wk; // (layer, dim, n_kv_heads * head_size)\n    final FloatBuffer[] wv; // (layer, dim, n_kv_heads * head_size)\n    final FloatBuffer[] wo; // (layer, n_heads * head_size, dim)\n    final FloatBuffer[] rms_ffn_weight; // (layer, dim)\n    // weights for ffn\n    final FloatBuffer[] w1; // (layer, hidden_dim, dim)\n    final FloatBuffer[] w2; // (layer, dim, hidden_dim)\n    final FloatBuffer[] w3; // (layer, hidden_dim, dim)\n    // final rmsnorm\n    final FloatBuffer rms_final_weight; // (dim,)\n    // (optional) classifier weights for the logits, on the last layer\n    final FloatBuffer wcls; // (vocab_size, dim)\n\n    static FloatBuffer takeFloats(MemorySegment memorySegment, long[] position, int... dims) {\n        long totalBytes = 1;\n        for (int d : dims) {\n            totalBytes *= d;\n        }\n        totalBytes *= Float.BYTES;\n        MemorySegment slice = memorySegment.asSlice(position[0], totalBytes);\n        position[0] += totalBytes;\n        return slice.asByteBuffer().order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer();\n    }\n\n    static FloatBuffer[] takeArray(MemorySegment memorySegment, long[] position, int dim0, int... dims) {\n        FloatBuffer[] segments = new FloatBuffer[dim0];\n        for (int i = 0; i < dim0; ++i) {\n            segments[i] = takeFloats(memorySegment, position, dims);\n        }\n        return segments;\n    }\n\n// ----------------------------------------------------------------------------\n// initialization: read from checkpoint\n\n    Weights(Config config, MemorySegment memorySegment) {\n        long[] position = new long[]{0};\n        this.token_embedding_table = takeFloats(memorySegment, position, config.vocab_size, config.dim);\n        this.rms_att_weight = takeArray(memorySegment, position, config.n_layers, config.dim);\n        this.wq = takeArray(memorySegment, position, config.n_layers, config.dim, config.n_heads * config.head_size);\n        this.wk = takeArray(memorySegment, position, config.n_layers, config.dim, config.n_kv_heads * config.head_size);\n        this.wv = takeArray(memorySegment, position, config.n_layers, config.dim, config.n_kv_heads * config.head_size);\n        this.wo = takeArray(memorySegment, position, config.n_layers, config.n_heads * config.head_size, config.dim);\n        this.rms_ffn_weight = takeArray(memorySegment, position, config.n_layers, config.dim);\n        this.w1 = takeArray(memorySegment, position, config.n_layers, config.hidden_dim, config.dim);\n        this.w2 = takeArray(memorySegment, position, config.n_layers, config.dim, config.hidden_dim);\n        this.w3 = takeArray(memorySegment, position, config.n_layers, config.hidden_dim, config.dim);\n        this.rms_final_weight = takeFloats(memorySegment, position, config.dim);\n        position[0] += (config.seq_len * config.head_size / 2) * Float.BYTES; // skip what used to be freq_cis_real (for RoPE)\n        position[0] += (config.seq_len * config.head_size / 2) * Float.BYTES; // skip what used to be freq_cis_imag (for RoPE)\n        this.wcls = config.shared_weights\n                ? this.token_embedding_table\n                : takeFloats(memorySegment, position, config.vocab_size, config.dim);\n    }\n}\n\nfinal class RunState {\n    // current wave of activations\n    final float[] x; // activation at current time stamp (dim,)\n    final float[] xb; // same, but inside a residual branch (dim,)\n    final float[] xb2; // an additional buffer just for convenience (dim,)\n    final float[] hb; // buffer for hidden dimension in the ffn (hidden_dim,)\n    final float[] hb2; // buffer for hidden dimension in the ffn (hidden_dim,)\n    final float[] q; // query (dim,)\n    final float[] k; // key (dim,)\n    final float[] v; // value (dim,)\n    final float[] att; // buffer for scores/attention values (n_heads, seq_len)\n    final float[] logits; // output logits\n    // kv cache\n    final float[][] key_cache;   // (layer, seq_len, dim)\n    final float[][] value_cache; // (layer, seq_len, dim)\n\n    RunState(Config config) {\n        int kv_dim = (config.dim * config.n_kv_heads) / config.n_heads;\n        this.x = new float[config.dim];\n        this.xb = new float[config.dim];\n        this.xb2 = new float[config.dim];\n        this.hb = new float[config.hidden_dim];\n        this.hb2 = new float[config.hidden_dim];\n        this.q = new float[config.dim];\n        this.k = new float[kv_dim];\n        this.v = new float[kv_dim];\n        this.att = new float[config.n_heads * config.seq_len];\n        this.logits = new float[config.vocab_size];\n        this.key_cache = new float[config.n_layers][config.seq_len * kv_dim];\n        this.value_cache = new float[config.n_layers][config.seq_len * kv_dim];\n    }\n}\n\nfinal class Transformer {\n    final Config config; // the hyperparameters of the architecture (the blueprint)\n    final Weights weights; // the weights of the model\n    final RunState state; // buffers for the \"wave\" of activations in the forward pass\n    // some more state needed to properly clean up the memory mapping (sigh)\n    final Arena memoryArena; // scope of the memory mapping\n    final MemorySegment data; // memory mapped data pointer\n    final long file_size; // size of the checkpoint file in bytes\n\n    Transformer(String checkpoint_path) throws IOException {\n        try (FileChannel fileChannel = FileChannel.open(Paths.get(checkpoint_path), StandardOpenOption.READ)) {\n            this.file_size = fileChannel.size();\n            this.memoryArena = Arena.ofAuto();\n            MemorySegment mappedFile = fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, this.file_size, this.memoryArena);\n            this.data = mappedFile;\n            int configSize = 7 * Integer.BYTES;\n            // read in the config header\n            ByteBuffer configBuffer = mappedFile.asSlice(0, configSize).asByteBuffer().order(ByteOrder.LITTLE_ENDIAN);\n            this.config = new Config(configBuffer);\n            System.out.println(config);\n            this.state = new RunState(config);\n            this.weights = new Weights(config, mappedFile.asSlice(configSize));\n        }\n    }\n}\n\nfinal class Tokenizer {\n    final String[] vocab;\n    final float[] vocab_scores;\n    final int vocab_size;\n    final int max_token_length;\n    Map<String, Integer> sorted_vocab;\n\n    Tokenizer(String tokenizer_path, int vocab_size) throws IOException {\n        // i should have written the vocab_size into the tokenizer file... sigh\n        this.vocab_size = vocab_size;\n        // malloc space to hold the scores and the strings\n        this.vocab = new String[vocab_size];\n        this.vocab_scores = new float[vocab_size];\n\n        // read in the file\n        try (FileChannel channel = FileChannel.open(Paths.get(tokenizer_path), StandardOpenOption.READ)) {\n            ByteBuffer tokBuffer = channel.map(FileChannel.MapMode.READ_ONLY, 0, channel.size());\n            tokBuffer.order(ByteOrder.LITTLE_ENDIAN);\n            this.max_token_length = tokBuffer.getInt();\n            for (int i = 0; i < vocab_size; i++) {\n                this.vocab_scores[i] = tokBuffer.getFloat();\n                int len = tokBuffer.getInt();\n                byte[] bytes = new byte[len];\n                tokBuffer.get(bytes);\n                this.vocab[i] = new String(bytes, StandardCharsets.UTF_8);\n            }\n        }\n    }\n}\n\nfinal class Sampler {\n    final int vocab_size;\n    final int[] probindex; // buffer used in top-p sampling\n    final float temperature;\n    final float topp;\n    long rng_seed;\n\n    Sampler(int vocab_size, float temperature, float topp, long rng_seed) {\n        this.vocab_size = vocab_size;\n        this.temperature = temperature;\n        this.topp = topp;\n        this.rng_seed = rng_seed;\n        // buffer only used with nucleus sampling; may not need but it's ~small\n        this.probindex = new int[vocab_size];\n    }\n\n    int random_u32() {\n        // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A\n        rng_seed ^= rng_seed >> 12;\n        rng_seed ^= rng_seed << 25;\n        rng_seed ^= rng_seed >> 27;\n        return (int) ((rng_seed * 0x2545F4914F6CDD1DL) >> 32);\n    }\n\n    float random_f32() { // random float32 in [0,1)\n        return (random_u32() >>> 8) / 16777216.0f;\n    }\n}\n\nclass Llama2 {\n\n// ----------------------------------------------------------------------------\n// neural net blocks; the dynamics of the Transformer\n\n    static void rmsnorm(float[] o, float[] x, FloatBuffer weight, int size) {\n        // calculate sum of squares\n        float ss = 0.0f;\n        for (int j = 0; j < size; j++) {\n            ss += x[j] * x[j];\n        }\n        ss /= size;\n        ss += 1e-5f;\n        ss = 1.0f / (float) Math.sqrt(ss);\n        // normalize and scale\n        for (int j = 0; j < size; j++) {\n            o[j] = weight.get(j) * (ss * x[j]);\n        }\n    }\n\n    static void softmax(float[] x, int xOffset, int size) {\n        // find max value (for numerical stability)\n        float max_val = x[0 + xOffset];\n        for (int i = 1; i < size; i++) {\n            if (x[i + xOffset] > max_val) {\n                max_val = x[i + xOffset];\n            }\n        }\n        // exp and sum\n        float sum = 0.0f;\n        for (int i = 0; i < size; i++) {\n            x[i + xOffset] = (float) Math.exp(x[i + xOffset] - max_val);\n            sum += x[i + xOffset];\n        }\n        // normalize\n        for (int i = 0; i < size; i++) {\n            x[i + xOffset] /= sum;\n        }\n    }\n\n    static final boolean USE_VECTOR_API = \"true\".equalsIgnoreCase(System.getProperty(\"llama2.VectorAPI\", \"true\"));\n\n    static void matmul(float[] xout, float[] x, FloatBuffer w, int n, int d) {\n        // W (d,n) @ x (n,) -> xout (d,)\n        // by far the most amount of time is spent inside this little function\n        MemorySegment wSegment = MemorySegment.ofBuffer(w);\n        IntStream.range(0, d).parallel().forEach(i -> {\n            float val = 0f;\n            int j = 0;\n            if (USE_VECTOR_API) {\n                VectorSpecies<Float> species = FloatVector.SPECIES_256;\n                FloatVector sum0 = FloatVector.zero(species);\n                FloatVector sum1 = FloatVector.zero(species);\n                FloatVector sum2 = FloatVector.zero(species);\n                FloatVector sum3 = FloatVector.zero(species);\n                int width = species.length();\n                int upperBound = n - n % (4 * width);\n                for (; j < upperBound; j += 4 * width) {\n                    var wj0 = FloatVector.fromMemorySegment(species, wSegment, (i * n + j + 0 * width) * Float.BYTES, ByteOrder.LITTLE_ENDIAN);\n                    var wj1 = FloatVector.fromMemorySegment(species, wSegment, (i * n + j + 1 * width) * Float.BYTES, ByteOrder.LITTLE_ENDIAN);\n                    var wj2 = FloatVector.fromMemorySegment(species, wSegment, (i * n + j + 2 * width) * Float.BYTES, ByteOrder.LITTLE_ENDIAN);\n                    var wj3 = FloatVector.fromMemorySegment(species, wSegment, (i * n + j + 3 * width) * Float.BYTES, ByteOrder.LITTLE_ENDIAN);\n                    var xj0 = FloatVector.fromArray(species, x, j + 0 * width);\n                    var xj1 = FloatVector.fromArray(species, x, j + 1 * width);\n                    var xj2 = FloatVector.fromArray(species, x, j + 2 * width);\n                    var xj3 = FloatVector.fromArray(species, x, j + 3 * width);\n                    sum0 = wj0.fma(xj0, sum0);\n                    sum1 = wj1.fma(xj1, sum1);\n                    sum2 = wj2.fma(xj2, sum2);\n                    sum3 = wj3.fma(xj3, sum3);\n                }\n                val = sum0.add(sum1).add(sum2).add(sum3).reduceLanes(VectorOperators.ADD);\n            }\n\n            // Graal's auto-vectorization.\n            int upperBound = n & ~3;\n            float[] sum = new float[4];\n            for (; j < upperBound; j += sum.length) {\n                sum[0] += w.get(i * n + j + 0) * x[j + 0];\n                sum[1] += w.get(i * n + j + 1) * x[j + 1];\n                sum[2] += w.get(i * n + j + 2) * x[j + 2];\n                sum[3] += w.get(i * n + j + 3) * x[j + 3];\n            }\n            val += sum[0] + sum[1] + sum[2] + sum[3];\n\n            for (; j < n; j++) {\n                val += w.get(i * n + j) * x[j];\n            }\n            xout[i] = val;\n        });\n    }\n\n    static float[] forward(Transformer transformer, int token, int pos) {\n        // a few convenience variables\n        Config p = transformer.config;\n        Weights w = transformer.weights;\n        RunState s = transformer.state;\n        int dim = p.dim;\n        int hidden_dim = p.hidden_dim;\n        int head_size = p.head_size;\n        int kv_dim = (p.dim * p.n_kv_heads) / p.n_heads;\n        int kv_mul = p.n_heads / p.n_kv_heads; // integer multiplier of the kv sharing in multiquery\n\n        // copy the token embedding into x\n        w.token_embedding_table.get(token * dim, s.x, 0, dim);\n\n        // forward all the layers\n        for (int l = 0; l < p.n_layers; l++) {\n\n            // attention rmsnorm\n            rmsnorm(s.xb, s.x, w.rms_att_weight[l], dim);\n\n            // qkv matmuls for this position\n            matmul(s.q, s.xb, w.wq[l], dim, dim);\n            matmul(s.k, s.xb, w.wk[l], dim, kv_dim);\n            matmul(s.v, s.xb, w.wv[l], dim, kv_dim);\n\n            // RoPE relative positional encoding: complex-valued rotate q and k in each head\n            for (int i = 0; i < dim; i+=2) {\n                int head_dim = i % head_size;\n                float freq = (float) (1.0 / Math.pow(10000.0f, head_dim / (float) head_size));\n                float val = pos * freq;\n                float fcr = (float) Math.cos(val);\n                float fci = (float) Math.sin(val);\n                int rotn = i < kv_dim ? 2 : 1; // how many vectors? 2 = q & k, 1 = q only\n                for (int v = 0; v < rotn; v++) {\n                    float[] vec = v == 0 ? s.q : s.k; // the vector to rotate (query or key)\n                    float v0 = vec[i];\n                    float v1 = vec[i + 1];\n                    vec[i] = v0 * fcr - v1 * fci;\n                    vec[i + 1] = v0 * fci + v1 * fcr;\n                }\n            }\n\n            // save key,value at this time step (pos) to our kv cache\n            //int loff = l * p.seq_len * kv_dim; // kv cache layer offset for convenience\n            System.arraycopy(s.k, 0, s.key_cache[l], pos * kv_dim, kv_dim);\n            System.arraycopy(s.v, 0, s.value_cache[l], pos * kv_dim, kv_dim);\n\n\n            final int curLayer = l;\n\n            // multihead attention. iterate over all heads\n            IntStream.range(0, p.n_heads).parallel().forEach(h -> {\n                // get the query vector for this head\n                // float* q = s.q + h * head_size;\n                int qOffset = h * head_size;\n\n                // attention scores for this head\n                // float* att = s.att + h * p.seq_len;\n                int attOffset = h * p.seq_len;\n\n                // iterate over all timesteps, including the current one\n                for (int t = 0; t <= pos; t++) {\n                    // get the key vector for this head and at this timestep\n                    // float* k = s->key_cache + loff + t * kv_dim + (h / kv_mul) * head_size;\n                    int keyCacheOffset = t * kv_dim + (h / kv_mul) * head_size;\n                    // calculate the attention score as the dot product of q and k\n                    float score = 0.0f;\n                    for (int i = 0; i < head_size; i++) {\n                        score += s.q[qOffset + i] * s.key_cache[curLayer][keyCacheOffset + i];\n                    }\n                    score /= (float) Math.sqrt(head_size);\n                    // save the score to the attention buffer\n                    s.att[attOffset + t] = score;\n                }\n\n                // softmax the scores to get attention weights, from 0..pos inclusively\n                softmax(s.att, attOffset, pos + 1);\n\n                // weighted sum of the values, store back into xb\n                // float* xb = s.xb + h * head_size;\n                int xbOffset = h * head_size;\n                // memset(xb, 0, head_size * sizeof(float));\n                Arrays.fill(s.xb, xbOffset, xbOffset + head_size, 0f);\n\n                for (int t = 0; t <= pos; t++) {\n                    // get the value vector for this head and at this timestep\n                    // float* v = s->value_cache + loff + t * kv_dim + (h / kv_mul) * head_size;\n                    int vOffset = t * kv_dim + (h / kv_mul) * head_size;\n                    // get the attention weight for this timestep\n                    float a = s.att[attOffset + t];\n                    // accumulate the weighted value inconfigto xb\n                    for (int i = 0; i < head_size; i++) {\n                        s.xb[xbOffset + i] += a * s.value_cache[curLayer][vOffset + i];\n                    }\n                }\n            });\n\n            // final matmul to get the output of the attention\n            matmul(s.xb2, s.xb, w.wo[l], dim, dim);\n\n            // residual connection back into x\n            for (int i = 0; i < dim; i++) {\n                s.x[i] += s.xb2[i];\n            }\n\n            // ffn rmsnorm\n            rmsnorm(s.xb, s.x, w.rms_ffn_weight[l], dim);\n\n            // Now for FFN in PyTorch we have: self.w2(F.silu(self.w1(x)) * self.w3(x))\n            // first calculate self.w1(x) and self.w3(x)\n            matmul(s.hb, s.xb, w.w1[l], dim, p.hidden_dim);\n            matmul(s.hb2, s.xb, w.w3[l], dim, p.hidden_dim);\n\n            // SwiGLU non-linearity\n            for (int i = 0; i < hidden_dim; i++) {\n                float val = s.hb[i];\n                // silu(x)=x*σ(x), where σ(x) is the logistic sigmoid\n                val *= (1.0f / (1.0f + Math.exp(-val)));\n                // elementwise multiply with w3(x)\n                s.hb[i] = val;\n            }\n\n            // elementwise multiply with w3(x)\n            for (int i = 0; i < hidden_dim; i++) {\n                s.hb[i] = s.hb[i] * s.hb2[i];\n            }\n\n            // final matmul to get the output of the ffn\n            matmul(s.xb, s.hb, w.w2[l], p.hidden_dim, dim);\n\n            // residual connection\n            for (int i = 0; i < dim; i++) {\n                s.x[i] += s.xb[i];\n            }\n        }\n\n        // final rmsnorm\n        rmsnorm(s.x, s.x, w.rms_final_weight, dim);\n\n        // classifier into logits\n        matmul(s.logits, s.x, w.wcls, dim, p.vocab_size);\n        return s.logits;\n    }\n\n// ----------------------------------------------------------------------------\n\n// The Byte Pair Encoding (BPE) Tokenizer that translates strings <-> tokens\n\n    static String decode(Tokenizer t, int prev_token, int token) {\n        String piece = t.vocab[token];\n        // following BOS (1) token, sentencepiece decoder strips any leading whitespace (see PR #89)\n        if (prev_token == 1 && piece.charAt(0) == ' ') {\n            piece = piece.substring(1);\n        }\n        // careful, some tokens designate raw bytes, and look like e.g. '<0x01>'\n        String prefix = \"<0x\";\n        String suffix = \">\";\n        if (piece.length() == 6 && piece.startsWith(prefix) && piece.endsWith(suffix)) {\n            String hex2 = piece.substring(prefix.length(), prefix.length() + 2);\n            char ch = (char) Integer.parseInt(hex2, 16);\n            // ok this token is a raw byte token, carefuly to only print printable chars or whitespace\n            // some of the other bytes can be various control codes, backspace, etc. => skip\n            piece = Character.toString(ch);\n        }\n        return piece;\n    }\n\n    static void safe_printf(String piece) {\n        // piece might be a raw byte token, and we only want to print printable chars or whitespace\n        // because some of the other bytes can be various control codes, backspace, etc.\n        if (piece == null) { return; }\n        if (piece.isEmpty()) { return; }\n        if (piece.length() == 1) {\n            char ch = piece.charAt(0);\n            boolean isPrintable = (32 <= ch && ch < 127);\n            if (!(isPrintable || Character.isWhitespace(ch))) {\n                return ;\n            }\n        }\n        System.out.print(piece);\n    }\n\n    static int str_lookup(String str, Map<String, Integer> sorted_vocab) {\n        // efficiently find the perfect match for str in vocab, return its index or -1 if not found\n        return sorted_vocab.getOrDefault(str, -1);\n    }\n\n    static int encode(Tokenizer t, String text, boolean bos, boolean eos, int[] tokens) {\n        // encode the string text (input) into an upper-bound preallocated tokens[] array\n        // bos != 0 means prepend the BOS token (=1), eos != 0 means append the EOS token (=2)\n        if (text == null) {\n            System.err.println(\"cannot encode NULL text\");\n            System.exit(1);\n        }\n\n        if (t.sorted_vocab == null) {\n            // sort vocabulary\n            t.sorted_vocab = new HashMap<>();\n            for (int i = 0; i < t.vocab_size; i++) {\n                assert !t.sorted_vocab.containsKey(t.vocab[i]);\n                t.sorted_vocab.put(t.vocab[i], i);\n            }\n        }\n\n        // start at 0 tokens\n        int n_tokens = 0; // the number of tokens\n\n        // add optional BOS (=1) token, if desired\n        if (bos) {\n            tokens[n_tokens++] = 1;\n        }\n\n        // so prepend a dummy prefix token to the input string, but only if text != \"\"\n        // TODO: pretty sure this isn't correct in the general case but I don't have the\n        // energy to read more of the sentencepiece code to figure out what it's doing\n        if (!\"\".equals(text)) {\n            int dummy_prefix = str_lookup(\" \", t.sorted_vocab);\n            tokens[n_tokens++] = dummy_prefix;\n        }\n\n        // first encode every individual codepoint in the input string\n        for (int i = 0, cpi; i < text.length(); i += Character.charCount(cpi)) {\n            cpi = text.codePointAt(i);\n\n            String singleCodepoint = Character.toString(cpi);\n            int id = str_lookup(singleCodepoint, t.sorted_vocab);\n\n            if (id != -1) {\n                // we found this codepoint in vocab, add it as a token\n                tokens[n_tokens++] = id;\n            } else {\n                // byte_fallback encoding: just encode each byte as a token\n                // +3 is here because the first 3 vocab elements are <unk>, <s>, </s>\n                // so the individual bytes only start at index 3\n                for (byte b : singleCodepoint.getBytes(StandardCharsets.UTF_8)) {\n                    tokens[n_tokens++] = Byte.toUnsignedInt(b) + 3;\n                }\n            }\n        }\n\n        // merge the best consecutive pair each iteration, according the scores in vocab_scores\n        while (true) {\n            float best_score = -1e10f;\n            int best_id = -1;\n            int best_idx = -1;\n\n            for (int i = 0; i < n_tokens - 1; ++i) {\n                // check if we can merge the pair (tokens[i], tokens[i+1])\n                String str_buffer = t.vocab[tokens[i]] + t.vocab[tokens[i + 1]];\n                int id = str_lookup(str_buffer, t.sorted_vocab);\n                if (id != -1 && t.vocab_scores[id] > best_score) {\n                    // this merge pair exists in vocab! record its score and position\n                    best_score = t.vocab_scores[id];\n                    best_id = id;\n                    best_idx = i;\n                }\n            }\n\n            if (best_idx == -1) {\n                break; // we couldn't find any more pairs to merge, so we're done\n            }\n\n            // merge the consecutive pair (best_idx, best_idx+1) into new token best_id\n            tokens[best_idx] = best_id;\n            // delete token at position best_idx+1, shift the entire sequence back 1\n            for (int i = best_idx + 1; i < n_tokens - 1; i++) {\n                tokens[i] = tokens[i + 1];\n            }\n            n_tokens--; // token length decreased\n        }\n\n        // add optional EOS (=2) token, if desired\n        if (eos) {\n            tokens[n_tokens++] = 2;\n        }\n\n        return n_tokens;\n    }\n\n// ----------------------------------------------------------------------------\n// utilities: time / rng\n\n    static long time_in_ms() {\n        // return time in milliseconds, for benchmarking the model speed\n        return System.nanoTime() / 1_000_000;\n    }\n\n\n// ----------------------------------------------------------------------------\n// generation loop\n\n    static void generate(Transformer transformer, Tokenizer tokenizer, Sampler sampler, String prompt, int steps) {\n        String empty_prompt = \"\";\n        if (prompt == null) {\n            prompt = empty_prompt;\n        }\n\n        // encode the (string) prompt into tokens sequence\n        int num_prompt_tokens = 0; // the total number of prompt tokens\n        int[] prompt_tokens = new int[prompt.length() * 2 + 3]; // +3 for '\\0', ?BOS, ?EOS\n        num_prompt_tokens = encode(tokenizer, prompt, true, false, prompt_tokens);\n        if (num_prompt_tokens < 1) {\n            System.err.println(\"something is wrong, expected at least 1 prompt token\");\n            System.exit(1);\n        }\n\n        // start the main loop\n        long start = 0;  // used to time our code, only initialized after first iteration\n        int next;        // will store the next token in the sequence\n        int token = prompt_tokens[0]; // kick off with the first token in the prompt\n        int pos = 0;     // position in the sequence\n        while (pos < steps) {\n            // forward the transformer to get logits for the next token\n            float[] logits = forward(transformer, token, pos);\n\n            // advance the state machine\n            if (pos < num_prompt_tokens - 1) {\n                // if we are still processing the input prompt, force the next prompt token\n                next = prompt_tokens[pos + 1];\n            } else {\n                // otherwise sample the next token from the logits\n                next = sample(sampler, logits);\n            }\n            pos++;\n\n            // data-dependent terminating condition: the BOS (=1) token delimits sequences\n            if (next == 1) {\n                break;\n            }\n\n            // print the token as string, decode it with the Tokenizer object\n            String piece = decode(tokenizer, token, next);\n            safe_printf(piece);\n\n            System.out.flush();\n            token = next;\n\n            // init the timer here because the first iteration can be slower\n            if (start == 0) {\n                start = time_in_ms();\n            }\n        }\n\n        System.out.println();\n\n        // report achieved tok/s (pos-1 because the timer starts after first iteration)\n        if (pos > 1) {\n            long end = time_in_ms();\n            System.err.printf(\"\\nachieved tok/s: %f\\n\", (pos - 1) / (double) (end - start) * 1000);\n        }\n    }\n\n\n// ----------------------------------------------------------------------------\n// sampling can be done in a few ways: greedy argmax, sampling, top-p sampling\n\n    static int sample_argmax(float[] probabilities, int n) {\n        // return the index that has the highest probability\n        int max_i = 0;\n        float max_p = probabilities[0];\n        for (int i = 1; i < n; i++) {\n            if (probabilities[i] > max_p) {\n                max_i = i;\n                max_p = probabilities[i];\n            }\n        }\n        return max_i;\n    }\n\n    static int sample_mult(float[] probabilities, int n, float coin) {\n        // sample index from probabilities (they must sum to 1!)\n        float cdf = 0.0f;\n        for (int i = 0; i < n; i++) {\n            cdf += probabilities[i];\n            if (coin < cdf) {\n                return i;\n            }\n        }\n        return n - 1; // in case of rounding errors\n    }\n\n    static void swap(int[] array, int from, int to) {\n        int tmp = array[from];\n        array[from] = array[to];\n        array[to] = tmp;\n    }\n\n    static void siftDown(int[] array, int from, int n, Comparator<Integer> comparator) {\n        int prev = from, next;\n        while ((next = 2 * prev + 1) < n) {\n            int r = 2 * prev + 2;\n            if (r < n && comparator.compare(array[r], array[next]) < 0) {\n                next = r;\n            }\n            if (comparator.compare(array[next], array[prev]) < 0) {\n                swap(array, prev, next);\n                prev = next;\n            } else {\n                break;\n            }\n        }\n    }\n\n    static int sample_topp(float[] probabilities, int n, float topp, int[] indices, float coin) {\n        // top-p sampling (or \"nucleus sampling\") samples from the smallest set of\n        // tokens that exceed probability topp. This way we never sample tokens that\n        // have very low probabilities and are less likely to go \"off the rails\".\n        // coin is a random number in [0, 1), usually from random_f32()\n        Comparator<Integer> comparator = Comparator.<Integer>comparingDouble(i -> probabilities[i]).reversed();\n\n        int head = 0;\n        int tail = n - 1;\n        // values smaller than (1 - topp) / (n - 1) cannot be part of the result\n        // so for efficiency we crop these out as candidates before sorting\n        float cutoff = (1.0f - topp) / (n - 1);\n        for (int i = 0; i < indices.length; i++) {\n            if (probabilities[i] >= cutoff) {\n                indices[head++] = i;\n            } else {\n                indices[tail--] = i;\n            }\n        }\n\n        int n0 = head;\n        // build heap O(n0)\n        for (int i = n0 / 2 - 1; i >= 0; --i) {\n            siftDown(indices, i, n0, comparator);\n        }\n\n        // truncate the list where cumulative probability of the largest k elements exceeds topp\n        // O(k lg n0)\n        float cumulative_prob = 0.0f;\n        int last_idx = 0;\n        for (int i = n0 - 1; i >= 0; i--) {\n            swap(indices, 0, i);\n            cumulative_prob += probabilities[indices[i]];\n            if (cumulative_prob > topp) {\n                last_idx = i;\n                break; // we've exceeded topp by including last_idx\n            }\n            siftDown(indices, 0, i - 1, comparator);\n        }\n\n        // sample from the truncated list\n        float r = coin * cumulative_prob;\n        float cdf = 0.0f;\n        for (int i = n0 - 1; i >= last_idx; i--) {\n            cdf += probabilities[indices[i]];\n            if (r < cdf) {\n                return indices[i];\n            }\n        }\n\n        return indices[last_idx]; // in case of rounding errors\n    }\n\n    static int sample(Sampler sampler, float[] logits) {\n        // sample the token given the logits and some hyperparameters\n        int next;\n        if (sampler.temperature == 0.0f) {\n            // greedy argmax sampling: take the token with the highest probability\n            next = sample_argmax(logits, sampler.vocab_size);\n        } else {\n            // apply the temperature to the logits\n            for (int q = 0; q < sampler.vocab_size; q++) {\n                logits[q] /= sampler.temperature;\n            }\n            // apply softmax to the logits to get the probabilities for next token\n            softmax(logits, 0, sampler.vocab_size);\n            // flip a (float) coin (this is our source of entropy for sampling)\n            float coin = sampler.random_f32();\n            // we sample from this distribution to get the next token\n            if (sampler.topp <= 0 || sampler.topp >= 1) {\n                // simply sample from the predicted probability distribution\n                next = sample_mult(logits, sampler.vocab_size, coin);\n            } else {\n                // top-p (nucleus) sampling, clamping the least likely tokens to zero\n                next = sample_topp(logits, sampler.vocab_size, sampler.topp, sampler.probindex, coin);\n            }\n        }\n        return next;\n    }\n\n    static String read_stdin(String guide) {\n        // read a line from stdin, up to but not including \\n\n        System.out.print(guide);\n        Scanner scanner = new Scanner(System.in);\n        if (scanner.hasNextLine()) {\n            return scanner.nextLine();\n        }\n        return null;\n    }\n\n// ----------------------------------------------------------------------------\n// chat loop\n// I manually inspected the tokens for a few chat conversations compared to\n// python reference and that seemed ok, but this was not thoroughly tested and\n// is not safely implemented, it's more a proof of concept atm.\n\n    static void chat(Transformer transformer, Tokenizer tokenizer, Sampler sampler,\n              String cli_user_prompt, String cli_system_prompt, int steps) {\n\n        // buffers for reading the system prompt and user prompt from stdin\n        String system_prompt = null;\n        String user_prompt = null;\n        String rendered_prompt = null;\n        int num_prompt_tokens = 0;\n        int[] prompt_tokens = new int[512];\n        int user_idx = 0;\n\n        // start the main loop\n        boolean user_turn = true; // user starts\n        int next = 0;        // will store the next token in the sequence\n        int token = 0;       // stores the current token to feed into the transformer\n        int prev_token;\n        int pos = 0;     // position in the sequence\n        while (pos < steps) {\n\n            // when it is the user's turn to contribute tokens to the dialog...\n            if (user_turn) {\n                // get the (optional) system prompt at position 0\n                if (pos == 0) {\n                    // at position 0, the user can also contribute a system prompt\n                    if (cli_system_prompt == null) {\n                        // system prompt was not passed in, attempt to get it from stdin\n                        system_prompt = read_stdin(\"Enter system prompt (optional): \");\n                    } else {\n                        // system prompt was passed in, use it\n                        system_prompt = cli_system_prompt;\n                    }\n                }\n                // get the user prompt\n                if (pos == 0 && cli_user_prompt != null) {\n                    // user prompt for position 0 was passed in, use it\n                    user_prompt = cli_user_prompt;\n                } else {\n                    // otherwise get user prompt from stdin\n                    user_prompt = read_stdin(\"User: \");\n                }\n                // render user/system prompts into the Llama 2 Chat schema\n                if (pos == 0 && system_prompt.isEmpty()) {\n                    String system_template = \"[INST] <<SYS>>\\n%s\\n<</SYS>>\\n\\n%s [/INST]\";\n                    rendered_prompt = system_template.formatted(system_prompt, user_prompt);\n                } else {\n                    String user_template = \"[INST] %s [/INST]\";\n                    rendered_prompt = user_template.formatted(user_prompt);\n                }\n                // encode the rendered prompt into tokens\n                num_prompt_tokens = encode(tokenizer, rendered_prompt, true, false, prompt_tokens);\n                user_idx = 0; // reset the user index\n                user_turn = false;\n                System.out.print(\"Assistant: \");\n            }\n\n            // determine the token to pass into the transformer next\n            if (user_idx < num_prompt_tokens) {\n                // if we are still processing the input prompt, force the next prompt token\n                token = prompt_tokens[user_idx++];\n            } else {\n                // otherwise use the next token sampled from previous turn\n                token = next;\n            }\n            // EOS (=2) token ends the Assistant turn\n            if (token == 2) {\n                user_turn = true;\n            }\n\n            // forward the transformer to get logits for the next token\n            float[] logits = forward(transformer, token, pos);\n            next = sample(sampler, logits);\n            pos++;\n\n            if (user_idx >= num_prompt_tokens && next != 2) {\n                // the Assistant is responding, so print its output\n                String piece = decode(tokenizer, token, next);\n                safe_printf(piece); // same as printf(\"%s\", piece), but skips \"unsafe\" bytes\n                System.out.flush();\n            }\n            if (next == 2) {\n                System.out.println();\n            }\n        }\n        System.out.println();\n    }\n\n\n// ----------------------------------------------------------------------------\n// int main\n\n    static void error_usage() {\n        System.err.println(\"Usage:   java Llama2 <checkpoint> [options]\");\n        System.err.println(\"Example: java Lamma2 model.bin -n 256 -i \\\"Once upon a time\\\"\");\n        System.err.println(\"Options:\");\n        System.err.println(\"  -t <float>  temperature in [0,inf], default 1.0\");\n        System.err.println(\"  -p <float>  p value in top-p (nucleus) sampling in [0,1] default 0.9\");\n        System.err.println(\"  -s <int>    random seed, default time(NULL)\");\n        System.err.println(\"  -n <int>    number of steps to run for, default 256. 0 = max_seq_len\");\n        System.err.println(\"  -i <string> input prompt\");\n        System.err.println(\"  -z <string> optional path to custom tokenizer\");\n        System.err.println(\"  -m <string> mode: generate|chat, default: generate\");\n        System.err.println(\"  -y <string> (optional) system prompt in chat mode\");\n        System.exit(1);\n    }\n\n    public static void main(String[] args) throws IOException {\n        // default parameters\n        String checkpoint_path = null; // e.g. out/model.bin\n        String tokenizer_path = \"tokenizer.bin\";\n        float temperature = 1.0f;    // 0.0 = greedy deterministic. 1.0 = original. don't set higher\n        float topp = 0.9f;           // top-p in nucleus sampling. 1.0 = off. 0.9 works well, but slower\n        long rng_seed = 0;           // seed rng with time by default\n        int steps = 256;             // max number of steps to run for, 0: use seq_len\n        String prompt = null;        // prompt string\n        String mode = \"generate\";    // generate|chat\n        String system_prompt = null; // the (optional) system prompt to use in chat mode\n\n        // poor man's C argparse so we can override the defaults above from the command line\n        if (args.length >= 1) {\n            checkpoint_path = args[0];\n        } else {\n            error_usage();\n        }\n        for (int i = 1; i < args.length; i += 2) {\n            // do some basic validation\n            if (i + 1 >= args.length) { error_usage(); } // must have arg after flag\n            if (args[i].charAt(0) != '-') { error_usage(); } // must start with dash\n            if (args[i].length() != 2) { error_usage(); } // must be -x (one dash, one letter)\n            // read in the args\n            switch (args[i].charAt(1)) {\n                case 't' -> temperature = Float.parseFloat(args[i + 1]);\n                case 'p' -> topp = Float.parseFloat(args[i + 1]);\n                case 's' -> rng_seed = Integer.parseInt(args[i + 1]);\n                case 'n' -> steps = Integer.parseInt(args[i + 1]);\n                case 'i' -> prompt = args[i + 1];\n                case 'z' -> tokenizer_path = args[i + 1];\n                case 'm' -> mode = args[i + 1];\n                case 'y' -> system_prompt = args[i + 1];\n                default -> error_usage();\n            }\n        }\n\n        // parameter validation/overrides\n        if (rng_seed <= 0) {\n            rng_seed = System.currentTimeMillis();\n        }\n        if (temperature < 0.0) {\n            temperature = 0.0f;\n        }\n        if (topp < 0.0 || 1.0 < topp) {\n            topp = 0.9f;\n        }\n        if (steps <= 0) {\n            steps = 0;\n        }\n\n        // build the Transformer via the model .bin file\n        Transformer transformer = new Transformer(checkpoint_path);\n        if (steps == 0 || steps > transformer.config.seq_len) {\n            steps = transformer.config.seq_len; // ovrerride to ~max length\n        }\n\n        // build the Tokenizer via the tokenizer .bin file\n        Tokenizer tokenizer = new Tokenizer(tokenizer_path, transformer.config.vocab_size);\n\n        // build the Sampler\n        Sampler sampler = new Sampler(transformer.config.vocab_size, temperature, topp, rng_seed);\n\n        // run!\n        switch (mode) {\n            case \"generate\" -> generate(transformer, tokenizer, sampler, prompt, steps);\n            case \"chat\" -> chat(transformer, tokenizer, sampler, prompt, system_prompt, steps);\n            default -> {\n                System.err.println(\"unknown mode: \" + mode);\n                error_usage();\n            }\n        }\n    }\n}\n"
  },
  {
    "path": "Makefile",
    "content": "ifdef JAVA_HOME\n\tJAVAC ?= ${JAVA_HOME}/bin/javac\n\tJAVA ?= ${JAVA_HOME}/bin/java\n\tJAR ?= ${JAVA_HOME}/bin/jar\n\tNATIVE_IMAGE ?= ${JAVA_HOME}/bin/native-image\nendif\n\nJAVAC ?= javac\nJAVA ?= java\nJAR ?= jar\nNATIVE_IMAGE ?= native-image\n\nJAVA_COMPILE_OPTIONS = --enable-preview -source 21 -g --add-modules jdk.incubator.vector\nJAVA_RUNTIME_OPTIONS += --enable-preview --add-modules jdk.incubator.vector\nNATIVE_IMAGE_OPTIONS += --enable-preview --add-modules jdk.incubator.vector\n\nJAVA_MAIN_CLASS = Llama2\nJAR_FILE = llama2.jar\n\nJAVA_SOURCES = $(wildcard *.java)\nJAVA_CLASSES = $(patsubst %.java, target/classes/%.class, $(JAVA_SOURCES))\n\n# Bundle all classes in a jar\n$(JAR_FILE): $(JAVA_CLASSES) target/META-INF/MANIFEST.MF\n\t$(JAR) -cvfm $(JAR_FILE) target/META-INF/MANIFEST.MF -C target/classes .\n\njar: $(JAR_FILE)\n\n# Compile the Java source files\ncompile: $(JAVA_CLASSES)\n\t$(info Java source files: $(JAVA_SOURCES))\n\t$(info Java .class files: $(JAVA_CLASSES))\n\n# Prints the command to run the Java main class\nrun-command:\n\t@echo $(JAVA) $(JAVA_RUNTIME_OPTIONS) -cp target/classes $(JAVA_MAIN_CLASS)\n\n# Prints the command to run the $(JAR_FILE)\nrun-jar-command:\n\t@echo $(JAVA) $(JAVA_RUNTIME_OPTIONS) -jar $(JAR_FILE)\n\n# Clean the target directory\nclean:\n\trm -rf target\n\trm $(JAR_FILE)\n\trm default.iprof\n\trm llama2\n\n# Creates the manifest for the .jar file\ntarget/META-INF/MANIFEST.MF:\n\tmkdir -p target/META-INF\n\t@echo \"Manifest-Version: 1.0\" > target/META-INF/MANIFEST.MF\n\t@echo \"Class-Path: .\" >> target/META-INF/MANIFEST.MF\n\t@echo \"Main-Class: $(JAVA_MAIN_CLASS)\" >> target/META-INF/MANIFEST.MF\n\t@echo \"\" >> target/META-INF/MANIFEST.MF\n\n# Create a standalone executable of the llama2.jar using GraalVM\nnative-image: $(JAR_FILE)\n\t$(NATIVE_IMAGE) $(NATIVE_IMAGE_OPTIONS) -jar $(JAR_FILE)\n\n# Compile the Java source files\ntarget/classes/%.class: %.java\n\t$(JAVAC) $(JAVA_COMPILE_OPTIONS) -d target/classes $<\n\n# Create the target directory\ntarget/classes:\n\tmkdir -p target/classes\n\n# Make the target directory a dependency of the Java class files\n$(JAVA_CLASSES): target/classes\ncompile: target/classes\ndefault: target/classes\n\n.PHONY: compile clean jar run-command run-jar-command\n.SUFFIXES: .java .class .jar .MF\n"
  },
  {
    "path": "README.md",
    "content": "# A Java port of Andrej Karpathy's llama2.c\n\n****Check the successor of this project: [Llama3.java](https://github.com/mukel/llama3.java): Practical Llama (3) inference in a single Java file, with additional features, including a `--chat` mode.**\n\nThis is a pure Java port of Andrej Karpathy's awesome [llama2.c](https://github.com/karpathy/llama2.c), a very simple implementation\nto run inference of models with a [Llama2](https://arxiv.org/pdf/2302.13971.pdf)-like transformer-based LLM architecture.  \n\n<p align=\"center\">\n  <img width=\"600\" src=\"https://github.com/mukel/llama2.java/assets/1896283/66a8a650-f1a9-4540-9587-b112294e5e6b\">\n</p>\n\nCurrently, there isn't anything really original here, but I'll continue polishing it while keeping it in sync with the original.  \nBesides the educational value, this project will be used to test and tune compiler optimizations on the JVM, particularly for the [Graal compiler](https://www.graalvm.org/latest/reference-manual/java/compiler).\nThis port used [llama2.scala](https://github.com/jrudolph/llama2.scala) initially as a reference.\n\n## Build\nJava 21+ is required, in particular the [`MemorySegment` mmap-ing feature](https://docs.oracle.com/en/java/javase/21/docs/api/java.base/java/nio/channels/FileChannel.html#map(java.nio.channels.FileChannel.MapMode,long,long,java.lang.foreign.Arena)).  \n\nThe code expects [`tokenizer.bin`](https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin) in the current directory.\nYou can use [TinyStories](https://huggingface.co/karpathy/tinyllamas/tree/main) checkpoints or get LLama2 models by [following instructions](https://github.com/karpathy/llama2.c#metas-llama-2-models).\n\n```bash\nwget https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin\nwget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin\n```\n\nTo build and run manually:\n```bash\njavac --enable-preview -source 21 --add-modules=jdk.incubator.vector Llama2.java\njava --enable-preview --add-modules=jdk.incubator.vector Llama2 stories15M.bin\n```\n\nOr run it directly with [JBang](https://www.jbang.dev/):\n```bash\njbang Llama2.java stories15M.bin\n# With additional -D options and custom Java home.\nJAVA_HOME=/path/to/java/home jbang -Djava.util.concurrent.ForkJoinPool.common.parallelism=0 -Dllama2.VectorAPI=false Llama2.java stories15M.bin\n```\n\nA `Makefile` and a `run.sh` script are also provided:\n\n```bash\nmake # optional, run.sh already runs make\n\nJAVA_HOME=$GRAALVM_HOME \\\nJAVA_RUNTIME_OPTIONS=-Djava.util.concurrent.ForkJoinPool.common.parallelism=8 \\\n./run.sh stories15M.bin\n```\n\n#### Native image\n\nA standalone native image can be created with [GraalVM](https://www.graalvm.org/)\n```bash\nJAVA_HOME=$GRAALVM_HOME NATIVE_IMAGE_OPTIONS=\"-march=native\" make native-image\n./llama2 stories15M.bin\n```\n\nOr can also be built with [Profile-Guided Optimizations (PGO)](https://www.graalvm.org/dev/reference-manual/native-image/guides/optimize-native-executable-with-pgo), on Oracle GaaalVM:\n```bash\nJAVA_HOME=$GRAALVM_HOME \\\nNATIVE_IMAGE_OPTIONS=\"--pgo-instrument -march=native --initialize-at-build-time=Llama2 -Dllama2.VectorAPI=false\" \\\nmake native-image\n\n# Profile run to generate default.iprof, with no parallelism to speedup profiling.\n./llama2 -Djava.util.concurrent.ForkJoinPool.common.parallelism=0 stories15M.bin\n\n# Build optimized image\nJAVA_HOME=$GRAALVM_HOME \\\nNATIVE_IMAGE_OPTIONS=\"--pgo -march=native --initialize-at-build-time=Llama2 -Dllama2.VectorAPI=false\" \\\nmake native-image\n\n# Should run ~2X faster than regular image.\n./llama2 stories15M.bin\n```\n\n## Performance\n\nQuick numbers on an AMD Ryzen 3950X 64GB, Arch Linux.  \n`llama2.java` executed on OpenJDK 20.0.2+9.  \nTo make things fair w.r.t. to vectorization, the Java version has a matmul implementation using the [Vector API](https://openjdk.org/jeps/448).  \nIn these measurements the JVM is warmed up enough to reach peak tokens/s.  \nOn GraalVM, please note that the Graal compiler doesn't support the Vector API yet, to avoid unexpected performance degradation, run with `-Dllama2.VectorAPI=false`.\n\n****Notes**  \n*The numbers below were collected using aggressive (gcc) compiler flags e.g. regular `gcc -O2 ...` wouldn't be as fast.*\n\n### Single-threaded\n\n`llama2.c` compiled with `gcc -Ofast -march=native run.c -lm -o run -march=native`  \n`llama2.java` executed with `-Djava.util.concurrent.ForkJoinPool.common.parallelism=0`\n\n| Model | Tokens per second | Speedup vs. llama2.c | Implementation |  \n| ------|------------------ | -------------------- | -------------- | \n| stories15M.bin  |   363 |  1.0 | llama2.c    |\n| stories15M.bin  |   237 | 0.65 | llama2.java |\n| stories110M.bin | 51.71 |  1.0 | llama2.c    |\n| stories110M.bin | 42.20 | 0.81 | llama2.java |\n| llama2_7B.bin   |  0.92 |  1.0 | llama2.c    |\n| llama2_7B.bin   |  0.88 | 0.95 | llama2.java |\n\n### Multi-threaded\n\n`llama2.c` compiled with `gcc -Ofast -fopenmp -march=native run.c -lm -o run -march=native`  \n`llama2.c` executed with `OMP_NUM_THREADS=8`  \n`llama2.java` executed with `-Djava.util.concurrent.ForkJoinPool.common.parallelism=8`  \n\n| Model | Tokens per second | Speedup vs. llama2.c | Implementation |  \n| ------|------------------ | -------------------- | -------------- |\n|  stories15M.bin |  1233 |  1.0 | llama2.c    |\n|  stories15M.bin |   438 | 0.35 | llama2.java |\n| stories110M.bin |    90 |  1.0 | llama2.c    |\n| stories110M.bin |    80 | 0.88 | llama2.java |\n|   llama2_7B.bin |  1.68 |  1.0 | llama2.c    |\n|   llama2_7B.bin |  1.65 | 0.98 | llama2.java |\n\n****Notes**  \n*In `stories15M.bin`, the C version shows a huge speedup, very likely a cache effect, this is considered an outlier.\nRunning with 16/32 threads may actually cause a slowdown; the performance is, in most cases, U-shaped w.r.t to the # of threads.\nWith that many threads, vectorization does not give any advantage, since throughput is limited by memory bandwidth.*\n\nPerformance is already comparable to the original C code, bar vectorization, even if the Java code has not been optimized yet.\n\n## License\n\nMIT\n"
  },
  {
    "path": "run.sh",
    "content": "#!/bin/bash\nmake compile\n`make run-command` \"$@\"\n"
  }
]