PyTorch 프로파일링 (Part 3): 어텐션은 모든 것을 아우른다

'PyTorch에서 프로파일링하기(Profiling in PyTorch)' 시리즈는 사용자가 프로파일러 트레이스 및 테이블을 편안하게 읽도록 하는 것이 목적입니다. Part 1에서는 덧셈과 곱셈 같은 기본적인 수학 연산을 프로파일링했습니다. 우리는 프로파일러 테이블이 어떻게 핫스팟(hotspots)을 밝혀내는지, 그리고 프로파일러 트레이스가 시간이 지남에 따라 알고리즘이 어떤 순서로 실행되는지 확인했습니다.

Part 2에서는 그 덧셈과 곱셈 연산을 torch linear layer로 감쌌습니다. 그런 다음 여러 개의 선형 레이어(multilayer perceptron)를 쌓아 올리고 이를 프로파일링했습니다. 이 과정에서 우리는 또한 결합된 커널(fused kernels) 및 수동으로 조정된 커널(hand-tuned kernels)도 프로파일링했습니다.

Transformer 아키텍처의 관점에서 볼 때, 우리가 다음으로 논리적으로 프로파일링해야 할 것은 또 다른 기본적인 알고리즘인 어텐션(attention)입니다. 이차 시간 복잡도(quadratic-time complexity)로 악명이 높지만, 이 문제를 완화하고 빠르게 만들 수 있는 많은 영리한 트릭들이 존재합니다. 여기서 우리의 목표는 모든 트릭을 자세히 다루는 것이 아닙니다. 대신, 프로파일러에서 각각의 트릭이 어떻게 다르게 보이는지 확인하는 것입니다.

이 블로그 게시물의 스크립트는 다음 위치에 있습니다:

04_a_naive_attention.py<br>,
04_b_inplace_ops_attention.py<br>,
04_c_sdpa_attention.py<br>,
그리고 04_d_kernels_attention.py.

이전과 마찬가지로, 이 파일들을 별도의 탭에서 열고 읽으면서 코드를 따라가 보는 것이 도움이 됩니다. 우리는 스크립트를 실행하기 위해 NVIDIA A100-SXM4-80GB GPU를 사용합니다. Hugging Face 인프라에 GPU를 설정하고 Dev Mode와 Spaces를 사용하여 스크립트를 실험하는 것은 매우 쉽습니다. 또한 Hugging Face Jobs 파이프라인으로도 스크립트를 실행할 수 있습니다.

어텐션은 쿼리(Queries, q), 키(Keys, k), 그리고 값(Values, v)을 사용합니다. 이들 간의 상호작용은 다음과 같은 짧은 단계들의 순서로 작성될 수 있습니다:

• 어텐션 점수 구축 (Build the attention scores):<br>scores := matmul(q, k.T)<br>- 점수 스케일링 (Scale the scores):<br>scores * scale<br>- 점수에 인과 마스크 적용 (Apply a causal mask to the scores):<br>scores.masked_fill(mask, "-inf")<br>- 소프트맥스를 사용하여 점수를 정규화하여 어텐션 가중치 얻기 (Normalize the scores with softmax to get the attention weights):<br>attn := softmax(scores)<br>- 이 가중치로 값 재가중치 부여 (Reweight the values with those weights):<br>matmul(attn, v)

따라서 어텐션은 실제로는 기본 연산들의 집합체입니다. 그중 일부는 이미 알고 있습니다(matmul). 나머지 부분은 쉽게 파악할 수 있습니다. PyTorch로 단순한 어텐션 모듈을 작성하고 이를 프로파일링해 봅시다.

class NaiveCausalAttention(nn.Module):
def __init__(self, head_dim):
super().__init__()
...

트레이스를 열기 전에, 평소처럼 우리가 무엇을 볼지 추측하는 연습을 해봅시다. 이 모듈의 forward를 트레이싱하면 다음이 예상됩니다:

• matmul 커널 (q . k.T) - mul 커널 (스케일링)
• 마스킹을 위한 연산
• softmax 커널
• matmul 커널 (atten . v)

uv run 04_a_naive_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

그림 1은 프로파일의 CPU 레인을 보여줍니다(GPU 레인은 너무 과부하가 걸리지 않도록 접혀 있습니다). attn_fwd (우리가 주석을 단 forward 호출) 내부에서 우리는 정확히 추측했던 연산들을 볼 수 있습니다. matmul은 이제 익숙한 친구이며, 새로운 연산들은 쉽게 파악할 수 있습니다:

mul: 스케일링
masked_fill: 인과 마스킹
softmax: softmax 커널

이제 GPU 레인을 펼쳐서 실제로 어떤 커널들이 실행되었는지 봅시다.

그림 2는 CPU 레인 옆에 GPU 레인을 보여줍니다. GPU 레인의 단일 attn_fwd 블록을 확대하여 커널들을 하나씩 살펴보겠습니다.

그림 3은 하나의 프로파일러 단계에 대한 개별 커널들을 읽어낼 수 있게 해줍니다:

• matmul (쿼리와 키)
• mul (스케일링)
• 메모리 복사 🤔
• 인과 마스킹
• softmax (어텐션 가중치를 생성합니다)
• matmul (어텐션 가중치와 값)

이 중 5가지가 예상되는 결과입니다. 메모리 복사(memory copy)는 예외적인 경우인데, 이것은 어디서 오는 것일까요? 단서는 PyTorch에 인플레이스 연산(in-place operations)이 있다는 것입니다. 일반적인 (아웃-오프레이스, out-of-place) 방식으로 텐서에 연산을 수행할 때, PyTorch는 종종 복사본을 만들고 요청된 연산을 그 복사본에 적용한 다음 복사본을 반환합니다. 일련의 연산 흐름을 따라가 보면, 여기서 문제의 원인은 우리의 masked_fill입니다.

만약 이것을 인플레이스 연산으로 대체한다면 어떨까요?

우리가 바꿀 것은 masked_fill을 masked_fill_로 바꾸는 것뿐입니다. (뒤에 언더스코어(underscore)가 붙은 것을 주목하세요. 이는 PyTorch의 인플레이스 연산 관례입니다.) 그리고 동일한 스크립트를 실행해 봅니다.

def forward(self, q, k, v, mask):
# q, k, v: [batch, heads, seq, head_dim]
scores = torch.matmul(q, k.transpose(-2, -1)) # [batch, heads, seq, seq]
...

트레이스(trace)를 살펴보고 무언가 달라졌는지 확인해 봅시다.

uv run 04_b_inplace_ops_attention.py
uvx trace-util -f traces/ -b <hf_uname>/traces

인플레이스 버전(Figure 5)은 마스킹 단계에서 아웃-오프레이스 버전(Figure 4)보다 훨씬 적은 CPU 연산(CPU ops)을 감싸고 있습니다. 이는 고무적인 신호입니다. GPU 레인을 펼쳐서 거기서 무슨 일이 일어났는지 확인해 봅시다.

GPU 레인에서는 Memcpy 커널이 완전히 사라졌습니다 (Figure 6 및 Figure 7). 단 한 줄의 변경으로 각 포워드 패스(forward pass)에서 전체 커널 하나를 제거했습니다. 이것 자체만으로는 크지 않아 보일 수 있지만, 이것이 단 하나의 어텐션 연산이라는 것을 기억하세요. 트랜스포머 기반 대규모 모델(LLMs, Diffusion models 등)의 맥락에서는 레이어마다 반복되며, 레이어가 많기 때문에 절약되는 양은 빠르게 쌓입니다 (그리고 만약 이로 인해 급여 인상을 받게 된다면, 최소 10%는 저희와 공유해 주시는 것이 공평할 것 같습니다).

아웃-오프레이스 방식이 PyTorch의 기본값인 데에는 이유가 있습니다. 그래디언트(gradient)를 계산하기 위해 autograd는 포워드 패스에서 본 텐서 값들을 기억해야 하는데, 많은 백워드 공식들이 이 값들을 재사용하기 때문입니다. 인플레이스 연산은 메모리 내에서 그 값들을 덮어쓰기 때문에, 백워드 패스가 잘못된 숫자를 읽게 됩니다. 따라서 우리는 forward를 torch.no_grad 아래에서 실행합니다.

, in-place는 우리에게 안전합니다. 역전파(backward pass)가 없고 손상시킬 것이 없기 때문입니다. 또한 인플레이스 연산은 시간만 절약하는 것(우리가 사례에서 본 것처럼)뿐만 아니라 메모리도 절약한다는 점이 주목할 만합니다 (추가 복사본이 없으므로). 이는 logits와 같은 대형 텐서에 매우 유용합니다!

