cuDNN: Deep Learning Primitives in a Java TaskGraph

TornadoVM’s cuDNN integration exposes NVIDIA’s deep-learning primitive library through the familiar TaskGraph API. Convolution layers, activation functions, pooling operations, and fused scaled-dot-product attention (flash attention) become library tasks that share TornadoVM-managed device buffers with adjacent JIT kernels on one CUDA stream. You can build a complete CNN forward pass — convolution, bias addition, activation, and pooling — as a single TaskGraph where each primitive runs with cuDNN’s highly-tuned implementation and the intermediate feature maps never leave the GPU.

Prerequisite: cuDNN 9 must be installed separately. It is not bundled with the CUDA Toolkit.

apt install libcudnn9-cuda-12 libcudnn9-dev-cuda-12

Build TornadoVM afterwards with make BACKEND=cuda.

Factory Methods

All factories are static methods on uk.ac.manchester.tornado.cudnn.CuDnn.

Activations

Factory	cuDNN Call	Notes
`cudnnSoftmax(in, out, rows, cols)`	`cudnnSoftmaxForward` (ACCURATE, INSTANCE)	Numerically stable per-row softmax; the standard attention-score shape
`cudnnRelu(in, out, size)`	`cudnnActivationForward` (RELU)	Element-wise ReLU
`cudnnSigmoid(in, out, size)`	`cudnnActivationForward` (SIGMOID)	Element-wise sigmoid
`cudnnTanh(in, out, size)`	`cudnnActivationForward` (TANH)	Element-wise tanh

Pooling

Factory	cuDNN Call	Notes
`cudnnMaxPool2d(in, out, n, c, h, w, window, stride)`	`cudnnPoolingForward` (MAX)	Square-window max pooling; no padding; input/output in NCHW layout

Convolution

Factory	cuDNN Call	Notes
`cudnnConv2d(in, filter, out, n, c, h, w, k, r, s, pad, stride)`	`cudnnConvolutionForward` (IMPLICIT_PRECOMP_GEMM)	2D cross-correlation; input NCHW; filter KCRS; square pad/stride

Flash Attention (SDPA)

Factory	cuDNN Call	Notes
`sdpaForward(q, k, v, o, b, h, sQ, sKv, d, scale, causal)`	cuDNN graph API fused SDPA	FP16 `HalfFloatArray`, packed BHSD layout, FP32 accumulate, inference only; `d` must be a multiple of 8 and ≤ 256; requires Ampere or newer

Data Layout

FP32 / NCHW

Activations, pooling, and convolution all use FP32 FloatArray with NCHW layout: batch dimension (N), then channels (C), then spatial height (H), then width (W). TornadoVM flat arrays map directly to this layout — element (n, c, h, w) is at index n*C*H*W + c*H*W + h*W + w.

FP16 / BHSD (Flash Attention)

The sdpaForward factory uses FP16 HalfFloatArray with BHSD layout: batch (B), heads (H), sequence length (S), head dimension (D). The head dimension d must be a multiple of 8 and no larger than 256. Requires an Ampere or newer GPU.

CNN Layer Example

The following example builds a complete CNN block — convolution, JIT bias addition, ReLU activation, and max pooling — as a single TaskGraph. All intermediate feature maps stay on the device.

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.cudnn.CuDnn;

// Input: 8 images, 64 channels, 56×56 spatial
int n = 8, c = 64, h = 56, w = 56;
// Filter: 64 output channels, 3×3 kernel
int k = 64, r = 3, s = 3, pad = 1, stride = 1;
int outH = (h + 2*pad - r) / stride + 1;  // = 56 with pad=1, stride=1
int outW = (w + 2*pad - s) / stride + 1;  // = 56

FloatArray input   = new FloatArray(n * c * h * w);
FloatArray filter  = new FloatArray(k * c * r * s);
FloatArray convOut = new FloatArray(n * k * outH * outW);
FloatArray bias    = new FloatArray(k);
FloatArray reluOut = new FloatArray(n * k * outH * outW);
FloatArray pooled  = new FloatArray(n * k * (outH/2) * (outW/2));

// ... fill input, filter, bias ...

TaskGraph graph = new TaskGraph("cnn")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, input, filter, bias)
    .libraryTask("conv",  CuDnn::cudnnConv2d,
            input, filter, convOut,
            n, c, h, w, k, r, s, pad, stride)
    .task("bias",         Layers::addBias, convOut, bias)              // JIT kernel
    .libraryTask("relu",  CuDnn::cudnnRelu,
            convOut, reluOut, n * k * outH * outW)
    .libraryTask("pool",  CuDnn::cudnnMaxPool2d,
            reluOut, pooled, n, k, outH, outW, 2, 2)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, pooled);

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

Flash Attention (SDPA) Example

The sdpaForward factory implements fused scaled-dot-product attention (flash attention) via the cuDNN graph API. It takes query, key, and value tensors in BHSD layout and produces the attention output in the same layout.

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

// Shape: 4 batches, 16 heads, 1024 query tokens, 128 head dim
int b = 4, heads = 16, sQ = 1024, sKv = 1024, d = 128;
float scale = 1.0f / (float) Math.sqrt(d);
boolean causal = true;

HalfFloatArray q = new HalfFloatArray(b * heads * sQ  * d);
HalfFloatArray k = new HalfFloatArray(b * heads * sKv * d);
HalfFloatArray v = new HalfFloatArray(b * heads * sKv * d);
HalfFloatArray o = new HalfFloatArray(b * heads * sQ  * d);

// ... fill q, k, v ...

