Provide an optimized rmsnorm kernel that fuses steps 1 and 2

Author: orionpapadakis

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

@@ -119,6 +119,49 @@ public static void rmsnormMapIndexInPlaceWithParallelOffset(
         //        }
     }
 
+    /**
+     * RmsNorm with parallel offset:
+     *
+     * Optimized kernel that combines Step 1 (Reduction) and Step 2 (Normalization).
+     */
+    public static void rmsnormWithParallelOffset(
+            KernelContext context,
+            FloatArray output,
+            FloatArray x,
+            int localMemSize,
+            int size,
+            float ermsNorm) {
+
+        int gid = context.globalIdx;
+        int lid = context.localIdx;
+        int groupId = context.groupIdx;
+        int groupSize = context.localGroupSizeX;
+
+        // Allocate local memory with the provided size
+        float[] localX = context.allocateFloatLocalArray(localMemSize);
+
+        // Load input value and compute square
+        localX[lid] = x.get(gid);
+        localX[lid] = localX[lid] * localX[lid];
+
+        // Perform parallel reduction within the work group
+        for (int stride = (groupSize / 2); stride > 0; stride /= 2) {
+            context.localBarrier();
+            if (lid < stride) {
+                localX[lid] += localX[lid + stride];
+            }
+        }
+
+        // Each workgroup performs the normalization
+        if (lid == 0) {
+            // Store the partial sum from each workgroup
+            localX[0] /= size;
+            localX[0] += ermsNorm;
+            localX[0] = 1.0f / TornadoMath.sqrt(localX[0]);
+            output.set(groupId, localX[0]);
+        }
+    }
+
     public static void reductionOneBlockWithLayerWithOffset(
             KernelContext context,
             FloatArray output,

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

@@ -166,18 +166,20 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
             //rmsnorm(state.q, state.q, weights.attnQNorm[curLayer], i * nEmbdHead, nEmbdHead, config.rmsNormEps());
             unifiedLayer
                     .task("rmsnormReduction_Qcur",
-                            Qwen3Kernels::rmsnormReductionWithOffset,
+                            Qwen3Kernels::rmsnormWithParallelOffset,
                             context,
                             state.tempQcur,         // output
                             state.wrapQ,            // input
-                            state.localSize) // currently 128, should be variable of global nEmbHead
-                    .task("rmsnormFinalNormalization_Qcur",
-                            Qwen3Kernels::rmsnormFinalNormalizationWithParallelOffset,
-                            context,
-                            state.tempQcur,     // output
-                            config.numberOfHeads(),
-                            nEmbdHead,
-                            config.rmsNormEps())
+                            state.localSize,        // currently 128, should be variable of global nEmbHead
+                            nEmbdHead,              // for normalization
+                            config.rmsNormEps())    // for normalization
+//                    .task("rmsnormFinalNormalization_Qcur",
+//                            Qwen3Kernels::rmsnormFinalNormalizationWithParallelOffset,
+//                            context,
+//                            state.tempQcur,     // output
+//                            config.numberOfHeads(),
+//                            nEmbdHead,
+//                            config.rmsNormEps())
                     .task("rmsnormMapIndexInPlace_Qcur",
                             Qwen3Kernels::rmsnormMapIndexInPlaceWithParallelOffset,
                             context,
@@ -192,18 +194,20 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
             //rmsnorm(state.k, state.k, weights.attnKNorm[curLayer], i * nEmbdHead, nEmbdHead, config.rmsNormEps());
             unifiedLayer
                     .task("rmsnormReduction_Kcur",
-                            Qwen3Kernels::rmsnormReductionWithOffset,
+                            Qwen3Kernels::rmsnormWithParallelOffset,
                             context,
                             state.tempKcur,         // output
                             state.wrapK,            // input
-                            state.localSize) // currently 128, should be variable of global nEmbHead
-                    .task("rmsnormFinalNormalization_Kcur",
-                            Qwen3Kernels::rmsnormFinalNormalizationWithParallelOffset,
-                            context,
-                            state.tempKcur,     // output
-                            config.numberOfKeyValueHeads(),
-                            nEmbdHead,
-                            config.rmsNormEps())
+                            state.localSize,        // currently 128, should be variable of global nEmbHead
+                            nEmbdHead,              // for normalization
+                            config.rmsNormEps())    // for normalization
+//                    .task("rmsnormFinalNormalization_Kcur",
+//                            Qwen3Kernels::rmsnormFinalNormalizationWithParallelOffset,
+//                            context,
+//                            state.tempKcur,     // output
+//                            config.numberOfKeyValueHeads(),
+//                            nEmbdHead,
+//                            config.rmsNormEps())
                     .task("rmsnormMapIndexInPlace_Kcur",
                             Qwen3Kernels::rmsnormMapIndexInPlaceWithParallelOffset,
                             context,
@@ -359,8 +363,8 @@ private GridScheduler setupQwen3GridSchedulersLayeredNonNvidia() {
         WorkerGrid qCurWorker = new WorkerGrid1D(config.numberOfHeads() * nEmbdHead);
         qCurWorker.setLocalWork(nEmbdHead, 1, 1);
 
-        WorkerGrid qCurWorker2 = new WorkerGrid1D(config.numberOfHeads());
-        qCurWorker2.setLocalWork(1, 1, 1);
+//        WorkerGrid qCurWorker2 = new WorkerGrid1D(config.numberOfHeads());
+//        qCurWorker2.setLocalWork(1, 1, 1);
 
         // Kcur
         // config.numberOfKeyValueHeads() = 8
@@ -369,8 +373,8 @@ private GridScheduler setupQwen3GridSchedulersLayeredNonNvidia() {
         WorkerGrid kCurWorker = new WorkerGrid1D(config.numberOfKeyValueHeads() * nEmbdHead);
         kCurWorker.setLocalWork(nEmbdHead, 1, 1);
 
-        WorkerGrid kCurWorker2 = new WorkerGrid1D(config.numberOfKeyValueHeads());
-        kCurWorker2.setLocalWork(1, 1, 1);
+//        WorkerGrid kCurWorker2 = new WorkerGrid1D(config.numberOfKeyValueHeads());
+//        kCurWorker2.setLocalWork(1, 1, 1);
 
         int h = config.numberOfHeads();
         int ic = nEmbdHead / 2;
@@ -413,12 +417,12 @@ private GridScheduler setupQwen3GridSchedulersLayeredNonNvidia() {
 
             // Qcur
             gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormReduction_Qcur", qCurWorker);
-            gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormFinalNormalization_Qcur", qCurWorker2);
+            //gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormFinalNormalization_Qcur", qCurWorker2);
             gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormMapIndexInPlace_Qcur", qCurWorker);
 
             // Kcur
             gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormReduction_Kcur", kCurWorker);
-            gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormFinalNormalization_Kcur", kCurWorker2);
+            //gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormFinalNormalization_Kcur", kCurWorker2);
             gridScheduler.addWorkerGrid("layer_" + i + ".rmsnormMapIndexInPlace_Kcur", kCurWorker);
 
             gridScheduler.addWorkerGrid("layer_" + i + ".ropeRotation", ropeWorker);