From a5889a6b10d0146b574a31669c00cec0271b7861 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Fri, 14 Aug 2026 16:24:13 +0100 Subject: [PATCH 1/4] Add Qwen2-MoE batch prefill and decode --- .../InferenceCoreBatchPrefillDecode.java | 291 ++++---- ...InferenceEngineWithBatchPrefillDecode.java | 38 +- .../inference/state/Qwen2MoEState.java | 77 ++- .../gpullama3/model/qwen2/Qwen2MoE.java | 77 ++- .../kernels/Qwen2MoEBatchKernels.java | 634 ++++++++++++++++++ .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 331 ++++++--- .../decode/Qwen2MoEQ8_0FFNLayersDecode.java | 105 +++ .../Qwen2MoEQ8_0LayersBatchPrefill.java | 493 ++++++++++++++ .../tornadovm/plan/ForwardPlanFactory.java | 161 +++-- .../q8_0/Qwen2MoEQ8_0PlanComponents.java | 51 +- 10 files changed, 1954 insertions(+), 304 deletions(-) create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/decode/Qwen2MoEQ8_0FFNLayersDecode.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java index c1458e42..ffed2941 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java @@ -1,6 +1,7 @@ package org.beehive.gpullama3.inference; import org.beehive.gpullama3.auxiliary.Parallel; +import org.beehive.gpullama3.inference.state.Qwen2MoEState; import org.beehive.gpullama3.inference.state.State; import org.beehive.gpullama3.inference.weights.standard.StandardWeights; import org.beehive.gpullama3.inference.weights.tornado.TornadoWeights; @@ -9,6 +10,7 @@ import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode; + import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import java.lang.foreign.MemorySegment; @@ -16,50 +18,45 @@ /** * Low-level forward passes for the batched prefill/decode inference path (Phase 3/4). * - *

Parallel to {@link InferenceCoreWithPrefillDecode} — does NOT modify it.

+ *

Parallel to {@link InferenceCoreWithPrefillDecode} — does NOT modify it. + * + *

Provides three operations: * - *

Provides three operations:

* */ public final class InferenceCoreBatchPrefillDecode { - private InferenceCoreBatchPrefillDecode() { - } + private InferenceCoreBatchPrefillDecode() {} /** * CPU batched prefill forward pass for LLaMA (Phase 3). * - *

Processes {@code batchSize} prompt tokens simultaneously through all - * transformer layers. For each layer, Q/K/V projections, output projection, - * and FFN projections are computed via batch matmul - * ({@link FloatTensor#matmul(int, FloatTensor[], FloatTensor[], int, int)}), - * which parallelises over both output dimension and batch simultaneously. - * Attention reuses {@code state.att} sequentially per token (parallel per - * head within each token), keeping memory overhead minimal.

+ *

Processes {@code batchSize} prompt tokens simultaneously through all transformer layers. + * For each layer, Q/K/V projections, output projection, and FFN projections are computed via + * batch matmul ({@link FloatTensor#matmul(int, FloatTensor[], FloatTensor[], int, int)}), which + * parallelises over both output dimension and batch simultaneously. Attention reuses {@code + * state.att} sequentially per token (parallel per head within each token), keeping memory + * overhead minimal. * - *

The logits layer is intentionally omitted — only the KV cache matters - * for prefill positions.

+ *

The logits layer is intentionally omitted — only the KV cache matters for prefill + * positions. * - * @param model - * the LLaMA model (must carry {@link StandardWeights}) - * @param state - * mutable inference state (KV cache, att buffer …) - * @param tokens - * input token ids, {@code tokens[b]} at position {@code startPos+b} - * @param startPos - * sequence position of {@code tokens[0]} - * @param batchSize - * number of tokens in this chunk ({@code tokens.length}) + * @param model the LLaMA model (must carry {@link StandardWeights}) + * @param state mutable inference state (KV cache, att buffer …) + * @param tokens input token ids, {@code tokens[b]} at position {@code startPos+b} + * @param startPos sequence position of {@code tokens[0]} + * @param batchSize number of tokens in this chunk ({@code tokens.length}) */ - public static void batchForwardJavaPrefill(Model model, State state, int[] tokens, int startPos, int batchSize) { + public static void batchForwardJavaPrefill( + Model model, State state, int[] tokens, int startPos, int batchSize) { final Configuration config = model.configuration(); final StandardWeights weights = (StandardWeights) model.weights(); int dim = config.dim(); @@ -89,75 +86,117 @@ public static void batchForwardJavaPrefill(Model model, State state, int[] token } // ── Token embeddings ────────────────────────────────────────────────── - Parallel.parallelFor(0, batchSize, b -> weights.token_embedding_table.copyTo(tokens[b] * dim, x[b], 0, dim)); + Parallel.parallelFor( + 0, + batchSize, + b -> weights.token_embedding_table.copyTo(tokens[b] * dim, x[b], 0, dim)); // ── Transformer layers ──────────────────────────────────────────────── for (int l = 0; l < config.numberOfLayers(); l++) { final int layer = l; - Parallel.parallelFor(0, batchSize, b -> InferenceCore.rmsnorm(xb[b], x[b], weights.rms_att_weight[layer], 0, dim, config.rmsNormEps())); + Parallel.parallelFor( + 0, + batchSize, + b -> + InferenceCore.rmsnorm( + xb[b], + x[b], + weights.rms_att_weight[layer], + 0, + dim, + config.rmsNormEps())); weights.wq[l].matmul(batchSize, xb, q, dim, dim); weights.wk[l].matmul(batchSize, xb, k, kvDim, dim); weights.wv[l].matmul(batchSize, xb, v, kvDim, dim); - Parallel.parallelFor(0, batchSize, b -> { - int pos = startPos + b; - for (int i = 0; i < dim; i += 2) { - int head_dim = i % headSize; - float fcr = weights.freq_cis_real.getFloat(pos * (headSize / 2) + (head_dim / 2)); - float fci = weights.freq_cis_imag.getFloat(pos * (headSize / 2) + (head_dim / 2)); - int rotn = i < kvDim ? 2 : 1; - for (int vv = 0; vv < rotn; vv++) { - FloatTensor vec = vv == 0 ? q[b] : k[b]; - float v0 = vec.getFloat(i); - float v1 = vec.getFloat(i + 1); - vec.setFloat(i, v0 * fcr - v1 * fci); - vec.setFloat(i + 1, v0 * fci + v1 * fcr); - } - } - k[b].copyTo(0, state.keyCache[layer], pos * kvDim, kvDim); - v[b].copyTo(0, state.valueCache[layer], pos * kvDim, kvDim); - }); + Parallel.parallelFor( + 0, + batchSize, + b -> { + int pos = startPos + b; + for (int i = 0; i < dim; i += 2) { + int head_dim = i % headSize; + float fcr = + weights.freq_cis_real.getFloat( + pos * (headSize / 2) + (head_dim / 2)); + float fci = + weights.freq_cis_imag.getFloat( + pos * (headSize / 2) + (head_dim / 2)); + int rotn = i < kvDim ? 2 : 1; + for (int vv = 0; vv < rotn; vv++) { + FloatTensor vec = vv == 0 ? q[b] : k[b]; + float v0 = vec.getFloat(i); + float v1 = vec.getFloat(i + 1); + vec.setFloat(i, v0 * fcr - v1 * fci); + vec.setFloat(i + 1, v0 * fci + v1 * fcr); + } + } + k[b].copyTo(0, state.keyCache[layer], pos * kvDim, kvDim); + v[b].copyTo(0, state.valueCache[layer], pos * kvDim, kvDim); + }); for (int b = 0; b < batchSize; b++) { final int pos_b = startPos + b; final int bFinal = b; - Parallel.parallelFor(0, config.numberOfHeads(), h -> { - int qOffset = h * headSize; - int attOffset = h * config.contextLength(); + Parallel.parallelFor( + 0, + config.numberOfHeads(), + h -> { + int qOffset = h * headSize; + int attOffset = h * config.contextLength(); - for (int t = 0; t <= pos_b; t++) { - int keyCacheOffset = t * kvDim + (h / kvMul) * headSize; - float score = q[bFinal].dot(qOffset, state.keyCache[layer], keyCacheOffset, headSize) / sqrtHeadSize; - state.att.setFloat(attOffset + t, score); - } - state.att.softmaxInPlace(attOffset, pos_b + 1); + for (int t = 0; t <= pos_b; t++) { + int keyCacheOffset = t * kvDim + (h / kvMul) * headSize; + float score = + q[bFinal].dot( + qOffset, + state.keyCache[layer], + keyCacheOffset, + headSize) + / sqrtHeadSize; + state.att.setFloat(attOffset + t, score); + } + state.att.softmaxInPlace(attOffset, pos_b + 1); - int xbOffset = h * headSize; - xb[bFinal].fillInPlace(xbOffset, headSize, 0f); - for (int t = 0; t <= pos_b; t++) { - int vOffset = t * kvDim + (h / kvMul) * headSize; - float a = state.att.getFloat(attOffset + t); - xb[bFinal].saxpyInPlace(xbOffset, state.valueCache[layer], vOffset, headSize, a); - } - }); + int xbOffset = h * headSize; + xb[bFinal].fillInPlace(xbOffset, headSize, 0f); + for (int t = 0; t <= pos_b; t++) { + int vOffset = t * kvDim + (h / kvMul) * headSize; + float a = state.att.getFloat(attOffset + t); + xb[bFinal].saxpyInPlace( + xbOffset, state.valueCache[layer], vOffset, headSize, a); + } + }); } weights.wo[l].matmul(batchSize, xb, xb2, dim, dim); - Parallel.parallelFor(0, batchSize, b -> { - x[b].addInPlace(xb2[b]); - InferenceCore.rmsnorm(xb[b], x[b], weights.rms_ffn_weight[layer], 0, dim, config.rmsNormEps()); - }); + Parallel.parallelFor( + 0, + batchSize, + b -> { + x[b].addInPlace(xb2[b]); + InferenceCore.rmsnorm( + xb[b], + x[b], + weights.rms_ffn_weight[layer], + 0, + dim, + config.rmsNormEps()); + }); weights.w1[l].matmul(batchSize, xb, hb, config.hiddenDim(), dim); weights.w3[l].matmul(batchSize, xb, hb2, config.hiddenDim(), dim); - Parallel.parallelFor(0, batchSize, b -> { - hb[b].mapInPlace(value -> value / (float) (1.0 + Math.exp(-value))); - hb[b].multiplyInPlace(hb2[b]); - }); + Parallel.parallelFor( + 0, + batchSize, + b -> { + hb[b].mapInPlace(value -> value / (float) (1.0 + Math.exp(-value))); + hb[b].multiplyInPlace(hb2[b]); + }); weights.w2[l].matmul(batchSize, hb, xb, dim, config.hiddenDim()); @@ -173,34 +212,43 @@ public static void batchForwardJavaPrefill(Model model, State state, int[] token /** * GPU batched prefill forward pass (Phase 4). * - *

Copies {@code chunkSize} token embeddings into device-visible state buffers, - * then delegates graph execution to the plan.

+ *

Copies {@code chunkSize} token embeddings into device-visible state buffers, then + * delegates graph execution to the plan. * - * @param model - * the LLaMA model - * @param state - * mutable inference state - * @param tokens - * token ids for this chunk - * @param startPos - * sequence position of {@code tokens[0]} - * @param chunkSize - * number of tokens in this chunk - * @param plan - * the batched prefill/decode GPU plan + * @param model the LLaMA model + * @param state mutable inference state + * @param tokens token ids for this chunk + * @param startPos sequence position of {@code tokens[0]} + * @param chunkSize number of tokens in this chunk + * @param plan the batched prefill/decode GPU plan */ - public static void batchForwardTornadoVMPrefill(Model model, State state, int[] tokens, int startPos, int chunkSize, TornadoVMMasterPlanBatchPrefillDecode plan) { + public static void batchForwardTornadoVMPrefill( + Model model, + State state, + int[] tokens, + int startPos, + int chunkSize, + TornadoVMMasterPlanBatchPrefillDecode plan) { final Configuration config = model.configuration(); final TornadoWeights weights = (TornadoWeights) model.weights(); state.batchStartPosHolder.set(0, startPos); + if (state instanceof Qwen2MoEState moeState && moeState.activeBatchSizeHolder != null) { + moeState.activeBatchSizeHolder.set(0, chunkSize); + } switch (weights.getWeightType()) { case F16 -> { - MemorySegment embTable = weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); + MemorySegment embTable = + weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); long dimBytes = (long) config.dim() * Short.BYTES; for (int b = 0; b < chunkSize; b++) { - MemorySegment.copy(embTable, (long) tokens[b] * dimBytes, state.embeddingXBatch.getSegment(), (long) b * dimBytes, dimBytes); + MemorySegment.copy( + embTable, + (long) tokens[b] * dimBytes, + state.embeddingXBatch.getSegment(), + (long) b * dimBytes, + dimBytes); } } case Q8_0 -> { @@ -210,14 +258,17 @@ public static void batchForwardTornadoVMPrefill(Model model, State state, int[] for (int b = 0; b < chunkSize; b++) { int tokenId = tokens[b]; for (int j = 0; j < dim; j++) { - int blockByteOffset = (tokenId * blocksPerRow + j / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + int blockByteOffset = + (tokenId * blocksPerRow + j / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; float scale = embTable.getHalfFloat(blockByteOffset).getFloat32(); float quant = embTable.get(blockByteOffset + 2 + j % Q8_0_BLOCK_SIZE); state.wrapXBatch.set(b * dim + j, quant * scale); } } } - default -> throw new IllegalArgumentException("Unsupported weight type: " + weights.getWeightType()); + default -> + throw new IllegalArgumentException( + "Unsupported weight type: " + weights.getWeightType()); } plan.tornadoVMForwardBatchPrefill(); @@ -226,37 +277,51 @@ public static void batchForwardTornadoVMPrefill(Model model, State state, int[] /** * GPU decode forward pass (Phase 4). * - *

Copies the token embedding into device-visible state, then delegates - * graph execution to the plan.

+ *

Copies the token embedding into device-visible state, then delegates graph execution to + * the plan. * - * @param model - * the LLaMA model - * @param state - * mutable inference state - * @param token - * current token id - * @param position - * sequence position - * @param plan - * the batched prefill/decode GPU plan + * @param model the LLaMA model + * @param state mutable inference state + * @param token current token id + * @param position sequence position + * @param plan the batched prefill/decode GPU plan * @return logits array for token sampling */ - public static FloatArray forwardTornadoVMDecode(Model model, State state, int token, int position, TornadoVMMasterPlanBatchPrefillDecode plan) { + public static FloatArray forwardTornadoVMDecode( + Model model, + State state, + int token, + int position, + TornadoVMMasterPlanBatchPrefillDecode plan) { final Configuration config = model.configuration(); final TornadoWeights weights = (TornadoWeights) model.weights(); switch (weights.getWeightType()) { case F16 -> { - MemorySegment embTable = weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); - MemorySegment.copy(embTable, (long) token * config.dim() * Short.BYTES, state.embeddingX.getSegment(), 0L, (long) config.dim() * Short.BYTES); + MemorySegment embTable = + weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); + MemorySegment.copy( + embTable, + (long) token * config.dim() * Short.BYTES, + state.embeddingX.getSegment(), + 0L, + (long) config.dim() * Short.BYTES); } case Q8_0 -> { - MemorySegment embTable = weights.getTokenEmbeddingTable().asByteArray().getSegment(); + MemorySegment embTable = + weights.getTokenEmbeddingTable().asByteArray().getSegment(); int blocksPerToken = (config.dim() + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; long bytesPerToken = (long) blocksPerToken * Q8_0_BLOCK_BYTES; - MemorySegment.copy(embTable, (long) token * bytesPerToken, state.embeddingX.getSegment(), 0L, bytesPerToken); + MemorySegment.copy( + embTable, + (long) token * bytesPerToken, + state.embeddingX.getSegment(), + 0L, + bytesPerToken); } - default -> throw new IllegalArgumentException("Unsupported weight type: " + weights.getWeightType()); + default -> + throw new IllegalArgumentException( + "Unsupported weight type: " + weights.getWeightType()); } return plan.tornadoVMForwardDecode(position); diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java index 4493340b..76f1e9bf 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java @@ -187,17 +187,30 @@ public static List generateTokensGPULlama(Model model, int N = promptTokens.size(); // ── Prefill ─────────────────────────────────────────────────────────── - // Build the token sequence at positions [startPosition .. startPosition+N-1]: - // position startPosition+0 : currentToken (BOS/previous token) - // position startPosition+k : promptTokens[k-1] - int[] prefillSeq = new int[N]; - prefillSeq[0] = currentToken; - for (int i = 1; i < N; i++) { - prefillSeq[i] = promptTokens.get(i - 1); + // Qwen's regular path forwards promptTokens[0] directly at position 0. + // Keep the final prompt token for the B1 decode graph, which produces the + // first generation logits without duplicating the ChatML start token. + boolean qwen2MoE = model.getModelType() + == org.beehive.gpullama3.model.ModelType.QWEN_2_MOE; + int prefillTokenCount = qwen2MoE ? Math.max(0, N - 1) : N; + int[] prefillSeq = new int[prefillTokenCount]; + if (qwen2MoE) { + for (int i = 0; i < prefillTokenCount; i++) { + prefillSeq[i] = promptTokens.get(i); + } + } else { + prefillSeq[0] = currentToken; + for (int i = 1; i < N; i++) { + prefillSeq[i] = promptTokens.get(i - 1); + } } - for (int chunkStart = 0; chunkStart < N && pos + chunkStart < actualMaxTokens; chunkStart += batchSize) { - int chunkEnd = Math.min(Math.min(chunkStart + batchSize, N), actualMaxTokens - pos); + for (int chunkStart = 0; + chunkStart < prefillTokenCount && pos + chunkStart < actualMaxTokens; + chunkStart += batchSize) { + int chunkEnd = Math.min( + Math.min(chunkStart + batchSize, prefillTokenCount), + actualMaxTokens - pos); int chunkSize = chunkEnd - chunkStart; int[] chunk = Arrays.copyOfRange(prefillSeq, chunkStart, chunkEnd); @@ -213,12 +226,15 @@ public static List generateTokensGPULlama(Model model, } currentToken = promptTokens.get(N - 1); - pos = startPosition + N; + pos = startPosition + (qwen2MoE ? N - 1 : N); state.latestToken = currentToken; long decodeStartNanos = System.nanoTime(); + int generatedTokenBudget = qwen2MoE + ? Math.max(0, actualMaxTokens - N) + : Integer.MAX_VALUE; // ── Decode ──────────────────────────────────────────────────────────── - while (pos < actualMaxTokens) { + while (pos < actualMaxTokens && generatedTokens.size() < generatedTokenBudget) { var logits = InferenceCoreBatchPrefillDecode.forwardTornadoVMDecode(model, state, currentToken, pos, plan); int nextToken = sampler.sampleToken(logits); 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 d7bdc92d..5ce078c3 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -4,6 +4,7 @@ import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; + import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; @@ -40,6 +41,20 @@ public class Qwen2MoEState extends Qwen2State { public final FloatArray wrapSharedGate; public final FloatArray wrapSharedOutput; + // TornadoVM buffers for the batch-prefill MoE path. + // Their shapes use the configured maximum batch size so TaskGraphs stay fixed. + public final FloatArray wrapRouterLogitsBatch; + public final IntArray activeBatchSizeHolder; + public final IntArray wrapSelectedExpertsBatch; + public final FloatArray wrapRoutingWeightsBatch; + public final IntArray wrapGroupedAssignmentIds; + public final IntArray wrapGroupedPositionByAssignment; + public final IntArray wrapExpertOffsets; + public final FloatArray wrapGroupedExpertHidden; + public final FloatArray wrapGroupedExpertDown; + public final FloatArray wrapSharedHiddenBatch; + public final FloatArray wrapSharedWeightBatch; + public Qwen2MoEState(Configuration config, int batchsize) { super(config, batchsize); Qwen2MoEConfiguration c = (Qwen2MoEConfiguration) config; @@ -56,6 +71,35 @@ public Qwen2MoEState(Configuration config, int batchsize) { this.wrapExpertGate = new FloatArray(c.moeHiddenDim() * c.numberOfExpertsUsed()); this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim()); this.wrapSharedOutput = new FloatArray(c.dim()); + + int gpuBatchSize = Integer.getInteger("llama.prefillBatchSize", 1); + if (gpuBatchSize > 1) { + int assignments = gpuBatchSize * c.numberOfExpertsUsed(); + this.wrapRouterLogitsBatch = new FloatArray(gpuBatchSize * c.numberOfExperts()); + this.activeBatchSizeHolder = new IntArray(1); + this.activeBatchSizeHolder.init(gpuBatchSize); + this.wrapSelectedExpertsBatch = new IntArray(assignments); + this.wrapRoutingWeightsBatch = new FloatArray(assignments); + this.wrapGroupedAssignmentIds = new IntArray(assignments); + this.wrapGroupedPositionByAssignment = new IntArray(assignments); + this.wrapExpertOffsets = new IntArray(c.numberOfExperts() + 1); + this.wrapGroupedExpertHidden = new FloatArray(assignments * c.moeHiddenDim()); + this.wrapGroupedExpertDown = new FloatArray(assignments * c.dim()); + this.wrapSharedHiddenBatch = new FloatArray(gpuBatchSize * c.sharedExpertHiddenDim()); + this.wrapSharedWeightBatch = new FloatArray(gpuBatchSize); + } else { + this.wrapRouterLogitsBatch = null; + this.activeBatchSizeHolder = null; + this.wrapSelectedExpertsBatch = null; + this.wrapRoutingWeightsBatch = null; + this.wrapGroupedAssignmentIds = null; + this.wrapGroupedPositionByAssignment = null; + this.wrapExpertOffsets = null; + this.wrapGroupedExpertHidden = null; + this.wrapGroupedExpertDown = null; + this.wrapSharedHiddenBatch = null; + this.wrapSharedWeightBatch = null; + } } @Override @@ -77,17 +121,21 @@ protected StateFields createStateFields(Configuration configuration) { fields.att = ArrayFloatTensor.allocate(config.numberOfHeads(), config.contextLength()); fields.logits = ArrayFloatTensor.allocate(config.vocabularySize()); - fields.keyCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) - .limit(config.numberOfLayers()) - .toArray(FloatTensor[]::new); - fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) - .limit(config.numberOfLayers()) - .toArray(FloatTensor[]::new); + fields.keyCache = + Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) + .limit(config.numberOfLayers()) + .toArray(FloatTensor[]::new); + fields.valueCache = + Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) + .limit(config.numberOfLayers()) + .toArray(FloatTensor[]::new); switch (config.quantization()) { case "FP16" -> fields.createActivationFP16(config.dim()); case "Q8_0" -> fields.createActivationQ8_0(config.dim()); - default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization()); + default -> + throw new UnsupportedOperationException( + "Unsupported quantization format: " + config.quantization()); } fields.wrapX = new FloatArray(config.dim()); fields.wrapXb = new FloatArray(config.dim()); @@ -101,8 +149,10 @@ protected StateFields createStateFields(Configuration configuration) { fields.wrapK = new FloatArray(config.kvDim()); fields.wrapV = new FloatArray(config.kvDim()); - fields.wrapKeyCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); - fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapKeyCache = + new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapValueCache = + new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); fields.wrapValueCache.init(0.f); fields.wrapKeyCache.init(0.f); fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); @@ -110,9 +160,12 @@ protected StateFields createStateFields(Configuration configuration) { // State invokes this override before the Qwen2State constructor body runs, // so use the Qwen2 work-group size directly instead of State.localSize. - fields.temp = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); - fields.tempFFN = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); - fields.tempLogits = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.temp = + new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.tempFFN = + new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.tempLogits = + new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); return fields; } diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java index 0a013e52..beb70562 100644 --- a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java @@ -1,7 +1,10 @@ package org.beehive.gpullama3.model.qwen2; +import static org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan.WITH_PREFILL_DECODE; + import org.beehive.gpullama3.inference.InferenceCore; import org.beehive.gpullama3.inference.InferenceEngine; +import org.beehive.gpullama3.inference.InferenceEngineWithBatchPrefillDecode; import org.beehive.gpullama3.inference.sampler.Sampler; import org.beehive.gpullama3.inference.state.Qwen2MoEState; import org.beehive.gpullama3.inference.state.State; @@ -17,13 +20,15 @@ import java.util.Set; import java.util.function.IntConsumer; -import static org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan.WITH_PREFILL_DECODE; - public class Qwen2MoE extends AbstractModel { Qwen2MoEConfiguration configuration; - public Qwen2MoE(Qwen2MoEConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) { + public Qwen2MoE( + Qwen2MoEConfiguration configuration, + Tokenizer tokenizer, + Weights weights, + ChatFormat chatFormat) { super(tokenizer, weights, chatFormat, null); this.configuration = configuration; } @@ -45,14 +50,16 @@ public ModelType getModelType() { @Override public State createNewState() { State state = new Qwen2MoEState(configuration(), -1); - state.latestToken = tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); + state.latestToken = + tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); return state; } @Override public State createNewState(int batchsize) { State state = new Qwen2MoEState(configuration(), batchsize); - state.latestToken = tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); + state.latestToken = + tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); return state; } @@ -81,21 +88,65 @@ public void forward(State state, int token, int position) { } @Override - public List generateTokens(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, + public List generateTokens( + State state, + int startPosition, + List promptTokens, + Set stopTokens, + int maxTokens, + Sampler sampler, + boolean echo, IntConsumer onTokenGenerated) { - return InferenceEngine.generateTokensQwen3(this, state, startPosition, promptTokens, stopTokens, maxTokens, sampler, echo, onTokenGenerated); + return InferenceEngine.generateTokensQwen3( + this, + state, + startPosition, + promptTokens, + stopTokens, + maxTokens, + sampler, + echo, + onTokenGenerated); } @Override - public List generateTokensGPU(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, - IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) { + public List generateTokensGPU( + State state, + int startPosition, + List promptTokens, + Set 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"); + return InferenceEngineWithBatchPrefillDecode.generateTokensGPULlama( + this, + state, + startPosition, + promptTokens, + stopTokens, + maxTokens, + sampler, + echo, + onTokenGenerated, + tornadoVMPlan); } if (WITH_PREFILL_DECODE) { - throw new UnsupportedOperationException("Prefill/decode on GPU not yet implemented for Qwen2-MoE"); + 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); + return InferenceEngine.generateTokensGPUQwen3( + this, + state, + startPosition, + promptTokens, + stopTokens, + maxTokens, + sampler, + echo, + onTokenGenerated, + tornadoVMPlan); } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java new file mode 100644 index 00000000..eab69eac --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java @@ -0,0 +1,634 @@ +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; + +/** GPU kernels used by Qwen2-MoE batch prefill. */ +public final class Qwen2MoEBatchKernels { + + private static final int Q8_0_BLOCK_SIZE = 32; + private static final int Q8_0_BLOCK_BYTES = 34; + + private Qwen2MoEBatchKernels() {} + + /** Computes one router score for every token-expert pair. */ + public static void batchedRouterProjection( + KernelContext context, + FloatArray input, + FloatArray routerLogits, + FloatArray routerWeights, + IntArray activeBatchSizeHolder, + int dim, + int numberOfExperts, + int localWorkGroupSize) { + + int groupId = context.groupIdx; + int localId = context.localIdx; + + // One work-group handles one (token, expert) pair. + // For 60 experts, groups 0..59 process token 0, groups 60..119 + // process token 1, and so on. + + int token = groupId / numberOfExperts; + int expert = groupId % numberOfExperts; + if (token >= activeBatchSizeHolder.get(0)) { + return; + } + + int inputOffset = token * dim; + int weightOffset = expert * dim; + + // Each local thread computes a strided part of the dot product. + float partialSum = 0.0f; + + for (int column = localId; column < dim; column += localWorkGroupSize) { + + partialSum += + input.get(inputOffset + column) * routerWeights.get(weightOffset + column); + } + + // Store every thread's partial sum in local memory. + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + + localSums[localId] = partialSum; + context.localBarrier(); + + // Reduce the partial sums inside this work-group. + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + + context.localBarrier(); + } + + // Local thread 0 writes routerLogits[token, expert]. + if (localId == 0) { + int outputOffset = token * numberOfExperts + expert; + routerLogits.set(outputOffset, localSums[0]); + } + } + + /** Adds Qwen2's Q, K, and V biases independently for every active token. */ + public static void batchedQKVBias( + KernelContext context, + FloatArray qBatch, + FloatArray kBatch, + FloatArray vBatch, + FloatArray qBias, + FloatArray kBias, + FloatArray vBias, + IntArray activeBatchSizeHolder, + int dim, + int kvDim) { + + int index = context.globalIdx; + int rowsPerToken = dim + 2 * kvDim; + int token = index / rowsPerToken; + int row = index % rowsPerToken; + + if (token >= activeBatchSizeHolder.get(0)) { + return; + } + + if (row < dim) { + int qIndex = token * dim + row; + qBatch.set(qIndex, qBatch.get(qIndex) + qBias.get(row)); + } else if (row < dim + kvDim) { + int kRow = row - dim; + int kIndex = token * kvDim + kRow; + kBatch.set(kIndex, kBatch.get(kIndex) + kBias.get(kRow)); + } else { + int vRow = row - dim - kvDim; + int vIndex = token * kvDim + vRow; + vBatch.set(vIndex, vBatch.get(vIndex) + vBias.get(vRow)); + } + } + + /** Applies Qwen2 split-half RoPE and writes active K/V rows to the cache. */ + public static void batchedRopeWithKVCacheQwen2( + KernelContext context, + IntArray batchStartPosHolder, + IntArray activeBatchSizeHolder, + FloatArray qBatch, + FloatArray kBatch, + FloatArray vBatch, + FloatArray keyCache, + FloatArray valueCache, + int kvDim, + int headSize, + int layerIndex, + int contextLength, + int dim, + float ropeTheta) { + + int index = context.globalIdx; + int pairsPerToken = dim / 2; + int token = index / pairsPerToken; + int pair = index % pairsPerToken; + + if (token >= activeBatchSizeHolder.get(0)) { + return; + } + + int position = batchStartPosHolder.get(0) + token; + int halfHeadSize = headSize / 2; + int component = pair % halfHeadSize; + int head = pair / halfHeadSize; + + float frequency = 1.0f / TornadoMath.pow(ropeTheta, 2.0f * component / (float) headSize); + float angle = position * frequency; + float cosine = TornadoMath.cos(angle); + float sine = TornadoMath.sin(angle); + + int qHeadOffset = token * dim + head * headSize; + float q0 = qBatch.get(qHeadOffset + component); + float q1 = qBatch.get(qHeadOffset + component + halfHeadSize); + qBatch.set(qHeadOffset + component, q0 * cosine - q1 * sine); + qBatch.set(qHeadOffset + component + halfHeadSize, q0 * sine + q1 * cosine); + + if (pair < kvDim / 2) { + int kvHead = pair / halfHeadSize; + int kHeadOffset = token * kvDim + kvHead * headSize; + float k0 = kBatch.get(kHeadOffset + component); + float k1 = kBatch.get(kHeadOffset + component + halfHeadSize); + float rotatedK0 = k0 * cosine - k1 * sine; + float rotatedK1 = k0 * sine + k1 * cosine; + + kBatch.set(kHeadOffset + component, rotatedK0); + kBatch.set(kHeadOffset + component + halfHeadSize, rotatedK1); + + int cacheOffset = + layerIndex * contextLength * kvDim + position * kvDim + kvHead * headSize; + keyCache.set(cacheOffset + component, rotatedK0); + keyCache.set(cacheOffset + component + halfHeadSize, rotatedK1); + valueCache.set(cacheOffset + component, vBatch.get(kHeadOffset + component)); + valueCache.set( + cacheOffset + component + halfHeadSize, + vBatch.get(kHeadOffset + component + halfHeadSize)); + } + } + + /** Applies softmax and selects Top-K experts independently for each token. */ + public static void batchedSoftmaxAndTopK( + KernelContext context, + FloatArray routerLogits, + IntArray selectedExperts, + FloatArray routingWeights, + IntArray activeBatchSizeHolder, + int numberOfExperts, + int topK) { + + // One GPU thread handles the complete routing result for one token. + int token = context.globalIdx; + if (token >= activeBatchSizeHolder.get(0)) { + return; + } + + int logitsOffset = token * numberOfExperts; + int assignmentOffset = token * topK; + + // Find this token's maximum router logit. + float maxLogit = Float.NEGATIVE_INFINITY; + for (int i = logitsOffset; i < logitsOffset + numberOfExperts; i++) { + maxLogit = Math.max(maxLogit, routerLogits.get(i)); + } + + // Compute the stable softmax denominator. + float sumExp = 0.0f; + + for (int expert = 0; expert < numberOfExperts; expert++) { + float logit = routerLogits.get(logitsOffset + expert); + sumExp += TornadoMath.exp(logit - maxLogit); + } + + // Convert this token's logits to probabilities. + for (int expert = 0; expert < numberOfExperts; expert++) { + int index = logitsOffset + expert; + float logit = routerLogits.get(index); + + float probability = TornadoMath.exp(logit - maxLogit) / sumExp; + + routerLogits.set(index, probability); + } + + // Select Top-K expert IDs and their routing weights. + for (int slot = 0; slot < topK; slot++) { + float currentMax = Float.NEGATIVE_INFINITY; + int selectedExpert = -1; + + for (int expert = 0; expert < numberOfExperts; expert++) { + float probability = routerLogits.get(logitsOffset + expert); + + if (probability > currentMax) { + currentMax = probability; + selectedExpert = expert; + } + } + + selectedExperts.set(assignmentOffset + slot, selectedExpert); + routingWeights.set(assignmentOffset + slot, currentMax); + + routerLogits.set(logitsOffset + selectedExpert, Float.NEGATIVE_INFINITY); + } + } + + /** Groups token-expert assignments by expert and records each expert's range. */ + public static void groupAssignmentsByExpert( + KernelContext context, + IntArray selectedExperts, + IntArray groupedAssignmentIds, + IntArray groupedPositionByAssignment, + IntArray expertOffsets, + IntArray activeBatchSizeHolder, + int numberOfExperts, + int topK) { + + // Start with one thread for a simple, deterministic implementation. + if (context.globalIdx != 0) { + return; + } + + int numberOfAssignments = activeBatchSizeHolder.get(0) * topK; + int groupedPosition = 0; + + // Visit each expert in expert-ID order. + for (int expert = 0; expert < numberOfExperts; expert++) { + // Record where this expert's assignments begin. + expertOffsets.set(expert, groupedPosition); + + // Find all active assignments that selected this expert. + for (int assignment = 0; assignment < numberOfAssignments; assignment++) { + + int selectedExpert = selectedExperts.get(assignment); + + if (selectedExpert == expert) { + // grouped position -> original assignment ID + groupedAssignmentIds.set(groupedPosition, assignment); + + // original assignment ID -> grouped position + groupedPositionByAssignment.set(assignment, groupedPosition); + + groupedPosition++; + } + } + } + + // Record the end of the final expert's assignment range. + expertOffsets.set(numberOfExperts, groupedPosition); + } + + /** Computes routed Gate/Up projections in expert-grouped assignment order. */ + public static void groupedRoutedExpertsGateUpSwiGLUQ8_0( + KernelContext context, + FloatArray inputBatch, + IntArray selectedExperts, + IntArray groupedAssignmentIds, + IntArray activeBatchSizeHolder, + ByteArray gateExperts, + ByteArray upExperts, + FloatArray groupedExpertHidden, + int dim, + int moeHiddenDim, + int numberOfExperts, + int topK, + int localWorkGroupSize) { + + int flatGroupId = context.groupIdx; + int localId = context.localIdx; + + int groupedPosition = flatGroupId / moeHiddenDim; + int rowId = flatGroupId % moeHiddenDim; + int numberOfAssignments = activeBatchSizeHolder.get(0) * topK; + + boolean active = groupedPosition < numberOfAssignments; + int assignment = 0; + int token = 0; + int expert = 0; + if (active) { + assignment = groupedAssignmentIds.get(groupedPosition); + token = assignment / topK; + expert = selectedExperts.get(assignment); + active = expert >= 0 && expert < numberOfExperts; + } + + int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = (expert * moeHiddenDim + rowId) * blocksPerRow; + int inputOffset = token * dim; + + 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 = inputBatch.get(inputOffset + 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)); + int outputOffset = groupedPosition * moeHiddenDim + rowId; + groupedExpertHidden.set(outputOffset, siluGate * up); + } + } + + /** Down-projects every routed assignment in expert-grouped order. */ + public static void groupedRoutedExpertsDownQ8_0( + KernelContext context, + FloatArray groupedExpertHidden, + IntArray selectedExperts, + IntArray groupedAssignmentIds, + IntArray activeBatchSizeHolder, + ByteArray downExperts, + FloatArray groupedExpertDown, + int dim, + int moeHiddenDim, + int numberOfExperts, + int topK, + int localWorkGroupSize) { + + int flatGroupId = context.groupIdx; + int localId = context.localIdx; + + int groupedPosition = flatGroupId / dim; + int rowId = flatGroupId % dim; + int numberOfAssignments = activeBatchSizeHolder.get(0) * topK; + + boolean active = groupedPosition < numberOfAssignments; + int assignment = 0; + int expert = 0; + if (active) { + assignment = groupedAssignmentIds.get(groupedPosition); + expert = selectedExperts.get(assignment); + active = expert >= 0 && expert < numberOfExperts; + } + + int blocksPerRow = (moeHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = (expert * dim + rowId) * blocksPerRow; + int hiddenOffset = groupedPosition * moeHiddenDim; + + float partialSum = 0.0f; + if (active) { + 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 * groupedExpertHidden.get(hiddenOffset + column); + } + } + + 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 && active) { + int outputOffset = groupedPosition * dim + rowId; + groupedExpertDown.set(outputOffset, localSums[0]); + } + } + + /** Adds the weighted routed-expert results back to each token's residual. */ + public static void accumulateGroupedRoutedExperts( + KernelContext context, + FloatArray groupedExpertDown, + IntArray groupedPositionByAssignment, + FloatArray routingWeights, + FloatArray residualBatch, + IntArray activeBatchSizeHolder, + int dim, + int topK) { + + int index = context.globalIdx; + int token = index / dim; + int rowId = index % dim; + + if (token >= activeBatchSizeHolder.get(0)) { + return; + } + + float result = residualBatch.get(index); + int assignmentOffset = token * topK; + for (int slot = 0; slot < topK; slot++) { + int assignment = assignmentOffset + slot; + int groupedPosition = groupedPositionByAssignment.get(assignment); + int downOffset = groupedPosition * dim + rowId; + result += routingWeights.get(assignment) * groupedExpertDown.get(downOffset); + } + + residualBatch.set(index, result); + } + + /** Computes the shared expert Gate/Up result independently for each token. */ + public static void batchedSharedExpertGateUpSwiGLUQ8_0( + KernelContext context, + FloatArray inputBatch, + IntArray activeBatchSizeHolder, + ByteArray sharedGate, + ByteArray sharedUp, + FloatArray sharedHiddenBatch, + int dim, + int sharedExpertHiddenDim, + int localWorkGroupSize) { + + int flatGroupId = context.groupIdx; + int localId = context.localIdx; + + int token = flatGroupId / sharedExpertHiddenDim; + int rowId = flatGroupId % sharedExpertHiddenDim; + boolean active = token < activeBatchSizeHolder.get(0); + + int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = rowId * blocksPerRow; + int inputOffset = token * dim; + + 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 = inputBatch.get(inputOffset + column); + float gateScale = sharedGate.getHalfFloat(blockByteOffset).getFloat32(); + float upScale = sharedUp.getHalfFloat(blockByteOffset).getFloat32(); + + gatePartialSum += sharedGate.get(quantOffset) * gateScale * inputValue; + upPartialSum += sharedUp.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)); + int outputOffset = token * sharedExpertHiddenDim + rowId; + sharedHiddenBatch.set(outputOffset, siluGate * up); + } + } + + /** Computes the sigmoid gate that scales each token's shared-expert output. */ + public static void batchedSharedExpertGateWeight( + KernelContext context, + FloatArray inputBatch, + FloatArray sharedGateInput, + FloatArray sharedWeightBatch, + IntArray activeBatchSizeHolder, + int dim, + int localWorkGroupSize) { + + int token = context.groupIdx; + int localId = context.localIdx; + boolean active = token < activeBatchSizeHolder.get(0); + int inputOffset = token * dim; + + float partialScore = 0.0f; + if (active) { + for (int column = localId; column < dim; column += localWorkGroupSize) { + partialScore += sharedGateInput.get(column) * inputBatch.get(inputOffset + 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(); + } + + if (localId == 0 && active) { + float sharedWeight = 1.0f / (1.0f + TornadoMath.exp(-localSums[0])); + sharedWeightBatch.set(token, sharedWeight); + } + } + + /** Down-projects the shared expert and adds it to each token's residual. */ + public static void batchedSharedExpertDownAndAccumulateQ8_0( + KernelContext context, + FloatArray sharedHiddenBatch, + FloatArray sharedWeightBatch, + IntArray activeBatchSizeHolder, + ByteArray sharedDown, + FloatArray residualBatch, + int dim, + int sharedExpertHiddenDim, + int localWorkGroupSize) { + + int flatGroupId = context.groupIdx; + int localId = context.localIdx; + + int token = flatGroupId / dim; + int rowId = flatGroupId % dim; + boolean active = token < activeBatchSizeHolder.get(0); + + int blocksPerRow = (sharedExpertHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = rowId * blocksPerRow; + int hiddenOffset = token * sharedExpertHiddenDim; + + float partialSum = 0.0f; + if (active) { + for (int column = localId; + column < sharedExpertHiddenDim; + 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 = + sharedDown.get(quantOffset) + * sharedDown.getHalfFloat(blockByteOffset).getFloat32(); + partialSum += weight * sharedHiddenBatch.get(hiddenOffset + column); + } + } + + 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 && active) { + int outputOffset = token * dim + rowId; + float weightedOutput = sharedWeightBatch.get(token) * localSums[0]; + residualBatch.set(outputOffset, residualBatch.get(outputOffset) + weightedOutput); + } + } +} 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 ee7a07e0..848376c2 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 @@ -10,6 +10,7 @@ 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; @@ -20,20 +21,21 @@ /** * 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}.

+ *

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 +public class Qwen2MoEQ8_0FFNLayers extends AbstractTransformerLayerTaskGraphs { - private final Qwen2MoEState moeState; + protected final Qwen2MoEState moeState; - public Qwen2MoEQ8_0FFNLayers(String taskGraphName, - Qwen2MoEState state, - Qwen2MoETornadoWeights weights, - Qwen2MoEConfiguration config, - SchedulerType schedulerType) { + public Qwen2MoEQ8_0FFNLayers( + String taskGraphName, + Qwen2MoEState state, + Qwen2MoETornadoWeights weights, + Qwen2MoEConfiguration config, + SchedulerType schedulerType) { super(taskGraphName, state, weights, config, schedulerType); this.moeState = state; setupFFNLayers(); @@ -42,8 +44,8 @@ public Qwen2MoEQ8_0FFNLayers(String taskGraphName, /** 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 rmsNormWorker = + WorkerGridFactory.createRmsNormWorker(moeState.localSize, moeState.localSize); WorkerGrid qkvWorker = workerForRows(config.dim() + 2 * config.kvDim()); WorkerGrid qkvBiasWorker = new WorkerGrid1D(config.dim()); @@ -98,16 +100,37 @@ private WorkerGrid workerForRows(int rows) { } /** - * Creates the complete GPU TaskGraph for one Transformer layer. - * {@code layerIndex} selects that layer's weights. + * 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, + String predecessor = predecessorGraphName(layerIndex); + if (predecessor == null) { + layer.consumeFromDevice(moeState.wrapX); + } else { + layer.consumeFromDevice(predecessor, moeState.wrapX); + } + layer = configureLayerWeights(layer, layerIndex); + layer = configureLayerDataTransfers(layer, layerIndex); + + configureAttention(layer, layerIndex); + configureRoutedExperts(layer, layerIndex); + layer.persistOnDevice(moeState.wrapX, moeState.wrapKeyCache, moeState.wrapValueCache); + return layer; + } + + /** Returns an explicit predecessor for plans that connect multiple graph chains. */ + protected String predecessorGraphName(int layerIndex) { + return null; + } + + /** Uploads this layer's read-only weights on its first execution. */ + protected TaskGraph configureLayerWeights(TaskGraph layer, int layerIndex) { + return layer.transferToDevice( + DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -125,136 +148,252 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { 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", + layer.task( + "attn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, - context, moeState.temp, moeState.wrapX, - config.dim(), config.rmsNormEps(), moeState.localSize); + context, + moeState.temp, + moeState.wrapX, + config.dim(), + config.rmsNormEps(), + moeState.localSize); - layer.task("attn_rms_qkv_projection", + 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, + 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); + config.dim(), + config.dim(), + config.kvDim(), + LOCAL_WORK_GROUP_SIZE_ALLOC); - layer.task("fused_qkv_bias", + layer.task( + "fused_qkv_bias", TransformerComputeKernelsLayered::fusedQKvBiasAddition, - context, moeState.wrapQ, moeState.wrapK, - weights.q_biasLayered[layerIndex].asFloatArray(), moeState.wrapV, + 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()); + 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( + "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( + "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", + layer.task( + "attn_output_proj", TransformerComputeKernelsLayered::matrixVectorGenericWithResidualQ8_0Byte, - context, moeState.wrapXb, moeState.wrapX, + context, + moeState.wrapXb, + moeState.wrapX, weights.woLayered[layerIndex].asByteArray(), - config.dim(), config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + 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. + * 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", + layer.task( + "ffn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, - context, moeState.tempFFN, moeState.wrapX, - config.dim(), config.rmsNormEps(), moeState.localSize); + context, + moeState.tempFFN, + moeState.wrapX, + config.dim(), + config.rmsNormEps(), + moeState.localSize); - layer.task("ffn_rms_apply", + layer.task( + "ffn_rms_apply", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, - context, moeState.wrapXb, moeState.wrapX, - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), moeState.tempFFN); + context, + moeState.wrapXb, + moeState.wrapX, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + moeState.tempFFN); - layer.task("router_projection", + layer.task( + "router_projection", TransformerComputeKernelsLayered::matrixVectorGeneric, - context, moeState.wrapXb, moeState.wrapRouterLogits, + context, + moeState.wrapXb, + moeState.wrapRouterLogits, weights.routerGateLayered[layerIndex].asFloatArray(), - config.dim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); + 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()); + layer.task( + "router_softmax_topk", + Qwen2MoEKernels::softmaxAndTopK, + context, + moeState.wrapRouterLogits, + moeState.wrapSelectedExperts, + moeState.wrapRoutingWeights, + config.numberOfExperts(), + config.numberOfExpertsUsed()); // 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", + layer.task( + "routed_experts_gate_up", Qwen2MoEKernels::fusedRoutedExpertsGateUpSwiGLUQ8_0, - context, moeState.wrapXb, moeState.wrapSelectedExperts, config.numberOfExpertsUsed(), + 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); + weights.upExpertsLayered[layerIndex].asByteArray(), + moeState.wrapExpertGate, + config.dim(), + config.moeHiddenDim(), + config.numberOfExperts(), + LOCAL_WORK_GROUP_SIZE_ALLOC); - layer.task("routed_experts_down", + layer.task( + "routed_experts_down", Qwen2MoEKernels::routedExpertsDownProjectAndAccumulateQ8_0, - context, moeState.wrapExpertGate, moeState.wrapX, - moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, config.numberOfExpertsUsed(), + 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); + 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, + 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); + weights.sharedUpLayered[layerIndex].asByteArray(), + moeState.wrapSharedGate, + 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); + 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. - */ + /** 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); + 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, + 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/layers/type/q8_0/decode/Qwen2MoEQ8_0FFNLayersDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/decode/Qwen2MoEQ8_0FFNLayersDecode.java new file mode 100644 index 00000000..6903e585 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/decode/Qwen2MoEQ8_0FFNLayersDecode.java @@ -0,0 +1,105 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode; + +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.layers.type.q8_0.Qwen2MoEQ8_0FFNLayers; +import org.beehive.gpullama3.tornadovm.scheduling.SchedulerType; + +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +/** Single-token decode layers that continue from the batch-prefill KV cache. */ +public final class Qwen2MoEQ8_0FFNLayersDecode extends Qwen2MoEQ8_0FFNLayers { + + public Qwen2MoEQ8_0FFNLayersDecode( + String taskGraph, + Qwen2MoEState state, + Qwen2MoETornadoWeights weights, + Qwen2MoEConfiguration config, + SchedulerType schedulerType) { + super(taskGraph, state, weights, config, schedulerType); + } + + /** Connects decode layer 0 to the decode activation and later layers in order. */ + @Override + protected String predecessorGraphName(int layerIndex) { + return layerIndex == 0 ? "decodeActivation" : "layer_" + (layerIndex - 1); + } + + /** Reuses the weights already uploaded by the corresponding batch-prefill layer. */ + @Override + protected TaskGraph configureLayerWeights(TaskGraph layer, int layerIndex) { + return layer.consumeFromDevice( + "batchPrefillLayer_" + layerIndex, + 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()); + } + + /** Reuses the KV cache produced by batch prefill instead of allocating a new cache. */ + @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.wrapAtt, + moeState.wrapRouterLogits, + moeState.wrapSelectedExperts, + moeState.wrapRoutingWeights, + moeState.wrapExpertGate, + moeState.wrapSharedGate, + moeState.wrapSharedOutput); + layer.consumeFromDevice( + "decodeActivation", moeState.wrapKeyCache, moeState.wrapValueCache); + } else { + String predecessor = "layer_" + (layerIndex - 1); + layer.consumeFromDevice( + predecessor, + 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, + moeState.temp, + moeState.tempFFN); + } + return layer; + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java new file mode 100644 index 00000000..e973937d --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java @@ -0,0 +1,493 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +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.Qwen2MoEBatchKernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; + +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** Batched-prefill Transformer-layer TaskGraphs for Qwen2-MoE Q8_0. */ +public final class Qwen2MoEQ8_0LayersBatchPrefill + implements BatchPrefillTransformerLayerTaskGraphs { + + private static final int LOCAL_WORK_GROUP_SIZE = 32; + + private final Qwen2MoEState state; + private final Qwen2MoETornadoWeights weights; + private final Qwen2MoEConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + private final int dim; + private final int kvDim; + private final int topK; + private final int numberOfAssignments; + private final List layerTaskGraphs; + private String lastLayerTaskGraphID; + + public Qwen2MoEQ8_0LayersBatchPrefill( + Qwen2MoEState state, + Qwen2MoETornadoWeights weights, + Qwen2MoEConfiguration config, + int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.dim = config.dim(); + this.kvDim = config.kvDim(); + this.topK = config.numberOfExpertsUsed(); + this.numberOfAssignments = batchSize * topK; + this.layerTaskGraphs = + IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + /** Creates the complete batch-prefill TaskGraph for one Transformer layer. */ + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) { + lastLayerTaskGraphID = graphName; + } + + TaskGraph layer = new TaskGraph(graphName); + configureDataTransfers(layer, layerIndex); + configureAttention(layer, layerIndex); + configureMoE(layer, layerIndex); + layer.persistOnDevice( + state.wrapXBatch, + state.wrapKeyCache, + state.wrapValueCache, + 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()); + return layer; + } + + /** Declares layer weights and the batch buffers that remain on the GPU. */ + private void configureDataTransfers(TaskGraph layer, int layerIndex) { + if (layerIndex == 0) { + layer.transferToDevice( + DataTransferMode.EVERY_EXECUTION, + state.batchStartPosHolder, + state.activeBatchSizeHolder); + layer.transferToDevice( + DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, + state.ffnScaleBatch, + state.wrapXbBatch, + state.wrapQBatch, + state.wrapKBatch, + state.wrapVBatch, + state.wrapKeyCache, + state.wrapValueCache, + state.wrapRouterLogitsBatch, + state.wrapSelectedExpertsBatch, + state.wrapRoutingWeightsBatch, + state.wrapGroupedAssignmentIds, + state.wrapGroupedPositionByAssignment, + state.wrapExpertOffsets, + state.wrapGroupedExpertHidden, + state.wrapGroupedExpertDown, + state.wrapSharedHiddenBatch, + state.wrapSharedWeightBatch); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String predecessor = "batchPrefillLayer_" + (layerIndex - 1); + layer.consumeFromDevice( + predecessor, + context, + state.wrapXBatch, + state.wrapXbBatch, + state.wrapQBatch, + state.wrapKBatch, + state.wrapVBatch, + state.wrapKeyCache, + state.wrapValueCache, + state.batchStartPosHolder, + state.activeBatchSizeHolder, + state.attnScaleBatch, + state.ffnScaleBatch, + state.wrapRouterLogitsBatch, + state.wrapSelectedExpertsBatch, + state.wrapRoutingWeightsBatch, + state.wrapGroupedAssignmentIds, + state.wrapGroupedPositionByAssignment, + state.wrapExpertOffsets, + state.wrapGroupedExpertHidden, + state.wrapGroupedExpertDown, + state.wrapSharedHiddenBatch, + state.wrapSharedWeightBatch); + } + + layer.transferToDevice( + DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.q_biasLayered[layerIndex].asFloatArray(), + weights.k_biasLayered[layerIndex].asFloatArray(), + weights.v_biasLayered[layerIndex].asFloatArray(), + weights.woLayered[layerIndex].asByteArray(), + 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()); + } + + /** Adds Qwen2 attention tasks for every token in the prefill batch. */ + private void configureAttention(TaskGraph layer, int layerIndex) { + layer.task( + "batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, + state.wrapXBatch, + state.attnScaleBatch, + dim, + config.rmsNormEps(), + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, + state.wrapXbBatch, + state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, + dim); + + layer.task( + "batch_qkv", + TransformerBatchPrefillKernels::batchedFusedQKVMatmulQ8, + context, + state.wrapXbBatch, + state.wrapQBatch, + state.wrapKBatch, + state.wrapVBatch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + dim, + kvDim, + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_qkv_bias", + Qwen2MoEBatchKernels::batchedQKVBias, + context, + state.wrapQBatch, + state.wrapKBatch, + state.wrapVBatch, + weights.q_biasLayered[layerIndex].asFloatArray(), + weights.k_biasLayered[layerIndex].asFloatArray(), + weights.v_biasLayered[layerIndex].asFloatArray(), + state.activeBatchSizeHolder, + dim, + kvDim); + + layer.task( + "batch_rope_kv", + Qwen2MoEBatchKernels::batchedRopeWithKVCacheQwen2, + context, + state.batchStartPosHolder, + state.activeBatchSizeHolder, + state.wrapQBatch, + state.wrapKBatch, + state.wrapVBatch, + state.wrapKeyCache, + state.wrapValueCache, + kvDim, + config.headSize(), + layerIndex, + config.contextLength(), + dim, + config.ropeTheta()); + + layer.task( + "batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, + context, + state.batchStartPosHolder, + state.wrapQBatch, + state.wrapKeyCache, + state.wrapValueCache, + state.wrapXbBatch, + config.numberOfHeads(), + config.headSize(), + kvDim, + config.kvMul(), + layerIndex, + config.contextLength(), + dim); + + layer.task( + "batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, + state.wrapXbBatch, + state.wrapXBatch, + weights.woLayered[layerIndex].asByteArray(), + dim, + dim, + LOCAL_WORK_GROUP_SIZE); + } + + /** Adds batched routing, routed experts, and the shared expert. */ + private void configureMoE(TaskGraph layer, int layerIndex) { + layer.task( + "batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, + state.wrapXBatch, + state.ffnScaleBatch, + dim, + config.rmsNormEps(), + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, + state.wrapXbBatch, + state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, + dim); + + layer.task( + "batch_router", + Qwen2MoEBatchKernels::batchedRouterProjection, + context, + state.wrapXbBatch, + state.wrapRouterLogitsBatch, + weights.routerGateLayered[layerIndex].asFloatArray(), + state.activeBatchSizeHolder, + dim, + config.numberOfExperts(), + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_topk", + Qwen2MoEBatchKernels::batchedSoftmaxAndTopK, + context, + state.wrapRouterLogitsBatch, + state.wrapSelectedExpertsBatch, + state.wrapRoutingWeightsBatch, + state.activeBatchSizeHolder, + config.numberOfExperts(), + topK); + + layer.task( + "batch_group_assignments", + Qwen2MoEBatchKernels::groupAssignmentsByExpert, + context, + state.wrapSelectedExpertsBatch, + state.wrapGroupedAssignmentIds, + state.wrapGroupedPositionByAssignment, + state.wrapExpertOffsets, + state.activeBatchSizeHolder, + config.numberOfExperts(), + topK); + + layer.task( + "batch_routed_gate_up", + Qwen2MoEBatchKernels::groupedRoutedExpertsGateUpSwiGLUQ8_0, + context, + state.wrapXbBatch, + state.wrapSelectedExpertsBatch, + state.wrapGroupedAssignmentIds, + state.activeBatchSizeHolder, + weights.gateExpertsLayered[layerIndex].asByteArray(), + weights.upExpertsLayered[layerIndex].asByteArray(), + state.wrapGroupedExpertHidden, + dim, + config.moeHiddenDim(), + config.numberOfExperts(), + topK, + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_routed_down", + Qwen2MoEBatchKernels::groupedRoutedExpertsDownQ8_0, + context, + state.wrapGroupedExpertHidden, + state.wrapSelectedExpertsBatch, + state.wrapGroupedAssignmentIds, + state.activeBatchSizeHolder, + weights.downExpertsLayered[layerIndex].asByteArray(), + state.wrapGroupedExpertDown, + dim, + config.moeHiddenDim(), + config.numberOfExperts(), + topK, + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_routed_accumulate", + Qwen2MoEBatchKernels::accumulateGroupedRoutedExperts, + context, + state.wrapGroupedExpertDown, + state.wrapGroupedPositionByAssignment, + state.wrapRoutingWeightsBatch, + state.wrapXBatch, + state.activeBatchSizeHolder, + dim, + topK); + + layer.task( + "batch_shared_gate_up", + Qwen2MoEBatchKernels::batchedSharedExpertGateUpSwiGLUQ8_0, + context, + state.wrapXbBatch, + state.activeBatchSizeHolder, + weights.sharedGateLayered[layerIndex].asByteArray(), + weights.sharedUpLayered[layerIndex].asByteArray(), + state.wrapSharedHiddenBatch, + dim, + config.sharedExpertHiddenDim(), + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_shared_weight", + Qwen2MoEBatchKernels::batchedSharedExpertGateWeight, + context, + state.wrapXbBatch, + weights.sharedGateInputLayered[layerIndex].asFloatArray(), + state.wrapSharedWeightBatch, + state.activeBatchSizeHolder, + dim, + LOCAL_WORK_GROUP_SIZE); + + layer.task( + "batch_shared_down", + Qwen2MoEBatchKernels::batchedSharedExpertDownAndAccumulateQ8_0, + context, + state.wrapSharedHiddenBatch, + state.wrapSharedWeightBatch, + state.activeBatchSizeHolder, + weights.sharedDownLayered[layerIndex].asByteArray(), + state.wrapXBatch, + dim, + config.sharedExpertHiddenDim(), + LOCAL_WORK_GROUP_SIZE); + } + + /** Configures the fixed WorkerGrid used by every task in every layer. */ + @Override + public void updateGridScheduler(GridScheduler scheduler) { + WorkerGrid rmsWorker = groupedRowsWorker(batchSize); + WorkerGrid batchScalarWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid batchElementWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid qkvWorker = groupedRowsWorker(batchSize * (dim + 2 * kvDim)); + int qkvBiasGlobalWork = batchSize * (dim + 2 * kvDim); + WorkerGrid qkvBiasWorker = + WorkerGridFactory.genericWorker( + qkvBiasGlobalWork, + validLocalSize(qkvBiasGlobalWork, 256)); + + int ropeGlobalWork = batchSize * (dim / 2); + int ropeLocalWork = validLocalSize(ropeGlobalWork, 512); + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobalWork, ropeLocalWork); + + int attentionLocalWork = validLocalSize(config.headSize(), 64); + WorkerGrid attentionWorker = + WorkerGridFactory.genericWorker( + batchSize * config.numberOfHeads() * attentionLocalWork, + attentionLocalWork); + + WorkerGrid batchDimRowsWorker = groupedRowsWorker(batchSize * dim); + WorkerGrid routerWorker = groupedRowsWorker(batchSize * config.numberOfExperts()); + WorkerGrid groupingWorker = WorkerGridFactory.createSingleWorker(); + WorkerGrid routedHiddenWorker = + groupedRowsWorker(numberOfAssignments * config.moeHiddenDim()); + WorkerGrid routedDownWorker = groupedRowsWorker(numberOfAssignments * dim); + WorkerGrid sharedHiddenWorker = + groupedRowsWorker(batchSize * config.sharedExpertHiddenDim()); + WorkerGrid sharedWeightWorker = groupedRowsWorker(batchSize); + + for (int layer = 0; layer < config.numberOfLayers(); layer++) { + String prefix = "batchPrefillLayer_" + layer + "."; + scheduler.addWorkerGrid(prefix + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(prefix + "batch_attn_rms_apply", batchElementWorker); + scheduler.addWorkerGrid(prefix + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(prefix + "batch_qkv_bias", qkvBiasWorker); + scheduler.addWorkerGrid(prefix + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(prefix + "batch_attention", attentionWorker); + scheduler.addWorkerGrid(prefix + "batch_attn_out", batchDimRowsWorker); + scheduler.addWorkerGrid(prefix + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(prefix + "batch_ffn_rms_apply", batchElementWorker); + scheduler.addWorkerGrid(prefix + "batch_router", routerWorker); + scheduler.addWorkerGrid(prefix + "batch_topk", batchScalarWorker); + scheduler.addWorkerGrid(prefix + "batch_group_assignments", groupingWorker); + scheduler.addWorkerGrid(prefix + "batch_routed_gate_up", routedHiddenWorker); + scheduler.addWorkerGrid(prefix + "batch_routed_down", routedDownWorker); + scheduler.addWorkerGrid(prefix + "batch_routed_accumulate", batchElementWorker); + scheduler.addWorkerGrid(prefix + "batch_shared_gate_up", sharedHiddenWorker); + scheduler.addWorkerGrid(prefix + "batch_shared_weight", sharedWeightWorker); + scheduler.addWorkerGrid(prefix + "batch_shared_down", batchDimRowsWorker); + } + } + + /** Creates a 32-thread work-group for every logical output row. */ + private static WorkerGrid groupedRowsWorker(int rows) { + return WorkerGridFactory.genericWorker(rows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + } + + /** Finds a legal local size that divides the requested global dimension. */ + private static int validLocalSize(int size, int maximum) { + int localSize = Math.min(size, maximum); + while (localSize > 1 && size % localSize != 0) { + localSize--; + } + return localSize; + } + + @Override + public List getLayerImmutableTaskGraphs() { + return layerTaskGraphs; + } + + @Override + public String getLastLayerTaskGraphID() { + return lastLayerTaskGraphID; + } + + public KernelContext getContext() { + return context; + } +} 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 6e8ba29d..df4133df 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java @@ -24,8 +24,8 @@ import org.beehive.gpullama3.tornadovm.plan.components.q8_0.LlamaQ8_0PlanComponents; 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.Qwen2Q8_0PlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen3Q8_0PlanComponents; // @formatter:off @@ -33,46 +33,54 @@ * Factory for {@link ForwardPlan} instances. * *

Dispatches across three axes in order: + * *

    - *
  1. Quantization ({@link GGMLType})
    2. - *
  2. Model family ({@link org.beehive.gpullama3.model.ModelType})
    3. - *
  3. Execution mode ({@link ExecutionMode})
    4. + *
  4. Quantization ({@link GGMLType}) + *
  5. Model family ({@link org.beehive.gpullama3.model.ModelType}) + *
  6. Execution mode ({@link ExecutionMode}) *
* - *

Use the typed convenience methods when the execution mode is known at the call site:

+ *

Use the typed convenience methods when the execution mode is known at the call site: + * *

*/ // @formatter:on public class ForwardPlanFactory { - private ForwardPlanFactory() { - } + private ForwardPlanFactory() {} // ── Typed public API ────────────────────────────────────────────────────── - public static SingleTokenForwardPlan createSingleToken(GGMLType quantization, State state, Model model) { + public static SingleTokenForwardPlan createSingleToken( + GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.STANDARD, state, model); - if (plan instanceof SingleTokenForwardPlan singleToken) - return singleToken; - throw new IllegalStateException("Expected SingleTokenForwardPlan for STANDARD mode but got " + plan.getClass().getSimpleName()); + if (plan instanceof SingleTokenForwardPlan singleToken) return singleToken; + throw new IllegalStateException( + "Expected SingleTokenForwardPlan for STANDARD mode but got " + + plan.getClass().getSimpleName()); } - public static PrefillDecodeForwardPlan createPrefillDecode(GGMLType quantization, State state, Model model) { + public static PrefillDecodeForwardPlan createPrefillDecode( + GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.PREFILL_DECODE, state, model); - if (plan instanceof PrefillDecodeForwardPlan prefillDecode) - return prefillDecode; - throw new IllegalStateException("Expected PrefillDecodeForwardPlan for PREFILL_DECODE mode but got " + plan.getClass().getSimpleName()); + if (plan instanceof PrefillDecodeForwardPlan prefillDecode) return prefillDecode; + throw new IllegalStateException( + "Expected PrefillDecodeForwardPlan for PREFILL_DECODE mode but got " + + plan.getClass().getSimpleName()); } - public static BatchPrefillDecodeForwardPlan createBatchPrefillDecode(GGMLType quantization, State state, Model model) { + public static BatchPrefillDecodeForwardPlan createBatchPrefillDecode( + GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.BATCH_PREFILL_DECODE, state, model); if (plan instanceof BatchPrefillDecodeForwardPlan batchPrefillDecode) return batchPrefillDecode; - throw new IllegalStateException("Expected BatchPrefillDecodeForwardPlan for BATCH_PREFILL_DECODE mode but got " + plan.getClass().getSimpleName()); + throw new IllegalStateException( + "Expected BatchPrefillDecodeForwardPlan for BATCH_PREFILL_DECODE mode but got " + + plan.getClass().getSimpleName()); } // ── Generic dispatch ────────────────────────────────────────────────────── @@ -83,7 +91,9 @@ static ForwardPlan create(GGMLType quantization, ExecutionMode mode, State state case Q8_0 -> createQ8_0Plan(mode, state, model); case F32 -> throw new UnsupportedOperationException("F32 plans not yet implemented"); case Q4_0 -> throw new UnsupportedOperationException("Q4_0 plans not yet implemented"); - default -> throw new UnsupportedOperationException("Quantization not supported: " + quantization); + default -> + throw new UnsupportedOperationException( + "Quantization not supported: " + quantization); }; } @@ -99,7 +109,9 @@ private static ForwardPlan createFP16Plan(ExecutionMode mode, State state, Model case PHI_3 -> createPhi3FP16Plan(mode, (Phi3State) state, model); case GRANITE -> createGraniteFP16Plan(mode, (GraniteState) state, model); case DEEPSEEK_R1_DISTILL_QWEN -> createQwen2FP16Plan(mode, (Qwen2State) state, model); - default -> throw new UnsupportedOperationException("F16 not supported for model: " + model.getModelType()); + default -> + throw new UnsupportedOperationException( + "F16 not supported for model: " + model.getModelType()); }; } @@ -116,111 +128,150 @@ private static ForwardPlan createQ8_0Plan(ExecutionMode mode, State state, Model case PHI_3 -> createPhi3Q8_0Plan(mode, (Phi3State) state, model); case GRANITE -> createGraniteQ8_0Plan(mode, (GraniteState) state, model); case DEEPSEEK_R1_DISTILL_QWEN -> createQwen2Q8_0Plan(mode, (Qwen2State) state, model); - default -> throw new UnsupportedOperationException("Q8_0 not supported for model: " + model.getModelType()); + default -> + throw new UnsupportedOperationException( + "Q8_0 not supported for model: " + model.getModelType()); }; } // ── Model+quant helpers — Llama (all 3 modes supported) ────────────────── - private static ForwardPlan createLlamaFP16Plan(ExecutionMode mode, LlamaState state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = new LlamaFP16PlanComponents(state, model); + private static ForwardPlan createLlamaFP16Plan( + ExecutionMode mode, LlamaState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = + new LlamaFP16PlanComponents(state, model); return switch (mode) { case STANDARD -> new SingleTokenForwardPlan(model, components); case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case BATCH_PREFILL_DECODE -> + new BatchPrefillDecodeForwardPlan( + model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createLlamaQ8_0Plan(ExecutionMode mode, LlamaState state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = new LlamaQ8_0PlanComponents(state, model); + private static ForwardPlan createLlamaQ8_0Plan( + ExecutionMode mode, LlamaState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = + new LlamaQ8_0PlanComponents(state, model); return switch (mode) { case STANDARD -> new SingleTokenForwardPlan(model, components); case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case BATCH_PREFILL_DECODE -> + new BatchPrefillDecodeForwardPlan( + model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } // ── Model+quant helpers — STANDARD only ────────────────────────────────── - private static ForwardPlan createMistralFP16Plan(ExecutionMode mode, LlamaState state, Model model) { + private static ForwardPlan createMistralFP16Plan( + ExecutionMode mode, LlamaState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for MISTRAL + F16"); return new SingleTokenForwardPlan(model, new MistralFP16PlanComponents(state, model)); } - private static ForwardPlan createMistralQ8_0Plan(ExecutionMode mode, LlamaState state, Model model) { + private static ForwardPlan createMistralQ8_0Plan( + ExecutionMode mode, LlamaState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for MISTRAL + Q8_0"); return new SingleTokenForwardPlan(model, new MistralQ8_0PlanComponents(state, model)); } - private static ForwardPlan createDevstralFP16Plan(ExecutionMode mode, DevstralState state, Model model) { + private static ForwardPlan createDevstralFP16Plan( + ExecutionMode mode, DevstralState state, Model model) { if (mode != ExecutionMode.STANDARD) - throw new UnsupportedOperationException(mode + " not yet supported for DEVSTRAL_2 + F16"); + throw new UnsupportedOperationException( + mode + " not yet supported for DEVSTRAL_2 + F16"); return new SingleTokenForwardPlan(model, new DevstralFP16PlanComponents(state, model)); } - private static ForwardPlan createDevstralQ8_0Plan(ExecutionMode mode, DevstralState state, Model model) { + private static ForwardPlan createDevstralQ8_0Plan( + ExecutionMode mode, DevstralState state, Model model) { if (mode != ExecutionMode.STANDARD) - throw new UnsupportedOperationException(mode + " not yet supported for DEVSTRAL_2 + Q8_0"); + throw new UnsupportedOperationException( + mode + " not yet supported for DEVSTRAL_2 + Q8_0"); return new SingleTokenForwardPlan(model, new DevstralQ8_0PlanComponents(state, model)); } - private static ForwardPlan createQwen2FP16Plan(ExecutionMode mode, Qwen2State state, Model model) { + private static ForwardPlan createQwen2FP16Plan( + ExecutionMode mode, Qwen2State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2 + F16"); return new SingleTokenForwardPlan(model, new Qwen2FP16PlanComponents(state, model)); } - private static ForwardPlan createQwen2Q8_0Plan(ExecutionMode mode, Qwen2State state, Model model) { + private static ForwardPlan createQwen2Q8_0Plan( + ExecutionMode mode, Qwen2State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2 + Q8_0"); 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 createQwen2MoEQ8_0Plan( + ExecutionMode mode, Qwen2MoEState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = + new Qwen2MoEQ8_0PlanComponents(state, model); + return switch (mode) { + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> + throw new UnsupportedOperationException( + mode + " not yet supported for QWEN_2_MOE + Q8_0"); + case BATCH_PREFILL_DECODE -> + new BatchPrefillDecodeForwardPlan( + model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + }; } - private static ForwardPlan createQwen3FP16Plan(ExecutionMode mode, Qwen3State state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = new Qwen3FP16PlanComponents(state, model); + private static ForwardPlan createQwen3FP16Plan( + ExecutionMode mode, Qwen3State state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = + new Qwen3FP16PlanComponents(state, model); return switch (mode) { - case STANDARD -> new SingleTokenForwardPlan(model, components); - case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); + case BATCH_PREFILL_DECODE -> + new BatchPrefillDecodeForwardPlan( + model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createQwen3Q8_0Plan(ExecutionMode mode, Qwen3State state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = new Qwen3Q8_0PlanComponents(state, model); + private static ForwardPlan createQwen3Q8_0Plan( + ExecutionMode mode, Qwen3State state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = + new Qwen3Q8_0PlanComponents(state, model); return switch (mode) { - case STANDARD -> new SingleTokenForwardPlan(model, components); - case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); + case BATCH_PREFILL_DECODE -> + new BatchPrefillDecodeForwardPlan( + model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createPhi3FP16Plan(ExecutionMode mode, Phi3State state, Model model) { + private static ForwardPlan createPhi3FP16Plan( + ExecutionMode mode, Phi3State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for PHI_3 + F16"); return new SingleTokenForwardPlan(model, new Phi3FP16PlanComponents(state, model)); } - private static ForwardPlan createPhi3Q8_0Plan(ExecutionMode mode, Phi3State state, Model model) { + private static ForwardPlan createPhi3Q8_0Plan( + ExecutionMode mode, Phi3State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for PHI_3 + Q8_0"); return new SingleTokenForwardPlan(model, new Phi3Q8_0PlanComponents(state, model)); } - private static ForwardPlan createGraniteFP16Plan(ExecutionMode mode, GraniteState state, Model model) { + private static ForwardPlan createGraniteFP16Plan( + ExecutionMode mode, GraniteState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for GRANITE + F16"); return new SingleTokenForwardPlan(model, new GraniteFP16PlanComponents(state, model)); } - private static ForwardPlan createGraniteQ8_0Plan(ExecutionMode mode, GraniteState state, Model model) { + private static ForwardPlan createGraniteQ8_0Plan( + ExecutionMode mode, GraniteState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for GRANITE + Q8_0"); return new SingleTokenForwardPlan(model, new GraniteQ8_0PlanComponents(state, model)); 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 index 93a1443f..6afa7719 100644 --- 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 @@ -7,15 +7,21 @@ 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.BatchPrefillTransformerLayerTaskGraphs; 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.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen2MoEQ8_0FFNLayersDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen2MoEQ8_0LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; +import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; +import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; 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 { +/** Assembles the single-token and batch-prefill Q8_0 GPU components for Qwen2-MoE. */ +public final class Qwen2MoEQ8_0PlanComponents implements BatchPrefillDecodeForwardPlanComponents { private final Qwen2MoEState state; private final Qwen2MoETornadoWeights weights; @@ -34,13 +40,50 @@ public ActivationTaskGraph singleTokenActivation() { return new Activation("activationUpdate", state, weights, config); } + @Override + public ActivationTaskGraph prefillDecodeActivation() { + return new Activation("decodeActivation", state, weights, config); + } + + @Override + public ActivationTaskGraph batchPrefillActivation(int batchSize) { + return new BatchPrefillActivation(state, config, batchSize, true); + } + + @Override + public ActivationTaskGraph batchDecodeActivation(String lastBatchLayerId) { + return new BatchDecodeActivation(state, config, lastBatchLayerId, true); + } + @Override public TransformerLayerTaskGraphs singleTokenTransformerLayers() { return new Qwen2MoEQ8_0FFNLayers("qwen2MoEFFN", state, weights, config, schedulerType); } + @Override + public TransformerLayerTaskGraphs prefillDecodeTransformerLayers() { + return new Qwen2MoEQ8_0FFNLayers("decode", state, weights, config, schedulerType); + } + + @Override + public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { + return new Qwen2MoEQ8_0FFNLayersDecode("decode", state, weights, config, schedulerType); + } + + @Override + public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { + return new Qwen2MoEQ8_0LayersBatchPrefill(state, weights, config, batchSize); + } + @Override public AbstractLogitsTaskGraph singleTokenLogits(String previousGraphId) { - return new LogitsQ8_0Layer("logits", state, weights, config, previousGraphId, schedulerType); + return new LogitsQ8_0Layer( + "logits", state, weights, config, previousGraphId, schedulerType); + } + + @Override + public AbstractLogitsTaskGraph decodeLogits(String previousGraphId) { + return new LogitsQ8_0LayerDecode( + "logits", state, weights, config, previousGraphId, schedulerType); } } From 6c7eb138a48198ed5a0be6ee954f6aeabbdc621b Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Fri, 14 Aug 2026 16:58:13 +0100 Subject: [PATCH 2/4] Reduce batch prefill formatting noise --- .../InferenceCoreBatchPrefillDecode.java | 287 +++++++--------- ...InferenceEngineWithBatchPrefillDecode.java | 15 +- .../inference/state/Qwen2MoEState.java | 34 +- .../gpullama3/model/qwen2/Qwen2MoE.java | 76 +---- .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 315 ++++++------------ .../tornadovm/plan/ForwardPlanFactory.java | 160 ++++----- .../q8_0/Qwen2MoEQ8_0PlanComponents.java | 6 +- 7 files changed, 296 insertions(+), 597 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java index ffed2941..9e3601a8 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java @@ -10,7 +10,6 @@ import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode; - import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import java.lang.foreign.MemorySegment; @@ -18,45 +17,50 @@ /** * Low-level forward passes for the batched prefill/decode inference path (Phase 3/4). * - *

Parallel to {@link InferenceCoreWithPrefillDecode} — does NOT modify it. - * - *

Provides three operations: + *

Parallel to {@link InferenceCoreWithPrefillDecode} — does NOT modify it.

* + *

Provides three operations:

* */ public final class InferenceCoreBatchPrefillDecode { - private InferenceCoreBatchPrefillDecode() {} + private InferenceCoreBatchPrefillDecode() { + } /** * CPU batched prefill forward pass for LLaMA (Phase 3). * - *

Processes {@code batchSize} prompt tokens simultaneously through all transformer layers. - * For each layer, Q/K/V projections, output projection, and FFN projections are computed via - * batch matmul ({@link FloatTensor#matmul(int, FloatTensor[], FloatTensor[], int, int)}), which - * parallelises over both output dimension and batch simultaneously. Attention reuses {@code - * state.att} sequentially per token (parallel per head within each token), keeping memory - * overhead minimal. + *

Processes {@code batchSize} prompt tokens simultaneously through all + * transformer layers. For each layer, Q/K/V projections, output projection, + * and FFN projections are computed via batch matmul + * ({@link FloatTensor#matmul(int, FloatTensor[], FloatTensor[], int, int)}), + * which parallelises over both output dimension and batch simultaneously. + * Attention reuses {@code state.att} sequentially per token (parallel per + * head within each token), keeping memory overhead minimal.

* - *

The logits layer is intentionally omitted — only the KV cache matters for prefill - * positions. + *

The logits layer is intentionally omitted — only the KV cache matters + * for prefill positions.

* - * @param model the LLaMA model (must carry {@link StandardWeights}) - * @param state mutable inference state (KV cache, att buffer …) - * @param tokens input token ids, {@code tokens[b]} at position {@code startPos+b} - * @param startPos sequence position of {@code tokens[0]} - * @param batchSize number of tokens in this chunk ({@code tokens.length}) + * @param model + * the LLaMA model (must carry {@link StandardWeights}) + * @param state + * mutable inference state (KV cache, att buffer …) + * @param tokens + * input token ids, {@code tokens[b]} at position {@code startPos+b} + * @param startPos + * sequence position of {@code tokens[0]} + * @param batchSize + * number of tokens in this chunk ({@code tokens.length}) */ - public static void batchForwardJavaPrefill( - Model model, State state, int[] tokens, int startPos, int batchSize) { + public static void batchForwardJavaPrefill(Model model, State state, int[] tokens, int startPos, int batchSize) { final Configuration config = model.configuration(); final StandardWeights weights = (StandardWeights) model.weights(); int dim = config.dim(); @@ -86,117 +90,75 @@ public static void batchForwardJavaPrefill( } // ── Token embeddings ────────────────────────────────────────────────── - Parallel.parallelFor( - 0, - batchSize, - b -> weights.token_embedding_table.copyTo(tokens[b] * dim, x[b], 0, dim)); + Parallel.parallelFor(0, batchSize, b -> weights.token_embedding_table.copyTo(tokens[b] * dim, x[b], 0, dim)); // ── Transformer layers ──────────────────────────────────────────────── for (int l = 0; l < config.numberOfLayers(); l++) { final int layer = l; - Parallel.parallelFor( - 0, - batchSize, - b -> - InferenceCore.rmsnorm( - xb[b], - x[b], - weights.rms_att_weight[layer], - 0, - dim, - config.rmsNormEps())); + Parallel.parallelFor(0, batchSize, b -> InferenceCore.rmsnorm(xb[b], x[b], weights.rms_att_weight[layer], 0, dim, config.rmsNormEps())); weights.wq[l].matmul(batchSize, xb, q, dim, dim); weights.wk[l].matmul(batchSize, xb, k, kvDim, dim); weights.wv[l].matmul(batchSize, xb, v, kvDim, dim); - Parallel.parallelFor( - 0, - batchSize, - b -> { - int pos = startPos + b; - for (int i = 0; i < dim; i += 2) { - int head_dim = i % headSize; - float fcr = - weights.freq_cis_real.getFloat( - pos * (headSize / 2) + (head_dim / 2)); - float fci = - weights.freq_cis_imag.getFloat( - pos * (headSize / 2) + (head_dim / 2)); - int rotn = i < kvDim ? 2 : 1; - for (int vv = 0; vv < rotn; vv++) { - FloatTensor vec = vv == 0 ? q[b] : k[b]; - float v0 = vec.getFloat(i); - float v1 = vec.getFloat(i + 1); - vec.setFloat(i, v0 * fcr - v1 * fci); - vec.setFloat(i + 1, v0 * fci + v1 * fcr); - } - } - k[b].copyTo(0, state.keyCache[layer], pos * kvDim, kvDim); - v[b].copyTo(0, state.valueCache[layer], pos * kvDim, kvDim); - }); + Parallel.parallelFor(0, batchSize, b -> { + int pos = startPos + b; + for (int i = 0; i < dim; i += 2) { + int head_dim = i % headSize; + float fcr = weights.freq_cis_real.getFloat(pos * (headSize / 2) + (head_dim / 2)); + float fci = weights.freq_cis_imag.getFloat(pos * (headSize / 2) + (head_dim / 2)); + int rotn = i < kvDim ? 2 : 1; + for (int vv = 0; vv < rotn; vv++) { + FloatTensor vec = vv == 0 ? q[b] : k[b]; + float v0 = vec.getFloat(i); + float v1 = vec.getFloat(i + 1); + vec.setFloat(i, v0 * fcr - v1 * fci); + vec.setFloat(i + 1, v0 * fci + v1 * fcr); + } + } + k[b].copyTo(0, state.keyCache[layer], pos * kvDim, kvDim); + v[b].copyTo(0, state.valueCache[layer], pos * kvDim, kvDim); + }); for (int b = 0; b < batchSize; b++) { final int pos_b = startPos + b; final int bFinal = b; - Parallel.parallelFor( - 0, - config.numberOfHeads(), - h -> { - int qOffset = h * headSize; - int attOffset = h * config.contextLength(); + Parallel.parallelFor(0, config.numberOfHeads(), h -> { + int qOffset = h * headSize; + int attOffset = h * config.contextLength(); - for (int t = 0; t <= pos_b; t++) { - int keyCacheOffset = t * kvDim + (h / kvMul) * headSize; - float score = - q[bFinal].dot( - qOffset, - state.keyCache[layer], - keyCacheOffset, - headSize) - / sqrtHeadSize; - state.att.setFloat(attOffset + t, score); - } - state.att.softmaxInPlace(attOffset, pos_b + 1); + for (int t = 0; t <= pos_b; t++) { + int keyCacheOffset = t * kvDim + (h / kvMul) * headSize; + float score = q[bFinal].dot(qOffset, state.keyCache[layer], keyCacheOffset, headSize) / sqrtHeadSize; + state.att.setFloat(attOffset + t, score); + } + state.att.softmaxInPlace(attOffset, pos_b + 1); - int xbOffset = h * headSize; - xb[bFinal].fillInPlace(xbOffset, headSize, 0f); - for (int t = 0; t <= pos_b; t++) { - int vOffset = t * kvDim + (h / kvMul) * headSize; - float a = state.att.getFloat(attOffset + t); - xb[bFinal].saxpyInPlace( - xbOffset, state.valueCache[layer], vOffset, headSize, a); - } - }); + int xbOffset = h * headSize; + xb[bFinal].fillInPlace(xbOffset, headSize, 0f); + for (int t = 0; t <= pos_b; t++) { + int vOffset = t * kvDim + (h / kvMul) * headSize; + float a = state.att.getFloat(attOffset + t); + xb[bFinal].saxpyInPlace(xbOffset, state.valueCache[layer], vOffset, headSize, a); + } + }); } weights.wo[l].matmul(batchSize, xb, xb2, dim, dim); - Parallel.parallelFor( - 0, - batchSize, - b -> { - x[b].addInPlace(xb2[b]); - InferenceCore.rmsnorm( - xb[b], - x[b], - weights.rms_ffn_weight[layer], - 0, - dim, - config.rmsNormEps()); - }); + Parallel.parallelFor(0, batchSize, b -> { + x[b].addInPlace(xb2[b]); + InferenceCore.rmsnorm(xb[b], x[b], weights.rms_ffn_weight[layer], 0, dim, config.rmsNormEps()); + }); weights.w1[l].matmul(batchSize, xb, hb, config.hiddenDim(), dim); weights.w3[l].matmul(batchSize, xb, hb2, config.hiddenDim(), dim); - Parallel.parallelFor( - 0, - batchSize, - b -> { - hb[b].mapInPlace(value -> value / (float) (1.0 + Math.exp(-value))); - hb[b].multiplyInPlace(hb2[b]); - }); + Parallel.parallelFor(0, batchSize, b -> { + hb[b].mapInPlace(value -> value / (float) (1.0 + Math.exp(-value))); + hb[b].multiplyInPlace(hb2[b]); + }); weights.w2[l].matmul(batchSize, hb, xb, dim, config.hiddenDim()); @@ -212,23 +174,23 @@ public static void batchForwardJavaPrefill( /** * GPU batched prefill forward pass (Phase 4). * - *

Copies {@code chunkSize} token embeddings into device-visible state buffers, then - * delegates graph execution to the plan. + *

Copies {@code chunkSize} token embeddings into device-visible state buffers, + * then delegates graph execution to the plan.

* - * @param model the LLaMA model - * @param state mutable inference state - * @param tokens token ids for this chunk - * @param startPos sequence position of {@code tokens[0]} - * @param chunkSize number of tokens in this chunk - * @param plan the batched prefill/decode GPU plan + * @param model + * the LLaMA model + * @param state + * mutable inference state + * @param tokens + * token ids for this chunk + * @param startPos + * sequence position of {@code tokens[0]} + * @param chunkSize + * number of tokens in this chunk + * @param plan + * the batched prefill/decode GPU plan */ - public static void batchForwardTornadoVMPrefill( - Model model, - State state, - int[] tokens, - int startPos, - int chunkSize, - TornadoVMMasterPlanBatchPrefillDecode plan) { + public static void batchForwardTornadoVMPrefill(Model model, State state, int[] tokens, int startPos, int chunkSize, TornadoVMMasterPlanBatchPrefillDecode plan) { final Configuration config = model.configuration(); final TornadoWeights weights = (TornadoWeights) model.weights(); @@ -239,16 +201,10 @@ public static void batchForwardTornadoVMPrefill( switch (weights.getWeightType()) { case F16 -> { - MemorySegment embTable = - weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); + MemorySegment embTable = weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); long dimBytes = (long) config.dim() * Short.BYTES; for (int b = 0; b < chunkSize; b++) { - MemorySegment.copy( - embTable, - (long) tokens[b] * dimBytes, - state.embeddingXBatch.getSegment(), - (long) b * dimBytes, - dimBytes); + MemorySegment.copy(embTable, (long) tokens[b] * dimBytes, state.embeddingXBatch.getSegment(), (long) b * dimBytes, dimBytes); } } case Q8_0 -> { @@ -258,17 +214,14 @@ public static void batchForwardTornadoVMPrefill( for (int b = 0; b < chunkSize; b++) { int tokenId = tokens[b]; for (int j = 0; j < dim; j++) { - int blockByteOffset = - (tokenId * blocksPerRow + j / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + int blockByteOffset = (tokenId * blocksPerRow + j / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; float scale = embTable.getHalfFloat(blockByteOffset).getFloat32(); float quant = embTable.get(blockByteOffset + 2 + j % Q8_0_BLOCK_SIZE); state.wrapXBatch.set(b * dim + j, quant * scale); } } } - default -> - throw new IllegalArgumentException( - "Unsupported weight type: " + weights.getWeightType()); + default -> throw new IllegalArgumentException("Unsupported weight type: " + weights.getWeightType()); } plan.tornadoVMForwardBatchPrefill(); @@ -277,51 +230,37 @@ public static void batchForwardTornadoVMPrefill( /** * GPU decode forward pass (Phase 4). * - *

Copies the token embedding into device-visible state, then delegates graph execution to - * the plan. + *

Copies the token embedding into device-visible state, then delegates + * graph execution to the plan.

* - * @param model the LLaMA model - * @param state mutable inference state - * @param token current token id - * @param position sequence position - * @param plan the batched prefill/decode GPU plan + * @param model + * the LLaMA model + * @param state + * mutable inference state + * @param token + * current token id + * @param position + * sequence position + * @param plan + * the batched prefill/decode GPU plan * @return logits array for token sampling */ - public static FloatArray forwardTornadoVMDecode( - Model model, - State state, - int token, - int position, - TornadoVMMasterPlanBatchPrefillDecode plan) { + public static FloatArray forwardTornadoVMDecode(Model model, State state, int token, int position, TornadoVMMasterPlanBatchPrefillDecode plan) { final Configuration config = model.configuration(); final TornadoWeights weights = (TornadoWeights) model.weights(); switch (weights.getWeightType()) { case F16 -> { - MemorySegment embTable = - weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); - MemorySegment.copy( - embTable, - (long) token * config.dim() * Short.BYTES, - state.embeddingX.getSegment(), - 0L, - (long) config.dim() * Short.BYTES); + MemorySegment embTable = weights.getTokenEmbeddingTable().asHalfFloatArray().getSegment(); + MemorySegment.copy(embTable, (long) token * config.dim() * Short.BYTES, state.embeddingX.getSegment(), 0L, (long) config.dim() * Short.BYTES); } case Q8_0 -> { - MemorySegment embTable = - weights.getTokenEmbeddingTable().asByteArray().getSegment(); + MemorySegment embTable = weights.getTokenEmbeddingTable().asByteArray().getSegment(); int blocksPerToken = (config.dim() + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; long bytesPerToken = (long) blocksPerToken * Q8_0_BLOCK_BYTES; - MemorySegment.copy( - embTable, - (long) token * bytesPerToken, - state.embeddingX.getSegment(), - 0L, - bytesPerToken); + MemorySegment.copy(embTable, (long) token * bytesPerToken, state.embeddingX.getSegment(), 0L, bytesPerToken); } - default -> - throw new IllegalArgumentException( - "Unsupported weight type: " + weights.getWeightType()); + default -> throw new IllegalArgumentException("Unsupported weight type: " + weights.getWeightType()); } return plan.tornadoVMForwardDecode(position); diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java index 76f1e9bf..a797956e 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java @@ -190,8 +190,7 @@ public static List generateTokensGPULlama(Model model, // Qwen's regular path forwards promptTokens[0] directly at position 0. // Keep the final prompt token for the B1 decode graph, which produces the // first generation logits without duplicating the ChatML start token. - boolean qwen2MoE = model.getModelType() - == org.beehive.gpullama3.model.ModelType.QWEN_2_MOE; + boolean qwen2MoE = model.getModelType() == org.beehive.gpullama3.model.ModelType.QWEN_2_MOE; int prefillTokenCount = qwen2MoE ? Math.max(0, N - 1) : N; int[] prefillSeq = new int[prefillTokenCount]; if (qwen2MoE) { @@ -205,12 +204,8 @@ public static List generateTokensGPULlama(Model model, } } - for (int chunkStart = 0; - chunkStart < prefillTokenCount && pos + chunkStart < actualMaxTokens; - chunkStart += batchSize) { - int chunkEnd = Math.min( - Math.min(chunkStart + batchSize, prefillTokenCount), - actualMaxTokens - pos); + for (int chunkStart = 0; chunkStart < prefillTokenCount && pos + chunkStart < actualMaxTokens; chunkStart += batchSize) { + int chunkEnd = Math.min(Math.min(chunkStart + batchSize, prefillTokenCount), actualMaxTokens - pos); int chunkSize = chunkEnd - chunkStart; int[] chunk = Arrays.copyOfRange(prefillSeq, chunkStart, chunkEnd); @@ -229,9 +224,7 @@ public static List generateTokensGPULlama(Model model, pos = startPosition + (qwen2MoE ? N - 1 : N); state.latestToken = currentToken; long decodeStartNanos = System.nanoTime(); - int generatedTokenBudget = qwen2MoE - ? Math.max(0, actualMaxTokens - N) - : Integer.MAX_VALUE; + int generatedTokenBudget = qwen2MoE ? Math.max(0, actualMaxTokens - N) : Integer.MAX_VALUE; // ── Decode ──────────────────────────────────────────────────────────── while (pos < actualMaxTokens && generatedTokens.size() < generatedTokenBudget) { 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 5ce078c3..c41ac03e 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -4,7 +4,6 @@ import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; - import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; @@ -121,21 +120,17 @@ protected StateFields createStateFields(Configuration configuration) { fields.att = ArrayFloatTensor.allocate(config.numberOfHeads(), config.contextLength()); fields.logits = ArrayFloatTensor.allocate(config.vocabularySize()); - fields.keyCache = - Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) - .limit(config.numberOfLayers()) - .toArray(FloatTensor[]::new); - fields.valueCache = - Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) - .limit(config.numberOfLayers()) - .toArray(FloatTensor[]::new); + fields.keyCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) + .limit(config.numberOfLayers()) + .toArray(FloatTensor[]::new); + fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)) + .limit(config.numberOfLayers()) + .toArray(FloatTensor[]::new); switch (config.quantization()) { case "FP16" -> fields.createActivationFP16(config.dim()); case "Q8_0" -> fields.createActivationQ8_0(config.dim()); - default -> - throw new UnsupportedOperationException( - "Unsupported quantization format: " + config.quantization()); + default -> throw new UnsupportedOperationException("Unsupported quantization format: " + config.quantization()); } fields.wrapX = new FloatArray(config.dim()); fields.wrapXb = new FloatArray(config.dim()); @@ -149,10 +144,8 @@ protected StateFields createStateFields(Configuration configuration) { fields.wrapK = new FloatArray(config.kvDim()); fields.wrapV = new FloatArray(config.kvDim()); - fields.wrapKeyCache = - new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); - fields.wrapValueCache = - new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapKeyCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); fields.wrapValueCache.init(0.f); fields.wrapKeyCache.init(0.f); fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); @@ -160,12 +153,9 @@ protected StateFields createStateFields(Configuration configuration) { // State invokes this override before the Qwen2State constructor body runs, // so use the Qwen2 work-group size directly instead of State.localSize. - fields.temp = - new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); - fields.tempFFN = - new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); - fields.tempLogits = - new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.temp = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.tempFFN = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); + fields.tempLogits = new FloatArray(1 + ((config.dim() + QWEN2_LOCAL_SIZE - 1) / QWEN2_LOCAL_SIZE)); return fields; } diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java index beb70562..49967115 100644 --- a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java @@ -1,7 +1,5 @@ package org.beehive.gpullama3.model.qwen2; -import static org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan.WITH_PREFILL_DECODE; - import org.beehive.gpullama3.inference.InferenceCore; import org.beehive.gpullama3.inference.InferenceEngine; import org.beehive.gpullama3.inference.InferenceEngineWithBatchPrefillDecode; @@ -20,15 +18,13 @@ import java.util.Set; import java.util.function.IntConsumer; +import static org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan.WITH_PREFILL_DECODE; + public class Qwen2MoE extends AbstractModel { Qwen2MoEConfiguration configuration; - public Qwen2MoE( - Qwen2MoEConfiguration configuration, - Tokenizer tokenizer, - Weights weights, - ChatFormat chatFormat) { + public Qwen2MoE(Qwen2MoEConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) { super(tokenizer, weights, chatFormat, null); this.configuration = configuration; } @@ -50,16 +46,14 @@ public ModelType getModelType() { @Override public State createNewState() { State state = new Qwen2MoEState(configuration(), -1); - state.latestToken = - tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); + state.latestToken = tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); return state; } @Override public State createNewState(int batchsize) { State state = new Qwen2MoEState(configuration(), batchsize); - state.latestToken = - tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); + state.latestToken = tokenizer.getSpecialTokens().get(chatFormat.chatTokens().tStartHeader()); return state; } @@ -88,65 +82,23 @@ public void forward(State state, int token, int position) { } @Override - public List generateTokens( - State state, - int startPosition, - List promptTokens, - Set stopTokens, - int maxTokens, - Sampler sampler, - boolean echo, + public List generateTokens(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated) { - return InferenceEngine.generateTokensQwen3( - this, - state, - startPosition, - promptTokens, - stopTokens, - maxTokens, - sampler, - echo, - onTokenGenerated); + return InferenceEngine.generateTokensQwen3(this, state, startPosition, promptTokens, stopTokens, maxTokens, sampler, echo, onTokenGenerated); } @Override - public List generateTokensGPU( - State state, - int startPosition, - List promptTokens, - Set stopTokens, - int maxTokens, - Sampler sampler, - boolean echo, - IntConsumer onTokenGenerated, - TornadoVMMasterPlan tornadoVMPlan) { + public List generateTokensGPU(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, + IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) { if (WITH_PREFILL_DECODE && TornadoVMMasterPlan.PREFILL_BATCH_SIZE > 1) { return InferenceEngineWithBatchPrefillDecode.generateTokensGPULlama( - this, - state, - startPosition, - promptTokens, - stopTokens, - maxTokens, - sampler, - echo, - onTokenGenerated, - tornadoVMPlan); + this, state, startPosition, promptTokens, stopTokens, maxTokens, + sampler, echo, onTokenGenerated, tornadoVMPlan); } if (WITH_PREFILL_DECODE) { - throw new UnsupportedOperationException( - "Prefill/decode on GPU not yet implemented for Qwen2-MoE"); + 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); + return InferenceEngine.generateTokensGPUQwen3(this, state, startPosition, promptTokens, + stopTokens, maxTokens, sampler, echo, onTokenGenerated, tornadoVMPlan); } } 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 848376c2..9b7a250d 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 @@ -10,7 +10,6 @@ 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; @@ -21,21 +20,20 @@ /** * 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}. + *

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 class Qwen2MoEQ8_0FFNLayers extends AbstractTransformerLayerTaskGraphs { protected final Qwen2MoEState moeState; - public Qwen2MoEQ8_0FFNLayers( - String taskGraphName, - Qwen2MoEState state, - Qwen2MoETornadoWeights weights, - Qwen2MoEConfiguration config, - SchedulerType schedulerType) { + public Qwen2MoEQ8_0FFNLayers(String taskGraphName, + Qwen2MoEState state, + Qwen2MoETornadoWeights weights, + Qwen2MoEConfiguration config, + SchedulerType schedulerType) { super(taskGraphName, state, weights, config, schedulerType); this.moeState = state; setupFFNLayers(); @@ -44,8 +42,8 @@ public Qwen2MoEQ8_0FFNLayers( /** 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 rmsNormWorker = WorkerGridFactory.createRmsNormWorker( + moeState.localSize, moeState.localSize); WorkerGrid qkvWorker = workerForRows(config.dim() + 2 * config.kvDim()); WorkerGrid qkvBiasWorker = new WorkerGrid1D(config.dim()); @@ -100,8 +98,8 @@ private WorkerGrid workerForRows(int rows) { } /** - * Creates the complete GPU TaskGraph for one Transformer layer. {@code layerIndex} selects that - * layer's weights. + * Creates the complete GPU TaskGraph for one Transformer layer. + * {@code layerIndex} selects that layer's weights. */ @Override protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { @@ -129,8 +127,7 @@ protected String predecessorGraphName(int layerIndex) { /** Uploads this layer's read-only weights on its first execution. */ protected TaskGraph configureLayerWeights(TaskGraph layer, int layerIndex) { - return layer.transferToDevice( - DataTransferMode.FIRST_EXECUTION, + return layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -152,248 +149,126 @@ protected TaskGraph configureLayerWeights(TaskGraph layer, int layerIndex) { /** Adds the normal Qwen2 attention tasks to this layer's TaskGraph. */ private void configureAttention(TaskGraph layer, int layerIndex) { - layer.task( - "attn_rms_reduce", + layer.task("attn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, - context, - moeState.temp, - moeState.wrapX, - config.dim(), - config.rmsNormEps(), - moeState.localSize); - - layer.task( - "attn_rms_qkv_projection", + 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, + 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); + config.dim(), config.dim(), config.kvDim(), LOCAL_WORK_GROUP_SIZE_ALLOC); - layer.task( - "fused_qkv_bias", + layer.task("fused_qkv_bias", TransformerComputeKernelsLayered::fusedQKvBiasAddition, - context, - moeState.wrapQ, - moeState.wrapK, - weights.q_biasLayered[layerIndex].asFloatArray(), - moeState.wrapV, + 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", + 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, + context, moeState.wrapXb, moeState.wrapX, weights.woLayered[layerIndex].asByteArray(), - config.dim(), - config.dim(), - LOCAL_WORK_GROUP_SIZE_ALLOC); + 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. + * 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", + layer.task("ffn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, - context, - moeState.tempFFN, - moeState.wrapX, - config.dim(), - config.rmsNormEps(), - moeState.localSize); - - layer.task( - "ffn_rms_apply", + 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); + context, moeState.wrapXb, moeState.wrapX, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), moeState.tempFFN); - layer.task( - "router_projection", + layer.task("router_projection", TransformerComputeKernelsLayered::matrixVectorGeneric, - context, - moeState.wrapXb, - moeState.wrapRouterLogits, + 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()); + 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()); // 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", + layer.task("routed_experts_gate_up", Qwen2MoEKernels::fusedRoutedExpertsGateUpSwiGLUQ8_0, - context, - moeState.wrapXb, - moeState.wrapSelectedExperts, - config.numberOfExpertsUsed(), + 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", + weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate, + config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + layer.task("routed_experts_down", Qwen2MoEKernels::routedExpertsDownProjectAndAccumulateQ8_0, - context, - moeState.wrapExpertGate, - moeState.wrapX, - moeState.wrapSelectedExperts, - moeState.wrapRoutingWeights, - config.numberOfExpertsUsed(), + 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); + 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, + 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); + 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. */ + /** + * 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); + 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, + 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 df4133df..d96e83e5 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java @@ -24,8 +24,8 @@ import org.beehive.gpullama3.tornadovm.plan.components.q8_0.LlamaQ8_0PlanComponents; 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.Qwen2MoEQ8_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 @@ -33,54 +33,46 @@ * Factory for {@link ForwardPlan} instances. * *

Dispatches across three axes in order: - * *

    - *
  1. Quantization ({@link GGMLType}) - *
  2. Model family ({@link org.beehive.gpullama3.model.ModelType}) - *
  3. Execution mode ({@link ExecutionMode}) + *
  4. Quantization ({@link GGMLType})
    5. + *
  5. Model family ({@link org.beehive.gpullama3.model.ModelType})
    6. + *
  6. Execution mode ({@link ExecutionMode})
    7. *
* - *

Use the typed convenience methods when the execution mode is known at the call site: - * + *

Use the typed convenience methods when the execution mode is known at the call site:

*
    - *
  • {@link #createSingleToken} — returns {@link SingleTokenForwardPlan} - *
  • {@link #createPrefillDecode} — returns {@link PrefillDecodeForwardPlan} - *
  • {@link #createBatchPrefillDecode} — returns {@link BatchPrefillDecodeForwardPlan} + *
  • {@link #createSingleToken} — returns {@link SingleTokenForwardPlan}
    • + *
  • {@link #createPrefillDecode} — returns {@link PrefillDecodeForwardPlan}
    • + *
  • {@link #createBatchPrefillDecode} — returns {@link BatchPrefillDecodeForwardPlan}
    • *
*/ // @formatter:on public class ForwardPlanFactory { - private ForwardPlanFactory() {} + private ForwardPlanFactory() { + } // ── Typed public API ────────────────────────────────────────────────────── - public static SingleTokenForwardPlan createSingleToken( - GGMLType quantization, State state, Model model) { + public static SingleTokenForwardPlan createSingleToken(GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.STANDARD, state, model); - if (plan instanceof SingleTokenForwardPlan singleToken) return singleToken; - throw new IllegalStateException( - "Expected SingleTokenForwardPlan for STANDARD mode but got " - + plan.getClass().getSimpleName()); + if (plan instanceof SingleTokenForwardPlan singleToken) + return singleToken; + throw new IllegalStateException("Expected SingleTokenForwardPlan for STANDARD mode but got " + plan.getClass().getSimpleName()); } - public static PrefillDecodeForwardPlan createPrefillDecode( - GGMLType quantization, State state, Model model) { + public static PrefillDecodeForwardPlan createPrefillDecode(GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.PREFILL_DECODE, state, model); - if (plan instanceof PrefillDecodeForwardPlan prefillDecode) return prefillDecode; - throw new IllegalStateException( - "Expected PrefillDecodeForwardPlan for PREFILL_DECODE mode but got " - + plan.getClass().getSimpleName()); + if (plan instanceof PrefillDecodeForwardPlan prefillDecode) + return prefillDecode; + throw new IllegalStateException("Expected PrefillDecodeForwardPlan for PREFILL_DECODE mode but got " + plan.getClass().getSimpleName()); } - public static BatchPrefillDecodeForwardPlan createBatchPrefillDecode( - GGMLType quantization, State state, Model model) { + public static BatchPrefillDecodeForwardPlan createBatchPrefillDecode(GGMLType quantization, State state, Model model) { ForwardPlan plan = create(quantization, ExecutionMode.BATCH_PREFILL_DECODE, state, model); if (plan instanceof BatchPrefillDecodeForwardPlan batchPrefillDecode) return batchPrefillDecode; - throw new IllegalStateException( - "Expected BatchPrefillDecodeForwardPlan for BATCH_PREFILL_DECODE mode but got " - + plan.getClass().getSimpleName()); + throw new IllegalStateException("Expected BatchPrefillDecodeForwardPlan for BATCH_PREFILL_DECODE mode but got " + plan.getClass().getSimpleName()); } // ── Generic dispatch ────────────────────────────────────────────────────── @@ -91,9 +83,7 @@ static ForwardPlan create(GGMLType quantization, ExecutionMode mode, State state case Q8_0 -> createQ8_0Plan(mode, state, model); case F32 -> throw new UnsupportedOperationException("F32 plans not yet implemented"); case Q4_0 -> throw new UnsupportedOperationException("Q4_0 plans not yet implemented"); - default -> - throw new UnsupportedOperationException( - "Quantization not supported: " + quantization); + default -> throw new UnsupportedOperationException("Quantization not supported: " + quantization); }; } @@ -109,9 +99,7 @@ private static ForwardPlan createFP16Plan(ExecutionMode mode, State state, Model case PHI_3 -> createPhi3FP16Plan(mode, (Phi3State) state, model); case GRANITE -> createGraniteFP16Plan(mode, (GraniteState) state, model); case DEEPSEEK_R1_DISTILL_QWEN -> createQwen2FP16Plan(mode, (Qwen2State) state, model); - default -> - throw new UnsupportedOperationException( - "F16 not supported for model: " + model.getModelType()); + default -> throw new UnsupportedOperationException("F16 not supported for model: " + model.getModelType()); }; } @@ -128,150 +116,114 @@ private static ForwardPlan createQ8_0Plan(ExecutionMode mode, State state, Model case PHI_3 -> createPhi3Q8_0Plan(mode, (Phi3State) state, model); case GRANITE -> createGraniteQ8_0Plan(mode, (GraniteState) state, model); case DEEPSEEK_R1_DISTILL_QWEN -> createQwen2Q8_0Plan(mode, (Qwen2State) state, model); - default -> - throw new UnsupportedOperationException( - "Q8_0 not supported for model: " + model.getModelType()); + default -> throw new UnsupportedOperationException("Q8_0 not supported for model: " + model.getModelType()); }; } // ── Model+quant helpers — Llama (all 3 modes supported) ────────────────── - private static ForwardPlan createLlamaFP16Plan( - ExecutionMode mode, LlamaState state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = - new LlamaFP16PlanComponents(state, model); + private static ForwardPlan createLlamaFP16Plan(ExecutionMode mode, LlamaState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = new LlamaFP16PlanComponents(state, model); return switch (mode) { case STANDARD -> new SingleTokenForwardPlan(model, components); case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> - new BatchPrefillDecodeForwardPlan( - model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createLlamaQ8_0Plan( - ExecutionMode mode, LlamaState state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = - new LlamaQ8_0PlanComponents(state, model); + private static ForwardPlan createLlamaQ8_0Plan(ExecutionMode mode, LlamaState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = new LlamaQ8_0PlanComponents(state, model); return switch (mode) { case STANDARD -> new SingleTokenForwardPlan(model, components); case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> - new BatchPrefillDecodeForwardPlan( - model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } // ── Model+quant helpers — STANDARD only ────────────────────────────────── - private static ForwardPlan createMistralFP16Plan( - ExecutionMode mode, LlamaState state, Model model) { + private static ForwardPlan createMistralFP16Plan(ExecutionMode mode, LlamaState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for MISTRAL + F16"); return new SingleTokenForwardPlan(model, new MistralFP16PlanComponents(state, model)); } - private static ForwardPlan createMistralQ8_0Plan( - ExecutionMode mode, LlamaState state, Model model) { + private static ForwardPlan createMistralQ8_0Plan(ExecutionMode mode, LlamaState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for MISTRAL + Q8_0"); return new SingleTokenForwardPlan(model, new MistralQ8_0PlanComponents(state, model)); } - private static ForwardPlan createDevstralFP16Plan( - ExecutionMode mode, DevstralState state, Model model) { + private static ForwardPlan createDevstralFP16Plan(ExecutionMode mode, DevstralState state, Model model) { if (mode != ExecutionMode.STANDARD) - throw new UnsupportedOperationException( - mode + " not yet supported for DEVSTRAL_2 + F16"); + throw new UnsupportedOperationException(mode + " not yet supported for DEVSTRAL_2 + F16"); return new SingleTokenForwardPlan(model, new DevstralFP16PlanComponents(state, model)); } - private static ForwardPlan createDevstralQ8_0Plan( - ExecutionMode mode, DevstralState state, Model model) { + private static ForwardPlan createDevstralQ8_0Plan(ExecutionMode mode, DevstralState state, Model model) { if (mode != ExecutionMode.STANDARD) - throw new UnsupportedOperationException( - mode + " not yet supported for DEVSTRAL_2 + Q8_0"); + throw new UnsupportedOperationException(mode + " not yet supported for DEVSTRAL_2 + Q8_0"); return new SingleTokenForwardPlan(model, new DevstralQ8_0PlanComponents(state, model)); } - private static ForwardPlan createQwen2FP16Plan( - ExecutionMode mode, Qwen2State state, Model model) { + private static ForwardPlan createQwen2FP16Plan(ExecutionMode mode, Qwen2State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2 + F16"); return new SingleTokenForwardPlan(model, new Qwen2FP16PlanComponents(state, model)); } - private static ForwardPlan createQwen2Q8_0Plan( - ExecutionMode mode, Qwen2State state, Model model) { + private static ForwardPlan createQwen2Q8_0Plan(ExecutionMode mode, Qwen2State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2 + Q8_0"); return new SingleTokenForwardPlan(model, new Qwen2Q8_0PlanComponents(state, model)); } - private static ForwardPlan createQwen2MoEQ8_0Plan( - ExecutionMode mode, Qwen2MoEState state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = - new Qwen2MoEQ8_0PlanComponents(state, model); + private static ForwardPlan createQwen2MoEQ8_0Plan(ExecutionMode mode, Qwen2MoEState state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = new Qwen2MoEQ8_0PlanComponents(state, model); return switch (mode) { - case STANDARD -> new SingleTokenForwardPlan(model, components); - case PREFILL_DECODE -> - throw new UnsupportedOperationException( - mode + " not yet supported for QWEN_2_MOE + Q8_0"); - case BATCH_PREFILL_DECODE -> - new BatchPrefillDecodeForwardPlan( - model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2_MOE + Q8_0"); + case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createQwen3FP16Plan( - ExecutionMode mode, Qwen3State state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = - new Qwen3FP16PlanComponents(state, model); + private static ForwardPlan createQwen3FP16Plan(ExecutionMode mode, Qwen3State state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = new Qwen3FP16PlanComponents(state, model); return switch (mode) { - case STANDARD -> new SingleTokenForwardPlan(model, components); - case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> - new BatchPrefillDecodeForwardPlan( - model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); + case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createQwen3Q8_0Plan( - ExecutionMode mode, Qwen3State state, Model model) { - BatchPrefillDecodeForwardPlanComponents components = - new Qwen3Q8_0PlanComponents(state, model); + private static ForwardPlan createQwen3Q8_0Plan(ExecutionMode mode, Qwen3State state, Model model) { + BatchPrefillDecodeForwardPlanComponents components = new Qwen3Q8_0PlanComponents(state, model); return switch (mode) { - case STANDARD -> new SingleTokenForwardPlan(model, components); - case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); - case BATCH_PREFILL_DECODE -> - new BatchPrefillDecodeForwardPlan( - model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); + case STANDARD -> new SingleTokenForwardPlan(model, components); + case PREFILL_DECODE -> new PrefillDecodeForwardPlan(model, components); + case BATCH_PREFILL_DECODE -> new BatchPrefillDecodeForwardPlan(model, components, TornadoVMMasterPlan.PREFILL_BATCH_SIZE); }; } - private static ForwardPlan createPhi3FP16Plan( - ExecutionMode mode, Phi3State state, Model model) { + private static ForwardPlan createPhi3FP16Plan(ExecutionMode mode, Phi3State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for PHI_3 + F16"); return new SingleTokenForwardPlan(model, new Phi3FP16PlanComponents(state, model)); } - private static ForwardPlan createPhi3Q8_0Plan( - ExecutionMode mode, Phi3State state, Model model) { + private static ForwardPlan createPhi3Q8_0Plan(ExecutionMode mode, Phi3State state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for PHI_3 + Q8_0"); return new SingleTokenForwardPlan(model, new Phi3Q8_0PlanComponents(state, model)); } - private static ForwardPlan createGraniteFP16Plan( - ExecutionMode mode, GraniteState state, Model model) { + private static ForwardPlan createGraniteFP16Plan(ExecutionMode mode, GraniteState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for GRANITE + F16"); return new SingleTokenForwardPlan(model, new GraniteFP16PlanComponents(state, model)); } - private static ForwardPlan createGraniteQ8_0Plan( - ExecutionMode mode, GraniteState state, Model model) { + private static ForwardPlan createGraniteQ8_0Plan(ExecutionMode mode, GraniteState state, Model model) { if (mode != ExecutionMode.STANDARD) throw new UnsupportedOperationException(mode + " not yet supported for GRANITE + Q8_0"); return new SingleTokenForwardPlan(model, new GraniteQ8_0PlanComponents(state, model)); 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 index 6afa7719..ea17283c 100644 --- 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 @@ -77,13 +77,11 @@ public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int @Override public AbstractLogitsTaskGraph singleTokenLogits(String previousGraphId) { - return new LogitsQ8_0Layer( - "logits", state, weights, config, previousGraphId, schedulerType); + return new LogitsQ8_0Layer("logits", state, weights, config, previousGraphId, schedulerType); } @Override public AbstractLogitsTaskGraph decodeLogits(String previousGraphId) { - return new LogitsQ8_0LayerDecode( - "logits", state, weights, config, previousGraphId, schedulerType); + return new LogitsQ8_0LayerDecode("logits", state, weights, config, previousGraphId, schedulerType); } } From d450f18b9d7e1d6c45e80c9d6393830a3765c7a8 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Fri, 14 Aug 2026 18:42:03 +0100 Subject: [PATCH 3/4] Fix Qwen decode position and token budget --- .../gpullama3/inference/InferenceEngine.java | 22 +++++++++++++------ ...InferenceEngineWithBatchPrefillDecode.java | 2 +- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index 85ca2ec5..049d9b78 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java @@ -166,6 +166,7 @@ public static List generateTokensQwen3(Model model, State state, int st // Storage for generated tokens List generatedTokens = new ArrayList<>(); + int generatedTokenBudget = Math.max(0, maxTokens - startPosition - promptTokens.size()); // Initialize token variables int currentToken = state.latestToken; // BOS? @@ -188,8 +189,11 @@ public static List generateTokensQwen3(Model model, State state, int st if (echo) { System.err.print(Tokenizer.replaceControlCharacters(model.tokenizer().decode(List.of(nextToken)))); } - // We have reached the last prompt token and computed the first response-token. - position++; // The current logit belongs to the next position + // The last prompt token produced the first response-token logits. + // The for-loop advances to the next sequence position. + if (generatedTokenBudget == 0) { + break; + } } else { // Mark the start of actual generation (after prompt processing) if (inferenceStartNanos == 0) { @@ -216,7 +220,7 @@ public static List generateTokensQwen3(Model model, State state, int st } // Check for stop condition - if (stopTokens.contains(nextToken)) { + if (generatedTokens.size() >= generatedTokenBudget || stopTokens.contains(nextToken)) { break; } @@ -393,6 +397,7 @@ public static List generateTokensGPUQwen3(Model model, State state, int // prompt is longer than the token budget (actualMaxTokens), the difference is // negative and would throw IllegalArgumentException("Illegal Capacity"). List generatedTokens = new ArrayList<>(Math.max(0, Math.min(256, actualMaxTokens - promptTokens.size()))); // Conservative estimate + int generatedTokenBudget = Math.max(0, actualMaxTokens - startPosition - promptTokens.size()); // Initialize token variables int currentToken = state.latestToken; // BOS? @@ -428,8 +433,11 @@ public static List generateTokensGPUQwen3(Model model, State state, int if (echo) { System.err.print(Tokenizer.replaceControlCharacters(model.tokenizer().decode(List.of(nextToken)))); } - // We have reached the last prompt token and computed the first response-token. - position++; // The current logit belongs to the next position + // The last prompt token produced the first response-token logits. + // The for-loop advances to the next sequence position. + if (generatedTokenBudget == 0) { + break; + } } else { // Mark the start of actual generation (after prompt processing) if (inferenceStartNanos == 0) { @@ -456,7 +464,7 @@ public static List generateTokensGPUQwen3(Model model, State state, int } // Check for stop condition - if (stopTokens.contains(nextToken)) { + if (generatedTokens.size() >= generatedTokenBudget || stopTokens.contains(nextToken)) { break; } @@ -678,4 +686,4 @@ public static List generateTokensGPUGranite(Model model, State state, i return generatedTokens; } -} \ No newline at end of file +} diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java index a797956e..bd09e6b2 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java @@ -224,7 +224,7 @@ public static List generateTokensGPULlama(Model model, pos = startPosition + (qwen2MoE ? N - 1 : N); state.latestToken = currentToken; long decodeStartNanos = System.nanoTime(); - int generatedTokenBudget = qwen2MoE ? Math.max(0, actualMaxTokens - N) : Integer.MAX_VALUE; + int generatedTokenBudget = Math.max(0, actualMaxTokens - startPosition - N); // ── Decode ──────────────────────────────────────────────────────────── while (pos < actualMaxTokens && generatedTokens.size() < generatedTokenBudget) { From 38c85922ce0c2aa5e1b82f577c09ff00b525a761 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Fri, 14 Aug 2026 21:48:55 +0100 Subject: [PATCH 4/4] Clean up Qwen2-MoE batch inference code --- .../InferenceCoreBatchPrefillDecode.java | 4 ++-- .../InferenceEngineWithBatchPrefillDecode.java | 4 +--- .../inference/state/Qwen2MoEState.java | 3 --- .../tornadovm/kernels/Qwen2MoEBatchKernels.java | 17 +---------------- .../prefill/Qwen2MoEQ8_0LayersBatchPrefill.java | 3 --- 5 files changed, 4 insertions(+), 27 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java index 9e3601a8..6f78fd2d 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java @@ -178,7 +178,7 @@ public static void batchForwardJavaPrefill(Model model, State state, int[] token * then delegates graph execution to the plan.

* * @param model - * the LLaMA model + * the model * @param state * mutable inference state * @param tokens @@ -234,7 +234,7 @@ public static void batchForwardTornadoVMPrefill(Model model, State state, int[] * graph execution to the plan.

* * @param model - * the LLaMA model + * the model * @param state * mutable inference state * @param token diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java index bd09e6b2..d811c214 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java @@ -147,9 +147,7 @@ public static List generateTokensLlama(Model model, // @formatter:on /** - * LLaMA batched GPU prefill token generation (GPU, Phase 4). - * - *

FP16 only; Q8_0 throws {@link UnsupportedOperationException}.

+ * Batched GPU prefill followed by single-token decode. * *

Split loop:

*
    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 c41ac03e..234978db 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -48,7 +48,6 @@ public class Qwen2MoEState extends Qwen2State { public final FloatArray wrapRoutingWeightsBatch; public final IntArray wrapGroupedAssignmentIds; public final IntArray wrapGroupedPositionByAssignment; - public final IntArray wrapExpertOffsets; public final FloatArray wrapGroupedExpertHidden; public final FloatArray wrapGroupedExpertDown; public final FloatArray wrapSharedHiddenBatch; @@ -81,7 +80,6 @@ public Qwen2MoEState(Configuration config, int batchsize) { this.wrapRoutingWeightsBatch = new FloatArray(assignments); this.wrapGroupedAssignmentIds = new IntArray(assignments); this.wrapGroupedPositionByAssignment = new IntArray(assignments); - this.wrapExpertOffsets = new IntArray(c.numberOfExperts() + 1); this.wrapGroupedExpertHidden = new FloatArray(assignments * c.moeHiddenDim()); this.wrapGroupedExpertDown = new FloatArray(assignments * c.dim()); this.wrapSharedHiddenBatch = new FloatArray(gpuBatchSize * c.sharedExpertHiddenDim()); @@ -93,7 +91,6 @@ public Qwen2MoEState(Configuration config, int batchsize) { this.wrapRoutingWeightsBatch = null; this.wrapGroupedAssignmentIds = null; this.wrapGroupedPositionByAssignment = null; - this.wrapExpertOffsets = null; this.wrapGroupedExpertHidden = null; this.wrapGroupedExpertDown = null; this.wrapSharedHiddenBatch = null; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java index eab69eac..ee685e7b 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEBatchKernels.java @@ -237,13 +237,12 @@ public static void batchedSoftmaxAndTopK( } } - /** Groups token-expert assignments by expert and records each expert's range. */ + /** Groups token-expert assignments by expert. */ public static void groupAssignmentsByExpert( KernelContext context, IntArray selectedExperts, IntArray groupedAssignmentIds, IntArray groupedPositionByAssignment, - IntArray expertOffsets, IntArray activeBatchSizeHolder, int numberOfExperts, int topK) { @@ -256,30 +255,16 @@ public static void groupAssignmentsByExpert( int numberOfAssignments = activeBatchSizeHolder.get(0) * topK; int groupedPosition = 0; - // Visit each expert in expert-ID order. for (int expert = 0; expert < numberOfExperts; expert++) { - // Record where this expert's assignments begin. - expertOffsets.set(expert, groupedPosition); - - // Find all active assignments that selected this expert. for (int assignment = 0; assignment < numberOfAssignments; assignment++) { - int selectedExpert = selectedExperts.get(assignment); - if (selectedExpert == expert) { - // grouped position -> original assignment ID groupedAssignmentIds.set(groupedPosition, assignment); - - // original assignment ID -> grouped position groupedPositionByAssignment.set(assignment, groupedPosition); - groupedPosition++; } } } - - // Record the end of the final expert's assignment range. - expertOffsets.set(numberOfExperts, groupedPosition); } /** Computes routed Gate/Up projections in expert-grouped assignment order. */ diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java index e973937d..1b159f0b 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen2MoEQ8_0LayersBatchPrefill.java @@ -114,7 +114,6 @@ private void configureDataTransfers(TaskGraph layer, int layerIndex) { state.wrapRoutingWeightsBatch, state.wrapGroupedAssignmentIds, state.wrapGroupedPositionByAssignment, - state.wrapExpertOffsets, state.wrapGroupedExpertHidden, state.wrapGroupedExpertDown, state.wrapSharedHiddenBatch, @@ -141,7 +140,6 @@ private void configureDataTransfers(TaskGraph layer, int layerIndex) { state.wrapRoutingWeightsBatch, state.wrapGroupedAssignmentIds, state.wrapGroupedPositionByAssignment, - state.wrapExpertOffsets, state.wrapGroupedExpertHidden, state.wrapGroupedExpertDown, state.wrapSharedHiddenBatch, @@ -319,7 +317,6 @@ private void configureMoE(TaskGraph layer, int layerIndex) { state.wrapSelectedExpertsBatch, state.wrapGroupedAssignmentIds, state.wrapGroupedPositionByAssignment, - state.wrapExpertOffsets, state.activeBatchSizeHolder, config.numberOfExperts(), topK);