NVIDIA CUDA Tile의 고성능 순전파 커널로 Mamba-2 가속하기

NVIDIA CUDA Tile(cuTile)은 NVIDIA GPU용 Python 기반 도메인 특화 언어(DSL)이자 병렬 프로그래밍 모델입니다. 데이터를 타일 단위로 처리하는 고수준 연산을 제공하므로, 개발자가 스레드와 워프의 스케줄링을 직접 다루는 부담을 줄여 줍니다. Python으로 GPU 커널을 작성하면 컴파일러가 Tensor Core를 비롯한 최신 하드웨어 기능을 활용하도록 최적화해 높은 성능을 끌어냅니다.

기술 중심의 글로벌 트레이딩 기업 Jane Street는 딥러닝을 활용해 퀀트 리서치를 발전시키고 있습니다. 이 회사의 머신러닝 엔지니어들은 cuTile을 초기에 도입해 시퀀스 모델인 Mamba-2를 구현하고, NVIDIA와 협력해 커널을 최적화했습니다. 이 글에서는 Mamba-2의 순전파 연산을 구성하는 세 가지 커스텀 커널을 살펴보고, 각 커널에 적용한 cuTile 최적화 기법을 설명합니다.

개요와 사전 요구 사항

cuTile을 사용하려면 먼저 cuTile Python Quickstart 가이드의 요구 사항을 환경이 충족하는지 확인하세요. 패키지는 cuda.tile로 임포트하며, 이 글의 모든 커널 구현에서는 ct라는 이름으로 사용합니다.

시퀀스 모델링의 연산 트레이드오프

시퀀스 모델링 아키텍처에는 연산 측면의 트레이드오프가 있습니다. 기존 소프트맥스 어텐션은 Tensor Core로 가속되는 행렬 곱셈 덕분에 학습이 빠르지만, 추론에서는 시퀀스 길이에 비례해 연산량과 메모리 요구량이 늘어납니다. 반면 순환 신경망인 RNN은 학습이 느리고 병렬화할 수 없지만, 시퀀스 길이와 관계없이 연산량이 일정하고 메모리 사용량도 고정되어 있습니다.

Mamba-2 아키텍처

Mamba-2는 두 방식의 장점을 모두 취하려는 모델 계열에서 나왔습니다. RNN처럼 고정 크기의 은닉 상태와 순환 업데이트를 사용하면서도, 이 업데이트를 Tensor Core로 청크 단위로 효율적으로 계산할 수 있습니다.

Mamba-2의 핵심 순환식은 다음과 같습니다.

ℎ𝑡 = exp(𝐴 ⋅ 𝑑𝑡𝑡) ℎ𝑡−1 + 𝑑𝑡𝑡 𝐵𝑡 𝑥𝑡

𝑦𝑡 = 𝑐𝑡ᵀ ℎ𝑡 + 𝐷 𝑥𝑡

각 항의 의미는 다음과 같습니다.

• 𝑥𝑡, 𝑑𝑡𝑡, 𝐵𝑡, 𝑐𝑡는 토큰 𝑡에서 Mamba-2 레이어의 입력입니다. 어텐션 레이어의 쿼리, 키, 값처럼 블록의 입력을 투영해 얻습니다.
• 𝑦𝑡는 토큰 𝑡의 출력입니다.
• ℎ𝑡는 토큰 간에 전달되는 고정 크기의 은닉 상태입니다. RNN처럼 벡터가 아니라 행렬이므로, 더 많은 수치를 담아 표현 용량이 큽니다.
• 𝐴와 𝐷는 학습되는 모델 파라미터입니다.

직관적으로 이해하면 다음과 같습니다.

• 𝑑𝑡𝑡는 헤드당 스칼라 하나로, 상태를 얼마나 갱신할지 정합니다. 이전 상태에 exp(𝐴 𝑑𝑡𝑡)를 곱하므로 𝑑𝑡𝑡 = 0이면 상태가 그대로 유지됩니다. 𝐴의 모든 값은 음수입니다.
• 𝑥𝑡는 크기 H의 벡터이며, 어텐션의 값과 비슷한 역할을 합니다.
• 𝐵𝑡는 크기 S의 벡터이며, 어텐션의 키와 비슷한 역할을 합니다.
• 𝑐𝑡는 크기 S의 벡터이며, 상태를 섞어 출력을 만드는 방법을 나타냅니다. 어텐션의 쿼리와 비슷한 역할을 합니다.

