NVIDIA Warp 리뷰: 시뮬레이션, 로보틱스 및 미분 가능한 ML을 위한 GPU 가속 Python

NVIDIA는 런타임에 CUDA로 컴파일되는 GPU 커널 (kernel) 작성을 위한 오픈 소스 Python 프레임워크인 Warp를 출시했습니다. 우리는 Warp가 실제로 설계된 워크로드인 입자 시뮬레이션 (particle simulation), 접촉이 빈번한 로보틱스 (contact-rich robotics), 그리고 미분 가능한 물리 (differentiable physics)를 대상으로 일주일간 테스트를 진행했으며, Warp가 진정으로 누구를 위한 것인지 알아내기 위해 JAX 및 Taichi와 나란히 비교했습니다.

요약하자면: Warp는 PyTorch를 대체하려는 것이 아니며, 대부분의 리뷰가 수행하는 JAX와의 비교는 핵심을 놓치고 있습니다. Warp는 "텐서 프로그램이 있다" (JAX가 빛을 발하는 영역)와 "직접 작성한 CUDA 커널이 있다" (Python을 포기해야 하는 영역) 사이의 간극에 존재합니다. 만약 당신의 작업이 불규칙한 제어 흐름 (irregular control flow), 희소 데이터 구조 (sparse data structures), 접촉 해상 (contact resolution), 레이 트레이싱 (ray tracing), 기하학적 처리 (geometry processing)와 같이 그 중간 어디쯤에 있다면, Warp는 다른 어떤 것도 깔끔하게 수행하지 못하는 빈자리를 채워줍니다.

Warp가 Python을 CUDA 커널로 컴파일하는 방법

Warp의 프로그래밍 모델은 두 가지 데코레이터 (decorator)를 중심으로 구축되었습니다: GPU에서 실행되는 병렬 함수를 위한 @wp.kernel과 디바이스 측 헬퍼 (device-side helpers)를 위한 @wp.func입니다. 사용자는 스칼라 (scalars), 벡터 (vectors), 행렬 (matrices), 구조체 (structs), 배열 (arrays)과 같은 제한된 타입 시스템 (type system)을 사용하여 Python을 작성하며, Warp의 트레이서 (tracer)가 함수를 C++/CUDA 소스로 변환하고, NVRTC로 컴파일한 뒤 결과물인 PTX를 캐시 (cache)합니다.

컴파일은 지연 방식 (lazy)으로 이루어지며 함수 시그니처 (function signature)별로 디스크에 캐시됩니다. 따라서 커널의 첫 호출에는 100-300ms가 소요되지만, 이후 실행은 캐시를 사용하여 마이크로초 (microseconds) 단위로 디스패치 (dispatch)됩니다. 이는 유사한 워크로드에서 JAX의 jit 웜업 (warmup)보다 눈에 띄게 빠른데, 이는 주로 Warp가 XLA 로워링 파이프라인 (lowering pipeline)을 건너뛰고 NVRTC로 바로 진행하기 때문입니다.

커널 (kernel) 내부에서 작성할 수 있는 내용은 의도적으로 제한되어 있습니다. 루프 (loops), 조건문 (conditionals), 함수 호출 (function calls)은 모두 작동하지만, Python의 동적 타이핑 (dynamic typing)은 지원되지 않습니다. 즉, 모든 변수는 Warp가 커널 시그니처 (kernel signature)로부터 추론하는 정적 타입 (static type)을 가집니다. 배열 (arrays)은 참조 (reference)로 전달되며, 스레드 ID를 위해 wp.tid()로 인덱싱됩니다. 가비지 컬렉션 (garbage collection), 예외 (exceptions), Python 객체 (Python objects)는 존재하지 않습니다. 이러한 제한 덕분에 Warp는 직접 작성한 CUDA와 동일한 속도로 실행되는 코드를 생성할 수 있습니다. 사실상 타입 추론 (type inference) 기능이 포함된 Python 문법으로 CUDA를 작성하는 것과 같습니다.

