Add Qwen2-MoE batch prefill and single-token decode

src/main/java/org/beehive/gpullama3/inference/InferenceCoreBatchPrefillDecode.java

@@ -1,6 +1,7 @@
 package org.beehive.gpullama3.inference;
 
 import org.beehive.gpullama3.auxiliary.Parallel;
+import org.beehive.gpullama3.inference.state.Qwen2MoEState;
 import org.beehive.gpullama3.inference.state.State;
 import org.beehive.gpullama3.inference.weights.standard.StandardWeights;
 import org.beehive.gpullama3.inference.weights.tornado.TornadoWeights;
@@ -177,7 +178,7 @@ public static void batchForwardJavaPrefill(Model model, State state, int[] token
      * then delegates graph execution to the plan.</p>
      *
      * @param model
-     *     the LLaMA model
+     *     the model
      * @param state
      *     mutable inference state
      * @param tokens
@@ -194,6 +195,9 @@ public static void batchForwardTornadoVMPrefill(Model model, State state, int[]
         final TornadoWeights weights = (TornadoWeights) model.weights();
 
         state.batchStartPosHolder.set(0, startPos);
+        if (state instanceof Qwen2MoEState moeState && moeState.activeBatchSizeHolder != null) {
+            moeState.activeBatchSizeHolder.set(0, chunkSize);
+        }
 
         switch (weights.getWeightType()) {
             case F16 -> {
@@ -230,7 +234,7 @@ public static void batchForwardTornadoVMPrefill(Model model, State state, int[]
      * graph execution to the plan.</p>
      *
      * @param model
-     *     the LLaMA model
+     *     the model
      * @param state
      *     mutable inference state
      * @param token

src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java

@@ -166,6 +166,7 @@ public static List<Integer> generateTokensQwen3(Model model, State state, int st
 
         // Storage for generated tokens
         List<Integer> generatedTokens = new ArrayList<>();
+        int generatedTokenBudget = Math.max(0, maxTokens - startPosition - promptTokens.size());
 
         // Initialize token variables
         int currentToken = state.latestToken; // BOS?
@@ -188,8 +189,11 @@ public static List<Integer> generateTokensQwen3(Model model, State state, int st
                 if (echo) {
                     System.err.print(Tokenizer.replaceControlCharacters(model.tokenizer().decode(List.of(nextToken))));
                 }
-                // We have reached the last prompt token and computed the first response-token.
-                position++; // The current logit belongs to the next position
+                // The last prompt token produced the first response-token logits.
+                // The for-loop advances to the next sequence position.
+                if (generatedTokenBudget == 0) {
+                    break;
+                }
             } else {
                 // Mark the start of actual generation (after prompt processing)
                 if (inferenceStartNanos == 0) {
@@ -216,7 +220,7 @@ public static List<Integer> generateTokensQwen3(Model model, State state, int st
             }
 
             // Check for stop condition
-            if (stopTokens.contains(nextToken)) {
+            if (generatedTokens.size() >= generatedTokenBudget || stopTokens.contains(nextToken)) {
                 break;
             }
 
@@ -393,6 +397,7 @@ public static List<Integer> generateTokensGPUQwen3(Model model, State state, int
         // prompt is longer than the token budget (actualMaxTokens), the difference is
         // negative and would throw IllegalArgumentException("Illegal Capacity").
         List<Integer> generatedTokens = new ArrayList<>(Math.max(0, Math.min(256, actualMaxTokens - promptTokens.size()))); // Conservative estimate
+        int generatedTokenBudget = Math.max(0, actualMaxTokens - startPosition - promptTokens.size());
 
         // Initialize token variables
         int currentToken = state.latestToken; // BOS?
@@ -428,8 +433,11 @@ public static List<Integer> generateTokensGPUQwen3(Model model, State state, int
                 if (echo) {
                     System.err.print(Tokenizer.replaceControlCharacters(model.tokenizer().decode(List.of(nextToken))));
                 }
-                // We have reached the last prompt token and computed the first response-token.
-                position++; // The current logit belongs to the next position
+                // The last prompt token produced the first response-token logits.
+                // The for-loop advances to the next sequence position.
+                if (generatedTokenBudget == 0) {
+                    break;
+                }
             } else {
                 // Mark the start of actual generation (after prompt processing)
                 if (inferenceStartNanos == 0) {
@@ -456,7 +464,7 @@ public static List<Integer> generateTokensGPUQwen3(Model model, State state, int
             }
 
             // Check for stop condition
-            if (stopTokens.contains(nextToken)) {
+            if (generatedTokens.size() >= generatedTokenBudget || stopTokens.contains(nextToken)) {
                 break;
             }
 
@@ -678,4 +686,4 @@ public static List<Integer> generateTokensGPUGranite(Model model, State state, i
 
         return generatedTokens;
     }
-}
+}

src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java

@@ -147,9 +147,7 @@ public static List<Integer> generateTokensLlama(Model model,
     // @formatter:on
 
     /**
-     * LLaMA batched GPU prefill token generation (GPU, Phase 4).
-     *
-     * <p>FP16 only; Q8_0 throws {@link UnsupportedOperationException}.</p>
+     * Batched GPU prefill followed by single-token decode.
      *
      * <p>Split loop:</p>
      * <ul>
@@ -187,17 +185,25 @@ public static List<Integer> generateTokensGPULlama(Model model,
         int N = promptTokens.size();
 
         // ── Prefill ───────────────────────────────────────────────────────────
-        // Build the token sequence at positions [startPosition .. startPosition+N-1]:
-        //   position startPosition+0 : currentToken (BOS/previous token)
-        //   position startPosition+k : promptTokens[k-1]
-        int[] prefillSeq = new int[N];
-        prefillSeq[0] = currentToken;
-        for (int i = 1; i < N; i++) {
-            prefillSeq[i] = promptTokens.get(i - 1);
+        // Qwen's regular path forwards promptTokens[0] directly at position 0.
+        // Keep the final prompt token for the B1 decode graph, which produces the
+        // first generation logits without duplicating the ChatML start token.
+        boolean qwen2MoE = model.getModelType() == org.beehive.gpullama3.model.ModelType.QWEN_2_MOE;
+        int prefillTokenCount = qwen2MoE ? Math.max(0, N - 1) : N;
+        int[] prefillSeq = new int[prefillTokenCount];
+        if (qwen2MoE) {
+            for (int i = 0; i < prefillTokenCount; i++) {
+                prefillSeq[i] = promptTokens.get(i);
+            }
+        } else {
+            prefillSeq[0] = currentToken;
+            for (int i = 1; i < N; i++) {
+                prefillSeq[i] = promptTokens.get(i - 1);
+            }
         }
 
-        for (int chunkStart = 0; chunkStart < N && pos + chunkStart < actualMaxTokens; chunkStart += batchSize) {
-            int chunkEnd = Math.min(Math.min(chunkStart + batchSize, N), actualMaxTokens - pos);
+        for (int chunkStart = 0; chunkStart < prefillTokenCount && pos + chunkStart < actualMaxTokens; chunkStart += batchSize) {
+            int chunkEnd = Math.min(Math.min(chunkStart + batchSize, prefillTokenCount), actualMaxTokens - pos);
             int chunkSize = chunkEnd - chunkStart;
             int[] chunk = Arrays.copyOfRange(prefillSeq, chunkStart, chunkEnd);
 
@@ -213,12 +219,13 @@ public static List<Integer> generateTokensGPULlama(Model model,
         }
 
         currentToken = promptTokens.get(N - 1);
-        pos = startPosition + N;
+        pos = startPosition + (qwen2MoE ? N - 1 : N);
         state.latestToken = currentToken;
         long decodeStartNanos = System.nanoTime();
+        int generatedTokenBudget = Math.max(0, actualMaxTokens - startPosition - N);
 
         // ── Decode ────────────────────────────────────────────────────────────
-        while (pos < actualMaxTokens) {
+        while (pos < actualMaxTokens && generatedTokens.size() < generatedTokenBudget) {
             var logits = InferenceCoreBatchPrefillDecode.forwardTornadoVMDecode(model, state, currentToken, pos, plan);
             int nextToken = sampler.sampleToken(logits);
 

src/main/java/org/beehive/gpullama3/inference/state/Qwen2MoEState.java

@@ -40,6 +40,19 @@ public class Qwen2MoEState extends Qwen2State {
     public final FloatArray wrapSharedGate;
     public final FloatArray wrapSharedOutput;
 
+    // TornadoVM buffers for the batch-prefill MoE path.
+    // Their shapes use the configured maximum batch size so TaskGraphs stay fixed.
+    public final FloatArray wrapRouterLogitsBatch;
+    public final IntArray activeBatchSizeHolder;
+    public final IntArray wrapSelectedExpertsBatch;
+    public final FloatArray wrapRoutingWeightsBatch;
+    public final IntArray wrapGroupedAssignmentIds;
+    public final IntArray wrapGroupedPositionByAssignment;
+    public final FloatArray wrapGroupedExpertHidden;
+    public final FloatArray wrapGroupedExpertDown;
+    public final FloatArray wrapSharedHiddenBatch;
+    public final FloatArray wrapSharedWeightBatch;
+
     public Qwen2MoEState(Configuration config, int batchsize) {
         super(config, batchsize);
         Qwen2MoEConfiguration c = (Qwen2MoEConfiguration) config;
@@ -56,6 +69,33 @@ public Qwen2MoEState(Configuration config, int batchsize) {
         this.wrapExpertGate = new FloatArray(c.moeHiddenDim() * c.numberOfExpertsUsed());
         this.wrapSharedGate = new FloatArray(c.sharedExpertHiddenDim());
         this.wrapSharedOutput = new FloatArray(c.dim());
+
+        int gpuBatchSize = Integer.getInteger("llama.prefillBatchSize", 1);
+        if (gpuBatchSize > 1) {
+            int assignments = gpuBatchSize * c.numberOfExpertsUsed();
+            this.wrapRouterLogitsBatch = new FloatArray(gpuBatchSize * c.numberOfExperts());
+            this.activeBatchSizeHolder = new IntArray(1);
+            this.activeBatchSizeHolder.init(gpuBatchSize);
+            this.wrapSelectedExpertsBatch = new IntArray(assignments);
+            this.wrapRoutingWeightsBatch = new FloatArray(assignments);
+            this.wrapGroupedAssignmentIds = new IntArray(assignments);
+            this.wrapGroupedPositionByAssignment = new IntArray(assignments);
+            this.wrapGroupedExpertHidden = new FloatArray(assignments * c.moeHiddenDim());
+            this.wrapGroupedExpertDown = new FloatArray(assignments * c.dim());
+            this.wrapSharedHiddenBatch = new FloatArray(gpuBatchSize * c.sharedExpertHiddenDim());
+            this.wrapSharedWeightBatch = new FloatArray(gpuBatchSize);
+        } else {
+            this.wrapRouterLogitsBatch = null;
+            this.activeBatchSizeHolder = null;
+            this.wrapSelectedExpertsBatch = null;
+            this.wrapRoutingWeightsBatch = null;
+            this.wrapGroupedAssignmentIds = null;
+            this.wrapGroupedPositionByAssignment = null;
+            this.wrapGroupedExpertHidden = null;
+            this.wrapGroupedExpertDown = null;
+            this.wrapSharedHiddenBatch = null;
+            this.wrapSharedWeightBatch = null;
+        }
     }
 
     @Override

src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2MoE.java

@@ -2,6 +2,7 @@
 
 import org.beehive.gpullama3.inference.InferenceCore;
 import org.beehive.gpullama3.inference.InferenceEngine;
+import org.beehive.gpullama3.inference.InferenceEngineWithBatchPrefillDecode;
 import org.beehive.gpullama3.inference.sampler.Sampler;
 import org.beehive.gpullama3.inference.state.Qwen2MoEState;
 import org.beehive.gpullama3.inference.state.State;
@@ -90,7 +91,9 @@ public List<Integer> generateTokens(State state, int startPosition, List<Integer
     public List<Integer> generateTokensGPU(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo,
             IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) {
         if (WITH_PREFILL_DECODE && TornadoVMMasterPlan.PREFILL_BATCH_SIZE > 1) {
-            throw new UnsupportedOperationException("Batch prefill/decode on GPU not yet implemented for Qwen2-MoE");
+            return InferenceEngineWithBatchPrefillDecode.generateTokensGPULlama(
+                    this, state, startPosition, promptTokens, stopTokens, maxTokens,
+                    sampler, echo, onTokenGenerated, tornadoVMPlan);
         }
         if (WITH_PREFILL_DECODE) {
             throw new UnsupportedOperationException("Prefill/decode on GPU not yet implemented for Qwen2-MoE");