상태 ℎ𝑡는 크기가 (S, H)인 행렬입니다(𝐵𝑡 𝑥𝑡ᵀ는 외적입니다). 계산을 쉽게 쓰기 위해 𝑎𝑡 = 𝐴 𝑑𝑡𝑡로 두고, 곱 𝑑𝑡𝑡 𝐵𝑡를 𝑏𝑡라 하겠습니다. 그러면 식은 다음과 같이 간단해집니다.

ℎ𝑡 = exp(𝑎𝑡) ℎ𝑡−1 + 𝑏𝑡 𝑥𝑡ᵀ, 𝑦𝑡 = 𝑐𝑡ᵀ ℎ𝑡−1 + 𝐷 𝑥𝑡

이 순환 형태는 순차 추론에 이상적입니다. 하지만 학습과 프리필에는 for 루프보다 빠른 방법이 필요합니다. 선형 상태 업데이트에는 병렬 스캔 알고리즘을 쓸 수 있지만, 이 방식은 엄청난 크기의 은닉 상태 텐서를 메모리에 실제로 만들어야 합니다. 그 결과 대부분 메모리 바운드가 되어 Tensor Core를 활용하지 못합니다.

다른 방법은 행렬 곱셈으로 임의의 미래 은닉 상태를 곧바로 계산해 Tensor Core를 활용하는 것입니다. 단순하게 설명하기 위해 exp(𝑎𝑡) 항을 잠시 무시하면, ℎ𝑡에서 ℎ𝑡+𝑇로 가는 식은 다음과 같습니다.

ℎ𝑡+𝑇 = ℎ𝑡 + 𝑏𝑡+1 𝑥ᵀ𝑡+1 + 𝑏𝑡+2 𝑥ᵀ𝑡+2 + ⋯ + 𝑏𝑡+𝑇 𝑥ᵀ𝑡+𝑇

모든 𝑏𝑡+𝑖를 쌓아 shape (T, S)의 행렬 𝐵를, 모든 𝑥𝑡+𝑖를 쌓아 shape (T, H)의 행렬 𝑋를 만들면 이 식은 ℎ𝑡+𝑇 = ℎ𝑡 + 𝐵ᵀ𝑋로 정리됩니다.

이 점프 공식을 쓰면 시퀀스 곳곳의 청크 “체크포인트”에서 은닉 상태를 효율적으로 계산할 수 있습니다. 다음으로 할 일은 ℎ0만으로 출력 𝑦1, 𝑦2, 𝑦3, 𝑦4를 계산하는 방법을 찾는 것입니다. ℎ1, ℎ2, ℎ3, ℎ4, …를 모두 계산하면 쉽겠지만, 그러면 점프 공식을 쓰는 의미가 사라집니다.

감쇠 항 exp(𝑎𝑡)를 계속 제외한 채 계산하려는 값은 다음과 같습니다.

𝑦𝑡+𝑖 = 𝑐ᵀ𝑡+𝑖 ℎ𝑡+𝑖 + 𝐷 𝑥𝑡+𝑖 = 𝑐ᵀ𝑡+𝑖 (ℎ𝑡 + 𝑏𝑡+1 𝑥ᵀ𝑡+1 + ⋯ + 𝑏𝑡+𝑖 𝑥ᵀ𝑡+𝑖) + 𝐷 𝑥𝑡+𝑖

앞의 공식으로 은닉 상태를 전개한 식입니다. 스킵 연결 항 𝐷 𝑥𝑡+𝑖는 따로 생각하기 쉽습니다. 나머지 항을 그림으로 나타내면 다음과 같습니다.

