Roblox에서 3D 제작을 위한 AI 추론 가속화
3D 객체 생성 속도 7.8배 향상 및 반응성 개선
- Roblox는 더 빠른 반복 처리를 위해 3D 메쉬 생성을 가속화하기 위해 CUDA Graphs와 KV 캐싱을 도입했습니다.
- 출시 당시 Cube 3D 모델은 토큰을 7.8밀리초(기존 60.5밀리초에서 단축) 만에, 전체 오브젝트를 4초(기존 31초에서 단축) 만에 생성할 수 있었습니다.
올해 초, Roblox는 Cube 3D 파운데이션 모델의 첫 번째 기능을 공개했습니다. Cube 3D를 통해 크리에이터는 텍스트 프롬프트만으로 3D 모델과 환경을 직접 생성할 수 있습니다. 우리는 처음부터 생성 속도가 느리면 본질적으로 반복적인 과정인 작업 흐름이 방해된다는 점을 인식하고, 지연 시간 최적화를 최우선 과제로 삼았습니다. 지난 3월 Cube 3D가 출시되기 전, 우리는 이미 추론 단계를 7.8배 더 빠르게 만들어 개발자와 사용자 모두에게 더 빠른 응답성을 제공했습니다.
출시 이후, 여러 주목할 만한 경험(experience)을 통해 578,000개 이상의 오브젝트가 생성되었습니다. 개발자들은 또한 사용자가 "고양이", "버거" 등과 같은 텍스트 프롬프트를 통해 경험 내에서 3D 오브젝트를 생성할 수 있도록 하는 데 관심을 보였습니다. 특히 음성 채팅을 사용하는 인기 있는 사교 게임인 'Mic Up'은 Cube 3D를 활용하여 플레이어에게 오브젝트를 생성하는 재미있고 상호작용적인 방식을 제공했습니다. 이 구현 방식에서 플레이어는 AI 아이콘을 포함한 추가 기능 메뉴를 왼쪽에서 열 수 있습니다. 해당 아이콘을 클릭한 후 텍스트 프롬프트를 입력하여 3D 오브젝트를 생성할 수 있습니다. 사용자에게는 생성 시간이 길어질수록 불편함이 가중되어, 자신의 아이디어가 실시간으로 3D로 구현되는 마법 같은 경험을 누리지 못하게 됩니다.
우리는 3D 생성 경험을 '멈추고 기다리는' 상호작용에서 반응성이 뛰어나고 자연스러우며, 신속한 실험이 가능한 경험으로 바꾸고자 했습니다. 씬에 오브젝트를 빠르게 추가할 수 있는 능력은 개발자에게 매우 중요합니다. Cube 3D의 속도를 높이기 위해, 우리는 먼저 추론 파이프라인을 프로파일링하여 성능 병목 현상을 파악했습니다. 강력한 GPU를 사용했음에도 불구하고, 연산 사이의 상당한 유휴 시간을 발견했습니다.
CPU-GPU 스케줄링 병목 현상 해결
현대 딥러닝 프레임워크는 GPU에서 연산(또는 커널)을 스케줄링하고 실행하는 데 CPU에 의존합니다. CPU는 각 연산을 준비하여 GPU로 전송한 후, 다음 연산을 준비하기 전에 확인 신호를 기다립니다. 이러한 대기 과정은 CPU가 다음 작업 배치를 준비하는 동안 GPU가 유휴 상태로 남아 있을 수 있는 스케줄링 병목 현상을 유발합니다. 이상적으로는 CPU가 GPU보다 앞서 연산을 준비하고 대기열에 넣어, GPU가 항상 처리할 작업을 갖도록 하는 것이 바람직합니다.
이는 Cube 3D와 같은 트랜스포머형 모델의 자기회귀 디코더에 특히 문제가 됩니다. 이러한 모델은 입력을 처리하고 토큰을 순차적으로 생성해야 하기 때문입니다. 이 모델들은 단 한 번의 생성 작업에도 수천 개의 개별 연산이 필요하며, 순서의 각 단계마다 계산 오버헤드가 누적됩니다.
“우리는 4차원 상호작용을 가능하게 하는 무언가를 만들고 싶었습니다.”라고 엔지니어링 부사장 아누팜 싱(Anupam Singh)은 로블록스가 자기회귀적 접근 방식을 선택한 이유를 설명하며 말했습니다. “우리는 단순히 자동차를 만드는 데 그치고 싶지 않습니다. 자동차 문을 열고 안으로 들어갈 수도 있기를 원합니다.”
각 연산에 소요된:
- 각 커널을 준비하는 데 소요되는 CPU 시간
- 커널 실행으로 인한 오버헤드
- GPU 실행 시간(실제 계산)
- 완료 여부 확인 시 발생하는 동기화 오버헤드
GPU에서 빠르게 실행되는 소규모 연산의 경우, 이러한 오버헤드가 추론 시간을 대부분 차지할 수 있습니다. GPU가 전체 추론 시간 중 극히 일부 시간 동안만 실제로 연산을 수행할 수 있습니다.
CUDA 그래프 구현: 중개자 제거
이러한 병목 현상을 해결하기 위해, 우리는 CPU의 개입 없이 GPU 연산 시퀀스를 기록하고 재실행할 수 있는 기능인 CUDA Graphs를 활용했습니다. Cube 3D 아키텍처의 자기회귀 디코더 구성 요소는 텍스트 프롬프트를 처리하고 고정 길이 벡터를 통해 형상 토큰을 생성합니다.
기능적으로는 기존의 대규모 언어 모델(LLM)과 유사하지만, 당사의 듀얼 스트림 디코더 아키텍처는 두 개의 병렬 어텐션 스트림을 사용한다는 중요한 차이점이 있습니다. 한 스트림은 조건 토큰 전용이고, 다른 스트림은 형상 토큰 전용입니다. 시중의 LLM 추론 엔진은 당사의 요구 사항에 적합하지 않았으며, 당사의 특정 아키텍처에 맞춤화된 구현이 필요했습니다.
CUDA 그래프를 GPU용 매크로 기록과 비슷하게 생각하면 됩니다. CPU가 각 명령을 개별적으로 내리는 대신, GPU 연산 전체 시퀀스(그래프)를 기록하고 단일 CPU 명령어로 전체 그래프를 실행합니다. 이 접근 방식은 추론 중에 CPU가 각 연산을 개별적으로 스케줄링할 필요가 없도록 하여 커널 실행 오버헤드를 획기적으로 줄여줍니다. 그래프가 실행되면 GPU는 추가 명령을 기다리지 않고 전체 시퀀스를 자율적으로 실행합니다.
CUDA 그래프에는 몇 가지 한계가 있습니다. 그래프 구조를 사전에 결정해야 하므로 고정된 배치 크기와 입력 차원이 필요합니다. 즉, 배치 크기나 입력 형식에 따라 별도의 그래프를 생성해야 합니다. Cube 3D를 활용한 우리의 사용 사례에서는 일반적인 입력 차원을 기준으로 추론 프로세스를 표준화할 수 있었기 때문에 이 제한 사항은 수용 가능했습니다.
Cube 3D 모델에 CUDA 그래프를 구현하기 위해 접근 방식을 조정해야 했습니다. 기존 대규모 언어 모델(LLM)에서는 어텐션 연산이 항상 동일한 시퀀스 길이로 수행되어 정적인 형식을 제공합니다. 그러나 당사의 맞춤형 듀얼 스트림 아키텍처에서는 일부 어텐션 레이어가 시퀀스 길이만으로 작동하는 반면, 다른 레이어는 시퀀스 길이와 조건 길이의 조합을 기반으로 작동합니다.
이러한 어려움에도 불구하고, CUDA Graphs를 구현한 후 놀라운 결과를 얻었습니다. 우리는 추론 중 각 토큰의 생성 시간을 측정하기 위해 토큰당 처리 시간(TPOT)을 사용합니다. CUDA Graphs를 구현한 후, TPOT은 60.5밀리초에서 20.5밀리초로 개선되어 2.9배의 성능 향상을 보였습니다. 전체 생성 시간은 31초에서 10.5초로 66% 감소했습니다.
KV 캐싱: 성공을 발판으로
지연 시간을 더욱 단축하기 위해, 업계 전반에서 매우 효과적인 것으로 입증된 LLM 추론의 표준 관행인 KV 캐싱을 구현했습니다.
Cube 3D와 같은 트랜스포머 기반 모델에서는 각 토큰을 생성할 때마다 이전에 생성된 모든 토큰을 기반으로 키(K) 행렬과 값(V) 행렬을 계산해야 합니다. 시퀀스가 길어질수록 모든 토큰에 대해 이러한 행렬을 재계산하는 것은 점점 더 비효율적이 됩니다.
KV 캐싱은 다음을 통해 이 문제를 해결합니다:
- 이전에 생성된 모든 토큰에 대한 K 및 V 행렬을 저장
- 새로운 토큰에 대해서만 K 및 V 행렬을 계산
- 이 새로운 행렬들을 캐시된 값들에 추가
이 접근 방식은 중복 계산을 제거하여 각 새로운 토큰에 필요한 작업량을 줄여줍니다. 이는 생성된 시퀀스가 길어질수록 특히 큰 효과를 발휘합니다.
KV 캐싱을 CUDA 그래프 구현에 통합하는 우리의 접근 방식은 기존의 LLM 추론 방식과 유사했습니다. KV 캐싱을 추가함으로써 TPOT(토큰당 처리 시간)을 7.8밀리초로 단축했습니다. 전체 생성 시간은 기존 31초에서 4초로 87% 감소했습니다. 이러한 상당한 시간 단축은 이 도구를 사용하는 크리에이터들에게 훨씬 더 큰 효율성을 제공합니다.
개발자와 사용자에게 미치는 실제적 영향 평가
저희는 최적화된 커널, 더 빠른 추론을 위한 모델 양자화, 하드웨어별 최적화, 병렬 토큰 생성 등을 포함하여 지연 시간을 더욱 줄이고 사용자 경험을 개선하는 기술을 탐구하고 있습니다.
이 작업은 전체 장면 생성 및 이해로 확장될 때 더욱 복잡해집니다. 이때는 레이아웃 내에서 수많은 3D 요소들이 서로 맥락을 고려하며 협력해야 하기 때문입니다. 또한 우리가 생성하는 3D 객체와 세계가 문이 열리고 닫히거나 바퀴가 회전하는 등 완벽하게 작동하기를 원합니다. 이를 달성하기 위해서는 전체 장면, 완벽하게 작동하는 객체, 아바타로 확장할 수 있도록 신속한 생성 및 반복 작업이 필요합니다. Cube 3D 파운데이션 모델을 확장해 나가며 추가적인 개선 사항과 새로운 기능을 공유하게 되어 기쁘며, 크리에이터 커뮤니티가 이를 활용해 만들어낼 몰입감 넘치는 세계를 기대합니다.
