🤖 machine learning

Hardware-Aware FP4 FlashAttention-4

이 논문은 비인과적 추론을 위한 Direct-P와 FP8 그래디언트를 사용하는 인과적 경로를 채택함으로써 Blackwell의 FP4 텐서 코어의 한계를 극복하는 하드웨어 인식형 FP4 FlashAttention-4를 소개하며, 이를 통해 MXFP4 훈련에서 관찰되는 발산 문제를 피하면서 포워드 처리량을 최대 2.13배, 단일 GPU 훈련 업데이트 속도를 최대 1.14배 향상시켰습니다.

원저자: Robert Hu

게시일 2026-09-04✓ Author reviewed ⓘ

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Robert Hu

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

인공지능의 세계에서 가장 강력한 모델들은 문맥을 이해하기 위해 "어텐션(attention)"이라 불리는 메커니즘에 의존합니다. 독자가 긴 문서를 훑어보는 모습을 상상해 보십시오. 어텐션은 모델이 나머지 부분은 무시하면서 문장 내에서 가장 관련 있는 단어들에 집중하고, 먼 거리의 아이디어들을 연결할 수 있게 해줍니다. 이를 위해 컴퓨터는 모든 단어를 다른 모든 단어와 비교하여 그 관계를 결정하는 방대한 일련의 계산을 수행합니다. 수년 동안 이러한 계산은 건물을 측정하기 위해 아주 미세하고 정밀한 눈금이 있는 자를 사용하는 것과 유사하게, 고정밀 숫자를 사용하여 수행되어 왔습니다. 이는 모델이 올바르게 학습하도록 보장하지만, 속도가 느리고 엄청난 양의 에너지를 소비합니다. 모델이 커짐에 따라 속도의 필요성은 매우 중요해집니다. 연구자들은 최근 훨씬 더 작고 거친 숫자—4비트 부동 소수점 값—를 사용하여 이러한 비교를 훨씬 빠르게 수행할 수 있는 칩을 개발했습니다. 하지만 단순히 더 작은 숫자로 전환하는 것만으로는 충분하지 않습니다. 데이터의 흐름을 관리하는 소프트웨어 또한 변해야 하며, 그렇지 않으면 속도 이점이 사라집니다.

Graphcore의 한 연구자는 "Direct-P"라고 부르는 새로운 방법으로 이 특정 병목 현상을 해결했습니다. 그들의 연구는 모델이 정보를 처리하여 답을 생성하는 "순방향(forward)" 패스와, 실수를 통해 학습하는 "역방향(backward)" 패스에 초점을 맞춥니다. 연구자는 새로운 칩이 숫자를 믿기 힘들 정도로 빠르게 곱할 수 있는 반면, 그 중간 단계인 원시 점수(raw scores)를 확률로 변환하는 과정이 모든 것을 느리게 만들고 있다는 것을 발견했습니다. 이것은 마치 모든 부품을 전달하기 전에 작업자가 모든 부품을 주의 깊게 측정해야 하기 때문에 계속 멈춰 서는 초고속 조립 라인과 같았습니다. 연구자는 복잡한 스케일링과 반올림을 포함하는 표준적인 방식의 이 변환 처리가 새로운 하드웨어의 속도 이점을 상쇄하는 교통 체증을 유발한다는 것을 발견했습니다.

이를 해결하기 위해, 연구자는 변환 과정을 직접적이고 간소화되도록 재설계했습니다. 정밀한 확률을 계산한 다음 반올림하는 대신, 그들의 방법은 원시 점수를 하드웨어가 사용하는 특정 코드로 직접 매핑합니다. 이는 지연을 일으켰던 중간 단계의 중간 관리자를 제거합니다. 이 접근 방식을 최신 세대의 데이터 센터 칩에서 테스트했을 때, 결과는 놀라웠습니다. 특정 형태의 데이터에 대해, 새로운 방법은 고정밀 숫자에 의존했던 이전 표준보다 정보를 2배 이상 빠르게 처리했습니다. 이는 비디오 생성 및 언어 이해와 같은 복잡한 작업에 충분히 정확한 출력을 유지하면서도 달성되었습니다. 비디오 생성 모델을 포함한 테스트에서, 새로운 방법은 표준 방식보다 거의 2배 더 빨랐으며, 결과물은 유한하고 사용 가능했지만, 더 느리고 정밀한 버전과 비교했을 때 측정 가능한 편차와 낮은 유사도를 보였습니다.

