From 3940767bff6bdafc76ebb84138df524a2cd91221 Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Thu, 6 Aug 2026 11:47:31 +0100 Subject: [PATCH 1/2] [moe] Fuse per-slot routed-expert launches: 8 kernel launches per layer to 2 --- .../inference/state/Qwen2MoEState.java | 2 +- .../tornadovm/kernels/Qwen2MoEKernels.java | 160 ++++++++++++++++++ .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 40 ++--- 3 files changed, 182 insertions(+), 20 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java index d0872883..d7bdc92d 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/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()); } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java index d16bc5b0..cbcc4867 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java @@ -186,6 +186,166 @@ public static void fusedRoutedExpertGateUpSwiGLUQ8_0( } } + /** + * Gate/Up + SiLU for all routed slots in a single launch. + * + *

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 all routed slots and accumulates them into the residual in one launch. + * + *

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}. diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java index a845c8b0..c085f09c 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java +++ b/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, From d022b75d7d44a45493c9b58eee599e55c9cdb150 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Thu, 6 Aug 2026 20:11:43 +0100 Subject: [PATCH 2/2] Clean up fused routed-expert kernels --- .../tornadovm/kernels/Qwen2MoEKernels.java | 184 +----------------- .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 5 +- 2 files changed, 9 insertions(+), 180 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java index cbcc4867..db4eac56 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java @@ -87,114 +87,14 @@ public static void softmaxAndTopK( } } - /** - * Computes the routed expert's gated activation for one top-K slot: - * {@code SiLU(W_gate[expert] * x) * (W_up[expert] * x)}. - * - *

The expert matrices are stacked in one Q8_0 tensor per layer. The - * selected expert id determines which matrix slice this kernel reads.

- */ - public static void fusedRoutedExpertGateUpSwiGLUQ8_0( - KernelContext context, - FloatArray input, - IntArray selectedExperts, - int slot, - ByteArray gateExperts, - ByteArray upExperts, - FloatArray expertHidden, - int dim, - int moeHiddenDim, - int numberOfExperts, - int localWorkGroupSize) { - - int rowId = context.groupIdx; - int localId = context.localIdx; - - int expert = selectedExperts.get(slot); - if (rowId >= moeHiddenDim || expert < 0 || expert >= numberOfExperts) { - return; - } - - // Locate this output row within the selected expert's stacked matrix. - int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; - int rowBlockOffset = - (expert * moeHiddenDim + rowId) * blocksPerRow; - - // One workgroup cooperates on the Gate and Up dot products for this row. - float gatePartialSum = 0.0f; - float upPartialSum = 0.0f; - - for (int column = localId; - column < dim; - column += localWorkGroupSize) { - - // Byte offset of the first byte of the Q8_0 block that contains this column. - // Each block occupies 34 bytes: a 2-byte FP16 scale plus 32 int8 quants. - int blockByteOffset = - (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; - - // Skip the 2-byte scale at the block start and locate this column's int8 quant. - int quantOffset = - blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; - - float inputValue = input.get(column); - - // getHalfFloat reads the FP16 scale from the first two block bytes, then converts it to FP32. - float gateScale = - gateExperts.getHalfFloat(blockByteOffset).getFloat32(); - float upScale = - upExperts.getHalfFloat(blockByteOffset).getFloat32(); - - byte gateQuant = gateExperts.get(quantOffset); - byte upQuant = upExperts.get(quantOffset); - - float gateWeight = gateQuant * gateScale; - float upWeight = upQuant * upScale; - - gatePartialSum += gateWeight * inputValue; - upPartialSum += upWeight * inputValue; - - } - - // Sum the partial gate values from all threads in this workgroup. - 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]; - - // Reuse local memory to sum the partial up values. - localSums[localId] = upPartialSum; - context.localBarrier(); - for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { - if (localId < stride) { - localSums[localId] += localSums[localId + stride]; - } - context.localBarrier(); - } - - // One thread writes this output row after both reductions are complete. - if (localId == 0) { - float up = localSums[0]; - float siluGate = gate / (1.0f + TornadoMath.exp(-gate)); - expertHidden.set(rowId, siluGate * up); - } - } - /** * Gate/Up + SiLU for all routed slots in a single launch. * - *

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. + *

The slot index is folded into the work-group id, so all top-K slots execute in one + * launch. Each slot writes its own {@code moeHiddenDim}-sized window of {@code expertHidden}, + * so the slots never alias. */ - public static void fusedRoutedExpertsGateUpSwiGLUQ8_0All( + public static void fusedRoutedExpertsGateUpSwiGLUQ8_0( KernelContext context, FloatArray input, IntArray selectedExperts, @@ -273,11 +173,10 @@ public static void fusedRoutedExpertsGateUpSwiGLUQ8_0All( /** * Down-projects all routed slots and accumulates them into the residual in one launch. * - *

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. + *

Each slot retains the original reduction tree and is accumulated in slot order. The + * optimization removes the per-slot kernel launches and writes the final residual once. */ - public static void routedExpertsDownProjectAndAccumulateQ8_0All( + public static void routedExpertsDownProjectAndAccumulateQ8_0( KernelContext context, FloatArray expertHidden, FloatArray residual, @@ -346,75 +245,6 @@ public static void routedExpertsDownProjectAndAccumulateQ8_0All( } } - /** - * Down-projects one selected expert and accumulates its routed contribution: - * {@code residual += routingWeight[slot] * W_down[expert] * expertHidden}. - */ - public static void routedExpertDownProjectAndAccumulateQ8_0( - KernelContext context, - FloatArray expertHidden, - FloatArray residual, - IntArray selectedExperts, - FloatArray routingWeights, - int slot, - ByteArray downExperts, - int dim, - int moeHiddenDim, - int numberOfExperts, - int localWorkGroupSize) { - - // One workgroup produces one element of the down-projected vector. - int rowId = context.groupIdx; - int localId = context.localIdx; - if (rowId >= dim) { - return; - } - - int expert = selectedExperts.get(slot); - if (expert < 0 || expert >= numberOfExperts) { - return; - } - float routingWeight = routingWeights.get(slot); - - // downExperts has the logical shape [experts, dim, moeHiddenDim]. - int blocksPerRow = (moeHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; - int rowBlockOffset = (expert * dim + rowId) * blocksPerRow; - - // Every thread accumulates a different subset of this row's dot product. - float partialSum = 0.0f; - for (int column = localId; - column < moeHiddenDim; - column += localWorkGroupSize) { - // The start byte of the Q8_0 block holding this down-projection weight. - // Block layout: a 2-byte FP16 scale followed by 32 int8 quants. - int blockByteOffset = - (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; - - // Quants begin immediately after the scale; column % 32 is the index within this block. - int quantOffset = blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; - - float weight = downExperts.get(quantOffset) - * downExperts.getHalfFloat(blockByteOffset).getFloat32(); - partialSum += weight * expertHidden.get(column); - } - - // Combine all thread-local partial sums into the completed output row. - float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); - 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) { - float outputValue = localSums[0]; - residual.set(rowId, residual.get(rowId) + routingWeight * outputValue); - } - } - /** Computes {@code SiLU(sharedGate * input) * (sharedUp * input)}. */ public static void sharedExpertGateUpSwiGLUQ8_0( KernelContext context, diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java index 4fd68d28..ee7a07e0 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java @@ -64,7 +64,6 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) { WorkerGrid routerWorker = workerForRows(config.numberOfExperts()); 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()); @@ -205,14 +204,14 @@ private void configureRoutedExperts(TaskGraph layer, int layerIndex) { // 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, + Qwen2MoEKernels::fusedRoutedExpertsGateUpSwiGLUQ8_0, 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, + Qwen2MoEKernels::routedExpertsDownProjectAndAccumulateQ8_0, context, moeState.wrapExpertGate, moeState.wrapX, moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, config.numberOfExpertsUsed(), weights.downExpertsLayered[layerIndex].asByteArray(),