Replace splitQKV kernel with fusedRmsNormQKVMatmulDirect, refactor Phi3 FP16 FFN layers to consolidate QKV projection tasks, and update worker grid/task configurations.

Author: mikepapadim

src/main/java/org/beehive/gpullama3/tornadovm/kernels/Phi3Kernels.java

@@ -192,4 +192,95 @@ public static void ropeRotationWithCacheCopyPhi3(
             valueCache.set(cacheOffset + base + idx + dimHalf, sv.get(base + idx + dimHalf));
         }
     }
+
+    /**
+     * Fused RMSNorm apply + QKV projection with direct output to separate Q, K, V buffers.
+     *
+     * <p>Eliminates the need for a separate splitQKV kernel by routing outputs
+     * directly based on row index:</p>
+     * <ul>
+     *   <li>Rows [0, dim): Q projection</li>
+     *   <li>Rows [dim, dim+kvDim): K projection</li>
+     *   <li>Rows [dim+kvDim, dim+2*kvDim): V projection</li>
+     * </ul>
+     *
+     * <p>Formula: output[row] = sum_j(Wqkv[row,j] * rmsWeight[j] * scale * x[j])</p>
+     *
+     * @param context           Kernel execution context
+     * @param x                 Input hidden state (FP32) [dim]
+     * @param q                 Output Q buffer (FP32) [dim]
+     * @param k                 Output K buffer (FP32) [kvDim]
+     * @param v                 Output V buffer (FP32) [kvDim]
+     * @param rmsWeights        RMS normalization weights (FP32) [dim]
+     * @param rmsScale          Precomputed RMS scale factor [1]
+     * @param wqkv              Combined QKV weight matrix (FP16) [opSize × dim]
+     * @param dim               Model dimension (Q output size)
+     * @param kvDim             KV dimension (K/V output size)
+     * @param localWorkGroupSize Local work group size for reduction
+     */
+    public static void fusedRmsNormQKVMatmulDirect(
+            KernelContext context,
+            FloatArray x,               // input (FP32)
+            FloatArray q,               // output Q (FP32)
+            FloatArray k,               // output K (FP32)
+            FloatArray v,               // output V (FP32)
+            FloatArray rmsWeights,      // RMS norm weights
+            FloatArray rmsScale,        // temp[0] = scale factor
+            HalfFloatArray wqkv,        // combined QKV weight matrix
+            int dim,                    // input dim and Q output dim
+            int kvDim,                  // K/V output dim
+            int localWorkGroupSize) {
+
+        int rowId = context.groupIdx;
+        int localId = context.localIdx;
+
+        // Total rows = dim (Q) + kvDim (K) + kvDim (V)
+        int totalRows = dim + 2 * kvDim;
+        if (rowId >= totalRows) {
+            return;
+        }
+
+        float scale = rmsScale.get(0);
+
+        // Allocate shared memory for reduction
+        float[] localSum = context.allocateFloatLocalArray(localWorkGroupSize);
+
+        int rowOffset = rowId * dim;
+
+        // Each thread computes partial dot product with inline normalization
+        float partialSum = 0.0f;
+        for (int j = localId; j < dim; j += localWorkGroupSize) {
+            float normalized = rmsWeights.get(j) * scale * x.get(j);
+            partialSum += wqkv.get(rowOffset + j).getFloat32() * normalized;
+        }
+
+        localSum[localId] = partialSum;
+        context.localBarrier();
+
+        // Parallel reduction within workgroup
+        for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
+            if (localId < stride) {
+                localSum[localId] += localSum[localId + stride];
+            }
+            context.localBarrier();
+        }
+
+        // Thread 0 writes to appropriate output buffer
+        if (localId == 0) {
+            float result = localSum[0];
+
+            if (rowId < dim) {
+                // Q projection: rows [0, dim)
+                q.set(rowId, result);
+            } else if (rowId < dim + kvDim) {
+                // K projection: rows [dim, dim+kvDim)
+                int kIdx = rowId - dim;
+                k.set(kIdx, result);
+            } else {
+                // V projection: rows [dim+kvDim, dim+2*kvDim)
+                int vIdx = rowId - dim - kvDim;
+                v.set(vIdx, result);
+            }
+        }
+    }
 }

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

