행렬 곱의 연산 방식

행렬곱은 대규모 AI 모델에서 계산 비용의 주요 병목 지점으로, 이를 효율적으로 처리하기 위해 병렬화가 필수적입니다. 이번 포스팅에서는 블록 단위로 행렬곱을 수행하는 두 가지 방식을 살펴봅니다. 일반적은 행렬 곱은 다음과 같습니다.

def dot_spec(x: Float32[Tensor, "4 32 32"], y: Float32[Tensor, "4 32 32"]) -> Float32[Tensor, "4 32 32"]:
    return x @ y

Blocked Matrix Multiplication

블록 행렬 곱은 행렬곱을 계산하는 데 있어 효율성을 높이기 위해 행렬을 작은 블록(또는 타일)으로 나누어 처리하는 기법입니다. 블록 행렬 곱은 블록 단위 계산과 K 차원의 순차적 누적을 핵심으로 하며, 캐시 최적화와 $B$, $ M $, $ N $ 차원의 병렬화에 초점을 맞춥니다.

@triton.jit
def dot_kernel(x_ptr, y_ptr, z_ptr, 
               N0, N1, N2, MID, 
               B0: tl.constexpr, B1: tl.constexpr, B2: tl.constexpr, B_MID: tl.constexpr):
    
    pid_0 = tl.program_id(0)  
    pid_1 = tl.program_id(1)  
    pid_2 = tl.program_id(2)  

    offs_m = pid_0 * B0 + tl.arange(0, B0)  
    offs_n = pid_1 * B1 + tl.arange(0, B1)   
    offs_k = tl.arange(0, B_MID)            

    x_batch_offset = pid_2 * N0 * MID  
    y_batch_offset = pid_2 * MID * N1  
    z_batch_offset = pid_2 * N0 * N1  

    acc = tl.zeros((B0, B1), dtype=tl.float32)

    for k in range(0, MID, B_MID):

        x_ptrs = x_ptr + x_batch_offset + offs_m[:, None] * MID + (k + offs_k)[None, :]
        y_ptrs = y_ptr + y_batch_offset + (k + offs_k)[:, None] * N1 + offs_n[None, :]

        x_mask = (offs_m[:, None] < N0) & (k + offs_k[None, :] < MID)
        y_mask = (k + offs_k[:, None] < MID) & (offs_n[None, :] < N1)

        x = tl.load(x_ptrs, mask=x_mask, other=0.0)
        y = tl.load(y_ptrs, mask=y_mask, other=0.0)

        acc += tl.dot(x, y)

    z_ptrs = z_ptr + z_batch_offset + offs_m[:, None] * N1 + offs_n[None, :]
    z_mask = (offs_m[:, None] < N0) & (offs_n[None, :] < N1)
    tl.store(z_ptrs, acc, mask=z_mask)

Parallel Reduction Matrix Multiplication

병렬 리덕션 기반의 행렬 곱 방식은 $M$, $N$, $K$ 차원을 모두 병렬화하고, atomic_add를 통해 결과를 누적하는 구조입니다. 이 방식은 현재 실험적인 수준에서만 사용되며, 실제 실무에서는 널리 활용되지 않습니다.

1. 배치 처리 시 $K$ 병렬화는 비효율적입니다.
   • 실무에서는 거의 대부분의 딥러닝 연산이 배치 단위($B$)로 처리됩니다.
   • $B$ 차원을 병렬화하는 것 만으로 높은 병렬성을 확보할 수 있고, $K$ 를 병렬화하는 이점이 상대적으로 감소합니다.
   • $K$ 차원을 병렬화하면 누적(reduction)이 필요한 구조로 바뀌므로, 단순 병렬처리 보다 복잡도가 증가합니다.
2. 부동 소수점 연산의 비결합성(Floating Point Non-Associativity) 문제가 발생합니다.
   • $K$ 차원을 나눠 여러 스레드에서 병렬로 곱하고, atomic_add로 누적하는 방식은 덧셈 순서가 달라지므로 부동소수점 오차가 누적됩니다.
   • 결과적으로 deterministic하지 않은 결과가 나와 재현성과 수치안정성을 해칩니다.
3. atomic_add 성능 병목
   • $K$ 차원을 병렬화하려면 여러 쓰레드가 동일한 출력 위치에 값을 쓰기 때문에 atomic_add가 필요합니다.
   • 하지만 atomic 연산은 동기화 비용이 크고, 특히 출력 차원이 작을 경우 심각한 성능 병목으로 이어집니다.

@triton.jit
def dot_kernel(x_ptr, y_ptr, z_ptr,
               N0, N1, N2, 
               B0: tl.constexpr, B1: tl.constexpr, B2: tl.constexpr):

    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    pid_k = tl.program_id(2)  

    offs_m = pid_m * B0 + tl.arange(0, B0)
    offs_n = pid_n * B1 + tl.arange(0, B1)
    offs_k = pid_k * B2 + tl.arange(0, B2)

    x_ptrs = x_ptr + offs_m[:, None] * N2 + offs_k[None, :]
    y_ptrs = y_ptr + offs_k[:, None] * N1 + offs_n[None, :]

    x = tl.load(x_ptrs, mask=(offs_m[:, None] < N0) & (offs_k[None, :] < N2), other=0.0)
    y = tl.load(y_ptrs, mask=(offs_k[:, None] < N2) & (offs_n[None, :] < N1), other=0.0)

    acc = tl.dot(x, y)

    z_ptrs = z_ptr + offs_m[:, None] * N1 + offs_n[None, :]
    z_mask = (offs_m[:, None] < N0) & (offs_n[None, :] < N1)
    tl.atomic_add(z_ptrs, acc, mask=z_mask)