From 678634128680ebff79db04a7ca75d79abb71fd74 Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Tue, 14 Jul 2026 14:35:08 +0100 Subject: [PATCH] Add on-device greedy sampling (-Dllama.deviceSample): GPU argmax over the logits appends to the FP16 logits graph so only the sampled token id (1 int) crosses to the host instead of the full vocab logits row; decode loop reads state.sampledToken. Guarded to GPU+greedy+FP16+Llama/Mistral/Qwen3, else the host still gets logits --- .gitignore | 1 + .../java/org/beehive/gpullama3/LlamaApp.java | 24 +++++++++++ .../gpullama3/inference/InferenceEngine.java | 18 +++++++- .../gpullama3/inference/state/State.java | 3 ++ .../kernels/TransformerComputeKernels.java | 42 +++++++++++++++++++ .../layers/type/fp16/LogitsFP16Layer.java | 34 ++++++++++++++- 6 files changed, 118 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index c8532c71..0cd3929f 100644 --- a/.gitignore +++ b/.gitignore @@ -67,3 +67,4 @@ logs/target/ .project .classpath .vscode/ +external/ diff --git a/src/main/java/org/beehive/gpullama3/LlamaApp.java b/src/main/java/org/beehive/gpullama3/LlamaApp.java index ef64ca23..f8fa495f 100644 --- a/src/main/java/org/beehive/gpullama3/LlamaApp.java +++ b/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()) { diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index d653ba58..14bb6688 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/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 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 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) { diff --git a/src/main/java/org/beehive/gpullama3/inference/state/State.java b/src/main/java/org/beehive/gpullama3/inference/state/State.java index 1f884b98..23055d7e 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/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; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java index 80617400..8a34e6f9 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernels.java +++ b/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. */ diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java index 1295db09..a97a37f2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LogitsFP16Layer.java +++ b/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; }