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

Long-context training of large language and multimodal models couples structured attention masks (document, prefix-LM, share-question, sliding-window, block-sparse) with multi-node context parallelism, making the mask a dominant bottleneck. FlashMask's column-wise sparse mask compresses them from O(n^2) to O(n) via four query-axis boundary vectors, covering broad structured patterns. Implemented for Ampere GPUs, it neither exploits FlashAttention-3's TMA/WGMMA warp-specialized pipeline on Hopper nor supports distribution. The column-wise sparse mask is indexed by key positions but stores query-axis interval boundaries. Under any multi-chunk context-parallel assignment, it reduces per-rank mask localization to a provably zero-communication clip--shift--max on the boundary vectors, compatible with the single-GPU kernel. FlashMask-3 introduces a mask-aware Hopper kernel with three-role warp specialization atop FlashAttention-3's TMA/WGMMA pipeline, plus a mask-aware distributed extension: zero-communication mask localization, a sparsity-aware load balancer (HCC-IPO), and a mask-aware compute-communication overlap with mask-driven communication pruning, whose forward all-gather uses topology-aware hierarchical routing. On a single H100 across 12 masks at 4K--128K context, it improves achieved TFLOPs over FlashMask, FlexAttention, and MagiAttention by 45--141%, 34--71%, and 2--44%. At context-parallel degrees 4--32 with 8K sequence per device, it attains 31--75% compute-communication overlap efficiency and improves throughput over Megatron-Core CP ring, Megatron-Core CP all-gather, and MagiAttention by 1.80times--3.71times, 1.71times--2.92times, and 1.09times--1.46times. In end-to-end training of a 118B-A14B MoE at 128K on 128 H100 GPUs, it raises throughput by 7.3--11.0% over a Dual Chunk all-gather baseline, with attention forward+backward 24.5--44.9% faster.

한 줄 요약

FlashMask-3는 Hopper GPU 기반의 분산 환경에서 효율적인 마스크 인식 어텐션 연산을 구현한 새로운 기술로, TFLOPs 성능과 통신 효율성을 크게 향상시킨다.

핵심 기여도

**Mask-aware Hopper kernel**: 세 가지 역할의 워프 전용화(KV Index Generator, KV Loader, MMA Consumer)를 통해 TMA/WGMMA 파이프라인에 마스크 인식 연산을 통합, 거의 무시할 수준의 오버헤드로 처리.
**Zero-communication mask localization**: `clip–shift–max` 연산을 통해 분산 환경에서 마스크 로컬라이제이션을 통신 없이 처리.
**Sparsity-aware load balancer (HCC-IPO)**: 희소성 기반의 계층적 로드 밸런싱을 통해 불균형 마스크를 처리.
**Mask-aware compute-communication overlap**: 마스크 기반 통신 최적화 및 계산-통신 겹침을 통해 31–75%의 겹침 효율성 달성.

핵심 아이디어

기존 FlashMask는 Ampere GPU에서 `O(n)` 크기의 4개 query-axis 경계 벡터를 사용해 `O(n²)` 마스크를 압축했으나, Hopper GPU의 TMA/WGMMA 파이프라인과 분산 환경을 지원하지 않았다. FlashMask-3은 이러한 제한을 극복하기 위해, query-axis 구조를 활용한 `column-wise sparse mask`가 key position 기반으로 저장되므로, `clip–shift–max` 연산을 통해 분산 환경에서도 마스크를 로컬로 처리할 수 있음을 증명(정리 1). 이를 기반으로 Hopper GPU의 TMA/WGMMA 파이프라인에 마스크 인식 연산을 통합하고, 분산 환경에서는 계산-통신 겹침과 통신 최적화를 통해 성능을 향상시킨다.

기술적 접근법

**Mask-aware Hopper kernel**: FlashAttention-3의 TMA/WGMMA 파이프라인에 세 가지 역할의 워프 전용화를 도입.
**Zero-communication mask localization**: `clip–shift–max` 연산을 통해 마스크 로컬라이제이션을 통신 없이 처리.
**Sparsity-aware load balancing**: HCC-IPO 알고리즘을 사용해 희소 마스크를 기반으로 로드 밸런싱.
**Compute-communication overlap**: 마스크 기반 통신 최적화 및 계산-통신 겹침을 통해 효율성 향상.
**Forward all-gather**: 토폴로지 인식 계층 라우팅을 사용.

주요 결과

**단일 H100 GPU**: 4K–128K 컨텍스트 길이에서 12개 마스크에 대해 FlashMask 대비 45–141%, FlexAttention 대비 34–71%, MagiAttention 대비 2–44% TFLOPs 향상.
**분산 환경 (CP degree 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배 성능 향상.
**118B-A14B MoE 모델 (128K 컨텍스트, 128 H100 GPU)**: Dual Chunk all-gather 기반 대비 7.3–11.0%의 투수율 향상, 어텐션 forward+backward 24.5–44.9% 가속.

의의 및 한계

FlashMask-3는 Hopper GPU의 고성능 파이프라인과 분산 환경을 모두 활용하는 첫 마스크 인식 어텐션 기술로, 기존 Ampere 기반 FlashMask와 비교해 TFLOPs, 계산-통신 겹침, 분산 투수율에서 모두 우수한 성능을 보인다. 특히, `column-wise sparse mask`의 구조적 특성을 기반으로 통신 없는 마스크 로컬라이제이션을 가능하게 한 점이 학술적·실용적 가치를 높인다. 그러나 본 연구는 H100 GPU에만 적용되었으며, 다른 하드웨어 아키텍처에서의 확장성은 명시되지 않음.

실용적 활용

FlashMask-3는 대규모 언어 및 멀티모달 모델의 긴 컨텍스트 학습에 적합하며, 특히 Hopper GPU 기반의 클러스터 환경에서 높은 성능 향상을 기대할 수 있다. 연구실 및 산업 내 대규모 모델 학습, 멀티노드 분산 학습 플랫폼에 즉시 적용 가능.