cuLA(CUDA Linear Attention): NVIDIA Blackwell, Hopper GPU를 위한 선형 어텐션 고성능 CUDA 커널 라이브러리

cuLA 소개

트랜스포머(Transformer) 기반 언어 모델이 점점 더 긴 컨텍스트를 다루게 되면서, 기존의 소프트맥스 어텐션(Softmax Attention)이 가진 O(n^2) 계산 복잡도 문제는 핵심적인 병목이 되었습니다. 이를 해결하기 위해 등장한 선형 어텐션(Linear Attention)은 상태 업데이트를 선형 시간으로 수행함으로써 긴 컨텍스트 처리에 훨씬 유리한 구조를 제공합니다. 그러나 이론적인 효율성을 실제 하드웨어 성능으로 연결하려면, GPU의 특성에 맞게 세심하게 최적화된 커널 구현이 반드시 필요합니다.

cuLA는 선형 어텐션 변형들을 위한 고성능 CUDA 커널 라이브러리입니다. inclusionAI가 개발한 이 프로젝트는 CuTe DSL과 CUTLASS C++를 활용하여 NVIDIA Blackwell(SM10X) 및 Hopper(SM90) GPU 아키텍처에 특화된 핸드튜닝(hand-tuned) 구현을 제공합니다. cuLA는 flash-linear-attention(FLA)의 서브모듈로 설계되어 있으며, 기존 FLA 사용자라면 단 한 줄의 임포트 변경만으로 cuLA의 가속화된 커널을 활용할 수 있습니다. 현재 v0.1.0(2026년 4월 2일 출시)으로 초기 개발 단계에 있으며, 커뮤니티의 기여를 적극적으로 환영하고 있습니다.

cuLA가 지원하는 선형 어텐션 알고리즘

cuLA는 최근 연구에서 제안된 여러 선형 어텐션 변형들을 구현합니다. 각 알고리즘은 표준 어텐션의 한계를 다양한 방식으로 보완하며, cuLA는 이들을 Blackwell 및 Hopper GPU 위에서 최대한 효율적으로 동작하도록 최적화합니다.

KDA(Kimi Delta Attention) 는 Moonshot AI의 Kimi 모델에서 사용되는 어텐션 메커니즘으로, 델타 업데이트 방식으로 선형 상태를 갱신합니다. cuLA는 모듈식 KDA 포워드(Modular KDA Forward)와 퓨전 KDA 포워드(Fused KDA Forward) 두 가지 변형을 지원합니다. 관련 논문은 arXiv:2510.26692에서 확인할 수 있습니다.

GLA(Gated Linear Attention) 는 게이팅 메커니즘을 선형 어텐션에 결합하여 표현력을 향상시킵니다. arXiv:2312.06635에 발표된 연구로, 하드웨어 효율적인 학습을 가능하게 합니다.

GDN(Gated Delta Networks) 는 게이팅과 델타 업데이트를 조합한 변형으로, arXiv:2412.06464에서 제안되었습니다.

Lightning Attention은 청크(chunk) 기반 분해를 통해 선형 어텐션의 훈련 효율성을 크게 개선한 방법으로, arXiv:2405.17381에 발표되었습니다. cuLA는 프리필(Prefill)과 변길이(Varlen) 두 가지 모드를 모두 지원합니다.

cuLA의 벤치마크 성능

cuLA는 flash-linear-attention v0.4.2를 기준으로 성능을 측정하였습니다. 벤치마크는 NVIDIA GB300/GB200/H200 GPU, CUDA Toolkit 12.9, PyTorch 2.9.1, Triton 3.5.1 환경에서 수행되었습니다.

커널 하드웨어 평균 속도 향상
KDA Modular Forward (고정 길이) Blackwell (SM10X) avg 1.45x
KDA Modular Forward (가변 길이) Blackwell (SM10X) avg 1.32x (18개 설정)
Lightning Attention Prefill Blackwell (SM10X) 최대 1.86x (B=2)
Lightning Attention Varlen Blackwell (SM10X) avg 1.54x (126개 설정)
KDA Fused Forward Hopper (SM90) avg 1.52x

특히 Lightning Attention의 최대 1.86배 가속은 배치 크기가 작은 추론 시나리오에서 cuLA의 Blackwell 최적화가 상당한 이점을 제공함을 보여줍니다. KDA 퓨전 포워드 커널의 Hopper(H200) 상의 1.52배 향상은 현재 가장 널리 배포된 데이터센터 GPU에서도 실질적인 성능 개선을 얻을 수 있음을 의미합니다.

cuLA 설치 및 사용법

cuLA는 Python 3.12+, CUDA Toolkit 12.9+, NVCC 12.9+, PyTorch 2.9.1+ 환경이 필요합니다. PyTorch의 CUDA 버전과 시스템 CUDA Toolkit 버전이 반드시 일치해야 합니다.

# 저장소 클론 및 서브모듈 초기화
git clone https://github.com/inclusionAI/cuLA.git
git submodule update --init --recursive

# PyTorch 설치 (CUDA 12.9)
pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cu129

# 의존성 및 cuLA 설치
pip install -e third_party/flash-linear-attention
pip install -e . --no-build-isolation

기존에 FLA(flash-linear-attention)을 사용하고 있다면, 단 한 줄의 임포트 변경으로 cuLA 커널로 교체할 수 있습니다.

import torch
from cula.kda import chunk_kda  # fla.ops.kda에서 이 한 줄만 변경

B, T, H, K, V = 2, 2048, 32, 128, 128
device = 'cuda'

q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16, requires_grad=True)
g = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1
beta = torch.randn(B, T, H, device=device, dtype=torch.bfloat16).sigmoid()
A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.01
dt_bias = torch.zeros(H * K, device=device, dtype=torch.float32)
init_state = torch.zeros(B, H, K, V, device=device, dtype=torch.float32)

# 포워드 패스
o, final_state = chunk_kda(
    q=q, k=k, v=v, g=g, beta=beta,
    A_log=A_log,
    dt_bias=dt_bias,
    initial_state=init_state,
    output_final_state=True,
    use_qk_l2norm_in_kernel=True,
    use_gate_in_kernel=True,
    safe_gate=True,      # TensorCore 가속을 위해 필수
    lower_bound=-5.0,
)

# 백워드 패스
do = torch.randn_like(o)
o.backward(do)

print(f'출력 형태: {o.shape}')            # [2, 2048, 32, 128]
print(f'최종 상태 형태: {final_state.shape}')  # [2, 32, 128, 128]

주의할 점은 safe_gate=True가 TensorCore 가속을 위해 반드시 필요하며, beta는 float32 또는 bfloat16을 지원하지만 initial_state는 반드시 float32여야 한다는 것입니다. 가변 길이 시퀀스를 위한 cu_seqlens는 int32 타입을 사용해야 합니다.

라이선스

cuLA는 Apache-2.0 라이선스로 공개되어 있어 연구 및 상업적 목적으로 자유롭게 사용할 수 있습니다.

:github: cuLA 프로젝트 GitHub 저장소

더 읽어보기




이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:south_korea:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:wrapped_gift: 아래:down_right_arrow:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck:

1개의 좋아요