Jane Street·NVIDIA, cuTile 기반 Mamba-2 순전파 커널 구현 공개

Jane Street와 NVIDIA가 Python 기반 GPU DSL인 cuTile로 구현한 Mamba-2 순전파 커널 3종과 최적화 방법을 공개했다. 상태 갱신에 Tensor Core를 쓰고, 출력 계산은 하나의 커널로 합쳐 중간 행렬을 레지스터에 유지한다. 입력 형태별로 메모리 배치·블록 크기 등을 자동 탐색하는 방식도 설명했다.

내용·출처 업데이트 · 2026. 10. 20. 오전 8:10

chunk_cumsum은 청크 내부 감쇠항의 누적합을 계산해 뒤의 두 커널에서 재사용한다.

fwd_pass_state는 청크 사이의 은닉 상태를 행렬 곱으로 갱신하며 ct.mma를 통해 Tensor Core를 활용한다.

compute_out은 이전 상태의 기여분·감쇠 및 인과 마스크·스킵 연결 계산을 합치고 중간 상호작용 행렬을 레지스터에 유지한다.

Blackwell Tensor Core의 피연산자별 메모리 제약에 맞춰 곱셈 순서를 바꾸고, 메모리 배치를 전치해 TMA 접근을 유도한다.

입력 형태마다 병목이 달라 지연시간 힌트·메모리 배치·블록 크기·점유율을 함께 탐색한다. 제공된 본문에는 완결된 성능 비교 수치가 없다.

중간 상태와 상호작용 행렬을 전역 메모리에 펼치는 비용을 줄이는 커널 설계 사례로 참고할 수 있다.