Refactor Phi3TornadoVMLayerPlanner by reverting fusedGateUpSiLUDown to individual FFN tasks, updating worker grid mappings, and removing the fused kernel from TransformerComputeKernelsLayered for improved maintainability and clarity.

mikepapadim · mikepapadim · commit 3ce5f368a4af · 2025-08-03T22:45:24.000+03:00
diff --git a/src/main/java/com/example/tornadovm/Phi3TornadoVMLayerPlanner.java b/src/main/java/com/example/tornadovm/Phi3TornadoVMLayerPlanner.java
@@ -81,22 +81,12 @@ public Tuple2<List<ImmutableTaskGraph>, GridScheduler> setupTornadoForwardPlanLa
                             state.wrapX, config.dim(), config.rmsNormEps(), state.localSize)
                     .task("mapContextFFN", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, state.wrapXb,
                             state.wrapX, weights.rms_ffn_weightLayered[layerIndex], state.tempFFN)
-                    // Before (3 tasks):
-                    // .task("wGateUp", ...)
-                    // .task("gateUpSiLU", ...)
-                    // .task("wDown", ...)
-
-                    // After (1 fused task):
-                    .task("fusedFFN", TransformerComputeKernelsLayered::fusedGateUpSiLUDown, context,
-                            state.wrapXb, state.wrapX, weights.wUpLayered[layerIndex],
-                            weights.wDownLayered[layerIndex], config.dim(), config.hiddenDim(),
-                            LOCAL_WORK_GROUP_SIZE_ALLOC)
-//                    .task("wGateUp", TransformerComputeKernelsLayered::matrixVectorGeneric, context,
-//                            state.wrapXb,   state.wrapHb, weights.wUpLayered[layerIndex],  config.dim(), 2 * config.hiddenDim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
-//                    .task("gateUpSiLU", TransformerComputeKernelsLayered::splitGateUpAndSiLU,
-//                            state.wrapHb, state.wrapHbG, state.wrapHbU, config.hiddenDim())
-//                    .task("wDown", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context,
-//                            state.wrapHbU, state.wrapX, weights.wDownLayered[layerIndex], config.hiddenDim(), config.dim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
+                    .task("wGateUp", TransformerComputeKernelsLayered::matrixVectorGeneric, context,
+                            state.wrapXb,   state.wrapHb, weights.wUpLayered[layerIndex],  config.dim(), 2 * config.hiddenDim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
+                    .task("gateUpSiLU", TransformerComputeKernelsLayered::splitGateUpAndSiLU,
+                            state.wrapHb, state.wrapHbG, state.wrapHbU, config.hiddenDim())
+                    .task("wDown", TransformerComputeKernelsLayered::matrixVectorGenericWithResidual, context,
+                            state.wrapHbU, state.wrapX, weights.wDownLayered[layerIndex], config.hiddenDim(), config.dim(),  LOCAL_WORK_GROUP_SIZE_ALLOC)
                     .persistOnDevice(
                             state.wrapX
                     );
@@ -348,9 +338,7 @@ private GridScheduler setupGridSchedulersLayered() {
             tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".parallel-attention", parallelAttentionWorker);
             tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".copyToCaches", copyToCachesWorker);
             // New FFN tasks
-//            tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".gateUpSiLU", splitGateUpSiLUWorker);
-            tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".fusedFFN", configDimRowMajorGlobalWorker);
-
+            tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".gateUpSiLU", splitGateUpSiLUWorker);
         }
 
         // Vocabulary worker configuration
diff --git a/src/main/java/com/example/tornadovm/TransformerComputeKernelsLayered.java b/src/main/java/com/example/tornadovm/TransformerComputeKernelsLayered.java
@@ -962,59 +962,5 @@ public static void splitGateUpAndSiLU(FloatArray hb, FloatArray hbG, FloatArray
             hbU.set(i, siluGate * upVal);
         }
     }
-    public static void fusedGateUpSiLUDown(KernelContext context,
-            FloatArray input,        // state.wrapXb
-            FloatArray output,       // state.wrapX (with residual)
-            HalfFloatArray wUp,      // weights.wUpLayered[layerIndex]
-            HalfFloatArray wDown,    // weights.wDownLayered[layerIndex]
-            int dim,                 // config.dim()
-            int hiddenDim,          // config.hiddenDim()
-            int localWorkGroupSize) {
-
-        int rowId = context.groupIdx;  // Each workgroup computes one output dimension
-        int localId = context.localIdx;
-
-        if (rowId >= dim) return;
-
-        float[] localSum = context.allocateFloatLocalArray(localWorkGroupSize);
-        float accumulator = 0.0f;
-
-        // Process hidden dimensions in chunks to maintain numerical stability
-        for (int h = localId; h < hiddenDim; h += localWorkGroupSize) {
-            // Step 1: Compute gate value (first half of wUp)
-            float gateValue = 0.0f;
-            for (int i = 0; i < dim; i++) {
-                gateValue += wUp.get(h * dim + i).getFloat32() * input.get(i);
-            }
-
-            // Step 2: Compute up value (second half of wUp)
-            float upValue = 0.0f;
-            for (int i = 0; i < dim; i++) {
-                upValue += wUp.get((h + hiddenDim) * dim + i).getFloat32() * input.get(i);
-            }
-
-            // Step 3: Apply SiLU to gate and multiply with up
-            float siluGate = gateValue / (1.0f + TornadoMath.exp(-gateValue));
-            float activated = siluGate * upValue;
-
-            // Step 4: Apply down projection for this row
-            accumulator += wDown.get(rowId * hiddenDim + h).getFloat32() * activated;
-        }
 
-        // Reduce across workgroup
-        localSum[localId] = accumulator;
-        context.localBarrier();
-
-        for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
-            if (localId < stride) {
-                localSum[localId] += localSum[localId + stride];
-            }
-            context.localBarrier();
-        }
-
-        // Add residual connection
-        if (localId == 0) {
-            output.set(rowId, output.get(rowId) + localSum[0]);
-        }
-    }
 }
