FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention

긴 문서를 여러 GPU에 나누어 학습할 때는 계산량만 문제가 되는 게 아니다. 각 토큰이 어떤 토큰을 참고할 수 있는지 정하는 어텐션 마스크도 GPU마다 맞게 다뤄야 한다. 마스크 때문에 어떤 GPU는 할 일이 많고 다른 GPU는 먼저 끝나면, 전체 학습은 가장 늦은 GPU를 기다리게 된다.

FlashMask-3는 마스크를 저장하는 방식에서 출발해, 계산할 작업을 고르고 GPU에 나누고 통신과 겹쳐 실행하는 과정까지 함께 설계한다. 논문이 보고한 실제 대규모 학습의 처리량 향상은 **7.3~11.0%**다. 이 수치는 128K 컨텍스트의 MoE 모델을 H100 128개로 학습한 결과이며, 어텐션 커널만의 속도 향상과는 구분해야 한다(§4.3, Table 1).

논문 · 코드

긴 문서를 나누면, 마스크도 문제가 된다

어텐션은 각 토큰이 다른 토큰의 정보를 얼마나 참고할지 계산한다. 마스크는 그중 허용하지 않는 관계를 막는 규칙이다. 예를 들어 여러 문서를 한 입력으로 묶어 학습한다면, 서로 다른 문서의 토큰끼리는 보지 못하게 할 수 있다. 가까운 토큰만 보게 하는 sliding-window도 이런 규칙의 하나다.

문제는 마스크를 거대한 표로 만들 때 생긴다. 토큰이 $n$개라면 토큰 쌍은 $n^2$개이므로, 입력이 길수록 마스크를 저장하고 처리하는 비용이 커진다. 기존 FlashMask는 규칙적인 마스크를 몇 개의 경계값으로 표현해 저장량을 $O(n)$으로 줄였다(§2.2).

여러 GPU를 쓰면 문제가 하나 더 생긴다. **컨텍스트 병렬화(Context Parallelism, CP)**는 긴 시퀀스의 어텐션 작업을 GPU들에 나누는 방식이다. 각 GPU는 자신이 맡은 토큰에 맞는 마스크가 필요하고, 토큰 수가 같아도 마스크가 허용하는 계산량은 다를 수 있다. 단순히 같은 크기로 나누는 것만으로는 부하가 고르게 분산되지 않는다(§2.3).

FlashMask-3가 풀려는 것은 이 두 문제의 연결이다. 압축한 마스크를 실제 GPU 커널에서 빠르게 쓰면서, 여러 GPU에 작업을 나눠도 마스크 처리와 통신이 발목을 잡지 않게 하려는 것이다.

핵심은 큰 표 대신 경계를 저장하는 것

어텐션 행렬에서 행은 정보를 요청하는 토큰인 쿼리(query), 열은 참고 대상인 **키(key)**에 대응한다. FlashMask 계열은 각 키에 대해 “어느 쿼리 구간을 가릴 것인가”를 경계값으로 저장한다. 표현할 수 있는 마스크에서는 키 하나당 최대 두 개의 가려진 구간, 즉 네 개의 경계값이면 된다(§2.2).

이 방식이 GPU별 마스크를 만드는 데 왜 유리한지 작은 예시로 보자. 아래는 논문의 측정값이 아닌 설명용 예시다.

토큰 8개를 GPU 두 개에 나누고, 어떤 키를 참고하지 못하는 쿼리 구간이 전역 좌표에서 $[6,8)$이라고 하자. GPU A는 쿼리 0~3, GPU B는 쿼리 4~7을 맡는다.

GPU B는 자신의 시작 위치인 4를 빼기만 하면 된다. GPU A는 자기 범위 밖의 경계를 끝으로 잘라내면 빈 구간을 얻는다. 둘 다 전역 경계와 자신이 맡은 위치를 알고 있다면, 이 변환을 위해 서로 정보를 주고받을 필요가 없다.

논문은 한 GPU가 떨어져 있는 여러 청크를 맡는 경우에도 이 아이디어를 확장한다. 경계를 담당 구간에 맞게 자르고(clip), 로컬 좌표로 옮긴 뒤(shift), 합치는(max) 변환이다. 이 결과가 원래 마스크와 같은 가림 관계를 보존한다는 것이 Theorem 1의 내용이다(§3.2, Appendix C.1).

여기서 통신이 없다는 말은 마스크의 로컬 좌표 변환에 한정된다. 어텐션 계산에 필요한 키·값 데이터를 GPU 사이에서 전송하는 일까지 없어지는 것은 아니다.

경계 정보를 계산과 작업 배정에도 쓴다

마스크를 작게 저장하는 것만으로 GPU가 빨라지지는 않는다. FlashMask-3는 같은 경계 정보를 이용해 실제로 할 일을 줄이고, 남은 일을 고르게 배치한다.