저희는 방금 기본 요소들로부터 어텐션을 구축했고, 심지어 Memcpy까지 줄였습니다.

좋은 소식은 PyTorch 팀이 이 모든 것을 우리를 위해 처리했으며, 전체 파이프라인을 단일 함수로 패키징했다는 것입니다:

from torch.nn import functional as F
F.scaled_dot_product_attention(q, k, v, is_causal=True)

이 한 줄 코드가 우리가 직접 작성한 모듈을 대체하며, is_causal=True는 심지어 수동으로 마스크를 구축하는 것까지 절약해 줍니다. 이 단일 호출이 얼마나 많은 것을 숨기고 있는지 감상할 가치가 있습니다. 그리고 이것은 단순히 코드 라인 이상의 것을 숨깁니다. Scaled Dot Product Attention (SDPA)는 단일 구현을 가지고 있지 않습니다. 내부적으로 여러 백엔드 중 하나로 *디스패치(dispatches)*되며, 입력값(dtype, head dimension, mask, 하드웨어 등)을 지원하는 가장 빠른 백엔드를 선택합니다.

공식 SDPA 튜토리얼은 이 선택 과정을 안내하며, 백엔드 자체는 torch.nn.attention.SDPBackend enum에 나열되어 있습니다:

from torch.nn.attention import SDPBackend
BACKENDS = {
"math": SDPBackend.MATH,
...

일반적으로 SDPA가 우리를 위해 선택하지만, torch.nn.attention.sdpa_kernel context manager를 사용하여 특정 백엔드를 고정할 수 있습니다. 이것이 우리가 스크립트에서 하는 일입니다. 이를 통해 각 백엔드를 개별적으로 프로파일링하고 트레이스(trace)에서 얼마나 다르게 나타나는지 확인할 수 있습니다. 하나씩 살펴보겠습니다.

uv run 04_c_sdpa_attention.py --backend math
uvx trace-util -f traces/ -b <hf_uname>/traces

무엇이든 열어보기 전에 추측해 봅시다. 우리는 직접 작성한 어텐션(matmul, mul, mask, softmax, matmul)을 단일 한 줄 코드로 대체했으므로, 트레이스가 더 단순하고 빨라질 것으로 예상해야 합니다. 커널(kernels)이 적고, CPU 디스패치도 적으며, 심지어 융합된 커널(fused kernel)까지 있을 수 있습니다. 먼저 프로파일러 테이블을 확인해 봅시다.

이것이 우리의 첫 번째 놀라움입니다. 한 줄짜리 구현(one liner)이 3.7배 더 느립니다.

트레이스(Figure 9)를 열어보면 왜 경고등이 켜졌는지 알 수 있습니다. math 백엔드는 naive 어텐션 구현(Figure 8)에서 사용된 5개 대신, 포워드 패스당 20개의 GPU 커널을 실행합니다. 이는 우리가 예상했던 것과 정반대입니다. 왜 이런 일이 발생하는지 알아봅시다.

Part 2에서는 커널 이름 같은 것을 지문처럼 읽는 법을 배웠습니다. 그 습관을 여기서 사용해 봅시다:

이 트레이스들을 포착하는 데 사용된 A100은 Tensor Cores를 탑재하고 출고됩니다. Tensor Cores는 matmul(행렬 곱셈)을 가속화하기 위해 특화된 하드웨어로, 일반적인 CUDA 코어보다 훨씬 빠르다는 것으로 알려져 있습니다. 이것이 왜 중요한지 이해하려면 GPU 내부에 무엇이 있는지 아는 것이 도움이 됩니다. 스트리밍 멀티프로세서(Streaming Multiprocessor, SM)는 GPU의 컴퓨팅 유닛이며, 각 SM에는 두 가지 종류의 산술 유닛인 CUDA 코어와 Tensor Cores가 있습니다. CUDA 코어는 범용적이며 한 번에 소수의 요소를 처리하는 반면, Tensor Cores는 단일 명령어(single instruction)로 전체 작은 행렬 타일(small matrix tile)을 곱하고 누적합니다. 따라서 질문은 간단합니다. "각 백엔드가 실제로 빠른 경로를 사용하고 있는가?"

커널 이름이 답을 알려줍니다. naive 커널(Figure 10)에 있는 s16816는 bfloat16 Tensor Core matmul(즉, 16x8x16 Tensor Core 명령어)의 시그니처이며, 따라서 naive 버전은 빠른 경로를 사용하고 있습니다. sgemm(Figure 11)은 일반적인 CUDA 코어에서 실행되는 고전적인 단정밀도(FP32) matmul입니다. 다시 말해, math 백엔드는 Tensor Cores에 전혀 접근하지 않습니다. 속도를 위해 수치적 정확성을 얻으려면 텐서를 FP32로 업캐스팅(upcast)합니다(입력이 bf16일 때도 데이터 전송량이 두 배가 됩니다). 그리고 더 느린 CUDA 코어로 폴백(fallback)합니다.

naive 버전에서는 인과 마스크(causal mask)를 한 번 만들고 재사용했습니다. 여기서는 is_causal=True로 전달했습니다.

그리고 수학 백엔드는 모든 호출마다 우리를 위해 하나를 구현했습니다. CPU 레인에서 일어나는 것을 볼 수 있습니다:

Figure 12에서 보이는 것이 바로 이것입니다.

aten::ones -> aten::tril build a [seq, seq] lower-triangular matrix
aten::scalar_tensor -> aten::fill_ make the -inf fill value
aten::where turn it into an additive bias (0 or -inf)

GPU에서는 이것이 triu_tril_kernel, 여러 개의 where 커널, 그리고 add_ 로 나타납니다. 마스크에 대해 생각하는 것을 멈추게 해준 편리한 플래그는 작업을 제거하지 않았고, 단지 그 작업을 한 단계 아래로 옮겨서 매 순전파(forward)마다 마스크가 처음부터 다시 구축되도록 했습니다.

저희가 직접 작성한 버전은 단순히 aten::softmax 라고 불렸습니다. 수학 백엔드는 aten::_safe_softmax 를 호출하며, 차이점은 다시 추가적인 커널로 보입니다 (Figure 13):

모든 항목이 마스크된(각 항목 -inf) 행은 일반적인 softmax가 exp(-inf)/sum(exp(-inf)) = 0/0 = NaN 을 계산하게 만들 것입니다.
_safe_softmax는 정확히 그것을 방지합니다. 저희의 순진한 커널은 결코 신경 쓰지 않았고, 그 경우에 조용히 NaN 을 생성했을 것입니다.

종합적으로 볼 때, 수학 백엔드가 참조 구현입니다. 이는 어텐션을 원시 ATen 연산(primitive ATen ops)으로 분해하는 간단하고, dtype-안전하며, NaN-안전한 방식입니다. 본질적으로 저희가 손으로 작성했던 순진한 어텐션과 같지만 더 신중합니다. 그 신중함이 바로 그것을 극도로 느리게 만드는 이유입니다.

그것의 임무는 빠르다는 것이 아니라 항상 작동하는 것입니다. 이것이 완벽한 기준선(baseline)을 만듭니다. 우리가 다음에 프로파일링할 모든 백엔드(flash, efficient, cudnn)는 20 개의 GPU 커널을 본질적으로 하나의 융합된 커널로 축소하여 bf16에 머무르면서 중간 행렬을 전혀 구현하지 않으려고 합니다.

uv run 04_c_sdpa_attention.py --backend efficient
uvx trace-util -f traces -b <hf_uname>/traces

수학 백엔드가 하나의 프로파일러 단계에 걸쳐 20개의 커널을 실행한 반면, 효율적인 백엔드는 단지 fmha_cutlassF_bf16_aligned_64x64_rf_sm80 만을 실행합니다 (Figure 14에서 볼 수 있듯이).

이 커널의 이름을 해독해 봅시다:

fmha

(fused multi-head attention): 어텐션의 모든 기본 연산이 이제 하나의 연산으로 '융합(fused)'되었습니다. cutlassF<br><br>:cutlassF는 CUTLASS (NVIDIA의 텐서 코어 GEMM용 오픈소스 템플릿)를 기반으로 하며, F는 순전파(forward)에 사용됩니다. bf16_aligned: bfloat16으로 실행되며 (math와 달리 FP32 업캐스트 없음), 64x64: 타일 크기이며, rf (레지스터 파일): 작업 집합이 레지스터에 유지되어 칩에서 가장 빠른 메모리입니다. sm80: Ampere(A100의 컴퓨팅 능력 8.0)용으로 컴파일되었습니다.<br><br>이는 Meta의 xformers 라이브러리에서 파생되어 PyTorch로 상위 호환된, 메모리 효율적인 어텐션 커널입니다. 사람들이