On-device greedy sampling (-Dllama.deviceSample): GPU argmax, keep logits device-side (~500x less D2H)

.gitignore

@@ -67,3 +67,4 @@ logs/target/
 .project
 .classpath
 .vscode/
+external/

src/main/java/org/beehive/gpullama3/LlamaApp.java

@@ -14,6 +14,29 @@ public class LlamaApp {
     public static final boolean USE_VECTOR_API = Boolean.parseBoolean(System.getProperty("llama.VectorAPI", "true"));   // Enable Java Vector API for CPU acceleration
     public static final boolean SHOW_PERF_INTERACTIVE = Boolean.parseBoolean(System.getProperty("llama.ShowPerfInteractive", "true")); // Show performance metrics in interactive mode
 
+    /**
+     * On-device greedy sampling ({@code -Dllama.deviceSample=true}) keeps the logits on the GPU
+     * and returns only the argmax token id. It is only valid on the GPU FP16 greedy path for the
+     * models whose decode loop reads {@code state.sampledToken} (Llama / Mistral / Qwen3). For any
+     * other configuration the host still needs the full logits row, so clear the flag here —
+     * before the logits task graph (which reads it as a static-final) is first built.
+     */
+    private static void guardDeviceSample(Model model, Options options) {
+        if (!Boolean.getBoolean("llama.deviceSample")) {
+            return;
+        }
+        boolean greedy = options.temperature() == 0.0f;
+        boolean fp16 = "FP16".equals(model.configuration().quantization());
+        var mt = model.getModelType();
+        boolean wiredLoop = mt == org.beehive.gpullama3.model.ModelType.LLAMA_3
+                || mt == org.beehive.gpullama3.model.ModelType.MISTRAL
+                || mt == org.beehive.gpullama3.model.ModelType.QWEN_3;
+        if (!(options.useTornadovm() && greedy && fp16 && wiredLoop)) {
+            System.err.println("[deviceSample] ignored — requires GPU + greedy (temperature 0) + FP16 + Llama/Mistral/Qwen3");
+            System.clearProperty("llama.deviceSample");
+        }
+    }
+
     private static void runSingleInstruction(Model model, Sampler sampler, Options options) {
         String response = model.runInstructOnce(sampler, options);
         System.out.println(response);
@@ -36,6 +59,7 @@ private static void runSingleInstruction(Model model, Sampler sampler, Options o
     static void main(String[] args) throws IOException {
         Options options = Options.parseOptions(args);
         Model model = loadModel(options);
+        guardDeviceSample(model, options);
         Sampler sampler = createSampler(model, options);
 
         if (options.interactive()) {

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

@@ -7,6 +7,7 @@
 import org.beehive.gpullama3.model.Model;
 import org.beehive.gpullama3.tokenizer.Tokenizer;
 import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan;
+import org.beehive.gpullama3.tornadovm.layers.type.fp16.LogitsFP16Layer;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 
 import java.io.ByteArrayOutputStream;
@@ -38,6 +39,19 @@ private InferenceEngine() {
         //prevent instantiation
     }
 
+    /**
+     * Greedy next-token pick for the GPU FP16 path. With {@code -Dllama.deviceSample=true}
+     * the argmax already ran on the device and the token id sits in {@code state.sampledToken};
+     * read it directly (the full logits row never left the GPU). Otherwise fall back to the
+     * host sampler over the transferred logits.
+     */
+    private static int sampleTokenGpu(State state, Sampler sampler, Object logits) {
+        if (LogitsFP16Layer.DEVICE_SAMPLE) {
+            return state.sampledToken.get(0);
+        }
+        return sampler.sampleToken(logits);
+    }
+
     /**
      * LLM generation entry point, ingest prompt tokens and generates new tokens.
      *
@@ -328,7 +342,7 @@ public static List<Integer> generateTokensGPULlama(Model model, State state, int
                 }
 
                 // Sample next token - use GPU sampling if available
-                nextToken = sampler.sampleToken(logits);
+                nextToken = sampleTokenGpu(state, sampler, logits);
 
                 // Add token consumer support
                 if (onTokenGenerated != null) {
@@ -422,7 +436,7 @@ public static List<Integer> generateTokensGPUQwen3(Model model, State state, int
             }
 
             // Sample the next token
-            nextToken = sampler.sampleToken(state.wrapLogits);
+            nextToken = sampleTokenGpu(state, sampler, state.wrapLogits);
 
             // Output the token if echo is enabled
             if (echo) {

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

@@ -59,6 +59,9 @@ public abstract class State {
     public final FloatArray wrapKeyCache;   // FloatArray wrapper for the key cache, optimized for TornadoVM.
     public final FloatArray wrapValueCache; // FloatArray wrapper for the value cache, optimized for TornadoVM.
     public final IntArray positionHolder;
+    // On-device greedy sampling: the GPU argmax kernel writes the sampled token id here
+    // (element 0), so only 1 int crosses to the host instead of the full vocab logits row.
+    public final IntArray sampledToken = new IntArray(1);
 
     public TornadoNativeArray embeddingX;
 

src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java

@@ -8,11 +8,53 @@
 import uk.ac.manchester.tornado.api.types.arrays.ByteArray;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
 
 public class TransformerComputeKernels {
 
+    /**
+     * On-device greedy sampling: single-workgroup argmax over the {@code vocab} logits.
+     * Each thread scans a strided slice tracking its local (max value, index); a local-memory
+     * tree reduction picks the global argmax and lane 0 writes the token id to {@code out[0]}.
+     * Lets the decode step transfer only that one int to the host instead of the whole
+     * vocab-sized logits row (D2H copy + host scan removed). Launch with one workgroup:
+     * {@code global == local == localMemSize}.
+     */
+    public static void argmaxLogits(KernelContext context, FloatArray logits, IntArray out, int vocab, int localMemSize) {
+        int tid = context.localIdx;
+        int localSz = context.localGroupSizeX;
+        float[] vals = context.allocateFloatLocalArray(localMemSize);
+        int[] idxs = context.allocateIntLocalArray(localMemSize);
+
+        float best = Float.NEGATIVE_INFINITY;
+        int bestIdx = 0;
+        for (int i = tid; i < vocab; i += localSz) {
+            float v = logits.get(i);
+            if (v > best) {
+                best = v;
+                bestIdx = i;
+            }
+        }
+        vals[tid] = best;
+        idxs[tid] = bestIdx;
+        context.localBarrier();
+
+        for (int s = localSz / 2; s > 0; s >>= 1) {
+            if (tid < s) {
+                if (vals[tid + s] > vals[tid]) {
+                    vals[tid] = vals[tid + s];
+                    idxs[tid] = idxs[tid + s];
+                }
+            }
+            context.localBarrier();
+        }
+        if (tid == 0) {
+            out.set(0, idxs[0]);
+        }
+    }
+
     /**
      * Default constructor for the TransformerComputeKernels class.
      */

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java

@@ -17,6 +17,17 @@
 
 public class LogitsFP16Layer extends AbstractLogitsTaskGraph {
 
+    /** Local workgroup size for the single-workgroup on-device argmax. */
+    private static final int SAMPLE_LOCAL = 256;
+
+    /**
+     * On-device greedy sampling: append a GPU argmax over the logits and transfer only the
+     * sampled token id (1 int) to the host instead of the full vocab logits row. Opt-in via
+     * {@code -Dllama.deviceSample=true}; only valid for greedy decoding (LlamaApp disables it
+     * for temperature > 0 and non-FP16 models, which still need the full logits host-side).
+     */
+    public static final boolean DEVICE_SAMPLE = Boolean.getBoolean("llama.deviceSample");
+
     public LogitsFP16Layer(String name, State state, Weights weights, Configuration config,
             String lastTaskGraphID, SchedulerType schedulerType) {
         super(name, state, weights, config, lastTaskGraphID, schedulerType);
@@ -92,8 +103,22 @@ protected TaskGraph setupLogitsTaskGraph(TornadoWeights weights, Configuration c
                 config.vocabularySize(),                       // output dimension
                 LOCAL_WORK_GROUP_SIZE_ALLOC * THREAD_SCALE_FOR_LOGITS);
 
-        // === Transfer Results to Host ===
-        logits.transferToHost(DataTransferMode.EVERY_EXECUTION, state.wrapLogits);
+        // === Sampling / result transfer ===
+        if (DEVICE_SAMPLE) {
+            // Greedy argmax on the GPU; only the token id crosses to the host (the full
+            // vocab logits row stays device-side — no big D2H copy, no host scan).
+            logits.transferToDevice(DataTransferMode.FIRST_EXECUTION, state.sampledToken);
+            logits.task("argmax_sample",
+                    TransformerComputeKernels::argmaxLogits,
+                    context,
+                    state.wrapLogits,
+                    state.sampledToken,
+                    config.vocabularySize(),
+                    SAMPLE_LOCAL);
+            logits.transferToHost(DataTransferMode.EVERY_EXECUTION, state.sampledToken);
+        } else {
+            logits.transferToHost(DataTransferMode.EVERY_EXECUTION, state.wrapLogits);
+        }
         configureAdditionalPersists(logits);
         return logits;
     }
@@ -108,6 +133,11 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler)
         tornadoForwardScheduler.addWorkerGrid("logits.rms_reduce", logitsRMS);
         tornadoForwardScheduler.addWorkerGrid("logits.rms_apply_fp16", logitsRMS);
         tornadoForwardScheduler.addWorkerGrid("logits.vocab_proj", vocabWorker);
+        if (DEVICE_SAMPLE) {
+            var argmaxWorker = new WorkerGrid1D(SAMPLE_LOCAL);   // one workgroup
+            argmaxWorker.setLocalWork(SAMPLE_LOCAL, 1, 1);
+            tornadoForwardScheduler.addWorkerGrid("logits.argmax_sample", argmaxWorker);
+        }
         return tornadoForwardScheduler;
     }