Triton으로 FlashAttention 1·2 forward 구현하기

17 Oct, 2026

앞 글에서 정리한 FlashAttention 1(FA1)과 FlashAttention 2(FA2)의 forward를 Triton으로 직접 구현해 보았다. 이번 글에서는 구현 과정에서 드러난 두 알고리즘의 구조적 차이와 실제 성능을 비교한다.

병렬화 단위: head 하나 vs Q 블록 하나

FA1에서는 program 하나가 head 하나를 맡도록 해서 grid 크기가 B × H다. 각 program은 K/V 블록을 순회하면서, 블록마다 모든 Q 블록의 출력을 갱신한다. 주소 계산과 causal 경계 처리를 빼고 뼈대만 남기면 다음과 같다.

# grid = (B * H,)
head = tl.program_id(0)

for start_n in tl.range(0, N_CTX, BLOCK_N):
    # K/V tile load
    for start_m in tl.range(0, N_CTX, BLOCK_M):
        # Q와 이전 출력·통계 load
        # attention 계산 후 출력·통계 store
        ...

FA2에서는 program 하나가 Q 블록 하나를 맡는다.

# grid = (triton.cdiv(N_CTX, BLOCK_M), B * H)
query_block = tl.program_id(0)
head = tl.program_id(1)

# 현재 program의 Q tile load
for start_n in tl.range(0, N_CTX, BLOCK_N):
    # K/V tile load
    # online softmax와 출력 누적
    ...

# 마지막에 정규화하고 출력 store

구현하면서 가장 눈에 띈 점은 FA2 알고리즘의 바깥 Q 루프가 그대로 grid로 옮겨간다는 것이었다. Q 블록마다 담당하는 출력 행이 다르므로 서로 독립적으로 처리할 수 있다. 그래서 커널 안에는 K/V를 순회하는 루프 하나만 남는다.

이 차이는 실행되는 program 수에서 바로 드러난다. B=4, H=32, N=2048, BLOCK_M=64일 때를 보자. FA1은 4 × 32 = 128개 program을 실행한다. FA2는 여기에 Q 블록 수 32를 곱한 4,096개를 실행한다. 이 차이가 성능에 어떤 영향을 주는지는 뒤의 batch 실험에서 확인한다.

출력 누적 방식

출력을 유지하는 방식도 다르다. FA1은 K/V 블록 하나를 처리할 때마다 정규화까지 마친 출력과 통계(m, l)를 global memory에 저장한다. 다음 K/V 블록에서는 이 값을 다시 읽어 와 보정한다. 결국 출력이 K/V 블록 수만큼 global memory를 오가게 된다. 중간 출력은 FP32로 유지하고, 마지막에만 FP16으로 변환했다.

FA2는 m_i, l_i, acc를 K/V 루프가 끝날 때까지 on-chip에 유지한다. acc는 정규화하기 전의 출력이다. 새 최댓값이 나올 때마다 스케일을 보정하면서 다음처럼 누적한다.

acc = alpha[:, None] * acc
acc += tl.dot(p.to(v.dtype), v)

모든 K/V 블록을 처리한 뒤에야 acc / l_i[:, None]로 한 번 나눠서 저장한다. 중간에 global memory를 거치지 않고, 정규화도 마지막에 한 번만 한다.

성능 비교

측정 방법

측정 환경은 RTX 5090이고 입력은 FP16이다. 설정 탐색은 tile 크기와 warp/stage 수가 다른 후보 6개 중에서 골랐다.

기준 shape: [4, 32, 2048, 64]

먼저 [B, H, N, d] = [4, 32, 2048, 64]에서 두 구현을 비교했다. 처음에는 둘 다 64×64 tile, 4 warps, 2 stages라는 같은 설정을 썼고, 그다음 구현별로 가장 빠른 설정을 골랐다.

FA1은 128×128 tile, 8 warps, 3 stages에서 가장 빨랐다. FA2는 탐색한 후보 중 기본 설정이 가장 빨랐다.

