cuBLAS and cuBLASLt: GPU Linear Algebra Tasks in Java

TornadoVM’s cuBLAS integration brings NVIDIA’s production-grade dense linear algebra library into the Java TaskGraph model. Rather than calling cuBLAS through a separate JNI layer with manually managed device memory, you express GEMM and GEMV operations as library tasks — first-class TaskGraph nodes that share TornadoVM-managed device buffers with your JIT-compiled kernels on one CUDA stream. The cuBLASLt variant extends this with fused epilogues (BIAS, GELU_BIAS) that collapse a GEMM plus a separate activation kernel into a single highly-optimized dispatch, which is especially valuable in the launch-bound regime of LLM token generation.

Prerequisite: CUDA Toolkit with cuBLAS (libcublas). cuBLAS ships inside the CUDA Toolkit — no separate package is needed. Build TornadoVM with make BACKEND=cuda.

cuBLAS Factory Methods

All factories are static methods on uk.ac.manchester.tornado.cublas.CuBlas. Pass them as the second argument to taskGraph.libraryTask(id, factory, args...).

Factory	cuBLAS Function	Semantics
`cublasSgemv(op, m, n, α, A, lda, x, incx, β, y, incy)`	`cublasSgemv`	`y = α·op(A)·x + β·y`
`cublasSgemm(opA, opB, m, n, k, α, A, lda, B, ldb, β, C, ldc)`	`cublasSgemm`	`C = α·op(A)·op(B) + β·C`
`cublasSgemmTF32(...)`	`cublasSgemm` + TF32 math mode	Same as SGEMM, executed on TF32 Tensor Cores (~1e-4 rel. error)
`cublasGemmExFP16(...)`	`cublasGemmEx`	FP16 inputs, FP16 output, FP32 Tensor Core accumulation
`cublasGemmExFP16FP32(...)`	`cublasGemmEx`	FP16 inputs, FP32 output, FP32 Tensor Core accumulation
`cublasGemmExBF16(...)`	`cublasGemmEx`	BF16 inputs and output, FP32 accumulation (`BFloat16Array`)
`cublasSgemmStridedBatched(...)`	`cublasSgemmStridedBatched`	`C[i] = α·op(A[i])·op(B[i]) + β·C[i]` over a flat-array batch

cuBLASLt Factory Methods

Fused-epilogue GEMM is available through uk.ac.manchester.tornado.cublas.CuBlasLt. Plans (descriptors + heuristic-selected algorithm) are created once per problem shape and cached with a 32 MiB device workspace.

Factory	Epilogue	Notes
`ltMatmulFP32(opA, opB, m, n, k, α, A, lda, B, ldb, β, C, ldc)`	None	FP32 matmul with plan caching
`ltMatmulFP16(...)`	None	FP16 matmul, FP32 Tensor Core accumulation
`ltMatmulFP8(...)`	None	E4M3 operands, FP16 output, TN layout, `ld` must be multiple of 16 B
`ltMatmulBiasFP16(..., bias)`	`BIAS`	`C = op(A)·op(B) + bias`, fused
`ltMatmulGeluBiasFP16(..., bias)`	`GELU_BIAS`	`C = GELU(op(A)·op(B) + bias)`, tanh approximation, fully fused

Column-Major Layout

cuBLAS is column-major. TornadoVM arrays are row-major. The standard tricks used throughout the tests:

SGEMV (row-major W·x)
SGEMM (row-major C = A·B)

A row-major weight matrix W of shape (d × n) is the column-major matrix (n × d) with lda = n. Compute y = W·x using the transpose op:

// W is stored row-major as (d rows × n cols)
// cuBLAS sees it as column-major (n × d) — pass CUBLAS_OP_T and swap m/n
taskGraph.libraryTask("sgemv", CuBlas::cublasSgemv,
        CuBlasOperation.CUBLAS_OP_T.operation(),
        n,           // rows of op(A) = cols of the row-major W
        d,           // cols of op(A) = rows of the row-major W
        1.0f,
        W, n,        // lda = n (leading dimension of the row-major layout)
        x, 1,
        0.0f, y, 1);

