Additional cleanup

Author: orionpapadakis

src/main/java/com/example/inference/state/Qwen3State.java

@@ -11,19 +11,16 @@
 
 public final class Qwen3State extends State {
 
-    // Qwen3-specific field
-    public final FloatTensor kq;
-
-    // Qwen3 temporary buffer for intermediate calculations, size adjusted for local workgroup size.
+    // Qwen3 specific fields
+    // Temporary buffers for intermediate calculations.
     public FloatArray tempQcur;
     public FloatArray tempKcur;
 
     public Qwen3State(Configuration config, int batchsize) {
         super(config, batchsize);
-        // Initialize Qwen3-specific field
+        // Initialize Qwen3-specific fields
         Qwen3Configuration qwen3config = (Qwen3Configuration) config;
         int nEmbdHead = qwen3config.numberOfHeads();
-        this.kq = ArrayFloatTensor.allocate(config.numberOfHeads(), 32, 15);
         this.tempQcur = new FloatArray(nEmbdHead);
         this.tempKcur = new FloatArray(nEmbdHead);
     }
@@ -34,9 +31,7 @@ protected StateFields createStateFields(Configuration configuration) {
 
         Qwen3Configuration config = (Qwen3Configuration) configuration;
 
-        //localSize = 128;
-
-        // Qwen3-specific calculations
+        // Qwen3-specific sizes
         int nHeadKv = config.numberOfKeyValueHeads();
         int nEmbdHeadK = config.numberOfHeadsKey();
         int nEmbdKGqa = nEmbdHeadK * nHeadKv;
@@ -51,8 +46,8 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.hb = ArrayFloatTensor.allocate(config.hiddenDim());
         fields.hb2 = ArrayFloatTensor.allocate(config.hiddenDim());
         fields.q = ArrayFloatTensor.allocate(nEmbdHeadK * config.numberOfHeads());
-        fields.k = ArrayFloatTensor.allocate(nEmbdKGqa);  // Different from Llama!
-        fields.v = ArrayFloatTensor.allocate(nEmbdKGqa);  // Different from Llama!
+        fields.k = ArrayFloatTensor.allocate(nEmbdKGqa);
+        fields.v = ArrayFloatTensor.allocate(nEmbdKGqa);
         fields.att = ArrayFloatTensor.allocate(config.numberOfHeads(), config.contextLength());
         fields.logits = ArrayFloatTensor.allocate(config.vocabularySize());
 
@@ -64,14 +59,14 @@ protected StateFields createStateFields(Configuration configuration) {
 
         // TornadoVM wrappers with Qwen3-specific sizes
         fields.wrapX = new FloatArray(config.dim());
-        fields.wrapXb = new FloatArray(nEmbdHeadK * config.numberOfHeads());  // Different from Llama!
+        fields.wrapXb = new FloatArray(nEmbdHeadK * config.numberOfHeads());
         fields.wrapXb2 = new FloatArray(config.dim());
         fields.wrapHb = new FloatArray(config.hiddenDim());
         fields.wrapHb2 = new FloatArray(config.hiddenDim());
         fields.wrapLogits = new FloatArray(config.vocabularySize());
-        fields.wrapQ = new FloatArray(nEmbdHeadK * config.numberOfHeads());   // Different from Llama!
-        fields.wrapK = new FloatArray(nEmbdKGqa);  // Different from Llama!
-        fields.wrapV = new FloatArray(nEmbdKGqa);  // Different from Llama!
+        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());

src/main/java/com/example/tornadovm/TornadoVMMasterPlan.java

@@ -96,7 +96,6 @@ public static TornadoVMMasterPlan initializeTornadoVMPlan(State state, Model mod
      * Dispatcher method to select the TornadoVMLayerPlanner for the model.
      */
     TornadoVMLayerPlanner createPlanner(State state, Model model) {
-        System.out.println("Creating TornadoVM layer planner : " + model.getModelType() );
         return switch (model.getModelType()) {
             case LLAMA_3, MISTRAL -> new TornadoVMLayerPlanner(state, model);
             case QWEN_3 -> new Qwen3TornadoVMLayerPlanner((Qwen3State) state, model);