FlashMask-3가 Hopper 분산 어텐션의 마스크 병목을 없앤다

무엇인가

긴 문맥 학습에서는 문서 마스크, prefix-LM 마스크, share-question 마스크, 슬라이딩 윈도, 블록 스파스 같은 구조적 어텐션 마스크를 하나의 패킹된 배치 안에 넣어 학습 데이터의 의미 구조를 표현한다. 문맥이 수십만 토큰으로 커지고 학습이 다중 노드 컨텍스트 병렬(CP)로 확장되면, 마스크 물질화와 마스크 인지 커널 스케줄링, 마스크가 유발하는 부하 불균형이 스텝 시간을 지배하는 1급 병목이 된다. 기존 FlashMask는 컬럼 단위 스파스 표현으로 마스크를 O(n²)에서 O(n)으로 압축했지만 Ampere GPU용으로 구현되어 Hopper의 TMA/WGMMA 워프 특화 파이프라인을 쓰지 못했고, 분산 지원도 없었다. FlexAttention은 임의 마스크를 컴파일해 처리하지만 FA3와 다른 코드 경로를 타고 마스크마다 컴파일 비용을 낸다.

어떻게 동작하나

논문의 핵심 관찰은 컬럼 단위 스파스 마스크가 키 위치로 인덱싱되면서도 쿼리 축 구간 경계를 저장한다는 구조적 비대칭성이다. 이 성질 덕분에 어떤 다중 청크 CP 할당에서도 랭크별 마스크 국소화가 경계 벡터에 대한 원소별 clip–shift–max로 환원되고, 랭크 간 통신이 전혀 필요 없다(정리 1). 수식으로는 각 청크 i에 대해 L_i(x)_k = clip(x_k − O_i, 0, C)로 전역 경계를 청크 로컬 좌표로 옮기고, L_i(x)_k + 1[L_i(x)_k > 0]·iC로 랭크 로컬 연결 공간에 삽입한 뒤, max_i로 집계해 x_k^(r)를 얻는다. 이 변환은 전역 마스크와 랭크가 소유한 오프셋만 있으면 되고, 출력은 단일 GPU 커널 인터페이스를 그대로 유지한다.

무엇과 다른가

Hopper 커널 쪽에서는 FA3의 Producer 워프 그룹이 TMA 발행에 128스레드를 쓰지만 실제로는 1스레드면 충분해 매 반복 96스레드가 놀고 있다는 점을 이용한다. 이 유휴 스레드를 KV Index Generator라는 세 번째 워프 특화 역할로 회수해 LTS/LTE/UTS/UTE의 블록별 max/min으로 각 타일을 완전 마스킹·부분 마스킹·비마스킹으로 분류한 큐를 만들고, KV Loader는 나열된 블록만 TMA로 로드하며, MMA Consumer는 부분 마스킹 블록에만 행 구간 마스크를 적용한다. 불규칙 마스크로 Q 블록별 유효 K 블록 수가 자릿수 단위로 달라지는 문제는 CTA가 전역 카운터에서 다음 Q 블록을 원자적으로 가져오는 Preemptive Persistent Tile(PPT) 스케줄러로 완화한다.

어떻게 쓰나

분산 확장은 세 부품으로 구성된다. 먼저 HCC-IPO 로드 밸런서는 랭크 간 연산 균형(카디널리티 제약 하의 min-max 분할), 노드 간 KV 트래픽 최소화(연산 균형을 ε 허용치 안에서 유지하는 청크 교환), 노드 내 연산 재균형을 교대로 수행한다. 다음으로 마스크 인지 연산-통신 오버랩은 NVSHMEM 대칭 메모리의 단일 SR 버퍼를 재사용하고, 스파시티 검사 커널이 완전히 마스킹된 KV 청크를 표시해 통신 커널이 이를 건너뛰게 한다. 실험에서 전형적인 causal document 및 prefix-LM document 마스크에서 AG/RS 트래픽의 약 50%가 제거됐다. 순방향 all-gather는 congruence group C(r) = {r mod G + k·G}로 라우팅해 1단계에서 노드 간 RoCE get을 N−1회, 2단계에서 노드 내 NVLink로 (G−1)·N개 청크를 가져온다. P=32일 때 노드 간 풀 횟수가 24회에서 3회로 약 90% 줄어든다. 역방향은 KV 컬럼 병렬로 청크별 work_done 플래그를 쓰고, reduce-scatter를 NVSHMEM one-sided put 기반 스파스 all-to-all로 구현하며 bitmask-AMO로 동기화한다.

전제와 한계

