CUDA.jl 6.3: Compiler caching, a new cuDNN, and dependent launches


Tim Besard

CUDA.jl 6.3 features better integration with Julia's compiler caches, so that GPU-side inference done while a package precompiles survives across sessions. The cuDNN wrappers have been rebuilt on cuDNN 9's backend graph API, and there is also support for programmatic dependent launch.

Kernel compilation with CompilerCaching.jl

When you launch a kernel, CUDA.jl has to find the compiled code for it. Until now it kept that mapping itself: a dictionary per CUDA context, from (method instance, world age, compiler configuration) to a CuFunction. It worked, but it duplicated bookkeeping Julia already does for the same method instances, and the cached entries did not survive across Julia sessions.

CUDA.jl 6.3 adopts GPUCompiler 2, which builds on CompilerCaching.jl, and drops that dictionary. Compilation results are now stored in the CodeInstance that Julia caches anyway. This makes it possible to cache on disk, by saving into system or package images.

Right now, we only store inferred code. Work is underway to make the generated LLVM IR and machine code relocatable, which will enable caching those as well. However, just caching the inference results is already a big win. Let's demonstrate using a simple package:

module Blur

using CUDA
using PrecompileTools

function blur_kernel!(dst, src, ::Val{R}) where R
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if i <= length(dst)
        acc = zero(eltype(src))
        for k in -R:R
            @inbounds acc += src[clamp(i + k, 1, length(src))] / (1 + abs(k))
        end
        @inbounds dst[i] = sqrt(abs(acc))
    end
    return
end

function blur(src, ::Val{R} = Val(4)) where R
    dst = similar(src)
    @cuda threads=256 blocks=cld(length(dst), 256) blur_kernel!(dst, src, Val(R))
    return dst
end

@setup_workload begin
    @compile_workload begin
        blur(CUDA.zeros(Float32, 1024))
    end
end

end

Timing the first call in a fresh session, on an RTX 5080 with Julia 1.12 and CUDA 13.3:

julia> using Blur, CUDA

julia> src = CUDA.rand(Float32, 1024);

julia> @time Blur.blur(src);
  0.129540 seconds (14.94 k allocations: 2.432 MiB, 71.87% compilation time: <1% of which was recompilation)

Delete the @setup_workload block, precompile again, and the same call in a fresh session costs this instead:

julia> @time Blur.blur(src);
  1.292274 seconds (5.43 M allocations: 262.556 MiB, 14.88% gc time, 95.62% compilation time: 10% of which was recompilation)

Note that this requires Julia 1.11 or later.

cuDNN, rebuilt on the graph API

cuDNN has two programming models: the legacy API, a fixed set of fixed-function operations and fusion patterns with a C entry point each, and the graph API, where you describe a computation as a graph of operations and let cuDNN pick an engine for the whole thing. The graph API can be reached two ways: directly through the C back-end API, or through NVIDIA's cudnn-frontend, whose C++ and Python layers provide a simplified programming model that covers most use cases.

cuDNN.jl was written against the fixed-function API. In version 6.3, it is rebuilt on the back-end API, with a front-end mimicking cudnn-frontend: a graph API and a set of operations implemented on top of it. The fixed-function wrappers are unchanged and remain available.

Graph front-end

Graph and Tensor describe a computation, build! lowers it, runs cuDNN's heuristics and selects an execution plan, and execute! binds arrays and runs it. Intermediate tensors are marked virtual, which is how the engine knows it may fuse instead of materializing them. As an example, a batched matrix multiply followed by a bias add and a ReLU, as one plan:

using CUDA, cuDNN
using cuDNN: Graph, tensor!, matmul!, pointwise!, build!, execute!

A    = CUDA.rand(Float16, 256, 256, 8)
B    = CUDA.rand(Float16, 256, 256, 8)
bias = CUDA.rand(Float16, 256, 1, 8)
C    = CUDA.zeros(Float16, 256, 256, 8)

g = Graph(io_dtype=Float16, intermediate_dtype=Float32, compute_dtype=Float32)
ta, tb = tensor!(g, A; name="A"), tensor!(g, B; name="B")
tbias  = tensor!(g, bias; name="Bias")
tc     = tensor!(g, C; name="C")

tmm  = matmul!(g, ta, tb; name="MM")    # virtual
tsum = pointwise!(g, :add, tmm, tbias)  # virtual
pointwise!(g, :relu, tsum; y=tc)        # writes C

build!(g)
execute!(g, Dict(ta => A, tb => B, tbias => bias, tc => C))

Operations layer

On top of the frontend sits a higher-level API that's easier to use: attention! and attention_backward!, convolution! with its two gradients, maxpool!/meanpool! and their gradients, and the batchnorm_* family. These take CuArrays in Julia memory order and hide the graph entirely.

Both of these APIs are very new, and minor changes to the design or implementation are to be expected in future releases. Feedback is very welcome, so please report issues or missing features on the CUDA.jl bug tracker.

Programmatic dependent launch

Two kernels back-to-back in the same stream are fully serialized: the second one does not start until the last block of the first one retires. That is often more ordering than needed. If the consumer starts with work that does not touch the producer's output, like loading weights or zeroing an accumulator, that work could have been running while the producer's last few blocks were still draining.

CUDA calls the escape hatch programmatic dependent launch, and CUDA.jl 6.3 supports it. The producer signals when its dependents may start, the consumer is launched with dependent=true, and the consumer waits before it touches anything the producer wrote:

@inline function busy(x::Float32, n::Int)   # stand-in for real work
    for _ in 1:n
        x = fma(x, 1.0000001f0, 1f-7)
    end
    return x
end

function producer!(out, n)
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    trigger_programmatic_launch_completion()
    @inbounds out[i] = busy(Float32(i), n)   # the tail
    return
end

function consumer!(out, in, n)
    i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    pre = busy(Float32(i) * 0.5f0, n)        # independent preamble
    grid_dependency_synchronize()
    @inbounds out[i] = pre + in[i]
    return
end

@cuda threads=256 blocks=32 producer!(a, 20_000)
@cuda threads=256 blocks=32 dependent=true consumer!(b, a, 20_000)

Each of these kernels takes about 36 µs on its own, and the grid is small enough that both fit on the device at once. Run back to back they cost 69 µs; with the trigger, the wait and dependent=true they cost 38 µs, so the consumer's preamble hides almost entirely behind the producer.

The trigger belongs at the point in the producer after which nothing else has to run before dependents may start, which is usually the top; a block that exits without calling it triggers completion implicitly. grid_dependency_synchronize is what makes the producer's writes visible, so the consumer needs it even when the trigger has already run. And the overlap is opportunistic: code whose correctness depends on the two kernels running concurrently can deadlock. Programmatic dependent launch requires compute capability 9.0 or higher.

Other changes

The full list is in NEWS.md and the release notes. If something in here breaks for you, please file an issue.