@@ -79,13 +79,17 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) {
         // FFN down projection worker
         int ffnDownGlobal = config.dim() * LOCAL_WORK_GROUP_SIZE_ALLOC;
         WorkerGrid ffnDownWorker = WorkerGridFactory.genericWorker(ffnDownGlobal, LOCAL_WORK_GROUP_SIZE_ALLOC);
+        // Same worker as before - total rows = dim + 2*kvDim = opSize
 
+        // Remove: gridScheduler.addWorkerGrid("layer_" + i + ".splitQKV", splitQKVWorker);
         // Map workers to tasks for each layer (in task execution order)
         for (int i = 0; i < config.numberOfLayers(); i++) {
             // === Attention Block ===
             gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_reduce", rmsNormWorker);
-            gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_qkv_matmul", fusedQkvWorker);
-            gridScheduler.addWorkerGrid("layer_" + i + ".splitQKV", splitQKVWorker);
+//            gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_qkv_matmul", fusedQkvWorker);
+            gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_qkv_projection", fusedQkvWorker);
+
+//            gridScheduler.addWorkerGrid("layer_" + i + ".splitQKV", splitQKVWorker);
             gridScheduler.addWorkerGrid("layer_" + i + ".rope_and_kv_cache", ropeWorker);
             gridScheduler.addWorkerGrid("layer_" + i + ".attention", parallelAttentionWorker);
             gridScheduler.addWorkerGrid("layer_" + i + ".attn_output_proj", matmul1Worker);
@@ -257,29 +261,44 @@ TaskGraph setupSinglePhi3FFNLayer(Phi3TornadoWeights weights, int layerIndex) {
                 phi3Config.rmsNormEps(),      // epsilon
                 phi3State.localSize);         // local memory size
 
-        // Fused RMS Apply + QKV Projection (combined matrix)
-        unifiedLayer.task("attn_rms_qkv_matmul",
-                Phi3Kernels::fusedRmsNormMatmul,
+//        // Fused RMS Apply + QKV Projection (combined matrix)
+//        unifiedLayer.task("attn_rms_qkv_matmul",
+//                Phi3Kernels::fusedRmsNormMatmul,
+//                context,
+//                phi3State.wrapX,              // input: raw hidden state (FP32)
+//                phi3State.wrapQkv,            // output: combined Q+K+V
+//                weights.rms_att_weightLayered[layerIndex].asFloatArray(),  // RMS weights
+//                phi3State.temp,               // RMS scale factor from reduction
+//                weights.wqkvLayered[layerIndex].asHalfFloatArray(),        // Wqkv [opSize × dim]
+//                phi3Config.dim(),             // input dimension
+//                opSize,                       // output dimension (Q + K + V)
+//                LOCAL_WORK_GROUP_SIZE_ALLOC);
+//
+//        // Split combined QKV into separate Q, K, V buffers
+//        unifiedLayer.task("splitQKV",
+//                TransformerComputeKernelsLayered::splitQKV,
+//                phi3State.wrapQkv,
+//                phi3State.wrapQ,
+//                phi3State.wrapK,
+//                phi3State.wrapV,
+//                phi3Config.dim(),
+//                phi3Config.headSize() * phi3Config.numberOfKeyValueHeads());
+
+        // AFTER: 1 task
+        unifiedLayer.task("attn_rms_qkv_projection",
+                Phi3Kernels::fusedRmsNormQKVMatmulDirect,
                 context,
-                phi3State.wrapX,              // input: raw hidden state (FP32)
-                phi3State.wrapQkv,            // output: combined Q+K+V
-                weights.rms_att_weightLayered[layerIndex].asFloatArray(),  // RMS weights
-                phi3State.temp,               // RMS scale factor from reduction
-                weights.wqkvLayered[layerIndex].asHalfFloatArray(),        // Wqkv [opSize × dim]
-                phi3Config.dim(),             // input dimension
-                opSize,                       // output dimension (Q + K + V)
+                phi3State.wrapX,              // input
+                phi3State.wrapQ,              // output Q
+                phi3State.wrapK,              // output K
+                phi3State.wrapV,              // output V
+                weights.rms_att_weightLayered[layerIndex].asFloatArray(),
+                phi3State.temp,               // RMS scale
+                weights.wqkvLayered[layerIndex].asHalfFloatArray(),
+                phi3Config.dim(),             // dim
+                phi3Config.kvDim(),           // kvDim
                 LOCAL_WORK_GROUP_SIZE_ALLOC);
 
-        // Split combined QKV into separate Q, K, V buffers
-        unifiedLayer.task("splitQKV",
-                TransformerComputeKernelsLayered::splitQKV,
-                phi3State.wrapQkv,
-                phi3State.wrapQ,
-                phi3State.wrapK,
-                phi3State.wrapV,
-                phi3Config.dim(),
-                phi3Config.headSize() * phi3Config.numberOfKeyValueHeads());
-
         // Fused Phi3 RoPE Rotation + KV Cache Write
         unifiedLayer.task("rope_and_kv_cache",
                 Phi3Kernels::ropeRotationWithCacheCopyPhi3,