🤖 machine learning

DASH: Deterministic Attention Scheduling for High-throughput Reproducible LLM Training

DASH(Deterministic Attention Scheduling for High-Throughput)는 역전파 과정을 DAG 스케줄링 문제로 정식화하고, 파이프라인 스톨을 줄이고 NVIDIA H800 GPU에서 처리량을 최대 1.28배 향상시키는 Descending Q-Tile Iteration 및 Shift Scheduling과 같은 새로운 전략을 도입함으로써 LLM 학습 시 결정론적 어텐션의 상당한 성능 오버헤드를 해결합니다.

원저자: Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

게시일 2026-06-11

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Xinwei Qiang, Hongmin Chen, Shixuan Sun, Jingwen Leng, Xin Liu, Minyi Guo

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 문제: "재현성"이라는 병목 현상

당신이 수백 명의 요리사(프로세싱 유닛)가 협력하여 거대한 식사(대규모 언어 모델 학습)를 요리하는 거대한 주방(GPU)을 운영하고 있다고 상상해 보세요.

AI의 세계에서 과학자들은 똑같은 레시피로 두 번 요리했을 때 정확히 똑같은 결과물을 얻을 수 있어야 합니다. 이것을 **재현성(reproducibility)**이라고 부릅니다. 만약 레시피를 아주 미세하게 수정했다면, 맛이 정확히 어떻게 변했는지 알아야 하기 때문입니다.

하지만 컴퓨터에는 특이한 점이 있습니다. 숫자를 더할 때 그 순서가 중요하다는 것입니다. 만약 요리사 A가 솥에 소금을 넣고 나서 요리사 B가 후추를 넣는 것과, 요리사 B가 후추를 먼저 넣고 나서 요리사 A가 소금을 넣는 것은 결과가 미세하게 다릅니다. 요리사들이 무작위로 주문을 외치며 돌아가는 혼란스러운 주방에서는, 요리할 때마다 최종적인 맛이 매번 조금씩 달라집니다. 이것이 바로 **비결정론(non-determinism)**입니다.

이를 해결하기 위해 현재의 표준 방식(FlashAttention-3)은 요리사들이 줄을 서서 정해진 순서대로 재료를 넣도록 강제합니다. 요리사 1번이 가고, 그다음 요리사 2번이 가고, 그다음 요리사 3번이 가는 식입니다. 이는 매번 요리할 때마다 정확히 똑같은 맛을 보장합니다.

문제는 이 점입니다: 이 엄격한 줄 세우기는 느립니다. 요리사 1번이 소금을 넣는 동안, 요리사 2번은 아무것도 하지 못한 채 서서 기다려야 합니다. 요리사 3번은 훨씬 더 오래 기다려야 하죠. 주방은 자기 차례를 기다리며 아무것도 하지 않고 서 있는 요리사들로 가득 차 있습니다. 논문은 이 "기다림"이 전체 학습 과정을 거의 **38%**나 느리게 만든다고 말합니다. 이는 엄청난 시간과 비용의 낭비입니다.

해결책: DASH (결정론적 어텐션 스케줄링)

저자들은 DASH라고 불리는 새로운 시스템을 만들었습니다. 단순히 모두를 지루한 줄에 세우는 대신, 레시피에 필요한 엄격한 순서를 따르면서도 요리사들이 계속 일할 수 있도록 주방의 워크플로우를 재설계했습니다.

그들은 이 문제를 교통 퍼즐처럼 다루었습니다. 요리사들을 고속도로에 합류하려는 자동차라고 상상해 보세요. 기존 방식은 차들을 한 대씩 합류하게 만들어 거대한 교통 체증을 일으켰습니다. DASH는 차들이 멈추지 않고 부드럽게 합류할 수 있도록 완벽한 타이밍을 찾아냅니다.

그들은 이 문제를 해결하기 위해 두 가지 주요 기술을 사용했습니다:

기술 1: "역순 줄 세우기" (Descending Q-Tile Iteration)

방에 들어가기 위해 줄을 서 있는 사람들을 상상해 보세요. 보통은 첫 번째 사람을 들여보내고, 그다음 두 번째, 그다음 세 번째 사람을 들여보냅니다. 하지만 이 특정 유형의 요리(Causal Attention이라 불림)에서는, 첫 번째 사람은 뒤에 있는 모든 사람이 작은 작업을 마칠 때까지 기다려야 합니다. 이로 인해 주방에 긴 공백이 생깁니다.

DASH의 해결책: 줄을 순서대로(1, 2, 3...) 부르는 대신, 역순으로(3, 2, 1...) 호출합니다.

