Hugging Face Diffusers 팀이 Flash Attention 3, torch.compile, FP8 양자화를 조합해 Flux.1-Dev 모델의 LoRA 추론 속도를 2.23배 향상

Context

Flux.1-Dev 모델에서 LoRA 어댑터를 hotswap할 때마다 torch.compile이 재컴파일되어 추론 지연이 발생했다. LoRA가 서로 다른 rank와 대상 레이어를 가질 수 있어 동적 속성에 대응하는 최적화가 필요했다.

Technical Solution

• Flash Attention 3(FA3) 프로세서 적용: DiffusionPipeline의 어텐션 처리를 FlashFluxAttnProcessor3_0으로 교체
• FP8 양자화 도입: TorchAoConfig를 이용해 transformer 레이어를 float8dq_e4m3_row로 양자화
• torch.compile 활성화: fullgraph=True, mode="max-autotune" 옵션으로 모델 그래프 컴파일
• Hotswap 메커니즘 구현: pipe.enable_lora_hotswap(target_rank=max_rank)로 아키텍처 변경 없이 가중치만 교체하여 재컴파일 회피
• 가중치 패딩 처리: LoRA 가중치를 최대 rank 크기로 패딩해 속성 재할당 없이 데이터 교체

Impact

NVIDIA H100 GPU에서 2.23배 추론 속도 향상. RTX 4090 같은 소비자용 GPU에서는 T5 텍스트 인코더 NF4 양자화와 영역별 컴파일 추가 적용으로 2.04배 속도 향상 달성.

Key Takeaway

동적 모델 구성(hotswap)을 요구하는 환경에서 컴파일 재실행을 방지하기 위해 아키텍처 불변성을 유지하면서 최적화 기법을 적층 적용할 수 있다. 최대 rank 사전 지정과 가중치 패딩은 런타임 재컴파일 오버헤드 제거의 핵심이다.