Question: GPU memory leak with CuArray inputs in iterative time stepping

[jakiesumrain]

Hi,

I'm using KrylovKit.exponentiate() with a custom GPU-accelerated operator for large-scale time propagation (>4000 iterations with time step 0.1 a.u.). I noticed different memory behaviors depending on the input vector type.

Observation

When passing CuArray inputs, GPU memory starts to grow at a certain iteration and quickly exausts in the next few time steps (I do not know exactly what happened, but the GPU usage quickly reached 100%, even though no OOM error occurred.), even though my operator uses pre-allocated buffers. With CPU Vector inputs (hybrid approach), memory stays stable.

My Implementation

I have two callables for the same operator:

# Approach 1: Pure GPU (CuArray input → CuArray output)
function (action::GPUBlockDenseAction{T})(v::CuArray{ComplexF64,1}) where T
    # [pre-allocated GPU buffers, no new allocations in this code]
    result = _matvec_batched!(action, v)
    return result  # CuArray
end

# Approach 2: Hybrid CPU/GPU (CPU Vector input → CPU Vector output)
function (action::GPUBlockDenseAction{T})(v_cpu::Vector{ComplexF64}) where T
    copyto!(action.v_gpu_input, v_cpu)        # Upload to GPU
    _matvec_batched!(action, action.v_gpu_input)  # GPU computation
    CUDA.synchronize()
    copyto!(action.result_cpu, action.temp_gpu)   # Download from GPU
    return action.result_cpu  # CPU Vector
end

Result

• Approach 1 (CuArray): VRAM grows from 4.8 GB → 8.6 GB starting at about 500 iterations and exausts (Adding CUDA.synchronize() in the CuArray version function does not solve memory leak)
• Approach 2 (Hybrid): VRAM stable at 4.8 GB throughout

Plus

• I did a simple timing of the iterative time stepping, when using the CuArray version function, the overhead of KrylovKit.exponentiate() alone (excluding matrix-vector multiplication on the GPU) slowly increases.

Question

Is this expected behavior? My hypothesis is that KrylovKit's internal allocations (basis vectors, orthogonalization temporaries) follow the input type, and Julia's GC doesn't trigger based on GPU memory pressure.