작동 원리: 줄의 끝에 있는 사람들(기다려야 하는 일이 가장 적은 사람들)이 즉시 요리를 시작할 수 있습니다. 이들이 작업을 마치면 다음 사람을 위한 공간이 확보됩니다. 이는 트럭의 짐을 뒤에서부터 내리는 것과 같습니다. 뒤쪽부터 비워내면 경로가 더 빨리 확보되어, 맨 앞쪽에서 발생하는 "교통 체증" 없이 전체 줄이 매끄럽게 움직이게 됩니다.

기술 2: "엇갈린 교대" (Shift Scheduling)

다른 유형의 요리(Full Attention이라 불림)의 경우, 문제는 모든 사람이 정확히 같은 시간에 같은 조리대를 사용하고 싶어 한다는 점입니다. 만약 그들이 동시에 같은 솥에 재료를 넣으려고 하면 충돌이 발생합니다.

DASH의 해결책: 그들은 **순환 이동(cyclic shift)**을 사용합니다. 주자들이 모두 동시에 출발하지 않는 계주 경기를 상상해 보세요.

요리사 1은 재료 A로 시작합니다.
요리사 2는 재료 B(요리사 1이 나중에 사용할 재료)로 시작합니다.
요리사 3은 재료 C로 시작합니다.
요리사 1이 A를 다 쓸 때쯤이면, 요리사 2가 준비를 마친 상태가 됩니다.

이것은 완벽한 "엇갈린" 리듬을 만들어냅니다. 모두가 동시에 서로 다른 부분의 작업을 수행하므로 조리대가 비기를 기다릴 필요가 없지만, 최종 조립은 레시피가 요구하는 엄격한 순서에 따라 이루어집니다.

결과: 빠르지만, 마법은 아니다

저자들은 강력한 NVIDIA H800 GPU(AI에 사용되는 슈퍼컴퓨터)에서 이 시스템을 테스트했습니다.

승리: 이들의 새로운 시스템은 기존의 느린 방식보다 "엄격한 순서" 요리를 1.28배 더 빠르게 만들었습니다. "빠르지만 엉망인 방식"과 "느리지만 완벽한 방식" 사이의 간극을 좁혔습니다.
현실적인 점검: 논문은 또한 현실 세계에서는 "완벽함"이 항상 "최선"은 아니라는 점도 발견했습니다.
- 매우 크고 복합적인 작업의 경우, "엇갈린 교대"(기술 2)가 기존 방식보다 오히려 조금 더 느려지는 현상이 나타났습니다.
- 이유는 무엇일까요? 새로운 방식이 너무 복잡해서 요리사들(GPU 코어)이 서로 다른 단계들을 기억하는 데 과부하가 걸렸기 때문입니다. 그들은 "메모지(registers)"가 부족해졌고, 결국 메모리 바닥에 노트를 떨어뜨려야 했으며, 이는 속도를 늦추는 결과를 초래했습니다.
- 교훈: 주방이 얼마나 큰지에 따라, 수학적으로 완벽하지만 복잡한 기술보다는 단순한 기술(예: 역순 줄 세우기)이 더 나을 수도 있습니다.

요약

이 논문은 AI 컴퓨터의 "요리사"들을 조직하는 더 스마트한 방법인 DASH를 소개합니다. DASH는 AI 학습이 완벽하게 재현 가능하도록(비트 단위까지 동일하도록) 보장하면서도, 컴퓨터가 아무것도 하지 않고 기다리게 만들지 않습니다. 작업 순서를 재배치하거나(때로는 줄을 거꾸로 세우고, 때로는 시작 시간을 엇갈리게 함으로써), 그들은 프로세스를 크게 가속화하여 더 신뢰할 수 있는 AI 모델을 더 저렴하고 빠르게 훈련할 수 있게 만들었습니다.

기술 요약: DASH – 고처리량 재현 가능한 LLM 학습을 위한 결정론적 어텐션 스케줄링 (Deterministic Attention Scheduling for High-Throughput Reproducible LLM Training)

1. 문제 정의

대규모 언어 모델(LLM) 학습에서 재현성은 불안정성을 진단하고 아키텍처 변경 사항을 검증하는 데 필수적이다. 그러나 결정론적(determinism)인 어텐션 메커니즘을 구현하는 것은 종종 심각한 성능 저하를 초래한다.

