stopTokens, int maxTokens, Sampler sampler, boolean echo,
+ IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) {
+ if (WITH_PREFILL_DECODE && TornadoVMMasterPlan.PREFILL_BATCH_SIZE > 1) {
+ throw new UnsupportedOperationException("Batch prefill/decode on GPU not yet implemented for Qwen2-MoE");
+ }
+ if (WITH_PREFILL_DECODE) {
+ throw new UnsupportedOperationException("Prefill/decode on GPU not yet implemented for Qwen2-MoE");
+ }
+ return InferenceEngine.generateTokensGPUQwen3(this, state, startPosition, promptTokens,
+ stopTokens, maxTokens, sampler, echo, onTokenGenerated, tornadoVMPlan);
+ }
+}
diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java
new file mode 100644
index 00000000..9b4a0ed5
--- /dev/null
+++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java
@@ -0,0 +1,47 @@
+package org.beehive.gpullama3.model.qwen2;
+
+import org.beehive.gpullama3.model.Configuration;
+
+public record Qwen2MoEConfiguration(String quantization,
+ int dim,
+ int hiddenDim,
+ int numberOfLayers,
+ int numberOfHeads,
+ int numberOfKeyValueHeads,
+ int numberOfHeadsKey,
+ int numberOfHeadsValue,
+ int vocabularySize,
+ int contextLengthModel,
+ int contextLength,
+ int numberOfExperts,
+ int numberOfExpertsUsed,
+ int moeHiddenDim,
+ int sharedExpertHiddenDim,
+ boolean sharedWeights,
+ float rmsNormEps,
+ float ropeTheta) implements Configuration {
+ @Override
+ public String quantization() {
+ return quantization;
+ }
+
+ @Override
+ public int headSize() {
+ return dim / numberOfHeads;
+ }
+
+ @Override
+ public int kvDim() {
+ return (dim * numberOfKeyValueHeads) / numberOfHeads;
+ }
+
+ @Override
+ public int kvMul() {
+ return numberOfHeads / numberOfKeyValueHeads;
+ }
+
+ @Override
+ public int contextLengthModel() {
+ return contextLengthModel;
+ }
+}
diff --git a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java
index 9067bde0..d3304d40 100644
--- a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java
+++ b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java
@@ -65,8 +65,15 @@ public float getFloat(int index) {
public static final ValueLayout.OfShort JAVA_SHORT_LE = ValueLayout.JAVA_SHORT.withOrder(ByteOrder.LITTLE_ENDIAN);
+ // When enabled, quantize the activation to Q8_0 and do an int8·int8 dot,
+ // matching llama.cpp's Q8_0 matmul quantization scheme.
+ static final boolean QUANTIZE_ACTIVATION = Boolean.parseBoolean(System.getProperty("llama.quantizeActivation", "true"));
+
@Override
public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) {
+ if (QUANTIZE_ACTIVATION) {
+ return dotQ8Activation(thisOffset, that, thatOffset, size);
+ }
if (USE_VECTOR_API) {
return vectorDot(this, thisOffset, (ArrayFloatTensor) that, thatOffset, size);
} else {
@@ -74,6 +81,47 @@ public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) {
}
}
+ /**
+ * Q8_0 weight · activation, where the activation is first quantized to Q8_0 (per 32-element
+ * block) and the dot is accumulated as int8·int8 -> int32, then scaled. This mirrors
+ * llama.cpp's ggml Q8_0 matmul path. Assumes thisOffset and size are multiples of the
+ * Q8_0 block size (32), which holds for all matmul callers.
+ */
+ private float dotQ8Activation(int thisOffset, FloatTensor that, int thatOffset, int size) {
+ final int BS = GGMLType.Q8_0.getBlockSize(); // 32
+ final int TS = GGMLType.Q8_0.getTypeSize(); // 34
+ float result = 0f;
+ int nBlocks = size / BS;
+ for (int b = 0; b < nBlocks; b++) {
+ int elemBase = b * BS;
+ int wBlockOffset = (thisOffset + elemBase) / BS * TS;
+ float wScale = Float.float16ToFloat(readShort(memorySegment, wBlockOffset));
+
+ // find the max abs of this activation block -> activation scale
+ float amax = 0f;
+ for (int i = 0; i < BS; i++) {
+ float av = Math.abs(that.getFloat(thatOffset + elemBase + i));
+ if (av > amax) amax = av;
+ }
+ // Match ggml's Q8_0 quantization order: derive the int8 values using the
+ // full-precision scale, but store/use the scale itself as f16.
+ float quantizationScale = amax / 127f;
+ float aScale = Float.float16ToFloat(Float.floatToFloat16(quantizationScale));
+ float aInv = quantizationScale != 0f ? 1f / quantizationScale : 0f;
+
+ // int8 · int8 accumulation (round-half-away-from-zero, matching ggml roundf)
+ int isum = 0;
+ for (int i = 0; i < BS; i++) {
+ float s = that.getFloat(thatOffset + elemBase + i) * aInv;
+ int aq = (int) (s + Math.copySign(0.5f, s));
+ byte wq = readByte(memorySegment, wBlockOffset + Float16.BYTES + i);
+ isum += aq * wq;
+ }
+ result += isum * (wScale * aScale);
+ }
+ return result;
+ }
+
private static float vectorDot(Q8_0FloatTensor thiz, int thisOffset, ArrayFloatTensor that, int thatOffset, int size) {
float result = 0f;
int j = 0;
diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java
new file mode 100644
index 00000000..d16bc5b0
--- /dev/null
+++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java
@@ -0,0 +1,447 @@
+package org.beehive.gpullama3.tornadovm.kernels;
+
+import uk.ac.manchester.tornado.api.KernelContext;
+import uk.ac.manchester.tornado.api.math.TornadoMath;
+import uk.ac.manchester.tornado.api.types.arrays.ByteArray;
+import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.IntArray;
+
+/**
+ * TornadoVM kernels specific to Qwen2-MoE / Qwen1.5-MoE inference.
+ */
+public final class Qwen2MoEKernels {
+
+ private static final int Q8_0_BLOCK_SIZE = 32;
+ private static final int Q8_0_BLOCK_BYTES = 34;
+
+ private Qwen2MoEKernels() {
+ }
+
+ /**
+ * Converts router scores to probabilities and selects the highest-scoring
+ * {@code topK} experts for one token.
+ *
+ * Inputs and outputs are GPU-resident TornadoVM arrays:
+ *
+ * - {@code routerLogits}: one raw score per expert
+ * - {@code selectedExperts}: output expert indices, length {@code topK}
+ * - {@code routingWeights}: output probabilities, length {@code topK}
+ *
+ *
+ * The first implementation deliberately uses one GPU thread for this
+ * small operation (60 experts in the target model). A later optimization
+ * can parallelize the reductions within one workgroup.
+ */
+ public static void softmaxAndTopK(
+ KernelContext context,
+ FloatArray routerLogits,
+ IntArray selectedExperts,
+ FloatArray routingWeights,
+ int numberOfExperts,
+ int topK) {
+
+ // The first correctness-oriented implementation is serial. Without
+ // this guard, every thread would race to overwrite the same buffers.
+ if (context.groupIdx != 0 || context.localIdx != 0) {
+ return;
+ }
+
+ // Find the maximum first to keep the softmax numerically stable.
+ float maxLogit = Float.NEGATIVE_INFINITY;
+ int maxIndex = -1;
+ for (int expert = 0; expert < numberOfExperts; expert++) {
+ float logit = routerLogits.get(expert);
+ if (logit > maxLogit) {
+ maxLogit = logit;
+ maxIndex = expert;
+ }
+ }
+
+ // Convert every router score to a probability over all experts.
+ float sumExp = 0.0f;
+ for (int expert = 0; expert < numberOfExperts; expert++) {
+ sumExp += TornadoMath.exp(routerLogits.get(expert) - maxLogit);
+ }
+ for (int expert = 0; expert < numberOfExperts; expert++) {
+ float probability = TornadoMath.exp(routerLogits.get(expert) - maxLogit) / sumExp;
+ routerLogits.set(expert, probability);
+ }
+
+ // Select the top-K probabilities without renormalizing their sum.
+ selectedExperts.set(0, maxIndex);
+ routingWeights.set(0, routerLogits.get(maxIndex));
+ routerLogits.set(maxIndex, Float.NEGATIVE_INFINITY);
+
+ for (int slot = 1; slot < topK; slot++) {
+ maxLogit = Float.NEGATIVE_INFINITY;
+ maxIndex = -1;
+ for (int expert = 0; expert < numberOfExperts; expert++) {
+ if (routerLogits.get(expert) > maxLogit) {
+ maxLogit = routerLogits.get(expert);
+ maxIndex = expert;
+ }
+ }
+ selectedExperts.set(slot, maxIndex);
+ routingWeights.set(slot, routerLogits.get(maxIndex));
+ routerLogits.set(maxIndex, Float.NEGATIVE_INFINITY);
+ }
+ }
+
+ /**
+ * 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);
+ }
+ }
+
+ /**
+ * 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,
+ FloatArray input,
+ ByteArray sharedGate,
+ ByteArray sharedUp,
+ FloatArray sharedHidden,
+ int dim,
+ int sharedExpertHiddenDim,
+ int localWorkGroupSize) {
+
+ // One workgroup computes one shared-expert hidden output row.
+ int rowId = context.groupIdx;
+ int localId = context.localIdx;
+
+ if (rowId >= sharedExpertHiddenDim) {
+ return;
+ }
+
+ // Number of Q8_0 blocks needed for one row of dim weights.
+ int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE;
+
+ // Index of this row's first Q8_0 block in the shared weight array.
+ int rowBlockOffset = rowId * blocksPerRow;
+
+ 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 =
+ sharedGate.getHalfFloat(blockByteOffset).getFloat32();
+ float upScale =
+ sharedUp.getHalfFloat(blockByteOffset).getFloat32();
+
+ byte gateQuant = sharedGate.get(quantOffset);
+ byte upQuant = sharedUp.get(quantOffset);
+
+ float gateWeight = gateQuant * gateScale;
+ float upWeight = upQuant * upScale;
+
+ gatePartialSum += gateWeight * inputValue;
+ upPartialSum += upWeight * 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];
+
+ // 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));
+ sharedHidden.set(rowId, siluGate * up);
+ }
+ }
+
+ /** Down-projects the shared expert hidden vector into the model dimension. */
+ public static void sharedExpertDownProjectQ8_0(
+ KernelContext context,
+ FloatArray sharedHidden,
+ ByteArray sharedDown,
+ FloatArray sharedOutput,
+ int dim,
+ int sharedExpertHiddenDim,
+ int localWorkGroupSize) {
+
+
+ int rowId = context.groupIdx;
+ int localId = context.localIdx;
+ if (rowId >= dim) {
+ return;
+ }
+
+ // Number of Q8_0 blocks needed for one row of dim weights.
+ int blocksPerRow = (sharedExpertHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE;
+
+ // Index of this row's first Q8_0 block in the shared weight array.
+ int rowBlockOffset = rowId * blocksPerRow;
+
+ float partialSum = 0.0f;
+ for (int column = localId;
+ column < sharedExpertHiddenDim;
+ 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 = sharedDown.get(quantOffset)
+ * sharedDown.getHalfFloat(blockByteOffset).getFloat32();
+ partialSum += weight * sharedHidden.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];
+ sharedOutput.set(rowId, outputValue);
+ }
+ }
+
+ /** Computes the shared gate sigmoid and adds the weighted shared output to the residual. */
+ public static void sharedExpertGateAndAccumulate(
+ KernelContext context,
+ FloatArray input,
+ FloatArray sharedGateInput,
+ FloatArray sharedOutput,
+ FloatArray residual,
+ int dim,
+ int localWorkGroupSize) {
+ int localId = context.localIdx;
+
+ float partialScore = 0.0f;
+
+ for (int column = localId;
+ column < dim;
+ column += localWorkGroupSize) {
+ partialScore += sharedGateInput.get(column) * input.get(column);
+ }
+
+ float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize);
+ localSums[localId] = partialScore;
+ context.localBarrier();
+ for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
+ if (localId < stride) {
+ localSums[localId] += localSums[localId + stride];
+ }
+ context.localBarrier();
+ }
+ float gateScore = localSums[0];
+ float sharedWeight =
+ 1.0f / (1.0f + TornadoMath.exp(-gateScore));
+
+ for (int index = localId;
+ index < dim;
+ index += localWorkGroupSize) {
+ residual.set(index,
+ residual.get(index) + sharedWeight * sharedOutput.get(index));
+ }
+ }
+}
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
new file mode 100644
index 00000000..0a439d80
--- /dev/null
+++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java
@@ -0,0 +1,261 @@
+package org.beehive.gpullama3.tornadovm.layers.type.q8_0;
+
+import org.beehive.gpullama3.inference.state.Qwen2MoEState;
+import org.beehive.gpullama3.inference.weights.tornado.Qwen2MoETornadoWeights;
+import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration;
+import org.beehive.gpullama3.tornadovm.kernels.Qwen2Kernels;
+import org.beehive.gpullama3.tornadovm.kernels.Qwen2MoEKernels;
+import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels;
+import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered;
+import org.beehive.gpullama3.tornadovm.layers.AbstractTransformerLayerTaskGraphs;
+import org.beehive.gpullama3.tornadovm.scheduling.SchedulerType;
+import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory;
+import uk.ac.manchester.tornado.api.GridScheduler;
+import uk.ac.manchester.tornado.api.TaskGraph;
+import uk.ac.manchester.tornado.api.WorkerGrid;
+import uk.ac.manchester.tornado.api.WorkerGrid1D;
+import uk.ac.manchester.tornado.api.WorkerGrid2D;
+import uk.ac.manchester.tornado.api.enums.DataTransferMode;
+
+/**
+ * Single-token Q8_0 TaskGraphs for Qwen2-MoE / Qwen1.5-MoE.
+ *
+ * The attention block follows Qwen2. Its dense FFN is replaced by the
+ * routed-expert pipeline: normalize, route, choose top-K experts, execute each
+ * selected expert, and accumulate its weighted output into {@code wrapX}.
+ */
+public final class Qwen2MoEQ8_0FFNLayers
+ extends AbstractTransformerLayerTaskGraphs {
+
+ private final Qwen2MoEState moeState;
+
+ public Qwen2MoEQ8_0FFNLayers(String taskGraphName,
+ Qwen2MoEState state,
+ Qwen2MoETornadoWeights weights,
+ Qwen2MoEConfiguration config,
+ SchedulerType schedulerType) {
+ super(taskGraphName, state, weights, config, schedulerType);
+ this.moeState = state;
+ setupFFNLayers();
+ }
+
+ /** Sets the GPU worker grid for each task in each Transformer layer. */
+ @Override
+ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
+ WorkerGrid rmsNormWorker = WorkerGridFactory.createRmsNormWorker(
+ moeState.localSize, moeState.localSize);
+
+ WorkerGrid qkvWorker = workerForRows(config.dim() + 2 * config.kvDim());
+ WorkerGrid qkvBiasWorker = new WorkerGrid1D(config.dim());
+ qkvBiasWorker.setGlobalWork(config.dim(), 1, 1);
+ qkvBiasWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1);
+
+ WorkerGrid ropeWorker = new WorkerGrid2D(config.numberOfHeads(), config.headSize() / 2);
+ ropeWorker.setGlobalWork(config.numberOfHeads(), config.headSize() / 2, 1);
+ ropeWorker.setLocalWork(1, 1, 1);
+
+ int attentionLocalSize = Math.min(config.headSize(), 64);
+ WorkerGrid attentionWorker = new WorkerGrid1D(config.numberOfHeads() * attentionLocalSize);
+ attentionWorker.setLocalWork(attentionLocalSize, 1, 1);
+
+ WorkerGrid dimElementWorker = new WorkerGrid1D(config.dim());
+ dimElementWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1);
+ WorkerGrid dimWorker = workerForRows(config.dim());
+ 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());
+ WorkerGrid sharedHiddenWorker = workerForRows(config.sharedExpertHiddenDim());
+
+ for (int layer = 0; layer < config.numberOfLayers(); layer++) {
+ String prefix = "layer_" + layer + ".";
+ scheduler.addWorkerGrid(prefix + "attn_rms_reduce", rmsNormWorker);
+ scheduler.addWorkerGrid(prefix + "attn_rms_qkv_projection", qkvWorker);
+ scheduler.addWorkerGrid(prefix + "fused_qkv_bias", qkvBiasWorker);
+ scheduler.addWorkerGrid(prefix + "rope_and_kv_cache", ropeWorker);
+ scheduler.addWorkerGrid(prefix + "attention", attentionWorker);
+ scheduler.addWorkerGrid(prefix + "attn_output_proj", dimWorker);
+ scheduler.addWorkerGrid(prefix + "ffn_rms_reduce", rmsNormWorker);
+ 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 + "shared_expert_gate_up", sharedHiddenWorker);
+ scheduler.addWorkerGrid(prefix + "shared_expert_down", dimWorker);
+ scheduler.addWorkerGrid(prefix + "shared_expert_gate_and_accumulate", topKWorker);
+ }
+ return scheduler;
+ }
+
+ /** Creates a grid where one GPU work-group computes one output row. */
+ private WorkerGrid workerForRows(int rows) {
+ WorkerGrid worker = new WorkerGrid1D(rows * LOCAL_WORK_GROUP_SIZE_ALLOC);
+ worker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1);
+ return worker;
+ }
+
+ /**
+ * Creates the complete GPU TaskGraph for one Transformer layer.
+ * {@code layerIndex} selects that layer's weights.
+ */
+ @Override
+ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) {
+ TaskGraph layer = new TaskGraph("layer_" + layerIndex);
+ // Reuse wrapX produced by the previous TaskGraph on the GPU.
+ layer.consumeFromDevice(moeState.wrapX);
+ // Upload this layer's read-only weights from CPU to GPU on the first execution.
+ layer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
+ weights.rms_att_weightLayered[layerIndex].asFloatArray(),
+ weights.wqLayered[layerIndex].asByteArray(),
+ weights.wkLayered[layerIndex].asByteArray(),
+ weights.wvLayered[layerIndex].asByteArray(),
+ weights.woLayered[layerIndex].asByteArray(),
+ weights.q_biasLayered[layerIndex].asFloatArray(),
+ weights.k_biasLayered[layerIndex].asFloatArray(),
+ weights.v_biasLayered[layerIndex].asFloatArray(),
+ weights.rms_ffn_weightLayered[layerIndex].asFloatArray(),
+ weights.routerGateLayered[layerIndex].asFloatArray(),
+ weights.gateExpertsLayered[layerIndex].asByteArray(),
+ weights.upExpertsLayered[layerIndex].asByteArray(),
+ weights.downExpertsLayered[layerIndex].asByteArray(),
+ weights.sharedGateLayered[layerIndex].asByteArray(),
+ weights.sharedUpLayered[layerIndex].asByteArray(),
+ weights.sharedDownLayered[layerIndex].asByteArray(),
+ weights.sharedGateInputLayered[layerIndex].asFloatArray());
+ layer = configureLayerDataTransfers(layer, layerIndex);
+
+ configureAttention(layer, layerIndex);
+ configureRoutedExperts(layer, layerIndex);
+ layer.persistOnDevice(moeState.wrapX);
+ return layer;
+ }
+
+ /** Adds the normal Qwen2 attention tasks to this layer's TaskGraph. */
+ private void configureAttention(TaskGraph layer, int layerIndex) {
+ layer.task("attn_rms_reduce",
+ TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup,
+ context, moeState.temp, moeState.wrapX,
+ config.dim(), config.rmsNormEps(), moeState.localSize);
+
+ layer.task("attn_rms_qkv_projection",
+ Qwen3Kernels::fusedRmsNormQKVMatmulQ8_0,
+ context, moeState.wrapX, moeState.wrapQ, moeState.wrapK, moeState.wrapV,
+ weights.rms_att_weightLayered[layerIndex].asFloatArray(), moeState.temp,
+ weights.wqLayered[layerIndex].asByteArray(),
+ weights.wkLayered[layerIndex].asByteArray(),
+ weights.wvLayered[layerIndex].asByteArray(),
+ config.dim(), config.dim(), config.kvDim(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+ layer.task("fused_qkv_bias",
+ TransformerComputeKernelsLayered::fusedQKvBiasAddition,
+ context, moeState.wrapQ, moeState.wrapK,
+ weights.q_biasLayered[layerIndex].asFloatArray(), moeState.wrapV,
+ weights.k_biasLayered[layerIndex].asFloatArray(),
+ weights.v_biasLayered[layerIndex].asFloatArray(),
+ config.dim(), config.kvDim());
+
+ layer.task("rope_and_kv_cache", Qwen3Kernels::ropeRotationWithCacheCopy,
+ context, moeState.positionHolder, moeState.wrapQ, moeState.wrapK, moeState.wrapV,
+ moeState.wrapKeyCache, moeState.wrapValueCache,
+ config.numberOfKeyValueHeads(), config.headSize(), config.kvDim(),
+ layerIndex, config.contextLength());
+
+ layer.task("attention", Qwen2Kernels::processHeadsFlashAttention,
+ context, moeState.wrapQ, moeState.wrapKeyCache, moeState.wrapValueCache,
+ moeState.wrapXb, config.numberOfHeads(), config.headSize(), config.kvDim(),
+ config.kvMul(), moeState.positionHolder, layerIndex, config.contextLength());
+
+ layer.task("attn_output_proj",
+ TransformerComputeKernelsLayered::matrixVectorGenericWithResidualQ8_0Byte,
+ context, moeState.wrapXb, moeState.wrapX,
+ weights.woLayered[layerIndex].asByteArray(),
+ config.dim(), config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+ }
+
+ /**
+ * Adds router, top-K, and selected-expert FFN tasks to this layer's TaskGraph.
+ * Their weighted outputs are added to the residual vector.
+ */
+ private void configureRoutedExperts(TaskGraph layer, int layerIndex) {
+ layer.task("ffn_rms_reduce",
+ TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup,
+ context, moeState.tempFFN, moeState.wrapX,
+ config.dim(), config.rmsNormEps(), moeState.localSize);
+
+ layer.task("ffn_rms_apply",
+ TransformerComputeKernelsLayered::reductionOneBlock2WithLayer,
+ context, moeState.wrapXb, moeState.wrapX,
+ weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), moeState.tempFFN);
+
+ layer.task("router_projection",
+ TransformerComputeKernelsLayered::matrixVectorGeneric,
+ context, moeState.wrapXb, moeState.wrapRouterLogits,
+ weights.routerGateLayered[layerIndex].asFloatArray(),
+ config.dim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+ layer.task("router_softmax_topk", Qwen2MoEKernels::softmaxAndTopK,
+ 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);
+ }
+
+ // The shared expert always runs; it does not depend on router top-K selection.
+ layer.task("shared_expert_gate_up", Qwen2MoEKernels::sharedExpertGateUpSwiGLUQ8_0,
+ context, moeState.wrapXb,
+ weights.sharedGateLayered[layerIndex].asByteArray(),
+ weights.sharedUpLayered[layerIndex].asByteArray(), moeState.wrapSharedGate,
+ config.dim(), config.sharedExpertHiddenDim(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+ layer.task("shared_expert_down", Qwen2MoEKernels::sharedExpertDownProjectQ8_0,
+ context, moeState.wrapSharedGate,
+ weights.sharedDownLayered[layerIndex].asByteArray(), moeState.wrapSharedOutput,
+ config.dim(), config.sharedExpertHiddenDim(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+ layer.task("shared_expert_gate_and_accumulate", Qwen2MoEKernels::sharedExpertGateAndAccumulate,
+ context, moeState.wrapXb, weights.sharedGateInputLayered[layerIndex].asFloatArray(),
+ moeState.wrapSharedOutput, moeState.wrapX, config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC);
+
+ }
+
+ /**
+ * Configures which TaskGraph data is uploaded from the CPU or reused on the GPU.
+ */
+ @Override
+ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) {
+ if (layerIndex == 0) {
+ layer.transferToDevice(DataTransferMode.EVERY_EXECUTION,
+ moeState.positionHolder, moeState.temp, moeState.tempFFN);
+ layer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
+ context, moeState.wrapXb, moeState.wrapXb2, moeState.wrapQ,
+ moeState.wrapK, moeState.wrapV, moeState.wrapKeyCache,
+ moeState.wrapValueCache, moeState.wrapAtt, moeState.wrapRouterLogits,
+ moeState.wrapSelectedExperts, moeState.wrapRoutingWeights,
+ moeState.wrapExpertGate, moeState.wrapSharedGate, moeState.wrapSharedOutput);
+ } else {
+ layer.consumeFromDevice(context, moeState.wrapXb, moeState.wrapXb2,
+ moeState.wrapQ, moeState.wrapK, moeState.wrapV, moeState.wrapKeyCache,
+ moeState.wrapValueCache, moeState.wrapAtt, moeState.wrapRouterLogits,
+ moeState.wrapSelectedExperts, moeState.wrapRoutingWeights,
+ moeState.wrapExpertGate, moeState.wrapSharedGate, moeState.wrapSharedOutput,
+ moeState.positionHolder);
+ }
+ return layer;
+ }
+}
diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java
index 042dd354..6e8ba29d 100644
--- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java
+++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java
@@ -4,6 +4,7 @@
import org.beehive.gpullama3.inference.state.GraniteState;
import org.beehive.gpullama3.inference.state.LlamaState;
import org.beehive.gpullama3.inference.state.Phi3State;
+import org.beehive.gpullama3.inference.state.Qwen2MoEState;
import org.beehive.gpullama3.inference.state.Qwen2State;
import org.beehive.gpullama3.inference.state.Qwen3State;
import org.beehive.gpullama3.inference.state.State;
@@ -24,6 +25,7 @@
import org.beehive.gpullama3.tornadovm.plan.components.q8_0.MistralQ8_0PlanComponents;
import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Phi3Q8_0PlanComponents;
import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen2Q8_0PlanComponents;
+import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen2MoEQ8_0PlanComponents;
import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen3Q8_0PlanComponents;
// @formatter:off
@@ -109,6 +111,7 @@ private static ForwardPlan createQ8_0Plan(ExecutionMode mode, State state, Model
case MISTRAL -> createMistralQ8_0Plan(mode, (LlamaState) state, model);
case DEVSTRAL_2 -> createDevstralQ8_0Plan(mode, (DevstralState) state, model);
case QWEN_2 -> createQwen2Q8_0Plan(mode, (Qwen2State) state, model);
+ case QWEN_2_MOE -> createQwen2MoEQ8_0Plan(mode, (Qwen2MoEState) state, model);
case QWEN_3 -> createQwen3Q8_0Plan(mode, (Qwen3State) state, model);
case PHI_3 -> createPhi3Q8_0Plan(mode, (Phi3State) state, model);
case GRANITE -> createGraniteQ8_0Plan(mode, (GraniteState) state, model);
@@ -175,6 +178,12 @@ private static ForwardPlan createQwen2Q8_0Plan(ExecutionMode mode, Qwen2State st
return new SingleTokenForwardPlan(model, new Qwen2Q8_0PlanComponents(state, model));
}
+ private static ForwardPlan createQwen2MoEQ8_0Plan(ExecutionMode mode, Qwen2MoEState state, Model model) {
+ if (mode != ExecutionMode.STANDARD)
+ throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2_MOE + Q8_0");
+ return new SingleTokenForwardPlan(model, new Qwen2MoEQ8_0PlanComponents(state, model));
+ }
+
private static ForwardPlan createQwen3FP16Plan(ExecutionMode mode, Qwen3State state, Model model) {
BatchPrefillDecodeForwardPlanComponents components = new Qwen3FP16PlanComponents(state, model);
return switch (mode) {
diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java
new file mode 100644
index 00000000..93a1443f
--- /dev/null
+++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java
@@ -0,0 +1,46 @@
+package org.beehive.gpullama3.tornadovm.plan.components.q8_0;
+
+import org.beehive.gpullama3.inference.state.Qwen2MoEState;
+import org.beehive.gpullama3.inference.weights.tornado.Qwen2MoETornadoWeights;
+import org.beehive.gpullama3.model.Model;
+import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration;
+import org.beehive.gpullama3.tornadovm.layers.AbstractLogitsTaskGraph;
+import org.beehive.gpullama3.tornadovm.layers.Activation;
+import org.beehive.gpullama3.tornadovm.layers.ActivationTaskGraph;
+import org.beehive.gpullama3.tornadovm.layers.TransformerLayerTaskGraphs;
+import org.beehive.gpullama3.tornadovm.layers.type.q8_0.LogitsQ8_0Layer;
+import org.beehive.gpullama3.tornadovm.layers.type.q8_0.Qwen2MoEQ8_0FFNLayers;
+import org.beehive.gpullama3.tornadovm.plan.components.SingleTokenForwardPlanComponents;
+import org.beehive.gpullama3.tornadovm.scheduling.SchedulerDetectionService;
+import org.beehive.gpullama3.tornadovm.scheduling.SchedulerType;
+
+/** Assembles the single-token Q8_0 GPU components for Qwen2-MoE. */
+public final class Qwen2MoEQ8_0PlanComponents implements SingleTokenForwardPlanComponents {
+
+ private final Qwen2MoEState state;
+ private final Qwen2MoETornadoWeights weights;
+ private final Qwen2MoEConfiguration config;
+ private final SchedulerType schedulerType;
+
+ public Qwen2MoEQ8_0PlanComponents(Qwen2MoEState state, Model model) {
+ this.state = state;
+ this.config = (Qwen2MoEConfiguration) model.configuration();
+ this.weights = (Qwen2MoETornadoWeights) model.weights();
+ this.schedulerType = SchedulerDetectionService.determineSchedulerType(model);
+ }
+
+ @Override
+ public ActivationTaskGraph singleTokenActivation() {
+ return new Activation("activationUpdate", state, weights, config);
+ }
+
+ @Override
+ public TransformerLayerTaskGraphs singleTokenTransformerLayers() {
+ return new Qwen2MoEQ8_0FFNLayers("qwen2MoEFFN", state, weights, config, schedulerType);
+ }
+
+ @Override
+ public AbstractLogitsTaskGraph singleTokenLogits(String previousGraphId) {
+ return new LogitsQ8_0Layer("logits", state, weights, config, previousGraphId, schedulerType);
+ }
+}