Warp는 어드조인트 코드 (adjoint code)를 자동으로 생성하는 내장 미분 가능 모드 (differentiable mode)를 제공합니다. 데코레이터 (decorator)를 적용한 모든 커널에 대해, Warp는 소스 투 소스 자동 미분 (source-to-source autodiff)을 통해 역전파 (backward pass)를 생성할 수 있습니다. 이는 물리 시뮬레이션이 엔드 투 엔드 (end-to-end)로 미분 가능함을 의미하며, Warp가 접촉이 빈번한 조작 (contact-rich manipulation) 및 소프트 바디 역학 (soft-body dynamics) 관련 연구 논문에 자주 등장하는 주요 이유입니다.

Warp vs JAX vs Taichi: 어떤 컴파일러가 귀하의 워크로드에 적합한가

이 세 프레임워크는 서로 다른 목적을 위해 최적화되어 있기 때문에, 정직한 비교는 보기보다 어렵습니다.

JAX는 귀하의 연산이 텐서 프로그램 (tensor program)이라고 가정합니다. NumPy 스타일로 코드를 작성하면, XLA가 연산들을 효율적인 커널로 융합 (fuse)합니다. JAX는 밀집 선형 대수 (dense linear algebra), 트랜스포머 (transformers), 그리고 미분 가능한 손실 함수 (differentiable losses)에 대한 경사 기반 최적화 (gradient-based optimization)에 탁월합니다. 하지만 불규칙한 메모리 액세스 (irregular memory access), 입자 상호작용 (particle interactions), 또는 깔끔하게 벡터화 (vectorize)되지 않는 작업에는 어려움을 겪습니다. JAX로 광역 단계 충돌 탐지기 (broad-phase collision detector)를 작성하려고 시도해 본다면 그 고충을 느끼게 될 것입니다.

Taichi는 개념적으로 Warp와 가장 유사합니다. Taichi 역시 데코레이터 기반의 DSL을 통해 Python을 GPU 코드로 컴파일하며, Warp가 채택한 많은 패턴의 선구자 역할을 했습니다. Taichi는 더 넓은 크로스 플랫폼 지원 (Vulkan, Metal, OpenGL 백엔드)을 갖추고 있는 반면, Warp는 NVIDIA 하드웨어와 더 밀접하게 결합되어 있으며 Omniverse, MuJoCo XLA, 그리고 Isaac Sim과의 직접적인 통합을 통해 이점을 얻습니다.

Warp는 세 가지 요소를 동시에 필요로 할 때 선택해야 할 최적의 도구입니다: NVIDIA GPU 성능, 불규칙한 제어 흐름 (irregular control flow)에 대한 자동 미분 (autodiff) 기반의 미분 가능한 물리 (differentiable physics), 그리고 wp.from_torch()를 통한 PyTorch 텐서와의 제로 카피 상호 운용성 (zero-copy interop)입니다. 이 PyTorch 상호 운용성은 많은 머신러닝 (ML) 연구자들이 과소평가하는 기능입니다. Warp 시뮬레이션을 미분 가능한 레이어 (differentiable layer)로 감싸고, PyTorch에서 PPO로 정책 (policy)을 학습시키면, 그래디언트 (gradients)가 엔드투엔드 (end-to-end)로 흐르게 됩니다.

100만 개의 입자 (particles)와 이웃 쿼리를 위한 공간 해시 그리드 (spatial hash grid)를 사용하여 실행한 입자 시뮬레이션 벤치마크에서, Warp는 RTX 4090에서 스텝당 1.8ms의 속도로 적분 커널 (integration kernel)을 디스패치 (dispatch)했습니다. jit 및 vmap을 사용한 동일한 JAX 구현은 희소 이웃 조회 (sparse neighbor lookup)로 인해 스캐터 연산 (scatter operations)으로 폴백 (fallback)되어야 했기 때문에 14ms가 소요되었습니다. Taichi는 동일한 하드웨어에서 Warp 수치의 5% 이내 성능을 보여주었습니다.

PyTorch 대신 Warp를 선택해야 하는 경우

밀집 행렬 곱셈 (dense matrix multiplies), 컨볼루션 (convolutions), 어텐션 (attention)과 같이 신경망 (neural networks) 형태를 띠는 워크로드에서는 항상 PyTorch가 승리할 것입니다. Warp는 해당 작업을 위해 경쟁하는 것이 아닙니다. Warp가 여러분의 스택에서 자리를 잡는 경우는 다음과 같습니다:

1. 학습 루프 내의 물리 시뮬레이터 (Physics simulators). 로봇 정책을 학습시키고 있는데 시뮬레이터 스텝이 병목 현상 (bottleneck)이 되는 경우입니다. Warp를 사용하면 시뮬레이터를 CUDA에 근접한 속도로 Python에서 작성할 수 있으며, 이를 통해 미분이 가능합니다.
2. 3D ML을 위한 기하학적 처리 (Geometry processing). 메시 연산 (Mesh operations), 부호 거리 함수 (signed distance fields), 마칭 큐브 (marching cubes) 등 PyTorch가 커스텀 CUDA 연산 (custom CUDA ops)을 통해 처리하기 까다로운 불규칙한 워크로드들입니다.
3. 입자 및 유체 역학 (Particle and fluid dynamics). SPH, MPM, FLIP 솔버 (solvers)는 Warp의 공간 데이터 구조 (spatial data structures)와 어드조인트 지원 (adjoint support)으로부터 엄청난 이점을 얻습니다.
4. 역 렌더링 (Inverse rendering) 및 레이 트레이싱 (ray tracing). Warp는 이러한 작업들을 순수 Python에서도 다룰 수 있게 해주는 BVH 및 레이 트레이싱 프리미티브 (ray-tracing primitives)를 포함하고 있습니다.

도입하기 전에 알아두어야 할 한계점

Warp는 완성된 제품이 아니며, README 파일에서도 이 점을 솔직하게 밝히고 있습니다.

• 실질적으로 NVIDIA 전용입니다. 디버깅을 위한 CPU 폴백(fallback) 기능이 존재하지만, 프로덕션 수준은 아닙니다. AMD 및 Apple Silicon 사용자는 Taichi를 살펴보는 것이 좋습니다.
• Python 서브셋(subset)이 엄격합니다. 커널(kernel) 내부에서 리스트 컴프리헨션(list comprehensions), 데코레이터(decorators) 또는 임의의 라이브러리 호출을 사용할 수 없습니다. 처음 포팅하는 몇 개의 커널은 리팩토링(refactor)이 필요할 것으로 예상하십시오.
• 디버깅 도구가 빈약합니다. 커널이 충돌하면 Python 소스 코드가 아닌, 생성된 C++의 줄 번호와 CUDA 에러 코드를 받게 됩니다. 소스 맵(Source-map) 지원은 로드맵에 포함되어 있으나 아직 출시되지 않았습니다.
• 문서가 API 업데이트를 따라가지 못합니다. 저장소(repo)에 있는 예제들이 가장 신뢰할 수 있는 참조 자료입니다. 공식 문서는 종종 한 단계 이전 버전에 머물러 있는 경우가 많으므로, 문서화되지 않은 문제에 부딪히면 테스트 파일들을 읽어보십시오.

만약 프로덕션용 시뮬레이터를 구축하고 있다면, Warp 버전을 고정(pin)하십시오. 현재 마이너 릴리스 간의 파괴적 변경(Breaking changes)이 흔하며, 한 릴리스에서 컴파일되었던 커널 시그니처(kernel signatures)가 다음 릴리스에서는 컴파일되지 않을 수 있습니다. 프로젝트가 빠르게 진행되고 있습니다. 이는 기능 측면에서는 좋지만, 안정성 측면에서는 비용이 따릅니다.

여러분의 작업이 시뮬레이션이나 기하학(geometry)과 관련이 있다면, 이 프레임워크는 시간을 투자할 가치가 있습니다. 만약 여러분이 트랜스포머(transformers)를 학습시키는 순수 ML 엔지니어라면, Warp는 아마 적합하지 않을 것입니다. 그리고 NVIDIA도 그 점을 알고 있는 듯합니다. 그들은 여러분의 학습 프레임워크를 대체하려는 것이 아닙니다. 그들은 텐서 프로그램(tensor-program) 형태에 맞지 않는 작업을 수행하는 모든 이들을 위해 Python에서 GPU로 이어지는 간극을 메우려 하고 있습니다.