Refactor tensor loading and introduce support for Half-Float precision in TornadoVM acceleration.

mikepapadim · mikepapadim · commit d0966eb62ec6 · 2025-11-26T16:05:01.000+02:00
diff --git a/set_paths b/set_paths
@@ -6,10 +6,10 @@
 
 # Resolve root of this project (LLaMA3) and TornadoVM
 export LLAMA_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-export TORNADO_ROOT="${LLAMA_ROOT}/external/tornadovm"
+#export TORNADO_ROOT="${LLAMA_ROOT}/external/tornadovm"
 
 # Set the path to TornadoVM SDK binaries
-export TORNADO_SDK="${TORNADO_ROOT}/bin/sdk"
+#export TORNADO_SDK="${TORNADO_ROOT}/bin/sdk"
 
 # Add TornadoVM and LLaMA bin directories to PATH
 export PATH="${PATH}:${TORNADO_SDK}:${LLAMA_ROOT}"
diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java
@@ -583,7 +583,7 @@ 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);
+        MemorySegment.copy(weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(), (long) token * configuration.dim() * Short.BYTES, state.embeddingX.getSegment(), 0, configuration.dim() * Short.BYTES);
 
         return tornadoVMMasterPlan.tornadoVMForwardExecuteLayered(position);
     }
diff --git a/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java b/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java
@@ -4,6 +4,7 @@
 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.HalfFloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