계산하지 않을 블록을 먼저 거른다

GPU 커널은 어텐션 행렬을 작은 블록으로 나눠 처리한다. 어떤 블록이 전부 가려져 있다면 데이터를 읽거나 행렬곱을 수행할 필요가 없다. 일부만 가려졌다면 계산하면서 마스크를 적용해야 한다.

FlashMask-3의 KV Index Generator는 이 구분을 먼저 수행해 계산할 블록 목록을 만든다. 데이터 로더는 목록에 있는 블록을 가져오고, 행렬곱 담당 스레드는 계산에 집중한다. Hopper GPU의 비동기 데이터 전송(TMA)과 행렬곱 명령(WGMMA)을 활용하면서, 마스크 검사 때문에 주 계산 흐름이 막히는 것을 피하는 설계다(§3.1).

이 역할 분리가 중요하다는 근거는 단순 포팅 실험에서 드러난다. 128K Document 마스크의 forward 처리량은 KV Index Generator가 없는 FA3 스타일 Hopper 포트에서 19.0 TFLOPs/s, 전체 FlashMask-3에서 577.6 TFLOPs/s였다. 다만 전체 구현에는 다른 최적화도 포함되므로, 차이 전부를 이 구성요소 하나의 효과로 볼 수는 없다(Appendix D.4, Table 7).

토큰 수보다 실제 계산량을 기준으로 나눈다

같은 길이의 청크라도 가려진 블록이 많으면 계산이 적다. FlashMask-3의 HCC-IPO는 이런 차이를 반영해 작업을 배정하는 알고리즘이다. 먼저 GPU별 계산량을 맞추고, 노드 간 데이터 전송을 줄이는 방향으로 배정을 조정한 다음, 노드 안에서 다시 균형을 맞춘다(§3.3).

한 GPU 안에서도 같은 문제가 생긴다. 일부 작업만 오래 걸리면 다른 연산 장치가 먼저 쉬게 된다. 이를 줄이기 위해 PPT 스케줄러는 작업을 끝낸 GPU 스레드 블록이 다음 작업을 가져가게 한다. GPU 사이의 배정과 GPU 내부의 배정에서 각각 기다리는 시간을 줄이는 셈이다(§3.1).

필요한 통신을 계산 뒤에 겹친다

컨텍스트를 나눠도 각 GPU는 다른 GPU가 가진 키·값 데이터가 필요하다. FlashMask-3는 마스크상 전혀 쓰지 않을 청크의 전송을 건너뛰고, 노드 간 통신과 노드 안 통신의 경로를 나눈다. 역전파에서는 작업을 여러 구간으로 나눠 한 구간의 통신이 다음 구간의 계산과 겹치도록 한다(§3.4).

전체 파이프라인

“마스크 때문에 느리다”는 말에는 세 가지 일이 섞여 있다. 가릴 칸을 표시한 큰 표를 만들고, 계산할 블록을 골라내고, GPU마다 할 일을 나누는 일이다. 이 준비와 배정이 오래 걸리거나 한 GPU에 일이 몰리면, 다음 학습 단계로 넘어가는 시간도 늦어진다.

아래 그림은 FlashMask-3가 다루는 작업을 작은 입력으로 풀어낸 설명용 예시다. 세 패널은 각각 독립된 예시이며 논문 성능 측정값이 아니다.

첫 패널에서 초록 칸은 계산할 토큰 쌍이고 회색 칸은 건너뛸 토큰 쌍이다. 키 6을 기준으로 보면 0~3행을 가리므로, 그 범위를 [0, 4)라는 경계로 적을 수 있다. 실제 커널은 이런 경계를 이용해 전부 가려진 블록을 읽거나 계산하지 않는다(§2.2, §3.1).

둘째 패널은 그 경계를 GPU 안의 좌표로 바꾸는 과정이다. 전역의 [6, 8)은 GPU B 안에서는 [2, 4)가 된다. 이 변환에는 GPU끼리 정보를 주고받을 필요가 없지만, 어텐션에 필요한 키·값 데이터의 전송은 여전히 남는다(§3.2).

셋째 패널에서는 작업 네 개를 같은 개수로 나눠도 계산량이 14와 6으로 벌어진다. 작업을 다시 조합하면 10씩 나눌 수 있다. 실제 HCC-IPO는 이런 계산량뿐 아니라 노드 간 전송 비용도 고려한다. 그림의 14→10은 원리를 보여주는 임의의 작업 비용이며 실측 향상률이 아니다(§3.3).