표준 어텐션 구현(예: FlashAttention-3)에서 비결정론성이 발생하는 근본 원인은 부동 소수점 산술의 비결합성(non-associativity)에 있다. 재현성을 보장하기 위해, 역전파(backward pass) 과정에서 그래디언트 텐서(특히 $dQ$)에 대한 고정된 누적 순서를 강제해야 한다. 현재의 결정론적 구현(FlashAttention-3 등)은 동기화 배리어(synchronization barriers)를 사용하여 그래디언트 리덕션(reduction) 작업을 직렬화함으로써 이를 달성한다. 이는 일관성을 보장하지만, 이전 리덕션이 완료될 때까지 스트리밍 멀티프로세서(SM)가 유휴 상태로 대기하게 만드는 심각한 파이프라인 스톨(pipeline stalls, 버블)을 생성한다. 경험적 데이터에 따르면, 타일 실행 스케줄과 엄격한 누적 순서 사이의 불일치로 인해 이로 인한 처리량 감소는 비결정론적 대응 방식 대비 최대 **37.9%**에 달할 수 있다.

2. 방법론

저자들은 DASH(Deterministic Attention Scheduling for High-Throughput)를 제안한다. 이 프레임워크는 결정론적 역전파를 유향 비순환 그래프(DAG) 상의 스케줄링 최적화 문제로 재정의한다.

2.1 문제 정식화

역전파는 다음과 같은 DAG로 모델링된다:

**노드(Nodes)**는 계산 단계(타일 프로세싱)와 리덕션 단계를 나타낸다.
**엣지(Edges)**는 데이터 의존성과 필수적인 누적 순서를 나타낸다.
목표: DAG의 임계 경로(critical path) 길이를 최소화하여 종단 간(end-to-end) 지연 시간을 줄이는 것이다.
제약 조건: 특정 키-값(KV) 타일에 대한 연산은 빠른 레지스터 기반 누적을 활용하기 위해 단일 SM에서 연속적으로 실행되어야 하며, 이는 그래프 내에서 끊기지 않는 체인을 형성해야 한다.

2.2 스케줄링 전략

DASH는 결정론적 순서를 준수하면서 파이프라인 버블을 최소화하기 위해 두 가지 상호 보완적인 전략을 도입한다.

A. 하향식 Q-타일 반복 (Descending Q-Tile Iteration, 휴리스틱)

메커니즘: 쿼리(Q) 블록의 처리 순서를 역전시킨다. Q-타일을 $0 $부터$ N $까지 처리하는 대신, 스케줄러는$ N $부터$ 0$까지 역순으로 처리한다.
근거: 인과적 어텐션(causal attention)에서 의존성은 삼각형 형태를 띤다. Q-타일을 역순으로 처리하면 SM이 의존성을 더 일찍 해결할 수 있어, 후속 어텐션 헤드를 위해 자원을 더 빨리 해제할 수 있다. 이는 헤드 사이의 유휴 간격을 크게 줄이는 긴밀하게 결합된 파이프라인을 생성한다.
적용 가능성: 풀 마스크(full mask)와 인과적 마스크(causal mask) 모두에 효과적이지만, 주로 인과적 스톨(causal stalls)을 완화하기 위해 설계되었다.

B. 시프트 스케줄링 (Shift Scheduling, 이론적 최적해)

메커니즘: 계산 작업을 SM에 순환적으로 할당하는 순환 시프트 할당을 채택한다.
- 풀 마스크 (Full Masks): SM은 KV 블록을 순환 순서로 처리한다 (예: SM $i$ 는 $i, i+1, \dots, N-1, 0, \dots, i-1$ 을 처리함). 이를 통해 특정 $dQ$ 블록에 대해 리덕션 작업이 자연스럽게 순서대로 수행되도록 하여, 어떤 의존성 엣지도 DAG의 깊이 단조성(depth monotonicity)을 위반하지 않음을 보장한다 (Lemma 1).
- 인과적 마스크 (Causal Masks, 대칭 시프트): 인과적 마스크의 불균형한 워크로드(초기 KV 블록이 후기 블록보다 더 많은 상호작 작용을 가짐) 문제를 해결한다. 이는 SM을 쌍으로 묶어 대칭적인 KV 블록( $i$ 와 $N-1-i$ )을 처리하게 함으로써, 삼각형 형태의 워크로드를 균형 잡힌 사각형 형태로 접는다(folding). 이는 밀집된 좌하단 직사각형과 접힌 우상단 삼각형의 두 단계로 실행된다.
이론적 기초: DAG에 0-가중치 의존성 엣지를 추가해도 임계 경로 길이가 보존되기 위한 필요충분조건은 소스 노드의 깊이가 목적지 노드의 깊이보다 작거나 같은 경우( $depth(u) \le depth(v)$ )라는 렘마(lemma)에 기반한다. 시프트 스케줄링은 이 제약 조건을 준수하는 충돌 없는 리덕션 순서를 구축한다.

