Refactor logits task graph to optimize kernel setup, update worker gr…

…ids, and deprecate redundant tasks in FP16 layer.

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

@@ -2,8 +2,8 @@
 
 import org.beehive.gpullama3.inference.state.State;
 import org.beehive.gpullama3.inference.weights.Weights;
-import org.beehive.gpullama3.inference.weights.tornado.TornadoWeights;
 import org.beehive.gpullama3.inference.weights.tornado.Qwen2TornadoWeights;
+import org.beehive.gpullama3.inference.weights.tornado.TornadoWeights;
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernels;
 import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered;
@@ -28,7 +28,7 @@ public class LogitsFP16Layer extends AbstractLayer {
     public LogitsFP16Layer(String name, State state, Weights weights, Configuration config, String lastTaskGraphID, SchedulerType schedulerType) {
         super(name, state, weights, config);
         this.lastTaskGraphID = lastTaskGraphID;
-        state.tempLogits.init(0.0f);
+        state.tempLogits.clear();
 
         var tornadoWeights = requireWeightsType(weights, TornadoWeights.class, "LogitsFP16Layer", "TornadoTensor");
         this.logitsTaskGraph = setupLogitsTaskGraph(tornadoWeights, config);
@@ -40,18 +40,20 @@ public LogitsFP16Layer(String name, State state, Weights weights, Configuration
      */
     private TaskGraph setupLogitsTaskGraph(TornadoWeights weights, Configuration config) {
         TaskGraph logits = new TaskGraph("logits");
-        logits.consumeFromDevice(lastTaskGraphID, state.wrapX).transferToDevice(DataTransferMode.EVERY_EXECUTION, state.tempLogits, state.wrapXFP16)
-                .transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.wrapLogits, weights.wclsByteArray.asHalfFloatArray(), weights.rms_final_weight_as_floatArray.asFloatArray())
-                .task("reductionsOneBlockLogits", TransformerComputeKernels::reductionOneBlockWithLayer, context, state.tempLogits, state.wrapX, config.dim(), config.rmsNormEps(), state.localSize);
-                if (schedulerType == SchedulerType.NON_NVIDIA) {
-                    logits.task("reductionFinalNormalizationLogits", TransformerComputeKernelsLayered::reductionFinalNormalization, context, state.tempLogits, config.dim(), config.rmsNormEps());
-                }
-                logits.task("mapContextLogits", TransformerComputeKernels::reductionOneBlock2WithLogits, context, state.wrapX, weights.rms_final_weight_as_floatArray.asFloatArray(), state.tempLogits)
-                .task("dequantizeX", TransformerComputeKernels::convertFP32toFP16v2, context, state.wrapX, state.wrapXFP16)
-                .task("projection", TransformerComputeKernelsLayered::matrixVectorGeneric, //
-                        context, state.wrapXFP16, state.wrapLogits,  //
-                        weights.wclsByteArray.asHalfFloatArray(), config.dim(), config.vocabularySize(), //
-                        LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS); //
+        logits.consumeFromDevice(lastTaskGraphID, state.wrapX) //
+                .transferToDevice(DataTransferMode.FIRST_EXECUTION, context,  //
+                        state.wrapLogits, state.wrapXbFP16,  //
+                        weights.wclsByteArray.asHalfFloatArray(),             //
+                        weights.rms_final_weight_as_floatArray.asFloatArray())  //
+                .task("rms_reduce", TransformerComputeKernels::reductionOneBlockWithLayer, context, state.tempLogits, state.wrapX, config.dim(), config.rmsNormEps(), state.localSize);
+        if (schedulerType == SchedulerType.NON_NVIDIA) {
+            logits.task("rms_finalize", TransformerComputeKernelsLayered::reductionFinalNormalization, context, state.tempLogits, config.dim(), config.rmsNormEps());
+        }
+        logits.task("rms_apply_fp16", TransformerComputeKernels::mapContextWithQuantizeLogits, context, state.wrapXbFP16, state.wrapX, weights.rms_final_weight_as_floatArray.asFloatArray(), state.tempLogits)
+                .task("vocab_proj", TransformerComputeKernelsLayered::matrixVectorGeneric, //
+                context, state.wrapXbFP16, state.wrapLogits,  //
+                weights.wclsByteArray.asHalfFloatArray(), config.dim(), config.vocabularySize(), //
+                LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS); //
         logits.transferToHost(DataTransferMode.EVERY_EXECUTION, state.wrapLogits);
         return logits;
     }
@@ -69,10 +71,9 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
         WorkerGrid vocabWorker = new WorkerGrid1D(vocabSizeRowMajor);
         vocabWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS, 1, 1);
 
-        tornadoForwardScheduler.addWorkerGrid("logits.dequantizeX", logitsRMS);
-        tornadoForwardScheduler.addWorkerGrid("logits.projection", vocabWorker);
-        tornadoForwardScheduler.addWorkerGrid("logits.reductionsOneBlockLogits", logitsRMS);
-        tornadoForwardScheduler.addWorkerGrid("logits.mapContextLogits", logitsRMS);
+        tornadoForwardScheduler.addWorkerGrid("logits.vocab_proj", vocabWorker);
+        tornadoForwardScheduler.addWorkerGrid("logits.rms_reduce", logitsRMS);
+        tornadoForwardScheduler.addWorkerGrid("logits.rms_apply_fp16", logitsRMS);
         return tornadoForwardScheduler;
     }