Merge pull request #79 from orionpapadakis/opt/q8-load-bytearray

[Opt] Manipulation of Q8_0 tensors with Tornado `ByteArray`s

Author: mikepapadim

pom.xml

@@ -54,12 +54,12 @@
             <dependency>
                 <groupId>io.github.beehive-lab</groupId>
                 <artifactId>tornado-api</artifactId>
-                <version>2.0.0</version>
+                <version>2.0.1-dev</version>
             </dependency>
             <dependency>
                 <groupId>io.github.beehive-lab</groupId>
                 <artifactId>tornado-runtime</artifactId>
-                <version>2.0.0</version>
+                <version>2.0.1-dev</version>
             </dependency>
         </dependencies>
 

src/main/java/org/beehive/gpullama3/inference/InferenceCore.java

@@ -583,7 +583,24 @@ public static FloatArray forwardTornadoVM(Model model, State state, int token, i
         final Configuration configuration = model.configuration();
         final TornadoWeights weights = (TornadoWeights) model.weights();
 
-        MemorySegment.copy(weights.getTokenEmbeddingTable().asFloatArray().getSegment(), (long) token * configuration.dim() * Float.BYTES, state.wrapX.getSegment(), 0, configuration.dim() * Float.BYTES);
+        switch (weights.getWeightType()) {
+            case F16 -> {
+                MemorySegment tokenEmbeddings = weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment();
+                int bytes = Short.BYTES;
+                MemorySegment.copy(tokenEmbeddings, (long) token * configuration.dim() * bytes, state.embeddingX.getSegment(), 0, (long) configuration.dim() * bytes);
+            }
+            case Q8_0 -> {
+                MemorySegment tokenEmbeddings = weights.getTokenEmbeddingTable().asByteArray().getSegment();
+                int blockSize = 32;
+                int Q8_0_BLOCK_BYTES = 34; // 2 bytes scale + 32 bytes quants
+                int blocksPerToken = (configuration.dim() + blockSize - 1) / blockSize; // Ceiling division
+                long bytesPerToken = (long) blocksPerToken * Q8_0_BLOCK_BYTES;
+
+                MemorySegment.copy(tokenEmbeddings, (long) token * bytesPerToken, state.embeddingX.getSegment(), 0, bytesPerToken);
+
+            }
+            default -> throw new IllegalArgumentException("Unsupported weight type: " + weights.getWeightType());
+        }
 
         return tornadoVMMasterPlan.tornadoVMForwardExecuteLayered(position);
     }

src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java

@@ -3,7 +3,9 @@
 import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor;
 import org.beehive.gpullama3.tensor.standard.FloatTensor;
 import org.beehive.gpullama3.model.Configuration;
+import uk.ac.manchester.tornado.api.types.arrays.ByteArray;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
@@ -52,6 +54,11 @@ protected StateFields createStateFields(Configuration config) {
         fields.wrapHb = new FloatArray(config.hiddenDim());
         fields.wrapHb2 = new FloatArray(config.hiddenDim());
 
+        switch (config.quantization()) {
+            case "FP16" -> fields.createActivationFP16(config.dim());
+            case "Q8_0" -> fields.createActivationQ8_0(config.dim());
+            default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization());
+        }
         fields.wrapLogits = new FloatArray(config.vocabularySize());
         fields.wrapQ = new FloatArray(config.dim());
         fields.wrapK = new FloatArray(config.dim());

src/main/java/org/beehive/gpullama3/inference/state/Phi3State.java

@@ -5,6 +5,7 @@
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.model.phi3.Phi3Configuration;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
@@ -79,6 +80,11 @@ protected StateFields createStateFields(Configuration config) {
         fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(contextLength, kvDim)).limit(nLayers).toArray(FloatTensor[]::new);
 
         // TornadoVM wrapper arrays for GPU acceleration
+        switch (config.quantization()) {
+            case "FP16" -> fields.createActivationFP16(config.dim());
+            case "Q8_0" -> fields.createActivationQ8_0(config.dim());
+            default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization());
+        }
         fields.wrapX = new FloatArray(dim);
         fields.wrapXb = new FloatArray(dim);
         fields.wrapXb2 = new FloatArray(dim);

src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java

@@ -5,6 +5,7 @@
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.model.qwen2.Qwen2Configuration;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
@@ -40,6 +41,11 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)).limit(config.numberOfLayers()).toArray(FloatTensor[]::new);
 
         // TornadoVM wrappers with Qwen2 dimensions
+        switch (config.quantization()) {
+            case "FP16" -> fields.createActivationFP16(config.dim());
+            case "Q8_0" -> fields.createActivationQ8_0(config.dim());
+            default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization());
+        }
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(config.dim());
         fields.wrapXb2 = new FloatArray(config.dim());

src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java

@@ -5,6 +5,7 @@
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.model.qwen3.Qwen3Configuration;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
@@ -65,6 +66,13 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)).limit(config.numberOfLayers()).toArray(FloatTensor[]::new);
 
         // TornadoVM wrappers with Qwen3-specific sizes
