행렬 곱셈과 어텐션

2026년 5월 1일 · 태그: ai, python, llms, compilers, mlir · 원문 · 한국어 번역(원문에서 변경됨)

GELU, 소프트맥스, 레이어 정규화는 이미 텐서 언어 위의 Python 함수로 표현했습니다. 그러나 트랜스포머 산술 연산의 대부분은 행렬 곱셈에서 나옵니다. 행렬 곱셈을 추가하면 컴파일러에 특수한 어텐션 커널을 도입하지 않고도 어텐션을 조립할 수 있습니다.

시퀀스 길이는 $T$ , 임베딩 폭은 $C$ , 헤드 수는 $H$ , 헤드 하나의 폭은 $D = C / H$ 로 표기하겠습니다. 모델은 한 번에 시퀀스 하나를 처리합니다. 헤드는 어텐션 행렬 곱에서 배치 차원이 됩니다.

행렬 연산 보존하기

형태가 $(M, K)$ 인 $A$ 와 형태가 $(K, N)$ 인 $B$ 에 대해:

$Y_{ij} = k = 0 \sum K - 1 A_{ik} B_{kj}$

프런트엔드의 __matmul__은 축약 차원을 검사하고, 출력 형태를 결정하며, 필요하면 dtype 변환을 삽입합니다. 그 결과는 matmul 표현식입니다. 이 표현식의 빌더는 이름 있는 linalg.matmul 연산을 생성합니다:

initial = linalg.fill(
    zero, outs=[tensor.EmptyOp(output_shape, element_type).result]
)
result = linalg.matmul(lhs, rhs, outs=[initial])

zero, 피연산자, 출력 타입은 타입이 지정된 표현식에서 가져옵니다. 행렬 곱셈은 결과를 대상에 누적하므로 0으로 초기화하는 것이 필수적입니다.

이 시점에서는 아직 루프 순서, 타일 크기, 벡터 폭, 스레드 매핑을 지정하지 않았습니다. IR은 여전히 이것이 행렬 곱셈이라고 말하고 있습니다. 따라서 백엔드가 이후에 구현 방식을 선택할 수 있습니다. CPU 파이프라인은 시스템 BLAS를 사용할 수 있으면 연산을 BLAS에 넘기고, 그렇지 않으면 MLIR의 루프 로어링을 사용합니다. GPU 타깃은 타일링된 커널로 대체합니다. 두 대체 모두 실행 시점에 연산을 여전히 알아볼 수 있어야 가능합니다.

작은 직사각형 예제로 형태 규칙과 생성된 연산을 확인합니다:

import numpy as np
from tinymlir import jit

@jit
def multiply(a, b):
    return a @ b

rng = np.random.default_rng(7)
a = rng.normal(size=(3, 5)).astype(np.float32)
b = rng.normal(size=(5, 7)).astype(np.float32)
np.testing.assert_allclose(multiply(a, b), a @ b, rtol=2e-6, atol=2e-6)
print(multiply.compile(a, b).module)

직사각형 입력은 뒤바뀐 차원과 잘못된 스트라이드를 드러내는 데 도움이 됩니다. 정사각형 예제만 모아 두면 이런 오류가 놀라울 만큼 잘 숨겨질 수 있습니다.

배치 차원이 일치하는 3차원 피연산자에 대해서는 빌더가 linalg.batch_matmul을 선택합니다. 현재 프런트엔드는 이 배치 차원들이 정확히 일치할 것을 요구합니다. NumPy의 배치 브로드캐스팅 규칙 전체를 구현하지는 않습니다.

선형 레이어

선형 레이어는 다음과 같이 간단합니다:

def linear(x, weight, bias):
    return x @ weight + bias

행렬 곱은 형태 (T, N)을 만들고, 벡터 편향은 그 행들에 걸쳐 브로드캐스트됩니다. GPT-2의 프로젝션 가중치는 형태가 (input_width, output_width)이므로 이 표현식에 그대로 들어맞습니다.

이 정의는 프런트엔드가 왜 중요한지도 보여 줍니다. gelu(linear(x, weight, bias))를 합성하면 전체 계산이 컴파일러에 드러납니다. 별도로 유지 관리하는 linear_gelu 구현은 없습니다. 텐서 연산은 MLIR의 융합 패스가 작업할 재료를 제공합니다.

쿼리, 키, 값

