Skip to content

Optimize routed-expert kernel launches and clean up code #145

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 3 commits into from
Aug 6, 2026
Merged

Optimize routed-expert kernel launches and clean up code #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3940767
[moe] Fuse per-slot routed-expert launches: 8 kernel launches per lay…
mikepapadim Aug 6, 2026
c097edf
Merge pull request #1 from mikepapadim/opt/moe-fused-expert-launches
MRPRESIDENT66 Aug 6, 2026
d022b75
Clean up fused routed-expert kernels
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
2 changes: 1 addition & 1 deletion 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
Expand Up @@ -53,7 +53,7 @@ public Qwen2MoEState(Configuration config, int batchsize) {
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.wrapExpertGate = new FloatArray(c.moeHiddenDim() * c.numberOfExpertsUsed());
this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim());
this.wrapSharedOutput = new FloatArray(c.dim());
}
Expand Down
174 changes: 82 additions & 92 deletions src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen2MoEKernels.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,17 +88,17 @@ public static void softmaxAndTopK(
}

/**
* Computes the routed expert's gated activation for one top-K slot:
* {@code SiLU(W_gate[expert] * x) * (W_up[expert] * x)}.
* Gate/Up + SiLU for <b>all</b> routed slots in a single launch.
*
* <p>The expert matrices are stacked in one Q8_0 tensor per layer. The
* selected expert id determines which matrix slice this kernel reads.</p>
* <p>The slot index is folded into the work-group id, so all top-K slots execute in one
* launch. Each slot writes its own {@code moeHiddenDim}-sized window of {@code expertHidden},
* so the slots never alias.
*/
public static void fusedRoutedExpertGateUpSwiGLUQ8_0(
public static void fusedRoutedExpertsGateUpSwiGLUQ8_0(
KernelContext context,
FloatArray input,
IntArray selectedExperts,
int slot,
int expertsUsed,
ByteArray gateExperts,
ByteArray upExperts,
FloatArray expertHidden,
Expand All @@ -107,56 +107,42 @@ public static void fusedRoutedExpertGateUpSwiGLUQ8_0(
int numberOfExperts,
int localWorkGroupSize) {

int rowId = context.groupIdx;
int flatGroupId = context.groupIdx;
int localId = context.localIdx;

int expert = selectedExperts.get(slot);
if (rowId >= moeHiddenDim || expert < 0 || expert >= numberOfExperts) {
return;
int slot = flatGroupId / moeHiddenDim;
int rowId = flatGroupId - slot * moeHiddenDim;

// A work-group whose slot or row falls outside the launch still has to reach the
// barriers below, so the guard only suppresses the memory accesses and the store.
boolean active = slot < expertsUsed && rowId < moeHiddenDim;
int expert = 0;
if (active) {
expert = selectedExperts.get(slot);
active = expert >= 0 && expert < numberOfExperts;
}

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

if (active) {
for (int column = localId; column < dim; column += localWorkGroupSize) {
int blockByteOffset =
(rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES;
int quantOffset =
blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE;

float inputValue = input.get(column);
float gateScale = gateExperts.getHalfFloat(blockByteOffset).getFloat32();
float upScale = upExperts.getHalfFloat(blockByteOffset).getFloat32();

gatePartialSum += (gateExperts.get(quantOffset) * gateScale) * inputValue;
upPartialSum += (upExperts.get(quantOffset) * upScale) * inputValue;
}
}

// Sum the partial gate values from all threads in this workgroup.
float[] localSums = context.allocateFloatLocalArray(localWorkGroupSize);
localSums[localId] = gatePartialSum;
context.localBarrier();
Expand All @@ -168,7 +154,6 @@ public static void fusedRoutedExpertGateUpSwiGLUQ8_0(
}
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) {
Expand All @@ -178,80 +163,85 @@ public static void fusedRoutedExpertGateUpSwiGLUQ8_0(
context.localBarrier();
}

// One thread writes this output row after both reductions are complete.
if (localId == 0) {
if (localId == 0 && active) {
float up = localSums[0];
float siluGate = gate / (1.0f + TornadoMath.exp(-gate));
expertHidden.set(rowId, siluGate * up);
expertHidden.set(slot * moeHiddenDim + rowId, siluGate * up);
}
}

/**
* Down-projects one selected expert and accumulates its routed contribution:
* {@code residual += routingWeight[slot] * W_down[expert] * expertHidden}.
* Down-projects <b>all</b> routed slots and accumulates them into the residual in one launch.
*
* <p>Each slot retains the original reduction tree and is accumulated in slot order. The
* optimization removes the per-slot kernel launches and writes the final residual once.
*/
public static void routedExpertDownProjectAndAccumulateQ8_0(
public static void routedExpertsDownProjectAndAccumulateQ8_0(
KernelContext context,
FloatArray expertHidden,
FloatArray residual,
IntArray selectedExperts,
FloatArray routingWeights,
int slot,
int expertsUsed,
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;
}
boolean active = rowId < dim;

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[] localSums = context.allocateFloatLocalArray(localWorkGroupSize);

float weight = downExperts.get(quantOffset)
* downExperts.getHalfFloat(blockByteOffset).getFloat32();
partialSum += weight * expertHidden.get(column);
// Lane 0 carries the running residual across slots. Each slot is reduced with the same
// tree and folded in the same order as the per-slot kernels, so the result is bit-identical
// to launching them separately - only the launch count and the residual write change.
float running = 0.0f;
if (localId == 0 && active) {
running = residual.get(rowId);
}

// 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];
for (int slot = 0; slot < expertsUsed; slot++) {
int expert = selectedExperts.get(slot);
boolean slotActive = active && expert >= 0 && expert < numberOfExperts;

float partialSum = 0.0f;
if (slotActive) {
int rowBlockOffset = (expert * dim + rowId) * blocksPerRow;
int hiddenBase = slot * moeHiddenDim;
for (int column = localId;
column < moeHiddenDim;
column += localWorkGroupSize) {
int blockByteOffset =
(rowBlockOffset + column / Q8_0_BLOCK_SIZE) * Q8_0_BLOCK_BYTES;
int quantOffset = blockByteOffset + 2 + column % Q8_0_BLOCK_SIZE;

float weight = downExperts.get(quantOffset)
* downExperts.getHalfFloat(blockByteOffset).getFloat32();
partialSum += weight * expertHidden.get(hiddenBase + column);
}
}

context.localBarrier();
localSums[localId] = partialSum;
context.localBarrier();
for (int stride = localWorkGroupSize / 2; stride > 0; stride >>= 1) {
if (localId < stride) {
localSums[localId] += localSums[localId + stride];
}
context.localBarrier();
}

if (localId == 0 && slotActive) {
running += routingWeights.get(slot) * localSums[0];
}
}

if (localId == 0) {
float outputValue = localSums[0];
residual.set(rowId, residual.get(rowId) + routingWeight * outputValue);
if (localId == 0 && active) {
residual.set(rowId, running);
}
}

Expand Down
39 changes: 20 additions & 19 deletions src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen2MoEQ8_0FFNLayers.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,9 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
WorkerGrid routerWorker = workerForRows(config.numberOfExperts());
WorkerGrid topKWorker = new WorkerGrid1D(LOCAL_WORK_GROUP_SIZE_ALLOC);
topKWorker.setLocalWork(LOCAL_WORK_GROUP_SIZE_ALLOC, 1, 1);
WorkerGrid expertHiddenWorker = workerForRows(config.moeHiddenDim());
// Fused routed-expert launch: the slot index is folded into the work-group id.
WorkerGrid allExpertsHiddenWorker =
workerForRows(config.moeHiddenDim() * config.numberOfExpertsUsed());
WorkerGrid sharedHiddenWorker = workerForRows(config.sharedExpertHiddenDim());

for (int layer = 0; layer < config.numberOfLayers(); layer++) {
Expand All @@ -79,10 +81,8 @@ public GridScheduler updateGridScheduler(GridScheduler scheduler) {
scheduler.addWorkerGrid(prefix + "ffn_rms_apply", dimElementWorker);
scheduler.addWorkerGrid(prefix + "router_projection", routerWorker);
scheduler.addWorkerGrid(prefix + "router_softmax_topk", topKWorker);
for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) {
scheduler.addWorkerGrid(prefix + "routed_expert_gate_up_" + slot, expertHiddenWorker);
scheduler.addWorkerGrid(prefix + "routed_expert_down_" + slot, dimWorker);
}
scheduler.addWorkerGrid(prefix + "routed_experts_gate_up", allExpertsHiddenWorker);
scheduler.addWorkerGrid(prefix + "routed_experts_down", dimWorker);
scheduler.addWorkerGrid(prefix + "shared_expert_gate_up", sharedHiddenWorker);
scheduler.addWorkerGrid(prefix + "shared_expert_down", dimWorker);
scheduler.addWorkerGrid(prefix + "shared_expert_gate_and_accumulate", topKWorker);
Expand Down Expand Up @@ -200,21 +200,22 @@ private void configureRoutedExperts(TaskGraph layer, int layerIndex) {
context, moeState.wrapRouterLogits, moeState.wrapSelectedExperts,
moeState.wrapRoutingWeights, config.numberOfExperts(), config.numberOfExpertsUsed());

for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) {
layer.task("routed_expert_gate_up_" + slot,
Qwen2MoEKernels::fusedRoutedExpertGateUpSwiGLUQ8_0,
context, moeState.wrapXb, moeState.wrapSelectedExperts, slot,
weights.gateExpertsLayered[layerIndex].asByteArray(),
weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate,
config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);
// All routed slots in two launches instead of two per slot: at top-4 this is 2 kernel
// launches per layer rather than 8, and the residual is accumulated once instead of
// four times.
layer.task("routed_experts_gate_up",
Qwen2MoEKernels::fusedRoutedExpertsGateUpSwiGLUQ8_0,
context, moeState.wrapXb, moeState.wrapSelectedExperts, config.numberOfExpertsUsed(),
weights.gateExpertsLayered[layerIndex].asByteArray(),
weights.upExpertsLayered[layerIndex].asByteArray(), moeState.wrapExpertGate,
config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);

layer.task("routed_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);
}
layer.task("routed_experts_down",
Qwen2MoEKernels::routedExpertsDownProjectAndAccumulateQ8_0,
context, moeState.wrapExpertGate, moeState.wrapX,
moeState.wrapSelectedExperts, moeState.wrapRoutingWeights, config.numberOfExpertsUsed(),
weights.downExpertsLayered[layerIndex].asByteArray(),
config.dim(), config.moeHiddenDim(), config.numberOfExperts(), LOCAL_WORK_GROUP_SIZE_ALLOC);

// The shared expert always runs; it does not depend on router top-K selection.
layer.task("shared_expert_gate_up", Qwen2MoEKernels::sharedExpertGateUpSwiGLUQ8_0,
Expand Down
Loading