NVIDIA CUDA Tile로 Mamba-2 순방향 전달 커널 가속화

Mamba-2 순방향 전달 최적화 개요

NVIDIA와 글로벌 트레이딩 기업 Jane Street은 최신 시퀀스 모델인 Mamba-2의 순방향 전달(forward pass) 커널을 NVIDIA CUDA Tile(cuTile)을 활용하여 최적화했다. cuTile은 NVIDIA GPU를 위한 Python 기반 DSL(Domain-Specific Language) 및 병렬 프로그래밍 모델로, 스레드 및 워프 수준의 스케줄링을 추상화하여 타일 중심의 고수준 프리미티브를 통해 GPU 커널을 Python으로 작성할 수 있도록 지원한다. 이를 통해 컴파일러 최적화가 Tensor Core와 같은 고급 하드웨어 기능을 자동으로 활용한다.

Mamba-2는 기존 시퀀스 모델 아키텍처의 계산적 한계를 극복하기 위해 설계된 모델이다. 전통적인 소프트맥스 어텐션은 Tensor Core를 활용한 행렬 곱셈으로 학습은 빠르지만, 추론 시 계산 및 메모리 요구 사항이 시퀀스 길이에 비례하여 증가한다. 반면, 순환 신경망(RNN)은 학습이 느리고 병렬화가 어렵지만, 추론 시 시퀀스 길이에 관계없이 일정한 계산량과 고정된 메모리 공간을 요구한다. Mamba-2는 이 두 가지 접근 방식의 장점을 결합하여 고정된 크기의 은닉 상태와 Tensor Core를 효율적으로 사용하는 청크 단위의 재귀적 업데이트를 특징으로 한다.

Mamba-2 아키텍처의 핵심 연산

Mamba-2의 핵심 재귀식은 \(y_t = c_t^\top h_t + D x_t\)로 표현된다. 여기서 \(h_t\)는 토큰 간에 전달되는 고정 크기의 은닉 상태 행렬이며, \(x_t\)는 어텐션의 '값'과 유사한 역할을, \(c_t\)는 어텐션의 '쿼리'와 유사하게 상태를 혼합하여 출력을 생성하는 역할을 한다. 이러한 재귀적 공식은 순차적 추론에 적합하지만, 학습 및 사전 채우기(prefill) 단계에서는 더 빠른 처리가 필요하다.

병렬 스캔 알고리즘은 선형 상태 업데이트를 가능하게 하지만, 방대한 은닉 상태 텐서를 구체화해야 하므로 주로 메모리 바운드(memory-bound)가 되어 Tensor Core를 활용하기 어렵다. 대안으로, Tensor Core를 활용하기 위해 미래의 은닉 상태를 행렬 곱셈을 통해 직접 계산하는 방식이 도입되었다. 특히, 로그 공간으로 재구성하여 순차적인 곱셈을 누적 합으로 대체하고, \(a_t\)의 누적 합을 텐서 \(G\)로 추적함으로써 새로운 은닉 상태 및 출력 계산식을 도출한다. 이 과정에서 지수 함수 계산 시 안정성을 위해 인과적 마스킹(causal masking)이 적용된다.

세 가지 핵심 커널 구현

Mamba-2 순방향 전달은 세 가지 주요 커널로 구성된다. 첫 번째 커널인 `chunk_cumsum`은 각 청크 내에서 \(a_t\)의 누적 합을 계산한다. 이 값은 이후 단계에서 여러 번 사용되므로 한 번만 계산하여 저장한다. 이 커널은 cuTile의 `ct.bid`, `ct.load`, `ct.store`, `ct.cumsum` 프리미티브를 활용하며, 로그 공간에서 누적 합을 계산하고 `exp2` 함수를 사용하기 위해 `log2(e)`로 스케일링한다.

두 번째 커널인 `fwd_pass_state`는 각 청크의 시작 부분에서 초기 은닉 상태와 마지막 청크 이후의 최종 은닉 상태를 계산한다. 이 커널은 출력을 직접 계산하지 않고, `compute_out` 커널이 나중에 출력으로 투영할 전달 상태를 준비한다. 핵심적으로 이 커널은 청크에 대한 순차적 루프를 포함하며, `ct.extract`로 마지막 토큰을 추출하고, `exp2`를 사용하여 성능을 최적화하며, `ct.mma` API를 통해 Tensor Core를 활용한 행렬 곱셈-덧셈 연산을 수행한다.

