Add Qwen1.5-MoE Q8_0 CPU and single-token TornadoVM GPU inference - #144
Conversation
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.
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.
|
@MRPRESIDENT66 thank you for your contribution. I wander if you tried the cuda backend instead of ptx? Can you share prompts you used? |
|
@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. |
|
@MRPRESIDENT66 Confirmed on my side too — CUDA backend, RTX 4090, JDK 21, TornadoVM One thing worth adding to your numbers: turn on CUDA graphs. It is off by default ( 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-graphsOn this model, medians of 5 runs at 256 tokens:
+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 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. |
|
@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: What it doesToday 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
8 launches per layer → 2. NumbersRTX 4090, JDK 21, TornadoVM
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 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 mineMy 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, Three different outputs. It reproduces with and without 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
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. |
|
@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 The change in one lineTwo 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 Backing numbersRTX 4090, JDK 21, TornadoVM
Spreads: baseline 61.73–62.31 and 85.48–85.59; fused 71.42–71.70 and 91.70–92.22. Combined with On the batched-prefill PR you mentionedWorth 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,
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:
Happy to open the fusion PR myself if you would rather not carry it, just say which you prefer. |
|
Opened it as offered: MRPRESIDENT66#1, targeting your 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 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 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. |
|
|
|
@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. |
|
@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. |
3 participants
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:
Key implementation notes
Qwen3ChatFormatnow 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.Q8_0FloatTensorpath 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 packagepassed.temperature=0, andseed=42.Qwen1.5-MoE-A2.7B-Chat.Q8_0.ggufCUDA backend
The CUDA backend was also validated on an RTX 4090 using TornadoVM 5.2.1 development build. For CUDA inference,
--cuda-graphsis 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-graphsResults after fixing the RMSNorm cross-workgroup reduction race:
--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.