Skip to content

Add Qwen1.5-MoE Q8_0 CPU and single-token TornadoVM GPU inference #144

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mikepapadim merged 10 commits into from
Aug 6, 2026
Merged

Add Qwen1.5-MoE Q8_0 CPU and single-token TornadoVM GPU inference #144

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
97ef9d2
Add Qwen1.5-MoE (qwen2moe) CPU inference support
MRPRESIDENT66 Jul 4, 2026
46cea94
Fix Qwen2MoEState to override createStateFields with correct config type
MRPRESIDENT66 Jul 4, 2026
973cdf3
Fix Qwen2-MoE routing probabilities
MRPRESIDENT66 Jul 5, 2026
8ddca7d
Add ChatML turn separator for Qwen messages
MRPRESIDENT66 Jul 5, 2026
25dd61b
Optimize Q8_0 dot products with quantized activations
MRPRESIDENT66 Jul 5, 2026
71d49d3
Clean up Qwen2-MoE CPU implementation
MRPRESIDENT66 Jul 12, 2026
c2e9583
Add Qwen2-MoE single-token GPU baseline
MRPRESIDENT66 Aug 2, 2026
2e843c6
Update baseline branch name
MRPRESIDENT66 Aug 2, 2026
b227bc2
chore(qwen2-moe): remove local validation artifacts
MRPRESIDENT66 Aug 5, 2026
7a4f197
Fix Qwen2-MoE RMSNorm reduction race
MRPRESIDENT66 Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 175 additions & 4 deletions src/main/java/org/beehive/gpullama3/inference/InferenceCore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -260,6 +260,177 @@ 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);

// MoE FFN pre-normalization
rmsnorm(state.xb, state.x, weights.rms_ffn_weight[curLayer], 0, dim, config.rmsNormEps());

int numberOfExperts = config.numberOfExperts();
int topK = config.numberOfExpertsUsed();
int expertHiddenDim = config.moeHiddenDim();

// 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 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 < numberOfExperts; j++) {
if (moeState.routerLogits.getFloat(j) > best) {
best = moeState.routerLogits.getFloat(j);
index = j;
}
}
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]);
}

// 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, sharedExpertHiddenDim);

// Gate the shared expert output.
float gateScore = weights.sharedGateInp[l].dot(0, state.xb, 0, dim);
float sharedExpertWeight = 1f / (1f + (float) Math.exp(-gateScore));
state.x.saxpyInPlace(0, moeState.yTmp, 0, dim, sharedExpertWeight);
}

// 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();
final Qwen2StandardWeights weights = (Qwen2StandardWeights) model.weights();
Expand Down
119 changes: 119 additions & 0 deletions src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
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;
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 {

// 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;

// 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 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;
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());

this.wrapRouterLogits = new FloatArray(c.numberOfExperts());
this.wrapSelectedExperts = new IntArray(c.numberOfExpertsUsed());
this.wrapRoutingWeights = new FloatArray(c.numberOfExpertsUsed());
this.wrapExpertGate = new FloatArray(c.moeHiddenDim());
this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim());
this.wrapSharedOutput = new FloatArray(c.dim());
}

@Override
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);

// 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;
}
}
12 changes: 8 additions & 4 deletions src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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;

Expand Down
Loading
Loading