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,