튜닝 효과는 FA1 쪽이 컸다. non-causal 기준으로 격차가 3.68배에서 2.16배로 줄었다. BLOCK_N이 두 배가 되면 K/V 루프 횟수가 절반으로 줄고, 그만큼 출력과 통계를 global memory에 읽고 쓰는 횟수도 줄기 때문으로 보인다.

튜닝 후 FA2는 non-causal에서 약 2.16배, causal에서 약 1.79배 빨랐다. FA2 논문이 밝힌 약 2배와 비슷한 수준이다. 다만 하드웨어와 비교 대상 구현이 다르기 때문에 참고 정도로 보면 될 것 같다

batch 크기에 따른 변화

H=32, N=2048, d=64는 고정하고 B만 바꿔 보았다. 각 구현을 batch마다 quick 튜닝했는데, 모든 batch에서 같은 설정이 선택되었다.

• FA1: 128×128 tile, 8 warps, 3 stages
• FA2: 64×64 tile, 4 warps, 2 stages

B·H가 작을 때 FA2의 이점이 특히 컸다. FA1의 시간은 B=1에서 1.144 ms, B=2에서 1.155 ms로 거의 같다. 일을 두 배로 늘려도 시간이 그대로라는 것은 GPU에 남는 자원이 있었다는 뜻이다. FA1은 program이 B × H개뿐이라, 32개나 64개로는 RTX 5090의 SM 170개를 채우지 못한다. 반면 FA2는 Q 블록까지 program으로 나누므로 B=1에서도 32 × 32 = 1,024개 program을 실행한다. 앞에서 본 "바깥 Q 루프가 grid로 옮겨간다"는 구조 차이가 가장 직접적으로 드러나는 구간이다.

B·H가 SM 수를 넘어도 차이는 유지됐다. B·H=256, 512에서는 FA1도 SM을 채울 만큼 program이 생긴다. 그런데도 FA2는 여전히 2.19~3.06배 빨랐다. 오히려 B=4에서 8로 갈 때 FA1의 시간은 약 2.5배 늘어, batch 증가 폭보다 더 빨리 늘었다. 같은 구간에서 FA2는 약 1.9배 늘었다. FA1은 K/V 블록마다 출력을 global memory에 왕복시키므로 메모리 트래픽이 유력한 후보다.

같은 설정으로 맞추면

두 구현을 모두 FA1의 최적 설정(128×128 tile, 8 warps, 3 stages)으로 맞춰서도 측정했다.

이 설정에서도 FA2가 모든 batch에서 빨랐다. 하지만 FA2 자체는 64×64 설정보다 약 5~22% 느려졌다. FA2도 타일이 크면 빨라질 줄 알았는데 의외였다. B=4 non-causal 기준으로 FA2의 자원 사용량을 비교하면 다음과 같다.

큰 tile에서는 register가 상한인 255까지 차고 spill이 발생했다. FA2는 K/V 루프 내내 acc, m_i, l_i를 들고 있는 구조라, tile이 커지면 register 부담이 같이 커진다. 이것이 성능 저하에 영향을 줬을 가능성이 있다.

같은 tile 크기가 FA1에는 이득이었고 FA2에는 손해였다. 두 구현의 최적 설정이 다를 수밖에 없는 이유가 여기서 보인다.

마치며

직접 구현하면서 가장 확실하게 이해한 것은 Q 블록을 program 단위로 나눠 맡기는 방식이었다. 논문에서 바깥 루프로 표현된 부분이 Triton에서는 grid가 되고, 각 program은 자기가 맡은 출력의 누적 상태를 끝까지 들고 간다.

이 구조 차이는 측정에서도 그대로 드러났다. batch가 작을 때 FA1은 program 수가 부족해 GPU를 다 쓰지 못했고, FA2는 Q 블록 단위로 일을 나눠 GPU를 채웠다. 각자 튜닝한 기준 shape에서는 FA2가 약 2배 빨랐다. 논문에서 읽기만 했던 속도 향상을 직접 짠 코드로 비슷한 폭으로 확인할 수 있어서 좋았다.

구현

• fa1
• fa2

#FlashAttention #attention #blog #cuda #fa1 #fa2 #gpu #triton