Add TornadoVM CUDA backend support and tensor-core (MMA) accelerated batch prefill (FP16 & Q8_0, Llama & Qwen3)

.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)'

.github/actions/setup-tornadovm/action.yml

@@ -3,7 +3,7 @@ description: Build TornadoVM develop once per backend and reuse across runs via
 
 inputs:
   backend:
-    description: 'TornadoVM backend to build (opencl or ptx)'
+    description: 'TornadoVM backend to build (opencl, ptx, or cuda)'
     required: true
 
 runs:

.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]
     timeout-minutes: 30
 
     steps:
@@ -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
@@ -106,6 +107,7 @@ jobs:
         backend:
           - name: opencl
           - name: ptx
+          - name: cuda
 
     steps:
       - name: Checkout GPULlama3
@@ -530,6 +532,7 @@ jobs:
         backend:
           - name: opencl
           - name: ptx
+          - name: cuda
 
     steps:
       - name: Checkout GPULlama3

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 `5.0.0-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)
 

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",

pom.xml

@@ -39,8 +39,9 @@
         <properties>
             <!-- CI-friendly version: resolved by flatten-maven-plugin at build time -->
             <revision>0.5.0</revision>
-            <tornadovm.base.version>4.0.1</tornadovm.base.version>
+            <tornadovm.base.version>5.0.0</tornadovm.base.version>
             <jdk.version.suffix>-jdk21</jdk.version.suffix>
+            <!-- CUDA backend is only available after 5.0.0 TornadoVM version -->
             <tornadovm.version>${tornadovm.base.version}${jdk.version.suffix}</tornadovm.version>
             <!-- Compiler defaults (overridden by JDK profiles below) -->
             <maven.compiler.source>25</maven.compiler.source>

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

@@ -85,6 +85,16 @@ 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;
+    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;
@@ -135,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);
@@ -148,6 +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(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(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 * (qDim + 2 * kvDim));
+            this.gateUpResultBatch = new FloatArray(paddedGpuBatch * 2 * config.hiddenDim());
         } else {
             this.embeddingXBatch = null;
             this.wrapXBatch = null;
@@ -160,6 +186,16 @@ 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;
+            this.qkvResultBatch = null;
+            this.gateUpResultBatch = null;
         }
     }
 

src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java

@@ -0,0 +1,20 @@
+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
+ * CUDA backend.
+ */
+public final class TensorCoreSupport {
+
+    public static boolean isTensorCoreCapableBackend() {
+        TornadoVMBackendType backendType = TornadoRuntimeProvider.getTornadoRuntime()
+                .getBackend(0)
+                .getBackendType();
+        return backendType == TornadoVMBackendType.CUDA;
+    }
+}