Reactor·Amazon, Trainium용 NKI 최적화로 실시간 영상 생성 실현

Reactor와 Amazon Neuron Science 팀은 AWS Trainium에서 실시간 인터랙티브 영상 생성을 가능하게 하는 커널 중심 최적화 전략을 개발했다. **Neuron Kernel Interface(NKI)**와 하이브리드 샤딩 방식을 활용해 Rolling Forcing 자기회귀 확산 모델을 실행했으며, 부드러운 재생을 위한 기준인 16 fps를 상회하는 프레임 속도를 달성했다.

커널 수준 최적화

실시간 영상 작업에서 범용 컴파일러가 처리하기 어려운 동적 형태, 비정상적인 메모리 접근 패턴, 무거운 캐시 관리 등 세 가지 주요 병목 현상을 해결했다. NKI를 사용해 비효율적인 연산을 하드웨어에 맞춘 커널로 교체했다.

• 3D-RoPE: 지연 시간이 5초에서 1.8밀리초로 단축됨.
• 캐시 복사: 레이어당 23밀리초에서 1.9밀리초로 감소.
• 어텐션 전치: 어텐션 커널에 융합하여 완전히 제거.

이러한 최적화로 파이프라인이 11 GB의 고대역폭 메모리 내에서 실행되었으며, 표준 eager 모드 경로에서는 메모리 부족 오류가 발생했다.

하이브리드 샤딩 전략

영상 확산 모델은 긴 토큰 시퀀스를 처리해야 하며, 셀프 어텐션이 연산 시간의 약 **70%**를 차지한다. 단일 코드로는 효율적인 처리가 불가능하며, 표준 병렬화 방식에는 문제가 있었다.

• 텐서 병렬화(TP) 단독: 모델의 어텐션 헤드가 12개인데 칩당 Neuron 코어가 8개라 나누어떨어지지 않아 패딩으로 인한 연산 낭비가 발생했다.
• 시퀀스 병렬화(SP) 단독: 파티션이 프레임 경계와 일치하지 않으면 3D 영상 토큰 구조가 깨질 위험이 있었다.

팀이 구현한 하이브리드 전략은 헤드를 4개 코어에, 시퀀스를 2개 코어에 분할한다. 이 방식으로 데이터 레이아웃의 정확성을 유지했으며, 공간 W축 샤딩을 사용해 VAE 디코더에서 8.25배의 초선형 속도 향상을 달성했다.

모델 구조 변경

하드웨어 활용률을 높이기 위해 확산 단계와 캐시 업데이트 단계 간 공통 컴포넌트를 배치 처리하도록 모델 코드를 최적화했다. 캐시 업데이트 단계는 연산량이 훨씬 적어 별도로 실행하면 활용률이 낮았다. 공유 컴포넌트를 배치하고 필요한 부분만 분할함으로써 인스턴스당 16개 칩과 8개 코어를 더 효과적으로 사용했다.

최종 파이프라인은 첫 번째 엔드투엔드 실행에서 올바른 영상을 생성하는 데 성공했다. 팀은 이러한 기술이 실시간 인터랙션을 요구하는 다른 자기회귀 확산 모델에도 일반화될 수 있다고 강조하며, Trainium을 확장 가능한 인터랙티브 월드 모델 플랫폼으로 자리매김하고 있다.