[moe] Fuse per-slot routed-expert launches: 8 kernel launches per layer to 2 (+15.4%)

[mikepapadim]

Stacked on top of moe-single-token (beehive-lab#144) — exactly one commit on top of b227bc2, so the diff here is only the optimisation.

Opening it as offered in the discussion on #144. Merge it, rewrite it, or close it — no attribution needed. It is your subsystem; this is just the prototype with the numbers attached.

What changes

The layer graph currently emits two tasks per routed slot:

for (int slot = 0; slot < config.numberOfExpertsUsed(); slot++) {
    layer.task("routed_expert_gate_up_" + slot, ...);
    layer.task("routed_expert_down_" + slot,    ...);
}

At expert_used_count = 4 that is 8 kernel launches per layer — 192 per token across 24 layers. This collapses them to two:

• fusedRoutedExpertsGateUpSwiGLUQ8_0All folds the slot index into the work-group id (slot = groupIdx / moeHiddenDim), so one launch covers every slot. Qwen2MoEState.wrapExpertGate grows from moeHiddenDim to moeHiddenDim * numberOfExpertsUsed so the slots write disjoint windows — 22 KB at this configuration.
• routedExpertsDownProjectAndAccumulateQ8_0All loops the slots inside one launch, with lane 0 carrying a running residual accumulator. Besides the launch saving, the residual is read and written once per row instead of once per slot.

8 launches per layer → 2. The per-slot kernels are left in place, untouched.

Numerics are preserved deliberately

My first version summed all slots into one partial before the reduction. That reorders the arithmetic, and the generated text diverged from baseline after ~40 tokens.

This version reduces each slot with the same tree and folds it into the residual in the same order as the per-slot kernels, so the result should be arithmetically identical rather than merely close. Lane 0 seeds its accumulator from residual.get(rowId) and writes back once, which reproduces the per-slot read-modify-write sequence exactly.

Numbers

RTX 4090, JDK 21.0.2, TornadoVM develop @ba28f1526 (CUDA backend), Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf, 256 tokens, --temperature 0 --seed 42. Medians of 5 runs, with the two jars interleaved run-by-run so GPU clock drift hits both arms equally:

	moe-single-token	+ this PR	gain
default	61.98 tok/s	71.54	+15.4 %
--cuda-graphs	85.52 tok/s	92.00	+7.6 %

Spreads: baseline 61.73–62.31 and 85.48–85.59; this PR 71.42–71.70 and 91.70–92.22.

Combined with --cuda-graphs (off by default, -Dllama.cudaGraphs): 61.98 → 92.00 tok/s, +48.4 % over today's default.

The gain is smaller with graphs enabled, as expected — graphs already amortise launch cost, so what remains is the single residual accumulation instead of four.

Correctness

At 24 tokens the output is identical between baseline and this PR, and stable across 3 runs of each.

Longer generations cannot be compared this way: the existing MoE path is non-deterministic beyond roughly 40 tokens — three runs of the unmodified moe-single-token branch, same jar, --temperature 0 --seed 42, 200 tokens, produce three different outputs, with and without --cuda-graphs. That is pre-existing and not introduced here, but it is worth fixing before this lands, since it makes any subsequent A/B harder to trust. Details in the #144 thread.

Reproducing

export TORNADOVM_HOME=/path/to/tornadovm-sdk   # make BACKEND=cuda
export JAVA_HOME=/path/to/jdk21
export LLAMA_ROOT=$PWD

./llama-tornado --gpu --model Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf \
  --prompt "Explain briefly how mixture of experts routing works." \
  --max-tokens 256 --temperature 0 --seed 42 --cuda-graphs

Note --max-tokens 1024 does not fit on a 24 GB card with this model: TornadoOutOfMemoryException: Unable to allocate 330612752 bytes.