Add disk or memory kv buffers

tjake · Jun 16, 2024 · 35ffb08 · 35ffb08
1 parent ad9183a
commit 35ffb08
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/math/VectorMath.java b/jlama-core/src/main/java/com/github/tjake/jlama/math/VectorMath.java
@@ -54,7 +54,7 @@ public static void pchunk(int offset, int length, BiIntConsumer action) {
         PhysicalCoreExecutor.instance.get().execute(() -> IntStream.range(0, fsplits)
                 .parallel()
                 .forEach(i -> action.accept(
-                        offset + (i * fchunkSize), fremainder > 0 ? fchunkSize + fremainder : fchunkSize)));
+                        offset + (i * fchunkSize), fremainder > 0 && i == fsplits - 1 ? fchunkSize + fremainder : fchunkSize)));
     }
 
     public static void softMax(AbstractTensor x, int offset, int length) {

diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java b/jlama-core/src/main/java/com/github/tjake/jlama/tensor/KvBufferCache.java
@@ -68,8 +68,13 @@ private Pair<RandomAccessFile, AbstractTensor> makeKvBuffer(UUID session) {
         } else {
             s = TensorShape.of(rawShape);
         }
-        Preconditions.checkArgument(model.getConfig().workingDirectory().isPresent());
 
+        // If we don't have a working directory, just use a FloatBufferTensor
+        if (model.getConfig().workingDirectory().isEmpty()) {
+            return Pair.create(null, new FloatBufferTensor(s));
+        }
+
+        // Otherwise, create a file-backed tensor
         try {
             RandomAccessFile raf = new RandomAccessFile(
                     Paths.get(model.getConfig().workingDirectory().get().toString(), session.toString())

diff --git a/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java b/jlama-tests/src/test/java/com/github/tjake/jlama/model/TestModels.java
@@ -133,7 +133,7 @@ public void MistralRun() throws Exception {
             Config c = om.readValue(new File(modelPrefix + "/config.json"), LlamaConfig.class);
             LlamaModel model = new LlamaModel(c, weights, tokenizer, DType.F32, DType.I8, Optional.empty());
             String prompt = "Simply put, the theory of relativity states that";
-            model.generate(UUID.randomUUID(), prompt, 0.7f, 256, false, makeOutHandler());
+            model.generate(UUID.randomUUID(), prompt, 0.7f, 64, false, makeOutHandler());
         }
     }