마찬가지로 𝑥𝑡, 𝑏𝑡, 𝑐𝑡, 𝑦𝑡를 청크 단위로 행렬에 모으면, 이를 Tensor Core에 친화적인 행렬 곱으로 다시 표현할 수 있습니다(그림 3). 여기서 C와 B의 곱에 나타나는 삼각형은 인과적 소프트맥스 어텐션과 같은 인과 마스크입니다.

출력 𝑦𝑡+𝑖를 계산할 때는 𝑐𝑡+𝑖와 𝑗 ≤ 𝑡+𝑖인 모든 𝑏𝑗의 곱이 등장하므로, 전체 행렬 곱을 계산할 때는 미래에 해당하는 항(대각선 위쪽)을 마스킹해야 합니다. 이를 수식으로 쓰면 다음과 같습니다(스킵 연결을 다시 추가했습니다).

𝑌 = 𝐶 ℎ𝑡 + ((𝐶 𝐵ᵀ) ⋅ 𝑀) 𝑋 + 𝐷 𝑋

여기서 𝑀은 일반적인 인과 마스크입니다.

이 식은 소프트맥스를 뺀 어텐션과 매우 비슷해 보입니다. 감쇠 항을 제거하면 Mamba-2는 이름만 바뀐 선형 어텐션이기 때문입니다. 감쇠를 다시 넣으면 상황이 조금 복잡해집니다. 순환식을 전개하면 다음과 같습니다.

ℎ(𝑡+𝑇) = exp(𝑎(𝑡+1) + 𝑎(𝑡+2) + ⋯ + 𝑎(𝑡+𝑇)) ℎ𝑡 + exp(𝑎(𝑡+2) + ⋯ + 𝑎(𝑡+𝑇)) 𝑏(𝑡+1) 𝑥(𝑡+1) + exp(𝑎(𝑡+3) + ⋯ + 𝑎(𝑡+𝑇)) 𝑏(𝑡+2) 𝑥(𝑡+2) + ⋯ + 𝑏(𝑡+𝑇) 𝑥(𝑡+𝑇)

로그 공간으로 다시 쓰면 순차 곱셈을 누적 합으로 바꿀 수 있습니다. 𝐵(또는 𝑋)의 행에 지수 값을 미리 곱해 두면 여전히 행렬 곱셈으로 계산됩니다. 𝑎𝑡의 누적 합을 텐서 𝐺에 기록하면 새 식은 다음과 같습니다.

ℎ𝑡+𝑇 = exp(𝐺last) ℎ𝑡 + (exp(𝐺last − 𝐺) ⋅ 𝐵) 𝑋ᵀ

새로운 은닉 상태 점프를 출력 계산에 대입하면 다음 식을 얻습니다.

𝑌 = exp(𝐺) 𝐶 ℎ𝑘𝑇 + (exp(𝐺 − 𝐺ᵀ) ⋅ (𝐶 𝐵ᵀ) ⋅ 𝑀) 𝑋 + 𝐷 𝑋

행렬 exp(𝐺 − 𝐺ᵀ)의 (𝑖, 𝑗) 성분은 exp(𝑔𝑖 − 𝑔𝑗)입니다. 𝑖 < 𝑗인 항은 모두 인과 마스크로 가려지므로, 𝑔𝑖가 𝑔𝑗보다 더 많은 𝑎𝑡의 합인 항만 남습니다. 구조상 모든 𝑎𝑡 값은 음수이므로 𝑔𝑖 − 𝑔𝑗도 음수가 됩니다. 양수의 지수는 남지 않으므로 수치 안정성이 보장됩니다.

이 계산은 다음 세 커널로 구현됩니다.

• chunk_cumsum: 각 청크 안에서 𝑎𝑡의 누적 합을 계산합니다.
• fwd_pass_state: 위의 점프 공식으로 각 청크 시작 시점의 은닉 상태를 계산합니다.
• compute_out: 이전 단계에서 계산한 청크별 초기 은닉 상태로 각 청크의 출력을 계산합니다.

커널 1: 누적 합 계산

𝑎𝑡 = 𝐴 𝑑𝑡𝑡의 누적 합 (G)는 두 번째와 세 번째 커널에서 모두 사용합니다. 같은 값을 반복해서 계산하지 않도록 첫 번째 커널에서 누적 합을 구해 저장합니다. 이때 누적 합은 전체 시퀀스가 아니라 각 청크 안에서 계산합니다.

