GPU에서의 두려움 없는 동시성: Rust를 이용한 안전한 GPU 추론, vLLM/SGLang과 경쟁 가능한 수준

저는 cuTile Rust를 유지 관리하고 있으며, 방금 "Fearless Concurrency on the GPU"라는 논문을 게시했습니다. 더 많은 GPU 코드가 AI에 의해 생성됨에 따라, 병목 현상은 코드를 작성하는 것에서 코드를 신뢰하는 것으로 이동하고 있습니다. cuTile Rust를 사용하면 Rust의 소유권(ownership) 및 빌림 검사(borrow checking)를 통해 컴파일러가 메모리 안전성(memory safety)과 데이터 경합(data-race)으로부터의 자유를 검증하는 GPU 커널(kernel)을 직접 작성하거나 생성할 수 있습니다. 이러한 보장(guarantees)은 구조적으로 얻어집니다. 이는 CUDA Tile IR로 낮아지는 타일 기반 프로그래밍 모델(tile-based programming model)로, Rust의 소유권 모델을 런치 경계(launch boundary) 너머로 전달합니다. 가변(mutable) 출력을 서로소인 가변 서브 텐서(disjoint mutable sub-tensors)로 분할하고, 입력을 공유 참조(shared references)로 전달하며, 컴파일러가 스레드 블록(thread blocks)으로 매핑하는 단일 스레드 의미론(single-threaded semantics)을 가진 타일 커널(tile kernels)을 작성합니다.

엔드 투 엔드(End to end)로, 저희는 Hugging Face와 함께 cuTile Rust를 기반으로 Qwen3 추론 엔진인 Grout를 구축했습니다. Batch-1 디코딩(decode) 시, RTX 5090에서 Qwen3-4B 기준 171 tok/s, B200에서 Qwen3-32B 기준 82 tok/s에 도달하며, 이는 vLLM 및 SGLang과 경쟁할 만한 수준입니다. Batch-1 디코딩은 메모리 대역폭 제한(memory-bandwidth-bound) 상황이며, Grout의 처리량(throughput)은 저희의 HBM 루프라인 분석(roofline analysis)과 일치합니다.

현재 Grout의 많은 커널은 여전히 unsafe 경로를 사용하고 있지만, 이를 safe 변체(variants)로 마이그레이션할 수 있으며, 이는 생성된 커널을 위한 검증 가능한 대상(verifiable target)을 제공합니다. 저희는 저장소(repo) 내의 cutile-kernels 크레이트(crate)에 이러한 커널 모음을 시작했습니다. 이것이 관심 분야라면, safe 변체를 기여하는 것이 향후 커널 합성(kernel synthesis)이 활용할 수 있는 안전하고 고성능인 커널 라이브러리를 성장시키는 데 도움이 됩니다.

커널 측면에서 안전성(safety)은 사실상 비용이 들지 않습니다. B200에서 safe GEMM은 직접 작성한 저수준 버전의 0.3% 이내 성능을 보여주며(~dense f16 peak의 약 92%), 요소별(element-wise) 연산은 측정 오차 범위 내에서 cuTile Python과 일치하는 약 7 TB/s에 도달합니다.

주의해야 할 몇 가지 추가 사항은 다음과 같습니다: Grout는 소수의 지원 모델을 가진 batch-1 모델이며(연구용 사례 연구이지, 즉시 교체 가능한 서버가 아님), NVIDIA 전용(Tile IR로 낮아짐)이며, GEMM은 일부 크기에서 여전히 cuBLAS보다 약간 뒤처집니다.

• 논문 (Paper): https://arxiv.org/abs/2606.15991
• 코드 (Code): https://github.com/nvlabs/cutile-rs
• Grout: https://github.com/huggingface/grout
  이 논문을 즐겁게 읽고 새로운 것을 배우시길 바랍니다! 질문이 있다면 기꺼이 답변해 드리겠습니다 :)
  /u/Exciting_Suspect9088 님이 r/MachineLearning 에 제출함
  [link] [comments]