diff --git a/README.md b/README.md index 6c89b398..d9bb49c6 100644 --- a/README.md +++ b/README.md @@ -493,6 +493,41 @@ Any OpenAI SDK works by pointing `base_url` at the server. Generation is seriali context (requests queue); the reusable core is `server/InferenceService`. Smoke test: `scripts/server-smoke-test.sh http://localhost:8080`. +## 📊 Benchmarking (`--bench`, llama-bench style) + +A [llama-bench](https://github.com/ggml-org/llama.cpp/tree/master/tools/llama-bench)-equivalent +benchmark for the GPU forward path: a cartesian matrix of tests over one or more models — +prompt processing (`pp N`), token generation (`tg N`) and combined (`pg`) — each repeated +`-r` times after an untimed warmup, reported as **average tokens/s ± stddev**. Timings cover +the forward pass only (no tokenization, no sampling), matching llama-bench methodology. + +```bash +# defaults: -p 512 -n 128 -r 5, markdown output +llama-tornado --gpu --cuda --model model1.gguf --bench + +# multiple models, custom matrix, CSV +llama-tornado --gpu --cuda --model model1.gguf --bench \ + --bench-args="-m model2.gguf -p 256,512 -n 64,128 -pg 512,128 -d 0,4096 -r 5 -o csv" +``` + +Benchmark options (`--bench-args="..."` — use the `=` form): `-m` extra models +(comma/repeatable), `-p` prompt sizes, `-n` generation lengths, `-pg pp,tg` combined tests, +`-b` prompt-processing batch size — when >1, `pp` tokens run through the batched-prefill +tensor-core (MMA) path, `b` tokens per forward (compute-bound, the llama-bench `-b`); +generation stays single-token decode. Supported models: Llama / Qwen3 / Mistral (FP16). +`-d` context depths (untimed KV prefill of `d` positions before each timed test, e.g. +`tg128@d4096`), `-r` repetitions, `-o md|csv|json|jsonl|sql`, `-oe ` second format to +stderr, `--delay ` pause between tests, `--no-warmup`. + +Example output (RTX 4090, CUDA backend): + +| model | quant | size | params | backend | test | t/s | +| ----- | ----- | ---: | -----: | ------- | ---- | --: | +| beehive-llama-3.2-1b-instruct-fp16 | FP16 | 2.31 GiB | 1.24 B | TornadoVM CUDA | pp512 | 86.62 ± 2.18 | +| beehive-llama-3.2-1b-instruct-fp16 | FP16 | 2.31 GiB | 1.24 B | TornadoVM CUDA | tg128 | 101.72 ± 0.43 | +| Qwen3-1.7B-f16 | FP16 | 3.79 GiB | 2.03 B | TornadoVM CUDA | pp512 | 57.29 ± 0.18 | +| Qwen3-1.7B-f16 | FP16 | 3.79 GiB | 2.03 B | TornadoVM CUDA | tg128 | 57.71 ± 0.10 | + ## Debug & Profiling Options View TornadoVM's internal behavior: ```bash diff --git a/llama-tornado b/llama-tornado index 581caba6..9f5367c6 100755 --- a/llama-tornado +++ b/llama-tornado @@ -196,13 +196,18 @@ class LlamaRunner: ] ) + if getattr(args, "server", False): + main_class = "org.beehive.gpullama3.server.OpenAIServer" + elif getattr(args, "bench", False): + main_class = "org.beehive.gpullama3.bench.LlamaBench" + else: + main_class = "org.beehive.gpullama3.LlamaApp" + module_config.extend( [ "-cp", self._find_llama_jar(), - "org.beehive.gpullama3.server.OpenAIServer" - if getattr(args, "server", False) - else "org.beehive.gpullama3.LlamaApp", + main_class, ] ) cmd.extend(module_config) @@ -235,6 +240,12 @@ class LlamaRunner: if args.use_gpu: srv.append("--gpu") return cmd + srv + if getattr(args, "bench", False): + # llama-bench mode: model(s) + benchmark options only. If --bench-args already + # supplies -m, use those verbatim (don't also prepend --model → duplicate rows). + extra = args.bench_args.split() if args.bench_args else [] + bench_args = extra if ("-m" in extra or "--model" in extra) else ["-m", args.model_path] + extra + return cmd + bench_args llama_args = [ "-m", args.model_path, @@ -426,6 +437,24 @@ def create_parser() -> argparse.ArgumentParser: help="Run in instruction mode (default)", ) + # OpenAI-compatible server + server_group = parser.add_argument_group("OpenAI-compatible server") + server_group.add_argument("--server", action="store_true", help="Run the OpenAI-compatible HTTP server instead of inference") + server_group.add_argument("--port", type=int, default=8080, help="Server port (default 8080)") + + # Benchmark (llama-bench style) + bench_group = parser.add_argument_group("Benchmark (llama-bench style)") + bench_group.add_argument( + "--bench", + action="store_true", + help="Run the llama-bench-style benchmark (bench.LlamaBench) instead of inference", + ) + bench_group.add_argument( + "--bench-args", + default="", + help='Extra benchmark options (use --bench-args="..."), e.g. "-p 512 -n 128 -d 0,4096 -r 5 -o md" (see LlamaBench)', + ) + # Hardware configuration hw_group = parser.add_argument_group("Hardware Configuration") hw_group.add_argument( @@ -475,11 +504,6 @@ def create_parser() -> argparse.ArgumentParser: help="Directory for profiler output", ) - # OpenAI-compatible server - server_group = parser.add_argument_group("OpenAI-compatible server") - server_group.add_argument("--server", action="store_true", help="Run the OpenAI-compatible HTTP server instead of inference") - server_group.add_argument("--port", type=int, default=8080, help="Server port (default 8080)") - # TornadoVM Execution Verbose options verbose_group = parser.add_argument_group("TornadoVM Execution Verbose") verbose_group.add_argument( diff --git a/llamaTornado b/llamaTornado index 068c7946..c1c28df6 100755 --- a/llamaTornado +++ b/llamaTornado @@ -17,7 +17,8 @@ record Config( boolean printBytecodes, boolean threads, boolean printKernel, boolean fullDump, boolean verboseInit, boolean showCommand, boolean executeAfterShow, - String openclFlags, int maxWaitEvents, boolean verbose + String openclFlags, int maxWaitEvents, boolean verbose, + boolean bench, String benchArgs ) {} Config parseArgs(String[] args) { @@ -51,6 +52,8 @@ Config parseArgs(String[] args) { String openclFlags = "-cl-denorms-are-zero -cl-no-signed-zeros -cl-finite-math-only"; int maxWaitEvents = 32000; boolean verbose = false; + boolean bench = false; + String benchArgs = ""; for (int i = 0; i < args.length; i++) { switch (args[i]) { @@ -86,6 +89,8 @@ Config parseArgs(String[] args) { case "--opencl-flags" -> openclFlags = args[++i]; case "--max-wait-events" -> maxWaitEvents = Integer.parseInt(args[++i]); case "--verbose", "-v" -> verbose = true; + case "--bench" -> bench = true; + case "--bench-args" -> benchArgs = args[++i]; default -> { System.err.println("Unknown option: " + args[i]); System.exit(1); @@ -111,7 +116,8 @@ Config parseArgs(String[] args) { return new Config(modelPath, prompt, systemPrompt, temperature, topP, seed, maxTokens, stream, echo, interactive, instruct, useGpu, backend, gpuMemory, heapMin, heapMax, directMemory, debug, profiler, profilerDumpDir, printBytecodes, threads, printKernel, fullDump, - verboseInit, showCommand, executeAfterShow, openclFlags, maxWaitEvents, verbose); + verboseInit, showCommand, executeAfterShow, openclFlags, maxWaitEvents, verbose, + bench, benchArgs); } String parseAndScale(String memoryValue, int multiplier) { @@ -170,6 +176,10 @@ void printUsage() { --execute-after-show Execute after showing command --verbose, -v Verbose output + Benchmark (llama-bench style): + --bench Run the llama-bench-style benchmark instead of inference + --bench-args Extra benchmark options, e.g. "-p 512 -n 128 -r 5 -o md" + -help Show this help -version Show version """.formatted(name)); @@ -281,7 +291,25 @@ List buildCommand(Config cfg, String javaHome, String tornadoSdk, String } } - cmd.addAll(List.of("-cp", findLlamaJar(llamaRoot), "org.beehive.gpullama3.LlamaApp")); + String mainClass = cfg.bench() + ? "org.beehive.gpullama3.bench.LlamaBench" + : "org.beehive.gpullama3.LlamaApp"; + cmd.addAll(List.of("-cp", findLlamaJar(llamaRoot), mainClass)); + + if (cfg.bench()) { + // llama-bench mode: model(s) + benchmark options only. If --bench-args already + // supplies -m, use those verbatim (don't also prepend --model → duplicate rows). + var extra = cfg.benchArgs() == null || cfg.benchArgs().isBlank() + ? List.of() + : List.of(cfg.benchArgs().trim().split("\\s+")); + if (extra.contains("-m") || extra.contains("--model")) { + cmd.addAll(extra); + } else { + cmd.addAll(List.of("-m", cfg.modelPath())); + cmd.addAll(extra); + } + return cmd; + } // LLaMA arguments cmd.addAll(List.of( diff --git a/src/main/java/org/beehive/gpullama3/bench/LlamaBench.java b/src/main/java/org/beehive/gpullama3/bench/LlamaBench.java new file mode 100644 index 00000000..f9de3631 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/bench/LlamaBench.java @@ -0,0 +1,354 @@ +package org.beehive.gpullama3.bench; + +import org.beehive.gpullama3.Options; +import org.beehive.gpullama3.inference.InferenceCore; +import org.beehive.gpullama3.inference.InferenceCoreBatchPrefillDecode; +import org.beehive.gpullama3.inference.state.State; +import org.beehive.gpullama3.model.Model; +import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlan; +import org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode; + +import java.util.Arrays; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.ArrayList; +import java.util.List; +import java.util.Locale; +import java.util.Random; + +import static org.beehive.gpullama3.model.loader.ModelLoader.loadModel; + +/** + * llama-bench-style performance benchmark for GPULlama3 (GPU forward path). + * + *

