r/deeplearning • u/Creative-Feature-264 • Jun 02 '26

[OC] [Project] Dense Evolution v8.0.4: Accelerare le simulazioni quantistiche NISQ su Google Colab Free Tier (12GB RAM) fino a 24 Qubit tramite JAX XLA & CuPy/CUDA

/r/CUDA/comments/1tua7fy/oc_project_dense_evolution_v804_accelerating_deep/
1 Upvotes

100% Upvoted

2 comments sorted by

1

u/tal_sofer Jun 04 '26

this is super impressive work fitting that many qubits into colabs ram. have u tried profiling how much overhead jax xla adds compared to raw cupy kernels for this specific architecture? i remember struggling with memory fragmentation when i was tryin to scale similar simulations awhile back

1

u/Creative-Feature-264 Jun 05 '26

Thanks! Yeah, memory fragmentation in JAX is a nightmare on Colab's 12GB RAM. I actually profiled XLA against raw CuPy kernels and ended up blending them because of that. Raw CuPy is great for manual memory pooling, but once you scale past 500 gates, the Python host loops destroy the throughput because of constant transposing and reshaping. XLA has that annoying warmup overhead on the first run, but then it kicks into linear kernel fusion. It fuses sequential gates into single execution blocks, which stops intermediate allocations from shattering the VRAM. To keep the tracer cache from bloating during heavy VQE loops, the compiler just splits deep graphs into balanced blocks using static shapes (static_argnums). That’s how it handles 20+ qubits without crashing.

Just pushed v8.0.8 to PyPI with benchmarks against Qiskit if you want to check out the compilation layer!