FastUSP 분산 확산 모델 추론을 위한 다중 레벨 가속 프레임워크

초록

본 논문은 대규모 텍스트‑투‑이미지 확산 모델(FLUX 12B, Qwen‑Image 등)의 멀티GPU 추론에서 핵심 병목인 커널 런치 오버헤드를 그래프 컴파일과 CUDA Graphs로 제거하고, FP8 양자화 통신 및 파이프라인 Ring Attention을 추가 최적화함으로써 2‑8 GPU 환경에서 1.09×‑1.16×의 일관된 속도 향상을 달성한 FastUSP 프레임워크를 제안한다.

상세 분석

FastUSP는 기존 USP(통합 시퀀스 병렬성) 구현이 안고 있던 세 가지 주요 비효율성을 체계적으로 해결한다. 첫 번째는 수백 개에 달하는 작은 CUDA 커널이 연속적으로 호출되면서 발생하는 CPU‑측 커널 런치 오버헤드이다. RTX 5090과 같은 최신 GPU에서는 개별 커널 실행 시간이 수십 마이크로초 수준으로 짧아, 런치 비용이 전체 단계 지연의 5‑10 %를 차지한다. FastUSP는 torch.compile의 “reduce‑overhead” 모드와 CUDA Graphs를 활용해 전체 연산 그래프를 캡처하고 재생함으로써 커널 수를 크게 줄이고, 런치 오버헤드를 거의 없앤다. 이 컴파일‑레벨 최적화가 전체 성능 향상의 9‑16 %를 담당한다는 실험 결과는 커널 런치 비용이 현재 고대역폭 NVLink 환경에서 가장 큰 병목임을 재확인한다.

두 번째는 통신 비용이다. USP는 Ulysses 단계에서 All‑to‑All을 두 번, Ring Attention 단계에서 N‑1번의 포인트‑투‑포인트 전송을 수행한다. NVLink의 900 GB/s 양방향 대역폭을 고려하면, 전체 연산 시간 대비 통신 비중은 5‑10 % 수준에 머문다. 따라서 통신을 최적화해도 전체 지연 감소 효과는 제한적이다. 그럼에도 불구하고 FastUSP는 FP8(E4M3) 양자화를 적용해 K/V 텐서의 전송량을 절반으로 줄이고, 정밀도 손실을 0.1 % 이하로 억제한다. 이는 클라우드 혹은 크로스‑노드 환경처럼 대역폭이 제한된 상황에서 유용하다.

세 번째는 연산‑통신 겹침이다. 기존 Ring Attention은 KV 청크를 순차적으로 받아가며 로컬 Query와 매칭해 계산하므로, 통신 지연이 그대로 누적된다. FastUSP는 이 과정을 이중 버퍼링된 CUDA 스트림으로 파이프라인화해, KV 청크를 비동기 수신하면서 동시에 이전 청크에 대한 attention을 수행한다. 실험 마이크로벤치마크에서는 1.25‑1.27×의 속도 향상을 보였지만, 전체 추론에서 차지하는 비중이 작아 최종 엔드‑투‑엔드 향상은 1‑2 % 수준에 머문다.

종합하면 FastUSP는 세 계층(컴파일, 통신, 연산) 최적화를 orthogonal하게 적용하면서도, 가장 큰 이득을 주는 컴파일‑레벨 최적화에 집중한다. 또한 Ring Attention 패턴이 현재 PyTorch Inductor와 완전 호환되지 않아 4‑8 GPU 규모에서 Qwen‑Image에 적용되지 못하는 한계를 명시하고, 향후 컴파일러‑런타임 공동 설계의 필요성을 제시한다.