TornadoVM Transformer Optimizations for GPU Inference in gpullama3.java

TornadoVM accelerates transformer inference by fusing kernels, quantizing weights to 8-bit, applying Flash Attention on NVIDIA GPUs, and orchestrating execution through a reusable GridScheduler, yielding multi-fold speedups over standard Java execution.

The gpullama3.java repository implements a high-performance LLM inference engine that leverages TornadoVM to compile Java code directly to OpenCL and CUDA kernels. These transformer optimizations eliminate memory bottlenecks and kernel launch overhead by collapsing multiple logical operations into single GPU passes.

Kernel Fusion and Memory Optimization

TornadoVM minimizes device memory traffic by fusing separate computational steps into unified kernels. This approach eliminates intermediate buffers and reduces PCIe transfer overhead.

Fused RMS-Norm and QKV Projection

Instead of launching separate kernels for normalization, dequantization, and matrix multiplication, TransformerComputeKernelsLayered.java combines these into a single operation:

public static void fusedRmsNormQKVMatmulQ8_0(
        KernelContext ctx, FloatArray x, FloatArray hb,
        FloatArray rmsWeights, FloatArray rmsScale,
        ByteArray w1, ByteArray w3,
        int inputDim, int hiddenDim, int localWorkGroupSize) {
    // RMS-norm, de-quantise Q8_0 weights, mat-mul in one pass
}

The Qwen3Kernels.java file provides a similar implementation in fusedRmsNormQKVMatmul for FP16 models. This fusion eliminates three separate kernel launches and associated global memory writes.

Fused FFN Gate and Up Projection

The feed-forward network (FFN) layer combines RMS-normalization, two parallel matrix multiplications (gate and up projections), SiLU activation, and GLU gating into one kernel:

// TransformerComputeKernelsLayered.java:fusedRmsNormFFNGateUpQ8_0
// Handles Q8_0 quantized weights

For FP16 execution, TransformerComputeKernelsLayered.java provides fusedRmsNormFFNGateUp, which uses HalfFloatArray to maintain 16-bit precision throughout the computation while accumulating in FP32 for numerical stability.

Fused Bias Addition and RoPE Rotation

The repository further optimizes attention mechanisms by fusing bias vector addition with rotary positional embedding (RoPE) calculations:

• TransformerComputeKernelsLayered.java:fusedQKvBiasAddition adds Q, K, and V bias vectors in a single pass before projection.
• TransformerComputeKernelsLayered.java:ropeRotationWithCacheCopy rotates queries and keys, writes rotated keys directly into the KV-cache, and copies values, eliminating a separate copy-to-cache kernel.

Quantization and Precision Optimization

Memory bandwidth often bottlenecks transformer inference. The implementation addresses this through aggressive quantization and mixed-precision arithmetic.

Q8_0 On-the-Fly Dequantization

The Q8_0 format stores weights in 34-byte blocks containing a 2-byte FP16 scale followed by 32 quantized bytes. During inference, TransformerComputeKernelsLayered.java:convertQ8_0toFP32 dequantizes values inside the mat-mul kernel:

HalfFloat scale = x.getHalfFloat(blockByteOffset);
byte quant = x.get(blockByteOffset + 2 + withinBlockIdx);
float dequant = ((float) quant) * scale.getFloat32();

This approach halves memory bandwidth requirements compared to FP16 storage while maintaining kernel fusion benefits.

FP16 Half-Precision Tensors

For models requiring higher precision, the FP16 path uses HalfFloatArray to store activations in 16 bits. Conversion kernels convertFP32toFP16 and convertFP16toFP32 handle host-device transfers, while compute kernels like those in LlamaFP16FFNLayers.java perform accumulation in FP32 for numerical safety before storing results back as FP16.

Attention Optimization

Flash Attention for NVIDIA GPUs

On NVIDIA hardware, the implementation activates Flash Attention kernels that avoid materializing the full Q·Kᵀ attention matrix. The AbstractFFNLayers.java file controls this optimization:

protected boolean shouldUseFinalNormalization() {
    // NVIDIA → Flash-Attention (skip final norm)
    return schedulerType == SchedulerType.NON_NVIDIA;
}

