cuTile Rust: 안전하고 데이터 경합이 없는 GPU 커널

cuTile Rust (cutile-rs)는 Rust의 소유권 및 빌림 규칙을 GPU 실행 경계 너머로 확장하여 메모리 안전하고 데이터 경합이 없는 GPU 커널 개발을 가능하게 합니다. 호스트 측에서 가변 텐서를 서로 겹치지 않는 조각으로 분할함으로써, 시스템은 각 GPU 타일 프로그램이 특정 데이터 세그먼트에 대한 독점적인 가변 참조를 받도록 보장하여 컴파일 타임에 데이터 경합을 제거합니다.

메모리 안전성과 타일 기반 모델

cuTile Rust는 수동 공유 메모리 관리와 unsafe 코드가 자주 필요한 전통적인 SIMT (Single Instruction, Multiple Threads) 프로그래밍을 안전한 타일 기반 프로그래밍 모델로 대체합니다.

실행 경계를 가로지르는 소유권

시스템은 엄격한 액세스 규칙을 강제하여 Rust의 안전성 보장을 유지합니다:

  • 가변 텐서 (Mutable Tensors): 커널이 실행되기 전에 서로 겹치지 않는 조각으로 분할됩니다. 각 타일은 할당된 분할 영역에 대한 독점적인 &mut 뷰를 부여받습니다.
  • 불변 텐서 (Immutable Tensors): 모든 타일에 읽기 전용 & 참조로 공유됩니다.
  • 런처 (Launchers): 생성된 런처는 GPU 작업이 진행되는 동안 소유권을 유지하며, 동기식 실행, 비동기 파이프라인 및 CUDA graph replay를 지원합니다.

CUDA Tile IR을 통한 JIT 컴파일

커널은 #[cutile::module] 매크로를 사용하여 정의되며, 이 매크로는 Rust AST를 캡처하여 호스트 바이너리에 내장합니다. 런타임에 cuTile Rust는 이 AST를 CUDA Tile IR을 통해 GPU cubin으로 JIT 컴파일합니다. 안전한 표면 API가 대부분의 작업을 처리하지만, Tile IR 작업 표면에 직접 접근해야 하는 사용자를 위해 unsafe 인트린직을 통한 저수준 제어 기능도 제공됩니다.

성능 벤치마크

NVIDIA B200에서 수행된 평가 결과, cuTile Rust는 측정 가능한 런타임 오버헤드 없이 안전성을 제공하며, 수동으로 최적화된 저수준 구현과 경쟁할 수 있는 성능 수준에 도달함을 보여줍니다.

  • Element-wise Operations: 7 TB/s에 도달하였으며, 이는 피크 메모리 대역폭의 약 91%입니다.
  • GEMM (General Matrix Multiply): 2 PFlop/s에 도달하였으며 (밀집 f16 피크의 92%), 저수준 Tile IR 변형과 0.3% 이내의 성능 차이를 보이며 cuBLAS와 경쟁할 수 있는 수준을 유지합니다.
  • 추론 성능 (Inference Performance): Hugging Face와의 협업을 통해, Grout 추론 엔진(cuTile Rust로 구축됨)은 RTX 5090에서 Qwen3-4B에 대해 171 tokens/s, B200에서 Qwen3-32B에 대해 batch-1 decode 시 82 tokens/s를 달성했습니다.

기술적 요구 사항 및 호환성

cuTile Rust는 compute capability sm_80 (Ampere) 이상의 NVIDIA GPU를 위해 설계되었습니다.

하드웨어 및 소프트웨어 종속성

  • GPU 아키텍처: 최소 sm_80. sm_100+ 지원을 위해서는 CUDA 13.1+가 필요하며, sm_8x는 CUDA 13.2, sm_90은 CUDA 13.3이 필요합니다.
  • CUDA Toolkit: FP4 packing 및 block-scaled MMA와 같은 기능을 활용하려면 버전 13.3을 권장합니다.
  • Rust 버전: 1.89+.
  • OS: Linux (Ubuntu 24.04에서 테스트됨).

다른 Rust GPU 프로젝트와의 비교

cuda-oxide와 같은 다른 프로젝트가 SIMT 스타일 커널을 위한 Rust-to-CUDA 컴파일을 제공하는 반면, cuTile Rust는 타일 기반 접근 방식을 취합니다. 이러한 트레이드오프는 안전성을 컴파일 타임에 확인할 수 있는 시맨틱 모델을 얻는 대신 명시적인 warp 프림티브와 수동 공유 메모리 제어를 포기하는 것입니다.

프로젝트 상태 및 생태계

cuTile Rust는 현재 활발한 개발 단계에 있는 초기 연구 프로젝트입니다. 사용자는 API 변경이 발생할 수 있으며 기능이 미완성일 수 있음을 유념해야 합니다. 프로젝트는 crates.io에서 cutile로 제공되며, Nix flake를 통해 간편한 환경 설정을 지원합니다.

워크스페이스 아키텍처

  • cutile: 커널 작성 및 실행을 위한 주요 사용자 대상 크레이트.
  • cutile-compiler: cutile-ir를 통해 Rust 커널을 실행 파일로 컴파일하는 작업을 처리합니다.
  • cuda-async / cuda-core: 각각 비동기 CUDA 실행 및 관용적인 안전한 CUDA API를 제공합니다.

Sources