+
+        switch (config.quantization()) {
+            case "FP16" -> fields.createActivationFP16(config.dim());
+            case "Q8_0" -> fields.createActivationQ8_0(config.dim());
+            default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization());
+        }
+
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(nEmbdHeadK * config.numberOfHeads());
         fields.wrapXb2 = new FloatArray(config.dim());
@@ -74,7 +82,6 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.wrapQ = new FloatArray(nEmbdHeadK * config.numberOfHeads());
         fields.wrapK = new FloatArray(nEmbdKGqa);
         fields.wrapV = new FloatArray(nEmbdKGqa);
-
         fields.wrapKeyCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
         fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
         fields.wrapValueCache.init(0.f);

src/main/java/org/beehive/gpullama3/inference/state/State.java

@@ -2,8 +2,8 @@
 
 import org.beehive.gpullama3.tensor.standard.FloatTensor;
 import org.beehive.gpullama3.model.Configuration;
-import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
-import uk.ac.manchester.tornado.api.types.arrays.IntArray;
+import uk.ac.manchester.tornado.api.types.HalfFloat;
+import uk.ac.manchester.tornado.api.types.arrays.*;
 
 /**
  * Represents the base state structure used during LLM inference.
@@ -57,6 +57,7 @@ public abstract class State {
     public final FloatArray wrapValueCache; // FloatArray wrapper for the value cache, optimized for TornadoVM.
     public final IntArray positionHolder;
 
+    public TornadoNativeArray embeddingX;
     // store inter
     public int localSize;
     public FloatArray temp;         // Temporary buffer for intermediate calculations, size adjusted for local workgroup size.
@@ -88,6 +89,7 @@ protected State(Configuration config, int batchsize) {
         this.keyCache = fields.keyCache;
         this.valueCache = fields.valueCache;
 
+        this.embeddingX = fields.embeddingX;
         this.wrapX = fields.wrapX;
         this.wrapXb = fields.wrapXb;
         this.wrapXb2 = fields.wrapXb2;
@@ -121,6 +123,19 @@ protected static class StateFields {
         public FloatArray wrapQ, wrapK, wrapV, wrapAtt, wrapKeyCache, wrapValueCache;
         public IntArray positionHolder;
         public FloatArray temp, tempFFN, tempLogits;
+        public TornadoNativeArray embeddingX;
+
+        public void createActivationFP16(int size) {
+            this.embeddingX = new HalfFloatArray(size);
+        }
+
+        public void createActivationQ8_0(int size) {
+            int blockSize = 32;
+            int Q8_0_BLOCK_BYTES = 34; // 2 bytes scale + 32 bytes quants
+            int blocksNeeded = (size + blockSize - 1) / blockSize;
+            int q8BytesNeeded = blocksNeeded * Q8_0_BLOCK_BYTES;
+            this.embeddingX = new ByteArray(q8BytesNeeded);
+        }
     }
 
     @Override

src/main/java/org/beehive/gpullama3/model/Configuration.java

@@ -2,6 +2,8 @@
 
 public interface Configuration {
 
+    String quantization();
+
     /** Transformer embedding dimension */
     int dim();
 

src/main/java/org/beehive/gpullama3/model/llama/LlamaConfiguration.java

@@ -3,7 +3,8 @@
 import org.beehive.gpullama3.model.Configuration;
 
 // @formatter:off
-public record LlamaConfiguration(int dim,
+public record LlamaConfiguration(String quantization,
+                                 int dim,
                                  int hiddenDim,
                                  int numberOfLayers,
                                  int numberOfHeads,
@@ -13,6 +14,11 @@ public record LlamaConfiguration(int dim,
                                  float rmsNormEps,
                                  float ropeTheta) implements Configuration {
 
+    @Override
+    public String quantization() {
+        return quantization;
+    }
+
     @Override
     public int numberOfHeadsKey() {
         throw new UnsupportedOperationException("Not supported for Llama.");
@@ -51,6 +57,7 @@ public LlamaConfiguration withContextLength(int newContextLength) {
             return this; // no change
         }
         return new LlamaConfiguration(
+                this.quantization,
                 this.dim,
                 this.hiddenDim,
                 this.numberOfLayers,

src/main/java/org/beehive/gpullama3/model/loader/AbstractModelLoader.java

@@ -35,6 +35,15 @@ protected AbstractModelLoader(FileChannel fileChannel, GGUF gguf, int contextLen
         this.useTornadovm = useTornadovm;
     }
 
+    protected String getModelQuantization(Map<String, Object> metadata) {
+        int modelQuantizationAsInt = (int) metadata.get("general.file_type");
+        return switch (modelQuantizationAsInt) {
+            case 1 -> "FP16";
+            case 7 -> "Q8_0";
+            default -> throw new UnsupportedOperationException("Unsupported quantization format: " + modelQuantizationAsInt + " (as int).");
+        };
+    }
+
     /**
      * Template method that defines the model loading workflow. Subclasses should not override this method.
      *