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

[MRPRESIDENT66]

This PR adds CPU and single-token TornadoVM GPU inference support for Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf.

Although the public model name is Qwen1.5-MoE, its GGUF architecture identifier is qwen2moe, which is used internally for model detection and class naming.

The implementation adds:

• Qwen2-MoE model detection, configuration, loader, state, and Q8_0 weight support.
• CPU MoE forward inference: router projection, softmax/Top-K routing, routed experts, shared expert, and residual accumulation.
• Single-token TornadoVM GPU inference for Q8_0 weights, including TaskGraphs and WorkerGrids for the router, routed experts, and shared expert.

Key implementation notes

• Qwen3ChatFormat now adds the standard newline after <|im_end|>. This affects Qwen-family models that use this shared ChatML formatter, but does not change Llama, Mistral, Granite, or Phi prompt formatting.
• The CPU Q8_0FloatTensor path now quantizes each FP32 activation block to int8 before the dot product, then performs an int8 × int8 accumulation with the Q8_0 scales. This follows the llama.cpp/ggml Q8_0 matmul approach and affects CPU Q8_0 inference generally, not only Qwen1.5-MoE. I added a switch, so the previous path remains available with -Dllama.quantizeActivation=false. Activation quantization is enabled by default.

Validation

Validated on an RTX 4090 (24 GB), JDK 21, and TornadoVM 5.1 with the PTX backend.

• ./mvnw -q -DskipTests package passed.
• CPU and GPU single-token smoke tests produced the same deterministic output prefix with the same prompt, temperature=0, and seed=42.
• Model: Qwen1.5-MoE-A2.7B-Chat.Q8_0.gguf
• 32 generated tokens:
  • CPU: 5.08 tok/s
  • GPU PTX: 24.59 tok/s average over three runs (24.89, 23.71, 25.17 tok/s)

CUDA backend

The CUDA backend was also validated on an RTX 4090 using TornadoVM 5.2.1 development build. For CUDA inference, --cuda-graphs is recommended because it significantly reduces the kernel-launch overhead of the MoE execution path.

The following prompt and settings were used:

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

Results after fixing the RMSNorm cross-workgroup reduction race:

• Default CUDA execution: 66.24 tok/s average over three runs (66.19, 66.00, 66.54 tok/s).
• With --cuda-graphs: 94.16 tok/s average over two runs (94.33, 93.98 tok/s).

All five runs produced the same token-ID hash.

Scope

This PR supports CPU and single-token GPU inference only. Batch prefill/decode support is intentionally out of scope and will be submitted separately.

[mikepapadim]

@MRPRESIDENT66 thank you for your contribution. I wander if you tried the cuda backend instead of ptx?
I am keen on merging it, but i think running it with the cuda backend will be faster.

Can you share prompts you used?

[MRPRESIDENT66]

@mikepapadim Thanks! I tested the CUDA backend too, and it works fine. It achieved 88.64 tok/s on the RTX 4090.

The prompt I used was: Explain briefly how mixture of experts routing works.

[mikepapadim]

@MRPRESIDENT66 Confirmed on my side too — CUDA backend, RTX 4090, JDK 21, TornadoVM develop @ba28f1526, this branch at b227bc2. With your prompt at 64 tokens I get 84.40 tok/s, right next to your 88.64. Nice work getting a MoE architecture running end to end; the routed/shared expert split in Qwen2MoEQ8_0FFNLayers reads cleanly, and having the CPU and GPU paths agree deterministically is exactly the right thing to have validated first.

One thing worth adding to your numbers: turn on CUDA graphs. It is off by default (llama.cudaGraphs=false) and needs the --cuda-graphs flag.

export JAVA_HOME=/path/to/jdk21
export TORNADOVM_HOME=/path/to/tornadovm-sdk    # built with: make BACKEND=cuda
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

On this model, medians of 5 runs at 256 tokens:

	tok/s
default	62.01
--cuda-graphs	85.42