그러나 이야기는 속도에 관한 것뿐만이 아닙니다. 모델이 학습하려고 할 때 어떤 일이 일어나는지에 대한 것이기도 합니다. 연구자는 모델이 지식을 업데이트하는 역방향 패스를 포함하여 전체 학습 과정에 이 초고속 저정밀 숫자를 사용할 수 있는지 탐구했습니다. 그들은 순방향 패스는 풀 스피드로 실행될 수 있는 반면, 역방향 패스는 신중한 타협이 필요하다는 것을 발견했습니다. 학습 중에 확률 값을 위해 가장 빠른 4비트 형식을 사용하려고 시도했을 때, 학습 과정은 불안정해졌고 모델은 개선되지 못했습니다. 숫자들이 너무 거칠어져서 학습 신호가 무너졌기 때문입니다. 결과적으로, 연구자는 학습이 안정적으로 유지되려면 나머지 계산이 더 빠른 4비트 형식을 사용하더라도 확률 값에는 약간 더 정밀한 8비트 형식을 사용해야 한다고 결정했습니다. 이 하이브리드 접근 방식은 단일 강력한 칩에서 전체 학습 주기를 약 14% 가속화했으며, 이는 대규모 모델에 있어 상당한 이득입니다.

연구자는 또한 컴퓨터 칩 자체의 물리적 한계도 살펴보았습니다. 그들은 계산의 속도가 더 이상 주요한 문제가 아니라는 것을 발견했습니다. 문제는 칩 내부에서 데이터가 어떻게 저장되고 이동하느냐였습니다. 칩에는 숫자를 다루는 동안 보유하는 작고 초고속인 메모리 영역이 있습니다. 연구자는 이 메모리 공간이 완전히 가득 차서, 계산의 서로 다른 단계들을 중첩(overlap)시킬 여유가 없다는 것을 발견했습니다. 그것은 마치 요리사가 칼이 아무리 날카롭더라도 현재의 채소를 다 썰기 전까지는 다음 채소를 썰기 시작할 수 없을 정도로 조리대가 꽉 찬 주방과 같았습니다. 메모리 공간이 완전히 점유되었기 때문에, 칩은 여러 단계를 동시에 작업할 수 없었으며, 이는 전체 프로세스가 얼마나 더 빨라질 수 있는지에 제한을 두었습니다.

이 연구는 인공지능을 더 빠르게 만드는 방법에 대해 우리가 생각하는 방식의 중요한 변화를 강조합니다. 단순히 더 빠른 계산기를 만드는 것만으로는 부족합니다. 하드웨어의 역량을 활용할 수 있도록 전체 워크플로우를 재설계해야 합니다. 연구자는 확률을 계산하는 방식을 바꾸고 데이터의 흐름을 신중하게 관리함으로써 막대한 속도 향상을 끌어낼 수 있음을 보여주었습니다. 그러나 그들은 또한 엄연한 한계가 존재함을 증명했습니다. 칩의 메모리 제약은 가장 빠른 산술 연산을 사용하더라도 중첩이 가능한 수준에 천장이 있음을 의미합니다. 그들이 제시하는 미래의 길은 단순히 더 빠른 산술이 아니라, 칩의 자원이 결코 기다리는 일이 없도록 데이터를 조직하는 더 스마트한 방식에 있습니다. 이 원시적인 속도와 신중한 데이터 관리 사이의 균형이 차세대 인공지능을 효율적으로 구동하게 할 것입니다.

=== 기술 요약: 하드웨어 인지형 FP4 FlashAttention-4 ===

문제 정의

NVIDIA Blackwell GPU는 행렬 곱셈을 크게 가속화하는 4비트 부동 소수점(FP4) 텐서 코어를 도입했습니다. 그러나 어텐션 메커니즘에서는 속도 저하가 행렬 곱( $QK^T$ 및 $PV$ )이 아니라, 중간 단계인 softmax에서 발생하는 경우가 많습니다. 이 단계는 스코어 리덕션(score reduction), 지수 함수 평가, 확률 생성, 그리고 스케일링 과정을 포함하며, 두 번째 행렬 곱이 시작되기 전에 반드시 완료되어야 합니다.

FP4를 어텐션 블록 전체에 활용하려는 기존의 시도들은 두 가지 주요 장애물에 직면해 있습니다:

타이밍(Timing): 스코어( $S$ )로부터 확률 텐서( $P$ )를 생성하는 과정은 임계 경로(critical path) 의존성을 만들어 $PV$ 곱의 실행을 지연시킵니다.
수치 범위(Numerical Range): FP4 포맷(특히 NVFP4 및 MXFP4)은 제한된 동적 범위를 가진 4비트 페이로드(E2M1)를 사용합니다. 표준 softmax 확률은 이 페이로드에 적합하지 않은 범위를 가지며, 이를 정밀하게 관리하지 않으면 언더플로우, "제로 스케일"(전체 블록이 사라지는 현상), 또는 심각한 정밀도 손실을 초래합니다.

나아가, **인과적 학습(causal training)**의 경우, 역전파(backward pass) 시 일반적으로 저장된 상태로부터 확률을 재구성해야 합니다. BF16으로 스코어를 전체 정밀도로 저장하거나 재계산하는 방식은 저정밀도 학습의 잠재적인 속도 향상을 상쇄합니다.

방법론: Direct-P 및 Causal Backward

본 논문은 Blackwell 하드웨어(GB200/B300)에서 풀-FP4 포워드 추론과 양자화된 인과적 역전파를 가능하게 하는 방법인 Direct-P를 제안합니다.

1. Forward Inference를 위한 Direct-P

Direct-P는 $QK^T$ 스코어 단편(fragment)의 완료와 $PV$ 명령 발행 사이의 간격을 수정합니다. 기존의 HAO AI Lab의 two-query 파이프라인과 텐서 메모리(TMEM) 소유권 프로토콜을 유지하면서 확률 생성을 최적화합니다:

직접 매핑(Direct Mapping): 고정밀 지수 함수를 계산하고 반올림하는 대신, Direct-P는 확률 생성을 코드 분류 문제로 취급합니다. 정규화된 로그 스코어를 아핀 분류기(affine classifier) 또는 선택적 네이티브 지수 함수를 사용하여 8개의 비음수 E2M1 크기로 직접 매핑합니다. 이는 임계 경로를 단축합니다.
일관된 정규화(Consistent Normalization): softmax의 분모(normalizer)는 $PV$ 명령에 의해 소비될 FP4 확률 코드와 블록 스케일로부터 직접 누적됩니다. 이를 통해 분자와 분모가 동일한 근사 연산자를 설명하도록 하여, 고정밀 합계와 저정밀 페이로드 간의 불일치를 방지합니다.
범위 가드(Range Guarding): 극단적인 로짓(예: >500)이 있는 레이어의 경우, 샘플링된 앵커 메커니즘을 사용하여 전체 스코어를 재스캔하지 않고도 서브노멀 언더플로우를 방지합니다.
포맷 선택: $Q$ 와 $K$ 에는 미세한 스케일링이 가능한 NVFP4를 사용하고, $P$ 와 $V$ 에는 ( $N_{32}$ 프로듀서 단편과 잘 정렬되는 거듭제곱 2 기반 블록 스케일을 갖춘) MXFP4를 사용합니다.

2. 양자화된 Causal Backward

학습을 위해, 본 논문은 포워드 패스에서 양자화된 $Q/K$ 페이로드, 스케일, 그리고 log-sum-exp (LSE) 정규화 값을 저장하는 경로를 도입합니다.

재구성(Reconstruction): 역전파는 별도의 BF16 스코어 경로를 생성하는 대신, 저장된 양자화된 피연산자로부터 확률을 직접 재구성합니다.
그래디언트 포맷: 프로젝션 에필로그(projection epilogues)는 그래디언트 행렬 곱에 피딩하기 위해 $Q, K, V$ 의 행/열 지향적 FP8 뷰를 발행합니다.
안정성 제약(Stability Constraint): 포워드 패스에서는 $P/V$ 에 MXFP4를 사용하지만, 분산 학습 실험 결과 MXFP4 P/V 궤적이 발산함이 밝혀졌습니다. 따라서 안정성을 보장하기 위해 유지된 학습 경로는 역전파 시 $P$ 와 $V$ 에 대해 FP8를 사용하며, $Q/K$ 는 NVFP4를 유지합니다.

주요 기여