Row-major C = A·B is mathematically equivalent to column-major C_cm = B_cm · A_cm. Swap the operands and pass CUBLAS_OP_N for both:

// Row-major A (m × k), B (k × n) → C (m × n)
// cuBLAS column-major trick: swap A↔B, pass OP_N, OP_N
taskGraph.libraryTask("sgemm", CuBlas::cublasSgemm,
        CuBlasOperation.CUBLAS_OP_N.operation(),
        CuBlasOperation.CUBLAS_OP_N.operation(),
        n, m, k,         // note: n first (cuBLAS convention)
        1.0f,
        matrixB, n,      // B comes first
        matrixA, k,      // then A
        0.0f, output, n);

When β != 0, cuBLAS reads the output operand before writing. The factory marks it READ_WRITE automatically. Include it in transferToDevice if its initial values come from the host (see TestCuBlasSgemvBeta).

Complete SGEMM Example

The following example runs a complete matrix multiply through a cuBLAS library task inside a TaskGraph, cross-validates against a sequential Java result, and measures throughput.

import uk.ac.manchester.tornado.api.*;
import uk.ac.manchester.tornado.api.enums.DataTransferMode;
import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
import uk.ac.manchester.tornado.cublas.CuBlas;
import uk.ac.manchester.tornado.cublas.enums.CuBlasOperation;

int size = 2048;
FloatArray matrixA = new FloatArray(size * size);
FloatArray matrixB = new FloatArray(size * size);
FloatArray output  = new FloatArray(size * size);

// ... fill matrixA and matrixB with data ...

// Row-major C = A·B using the operand-swap trick
TaskGraph graph = new TaskGraph("sgemm")
    .transferToDevice(DataTransferMode.FIRST_EXECUTION, matrixA, matrixB)
    .libraryTask("sgemm", CuBlas::cublasSgemm,
            CuBlasOperation.CUBLAS_OP_N.operation(),
            CuBlasOperation.CUBLAS_OP_N.operation(),
            size, size, size,
            1.0f, matrixB, size, matrixA, size,
            0.0f, output,  size)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, output);

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(graph.snapshot())) {
    plan.execute();
}

Run the built-in benchmark (RTX 4090 reference numbers below):

# args: [size] [iterations]
tornado -m tornado.cublas/uk.ac.manchester.tornado.cublas.tests.BenchmarkSgemm 2048 50

Performance Reference

All numbers below are from an NVIDIA GeForce RTX 4090 with CUDA 12.6.

SGEMM Throughput

Size	JIT Kernel	cuBLAS FP32	cuBLAS TF32	cuBLAS FP16
1024	3.9 TFLOP/s	26.1 TFLOP/s	28.2 TFLOP/s	44.1 TFLOP/s
2048	6.0 TFLOP/s	48.3 TFLOP/s	62.8 TFLOP/s	120.6 TFLOP/s
4096	5.8 TFLOP/s	57.0 TFLOP/s	81.1 TFLOP/s	160.6 TFLOP/s

160.6 TFLOP/s is the RTX 4090’s peak FP16 Tensor Core throughput — reached from Java by changing only the factory name.

Fused Epilogue Speedup (cuBLASLt)

BenchmarkLtFusedMlp: fused GELU_BIAS vs. unfused GemmEx FP16 + JIT kernel.

Size	Unfused	Fused	Speedup
1024	17.5 TFLOP/s	29.4 TFLOP/s	1.68×
2048	73.8 TFLOP/s	98.3 TFLOP/s	1.33×
4096	132.1 TFLOP/s	147.3 TFLOP/s	1.11×

Fusion pays most at small/medium sizes — the launch-bound regime of LLM decoding.

TF32 and FP16 Tensor Cores

TornadoVM exposes Tensor Core acceleration through named factory variants — no code restructuring required.

TF32 Tensor Cores
FP16 Tensor Cores
FP16 cuBLASLt (fused)

TF32 uses the same FloatArray inputs as FP32 but routes through Tensor Cores via the CUBLAS_TF32_TENSOR_OP_MATH math mode. Expect ~1e-4 relative error compared to FP32.

