Python에서 CUDA 심층 신경망(cuDNN) 사용하기
cuDNN Python API를 사용해 스케일드 닷 프로덕트 어텐션(scaled dot product attention)을 구현하는 방법을 살펴봅시다. 이는 트랜스포머 계열 모델의 추론에서 계산 비용이 가장 큰 부분이면서도 부분적으로 병렬화할 수 있어서, 보통 GPU로 오프로드합니다. 내부적으로는 FlashAttention-2 알고리즘을 사용하지만, 단순한 원시 커널 구현보다 더 높은 수준의 API를 제공합니다.
먼저 필요한 라이브러리를 임포트합니다. cuDNN API에는 cudnn을, 텐서 연산과 비교에는 torch를 사용합니다. 최적화된 커널을 사용하려면 컴퓨트 성능(compute capability) SM80 아키텍처(Ampere) 이상의 GPU가 필요하다는 점에 유의하십시오.
import cudnn
import torch
import math
torch.manual_seed(0xDEADBEEF)
handle = cudnn.create_handle()
assert torch.cuda.is_available()
다음으로 GPT-2와 비슷한 구성을 바탕으로 어텐션 메커니즘의 문제 차원을 정의합니다. 배치 크기(b), 헤드 수(h), 최대 시퀀스 길이(s), 헤드당 임베딩 차원(d)이 그것입니다. 어텐션 스케일링 계수 attn_scale도 계산합니다.
b = 4
h = 12
s = 1024
d = 64
attn_scale = 1.0 / math.sqrt(d)
그런 다음 PyTorch를 사용해 GPU에 쿼리(Q), 키(K), 값(V), 출력(O) 텐서를 생성합니다. 이 텐서들은 반정밀도(.half())로 초기화합니다. 논리적 차원은 (b, h, s, d)(BHSD)로 정의합니다. 특히 명시적인 레이아웃 제어를 보여 주기 위해 strides를 사용하여 BSHD(배치, 시퀀스, 헤드, 헤드당 차원)에 해당하는 물리적 메모리 레이아웃을 지정합니다. as_strided 메서드는 데이터를 복사하지 않고 이 레이아웃을 적용합니다.
dims = (b, h, s, d)
strides = (s * h * d, d, h * d, 1) # BSHD physical layout
q_gpu = torch.randn(b * s * h * d).half().cuda().as_strided(dims, strides)
k_gpu = torch.randn(b * s * h * d).half().cuda().as_strided(dims, strides)
v_gpu = torch.randn(b * s * h * d).half().cuda().as_strided(dims, strides)
o_gpu = torch.empty(b * s * h * d).half().cuda().as_strided(dims, strides)
이제 cuDNN 계산 그래프를 구성합니다. cuDNN 그래프 시스템은 텐서 연산을 GPU에서 효율적으로 실행할 수 있는 데이터 흐름 그래프로 기술하기 위한 고수준 인터페이스를 제공합니다. 사용자는 그래프에 연산을 추가하는 방식으로 계산을 구성하며, 이 그래프는 일반적으로 전체 신경망 중에서 최적화하여 특화된 GPU 커널로 오프로드하려는 일부분을 나타냅니다. 그래프가 확정되면 cuDNN은 서로 다른 트레이드오프를 가진 여러 실행 엔진을 제공합니다. 어떤 엔진은 사용 편의성을 우선시하고, 어떤 엔진은 런타임 오버헤드를 최소화하거나 성능을 극대화합니다.
그래프는 입출력에는 반정밀도를, 중간 연산과 계산 연산에는 float 정밀도를 지정하여 초기화합니다. PyTorch GPU 텐서에 대응하는 그래프 텐서 디스크립터(q, k, v)를 생성합니다. SDPA 연산을 그래프에 추가하며, 추론용으로 구성하고(is_inference=True), 계산된 attn_scale을 사용하며, 인과 마스킹을 활성화합니다(use_causal_mask=True). 두 번째 반환값(stats 텐서)은 학습용이므로 무시합니다. 마지막으로 출력 그래프 텐서 o를 출력으로 표시하고 차원과 스트라이드를 설정합니다.
graph = cudnn.pygraph(
io_data_type=cudnn.data_type.HALF,
intermediate_data_type=cudnn.data_type.FLOAT,
compute_data_type=cudnn.data_type.FLOAT,
)
q = graph.tensor_like(q_gpu)
k = graph.tensor_like(k_gpu)
v = graph.tensor_like(v_gpu)
o, _ = graph.sdpa(
name="sdpa",
q=q,
k=k,
v=v,
is_inference=True,
attn_scale=attn_scale,
use_causal_mask=True,
)
o.set_output(True).set_dim(dims).set_stride(strides)
그래프를 정의했으면 이를 빌드합니다. 이 과정에는 그래프 구조 검증, 내부 연산 그래프 구축, cuDNN의 휴리스틱(Mode A와 Fallback)을 사용한 실행 계획 생성, 선택된 계획의 지원 여부 확인, 그리고 마지막으로 이 계획들의 빌드가 포함되며, 빌드 과정에서 커널의 JIT 컴파일이 일어날 수 있습니다. 휴리스틱 Mode A는 빠르면서도 대부분의 연산 그래프 패턴을 처리할 수 있도록 설계되었습니다. 휴리스틱 Mode B는 더 정확하지만 더 느리도록 설계되었습니다.
graph.validate()
graph.build_operation_graph()
graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
graph.check_support()
graph.build_plans()
그래프를 실행하기 위해 variant_pack 딕셔너리를 생성합니다. 이 딕셔너리는 (앞서 정의한) 심볼릭 그래프 텐서를 실제 데이터를 담고 있는 PyTorch GPU 텐서에 매핑합니다. 또한 실행 계획에 필요한 워크스페이스 크기를 조회하여 GPU에 할당합니다. 그런 다음 그래프를 실행하며, torch.cuda.synchronize()는 다음 단계로 넘어가기 전에 계산이 완료되도록 보장합니다.
variant_pack = {
q: q_gpu,
k: k_gpu,
v: v_gpu,
o: o_gpu,
}
workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8)
graph.execute(variant_pack, workspace)
torch.cuda.synchronize()
마지막으로, cuDNN SDPA 구현의 출력을 PyTorch의 네이티브 scaled_dot_product_attention과 비교하여 검증합니다. 원래의 GPU 텐서로부터 참조 텐서 q_ref, k_ref, v_ref를 준비하고, 이를 float로 변환합니다. cuDNN 텐서는 BSHD 물리적 레이아웃을 사용했고 PyTorch의 네이티브 함수는 BHSD 형식의 4D 입력을 기대하므로, 참조 텐서를 BSHD에서 BHSD로 순열 변환(permute)합니다. 그런 다음 참조 출력 o_ref를 계산합니다.
간단히 말해, permute(0, 2, 1, 3)은 4D 텐서의 차원을 재배열합니다. 원래 텐서의 차원이 (배치, 시퀀스, 헤드, 헤드당 차원)(BSHD)을 나타낸다면, 이 연산은 두 번째와 세 번째 차원(시퀀스와 헤드)을 맞바꾸어 (배치, 헤드, 시퀀스, 헤드당 차원)(BHSD) 차원을 갖는 텐서를 만듭니다. 이렇게 하는 이유는 torch.nn.functional.scaled_dot_product_attention 함수가 4D 입력 텐서를 BHSD 형식으로 기대하기 때문이며, 따라서 이 순열 변환은 텐서의 레이아웃이 함수의 요구 사항과 일치하도록 보장합니다.
q_ref = q_gpu.detach().clone().float()
k_ref = k_gpu.detach().clone().float()
v_ref = v_gpu.detach().clone().float()
# Permute from BSHD to BHSD for PyTorch native SDPA
q_ref = q_ref.permute(0, 2, 1, 3).contiguous()
k_ref = k_ref.permute(0, 2, 1, 3).contiguous()
v_ref = v_ref.permute(0, 2, 1, 3).contiguous()
o_ref = torch.nn.functional.scaled_dot_product_attention(
q_ref, k_ref, v_ref, is_causal=True, scale=attn_scale
)
o_gpu_for_comparison = o_gpu.float()
# Permute cuDNN output from BSHD to BHSD for comparison
o_gpu_for_comparison = o_gpu_for_comparison.permute(0, 2, 1, 3).contiguous()
torch.testing.assert_close(o_ref, o_gpu_for_comparison, atol=5e-3, rtol=3e-3)
print("PyTorch and cuDNN SDPA implementations match!")
현대 AI 도구 스택에서 cuDNN 프런트엔드 API 같은 도구는 고수준 딥러닝 프레임워크(PyTorch나 Jax 등)보다는 아래에, 하지만 원시 CUDA 프로그래밍보다는 위에 자리합니다. SDPA와 같이 고도로 최적화된 커널에 대해 이 수준의 제어를 직접 노출하면, 최적화할 때마다 C++까지 파고들 필요 없이 모델 구성 요소의 신속한 프로토타이핑과 성능 튜닝이 가능해집니다.