[moe] Fuse per-slot routed-expert launches: 8 kernel launches per layer to 2 (+15.4%)

src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java

@@ -53,7 +53,7 @@ public Qwen2MoEState(Configuration config, int batchsize) {
         this.wrapRouterLogits = new FloatArray(c.numberOfExperts());
         this.wrapSelectedExperts = new IntArray(c.numberOfExpertsUsed());
         this.wrapRoutingWeights = new FloatArray(c.numberOfExpertsUsed());
-        this.wrapExpertGate = new FloatArray(c.moeHiddenDim());
+        this.wrapExpertGate = new FloatArray(c.moeHiddenDim() * c.numberOfExpertsUsed());
         this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim());
         this.wrapSharedOutput = new FloatArray(c.dim());
     }

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

@@ -186,6 +186,166 @@ public static void fusedRoutedExpertGateUpSwiGLUQ8_0(
         }
     }
 
+    /**
+     * Gate/Up + SiLU for <b>all</b> routed slots in a single launch.
+     *
+     * <p>Functionally identical to calling {@link #fusedRoutedExpertGateUpSwiGLUQ8_0} once per
+     * slot; the slot index is folded into the work-group id instead, so top-K launches collapse
+     * into one. Each slot writes its own {@code moeHiddenDim}-sized window of
+     * {@code expertHidden}, so the slots never alias.
+     */
+    public static void fusedRoutedExpertsGateUpSwiGLUQ8_0All(
+            KernelContext context,
+            FloatArray input,
+            IntArray selectedExperts,
+            int expertsUsed,
+            ByteArray gateExperts,
+            ByteArray upExperts,
+            FloatArray expertHidden,
+            int dim,
+            int moeHiddenDim,
+            int numberOfExperts,
+            int localWorkGroupSize) {
+
+        int flatGroupId = context.groupIdx;
+        int localId = context.localIdx;
+
+        int slot = flatGroupId / moeHiddenDim;
+        int rowId = flatGroupId - slot * moeHiddenDim;
+
+        // A work-group whose slot or row falls outside the launch still has to reach the
+        // barriers below, so the guard only suppresses the memory accesses and the store.
+        boolean active = slot < expertsUsed && rowId < moeHiddenDim;
+        int expert = 0;
+        if (active) {
+            expert = selectedExperts.get(slot);
+            active = expert >= 0 && expert < numberOfExperts;
+        }
+
+        int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE;
+        int rowBlockOffset = (expert * moeHiddenDim + rowId) * blocksPerRow;
+
+        float gatePartialSum = 0.0f;
+        float upPartialSum = 0.0f;
+        if (active) {
+            for (int column = localId; column < dim; column += localWorkGroupSize) {
+                int blockByteOffset =
+                        (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES;
+                int quantOffset =
+                        blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE;
+
+                float inputValue = input.get(column);
+                float gateScale = gateExperts.getHalfFloat(blockByteOffset).getFloat32();
+                float upScale = upExperts.getHalfFloat(blockByteOffset).getFloat32();
+
+                gatePartialSum += (gateExperts.get(quantOffset) * gateScale) * inputValue;
+                upPartialSum += (upExperts.get(quantOffset) * upScale) * inputValue;
+            }
+        }
+
+        float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize);
+        localSums[localId] = gatePartialSum;
+        context.localBarrier();
+        for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
+            if (localId < stride) {
+                localSums[localId] += localSums[localId + stride];
+            }
+            context.localBarrier();
+        }
+        float gate = localSums[0];
+
+        localSums[localId] = upPartialSum;
+        context.localBarrier();
+        for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
+            if (localId < stride) {
+                localSums[localId] += localSums[localId + stride];
+            }
+            context.localBarrier();
+        }
+
+        if (localId == 0 && active) {
+            float up = localSums[0];
+            float siluGate = gate / (1.0f + TornadoMath.exp(-gate));
+            expertHidden.set(slot * moeHiddenDim + rowId, siluGate * up);
+        }
+    }
+
+    /**
+     * Down-projects <b>all</b> routed slots and accumulates them into the residual in one launch.
+     *
+     * <p>Beyond collapsing top-K launches into one, this also folds the per-slot partial sums
+     * before the reduction, so the work-group reduces once instead of K times and the residual is
+     * read-modify-written once instead of K times.
+     */
+    public static void routedExpertsDownProjectAndAccumulateQ8_0All(
+            KernelContext context,
+            FloatArray expertHidden,
+            FloatArray residual,
+            IntArray selectedExperts,
+            FloatArray routingWeights,
+            int expertsUsed,
+            ByteArray downExperts,
+            int dim,
+            int moeHiddenDim,
+            int numberOfExperts,
+            int localWorkGroupSize) {
+
+        int rowId = context.groupIdx;
+        int localId = context.localIdx;
+        boolean active = rowId < dim;
+
+        int blocksPerRow = (moeHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE;
+        float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize);
+
+        // Lane 0 carries the running residual across slots. Each slot is reduced with the same
+        // tree and folded in the same order as the per-slot kernels, so the result is bit-identical
+        // to launching them separately - only the launch count and the residual write change.
+        float running = 0.0f;
+        if (localId == 0 && active) {
+            running = residual.get(rowId);
+        }
+
+        for (int slot = 0; slot < expertsUsed; slot++) {
+            int expert = selectedExperts.get(slot);
+            boolean slotActive = active && expert >= 0 && expert < numberOfExperts;
+
+            float partialSum = 0.0f;
+            if (slotActive) {
+                int rowBlockOffset = (expert * dim + rowId) * blocksPerRow;
+                int hiddenBase = slot * moeHiddenDim;
+                for (int column = localId;
+                     column < moeHiddenDim;
+                     column += localWorkGroupSize) {
+                    int blockByteOffset =
+                            (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES;
+                    int quantOffset = blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE;
+
+                    float weight = downExperts.get(quantOffset)
+                            * downExperts.getHalfFloat(blockByteOffset).getFloat32();
+                    partialSum += weight * expertHidden.get(hiddenBase + column);
+                }
+            }
+
+            context.localBarrier();
+            localSums[localId] = partialSum;
+            context.localBarrier();
+            for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
+                if (localId < stride) {
+                    localSums[localId] += localSums[localId + stride];
+                }
+                context.localBarrier();
+            }
+
+            if (localId == 0 && slotActive) {
+                running += routingWeights.get(slot) * localSums[0];
+            }
+        }
+
+        if (localId == 0 && active) {
+            residual.set(rowId, running);
+        }
+    }
+
     /**
      * Down-projects one selected expert and accumulates its routed contribution:
      * {@code residual += routingWeight[slot] * W_down[expert] * expertHidden}.

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java

@@ -64,6 +64,9 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
         WorkerGrid topKWorker = new WorkerGrid1D(LOCAL_WORK_GROUP_SIZE_ALLOC);
         topKWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1);
         WorkerGrid expertHiddenWorker = workerForRows(config.moeHiddenDim());
+        // Fused routed-expert launch: the slot index is folded into the work-group id.
+        WorkerGrid allExpertsHiddenWorker =
+                workerForRows(config.moeHiddenDim() * config.numberOfExpertsUsed());
         WorkerGrid sharedHiddenWorker = workerForRows(config.sharedExpertHiddenDim());
 
         for (int layer = 0; layer < config.numberOfLayers(); layer++) {
@@ -78,10 +81,8 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
             scheduler.addWorkerGrid(prefix + "ffn_rms_apply", dimElementWorker);
             scheduler.addWorkerGrid(prefix + "router_projection", routerWorker);
             scheduler.addWorkerGrid(prefix + "router_softmax_topk", topKWorker);
-            for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) {
-                scheduler.addWorkerGrid(prefix + "routed_expert_gate_up_" + slot, expertHiddenWorker);
-                scheduler.addWorkerGrid(prefix + "routed_expert_down_" + slot, dimWorker);
-            }
+            scheduler.addWorkerGrid(prefix + "routed_experts_gate_up", allExpertsHiddenWorker);
+            scheduler.addWorkerGrid(prefix + "routed_experts_down", dimWorker);
             scheduler.addWorkerGrid(prefix + "shared_expert_gate_up", sharedHiddenWorker);
             scheduler.addWorkerGrid(prefix + "shared_expert_down", dimWorker);
             scheduler.addWorkerGrid(prefix + "shared_expert_gate_and_accumulate", topKWorker);
@@ -211,21 +212,22 @@ private void configureRoutedExperts(TaskGraph layer, int layerIndex) {
                 context, moeState.wrapRouterLogits, moeState.wrapSelectedExperts,
                 moeState.wrapRoutingWeights, config.numberOfExperts(), config.numberOfExpertsUsed());
 
-        for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) {
-            layer.task("routed_expert_gate_up_" + slot,
-                    Qwen2MoEKernels::fusedRoutedExpertGateUpSwiGLUQ8_0,
-                    context, moeState.wrapXb, moeState.wrapSelectedExperts, slot,
-                    weights.gateExpertsLayered[layerIndex].asByteArray(),
-                    weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate,
-                    config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
-
-            layer.task("routed_expert_down_" + slot,
-                    Qwen2MoEKernels::routedExpertDownProjectAndAccumulateQ8_0,
-                    context, moeState.wrapExpertGate, moeState.wrapX,
-                    moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, slot,
-                    weights.downExpertsLayered[layerIndex].asByteArray(),
-                    config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
-        }
+        // All routed slots in two launches instead of two per slot: at top-4 this is 2 kernel
+        // launches per layer rather than 8, and the residual is accumulated once instead of
+        // four times.
+        layer.task("routed_experts_gate_up",
+                Qwen2MoEKernels::fusedRoutedExpertsGateUpSwiGLUQ8_0All,
+                context, moeState.wrapXb, moeState.wrapSelectedExperts, config.numberOfExpertsUsed(),
+                weights.gateExpertsLayered[layerIndex].asByteArray(),
+                weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate,
+                config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+        layer.task("routed_experts_down",
+                Qwen2MoEKernels::routedExpertsDownProjectAndAccumulateQ8_0All,
+                context, moeState.wrapExpertGate, moeState.wrapX,
+                moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, config.numberOfExpertsUsed(),
+                weights.downExpertsLayered[layerIndex].asByteArray(),
+                config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
 
         // The shared expert always runs; it does not depend on router top-K selection.
         layer.task("shared_expert_gate_up", Qwen2MoEKernels::sharedExpertGateUpSwiGLUQ8_0,