When SchedulerType.NVIDIA is detected, the system selects processHeadsFlashAttention* kernels that operate on tiles of 16–32 elements, loading only required Q/K blocks into shared memory. This reduces memory traffic from O(N²) to O(N·tile), yielding 2–3× speedups on RTX-class GPUs.

Scheduling and Runtime Optimization

GridScheduler Task-Graph Caching

The Q8_0LayerPlanner.java and FP16LayerPlanner.java classes eliminate per-inference scheduling overhead by building a master GridScheduler once during model initialization:

GridScheduler masterScheduler = new GridScheduler();
// ... add worker grids for each layer ...
allTaskGraphs.add(activationLayer.getImmutableTaskGraph());
this.cachedTaskGraphs = allTaskGraphs;
this.cachedScheduler = masterScheduler;

Subsequent token generations reuse the cached scheduler and immutable task graphs, removing the >100ms graph construction penalty typically incurred by dynamic GPU programming frameworks.

Parallel Reduction for RMS-Norm

The RMS-normalization implementation uses a two-phase parallel reduction to compute the scaling factor with minimal synchronization. The TransformerComputeKernelsLayered.java:reductionOneBlockWithLayer method implements a tree reduction:

for (int stride = (groupSize/2); stride > 0; stride /= 2) {
    context.localBarrier();
    if (lid < stride) localX[lid] += localX[lid+stride];
}

Local work groups first reduce values internally, then the first thread aggregates partial sums across groups. This approach maximizes occupancy while ensuring accurate RMS statistics.

Java Vector API Fallback

When GPU acceleration is unavailable, the system falls back to the Java Vector API for SIMD acceleration. The LlamaApp.java entry point checks the USE_VECTOR_API flag:

public static final boolean USE_VECTOR_API =
    Boolean.parseBoolean(System.getProperty("llama.VectorAPI", "true"));

This fallback path utilizes FloatVector operations to exploit AVX-512 or NEON instructions on modern CPUs, ensuring acceptable performance even without discrete GPU hardware.

Summary

• Kernel fusion collapses RMS-norm, mat-mul, activation, and caching into single GPU kernels, eliminating intermediate memory copies in TransformerComputeKernelsLayered.java.
• Flash Attention activates on NVIDIA GPUs via AbstractFFNLayers.java, reducing attention complexity from O(N²) to O(N·tile) through tiled shared-memory operations.
• Q8_0 quantization halves memory bandwidth by storing 8-bit quantized weights with FP16 scales, dequantized on-the-fly during matrix multiplication.
• GridScheduler caching in Q8_0LayerPlanner.java and FP16LayerPlanner.java removes per-inference scheduling overhead by reusing immutable task graphs.
• Java Vector API fallback in LlamaApp.java provides SIMD acceleration on CPU-only systems when TornadoVM GPU targets are unavailable.

Frequently Asked Questions

What is TornadoVM and how does it optimize transformer inference?

TornadoVM is a Java framework that compiles Java bytecode to OpenCL, CUDA, or SPIR-V kernels for heterogeneous hardware acceleration. In gpullama3.java, TornadoVM applies transformer optimizations by fusing multiple neural network operations into single GPU kernels, quantizing weights to reduce memory bandwidth, and using hardware-specific schedulers to minimize kernel launch overhead.

How does kernel fusion improve performance compared to standard layer execution?

Kernel fusion eliminates the need to write intermediate results back to global memory between operations. For example, the fusedRmsNormQKVMatmulQ8_0 kernel in TransformerComputeKernelsLayered.java performs RMS-normalization, weight dequantization, and matrix multiplication in a single pass. This reduces memory traffic by 3–4× compared to executing separate kernels for each step, directly improving inference latency on bandwidth-constrained GPUs.

What is Flash Attention and why is it limited to NVIDIA GPUs?

Flash Attention is an algorithm that computes the attention mechanism in blocks without materializing the full N×N attention matrix, reducing memory complexity from O(N²) to O(N). The gpullama3.java implementation activates Flash Attention via AbstractFFNLayers.java when SchedulerType.NVIDIA is detected because it relies on specific shared-memory optimizations and warp primitives that are currently optimized for NVIDIA's CUDA architecture in this codebase. Support for other GPU vendors may require different tile size optimizations or memory barrier strategies.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}