Dispatch-Aware Ragged Attention for Pruned Vision Transformers

작성자

Haebom

카테고리

Empty

저자

Seifeldin Abdellatif, Ahmad Almasri

💡 개요

본 논문은 Vision Transformer(ViT)에서 토큰 가지치기(pruning)를 통해 연산량을 줄이더라도, 짧은 시퀀스 길이에서 발생하는 커널 디스패치 오버헤드로 인해 실제 속도 향상이 나타나지 않는 문제를 해결합니다. 이를 위해 기존 FlashAttention-2보다 낮은 디스패치 오버헤드를 가지는 경량화된 양방향 Triton 어텐션 커널을 제안합니다. 제안된 방법은 패딩 기반 PyTorch SDPA 대비 최대 2.51배의 처리량 향상을 달성하며, 가지치기된 ViT의 효율성을 높입니다.

🔑 시사점 및 한계

•

ViT의 토큰 가지치기에서 발생하는 실제 속도 향상은 커널 디스패치 오버헤드에 의해 제한될 수 있으며, 이를 해결하기 위한 효율적인 커널 구현이 중요합니다.

•

제안된 Dispatch-Aware Ragged Attention 커널은 짧은 시퀀스 길이에서도 가지치기의 이점을 살려 효율성을 크게 향상시킬 수 있습니다.

•

실험 결과는 제안된 시스템이 다양한 입력 크기 및 가지치기율에서 우수한 성능을 보이며, 수치적 정확성 또한 검증되었음을 보여줍니다.

•

제안된 커널이 특정 하드웨어(NVIDIA RTX 4000 Ada Generation GPU)에 최적화되어 있어, 다른 아키텍처에서의 성능 검증 및 일반화 가능성에 대한 추가 연구가 필요할 수 있습니다.

PDF 보기

Made with Slashpage