+37.8%, and the generated text is byte-identical (checked at temperature 0, seed 42).

The reason it pays off so well here is the launch count. Your FFN layer builds 15 fixed tasks plus 2 per routed-expert slot, and the GGUF says expert_used_count = 4, block_count = 24 — so roughly 560 kernel launches per token. Without graphs each one is dispatched individually; with graphs the whole per-layer graph is captured once and replayed. Dense models see maybe +28% from the same flag, MoE more, precisely because of that launch count.

So the honest headline for this PR is closer to ~85 tok/s on CUDA rather than 62, with no code change. Might be worth quoting that in the description.

Full details, plus a negative result and why FP16 does not fit on a 24 GB card, in my other comment above.

[mikepapadim]

@MRPRESIDENT66 I prototyped suggestion 2 from my earlier comment — collapsing the per-slot routed-expert launches — and it is worth more than I expected. Branch: mikepapadim/GPULlama3.javaF:opt/moe-fused-expert-launches (3940767, +182/−20 over this PR). Take it or leave it, it is yours if useful.

What it does

Today the layer graph 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. The change turns them into two:

• fusedRoutedExpertsGateUpSwiGLUQ8_0All — folds the slot index into the work-group id (slot = groupIdx / moeHiddenDim), so one launch covers all slots. wrapExpertGate grows from moeHiddenDim to moeHiddenDim * numberOfExpertsUsed so slots write disjoint windows (22 KB total, nothing).
• routedExpertsDownProjectAndAccumulateQ8_0All — loops the slots inside one launch, and folds each into the residual through lane 0's running accumulator, so the residual is read and written once per row instead of four times.

8 launches per layer → 2.

Numbers

RTX 4090, JDK 21, TornadoVM develop @ba28f1526, 256 tokens, medians of 5 runs, the two jars interleaved run-by-run so clock drift hits both arms equally:

	baseline (this PR)	fused	gain
default	61.98 tok/s	71.54	+15.4 %
--cuda-graphs	85.52 tok/s	92.00	+7.6 %

Variance was tight — baseline 61.73–62.31 and 85.48–85.59, fused 71.42–71.70 and 91.70–92.22.

Stacked with the --cuda-graphs flag from my earlier comment: 61.98 → 92.00 tok/s, +48.4 % over the current default.

The gain is smaller with graphs on, exactly as you would expect — graphs already amortise launch cost, so what is left is the single residual accumulation instead of four.

On correctness — and a bug that is not mine

My first version accumulated all slots into one partial sum before the reduction. That changes the summation order, and the text diverged from baseline after ~40 tokens. I rewrote the down-projection so each slot is reduced with the same tree and folded in the same order as the per-slot kernels, which should be arithmetically identical.

Checking that turned up something more important. The MoE GPU path is not deterministic at longer generations. Three runs of the unmodified PR branch, same jar, --temperature 0 --seed 42, 200 tokens:

b6c032c4a9619683060ee14168bdbc83  run 1
d1e4e4682ea7ef266292fa35ba8369c3  run 2
61645db0413d4d27021953c385f6b4bf  run 3

Three different outputs. It reproduces with and without --cuda-graphs, so graphs are not the cause, and it is present without my patch — this is on the existing path, not something the fusion introduced.

Short generations are stable: at 24 tokens, baseline is identical across 3 runs, fused is identical across 3 runs, and baseline == fused. So the fusion is sound on the deterministic prefix, and your PR's "same deterministic output prefix" validation holds — it is the longer tail that drifts.

Worth chasing before this lands, since it undercuts CPU/GPU comparison at realistic lengths. My first suspect would be a buffer that is read before the producing task's write is visible, or a reduction whose work-group count does not match its assumption; but I have not localised it.

One more constraint

--max-tokens 1024 does not run on a 4090:

TornadoOutOfMemoryException: Unable to allocate 330612752 bytes of memory.
  at TornadoVMMasterPlanSingleToken.forceCopyInReadOnlyData(...:116)

