트랜스포머 커널 컴파일하기

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

이제 우리의 프런트엔드는 Python 텐서 표현식을 MLIR 연산으로 변환할 수 있습니다. 이를 사용해 NumPy GPT-2 모델의 GELU, 소프트맥스, 레이어 정규화를 구현하겠습니다. 이 순서는 단계적으로 나아가기에 유용합니다. GELU는 원소별 연산이고, 소프트맥스는 리덕션을 도입하며, 정규화는 리덕션과 브로드캐스팅을 결합합니다.

라이브러리 함수와 컴파일러 프리미티브를 구분하면 구현을 작게 유지할 수 있습니다. GELU는 라이브러리 함수입니다. 덧셈, 곱셈, 거듭제곱, tanh는 컴파일러 프리미티브입니다. 컴파일러는 GELU의 공식을 컴파일하기 위해 gelu라는 이름을 인식할 필요가 없습니다.

아래 예제는 Part 9의 환경을 사용하여 tiny-mlir 동반 프로젝트에서 실행됩니다.

GELU

GPT-2는 tanh 근사를 사용합니다.

$$\mathrm{GELU}(x)=\frac{x}{2}\left(1+\tanh \left[\sqrt{\frac{2}{\pi }}(x+0.044715x^3)\right]\right)$$

우리의 구현은 레퍼런스를 충실히 따릅니다.

import numpy as np
from tinymlir import tanh


def gelu(x):
    return 0.5 * x * (1 + tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))

np.sqrt(2 / np.pi)는 인자가 상수이므로 트레이싱 중에 평가됩니다. x에 의존하는 모든 것은 텐서 그래프의 일부가 됩니다. MLIR 빌더에서 산술 연산은 linalg.generic의 스칼라 영역으로 들어가고, tanh는 math.tanh가 되며, 거듭제곱은 math.powf가 됩니다.

from tinymlir import jit

compiled_gelu = jit(gelu)
x = np.linspace(-5, 5, 31, dtype=np.float32)
y = compiled_gelu(x)
expected = 0.5 * x * (
    1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3))
)
np.testing.assert_allclose(y, expected, rtol=2e-6, atol=2e-6)

여기에는 미묘한 세부 사항이 있습니다. np.sqrt가 생성하는 NumPy 스칼라의 dtype은 float64입니다. 레퍼런스 환경에서는 이 값을 float32 표현식에 곱하면 해당 부분의 계산이 승격됩니다. 그 결과 GELU 출력은 float64가 됩니다. 우리의 프런트엔드는 명시적인 캐스트 노드를 통해 이 승격을 유지합니다.

모든 연산을 float32로 하는 모델을 설계하는 것도 합리적이지만, 그것은 레퍼런스 계산을 변경하는 것이 됩니다. 이 시리즈에서 우리가 재현하려는 것은 레퍼런스의 동작입니다. 참조 구현의 공식을 맞춘다는 것에는 그 dtype과 평가 순서까지 맞추는 것이 포함됩니다.

리덕션

원소별 연산은 각 출력을 독립적으로 계산할 수 있습니다. 합계는 여러 입력 원소를 결합해야 합니다.

$$y_i=\sum _{j=0}^{N-1}{x}_{ij}$$

프런트엔드는 이를 축을 가진 리덕션 노드로 표현합니다. 이 노드의 타입 규칙은 리덕션된 차원을 제거합니다. (3, 7) 텐서를 축 1에 대해 리덕션하면 (3,) 형태가 생성됩니다. keepdims=True를 사용하면 그 결과를 (3, 1)로 reshape하여 이후의 브로드캐스팅을 간단하게 만듭니다.

우리의 빌더는 linalg.reduce를 사용합니다. 다음은 핵심 구성으로, operand, output_type, axis와 원소 타입은 표현식이 제공합니다.