입력은 헤드별 스칼라 감쇠 파라미터 a와 토큰별 타임스텝 dt입니다. 헤드 h, 토큰 t의 유효 감쇠는 a[h] * dt[t, h]이며, 이를 청크 내 토큰 순서에 따라 누적합니다. 이 누적 합은 로그 공간의 값으로, 이후 다른 값과 선형 결합하거나 그대로 지수 함수에 입력합니다. 이때 계산 비용이 더 낮은 exp2를 exp 대신 사용할 수 있도록, 누적 합에 (\log_2(e))를 곱해 저장합니다.

이 커널에서는 그리드 내 위치를 확인하는 ct.bid, 데이터를 읽고 쓰는 ct.load와 ct.store, 누적 합을 계산하는 ct.cumsum을 사용합니다.

전체 커널은 짧습니다.

@ct.kernel
def chunk_cumsum_kernel(
    a: torch.Tensor,        # [num_heads]
    dt: torch.Tensor,       # [batch_size, num_tokens, num_heads]
    output: torch.Tensor,   # [batch_size, num_tokens, num_heads]
    chunk_size: ct.Constant[int],
    BLOCK_H: ct.Constant[int],
) -> None:
    chunk_idx, head_tile_idx, batch = ct.bid(0), ct.bid(1), ct.bid(2)
    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))

커널은 (num_chunks, num_head_tiles, batch_size) 그리드로 실행되며, 첫 줄에서 ct.bid로 그리드 안의 위치를 파악합니다. a_tile(shape [BLOCK_H])과 dt_tile(shape [1, chunk_size, BLOCK_H])을 브로드캐스트해 곱하면 타일 전체의 토큰별, 헤드별 유효 감쇠가 나옵니다. 이어서 ct.cumsum(..., axis=1)이 청크 내 토큰 차원을 따라 스캔하고, 결과는 fwd_pass_state와 compute_out이 기대하는 레이아웃으로 저장됩니다.

커널 2: fwd_pass_state

fwd_pass_state는 각 청크의 초기 은닉 상태와 마지막 청크 이후의 최종 은닉 상태를 계산합니다. 출력을 직접 계산하지는 않으며, 나중에 compute_out이 출력으로 투영할 전달 상태를 준비합니다. 핵심은 앞서 설명한 대로 청크를 순차적으로 도는 루프입니다. 커널을 한 번 실행할 때마다 아래와 같이 은닉 상태의 (batch, head, head_tile, state_tile) 슬라이스 하나를 로드합니다. 이 슬라이스 안에서 hid_tile은 shape이 [BLOCK_HEAD, BLOCK_STATE]이며 각 청크의 초기 은닉 상태 타일을 담습니다.

head_tile_idx, state_tile_idx, batch_head = ct.bid(0), ct.bid(1), ct.bid(2)
batch = batch_head // x.shape[2]
head = batch_head % x.shape[2]
group = head // num_heads_per_group
 
hid_tile = ct.load(
    init_state,
    index=(batch, head, head_tile_idx, state_tile_idx),
    shape=(1, 1, BLOCK_HEAD, BLOCK_STATE),
).reshape((BLOCK_HEAD, BLOCK_STATE))

루프를 돌 때마다 먼저 은닉 상태를 저장합니다.

num_chunks = ct.cdiv(b.shape[1], chunk_size)
for chunk_idx in range(num_chunks):
    ct.store(
        hiddens,
        index=(batch, head, chunk_idx, head_tile_idx, state_tile_idx),
        tile=hid_tile.reshape((1, 1, 1, BLOCK_HEAD, BLOCK_STATE)).astype(hiddens.dtype),
    )

루프의 나머지 부분은 x, b, dt, decay_cumsum의 청크 하나를 로드해 다음 청크의 초기 은닉 상태를 계산합니다. 앞서 본 수식은 다음과 같습니다.

ℎ𝑡+𝑇 = exp(𝐺last) ℎ𝑡 + (exp(𝐺last − 𝐺) ⋅ 𝐵) 𝑋ᵀ

