Add Qwen2-MoE batch prefill and single-token decode

[MRPRESIDENT66]

Summary

This PR adds batch prefill with B1 decode for Qwen2-MoE Q8_0 models.

The prompt is processed in batches, while autoregressive generation continues one token at a time. Both phases reuse the same GPU model weights and KV cache.

Main changes:

• Added batched Attention and MoE kernels.
• Added GPU-side token-expert grouping.
• Added a Qwen2-MoE batch-prefill TaskGraph.
• Connected B16 prefill to B1 decode.
• Correctly handled the final incomplete prefill batch.
• Fixed the Qwen decode position and generated-token budget.

Validation

The Single-token and B16 paths generated identical output with the same prompt, temperature, and seed.

Each run processed 64 tokens:

18 prompt tokens + 46 generated tokens

Performance

Tested on an RTX 5070 Ti with JDK 21, TornadoVM 5.2.1-dev, the CUDA backend, and CUDA Graphs.

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

The Single-token baseline used the same parameters without:

--with-prefill-decode --batch-prefill-size 16

One warm-up run was excluded, followed by five interleaved measured runs.

Mode	Run 1	Run 2	Run 3	Run 4	Run 5	Mean
Single-token baseline	117.19	118.44	116.83	116.98	117.81	117.45 tok/s
B16 prefill + B1 decode	128.20	128.70	128.35	128.25	128.56	128.41 tok/s

This is a 9.33% improvement.

throughput was calculated as:

(prompt tokens + generated tokens) / timed inference duration

Phase	Tokens	Mean throughput
B16 prefill	18 prompt tokens	246.26 prompt tok/s
B1 decode	46 generated tokens	108.16 generated tok/s

Profiling

The B16 prefill and B1 decode path was profiled on an RTX 5070 Ti using JFR and Nsight Systems.

JFR showed no major Java-side bottleneck. GC pauses accounted for approximately 0.25% of the 21-second recording.

For one full B16 prefill chunk across 24 layers:

Component	Time	GPU kernel share
Attention	9.375 ms	20.26%
FFN RMSNorm	0.094 ms	0.20%
Router and Top-K	0.291 ms	0.63%
Assignment grouping	2.281 ms	4.93%
Routed experts	19.337 ms	41.78%
Shared expert	14.905 ms	32.20%
Total	46.284 ms	100%

Within the routed-expert computation, Gate/Up took 10.949 ms, Down took 8.352 ms, and accumulation took 0.036 ms. Routed and shared experts together accounted for 73.98% of GPU kernel time.

Task gaps inside each layer were approximately 0.18–0.19 μs. Gaps between layer TaskGraphs totalled approximately 2.23 ms per B16 chunk.

[mikepapadim]

thanks @MRPRESIDENT66