From 97ef9d29d27a78e9a5388393571a3d1219680ab5 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sat, 4 Jul 2026 01:17:05 +0100 Subject: [PATCH 01/10] Add Qwen1.5-MoE (qwen2moe) CPU inference support Adds config/weights/state/model/loader classes for Qwen2-MoE and implements the router + top-k expert + shared-expert FFN block in InferenceCore.forwardJavaQwen2MoE, plus architecture detection for "qwen2moe" GGUF files. --- .../gpullama3/inference/InferenceCore.java | 206 +++++++++++++++++- .../inference/state/Qwen2MoEState.java | 43 ++++ .../standard/Qwen2MoEStandardWeights.java | 80 +++++++ .../beehive/gpullama3/model/ModelType.java | 8 + .../gpullama3/model/loader/ModelLoader.java | 6 + .../model/loader/Qwen2MoEModelLoader.java | 123 +++++++++++ .../gpullama3/model/qwen2/Qwen2MoE.java | 88 ++++++++ .../model/qwen2/Qwen2MoEConfiguration.java | 48 ++++ 8 files changed, 598 insertions(+), 4 deletions(-) create mode 100644 src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java create mode 100644 src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java create mode 100644 src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java create mode 100644 src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java create mode 100644 src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java index 25e46972..edd383b6 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java @@ -1,13 +1,13 @@ package org.beehive.gpullama3.inference; import org.beehive.gpullama3.auxiliary.Parallel; +import org.beehive.gpullama3.inference.state.Qwen2MoEState; +import org.beehive.gpullama3.inference.weights.standard.*; +import org.beehive.gpullama3.model.qwen2.Qwen2MoE; +import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; import org.beehive.gpullama3.tensor.standard.FloatTensor; import org.beehive.gpullama3.inference.state.Phi3State; import org.beehive.gpullama3.inference.state.State; -import org.beehive.gpullama3.inference.weights.standard.Phi3StandardWeights; -import org.beehive.gpullama3.inference.weights.standard.Qwen2StandardWeights; -import org.beehive.gpullama3.inference.weights.standard.Qwen3StandardWeights; -import org.beehive.gpullama3.inference.weights.standard.StandardWeights; import org.beehive.gpullama3.inference.weights.tornado.TornadoWeights; import org.beehive.gpullama3.model.Configuration; import org.beehive.gpullama3.model.Model; @@ -259,6 +259,204 @@ public static FloatTensor forwardJavaDevstral(Model model, State state, int toke return state.logits; } + public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int token, int position) { + final Qwen2MoEConfiguration config = (Qwen2MoEConfiguration) model.configuration(); + final Qwen2MoEStandardWeights weights = (Qwen2MoEStandardWeights) model.weights(); + final Qwen2MoEState moeState = (Qwen2MoEState) state; + int dim = config.dim(); + int headSize = config.headSize(); + int kvDim = (config.dim() * config.numberOfKeyValueHeads()) / config.numberOfHeads(); + int kvMul = config.numberOfHeads() / config.numberOfKeyValueHeads(); // integer multiplier of the kv sharing in multiquery + float sqrtHeadSize = (float) Math.sqrt(headSize); + + weights.token_embedding_table.copyTo(token * dim, state.x, 0, dim); + + // forward all the layers + for (int l = 0; l < config.numberOfLayers(); l++) { + // attention rmsnorm + final int curLayer = l; + rmsnorm(state.xb, state.x, weights.rms_att_weight[curLayer], 0, dim, config.rmsNormEps()); + + // qkv matmuls for this position + weights.wq[l].matmul(state.xb, state.q, dim, dim); + weights.wk[l].matmul(state.xb, state.k, kvDim, dim); + weights.wv[l].matmul(state.xb, state.v, kvDim, dim); + + // qkv additions with qkv bias + state.q.addInPlace(weights.q_bias[curLayer]); + state.k.addInPlace(weights.k_bias[curLayer]); + state.v.addInPlace(weights.v_bias[curLayer]); + + // RoPE relative positional encoding: complex-valued rotate q and k in each head + // GPT-NeoX style RoPE, real/imaginary components are stored with a headSize/2 offset per head, instead of consecutive. + for (int h = 0; h < config.numberOfHeads(); ++h) { + int rotn = h < config.numberOfKeyValueHeads() ? 2 : 1; // how many vectors? 2 = q & k, 1 = q only + int poffset = h * headSize; + for (int i0 = 0; i0 < headSize; i0 += 2) { + int ic = i0 / 2; + float fcr = weights.freq_cis_real.getFloat((position) * (headSize / 2) + ic); + float fci = weights.freq_cis_imag.getFloat((position) * (headSize / 2) + ic); + for (int vi = 0; vi < rotn; vi++) { + FloatTensor vec = (vi == 0) ? state.q : state.k; // the vector to rotate (query or key) + float v0 = vec.getFloat(poffset + ic); + float v1 = vec.getFloat(poffset + ic + headSize / 2); + vec.setFloat(poffset + ic, v0 * fcr - v1 * fci); + vec.setFloat(poffset + ic + headSize / 2, v0 * fci + v1 * fcr); + } + } + } + + // save key,value at this time step (position) to our kv cache + //int loff = l * config.seq_len * kvDim; // kv cache layer offset for convenience + state.k.copyTo(0, state.keyCache[curLayer], position * kvDim, kvDim); + state.v.copyTo(0, state.valueCache[curLayer], position * kvDim, kvDim); + + // multihead attention. iterate over all heads + Parallel.parallelFor(0, config.numberOfHeads(), h -> { + // get the query vector for this head + // float* q = s.q + h * headSize; + int qOffset = h * headSize; + + // attention scores for this head + // float* att = s.att + h * config.seq_len; + int attOffset = h * config.contextLength(); + + // iterate over all timesteps, including the current one + for (int t = 0; t <= position; t++) { + // get the key vector for this head and at this timestep + // float* k = s.key_cache + loff + t * dim + h * headSize; + int keyCacheOffset = /* loff + */ t * kvDim + (h / kvMul) * headSize; + // calculate the attention score as the dot product of q and k + float score = state.q.dot(qOffset, state.keyCache[curLayer], keyCacheOffset, headSize); + score /= sqrtHeadSize; + // save the score to the attention buffer + state.att.setFloat(attOffset + t, score); + } + + // softmax the scores to get attention weights, from 0..position inclusively + state.att.softmaxInPlace(attOffset, position + 1); + + // weighted sum of the values, store back into xb + // float* xb = s.xb + h * headSize; + int xbOffset = h * headSize; + // memset(xb, 0, headSize * sizeof(float)); + state.xb.fillInPlace(xbOffset, headSize, 0f); + + for (int t = 0; t <= position; t++) { + // get the value vector for this head and at this timestep + // float* v = s.value_cache + loff + t * dim + h * headSize;C + int vOffset = /* loff + */ t * kvDim + (h / kvMul) * headSize; + // get the attention weight for this timestep + float a = state.att.getFloat(attOffset + t); + // accumulate the weighted value into xb + state.xb.saxpyInPlace(xbOffset, state.valueCache[curLayer], vOffset, headSize, a); + } + }); + + // final matmul to get the output of the attention + weights.wo[l].matmul(state.xb, state.xb2, dim, dim); + + // residual connection back into x + state.x.addInPlace(state.xb2); + + // ========================= FFN block: MoE ========================= + // NOTE: this scaffold references fields you still need to create: + // Qwen2MoEConfiguration: numberOfExperts(), numberOfExpertsUsed(), + // moeHiddenDim(), sharedExpertHiddenDim() + // Qwen2MoEStandardWeights: routerGate[], gateExps[], upExps[], downExps[], + // sharedGate[], sharedUp[], sharedDown[], sharedGateInp[] + // Qwen2MoEState buffers: routerLogits, hbE, hbE2, hbS, hbS2, yTmp + // Once those exist, change the two casts at the top of this method to the MoE + // types, then replace the TODO lines below with the real calls shown. + + // FFN pre-norm (same as dense): state.xb = rmsnorm(state.x) + rmsnorm(state.xb, state.x, weights.rms_ffn_weight[curLayer], 0, dim, config.rmsNormEps()); + + int nExp = config.numberOfExperts(); // 60 + int topK = config.numberOfExpertsUsed(); // 4 + int moeHid = config.moeHiddenDim(); // 1408 + + // --- Step A: router scores --- + // routerLogits[nExp] = routerGate[l] · xb + weights.routerGate[l].matmul(state.xb, moeState.routerLogits, nExp, dim); + // --- Step B: pick top-K experts + normalize their weights --- + int[] idx = new int[topK]; float[] wgt = new float[topK]; + for (int i = 0; i < topK; i++) { + float best = Float.NEGATIVE_INFINITY; + int index = -1; + for (int j = 0; j < nExp; j++) { + if (moeState.routerLogits.getFloat(j) > best) { + best = moeState.routerLogits.getFloat(j); + index = j; + } + } + idx[i] = index; + wgt[i] = best; + moeState.routerLogits.setFloat(index, Float.NEGATIVE_INFINITY); + } + + float maxVal = wgt[0]; + for (int i = 1; i < topK; i++) { + if (wgt[i] > maxVal) maxVal = wgt[i]; + } + + float sum = 0f; + for (int i = 0; i < topK; i++) { + wgt[i] = (float) Math.exp(wgt[i] - maxVal); + sum += wgt[i]; + } + + for (int i = 0; i < topK; i++) { + wgt[i] /= sum; + } + + +// --- Step C: compute each selected expert and accumulate into x --- + for (int j = 0; j < topK; j++) { + int e = idx[j]; + int baseGU = e * moeHid * dim; // gate/up expert offset + int baseD = e * dim * moeHid; // down expert offset (rows/cols swapped) + matmulExpert(weights.gateExps[l], baseGU, state.xb, moeState.hbE, moeHid, dim); + matmulExpert(weights.upExps[l], baseGU, state.xb, moeState.hbE2, moeHid, dim); + moeState.hbE.mapInPlace(v -> v / (float)(1.0 + Math.exp(-v))); // silu + moeState.hbE.multiplyInPlace(moeState.hbE2); // gate ⊙ up + matmulExpert(weights.downExps[l], baseD, moeState.hbE, moeState.yTmp, dim, moeHid); + state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, wgt[j]); // x += wgt*y + } + + // --- Step D: shared expert (always-on, plain dense FFN) --- + // int sHid = config.sharedExpertHiddenDim(); // 5632 + int sHid = config.sharedExpertHiddenDim(); // 5632 + + weights.sharedGate[l].matmul(state.xb, moeState.hbS, sHid, dim); + weights.sharedUp[l].matmul(state.xb, moeState.hbS2, sHid, dim); + moeState.hbS.mapInPlace(v -> v / (float)(1.0 + Math.exp(-v))); // silu + moeState.hbS.multiplyInPlace(moeState.hbS2); + weights.sharedDown[l].matmul(moeState.hbS, moeState.yTmp, dim, sHid); + +// 算 shared expert 的门控值 g + float gateScore = weights.sharedGateInp[l].dot(0, state.xb, 0, dim); + float g = 1f / (1f + (float) Math.exp(-gateScore)); + + state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, g); + // ================================================================== + } + + // final rmsnorm + classifier (same as dense Qwen2) + rmsnorm(state.x, state.x, weights.rms_final_weight, 0, dim, config.rmsNormEps()); + weights.wcls.matmul(state.x, state.logits, config.vocabularySize(), dim); + + return state.logits; + } + + /** + * Like {@link FloatTensor#matmul}, but reads the weight matrix starting at element + * offset {@code base} instead of 0 — so it can target ONE expert inside a stacked + * {@code [nExpert × d0 × d1]} tensor. out[d0] = (d0×d1 sub-matrix at base) · in[d1]. + */ + private static void matmulExpert(FloatTensor w, int base, FloatTensor in, FloatTensor out, int d0, int d1) { + Parallel.parallelFor(0, d0, i -> out.setFloat(i, w.dot(base + i * d1, in, 0, d1))); + } public static FloatTensor forwardJavaQwen2(Model model, State state, int token, int position) { final Qwen2Configuration config = (Qwen2Configuration) model.configuration(); diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java new file mode 100644 index 00000000..f70084d5 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -0,0 +1,43 @@ +package org.beehive.gpullama3.inference.state; + +import org.beehive.gpullama3.model.Configuration; +import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; +import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; +import org.beehive.gpullama3.tensor.standard.FloatTensor; + +public class Qwen2MoEState extends Qwen2State { + + // Router output scores, one per expert (length = numberOfExperts). + // Used to pick the top-k experts for the current token. + public final FloatTensor routerLogits; + + // Scratch buffers for a single routed expert's internal FFN (length = moeHiddenDim). + // hbE holds gate_proj(xb), hbE2 holds up_proj(xb); combined via silu(hbE) * hbE2. + public final FloatTensor hbE; + public final FloatTensor hbE2; + + // Scratch buffers for the shared expert's internal FFN (length = sharedExpertHiddenDim). + // Same role as hbE/hbE2, but for the always-on shared expert instead of a routed one. + public final FloatTensor hbS; + public final FloatTensor hbS2; + + // Temporary holder for a single expert's down-projected output (length = dim), + // before it is weighted and accumulated into the residual stream (state.x). + public final FloatTensor yTmp; + + public Qwen2MoEState(Configuration config, int batchsize) { + super(config, batchsize); // allocate all the regular Qwen2State buffers first + Qwen2MoEConfiguration c = (Qwen2MoEConfiguration) config; + this.routerLogits = ArrayFloatTensor.allocate(c.numberOfExperts()); + this.hbE = ArrayFloatTensor.allocate(c.moeHiddenDim()); + this.hbE2 = ArrayFloatTensor.allocate(c.moeHiddenDim()); + this.hbS = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); + this.hbS2 = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); + this.yTmp = ArrayFloatTensor.allocate(c.dim()); + } + + @Override + protected StateFields createStateFields(Configuration config) { + return super.createStateFields(config); + } +} diff --git a/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java b/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java new file mode 100644 index 00000000..75f2e3c4 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java @@ -0,0 +1,80 @@ +package org.beehive.gpullama3.inference.weights.standard; + +import org.beehive.gpullama3.tensor.GGMLType; +import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; +import org.beehive.gpullama3.tensor.standard.FloatTensor; + +public class Qwen2MoEStandardWeights extends StandardWeights{ + // Qwen2-specific weights + public final FloatTensor[] q_bias, k_bias, v_bias; + public final FloatTensor[] routerGate; + public final FloatTensor[] gateExps; + public final FloatTensor[] upExps; // [层数] + public final FloatTensor[] downExps; // [层数] + public final FloatTensor[] sharedGate; // [层数]，shared expert + public final FloatTensor[] sharedUp; // [层数] + public final FloatTensor[] sharedDown; // [层数] + public final FloatTensor[] sharedGateInp; // [层数]，shared expert 的 sigmoid gate 输入 + + public Qwen2MoEStandardWeights( + FloatTensor token_embedding_table, + FloatTensor[] rms_att_weight, + FloatTensor[] wq, + FloatTensor[] wk, + FloatTensor[] wv, + FloatTensor[] q_bias, + FloatTensor[] k_bias, + FloatTensor[] v_bias, + FloatTensor[] wo, + FloatTensor[] rms_ffn_weight, + FloatTensor[] w1, + FloatTensor[] w2, + FloatTensor[] w3, + FloatTensor[] routerGate, + FloatTensor[] gateExps, + FloatTensor[] upExps, + FloatTensor[] downExps, + FloatTensor[] sharedGate, + FloatTensor[] sharedUp, + FloatTensor[] sharedDown, + FloatTensor[] sharedGateInp, + FloatTensor rms_final_weight, + ArrayFloatTensor freq_cis_real, + ArrayFloatTensor freq_cis_imag, + FloatTensor wcls, + GGMLType weightType) { + // call to StandardWeights constructor + super(token_embedding_table, + rms_att_weight, + wq, + wk, + wv, + wo, + rms_ffn_weight, + w1, + w2, + w3, + rms_final_weight, + freq_cis_real, + freq_cis_imag, + wcls, + weightType); + // init Qwen2-specific fields + this.q_bias = q_bias; + this.k_bias = k_bias; + this.v_bias = v_bias; + this.routerGate = routerGate; + this.gateExps = gateExps; + this.upExps = upExps; + this.downExps = downExps; + this.sharedGate = sharedGate; + this.sharedUp = sharedUp; + this.sharedDown = sharedDown; + this.sharedGateInp = sharedGateInp; + } + + @Override + public GGMLType getWeightType() { + return weightType; + } +} diff --git a/src/main/java/org/beehive/gpullama3/model/ModelType.java b/src/main/java/org/beehive/gpullama3/model/ModelType.java index 0659da7d..d3d7fa63 100644 --- a/src/main/java/org/beehive/gpullama3/model/ModelType.java +++ b/src/main/java/org/beehive/gpullama3/model/ModelType.java @@ -7,6 +7,7 @@ import org.beehive.gpullama3.model.loader.MistralModelLoader; import org.beehive.gpullama3.model.loader.Phi3ModelLoader; import org.beehive.gpullama3.model.loader.Qwen2ModelLoader; +import org.beehive.gpullama3.model.loader.Qwen2MoEModelLoader; import org.beehive.gpullama3.model.loader.Qwen3ModelLoader; import java.nio.channels.FileChannel; @@ -59,6 +60,13 @@ public Model loadModel(FileChannel fileChannel, GGUF gguf, int contextLength, bo } }, + QWEN_2_MOE { + @Override + public Model loadModel(FileChannel fileChannel, GGUF gguf, int contextLength, boolean useTornadovm) { + return new Qwen2MoEModelLoader(fileChannel, gguf, contextLength, useTornadovm).loadModel(); + } + }, + DEEPSEEK_R1_DISTILL_QWEN { @Override public Model loadModel(FileChannel fileChannel, GGUF gguf, int contextLength, boolean useTornadovm) { diff --git a/src/main/java/org/beehive/gpullama3/model/loader/ModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/ModelLoader.java index 353aea91..1e14e12b 100644 --- a/src/main/java/org/beehive/gpullama3/model/loader/ModelLoader.java +++ b/src/main/java/org/beehive/gpullama3/model/loader/ModelLoader.java @@ -45,6 +45,12 @@ public ModelLoader(FileChannel fileChannel, GGUF gguf, int contextLength, boolea } private static ModelType detectModelType(Map metadata) { + // Architecture key is authoritative (set by llama.cpp conversion) and doesn't + // depend on how the model happens to be named, unlike general.name below. + if ("qwen2moe".equals(metadata.get("general.architecture"))) { + return ModelType.QWEN_2_MOE; + } + String name = (String) metadata.get("general.name"); // Check by name first diff --git a/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java new file mode 100644 index 00000000..410c3a24 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java @@ -0,0 +1,123 @@ +package org.beehive.gpullama3.model.loader; + +import org.beehive.gpullama3.inference.operation.RoPE; +import org.beehive.gpullama3.inference.weights.Weights; +import org.beehive.gpullama3.inference.weights.standard.Qwen2MoEStandardWeights; +import org.beehive.gpullama3.model.format.ChatFormat; +import org.beehive.gpullama3.model.qwen2.Qwen2MoE; +import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; +import org.beehive.gpullama3.tensor.GGMLTensorEntry; +import org.beehive.gpullama3.tensor.GGUF; +import org.beehive.gpullama3.auxiliary.Pair; +import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; +import org.beehive.gpullama3.tokenizer.Qwen3Tokenizer; +import org.beehive.gpullama3.tokenizer.Tokenizer; +import org.beehive.gpullama3.tokenizer.Vocabulary; + +import static org.beehive.gpullama3.model.loader.ModelLoader.*; + +import java.nio.channels.FileChannel; +import java.util.Map; + +public class Qwen2MoEModelLoader extends AbstractModelLoader { + @Override + protected Vocabulary loadVocabulary(Map metadata) { + return Vocabulary.loadQwen3Vocabulary(metadata); + } + + @Override + protected Tokenizer createTokenizer(Map metadata, Vocabulary vocabulary) { + boolean isDeepSeekR1DistillQwen = "DeepSeek-R1-Distill-Qwen".equals(metadata.get("general.basename")); + return new Qwen3Tokenizer(metadata, vocabulary, isDeepSeekR1DistillQwen); + } + + @Override + protected Pair precomputeRopeFrequencies(Qwen2MoEConfiguration config) { + return RoPE.precomputeFreqsCis(config.contextLengthModel(), config.headSize(), config.ropeTheta(), false, 8, 1, 3, 8192); + } + + @Override + protected Qwen2MoE createModel(Qwen2MoEConfiguration config, Tokenizer tokenizer, Weights weights) { + ChatFormat.ChatTokens chatTokens = new ChatFormat.ChatTokens("<|im_start|>", "<|im_end|>", "", "<|end_of_text|>", "<|endoftext|>"); + return new Qwen2MoE(config, tokenizer, weights, ChatFormat.create(tokenizer, chatTokens)); + } + + public Qwen2MoEModelLoader(FileChannel fileChannel, GGUF gguf, int contextLength, boolean useTornadovm) { + super(fileChannel, gguf, contextLength, useTornadovm); + } + + @Override + protected Qwen2MoEConfiguration createConfiguration(Map metadata) { + int modelContextLength = (int) metadata.get("qwen2moe.context_length"); + int finalContextLength = (contextLength < 0 || modelContextLength < contextLength) ? modelContextLength : contextLength; + int numberOfKeyValueHeads = (int) metadata.get("qwen2moe.attention.head_count_kv"); + int vocabSize = vocabulary.size(); + + int moeHiddenDim = gguf.getTensorInfos().get("blk.0.ffn_down_exps.weight").dimensions()[0]; + + return new Qwen2MoEConfiguration( + getModelQuantization(metadata), // quantization + (int) metadata.get("qwen2moe.embedding_length"), // dim + 0, // hiddenDim + (int) metadata.get("qwen2moe.block_count"), // numberOfLayers + (int) metadata.get("qwen2moe.attention.head_count"), // numberOfHeads + numberOfKeyValueHeads, // numberOfKeyValueHeads + numberOfKeyValueHeads, // numberOfHeadsKey + numberOfKeyValueHeads, // numberOfHeadsValue + vocabSize, // vocabularySize + modelContextLength, // contextLengthModel + finalContextLength, // contextLength + (int) metadata.get("qwen2moe.expert_count"), // numberOfExperts + (int) metadata.get("qwen2moe.expert_used_count"), // numberOfExpertsUsed + moeHiddenDim, // moeHiddenDim + (int) metadata.get("qwen2moe.feed_forward_length"), // sharedExpertHiddenDim + false, // sharedWeights + (float) metadata.get("qwen2moe.attention.layer_norm_rms_epsilon"), // rmsNormEps + (float) metadata.get("qwen2moe.rope.freq_base") // ropeTheta +); + } + + @Override + protected Weights createStandardWeights(Map tensorEntries, Qwen2MoEConfiguration config, + Pair ropeFreqs, GGMLTensorEntry tokenEmbeddings, + GGMLTensorEntry outputWeight) { + + final int nl = config.numberOfLayers(); + + return new Qwen2MoEStandardWeights( + loadTensor(tokenEmbeddings), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_norm.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_v.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.bias")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.bias")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_v.bias")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_output.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_norm.weight")), + null, // w1: qwen2moe has no dense ffn_gate, only routed/shared experts below + null, // w2: qwen2moe has no dense ffn_down + null, // w3: qwen2moe has no dense ffn_up + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_inp.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_exps.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_up_exps.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_down_exps.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_shexp.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_up_shexp.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_down_shexp.weight")), + loadArrayOfTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_inp_shexp.weight")), + loadTensor(tensorEntries.get("output_norm.weight")), + new ArrayFloatTensor(ropeFreqs.first()), + new ArrayFloatTensor(ropeFreqs.second()), + loadTensor(outputWeight), + outputWeight.ggmlType() + ); + } + + @Override + protected Weights createTornadoVMWeights(Map tensorEntries, Qwen2MoEConfiguration config, + Pair ropeFreqs, GGMLTensorEntry tokenEmbeddings, + GGMLTensorEntry outputWeight) { + throw new UnsupportedOperationException("GPU MoE weights not yet supported"); + } +} \ No newline at end of file diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java new file mode 100644 index 00000000..5aaa1111 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java @@ -0,0 +1,88 @@ +package org.beehive.gpullama3.model.qwen2; + +import org.beehive.gpullama3.inference.InferenceCore; +import org.beehive.gpullama3.inference.InferenceEngine; +import org.beehive.gpullama3.inference.sampler.Sampler; +import org.beehive.gpullama3.inference.state.Qwen2MoEState; +import org.beehive.gpullama3.inference.state.State; +import org.beehive.gpullama3.inference.weights.Weights; +import org.beehive.gpullama3.model.AbstractModel; +import org.beehive.gpullama3.model.ModelType; +import org.beehive.gpullama3.model.format.ChatFormat; +import org.beehive.gpullama3.tokenizer.Qwen3Tokenizer; +import org.beehive.gpullama3.tokenizer.Tokenizer; +import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; + +import java.util.List; +import java.util.Set; +import java.util.function.IntConsumer; + +public class Qwen2MoE extends AbstractModel { + + Qwen2MoEConfiguration configuration; + + public Qwen2MoE(Qwen2MoEConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) { + super(tokenizer, weights, chatFormat, null); + this.configuration = configuration; + } + + public Qwen2MoEConfiguration configuration() { + return configuration; + } + + @Override + public Tokenizer tokenizer() { + return (Qwen3Tokenizer) tokenizer; + } + + @Override + public ModelType getModelType() { + return ModelType.QWEN_2_MOE; + } + + @Override + public State createNewState() { + State state = new Qwen2MoEState(configuration(), -1); + 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()); + return state; + } + + @Override + public boolean shouldAddBeginOfText() { + return false; + } + + @Override + public boolean shouldAddSystemPrompt() { + return true; + } + + @Override + public boolean shouldIncludeReasoning() { + return false; + } + + @Override + public void forward(State state, int token, int position) { + InferenceCore.forwardJavaQwen2MoE(this, state, token, position); + } + + @Override + 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); + } + + @Override + public List generateTokensGPU(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, + IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) { + throw new UnsupportedOperationException("GPU MoE not yet supported"); + } +} diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java new file mode 100644 index 00000000..b736361b --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java @@ -0,0 +1,48 @@ +package org.beehive.gpullama3.model.qwen2; + +import org.beehive.gpullama3.model.Configuration; + +public record Qwen2MoEConfiguration(String quantization, + int dim, + int hiddenDim, + int numberOfLayers, + int numberOfHeads, + int numberOfKeyValueHeads, + int numberOfHeadsKey, + int numberOfHeadsValue, + int vocabularySize, + int contextLengthModel, + int contextLength, + int numberOfExperts, + int numberOfExpertsUsed, + int moeHiddenDim, + int sharedExpertHiddenDim, + boolean sharedWeights, + float rmsNormEps, + float ropeTheta) implements Configuration { + @Override + public String quantization() { + return quantization; + } + + @Override + public int headSize() { + return dim / numberOfHeads; + } + + @Override + public int kvDim() { + return (dim * numberOfKeyValueHeads) / numberOfHeads; + } + + @Override + public int kvMul() { + return numberOfHeads / numberOfKeyValueHeads; + } + + @Override + public int contextLengthModel() { + return contextLengthModel; + } +} + From 46cea940a715f175108d4b179e44a4720587d899 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sat, 4 Jul 2026 01:43:26 +0100 Subject: [PATCH 02/10] Fix Qwen2MoEState to override createStateFields with correct config type Qwen2MoEState inherited Qwen2State.createStateFields, which casts the Configuration to Qwen2Configuration. Since Qwen2MoEConfiguration is a separate record (not a subtype of Qwen2Configuration), this cast threw a ClassCastException at runtime when creating a new state. Overriding with a Qwen2MoEConfiguration cast fixes state allocation. --- .../inference/state/Qwen2MoEState.java | 57 ++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java index f70084d5..7bc7ce84 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,11 @@ 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; + +import java.util.stream.Stream; public class Qwen2MoEState extends Qwen2State { @@ -37,7 +42,55 @@ public Qwen2MoEState(Configuration config, int batchsize) { } @Override - protected StateFields createStateFields(Configuration config) { - return super.createStateFields(config); + protected StateFields createStateFields(Configuration configuration) { + StateFields fields = new StateFields(); + + Qwen2MoEConfiguration config = (Qwen2MoEConfiguration) configuration; + + int nEmbdGqa = config.kvDim(); + + fields.x = ArrayFloatTensor.allocate(config.dim()); + fields.xb = ArrayFloatTensor.allocate(config.dim()); + fields.xb2 = ArrayFloatTensor.allocate(config.dim()); + fields.hb = ArrayFloatTensor.allocate(config.hiddenDim()); + fields.hb2 = ArrayFloatTensor.allocate(config.hiddenDim()); + fields.q = ArrayFloatTensor.allocate(config.dim()); + fields.k = ArrayFloatTensor.allocate(config.kvDim()); + fields.v = ArrayFloatTensor.allocate(config.kvDim()); + 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); + + 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()); + } + fields.wrapX = new FloatArray(config.dim()); + fields.wrapXb = new FloatArray(config.dim()); + fields.wrapXbFP16 = new HalfFloatArray(config.dim()); + fields.wrapXb2 = new FloatArray(config.dim()); + fields.wrapHb = new FloatArray(config.hiddenDim()); + fields.wrapHb2 = new FloatArray(config.hiddenDim()); + + fields.wrapLogits = new FloatArray(config.vocabularySize()); + fields.wrapQ = new FloatArray(config.dim()); + 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.wrapValueCache.init(0.f); + fields.wrapKeyCache.init(0.f); + fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); + fields.positionHolder = new IntArray(1); + + fields.temp = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); + fields.tempFFN = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); + fields.tempLogits = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); + + return fields; } } From 973cdf3918d7006c7a82850e4ef25c051bd352b8 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 5 Jul 2026 15:55:44 +0100 Subject: [PATCH 03/10] Fix Qwen2-MoE routing probabilities --- .../gpullama3/inference/InferenceCore.java | 24 +++++-------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java index edd383b6..d2bde75d 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java @@ -379,8 +379,11 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // --- Step A: router scores --- // routerLogits[nExp] = routerGate[l] · xb weights.routerGate[l].matmul(state.xb, moeState.routerLogits, nExp, dim); - // --- Step B: pick top-K experts + normalize their weights --- - int[] idx = new int[topK]; float[] wgt = new float[topK]; + // --- Step B: softmax over ALL experts, then pick top-K (no renormalization) --- + // Qwen1.5-MoE uses norm_topk_prob=false: each selected expert's routing weight is + // its softmax probability over all experts, WITHOUT rescaling the top-K to sum=1. + moeState.routerLogits.softmaxInPlace(0, nExp); + int[] idx = new int[topK]; float[] wgt = new float[topK]; for (int i = 0; i < topK; i++) { float best = Float.NEGATIVE_INFINITY; int index = -1; @@ -391,25 +394,10 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke } } idx[i] = index; - wgt[i] = best; + wgt[i] = best; // softmax probability over all experts, used directly moeState.routerLogits.setFloat(index, Float.NEGATIVE_INFINITY); } - float maxVal = wgt[0]; - for (int i = 1; i < topK; i++) { - if (wgt[i] > maxVal) maxVal = wgt[i]; - } - - float sum = 0f; - for (int i = 0; i < topK; i++) { - wgt[i] = (float) Math.exp(wgt[i] - maxVal); - sum += wgt[i]; - } - - for (int i = 0; i < topK; i++) { - wgt[i] /= sum; - } - // --- Step C: compute each selected expert and accumulate into x --- for (int j = 0; j < topK; j++) { From 8ddca7d48f09009b1d2a90abb7a549aca6924c0d Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 5 Jul 2026 15:55:44 +0100 Subject: [PATCH 04/10] Add ChatML turn separator for Qwen messages --- .../org/beehive/gpullama3/model/format/Qwen3ChatFormat.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/main/java/org/beehive/gpullama3/model/format/Qwen3ChatFormat.java b/src/main/java/org/beehive/gpullama3/model/format/Qwen3ChatFormat.java index 5d121c31..cd87a3d1 100644 --- a/src/main/java/org/beehive/gpullama3/model/format/Qwen3ChatFormat.java +++ b/src/main/java/org/beehive/gpullama3/model/format/Qwen3ChatFormat.java @@ -93,6 +93,8 @@ public List encodeMessage(Message message) { if (imEnd != -1 && !isFim) { // Add the end token directly tokens.add(imEnd); + // Standard ChatML: a newline follows <|im_end|> + tokens.addAll(this.tokenizer.encodeOrdinaryAsList("\n")); } return tokens; } From 25dd61b806f676cff14e201899dcf628df636c4a Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 5 Jul 2026 15:55:44 +0100 Subject: [PATCH 05/10] Optimize Q8_0 dot products with quantized activations --- .../tensor/standard/Q8_0FloatTensor.java | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java index 9067bde0..e908d540 100644 --- a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java +++ b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java @@ -65,8 +65,15 @@ public float getFloat(int index) { public static final ValueLayout.OfShort JAVA_SHORT_LE = ValueLayout.JAVA_SHORT.withOrder(ByteOrder.LITTLE_ENDIAN); + // When enabled, quantize the activation to Q8_0 and do an int8·int8 dot, + // matching llama.cpp's Q8_0 matmul numerics exactly. + static final boolean QUANTIZE_ACTIVATION = Boolean.parseBoolean(System.getProperty("llama.quantizeActivation", "true")); + @Override public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) { + if (QUANTIZE_ACTIVATION) { + return dotQ8Activation(thisOffset, that, thatOffset, size); + } if (USE_VECTOR_API) { return vectorDot(this, thisOffset, (ArrayFloatTensor) that, thatOffset, size); } else { @@ -74,6 +81,47 @@ public float dot(int thisOffset, FloatTensor that, int thatOffset, int size) { } } + /** + * Q8_0 weight · activation, where the activation is first quantized to Q8_0 (per 32-element + * block) and the dot is accumulated as int8·int8 -> int32, then scaled. This mirrors + * llama.cpp's ggml Q8_0 matmul path. Assumes thisOffset and size are multiples of the + * Q8_0 block size (32), which holds for all matmul callers. + */ + private float dotQ8Activation(int thisOffset, FloatTensor that, int thatOffset, int size) { + final int BS = GGMLType.Q8_0.getBlockSize(); // 32 + final int TS = GGMLType.Q8_0.getTypeSize(); // 34 + float result = 0f; + int nBlocks = size / BS; + for (int b = 0; b < nBlocks; b++) { + int elemBase = b * BS; + int wBlockOffset = (thisOffset + elemBase) / BS * TS; + float wScale = Float.float16ToFloat(readShort(memorySegment, wBlockOffset)); + + // find the max abs of this activation block -> activation scale + float amax = 0f; + for (int i = 0; i < BS; i++) { + float av = Math.abs(that.getFloat(thatOffset + elemBase + i)); + if (av > amax) amax = av; + } + // Match ggml's Q8_0 quantization order: derive the int8 values using the + // full-precision scale, but store/use the scale itself as f16. + float quantizationScale = amax / 127f; + float aScale = Float.float16ToFloat(Float.floatToFloat16(quantizationScale)); + float aInv = quantizationScale != 0f ? 1f / quantizationScale : 0f; + + // int8 · int8 accumulation (round-half-away-from-zero, matching ggml roundf) + int isum = 0; + for (int i = 0; i < BS; i++) { + float s = that.getFloat(thatOffset + elemBase + i) * aInv; + int aq = (int) (s + Math.copySign(0.5f, s)); + byte wq = readByte(memorySegment, wBlockOffset + Float16.BYTES + i); + isum += aq * wq; + } + result += isum * (wScale * aScale); + } + return result; + } + private static float vectorDot(Q8_0FloatTensor thiz, int thisOffset, ArrayFloatTensor that, int thatOffset, int size) { float result = 0f; int j = 0; From 71d49d35b25cb9d24d6f2479b89a148261e09611 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 12 Jul 2026 17:28:34 +0100 Subject: [PATCH 06/10] Clean up Qwen2-MoE CPU implementation --- .../gpullama3/inference/InferenceCore.java | 86 ++++++++----------- .../inference/state/Qwen2MoEState.java | 14 +-- .../standard/Qwen2MoEStandardWeights.java | 18 ++-- .../model/loader/Qwen2MoEModelLoader.java | 40 ++++----- .../model/qwen2/Qwen2MoEConfiguration.java | 27 +++--- .../tensor/standard/Q8_0FloatTensor.java | 2 +- 6 files changed, 87 insertions(+), 100 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java index d2bde75d..b78eb7ec 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java @@ -259,6 +259,7 @@ public static FloatTensor forwardJavaDevstral(Model model, State state, int toke return state.logits; } + public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int token, int position) { final Qwen2MoEConfiguration config = (Qwen2MoEConfiguration) model.configuration(); final Qwen2MoEStandardWeights weights = (Qwen2MoEStandardWeights) model.weights(); @@ -359,75 +360,60 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // residual connection back into x state.x.addInPlace(state.xb2); - // ========================= FFN block: MoE ========================= - // NOTE: this scaffold references fields you still need to create: - // Qwen2MoEConfiguration: numberOfExperts(), numberOfExpertsUsed(), - // moeHiddenDim(), sharedExpertHiddenDim() - // Qwen2MoEStandardWeights: routerGate[], gateExps[], upExps[], downExps[], - // sharedGate[], sharedUp[], sharedDown[], sharedGateInp[] - // Qwen2MoEState buffers: routerLogits, hbE, hbE2, hbS, hbS2, yTmp - // Once those exist, change the two casts at the top of this method to the MoE - // types, then replace the TODO lines below with the real calls shown. - - // FFN pre-norm (same as dense): state.xb = rmsnorm(state.x) + // MoE FFN pre-normalization rmsnorm(state.xb, state.x, weights.rms_ffn_weight[curLayer], 0, dim, config.rmsNormEps()); - int nExp = config.numberOfExperts(); // 60 - int topK = config.numberOfExpertsUsed(); // 4 - int moeHid = config.moeHiddenDim(); // 1408 + int numberOfExperts = config.numberOfExperts(); + int topK = config.numberOfExpertsUsed(); + int expertHiddenDim = config.moeHiddenDim(); - // --- Step A: router scores --- - // routerLogits[nExp] = routerGate[l] · xb - weights.routerGate[l].matmul(state.xb, moeState.routerLogits, nExp, dim); - // --- Step B: softmax over ALL experts, then pick top-K (no renormalization) --- + // Compute routing probabilities over all experts, then select top-k. // Qwen1.5-MoE uses norm_topk_prob=false: each selected expert's routing weight is - // its softmax probability over all experts, WITHOUT rescaling the top-K to sum=1. - moeState.routerLogits.softmaxInPlace(0, nExp); - int[] idx = new int[topK]; float[] wgt = new float[topK]; + // its probability over all experts without rescaling the top-k weights to sum to one. + weights.routerGate[l].matmul(state.xb, moeState.routerLogits, numberOfExperts, dim); + moeState.routerLogits.softmaxInPlace(0, numberOfExperts); + + int[] selectedExperts = new int[topK]; + float[] routingWeights = new float[topK]; for (int i = 0; i < topK; i++) { float best = Float.NEGATIVE_INFINITY; int index = -1; - for (int j = 0; j < nExp; j++) { + for (int j = 0; j < numberOfExperts; j++) { if (moeState.routerLogits.getFloat(j) > best) { best = moeState.routerLogits.getFloat(j); index = j; } } - idx[i] = index; - wgt[i] = best; // softmax probability over all experts, used directly + selectedExperts[i] = index; + routingWeights[i] = best; moeState.routerLogits.setFloat(index, Float.NEGATIVE_INFINITY); } + // Compute each selected expert and accumulate its weighted output. + for (int j = 0; j < topK; j++) { + int expert = selectedExperts[j]; + int gateUpOffset = expert * expertHiddenDim * dim; + int downOffset = expert * dim * expertHiddenDim; + matmulExpert(weights.gateExps[l], gateUpOffset, state.xb, moeState.hbE, expertHiddenDim, dim); + matmulExpert(weights.upExps[l], gateUpOffset, state.xb, moeState.hbE2, expertHiddenDim, dim); + moeState.hbE.mapInPlace(v -> v / (float) (1.0 + Math.exp(-v))); + moeState.hbE.multiplyInPlace(moeState.hbE2); + matmulExpert(weights.downExps[l], downOffset, moeState.hbE, moeState.yTmp, dim, expertHiddenDim); + state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, routingWeights[j]); + } -// --- Step C: compute each selected expert and accumulate into x --- - for (int j = 0; j < topK; j++) { - int e = idx[j]; - int baseGU = e * moeHid * dim; // gate/up expert offset - int baseD = e * dim * moeHid; // down expert offset (rows/cols swapped) - matmulExpert(weights.gateExps[l], baseGU, state.xb, moeState.hbE, moeHid, dim); - matmulExpert(weights.upExps[l], baseGU, state.xb, moeState.hbE2, moeHid, dim); - moeState.hbE.mapInPlace(v -> v / (float)(1.0 + Math.exp(-v))); // silu - moeState.hbE.multiplyInPlace(moeState.hbE2); // gate ⊙ up - matmulExpert(weights.downExps[l], baseD, moeState.hbE, moeState.yTmp, dim, moeHid); - state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, wgt[j]); // x += wgt*y - } - - // --- Step D: shared expert (always-on, plain dense FFN) --- - // int sHid = config.sharedExpertHiddenDim(); // 5632 - int sHid = config.sharedExpertHiddenDim(); // 5632 - - weights.sharedGate[l].matmul(state.xb, moeState.hbS, sHid, dim); - weights.sharedUp[l].matmul(state.xb, moeState.hbS2, sHid, dim); - moeState.hbS.mapInPlace(v -> v / (float)(1.0 + Math.exp(-v))); // silu + // Compute the always-on shared expert. + int sharedExpertHiddenDim = config.sharedExpertHiddenDim(); + weights.sharedGate[l].matmul(state.xb, moeState.hbS, sharedExpertHiddenDim, dim); + weights.sharedUp[l].matmul(state.xb, moeState.hbS2, sharedExpertHiddenDim, dim); + moeState.hbS.mapInPlace(v -> v / (float) (1.0 + Math.exp(-v))); moeState.hbS.multiplyInPlace(moeState.hbS2); - weights.sharedDown[l].matmul(moeState.hbS, moeState.yTmp, dim, sHid); + weights.sharedDown[l].matmul(moeState.hbS, moeState.yTmp, dim, sharedExpertHiddenDim); -// 算 shared expert 的门控值 g + // Gate the shared expert output. float gateScore = weights.sharedGateInp[l].dot(0, state.xb, 0, dim); - float g = 1f / (1f + (float) Math.exp(-gateScore)); - - state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, g); - // ================================================================== + float sharedExpertWeight = 1f / (1f + (float) Math.exp(-gateScore)); + state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, sharedExpertWeight); } // final rmsnorm + classifier (same as dense Qwen2) 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 7bc7ce84..c59a9923 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -31,12 +31,12 @@ public class Qwen2MoEState extends Qwen2State { public final FloatTensor yTmp; public Qwen2MoEState(Configuration config, int batchsize) { - super(config, batchsize); // allocate all the regular Qwen2State buffers first + super(config, batchsize); Qwen2MoEConfiguration c = (Qwen2MoEConfiguration) config; this.routerLogits = ArrayFloatTensor.allocate(c.numberOfExperts()); - this.hbE = ArrayFloatTensor.allocate(c.moeHiddenDim()); + this.hbE = ArrayFloatTensor.allocate(c.moeHiddenDim()); this.hbE2 = ArrayFloatTensor.allocate(c.moeHiddenDim()); - this.hbS = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); + this.hbS = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); this.hbS2 = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); this.yTmp = ArrayFloatTensor.allocate(c.dim()); } @@ -60,8 +60,12 @@ 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()); diff --git a/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java b/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java index 75f2e3c4..3620281c 100644 --- a/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java +++ b/src/main/java/org/beehive/gpullama3/inference/weights/standard/Qwen2MoEStandardWeights.java @@ -4,17 +4,17 @@ import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; -public class Qwen2MoEStandardWeights extends StandardWeights{ - // Qwen2-specific weights +public class Qwen2MoEStandardWeights extends StandardWeights { + // Qwen2-MoE-specific weights public final FloatTensor[] q_bias, k_bias, v_bias; public final FloatTensor[] routerGate; public final FloatTensor[] gateExps; - public final FloatTensor[] upExps; // [层数] - public final FloatTensor[] downExps; // [层数] - public final FloatTensor[] sharedGate; // [层数]，shared expert - public final FloatTensor[] sharedUp; // [层数] - public final FloatTensor[] sharedDown; // [层数] - public final FloatTensor[] sharedGateInp; // [层数]，shared expert 的 sigmoid gate 输入 + public final FloatTensor[] upExps; + public final FloatTensor[] downExps; + public final FloatTensor[] sharedGate; + public final FloatTensor[] sharedUp; + public final FloatTensor[] sharedDown; + public final FloatTensor[] sharedGateInp; public Qwen2MoEStandardWeights( FloatTensor token_embedding_table, @@ -43,7 +43,6 @@ public Qwen2MoEStandardWeights( ArrayFloatTensor freq_cis_imag, FloatTensor wcls, GGMLType weightType) { - // call to StandardWeights constructor super(token_embedding_table, rms_att_weight, wq, @@ -59,7 +58,6 @@ public Qwen2MoEStandardWeights( freq_cis_imag, wcls, weightType); - // init Qwen2-specific fields this.q_bias = q_bias; this.k_bias = k_bias; this.v_bias = v_bias; diff --git a/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java index 410c3a24..b21c1071 100644 --- a/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java +++ b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java @@ -1,5 +1,6 @@ package org.beehive.gpullama3.model.loader; +import org.beehive.gpullama3.auxiliary.Pair; import org.beehive.gpullama3.inference.operation.RoPE; import org.beehive.gpullama3.inference.weights.Weights; import org.beehive.gpullama3.inference.weights.standard.Qwen2MoEStandardWeights; @@ -8,18 +9,21 @@ import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; import org.beehive.gpullama3.tensor.GGMLTensorEntry; import org.beehive.gpullama3.tensor.GGUF; -import org.beehive.gpullama3.auxiliary.Pair; import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tokenizer.Qwen3Tokenizer; import org.beehive.gpullama3.tokenizer.Tokenizer; import org.beehive.gpullama3.tokenizer.Vocabulary; -import static org.beehive.gpullama3.model.loader.ModelLoader.*; - import java.nio.channels.FileChannel; import java.util.Map; -public class Qwen2MoEModelLoader extends AbstractModelLoader { +import static org.beehive.gpullama3.model.loader.ModelLoader.*; + +public class Qwen2MoEModelLoader extends AbstractModelLoader { + public Qwen2MoEModelLoader(FileChannel fileChannel, GGUF gguf, int contextLength, boolean useTornadovm) { + super(fileChannel, gguf, contextLength, useTornadovm); + } + @Override protected Vocabulary loadVocabulary(Map metadata) { return Vocabulary.loadQwen3Vocabulary(metadata); @@ -42,10 +46,6 @@ protected Qwen2MoE createModel(Qwen2MoEConfiguration config, Tokenizer tokenizer return new Qwen2MoE(config, tokenizer, weights, ChatFormat.create(tokenizer, chatTokens)); } - public Qwen2MoEModelLoader(FileChannel fileChannel, GGUF gguf, int contextLength, boolean useTornadovm) { - super(fileChannel, gguf, contextLength, useTornadovm); - } - @Override protected Qwen2MoEConfiguration createConfiguration(Map metadata) { int modelContextLength = (int) metadata.get("qwen2moe.context_length"); @@ -57,24 +57,24 @@ protected Qwen2MoEConfiguration createConfiguration(Map metadata return new Qwen2MoEConfiguration( getModelQuantization(metadata), // quantization - (int) metadata.get("qwen2moe.embedding_length"), // dim + (int) metadata.get("qwen2moe.embedding_length"), // dim 0, // hiddenDim - (int) metadata.get("qwen2moe.block_count"), // numberOfLayers + (int) metadata.get("qwen2moe.block_count"), // numberOfLayers (int) metadata.get("qwen2moe.attention.head_count"), // numberOfHeads - numberOfKeyValueHeads, // numberOfKeyValueHeads - numberOfKeyValueHeads, // numberOfHeadsKey - numberOfKeyValueHeads, // numberOfHeadsValue - vocabSize, // vocabularySize - modelContextLength, // contextLengthModel - finalContextLength, // contextLength + numberOfKeyValueHeads, // numberOfKeyValueHeads + numberOfKeyValueHeads, // numberOfHeadsKey + numberOfKeyValueHeads, // numberOfHeadsValue + vocabSize, // vocabularySize + modelContextLength, // contextLengthModel + finalContextLength, // contextLength (int) metadata.get("qwen2moe.expert_count"), // numberOfExperts (int) metadata.get("qwen2moe.expert_used_count"), // numberOfExpertsUsed - moeHiddenDim, // moeHiddenDim + moeHiddenDim, // moeHiddenDim (int) metadata.get("qwen2moe.feed_forward_length"), // sharedExpertHiddenDim - false, // sharedWeights + false, // sharedWeights (float) metadata.get("qwen2moe.attention.layer_norm_rms_epsilon"), // rmsNormEps (float) metadata.get("qwen2moe.rope.freq_base") // ropeTheta -); + ); } @Override @@ -120,4 +120,4 @@ protected Weights createTornadoVMWeights(Map tensorEntr GGMLTensorEntry outputWeight) { throw new UnsupportedOperationException("GPU MoE weights not yet supported"); } -} \ No newline at end of file +} diff --git a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java index b736361b..9b4a0ed5 100644 --- a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoEConfiguration.java @@ -3,23 +3,23 @@ import org.beehive.gpullama3.model.Configuration; public record Qwen2MoEConfiguration(String quantization, - int dim, - int hiddenDim, - int numberOfLayers, - int numberOfHeads, - int numberOfKeyValueHeads, - int numberOfHeadsKey, - int numberOfHeadsValue, - int vocabularySize, - int contextLengthModel, - int contextLength, + int dim, + int hiddenDim, + int numberOfLayers, + int numberOfHeads, + int numberOfKeyValueHeads, + int numberOfHeadsKey, + int numberOfHeadsValue, + int vocabularySize, + int contextLengthModel, + int contextLength, int numberOfExperts, int numberOfExpertsUsed, int moeHiddenDim, int sharedExpertHiddenDim, - boolean sharedWeights, - float rmsNormEps, - float ropeTheta) implements Configuration { + boolean sharedWeights, + float rmsNormEps, + float ropeTheta) implements Configuration { @Override public String quantization() { return quantization; @@ -45,4 +45,3 @@ public int contextLengthModel() { return contextLengthModel; } } - diff --git a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java index e908d540..d3304d40 100644 --- a/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java +++ b/src/main/java/org/beehive/gpullama3/tensor/standard/Q8_0FloatTensor.java @@ -66,7 +66,7 @@ public float getFloat(int index) { public static final ValueLayout.OfShort JAVA_SHORT_LE = ValueLayout.JAVA_SHORT.withOrder(ByteOrder.LITTLE_ENDIAN); // When enabled, quantize the activation to Q8_0 and do an int8·int8 dot, - // matching llama.cpp's Q8_0 matmul numerics exactly. + // matching llama.cpp's Q8_0 matmul quantization scheme. static final boolean QUANTIZE_ACTIVATION = Boolean.parseBoolean(System.getProperty("llama.quantizeActivation", "true")); @Override From c2e9583561093ea504988136045bdd895af7598a Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 2 Aug 2026 15:31:12 +0100 Subject: [PATCH 07/10] Add Qwen2-MoE single-token GPU baseline --- docs/qwen2-moe-gpu-baseline.md | 143 ++++++ scripts/compare_moe_correctness.py | 148 ++++++ scripts/summarize_tornado_profile.py | 170 +++++++ .../gpullama3/inference/InferenceCore.java | 13 + .../gpullama3/inference/InferenceEngine.java | 5 +- .../inference/state/Qwen2MoEState.java | 27 +- .../gpullama3/inference/state/Qwen2State.java | 12 +- .../tornado/Qwen2MoETornadoWeights.java | 78 +++ .../model/loader/Qwen2MoEModelLoader.java | 37 +- .../gpullama3/model/qwen2/Qwen2MoE.java | 17 +- .../TornadoVMMasterPlanSingleToken.java | 8 + .../tornadovm/kernels/Qwen2MoEKernels.java | 456 ++++++++++++++++++ .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 300 ++++++++++++ .../tornadovm/plan/ForwardPlanFactory.java | 9 + .../q8_0/Qwen2MoEQ8_0PlanComponents.java | 46 ++ .../validation/MoECorrectnessTrace.java | 192 ++++++++ 16 files changed, 1650 insertions(+), 11 deletions(-) create mode 100644 docs/qwen2-moe-gpu-baseline.md create mode 100755 scripts/compare_moe_correctness.py create mode 100644 scripts/summarize_tornado_profile.py create mode 100644 src/main/java/org/beehive/gpullama3/inference/weights/tornado/Qwen2MoETornadoWeights.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java create mode 100644 src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java diff --git a/docs/qwen2-moe-gpu-baseline.md b/docs/qwen2-moe-gpu-baseline.md new file mode 100644 index 00000000..e5611a3d --- /dev/null +++ b/docs/qwen2-moe-gpu-baseline.md @@ -0,0 +1,143 @@ +# Qwen2-MoE Q8_0 GPU Baseline + +This document records the working single-token GPU baseline before further +kernel and scheduling optimizations. The baseline is preserved on branch +`codex/qwen2-moe-gpu-baseline`. + +## Scope + +- Model: `Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf` +- GGUF file size: approximately 15 GB +- Weight format: Q8_0 (32 int8 values and one FP16 scale per block) +- Execution mode: single-token inference +- GPU backend: TornadoVM PTX +- Unsupported in this baseline: sequential prefill/decode and batch prefill/decode + +The GPU path includes Qwen2 attention, router projection, softmax and Top-K, +four routed experts, the shared expert, residual accumulation, and final logits. + +## Test Environment + +- Server: `storm` +- GPU: NVIDIA GeForce RTX 4090, 24 GB +- TornadoVM SDK: 5.1.0, JDK 21, PTX and OpenCL backends installed +- Model path: `/home/mingyi/models/Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf` +- GPU memory limit: 20 GB + +Representative command: + +```bash +./llama-tornado \ + --gpu --ptx \ + --gpu-memory 20GB \ + --heap-min 2g --heap-max 8g \ + --model /home/mingyi/models/Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf \ + --prompt "Hi" \ + --temperature 0 \ + --seed 42 \ + --max-tokens 128 +``` + +## Correctness Results + +The repository contains an optional JSONL trace and comparison script that +compare generated token IDs, per-layer router logits, Top-K expert IDs, +routing weights, and final logits between CPU and GPU executions. + +When CPU activation quantization was disabled so that the CPU arithmetic more +closely matched the current GPU kernels, the common trace prefix produced: + +- 3 compared generated token IDs with no mismatch +- 20 final Top-1 predictions with no mismatch +- 4 Top-K expert-set mismatches out of 495 layer comparisons +- mean absolute router-logit error: 0.001323 +- mean absolute final-logit error: 0.01094 + +The first expert-set mismatch occurred around an almost tied routing decision. +This indicates that the main remaining differences are numerical rather than a +large structural error in the MoE pipeline. Longer correctness traces are still +required before claiming full numerical equivalence. + +## Throughput Baseline + +Three interleaved 128-token PTX runs measured: + +| Run | Throughput | +|---:|---:| +| 1 | 17.75 tokens/s | +| 2 | 17.77 tokens/s | +| 3 | 17.58 tokens/s | +| **Mean** | **17.70 tokens/s** | + +A later 19-token smoke test reached 24.63 tokens/s. This short result is kept as +a health check, not as the main baseline, because short generations are more +sensitive to prompt length, warm-up, and measurement variance. + +## Kernel Profiling Results + +The TornadoVM profiler was run with 16 generated tokens. The first execution of +each TaskGraph was excluded because it includes initialization and initial +weight transfer, leaving 15 steady-state iterations. + +| Component | Time per token | Share of GPU kernel time | +|---|---:|---:| +| Attention | 5.222 ms | 28.62% | +| Shared expert | 4.796 ms | 26.29% | +| Routed Gate/Up | 3.539 ms | 19.40% | +| Routed Down | 2.536 ms | 13.90% | +| Router and Top-K | 1.179 ms | 6.46% | +| FFN RMSNorm | 0.536 ms | 2.94% | +| Other kernels | 0.435 ms | 2.39% | + +Routed and shared expert computation accounts for approximately 59.6% of GPU +kernel time. The most expensive individual tasks were: + +| Task | Time per token | Share of GPU kernel time | +|---|---:|---:| +| Attention kernel | 3.420 ms | 18.75% | +| Four routed Gate/Up kernels | 3.539 ms | 19.40% | +| Four routed Down kernels | 2.536 ms | 13.90% | +| Shared expert Gate/Up | 2.105 ms | 11.54% | +| Shared expert Down | 1.824 ms | 10.00% | +| Shared gate and accumulation | 0.867 ms | 4.75% | +| Router projection | 0.671 ms | 3.68% | +| Softmax and Top-K | 0.508 ms | 2.78% | + +Runtime-level profiler totals per token were: + +- GPU kernels: 18.244 ms +- copy-in: 4.962 ms +- runtime/profiler residual: 24.118 ms +- total TaskGraph time: 47.323 ms + +The residual is only an upper bound. It combines host dispatch, +synchronization, event handling, and profiler overhead, so it must not be +reported as pure kernel-launch time. + +## Q8_0 Activation Experiments + +Several experimental kernels quantized activations to Q8_0 and reused the +quantized values for integer dot products. These experiments were rolled back +from the working baseline: + +| Experiment | Result | +|---|---:| +| Routed Gate/Up only | approximately 2.5% faster | +| Routed and shared Gate/Up | approximately 3.4% faster | +| All tested Q8_0 matrix-vector paths | approximately 3.7% slower | + +The full conversion also increased numerical error. The current baseline +therefore keeps FP32 activations and reads Q8_0 weights by applying each block's +FP16 scale during the dot product. + +## Next Optimization Target + +Profiling shows that Top-K selection is too small to be the first target. The +next branch should investigate expert-task fusion and scheduling overhead: + +- reduce the four routed Gate/Up tasks to one task per layer where practical; +- reduce the four routed Down tasks to one task per layer where practical; +- preserve the current correctness trace as the regression oracle; +- compare steady-state throughput and profiler results against 17.70 tokens/s; +- keep any optimization only if it improves performance without unacceptable + correctness loss. diff --git a/scripts/compare_moe_correctness.py b/scripts/compare_moe_correctness.py new file mode 100755 index 00000000..7a16de73 --- /dev/null +++ b/scripts/compare_moe_correctness.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Compare Qwen2-MoE CPU and GPU JSONL correctness traces.""" + +import argparse +import json +import math +from pathlib import Path + + +def load_trace(path: Path): + records = {"router": {}, "logits": {}, "token": {}} + lines = path.read_text(encoding="utf-8").splitlines() + for line_number, line in enumerate(lines, 1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError: + if line_number == len(lines): + print(f"warning: ignoring incomplete final line in {path}") + break + raise + kind = record["type"] + if kind == "router": + key = (record["position"], record["layer"]) + elif kind == "logits": + key = record["position"] + elif kind == "token": + key = record["index"] + else: + raise ValueError(f"{path}:{line_number}: unknown record type {kind}") + if key in records[kind]: + raise ValueError(f"{path}:{line_number}: duplicate {kind} key {key}") + records[kind][key] = record + return records + + +def errors(left, right): + if len(left) != len(right): + raise ValueError(f"array lengths differ: {len(left)} != {len(right)}") + differences = [abs(a - b) for a, b in zip(left, right)] + if not all(math.isfinite(value) for value in differences): + return math.inf, math.inf + return max(differences, default=0.0), sum(differences) / max(1, len(differences)) + + +def argmax(values): + return max(range(len(values)), key=values.__getitem__) + + +def matching_keys(cpu, gpu, kind, common_only): + cpu_keys = set(cpu[kind]) + gpu_keys = set(gpu[kind]) + if cpu_keys != gpu_keys and not common_only: + missing_gpu = sorted(cpu_keys - gpu_keys) + missing_cpu = sorted(gpu_keys - cpu_keys) + raise ValueError( + f"{kind} keys differ; missing on GPU={missing_gpu[:10]}, " + f"missing on CPU={missing_cpu[:10]}" + ) + return sorted(cpu_keys & gpu_keys) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("cpu_trace", type=Path) + parser.add_argument("gpu_trace", type=Path) + parser.add_argument( + "--common-only", action="store_true", + help="compare only records completed by both traces", + ) + args = parser.parse_args() + + cpu = load_trace(args.cpu_trace) + gpu = load_trace(args.gpu_trace) + + token_keys = matching_keys(cpu, gpu, "token", args.common_only) + token_mismatches = [ + key for key in token_keys + if cpu["token"][key]["id"] != gpu["token"][key]["id"] + ] + + router_keys = matching_keys(cpu, gpu, "router", args.common_only) + expert_order_mismatches = [] + expert_set_mismatches = [] + router_max = router_mean_sum = routing_max = routing_mean_sum = 0.0 + for key in router_keys: + cpu_record = cpu["router"][key] + gpu_record = gpu["router"][key] + if cpu_record["experts"] != gpu_record["experts"]: + expert_order_mismatches.append(key) + if set(cpu_record["experts"]) != set(gpu_record["experts"]): + expert_set_mismatches.append(key) + maximum, mean = errors(cpu_record["logits"], gpu_record["logits"]) + router_max = max(router_max, maximum) + router_mean_sum += mean + maximum, mean = errors(cpu_record["weights"], gpu_record["weights"]) + routing_max = max(routing_max, maximum) + routing_mean_sum += mean + + logits_keys = matching_keys(cpu, gpu, "logits", args.common_only) + top1_mismatches = [] + logits_max = logits_mean_sum = 0.0 + for key in logits_keys: + cpu_values = cpu["logits"][key]["values"] + gpu_values = gpu["logits"][key]["values"] + if argmax(cpu_values) != argmax(gpu_values): + top1_mismatches.append(key) + maximum, mean = errors(cpu_values, gpu_values) + logits_max = max(logits_max, maximum) + logits_mean_sum += mean + + print(f"tokens compared: {len(token_keys)}") + print(f"token ID mismatches: {len(token_mismatches)} {token_mismatches[:10]}") + print(f"router layers compared: {len(router_keys)}") + print(f"Top-K order mismatches: {len(expert_order_mismatches)} {expert_order_mismatches[:10]}") + print(f"Top-K set mismatches: {len(expert_set_mismatches)} {expert_set_mismatches[:10]}") + print(f"router logits max abs error: {router_max:.8g}") + print(f"router logits mean abs err: {router_mean_sum / max(1, len(router_keys)):.8g}") + print(f"routing weight max abs err: {routing_max:.8g}") + print(f"routing weight mean abs err: {routing_mean_sum / max(1, len(router_keys)):.8g}") + print(f"logit vectors compared: {len(logits_keys)}") + print(f"final Top-1 mismatches: {len(top1_mismatches)} {top1_mismatches[:10]}") + print(f"final logits max abs error: {logits_max:.8g}") + print(f"final logits mean abs error: {logits_mean_sum / max(1, len(logits_keys)):.8g}") + + if expert_set_mismatches: + first = expert_set_mismatches[0] + cpu_record = cpu["router"][first] + gpu_record = gpu["router"][first] + print(f"first Top-K set mismatch {first}:") + print(f" CPU experts={cpu_record['experts']} weights={cpu_record['weights']}") + print(f" GPU experts={gpu_record['experts']} weights={gpu_record['weights']}") + for expert in sorted(set(cpu_record["experts"]) | set(gpu_record["experts"])): + cpu_logit = cpu_record["logits"][expert] + gpu_logit = gpu_record["logits"][expert] + print( + f" expert {expert:2d}: CPU raw={cpu_logit: .8f}, " + f"GPU raw={gpu_logit: .8f}, abs err={abs(cpu_logit - gpu_logit):.3g}" + ) + + passed = not token_mismatches and not expert_set_mismatches and not top1_mismatches + print("RESULT: " + ("PASS" if passed else "FAIL")) + raise SystemExit(0 if passed else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/summarize_tornado_profile.py b/scripts/summarize_tornado_profile.py new file mode 100644 index 00000000..6f5d6f85 --- /dev/null +++ b/scripts/summarize_tornado_profile.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Summarize concatenated TornadoVM profiler JSON objects by MoE task category.""" + +import argparse +import collections +import json +from pathlib import Path + + +ATTENTION_TASKS = { + "attn_rms_reduce", + "attn_rms_finalize", + "attn_rms_qkv_projection", + "fused_qkv_bias", + "rope_and_kv_cache", + "attention", + "attn_output_proj", +} + + +def load_objects(path: Path): + """Read TornadoVM's stream of adjacent top-level JSON objects.""" + text = path.read_text(encoding="utf-8") + decoder = json.JSONDecoder() + offset = 0 + objects = [] + while offset < len(text): + while offset < len(text) and text[offset].isspace(): + offset += 1 + if offset >= len(text): + break + value, offset = decoder.raw_decode(text, offset) + objects.append(value) + return objects + + +def task_category(full_name: str): + task = full_name.rsplit(".", 1)[-1] + if task.startswith("routed_expert_gate_up_"): + return "Routed Gate/Up" + if task.startswith("routed_expert_down_"): + return "Routed Down" + if task.startswith("shared_expert_"): + return "Shared expert" + if task.startswith("router_"): + return "Router + Top-K" + if task in ATTENTION_TASKS: + return "Attention" + if task.startswith("ffn_rms_"): + return "FFN RMSNorm" + return "Other" + + +def collect(objects, warmup_per_graph): + occurrences = collections.Counter() + categories = collections.Counter() + tasks = collections.Counter() + totals = collections.Counter() + graph_counts = collections.Counter() + task_count = 0 + + for record in objects: + graph_name = next(iter(record)) + graph = record[graph_name] + occurrence = occurrences[graph_name] + occurrences[graph_name] += 1 + if occurrence < warmup_per_graph: + continue + + graph_counts[graph_name] += 1 + totals["kernel"] += int(graph.get("TOTAL_KERNEL_TIME", 0)) + totals["copy_in"] += int(graph.get("COPY_IN_TIME", 0)) + totals["task_graph"] += int(graph.get("TOTAL_TASK_GRAPH_TIME", 0)) + totals["copy_bytes"] += int(graph.get("TOTAL_COPY_IN_SIZE_BYTES", 0)) + + for task_name, task in graph.items(): + if isinstance(task, dict): + task_time = int(task.get("TASK_KERNEL_TIME", 0)) + categories[task_category(task_name)] += task_time + tasks[task_name.rsplit(".", 1)[-1]] += task_time + task_count += 1 + + return occurrences, graph_counts, categories, tasks, totals, task_count + + +def milliseconds(nanoseconds): + return nanoseconds / 1_000_000.0 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("profile", type=Path) + parser.add_argument( + "--warmup-per-graph", + type=int, + default=1, + help="ignore this many initial executions of every TaskGraph (default: 1)", + ) + args = parser.parse_args() + + objects = load_objects(args.profile) + occurrences, graphs, categories, tasks, totals, task_count = collect( + objects, args.warmup_per_graph + ) + retained_iterations = min(graphs.values(), default=0) + kernel_time = totals["kernel"] + residual = totals["task_graph"] - kernel_time - totals["copy_in"] + + print(f"profile objects: {len(objects)}") + print(f"TaskGraphs: {len(occurrences)}") + print(f"executions per graph: {sorted(set(occurrences.values()))}") + print(f"warmups skipped per graph: {args.warmup_per_graph}") + print(f"steady iterations retained: {retained_iterations}") + print(f"task executions retained: {task_count}") + print() + + print("GPU kernel breakdown") + for name, time_ns in categories.most_common(): + percentage = 100.0 * time_ns / kernel_time if kernel_time else 0.0 + per_iteration = ( + milliseconds(time_ns) / retained_iterations + if retained_iterations + else 0.0 + ) + print( + f"{name:20s} {milliseconds(time_ns):9.3f} ms total" + f" {per_iteration:8.3f} ms/iter {percentage:6.2f}%" + ) + + print() + print("Top individual tasks") + for name, time_ns in tasks.most_common(15): + percentage = 100.0 * time_ns / kernel_time if kernel_time else 0.0 + per_iteration = ( + milliseconds(time_ns) / retained_iterations + if retained_iterations + else 0.0 + ) + print( + f"{name:38s} {per_iteration:8.3f} ms/iter {percentage:6.2f}%" + ) + + print() + print("Runtime-level totals") + for name, value in ( + ("GPU kernels", kernel_time), + ("copy-in", totals["copy_in"]), + ("runtime/profiler residual", residual), + ("TaskGraph total", totals["task_graph"]), + ): + per_iteration = ( + milliseconds(value) / retained_iterations + if retained_iterations + else 0.0 + ) + print( + f"{name:26s} {milliseconds(value):9.3f} ms total" + f" {per_iteration:8.3f} ms/iter" + ) + + print() + print( + "Note: runtime/profiler residual is an upper bound containing host-side " + "dispatch, synchronization, event handling, and profiler overhead; it is " + "not a pure kernel-launch measurement." + ) + + +if __name__ == "__main__": + main() diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java index b78eb7ec..4753764c 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java @@ -18,6 +18,7 @@ import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import org.beehive.gpullama3.validation.MoECorrectnessTrace; import java.lang.foreign.MemorySegment; @@ -371,6 +372,13 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // Qwen1.5-MoE uses norm_topk_prob=false: each selected expert's routing weight is // its probability over all experts without rescaling the top-k weights to sum to one. weights.routerGate[l].matmul(state.xb, moeState.routerLogits, numberOfExperts, dim); + float[] rawRouterLogits = null; + if (MoECorrectnessTrace.isEnabled()) { + rawRouterLogits = new float[numberOfExperts]; + for (int expert = 0; expert < numberOfExperts; expert++) { + rawRouterLogits[expert] = moeState.routerLogits.getFloat(expert); + } + } moeState.routerLogits.softmaxInPlace(0, numberOfExperts); int[] selectedExperts = new int[topK]; @@ -388,6 +396,10 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke routingWeights[i] = best; moeState.routerLogits.setFloat(index, Float.NEGATIVE_INFINITY); } + if (MoECorrectnessTrace.isEnabled()) { + MoECorrectnessTrace.recordCpuRouter(position, l, rawRouterLogits, + selectedExperts, routingWeights); + } // Compute each selected expert and accumulate its weighted output. for (int j = 0; j < topK; j++) { @@ -419,6 +431,7 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // final rmsnorm + classifier (same as dense Qwen2) rmsnorm(state.x, state.x, weights.rms_final_weight, 0, dim, config.rmsNormEps()); weights.wcls.matmul(state.x, state.logits, config.vocabularySize(), dim); + MoECorrectnessTrace.recordCpuLogits(position, state.logits); return state.logits; } diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index 85ca2ec5..d9549599 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java @@ -9,6 +9,7 @@ import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; import org.beehive.gpullama3.tornadovm.layers.type.fp16.LogitsFP16Layer; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import org.beehive.gpullama3.validation.MoECorrectnessTrace; import java.io.ByteArrayOutputStream; import java.util.ArrayList; @@ -209,6 +210,7 @@ public static List generateTokensQwen3(Model model, State state, int st // Track the generated token generatedTokens.add(nextToken); + MoECorrectnessTrace.recordToken(nextToken); // Notify via callback if provided if (onTokenGenerated != null) { @@ -449,6 +451,7 @@ public static List generateTokensGPUQwen3(Model model, State state, int // Track the generated token generatedTokens.add(nextToken); + MoECorrectnessTrace.recordToken(nextToken); // Notify via callback if provided if (onTokenGenerated != null) { @@ -678,4 +681,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/state/Qwen2MoEState.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java index c59a9923..0a47e77b 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -30,6 +30,17 @@ public class Qwen2MoEState extends Qwen2State { // before it is weighted and accumulated into the residual stream (state.x). public final FloatTensor yTmp; + // TornadoVM buffers for the single-token GPU MoE path. These are deliberately + // separate from the CPU FloatTensor fields above: TaskGraph kernels operate on + // TornadoVM arrays that can remain resident on the device between tasks. + public final FloatArray wrapRouterLogits; + public final FloatArray wrapRawRouterLogits; + public final IntArray wrapSelectedExperts; + public final FloatArray wrapRoutingWeights; + public final FloatArray wrapExpertGate; + public final FloatArray wrapSharedGate; + public final FloatArray wrapSharedOutput; + public Qwen2MoEState(Configuration config, int batchsize) { super(config, batchsize); Qwen2MoEConfiguration c = (Qwen2MoEConfiguration) config; @@ -39,6 +50,14 @@ public Qwen2MoEState(Configuration config, int batchsize) { this.hbS = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); this.hbS2 = ArrayFloatTensor.allocate(c.sharedExpertHiddenDim()); this.yTmp = ArrayFloatTensor.allocate(c.dim()); + + this.wrapRouterLogits = new FloatArray(c.numberOfExperts()); + this.wrapRawRouterLogits = new FloatArray(c.numberOfExperts()); + this.wrapSelectedExperts = new IntArray(c.numberOfExpertsUsed()); + this.wrapRoutingWeights = new FloatArray(c.numberOfExpertsUsed()); + this.wrapExpertGate = new FloatArray(c.moeHiddenDim()); + this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim()); + this.wrapSharedOutput = new FloatArray(c.dim()); } @Override @@ -91,9 +110,11 @@ protected StateFields createStateFields(Configuration configuration) { fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); fields.positionHolder = new IntArray(1); - fields.temp = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); - fields.tempFFN = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); - fields.tempLogits = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); + // 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)); return fields; } diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java index 266730ac..5f8213eb 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java @@ -12,9 +12,11 @@ public class Qwen2State extends State { + protected static final int QWEN2_LOCAL_SIZE = 32; + public Qwen2State(Configuration config, int batchsize) { super(config, batchsize); - this.localSize = 32; + this.localSize = QWEN2_LOCAL_SIZE; } @Override protected StateFields createStateFields(Configuration configuration) { @@ -66,9 +68,11 @@ protected StateFields createStateFields(Configuration configuration) { fields.positionHolder = new IntArray(1); // Temporary arrays - fields.temp = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); - fields.tempFFN = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); - fields.tempLogits = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize)); + // 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)); return fields; diff --git a/src/main/java/org/beehive/gpullama3/inference/weights/tornado/Qwen2MoETornadoWeights.java b/src/main/java/org/beehive/gpullama3/inference/weights/tornado/Qwen2MoETornadoWeights.java new file mode 100644 index 00000000..68f7b5b9 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/inference/weights/tornado/Qwen2MoETornadoWeights.java @@ -0,0 +1,78 @@ +package org.beehive.gpullama3.inference.weights.tornado; + +import org.beehive.gpullama3.tensor.GGMLType; +import org.beehive.gpullama3.tensor.tornado.TornadoTensor; + +/** + * TornadoVM weight container for Qwen2-MoE / Qwen1.5-MoE models. + * + *

The inherited fields provide the embedding, attention, RMSNorm and final + * classifier weights. MoE-specific tensors remain in their GGUF Q8_0 layout; + * GPU kernels will access their {@code ByteArray} representation through + * {@link TornadoTensor#asByteArray()}.

+ */ +public final class Qwen2MoETornadoWeights extends Qwen2TornadoWeights { + + public final TornadoTensor[] routerGateLayered; + public final TornadoTensor[] gateExpertsLayered; + public final TornadoTensor[] upExpertsLayered; + public final TornadoTensor[] downExpertsLayered; + public final TornadoTensor[] sharedGateLayered; + public final TornadoTensor[] sharedUpLayered; + public final TornadoTensor[] sharedDownLayered; + public final TornadoTensor[] sharedGateInputLayered; + + // @formatter:off + public Qwen2MoETornadoWeights( + TornadoTensor tokenEmbeddingTable, + TornadoTensor[] rmsAttWeightLayered, + TornadoTensor[] wqLayered, + TornadoTensor[] wkLayered, + TornadoTensor[] wvLayered, + TornadoTensor[] qBiasLayered, + TornadoTensor[] kBiasLayered, + TornadoTensor[] vBiasLayered, + TornadoTensor[] woLayered, + TornadoTensor[] rmsFfnWeightLayered, + TornadoTensor[] routerGateLayered, + TornadoTensor[] gateExpertsLayered, + TornadoTensor[] upExpertsLayered, + TornadoTensor[] downExpertsLayered, + TornadoTensor[] sharedGateLayered, + TornadoTensor[] sharedUpLayered, + TornadoTensor[] sharedDownLayered, + TornadoTensor[] sharedGateInputLayered, + TornadoTensor rmsFinalWeight, + TornadoTensor freqCisReal, + TornadoTensor freqCisImag, + TornadoTensor wCls, + GGMLType weightType) { + super(tokenEmbeddingTable, + rmsAttWeightLayered, + wqLayered, + wkLayered, + wvLayered, + qBiasLayered, + kBiasLayered, + vBiasLayered, + woLayered, + rmsFfnWeightLayered, + null, + null, + null, + rmsFinalWeight, + freqCisReal, + freqCisImag, + wCls, + weightType); + this.routerGateLayered = routerGateLayered; + this.gateExpertsLayered = gateExpertsLayered; + this.upExpertsLayered = upExpertsLayered; + this.downExpertsLayered = downExpertsLayered; + this.sharedGateLayered = sharedGateLayered; + this.sharedUpLayered = sharedUpLayered; + this.sharedDownLayered = sharedDownLayered; + this.sharedGateInputLayered = sharedGateInputLayered; + } + // @formatter:on +} diff --git a/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java index b21c1071..aed25d90 100644 --- a/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java +++ b/src/main/java/org/beehive/gpullama3/model/loader/Qwen2MoEModelLoader.java @@ -4,12 +4,15 @@ import org.beehive.gpullama3.inference.operation.RoPE; import org.beehive.gpullama3.inference.weights.Weights; import org.beehive.gpullama3.inference.weights.standard.Qwen2MoEStandardWeights; +import org.beehive.gpullama3.inference.weights.tornado.Qwen2MoETornadoWeights; import org.beehive.gpullama3.model.format.ChatFormat; import org.beehive.gpullama3.model.qwen2.Qwen2MoE; import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; import org.beehive.gpullama3.tensor.GGMLTensorEntry; +import org.beehive.gpullama3.tensor.GGMLType; import org.beehive.gpullama3.tensor.GGUF; import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; +import org.beehive.gpullama3.tensor.tornado.FP32TornadoTensor; import org.beehive.gpullama3.tokenizer.Qwen3Tokenizer; import org.beehive.gpullama3.tokenizer.Tokenizer; import org.beehive.gpullama3.tokenizer.Vocabulary; @@ -17,6 +20,8 @@ import java.nio.channels.FileChannel; import java.util.Map; +import uk.ac.manchester.tornado.api.types.arrays.FloatArray; + import static org.beehive.gpullama3.model.loader.ModelLoader.*; public class Qwen2MoEModelLoader extends AbstractModelLoader { @@ -118,6 +123,36 @@ protected Weights createStandardWeights(Map tensorEntri protected Weights createTornadoVMWeights(Map tensorEntries, Qwen2MoEConfiguration config, Pair ropeFreqs, GGMLTensorEntry tokenEmbeddings, GGMLTensorEntry outputWeight) { - throw new UnsupportedOperationException("GPU MoE weights not yet supported"); + GGMLType ggmlType = effectiveGpuWeightType(outputWeight.ggmlType()); + if (ggmlType != GGMLType.F16 && ggmlType != GGMLType.Q8_0) { + throw new UnsupportedOperationException( + "Type: " + ggmlType + " currently not supported for TornadoVM weights."); + } + + final int nl = config.numberOfLayers(); + return new Qwen2MoETornadoWeights( + loadTornadoTensor(tokenEmbeddings), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_norm.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_v.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_q.bias")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_k.bias")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_v.bias")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".attn_output.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_norm.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_inp.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_exps.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_up_exps.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_down_exps.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_shexp.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_up_shexp.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_down_shexp.weight")), + loadArrayOfTornadoTensors(nl, i -> tensorEntries.get("blk." + i + ".ffn_gate_inp_shexp.weight")), + loadTornadoTensor(tensorEntries.get("output_norm.weight")), + new FP32TornadoTensor(FloatArray.fromArray(ropeFreqs.first())), + new FP32TornadoTensor(FloatArray.fromArray(ropeFreqs.second())), + loadTornadoTensor(outputWeight), + ggmlType); } } 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 5aaa1111..0a013e52 100644 --- a/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java +++ b/src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java @@ -17,6 +17,8 @@ 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; @@ -71,7 +73,11 @@ public boolean shouldIncludeReasoning() { @Override public void forward(State state, int token, int position) { - InferenceCore.forwardJavaQwen2MoE(this, state, token, position); + if (plan == null) { + InferenceCore.forwardJavaQwen2MoE(this, state, token, position); + } else { + InferenceCore.forwardTornadoVM(this, state, token, position, tornadoVMPlan()); + } } @Override @@ -83,6 +89,13 @@ public List generateTokens(State state, int startPosition, List generateTokensGPU(State state, int startPosition, List promptTokens, Set stopTokens, int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) { - throw new UnsupportedOperationException("GPU MoE not yet supported"); + if (WITH_PREFILL_DECODE && TornadoVMMasterPlan.PREFILL_BATCH_SIZE > 1) { + throw new UnsupportedOperationException("Batch prefill/decode on GPU not yet implemented for Qwen2-MoE"); + } + if (WITH_PREFILL_DECODE) { + throw new UnsupportedOperationException("Prefill/decode on GPU not yet implemented for Qwen2-MoE"); + } + return InferenceEngine.generateTokensGPUQwen3(this, state, startPosition, promptTokens, + stopTokens, maxTokens, sampler, echo, onTokenGenerated, tornadoVMPlan); } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java b/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java index 5035d384..46f0e405 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java @@ -11,6 +11,8 @@ import uk.ac.manchester.tornado.api.ImmutableTaskGraph; import uk.ac.manchester.tornado.api.TornadoExecutionPlan; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import org.beehive.gpullama3.inference.state.Qwen2MoEState; +import org.beehive.gpullama3.validation.MoECorrectnessTrace; /** * Standard (single-token) GPU execution plan. @@ -81,6 +83,11 @@ public FloatArray tornadoVMForwardDecode(int position) { executionPlan.withGraph(taskGraphLayout.layerIdx(layer)) .withGridScheduler(tornadoVMForwardPlan.getGridScheduler()) .execute(); + if (MoECorrectnessTrace.isEnabled() && state instanceof Qwen2MoEState moeState) { + MoECorrectnessTrace.recordGpuRouter(position, layer, + moeState.wrapRawRouterLogits, moeState.wrapSelectedExperts, + moeState.wrapRoutingWeights); + } } state.tempLogits.clear(); state.wrapLogits.clear(); @@ -90,6 +97,7 @@ public FloatArray tornadoVMForwardDecode(int position) { logitsGraph.withCUDAGraph(); } logitsGraph.execute(); + MoECorrectnessTrace.recordGpuLogits(position, state.wrapLogits); return state.wrapLogits; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java new file mode 100644 index 00000000..91a7dd46 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java @@ -0,0 +1,456 @@ +package org.beehive.gpullama3.tornadovm.kernels; + +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.math.TornadoMath; +import uk.ac.manchester.tornado.api.types.arrays.ByteArray; +import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import uk.ac.manchester.tornado.api.types.arrays.IntArray; + +/** + * TornadoVM kernels specific to Qwen2-MoE / Qwen1.5-MoE inference. + */ +public final class Qwen2MoEKernels { + + /** Copies raw router scores before softmax/top-K modifies its input buffer. */ + public static void copyRouterLogits(KernelContext context, FloatArray source, + FloatArray destination, int numberOfExperts) { + int expert = context.globalIdx; + if (expert < numberOfExperts) { + destination.set(expert, source.get(expert)); + } + } + + private static final int Q8_0_BLOCK_SIZE = 32; + private static final int Q8_0_BLOCK_BYTES = 34; + + private Qwen2MoEKernels() { + } + + /** + * Converts router scores to probabilities and selects the highest-scoring + * {@code topK} experts for one token. + * + *

Inputs and outputs are GPU-resident TornadoVM arrays:

+ *
    + *
  • {@code routerLogits}: one raw score per expert
    • + *
  • {@code selectedExperts}: output expert indices, length {@code topK}
    • + *
  • {@code routingWeights}: output probabilities, length {@code topK}
    • + *
+ * + *

The first implementation deliberately uses one GPU thread for this + * small operation (60 experts in the target model). A later optimization + * can parallelize the reductions within one workgroup.

+ */ + public static void softmaxAndTopK( + KernelContext context, + FloatArray routerLogits, + IntArray selectedExperts, + FloatArray routingWeights, + int numberOfExperts, + int topK) { + + // The first correctness-oriented implementation is serial. Without + // this guard, every thread would race to overwrite the same buffers. + if (context.groupIdx != 0 || context.localIdx != 0) { + return; + } + + // Find the maximum first to keep the softmax numerically stable. + float maxLogit = Float.NEGATIVE_INFINITY; + int maxIndex = -1; + for (int expert = 0; expert < numberOfExperts; expert++) { + float logit = routerLogits.get(expert); + if (logit > maxLogit) { + maxLogit = logit; + maxIndex = expert; + } + } + + // Convert every router score to a probability over all experts. + float sumExp = 0.0f; + for (int expert = 0; expert < numberOfExperts; expert++) { + sumExp += TornadoMath.exp(routerLogits.get(expert) - maxLogit); + } + for (int expert = 0; expert < numberOfExperts; expert++) { + float probability = TornadoMath.exp(routerLogits.get(expert) - maxLogit) / sumExp; + routerLogits.set(expert, probability); + } + + // Select the top-K probabilities without renormalizing their sum. + selectedExperts.set(0, maxIndex); + routingWeights.set(0, routerLogits.get(maxIndex)); + routerLogits.set(maxIndex, Float.NEGATIVE_INFINITY); + + for (int slot = 1; slot < topK; slot++) { + maxLogit = Float.NEGATIVE_INFINITY; + maxIndex = -1; + for (int expert = 0; expert < numberOfExperts; expert++) { + if (routerLogits.get(expert) > maxLogit) { + maxLogit = routerLogits.get(expert); + maxIndex = expert; + } + } + selectedExperts.set(slot, maxIndex); + routingWeights.set(slot, routerLogits.get(maxIndex)); + routerLogits.set(maxIndex, Float.NEGATIVE_INFINITY); + } + } + + /** + * Computes the routed expert's gated activation for one top-K slot: + * {@code SiLU(W_gate[expert] * x) * (W_up[expert] * x)}. + * + *

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

+ */ + public static void fusedRoutedExpertGateUpSwiGLUQ8_0( + KernelContext context, + FloatArray input, + IntArray selectedExperts, + int slot, + ByteArray gateExperts, + ByteArray upExperts, + FloatArray expertHidden, + int dim, + int moeHiddenDim, + int numberOfExperts, + int localWorkGroupSize) { + + int rowId = context.groupIdx; + int localId = context.localIdx; + + int expert = selectedExperts.get(slot); + if (rowId >= moeHiddenDim || expert < 0 || expert >= numberOfExperts) { + return; + } + + // Locate this output row within the selected expert's stacked matrix. + int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = + (expert * moeHiddenDim + rowId) * blocksPerRow; + + // One workgroup cooperates on the Gate and Up dot products for this row. + float gatePartialSum = 0.0f; + float upPartialSum = 0.0f; + + for (int column = localId; + column < dim; + column += localWorkGroupSize) { + + // Byte offset of the first byte of the Q8_0 block that contains this column. + // Each block occupies 34 bytes: a 2-byte FP16 scale plus 32 int8 quants. + int blockByteOffset = + (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + + // Skip the 2-byte scale at the block start and locate this column's int8 quant. + int quantOffset = + blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; + + float inputValue = input.get(column); + + // getHalfFloat reads the FP16 scale from the first two block bytes, then converts it to FP32. + float gateScale = + gateExperts.getHalfFloat(blockByteOffset).getFloat32(); + float upScale = + upExperts.getHalfFloat(blockByteOffset).getFloat32(); + + byte gateQuant = gateExperts.get(quantOffset); + byte upQuant = upExperts.get(quantOffset); + + float gateWeight = gateQuant * gateScale; + float upWeight = upQuant * upScale; + + gatePartialSum += gateWeight * inputValue; + upPartialSum += upWeight * inputValue; + + } + + // Sum the partial gate values from all threads in this workgroup. + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + localSums[localId] = gatePartialSum; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + float gate = localSums[0]; + + // Reuse local memory to sum the partial up values. + localSums[localId] = upPartialSum; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + + // One thread writes this output row after both reductions are complete. + if (localId == 0) { + float up = localSums[0]; + float siluGate = gate / (1.0f + TornadoMath.exp(-gate)); + expertHidden.set(rowId, siluGate * up); + } + } + + /** + * Down-projects one selected expert and accumulates its routed contribution: + * {@code residual += routingWeight[slot] * W_down[expert] * expertHidden}. + */ + public static void routedExpertDownProjectAndAccumulateQ8_0( + KernelContext context, + FloatArray expertHidden, + FloatArray residual, + IntArray selectedExperts, + FloatArray routingWeights, + int slot, + ByteArray downExperts, + int dim, + int moeHiddenDim, + int numberOfExperts, + int localWorkGroupSize) { + + // One workgroup produces one element of the down-projected vector. + int rowId = context.groupIdx; + int localId = context.localIdx; + if (rowId >= dim) { + return; + } + + int expert = selectedExperts.get(slot); + if (expert < 0 || expert >= numberOfExperts) { + return; + } + float routingWeight = routingWeights.get(slot); + + // downExperts has the logical shape [experts, dim, moeHiddenDim]. + int blocksPerRow = (moeHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + int rowBlockOffset = (expert * dim + rowId) * blocksPerRow; + + // Every thread accumulates a different subset of this row's dot product. + float partialSum = 0.0f; + for (int column = localId; + column < moeHiddenDim; + column += localWorkGroupSize) { + // The start byte of the Q8_0 block holding this down-projection weight. + // Block layout: a 2-byte FP16 scale followed by 32 int8 quants. + int blockByteOffset = + (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + + // Quants begin immediately after the scale; column % 32 is the index within this block. + int quantOffset = blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; + + float weight = downExperts.get(quantOffset) + * downExperts.getHalfFloat(blockByteOffset).getFloat32(); + partialSum += weight * expertHidden.get(column); + } + + // Combine all thread-local partial sums into the completed output row. + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + localSums[localId] = partialSum; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + + if (localId == 0) { + float outputValue = localSums[0]; + residual.set(rowId, residual.get(rowId) + routingWeight * outputValue); + } + } + + /** Computes {@code SiLU(sharedGate * input) * (sharedUp * input)}. */ + public static void sharedExpertGateUpSwiGLUQ8_0( + KernelContext context, + FloatArray input, + ByteArray sharedGate, + ByteArray sharedUp, + FloatArray sharedHidden, + int dim, + int sharedExpertHiddenDim, + int localWorkGroupSize) { + + // One workgroup computes one shared-expert hidden output row. + int rowId = context.groupIdx; + int localId = context.localIdx; + + if (rowId >= sharedExpertHiddenDim) { + return; + } + + // Number of Q8_0 blocks needed for one row of dim weights. + int blocksPerRow = (dim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + + // Index of this row's first Q8_0 block in the shared weight array. + int rowBlockOffset = rowId * blocksPerRow; + + float gatePartialSum = 0.0f; + float upPartialSum = 0.0f; + + for (int column = localId; + column < dim; + column += localWorkGroupSize) { + + // Byte offset of the first byte of the Q8_0 block that contains this column. + // Each block occupies 34 bytes: a 2-byte FP16 scale plus 32 int8 quants. + int blockByteOffset = + (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + + // Skip the 2-byte scale at the block start and locate this column's int8 quant. + int quantOffset = + blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; + + float inputValue = input.get(column); + + // getHalfFloat reads the FP16 scale from the first two block bytes, then converts it to FP32. + float gateScale = + sharedGate.getHalfFloat(blockByteOffset).getFloat32(); + float upScale = + sharedUp.getHalfFloat(blockByteOffset).getFloat32(); + + byte gateQuant = sharedGate.get(quantOffset); + byte upQuant = sharedUp.get(quantOffset); + + float gateWeight = gateQuant * gateScale; + float upWeight = upQuant * upScale; + + gatePartialSum += gateWeight * inputValue; + upPartialSum += upWeight * inputValue; + + + } + + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + + localSums[localId] = gatePartialSum; + context.localBarrier(); + + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + + float gate = localSums[0]; + + // Reuse local memory to sum the partial up values. + localSums[localId] = upPartialSum; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + + // One thread writes this output row after both reductions are complete. + if (localId == 0) { + float up = localSums[0]; + float siluGate = gate / (1.0f + TornadoMath.exp(-gate)); + sharedHidden.set(rowId, siluGate * up); + } + } + + /** Down-projects the shared expert hidden vector into the model dimension. */ + public static void sharedExpertDownProjectQ8_0( + KernelContext context, + FloatArray sharedHidden, + ByteArray sharedDown, + FloatArray sharedOutput, + int dim, + int sharedExpertHiddenDim, + int localWorkGroupSize) { + + + int rowId = context.groupIdx; + int localId = context.localIdx; + if (rowId >= dim) { + return; + } + + // Number of Q8_0 blocks needed for one row of dim weights. + int blocksPerRow = (sharedExpertHiddenDim + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; + + // Index of this row's first Q8_0 block in the shared weight array. + int rowBlockOffset = rowId * blocksPerRow; + + float partialSum = 0.0f; + for (int column = localId; + column < sharedExpertHiddenDim; + column += localWorkGroupSize) { + // The start byte of the Q8_0 block holding this down-projection weight. + // Block layout: a 2-byte FP16 scale followed by 32 int8 quants. + int blockByteOffset = + (rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES; + + // Quants begin immediately after the scale; column % 32 is the index within this block. + int quantOffset = blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE; + + float weight = sharedDown.get(quantOffset) + * sharedDown.getHalfFloat(blockByteOffset).getFloat32(); + partialSum += weight * sharedHidden.get(column); + } + + // Combine all thread-local partial sums into the completed output row. + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + localSums[localId] = partialSum; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + + if (localId == 0) { + float outputValue = localSums[0]; + sharedOutput.set(rowId, outputValue); + } + } + + /** Computes the shared gate sigmoid and adds the weighted shared output to the residual. */ + public static void sharedExpertGateAndAccumulate( + KernelContext context, + FloatArray input, + FloatArray sharedGateInput, + FloatArray sharedOutput, + FloatArray residual, + int dim, + int localWorkGroupSize) { + int localId = context.localIdx; + + float partialScore = 0.0f; + + for (int column = localId; + column < dim; + column += localWorkGroupSize) { + partialScore += sharedGateInput.get(column) * input.get(column); + } + + float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize); + localSums[localId] = partialScore; + context.localBarrier(); + for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSums[localId] += localSums[localId + stride]; + } + context.localBarrier(); + } + float gateScore = localSums[0]; + float sharedWeight = + 1.0f / (1.0f + TornadoMath.exp(-gateScore)); + + for (int index = localId; + index < dim; + index += localWorkGroupSize) { + residual.set(index, + residual.get(index) + sharedWeight * sharedOutput.get(index)); + } + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java new file mode 100644 index 00000000..18215b7f --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java @@ -0,0 +1,300 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0; + +import org.beehive.gpullama3.inference.state.Qwen2MoEState; +import org.beehive.gpullama3.inference.weights.tornado.Qwen2MoETornadoWeights; +import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen2Kernels; +import org.beehive.gpullama3.tornadovm.kernels.Qwen2MoEKernels; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerComputeKernelsLayered; +import org.beehive.gpullama3.tornadovm.layers.AbstractTransformerLayerTaskGraphs; +import org.beehive.gpullama3.tornadovm.scheduling.SchedulerType; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.validation.MoECorrectnessTrace; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +/** + * Single-token Q8_0 TaskGraphs for Qwen2-MoE / Qwen1.5-MoE. + * + *

The attention block follows Qwen2. Its dense FFN is replaced by the + * routed-expert pipeline: normalize, route, choose top-K experts, execute each + * selected expert, and accumulate its weighted output into {@code wrapX}.

+ */ +public final class Qwen2MoEQ8_0FFNLayers + extends AbstractTransformerLayerTaskGraphs { + + private final Qwen2MoEState moeState; + + public Qwen2MoEQ8_0FFNLayers(String taskGraphName, + Qwen2MoEState state, + Qwen2MoETornadoWeights weights, + Qwen2MoEConfiguration config, + SchedulerType schedulerType) { + super(taskGraphName, state, weights, config, schedulerType); + this.moeState = state; + setupFFNLayers(); + } + + /** Sets the GPU worker grid for each task in each Transformer layer. */ + @Override + public GridScheduler updateGridScheduler(GridScheduler scheduler) { + WorkerGrid rmsNormWorker = WorkerGridFactory.createRmsNormWorker(config.dim(), moeState.localSize); + + WorkerGrid qkvWorker = workerForRows(config.dim() + 2 * config.kvDim()); + WorkerGrid qkvBiasWorker = new WorkerGrid1D(config.dim()); + qkvBiasWorker.setGlobalWork(config.dim(), 1, 1); + qkvBiasWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); + + WorkerGrid ropeWorker = new WorkerGrid2D(config.numberOfHeads(), config.headSize() / 2); + ropeWorker.setGlobalWork(config.numberOfHeads(), config.headSize() / 2, 1); + ropeWorker.setLocalWork(1, 1, 1); + + int attentionLocalSize = Math.min(config.headSize(), 64); + WorkerGrid attentionWorker = new WorkerGrid1D(config.numberOfHeads() * attentionLocalSize); + attentionWorker.setLocalWork(attentionLocalSize, 1, 1); + + WorkerGrid dimElementWorker = new WorkerGrid1D(config.dim()); + dimElementWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); + WorkerGrid dimWorker = workerForRows(config.dim()); + WorkerGrid routerWorker = workerForRows(config.numberOfExperts()); + int routerCopyGlobalSize = ((config.numberOfExperts() + LOCAL_WORK_GROUP_SIZE_ALLOC - 1) + / LOCAL_WORK_GROUP_SIZE_ALLOC) * LOCAL_WORK_GROUP_SIZE_ALLOC; + WorkerGrid routerCopyWorker = new WorkerGrid1D(routerCopyGlobalSize); + routerCopyWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); + WorkerGrid topKWorker = new WorkerGrid1D(LOCAL_WORK_GROUP_SIZE_ALLOC); + topKWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); + WorkerGrid expertHiddenWorker = workerForRows(config.moeHiddenDim()); + WorkerGrid sharedHiddenWorker = workerForRows(config.sharedExpertHiddenDim()); + + for (int layer = 0; layer < config.numberOfLayers(); layer++) { + String prefix = "layer_" + layer + "."; + scheduler.addWorkerGrid(prefix + "attn_rms_reduce", rmsNormWorker); + scheduler.addWorkerGrid(prefix + "attn_rms_qkv_projection", qkvWorker); + scheduler.addWorkerGrid(prefix + "fused_qkv_bias", qkvBiasWorker); + scheduler.addWorkerGrid(prefix + "rope_and_kv_cache", ropeWorker); + scheduler.addWorkerGrid(prefix + "attention", attentionWorker); + scheduler.addWorkerGrid(prefix + "attn_output_proj", dimWorker); + scheduler.addWorkerGrid(prefix + "ffn_rms_reduce", rmsNormWorker); + scheduler.addWorkerGrid(prefix + "ffn_rms_apply", dimElementWorker); + scheduler.addWorkerGrid(prefix + "router_projection", routerWorker); + if (MoECorrectnessTrace.isEnabled()) { + scheduler.addWorkerGrid(prefix + "router_trace_copy", routerCopyWorker); + } + scheduler.addWorkerGrid(prefix + "router_softmax_topk", topKWorker); + for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) { + scheduler.addWorkerGrid(prefix + "routed_expert_gate_up_" + slot, expertHiddenWorker); + scheduler.addWorkerGrid(prefix + "routed_expert_down_" + slot, dimWorker); + } + scheduler.addWorkerGrid(prefix + "shared_expert_gate_up", sharedHiddenWorker); + scheduler.addWorkerGrid(prefix + "shared_expert_down", dimWorker); + scheduler.addWorkerGrid(prefix + "shared_expert_gate_and_accumulate", topKWorker); + } + return scheduler; + } + + /** Creates a grid where one GPU work-group computes one output row. */ + private WorkerGrid workerForRows(int rows) { + WorkerGrid worker = new WorkerGrid1D(rows * LOCAL_WORK_GROUP_SIZE_ALLOC); + worker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); + return worker; + } + + /** + * Creates the complete GPU TaskGraph for one Transformer layer. + * {@code layerIndex} selects that layer's weights. + */ + @Override + protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { + TaskGraph layer = new TaskGraph("layer_" + layerIndex); + // Reuse wrapX produced by the previous TaskGraph on the GPU. + layer.consumeFromDevice(moeState.wrapX); + // Upload this layer's read-only weights from CPU to GPU on the first execution. + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.q_biasLayered[layerIndex].asFloatArray(), + weights.k_biasLayered[layerIndex].asFloatArray(), + weights.v_biasLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.routerGateLayered[layerIndex].asFloatArray(), + weights.gateExpertsLayered[layerIndex].asByteArray(), + weights.upExpertsLayered[layerIndex].asByteArray(), + weights.downExpertsLayered[layerIndex].asByteArray(), + weights.sharedGateLayered[layerIndex].asByteArray(), + weights.sharedUpLayered[layerIndex].asByteArray(), + weights.sharedDownLayered[layerIndex].asByteArray(), + weights.sharedGateInputLayered[layerIndex].asFloatArray()); + layer = configureLayerDataTransfers(layer, layerIndex); + + configureAttention(layer, layerIndex); + configureRoutedExperts(layer, layerIndex); + layer.persistOnDevice(moeState.wrapX); + return layer; + } + + /** Adds the normal Qwen2 attention tasks to this layer's TaskGraph. */ + private void configureAttention(TaskGraph layer, int layerIndex) { + layer.task("attn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayer, + context, moeState.temp, moeState.wrapX, + config.dim(), config.rmsNormEps(), moeState.localSize); + + if (shouldUseFinalNormalization()) { + layer.task("attn_rms_finalize", + TransformerComputeKernelsLayered::reductionFinalNormalization, + context, moeState.temp, config.dim(), config.rmsNormEps()); + } + + layer.task("attn_rms_qkv_projection", + Qwen3Kernels::fusedRmsNormQKVMatmulQ8_0, + context, moeState.wrapX, moeState.wrapQ, moeState.wrapK, moeState.wrapV, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), moeState.temp, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + config.dim(), config.dim(), config.kvDim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + layer.task("fused_qkv_bias", + TransformerComputeKernelsLayered::fusedQKvBiasAddition, + context, moeState.wrapQ, moeState.wrapK, + weights.q_biasLayered[layerIndex].asFloatArray(), moeState.wrapV, + weights.k_biasLayered[layerIndex].asFloatArray(), + weights.v_biasLayered[layerIndex].asFloatArray(), + config.dim(), config.kvDim()); + + layer.task("rope_and_kv_cache", Qwen3Kernels::ropeRotationWithCacheCopy, + context, moeState.positionHolder, moeState.wrapQ, moeState.wrapK, moeState.wrapV, + moeState.wrapKeyCache, moeState.wrapValueCache, + config.numberOfKeyValueHeads(), config.headSize(), config.kvDim(), + layerIndex, config.contextLength()); + + layer.task("attention", Qwen2Kernels::processHeadsFlashAttention, + context, moeState.wrapQ, moeState.wrapKeyCache, moeState.wrapValueCache, + moeState.wrapXb, config.numberOfHeads(), config.headSize(), config.kvDim(), + config.kvMul(), moeState.positionHolder, layerIndex, config.contextLength()); + + layer.task("attn_output_proj", + TransformerComputeKernelsLayered::matrixVectorGenericWithResidualQ8_0Byte, + context, moeState.wrapXb, moeState.wrapX, + weights.woLayered[layerIndex].asByteArray(), + config.dim(), config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + } + + /** + * Adds router, top-K, and selected-expert FFN tasks to this layer's TaskGraph. + * Their weighted outputs are added to the residual vector. + */ + private void configureRoutedExperts(TaskGraph layer, int layerIndex) { + layer.task("ffn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayer, + context, moeState.tempFFN, moeState.wrapX, + config.dim(), config.rmsNormEps(), moeState.localSize); + + if (shouldUseFinalNormalization()) { + layer.task("ffn_rms_finalize", + TransformerComputeKernelsLayered::reductionFinalNormalization, + context, moeState.tempFFN, config.dim(), config.rmsNormEps()); + } + + layer.task("ffn_rms_apply", + TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, + context, moeState.wrapXb, moeState.wrapX, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), moeState.tempFFN); + + layer.task("router_projection", + TransformerComputeKernelsLayered::matrixVectorGeneric, + context, moeState.wrapXb, moeState.wrapRouterLogits, + weights.routerGateLayered[layerIndex].asFloatArray(), + config.dim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + // In correctness mode, preserve raw scores before softmax/top-K mutates them. + if (MoECorrectnessTrace.isEnabled()) { + layer.task("router_trace_copy", Qwen2MoEKernels::copyRouterLogits, + context, moeState.wrapRouterLogits, moeState.wrapRawRouterLogits, + config.numberOfExperts()); + } + + layer.task("router_softmax_topk", Qwen2MoEKernels::softmaxAndTopK, + context, moeState.wrapRouterLogits, moeState.wrapSelectedExperts, + moeState.wrapRoutingWeights, config.numberOfExperts(), config.numberOfExpertsUsed()); + + for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) { + layer.task("routed_expert_gate_up_" + slot, + Qwen2MoEKernels::fusedRoutedExpertGateUpSwiGLUQ8_0, + context, moeState.wrapXb, moeState.wrapSelectedExperts, slot, + weights.gateExpertsLayered[layerIndex].asByteArray(), + weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate, + config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + layer.task("routed_expert_down_" + slot, + Qwen2MoEKernels::routedExpertDownProjectAndAccumulateQ8_0, + context, moeState.wrapExpertGate, moeState.wrapX, + moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, slot, + weights.downExpertsLayered[layerIndex].asByteArray(), + config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); + } + + // The shared expert always runs; it does not depend on router top-K selection. + layer.task("shared_expert_gate_up", Qwen2MoEKernels::sharedExpertGateUpSwiGLUQ8_0, + context, moeState.wrapXb, + weights.sharedGateLayered[layerIndex].asByteArray(), + weights.sharedUpLayered[layerIndex].asByteArray(), moeState.wrapSharedGate, + config.dim(), config.sharedExpertHiddenDim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + layer.task("shared_expert_down", Qwen2MoEKernels::sharedExpertDownProjectQ8_0, + context, moeState.wrapSharedGate, + weights.sharedDownLayered[layerIndex].asByteArray(), moeState.wrapSharedOutput, + config.dim(), config.sharedExpertHiddenDim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + layer.task("shared_expert_gate_and_accumulate", Qwen2MoEKernels::sharedExpertGateAndAccumulate, + context, moeState.wrapXb, weights.sharedGateInputLayered[layerIndex].asFloatArray(), + moeState.wrapSharedOutput, moeState.wrapX, config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC); + + if (MoECorrectnessTrace.isEnabled()) { + layer.transferToHost(DataTransferMode.EVERY_EXECUTION, + moeState.wrapRawRouterLogits, moeState.wrapSelectedExperts, + moeState.wrapRoutingWeights); + } + } + + /** + * Configures which TaskGraph data is uploaded from the CPU or reused on the GPU. + */ + @Override + protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) { + if (layerIndex == 0) { + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, + moeState.positionHolder, moeState.temp, moeState.tempFFN); + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, moeState.wrapXb, moeState.wrapXb2, moeState.wrapQ, + moeState.wrapK, moeState.wrapV, moeState.wrapKeyCache, + moeState.wrapValueCache, moeState.wrapAtt, moeState.wrapRouterLogits, + moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, + moeState.wrapExpertGate, moeState.wrapSharedGate, moeState.wrapSharedOutput); + } else { + layer.consumeFromDevice(context, moeState.wrapXb, moeState.wrapXb2, + moeState.wrapQ, moeState.wrapK, moeState.wrapV, moeState.wrapKeyCache, + moeState.wrapValueCache, moeState.wrapAtt, moeState.wrapRouterLogits, + moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, + moeState.wrapExpertGate, moeState.wrapSharedGate, moeState.wrapSharedOutput, + moeState.positionHolder); + } + if (MoECorrectnessTrace.isEnabled()) { + if (layerIndex == 0) { + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + moeState.wrapRawRouterLogits); + } else { + layer.consumeFromDevice(moeState.wrapRawRouterLogits); + } + } + return layer; + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java index 042dd354..6e8ba29d 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/ForwardPlanFactory.java @@ -4,6 +4,7 @@ import org.beehive.gpullama3.inference.state.GraniteState; import org.beehive.gpullama3.inference.state.LlamaState; import org.beehive.gpullama3.inference.state.Phi3State; +import org.beehive.gpullama3.inference.state.Qwen2MoEState; import org.beehive.gpullama3.inference.state.Qwen2State; import org.beehive.gpullama3.inference.state.Qwen3State; import org.beehive.gpullama3.inference.state.State; @@ -24,6 +25,7 @@ import org.beehive.gpullama3.tornadovm.plan.components.q8_0.MistralQ8_0PlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Phi3Q8_0PlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen2Q8_0PlanComponents; +import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen2MoEQ8_0PlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.q8_0.Qwen3Q8_0PlanComponents; // @formatter:off @@ -109,6 +111,7 @@ private static ForwardPlan createQ8_0Plan(ExecutionMode mode, State state, Model case MISTRAL -> createMistralQ8_0Plan(mode, (LlamaState) state, model); case DEVSTRAL_2 -> createDevstralQ8_0Plan(mode, (DevstralState) state, model); case QWEN_2 -> createQwen2Q8_0Plan(mode, (Qwen2State) state, model); + case QWEN_2_MOE -> createQwen2MoEQ8_0Plan(mode, (Qwen2MoEState) state, model); case QWEN_3 -> createQwen3Q8_0Plan(mode, (Qwen3State) state, model); case PHI_3 -> createPhi3Q8_0Plan(mode, (Phi3State) state, model); case GRANITE -> createGraniteQ8_0Plan(mode, (GraniteState) state, model); @@ -175,6 +178,12 @@ private static ForwardPlan createQwen2Q8_0Plan(ExecutionMode mode, Qwen2State st return new SingleTokenForwardPlan(model, new Qwen2Q8_0PlanComponents(state, model)); } + private static ForwardPlan createQwen2MoEQ8_0Plan(ExecutionMode mode, Qwen2MoEState state, Model model) { + if (mode != ExecutionMode.STANDARD) + throw new UnsupportedOperationException(mode + " not yet supported for QWEN_2_MOE + Q8_0"); + return new SingleTokenForwardPlan(model, new Qwen2MoEQ8_0PlanComponents(state, model)); + } + private static ForwardPlan createQwen3FP16Plan(ExecutionMode mode, Qwen3State state, Model model) { BatchPrefillDecodeForwardPlanComponents components = new Qwen3FP16PlanComponents(state, model); return switch (mode) { diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java new file mode 100644 index 00000000..93a1443f --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen2MoEQ8_0PlanComponents.java @@ -0,0 +1,46 @@ +package org.beehive.gpullama3.tornadovm.plan.components.q8_0; + +import org.beehive.gpullama3.inference.state.Qwen2MoEState; +import org.beehive.gpullama3.inference.weights.tornado.Qwen2MoETornadoWeights; +import org.beehive.gpullama3.model.Model; +import org.beehive.gpullama3.model.qwen2.Qwen2MoEConfiguration; +import org.beehive.gpullama3.tornadovm.layers.AbstractLogitsTaskGraph; +import org.beehive.gpullama3.tornadovm.layers.Activation; +import org.beehive.gpullama3.tornadovm.layers.ActivationTaskGraph; +import org.beehive.gpullama3.tornadovm.layers.TransformerLayerTaskGraphs; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.LogitsQ8_0Layer; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.Qwen2MoEQ8_0FFNLayers; +import org.beehive.gpullama3.tornadovm.plan.components.SingleTokenForwardPlanComponents; +import org.beehive.gpullama3.tornadovm.scheduling.SchedulerDetectionService; +import org.beehive.gpullama3.tornadovm.scheduling.SchedulerType; + +/** Assembles the single-token Q8_0 GPU components for Qwen2-MoE. */ +public final class Qwen2MoEQ8_0PlanComponents implements SingleTokenForwardPlanComponents { + + private final Qwen2MoEState state; + private final Qwen2MoETornadoWeights weights; + private final Qwen2MoEConfiguration config; + private final SchedulerType schedulerType; + + public Qwen2MoEQ8_0PlanComponents(Qwen2MoEState state, Model model) { + this.state = state; + this.config = (Qwen2MoEConfiguration) model.configuration(); + this.weights = (Qwen2MoETornadoWeights) model.weights(); + this.schedulerType = SchedulerDetectionService.determineSchedulerType(model); + } + + @Override + public ActivationTaskGraph singleTokenActivation() { + return new Activation("activationUpdate", state, weights, config); + } + + @Override + public TransformerLayerTaskGraphs singleTokenTransformerLayers() { + return new Qwen2MoEQ8_0FFNLayers("qwen2MoEFFN", state, weights, config, schedulerType); + } + + @Override + public AbstractLogitsTaskGraph singleTokenLogits(String previousGraphId) { + return new LogitsQ8_0Layer("logits", state, weights, config, previousGraphId, schedulerType); + } +} diff --git a/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java b/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java new file mode 100644 index 00000000..64ab1470 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java @@ -0,0 +1,192 @@ +package org.beehive.gpullama3.validation; + +import org.beehive.gpullama3.tensor.standard.FloatTensor; +import uk.ac.manchester.tornado.api.types.arrays.FloatArray; +import uk.ac.manchester.tornado.api.types.arrays.IntArray; + +import java.io.BufferedWriter; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Optional JSONL trace used to compare Qwen2-MoE CPU and GPU inference. + * + *

The trace is disabled unless {@code -Dllama.correctnessTrace=/path/file.jsonl} + * is supplied, so normal inference does not perform any file I/O.

+ */ +public final class MoECorrectnessTrace { + + private static final String TRACE_PATH = System.getProperty("llama.correctnessTrace"); + private static final boolean ENABLED = TRACE_PATH != null && !TRACE_PATH.isBlank(); + private static final AtomicInteger TOKEN_INDEX = new AtomicInteger(); + private static final BufferedWriter WRITER = createWriter(); + + private MoECorrectnessTrace() { + } + + public static boolean isEnabled() { + return ENABLED; + } + + public static void recordCpuRouter(int position, int layer, float[] logits, + int[] experts, float[] weights) { + if (!ENABLED) { + return; + } + writeRouterPrefix(position, layer); + writeFloatArray(logits); + write(",\"experts\":"); + writeIntArray(experts); + write(",\"weights\":"); + writeFloatArray(weights); + writeLine("}"); + } + + public static void recordGpuRouter(int position, int layer, FloatArray logits, + IntArray experts, FloatArray weights) { + if (!ENABLED) { + return; + } + writeRouterPrefix(position, layer); + writeFloatArray(logits); + write(",\"experts\":"); + writeIntArray(experts); + write(",\"weights\":"); + writeFloatArray(weights); + writeLine("}"); + } + + public static void recordCpuLogits(int position, FloatTensor logits) { + if (!ENABLED) { + return; + } + write("{\"type\":\"logits\",\"position\":" + position + ",\"values\":"); + writeFloatTensor(logits); + writeLine("}"); + } + + public static void recordGpuLogits(int position, FloatArray logits) { + if (!ENABLED) { + return; + } + write("{\"type\":\"logits\",\"position\":" + position + ",\"values\":"); + writeFloatArray(logits); + writeLine("}"); + } + + public static void recordToken(int tokenId) { + if (!ENABLED) { + return; + } + writeLine("{\"type\":\"token\",\"index\":" + TOKEN_INDEX.getAndIncrement() + + ",\"id\":" + tokenId + "}"); + } + + private static BufferedWriter createWriter() { + if (!ENABLED) { + return null; + } + try { + Path path = Path.of(TRACE_PATH); + Path parent = path.toAbsolutePath().getParent(); + if (parent != null) { + Files.createDirectories(parent); + } + BufferedWriter writer = Files.newBufferedWriter(path, + StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, + StandardOpenOption.WRITE); + Runtime.getRuntime().addShutdownHook(new Thread(() -> { + try { + writer.close(); + } catch (IOException ignored) { + // The process is already terminating. + } + })); + return writer; + } catch (IOException e) { + throw new UncheckedIOException("Cannot create correctness trace: " + TRACE_PATH, e); + } + } + + private static void writeRouterPrefix(int position, int layer) { + write("{\"type\":\"router\",\"position\":" + position + + ",\"layer\":" + layer + ",\"logits\":"); + } + + private static void writeFloatTensor(FloatTensor values) { + write("["); + for (int i = 0; i < values.size(); i++) { + if (i > 0) { + write(","); + } + write(Float.toString(values.getFloat(i))); + } + write("]"); + } + + private static void writeFloatArray(FloatArray values) { + write("["); + for (int i = 0; i < values.getSize(); i++) { + if (i > 0) { + write(","); + } + write(Float.toString(values.get(i))); + } + write("]"); + } + + private static void writeFloatArray(float[] values) { + write("["); + for (int i = 0; i < values.length; i++) { + if (i > 0) { + write(","); + } + write(Float.toString(values[i])); + } + write("]"); + } + + private static void writeIntArray(IntArray values) { + write("["); + for (int i = 0; i < values.getSize(); i++) { + if (i > 0) { + write(","); + } + write(Integer.toString(values.get(i))); + } + write("]"); + } + + private static void writeIntArray(int[] values) { + write("["); + for (int i = 0; i < values.length; i++) { + if (i > 0) { + write(","); + } + write(Integer.toString(values[i])); + } + write("]"); + } + + private static synchronized void write(String value) { + try { + WRITER.write(value); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + private static synchronized void writeLine(String value) { + try { + WRITER.write(value); + WRITER.newLine(); + WRITER.flush(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } +} From 2e843c64f1e3a0dc4903fb005fd43f7a40b21d14 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Sun, 2 Aug 2026 15:34:06 +0100 Subject: [PATCH 08/10] Update baseline branch name --- docs/qwen2-moe-gpu-baseline.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/qwen2-moe-gpu-baseline.md b/docs/qwen2-moe-gpu-baseline.md index e5611a3d..a74479ba 100644 --- a/docs/qwen2-moe-gpu-baseline.md +++ b/docs/qwen2-moe-gpu-baseline.md @@ -2,7 +2,7 @@ This document records the working single-token GPU baseline before further kernel and scheduling optimizations. The baseline is preserved on branch -`codex/qwen2-moe-gpu-baseline`. +`qwen2-moe-gpu-baseline`. ## Scope From b227bc2c7dec5c185d116f8fa7aaf12257ec944a Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Wed, 5 Aug 2026 22:57:04 +0100 Subject: [PATCH 09/10] chore(qwen2-moe): remove local validation artifacts --- docs/qwen2-moe-gpu-baseline.md | 143 ------------- scripts/compare_moe_correctness.py | 148 -------------- scripts/summarize_tornado_profile.py | 170 ---------------- .../gpullama3/inference/InferenceCore.java | 14 -- .../gpullama3/inference/InferenceEngine.java | 5 +- .../inference/state/Qwen2MoEState.java | 2 - .../TornadoVMMasterPlanSingleToken.java | 8 - .../tornadovm/kernels/Qwen2MoEKernels.java | 9 - .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 28 --- .../validation/MoECorrectnessTrace.java | 192 ------------------ 10 files changed, 1 insertion(+), 718 deletions(-) delete mode 100644 docs/qwen2-moe-gpu-baseline.md delete mode 100755 scripts/compare_moe_correctness.py delete mode 100644 scripts/summarize_tornado_profile.py delete mode 100644 src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java diff --git a/docs/qwen2-moe-gpu-baseline.md b/docs/qwen2-moe-gpu-baseline.md deleted file mode 100644 index a74479ba..00000000 --- a/docs/qwen2-moe-gpu-baseline.md +++ /dev/null @@ -1,143 +0,0 @@ -# Qwen2-MoE Q8_0 GPU Baseline - -This document records the working single-token GPU baseline before further -kernel and scheduling optimizations. The baseline is preserved on branch -`qwen2-moe-gpu-baseline`. - -## Scope - -- Model: `Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf` -- GGUF file size: approximately 15 GB -- Weight format: Q8_0 (32 int8 values and one FP16 scale per block) -- Execution mode: single-token inference -- GPU backend: TornadoVM PTX -- Unsupported in this baseline: sequential prefill/decode and batch prefill/decode - -The GPU path includes Qwen2 attention, router projection, softmax and Top-K, -four routed experts, the shared expert, residual accumulation, and final logits. - -## Test Environment - -- Server: `storm` -- GPU: NVIDIA GeForce RTX 4090, 24 GB -- TornadoVM SDK: 5.1.0, JDK 21, PTX and OpenCL backends installed -- Model path: `/home/mingyi/models/Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf` -- GPU memory limit: 20 GB - -Representative command: - -```bash -./llama-tornado \ - --gpu --ptx \ - --gpu-memory 20GB \ - --heap-min 2g --heap-max 8g \ - --model /home/mingyi/models/Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf \ - --prompt "Hi" \ - --temperature 0 \ - --seed 42 \ - --max-tokens 128 -``` - -## Correctness Results - -The repository contains an optional JSONL trace and comparison script that -compare generated token IDs, per-layer router logits, Top-K expert IDs, -routing weights, and final logits between CPU and GPU executions. - -When CPU activation quantization was disabled so that the CPU arithmetic more -closely matched the current GPU kernels, the common trace prefix produced: - -- 3 compared generated token IDs with no mismatch -- 20 final Top-1 predictions with no mismatch -- 4 Top-K expert-set mismatches out of 495 layer comparisons -- mean absolute router-logit error: 0.001323 -- mean absolute final-logit error: 0.01094 - -The first expert-set mismatch occurred around an almost tied routing decision. -This indicates that the main remaining differences are numerical rather than a -large structural error in the MoE pipeline. Longer correctness traces are still -required before claiming full numerical equivalence. - -## Throughput Baseline - -Three interleaved 128-token PTX runs measured: - -| Run | Throughput | -|---:|---:| -| 1 | 17.75 tokens/s | -| 2 | 17.77 tokens/s | -| 3 | 17.58 tokens/s | -| **Mean** | **17.70 tokens/s** | - -A later 19-token smoke test reached 24.63 tokens/s. This short result is kept as -a health check, not as the main baseline, because short generations are more -sensitive to prompt length, warm-up, and measurement variance. - -## Kernel Profiling Results - -The TornadoVM profiler was run with 16 generated tokens. The first execution of -each TaskGraph was excluded because it includes initialization and initial -weight transfer, leaving 15 steady-state iterations. - -| Component | Time per token | Share of GPU kernel time | -|---|---:|---:| -| Attention | 5.222 ms | 28.62% | -| Shared expert | 4.796 ms | 26.29% | -| Routed Gate/Up | 3.539 ms | 19.40% | -| Routed Down | 2.536 ms | 13.90% | -| Router and Top-K | 1.179 ms | 6.46% | -| FFN RMSNorm | 0.536 ms | 2.94% | -| Other kernels | 0.435 ms | 2.39% | - -Routed and shared expert computation accounts for approximately 59.6% of GPU -kernel time. The most expensive individual tasks were: - -| Task | Time per token | Share of GPU kernel time | -|---|---:|---:| -| Attention kernel | 3.420 ms | 18.75% | -| Four routed Gate/Up kernels | 3.539 ms | 19.40% | -| Four routed Down kernels | 2.536 ms | 13.90% | -| Shared expert Gate/Up | 2.105 ms | 11.54% | -| Shared expert Down | 1.824 ms | 10.00% | -| Shared gate and accumulation | 0.867 ms | 4.75% | -| Router projection | 0.671 ms | 3.68% | -| Softmax and Top-K | 0.508 ms | 2.78% | - -Runtime-level profiler totals per token were: - -- GPU kernels: 18.244 ms -- copy-in: 4.962 ms -- runtime/profiler residual: 24.118 ms -- total TaskGraph time: 47.323 ms - -The residual is only an upper bound. It combines host dispatch, -synchronization, event handling, and profiler overhead, so it must not be -reported as pure kernel-launch time. - -## Q8_0 Activation Experiments - -Several experimental kernels quantized activations to Q8_0 and reused the -quantized values for integer dot products. These experiments were rolled back -from the working baseline: - -| Experiment | Result | -|---|---:| -| Routed Gate/Up only | approximately 2.5% faster | -| Routed and shared Gate/Up | approximately 3.4% faster | -| All tested Q8_0 matrix-vector paths | approximately 3.7% slower | - -The full conversion also increased numerical error. The current baseline -therefore keeps FP32 activations and reads Q8_0 weights by applying each block's -FP16 scale during the dot product. - -## Next Optimization Target - -Profiling shows that Top-K selection is too small to be the first target. The -next branch should investigate expert-task fusion and scheduling overhead: - -- reduce the four routed Gate/Up tasks to one task per layer where practical; -- reduce the four routed Down tasks to one task per layer where practical; -- preserve the current correctness trace as the regression oracle; -- compare steady-state throughput and profiler results against 17.70 tokens/s; -- keep any optimization only if it improves performance without unacceptable - correctness loss. diff --git a/scripts/compare_moe_correctness.py b/scripts/compare_moe_correctness.py deleted file mode 100755 index 7a16de73..00000000 --- a/scripts/compare_moe_correctness.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env python3 -"""Compare Qwen2-MoE CPU and GPU JSONL correctness traces.""" - -import argparse -import json -import math -from pathlib import Path - - -def load_trace(path: Path): - records = {"router": {}, "logits": {}, "token": {}} - lines = path.read_text(encoding="utf-8").splitlines() - for line_number, line in enumerate(lines, 1): - if not line.strip(): - continue - try: - record = json.loads(line) - except json.JSONDecodeError: - if line_number == len(lines): - print(f"warning: ignoring incomplete final line in {path}") - break - raise - kind = record["type"] - if kind == "router": - key = (record["position"], record["layer"]) - elif kind == "logits": - key = record["position"] - elif kind == "token": - key = record["index"] - else: - raise ValueError(f"{path}:{line_number}: unknown record type {kind}") - if key in records[kind]: - raise ValueError(f"{path}:{line_number}: duplicate {kind} key {key}") - records[kind][key] = record - return records - - -def errors(left, right): - if len(left) != len(right): - raise ValueError(f"array lengths differ: {len(left)} != {len(right)}") - differences = [abs(a - b) for a, b in zip(left, right)] - if not all(math.isfinite(value) for value in differences): - return math.inf, math.inf - return max(differences, default=0.0), sum(differences) / max(1, len(differences)) - - -def argmax(values): - return max(range(len(values)), key=values.__getitem__) - - -def matching_keys(cpu, gpu, kind, common_only): - cpu_keys = set(cpu[kind]) - gpu_keys = set(gpu[kind]) - if cpu_keys != gpu_keys and not common_only: - missing_gpu = sorted(cpu_keys - gpu_keys) - missing_cpu = sorted(gpu_keys - cpu_keys) - raise ValueError( - f"{kind} keys differ; missing on GPU={missing_gpu[:10]}, " - f"missing on CPU={missing_cpu[:10]}" - ) - return sorted(cpu_keys & gpu_keys) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("cpu_trace", type=Path) - parser.add_argument("gpu_trace", type=Path) - parser.add_argument( - "--common-only", action="store_true", - help="compare only records completed by both traces", - ) - args = parser.parse_args() - - cpu = load_trace(args.cpu_trace) - gpu = load_trace(args.gpu_trace) - - token_keys = matching_keys(cpu, gpu, "token", args.common_only) - token_mismatches = [ - key for key in token_keys - if cpu["token"][key]["id"] != gpu["token"][key]["id"] - ] - - router_keys = matching_keys(cpu, gpu, "router", args.common_only) - expert_order_mismatches = [] - expert_set_mismatches = [] - router_max = router_mean_sum = routing_max = routing_mean_sum = 0.0 - for key in router_keys: - cpu_record = cpu["router"][key] - gpu_record = gpu["router"][key] - if cpu_record["experts"] != gpu_record["experts"]: - expert_order_mismatches.append(key) - if set(cpu_record["experts"]) != set(gpu_record["experts"]): - expert_set_mismatches.append(key) - maximum, mean = errors(cpu_record["logits"], gpu_record["logits"]) - router_max = max(router_max, maximum) - router_mean_sum += mean - maximum, mean = errors(cpu_record["weights"], gpu_record["weights"]) - routing_max = max(routing_max, maximum) - routing_mean_sum += mean - - logits_keys = matching_keys(cpu, gpu, "logits", args.common_only) - top1_mismatches = [] - logits_max = logits_mean_sum = 0.0 - for key in logits_keys: - cpu_values = cpu["logits"][key]["values"] - gpu_values = gpu["logits"][key]["values"] - if argmax(cpu_values) != argmax(gpu_values): - top1_mismatches.append(key) - maximum, mean = errors(cpu_values, gpu_values) - logits_max = max(logits_max, maximum) - logits_mean_sum += mean - - print(f"tokens compared: {len(token_keys)}") - print(f"token ID mismatches: {len(token_mismatches)} {token_mismatches[:10]}") - print(f"router layers compared: {len(router_keys)}") - print(f"Top-K order mismatches: {len(expert_order_mismatches)} {expert_order_mismatches[:10]}") - print(f"Top-K set mismatches: {len(expert_set_mismatches)} {expert_set_mismatches[:10]}") - print(f"router logits max abs error: {router_max:.8g}") - print(f"router logits mean abs err: {router_mean_sum / max(1, len(router_keys)):.8g}") - print(f"routing weight max abs err: {routing_max:.8g}") - print(f"routing weight mean abs err: {routing_mean_sum / max(1, len(router_keys)):.8g}") - print(f"logit vectors compared: {len(logits_keys)}") - print(f"final Top-1 mismatches: {len(top1_mismatches)} {top1_mismatches[:10]}") - print(f"final logits max abs error: {logits_max:.8g}") - print(f"final logits mean abs error: {logits_mean_sum / max(1, len(logits_keys)):.8g}") - - if expert_set_mismatches: - first = expert_set_mismatches[0] - cpu_record = cpu["router"][first] - gpu_record = gpu["router"][first] - print(f"first Top-K set mismatch {first}:") - print(f" CPU experts={cpu_record['experts']} weights={cpu_record['weights']}") - print(f" GPU experts={gpu_record['experts']} weights={gpu_record['weights']}") - for expert in sorted(set(cpu_record["experts"]) | set(gpu_record["experts"])): - cpu_logit = cpu_record["logits"][expert] - gpu_logit = gpu_record["logits"][expert] - print( - f" expert {expert:2d}: CPU raw={cpu_logit: .8f}, " - f"GPU raw={gpu_logit: .8f}, abs err={abs(cpu_logit - gpu_logit):.3g}" - ) - - passed = not token_mismatches and not expert_set_mismatches and not top1_mismatches - print("RESULT: " + ("PASS" if passed else "FAIL")) - raise SystemExit(0 if passed else 1) - - -if __name__ == "__main__": - main() diff --git a/scripts/summarize_tornado_profile.py b/scripts/summarize_tornado_profile.py deleted file mode 100644 index 6f5d6f85..00000000 --- a/scripts/summarize_tornado_profile.py +++ /dev/null @@ -1,170 +0,0 @@ -#!/usr/bin/env python3 -"""Summarize concatenated TornadoVM profiler JSON objects by MoE task category.""" - -import argparse -import collections -import json -from pathlib import Path - - -ATTENTION_TASKS = { - "attn_rms_reduce", - "attn_rms_finalize", - "attn_rms_qkv_projection", - "fused_qkv_bias", - "rope_and_kv_cache", - "attention", - "attn_output_proj", -} - - -def load_objects(path: Path): - """Read TornadoVM's stream of adjacent top-level JSON objects.""" - text = path.read_text(encoding="utf-8") - decoder = json.JSONDecoder() - offset = 0 - objects = [] - while offset < len(text): - while offset < len(text) and text[offset].isspace(): - offset += 1 - if offset >= len(text): - break - value, offset = decoder.raw_decode(text, offset) - objects.append(value) - return objects - - -def task_category(full_name: str): - task = full_name.rsplit(".", 1)[-1] - if task.startswith("routed_expert_gate_up_"): - return "Routed Gate/Up" - if task.startswith("routed_expert_down_"): - return "Routed Down" - if task.startswith("shared_expert_"): - return "Shared expert" - if task.startswith("router_"): - return "Router + Top-K" - if task in ATTENTION_TASKS: - return "Attention" - if task.startswith("ffn_rms_"): - return "FFN RMSNorm" - return "Other" - - -def collect(objects, warmup_per_graph): - occurrences = collections.Counter() - categories = collections.Counter() - tasks = collections.Counter() - totals = collections.Counter() - graph_counts = collections.Counter() - task_count = 0 - - for record in objects: - graph_name = next(iter(record)) - graph = record[graph_name] - occurrence = occurrences[graph_name] - occurrences[graph_name] += 1 - if occurrence < warmup_per_graph: - continue - - graph_counts[graph_name] += 1 - totals["kernel"] += int(graph.get("TOTAL_KERNEL_TIME", 0)) - totals["copy_in"] += int(graph.get("COPY_IN_TIME", 0)) - totals["task_graph"] += int(graph.get("TOTAL_TASK_GRAPH_TIME", 0)) - totals["copy_bytes"] += int(graph.get("TOTAL_COPY_IN_SIZE_BYTES", 0)) - - for task_name, task in graph.items(): - if isinstance(task, dict): - task_time = int(task.get("TASK_KERNEL_TIME", 0)) - categories[task_category(task_name)] += task_time - tasks[task_name.rsplit(".", 1)[-1]] += task_time - task_count += 1 - - return occurrences, graph_counts, categories, tasks, totals, task_count - - -def milliseconds(nanoseconds): - return nanoseconds / 1_000_000.0 - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("profile", type=Path) - parser.add_argument( - "--warmup-per-graph", - type=int, - default=1, - help="ignore this many initial executions of every TaskGraph (default: 1)", - ) - args = parser.parse_args() - - objects = load_objects(args.profile) - occurrences, graphs, categories, tasks, totals, task_count = collect( - objects, args.warmup_per_graph - ) - retained_iterations = min(graphs.values(), default=0) - kernel_time = totals["kernel"] - residual = totals["task_graph"] - kernel_time - totals["copy_in"] - - print(f"profile objects: {len(objects)}") - print(f"TaskGraphs: {len(occurrences)}") - print(f"executions per graph: {sorted(set(occurrences.values()))}") - print(f"warmups skipped per graph: {args.warmup_per_graph}") - print(f"steady iterations retained: {retained_iterations}") - print(f"task executions retained: {task_count}") - print() - - print("GPU kernel breakdown") - for name, time_ns in categories.most_common(): - percentage = 100.0 * time_ns / kernel_time if kernel_time else 0.0 - per_iteration = ( - milliseconds(time_ns) / retained_iterations - if retained_iterations - else 0.0 - ) - print( - f"{name:20s} {milliseconds(time_ns):9.3f} ms total" - f" {per_iteration:8.3f} ms/iter {percentage:6.2f}%" - ) - - print() - print("Top individual tasks") - for name, time_ns in tasks.most_common(15): - percentage = 100.0 * time_ns / kernel_time if kernel_time else 0.0 - per_iteration = ( - milliseconds(time_ns) / retained_iterations - if retained_iterations - else 0.0 - ) - print( - f"{name:38s} {per_iteration:8.3f} ms/iter {percentage:6.2f}%" - ) - - print() - print("Runtime-level totals") - for name, value in ( - ("GPU kernels", kernel_time), - ("copy-in", totals["copy_in"]), - ("runtime/profiler residual", residual), - ("TaskGraph total", totals["task_graph"]), - ): - per_iteration = ( - milliseconds(value) / retained_iterations - if retained_iterations - else 0.0 - ) - print( - f"{name:26s} {milliseconds(value):9.3f} ms total" - f" {per_iteration:8.3f} ms/iter" - ) - - print() - print( - "Note: runtime/profiler residual is an upper bound containing host-side " - "dispatch, synchronization, event handling, and profiler overhead; it is " - "not a pure kernel-launch measurement." - ) - - -if __name__ == "__main__": - main() diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java index 4753764c..b68e951e 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceCore.java @@ -18,7 +18,6 @@ import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; -import org.beehive.gpullama3.validation.MoECorrectnessTrace; import java.lang.foreign.MemorySegment; @@ -372,13 +371,6 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // Qwen1.5-MoE uses norm_topk_prob=false: each selected expert's routing weight is // its probability over all experts without rescaling the top-k weights to sum to one. weights.routerGate[l].matmul(state.xb, moeState.routerLogits, numberOfExperts, dim); - float[] rawRouterLogits = null; - if (MoECorrectnessTrace.isEnabled()) { - rawRouterLogits = new float[numberOfExperts]; - for (int expert = 0; expert < numberOfExperts; expert++) { - rawRouterLogits[expert] = moeState.routerLogits.getFloat(expert); - } - } moeState.routerLogits.softmaxInPlace(0, numberOfExperts); int[] selectedExperts = new int[topK]; @@ -396,11 +388,6 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke routingWeights[i] = best; moeState.routerLogits.setFloat(index, Float.NEGATIVE_INFINITY); } - if (MoECorrectnessTrace.isEnabled()) { - MoECorrectnessTrace.recordCpuRouter(position, l, rawRouterLogits, - selectedExperts, routingWeights); - } - // Compute each selected expert and accumulate its weighted output. for (int j = 0; j < topK; j++) { int expert = selectedExperts[j]; @@ -431,7 +418,6 @@ public static FloatTensor forwardJavaQwen2MoE(Model model, State state, int toke // final rmsnorm + classifier (same as dense Qwen2) rmsnorm(state.x, state.x, weights.rms_final_weight, 0, dim, config.rmsNormEps()); weights.wcls.matmul(state.x, state.logits, config.vocabularySize(), dim); - MoECorrectnessTrace.recordCpuLogits(position, state.logits); return state.logits; } diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index d9549599..85ca2ec5 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java @@ -9,7 +9,6 @@ import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; import org.beehive.gpullama3.tornadovm.layers.type.fp16.LogitsFP16Layer; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; -import org.beehive.gpullama3.validation.MoECorrectnessTrace; import java.io.ByteArrayOutputStream; import java.util.ArrayList; @@ -210,7 +209,6 @@ public static List generateTokensQwen3(Model model, State state, int st // Track the generated token generatedTokens.add(nextToken); - MoECorrectnessTrace.recordToken(nextToken); // Notify via callback if provided if (onTokenGenerated != null) { @@ -451,7 +449,6 @@ public static List generateTokensGPUQwen3(Model model, State state, int // Track the generated token generatedTokens.add(nextToken); - MoECorrectnessTrace.recordToken(nextToken); // Notify via callback if provided if (onTokenGenerated != null) { @@ -681,4 +678,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/state/Qwen2MoEState.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java index 0a47e77b..d0872883 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java @@ -34,7 +34,6 @@ public class Qwen2MoEState extends Qwen2State { // separate from the CPU FloatTensor fields above: TaskGraph kernels operate on // TornadoVM arrays that can remain resident on the device between tasks. public final FloatArray wrapRouterLogits; - public final FloatArray wrapRawRouterLogits; public final IntArray wrapSelectedExperts; public final FloatArray wrapRoutingWeights; public final FloatArray wrapExpertGate; @@ -52,7 +51,6 @@ public Qwen2MoEState(Configuration config, int batchsize) { this.yTmp = ArrayFloatTensor.allocate(c.dim()); this.wrapRouterLogits = new FloatArray(c.numberOfExperts()); - this.wrapRawRouterLogits = new FloatArray(c.numberOfExperts()); this.wrapSelectedExperts = new IntArray(c.numberOfExpertsUsed()); this.wrapRoutingWeights = new FloatArray(c.numberOfExpertsUsed()); this.wrapExpertGate = new FloatArray(c.moeHiddenDim()); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java b/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java index 46f0e405..5035d384 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanSingleToken.java @@ -11,8 +11,6 @@ import uk.ac.manchester.tornado.api.ImmutableTaskGraph; import uk.ac.manchester.tornado.api.TornadoExecutionPlan; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; -import org.beehive.gpullama3.inference.state.Qwen2MoEState; -import org.beehive.gpullama3.validation.MoECorrectnessTrace; /** * Standard (single-token) GPU execution plan. @@ -83,11 +81,6 @@ public FloatArray tornadoVMForwardDecode(int position) { executionPlan.withGraph(taskGraphLayout.layerIdx(layer)) .withGridScheduler(tornadoVMForwardPlan.getGridScheduler()) .execute(); - if (MoECorrectnessTrace.isEnabled() && state instanceof Qwen2MoEState moeState) { - MoECorrectnessTrace.recordGpuRouter(position, layer, - moeState.wrapRawRouterLogits, moeState.wrapSelectedExperts, - moeState.wrapRoutingWeights); - } } state.tempLogits.clear(); state.wrapLogits.clear(); @@ -97,7 +90,6 @@ public FloatArray tornadoVMForwardDecode(int position) { logitsGraph.withCUDAGraph(); } logitsGraph.execute(); - MoECorrectnessTrace.recordGpuLogits(position, state.wrapLogits); return state.wrapLogits; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java index 91a7dd46..d16bc5b0 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java @@ -11,15 +11,6 @@ */ public final class Qwen2MoEKernels { - /** Copies raw router scores before softmax/top-K modifies its input buffer. */ - public static void copyRouterLogits(KernelContext context, FloatArray source, - FloatArray destination, int numberOfExperts) { - int expert = context.globalIdx; - if (expert < numberOfExperts) { - destination.set(expert, source.get(expert)); - } - } - private static final int Q8_0_BLOCK_SIZE = 32; private static final int Q8_0_BLOCK_BYTES = 34; 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 18215b7f..a845c8b0 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 org.beehive.gpullama3.validation.MoECorrectnessTrace; import uk.ac.manchester.tornado.api.GridScheduler; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; @@ -62,10 +61,6 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) { dimElementWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); WorkerGrid dimWorker = workerForRows(config.dim()); WorkerGrid routerWorker = workerForRows(config.numberOfExperts()); - int routerCopyGlobalSize = ((config.numberOfExperts() + LOCAL_WORK_GROUP_SIZE_ALLOC - 1) - / LOCAL_WORK_GROUP_SIZE_ALLOC) * LOCAL_WORK_GROUP_SIZE_ALLOC; - WorkerGrid routerCopyWorker = new WorkerGrid1D(routerCopyGlobalSize); - routerCopyWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); WorkerGrid topKWorker = new WorkerGrid1D(LOCAL_WORK_GROUP_SIZE_ALLOC); topKWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1); WorkerGrid expertHiddenWorker = workerForRows(config.moeHiddenDim()); @@ -82,9 +77,6 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) { scheduler.addWorkerGrid(prefix + "ffn_rms_reduce", rmsNormWorker); scheduler.addWorkerGrid(prefix + "ffn_rms_apply", dimElementWorker); scheduler.addWorkerGrid(prefix + "router_projection", routerWorker); - if (MoECorrectnessTrace.isEnabled()) { - scheduler.addWorkerGrid(prefix + "router_trace_copy", routerCopyWorker); - } scheduler.addWorkerGrid(prefix + "router_softmax_topk", topKWorker); for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) { scheduler.addWorkerGrid(prefix + "routed_expert_gate_up_" + slot, expertHiddenWorker); @@ -215,13 +207,6 @@ private void configureRoutedExperts(TaskGraph layer, int layerIndex) { weights.routerGateLayered[layerIndex].asFloatArray(), config.dim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC); - // In correctness mode, preserve raw scores before softmax/top-K mutates them. - if (MoECorrectnessTrace.isEnabled()) { - layer.task("router_trace_copy", Qwen2MoEKernels::copyRouterLogits, - context, moeState.wrapRouterLogits, moeState.wrapRawRouterLogits, - config.numberOfExperts()); - } - layer.task("router_softmax_topk", Qwen2MoEKernels::softmaxAndTopK, context, moeState.wrapRouterLogits, moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, config.numberOfExperts(), config.numberOfExpertsUsed()); @@ -258,11 +243,6 @@ private void configureRoutedExperts(TaskGraph layer, int layerIndex) { context, moeState.wrapXb, weights.sharedGateInputLayered[layerIndex].asFloatArray(), moeState.wrapSharedOutput, moeState.wrapX, config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC); - if (MoECorrectnessTrace.isEnabled()) { - layer.transferToHost(DataTransferMode.EVERY_EXECUTION, - moeState.wrapRawRouterLogits, moeState.wrapSelectedExperts, - moeState.wrapRoutingWeights); - } } /** @@ -287,14 +267,6 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) moeState.wrapExpertGate, moeState.wrapSharedGate, moeState.wrapSharedOutput, moeState.positionHolder); } - if (MoECorrectnessTrace.isEnabled()) { - if (layerIndex == 0) { - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - moeState.wrapRawRouterLogits); - } else { - layer.consumeFromDevice(moeState.wrapRawRouterLogits); - } - } return layer; } } diff --git a/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java b/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java deleted file mode 100644 index 64ab1470..00000000 --- a/src/main/java/org/beehive/gpullama3/validation/MoECorrectnessTrace.java +++ /dev/null @@ -1,192 +0,0 @@ -package org.beehive.gpullama3.validation; - -import org.beehive.gpullama3.tensor.standard.FloatTensor; -import uk.ac.manchester.tornado.api.types.arrays.FloatArray; -import uk.ac.manchester.tornado.api.types.arrays.IntArray; - -import java.io.BufferedWriter; -import java.io.IOException; -import java.io.UncheckedIOException; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; -import java.util.concurrent.atomic.AtomicInteger; - -/** - * Optional JSONL trace used to compare Qwen2-MoE CPU and GPU inference. - * - *

The trace is disabled unless {@code -Dllama.correctnessTrace=/path/file.jsonl} - * is supplied, so normal inference does not perform any file I/O.

- */ -public final class MoECorrectnessTrace { - - private static final String TRACE_PATH = System.getProperty("llama.correctnessTrace"); - private static final boolean ENABLED = TRACE_PATH != null && !TRACE_PATH.isBlank(); - private static final AtomicInteger TOKEN_INDEX = new AtomicInteger(); - private static final BufferedWriter WRITER = createWriter(); - - private MoECorrectnessTrace() { - } - - public static boolean isEnabled() { - return ENABLED; - } - - public static void recordCpuRouter(int position, int layer, float[] logits, - int[] experts, float[] weights) { - if (!ENABLED) { - return; - } - writeRouterPrefix(position, layer); - writeFloatArray(logits); - write(",\"experts\":"); - writeIntArray(experts); - write(",\"weights\":"); - writeFloatArray(weights); - writeLine("}"); - } - - public static void recordGpuRouter(int position, int layer, FloatArray logits, - IntArray experts, FloatArray weights) { - if (!ENABLED) { - return; - } - writeRouterPrefix(position, layer); - writeFloatArray(logits); - write(",\"experts\":"); - writeIntArray(experts); - write(",\"weights\":"); - writeFloatArray(weights); - writeLine("}"); - } - - public static void recordCpuLogits(int position, FloatTensor logits) { - if (!ENABLED) { - return; - } - write("{\"type\":\"logits\",\"position\":" + position + ",\"values\":"); - writeFloatTensor(logits); - writeLine("}"); - } - - public static void recordGpuLogits(int position, FloatArray logits) { - if (!ENABLED) { - return; - } - write("{\"type\":\"logits\",\"position\":" + position + ",\"values\":"); - writeFloatArray(logits); - writeLine("}"); - } - - public static void recordToken(int tokenId) { - if (!ENABLED) { - return; - } - writeLine("{\"type\":\"token\",\"index\":" + TOKEN_INDEX.getAndIncrement() - + ",\"id\":" + tokenId + "}"); - } - - private static BufferedWriter createWriter() { - if (!ENABLED) { - return null; - } - try { - Path path = Path.of(TRACE_PATH); - Path parent = path.toAbsolutePath().getParent(); - if (parent != null) { - Files.createDirectories(parent); - } - BufferedWriter writer = Files.newBufferedWriter(path, - StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, - StandardOpenOption.WRITE); - Runtime.getRuntime().addShutdownHook(new Thread(() -> { - try { - writer.close(); - } catch (IOException ignored) { - // The process is already terminating. - } - })); - return writer; - } catch (IOException e) { - throw new UncheckedIOException("Cannot create correctness trace: " + TRACE_PATH, e); - } - } - - private static void writeRouterPrefix(int position, int layer) { - write("{\"type\":\"router\",\"position\":" + position - + ",\"layer\":" + layer + ",\"logits\":"); - } - - private static void writeFloatTensor(FloatTensor values) { - write("["); - for (int i = 0; i < values.size(); i++) { - if (i > 0) { - write(","); - } - write(Float.toString(values.getFloat(i))); - } - write("]"); - } - - private static void writeFloatArray(FloatArray values) { - write("["); - for (int i = 0; i < values.getSize(); i++) { - if (i > 0) { - write(","); - } - write(Float.toString(values.get(i))); - } - write("]"); - } - - private static void writeFloatArray(float[] values) { - write("["); - for (int i = 0; i < values.length; i++) { - if (i > 0) { - write(","); - } - write(Float.toString(values[i])); - } - write("]"); - } - - private static void writeIntArray(IntArray values) { - write("["); - for (int i = 0; i < values.getSize(); i++) { - if (i > 0) { - write(","); - } - write(Integer.toString(values.get(i))); - } - write("]"); - } - - private static void writeIntArray(int[] values) { - write("["); - for (int i = 0; i < values.length; i++) { - if (i > 0) { - write(","); - } - write(Integer.toString(values[i])); - } - write("]"); - } - - private static synchronized void write(String value) { - try { - WRITER.write(value); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - } - - private static synchronized void writeLine(String value) { - try { - WRITER.write(value); - WRITER.newLine(); - WRITER.flush(); - } catch (IOException e) { - throw new UncheckedIOException(e); - } - } -} From 7a4f19777da78b86e3a627b5e9c18cc57b4993d2 Mon Sep 17 00:00:00 2001 From: Mingyi Jin Date: Thu, 6 Aug 2026 13:40:43 +0100 Subject: [PATCH 10/10] Fix Qwen2-MoE RMSNorm reduction race --- .../type/q8_0/Qwen2MoEQ8_0FFNLayers.java | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java index a845c8b0..0a439d80 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 @@ -42,7 +42,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(config.dim(), moeState.localSize); + WorkerGrid rmsNormWorker = WorkerGridFactory.createRmsNormWorker( + moeState.localSize, moeState.localSize); WorkerGrid qkvWorker = workerForRows(config.dim() + 2 * config.kvDim()); WorkerGrid qkvBiasWorker = new WorkerGrid1D(config.dim()); @@ -135,16 +136,10 @@ protected TaskGraph createFFNLayerTaskGraph(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", - TransformerComputeKernelsLayered::reductionOneBlockWithLayer, + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, context, moeState.temp, moeState.wrapX, config.dim(), config.rmsNormEps(), moeState.localSize); - if (shouldUseFinalNormalization()) { - layer.task("attn_rms_finalize", - TransformerComputeKernelsLayered::reductionFinalNormalization, - context, moeState.temp, config.dim(), config.rmsNormEps()); - } - layer.task("attn_rms_qkv_projection", Qwen3Kernels::fusedRmsNormQKVMatmulQ8_0, context, moeState.wrapX, moeState.wrapQ, moeState.wrapK, moeState.wrapV, @@ -186,16 +181,10 @@ private void configureAttention(TaskGraph layer, int layerIndex) { */ private void configureRoutedExperts(TaskGraph layer, int layerIndex) { layer.task("ffn_rms_reduce", - TransformerComputeKernelsLayered::reductionOneBlockWithLayer, + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, context, moeState.tempFFN, moeState.wrapX, config.dim(), config.rmsNormEps(), moeState.localSize); - if (shouldUseFinalNormalization()) { - layer.task("ffn_rms_finalize", - TransformerComputeKernelsLayered::reductionFinalNormalization, - context, moeState.tempFFN, config.dim(), config.rmsNormEps()); - } - layer.task("ffn_rms_apply", TransformerComputeKernelsLayered::reductionOneBlock2WithLayer, context, moeState.wrapXb, moeState.wrapX,