TaskGraph graph = new TaskGraph("sdpa")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, q, k, v)
    .libraryTask("sdpa", CuDnn::sdpaForward,
            q, k, v, o,
            b, heads, sQ, sKv, d,
            scale, causal)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, o);

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

sdpaForward requires the head dimension d to be a multiple of 8 and no larger than 256. It also requires an Ampere or newer GPU (compute capability 8.0+). On older hardware the task will report UNSUPPORTED.

CUDA Graph Capture

Convolution descriptors, algorithm selection, and the grow-only workspace are created once per shape in the prepare() hook, before CUDA Graph capture begins. This makes all cuDNN library tasks CUDA Graph compatible.

try (TornadoExecutionPlan plan = new TornadoExecutionPlan(graph.snapshot())) {
    plan.withCUDAGraph();   // iteration 0 captures; subsequent iterations replay
    for (int i = 0; i < 100; i++) {
        plan.execute();
    }
}

Performance

Conv2d vs JIT Kernel

BenchmarkConv2d on RTX 4090, FP32, cuDNN 9.23:

Input (NCHW), Filter	JIT Direct Conv	cuDNN Task	Speedup
8×64×56×56, k=64	3.8 TFLOP/s	11.6 TFLOP/s	3.1×
16×128×28×28, k=128	4.0 TFLOP/s	21.1 TFLOP/s	5.3×

tornado -m tornado.cudnn/\
uk.ac.manchester.tornado.cudnn.tests.BenchmarkConv2d \
8 64 56 64 50

SDPA vs JIT Attention

BenchmarkSdpa on RTX 4090, FP16 cuDNN vs FP32 JIT attention:

Shape (b, h, s, d)	JIT Kernel	cuDNN SDPA	Speedup
4, 16, 1024, 128	90.7 ms (0.19 TFLOP/s)	0.268 ms (64.2 TFLOP/s)	339×
1, 32, 2048, 128	203.3 ms (0.17 TFLOP/s)	0.407 ms (84.3 TFLOP/s)	499×

tornado -m tornado.cudnn/\
uk.ac.manchester.tornado.cudnn.tests.BenchmarkSdpa \
4 16 1024 128 20

Mixing cuDNN with JIT Kernels

cuDNN library tasks can be freely mixed with JIT-compiled tasks in the same TaskGraph. The JIT bias task in the CNN example above writes back to the convOut buffer on the device; the subsequent cudnnRelu call reads that modified buffer directly from device memory — the CUDA stream guarantees ordering.

// Softmax → JIT temperature scaling → another cuDNN operation
TaskGraph graph = new TaskGraph("transformer_block")
    .transferToDevice(DataTransferMode.EVERY_EXECUTION, scores)
    .libraryTask("softmax",  CuDnn::cudnnSoftmax, scores, attnWeights, seqLen, seqLen)
    .task("temperature",     Ops::scale, attnWeights, invTemp)          // JIT kernel
    .libraryTask("relu",     CuDnn::cudnnRelu, attnWeights, out, size)
    .transferToHost(DataTransferMode.EVERY_EXECUTION, out);

Unit Tests and Runnable Examples

# Full cuDNN test suite (auto-skips without CUDA backend / libtornado-cudnn)
tornado-test -V uk.ac.manchester.tornado.unittests.cudnn.TestCuDnn

# CUDA Graph capture test for Conv2d
tornado-test -V uk.ac.manchester.tornado.unittests.cudnn.TestCuDnn#testConv2dWithCudaGraph

Results from BenchmarkSdpa are cross-validated against the FP32 JIT attention kernel on the same FP16-rounded inputs, confirming numerical correctness before reporting throughput numbers.

Known Limitations

The current cuDNN bindings use the cuDNN legacy compute API (deprecated since cuDNN 9.0, but still fully functional — validated on cuDNN 9.23). The ROADMAP tracks migration to the graph API for additional operations:

FP16/BF16 Tensor Core convolution
cudnnFindConvolutionForwardAlgorithm autotuning (fixed IMPLICIT_PRECOMP_GEMM algorithm is used today)
Backward/training operations
RMSNorm and LayerNorm via cuDNN graph API

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

cuDNN: Deep Learning Primitives in a Java TaskGraph

Factory Methods

Activations

Pooling

Convolution

Flash Attention (SDPA)

Data Layout

FP32 / NCHW

FP16 / BHSD (Flash Attention)

CNN Layer Example

Flash Attention (SDPA) Example

CUDA Graph Capture

Performance

Conv2d vs JIT Kernel

SDPA vs JIT Attention

Mixing cuDNN with JIT Kernels

Unit Tests and Runnable Examples

Known Limitations

Build docs developers (and LLMs) love

Get Started

Programming Guide

Backends

NVIDIA Integration

Guides

Contributing

Documentation Index

​Factory Methods

​Activations

​Pooling

​Convolution

​Flash Attention (SDPA)

​Data Layout

FP32 / NCHW

FP16 / BHSD (Flash Attention)

​CNN Layer Example

​Flash Attention (SDPA) Example

​CUDA Graph Capture

​Performance

Conv2d vs JIT Kernel

SDPA vs JIT Attention

​Mixing cuDNN with JIT Kernels

​Unit Tests and Runnable Examples

​Known Limitations

Build docs developers (and LLMs) love

Factory Methods

Activations

Pooling

Convolution

Flash Attention (SDPA)

Data Layout

CNN Layer Example

Flash Attention (SDPA) Example

CUDA Graph Capture

Performance

Mixing cuDNN with JIT Kernels

Unit Tests and Runnable Examples

Known Limitations