Merge pull request #78 from beehive-lab/feat/deq-n-compute

[FP16] Improved performance by fusing dequantize with compute  in kernels: 20-30% Inference Speedup

Author: mikepapadim

llama-tornado

@@ -422,7 +422,7 @@ def create_parser() -> argparse.ArgumentParser:
     )
     debug_group.add_argument(
         "--profiler-dump-dir",
-        default="/home/mikepapadim/repos/gpu-llama3.java/prof.json",
+        default=None,
         help="Directory for profiler output",
     )
 
@@ -498,6 +498,11 @@ def main():
     parser = create_parser()
     args = parser.parse_args()
 
+    # Set default profiler log path relative to LLAMA_ROOT
+    if args.profiler_dump_dir is None:
+        llama_root = os.environ.get("LLAMA_ROOT")
+        args.profiler_dump_dir = os.path.join(llama_root, "profiler-log.json")
+
     # Set default seed if not provided
     if args.seed is None:
         args.seed = int(time.time())

pom.xml

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

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

@@ -64,6 +64,8 @@ protected StateFields createStateFields(Configuration config) {
         fields.wrapK = new FloatArray(config.dim());
         fields.wrapV = new FloatArray(config.dim());
 
+        fields.wrapXFP16 = new HalfFloatArray(config.dim());
+        fields.wrapXbFP16 = new HalfFloatArray(config.dim());
         // dim vs kvdim
         fields.wrapKeyCache = new FloatArray(config.contextLength() * kvDim * config.numberOfLayers());
         fields.wrapValueCache = new FloatArray(config.contextLength() * kvDim * config.numberOfLayers());

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

@@ -87,6 +87,8 @@ protected StateFields createStateFields(Configuration config) {
         }
         fields.wrapX = new FloatArray(dim);
         fields.wrapXb = new FloatArray(dim);
+        fields.wrapXFP16 = new HalfFloatArray(dim);
+        fields.wrapXbFP16 = new HalfFloatArray(dim);
         fields.wrapXb2 = new FloatArray(dim);
         fields.wrapHb = new FloatArray(2 * hiddenDim);
         fields.wrapHb2 = new FloatArray(hiddenDim);

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

@@ -48,6 +48,7 @@ protected StateFields createStateFields(Configuration configuration) {
         }
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(config.dim());
+        fields.wrapXbFP16 = new HalfFloatArray(config.dim());
         fields.wrapXb2 = new FloatArray(config.dim());
         fields.wrapHb = new FloatArray(config.hiddenDim());
         fields.wrapHb2 = new FloatArray(config.hiddenDim());

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

@@ -75,6 +75,8 @@ protected StateFields createStateFields(Configuration configuration) {
 
         fields.wrapX = new FloatArray(config.dim());
         fields.wrapXb = new FloatArray(nEmbdHeadK * config.numberOfHeads());
+        fields.wrapXbFP16 = new HalfFloatArray(nEmbdHeadK * config.numberOfHeads());
+
         fields.wrapXb2 = new FloatArray(config.dim());
         fields.wrapHb = new FloatArray(config.hiddenDim());
         fields.wrapHb2 = new FloatArray(config.hiddenDim());

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

@@ -4,6 +4,9 @@
 import org.beehive.gpullama3.model.Configuration;
 import uk.ac.manchester.tornado.api.types.HalfFloat;
 import uk.ac.manchester.tornado.api.types.arrays.*;
+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;
 
 /**
  * Represents the base state structure used during LLM inference.
@@ -58,13 +61,17 @@ public abstract class State {
     public final IntArray positionHolder;
 
     public TornadoNativeArray embeddingX;
+
+    public final HalfFloatArray wrapXbFP16;         // FloatArray wrapper for xb (residual branch activation), optimized for TornadoVM usage.
+
     // store inter
     public int localSize;
     public FloatArray temp;         // Temporary buffer for intermediate calculations, size adjusted for local workgroup size.
     public FloatArray tempFFN;      // Temporary buffer for feed-forward network calculations, size adjusted for local workgroup size.
     public FloatArray tempLogits;   // Temporary buffer for logits calculations, size adjusted for local workgroup size.
     public int latestToken;         // Keeps track of the most recent token processed by the model. Useful for stateful or autoregressive models.
 
+    public HalfFloatArray wrapXFP16;
     /** last index in previous block */
 
     protected State(Configuration config, int batchsize) {
@@ -100,6 +107,9 @@ protected State(Configuration config, int batchsize) {
         this.wrapK = fields.wrapK;
         this.wrapV = fields.wrapV;
 
+        this.wrapXFP16 = fields.wrapXFP16;
+        this.wrapXbFP16 = fields.wrapXbFP16;
+
         // dim vs kvdim
         this.wrapKeyCache = fields.wrapKeyCache;
         this.wrapValueCache = fields.wrapValueCache;
@@ -136,6 +146,7 @@ public void createActivationQ8_0(int size) {
             int q8BytesNeeded = blocksNeeded * Q8_0_BLOCK_BYTES;
             this.embeddingX = new ByteArray(q8BytesNeeded);
         }
+        public HalfFloatArray wrapXFP16, wrapXbFP16;
     }
 
     @Override

src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlan.java

@@ -1,9 +1,9 @@
 package org.beehive.gpullama3.tornadovm;
 
-import org.beehive.gpullama3.tensor.GGMLType;
 import org.beehive.gpullama3.inference.state.State;
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.model.Model;
+import org.beehive.gpullama3.tensor.GGMLType;
 import org.beehive.gpullama3.tornadovm.layerplanner.base.QuantizationPlannerFactory;
 import uk.ac.manchester.tornado.api.ImmutableTaskGraph;
 import uk.ac.manchester.tornado.api.TornadoExecutionPlan;
@@ -133,6 +133,8 @@ public FloatArray tornadoVMForwardExecuteLayered(int position) {
 
         // Set the position in the state object (used by attention layers)
         state.positionHolder.set(0, position);
+        state.temp.clear();
+        state.tempFFN.clear();
 
         // 2. Execute each transformer layer graph sequentially
         // Each graph computes attention and feed-forward transformations for one layer
@@ -141,7 +143,8 @@ public FloatArray tornadoVMForwardExecuteLayered(int position) {
                     .withGridScheduler(tornadoVMLayerPlanner.getGridScheduler())
                     .execute();
         }
-
+        state.tempLogits.clear(); // Clear the intermediate logits tensor -> set to 0f
+        state.wrapLogits.clear(); // Clear the output logits tensor -> set to 0f
         // 3. Execute the final graph that projects the last hidden state to output logits
         executionPlan.withGraph(getFinalLogitsGraphIndex())
                 .withGridScheduler(tornadoVMLayerPlanner.getGridScheduler())
@@ -179,7 +182,7 @@ private int getFinalLogitsGraphIndex() {
     /// Execute the forward pass of the LLaMA transformer model using TornadoVM acceleration just once to copy the data into the read-only data layer.
     public void forceCopyInReadOnlyDataLayered() {
         // Execute all TornadoVM graphs
-        state.wrapX.init(0.0f);
+        state.wrapX.clear();
         state.positionHolder.init(0);
 
         // Execute activation update graph