// Switch from FP32 to TF32 by changing only the factory name
taskGraph.libraryTask("sgemm_tf32", CuBlas::cublasSgemmTF32,
        CuBlasOperation.CUBLAS_OP_N.operation(),
        CuBlasOperation.CUBLAS_OP_N.operation(),
        n, m, k, 1.0f, matrixB, n, matrixA, k, 0.0f, output, n);

cublasGemmExFP16 expects HalfFloatArray inputs. Use new HalfFloat(float) to populate and .getFloat32() to read back.

import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray;
import uk.ac.manchester.tornado.api.types.HalfFloat;

HalfFloatArray aFP16 = new HalfFloatArray(size * size);
HalfFloatArray bFP16 = new HalfFloatArray(size * size);
HalfFloatArray cFP16 = new HalfFloatArray(size * size);

// ... populate with HalfFloat values ...

taskGraph.libraryTask("hgemm", CuBlas::cublasGemmExFP16,
        CuBlasOperation.CUBLAS_OP_N.operation(),
        CuBlasOperation.CUBLAS_OP_N.operation(),
        n, m, k, 1.0f, bFP16, n, aFP16, k, 0.0f, cFP16, n);

The cuBLASLt path adds bias and activation fusion on top of FP16 Tensor Core GEMMs:

// C = GELU(A * B + bias), one fused cuBLASLt call
taskGraph.libraryTask("mlp", CuBlasLt::ltMatmulGeluBiasFP16,
        CuBlasOperation.CUBLAS_OP_N.operation(),
        CuBlasOperation.CUBLAS_OP_N.operation(),
        n, m, k,
        1.0f, bFP16, n, aFP16, k, 0.0f, cFP16, n,
        bias);  // bias vector (one entry per output feature)

Strided Batched GEMM

One call launches a whole batch of equally-shaped GEMMs stored as contiguous slices in flat arrays.

int batchCount = 32;
long stride = (long) size * size;   // elements between consecutive matrices in the flat array

taskGraph.libraryTask("batched", CuBlas::cublasSgemmStridedBatched,
        CuBlasOperation.CUBLAS_OP_N.operation(),
        CuBlasOperation.CUBLAS_OP_N.operation(),
        size, size, size,
        1.0f,
        matrixB, size, stride,
        matrixA, size, stride,
        0.0f,
        matrixC, size, stride,
        batchCount);

Unit Tests and Runnable Examples

# Full cuBLAS test suite (auto-skips without CUDA backend)
tornado-test -V uk.ac.manchester.tornado.unittests.cublas.TestCuBlas

# cuBLASLt fused epilogue tests
tornado-test -V uk.ac.manchester.tornado.unittests.cublas.TestCuBlasLt

TestCuBlas#testSharedBufferAcrossTaskGraphs demonstrates the cross-graph shared-buffer pattern: a JIT task graph produces the matrix on the device, and a second graph consumes it with a cuBLAS call — the data never returns to the host between the two graphs.

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

cuBLAS and cuBLASLt: GPU Linear Algebra Tasks in Java

cuBLAS Factory Methods

cuBLASLt Factory Methods

Column-Major Layout

Complete SGEMM Example

Performance Reference

SGEMM Throughput

Fused Epilogue Speedup (cuBLASLt)

TF32 and FP16 Tensor Cores

Strided Batched GEMM

Unit Tests and Runnable Examples

Build docs developers (and LLMs) love

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Documentation Index

​cuBLAS Factory Methods

​cuBLASLt Factory Methods

​Column-Major Layout

​Complete SGEMM Example

​Performance Reference

SGEMM Throughput

Fused Epilogue Speedup (cuBLASLt)

​TF32 and FP16 Tensor Cores

​Strided Batched GEMM

​Unit Tests and Runnable Examples

Build docs developers (and LLMs) love

cuBLAS Factory Methods

cuBLASLt Factory Methods

Column-Major Layout

Complete SGEMM Example

Performance Reference

TF32 and FP16 Tensor Cores

Strided Batched GEMM

Unit Tests and Runnable Examples