zero = arith.ConstantOp(element_type, 0.0).result
initial = tensor.EmptyOp(output_shape, element_type).result
filled = linalg.fill(zero, outs=[initial])
reduction = linalg.ReduceOp(
    [output_type], [operand], [filled], [axis]
)
block = reduction.regions[0].blocks.append(element_type, element_type)
with ir.InsertionPoint(block):
    total = arith.AddFOp(*block.arguments)
    linalg.YieldOp([total.result])

초기 텐서는 0으로 채워져 있어야 합니다. tensor.empty는 형태를 가진 목적지를 제공할 뿐, 누산기를 초기화하지는 않습니다. 이 차이는 리덕션과 행렬 곱셈 모두에서 중요합니다.

본문은 입력 원소 하나와 현재까지 누적된 값을 받습니다. 본문은 이 둘을 더하고 새 누산기를 내놓습니다. 최댓값 리덕션도 같은 구조를 사용하되, 음의 무한대로 초기화하고 영역 안에서 arith.maximumf를 사용합니다.

이 연산은 여전히 리덕션을 기술할 뿐이며, 각 덧셈을 어느 프로세서나 스레드가 계산할지는 지정하지 않습니다. 따라서 스케줄링 결정은 이후의 변환에 맡겨집니다. 우리의 초기 네이티브 백엔드는 결국 이 연산을 루프로 로어링합니다.

안정적인 소프트맥스

소프트맥스는 점수로 이루어진 행 하나를 확률로 바꿉니다.

$$p_j=\frac{e^{x_j}}{{\sum }_k{e}^{x_k}}$$

큰 점수에 지수 함수를 직접 적용하면 오버플로가 발생할 수 있습니다. 행의 최댓값을 빼면 실수 산술에서는 결과가 바뀌지 않으면서도 가장 큰 지수 값이 1로 유지됩니다.

$$m=\underset{k}{\max }{x}_k,p_j=\frac{e^{x_j-m}}{{\sum }_k{e}^{x_k-m}}$$

Python 정의는 지금까지 갖춘 연산들의 조합입니다.

from tinymlir import exp


def softmax(x):
    numerator = exp(x - x.max(axis=-1, keepdims=True))
    return numerator / numerator.sum(axis=-1, keepdims=True)

형태가 (H, T, T)인 입력에 대해 최댓값과 합의 형태는 (H, T, 1)입니다. 이들을 소비하는 원소별 연산은 마지막의 크기 1인 차원을 브로드캐스트합니다. 같은 함수는 일반적인 2차원 행렬에도 동작하며, 프런트엔드가 각 특수화에 맞는 타입과 맵을 도출합니다.

컴파일러는 공유된 numerator 식을 그대로 유지합니다. Python에서 사용될 때마다 지수 계산을 따로 다시 만들지 않습니다. 이후의 변환이 중간값을 다시 계산할지 실체화할지는 별개의 최적화 문제입니다.

지수 공식을 직접 적용하면 실패할 경우들을 테스트해 볼 수 있습니다.

x = np.array([
    [10000, 9999, -10000],
    [-10000, -9999, -10001],
    [0, -np.inf, -np.inf],
], dtype=np.float32)
probabilities = jit(softmax)(x)
np.testing.assert_allclose(probabilities.sum(axis=-1), 1, atol=1e-6)
assert probabilities[2, 0] == 1

음의 무한대로만 이루어진 행은 이 예제가 유용하게 다루는 범위를 벗어납니다. 그 최댓값을 빼면 정의되지 않은 값이 나오기 때문입니다. GPT-2의 인과 마스크는 항상 대각선 성분을 남겨 두므로, 비어 있지 않은 쿼리 행에는 허용된 키가 적어도 하나 있습니다. 이 성질은 어텐션을 구현할 때 명시적으로 드러내겠습니다.

레이어 정규화

레이어 정규화는 각 토큰 표현의 마지막 차원에 걸쳐 통계량을 계산합니다.

$${\mu }_i=\frac{1}{C}\sum _j{x}_{ij},{\sigma }_i^2=\frac{1}{C}\sum _j(x_{ij}-{\mu }_i{)}^2$$