마지막으로 필요한 데이터의 전송을 계산과 겹쳐 실행한다. 예를 들어 다음 구간의 데이터를 가져오는 동안 준비된 구간을 계산해 기다림을 줄인다. 위 그림은 저장·좌표 변환·작업 배정을 보여주며, 이 통신 중첩의 세부 구현까지 나타낸 것은 아니다(§3.4).

실제 학습은 얼마나 빨라졌나

이 논문에서 실무적으로 먼저 볼 결과는 전체 학습 처리량이다. 어텐션만 빨라져도 나머지 학습 단계가 그대로라면 전체 이득은 작아질 수 있기 때문이다.

실험은 총 118B·활성 14B 파라미터의 MoE 모델을 128K 컨텍스트로 학습했다. 하드웨어는 H100 128개이며 CP 크기는 16, GPU당 로컬 시퀀스는 8K다. 아래 기준선도 FlashMask-3 커널을 사용하지만, 작업 배정은 Dual Chunk, 통신은 all-gather를 쓰고 제안한 부하 균형·통신 중첩을 끈 설정이다. 따라서 이 표는 커널 자체의 교체 효과보다 분산 실행 최적화의 추가 이득을 보여준다(§4.3, Table 1).

표는 원문 평균값을 옮겼다. 표준편차와 구성요소별 실험은 아래 접힌 원문 표에서 확인할 수 있다.

해석에서 중요한 것은 마지막 두 열의 차이다. 어텐션 계산은 최대 44.9% 개선됐지만 전체 학습은 최대 11.0% 개선됐다. 그래디언트 동기화나 옵티마이저 업데이트처럼 이 방법이 바꾸지 않는 작업도 학습 시간에 포함되기 때문이다. “훈련이 44.9% 빨라졌다”라고 읽으면 범위를 잘못 이해하게 된다.

별도의 분산 어텐션 실험에서는 MagiAttention 대비 총 forward+backward 처리량이 1.09~1.46배였다. 이 실험은 CP 크기 4~32에서 GPU당 시퀀스 길이를 8K로 고정한 조건이므로, 위의 전체 학습 수치와 같은 지표로 비교하면 안 된다(§4.2, Table 5).

그렇다면 단일 GPU에서는 FlashAttention-3보다 빠를까? 원문 Table 6에서 비교한 다섯 마스크에서는 그렇지 않다. FlashMask-3의 처리량은 FA3의 0.71~1.00배였다. 이 설계의 가치는 모든 경우에서 FA3를 이기는 데 있기보다, 더 다양한 구조의 마스크와 분산 실행을 함께 지원하는 데 있다(Appendix D.4).

적용 전에 확인할 한계

가장 먼저 확인할 것은 마스크를 이 표현으로 담을 수 있는가다. 키 하나에 대해 가려지는 쿼리들이 두 구간으로 표현되지 않는다면 현재 형식을 그대로 쓸 수 없다. 또한 커널은 Hopper의 TMA·WGMMA에 맞춰 설계됐으므로 다른 GPU에서 같은 이득을 기대하려면 별도 구현과 검증이 필요하다. 둘 다 저자가 명시한 제약이다(Appendix A).

분산 통신 최적화도 하드웨어 구성을 탄다. 노드 안의 빠른 NVLink와 상대적으로 느린 노드 간 연결을 활용하는 설계이며, 계층적 라우팅은 현재 forward all-gather에 적용된다. 다른 토폴로지나 backward reduce-scatter까지 같은 결과가 나는지는 추가 검증이 필요하다(Appendix A, C.4).

리뷰어 관점에서는 작업 배정 비용도 살펴볼 필요가 있다. 논문은 스케줄링을 사전 계산하고 미리 불러와 실행 중 비용을 숨긴다. 하지만 마스크가 자주 달라져 이 준비를 재사용하기 어렵다면 이득이 줄 수 있다. 실제 Figure 7(b)에서 IPO와 더 단순한 LPT의 부하 균형 차이는 개선된 경우에도 약 0.01 수준이므로, 배정 알고리즘의 복잡도가 내 워크로드에서도 필요한지는 따로 판단해야 한다(§4.4, Appendix D.5).

품질 검증의 범위도 구분해야 한다. 마스크 변환은 원래 가림 관계를 보존하며, 논문은 1,581스텝의 학습 손실을 비교했다. 그러나 별도 다운스트림 벤치마크 결과를 제시한 것은 아니다(Theorem 1, Appendix D.3).

FlashMask-3는 긴 문서, 구조적 마스크, Hopper 기반의 다중 GPU 학습이 함께 필요한 환경에서 검토할 만하다. 마스크의 경계를 단순한 저장 형식으로 두지 않고, 실제 작업량과 통신을 결정하는 정보로 활용했다는 점이 핵심이다. 도입 여부를 판단할 때는 마스크 표현 가능성, GPU와 네트워크 구성, 전체 학습 처리량의 순서로 확인하는 편이 좋다.