Refactor TornadoVM Qwen layers and grid scheduler:

- Remove `GenericLayerPlanner` interface and consolidate layer planner logic.
- Standardize `qwen2State` and `qwen3State` usage across layers.
- Adjust local work sizes for better efficiency in FP16 and Q8_0 layers.
- Cleanup redundant code and comments for improved readability.

Author: mikepapadim

src/main/java/org/beehive/gpullama3/tornadovm/layerplanner/GenericLayerPlanner.java

@@ -6,6 +6,7 @@
 import org.beehive.gpullama3.inference.weights.tornado.FP16Weights.FP16Weights;
 import org.beehive.gpullama3.inference.weights.tornado.FP16Weights.LlamaTornadoWeights;
 import org.beehive.gpullama3.model.Configuration;
+import org.beehive.gpullama3.model.Model;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernels;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered;
 import org.beehive.gpullama3.tornadovm.layers.AbstractLayer;
@@ -133,9 +134,12 @@ private GridScheduler setupGridSchedulerForLogits(Configuration config) {
     @Override
     public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) {
             // RMSNorm operations
-            WorkerGrid rmsNormWorker = new WorkerGrid1D(config.dim());
-            rmsNormWorker.setGlobalWork(config.dim(), 1, 1);
-            rmsNormWorker.setLocalWork(256, 1, 1);
+        WorkerGrid rmsNormWorker = new WorkerGrid1D(config.dim());
+
+        rmsNormWorker.setGlobalWork(config.dim(), 1, 1);  // Set global work size to total dimension
+
+        //TODO: XXX
+        rmsNormWorker.setLocalWork(32, 1, 1);         // Set local work size to 256 (standard efficient size)
 
         // OpenCL equivalent: clEnqueueNDRangeKernel(globalWorkSize=[config.vocabularySize,1,1], localWorkSize=[16,1,1])
         // CUDA equivalent: kernel<<<dim3((config.vocabularySize+15)/16,1,1), dim3(16,1,1)>>>

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

@@ -53,10 +53,10 @@ public Qwen2FP16FFNLayers(String taskGraphName, Qwen2State state, Qwen2TornadoWe
         this.qwen2State = state;
         this.qwen2Config = config;
 
-//        state.temp.init(0.0f);
-//        state.tempFFN.init(0.0f);
-//        state.tempLogits.init(0.0f);
-//        state.wrapLogits.init(0.0f);
+//        qwen2State.temp.init(0.0f);
+//        qwen2State.tempFFN.init(0.0f);
+//        qwen2State.tempLogits.init(0.0f);
+//        qwen2State.wrapLogits.init(0.0f);
 
 
         // Ensure we have Qwen2-specific weights
@@ -71,7 +71,6 @@ public Qwen2FP16FFNLayers(String taskGraphName, Qwen2State state, Qwen2TornadoWe
 
     @Override
     public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) {
-
         // Single worker for tasks running with a single thread
         // OpenCL equivalent: clEnqueueNDRangeKernel(globalWorkSize=[1,1,1], localWorkSize=[1,1,1])
         // CUDA equivalent: kernel<<<dim3(1,1,1), dim3(1,1,1)>>>
@@ -209,9 +208,8 @@ private void setupLastID(String taskGraphID) {
     List<ImmutableTaskGraph> setupFFNLayered() {
         List<ImmutableTaskGraph> ffnGraphs = new ArrayList<>();
 
-        state.temp.init(0.0f);
-        qwen2State
-                .tempFFN.init(0.0f);
+        qwen2State.temp.init(0.0f);
+        qwen2State.tempFFN.init(0.0f);
 
 
         for (int layerIndex = 0; layerIndex < qwen2Config.numberOfLayers(); layerIndex++) {
@@ -229,59 +227,39 @@ List<ImmutableTaskGraph> setupFFNLayered() {
      * Setup a single transformer layer for Qwen2 with GQA
      */
     TaskGraph setupSingleQwen2FFNLayer(Qwen2TornadoWeights weights, int layerIndex) {
-     TaskGraph   unifiedLayer = new TaskGraph("layer_" + layerIndex);
+       TaskGraph unifiedLayer = new TaskGraph("layer_" + layerIndex);
         unifiedLayer.consumeFromDevice(state.wrapX);
         unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
                 //Copy-in weights per layer for batched-layered layout
-                weights.rms_att_weightLayered[layerIndex],
-                weights.wqLayered[layerIndex],
-                weights.wkLayered[layerIndex],
-                weights.wvLayered[layerIndex],
-                weights.woLayered[layerIndex],
-                weights.q_biasLayered[layerIndex],
-                weights.k_biasLayered[layerIndex],
-                weights.v_biasLayered[layerIndex],
-                weights.rms_ffn_weightLayered[layerIndex],
-                weights.w1Layered[layerIndex],
-                weights.w2Layered[layerIndex],
-                weights.w3Layered[layerIndex]
-        );
+                weights.rms_att_weightLayered[layerIndex], weights.wqLayered[layerIndex], weights.wkLayered[layerIndex], weights.wvLayered[layerIndex], weights.woLayered[layerIndex],
+                weights.q_biasLayered[layerIndex], weights.k_biasLayered[layerIndex], weights.v_biasLayered[layerIndex], weights.rms_ffn_weightLayered[layerIndex], weights.w1Layered[layerIndex],
+                weights.w2Layered[layerIndex], weights.w3Layered[layerIndex]);
         unifiedLayer = configureLayerDataTransfers(unifiedLayer, layerIndex);
 
-        unifiedLayer.task("reductionsOneBlock" , TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, state.temp,
-                        state.wrapX, config.dim(), config.rmsNormEps(), state.localSize)
-                .task("mapContext", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, state.wrapXb,
-                        state.wrapX, weights.rms_att_weightLayered[layerIndex], state.temp)
-                .task("qmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context,
-                        state.wrapXb,  state.wrapQ, weights.wqLayered[layerIndex], config.dim(), config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .task("kmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context,
-                        state.wrapXb,  state.wrapK, weights.wkLayered[layerIndex], config.dim(), config.kvDim(), LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .task("vmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context,
-                        state.wrapXb,   state.wrapV, weights.wvLayered[layerIndex], config.dim(), config.kvDim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .task("qbias", TransformerComputeKernelsLayered::addInPlace, state.wrapQ, weights.q_biasLayered[layerIndex], config.dim())
-                .task("kbias", TransformerComputeKernelsLayered::addInPlace, state.wrapK, weights.k_biasLayered[layerIndex], config.kvDim())
-                .task("vbias", TransformerComputeKernelsLayered::addInPlace, state.wrapV, weights.v_biasLayered[layerIndex], config.kvDim())
-                .task("rope", Qwen3Kernels::ropeRotation,context, state.positionHolder, state.wrapQ, state.wrapK, config.numberOfKeyValueHeads(),
-                        config.headSize())
-                .task("copyToCaches", TransformerComputeKernelsLayered::copyToCache,
-                        state.wrapKeyCache, state.wrapK,  state.wrapValueCache, state.wrapV, state.positionHolder, config.kvDim(), layerIndex, config.contextLength())
-                .task("parallel-attention", Qwen2Kernels::processHeadsFlashAttention, context,
-                        state.wrapQ, state.wrapKeyCache, state.wrapValueCache, state.wrapXb,
-                        config.numberOfHeads(), config.headSize(), config.kvDim(), config.kvMul(),
-                        state.positionHolder, layerIndex, config.contextLength())
-                .task("matmul1", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context,
-                        state.wrapXb,  state.wrapX, weights.woLayered[layerIndex], config.dim(), config.dim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .task("reductionsOneBlockFFN", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, state.tempFFN,
-                        state.wrapX, config.dim(), config.rmsNormEps(), state.localSize)
-                .task("mapContextFFN", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, state.wrapXb,
-                        state.wrapX, weights.rms_ffn_weightLayered[layerIndex], state.tempFFN)
-                .task("fused_ffn_w1_w3", TransformerComputeKernelsLayered::fusedFeedForwardWithSiLUAndGLUActivation, context,
-                        state.wrapXb,   state.wrapHb, weights.w1Layered[layerIndex], weights.w3Layered[layerIndex], config.dim(), config.hiddenDim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .task("projectionTwo", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context,
-                        state.wrapHb, state.wrapX, weights.w2Layered[layerIndex], config.hiddenDim(), config.dim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
-                .persistOnDevice(
-                        state.wrapX
-                );
+        unifiedLayer.task("reductionsOneBlock", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, qwen2State.temp, qwen2State.wrapX, config.dim(), config.rmsNormEps(), qwen2State.localSize)
+                .task("mapContext", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, qwen2State.wrapXb, qwen2State.wrapX, weights.rms_att_weightLayered[layerIndex], qwen2State.temp)
+                .task("qmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context, qwen2State.wrapXb, qwen2State.wrapQ, weights.wqLayered[layerIndex], config.dim(), config.dim(),
+                        LOCAL_WORK_GROUP_SIZE_ALLOC)
+                .task("kmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context, qwen2State.wrapXb, qwen2State.wrapK, weights.wkLayered[layerIndex], config.dim(), config.kvDim(),
+                        LOCAL_WORK_GROUP_SIZE_ALLOC)
+                .task("vmatmul", TransformerComputeKernelsLayered::matrixVectorGeneric, context, qwen2State.wrapXb, qwen2State.wrapV, weights.wvLayered[layerIndex], config.dim(), config.kvDim(),
+                        LOCAL_WORK_GROUP_SIZE_ALLOC).task("qbias", TransformerComputeKernelsLayered::addInPlace, qwen2State.wrapQ, weights.q_biasLayered[layerIndex], config.dim())
+                .task("kbias", TransformerComputeKernelsLayered::addInPlace, qwen2State.wrapK, weights.k_biasLayered[layerIndex], config.kvDim())
+                .task("vbias", TransformerComputeKernelsLayered::addInPlace, qwen2State.wrapV, weights.v_biasLayered[layerIndex], config.kvDim())
+                .task("rope", Qwen3Kernels::ropeRotation, context, qwen2State.positionHolder, qwen2State.wrapQ, qwen2State.wrapK, config.numberOfKeyValueHeads(), config.headSize())
+                .task("copyToCaches", TransformerComputeKernelsLayered::copyToCache, qwen2State.wrapKeyCache, qwen2State.wrapK, qwen2State.wrapValueCache, qwen2State.wrapV, qwen2State.positionHolder, config.kvDim(),
+                        layerIndex, config.contextLength())
+                .task("parallel-attention", Qwen2Kernels::processHeadsFlashAttention, context, qwen2State.wrapQ, qwen2State.wrapKeyCache, qwen2State.wrapValueCache, qwen2State.wrapXb, config.numberOfHeads(),
+                        config.headSize(), config.kvDim(), config.kvMul(), qwen2State.positionHolder, layerIndex, config.contextLength())
+                .task("matmul1", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context, qwen2State.wrapXb, qwen2State.wrapX, weights.woLayered[layerIndex], config.dim(), config.dim(),
+                        LOCAL_WORK_GROUP_SIZE_ALLOC)
+                .task("reductionsOneBlockFFN", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, qwen2State.tempFFN, qwen2State.wrapX, config.dim(), config.rmsNormEps(), qwen2State.localSize)
+                .task("mapContextFFN", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, qwen2State.wrapXb, qwen2State.wrapX, weights.rms_ffn_weightLayered[layerIndex], qwen2State.tempFFN)
+                .task("fused_ffn_w1_w3", TransformerComputeKernelsLayered::fusedFeedForwardWithSiLUAndGLUActivation, context, qwen2State.wrapXb, qwen2State.wrapHb, weights.w1Layered[layerIndex],
+                        weights.w3Layered[layerIndex], config.dim(), config.hiddenDim(), LOCAL_WORK_GROUP_SIZE_ALLOC)
+                .task("projectionTwo", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context, qwen2State.wrapHb, qwen2State.wrapX, weights.w2Layered[layerIndex], config.hiddenDim(),
+                        config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC).persistOnDevice(state.wrapX);
+
         return unifiedLayer;
     }
 
@@ -292,19 +270,19 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye
         // First layer: Transfer initial data to device (one-time transfer)
         if (layerIndex == 0) {
             // Transfer all attention-related data: query, key, value matrices and their caches
-            unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.positionHolder, state.temp, state.tempFFN); //
+            unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, qwen2State.positionHolder, qwen2State.temp, qwen2State.tempFFN); //
             unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, //
-                    context, state.wrapXb, state.wrapXb2, //
-                    state.wrapQ, state.wrapK, state.wrapV, //
-                    state.wrapKeyCache, state.wrapValueCache, //
-                    state.wrapAtt, state.wrapHb); //
+                    context, qwen2State.wrapXb, qwen2State.wrapXb2, //
+                    qwen2State.wrapQ, qwen2State.wrapK, qwen2State.wrapV, //
+                    qwen2State.wrapKeyCache, qwen2State.wrapValueCache, //
+                    qwen2State.wrapAtt, qwen2State.wrapHb); //
         } else {
             // Subsequent layers: Consume data already on device from previous layer
-            unifiedLayer.consumeFromDevice(context, state.wrapXb, state.wrapXb2, //
-                    state.wrapQ, state.wrapK, state.wrapV, //
-                    state.wrapKeyCache, state.wrapValueCache, //
-                    state.wrapAtt, state.wrapHb, //
-                    state.positionHolder //
+            unifiedLayer.consumeFromDevice(context, qwen2State.wrapXb, qwen2State.wrapXb2, //
+                    qwen2State.wrapQ, qwen2State.wrapK, qwen2State.wrapV, //
+                    qwen2State.wrapKeyCache, qwen2State.wrapValueCache, //
+                    qwen2State.wrapAtt, qwen2State.wrapHb, //
+                    qwen2State.positionHolder //
             );
         }
         return unifiedLayer;

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

@@ -37,7 +37,7 @@ public LogitsQ8_0Layer(String taskGraphName, State state, Weights weights, Confi
     public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) {
         WorkerGrid rmsNormWorker = new WorkerGrid1D(config.dim());
         rmsNormWorker.setGlobalWork(config.dim(), 1, 1);
-        rmsNormWorker.setLocalWork(256, 1, 1);
+        rmsNormWorker.setLocalWork(32, 1, 1);
         // RMSNorm operations
         int vocabSizeRowMajor = config.vocabularySize() * LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS;
         WorkerGrid vocabWorker = new WorkerGrid1D(vocabSizeRowMajor);