From 33b9835112076cac5f9a45a8d5f78dcbf9b9f6d5 Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Sat, 20 Jun 2026 14:46:52 +0300 Subject: [PATCH 01/18] feat(launcher): add CUDA backend path Add a --cuda flag to llama-tornado that selects the TornadoVM CUDA backend, mirroring the existing --opencl/--ptx/--metal plumbing: it loads the tornado.drivers.cuda module and the cuda-exports export list. Also disambiguate --ptx help text (was 'PTX/CUDA'). --- llama-tornado | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/llama-tornado b/llama-tornado index 1d6c3d23..78388295 100755 --- a/llama-tornado +++ b/llama-tornado @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """ llama-tornado: GPU-accelerated Java LLM runner with TornadoVM -Run LLM models using either OpenCL or PTX backends. +Run LLM models using OpenCL, PTX, CUDA, or Metal backends. """ import argparse @@ -19,6 +19,7 @@ from enum import Enum class Backend(Enum): OPENCL = "opencl" PTX = "ptx" + CUDA = "cuda" METAL = "metal" @@ -178,6 +179,14 @@ class LlamaRunner: "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.ptx", ] ) + elif args.backend == Backend.CUDA: + module_config.extend( + [ + f"@{self.tornado_sdk}/etc/exportLists/cuda-exports", + "--add-modules", + "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.cuda", + ] + ) elif args.backend == Backend.METAL: module_config.extend( [ @@ -426,7 +435,14 @@ def create_parser() -> argparse.ArgumentParser: dest="backend", action="store_const", const=Backend.PTX, - help="Use PTX/CUDA backend", + help="Use PTX backend", + ) + hw_group.add_argument( + "--cuda", + dest="backend", + action="store_const", + const=Backend.CUDA, + help="Use CUDA backend (requires TornadoVM built with the CUDA backend)", ) hw_group.add_argument( "--metal", From 2e2fa90b3fc3fa144afe76adfe68a1248c66b6f1 Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Sat, 20 Jun 2026 14:46:52 +0300 Subject: [PATCH 02/18] build(pom): build against TornadoVM 4.0.2-jdk21-dev (CUDA backend) The CUDA backend is only available in a dev build of TornadoVM (PR #861), so point the JDK21 build at 4.0.2-jdk21-dev. The project's own version is unchanged. --- pom.xml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pom.xml b/pom.xml index a83c7ecf..82e875e9 100644 --- a/pom.xml +++ b/pom.xml @@ -39,9 +39,10 @@ 0.4.0 - 4.0.1 + 4.0.2 -jdk21 - ${tornadovm.base.version}${jdk.version.suffix} + + ${tornadovm.base.version}${jdk.version.suffix}-dev 25 25 @@ -147,7 +148,8 @@ 21 21 -jdk21 - ${tornadovm.base.version}${jdk.version.suffix} + + ${tornadovm.base.version}${jdk.version.suffix}-dev From 74b88c209e5ddb2ed2c2c465177bfd4d82bb3aeb Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Sat, 20 Jun 2026 14:46:52 +0300 Subject: [PATCH 03/18] docs: document CUDA backend and TornadoVM PR #861 requirement List CUDA among the supported backends, add a --cuda usage example, and note that the CUDA backend requires a TornadoVM build with the CUDA backend from PR #861 (https://github.com/beehive-lab/TornadoVM/pull/861). --- README.md | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 2e2db217..9cb3a0a9 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,8 @@ GPULlama3ChatModel model = GPULlama3ChatModel.builder() Ensure you have the following installed and configured: - **Java 21**: Required for Vector API support & TornadoVM. -- [TornadoVM](https://github.com/beehive-lab/TornadoVM) with OpenCL or PTX backends. +- [TornadoVM](https://github.com/beehive-lab/TornadoVM) with OpenCL, PTX, or CUDA backends. + - The `--cuda` backend requires a TornadoVM build that includes the CUDA backend from [TornadoVM PR #861](https://github.com/beehive-lab/TornadoVM/pull/861). This project currently builds against TornadoVM `4.0.2-jdk21-dev`. - GCC/G++ 13 or newer: Required to build and run TornadoVM native components. ### Install, Build, and Run @@ -305,6 +306,12 @@ Run a model with a text prompt: ./llama-tornado --gpu --verbose-init --opencl --model beehive-llama-3.2-1b-instruct-fp16.gguf --prompt "Explain the benefits of GPU acceleration." ``` +Select a backend explicitly with `--opencl`, `--ptx`, or `--cuda` (NVIDIA), or `--metal` (Apple Silicon). For example, to run on the CUDA backend: + +```bash +./llama-tornado --gpu --cuda --model beehive-llama-3.2-1b-instruct-fp16.gguf --prompt "Explain the benefits of GPU acceleration." +``` + #### GPU Execution (FP16 Model) Enable GPU acceleration with Q8_0 quantization: ```bash @@ -393,7 +400,7 @@ Supported command-line options include: ```bash cmd ➜ llama-tornado --help usage: llama-tornado [-h] --model MODEL_PATH [--prompt PROMPT] [-sp SYSTEM_PROMPT] [--temperature TEMPERATURE] [--top-p TOP_P] [--seed SEED] [-n MAX_TOKENS] - [--stream STREAM] [--echo ECHO] [-i] [--instruct] [--gpu] [--opencl] [--ptx] [--gpu-memory GPU_MEMORY] [--heap-min HEAP_MIN] [--heap-max HEAP_MAX] + [--stream STREAM] [--echo ECHO] [-i] [--instruct] [--gpu] [--opencl] [--ptx] [--cuda] [--metal] [--gpu-memory GPU_MEMORY] [--heap-min HEAP_MIN] [--heap-max HEAP_MAX] [--debug] [--profiler] [--profiler-dump-dir PROFILER_DUMP_DIR] [--print-bytecodes] [--print-threads] [--print-kernel] [--full-dump] [--show-command] [--execute-after-show] [--opencl-flags OPENCL_FLAGS] [--max-wait-events MAX_WAIT_EVENTS] [--verbose] @@ -424,7 +431,9 @@ Mode Selection: Hardware Configuration: --gpu Enable GPU acceleration (default: False) --opencl Use OpenCL backend (default) (default: None) - --ptx Use PTX/CUDA backend (default: None) + --ptx Use PTX backend (default: None) + --cuda Use CUDA backend (requires TornadoVM built with the CUDA backend) (default: None) + --metal Use Apple Metal backend (macOS only) (default: None) --gpu-memory GPU_MEMORY GPU memory allocation (default: 7GB) --heap-min HEAP_MIN Minimum JVM heap size (default: 20g) @@ -480,9 +489,9 @@ View TornadoVM's internal behavior: - **Support for GGUF format models** with full FP16 and partial support for Q8_0 and Q4_0 quantization. - **Instruction-following and chat modes** for various use cases. - **Interactive CLI** with `--interactive` and `--instruct` modes. - - **Flexible backend switching** - choose OpenCL or PTX at runtime (need to build TornadoVM with both enabled). + - **Flexible backend switching** - choose OpenCL, PTX, or CUDA at runtime (need to build TornadoVM with the chosen backends enabled). - **Cross-platform compatibility**: - - ✅ NVIDIA GPUs (OpenCL & PTX ) + - ✅ NVIDIA GPUs (OpenCL, PTX & CUDA) - ✅ Intel GPUs (OpenCL) - ✅ Apple GPUs (OpenCL) From ed4db212e42977400eef82a3ee143e7945a504ec Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Sat, 20 Jun 2026 15:08:47 +0300 Subject: [PATCH 04/18] ci: add CUDA backend to build and inference matrices Add a cuda variant to the build, standalone-inference, and quarkus-integration backend matrices. The setup-tornadovm action now builds the CUDA backend from the cuda2 branch (TornadoVM PR #861) until it is merged to master; other backends still build from master. Shared inference steps run on CUDA via the matrix; the PTX-only CUDA-graph steps remain gated to ptx. --- .github/actions/run-inference/action.yml | 2 +- .github/actions/setup-tornadovm/action.yml | 20 ++++++++++++++++---- .github/workflows/build-and-run.yml | 5 ++++- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/.github/actions/run-inference/action.yml b/.github/actions/run-inference/action.yml index 314a8be5..fe3d574b 100644 --- a/.github/actions/run-inference/action.yml +++ b/.github/actions/run-inference/action.yml @@ -3,7 +3,7 @@ description: Run one llama-tornado inference pass and write the metrics + sideca inputs: backend: - description: 'GPU backend (opencl or ptx)' + description: 'GPU backend (opencl, ptx, or cuda)' required: true model_file: description: 'Model filename inside $MODELS_DIR (e.g. Llama-3.2-1B-Instruct-F16.gguf)' diff --git a/.github/actions/setup-tornadovm/action.yml b/.github/actions/setup-tornadovm/action.yml index 3b1c5070..01fc41ac 100644 --- a/.github/actions/setup-tornadovm/action.yml +++ b/.github/actions/setup-tornadovm/action.yml @@ -3,17 +3,29 @@ description: Build TornadoVM once per backend and reuse across runs via a local inputs: backend: - description: 'TornadoVM backend to build (opencl or ptx)' + description: 'TornadoVM backend to build (opencl, ptx, or cuda)' required: true runs: using: composite steps: + - name: Determine TornadoVM branch + id: branch + shell: bash + run: | + # The CUDA backend currently lives on the cuda2 branch (TornadoVM PR #861) + # until it is merged to master; all other backends build from master. + if [ "${{ inputs.backend }}" = "cuda" ]; then + echo "ref=cuda2" >> $GITHUB_OUTPUT + else + echo "ref=master" >> $GITHUB_OUTPUT + fi + - name: Get TornadoVM HEAD SHA id: tornado_sha shell: bash run: | - SHA=$(git ls-remote https://github.com/beehive-lab/TornadoVM HEAD | cut -f1) + SHA=$(git ls-remote https://github.com/beehive-lab/TornadoVM ${{ steps.branch.outputs.ref }} | cut -f1) echo "sha=$SHA" >> $GITHUB_OUTPUT - name: Check local build sentinel @@ -27,12 +39,12 @@ runs: echo "up-to-date=false" >> $GITHUB_OUTPUT fi - - name: Clone TornadoVM master + - name: Clone TornadoVM if: steps.sentinel.outputs.up-to-date != 'true' shell: bash run: | rm -rf $TORNADO_ROOT - git clone --depth 1 --branch master \ + git clone --depth 1 --branch ${{ steps.branch.outputs.ref }} \ https://github.com/beehive-lab/TornadoVM.git \ $TORNADO_ROOT diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index f60083cc..7b7da0ad 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -31,7 +31,7 @@ jobs: # ./mvnw -T12C -Pspotless spotless:check # Build: TornadoVM → GPULlama3 → Quarkus LangChain4j - # max-parallel: 1 ensures the opencl and ptx variants run sequentially so + # max-parallel: 1 ensures the opencl, ptx and cuda variants run sequentially so # there are no workspace conflicts between matrix jobs. build: if: github.repository == 'beehive-lab/GPULlama3.java' @@ -45,6 +45,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 @@ -99,6 +100,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 @@ -523,6 +525,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 From 3f13606766ea4273ad43078e674c1fedb8881c90 Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Mon, 22 Jun 2026 13:37:10 +0300 Subject: [PATCH 05/18] [hack] force ci to run on strix --- .github/workflows/build-and-run.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index 7b7da0ad..419b1698 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -18,7 +18,7 @@ env: jobs: code-quality: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: self-hosted + runs-on: [self-hosted, orion-strix] timeout-minutes: 30 steps: From bd4a70b679d50881eec421747668018c4079738f Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Mon, 22 Jun 2026 14:03:39 +0300 Subject: [PATCH 06/18] [hack] force all ci jobs to run on strix --- .github/workflows/build-and-run.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index 419b1698..f2ef0606 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -35,7 +35,7 @@ jobs: # there are no workspace conflicts between matrix jobs. build: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: [self-hosted, orion-strix] needs: code-quality timeout-minutes: 30 strategy: @@ -91,7 +91,7 @@ jobs: standalone-inference: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: [self-hosted, orion-strix] needs: build timeout-minutes: 30 strategy: @@ -516,7 +516,7 @@ jobs: # Test integration with Quarkus-langchain4j quarkus-langchain4j-integration: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: [self-hosted, orion-strix] needs: build timeout-minutes: 10 strategy: @@ -618,7 +618,7 @@ jobs: github.event_name == 'push' && github.ref == 'refs/heads/main' - runs-on: [self-hosted] + runs-on: [self-hosted, orion-strix] needs: standalone-inference timeout-minutes: 15 From 3efde927f567b1d9099738e645cdfbcbdad4024f Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Mon, 22 Jun 2026 14:12:16 +0300 Subject: [PATCH 07/18] [hack] use correct runner custom label --- .github/workflows/build-and-run.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index f2ef0606..bdb48336 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -18,7 +18,7 @@ env: jobs: code-quality: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, orion-strix] + runs-on: [self-hosted, 5090-laptop] timeout-minutes: 30 steps: @@ -35,7 +35,7 @@ jobs: # there are no workspace conflicts between matrix jobs. build: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, orion-strix] + runs-on: [self-hosted, 5090-laptop] needs: code-quality timeout-minutes: 30 strategy: @@ -91,7 +91,7 @@ jobs: standalone-inference: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, orion-strix] + runs-on: [self-hosted, 5090-laptop] needs: build timeout-minutes: 30 strategy: @@ -516,7 +516,7 @@ jobs: # Test integration with Quarkus-langchain4j quarkus-langchain4j-integration: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, orion-strix] + runs-on: [self-hosted, 5090-laptop] needs: build timeout-minutes: 10 strategy: @@ -618,7 +618,7 @@ jobs: github.event_name == 'push' && github.ref == 'refs/heads/main' - runs-on: [self-hosted, orion-strix] + runs-on: [self-hosted, 5090-laptop] needs: standalone-inference timeout-minutes: 15 From dcee6cc0ae80c975e831a0c8bc3ac847ed302ece Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Tue, 23 Jun 2026 16:57:37 +0300 Subject: [PATCH 08/18] Revert specific workflow runner labels --- .github/workflows/build-and-run.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index bdb48336..3b33059c 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -18,7 +18,7 @@ env: jobs: code-quality: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, 5090-laptop] + runs-on: [self-hosted] timeout-minutes: 30 steps: @@ -35,7 +35,7 @@ jobs: # there are no workspace conflicts between matrix jobs. build: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, 5090-laptop] + runs-on: [self-hosted] needs: code-quality timeout-minutes: 30 strategy: @@ -91,7 +91,7 @@ jobs: standalone-inference: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, 5090-laptop] + runs-on: [self-hosted] needs: build timeout-minutes: 30 strategy: @@ -516,7 +516,7 @@ jobs: # Test integration with Quarkus-langchain4j quarkus-langchain4j-integration: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted, 5090-laptop] + runs-on: [self-hosted] needs: build timeout-minutes: 10 strategy: @@ -618,7 +618,7 @@ jobs: github.event_name == 'push' && github.ref == 'refs/heads/main' - runs-on: [self-hosted, 5090-laptop] + runs-on: [self-hosted] needs: standalone-inference timeout-minutes: 15 From 4bb539dc94e5fbc5d2e6636451991069f51ed1bb Mon Sep 17 00:00:00 2001 From: MaryXek Date: Thu, 2 Jul 2026 14:21:30 +0300 Subject: [PATCH 09/18] Convert QVK, Wo and W2 prefill projections to use MMA GEMM --- .../gpullama3/inference/state/State.java | 25 +++ .../TransformerBatchPrefillKernels.java | 190 ++++++++++++++++++ .../prefill/LlamaFP16LayersBatchPrefill.java | 185 +++++++++++++---- 3 files changed, 364 insertions(+), 36 deletions(-) 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 0807b756..ae929e04 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -85,6 +85,14 @@ public abstract class State { public final FloatArray attnScaleBatch; // B (per-token RMS scale, attn) public final FloatArray ffnScaleBatch; // B (per-token RMS scale, FFN) public final IntArray batchStartPosHolder; // 1 (start position of chunk) + public final HalfFloatArray normedXFFNFP16; + public final FloatArray ffnGateResult; + public final FloatArray ffnUpResult; + public final HalfFloatArray xbFP16Batch; + public final HalfFloatArray attnOutFP16; + public final FloatArray woOut; + public final HalfFloatArray wrapHbFP16Batch; + public final FloatArray w2Out; protected State(Configuration config, int batchsize) { this.batchsize = batchsize; @@ -148,6 +156,15 @@ protected State(Configuration config, int batchsize) { this.attnScaleBatch = new FloatArray(gpuBatchSize); this.ffnScaleBatch = new FloatArray(gpuBatchSize); this.batchStartPosHolder = new IntArray(1); + this.normedXFFNFP16 = new HalfFloatArray(gpuBatchSize * config.dim()); + this.ffnGateResult = new FloatArray(gpuBatchSize * config.hiddenDim()); + this.ffnUpResult = new FloatArray(gpuBatchSize * config.hiddenDim()); + + this.xbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); + this.attnOutFP16 = new HalfFloatArray(gpuBatchSize * config.dim()); + this.woOut = new FloatArray(gpuBatchSize * config.dim()); + this.wrapHbFP16Batch = new HalfFloatArray(gpuBatchSize * config.hiddenDim()); + this.w2Out = new FloatArray(gpuBatchSize * config.dim()); } else { this.embeddingXBatch = null; this.wrapXBatch = null; @@ -160,6 +177,14 @@ protected State(Configuration config, int batchsize) { this.attnScaleBatch = null; this.ffnScaleBatch = null; this.batchStartPosHolder = null; + this.normedXFFNFP16 = null; + this.ffnGateResult = null; + this.ffnUpResult = null; + this.xbFP16Batch = null; + this.attnOutFP16 = null; + this.woOut = null; + this.wrapHbFP16Batch = null; + this.w2Out = null; } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 057e9f87..3f99bb0e 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -1,6 +1,7 @@ package org.beehive.gpullama3.tornadovm.kernels; import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.enums.MMAShape; import uk.ac.manchester.tornado.api.math.TornadoMath; import uk.ac.manchester.tornado.api.types.HalfFloat; import uk.ac.manchester.tornado.api.types.arrays.ByteArray; @@ -718,5 +719,194 @@ public static void batchedFusedRmsNormFFNGateUpQ8(KernelContext context, } } + /** + * RMS-apply for FFN, writing FP16. Mirrors batchedRmsApplyFP16 but pulls + * scale from ffnScaleBatch. Output is the A operand for the W1/W3 MMA tasks. + * + * Worker: B*dim global threads, localSize=256. + */ + public static void batchedFFNRmsApplyFP16(KernelContext context, + HalfFloatArray normedXFFNFP16, + FloatArray wrapXBatch, + FloatArray rmsFFNWeights, + FloatArray ffnScaleBatch, + int dim) { + int gid = context.globalIdx; + int b = gid / dim; + int i = gid % dim; + float scale = ffnScaleBatch.get(b); + float result = rmsFFNWeights.get(i) * scale * wrapXBatch.get(gid); + normedXFFNFP16.set(gid, new HalfFloat(result)); + } + + /** + * Fused SiLU(gate) * up after the two FFN matmuls. + * Operates on FP32 inputs (MMA writes FP32). + * + * Worker: B*hiddenDim global threads, localSize=256. + */ + public static void batchedFFNSwiGLU(KernelContext context, + FloatArray wrapHbBatch, + FloatArray ffnGateResult, + FloatArray ffnUpResult, + int hiddenDim) { + int gid = context.globalIdx; + float g = ffnGateResult.get(gid); + float u = ffnUpResult.get(gid); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbBatch.set(gid, silu * u); + } + + private static final int WARP_SIZE = 32; + private static final int BM = 128, BN = 128, BK = 16; + private static final int WARPS_M = 4, WARPS_N = 2; + private static final int WARPS_PER_BLOCK = WARPS_M * WARPS_N; + private static final int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + private static final int WM = BM / WARPS_M; + private static final int WN = BN / WARPS_N; + private static final int B_SUBTILE_BYTES = 256; + + public static void gemmMMA(KernelContext ctx, + HalfFloatArray A, HalfFloatArray B, FloatArray C, + int M, int N, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + int kBase = kStep * BK; + // A load: A is [M, K] row-major (unchanged from gemmMMA) + for (int idx = tid; idx < BM * BK / 2; idx += THREADS_PER_BLOCK) { + int m_row = idx / (BK / 2); + int k_pair = idx % (BK / 2); + int k_base = k_pair * 2; + int gA = (blockRow + m_row) * K + (kBase + k_base); + int lo = A.get(gA).getHalfFloatValue() & 0xFFFF; + int hi = A.get(gA + 1).getHalfFloatValue() & 0xFFFF; + aTile[m_row * (BK / 2) + k_pair] = lo | (hi << 16); + } + // B load: B is [N, K] row-major. + // The shared-memory layout in bTile is unchanged; we read from + // a different global layout and pack into the same shared positions + // so that mmaLoadB, mma fragments, and stores all operate identically. + for (int idx = tid; idx < BK * BN / 2; idx += THREADS_PER_BLOCK) { + int subTileId = idx / 64; + int intInSub = idx % 64; + int k_row = intInSub / 4; + int j_pair = intInSub % 4; + int j_base = j_pair * 2; + int col_in_block = subTileId * 8 + j_base; + // B[col, k] is at col * K + k in [N, K] row-major. + int gB0 = (blockCol + col_in_block) * K + (kBase + k_row); + int gB1 = (blockCol + col_in_block + 1) * K + (kBase + k_row); + int lo = B.get(gB0).getHalfFloatValue() & 0xFFFF; + int hi = B.get(gB1).getHalfFloatValue() & 0xFFFF; + bTile[idx] = lo | (hi << 16); + } + ctx.localBarrier(); + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, C, rBase + 0, cBase + 0, N); + ctx.mmaStore(c01, C, rBase + 0, cBase + 8, N); + ctx.mmaStore(c02, C, rBase + 0, cBase + 16, N); + ctx.mmaStore(c03, C, rBase + 0, cBase + 24, N); + ctx.mmaStore(c04, C, rBase + 0, cBase + 32, N); + ctx.mmaStore(c05, C, rBase + 0, cBase + 40, N); + ctx.mmaStore(c06, C, rBase + 0, cBase + 48, N); + ctx.mmaStore(c07, C, rBase + 0, cBase + 56, N); + ctx.mmaStore(c10, C, rBase + 16, cBase + 0, N); + ctx.mmaStore(c11, C, rBase + 16, cBase + 8, N); + ctx.mmaStore(c12, C, rBase + 16, cBase + 16, N); + ctx.mmaStore(c13, C, rBase + 16, cBase + 24, N); + ctx.mmaStore(c14, C, rBase + 16, cBase + 32, N); + ctx.mmaStore(c15, C, rBase + 16, cBase + 40, N); + ctx.mmaStore(c16, C, rBase + 16, cBase + 48, N); + ctx.mmaStore(c17, C, rBase + 16, cBase + 56, N); + } + + // ── Residual add (FP32) ─────────────────────────────────────────────── + // gemmMMA overwrites C, but Wo and W2 both need x = x + W·a. + // Worker: B*dim global threads (valid rows only), localSize=256. + public static void batchedResidualAddFP32(KernelContext context, + FloatArray residual, // x (in/out) + FloatArray delta) { // GEMM output + int gid = context.globalIdx; + residual.set(gid, residual.get(gid) + delta.get(gid)); + } + + // ── SwiGLU emitting FP16 ────────────────────────────────────────────── + // Replaces batchedFFNSwiGLU. Output is the A operand for the W2 GEMM. + // Worker: B*hiddenDim global threads, localSize=256. + public static void batchedFFNSwiGLUFP16(KernelContext context, + HalfFloatArray wrapHbFP16Batch, + FloatArray ffnGateResult, + FloatArray ffnUpResult, + int hiddenDim) { + int gid = context.globalIdx; + float g = ffnGateResult.get(gid); + float u = ffnUpResult.get(gid); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbFP16Batch.set(gid, new HalfFloat(silu * u)); + } + + // ── FP32 → FP16 cast (Option B only, see Wo below) ──────────────────── + // Worker: B*dim global threads, localSize=256. + public static void batchedConvertFP32toFP16(KernelContext context, + FloatArray in, + HalfFloatArray out) { + int gid = context.globalIdx; + out.set(gid, new HalfFloat(in.get(gid))); + } + // @formatter:on } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java index 6f0b3e4b..f590b7b9 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java @@ -11,6 +11,8 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; @@ -70,7 +72,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapXbBatch, state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.ffnGateResult, state.ffnUpResult, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); // wrapXBatch produced by the prefillActivation graph and persists in device memory // to consume it from there we should use the explicit uniqueTaskGraph name // the no-arg form would use current graph name, which causes NPE without CUDA Graphs @@ -81,13 +85,15 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapXbBatch, state.wrapHbBatch, state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); + state.normedXFFNFP16, state.ffnGateResult, state.ffnUpResult, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } // Per-layer weights: upload once @@ -118,15 +124,28 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - batchPrefillLayer.task("batch_qkv", - TransformerBatchPrefillKernels::batchedFusedQKVMatmul, - context, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - weights.wqLayered[layerIndex].asHalfFloatArray(), - weights.wkLayered[layerIndex].asHalfFloatArray(), - weights.wvLayered[layerIndex].asHalfFloatArray(), - dim, kvDim, LOCAL_WORK_GROUP_SIZE); +// batchPrefillLayer.task("batch_qkv", +// TransformerBatchPrefillKernels::batchedFusedQKVMatmul, +// context, +// state.wrapXbFP16Batch, +// state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, +// weights.wqLayered[layerIndex].asHalfFloatArray(), +// weights.wkLayered[layerIndex].asHalfFloatArray(), +// weights.wvLayered[layerIndex].asHalfFloatArray(), +// dim, kvDim, LOCAL_WORK_GROUP_SIZE); + + batchPrefillLayer.task("qProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + state.wrapQBatch, batchSize, dim, dim) + .task("kProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapXbFP16Batch, + weights.wkLayered[layerIndex].asHalfFloatArray(), + state.wrapKBatch, batchSize, kvDim, dim) + .task("vProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapXbFP16Batch, + weights.wvLayered[layerIndex].asHalfFloatArray(), + state.wrapVBatch, batchSize, kvDim, dim); batchPrefillLayer.task("batch_rope_kv", TransformerBatchPrefillKernels::batchedRopeWithKVCache, @@ -143,11 +162,19 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - batchPrefillLayer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapXbBatch, state.wrapXBatch, - weights.woLayered[layerIndex].asHalfFloatArray(), - dim, dim, LOCAL_WORK_GROUP_SIZE); +// batchPrefillLayer.task("batch_attn_out", +// TransformerBatchPrefillKernels::batchedMatVecWithResidual, +// context, state.wrapXbBatch, state.wrapXBatch, +// weights.woLayered[layerIndex].asHalfFloatArray(), +// dim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("attnCast", TransformerBatchPrefillKernels::batchedConvertFP32toFP16, + context, state.wrapXbBatch, state.attnOutFP16) + .task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asHalfFloatArray(), + state.woOut, batchSize, dim, dim) + .task("woResid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.woOut); // ── FFN Block ────────────────────────────────────────────────────────── batchPrefillLayer.task("batch_ffn_rms", @@ -155,20 +182,55 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { context, state.wrapXBatch, state.ffnScaleBatch, dim, config.rmsNormEps()); - batchPrefillLayer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, - context, state.wrapXBatch, state.wrapHbBatch, + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("batch_ffn_w1_mma", + TransformerBatchPrefillKernels::gemmMMA, + context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asHalfFloatArray(), - weights.w3Layered[layerIndex].asHalfFloatArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); + state.ffnGateResult, + batchSize, hidDim, dim); - batchPrefillLayer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("batch_ffn_w3_mma", + TransformerBatchPrefillKernels::gemmMMA, + context, state.normedXFFNFP16, + weights.w3Layered[layerIndex].asHalfFloatArray(), + state.ffnUpResult, + batchSize, hidDim, dim); + + +// batchPrefillLayer.task("batch_ffn_swiglu", +// TransformerBatchPrefillKernels::batchedFFNSwiGLU, +// context, state.wrapHbBatch, state.ffnGateResult, state.ffnUpResult, +// hidDim); + + +// batchPrefillLayer.task("batch_ffn_gate_up", +// TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, +// context, state.wrapXBatch, state.wrapHbBatch, +// weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), +// state.ffnScaleBatch, +// weights.w1Layered[layerIndex].asHalfFloatArray(), +// weights.w3Layered[layerIndex].asHalfFloatArray(), +// dim, hidDim, LOCAL_WORK_GROUP_SIZE); + +// batchPrefillLayer.task("batch_ffn_down", +// TransformerBatchPrefillKernels::batchedMatVecWithResidual, +// context, state.wrapHbBatch, state.wrapXBatch, +// weights.w2Layered[layerIndex].asHalfFloatArray(), +// hidDim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("swiglu", TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16, + context, state.wrapHbFP16Batch, state.ffnGateResult, state.ffnUpResult, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); // Persist wrapXBatch for the next layer, and KV cache so the decode // layers can consume it via the activation graph pass-through. @@ -178,6 +240,21 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { } // @formatter:on + // gemmMMA: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { int dim = config.dim(); @@ -214,17 +291,53 @@ public void updateGridScheduler(GridScheduler scheduler) { WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + // FFN RMS apply: B*dim threads, local=256 + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // MMA: 128x128 block tile, 256 threads/block, M=batchSize, N=hiddenDim, K=dim + // Global = (M/128)*256 in X, (N/128) in Y, local (256,1,1) + WorkerGrid mmaFFNWorker = new WorkerGrid2D((batchSize / 128) * 256, hidDim / 128); + mmaFFNWorker.setLocalWork(256, 1, 1); + + // SwiGLU: B*hiddenDim threads, local=256 + WorkerGrid swigluWorker = WorkerGridFactory.genericWorker(batchSize * hidDim, 256); + + WorkerGrid mmaDimWorker = mmaGrid(batchSize, dim); // qProj, woProj, w2Proj + WorkerGrid mmaKvWorker = mmaGrid(batchSize, kvDim); // kProj, vProj + WorkerGrid mmaHidWorker = mmaGrid(batchSize, hidDim); // w1, w3 (replaces mmaFFNWorker) + +// Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // attnCast, woResid, w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); +// scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); +// scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); +// scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); +// scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); +// scheduler.addWorkerGrid(p + "batch_attention", attnWorker); +// scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); +// scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); +// scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); +// scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + scheduler.addWorkerGrid(p + "qProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "kProj", mmaKvWorker); + scheduler.addWorkerGrid(p + "vProj", mmaKvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "attnCast", ewDimWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "woResid", ewDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "batch_ffn_w1_mma", mmaHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_w3_mma", mmaHidWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } } From f3c6f562b4aa8efbd7129a937d3b32366708a579 Mon Sep 17 00:00:00 2001 From: MaryXek Date: Thu, 2 Jul 2026 15:40:09 +0300 Subject: [PATCH 10/18] Fuse QKV/GateUp MMA GEMMs, pipeline gemmMMA loads, and fix flash attention accumulation in FP16 batch prefill --- .../gpullama3/inference/state/State.java | 6 + .../TransformerBatchPrefillKernels.java | 692 +++++++++++++++++- .../prefill/LlamaFP16LayersBatchPrefill.java | 240 +++--- 3 files changed, 751 insertions(+), 187 deletions(-) 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 ae929e04..1441e417 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -93,6 +93,8 @@ public abstract class State { public final FloatArray woOut; public final HalfFloatArray wrapHbFP16Batch; public final FloatArray w2Out; + public final FloatArray qkvResultBatch; // B × (dim + 2*kvDim), packed [q|k|v] rows + public final FloatArray gateUpResultBatch; // B × 2*hiddenDim, packed [gate|up] rows protected State(Configuration config, int batchsize) { this.batchsize = batchsize; @@ -165,6 +167,8 @@ protected State(Configuration config, int batchsize) { this.woOut = new FloatArray(gpuBatchSize * config.dim()); this.wrapHbFP16Batch = new HalfFloatArray(gpuBatchSize * config.hiddenDim()); this.w2Out = new FloatArray(gpuBatchSize * config.dim()); + this.qkvResultBatch = new FloatArray(gpuBatchSize * (config.dim() + 2 * config.kvDim())); + this.gateUpResultBatch = new FloatArray(gpuBatchSize * 2 * config.hiddenDim()); } else { this.embeddingXBatch = null; this.wrapXBatch = null; @@ -185,6 +189,8 @@ protected State(Configuration config, int batchsize) { this.woOut = null; this.wrapHbFP16Batch = null; this.w2Out = null; + this.qkvResultBatch = null; + this.gateUpResultBatch = null; } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 3f99bb0e..0e9692e7 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -766,6 +766,29 @@ public static void batchedFFNSwiGLU(KernelContext context, private static final int WN = BN / WARPS_N; private static final int B_SUBTILE_BYTES = 256; + /** + * Packs two consecutive FP16 values into one int (lo | hi<<16) for the + * shared-memory ldmatrix tiles. Leaf helper; inlined by the Tornado JIT. + */ + private static int packHalves(HalfFloatArray src, int idxLo, int idxHi) { + int lo = src.get(idxLo).getHalfFloatValue() & 0xFFFF; + int hi = src.get(idxHi).getHalfFloatValue() & 0xFFFF; + return lo | (hi << 16); + } + + /** + * Tensor-core GEMM: C[M,N] (FP32) = A[M,K] (FP16, row-major) × B[N,K] (FP16, row-major). + * + *

Software-pipelined: each thread stages the NEXT K-step's A/B elements in + * registers while the CURRENT step's ldmatrix+MMA execute, so global-memory + * latency is hidden behind tensor-core compute. Shared memory stays + * single-buffered; the two block barriers per step preserve correctness + * (read-complete before overwrite, write-complete before next ldmatrix).

+ * + *

Requires M % 128 == 0, N % 128 == 0, K % 16 == 0, SM 8.0+.

+ * + * Worker: WorkerGrid2D((M/128)*256, N/128), local (256,1,1). + */ public static void gemmMMA(KernelContext ctx, HalfFloatArray A, HalfFloatArray B, FloatArray C, int M, int N, int K) { @@ -788,38 +811,48 @@ public static void gemmMMA(KernelContext ctx, float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + // ── Per-thread staging index math (constant across K-steps) ────────── + // A tile: BM*BK/2 = 1024 ints; layout idx = m_row*(BK/2) + k_pair. + // m_row = idx >>> 3, k = (idx & 7)*2; global A element (blockRow+m_row, kBase+k). + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + // B tile: BK*BN/2 = 1024 ints; subTileId = idx >>> 6, k_row = (idx & 63) >>> 2, + // col = subTileId*8 + (idx & 3)*2; B[col, k] at col*K + k (pair at +K). + int bIdx0 = tid; int gB0 = (blockCol + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (blockCol + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (blockCol + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (blockCol + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + // ── Prologue: stage K-step 0 ───────────────────────────────────────── + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0 = packHalves(B, gB0, gB0 + K); + int bReg1 = packHalves(B, gB1, gB1 + K); + int bReg2 = packHalves(B, gB2, gB2 + K); + int bReg3 = packHalves(B, gB3, gB3 + K); + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + int numKSteps = K / BK; for (int kStep = 0; kStep < numKSteps; kStep++) { - int kBase = kStep * BK; - // A load: A is [M, K] row-major (unchanged from gemmMMA) - for (int idx = tid; idx < BM * BK / 2; idx += THREADS_PER_BLOCK) { - int m_row = idx / (BK / 2); - int k_pair = idx % (BK / 2); - int k_base = k_pair * 2; - int gA = (blockRow + m_row) * K + (kBase + k_base); - int lo = A.get(gA).getHalfFloatValue() & 0xFFFF; - int hi = A.get(gA + 1).getHalfFloatValue() & 0xFFFF; - aTile[m_row * (BK / 2) + k_pair] = lo | (hi << 16); - } - // B load: B is [N, K] row-major. - // The shared-memory layout in bTile is unchanged; we read from - // a different global layout and pack into the same shared positions - // so that mmaLoadB, mma fragments, and stores all operate identically. - for (int idx = tid; idx < BK * BN / 2; idx += THREADS_PER_BLOCK) { - int subTileId = idx / 64; - int intInSub = idx % 64; - int k_row = intInSub / 4; - int j_pair = intInSub % 4; - int j_base = j_pair * 2; - int col_in_block = subTileId * 8 + j_base; - // B[col, k] is at col * K + k in [N, K] row-major. - int gB0 = (blockCol + col_in_block) * K + (kBase + k_row); - int gB1 = (blockCol + col_in_block + 1) * K + (kBase + k_row); - int lo = B.get(gB0).getHalfFloatValue() & 0xFFFF; - int hi = B.get(gB1).getHalfFloatValue() & 0xFFFF; - bTile[idx] = lo | (hi << 16); + // Issue next step's global loads FIRST: independent of the MMAs below, + // so their latency overlaps ldmatrix + tensor-core compute. + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + bReg0 = packHalves(B, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(B, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(B, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(B, gB3 + kOff, gB3 + kOff + K); } - ctx.localBarrier(); int aOff0 = warpM * 1024; int aOff1 = warpM * 1024 + 512; @@ -834,6 +867,13 @@ public static void gemmMMA(KernelContext ctx, HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); // all shared reads for this K-step complete + + // Overwrite shared tiles for the next step; overlaps the MMAs below. + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); @@ -851,7 +891,7 @@ public static void gemmMMA(KernelContext ctx, c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); - ctx.localBarrier(); + ctx.localBarrier(); // shared writes visible before next step's ldmatrix } int rBase = blockRow + warpM * WM; @@ -908,5 +948,597 @@ public static void batchedConvertFP32toFP16(KernelContext context, out.set(gid, new HalfFloat(in.get(gid))); } + + // ── Fused MMA projections ───────────────────────────────────────────────── + // + // Q, K, V (and gate/up) share the same A operand and the same K dimension, + // so they are fused into ONE kernel launch each. The N grid spans the packed + // output [dim | kvDim | kvDim] (resp. [hidDim | hidDim]); each thread block + // selects its weight matrix from groupIdy — a block-uniform branch, so there + // is zero divergence and no weight duplication in memory. This restores the + // A-reuse of the old fused matvec kernels AND fixes grid starvation for the + // skinny GQA projections (kvDim/128 blocks alone cannot fill the SMs). + + /** + * Fused QKV tensor-core GEMM into a PACKED output: + * qkvOut[M, dim+2*kvDim] = A[M,K] × [Wq | Wk | Wv] (each [N_i, K] row-major). + * + *

Layout of a row of qkvOut: [ q(0..dim) | k(0..kvDim) | v(0..kvDim) ]. + * Requires dim % 128 == 0 and kvDim % 128 == 0.

+ * + * Worker: WorkerGrid2D((M/128)*256, (dim+2*kvDim)/128), local (256,1,1). + */ + public static void gemmMMAQKV(KernelContext ctx, + HalfFloatArray A, + HalfFloatArray wq, HalfFloatArray wk, HalfFloatArray wv, + FloatArray qkvOut, + int M, int dim, int kvDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; // column in the packed output + int qkvStride = dim + 2 * kvDim; + + // Column base inside the segment's own weight matrix (block-uniform). + int wColBase = blockCol; + if (blockCol >= dim) wColBase -= dim; + if (blockCol >= dim + kvDim) wColBase -= kvDim; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int gB0 = (wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + // ── Prologue: stage K-step 0 ── + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < dim) { + bReg0 = packHalves(wq, gB0, gB0 + K); + bReg1 = packHalves(wq, gB1, gB1 + K); + bReg2 = packHalves(wq, gB2, gB2 + K); + bReg3 = packHalves(wq, gB3, gB3 + K); + } else if (blockCol < dim + kvDim) { + bReg0 = packHalves(wk, gB0, gB0 + K); + bReg1 = packHalves(wk, gB1, gB1 + K); + bReg2 = packHalves(wk, gB2, gB2 + K); + bReg3 = packHalves(wk, gB3, gB3 + K); + } else { + bReg0 = packHalves(wv, gB0, gB0 + K); + bReg1 = packHalves(wv, gB1, gB1 + K); + bReg2 = packHalves(wv, gB2, gB2 + K); + bReg3 = packHalves(wv, gB3, gB3 + K); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < dim) { + bReg0 = packHalves(wq, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wq, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wq, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wq, gB3 + kOff, gB3 + kOff + K); + } else if (blockCol < dim + kvDim) { + bReg0 = packHalves(wk, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wk, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wk, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wk, gB3 + kOff, gB3 + kOff + K); + } else { + bReg0 = packHalves(wv, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wv, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wv, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wv, gB3 + kOff, gB3 + kOff + K); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + // Stores are uniform: packed column == blockCol-relative column, + // stride is the packed row width. + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, qkvOut, rBase + 0, cBase + 0, qkvStride); + ctx.mmaStore(c01, qkvOut, rBase + 0, cBase + 8, qkvStride); + ctx.mmaStore(c02, qkvOut, rBase + 0, cBase + 16, qkvStride); + ctx.mmaStore(c03, qkvOut, rBase + 0, cBase + 24, qkvStride); + ctx.mmaStore(c04, qkvOut, rBase + 0, cBase + 32, qkvStride); + ctx.mmaStore(c05, qkvOut, rBase + 0, cBase + 40, qkvStride); + ctx.mmaStore(c06, qkvOut, rBase + 0, cBase + 48, qkvStride); + ctx.mmaStore(c07, qkvOut, rBase + 0, cBase + 56, qkvStride); + ctx.mmaStore(c10, qkvOut, rBase + 16, cBase + 0, qkvStride); + ctx.mmaStore(c11, qkvOut, rBase + 16, cBase + 8, qkvStride); + ctx.mmaStore(c12, qkvOut, rBase + 16, cBase + 16, qkvStride); + ctx.mmaStore(c13, qkvOut, rBase + 16, cBase + 24, qkvStride); + ctx.mmaStore(c14, qkvOut, rBase + 16, cBase + 32, qkvStride); + ctx.mmaStore(c15, qkvOut, rBase + 16, cBase + 40, qkvStride); + ctx.mmaStore(c16, qkvOut, rBase + 16, cBase + 48, qkvStride); + ctx.mmaStore(c17, qkvOut, rBase + 16, cBase + 56, qkvStride); + } + + /** + * Fused W1/W3 (gate/up) tensor-core GEMM into a PACKED output: + * gateUpOut[M, 2*hidDim] = A[M,K] × [W1 | W3] (each [hidDim, K] row-major). + * + *

Layout of a row: [ gate(0..hidDim) | up(0..hidDim) ]. + * Requires hidDim % 128 == 0.

+ * + * Worker: WorkerGrid2D((M/128)*256, (2*hidDim)/128), local (256,1,1). + */ + public static void gemmMMAGateUp(KernelContext ctx, + HalfFloatArray A, + HalfFloatArray w1, HalfFloatArray w3, + FloatArray gateUpOut, + int M, int hidDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; // column in the packed output + int outStride = 2 * hidDim; + + int wColBase = (blockCol < hidDim) ? blockCol : (blockCol - hidDim); + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int gB0 = (wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < hidDim) { + bReg0 = packHalves(w1, gB0, gB0 + K); + bReg1 = packHalves(w1, gB1, gB1 + K); + bReg2 = packHalves(w1, gB2, gB2 + K); + bReg3 = packHalves(w1, gB3, gB3 + K); + } else { + bReg0 = packHalves(w3, gB0, gB0 + K); + bReg1 = packHalves(w3, gB1, gB1 + K); + bReg2 = packHalves(w3, gB2, gB2 + K); + bReg3 = packHalves(w3, gB3, gB3 + K); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < hidDim) { + bReg0 = packHalves(w1, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(w1, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(w1, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(w1, gB3 + kOff, gB3 + kOff + K); + } else { + bReg0 = packHalves(w3, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(w3, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(w3, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(w3, gB3 + kOff, gB3 + kOff + K); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, gateUpOut, rBase + 0, cBase + 0, outStride); + ctx.mmaStore(c01, gateUpOut, rBase + 0, cBase + 8, outStride); + ctx.mmaStore(c02, gateUpOut, rBase + 0, cBase + 16, outStride); + ctx.mmaStore(c03, gateUpOut, rBase + 0, cBase + 24, outStride); + ctx.mmaStore(c04, gateUpOut, rBase + 0, cBase + 32, outStride); + ctx.mmaStore(c05, gateUpOut, rBase + 0, cBase + 40, outStride); + ctx.mmaStore(c06, gateUpOut, rBase + 0, cBase + 48, outStride); + ctx.mmaStore(c07, gateUpOut, rBase + 0, cBase + 56, outStride); + ctx.mmaStore(c10, gateUpOut, rBase + 16, cBase + 0, outStride); + ctx.mmaStore(c11, gateUpOut, rBase + 16, cBase + 8, outStride); + ctx.mmaStore(c12, gateUpOut, rBase + 16, cBase + 16, outStride); + ctx.mmaStore(c13, gateUpOut, rBase + 16, cBase + 24, outStride); + ctx.mmaStore(c14, gateUpOut, rBase + 16, cBase + 32, outStride); + ctx.mmaStore(c15, gateUpOut, rBase + 16, cBase + 40, outStride); + ctx.mmaStore(c16, gateUpOut, rBase + 16, cBase + 48, outStride); + ctx.mmaStore(c17, gateUpOut, rBase + 16, cBase + 56, outStride); + } + + // ── Parallel RMS reductions ─────────────────────────────────────────────── + // Replace the localSize=1 sequential reductions (one thread walking `dim` + // elements alone) with one 256-thread workgroup per token and a shared-memory + // tree reduction. + + /** + * Parallel RMS square-sum reduction. One workgroup per batch token. + * + * Worker: B workgroups × localSize threads (localSize=256). + */ + public static void batchedRmsReduceParallel(KernelContext context, + FloatArray wrapXBatch, + FloatArray scaleBatch, + int dim, float eps, int localSize) { + int tid = context.localIdx; + int b = context.groupIdx; + int localSz = context.localGroupSizeX; + float[] partial = context.allocateFloatLocalArray(localSize); + + int base = b * dim; + float ss = 0.0f; + for (int i = tid; i < dim; i += localSz) { + float v = wrapXBatch.get(base + i); + ss += v * v; + } + partial[tid] = ss; + context.localBarrier(); + for (int s = localSz / 2; s > 0; s >>= 1) { + if (tid < s) { + partial[tid] += partial[tid + s]; + } + context.localBarrier(); + } + if (tid == 0) { + float m = partial[0] / dim + eps; + scaleBatch.set(b, 1.0f / TornadoMath.sqrt(m)); + } + } + + /** + * Parallel RMS reduction FUSED with the pending residual add: + * x[b,i] += delta[b,i] first, then square-sum over the updated row. + * Replaces the separate woResid task + FFN RMS reduce (each element is + * visited exactly once, so the in-place update is race-free). + * + * Worker: B workgroups × localSize threads (localSize=256). + */ + public static void batchedRmsReduceFusedResidual(KernelContext context, + FloatArray wrapXBatch, + FloatArray delta, + FloatArray scaleBatch, + int dim, float eps, int localSize) { + int tid = context.localIdx; + int b = context.groupIdx; + int localSz = context.localGroupSizeX; + float[] partial = context.allocateFloatLocalArray(localSize); + + int base = b * dim; + float ss = 0.0f; + for (int i = tid; i < dim; i += localSz) { + float v = wrapXBatch.get(base + i) + delta.get(base + i); + wrapXBatch.set(base + i, v); + ss += v * v; + } + partial[tid] = ss; + context.localBarrier(); + for (int s = localSz / 2; s > 0; s >>= 1) { + if (tid < s) { + partial[tid] += partial[tid + s]; + } + context.localBarrier(); + } + if (tid == 0) { + float m = partial[0] / dim + eps; + scaleBatch.set(b, 1.0f / TornadoMath.sqrt(m)); + } + } + + // ── RoPE + KV cache over the packed QKV buffer ──────────────────────────── + + /** + * Fused batched RoPE rotation + KV cache write, reading/writing the PACKED + * QKV buffer produced by {@link #gemmMMAQKV}. Row layout: [ q | k | v ]. + * Q is rotated in place (consumed from the packed buffer by attention). + * + * Worker: B*(dim/2) global threads, localSize=512 (or less). + */ + public static void batchedRopeWithKVCachePacked(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + int kvDim, int headSize, + int layerIndex, int contextLength, int dim) { + int globalIdx = context.globalIdx; + int halfDim = dim / 2; + int batchIdx = globalIdx / halfDim; + int pairIdx = globalIdx % halfDim; + int i = pairIdx * 2; + int qkvStride = dim + 2 * kvDim; + + int pos = batchStartPosHolder.get(0) + batchIdx; + int qOffset = batchIdx * qkvStride; + int kOffset = batchIdx * qkvStride + dim; + int vOffset = batchIdx * qkvStride + dim + kvDim; + + if (i + 1 < dim) { + int head_dim = i % headSize; + float freq = 1.0f / TornadoMath.pow(50000.0f, head_dim / (float) headSize); + float val = pos * freq; + float fcr = TornadoMath.cos(val); + float fci = TornadoMath.sin(val); + + // Rotate Q in place + float v0q = qkvBatch.get(qOffset + i); + float v1q = qkvBatch.get(qOffset + i + 1); + qkvBatch.set(qOffset + i, v0q * fcr - v1q * fci); + qkvBatch.set(qOffset + i + 1, v0q * fci + v1q * fcr); + + // Rotate K and write K,V to cache + if (i + 1 < kvDim) { + float v0k = qkvBatch.get(kOffset + i); + float v1k = qkvBatch.get(kOffset + i + 1); + float rotK0 = v0k * fcr - v1k * fci; + float rotK1 = v0k * fci + v1k * fcr; + + int cacheOff = layerIndex * contextLength * kvDim + pos * kvDim; + wrapKeyCache.set(cacheOff + i, rotK0); + wrapKeyCache.set(cacheOff + i + 1, rotK1); + wrapValueCache.set(cacheOff + i, qkvBatch.get(vOffset + i)); + wrapValueCache.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); + } + } + } + + // ── Flash attention (fixed accumulation, FP16 output) ──────────────────── + + /** + * Batched causal flash attention over the packed QKV buffer, writing FP16 + * directly (the A operand of the Wo GEMM — eliminates the attnCast pass). + * + *

Fixes the redundant accumulation of the previous version: each thread + * now OWNS output dims {tid, tid+localSz} and accumulates them in registers, + * instead of every thread redundantly computing the full headSize output + * vector into a (spilled) private array. K/V tile loads are flattened over + * (t, d) so consecutive threads issue coalesced reads.

+ * + *

Requires headSize <= 2*localSz (localSz = min(headSize, 128)).

+ * + * Worker: B*nHeads workgroups × min(headSize,128) threads. + */ + public static void batchedFlashAttentionFP16Out(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + HalfFloatArray attnOutFP16, + int nHeads, int headSize, + int kvDim, int kvMul, + int layerIndex, int contextLength, int dim) { + int tid = context.localIdx; + int groupId = context.groupIdx; + int localSz = context.localGroupSizeX; + + int batchIdx = groupId / nHeads; + int h = groupId % nHeads; + int pos = batchStartPosHolder.get(0) + batchIdx; + int loff = layerIndex * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + int BLOCK_C = 16; + int qkvStride = dim + 2 * kvDim; + + float[] qShared = context.allocateFloatLocalArray(headSize); + float[] kTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] vTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] sTile = context.allocateFloatLocalArray(BLOCK_C); + + // Load Q (rotated, from the packed QKV buffer) into shared memory + int qOffset = batchIdx * qkvStride + h * headSize; + for (int i = tid; i < headSize; i += localSz) { + qShared[i] = qkvBatch.get(qOffset + i); + } + context.localBarrier(); + + float maxScore = Float.NEGATIVE_INFINITY; + float sumExp = 0.0f; + // Each thread owns output dims d0 = tid and (if headSize > localSz) d1. + float acc0 = 0.0f; + float acc1 = 0.0f; + int d1 = tid + localSz; + + for (int tileC = 0; tileC <= pos; tileC += BLOCK_C) { + int tileEnd = Math.min(tileC + BLOCK_C - 1, pos); + int tileLen = tileEnd - tileC + 1; + + // Load K/V tile — flattened over (t, d) for coalescing + for (int idx = tid; idx < tileLen * headSize; idx += localSz) { + int tInTile = idx / headSize; + int d = idx % headSize; + int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + d; + kTile[tInTile * headSize + d] = wrapKeyCache.get(kvOff); + vTile[tInTile * headSize + d] = wrapValueCache.get(kvOff); + } + context.localBarrier(); + + // Scores: one thread per key position in the tile + for (int t = tileC + tid; t <= tileEnd; t += localSz) { + int tInTile = t - tileC; + float score = 0.0f; + for (int d = 0; d < headSize; d++) { + score += qShared[d] * kTile[tInTile * headSize + d]; + } + sTile[tInTile] = score / TornadoMath.sqrt(headSize); + } + context.localBarrier(); + + // Tile max: redundant per-thread scan over <= 16 shared values — + // deterministic across the workgroup, no broadcast needed. + float tileMax = Float.NEGATIVE_INFINITY; + for (int t = 0; t < tileLen; t++) { + if (sTile[t] > tileMax) { + tileMax = sTile[t]; + } + } + + float newMax = Math.max(maxScore, tileMax); + if (maxScore != Float.NEGATIVE_INFINITY && newMax != maxScore) { + float corr = TornadoMath.exp(maxScore - newMax); + sumExp *= corr; + acc0 *= corr; + acc1 *= corr; + } + maxScore = newMax; + + for (int t = 0; t < tileLen; t++) { + float p = TornadoMath.exp(sTile[t] - maxScore); + sumExp += p; + acc0 += p * vTile[t * headSize + tid]; + if (d1 < headSize) { + acc1 += p * vTile[t * headSize + d1]; + } + } + context.localBarrier(); + } + + float norm = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + int outOffset = batchIdx * dim + h * headSize; + attnOutFP16.set(outOffset + tid, new HalfFloat(acc0 * norm)); + if (d1 < headSize) { + attnOutFP16.set(outOffset + d1, new HalfFloat(acc1 * norm)); + } + } + + // ── SwiGLU over the packed gate/up buffer, emitting FP16 ───────────────── + + /** + * Fused SiLU(gate) * up over the PACKED [gate | up] GEMM output, + * emitting FP16 (the A operand of the W2 GEMM). + * + * Worker: B*hiddenDim global threads, localSize=256. + */ + public static void batchedFFNSwiGLUFP16Packed(KernelContext context, + HalfFloatArray wrapHbFP16Batch, + FloatArray gateUpResult, + int hiddenDim) { + int gid = context.globalIdx; + int b = gid / hiddenDim; + int i = gid % hiddenDim; + int rowBase = b * 2 * hiddenDim; + float g = gateUpResult.get(rowBase + i); + float u = gateUpResult.get(rowBase + hiddenDim + i); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbFP16Batch.set(gid, new HalfFloat(silu * u)); + } + // @formatter:on } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java index f590b7b9..2e7a4705 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java @@ -25,13 +25,29 @@ *

One {@link ImmutableTaskGraph} per transformer layer, each processing * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

* + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device * after every layer so the subsequent single-token decode layers can consume it.

*/ public class LlamaFP16LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - // Matches the local workgroup size used by the single-token kernels. - static final int LOCAL_WORK_GROUP_SIZE = 32; + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; private final LlamaState state; private final LlamaTornadoWeights weights; @@ -69,11 +85,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { context, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.ffnGateResult, state.ffnUpResult, + state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); // wrapXBatch produced by the prefillActivation graph and persists in device memory // to consume it from there we should use the explicit uniqueTaskGraph name @@ -88,11 +102,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.batchStartPosHolder, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.ffnGateResult, state.ffnUpResult, + state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } @@ -114,9 +126,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { // ── Attention Block ──────────────────────────────────────────────────── batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, + TransformerBatchPrefillKernels::batchedRmsReduceParallel, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); batchPrefillLayer.task("batch_attn_rms_apply", TransformerBatchPrefillKernels::batchedRmsApplyFP16, @@ -124,63 +136,45 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); -// batchPrefillLayer.task("batch_qkv", -// TransformerBatchPrefillKernels::batchedFusedQKVMatmul, -// context, -// state.wrapXbFP16Batch, -// state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, -// weights.wqLayered[layerIndex].asHalfFloatArray(), -// weights.wkLayered[layerIndex].asHalfFloatArray(), -// weights.wvLayered[layerIndex].asHalfFloatArray(), -// dim, kvDim, LOCAL_WORK_GROUP_SIZE); - - batchPrefillLayer.task("qProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.wrapXbFP16Batch, - weights.wqLayered[layerIndex].asHalfFloatArray(), - state.wrapQBatch, batchSize, dim, dim) - .task("kProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.wrapXbFP16Batch, - weights.wkLayered[layerIndex].asHalfFloatArray(), - state.wrapKBatch, batchSize, kvDim, dim) - .task("vProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.wrapXbFP16Batch, - weights.wvLayered[layerIndex].asHalfFloatArray(), - state.wrapVBatch, batchSize, kvDim, dim); + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + state.qkvResultBatch, batchSize, dim, kvDim, dim); batchPrefillLayer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCache, + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); -// batchPrefillLayer.task("batch_attn_out", -// TransformerBatchPrefillKernels::batchedMatVecWithResidual, -// context, state.wrapXbBatch, state.wrapXBatch, -// weights.woLayered[layerIndex].asHalfFloatArray(), -// dim, dim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("attnCast", TransformerBatchPrefillKernels::batchedConvertFP32toFP16, - context, state.wrapXbBatch, state.attnOutFP16) - .task("woProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.attnOutFP16, - weights.woLayered[layerIndex].asHalfFloatArray(), - state.woOut, batchSize, dim, dim) - .task("woResid", TransformerBatchPrefillKernels::batchedResidualAddFP32, - context, state.wrapXBatch, state.woOut); + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asHalfFloatArray(), + state.woOut, batchSize, dim, dim); // ── FFN Block ────────────────────────────────────────────────────────── + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); batchPrefillLayer.task("batch_ffn_rms_apply", TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, @@ -188,43 +182,17 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), state.ffnScaleBatch, dim); - batchPrefillLayer.task("batch_ffn_w1_mma", - TransformerBatchPrefillKernels::gemmMMA, + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asHalfFloatArray(), - state.ffnGateResult, - batchSize, hidDim, dim); - - batchPrefillLayer.task("batch_ffn_w3_mma", - TransformerBatchPrefillKernels::gemmMMA, - context, state.normedXFFNFP16, weights.w3Layered[layerIndex].asHalfFloatArray(), - state.ffnUpResult, - batchSize, hidDim, dim); - - -// batchPrefillLayer.task("batch_ffn_swiglu", -// TransformerBatchPrefillKernels::batchedFFNSwiGLU, -// context, state.wrapHbBatch, state.ffnGateResult, state.ffnUpResult, -// hidDim); - - -// batchPrefillLayer.task("batch_ffn_gate_up", -// TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, -// context, state.wrapXBatch, state.wrapHbBatch, -// weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), -// state.ffnScaleBatch, -// weights.w1Layered[layerIndex].asHalfFloatArray(), -// weights.w3Layered[layerIndex].asHalfFloatArray(), -// dim, hidDim, LOCAL_WORK_GROUP_SIZE); - -// batchPrefillLayer.task("batch_ffn_down", -// TransformerBatchPrefillKernels::batchedMatVecWithResidual, -// context, state.wrapHbBatch, state.wrapXBatch, -// weights.w2Layered[layerIndex].asHalfFloatArray(), -// hidDim, dim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("swiglu", TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16, - context, state.wrapHbFP16Batch, state.ffnGateResult, state.ffnUpResult, hidDim) + state.gateUpResultBatch, batchSize, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, context, state.wrapHbFP16Batch, weights.w2Layered[layerIndex].asHalfFloatArray(), @@ -240,7 +208,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { } // @formatter:on - // gemmMMA: 256 threads/block (1D within block), grid over M- and N-blocks. + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. static WorkerGrid mmaGrid(int paddedM, int N) { int mBlocks = paddedM / 128; // BM int nBlocks = N / 128; // BN @@ -263,16 +231,13 @@ public void updateGridScheduler(GridScheduler scheduler) { int nHeads = config.numberOfHeads(); int headSz = config.headSize(); - // RMS: one thread per batch token - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - // QKV: B*(dim+2*kvDim) workgroups × LOCAL_WORK_GROUP_SIZE - int qkvRows = dim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); // RoPE+KV cache: B*(dim/2) threads, local=512 int ropeGlobal = batchSize * (dim / 2); @@ -280,75 +245,36 @@ public void updateGridScheduler(GridScheduler scheduler) { while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - // Attention (flash): B*nHeads workgroups × optimalLocalSize - int optLocal = findOptimalLocalSize(headSz); + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * optLocal, optLocal); - - // Mat-vec (Wo, W2): B*d workgroups × LOCAL_WORK_GROUP_SIZE - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - // FFN RMS apply: B*dim threads, local=256 - WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - // MMA: 128x128 block tile, 256 threads/block, M=batchSize, N=hiddenDim, K=dim - // Global = (M/128)*256 in X, (N/128) in Y, local (256,1,1) - WorkerGrid mmaFFNWorker = new WorkerGrid2D((batchSize / 128) * 256, hidDim / 128); - mmaFFNWorker.setLocalWork(256, 1, 1); + batchSize * nHeads * attnLocal, attnLocal); - // SwiGLU: B*hiddenDim threads, local=256 - WorkerGrid swigluWorker = WorkerGridFactory.genericWorker(batchSize * hidDim, 256); + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(batchSize, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(batchSize, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(batchSize, 2 * hidDim); // fused W1/W3 - WorkerGrid mmaDimWorker = mmaGrid(batchSize, dim); // qProj, woProj, w2Proj - WorkerGrid mmaKvWorker = mmaGrid(batchSize, kvDim); // kProj, vProj - WorkerGrid mmaHidWorker = mmaGrid(batchSize, hidDim); // w1, w3 (replaces mmaFFNWorker) - -// Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) - WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // attnCast, woResid, w2Resid + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; -// scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); -// scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); -// scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); -// scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); -// scheduler.addWorkerGrid(p + "batch_attention", attnWorker); -// scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); -// scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); -// scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); -// scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "qProj", mmaDimWorker); - scheduler.addWorkerGrid(p + "kProj", mmaKvWorker); - scheduler.addWorkerGrid(p + "vProj", mmaKvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "attnCast", ewDimWorker); - scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); - scheduler.addWorkerGrid(p + "woResid", ewDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_ffn_w1_mma", mmaHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_w3_mma", mmaHidWorker); - scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); - scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); - scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } - return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } From f700c3c96dcced67afee99c83bc91053bab092bd Mon Sep 17 00:00:00 2001 From: MaryXek Date: Mon, 6 Jul 2026 16:54:17 +0300 Subject: [PATCH 11/18] Apply padding if batch size is not multiple of 128 --- .../gpullama3/inference/state/State.java | 21 +++++++++++------- .../prefill/LlamaFP16LayersBatchPrefill.java | 22 ++++++++++++++----- 2 files changed, 29 insertions(+), 14 deletions(-) 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 1441e417..e3092291 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -145,11 +145,16 @@ protected State(Configuration config, int batchsize) { int gpuBatchSize = Integer.getInteger("llama.prefillBatchSize", 1); if (gpuBatchSize > 1) { + // The tensor-core GEMM kernels operate on full 128-row M tiles + // (BM = 128). Pad the GEMM-adjacent activation buffers so any + // batch size launches whole tiles; rows >= gpuBatchSize hold + // garbage and are never read by the non-GEMM kernels. + int paddedGpuBatch = (gpuBatchSize + 127) & ~127; int qDim = batchQDim(config); int kvDim = batchKvDim(config); this.embeddingXBatch = new HalfFloatArray(gpuBatchSize * config.dim()); this.wrapXBatch = new FloatArray(gpuBatchSize * config.dim()); - this.wrapXbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); + this.wrapXbFP16Batch = new HalfFloatArray(paddedGpuBatch * config.dim()); this.wrapQBatch = new FloatArray(gpuBatchSize * qDim); this.wrapKBatch = new FloatArray(gpuBatchSize * kvDim); this.wrapVBatch = new FloatArray(gpuBatchSize * kvDim); @@ -158,17 +163,17 @@ protected State(Configuration config, int batchsize) { this.attnScaleBatch = new FloatArray(gpuBatchSize); this.ffnScaleBatch = new FloatArray(gpuBatchSize); this.batchStartPosHolder = new IntArray(1); - this.normedXFFNFP16 = new HalfFloatArray(gpuBatchSize * config.dim()); + this.normedXFFNFP16 = new HalfFloatArray(paddedGpuBatch * config.dim()); this.ffnGateResult = new FloatArray(gpuBatchSize * config.hiddenDim()); this.ffnUpResult = new FloatArray(gpuBatchSize * config.hiddenDim()); this.xbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); - this.attnOutFP16 = new HalfFloatArray(gpuBatchSize * config.dim()); - this.woOut = new FloatArray(gpuBatchSize * config.dim()); - this.wrapHbFP16Batch = new HalfFloatArray(gpuBatchSize * config.hiddenDim()); - this.w2Out = new FloatArray(gpuBatchSize * config.dim()); - this.qkvResultBatch = new FloatArray(gpuBatchSize * (config.dim() + 2 * config.kvDim())); - this.gateUpResultBatch = new FloatArray(gpuBatchSize * 2 * config.hiddenDim()); + this.attnOutFP16 = new HalfFloatArray(paddedGpuBatch * config.dim()); + this.woOut = new FloatArray(paddedGpuBatch * config.dim()); + this.wrapHbFP16Batch = new HalfFloatArray(paddedGpuBatch * config.hiddenDim()); + this.w2Out = new FloatArray(paddedGpuBatch * config.dim()); + this.qkvResultBatch = new FloatArray(paddedGpuBatch * (config.dim() + 2 * config.kvDim())); + this.gateUpResultBatch = new FloatArray(paddedGpuBatch * 2 * config.hiddenDim()); } else { this.embeddingXBatch = null; this.wrapXBatch = null; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java index 2e7a4705..fa04c0b8 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java @@ -54,6 +54,10 @@ public class LlamaFP16LayersBatchPrefill implements BatchPrefillTransformerLayer private final LlamaConfiguration config; private final KernelContext context = new KernelContext(); private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; private final List layerITGs; private String lastLayerTaskGraphID; @@ -63,6 +67,12 @@ public LlamaFP16LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) .map(TaskGraph::snapshot) @@ -145,7 +155,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), weights.wvLayered[layerIndex].asHalfFloatArray(), - state.qkvResultBatch, batchSize, dim, kvDim, dim); + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); batchPrefillLayer.task("batch_rope_kv", TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, @@ -167,7 +177,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, context, state.attnOutFP16, weights.woLayered[layerIndex].asHalfFloatArray(), - state.woOut, batchSize, dim, dim); + state.woOut, paddedBatch, dim, dim); // ── FFN Block ────────────────────────────────────────────────────────── // x += woOut is fused into the FFN RMS reduction (drops the woResid task). @@ -188,7 +198,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asHalfFloatArray(), weights.w3Layered[layerIndex].asHalfFloatArray(), - state.gateUpResultBatch, batchSize, hidDim, dim); + state.gateUpResultBatch, paddedBatch, hidDim, dim); batchPrefillLayer.task("swiglu", TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, @@ -252,9 +262,9 @@ public void updateGridScheduler(GridScheduler scheduler) { batchSize * nHeads * attnLocal, attnLocal); // MMA grids - WorkerGrid mmaQkvWorker = mmaGrid(batchSize, dim + 2 * kvDim); // fused QKV - WorkerGrid mmaDimWorker = mmaGrid(batchSize, dim); // woProj, w2Proj - WorkerGrid mmaGateUpWorker = mmaGrid(batchSize, 2 * hidDim); // fused W1/W3 + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid From 4b26422663fc3fabc8ce3bb0e2f66905994e50cd Mon Sep 17 00:00:00 2001 From: MaryXek Date: Mon, 6 Jul 2026 17:12:25 +0300 Subject: [PATCH 12/18] Add Q8_0 W8A16 tensor-core batch prefill: dequantize weights to FP16 in GEMM staging --- .../TransformerBatchPrefillKernels.java | 451 ++++++++++++++++++ .../prefill/LlamaQ8_0LayersBatchPrefill.java | 314 +++++++----- 2 files changed, 637 insertions(+), 128 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 0e9692e7..9d85c936 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -1540,5 +1540,456 @@ public static void batchedFFNSwiGLUFP16Packed(KernelContext context, wrapHbFP16Batch.set(gid, new HalfFloat(silu * u)); } + + // ── Q8_0 tensor-core GEMMs (W8A16) ─────────────────────────────────────── + // + // Q8_0 weights stay quantized in global memory (34-byte GGUF blocks: + // FP16 scale + 32 int8 quants) and are dequantized to FP16 *in the + // register-staging step* of the software pipeline, then flow through the + // identical ldmatrix + m16n8k16 path as the FP16 GEMMs. This halves the + // weight-side memory traffic relative to FP16 while reusing the proven + // FP16 tensor-core pipeline. A true INT8 (m16n8k32) path would need + // per-k-block accumulator rescaling — blocked on fragment-level access + // in the TornadoVM intrinsics; see paper future work. + // + // Note: BK = 16, so a K-step never straddles a Q8_0 block boundary + // (32 | K), and each staged pair reads exactly one scale per column. + + /** + * Dequantizes and packs two vertically-adjacent (col, col+1) Q8_0 weight + * elements at depth k into one int of two FP16 values, matching the + * bTile layout expected by mmaLoadB. Leaf helper; inlined by the JIT. + */ + private static int packQ8Halves(ByteArray w, int col, int k, int blocksPerRow) { + int kBlock = k >>> 5; // k / 32 + int kIn = k & 31; // k % 32 + int off0 = (col * blocksPerRow + kBlock) * 34; + int off1 = off0 + blocksPerRow * 34; // column col+1, same k-block + float v0 = w.getHalfFloat(off0).getFloat32() * w.get(off0 + 2 + kIn); + float v1 = w.getHalfFloat(off1).getFloat32() * w.get(off1 + 2 + kIn); + int lo = new HalfFloat(v0).getHalfFloatValue() & 0xFFFF; + int hi = new HalfFloat(v1).getHalfFloatValue() & 0xFFFF; + return lo | (hi << 16); + } + + /** + * Tensor-core GEMM with Q8_0 weights: + * C[M,N] (FP32) = A[M,K] (FP16) × B[N,K] (Q8_0 blocks, row-major). + * Same tiling, pipeline, and constraints as {@link #gemmMMA}. + */ + public static void gemmMMAQ8(KernelContext ctx, + HalfFloatArray A, ByteArray B, FloatArray C, + int M, int N, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int blocksPerRow = K / 32; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + // B staging keeps (col, k) coordinates explicit for block-offset math. + int bIdx0 = tid; int bCol0 = blockCol + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = blockCol + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = blockCol + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = blockCol + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0 = packQ8Halves(B, bCol0, bK0, blocksPerRow); + int bReg1 = packQ8Halves(B, bCol1, bK1, blocksPerRow); + int bReg2 = packQ8Halves(B, bCol2, bK2, blocksPerRow); + int bReg3 = packQ8Halves(B, bCol3, bK3, blocksPerRow); + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + bReg0 = packQ8Halves(B, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(B, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(B, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(B, bCol3, kOff + bK3, blocksPerRow); + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, C, rBase + 0, cBase + 0, N); + ctx.mmaStore(c01, C, rBase + 0, cBase + 8, N); + ctx.mmaStore(c02, C, rBase + 0, cBase + 16, N); + ctx.mmaStore(c03, C, rBase + 0, cBase + 24, N); + ctx.mmaStore(c04, C, rBase + 0, cBase + 32, N); + ctx.mmaStore(c05, C, rBase + 0, cBase + 40, N); + ctx.mmaStore(c06, C, rBase + 0, cBase + 48, N); + ctx.mmaStore(c07, C, rBase + 0, cBase + 56, N); + ctx.mmaStore(c10, C, rBase + 16, cBase + 0, N); + ctx.mmaStore(c11, C, rBase + 16, cBase + 8, N); + ctx.mmaStore(c12, C, rBase + 16, cBase + 16, N); + ctx.mmaStore(c13, C, rBase + 16, cBase + 24, N); + ctx.mmaStore(c14, C, rBase + 16, cBase + 32, N); + ctx.mmaStore(c15, C, rBase + 16, cBase + 40, N); + ctx.mmaStore(c16, C, rBase + 16, cBase + 48, N); + ctx.mmaStore(c17, C, rBase + 16, cBase + 56, N); + } + + /** + * Fused QKV tensor-core GEMM with Q8_0 weights into the PACKED output: + * qkvOut[M, dim+2*kvDim] = A[M,K] (FP16) × [Wq | Wk | Wv] (Q8_0, [N_i,K] row-major). + * Same layout, fusion, and constraints as {@link #gemmMMAQKV}. + */ + public static void gemmMMAQKVQ8(KernelContext ctx, + HalfFloatArray A, + ByteArray wq, ByteArray wk, ByteArray wv, + FloatArray qkvOut, + int M, int dim, int kvDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int qkvStride = dim + 2 * kvDim; + int blocksPerRow = K / 32; + + int wColBase = blockCol; + if (blockCol >= dim) wColBase -= dim; + if (blockCol >= dim + kvDim) wColBase -= kvDim; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int bCol0 = wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < dim) { + bReg0 = packQ8Halves(wq, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wq, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wq, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wq, bCol3, bK3, blocksPerRow); + } else if (blockCol < dim + kvDim) { + bReg0 = packQ8Halves(wk, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wk, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wk, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wk, bCol3, bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(wv, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wv, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wv, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wv, bCol3, bK3, blocksPerRow); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < dim) { + bReg0 = packQ8Halves(wq, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wq, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wq, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wq, bCol3, kOff + bK3, blocksPerRow); + } else if (blockCol < dim + kvDim) { + bReg0 = packQ8Halves(wk, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wk, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wk, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wk, bCol3, kOff + bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(wv, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wv, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wv, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wv, bCol3, kOff + bK3, blocksPerRow); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, qkvOut, rBase + 0, cBase + 0, qkvStride); + ctx.mmaStore(c01, qkvOut, rBase + 0, cBase + 8, qkvStride); + ctx.mmaStore(c02, qkvOut, rBase + 0, cBase + 16, qkvStride); + ctx.mmaStore(c03, qkvOut, rBase + 0, cBase + 24, qkvStride); + ctx.mmaStore(c04, qkvOut, rBase + 0, cBase + 32, qkvStride); + ctx.mmaStore(c05, qkvOut, rBase + 0, cBase + 40, qkvStride); + ctx.mmaStore(c06, qkvOut, rBase + 0, cBase + 48, qkvStride); + ctx.mmaStore(c07, qkvOut, rBase + 0, cBase + 56, qkvStride); + ctx.mmaStore(c10, qkvOut, rBase + 16, cBase + 0, qkvStride); + ctx.mmaStore(c11, qkvOut, rBase + 16, cBase + 8, qkvStride); + ctx.mmaStore(c12, qkvOut, rBase + 16, cBase + 16, qkvStride); + ctx.mmaStore(c13, qkvOut, rBase + 16, cBase + 24, qkvStride); + ctx.mmaStore(c14, qkvOut, rBase + 16, cBase + 32, qkvStride); + ctx.mmaStore(c15, qkvOut, rBase + 16, cBase + 40, qkvStride); + ctx.mmaStore(c16, qkvOut, rBase + 16, cBase + 48, qkvStride); + ctx.mmaStore(c17, qkvOut, rBase + 16, cBase + 56, qkvStride); + } + + /** + * Fused W1/W3 (gate/up) tensor-core GEMM with Q8_0 weights into the PACKED + * output gateUpOut[M, 2*hidDim]. Same layout and constraints as + * {@link #gemmMMAGateUp}. + */ + public static void gemmMMAGateUpQ8(KernelContext ctx, + HalfFloatArray A, + ByteArray w1, ByteArray w3, + FloatArray gateUpOut, + int M, int hidDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int outStride = 2 * hidDim; + int blocksPerRow = K / 32; + + int wColBase = (blockCol < hidDim) ? blockCol : (blockCol - hidDim); + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int bCol0 = wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < hidDim) { + bReg0 = packQ8Halves(w1, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(w1, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(w1, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(w1, bCol3, bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(w3, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(w3, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(w3, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(w3, bCol3, bK3, blocksPerRow); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < hidDim) { + bReg0 = packQ8Halves(w1, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(w1, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(w1, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(w1, bCol3, kOff + bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(w3, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(w3, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(w3, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(w3, bCol3, kOff + bK3, blocksPerRow); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, gateUpOut, rBase + 0, cBase + 0, outStride); + ctx.mmaStore(c01, gateUpOut, rBase + 0, cBase + 8, outStride); + ctx.mmaStore(c02, gateUpOut, rBase + 0, cBase + 16, outStride); + ctx.mmaStore(c03, gateUpOut, rBase + 0, cBase + 24, outStride); + ctx.mmaStore(c04, gateUpOut, rBase + 0, cBase + 32, outStride); + ctx.mmaStore(c05, gateUpOut, rBase + 0, cBase + 40, outStride); + ctx.mmaStore(c06, gateUpOut, rBase + 0, cBase + 48, outStride); + ctx.mmaStore(c07, gateUpOut, rBase + 0, cBase + 56, outStride); + ctx.mmaStore(c10, gateUpOut, rBase + 16, cBase + 0, outStride); + ctx.mmaStore(c11, gateUpOut, rBase + 16, cBase + 8, outStride); + ctx.mmaStore(c12, gateUpOut, rBase + 16, cBase + 16, outStride); + ctx.mmaStore(c13, gateUpOut, rBase + 16, cBase + 24, outStride); + ctx.mmaStore(c14, gateUpOut, rBase + 16, cBase + 32, outStride); + ctx.mmaStore(c15, gateUpOut, rBase + 16, cBase + 40, outStride); + ctx.mmaStore(c16, gateUpOut, rBase + 16, cBase + 48, outStride); + ctx.mmaStore(c17, gateUpOut, rBase + 16, cBase + 56, outStride); + } + // @formatter:on } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java index 1c739f4e..efbab660 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java @@ -11,44 +11,73 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan (Q8_0). + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan + * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). * - *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} - * but uses Q8_0 kernels with inline dequantization. Key differences from the FP16 path:

- *
    - *
  • {@code wrapXBatch} is filled with dequantized FP32 embeddings by the host before - * the activation graph runs (no on-device FP16→FP32 conversion).
    • - *
  • {@code wrapXbBatch} (FP32) is reused as the normalized xb intermediate: written - * by {@code batchedRmsApplyFP32}, read by {@code batchedFusedQKVMatmulQ8}, then - * overwritten by flash attention output.
    • - *
  • {@code wrapXbFP16Batch} is not used.
    • - *
  • Weight matrices are {@code ByteArray} (Q8_0 format).
    • - *
+ *

One {@link ImmutableTaskGraph} per transformer layer, each processing + * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

+ * + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA, Q8_0 weights + * dequantized to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device + * after every layer so the subsequent single-token decode layers can consume it.

*/ public class LlamaQ8_0LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - static final int LOCAL_WORK_GROUP_SIZE = 32; + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; private final LlamaState state; private final LlamaTornadoWeights weights; private final LlamaConfiguration config; private final KernelContext context = new KernelContext(); private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; private final List layerITGs; private String lastLayerTaskGraphID; - public LlamaQ8_0LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights, LlamaConfiguration config, int batchSize) { + public LlamaQ8_0LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { this.state = state; this.weights = weights; this.config = config; this.batchSize = batchSize; - this.layerITGs = IntStream.range(0, config.numberOfLayers()).mapToObj(this::createBatchPrefillLayerTaskGraph).map(TaskGraph::snapshot).toList(); + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); } // @formatter:off @@ -56,37 +85,42 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph layer = new TaskGraph(graphName); + TaskGraph batchPrefillLayer = new TaskGraph(graphName); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - // Allocate GPU-side batch intermediates once. - // wrapXBatch is filled with dequantized FP32 by the host, persisted by prefillActivation. - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate persistent GPU-side intermediates once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + // wrapXBatch produced by the prefillActivation graph and persists in device memory + // to consume it from there we should use the explicit uniqueTaskGraph name + // the no-arg form would use current graph name, which causes NPE without CUDA Graphs + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { + // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, + batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache, state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } - // Per-layer weights: upload once (Q8_0 format) - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -97,140 +131,164 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w2Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray()); - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - // Writes FP32 normalized xb into wrapXbBatch (reused later by flash attention) - layer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP32, - context, state.wrapXbBatch, state.wrapXBatch, + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - layer.task("batch_qkv", - TransformerBatchPrefillKernels::batchedFusedQKVMatmulQ8, - context, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), weights.wvLayered[layerIndex].asByteArray(), - dim, kvDim, LOCAL_WORK_GROUP_SIZE); + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); - layer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCache, + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, config.headSize(), layerIndex, config.contextLength(), dim); - // Overwrites wrapXbBatch with attention output - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapXbBatch, state.wrapXBatch, + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, weights.woLayered[layerIndex].asByteArray(), - dim, dim, LOCAL_WORK_GROUP_SIZE); + state.woOut, paddedBatch, dim, dim); // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, - context, state.wrapXBatch, state.wrapHbBatch, + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, + state.ffnScaleBatch, dim); + + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); + state.gateUpResultBatch, paddedBatch, hidDim, dim); - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asByteArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + // Persist wrapXBatch for the next layer, and KV cache so the decode + // layers can consume it via the activation graph pass-through. + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return layer; + return batchPrefillLayer; } // @formatter:on + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); - int nHeads = config.numberOfHeads(); - int headSz = config.headSize(); - - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - int qkvRows = dim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker(batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // RoPE+KV cache: B*(dim/2) threads, local=512 int ropeGlobal = batchSize * (dim / 2); - int ropeLocal = Math.min(512, ropeGlobal); - while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) { - ropeLocal--; - } + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - int optLocal = findOptimalLocalSize(headSz); - WorkerGrid attnWorker = WorkerGridFactory.genericWorker(batchSize * nHeads * optLocal, optLocal); - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker(batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker(batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { - optimal = s; - break; - } - } + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } - return optimal; - } - - public List getLayerImmutableTaskGraphs() { - return layerITGs; } - public String getLastLayerTaskGraphID() { - return lastLayerTaskGraphID; - } - - public KernelContext getContext() { - return context; - } + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } } From 3f9fba0e54adeb36de4832e12cb7d7004e776e2b Mon Sep 17 00:00:00 2001 From: MaryXek Date: Mon, 6 Jul 2026 23:59:25 +0300 Subject: [PATCH 13/18] Port Qwen3 batch prefill (FP16 and Q8_0) to the tensor-core pipeline with packed QK-norm and split-half RoPE --- .../gpullama3/inference/state/State.java | 4 +- .../tornadovm/kernels/Qwen3Kernels.java | 132 +++++++++ .../prefill/Qwen3FP16LayersBatchPrefill.java | 251 ++++++++++------- .../prefill/Qwen3Q8_0LayersBatchPrefill.java | 259 +++++++++++------- 4 files changed, 440 insertions(+), 206 deletions(-) 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 e3092291..1f884b98 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -168,11 +168,11 @@ protected State(Configuration config, int batchsize) { this.ffnUpResult = new FloatArray(gpuBatchSize * config.hiddenDim()); this.xbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); - this.attnOutFP16 = new HalfFloatArray(paddedGpuBatch * config.dim()); + this.attnOutFP16 = new HalfFloatArray(paddedGpuBatch * qDim); // qDim == dim for Llama; qDim = nHeads*headDim for Qwen3 this.woOut = new FloatArray(paddedGpuBatch * config.dim()); this.wrapHbFP16Batch = new HalfFloatArray(paddedGpuBatch * config.hiddenDim()); this.w2Out = new FloatArray(paddedGpuBatch * config.dim()); - this.qkvResultBatch = new FloatArray(paddedGpuBatch * (config.dim() + 2 * config.kvDim())); + this.qkvResultBatch = new FloatArray(paddedGpuBatch * (qDim + 2 * kvDim)); this.gateUpResultBatch = new FloatArray(paddedGpuBatch * 2 * config.hiddenDim()); } else { this.embeddingXBatch = null; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java index 26b2f2b0..afc2d730 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java @@ -1214,5 +1214,137 @@ public static void batchedRopeWithKVCacheQwen3( wrapValueCache.set(cacheOff + ic + halfEmbdHead, wrapVBatch.get(kHeadBase + ic + halfEmbdHead)); } } + + // ── Packed-QKV variants for the tensor-core batch prefill path ─────────── + // These mirror batchedFusedQKRmsNorm / batchedRopeWithKVCacheQwen3 but read + // and write the packed [q | k | v] buffer produced by gemmMMAQKV + // (row stride qDim + 2*kvDim), so no separate Q/K/V buffers are needed. + + public static void batchedFusedQKRmsNormPacked( + KernelContext context, + FloatArray qkvBatch, + FloatArray qWeights, + FloatArray kWeights, + int nHeads, + int nHeadKv, + int nEmbdHead, + int qDim, + int kvDim, + float rmsNormEps) { + + int groupId = context.globalIdx / nEmbdHead; + int localId = context.localIdx; + int localSize = context.localGroupSizeX; + int totalHeadsPerBatch = nHeads + nHeadKv; + int qkvStride = qDim + 2 * kvDim; + + int batchIdx = groupId / totalHeadsPerBatch; + int headSlot = groupId % totalHeadsPerBatch; + + float[] localSum = context.allocateFloatLocalArray(nEmbdHead); + + if (headSlot < nHeads) { + // Q head (packed offset 0) + int headOffset = batchIdx * qkvStride + headSlot * nEmbdHead; + float partialSum = 0.0f; + for (int i = localId; i < nEmbdHead; i += localSize) { + float val = qkvBatch.get(headOffset + i); + partialSum += val * val; + } + localSum[localId] = partialSum; + context.localBarrier(); + for (int stride = localSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSum[localId] += localSum[localId + stride]; + } + context.localBarrier(); + } + float ss = localSum[0] / nEmbdHead + rmsNormEps; + ss = 1.0f / TornadoMath.sqrt(ss); + context.localBarrier(); + for (int i = localId; i < nEmbdHead; i += localSize) { + qkvBatch.set(headOffset + i, qWeights.get(i) * ss * qkvBatch.get(headOffset + i)); + } + } else { + // K head (packed offset qDim) + int kHeadIdx = headSlot - nHeads; + int headOffset = batchIdx * qkvStride + qDim + kHeadIdx * nEmbdHead; + float partialSum = 0.0f; + for (int i = localId; i < nEmbdHead; i += localSize) { + float val = qkvBatch.get(headOffset + i); + partialSum += val * val; + } + localSum[localId] = partialSum; + context.localBarrier(); + for (int stride = localSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSum[localId] += localSum[localId + stride]; + } + context.localBarrier(); + } + float ss = localSum[0] / nEmbdHead + rmsNormEps; + ss = 1.0f / TornadoMath.sqrt(ss); + context.localBarrier(); + for (int i = localId; i < nEmbdHead; i += localSize) { + qkvBatch.set(headOffset + i, kWeights.get(i) * ss * qkvBatch.get(headOffset + i)); + } + } + } + + public static void batchedRopeWithKVCacheQwen3Packed( + KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + int kvDim, + int nEmbdHead, + int layerIndex, + int contextLength, + int qDim) { + + int globalIdx = context.globalIdx; + int halfQDim = qDim / 2; + int batchIdx = globalIdx / halfQDim; + int pairIdx = globalIdx % halfQDim; + int qkvStride = qDim + 2 * kvDim; + + int pos = batchStartPosHolder.get(0) + batchIdx; + + // Qwen3 uses split-half RoPE: pair element ic with ic + nEmbdHead/2 within each head. + int halfEmbdHead = nEmbdHead / 2; + int ic = pairIdx % halfEmbdHead; + int headIdx = pairIdx / halfEmbdHead; + + float freq = 1.0f / TornadoMath.pow(1000000.0f, 2.0f * ic / (float) nEmbdHead); + float val = pos * freq; + float fcr = TornadoMath.cos(val); + float fci = TornadoMath.sin(val); + + // Rotate Q in place (packed offset 0) + int qHeadBase = batchIdx * qkvStride + headIdx * nEmbdHead; + float v0q = qkvBatch.get(qHeadBase + ic); + float v1q = qkvBatch.get(qHeadBase + ic + halfEmbdHead); + qkvBatch.set(qHeadBase + ic, v0q * fcr - v1q * fci); + qkvBatch.set(qHeadBase + ic + halfEmbdHead, v0q * fci + v1q * fcr); + + // Rotate K (packed offset qDim) and write K,V to cache + if (pairIdx < kvDim / 2) { + int kHeadIdx = pairIdx / halfEmbdHead; + int kHeadBase = batchIdx * qkvStride + qDim + kHeadIdx * nEmbdHead; + int vHeadBase = batchIdx * qkvStride + qDim + kvDim + kHeadIdx * nEmbdHead; + float v0k = qkvBatch.get(kHeadBase + ic); + float v1k = qkvBatch.get(kHeadBase + ic + halfEmbdHead); + float rotK0 = v0k * fcr - v1k * fci; + float rotK1 = v0k * fci + v1k * fcr; + + int cacheOff = layerIndex * contextLength * kvDim + pos * kvDim + kHeadIdx * nEmbdHead; + wrapKeyCache.set(cacheOff + ic, rotK0); + wrapKeyCache.set(cacheOff + ic + halfEmbdHead, rotK1); + wrapValueCache.set(cacheOff + ic, qkvBatch.get(vHeadBase + ic)); + wrapValueCache.set(cacheOff + ic + halfEmbdHead, qkvBatch.get(vHeadBase + ic + halfEmbdHead)); + } + } + } // @formatter:on diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java index 85b116d2..55b54afd 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java @@ -12,30 +12,54 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Batched-prefill transformer-layer TaskGraphs for the Qwen3 FP16 unified batched prefill-decode plan. + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefill} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). * - *

Mirrors {@link LlamaFP16LayersBatchPrefill} but adapts to Qwen3's GQA layout and - * Qwen3-specific kernels (fused Q/K RMSNorm, RoPE theta = 1 000 000). Avoids any calls to - * {@code Qwen3Configuration.headSize()}, {@code kvDim()}, or {@code kvMul()} which throw.

+ *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

*/ public class Qwen3FP16LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - static final int LOCAL_WORK_GROUP_SIZE = 32; + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; private final Qwen3State state; private final Qwen3TornadoWeights weights; private final Qwen3Configuration config; private final KernelContext context = new KernelContext(); private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; private final int nHeadKv; - private final int nEmbdHeadK; - private final int nEmbdHeadV; private final int nEmbdHead; private final int qDim; private final int kvDim; @@ -49,12 +73,16 @@ public Qwen3FP16LayersBatchPrefill(Qwen3State state, Qwen3TornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHeadK = config.numberOfHeadsKey(); - this.nEmbdHeadV = config.numberOfHeadsValue(); - this.nEmbdHead = nEmbdHeadV; - this.qDim = nEmbdHeadK * config.numberOfHeads(); - this.kvDim = nEmbdHeadV * nHeadKv; + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; this.gqa = config.numberOfHeads() / nHeadKv; this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) @@ -67,38 +95,38 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph layer = new TaskGraph(graphName); + TaskGraph batchPrefillLayer = new TaskGraph(graphName); int dim = config.dim(); int hidDim = config.hiddenDim(); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, + batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } // Per-layer weights: upload once - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), @@ -112,139 +140,160 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w3Layered[layerIndex].asHalfFloatArray()); // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - layer.task("batch_attn_rms_apply", + batchPrefillLayer.task("batch_attn_rms_apply", TransformerBatchPrefillKernels::batchedRmsApplyFP16, context, state.wrapXbFP16Batch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - layer.task("batch_qkv", - Qwen3Kernels::batchedFusedQKVMatmulFP16, - context, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), weights.wvLayered[layerIndex].asHalfFloatArray(), - dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); - layer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNorm, - context, - state.wrapQBatch, state.wrapKBatch, + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, weights.rms_att_QNormLayered[layerIndex].asFloatArray(), weights.rms_att_KNormLayered[layerIndex].asFloatArray(), config.numberOfHeads(), nHeadKv, nEmbdHead, qDim, kvDim, config.rmsNormEps()); - layer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3, + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - // Reuses batchedFlashAttention: passes qDim as the 'dim' stride parameter. - // Valid because qDim == dim for all standard Qwen3 models (nEmbdHeadK = dim/nHeads). - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, config.numberOfHeads(), nEmbdHead, kvDim, gqa, layerIndex, config.contextLength(), qDim); - // Output projection: n=qDim (input), d=dim (output) - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapXbBatch, state.wrapXBatch, + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, weights.woLayered[layerIndex].asHalfFloatArray(), - qDim, dim, LOCAL_WORK_GROUP_SIZE); + state.woOut, paddedBatch, dim, qDim); // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, - context, state.wrapXBatch, state.wrapHbBatch, + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, + context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asHalfFloatArray(), weights.w3Layered[layerIndex].asHalfFloatArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); + state.gateUpResultBatch, paddedBatch, hidDim, dim); - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return layer; + return batchPrefillLayer; } // @formatter:on + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); - int qkvRows = qDim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + // Split-half RoPE: B*(qDim/2) threads int ropeGlobal = batchSize * (qDim / 2); int ropeLocal = Math.min(512, ropeGlobal); while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - int optLocal = findOptimalLocalSize(nEmbdHead); + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * config.numberOfHeads() * optLocal, optLocal); + batchSize * nHeads * attnLocal, attnLocal); - // Wo: B*dim output rows (n=qDim, d=dim) - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } - return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java index b3db3f41..a561538e 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java @@ -12,30 +12,55 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Batched-prefill transformer-layer TaskGraphs for the Qwen3 Q8_0 unified batched prefill-decode plan. + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). * - *

Q8_0 path: wrapXbBatch (FP32) holds normalized activations; wrapXbFP16Batch is not used. - * Mirrors {@link Qwen3FP16LayersBatchPrefill} but uses Q8_0 weights (ByteArray) and FP32 - * attention normalization path.

+ *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA, Q8_0 weights dequantized + * to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

*/ public class Qwen3Q8_0LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - static final int LOCAL_WORK_GROUP_SIZE = 32; + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; private final Qwen3State state; private final Qwen3TornadoWeights weights; private final Qwen3Configuration config; private final KernelContext context = new KernelContext(); private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; private final int nHeadKv; - private final int nEmbdHeadK; - private final int nEmbdHeadV; private final int nEmbdHead; private final int qDim; private final int kvDim; @@ -49,12 +74,16 @@ public Qwen3Q8_0LayersBatchPrefill(Qwen3State state, Qwen3TornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHeadK = config.numberOfHeadsKey(); - this.nEmbdHeadV = config.numberOfHeadsValue(); - this.nEmbdHead = nEmbdHeadV; - this.qDim = nEmbdHeadK * config.numberOfHeads(); - this.kvDim = nEmbdHeadV * nHeadKv; + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; this.gqa = config.numberOfHeads() / nHeadKv; this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) @@ -67,36 +96,38 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph layer = new TaskGraph(graphName); + TaskGraph batchPrefillLayer = new TaskGraph(graphName); int dim = config.dim(); int hidDim = config.hiddenDim(); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, + batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache, state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } - // Per-layer weights (Q8_0 format) - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -110,138 +141,160 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w3Layered[layerIndex].asByteArray()); // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - // FP32 normalize into wrapXbBatch (Q8_0 path: no FP16 quantize step) - layer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP32, - context, state.wrapXbBatch, state.wrapXBatch, + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - layer.task("batch_qkv", - Qwen3Kernels::batchedFusedQKVMatmulQ8_0, - context, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), weights.wvLayered[layerIndex].asByteArray(), - dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); - layer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNorm, - context, - state.wrapQBatch, state.wrapKBatch, + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, weights.rms_att_QNormLayered[layerIndex].asFloatArray(), weights.rms_att_KNormLayered[layerIndex].asFloatArray(), config.numberOfHeads(), nHeadKv, nEmbdHead, qDim, kvDim, config.rmsNormEps()); - layer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3, + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - // Reuses batchedFlashAttention; passes qDim as the 'dim' stride (valid: qDim==dim typically). - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, config.numberOfHeads(), nEmbdHead, kvDim, gqa, layerIndex, config.contextLength(), qDim); - // Output projection (Q8_0): n=qDim, d=dim - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapXbBatch, state.wrapXBatch, + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, weights.woLayered[layerIndex].asByteArray(), - qDim, dim, LOCAL_WORK_GROUP_SIZE); + state.woOut, paddedBatch, dim, qDim); // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, - context, state.wrapXBatch, state.wrapHbBatch, + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); + state.gateUpResultBatch, paddedBatch, hidDim, dim); - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asByteArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return layer; + return batchPrefillLayer; } // @formatter:on + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); - int qkvRows = qDim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + // Split-half RoPE: B*(qDim/2) threads int ropeGlobal = batchSize * (qDim / 2); int ropeLocal = Math.min(512, ropeGlobal); while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - int optLocal = findOptimalLocalSize(nEmbdHead); + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * config.numberOfHeads() * optLocal, optLocal); + batchSize * nHeads * attnLocal, attnLocal); - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } - return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } From 9c5f8ca361d86c9030f67fa013be0a9de08f478d Mon Sep 17 00:00:00 2001 From: MaryXek Date: Tue, 7 Jul 2026 00:31:33 +0300 Subject: [PATCH 14/18] Put a check to execute batch MMA kernels only for the CUDA and PTX backends --- .../tornadovm/TensorCoreSupport.java | 31 +++ .../LlamaFP16LayersBatchPrefillGeneric.java | 250 +++++++++++++++++ .../Qwen3FP16LayersBatchPrefillGeneric.java | 259 ++++++++++++++++++ .../LlamaQ8_0LayersBatchPrefillGeneric.java | 242 ++++++++++++++++ .../Qwen3Q8_0LayersBatchPrefillGeneric.java | 256 +++++++++++++++++ .../fp16/LlamaFP16PlanComponents.java | 7 +- .../fp16/Qwen3FP16PlanComponents.java | 7 +- .../q8_0/LlamaQ8_0PlanComponents.java | 7 +- .../q8_0/Qwen3Q8_0PlanComponents.java | 7 +- 9 files changed, 1062 insertions(+), 4 deletions(-) create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java new file mode 100644 index 00000000..d9f77cbc --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java @@ -0,0 +1,31 @@ +package org.beehive.gpullama3.tornadovm; + +import uk.ac.manchester.tornado.api.enums.TornadoVMBackendType; +import uk.ac.manchester.tornado.api.runtime.TornadoRuntimeProvider; + +/** + * Detects whether the active TornadoVM backend can execute the tensor-core + * (MMA) batch-prefill kernels. TornadoVM lowers the MMA intrinsics + * ({@code mmaLoadA/B}, {@code mma}, {@code mmaStore}) only on the NVIDIA + * PTX and CUDA backends; on OpenCL, SPIR-V, and Metal the batch-prefill + * planners fall back to the portable matvec pipeline ({@code *Generic} + * planner classes). + */ +public final class TensorCoreSupport { + + private static boolean notified = false; + + private TensorCoreSupport() { + } + + public static synchronized boolean isTensorCoreCapableBackend() { + TornadoVMBackendType backendType = TornadoRuntimeProvider.getTornadoRuntime() + .getBackend(0) + .getBackendType(); + boolean capable = backendType == TornadoVMBackendType.PTX || backendType == TornadoVMBackendType.CUDA; + if (!capable && !notified) { + notified = true; + } + return capable; + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java new file mode 100644 index 00000000..ed162f36 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java @@ -0,0 +1,250 @@ +package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; + +import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; +import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan + * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). + * + *

One {@link ImmutableTaskGraph} per transformer layer, each processing + * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

+ * + *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device + * after every layer so the subsequent single-token decode layers can consume it.

+ */ +/** + * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec + * pipeline, retained as the fallback for backends without MMA intrinsics + * Selected automatically when the active TornadoVM + * backend is not PTX or CUDA. + */ +public class LlamaFP16LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { + + // Matches the local workgroup size used by the single-token kernels. + static final int LOCAL_WORK_GROUP_SIZE = 32; + + private final LlamaState state; + private final LlamaTornadoWeights weights; + private final LlamaConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public LlamaFP16LayersBatchPrefillGeneric(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate persistent GPU-side intermediates once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + // wrapXBatch produced by the prefillActivation graph and persists in device memory + // to consume it from there we should use the explicit uniqueTaskGraph name + // the no-arg form would use current graph name, which causes NPE without CUDA Graphs + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + weights.woLayered[layerIndex].asHalfFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w2Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray()); + + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + batchPrefillLayer.task("batch_qkv", + TransformerBatchPrefillKernels::batchedFusedQKVMatmul, + context, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + dim, kvDim, LOCAL_WORK_GROUP_SIZE); + + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCache, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapXbBatch, state.wrapXBatch, + weights.woLayered[layerIndex].asHalfFloatArray(), + dim, dim, LOCAL_WORK_GROUP_SIZE); + + // ── FFN Block ────────────────────────────────────────────────────────── + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, + context, state.wrapXBatch, state.wrapHbBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray(), + dim, hidDim, LOCAL_WORK_GROUP_SIZE); + + batchPrefillLayer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); + + // Persist wrapXBatch for the next layer, and KV cache so the decode + // layers can consume it via the activation graph pass-through. + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + // RMS: one thread per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + + // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // QKV: B*(dim+2*kvDim) workgroups × LOCAL_WORK_GROUP_SIZE + int qkvRows = dim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + // RoPE+KV cache: B*(dim/2) threads, local=512 + int ropeGlobal = batchSize * (dim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention (flash): B*nHeads workgroups × optimalLocalSize + int optLocal = findOptimalLocalSize(headSz); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * optLocal, optLocal); + + // Mat-vec (Wo, W2): B*d workgroups × LOCAL_WORK_GROUP_SIZE + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } + } + return optimal; + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java new file mode 100644 index 00000000..e9074764 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java @@ -0,0 +1,259 @@ +package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the Qwen3 FP16 unified batched prefill-decode plan. + * + *

Mirrors {@link LlamaFP16LayersBatchPrefill} but adapts to Qwen3's GQA layout and + * Qwen3-specific kernels (fused Q/K RMSNorm, RoPE theta = 1 000 000). Avoids any calls to + * {@code Qwen3Configuration.headSize()}, {@code kvDim()}, or {@code kvMul()} which throw.

+ */ +/** + * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec + * pipeline, retained as the fallback for backends without MMA intrinsics. + * Selected automatically when the active TornadoVM + * backend is not PTX or CUDA. + */ +public class Qwen3FP16LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { + + static final int LOCAL_WORK_GROUP_SIZE = 32; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + private final int nHeadKv; + private final int nEmbdHeadK; + private final int nEmbdHeadV; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3FP16LayersBatchPrefillGeneric(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHeadK = config.numberOfHeadsKey(); + this.nEmbdHeadV = config.numberOfHeadsValue(); + this.nEmbdHead = nEmbdHeadV; + this.qDim = nEmbdHeadK * config.numberOfHeads(); + this.kvDim = nEmbdHeadV * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph layer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + layer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); + } + + // Per-layer weights: upload once + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + weights.woLayered[layerIndex].asHalfFloatArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w2Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + layer.task("batch_qkv", + Qwen3Kernels::batchedFusedQKVMatmulFP16, + context, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNorm, + context, + state.wrapQBatch, state.wrapKBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + layer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Reuses batchedFlashAttention: passes qDim as the 'dim' stride parameter. + // Valid because qDim == dim for all standard Qwen3 models (nEmbdHeadK = dim/nHeads). + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection: n=qDim (input), d=dim (output) + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapXbBatch, state.wrapXBatch, + weights.woLayered[layerIndex].asHalfFloatArray(), + qDim, dim, LOCAL_WORK_GROUP_SIZE); + + // ── FFN Block ────────────────────────────────────────────────────────── + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, + context, state.wrapXBatch, state.wrapHbBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray(), + dim, hidDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); + + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return layer; + } + // @formatter:on + + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + int qkvRows = qDim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); + + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + int optLocal = findOptimalLocalSize(nEmbdHead); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * config.numberOfHeads() * optLocal, optLocal); + + // Wo: B*dim output rows (n=qDim, d=dim) + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } + } + return optimal; + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java new file mode 100644 index 00000000..d0627ef0 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java @@ -0,0 +1,242 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; +import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan (Q8_0). + * + *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} + * but uses Q8_0 kernels with inline dequantization. Key differences from the FP16 path:

+ *
    + *
  • {@code wrapXBatch} is filled with dequantized FP32 embeddings by the host before + * the activation graph runs (no on-device FP16→FP32 conversion).
    • + *
  • {@code wrapXbBatch} (FP32) is reused as the normalized xb intermediate: written + * by {@code batchedRmsApplyFP32}, read by {@code batchedFusedQKVMatmulQ8}, then + * overwritten by flash attention output.
    • + *
  • {@code wrapXbFP16Batch} is not used.
    • + *
  • Weight matrices are {@code ByteArray} (Q8_0 format).
    • + *
+ */ +/** + * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec + * pipeline, retained as the fallback for backends without MMA intrinsics. + * Selected automatically when the active TornadoVM + * backend is not PTX or CUDA. + */ +public class LlamaQ8_0LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { + + static final int LOCAL_WORK_GROUP_SIZE = 32; + + private final LlamaState state; + private final LlamaTornadoWeights weights; + private final LlamaConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public LlamaQ8_0LayersBatchPrefillGeneric(LlamaState state, LlamaTornadoWeights weights, LlamaConfiguration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.layerITGs = IntStream.range(0, config.numberOfLayers()).mapToObj(this::createBatchPrefillLayerTaskGraph).map(TaskGraph::snapshot).toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph layer = new TaskGraph(graphName); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate GPU-side batch intermediates once. + // wrapXBatch is filled with dequantized FP32 by the host, persisted by prefillActivation. + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + layer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); + } + + // Per-layer weights: upload once (Q8_0 format) + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + + // ── Attention Block ──────────────────────────────────────────────────── + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps()); + + // Writes FP32 normalized xb into wrapXbBatch (reused later by flash attention) + layer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, state.wrapXbBatch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + layer.task("batch_qkv", + TransformerBatchPrefillKernels::batchedFusedQKVMatmulQ8, + context, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + dim, kvDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCache, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + // Overwrites wrapXbBatch with attention output + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapXbBatch, state.wrapXBatch, + weights.woLayered[layerIndex].asByteArray(), + dim, dim, LOCAL_WORK_GROUP_SIZE); + + // ── FFN Block ────────────────────────────────────────────────────────── + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, + context, state.wrapXBatch, state.wrapHbBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + dim, hidDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asByteArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); + + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return layer; + } + // @formatter:on + + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + int qkvRows = dim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker(batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + int ropeGlobal = batchSize * (dim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) { + ropeLocal--; + } + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + int optLocal = findOptimalLocalSize(headSz); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker(batchSize * nHeads * optLocal, optLocal); + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker(batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker(batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { + optimal = s; + break; + } + } + } + return optimal; + } + + public List getLayerImmutableTaskGraphs() { + return layerITGs; + } + + public String getLastLayerTaskGraphID() { + return lastLayerTaskGraphID; + } + + public KernelContext getContext() { + return context; + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java new file mode 100644 index 00000000..b30973c5 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java @@ -0,0 +1,256 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the Qwen3 Q8_0 unified batched prefill-decode plan. + * + *

Q8_0 path: wrapXbBatch (FP32) holds normalized activations; wrapXbFP16Batch is not used. + * Mirrors {@link Qwen3FP16LayersBatchPrefill} but uses Q8_0 weights (ByteArray) and FP32 + * attention normalization path.

+ */ +/** + * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec + * pipeline, retained as the fallback for backends without MMA intrinsics. + * Selected automatically when the active TornadoVM + * backend is not PTX or CUDA. + */ +public class Qwen3Q8_0LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { + + static final int LOCAL_WORK_GROUP_SIZE = 32; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + private final int nHeadKv; + private final int nEmbdHeadK; + private final int nEmbdHeadV; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3Q8_0LayersBatchPrefillGeneric(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHeadK = config.numberOfHeadsKey(); + this.nEmbdHeadV = config.numberOfHeadsValue(); + this.nEmbdHead = nEmbdHeadV; + this.qDim = nEmbdHeadK * config.numberOfHeads(); + this.kvDim = nEmbdHeadV * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph layer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + layer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); + } + + // Per-layer weights (Q8_0 format) + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps()); + + // FP32 normalize into wrapXbBatch (Q8_0 path: no FP16 quantize step) + layer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, state.wrapXbBatch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + layer.task("batch_qkv", + Qwen3Kernels::batchedFusedQKVMatmulQ8_0, + context, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNorm, + context, + state.wrapQBatch, state.wrapKBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + layer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Reuses batchedFlashAttention; passes qDim as the 'dim' stride (valid: qDim==dim typically). + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, + context, state.batchStartPosHolder, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection (Q8_0): n=qDim, d=dim + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapXbBatch, state.wrapXBatch, + weights.woLayered[layerIndex].asByteArray(), + qDim, dim, LOCAL_WORK_GROUP_SIZE); + + // ── FFN Block ────────────────────────────────────────────────────────── + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, + context, state.wrapXBatch, state.wrapHbBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + dim, hidDim, LOCAL_WORK_GROUP_SIZE); + + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asByteArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); + + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return layer; + } + // @formatter:on + + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + int qkvRows = qDim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); + + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + int optLocal = findOptimalLocalSize(nEmbdHead); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * config.numberOfHeads() * optLocal, optLocal); + + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } + } + return optimal; + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java index e5166844..75bb1490 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java @@ -15,6 +15,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LlamaFP16FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LogitsFP16LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillGeneric; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -81,7 +83,10 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { - return new LlamaFP16LayersBatchPrefill(state, weights, config, batchSize); + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new LlamaFP16LayersBatchPrefill(state, weights, config, batchSize); + } + return new LlamaFP16LayersBatchPrefillGeneric(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java index 4cb7a0aa..00521d7a 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java @@ -15,6 +15,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefillGeneric; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -76,7 +78,10 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { - return new Qwen3FP16LayersBatchPrefill(state, weights, config, batchSize); + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new Qwen3FP16LayersBatchPrefill(state, weights, config, batchSize); + } + return new Qwen3FP16LayersBatchPrefillGeneric(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java index 536a8789..35e498e9 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java @@ -15,6 +15,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LlamaQ8_0FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefillGeneric; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -85,7 +87,10 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { - return new LlamaQ8_0LayersBatchPrefill(state, weights, config, batchSize); + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new LlamaQ8_0LayersBatchPrefill(state, weights, config, batchSize); + } + return new LlamaQ8_0LayersBatchPrefillGeneric(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java index 6f067d1f..1baf0cf1 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java @@ -15,6 +15,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefillGeneric; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -76,7 +78,10 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { - return new Qwen3Q8_0LayersBatchPrefill(state, weights, config, batchSize); + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new Qwen3Q8_0LayersBatchPrefill(state, weights, config, batchSize); + } + return new Qwen3Q8_0LayersBatchPrefillGeneric(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── From 086584130d6fe07be99fa41ab45a25dc26e06826 Mon Sep 17 00:00:00 2001 From: MaryXek Date: Tue, 7 Jul 2026 13:31:36 +0300 Subject: [PATCH 15/18] Enable MMA path only on CUDA backend --- .../gpullama3/tornadovm/TensorCoreSupport.java | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java index d9f77cbc..0383e24b 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java @@ -7,25 +7,14 @@ * Detects whether the active TornadoVM backend can execute the tensor-core * (MMA) batch-prefill kernels. TornadoVM lowers the MMA intrinsics * ({@code mmaLoadA/B}, {@code mma}, {@code mmaStore}) only on the NVIDIA - * PTX and CUDA backends; on OpenCL, SPIR-V, and Metal the batch-prefill - * planners fall back to the portable matvec pipeline ({@code *Generic} - * planner classes). + * CUDA backend. */ public final class TensorCoreSupport { - private static boolean notified = false; - - private TensorCoreSupport() { - } - - public static synchronized boolean isTensorCoreCapableBackend() { + public static boolean isTensorCoreCapableBackend() { TornadoVMBackendType backendType = TornadoRuntimeProvider.getTornadoRuntime() .getBackend(0) .getBackendType(); - boolean capable = backendType == TornadoVMBackendType.PTX || backendType == TornadoVMBackendType.CUDA; - if (!capable && !notified) { - notified = true; - } - return capable; + return backendType == TornadoVMBackendType.CUDA; } } From 1e56c0c8664454f1f109f5fcee7591613bff3214 Mon Sep 17 00:00:00 2001 From: MaryXek Date: Tue, 7 Jul 2026 13:54:58 +0300 Subject: [PATCH 16/18] Refactor prefill layer class names --- ...atchPrefillTransformerLayerTaskGraphs.java | 2 +- .../prefill/LlamaFP16LayersBatchPrefill.java | 201 +++++------ ...va => LlamaFP16LayersBatchPrefillMMA.java} | 213 +++++++----- .../prefill/Qwen3FP16LayersBatchPrefill.java | 251 ++++++-------- .../Qwen3FP16LayersBatchPrefillGeneric.java | 259 --------------- .../Qwen3FP16LayersBatchPrefillMMA.java | 302 +++++++++++++++++ .../prefill/LlamaQ8_0LayersBatchPrefill.java | 314 +++++++----------- .../LlamaQ8_0LayersBatchPrefillGeneric.java | 242 -------------- .../LlamaQ8_0LayersBatchPrefillMMA.java | 294 ++++++++++++++++ .../prefill/Qwen3Q8_0LayersBatchPrefill.java | 259 ++++++--------- .../Qwen3Q8_0LayersBatchPrefillGeneric.java | 256 -------------- .../Qwen3Q8_0LayersBatchPrefillMMA.java | 304 +++++++++++++++++ .../fp16/LlamaFP16PlanComponents.java | 6 +- .../fp16/Qwen3FP16PlanComponents.java | 6 +- .../q8_0/LlamaQ8_0PlanComponents.java | 6 +- .../q8_0/Qwen3Q8_0PlanComponents.java | 6 +- 16 files changed, 1449 insertions(+), 1472 deletions(-) rename src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/{LlamaFP16LayersBatchPrefillGeneric.java => LlamaFP16LayersBatchPrefillMMA.java} (52%) delete mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java delete mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java delete mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java create mode 100644 src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java index cf406091..a46bbbae 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java @@ -8,7 +8,7 @@ /** * Interface for a group of N batched-prefill transformer-layer TornadoVM TaskGraphs. * - *

Implemented by {@code LlamaFP16LayersBatchPrefill} and {@code LlamaQ8_0LayersBatchPrefill}.

+ *

Implemented by {@code LlamaFP16LayersBatchPrefillMMA} and {@code LlamaQ8_0LayersBatchPrefillMMA}.

*/ public interface BatchPrefillTransformerLayerTaskGraphs { List getLayerImmutableTaskGraphs(); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java index fa04c0b8..6f0b3e4b 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java @@ -11,8 +11,6 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.WorkerGrid1D; -import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; @@ -25,39 +23,19 @@ *

One {@link ImmutableTaskGraph} per transformer layer, each processing * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

* - *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA):

- * 
- *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
- *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
- *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
- *   batch_rope_kv         RoPE over the packed buffer + KV cache write
- *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
- *   woProj                MMA GEMM → woOut
- *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
- *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
- *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
- *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
- *   w2Proj                MMA GEMM → w2Out
- *   w2Resid               x += w2Out
- *
- * *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device * after every layer so the subsequent single-token decode layers can consume it.

*/ public class LlamaFP16LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - // Local size for the parallel RMS reductions (one workgroup per token). - static final int RMS_LOCAL_SIZE = 256; + // Matches the local workgroup size used by the single-token kernels. + static final int LOCAL_WORK_GROUP_SIZE = 32; private final LlamaState state; private final LlamaTornadoWeights weights; private final LlamaConfiguration config; private final KernelContext context = new KernelContext(); private final int batchSize; - // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent - // buffers in State are allocated at this padded size; rows >= batchSize are - // computed but never consumed. All non-GEMM kernels use the true batchSize. - private final int paddedBatch; private final List layerITGs; private String lastLayerTaskGraphID; @@ -67,12 +45,6 @@ public LlamaFP16LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; - this.paddedBatch = (batchSize + 127) & ~127; - if (batchSize % 128 != 0) { - System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " - + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", - batchSize, paddedBatch, batchSize, paddedBatch); - } this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) .map(TaskGraph::snapshot) @@ -95,10 +67,10 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { context, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); // wrapXBatch produced by the prefillActivation graph and persists in device memory // to consume it from there we should use the explicit uniqueTaskGraph name // the no-arg form would use current graph name, which causes NPE without CUDA Graphs @@ -109,13 +81,13 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); } // Per-layer weights: upload once @@ -136,9 +108,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { // ── Attention Block ──────────────────────────────────────────────────── batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceParallel, + TransformerBatchPrefillKernels::batchedRmsReduce, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + dim, config.rmsNormEps()); batchPrefillLayer.task("batch_attn_rms_apply", TransformerBatchPrefillKernels::batchedRmsApplyFP16, @@ -146,69 +118,57 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. - // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the - // skinny GQA projections, and the A operand is read once, not thrice. - batchPrefillLayer.task("qkvProj", - TransformerBatchPrefillKernels::gemmMMAQKV, - context, state.wrapXbFP16Batch, + batchPrefillLayer.task("batch_qkv", + TransformerBatchPrefillKernels::batchedFusedQKVMatmul, + context, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), weights.wvLayered[layerIndex].asHalfFloatArray(), - state.qkvResultBatch, paddedBatch, dim, kvDim, dim); + dim, kvDim, LOCAL_WORK_GROUP_SIZE); batchPrefillLayer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, + TransformerBatchPrefillKernels::batchedRopeWithKVCache, context, state.batchStartPosHolder, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, config.headSize(), layerIndex, config.contextLength(), dim); - // Register-partitioned P·V accumulation + direct FP16 emission - // (replaces batchedFlashAttention + attnCast). batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + TransformerBatchPrefillKernels::batchedFlashAttention, context, state.batchStartPosHolder, - state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.attnOutFP16, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.attnOutFP16, + batchPrefillLayer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapXbBatch, state.wrapXBatch, weights.woLayered[layerIndex].asHalfFloatArray(), - state.woOut, paddedBatch, dim, dim); + dim, dim, LOCAL_WORK_GROUP_SIZE); // ── FFN Block ────────────────────────────────────────────────────────── - // x += woOut is fused into the FFN RMS reduction (drops the woResid task). batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, - context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); - batchPrefillLayer.task("batch_ffn_rms_apply", - TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, - context, state.normedXFFNFP16, state.wrapXBatch, + batchPrefillLayer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, + context, state.wrapXBatch, state.wrapHbBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, dim); - - // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. - batchPrefillLayer.task("gateUpProj", - TransformerBatchPrefillKernels::gemmMMAGateUp, - context, state.normedXFFNFP16, + state.ffnScaleBatch, weights.w1Layered[layerIndex].asHalfFloatArray(), weights.w3Layered[layerIndex].asHalfFloatArray(), - state.gateUpResultBatch, paddedBatch, hidDim, dim); - - batchPrefillLayer.task("swiglu", - TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, - context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) - .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, - context, state.wrapHbFP16Batch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - state.w2Out, batchSize, dim, hidDim) - .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, - context, state.wrapXBatch, state.w2Out); + dim, hidDim, LOCAL_WORK_GROUP_SIZE); + + batchPrefillLayer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); // Persist wrapXBatch for the next layer, and KV cache so the decode // layers can consume it via the activation graph pass-through. @@ -218,21 +178,6 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { } // @formatter:on - // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. - static WorkerGrid mmaGrid(int paddedM, int N) { - int mBlocks = paddedM / 128; // BM - int nBlocks = N / 128; // BN - WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); - g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) - return g; - } - - static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 - WorkerGrid1D g = new WorkerGrid1D(n); - g.setLocalWork(256, 1, 1); - return g; - } - /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { int dim = config.dim(); @@ -241,13 +186,16 @@ public void updateGridScheduler(GridScheduler scheduler) { int nHeads = config.numberOfHeads(); int headSz = config.headSize(); - // Parallel RMS reductions: one 256-thread workgroup per batch token - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( - batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + // RMS: one thread per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // QKV: B*(dim+2*kvDim) workgroups × LOCAL_WORK_GROUP_SIZE + int qkvRows = dim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); // RoPE+KV cache: B*(dim/2) threads, local=512 int ropeGlobal = batchSize * (dim / 2); @@ -255,36 +203,39 @@ public void updateGridScheduler(GridScheduler scheduler) { while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. - // The kernel requires headSize <= 2*localSize. - int attnLocal = Math.min(headSz, 128); + // Attention (flash): B*nHeads workgroups × optimalLocalSize + int optLocal = findOptimalLocalSize(headSz); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * attnLocal, attnLocal); + batchSize * nHeads * optLocal, optLocal); - // MMA grids - WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV - WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj - WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 - - // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) - WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid - WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + // Mat-vec (Wo, W2): B*d workgroups × LOCAL_WORK_GROUP_SIZE + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); - scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); - scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); - scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); - scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } } + return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java similarity index 52% rename from src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java rename to src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java index ed162f36..7d04117c 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillGeneric.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java @@ -11,6 +11,8 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; @@ -23,34 +25,54 @@ *

One {@link ImmutableTaskGraph} per transformer layer, each processing * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

* + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device * after every layer so the subsequent single-token decode layers can consume it.

*/ -/** - * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec - * pipeline, retained as the fallback for backends without MMA intrinsics - * Selected automatically when the active TornadoVM - * backend is not PTX or CUDA. - */ -public class LlamaFP16LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { +public class LlamaFP16LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { - // Matches the local workgroup size used by the single-token kernels. - static final int LOCAL_WORK_GROUP_SIZE = 32; + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; private final LlamaState state; private final LlamaTornadoWeights weights; private final LlamaConfiguration config; private final KernelContext context = new KernelContext(); private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; private final List layerITGs; private String lastLayerTaskGraphID; - public LlamaFP16LayersBatchPrefillGeneric(LlamaState state, LlamaTornadoWeights weights, - LlamaConfiguration config, int batchSize) { + public LlamaFP16LayersBatchPrefillMMA(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { this.state = state; this.weights = weights; this.config = config; this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) .map(TaskGraph::snapshot) @@ -73,10 +95,10 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { context, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); // wrapXBatch produced by the prefillActivation graph and persists in device memory // to consume it from there we should use the explicit uniqueTaskGraph name // the no-arg form would use current graph name, which causes NPE without CUDA Graphs @@ -87,13 +109,13 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { batchPrefillLayer.consumeFromDevice(pred, context, state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } // Per-layer weights: upload once @@ -114,9 +136,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { // ── Attention Block ──────────────────────────────────────────────────── batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, + TransformerBatchPrefillKernels::batchedRmsReduceParallel, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); batchPrefillLayer.task("batch_attn_rms_apply", TransformerBatchPrefillKernels::batchedRmsApplyFP16, @@ -124,57 +146,69 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - batchPrefillLayer.task("batch_qkv", - TransformerBatchPrefillKernels::batchedFusedQKVMatmul, - context, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), weights.wvLayered[layerIndex].asHalfFloatArray(), - dim, kvDim, LOCAL_WORK_GROUP_SIZE); + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); batchPrefillLayer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCache, + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - batchPrefillLayer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapXbBatch, state.wrapXBatch, + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, weights.woLayered[layerIndex].asHalfFloatArray(), - dim, dim, LOCAL_WORK_GROUP_SIZE); + state.woOut, paddedBatch, dim, dim); // ── FFN Block ────────────────────────────────────────────────────────── + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - batchPrefillLayer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, - context, state.wrapXBatch, state.wrapHbBatch, + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, + state.ffnScaleBatch, dim); + + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, + context, state.normedXFFNFP16, weights.w1Layered[layerIndex].asHalfFloatArray(), weights.w3Layered[layerIndex].asHalfFloatArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); - - batchPrefillLayer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); // Persist wrapXBatch for the next layer, and KV cache so the decode // layers can consume it via the activation graph pass-through. @@ -184,6 +218,21 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { } // @formatter:on + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { int dim = config.dim(); @@ -192,16 +241,13 @@ public void updateGridScheduler(GridScheduler scheduler) { int nHeads = config.numberOfHeads(); int headSz = config.headSize(); - // RMS: one thread per batch token - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - // QKV: B*(dim+2*kvDim) workgroups × LOCAL_WORK_GROUP_SIZE - int qkvRows = dim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); // RoPE+KV cache: B*(dim/2) threads, local=512 int ropeGlobal = batchSize * (dim / 2); @@ -209,39 +255,36 @@ public void updateGridScheduler(GridScheduler scheduler) { while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - // Attention (flash): B*nHeads workgroups × optimalLocalSize - int optLocal = findOptimalLocalSize(headSz); + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * optLocal, optLocal); + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 - // Mat-vec (Wo, W2): B*d workgroups × LOCAL_WORK_GROUP_SIZE - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); } - return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java index 55b54afd..85b116d2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefill.java @@ -12,54 +12,30 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.WorkerGrid1D; -import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) - * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefill} with the Qwen3 - * architectural additions: per-head Q/K RMS normalization between the QKV - * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention - * (qDim = nHeads * headDim, which may differ from the model dim). + * Batched-prefill transformer-layer TaskGraphs for the Qwen3 FP16 unified batched prefill-decode plan. * - *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA):

- * 
- *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
- *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
- *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
- *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
- *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
- *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
- *   woProj                MMA GEMM [dim × qDim] → woOut
- *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
- *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
- *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
- *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
- *   w2Proj                MMA GEMM → w2Out
- *   w2Resid               x += w2Out
- *
- * - *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 - * (holds for all standard Qwen3 checkpoints).

+ *

Mirrors {@link LlamaFP16LayersBatchPrefill} but adapts to Qwen3's GQA layout and + * Qwen3-specific kernels (fused Q/K RMSNorm, RoPE theta = 1 000 000). Avoids any calls to + * {@code Qwen3Configuration.headSize()}, {@code kvDim()}, or {@code kvMul()} which throw.

*/ public class Qwen3FP16LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - // Local size for the parallel RMS reductions (one workgroup per token). - static final int RMS_LOCAL_SIZE = 256; + static final int LOCAL_WORK_GROUP_SIZE = 32; private final Qwen3State state; private final Qwen3TornadoWeights weights; private final Qwen3Configuration config; private final KernelContext context = new KernelContext(); private final int batchSize; - // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama - // planner for the padding rationale. Non-GEMM kernels use the true batchSize. - private final int paddedBatch; private final int nHeadKv; + private final int nEmbdHeadK; + private final int nEmbdHeadV; private final int nEmbdHead; private final int qDim; private final int kvDim; @@ -73,16 +49,12 @@ public Qwen3FP16LayersBatchPrefill(Qwen3State state, Qwen3TornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; - this.paddedBatch = (batchSize + 127) & ~127; - if (batchSize % 128 != 0) { - System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " - + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", - batchSize, paddedBatch, batchSize, paddedBatch); - } this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHead = config.numberOfHeadsValue(); - this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); - this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.nEmbdHeadK = config.numberOfHeadsKey(); + this.nEmbdHeadV = config.numberOfHeadsValue(); + this.nEmbdHead = nEmbdHeadV; + this.qDim = nEmbdHeadK * config.numberOfHeads(); + this.kvDim = nEmbdHeadV * nHeadKv; this.gqa = config.numberOfHeads() / nHeadKv; this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) @@ -95,38 +67,38 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph batchPrefillLayer = new TaskGraph(graphName); + TaskGraph layer = new TaskGraph(graphName); int dim = config.dim(); int hidDim = config.hiddenDim(); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { - batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); - batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { String pred = "batchPrefillLayer_" + (layerIndex - 1); - batchPrefillLayer.consumeFromDevice(pred, + layer.consumeFromDevice(pred, context, state.wrapXBatch, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapXbBatch, + state.wrapHbBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); } // Per-layer weights: upload once - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), @@ -140,160 +112,139 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w3Layered[layerIndex].asHalfFloatArray()); // ── Attention Block ──────────────────────────────────────────────────── - batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceParallel, + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + dim, config.rmsNormEps()); - batchPrefillLayer.task("batch_attn_rms_apply", + layer.task("batch_attn_rms_apply", TransformerBatchPrefillKernels::batchedRmsApplyFP16, context, state.wrapXbFP16Batch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) - batchPrefillLayer.task("qkvProj", - TransformerBatchPrefillKernels::gemmMMAQKV, - context, state.wrapXbFP16Batch, + layer.task("batch_qkv", + Qwen3Kernels::batchedFusedQKVMatmulFP16, + context, + state.wrapXbFP16Batch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, weights.wqLayered[layerIndex].asHalfFloatArray(), weights.wkLayered[layerIndex].asHalfFloatArray(), weights.wvLayered[layerIndex].asHalfFloatArray(), - state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); - // Qwen3: per-head RMS norm on Q and K before RoPE - batchPrefillLayer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNormPacked, - context, state.qkvResultBatch, + layer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNorm, + context, + state.wrapQBatch, state.wrapKBatch, weights.rms_att_QNormLayered[layerIndex].asFloatArray(), weights.rms_att_KNormLayered[layerIndex].asFloatArray(), config.numberOfHeads(), nHeadKv, nEmbdHead, qDim, kvDim, config.rmsNormEps()); - batchPrefillLayer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + layer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3, context, state.batchStartPosHolder, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - // Register-partitioned flash attention over the packed buffer. - // The 'dim' parameter doubles as the packed-Q stride base and the - // attnOutFP16 row width — both are qDim for Qwen3. - batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + // Reuses batchedFlashAttention: passes qDim as the 'dim' stride parameter. + // Valid because qDim == dim for all standard Qwen3 models (nEmbdHeadK = dim/nHeads). + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, context, state.batchStartPosHolder, - state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.attnOutFP16, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, config.numberOfHeads(), nEmbdHead, kvDim, gqa, layerIndex, config.contextLength(), qDim); - // Output projection: [M=batch, N=dim, K=qDim] - batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, - context, state.attnOutFP16, + // Output projection: n=qDim (input), d=dim (output) + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapXbBatch, state.wrapXBatch, weights.woLayered[layerIndex].asHalfFloatArray(), - state.woOut, paddedBatch, dim, qDim); + qDim, dim, LOCAL_WORK_GROUP_SIZE); // ── FFN Block ────────────────────────────────────────────────────────── - batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, - context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - - batchPrefillLayer.task("batch_ffn_rms_apply", - TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, - context, state.normedXFFNFP16, state.wrapXBatch, + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, + context, state.wrapXBatch, state.wrapHbBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, dim); - - batchPrefillLayer.task("gateUpProj", - TransformerBatchPrefillKernels::gemmMMAGateUp, - context, state.normedXFFNFP16, + state.ffnScaleBatch, weights.w1Layered[layerIndex].asHalfFloatArray(), weights.w3Layered[layerIndex].asHalfFloatArray(), - state.gateUpResultBatch, paddedBatch, hidDim, dim); + dim, hidDim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("swiglu", - TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, - context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) - .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, - context, state.wrapHbFP16Batch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - state.w2Out, paddedBatch, dim, hidDim) - .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, - context, state.wrapXBatch, state.w2Out); + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidual, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return batchPrefillLayer; + return layer; } // @formatter:on - // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. - static WorkerGrid mmaGrid(int paddedM, int N) { - int mBlocks = paddedM / 128; // BM - int nBlocks = N / 128; // BN - WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); - g.setLocalWork(256, 1, 1); - return g; - } - - static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 - WorkerGrid1D g = new WorkerGrid1D(n); - g.setLocalWork(256, 1, 1); - return g; - } - - /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); - int nHeads = config.numberOfHeads(); + int dim = config.dim(); + int hidDim = config.hiddenDim(); - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( - batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + int qkvRows = qDim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); - // Split-half RoPE: B*(qDim/2) threads int ropeGlobal = batchSize * (qDim / 2); int ropeLocal = Math.min(512, ropeGlobal); while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads - int attnLocal = Math.min(nEmbdHead, 128); + int optLocal = findOptimalLocalSize(nEmbdHead); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * attnLocal, attnLocal); + batchSize * config.numberOfHeads() * optLocal, optLocal); - // MMA grids - WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV - WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj - WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 - - WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid - WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + // Wo: B*dim output rows (n=qDim, d=dim) + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); - scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); - scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); - scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); - scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } } + return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java deleted file mode 100644 index e9074764..00000000 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillGeneric.java +++ /dev/null @@ -1,259 +0,0 @@ -package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; - -import org.beehive.gpullama3.inference.state.Qwen3State; -import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; -import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; -import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; -import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; -import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; -import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; -import uk.ac.manchester.tornado.api.GridScheduler; -import uk.ac.manchester.tornado.api.ImmutableTaskGraph; -import uk.ac.manchester.tornado.api.KernelContext; -import uk.ac.manchester.tornado.api.TaskGraph; -import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.enums.DataTransferMode; - -import java.util.List; -import java.util.stream.IntStream; - -/** - * Batched-prefill transformer-layer TaskGraphs for the Qwen3 FP16 unified batched prefill-decode plan. - * - *

Mirrors {@link LlamaFP16LayersBatchPrefill} but adapts to Qwen3's GQA layout and - * Qwen3-specific kernels (fused Q/K RMSNorm, RoPE theta = 1 000 000). Avoids any calls to - * {@code Qwen3Configuration.headSize()}, {@code kvDim()}, or {@code kvMul()} which throw.

- */ -/** - * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec - * pipeline, retained as the fallback for backends without MMA intrinsics. - * Selected automatically when the active TornadoVM - * backend is not PTX or CUDA. - */ -public class Qwen3FP16LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { - - static final int LOCAL_WORK_GROUP_SIZE = 32; - - private final Qwen3State state; - private final Qwen3TornadoWeights weights; - private final Qwen3Configuration config; - private final KernelContext context = new KernelContext(); - private final int batchSize; - private final int nHeadKv; - private final int nEmbdHeadK; - private final int nEmbdHeadV; - private final int nEmbdHead; - private final int qDim; - private final int kvDim; - private final int gqa; - private final List layerITGs; - private String lastLayerTaskGraphID; - - public Qwen3FP16LayersBatchPrefillGeneric(Qwen3State state, Qwen3TornadoWeights weights, - Qwen3Configuration config, int batchSize) { - this.state = state; - this.weights = weights; - this.config = config; - this.batchSize = batchSize; - this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHeadK = config.numberOfHeadsKey(); - this.nEmbdHeadV = config.numberOfHeadsValue(); - this.nEmbdHead = nEmbdHeadV; - this.qDim = nEmbdHeadK * config.numberOfHeads(); - this.kvDim = nEmbdHeadV * nHeadKv; - this.gqa = config.numberOfHeads() / nHeadKv; - this.layerITGs = IntStream.range(0, config.numberOfLayers()) - .mapToObj(this::createBatchPrefillLayerTaskGraph) - .map(TaskGraph::snapshot) - .toList(); - } - - // @formatter:off - private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { - String graphName = "batchPrefillLayer_" + layerIndex; - if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - - TaskGraph layer = new TaskGraph(graphName); - int dim = config.dim(); - int hidDim = config.hiddenDim(); - - // ── Data Transfers ───────────────────────────────────────────────────── - if (layerIndex == 0) { - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - context, - state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); - } else { - String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, - context, - state.wrapXBatch, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapXbBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); - } - - // Per-layer weights: upload once - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - weights.wqLayered[layerIndex].asHalfFloatArray(), - weights.wkLayered[layerIndex].asHalfFloatArray(), - weights.wvLayered[layerIndex].asHalfFloatArray(), - weights.woLayered[layerIndex].asHalfFloatArray(), - weights.rms_att_QNormLayered[layerIndex].asFloatArray(), - weights.rms_att_KNormLayered[layerIndex].asFloatArray(), - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - weights.w1Layered[layerIndex].asHalfFloatArray(), - weights.w2Layered[layerIndex].asHalfFloatArray(), - weights.w3Layered[layerIndex].asHalfFloatArray()); - - // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, - context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP16, - context, state.wrapXbFP16Batch, state.wrapXBatch, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - state.attnScaleBatch, dim); - - layer.task("batch_qkv", - Qwen3Kernels::batchedFusedQKVMatmulFP16, - context, - state.wrapXbFP16Batch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - weights.wqLayered[layerIndex].asHalfFloatArray(), - weights.wkLayered[layerIndex].asHalfFloatArray(), - weights.wvLayered[layerIndex].asHalfFloatArray(), - dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNorm, - context, - state.wrapQBatch, state.wrapKBatch, - weights.rms_att_QNormLayered[layerIndex].asFloatArray(), - weights.rms_att_KNormLayered[layerIndex].asFloatArray(), - config.numberOfHeads(), nHeadKv, nEmbdHead, - qDim, kvDim, config.rmsNormEps()); - - layer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapKeyCache, state.wrapValueCache, - kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - - // Reuses batchedFlashAttention: passes qDim as the 'dim' stride parameter. - // Valid because qDim == dim for all standard Qwen3 models (nEmbdHeadK = dim/nHeads). - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, - config.numberOfHeads(), nEmbdHead, - kvDim, gqa, layerIndex, config.contextLength(), qDim); - - // Output projection: n=qDim (input), d=dim (output) - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapXbBatch, state.wrapXBatch, - weights.woLayered[layerIndex].asHalfFloatArray(), - qDim, dim, LOCAL_WORK_GROUP_SIZE); - - // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp, - context, state.wrapXBatch, state.wrapHbBatch, - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, - weights.w1Layered[layerIndex].asHalfFloatArray(), - weights.w3Layered[layerIndex].asHalfFloatArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidual, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asHalfFloatArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); - - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - - return layer; - } - // @formatter:on - - public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); - - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - int qkvRows = qDim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); - - int ropeGlobal = batchSize * (qDim / 2); - int ropeLocal = Math.min(512, ropeGlobal); - while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; - WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - - int optLocal = findOptimalLocalSize(nEmbdHead); - WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * config.numberOfHeads() * optLocal, optLocal); - - // Wo: B*dim output rows (n=qDim, d=dim) - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - for (int i = 0; i < config.numberOfLayers(); i++) { - String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } - } - return optimal; - } - - public List getLayerImmutableTaskGraphs() { return layerITGs; } - public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } - public KernelContext getContext() { return context; } -} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java new file mode 100644 index 00000000..f12b842d --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java @@ -0,0 +1,302 @@ +package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefillMMA} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). + * + *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

+ */ +public class Qwen3FP16LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final int nHeadKv; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3FP16LayersBatchPrefillMMA(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + weights.woLayered[layerIndex].asHalfFloatArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w2Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asHalfFloatArray(), + state.woOut, paddedBatch, dim, qDim); + + // ── FFN Block ────────────────────────────────────────────────────────── + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + + // Split-half RoPE: B*(qDim/2) threads + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java index efbab660..a3016cb7 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java @@ -11,73 +11,44 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.WorkerGrid1D; -import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan - * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan (Q8_0). * - *

One {@link ImmutableTaskGraph} per transformer layer, each processing - * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

- * - *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA, Q8_0 weights - * dequantized to FP16 in the GEMM staging registers — W8A16):

- * 
- *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
- *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
- *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
- *   batch_rope_kv         RoPE over the packed buffer + KV cache write
- *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
- *   woProj                MMA GEMM → woOut
- *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
- *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
- *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
- *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
- *   w2Proj                MMA GEMM → w2Out
- *   w2Resid               x += w2Out
- *
- * - *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device - * after every layer so the subsequent single-token decode layers can consume it.

+ *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA} + * but uses Q8_0 kernels with inline dequantization. Key differences from the FP16 path:

+ *
    + *
  • {@code wrapXBatch} is filled with dequantized FP32 embeddings by the host before + * the activation graph runs (no on-device FP16→FP32 conversion).
    • + *
  • {@code wrapXbBatch} (FP32) is reused as the normalized xb intermediate: written + * by {@code batchedRmsApplyFP32}, read by {@code batchedFusedQKVMatmulQ8}, then + * overwritten by flash attention output.
    • + *
  • {@code wrapXbFP16Batch} is not used.
    • + *
  • Weight matrices are {@code ByteArray} (Q8_0 format).
    • + *
*/ public class LlamaQ8_0LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - // Local size for the parallel RMS reductions (one workgroup per token). - static final int RMS_LOCAL_SIZE = 256; + static final int LOCAL_WORK_GROUP_SIZE = 32; private final LlamaState state; private final LlamaTornadoWeights weights; private final LlamaConfiguration config; private final KernelContext context = new KernelContext(); private final int batchSize; - // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent - // buffers in State are allocated at this padded size; rows >= batchSize are - // computed but never consumed. All non-GEMM kernels use the true batchSize. - private final int paddedBatch; private final List layerITGs; private String lastLayerTaskGraphID; - public LlamaQ8_0LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights, - LlamaConfiguration config, int batchSize) { + public LlamaQ8_0LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights, LlamaConfiguration config, int batchSize) { this.state = state; this.weights = weights; this.config = config; this.batchSize = batchSize; - this.paddedBatch = (batchSize + 127) & ~127; - if (batchSize % 128 != 0) { - System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " - + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", - batchSize, paddedBatch, batchSize, paddedBatch); - } - this.layerITGs = IntStream.range(0, config.numberOfLayers()) - .mapToObj(this::createBatchPrefillLayerTaskGraph) - .map(TaskGraph::snapshot) - .toList(); + this.layerITGs = IntStream.range(0, config.numberOfLayers()).mapToObj(this::createBatchPrefillLayerTaskGraph).map(TaskGraph::snapshot).toList(); } // @formatter:off @@ -85,42 +56,37 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph batchPrefillLayer = new TaskGraph(graphName); + TaskGraph layer = new TaskGraph(graphName); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION - batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - // Allocate persistent GPU-side intermediates once - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate GPU-side batch intermediates once. + // wrapXBatch is filled with dequantized FP32 by the host, persisted by prefillActivation. + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbFP16Batch, - state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); - // wrapXBatch produced by the prefillActivation graph and persists in device memory - // to consume it from there we should use the explicit uniqueTaskGraph name - // the no-arg form would use current graph name, which causes NPE without CUDA Graphs - batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { - // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume String pred = "batchPrefillLayer_" + (layerIndex - 1); - batchPrefillLayer.consumeFromDevice(pred, + layer.consumeFromDevice(pred, context, state.wrapXBatch, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbFP16Batch, - state.qkvResultBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); } - // Per-layer weights: upload once - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + // Per-layer weights: upload once (Q8_0 format) + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -131,164 +97,140 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w2Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray()); - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); // ── Attention Block ──────────────────────────────────────────────────── - batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceParallel, + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + dim, config.rmsNormEps()); - batchPrefillLayer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP16, - context, state.wrapXbFP16Batch, state.wrapXBatch, + // Writes FP32 normalized xb into wrapXbBatch (reused later by flash attention) + layer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, state.wrapXbBatch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. - // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the - // skinny GQA projections, and the A operand is read once, not thrice. - batchPrefillLayer.task("qkvProj", - TransformerBatchPrefillKernels::gemmMMAQKVQ8, - context, state.wrapXbFP16Batch, + layer.task("batch_qkv", + TransformerBatchPrefillKernels::batchedFusedQKVMatmulQ8, + context, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), weights.wvLayered[layerIndex].asByteArray(), - state.qkvResultBatch, paddedBatch, dim, kvDim, dim); + dim, kvDim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, + layer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCache, context, state.batchStartPosHolder, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, config.headSize(), layerIndex, config.contextLength(), dim); - // Register-partitioned P·V accumulation + direct FP16 emission - // (replaces batchedFlashAttention + attnCast). - batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + // Overwrites wrapXbBatch with attention output + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, context, state.batchStartPosHolder, - state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.attnOutFP16, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, config.numberOfHeads(), config.headSize(), kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, - context, state.attnOutFP16, + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapXbBatch, state.wrapXBatch, weights.woLayered[layerIndex].asByteArray(), - state.woOut, paddedBatch, dim, dim); + dim, dim, LOCAL_WORK_GROUP_SIZE); // ── FFN Block ────────────────────────────────────────────────────────── - // x += woOut is fused into the FFN RMS reduction (drops the woResid task). - batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, - context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - - batchPrefillLayer.task("batch_ffn_rms_apply", - TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, - context, state.normedXFFNFP16, state.wrapXBatch, + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, + context, state.wrapXBatch, state.wrapHbBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, dim); - - // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. - batchPrefillLayer.task("gateUpProj", - TransformerBatchPrefillKernels::gemmMMAGateUpQ8, - context, state.normedXFFNFP16, + state.ffnScaleBatch, weights.w1Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray(), - state.gateUpResultBatch, paddedBatch, hidDim, dim); + dim, hidDim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("swiglu", - TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, - context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) - .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, - context, state.wrapHbFP16Batch, - weights.w2Layered[layerIndex].asByteArray(), - state.w2Out, batchSize, dim, hidDim) - .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, - context, state.wrapXBatch, state.w2Out); + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asByteArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); - // Persist wrapXBatch for the next layer, and KV cache so the decode - // layers can consume it via the activation graph pass-through. - batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return batchPrefillLayer; + return layer; } // @formatter:on - // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. - static WorkerGrid mmaGrid(int paddedM, int N) { - int mBlocks = paddedM / 128; // BM - int nBlocks = N / 128; // BN - WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); - g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) - return g; - } - - static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 - WorkerGrid1D g = new WorkerGrid1D(n); - g.setLocalWork(256, 1, 1); - return g; - } - - /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); - int nHeads = config.numberOfHeads(); - int headSz = config.headSize(); - - // Parallel RMS reductions: one 256-thread workgroup per batch token - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( - batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); - - // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - // RoPE+KV cache: B*(dim/2) threads, local=512 + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + int qkvRows = dim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker(batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); int ropeGlobal = batchSize * (dim / 2); - int ropeLocal = Math.min(512, ropeGlobal); - while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) { + ropeLocal--; + } WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - - // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. - // The kernel requires headSize <= 2*localSize. - int attnLocal = Math.min(headSz, 128); - WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * attnLocal, attnLocal); - - // MMA grids - WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV - WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj - WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 - - // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) - WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid - WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + int optLocal = findOptimalLocalSize(headSz); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker(batchSize * nHeads * optLocal, optLocal); + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker(batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker(batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); - scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); - scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); - scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); - scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { + optimal = s; + break; + } + } } + return optimal; + } + + public List getLayerImmutableTaskGraphs() { + return layerITGs; } - public List getLayerImmutableTaskGraphs() { return layerITGs; } - public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } - public KernelContext getContext() { return context; } + public String getLastLayerTaskGraphID() { + return lastLayerTaskGraphID; + } + + public KernelContext getContext() { + return context; + } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java deleted file mode 100644 index d0627ef0..00000000 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillGeneric.java +++ /dev/null @@ -1,242 +0,0 @@ -package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; - -import org.beehive.gpullama3.inference.state.LlamaState; -import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; -import org.beehive.gpullama3.model.llama.LlamaConfiguration; -import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; -import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; -import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; -import uk.ac.manchester.tornado.api.GridScheduler; -import uk.ac.manchester.tornado.api.ImmutableTaskGraph; -import uk.ac.manchester.tornado.api.KernelContext; -import uk.ac.manchester.tornado.api.TaskGraph; -import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.enums.DataTransferMode; - -import java.util.List; -import java.util.stream.IntStream; - -/** - * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan (Q8_0). - * - *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} - * but uses Q8_0 kernels with inline dequantization. Key differences from the FP16 path:

- *
    - *
  • {@code wrapXBatch} is filled with dequantized FP32 embeddings by the host before - * the activation graph runs (no on-device FP16→FP32 conversion).
    • - *
  • {@code wrapXbBatch} (FP32) is reused as the normalized xb intermediate: written - * by {@code batchedRmsApplyFP32}, read by {@code batchedFusedQKVMatmulQ8}, then - * overwritten by flash attention output.
    • - *
  • {@code wrapXbFP16Batch} is not used.
    • - *
  • Weight matrices are {@code ByteArray} (Q8_0 format).
    • - *
- */ -/** - * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec - * pipeline, retained as the fallback for backends without MMA intrinsics. - * Selected automatically when the active TornadoVM - * backend is not PTX or CUDA. - */ -public class LlamaQ8_0LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { - - static final int LOCAL_WORK_GROUP_SIZE = 32; - - private final LlamaState state; - private final LlamaTornadoWeights weights; - private final LlamaConfiguration config; - private final KernelContext context = new KernelContext(); - private final int batchSize; - private final List layerITGs; - private String lastLayerTaskGraphID; - - public LlamaQ8_0LayersBatchPrefillGeneric(LlamaState state, LlamaTornadoWeights weights, LlamaConfiguration config, int batchSize) { - this.state = state; - this.weights = weights; - this.config = config; - this.batchSize = batchSize; - this.layerITGs = IntStream.range(0, config.numberOfLayers()).mapToObj(this::createBatchPrefillLayerTaskGraph).map(TaskGraph::snapshot).toList(); - } - - // @formatter:off - private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { - String graphName = "batchPrefillLayer_" + layerIndex; - if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - - TaskGraph layer = new TaskGraph(graphName); - - // ── Data Transfers ───────────────────────────────────────────────────── - if (layerIndex == 0) { - // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - // Allocate GPU-side batch intermediates once. - // wrapXBatch is filled with dequantized FP32 by the host, persisted by prefillActivation. - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - context, - state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); - } else { - String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, - context, - state.wrapXBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); - } - - // Per-layer weights: upload once (Q8_0 format) - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - weights.wqLayered[layerIndex].asByteArray(), - weights.wkLayered[layerIndex].asByteArray(), - weights.wvLayered[layerIndex].asByteArray(), - weights.woLayered[layerIndex].asByteArray(), - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - weights.w1Layered[layerIndex].asByteArray(), - weights.w2Layered[layerIndex].asByteArray(), - weights.w3Layered[layerIndex].asByteArray()); - - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); - - // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, - context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); - - // Writes FP32 normalized xb into wrapXbBatch (reused later by flash attention) - layer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP32, - context, state.wrapXbBatch, state.wrapXBatch, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - state.attnScaleBatch, dim); - - layer.task("batch_qkv", - TransformerBatchPrefillKernels::batchedFusedQKVMatmulQ8, - context, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - weights.wqLayered[layerIndex].asByteArray(), - weights.wkLayered[layerIndex].asByteArray(), - weights.wvLayered[layerIndex].asByteArray(), - dim, kvDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_rope_kv", - TransformerBatchPrefillKernels::batchedRopeWithKVCache, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapKeyCache, state.wrapValueCache, - kvDim, config.headSize(), layerIndex, config.contextLength(), dim); - - // Overwrites wrapXbBatch with attention output - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, - config.numberOfHeads(), config.headSize(), - kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); - - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapXbBatch, state.wrapXBatch, - weights.woLayered[layerIndex].asByteArray(), - dim, dim, LOCAL_WORK_GROUP_SIZE); - - // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, - context, state.wrapXBatch, state.wrapHbBatch, - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, - weights.w1Layered[layerIndex].asByteArray(), - weights.w3Layered[layerIndex].asByteArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asByteArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); - - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - - return layer; - } - // @formatter:on - - public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int kvDim = config.kvDim(); - int hidDim = config.hiddenDim(); - int nHeads = config.numberOfHeads(); - int headSz = config.headSize(); - - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - int qkvRows = dim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker(batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - int ropeGlobal = batchSize * (dim / 2); - int ropeLocal = Math.min(512, ropeGlobal); - while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) { - ropeLocal--; - } - WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - int optLocal = findOptimalLocalSize(headSz); - WorkerGrid attnWorker = WorkerGridFactory.genericWorker(batchSize * nHeads * optLocal, optLocal); - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker(batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker(batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - for (int i = 0; i < config.numberOfLayers(); i++) { - String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { - optimal = s; - break; - } - } - } - return optimal; - } - - public List getLayerImmutableTaskGraphs() { - return layerITGs; - } - - public String getLastLayerTaskGraphID() { - return lastLayerTaskGraphID; - } - - public KernelContext getContext() { - return context; - } -} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java new file mode 100644 index 00000000..24264400 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java @@ -0,0 +1,294 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; +import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan + * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). + * + *

One {@link ImmutableTaskGraph} per transformer layer, each processing + * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

+ * + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA, Q8_0 weights + * dequantized to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device + * after every layer so the subsequent single-token decode layers can consume it.

+ */ +public class LlamaQ8_0LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final LlamaState state; + private final LlamaTornadoWeights weights; + private final LlamaConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public LlamaQ8_0LayersBatchPrefillMMA(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate persistent GPU-side intermediates once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + // wrapXBatch produced by the prefillActivation graph and persists in device memory + // to consume it from there we should use the explicit uniqueTaskGraph name + // the no-arg form would use current graph name, which causes NPE without CUDA Graphs + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); + + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asByteArray(), + state.woOut, paddedBatch, dim, dim); + + // ── FFN Block ────────────────────────────────────────────────────────── + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + // Persist wrapXBatch for the next layer, and KV cache so the decode + // layers can consume it via the activation graph pass-through. + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // RoPE+KV cache: B*(dim/2) threads, local=512 + int ropeGlobal = batchSize * (dim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java index a561538e..b3db3f41 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefill.java @@ -12,55 +12,30 @@ import uk.ac.manchester.tornado.api.KernelContext; import uk.ac.manchester.tornado.api.TaskGraph; import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.WorkerGrid1D; -import uk.ac.manchester.tornado.api.WorkerGrid2D; import uk.ac.manchester.tornado.api.enums.DataTransferMode; import java.util.List; import java.util.stream.IntStream; /** - * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) - * pipeline. Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} with the Qwen3 - * architectural additions: per-head Q/K RMS normalization between the QKV - * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention - * (qDim = nHeads * headDim, which may differ from the model dim). + * Batched-prefill transformer-layer TaskGraphs for the Qwen3 Q8_0 unified batched prefill-decode plan. * - *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA, Q8_0 weights dequantized - * to FP16 in the GEMM staging registers — W8A16):

- * 
- *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
- *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
- *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
- *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
- *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
- *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
- *   woProj                MMA GEMM [dim × qDim] → woOut
- *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
- *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
- *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
- *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
- *   w2Proj                MMA GEMM → w2Out
- *   w2Resid               x += w2Out
- *
- * - *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 - * (holds for all standard Qwen3 checkpoints).

+ *

Q8_0 path: wrapXbBatch (FP32) holds normalized activations; wrapXbFP16Batch is not used. + * Mirrors {@link Qwen3FP16LayersBatchPrefill} but uses Q8_0 weights (ByteArray) and FP32 + * attention normalization path.

*/ public class Qwen3Q8_0LayersBatchPrefill implements BatchPrefillTransformerLayerTaskGraphs { - // Local size for the parallel RMS reductions (one workgroup per token). - static final int RMS_LOCAL_SIZE = 256; + static final int LOCAL_WORK_GROUP_SIZE = 32; private final Qwen3State state; private final Qwen3TornadoWeights weights; private final Qwen3Configuration config; private final KernelContext context = new KernelContext(); private final int batchSize; - // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama - // planner for the padding rationale. Non-GEMM kernels use the true batchSize. - private final int paddedBatch; private final int nHeadKv; + private final int nEmbdHeadK; + private final int nEmbdHeadV; private final int nEmbdHead; private final int qDim; private final int kvDim; @@ -74,16 +49,12 @@ public Qwen3Q8_0LayersBatchPrefill(Qwen3State state, Qwen3TornadoWeights weights this.weights = weights; this.config = config; this.batchSize = batchSize; - this.paddedBatch = (batchSize + 127) & ~127; - if (batchSize % 128 != 0) { - System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " - + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", - batchSize, paddedBatch, batchSize, paddedBatch); - } this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHead = config.numberOfHeadsValue(); - this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); - this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.nEmbdHeadK = config.numberOfHeadsKey(); + this.nEmbdHeadV = config.numberOfHeadsValue(); + this.nEmbdHead = nEmbdHeadV; + this.qDim = nEmbdHeadK * config.numberOfHeads(); + this.kvDim = nEmbdHeadV * nHeadKv; this.gqa = config.numberOfHeads() / nHeadKv; this.layerITGs = IntStream.range(0, config.numberOfLayers()) .mapToObj(this::createBatchPrefillLayerTaskGraph) @@ -96,38 +67,36 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { String graphName = "batchPrefillLayer_" + layerIndex; if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - TaskGraph batchPrefillLayer = new TaskGraph(graphName); + TaskGraph layer = new TaskGraph(graphName); int dim = config.dim(); int hidDim = config.hiddenDim(); // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { - batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, context, state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbFP16Batch, - state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); - batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, + state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("prefillActivation", state.wrapXBatch); } else { String pred = "batchPrefillLayer_" + (layerIndex - 1); - batchPrefillLayer.consumeFromDevice(pred, + layer.consumeFromDevice(pred, context, state.wrapXBatch, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbFP16Batch, - state.qkvResultBatch, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, + state.wrapHbBatch, state.wrapKeyCache, state.wrapValueCache, - state.normedXFFNFP16, state.gateUpResultBatch, - state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch); } - // Per-layer weights: upload once - batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + // Per-layer weights (Q8_0 format) + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, weights.rms_att_weightLayered[layerIndex].asFloatArray(), weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), @@ -141,160 +110,138 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { weights.w3Layered[layerIndex].asByteArray()); // ── Attention Block ──────────────────────────────────────────────────── - batchPrefillLayer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceParallel, + layer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduce, context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + dim, config.rmsNormEps()); - batchPrefillLayer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP16, - context, state.wrapXbFP16Batch, state.wrapXBatch, + // FP32 normalize into wrapXbBatch (Q8_0 path: no FP16 quantize step) + layer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP32, + context, state.wrapXbBatch, state.wrapXBatch, weights.rms_att_weightLayered[layerIndex].asFloatArray(), state.attnScaleBatch, dim); - // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) - batchPrefillLayer.task("qkvProj", - TransformerBatchPrefillKernels::gemmMMAQKVQ8, - context, state.wrapXbFP16Batch, + layer.task("batch_qkv", + Qwen3Kernels::batchedFusedQKVMatmulQ8_0, + context, + state.wrapXbBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, weights.wqLayered[layerIndex].asByteArray(), weights.wkLayered[layerIndex].asByteArray(), weights.wvLayered[layerIndex].asByteArray(), - state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); - // Qwen3: per-head RMS norm on Q and K before RoPE - batchPrefillLayer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNormPacked, - context, state.qkvResultBatch, + layer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNorm, + context, + state.wrapQBatch, state.wrapKBatch, weights.rms_att_QNormLayered[layerIndex].asFloatArray(), weights.rms_att_KNormLayered[layerIndex].asFloatArray(), config.numberOfHeads(), nHeadKv, nEmbdHead, qDim, kvDim, config.rmsNormEps()); - batchPrefillLayer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + layer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3, context, state.batchStartPosHolder, - state.qkvResultBatch, + state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, state.wrapKeyCache, state.wrapValueCache, kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - // Register-partitioned flash attention over the packed buffer. - // The 'dim' parameter doubles as the packed-Q stride base and the - // attnOutFP16 row width — both are qDim for Qwen3. - batchPrefillLayer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + // Reuses batchedFlashAttention; passes qDim as the 'dim' stride (valid: qDim==dim typically). + layer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttention, context, state.batchStartPosHolder, - state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, - state.attnOutFP16, + state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, + state.wrapXbBatch, config.numberOfHeads(), nEmbdHead, kvDim, gqa, layerIndex, config.contextLength(), qDim); - // Output projection: [M=batch, N=dim, K=qDim] - batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, - context, state.attnOutFP16, + // Output projection (Q8_0): n=qDim, d=dim + layer.task("batch_attn_out", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapXbBatch, state.wrapXBatch, weights.woLayered[layerIndex].asByteArray(), - state.woOut, paddedBatch, dim, qDim); + qDim, dim, LOCAL_WORK_GROUP_SIZE); // ── FFN Block ────────────────────────────────────────────────────────── - batchPrefillLayer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, - context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, - dim, config.rmsNormEps(), RMS_LOCAL_SIZE); - - batchPrefillLayer.task("batch_ffn_rms_apply", - TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, - context, state.normedXFFNFP16, state.wrapXBatch, + layer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedFFNRmsReduce, + context, state.wrapXBatch, state.ffnScaleBatch, + dim, config.rmsNormEps()); + + layer.task("batch_ffn_gate_up", + TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, + context, state.wrapXBatch, state.wrapHbBatch, weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, dim); - - batchPrefillLayer.task("gateUpProj", - TransformerBatchPrefillKernels::gemmMMAGateUpQ8, - context, state.normedXFFNFP16, + state.ffnScaleBatch, weights.w1Layered[layerIndex].asByteArray(), weights.w3Layered[layerIndex].asByteArray(), - state.gateUpResultBatch, paddedBatch, hidDim, dim); + dim, hidDim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.task("swiglu", - TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, - context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) - .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, - context, state.wrapHbFP16Batch, - weights.w2Layered[layerIndex].asByteArray(), - state.w2Out, paddedBatch, dim, hidDim) - .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, - context, state.wrapXBatch, state.w2Out); + layer.task("batch_ffn_down", + TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, + context, state.wrapHbBatch, state.wrapXBatch, + weights.w2Layered[layerIndex].asByteArray(), + hidDim, dim, LOCAL_WORK_GROUP_SIZE); - batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - return batchPrefillLayer; + return layer; } // @formatter:on - // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. - static WorkerGrid mmaGrid(int paddedM, int N) { - int mBlocks = paddedM / 128; // BM - int nBlocks = N / 128; // BN - WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); - g.setLocalWork(256, 1, 1); - return g; - } - - static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 - WorkerGrid1D g = new WorkerGrid1D(n); - g.setLocalWork(256, 1, 1); - return g; - } - - /** Registers all batch layer workers in the shared {@link GridScheduler}. */ public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); - int nHeads = config.numberOfHeads(); + int dim = config.dim(); + int hidDim = config.hiddenDim(); - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( - batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + int qkvRows = qDim + 2 * kvDim; + WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( + batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); - // Split-half RoPE: B*(qDim/2) threads int ropeGlobal = batchSize * (qDim / 2); int ropeLocal = Math.min(512, ropeGlobal); while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads - int attnLocal = Math.min(nEmbdHead, 128); + int optLocal = findOptimalLocalSize(nEmbdHead); WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * nHeads * attnLocal, attnLocal); + batchSize * config.numberOfHeads() * optLocal, optLocal); - // MMA grids - WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV - WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj - WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 - - WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid - WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( + batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); + WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( + batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); for (int i = 0; i < config.numberOfLayers(); i++) { String p = "batchPrefillLayer_" + i + "."; scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); - scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); - scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); - scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); - scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); + scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); + } + } + + private static int findOptimalLocalSize(int size) { + int optimal = Math.min(size, 64); + if (size % optimal != 0) { + for (int s = 64; s >= 1; s--) { + if (size % s == 0) { optimal = s; break; } + } } + return optimal; } public List getLayerImmutableTaskGraphs() { return layerITGs; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java deleted file mode 100644 index b30973c5..00000000 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillGeneric.java +++ /dev/null @@ -1,256 +0,0 @@ -package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; - -import org.beehive.gpullama3.inference.state.Qwen3State; -import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; -import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; -import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; -import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; -import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; -import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; -import uk.ac.manchester.tornado.api.GridScheduler; -import uk.ac.manchester.tornado.api.ImmutableTaskGraph; -import uk.ac.manchester.tornado.api.KernelContext; -import uk.ac.manchester.tornado.api.TaskGraph; -import uk.ac.manchester.tornado.api.WorkerGrid; -import uk.ac.manchester.tornado.api.enums.DataTransferMode; - -import java.util.List; -import java.util.stream.IntStream; - -/** - * Batched-prefill transformer-layer TaskGraphs for the Qwen3 Q8_0 unified batched prefill-decode plan. - * - *

Q8_0 path: wrapXbBatch (FP32) holds normalized activations; wrapXbFP16Batch is not used. - * Mirrors {@link Qwen3FP16LayersBatchPrefill} but uses Q8_0 weights (ByteArray) and FP32 - * attention normalization path.

- */ -/** - * Portable (non-MMA) batched-prefill planner: the pre-tensor-core matvec - * pipeline, retained as the fallback for backends without MMA intrinsics. - * Selected automatically when the active TornadoVM - * backend is not PTX or CUDA. - */ -public class Qwen3Q8_0LayersBatchPrefillGeneric implements BatchPrefillTransformerLayerTaskGraphs { - - static final int LOCAL_WORK_GROUP_SIZE = 32; - - private final Qwen3State state; - private final Qwen3TornadoWeights weights; - private final Qwen3Configuration config; - private final KernelContext context = new KernelContext(); - private final int batchSize; - private final int nHeadKv; - private final int nEmbdHeadK; - private final int nEmbdHeadV; - private final int nEmbdHead; - private final int qDim; - private final int kvDim; - private final int gqa; - private final List layerITGs; - private String lastLayerTaskGraphID; - - public Qwen3Q8_0LayersBatchPrefillGeneric(Qwen3State state, Qwen3TornadoWeights weights, - Qwen3Configuration config, int batchSize) { - this.state = state; - this.weights = weights; - this.config = config; - this.batchSize = batchSize; - this.nHeadKv = config.numberOfKeyValueHeads(); - this.nEmbdHeadK = config.numberOfHeadsKey(); - this.nEmbdHeadV = config.numberOfHeadsValue(); - this.nEmbdHead = nEmbdHeadV; - this.qDim = nEmbdHeadK * config.numberOfHeads(); - this.kvDim = nEmbdHeadV * nHeadKv; - this.gqa = config.numberOfHeads() / nHeadKv; - this.layerITGs = IntStream.range(0, config.numberOfLayers()) - .mapToObj(this::createBatchPrefillLayerTaskGraph) - .map(TaskGraph::snapshot) - .toList(); - } - - // @formatter:off - private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { - String graphName = "batchPrefillLayer_" + layerIndex; - if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; - - TaskGraph layer = new TaskGraph(graphName); - int dim = config.dim(); - int hidDim = config.hiddenDim(); - - // ── Data Transfers ───────────────────────────────────────────────────── - if (layerIndex == 0) { - layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - context, - state.attnScaleBatch, state.ffnScaleBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache); - layer.consumeFromDevice("prefillActivation", state.wrapXBatch); - } else { - String pred = "batchPrefillLayer_" + (layerIndex - 1); - layer.consumeFromDevice(pred, - context, - state.wrapXBatch, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapHbBatch, - state.wrapKeyCache, state.wrapValueCache, - state.batchStartPosHolder, - state.attnScaleBatch, state.ffnScaleBatch); - } - - // Per-layer weights (Q8_0 format) - layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - weights.wqLayered[layerIndex].asByteArray(), - weights.wkLayered[layerIndex].asByteArray(), - weights.wvLayered[layerIndex].asByteArray(), - weights.woLayered[layerIndex].asByteArray(), - weights.rms_att_QNormLayered[layerIndex].asFloatArray(), - weights.rms_att_KNormLayered[layerIndex].asFloatArray(), - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - weights.w1Layered[layerIndex].asByteArray(), - weights.w2Layered[layerIndex].asByteArray(), - weights.w3Layered[layerIndex].asByteArray()); - - // ── Attention Block ──────────────────────────────────────────────────── - layer.task("batch_attn_rms", - TransformerBatchPrefillKernels::batchedRmsReduce, - context, state.wrapXBatch, state.attnScaleBatch, - dim, config.rmsNormEps()); - - // FP32 normalize into wrapXbBatch (Q8_0 path: no FP16 quantize step) - layer.task("batch_attn_rms_apply", - TransformerBatchPrefillKernels::batchedRmsApplyFP32, - context, state.wrapXbBatch, state.wrapXBatch, - weights.rms_att_weightLayered[layerIndex].asFloatArray(), - state.attnScaleBatch, dim); - - layer.task("batch_qkv", - Qwen3Kernels::batchedFusedQKVMatmulQ8_0, - context, - state.wrapXbBatch, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - weights.wqLayered[layerIndex].asByteArray(), - weights.wkLayered[layerIndex].asByteArray(), - weights.wvLayered[layerIndex].asByteArray(), - dim, qDim, kvDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_qk_rmsnorm", - Qwen3Kernels::batchedFusedQKRmsNorm, - context, - state.wrapQBatch, state.wrapKBatch, - weights.rms_att_QNormLayered[layerIndex].asFloatArray(), - weights.rms_att_KNormLayered[layerIndex].asFloatArray(), - config.numberOfHeads(), nHeadKv, nEmbdHead, - qDim, kvDim, config.rmsNormEps()); - - layer.task("batch_rope_kv", - Qwen3Kernels::batchedRopeWithKVCacheQwen3, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKBatch, state.wrapVBatch, - state.wrapKeyCache, state.wrapValueCache, - kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); - - // Reuses batchedFlashAttention; passes qDim as the 'dim' stride (valid: qDim==dim typically). - layer.task("batch_attention", - TransformerBatchPrefillKernels::batchedFlashAttention, - context, state.batchStartPosHolder, - state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache, - state.wrapXbBatch, - config.numberOfHeads(), nEmbdHead, - kvDim, gqa, layerIndex, config.contextLength(), qDim); - - // Output projection (Q8_0): n=qDim, d=dim - layer.task("batch_attn_out", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapXbBatch, state.wrapXBatch, - weights.woLayered[layerIndex].asByteArray(), - qDim, dim, LOCAL_WORK_GROUP_SIZE); - - // ── FFN Block ────────────────────────────────────────────────────────── - layer.task("batch_ffn_rms", - TransformerBatchPrefillKernels::batchedFFNRmsReduce, - context, state.wrapXBatch, state.ffnScaleBatch, - dim, config.rmsNormEps()); - - layer.task("batch_ffn_gate_up", - TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUpQ8, - context, state.wrapXBatch, state.wrapHbBatch, - weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), - state.ffnScaleBatch, - weights.w1Layered[layerIndex].asByteArray(), - weights.w3Layered[layerIndex].asByteArray(), - dim, hidDim, LOCAL_WORK_GROUP_SIZE); - - layer.task("batch_ffn_down", - TransformerBatchPrefillKernels::batchedMatVecWithResidualQ8, - context, state.wrapHbBatch, state.wrapXBatch, - weights.w2Layered[layerIndex].asByteArray(), - hidDim, dim, LOCAL_WORK_GROUP_SIZE); - - layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); - - return layer; - } - // @formatter:on - - public void updateGridScheduler(GridScheduler scheduler) { - int dim = config.dim(); - int hidDim = config.hiddenDim(); - - WorkerGrid rmsWorker = WorkerGridFactory.genericWorker(batchSize, 1); - WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); - - int qkvRows = qDim + 2 * kvDim; - WorkerGrid qkvWorker = WorkerGridFactory.genericWorker( - batchSize * qkvRows * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( - batchSize * (config.numberOfHeads() + nHeadKv) * nEmbdHead, nEmbdHead); - - int ropeGlobal = batchSize * (qDim / 2); - int ropeLocal = Math.min(512, ropeGlobal); - while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; - WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); - - int optLocal = findOptimalLocalSize(nEmbdHead); - WorkerGrid attnWorker = WorkerGridFactory.genericWorker( - batchSize * config.numberOfHeads() * optLocal, optLocal); - - WorkerGrid matVecDimWorker = WorkerGridFactory.genericWorker( - batchSize * dim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - WorkerGrid matVecHidWorker = WorkerGridFactory.genericWorker( - batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE); - - for (int i = 0; i < config.numberOfLayers(); i++) { - String p = "batchPrefillLayer_" + i + "."; - scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); - scheduler.addWorkerGrid(p + "batch_qkv", qkvWorker); - scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); - scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); - scheduler.addWorkerGrid(p + "batch_attention", attnWorker); - scheduler.addWorkerGrid(p + "batch_attn_out", matVecDimWorker); - scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); - scheduler.addWorkerGrid(p + "batch_ffn_gate_up", matVecHidWorker); - scheduler.addWorkerGrid(p + "batch_ffn_down", matVecDimWorker); - } - } - - private static int findOptimalLocalSize(int size) { - int optimal = Math.min(size, 64); - if (size % optimal != 0) { - for (int s = 64; s >= 1; s--) { - if (size % s == 0) { optimal = s; break; } - } - } - return optimal; - } - - public List getLayerImmutableTaskGraphs() { return layerITGs; } - public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } - public KernelContext getContext() { return context; } -} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java new file mode 100644 index 00000000..1cd02ee4 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java @@ -0,0 +1,304 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefillMMA} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). + * + *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA, Q8_0 weights dequantized + * to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

+ */ +public class Qwen3Q8_0LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final int nHeadKv; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3Q8_0LayersBatchPrefillMMA(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asByteArray(), + state.woOut, paddedBatch, dim, qDim); + + // ── FFN Block ────────────────────────────────────────────────────────── + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + + // Split-half RoPE: B*(qDim/2) threads + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java index 75bb1490..a4490d88 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java @@ -14,8 +14,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LlamaFP16FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LlamaFP16FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LogitsFP16LayerDecode; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill; -import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillGeneric; import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; @@ -84,9 +84,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { if (TensorCoreSupport.isTensorCoreCapableBackend()) { - return new LlamaFP16LayersBatchPrefill(state, weights, config, batchSize); + return new LlamaFP16LayersBatchPrefillMMA(state, weights, config, batchSize); } - return new LlamaFP16LayersBatchPrefillGeneric(state, weights, config, batchSize); + return new LlamaFP16LayersBatchPrefill(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java index 00521d7a..7a8008f8 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java @@ -14,8 +14,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LogitsFP16LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersPrefillDecode; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefill; -import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefillGeneric; import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; @@ -79,9 +79,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { if (TensorCoreSupport.isTensorCoreCapableBackend()) { - return new Qwen3FP16LayersBatchPrefill(state, weights, config, batchSize); + return new Qwen3FP16LayersBatchPrefillMMA(state, weights, config, batchSize); } - return new Qwen3FP16LayersBatchPrefillGeneric(state, weights, config, batchSize); + return new Qwen3FP16LayersBatchPrefill(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java index 35e498e9..bd7402f2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java @@ -14,8 +14,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LlamaQ8_0FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LlamaQ8_0FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefill; -import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefillGeneric; import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; @@ -88,9 +88,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { if (TensorCoreSupport.isTensorCoreCapableBackend()) { - return new LlamaQ8_0LayersBatchPrefill(state, weights, config, batchSize); + return new LlamaQ8_0LayersBatchPrefillMMA(state, weights, config, batchSize); } - return new LlamaQ8_0LayersBatchPrefillGeneric(state, weights, config, batchSize); + return new LlamaQ8_0LayersBatchPrefill(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java index 1baf0cf1..e4d3b6e2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java @@ -14,8 +14,8 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersPrefillDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefill; -import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefillGeneric; import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; @@ -79,9 +79,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { if (TensorCoreSupport.isTensorCoreCapableBackend()) { - return new Qwen3Q8_0LayersBatchPrefill(state, weights, config, batchSize); + return new Qwen3Q8_0LayersBatchPrefillMMA(state, weights, config, batchSize); } - return new Qwen3Q8_0LayersBatchPrefillGeneric(state, weights, config, batchSize); + return new Qwen3Q8_0LayersBatchPrefill(state, weights, config, batchSize); } // ── Logits layers ───────────────────────────────────────────────────────── From f91cd59fbf59b9528c5b6f4056d6a6a7d68c80d7 Mon Sep 17 00:00:00 2001 From: MaryXek Date: Tue, 7 Jul 2026 14:00:26 +0300 Subject: [PATCH 17/18] Minor fixes in the javadocs --- .../layers/BatchPrefillTransformerLayerTaskGraphs.java | 2 +- .../layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java index a46bbbae..9818eaa5 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java @@ -8,7 +8,7 @@ /** * Interface for a group of N batched-prefill transformer-layer TornadoVM TaskGraphs. * - *

Implemented by {@code LlamaFP16LayersBatchPrefillMMA} and {@code LlamaQ8_0LayersBatchPrefillMMA}.

+ *

Implemented by {@code LlamaFP16LayersBatchPrefillMMA}, {@code LlamaFP16LayersBatchPrefill}, {@code LlamaQ8_0LayersBatchPrefillMMA} and {@code LlamaQ8_0LayersBatchPrefill}.

*/ public interface BatchPrefillTransformerLayerTaskGraphs { List getLayerImmutableTaskGraphs(); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java index a3016cb7..1c739f4e 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefill.java @@ -19,7 +19,7 @@ /** * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan (Q8_0). * - *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA} + *

Mirrors {@link org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill} * but uses Q8_0 kernels with inline dequantization. Key differences from the FP16 path:

*
    *
  • {@code wrapXBatch} is filled with dequantized FP32 embeddings by the host before From 3e676e9edde3312bd0501a5bd8b286fa9b01e095 Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Fri, 10 Jul 2026 13:40:26 +0100 Subject: [PATCH 18/18] Update TornadoVM version to latest release --- README.md | 2 +- pom.xml | 9 ++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 9cb3a0a9..c84b57dc 100644 --- a/README.md +++ b/README.md @@ -67,7 +67,7 @@ Ensure you have the following installed and configured: - **Java 21**: Required for Vector API support & TornadoVM. - [TornadoVM](https://github.com/beehive-lab/TornadoVM) with OpenCL, PTX, or CUDA backends. - - The `--cuda` backend requires a TornadoVM build that includes the CUDA backend from [TornadoVM PR #861](https://github.com/beehive-lab/TornadoVM/pull/861). This project currently builds against TornadoVM `4.0.2-jdk21-dev`. + - The `--cuda` backend requires a TornadoVM build that includes the CUDA backend from [TornadoVM PR #861](https://github.com/beehive-lab/TornadoVM/pull/861). This project currently builds against TornadoVM `5.0.0-jdk21-dev`. - GCC/G++ 13 or newer: Required to build and run TornadoVM native components. ### Install, Build, and Run diff --git a/pom.xml b/pom.xml index 82e875e9..10361b06 100644 --- a/pom.xml +++ b/pom.xml @@ -39,10 +39,10 @@ 0.4.0 - 4.0.2 + 5.0.0 -jdk21 - - ${tornadovm.base.version}${jdk.version.suffix}-dev + + ${tornadovm.base.version}${jdk.version.suffix} 25 25 @@ -148,8 +148,7 @@ 21 21 -jdk21 - - ${tornadovm.base.version}${jdk.version.suffix}-dev + ${tornadovm.base.version}${jdk.version.suffix}