Mirrors llama.cpp's {@code llama-bench}: a cartesian matrix of tests over one or more + * models — prompt processing ({@code pp N}: N sequential forwards from position 0), token + * generation ({@code tg N}: N single-token forwards over a growing KV cache) and combined + * ({@code pg pp+tg}) — each repeated {@code -r} times after an untimed warmup, reported as + * average tokens/s ± stddev in markdown (default), CSV or JSON. Timings cover the forward + * pass only: no tokenization, no sampling, no host argmax (llama-bench parity).

+ * + * 
+ * gpullama3-bench (via llama-tornado --bench):
+ *   -m  model.gguf[,model2.gguf]   models (repeatable / comma-separated)
+ *   -p  512[,1024]                 prompt-processing sizes       (default 512)
+ *   -n  128[,256]                  generation lengths            (default 128)
+ *   -b  1                          prompt-processing batch (>1 = batched-prefill MMA path,
+ *                                  compute-bound pp; Llama/Qwen3/Mistral FP16)
+ *   -pg 512,128                    combined prompt+gen test      (repeatable)
+ *   -d  0[,4096]                   context depths: untimed KV prefill of d positions
+ *                                  before each timed test (llama-bench -d)
+ *   -r  5                          repetitions                   (default 5)
+ *   -o  md|csv|json|jsonl|sql      output format                 (default md)
+ *   -oe fmt                        also print results to stderr in this format
+ *   --delay N                      sleep N s between tests (GPU thermals)
+ *   --no-warmup                    skip the untimed warmup rep
+ *
+ */ +public class LlamaBench { + + record TestSpec(int nPrompt, int nGen, int depth) { + String name() { + String base = (nPrompt > 0 && nGen > 0) ? "pp" + nPrompt + "+tg" + nGen + : nPrompt > 0 ? "pp" + nPrompt : "tg" + nGen; + return depth > 0 ? base + "@d" + depth : base; + } + + int tokens() { + return nPrompt + nGen; + } + } + + record Result(String model, String quant, double sizeGiB, double paramsB, String backend, String test, double avg, double stddev, double[] samples) {} + + public static void main(String[] args) throws Exception { + System.setProperty("llama.enableTornadoVM", "true"); + + // Pre-scan -b: the batched-prefill plan + state buffers are gated on these system + // properties, read once at class-init — set BEFORE any TornadoVM/State class loads. + int batch = 1; + for (int i = 0; i < args.length - 1; i++) { + if (args[i].equals("-b") || args[i].equals("--batch-size")) { + batch = Integer.parseInt(args[i + 1]); + } + } + if (batch > 1) { + System.setProperty("llama.withPrefillDecode", "true"); + System.setProperty("llama.prefillBatchSize", String.valueOf(batch)); + } + final int batchSize = batch; + + List models = new ArrayList<>(); + List pps = new ArrayList<>(); + List tgs = new ArrayList<>(); + List pgs = new ArrayList<>(); + List depths = new ArrayList<>(); + int reps = 5; + int delay = 0; + String out = "md"; + String outErr = null; + boolean warmup = true; + + for (int i = 0; i < args.length; i++) { + switch (args[i]) { + case "-m", "--model" -> { for (String m : args[++i].split(",")) models.add(m); } + case "-p", "--n-prompt" -> { for (String v : args[++i].split(",")) pps.add(Integer.parseInt(v)); } + case "-n", "--n-gen" -> { for (String v : args[++i].split(",")) tgs.add(Integer.parseInt(v)); } + case "-pg" -> { + String[] parts = args[++i].split("[,+]"); + pgs.add(new int[] { Integer.parseInt(parts[0]), Integer.parseInt(parts[1]) }); + } + case "-b", "--batch-size" -> i++; // consumed in pre-scan + case "-d", "--n-depth" -> { for (String v : args[++i].split(",")) depths.add(Integer.parseInt(v)); } + case "-r", "--repetitions" -> reps = Integer.parseInt(args[++i]); + case "-o", "--output" -> out = args[++i]; + case "-oe", "--output-err" -> outErr = args[++i]; + case "--delay" -> delay = Integer.parseInt(args[++i]); + case "--no-warmup" -> warmup = false; + default -> { + // Ignore the launcher's trailing single-run options (prompt etc.). + } + } + } + if (models.isEmpty()) { + System.err.println("usage: LlamaBench -m model.gguf [-m model2.gguf] [-p 512] [-n 128] [-pg 512,128] [-d 0] [-r 5] [-o md|csv|json|jsonl|sql] [-oe fmt] [--delay s] [--no-warmup]"); + System.exit(1); + } + if (pps.isEmpty() && tgs.isEmpty() && pgs.isEmpty()) { + pps.add(512); + tgs.add(128); + } + + if (depths.isEmpty()) { + depths.add(0); + } + List tests = new ArrayList<>(); + for (int d : depths) { + for (int p : pps) { + if (p > 0) { + tests.add(new TestSpec(p, 0, d)); + } + } + for (int n : tgs) { + if (n > 0) { + tests.add(new TestSpec(0, n, d)); + } + } + for (int[] pg : pgs) { + tests.add(new TestSpec(pg[0], pg[1], d)); + } + } + + List results = new ArrayList<>(); + for (String modelPath : models) { + try { + results.addAll(benchModel(modelPath, tests, reps, warmup, delay, batchSize)); + } catch (Throwable e) { + System.err.printf(Locale.ROOT, "[bench] %s FAILED (batch=%d unsupported for this model?): %s%n", + modelPath, batchSize, e); + } + } + + print(out, results, System.out); + if (outErr != null) { + print(outErr, results, System.err); + } + // TornadoVM daemon threads keep the JVM alive after plan teardown. + System.exit(0); + } + + static void print(String format, List results, java.io.PrintStream ps) { + switch (format) { + case "csv" -> printCsv(results, ps); + case "json" -> printJson(results, ps); + case "jsonl" -> printJsonl(results, ps); + case "sql" -> printSql(results, ps); + default -> printMarkdown(results, ps); + } + } + + static List benchModel(String modelPath, List tests, int reps, boolean warmup, int delay, int batch) throws Exception { + Path path = Paths.get(modelPath); + int maxCtx = tests.stream().mapToInt(t -> t.depth() + t.tokens()).max().orElse(1024) + 8; + Options options = new Options(path, "bench", null, null, false, 0.0f, 1.0f, 42, maxCtx, false, false, true, batch > 1, batch); + Model model = loadModel(options); + State state = model.createNewState(); + TornadoVMMasterPlan plan = TornadoVMMasterPlan.initializeTornadoVMPlan(state, model); + + int vocab = model.configuration().vocabularySize(); + String name = path.getFileName().toString().replaceAll("\\.gguf$", ""); + String quant = model.configuration().quantization(); + double sizeGiB = Files.size(path) / (1024.0 * 1024.0 * 1024.0); + double paramsB = estimateParamsB(Files.size(path), quant); + String backend = System.getProperty("tornado.backend.name", "TornadoVM " + backendName()); + + // Deterministic synthetic token stream (llama-bench uses random ids too). + Random rng = new Random(42); + int maxTokens = tests.stream().mapToInt(t -> t.depth() + t.tokens()).max().orElse(0); + int[] toks = new int[maxTokens]; + for (int i = 0; i < maxTokens; i++) { + toks[i] = rng.nextInt(vocab); + } + + List results = new ArrayList<>(); + for (TestSpec t : tests) { + if (delay > 0) { + Thread.sleep(delay * 1000L); + } + if (warmup) { + runTest(model, state, plan, toks, t, batch); // untimed + } + double[] samples = new double[reps]; + for (int r = 0; r < reps; r++) { + samples[r] = runTest(model, state, plan, toks, t, batch); + } + double avg = 0; + for (double s : samples) { + avg += s; + } + avg /= samples.length; + double var = 0; + for (double s : samples) { + var += (s - avg) * (s - avg); + } + double stddev = samples.length > 1 ? Math.sqrt(var / (samples.length - 1)) : 0.0; + String testName = batch > 1 ? t.name() + " b" + batch : t.name(); + results.add(new Result(name, quant, sizeGiB, paramsB, backend, testName, avg, stddev, samples)); + System.err.printf(Locale.ROOT, "[bench] %-28s %-14s %8.2f ± %.2f t/s%n", name, testName, avg, stddev); + } + plan.freeTornadoExecutionPlan(); + return results; + } + + /** + * One timed repetition: fresh sequence from position 0 (KV overwritten). With a depth d, + * d positions are prefilled untimed first and the timed window runs at positions + * d..d+tokens (llama-bench {@code -d}: measures throughput at context depth). + * + *

With {@code batch > 1} the prompt-processing tokens (nPrompt) run through the + * batched-prefill MMA path — {@code batch} tokens per forward, compute-bound (llama-bench + * {@code -b}) — while generation tokens (nGen) stay single-token decode. Returns tokens/s.

+ */ + static double runTest(Model model, State state, TornadoVMMasterPlan plan, int[] toks, TestSpec t, int batch) { + // Untimed depth prefill. + prefill(model, state, plan, toks, 0, t.depth(), batch); + + int base = t.depth(); + long t0 = System.nanoTime(); + // Prompt processing: batched-prefill chunks when batch>1, else single-token. + prefill(model, state, plan, toks, base, t.nPrompt(), batch); + // Generation: always single-token decode over the growing KV cache. + for (int i = 0; i < t.nGen(); i++) { + int pos = base + t.nPrompt() + i; + if (batch > 1) { + InferenceCoreBatchPrefillDecode.forwardTornadoVMDecode(model, state, toks[pos], pos, + (TornadoVMMasterPlanBatchPrefillDecode) plan); + } else { + InferenceCore.forwardTornadoVM(model, state, toks[pos], pos, plan); + } + } + long t1 = System.nanoTime(); + return t.tokens() / ((t1 - t0) / 1e9); + } + + /** Feed {@code count} tokens from {@code toks[start..]} at positions {@code start..} (prefill). */ + static void prefill(Model model, State state, TornadoVMMasterPlan plan, int[] toks, int start, int count, int batch) { + if (count <= 0) { + return; + } + if (batch > 1) { + var bp = (TornadoVMMasterPlanBatchPrefillDecode) plan; + for (int off = 0; off < count; off += batch) { + int chunkSize = Math.min(batch, count - off); + int[] chunk = Arrays.copyOfRange(toks, start + off, start + off + chunkSize); + InferenceCoreBatchPrefillDecode.batchForwardTornadoVMPrefill(model, state, chunk, start + off, chunkSize, bp); + } + } else { + for (int i = 0; i < count; i++) { + InferenceCore.forwardTornadoVM(model, state, toks[start + i], start + i, plan); + } + } + } + + /** Rough parameter count from file size + quantization (F16 = 2 B/param, Q8_0 = 34/32 B/param). */ + static double estimateParamsB(long bytes, String quant) { + double bytesPerParam = switch (quant) { + case "FP16" -> 2.0; + case "Q8_0" -> 34.0 / 32.0; + default -> 2.0; + }; + return bytes / bytesPerParam / 1e9; + } + + static String backendName() { + try { + return uk.ac.manchester.tornado.api.runtime.TornadoRuntimeProvider.getTornadoRuntime().getBackend(0).getBackendType().name(); + } catch (Throwable e) { + return "unknown"; + } + } + + static void printMarkdown(List rs, java.io.PrintStream ps) { + ps.println(); + ps.println("| model | quant | size | params | backend | test | t/s |"); + ps.println("| ----- | ----- | ---: | -----: | ------- | ---- | --: |"); + for (Result r : rs) { + ps.printf(Locale.ROOT, "| %s | %s | %.2f GiB | %.2f B | %s | %s | %.2f ± %.2f |%n", + r.model(), r.quant(), r.sizeGiB(), r.paramsB(), r.backend(), r.test(), r.avg(), r.stddev()); + } + } + + static void printCsv(List rs, java.io.PrintStream ps) { + ps.println("model,quant,size_gib,params_b,backend,test,avg_ts,stddev_ts,samples"); + for (Result r : rs) { + StringBuilder samples = new StringBuilder(); + for (double s : r.samples()) { + if (samples.length() > 0) { + samples.append(';'); + } + samples.append(String.format(Locale.ROOT, "%.2f", s)); + } + ps.printf(Locale.ROOT, "%s,%s,%.3f,%.3f,%s,%s,%.2f,%.2f,%s%n", + r.model(), r.quant(), r.sizeGiB(), r.paramsB(), r.backend(), r.test(), r.avg(), r.stddev(), samples); + } + } + + static String jsonRow(Result r) { + StringBuilder sb = new StringBuilder(String.format(Locale.ROOT, + "{\"model\": \"%s\", \"quant\": \"%s\", \"size_gib\": %.3f, \"params_b\": %.3f, \"backend\": \"%s\", \"test\": \"%s\", \"avg_ts\": %.2f, \"stddev_ts\": %.2f, \"samples_ts\": [", + r.model(), r.quant(), r.sizeGiB(), r.paramsB(), r.backend(), r.test(), r.avg(), r.stddev())); + for (int j = 0; j < r.samples().length; j++) { + if (j > 0) { + sb.append(", "); + } + sb.append(String.format(Locale.ROOT, "%.2f", r.samples()[j])); + } + sb.append("]}"); + return sb.toString(); + } + + static void printJson(List rs, java.io.PrintStream ps) { + StringBuilder sb = new StringBuilder("[\n"); + for (int i = 0; i < rs.size(); i++) { + sb.append(" ").append(jsonRow(rs.get(i))).append(i < rs.size() - 1 ? "," : "").append('\n'); + } + sb.append("]"); + ps.println(sb); + } + + static void printJsonl(List rs, java.io.PrintStream ps) { + for (Result r : rs) { + ps.println(jsonRow(r)); + } + } + + static void printSql(List rs, java.io.PrintStream ps) { + ps.println("CREATE TABLE IF NOT EXISTS llama_bench (model TEXT, quant TEXT, size_gib REAL, params_b REAL, backend TEXT, test TEXT, avg_ts REAL, stddev_ts REAL);"); + for (Result r : rs) { + ps.printf(Locale.ROOT, "INSERT INTO llama_bench VALUES ('%s', '%s', %.3f, %.3f, '%s', '%s', %.2f, %.2f);%n", + r.model(), r.quant(), r.sizeGiB(), r.paramsB(), r.backend(), r.test(), r.avg(), r.stddev()); + } + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java index 0b7482e3..1b400e61 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java @@ -2489,6 +2489,44 @@ public static float matrixVectorRowMajorOptimizedx(KernelContext context, int lo return localSum[0]; } + /** + * Race-free RMS-norm reduction in ONE workgroup: threads stride over the input, reduce in + * local memory, and lane 0 writes the final scale factor to {@code output[0]}. + * + *

The multi-workgroup variant ({@code reductionOneBlockWithLayer}) has workgroup 0 + * combining the other workgroups' partial sums with no inter-workgroup + * synchronization — a data race. On the CUDA backend the race outcome is + * schedule/compilation dependent: e.g. Qwen3-1.7B reads stale partials on every layer + * (wrong normalization scale → garbage output) while Llama-1B and Qwen3-4B happen to win + * the race. The NON_NVIDIA scheduler avoids the race with a separate + * {@code reductionFinalNormalization} task; this kernel is the NVIDIA-path equivalent. + * Launch with exactly one workgroup: {@code global == local == localMemSize}.

+ */ + public static void reductionOneBlockWithLayerSingleGroup(KernelContext context, FloatArray output, FloatArray x, int size, float ermsNorm, int localMemSize) { + int lid = context.localIdx; + int groupSize = context.localGroupSizeX; + float[] localX = context.allocateFloatLocalArray(localMemSize); + + float partial = 0.0f; + for (int j = lid; j < size; j += groupSize) { + float v = x.get(j); + partial += v * v; + } + localX[lid] = partial; + + for (int stride = groupSize / 2; stride > 0; stride /= 2) { + context.localBarrier(); + if (lid < stride) { + localX[lid] += localX[lid + stride]; + } + } + + if (lid == 0) { + float ss = localX[0] / size + ermsNorm; + output.set(0, 1.0f / TornadoMath.sqrt(ss)); + } + } + // Second kernel - Combines partial sums and computes final normalization public static void reductionFinalNormalization(KernelContext context, FloatArray output, int size, float ermsNorm) { int gid = context.globalIdx; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java index ee2faccb..f0da8b54 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java @@ -59,6 +59,8 @@ public Qwen3FP16FFNLayers(String taskGraphName, Qwen3State state, Qwen3TornadoWe @Override public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { WorkerGrid rmsNormWorker = WorkerGridFactory.createRmsNormWorker(config.dim(), state.localSize); + // Single-workgroup grid for the race-free NVIDIA-path RMS reduction (global == local). + WorkerGrid rmsSingleGroupWorker = WorkerGridFactory.createRmsNormWorker(state.localSize, state.localSize); WorkerGrid ropeWorker = WorkerGridFactory.createRoPEWorker(config.numberOfHeads(), nEmbdHead); // Split-KV attention launches nHeads*nSplits workgroups (one per head-split) followed by a combine // pass over nHeads workgroups. @@ -85,7 +87,8 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { // Map workers to tasks for each layer (in task execution order) for (int i = 0; i < config.numberOfLayers(); i++) { // === Attention Block === - gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_reduce", rmsNormWorker); + gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_reduce", + shouldUseFinalNormalization() ? rmsNormWorker : rmsSingleGroupWorker); gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_qkv_projection", fusedQKVWorker); gridScheduler.addWorkerGrid("layer_" + i + ".qk_rmsnorm", qkRmsNormWorker); gridScheduler.addWorkerGrid("layer_" + i + ".rope_and_kv_cache", ropeWorker); @@ -93,7 +96,8 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { gridScheduler.addWorkerGrid("layer_" + i + ".attention_combine", attentionCombineWorker); gridScheduler.addWorkerGrid("layer_" + i + ".attn_output_proj", matmul1Worker); // === FFN Block === - gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_reduce", rmsNormWorker); + gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_reduce", + shouldUseFinalNormalization() ? rmsNormWorker : rmsSingleGroupWorker); if (shouldUseFinalNormalization()) { gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_finalize", rmsNormWorker); } @@ -230,6 +234,18 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { // ═══════════════════════════════════════════════════════════════════════ // RMS Normalization - compute scale factor + if (!shouldUseFinalNormalization()) { + // NVIDIA path: single-workgroup reduction. The multi-workgroup kernel relies on a + // racy cross-workgroup combine (no finalize task on this path) — see kernel javadoc. + unifiedLayer.task("attn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, + context, + qwen3State.temp, + qwen3State.wrapX, + config.dim(), + config.rmsNormEps(), + qwen3State.localSize); + } else unifiedLayer.task("attn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, @@ -361,6 +377,17 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { // ═══════════════════════════════════════════════════════════════════════ // RMS Normalization - compute scale factor + if (!shouldUseFinalNormalization()) { + // NVIDIA path: single-workgroup reduction (race-free) — see attn_rms_reduce. + unifiedLayer.task("ffn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, + context, + qwen3State.tempFFN, + qwen3State.wrapX, + config.dim(), + config.rmsNormEps(), + qwen3State.localSize); + } else unifiedLayer.task("ffn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context,