단일 H100에서 12개 구조적 마스크, 4K~128K 문맥을 훑은 결과 FlashMask-3는 FlashMask 대비 45~141%, FlexAttention 대비 34~71%, MagiAttention 대비 2~44% 높은 TFLOPs/s를 달성했고 모든 마스크 패턴과 시퀀스 길이에서 세 베이스라인을 앞섰다. 기하평균으로는 FlashMask 대비 82.6%, FlexAttention 대비 55.3%, MagiAttention 대비 15.0% 빠르며, MagiAttention의 블록 스파스 표현이 맞지 않는 Random Eviction 마스크(50배 이상 느려짐)는 제외한 수치다. CP 4~32, 디바이스당 8K 시퀀스 조건에서는 31~75%의 연산-통신 오버랩 효율을 얻었고, 처리량은 Megatron-Core CP(ring) 대비 1.80~3.71배, Megatron-Core CP(all-gather) 대비 1.71~2.92배, MagiAttention 대비 1.09~1.46배였다. Causal Document Mask에서 처리량은 CP=4의 430.5 TFLOPs/s에서 CP=8/16/32에서 각각 471.8, 491.8, 501.6 TFLOPs/s로 증가했다.

엔드투엔드 실험은 128 H100(16노드), CP 16, 디바이스당 8K, Hq=64, Hkv=8, 문서 마스크로 118B-A14B MoE를 128K 문맥에서 학습한 것이다. 베이스라인은 밸런서와 오버랩을 끈 Dual Chunk + all-gather 구성이며, 마스크 스파시티 70%/80%/86% 데이터셋에서 처리량이 7.3%/7.7%/11.0% 올랐고 어텐션 순방향+역방향은 24.5%/42.1%/44.9% 빨라졌다. 70% 스파시티에서 스텝당 10900.5ms가 8757.6ms로 줄었다. 엔드투엔드 이득이 더 작은 이유는 어텐션이 전체 순방향+역방향 연산 시간의 43.5%/25.8%/24.4%만 차지하고, 측정에 그래디언트 동기화와 옵티마이저 갱신이 포함되기 때문이다. 구성 요소별로는 밸런싱이 2.8/4.4/5.7%, 오버랩이 5.4/4.2/5.9%, 청크 스킵 프루닝이 2.0/2.9/1.5%를 기여했다. 계층적 라우팅을 끄면 밸런싱+오버랩이 −1.7%/−1.5%/+1.9%로 떨어져 오버랩 단독보다도 나빠지는데, 밸런서가 청크를 노드 전반에 흩뿌리면서 트래픽이 congruence group에 갇히지 않기 때문이다. 역방향 파이프라인 단계는 기본 4단계이며 2로 줄이면 최대 3.6% 손해, 8로 늘려도 유의미한 추가 이득이 없었다. 어블레이션에서 PPT는 정적 스케줄러 대비 최대 12% 빨랐고, IPO는 세 마스크 패턴 모두에서 가장 낮은 peak-to-average 비율을 기록했다.

실무 관점에서 이 논문이 주는 신호는 마스크 표현을 커널과 분산 스케줄러가 공유하는 단일 인터페이스로 유지했다는 점이다. 하나의 컬럼 단위 스파스 마스크 추상화와 하나의 FA3 파생 커널 바이너리가 단일 GPU와 쿼리 샤딩 CP 실행을 모두 처리하므로, 마스크 종류가 늘어도 커널을 다시 컴파일하거나 랭크별 슬라이스를 물질화할 필요가 없다. 다만 이득은 마스크 스파시티에 크게 의존한다. 스파시티가 낮으면 어텐션 비중 자체가 작아 엔드투엔드 이득이 줄어들므로, 도입 전에 자신의 마스크 분포와 어텐션이 스텝 시간에서 차지하는 비율을 먼저 재야 한다. 또한 저자들은 FlashMask-3가 정확한 어텐션을 계산한다고 명시한다. 마스크 국소화가 전역 마스크와 일치하고(정리 1), 밸런서·오버랩·프루닝은 무엇을 계산하는지가 아니라 어디에 배치하는지만 바꾸므로 수렴과 품질에는 영향이 없다는 것이 논문의 주장이다.

한계와 전제도 분명하다. 구현은 Hopper(H100)와 NVSHMEM, 노드 내 NVLink·노드 간 RoCE 토폴로지, 노드당 8 GPU(G=8)라는 하드웨어 전제에 묶여 있고, 상류 스케줄러가 정한 청크 할당을 입력으로 받는다. 연산과 통신의 정확한 결합 최적화는 조합 폭발과 스파시티로 인한 불규칙성 때문에 다루기 어렵다고 인정하고, HCC-IPO라는 반복적 휴리스틱으로 근사한다. 성능 수치는 특정 마스크 집합과 시퀀스 길이, CP 차수에서 측정된 것이며, 논문은 별도의 한계 절을 두지 않고 이 전제들을 각 절에 흩어 서술한다.