세 번째 커널인 `compute_out`은 순방향 전달의 마지막 단계로, `fwd_pass_state`가 생성한 청크별 초기 은닉 상태를 받아 토큰 출력을 계산한다. 이 커널은 PyTorch의 직접 구현에서 발생할 수 있는 전체 상호작용 행렬의 전역 메모리 구체화를 방지하기 위해 세 가지 항을 하나의 타일 커널로 융합하고 해당 행렬을 레지스터에 유지한다. `ct.mma`를 사용하여 상호작용 행렬을 구축하고 전달 상태 항을 동시에 누적하며, `ct.arange`와 `ct.where`를 사용하여 하삼각 마스크를 생성하여 마스크 텐서의 구체화를 피한다. `ct.exp2`를 사용하여 감쇠를 적용하고, 잠재적인 오버플로우를 방지하기 위해 인과적 마스킹을 즉시 적용한다.

cuTile 기반 성능 최적화 기법

cuTile-Python DSL로 Mamba-2 워크로드를 정의하면서 여러 핵심 성능 최적화가 도입되었다. 첫째, cuTile은 `ct.load` 및 `ct.store` 연산에 지연 시간 힌트(latency hint)를 주석으로 달아 DRAM 트래픽을 모델링할 수 있게 한다. 컴파일러는 이 힌트를 활용하여 명령어 스케줄링 및 파이프라인 버퍼 깊이를 최적화한다.

둘째, NVIDIA Blackwell 5세대 Tensor Core에서 MMA(Matrix Multiply-Accumulate) 피연산자는 비대칭적이다. 피연산자 A는 Tensor Memory에 있을 수 있지만, 피연산자 B는 Shared Memory에 있어야 한다. `fwd_pass_state`와 같은 입력 융합 패턴의 경우, \(B^{\mathsf{T}}A^{\mathsf{T}} = (AB)^{\mathsf{T}}\) 항등식을 사용하여 융합된 피연산자를 피연산자 A로 라우팅함으로써 Shared Memory 경합을 완화한다.

셋째, GMEM(Global Memory) 레이아웃 최적화가 적용되었다. `decay_cumsum` 및 `dt`에 대한 기본 스트라이드 메모리 접근 방식은 캐시 친화적이지 않은 LDG 명령어로 이어진다. 이를 [batch, head, token] 레이아웃으로 전치하여 연속적인 접근을 보장하고 TMA(Tensor Memory Access) 사용을 가능하게 한다.

마지막으로, cuTile-Python 자동 튜너가 활용된다. 성능 병목 현상은 입력 형상에 따라 달라지므로, 모든 모델 구성에 맞는 단일 최적화 전략은 없다. 이러한 최적화는 cuTile 매개변수(block_size, occupancy)와 함께 통합된 탐색 공간에 통합되며, cuTile의 `exhaustive_search` 기능을 사용하여 각 입력 형상에 대한 최적의 설정을 자동으로 찾는다.

성능 벤치마크 결과 및 활용

구현된 Mamba-2 순방향 전달 커널의 성능은 Triton으로 작성된 오픈 소스 Mamba-2 커널 구현과 비교하여 평가되었다. 테스트는 단일 NVIDIA B300 SXM6 GPU(컴퓨팅 능력 10.3, Blackwell Ultra)에서 수행되었으며, 소프트웨어 환경은 mamba_ssm 2.3.2.post1, Triton 3.7.1, cuTile 1.5.0, PyTorch 2.13.0, CUDA 13.3 등으로 구성되었다.

벤치마크 결과, 최적화된 cuTile 순방향 전달 커널은 Triton 기반 구현 대비 1.14배에서 1.82배 더 빠른 성능을 보였다. 특히 시퀀스 길이와 상태 크기가 증가할수록 속도 향상이 더욱 두드러지는 경향을 나타냈다. 이 구현은 TileGym을 통해 사용할 수 있으며, `tilegym.ops.mamba2_chunk_forward`를 통해 접근 가능하다.

이러한 성능 향상은 커널 수준에서 측정된 것이며, 엔드투엔드 학습 결과는 역방향 전달(backward pass) 구현에도 영향을 받는다. cuTile은 Python으로 명확하고 가독성 높으며 고성능 커널을 작성할 수 있도록 지원하며, 사용자는 cuTile 문서, GitHub 저장소 및 샘플 코드를 통해 시작할 수 있다.