ct.extract로 로드한 감쇠 타일에서 마지막 토큰을 꺼내 𝐺last를 계산합니다. 앞서 설명했듯 성능을 위해 exp 대신 exp2를 씁니다. ct.mma는 Tensor Core를 활용해 행렬 곱셈-덧셈을 실행하는 API입니다.

x_tile = ct.load(
       x, index=(batch, chunk_idx, head, head_tile_idx),
       shape=(1, chunk_size, 1, BLOCK_HEAD),
       padding_mode=ct.PaddingMode.ZERO,
   ).reshape((chunk_size, BLOCK_HEAD))
   b_tile = ct.load(
       b, index=(batch, chunk_idx, group, state_tile_idx),
       shape=(1, chunk_size, 1, BLOCK_STATE),
       padding_mode=ct.PaddingMode.ZERO,
   ).reshape((chunk_size, BLOCK_STATE))
   dt_tile = ct.load(
       dt, index=(batch, chunk_idx, head),
       shape=(1, chunk_size, 1),
       padding_mode=ct.PaddingMode.ZERO,
   ).reshape((chunk_size,))
   d_tile = ct.load(
       decay_cumsum, index=(batch, chunk_idx, head),
       shape=(1, chunk_size, 1),
       padding_mode=ct.PaddingMode.ZERO,
   ).reshape((chunk_size,))
   d_last = ct.extract(d_tile, index=(chunk_size - 1,), shape=(1,))
 
   b_dt_d = b_tile * (dt_tile * ct.exp2(d_last - d_tile))[:, None]
   hid_tile = ct.exp2(d_last) * hid_tile
   hid_tile = ct.mma(x_tile.transpose(), b_dt_d.astype(b.dtype), hid_tile)

마지막 청크 이후 hid_tile에는 최종 은닉 상태가 들어 있으며, 이를 저장합니다. 추론에서는 이 값을 다음 호출의 init_state로 되먹일 수 있고, 학습에서는 보통 버립니다.

ct.store(
    final_state,
    index=(batch, head, head_tile_idx, state_tile_idx),
    tile=hid_tile.reshape((1, 1, BLOCK_HEAD, BLOCK_STATE)),
)

순차 의존성은 단일 그리드 포인트(프로그램 인스턴스) 안에 머물고, 서로 독립적인 (batch, head, head_tile, state_tile) 슬라이스는 그리드 전체에서 병렬로 실행됩니다.

커널 3: compute_out

compute_out은 순전파의 마지막 커널입니다. fwd_pass_state가 만든 청크별 초기 은닉 상태를 받아 토큰 출력을 계산하며, 앞서 유도한 출력식을 구현합니다.

𝑌 = exp(𝐺) 𝐶 ℎ𝑘𝑇 + (exp(𝐺 − 𝐺ᵀ) ⋅ (𝐶 𝐵ᵀ) ⋅ 𝑀) 𝑋 + 𝐷 𝑋

PyTorch로 직접 구현하면 이 세 항을 따로 계산하며, [B, n_chunks, chunk_size, chunk_size] 크기의 상호작용 행렬 C @ B.T 전체를 글로벌 메모리에 구체화해야 합니다. compute_out은 세 항을 하나의 타일 커널로 융합하고, 이 행렬을 레지스터에 유지합니다.

커널을 한 번 실행할 때마다 (batch, chunk, head, head_tile) 출력 타일 하나를 맡아 두 단계로 작업합니다. 1단계에서는 상태 타일을 순회하면서 ct.mma로 상호작용 행렬 cb_tile = C @ B.T를 만들고, 동시에 전달 상태 항 C @ H.T를 누적합니다.

chunk_idx, head_tile_idx, batch_head = ct.bid(0), ct.bid(1), ct.bid(2)
batch = batch_head // x.shape[2]
head  = batch_head %  x.shape[2]
group = head // num_heads_per_group
 
cb_tile  = ct.full((chunk_size, chunk_size), 0.0, dtype=ct.float32)
out_tile = ct.full((chunk_size, BLOCK_HEAD),  0.0, dtype=ct.float32)
 