Direct-P 구현: 스코어에서 MXFP4 확률 코드로 직접 매핑하고 동일한 표현된 값들을 사용하여 정규화함으로써, 유리한 형태(shape)에 대해 GB200에서 최대 2.13배의 BF16 포워드 처리량을 달리는 방법을 제시했습니다. 테스트된 D128 그리드 전반에서 기하 평균 속도 향상은 2.02배입니다.
양자화된 Causal Training 경로: 양자화된 $Q/K$ 페이로드, 스케일, LSE 정규화 값을 저장하여, 배치 크기 4에서 전체 80억 파라미터 업데이트를 최대 1.14배 가속화하는 역전파를 도입했습니다.
하드웨어 프로파일링: 텐서 메모리(TMEM) 소유권을 오버랩(overlap)의 주요 제한 요소로 식별했습니다. 본 연구는 단순히 더 빠른 텐서 코어를 갖는 것만으로는 부족하며, 저장 레이아웃(D128 기준 512개 컬럼)이 현재의 $PV$ 단계와 다음 $QK$ 연산을 오버랩하는 것을 방해할 수 있음을 보여줍니다.

결과

포워드 추론 성능

처리량(Throughput): GB200에서 "Fast" Direct-P 경로는 HAO의 BF16 베이스라인 대비 2.02배의 기하 평균 속도 향상을 달성하며, 최고 2998 TFLOP/s(S32768/H24)에 도달합니다.
정확도: Direct-P는 속도를 위해 일부 정밀도를 희생합니다. 가우시안 소프트맵 테스트에서 "Fast" 경로는 BF16 대비 약 0.944의 평균 코사인 유사도를 보였으며, 이는 HAO의 NV/FP8 경로의 약 0.990과 대조됩니다.
모델 영향:
- ViT/BERT: 고정 입력 평가 시, "Fast" 경로는 연산자 수준의 코사인 유사도가 낮음에도 불구하고 태스크 정확도(예: ViT S4096 top-1 정확도가 BF16과 동일한 88.5% 유지)를 유지합니다.
- Wan Video Diffusion: 이 방법은 20단계 디퓨전에서 유한하고 기능적으로 작동합니다. 그러나 단계가 진행됨에 따라 오차가 누적됩니다. 14B 모델의 20단계 결과는 BF16 대비 상대적 L2 오차가 ~0.53로 나타나, 드리프트(drift)는 있으나 생성은 기능적임을 보여줍니다.

학습 성능

단일 GPU: 양자화된 경로(NVFP4 프로젝션 + FP8 P/V)는 배치 크기 4에서 전체 어텐션 서브레이어를 1.25배 가속하고, 완전한 8B 모델 업데이트를 1.14배 가속합니다.
분산 학습: 1,000억 토큰에 걸친 64-GPU 실험에서, FP8 P/V 경로는 BF16 대비 1.112배의 처리량 개선을 달왔습니다.
안정성: 논문은 MXFP4 P/V 궤적이 분산 환경에서 발산함을 명시적으로 언급했습니다. 성공적인 학습 경로는 $P/V$ 제품에 FP8을 사용하는 데 의존하며, 이는 학습 경로가 엔드 투 엔드로 "순수한" FP4는 아님을 의미합니다.

의의 및 주장

본 논문은 Direct-P가 행렬 곱 사이의 임계 경로를 단축함으로써 풀-FP4 포워드 어텐션을 실용적으로 만든다고 주장합니다. 이는 텐서 코어가 더 빨라지더라도 병목 현상이 데이터 이동, 확률 생성, 그리고 메모리 소유권으로 이동했음을 입증합니다.

핵심 요약:

하드웨어 한계: 추가적인 속도 향상의 주요 장벽은 산술 처리량이 아니라 TMEM 소유권입니다. 현재의 D128 레이아웃은 512개의 TMEM 컬럼을 모두 점유하여, 다음 $QK$ 연산이 $PV$ 단계의 꼬리 부분과 오버랩되는 것을 방해합니다.
준비 지연(Readiness Stalls): 낮은 텐서 코어 활성도는 연산 능력의 부족보다는 피연산자의 준비 상태(확률 생성 대기) 때문인 경우가 많습니다.
학습 트레이드오프: FP4 포워드 추론은 실행 가능하지만, 안정적인 인과적 학습을 위해서는 현재 $P/V$ 피연산자에 대한 FP8 양보가 필요합니다. 논문은 향후 Unsigned E5M3(UE5M3)와 같은 포맷이 완전한 FP4 학습을 가능하게 할 수 있다고 제안하지만, 이는 본 연구에서 평가되지 않았습니다.

결론적으로, 더 빠른 행렬 산술만으로는 충분하지 않으며, 더 큰 할당 가능 오버랩 윈도우(예: 추가적인 스코어 목적지)를 허용하는 아키텍처적 변화와 최적화된 저정밀도 알고리즘이 병행되어야 높은 효율성을 달성할 수 있습니다.