LowRank-SSM: FPGA 상의 랭크 축소된 Mamba 가속을 위한 하드웨어-소프트웨어 공동 설계
요약
Mamba와 같은 상태 공간 모델(SSM)의 FPGA 배포 시 발생하는 메모리 및 대역폭 병목 현상을 해결하기 위한 LowRank-SSM 프레임워크를 제안합니다. SVD 기반의 가중치 분해 소프트웨어 알고리즘과 이중 경로 투영을 지원하는 하드웨어 가속기를 공동 설계하여 효율성을 극대화했습니다.
핵심 포인트
- SVD를 활용한 투영 레이어 가중치 분해 및 랭크 할당 알고리즘 도입
- FPGA 상의 이중 경로 투영 및 융합된 선택적 스캔 유닛 설계
- 혼합 랭크 실행을 통해 정확도와 처리량 간의 최적 트레이드오프 달성
- SOTA 대비 처리량 2.19배, 에너지 효율 2.03배 향상 입증
Mamba 및 Mamba-2와 같은 상태 공간 모델(State Space Models (SSMs))은 선형 시간 자기회귀 추론(linear-time autoregressive inference)을 달성하여, 지연 시간에 민감하고 자원이 제한된 배포 환경에 매력적입니다. 그러나 이들의 거대한 입력 및 출력 투영 레이어(projection layers)는 이차적인 가중치 메모리 및 오프칩 대역폭 비용을 발생시켜 실제 FPGA 배포의 병목 현상을 일으키며, 시퀀스 길이가 1,024 이상일 때 토큰당 실행 시간의 60% 이상을 차지합니다. 기존의 가속기들은 양자화(quantization)나 활성화 희소성(activation sparsity)을 통해 이러한 오버헤드를 줄이려 노력하지만, 투영 랭크(projection rank)를 명시적인 하드웨어 설계 변수로 다루는 방식은 없었으며, 이로 인해 체계적인 정확도-처리량 트레이드오프(accuracy-throughput trade-off)가 탐구되지 않은 상태로 남아 있었습니다. 본 논문에서는 이러한 격차를 해소하는 하드웨어-소프트웨어 공동 설계 프레임워크인 LowRank-SSM을 제시합니다. 소프트웨어 측면에서는, 사후 학습 절단된 SVD(post-training truncated SVD)를 통해 입력 및 출력 투영 가중치를 분해하고, 사용자가 지정한 정확도 제약 조건을 준수하면서 가중치 저장 공간을 최소화하는 밴드별 랭크 벡터를 탐색하는 탐욕적 밴드별 랭크 할당 알고리즘(greedy bandwise rank-allocation algorithm)을 도입합니다. 하드웨어 측면에서는, 결과적으로 인수분해된 투영(factored projections)을 FPGA 상의 완전 파이프라인 가속기에 매핑하며, 여기에는 이중 경로 투영(dual-path projection; 저랭크 경로 및 풀랭크 경로), 융합된 선택적 스캔 유닛(fused selective-scan unit), 그리고 버스 경합(bus contention) 없이 DDR4 대역폭을 포화시키는 5개의 독립적인 AXI 마스터 번들이 포함됩니다. 밴드별 실행 시간 랭크 마스크(per-band runtime rank mask)를 통해 아키텍처 오버헤드 없이 64개 모든 레이어에 걸쳐 혼합 랭크 실행(mixed-rank execution)이 가능합니다. Xilinx Versal VC1902 400 MHz 환경에서 배포된 혼합 랭크 INT8 설계는 7.89~tokens/s를 달성하였으며, 이는 유사한 전력 및 정확도 조건에서 SOTA 대비 ${2.19 imes}$의 처리량 향상과 ${2.03 imes}$의 에너지 효율 향상을 나타냅니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.AR의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기