15.23 GB of Q8_0 weights leaves little headroom for the KV cache. Fine at 256, worth knowing as a documented limit.

Happy to open this as a PR against your branch, or leave it here for you to cherry-pick.

[mikepapadim]

@MRPRESIDENT66 To make the offer concrete: please take the routed-expert fusion as your own separate PR if you like it. It is your subsystem, and it will review better as a focused follow-up than bolted onto this one. Branch is mikepapadim/GPULlama3.javaF:opt/moe-fused-expert-launches (3940767, 3 files, +182/−20) — cherry-pick it, rewrite it, or ignore it, no attribution needed.

The change in one line

Two tasks per routed slot become two tasks total: the slot index is folded into the work-group id for gate/up, and the down-projection loops slots internally with lane 0 carrying a running residual accumulator. At expert_used_count = 4 that is 8 kernel launches per layer → 2, i.e. 192 → 48 per token across 24 layers.

Backing numbers

RTX 4090, JDK 21, TornadoVM develop @ba28f1526, 256 tokens, medians of 5 runs, the two jars interleaved run-by-run so GPU clock drift hits both arms equally:

	this PR	+ fusion	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; fused 71.42–71.70 and 91.70–92.22.

Combined with --cuda-graphs: 61.98 → 92.00 tok/s, +48.4 % over today's default, no change to model quality — identical 24-token prefix, stable across 3 runs on both jars.

On the batched-prefill PR you mentioned

Worth doing for the API surface, but I would not expect a large decode-side win on this model, and I would rather tell you that now than have you build it expecting one. A rough probe, single-token path, --max-tokens 1, warm:

prompt	wall
280 words	7.34 s
1680 words	8.32 s

About 1 s for ~1400 extra words, i.e. roughly 0.5 ms per prompt token, against ~11–16 ms per token during decode. So prefill on this path is already far cheaper per token than decode — nothing like the ratio you see on dense models where batched prefill is the headline. Treat that as a rough probe rather than a measurement: model load dominates the wall clock, and I did not instrument prefill separately.

So my suggested ordering, for whatever it is worth:

1. This PR, plus a note about --cuda-graphs — biggest win, zero code.
2. The routed-expert fusion as its own PR — +15 % default, +7.6 % with graphs.
3. The non-determinism at longer generations (three different outputs from three identical runs, with and without graphs, on the unmodified branch) — this one I would fix before the fusion, since it makes every subsequent A/B harder to trust.
4. PREFILL_DECODE for QWEN_2_MOE — for API completeness and prompt-heavy workloads, with the caveat above about expected gain.

Happy to open the fusion PR myself if you would rather not carry it, just say which you prefer.

[mikepapadim]

Opened it as offered: MRPRESIDENT66#1, targeting your moe-single-token branch so it stacks on this PR rather than duplicating it — one commit on top of b227bc2, so the diff there is only the optimisation.

Merge, rewrite, or close it as you see fit; no attribution needed. If you would rather carry the change yourself, close it and cherry-pick 3940767 — that is equally good from my side.

Recap of the numbers, all interleaved run-by-run against your branch: 61.98 → 71.54 tok/s (+15.4 %) by default, 85.52 → 92.00 (+7.6 %) with --cuda-graphs, identical 24-token output.

The one thing I would sequence ahead of it is the non-determinism at longer generations, since it makes every later comparison harder to trust.

[CLAassistant]

All committers have signed the CLA.

[MRPRESIDENT66]

@mikepapadim Thanks for your work on this! I found the cause of the inconsistent output in longer generations. The RMSNorm reduction used multiple work-groups to calculate the sum of squares, so work-group 0 could read partial results before the other groups had finished. I have now changed it to use a single work-group and pushed the fix.

[MRPRESIDENT66]

@mikepapadim I’ll keep the routed-expert fusion separate. Once #144 is merged, I’ll create a new branch from the updated main, cherry-pick your commit, test it, and submit it as a follow-up PR.