ScanWeaver: 결합 법칙 스캔 로워링 (Associative Scan Lowering)을 통한 아핀 재귀 (Affine

요약

Mamba와 같은 선택적 상태 공간 모델의 순차적 의존성 문제를 해결하기 위한 컴파일러 프레임워크 ScanWeaver를 제안합니다. 재귀 계산을 결합 법칙 스캔 표현으로 변환하여 GPU에서 병렬 실행이 가능하도록 MLIR 기반의 엔드 투 엔드 로워링을 수행합니다.

핵심 포인트

재귀 기반 계산을 결합 법칙 스캔으로 변환하여 병렬성 확보
MLIR 기반의 체계적인 GPU 로워링 프레임워크 ScanWeaver 제안
Mamba 스타일의 선택적 스캔을 아핀 재귀 클래스로 일반화
PyTorch 및 CUDA 베이스라인 대비 효율적인 GPU 실행 검증

Mamba와 같은 선택적 상태 공간 모델 (Selective State-Space Models)은 선형 시간 시퀀스 모델링 (Linear-time sequence modeling) 능력을 유지하면서 언어 모델링 능력을 향상시키는, 입력 의존적 스캔 재귀 (Input-dependent scan recurrences)의 실질적인 중요성을 강조합니다. 그러나 이러한 재귀는 기존의 구조화된 SSM (Structured SSMs)보다 더 엄격한 순차적 의존성 (Sequential dependencies)을 도입하여, 현대적인 가속기 (Accelerators)에서의 병렬 실행을 제한합니다. 본 논문에서는 재귀 기반 계산을 결합 법칙 스캔 표현 (Associative scan representations)으로 변환하고, 이를 실행 가능한 GPU 프로그램으로 엔드 투 엔드 (End-to-end) 로워링 (Lowering)하는 컴파일러 프레임워크인 extbf{ScanWeaver}를 제안합니다. 우리는 현대적인 ML 워크로드에서 발생하는 광범위한 아핀 재귀 (Affine recurrences) 클래스의 동기 부여 사례로 Mamba 스타일의 선택적 스캔 (Selective scan)을 사용합니다. ScanWeaver는 단일 모델 제품군을 목표로 하는 대신, 이 재귀 구조를 일급 컴파일러 추상화 (First-class compiler abstraction)로 격상시켜, GPU 상에서 컴파일러가 생성한 Blelloch 스캔 실행으로의 체계적인 MLIR 기반 로워링을 가능하게 합니다. 일치하는 로컬 재귀 의미론 (Local recurrence semantics)을 가진 순방향 선택적 스캔 워크로드 전반에 걸쳐, 우리는 아핀 재귀 분해 (Affine recurrence decomposition), Blelloch 로워링, MLIR GPU 로워링, 실행 가능한 아티팩트 생성, 그리고 생성된 MLIR로부터의 실제 GPU 실행을 검증합니다. 우리는 결과물인 ScanWeaver GPU 실행을 PyTorch 및 CUDA 순차적 베이스라인 (Sequential baselines)과 벤치마킹하며, 시스템 맥락을 위해 Mamba 커널을 융합된 프로덕션 베이스라인 (Fused production baseline)으로 포함합니다.

AI 자동 생성 콘텐츠

원문 바로가기

ScanWeaver: 결합 법칙 스캔 로워링 (Associative Scan Lowering)을 통한 아핀 재귀 (Affine

요약

핵심 포인트

댓글