# Phase 1: build CB = C @ B.T and accumulate the carried-state output
for state_tile_idx in range(state_size // BLOCK_STATE):
    b_tile = ct.load(
        b, index=(batch, chunk_idx, group, state_tile_idx),
        shape=(1, chunk_size, 1, BLOCK_STATE), padding_mode=ct.PaddingMode.ZERO,
    ).reshape((chunk_size, BLOCK_STATE))
    c_tile = ct.load(
        c, index=(batch, chunk_idx, group, state_tile_idx),
        shape=(1, chunk_size, 1, BLOCK_STATE), padding_mode=ct.PaddingMode.ZERO,
    ).reshape((chunk_size, BLOCK_STATE))
    h_tile = ct.load(
        hiddens, index=(batch, head, chunk_idx, head_tile_idx, state_tile_idx),
        shape=(1, 1, 1, BLOCK_HEAD, BLOCK_STATE),
    ).reshape((BLOCK_HEAD, BLOCK_STATE))
 
    cb_tile  = ct.mma(c_tile, b_tile.transpose(), cb_tile)
    out_tile = ct.mma(c_tile, h_tile.transpose().astype(c_tile.dtype), out_tile)

2단계에서는 cb_tile에 감쇠와 인과 마스크를 적용하고(ct.arange와 ct.where로 하삼각 마스크 𝑀을 만들기 때문에 마스크 텐서를 따로 구체화하지 않습니다), 모든 항을 더해 최종 출력 타일을 얻습니다. 두 누산기는 모두 ct.full 레지스터 타일로 시작하며, 마지막 ct.mma는 ct.tfloat32로 캐스팅하여 FP32에 가까운 정확도로 Tensor Core에서 실행합니다.

# Phase 2: decay + causal mask then final computation
d_tile  = ct.load(decay_cumsum, index=(batch, chunk_idx, head),
                  shape=(1, chunk_size, 1), padding_mode=ct.PaddingMode.ZERO).reshape((chunk_size,))
dt_tile = ct.load(dt, index=(batch, chunk_idx, head),
                  shape=(1, chunk_size, 1), padding_mode=ct.PaddingMode.ZERO).reshape((chunk_size,))
x_tile  = ct.load(x, index=(batch, chunk_idx, head, head_tile_idx),
                  shape=(1, chunk_size, 1, BLOCK_HEAD), padding_mode=ct.PaddingMode.ZERO).reshape((chunk_size, BLOCK_HEAD))
d_val   = ct.load(d, index=(head,), shape=(1,))
 
cb_tile  = cb_tile * ct.exp2(d_tile[:, None] - d_tile[None, :])  # pairwise decay
idx      = ct.arange(chunk_size, dtype=ct.int32)
cb_tile  = ct.where(idx[:, None] >= idx[None, :], cb_tile, 0.0)  # causal mask
out_tile = out_tile * ct.exp2(d_tile[:, None])                   # decay carried state
 
dt_x     = dt_tile[:, None] * x_tile
out_tile = ct.mma(cb_tile.astype(ct.tfloat32), dt_x.astype(ct.tfloat32), out_tile)
out_tile = out_tile + d_val * x_tile                            # D * x skip
ct.store(out, index=(batch, chunk_idx, head, head_tile_idx),
         tile=out_tile.reshape((1, chunk_size, 1, BLOCK_HEAD)).astype(out.dtype))

앞서 언급했듯, ct.exp2(d_tile[:, None] - d_tile[None, :]) 줄에서는 수치적으로 불안정한 지수(일부는 양수)도 계산됩니다. 이 값들은 곧바로 마스킹하여, 오버플로 가능성이 있는 값이 계산을 오염시키지 않게 합니다.

Mamba-2 성능 최적화

Mamba-2 워크로드를 cuTile-Python DSL로 정의한 뒤, 몇 가지 핵심 성능 최적화를 도입했습니다.

• 지연 시간 힌트: cuTile에서는 ct.load와 ct.store 연산에 지연 시간 힌트를 달아 DRAM 트래픽을 모델링할 수 있습니다. 컴파일러는 이 힌트로 명령어 스케줄링과 파이프라인 버퍼 깊이를 최적화합니다.
• A/B 교환: NVIDIA Blackwell 5세대 Tensor Core에서는 MMA 피연산자가 비대칭입니다. 피연산자 A는 Tensor Memory에 둘 수 있지만 피연산자 B는 공유 메모리에 있어야 합니다. fwd_pass_state의 SIMT 연산과 Dot 연산처럼 입력 융합 패턴에서는 항등식 𝐵ᵀ𝐴ᵀ = (𝐴𝐵)ᵀ를 이용해 융합된 피연산자(b_dt_d)를 피연산자 A로 보내 공유 메모리 경합을 줄입니다.
• GMEM 레이아웃: decay_cumsum과 dt의 기본 스트라이드 메모리 접근은 캐시 효율이 낮은 LDG 명령어로 내려갑니다. [batch, head, token] 레이아웃으로 전치하면 연속적인 접근이 보장되어 TMA를 사용할 수 있습니다.
• cuTile-Python 오토튜너: 입력 shape에 따라 성능 병목이 달라지므로, 모든 모델 구성에 맞는 단일 최적화 전략은 없습니다. 이러한 최적화를 cuTile 파라미터(block_size, occupancy)와 함께 하나의 통합 탐색 공간에 넣고, cuTile의 exhaustive_search로 입력 shape마다 최적의 설정을 자동으로 찾습니다.

Mamba-2 순전파 성능 벤치마크

이 Mamba-2 순전파 커널 구현은 TileGym에서 이용할 수 있습니다. 전체 예제는 다음과 같습니다.

import torch
import tilegym
from tests import common
from tests.ops.test_mamba2 import Mamba2ProblemSpec
 
tilegym.set_backend("cutile")
num_tokens, state_size = 16384, 128
 
torch.manual_seed(0)
spec = Mamba2ProblemSpec(batch_size=4, num_tokens=num_tokens, num_groups=1, 
num_heads=8, state_size=state_size, head_size=64, 
chunk_size=128, dtype=torch.bfloat16, state_dtype=torch.float32)
fwd_inputs = spec.get_sample_inputs(with_reference=False, with_backward=False)
fwd = common.benchmark(
    lambda: tilegym.ops.mamba2_chunk_forward(**fwd_inputs),
    use_cupti=True,
)
print(f"Forward: {fwd['forward']['median'] * 1000:.2f} us")

구현한 Mamba-2 순전파의 성능은 Triton으로 작성된 오픈소스 Mamba-2 커널 구현과 비교해 다음 소프트웨어 및 하드웨어 구성에서 평가했습니다.

소프트웨어

• mamba_ssm 2.3.2.post1 (NVIDIA Dynamo-Triton reference), Triton 3.7.1
• cuTile: cuda-tile 1.5.0
• TileIR assembler (tileiras) 13.4
• PyTorch 2.13.0, CUDA 13.3 (Driver and Runtime API, NVCC V13.3.33), Driver 595.58.03
• cuda-python / cuda-bindings 13.3.1
• Python 3.12.3, Ubuntu 24.04.4 LTS, kernel 6.8.0-107

하드웨어

• NVIDIA B300 SXM6 1개, 컴퓨트 capability 10.3 (NVIDIA Blackwell Ultra)
• 최대 SM 클럭 1920 MHz, 전력 제한 1100 W, HBM 268 GiB

최적화한 cuTile 순전파 커널은 비교 대상보다 1.14~1.82배 빠른 성능을 보였으며, 시퀀스 길이와 상태 크기가 커질수록 성능 차이도 커졌습니다. 다만 이는 순전파 커널의 성능을 비교한 결과로, 전체 학습 속도는 역전파 구현의 성능에도 영향을 받습니다.

cuTile 시작하기

cuTile을 활용하면 Python으로 읽기 쉽고 성능도 뛰어난 GPU 커널을 작성할 수 있습니다. 문서와 GitHub 리포지토리에서 사용법을 살펴보고, 샘플 코드를 직접 실행해 보세요.