3. 주요 기여

근본 원인 식별: 본 논문은 결정론적 어텐션의 성능 저하가 단순히 직렬화 자체 때문이 아니라, 타일 스케줄링과 미리 정해진 누적 순서 사이의 불일치 때문임을 밝혀냈다.
DAG 정식화: 결정론적 어텐션 역전파 스케줄링 문제를 DAG 기반으로 최초로 정식화하여, 임계 경로 길이를 원칙적으로 최적화할 수 있게 하였다.
새로운 스케줄링 알고리즘:
- 하향식 Q-타일 반복: 쿼리 블록 순회를 역전시켜 인과적 어텐션 성능을 개선하는 강력한 휴리스틱을 제공한다.
- 시프트 스케줄링: 풀 마스크에 대해 워크로드를 균형 있게 맞추고 파이프라인 버블을 제거하는 이론적으로 최적인 알고리즘이다.
- 대칭 시프트 스케줄링: 대칭적 쌍 구성과 2단계 폴딩을 통해 불균형한 워크로드를 균형 있게 만드는 인과적 마스크를 위한 확장 모델이다.
경험적 검증: NVIDIA H800 GPU에서 이러한 전략들이 결정론적 어텐션의 성능 격차를 유의미하게 좁힐 수 있음을 입증하였다.

4. 실험 결과

실험은 NVIDIA H800 GPU에서 FlashAttention-3를 베이스라인으로 사용하여 수행되었다.

처리량 개선:
- 풀 마스크 (Full Masks): 시프트 스케줄링은 베이스라인을 지속적으로 능가하며, 어텐션 역전파에서 최대 1.28배의 속도 향상을 달성했다.
- 인과적 마스크 (Causal Masks): 하향식 Q-타일 반복과 대칭 시프트 스케줄링 모두 모든 설정에서 성능 향상을 보여주었다.
종단 간 성능 (End-to-End Performance):
- LLaMA3-8B, Qwen2.5-7B, Mistral-8x7B와 같은 모델의 전체 트랜스포머 블록(순전파 및 역전파 포함)에 적용했을 때, DASH는 평균 ~5% (시퀀스 길이와 모델에 따라 2%에서 10% 사이)의 속도 향상을 달성했다.
- 풀 마스크 모델(예: Stable Diffusion 3.5, SAM) 또한 약 4%의 속도 향상을 보였다.
하드웨어 제한 및 트레이드오프:
- 논문은 특정 시나리오에서 이론적 최적성과 실제 성능 사이의 괴리를 언급한다. 헤드 차원(head dimension)이 큰(예: 128) 인과적 마스크의 경우, 대칭 시프트 스케줄링이 단순한 하향식 Q-타일 반복보다 성능이 떨어진다.
- 이유: 대칭 시프트의 복잡한 상태 관리가 레지스터 압박(register pressure)을 증가시킨다. H800 아키텍처에서 이는 레지스터 스필링(register spilling)을 느린 로컬 메모리로 유발하여 알고리즘적 이점을 상쇄한다. 단순한 하향식 전략은 이러한 레지스터 압박을 피할 수 있다.
- 풀 마스크 환경의 극단적인 시퀀스 길이(16k)에서는, 높은 병렬성 영역에서 계산 이득보다 원격 L2 캐시 세그먼트에 대한 SM 간 통신 지연 시간이 더 커지기 때문에 시프트 스케줄링이 가끔 약간 저하되는 현상이 발생한다.

5. 의의 및 주장

DASH는 결정론성을 희생하지 않으면서도 재현 가능한 LLM 학습의 효율성을 크게 발전시킨다고 주장한다. 지능적인 DAG 기반 계획을 통해 누적 순서를 실행 스케줄로부터 분리함으로써, DASH는 결정론성의 성능 페널티가 고정된 비용이 아니라 해결 가능한 스케줄링 문제임을 입증한다.

저자들은 이론적 최적성(시프트 스케줄링)이 달성 가능하지만, 실제 배포 시에는 레지스터 압박이나 메모리 계층 지연 시간과 같은 하드웨어 제약 조건을 고려해야 한다고 강조한다. 결과적으로, DASH는 실무자가 자신의 하드웨어 및 모델 구성에 따라 최적의 접근 방식(휴리스틱 vs 최적화)을 선택할 수 있도록 일련의 전략을 제공한다. 연구를 촉진하기 위해 코드는 오픈 소스로 공개되었다.