Deterministic Attention Scheduling으로 고속 LLM 학습 구현

초록

**
본 논문은 FlashAttention‑3의 결정적 역전파에서 발생하는 37.9% 수준의 처리량 감소 원인을 DAG 기반 스케줄링 문제로 정의하고, 두 가지 최적화 전략인 Descending Q‑Tile Iteration과 Shift Scheduling을 제시한다. 이 기법들은 GPU SM 간의 파이프라인 버블을 최소화해 결정적 역전파의 임계 경로 길이를 단축시키며, NVIDIA H800 GPU 실험에서 최대 1.28배의 속도 향상을 달성한다.

**

상세 분석

**
FlashAttention‑3은 메모리 효율성을 위해 KV 차원을 여러 SM에 분산하고, 각 SM이 로컬 축소(local reduction)를 수행한 뒤 전역 축소(global reduction) 단계에서 dQ를 합산한다. 비결정적 모드에서는 원자적 add(atomicAdd)를 이용해 순서에 무관하게 누적하지만, 이는 부동소수점 비결합성 때문에 비트 단위 재현성을 보장하지 못한다. 결정적 모드에서는 CT‑A( cooperative thread array )별 고정 순서를 강제하기 위해 barrier와 순차적 축소를 도입한다. 이 과정에서 각 SM이 전역 축소를 기다리는 동안 유휴 상태가 발생해 파이프라인 버블이 크게 늘어나며, 특히 causal mask에서는 매 타일마다 의존성이 선형적으로 쌓여 전체 실행 시간이 n·(c+r)+(n‑1)·r 로 비효율적으로 증가한다.

논문은 이러한 현상을 DAG(Directed Acyclic Graph)로 모델링한다. 각 (i, j) 타일 작업은 Compute 단계 C(i,j)와 Reduction 단계 R(i,j)로 구성된 두 노드로 표현되며, 동일 KV 타일 내에서는 연속적인 의존성(edge)으로 연결된다. 전역 축소 순서를 보장하기 위한 zero‑weight 의존성(edge)도 삽입해 SM 간 순서를 강제한다. 최적화 목표는 이 DAG의 크리티컬 패스(critical path)를 최소화하는 것으로, 이는 전체 역전파 지연을 직접 감소시킨다.

첫 번째 전략인 Descending Q‑Tile Iteration은 causal mask 상황에서 Q‑tile을 역순으로 탐색한다. 기존 순방향 탐색에서는 앞쪽 Q‑tile이 먼저 완료돼야 뒤쪽 타일이 전역 축소를 시작할 수 있어 SM이 대기한다. 역순으로 처리하면 짧은 작업이 먼저 끝나고, 이후 타일이 즉시 전역 축소에 진입해 파이프라인 버블을 크게 줄인다. 수식적으로는 전체 실행 시간을 T_rev ≈ m·(n+1)(c+r)/2 + (n‑1)·r 로 근사할 수 있다.

두 번째 전략인 Shift Scheduling은 DAG 모델이 보장하는 최적성을 활용한다. 여기서는 각 SM에 할당된 Q‑tile 순서를 순환(cyclic) 시프트시켜, 모든 SM이 동일한 시간 간격으로 전역 축소를 수행하도록 만든다. 즉, SM0은 Q0→Q1→…→Qn‑1 순으로, SM1은 Q1→Q2→…→Q0 순으로 작업한다. 이렇게 하면 전역 축소 단계가 자연스럽게 충돌 없이 연속적으로 진행되며, 각 SM의 작업 부하가 완벽히 균형을 이룬다. 이론적으로는 크리티컬 패스 길이가 2c+2r 로 최소화되며, 실제 GPU 실행에서도 파이프라인 버블이 거의 사라진다.

실험은 NVIDIA H800 GPU(8 SM 기준)에서 수행되었으며, Full‑mask와 Causal‑mask 모두에 대해 기존 결정적 구현 대비 평균 1.12×~1.28×의 속도 향상을 기록했다. 특히 Causal‑mask에서는 Descending Q‑Tile Iteration이 1.18×, Shift Scheduling이 1.28×의 개선을 보였다. 코드와 스케줄링 로직은 모두 오픈소스로 공개돼 재현성과 확장성을 확보한다.

**