GPT-2는 세 가지 투영을 한 번에 계산합니다:

$[Q; K; V] = X W_{q k v} + b_{q k v}$

결과의 형상은 (T, 3*C)입니다. 세 개의 슬라이스를 취하여 각각을 (T, H, D)로 재형상화한 다음, (H, T, D)로 전치합니다:

packed = linear(x, qkv_weight, qkv_bias)
length, width = x.shape
depth = width // heads
q, k, v = [
    packed[:, i*width:(i+1)*width]
        .reshape((length, heads, depth))
        .transpose((1, 0, 2))
    for i in range(3)
]

이 문장들은 트레이싱 중에 실행됩니다. 루프의 범위는 정적이며, 모든 슬라이스와 형상이 알려져 있습니다. 배열 자체는 여전히 런타임 입력입니다.

각 연산에는 대응하는 구조화된 로어링이 있습니다. 슬라이스는 tensor.extract_slice가 됩니다. 재형상화는 상수 형상 피연산자를 가진 tensor.reshape가 됩니다. 전치는 순열 속성을 가진 linalg.transpose가 됩니다.

재형상화 연산의 선택은 의도적입니다. MLIR에는 tensor.collapse_shape와 tensor.expand_shape도 있습니다. 이 연산들은 명시적인 재결합 그룹을 지니며, 두 정적 형상 사이의 재형상화에 자연스럽게 들어맞습니다. 그러나 원소별 융합 패스는 이러한 연산을 인접한 linalg.matmul에 접어 넣는 방법을 알고 있으며, 명명된 연산을 재형상화를 흡수하는 인덱싱 맵을 가진 linalg.generic으로 다시 작성합니다. 이 변환은 적법하며 연산 하나를 절약합니다. 동시에 그 연산이 행렬 곱셈이었다는 사실도 지워 버리는데, 이것이 바로 다음 장들의 BLAS 및 GPU 대체가 찾는 정보입니다. 형상 피연산자를 가진 tensor.reshape는 융합 패스가 인식하는 재결합이 아니므로, matmul은 백엔드가 이를 가져갈 수 있는 지점까지 살아남습니다.

재형상화는 원소 개수를 보존해야 합니다. 전치는 순열을 지정해야 합니다. 슬라이스는 정적 형상에 들어맞아야 하며, 우리의 작은 언어에서는 단위 스트라이드를 가지고 비어 있지 않은 결과를 내야 합니다. 이러한 검사는 MLIR 연산을 생성하기 전에 이루어집니다.

레이아웃 연산은 버퍼화와 로어링 때까지 그 의미를 유지합니다. 특정 조합이 뷰가 될지 복사가 필요할지는 그 레이아웃과 사용 방식에 따라 달라집니다. Python 표현식이 짧다는 이유만으로 모든 전치가 비용 없이 이루어진다고 가정해서는 안 됩니다.

점수와 마스크

각 헤드에 대해 어텐션 점수는 다음과 같습니다:

$S_{h} = \frac{Q _{h} K _{h}^{⊤}}{D} + M$

배치 곱의 형상은 (H, T, T)입니다. 인과 마스크의 형상은 (T, T)이며 헤드 전체에 걸쳐 브로드캐스트됩니다. NumPy 참조 구현과 일치시키기 위해, 허용된 항목은 0이고 미래 위치에는 투영의 dtype으로 표현한 $- 1 0^{10}$ 이 들어갑니다.

우리의 마스크는 linalg.generic으로 만든 작은 프리미티브입니다. 그 영역은 linalg.index를 사용해 행과 열 인덱스를 얻고, arith.cmpi로 이를 비교한 다음, 0 또는 마스크 값을 선택합니다. 이는 인덱스에 의존하는 텐서 구성이며, 마스크를 담은 호스트 NumPy 배열을 필요로 하지 않습니다.

대각선은 여전히 허용되므로, 모든 쿼리 위치에는 보이는 키가 적어도 하나 있습니다. 10부의 소프트맥스를 적용하면 마지막 차원이 정규화됩니다:

scores = (q @ k.transpose((0, 2, 1))) / np.sqrt(depth)
probabilities = softmax(scores + causal_mask(length, packed.dtype))

NumPy 제곱근은 float64 스칼라를 반환합니다. 따라서 우리의 타입 규칙은 참조 구현과 마찬가지로 점수 계산을 승격합니다. 이 전파는 런타임 변환 속에 숨겨지지 않고, IR에서 명시적인 원소별 캐스트로 드러납니다.