@@ -52,6 +53,8 @@ protected StateFields createStateFields(Configuration config) {
         fields.wrapHb = new FloatArray(config.hiddenDim());
         fields.wrapHb2 = new FloatArray(config.hiddenDim());
 
+        fields.embeddingX = new HalfFloatArray(config.dim());
+
         fields.wrapLogits = new FloatArray(config.vocabularySize());
         fields.wrapQ = new FloatArray(config.dim());
         fields.wrapK = new FloatArray(config.dim());
diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Phi3State.java b/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,7 @@ protected StateFields createStateFields(Configuration config) {
         fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(contextLength, kvDim)).limit(nLayers).toArray(FloatTensor[]::new);
 
         // TornadoVM wrapper arrays for GPU acceleration
+        fields.embeddingX = new HalfFloatArray(config.dim());
         fields.wrapX = new FloatArray(dim);
         fields.wrapXb = new FloatArray(dim);
         fields.wrapXb2 = new FloatArray(dim);
diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java b/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,7 @@ 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
+        fields.embeddingX = new HalfFloatArray(config.dim());
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(config.dim());
         fields.wrapXb2 = new FloatArray(config.dim());
diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java b/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,8 @@ 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
+
+        fields.embeddingX = new HalfFloatArray(config.dim());
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(nEmbdHeadK * config.numberOfHeads());
         fields.wrapXb2 = new FloatArray(config.dim());
@@ -74,7 +77,7 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.wrapQ = new FloatArray(nEmbdHeadK * config.numberOfHeads());
         fields.wrapK = new FloatArray(nEmbdKGqa);
         fields.wrapV = new FloatArray(nEmbdKGqa);
-
+        fields.embeddingX = new HalfFloatArray(config.dim());
         fields.wrapKeyCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
         fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
         fields.wrapValueCache.init(0.f);
diff --git a/src/main/java/org/beehive/gpullama3/inference/state/State.java b/src/main/java/org/beehive/gpullama3/inference/state/State.java
@@ -2,7 +2,9 @@
 
 import org.beehive.gpullama3.tensor.standard.FloatTensor;
 import org.beehive.gpullama3.model.Configuration;
+import uk.ac.manchester.tornado.api.types.HalfFloat;
 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;
 
 /**
@@ -57,6 +59,7 @@ public abstract class State {
     public final FloatArray wrapValueCache; // FloatArray wrapper for the value cache, optimized for TornadoVM.
     public final IntArray positionHolder;
 
+    public HalfFloatArray embeddingX;
     // store inter
     public int localSize;
     public FloatArray temp;         // Temporary buffer for intermediate calculations, size adjusted for local workgroup size.
@@ -88,6 +91,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 +125,7 @@ protected static class StateFields {
         public FloatArray wrapQ, wrapK, wrapV, wrapAtt, wrapKeyCache, wrapValueCache;
         public IntArray positionHolder;
         public FloatArray temp, tempFFN, tempLogits;
+        public HalfFloatArray embeddingX;
     }
 
     @Override
diff --git a/src/main/java/org/beehive/gpullama3/model/loader/LlamaModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/LlamaModelLoader.java
@@ -120,7 +120,7 @@ protected Weights createTornadoVMWeights(Map<String, GGMLTensorEntry> tensorEntr
 
         // Load all tensors uniformly as TornadoTensor hierarchy
         return new LlamaTornadoWeights(
-                loadTornadoTensorAsFP32(tokenEmbeddings),
+                loadTornadoTensor(tokenEmbeddings),
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_norm.weight")),    // fp32
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.weight")),
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.weight")),
diff --git a/src/main/java/org/beehive/gpullama3/model/loader/MistralModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/MistralModelLoader.java
@@ -130,7 +130,7 @@ protected Weights createTornadoVMWeights(Map<String, GGMLTensorEntry> tensorEntr
 
         // Load all tensors uniformly as TornadoTensor hierarchy
         return new LlamaTornadoWeights(
-                loadTornadoTensorAsFP32(tokenEmbeddings),
+                loadTornadoTensor(tokenEmbeddings),
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_norm.weight")),    // fp32
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.weight")),
                 loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.weight")),
diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java
@@ -2,7 +2,9 @@
 
 import uk.ac.manchester.tornado.api.KernelContext;
 import uk.ac.manchester.tornado.api.math.TornadoMath;
+import uk.ac.manchester.tornado.api.types.HalfFloat;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 
 public class TransformerComputeKernels {
 
@@ -19,6 +21,18 @@ public static void emptyTaskToForceCopyIn(FloatArray buffer) {
         }
     }
 
+    public static void convertFP16toFP32(KernelContext context, HalfFloatArray x, FloatArray wrapX) {
+        int i = context.globalIdx;
+        wrapX.set(i, x.get(i).getFloat32());
+    }
+
+    public static void convertFP32toFP16(KernelContext context,  FloatArray wrapX, HalfFloatArray x) {
+        int i = context.globalIdx;
+        float valInput = wrapX.get(i);
+        HalfFloat val = new HalfFloat(valInput);
+        x.set(i,val);
+    }
+
     /**
      * Performs RMS (Root Mean Square) normalization using parallel reduction.
      * This is a two-phase reduction: first within work groups, then across work groups.
diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/Activation.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/Activation.java
@@ -7,8 +7,10 @@
 import org.beehive.gpullama3.tornadovm.layerplanner.WorkerGridFactory;
 import uk.ac.manchester.tornado.api.GridScheduler;
 import uk.ac.manchester.tornado.api.ImmutableTaskGraph;
+import uk.ac.manchester.tornado.api.KernelContext;
 import uk.ac.manchester.tornado.api.TaskGraph;
 import uk.ac.manchester.tornado.api.WorkerGrid;
+import uk.ac.manchester.tornado.api.WorkerGrid1D;
 import uk.ac.manchester.tornado.api.enums.DataTransferMode;
 
 public class Activation extends AbstractLayer {
@@ -17,16 +19,20 @@ public class Activation extends AbstractLayer {
     public Activation(String taskGraphHandle, State state, Weights weights, Configuration config) {
         super(taskGraphHandle, state, weights, config);
 
-        // formatter:off
-        this.activationUpdate = new TaskGraph(taskGraphHandle).transferToDevice(DataTransferMode.EVERY_EXECUTION, state.wrapX)
-                .task("updateX", TransformerComputeKernels::emptyTaskToForceCopyIn, state.wrapX).persistOnDevice(state.wrapX);
-        // formatter:on
+        KernelContext kernelContext = new KernelContext();
+        // @formatter:off
+        this.activationUpdate = new TaskGraph(taskGraphHandle)
+            .transferToDevice(DataTransferMode.EVERY_EXECUTION, state.embeddingX)
+            .task("updateX", TransformerComputeKernels::convertFP16toFP32, kernelContext, state.embeddingX, state.wrapX)
+            .persistOnDevice(state.wrapX);
+        // @formatter:on
     }
 
     @Override
     public GridScheduler updateGridScheduler(GridScheduler scheduler) {
-        WorkerGrid singleWorker = WorkerGridFactory.createSingleWorker();
-        scheduler.addWorkerGrid("activationUpdate.updateX", singleWorker);
+        WorkerGrid worker = new WorkerGrid1D(config.dim());
+        worker.setLocalWork(128, 1, 1);
+        scheduler.addWorkerGrid("activationUpdate.updateX", worker);
         return scheduler;
     }
 

Original file line number	Diff line number	Diff line change
`@@ -583,7 +583,7 @@ public static FloatArray forwardTornadoVM(Model model, State state, int token, i`
`583`	`583`	`final Configuration configuration = model.configuration();`
`584`	`584`	`final TornadoWeights weights = (TornadoWeights) model.weights();`
`585`	`585`
`586`		`- MemorySegment.copy(weights.getTokenEmbeddingTable().asFloatArray().getSegment(), (long) token * configuration.dim() * Float.BYTES, state.wrapX.getSegment(), 0, configuration.dim() * Float.BYTES);`
	`586`	`+ MemorySegment.copy(weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(), (long) token * configuration.dim() * Short.BYTES, state.embeddingX.getSegment(), 0, configuration.dim() * Short.BYTES);`
`587`	`587`
`588`	`588`	`return tornadoVMMasterPlan.tornadoVMForwardExecuteLayered(position);`
`589`	`589`	`}`