메모리 제약 GPU에서 블록 저랭크 파운데이션 모델 추론 가속화: Triton 커널로 최대 3.76배 속도 향상

Accelerating Block Low-Rank Foundation Model Inference on MemoryConstrained GPUs

블록 저랭크(BLR) 기법(Monarch, BLAST)은 파운데이션 모델의 가중치를 압축해 메모리와 연산량을 줄이지만, 다중 토큰 추론에서는 중간 데이터 이동과 PyTorch 컴파일러 한계로 인해 오히려 메모리 대역폭에 묶여 지연 시간이 늘어날 수 있습니다. 본 연구는 roofline 분석을 통해 이 문제를 규명하고, 부분 융합(partial fusion)과 메모리 레이아웃 최적화를 적용한 맞춤형 Triton 커널을 제안합니다. Jetson Orin Nano와 A40 같은 메모리 제약 GPU에서 Llama-7B/1B, GPT2-S, DiT-XL/2, ViT-B 모델을 대상으로, PyTorch CUDA dense 기준 대비 최대 3.76배 속도 향상과 3배 모델 압축을 달성했습니다.