NVIDIA와 제인 스트리트가 Mamba-2의 입력을 받아 출력을 만드는 순방향 계산을 가속하는 GPU 커널, GPU에서 특정 연산을 수행하는 코드, 세 개를 공개했다. 모델을 최적화하는 개발자는 파이썬으로 작성된 공개 코드를 활용할 수 있으며, 회사 발표 기준 기존 오픈소스 커널보다 1.14~1.82배 빠른 결과가 나왔다.

1. 순환 모델 계산을 청크 단위로 바꾸다

Mamba-2는 앞서 처리한 정보를 고정 크기의 내부 기억인 ‘은닉 상태’에 담아 갱신하는 순환형 모델이다. 학습이나 프롬프트 처리에서는 토큰을 하나씩 계산하는 대신 여러 토큰을 묶은 청크의 계산을 행렬 곱셈으로 바꾼다. GPU의 행렬 연산 전용 장치인 Tensor Core를 활용하기 위해서다.

NVIDIA의 cuTile은 GPU 커널을 파이썬으로 작성하는 도구다. 개발자는 데이터 묶음인 타일의 연산을 작성하고, 세부 실행 배치는 컴파일러에 맡긴다. 제인 스트리트 엔지니어들이 이를 이용해 Mamba-2를 구현했고, NVIDIA와 함께 커널을 최적화했다.

2. 세 커널이 상태와 출력을 나눠 계산

첫 커널 chunk_cumsum은 이전 정보가 줄어드는 정도를 나타내는 감쇠율과 시간 간격의 곱을 청크 안에서 누적해 저장한다. 뒤의 두 커널이 이 값을 재사용한다.

둘째 fwd_pass_state는 청크 시작의 은닉 상태와 마지막 청크 이후의 상태를 계산한다. 청크 순서는 지키되 서로 독립적인 상태 조각은 병렬로 처리한다.

셋째 compute_out은 청크 시작 상태로 토큰 출력을 만든다. 중간 행렬을 큰 전역 메모리에 따로 쓰는 대신 계산을 합쳐 연산 장치 가까이의 작은 저장 공간인 레지스터에 유지한다.

블로그의 누적 합 커널은 타일을 읽고 청크 안에서 누적해 저장하는 흐름을 다음처럼 보여준다.

    dt_tile = ct.load(dt, index=(batch, chunk_idx, head_tile_idx),
                          shape=(1, chunk_size, BLOCK_H),
                          padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
    a_tile = ct.load(a, index=(head_tile_idx,),
                        shape=(BLOCK_H,),
                        padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)

    d_tile = ct.cumsum(a_tile * dt_tile, axis=1) * LOG2E
    ct.store(output, index=(batch, chunk_idx, head_tile_idx), tile=d_tile.astype(output.dtype))

3. 하드웨어에 맞춘 메모리·행렬곱 최적화

구현은 메모리 읽기·쓰기에 예상 지연 정보를 붙여 컴파일러가 명령 순서를 조정하도록 했다. Blackwell 5세대 Tensor Core의 입력 배치 특성에 맞춰 곱셈 입력의 순서를 바꿔 공유 메모리 경합, 여러 작업이 저장 공간을 다투는 현상도 줄였다. 감쇠 누적값과 시간 간격의 메모리 배치를 바꿔 데이터를 연속으로 읽도록 한 것도 최적화 요소다.

입력 형태에 따라 병목이 달라지는 만큼 블록 크기, 한 번에 처리할 데이터 묶음의 크기와 GPU 자원 활용 설정을 자동 탐색한다. 실행 시간을 줄이는 이익을 얻는 대신, 입력 형태별로 설정 조합을 탐색하는 작업이 필요해지는 구조다.

4. 회사 발표 기준 커널 측정은 최대 1.82배

NVIDIA 블로그의 측정은 단일 B300 SXM6에서 cuTile 1.5.0을 사용했다. GPU의 최대 연산 클럭은 1920 MHz, 고대역폭 메모리 용량은 268 GiB다.

비교 대상은 Triton 3.7.1로 작성된 오픈소스 커널이며, mamba_ssm 2.3.2.post1의 엔비디아 다이너모 계열 기준 구현이다.

측정에는 PyTorch 2.13.0과 CUDA 13.3을 사용했다.

회사 발표 기준 cuTile 커널은 Triton보다 1.14~1.82배 빠르며, 시퀀스 길이와 상태 크기가 커질수록 속도 향상이 커진다. 그래프는 1K~256K 토큰과 상태 크기 64, 128, 256, 512를 다루지만 각 막대에 정확한 수치 라벨은 없다.

이 결과는 순방향 커널 단위의 비교이며, 학습 전체의 성능은 역방향 계산, 출력에서 입력 쪽으로 되짚는 계산, 구현에도 달려 있다고 블로그는 설명한다.

5. TileGym에서 구현을 확인

코드와 벤치마크는 NVIDIA의 TileGym 저장소에서 확인할 수 있다. 기본 실행 방식은 cuTile이며, 저장소의 설치 지침을 따라 연산 예제와 성능 측정을 살펴볼 수 있다.

저장소는 cuTile 1.5.0 이상과 실행 시 필요한 컴파일러를 명시한다. 설치 안내에는 Blackwell GPU에서 CUDA 13.1 이상, 앰프레 GPU에서는 CUDA 13.2 이상을 지원한다고 적혀 있다.

이번 공개는 Mamba-2 순방향 계산의 행렬화와 하드웨어별 튜닝이 실측 이득으로 이어질 수 있음을 코드와 함께 제시한다.