그런 다음 학습된 스케일 벡터와 바이어스 벡터를 적용합니다.

$$y_{ij}=g_j\frac{x_{ij}-{\mu }_i}{\sqrt{{\sigma }_i^2+ϵ}}+b_j$$

모분산과 GPT-2의 $ϵ=10^{-5}$를 사용합니다.

from tinymlir import sqrt


def layer_norm(x, gain, bias):
    centered = x - x.mean(axis=-1, keepdims=True)
    variance = (centered * centered).mean(axis=-1, keepdims=True)
    return gain * centered / sqrt(variance + 1e-5) + bias

mean 자체는 합을 정적으로 알려진 리덕션 길이로 나눈 것입니다. 백엔드에 별도의 정규화 루프는 없습니다. 이 함수는 앞선 예제들과 같은 로어링 규칙을 사용하여 리덕션, 원소별 연산, 브로드캐스트 맵으로 전개됩니다.

중심화된 차이를 사용하는 것은 수치적으로도 중요합니다. 대수적으로 동등한 식인 $E[x^2]-E[x{]}^2$는 큰 공통 오프셋을 가진 값에서 심각한 소거 오차를 겪을 수 있습니다. 저희는 더 편리해 보이는 항등식을 택하지 않고 레퍼런스의 중심화된 계산 방식을 따릅니다.

상수 행은 유용한 테스트입니다. 상수 행의 분산은 0이므로, 정규화 후의 출력은 학습된 편향과 같아야 합니다. 테스트 스위트는 특수한 경우가 일반 경로의 오류를 가리지 않도록 상수 행과 무작위 행을 함께 사용합니다.

옵티마이저가 보는 것

이 시점에서 컴파일러는 이름이 붙은 커널들을 연달아 호출하는 것 이상을 봅니다. Python 식 전체가 함수 내부의 텐서 그래프가 됩니다. 뺄셈 뒤에 지수 함수와 나눗셈이 이어지는 계산은 region과 인덱싱 맵을 합성할 수 있는 구조화된 연산으로 나타납니다.

융합 전후의 차이를 살펴볼 수 있습니다:

uv run python emit.py softmax
uv run python emit.py softmax --stage optimized

리덕션은 융합에 제약을 줍니다. 행의 합을 계산하려면 그 행에 기여하는 값들이 필요합니다. 패스는 생산자-소비자 관계와 리덕션 의미론을 보존해야 하며, 모든 연산을 임의의 원소별 region 안으로 옮길 수는 없습니다. 이러한 의미론을 명시적으로 유지하는 것이 애초에 구조화된 IR을 사용하는 이유입니다.

수치 검증

소규모 연산 테스트는 엄격한 허용 오차로 NumPy와 결과를 비교하고 출력 dtype을 확인합니다. 이 테스트들은 여러 리덕션 축, 유지된 차원, 브로드캐스트된 벡터, 큰 소프트맥스 오프셋, 상수 정규화 행, 그리고 GELU의 혼합 정밀도를 다룹니다.

uv run pytest

이 테스트들은 기본 구성 요소를 확립합니다. 최종 모델 테스트에서는 중간 잔차 출력과 로짓도 원본 tiny-gpt2 저장소와 비교하고, 이어서 생성된 토큰 ID가 정확히 일치하는지 확인합니다. 구현에 국소적인 반올림 차이가 있더라도 수치적으로는 여전히 올바를 수 있습니다. 필수 예제에서 다른 토큰을 선택하는 것은 이와 별개로 관찰 가능한 실패입니다.

다음으로 행렬 곱셈을 추가하고, 이와 같은 합성 가능한 텐서 연산을 사용해 어텐션 계산을 표현하겠습니다.

외부 자료

• GPT-2 모델 구현
• MLIR Linalg 다이얼렉트
• MLIR Math 다이얼렉트
• NumPy GPT-2 레퍼런스

이전: 9부 - MLIR을 위한 Python 프런트엔드 · 시리즈 목차 · 다음: 11부 - 행렬 곱셈과 어텐션