값 결합하기

확률은 값 벡터에 가중치를 부여합니다:

$O_{h} = P_{h} V_{h}$

결과의 형태는 (H, T, D)입니다. 이를 (T, H, D)로 전치하고 (T, C)로 reshape하면 헤드가 결합됩니다. 마지막 선형 투영은 잔차 스트림의 폭으로 되돌립니다:

context = (probabilities @ v).transpose((1, 0, 2)).reshape((length, width))
return linear(context, out_weight, out_bias)

이 조각들을 합치면 tinymlir/ops.py에 있는 어텐션 함수 전체가 됩니다. 이 함수의 백엔드 어휘는 행렬 곱, 원소별 산술, 축약(reduction), 슬라이스, reshape, 전치로 이루어져 있습니다. 컴파일러에는 모델의 어텐션 함수를 이름으로 인식하는 특별한 규칙이 필요하지 않습니다.

이 버전은 점수 텐서와 확률 텐서를 실체화합니다. 그 크기는 $H T^{2}$ 에 비례합니다. 코드는 간결하지만, 알고리즘은 여전히 이차 크기의 어텐션 저장 공간을 사용합니다. FlashAttention을 구현하려면 타일 단위로 처리하고, 그 전체 텐서를 실체화하지 않은 채 정규화 통계를 유지하는 다른 스케줄이 필요합니다. 이는 일반적인 원소별 융합으로 얻어지는 결과가 아닙니다.

잔차 연결

트랜스포머는 어텐션 전에 입력을 정규화한 다음, 어텐션 출력을 다시 입력에 더합니다:

@jit
def attend(x, gain, bias, qkv_weight, qkv_bias, out_weight, out_bias, *, heads):
    normalized = layer_norm(x, gain, bias)
    return x + attention(
        normalized, qkv_weight, qkv_bias, out_weight, out_bias, heads
    )

heads는 정적 키워드 매개변수입니다. 이 값을 바꾸면 다른 특수화가 선택됩니다. 텐서는 동적 데이터이며, 그 형태와 dtype이 특수화를 결정합니다.

피드포워드 절반도 같은 잔차 구조를 가집니다:

@jit
def feedforward(x, gain, bias, up_weight, up_bias, down_weight, down_bias):
    normalized = layer_norm(x, gain, bias)
    hidden = gelu(linear(normalized, up_weight, up_bias))
    return x + linear(hidden, down_weight, down_bias)

이것들이 최종 모델이 실행하는 함수입니다. 각 함수는 하나의 모듈 안에 조합 가능한 텐서 연산을 여러 개 담고 있어, 컴파일러가 변환에 활용할 수 있는 유용한 생산자-소비자 관계를 제공합니다.

인과성 확인하기

수치 비교는 필요하지만, 어텐션에는 유용한 구조적 성질도 있습니다. 미래의 입력을 바꾸더라도 이전 출력은 변하지 않아야 합니다.

테스트 스위트는 짧은 시퀀스를 만들고, 컴파일된 어텐션을 실행한 뒤, 마지막 두 입력 행을 바꾸고, 처음 세 출력 행이 그대로인지 확인합니다. 이를 통해 뒤집힌 마스크, 잘못된 전치, 그리고 형태 검사만으로는 잡아낼 수 없는 일부 헤드 레이아웃 실수를 잡아낼 수 있습니다.

또한 비정사각 행렬 곱, 배치 곱, 그리고 자명하지 않은 슬라이스에서의 레이아웃 변경도 테스트합니다. 이들은 인덱싱 오류가 모델에 유입될 수 있는 서로 독립적인 지점입니다. 이들을 따로 검사하면 최종 로짓에서 불일치가 발생했을 때 그 위치를 훨씬 쉽게 찾을 수 있습니다.

이제 GPT-2 블록에 필요한 산술 연산을 모두 갖추었습니다. 다음으로는 이 구조화된 텐서 프로그램을 호출 가능한 네이티브 코드로 바꾸는 변환 및 실행 파이프라인을 살펴보겠습니다.

외부 자료

이전: 10부 - 트랜스포머 커널 컴파일하기 · 시리즈 목차 · 다음: 12부 - 융합, 버퍼화, 실행

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역