JAX JIT 기반 TensorCircuit-NG, cuQuantum 대비 최대 147배 성능 우위 달성

Context

NVIDIA cuQuantum은 GPU 최적화 라이브러리를 통해 강력한 성능을 제공하나, end-to-end differentiable VQE 워크로드에서 높은 엔지니어링 공수와 verbose한 코드가 병목으로 작용함. 특히 단순 커널 성능과 전체 시스템의 실행 효율 사이의 간극으로 인한 최적화 한계가 존재함.

Technical Solution

• JAX JIT compilation 도입을 통한 반복적인 value-and-gradient 평가 비용의 amortize 구조 설계
• VQE 레이어의 scan 연산을 적용하여 JAX 컴파일 및 staging 시간을 단축하는 TC-JAX scan 모드 구현
• 전체 레이어를 직접 구축하여 traced program 크기를 늘리는 대신, 컴파일 후 peak throughput을 극대화한 TC-JAX unrolled 모드 제공
• high-level circuit-based programming model을 통해 low-level GPU 메모리 관리 및 binding signature 설정의 복잡성 제거
• 단순 GPU 커널 호출을 넘어 autodiff 통합과 graph-level optimization을 통한 end-to-end 실행 경로 최적화

Impact

• 24 qubits 기준 TC-JAX unrolled 모드가 cuStateVec adjoint 대비 11.83배, cuTensorNet full-state autograd 대비 90.46배 빠른 실행 속도 기록
• 20 qubits 기준 TC-JAX unrolled와 cuTensorNet 간 최대 147.93배의 성능 격차 확인
• 24 qubits 환경에서 unrolled 모드가 scan 모드 대비 약 2.50배 빠른 post-compilation 처리 속도 달성

Key Takeaway

하드웨어 벤더의 low-level 라이브러리가 반드시 최적의 end-to-end 성능을 보장하지 않으며, JIT 컴파일러를 통한 그래프 최적화와 고수준 추상화가 실제 런타임 효율성을 압도할 수 있음을 증명함.