custatevec OOM on DGX Spark (GB10 UMA) — cudaMalloc fails on unified memory architecture

[reinamora137]

Summary

custatevec-backed quantum simulation (via PennyLane lightning.gpu) fails with out of memory on NVIDIA DGX Spark (GB10 Superchip, SM 12.1) even for trivially small state vectors (4 qubits = 256 bytes). The failure originates from cudaMalloc encountering the unified memory architecture where cudaMemGetInfo reports near-zero free memory despite 128GB of shared DRAM being available.

Environment

• Hardware: NVIDIA DGX Spark, GB10 Superchip (SM 12.1)
• Memory: 128GB LPDDR5X unified CPU+GPU (no discrete VRAM)
• CUDA: 13.0, Driver 580.126.09
• OS: Ubuntu (aarch64)
• custatevec: 1.12.0 (via custatevec-cu12 pip package)

Reproducer

import pennylane as qml

# Uses custatevec under the hood
dev = qml.device('lightning.gpu', wires=4)

@qml.qnode(dev)
def circuit():
    qml.Hadamard(wires=0)
    return qml.expval(qml.PauliZ(0))

result = circuit()  # Fails with OOM

pennylane_lightning.lightning_gpu_ops.LightningException: 
[...DevTag.hpp][Line:65][Method:refresh]: Error in PennyLane Lightning: out of memory

Root Cause

DGX Spark implements Unified Memory Architecture (UMA) — the GPU and CPU share the same physical DRAM. Standard CUDA memory query APIs (cudaMemGetInfo) return misleading values on UMA systems, causing libraries that pre-check available memory to believe no GPU memory exists.

This is documented by NVIDIA:

• DGX Spark Known Issues: "cudaMemGetInfo does not account for memory that could potentially be reclaimed from SWAP"
• DGX Spark CUDA Porting Guide

Similar issues in other CUDA libraries:

• PyTorch #174358
• RAPIDS cuml #3979

Questions

1. Does custatevec internally use cudaMemGetInfo for pre-allocation checks? If so, is there a plan to support UMA platforms?
2. Would using cudaMallocManaged instead of cudaMalloc on UMA systems (detected via cudaDeviceProperties::integrated) be a viable fix?
3. Is there a custatevec configuration option or environment variable to bypass the memory pre-check?
4. What is the roadmap for cuQuantum support on DGX Spark / Grace-Blackwell UMA systems?

Context

The DGX Spark is shipping to quantum computing researchers and developers. GPU-accelerated quantum simulation is a natural use case for this hardware. Currently, custatevec is unusable on it, forcing fallback to CPU-only simulation.

We've also filed a related issue on PennyLaneAI/pennylane-lightning since the allocation code path goes through their DataBuffer.hpp, but the underlying question is whether custatevec itself has UMA-incompatible memory assumptions.

[reinamora137]

Update: Confirmed fix — custatevec works on DGX Spark with cudaMallocManaged

We built PennyLane-Lightning from source with a patch that replaces cudaMalloc with cudaMallocManaged on UMA systems (detected via cudaDeviceProp::integrated).

Results:

• 20-qubit circuits: 13.9x speedup over CPU (8.39ms vs 116.31ms)
• 24-qubit circuits: execute in 399ms (268MB state vector) — previously impossible

Also note: custatevec-cu13 (pip package) is needed for CUDA 13.0 systems. The custatevec-cu12 package fails due to SONAME mismatch (libcublas.so.12 not found on systems with libcublas.so.13).

The fix is purely in the calling code (PennyLane-Lightning's DataBuffer.hpp), not in custatevec itself. However, if custatevec internally performs any cudaMemGetInfo-based pre-allocation checks, those would also need updating for UMA support.

[ymagchi]

Hi @reinamora137,

Thank you for your interest in cuStateVec.
The custatevec-cu13 package is built for CUDA 13, while custatevec-cu12 is for CUDA 12. Neither package includes cudaMemGetInfo API call inside the library. The cuStateVec API sets doesn't allocate memory for the state vector internally; it performs computations using the device pointer provided by the user.

Based on the discussion above, it seems that this issue has been resolved. Would it be alright to close this thread?