Merge pull request beehive-lab#64 from orionpapadakis/fix/bug-fixes

Bug fixes in sizes and names of GridScheduler

Author: mikepapadim

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java

@@ -172,7 +172,7 @@ List<ImmutableTaskGraph> setupFFNLayered() {
      * Setup a single transformer layer for Qwen3 with GQA
      */
     TaskGraph setupSingleQwen3FFNLayer(Qwen3TornadoWeights weights, int layerIndex) {
-        var taskGraphName = "ffn_layer_" + layerIndex;
+        var taskGraphName = "layer_" + layerIndex;
         TaskGraph unifiedLayer = new TaskGraph(taskGraphName);
         unifiedLayer.consumeFromDevice(state.wrapX);
         unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, //

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/LogitsQ8_0Layer.java

@@ -3,7 +3,7 @@
 import org.beehive.gpullama3.inference.state.State;
 import org.beehive.gpullama3.inference.weights.Weights;
 import org.beehive.gpullama3.inference.weights.tornado.q8_0.LlamaTornadoWeightsQ8_0;
-import org.beehive.gpullama3.inference.weights.tornado.q8_0.Qwen3TornadoWeightsQ8_0;
+import org.beehive.gpullama3.inference.weights.tornado.q8_0.Qwen2TornadoWeightsQ8_0;
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernels;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered;
@@ -34,7 +34,7 @@ public LogitsQ8_0Layer(String taskGraphName, State state, Weights weights, Confi
     @Override
     public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) {
         WorkerGrid logitsRMS;
-        if (weights instanceof Qwen3TornadoWeightsQ8_0) {
+        if (weights instanceof Qwen2TornadoWeightsQ8_0) {
             logitsRMS = WorkerGridFactory.createRmsNormWorker(config.dim(), 32);
         } else {
             logitsRMS = WorkerGridFactory.createRmsNormWorker(config.dim(), 256);