cuBLAS 대비 96% 성능 구현 및 Rust Borrow Checker 기반의 GPU Memory Safety 확보

Context

기존 GPU 프로그래밍은 Host-to-GPU 경계에서 Rust의 Borrow Checker가 무효화되어 raw pointer 기반의 수동 메모리 관리가 필수적이었음. 이로 인해 Aliasing 및 Synchronization 오류로 인한 데이터 레이스 위험과 메모리 불안정성이 상존함.

Technical Solution

• Host 단에서 Mutable Output Tensor를 중첩되지 않는 disjoint tiles로 분할하여 데이터 레이스 원천 차단
• 분할된 각 Tile Program에 독점적 &mut view를 할당하고 Input은 Shared & reference로 제공하는 Type Construction 설계
• Const-generic shape 파라미터를 통해 Tile 크기를 Type 수준에서 인코딩하여 Borrow Checker의 정적 검증 가능케 함
• MLIR 기반의 CUDA Tile IR을 통한 JIT Compilation 구조로 Rust 코드를 고성능 GPU Binary로 변환
• Lazy DeviceOp 모델 및 CUDA Graph Capture를 통한 버퍼 재사용 최적화 및 실행 순서 강제
• #[cutile::module] 매크로를 통한 Host-side Launcher 자동 생성으로 수동 Dispatch 코드 제거

Impact

• NVIDIA B200 기준 GEMM 연산 시 cuBLAS 성능의 약 96% 달성 및 2 PFlop/s 기록
• Memory-bound element-wise 연산에서 7 TB/s의 처리량 확보
• Grout 추론 엔진 적용 시 RTX 5090(Qwen3-4B)에서 171 tokens/s, B200(Qwen3-32B)에서 82 tokens/s 달성

Key Takeaway

런타임 락 없이 Type System의 제약 조건(Ownership/Borrowing)을 데이터 분할 전략과 결합함으로써, 성능 저하 없는 메모리 안전성을 달성하는 구조적 해결책 제시