FLAMIA가 Blackwell 저정밀 attention 학습을 안전하게 가속한다

무엇인가

이 논문은 Transformer 학습에서 긴 문맥 attention 비용을 Blackwell의 FP8·FP4 tensor core로 끌어내리는 문제를 다룬다. Blackwell은 BF16 대비 FP8과 FP4 행렬곱에서 각각 2배, 4배의 피크 처리량을 제공하지만, attention에 그대로 적용하려면 양자화 오버헤드, causal mask 아래에서 공유 스케일을 통한 미래 토큰 정보 누수, forward와 backward의 수치 불일치, 그리고 사전학습에서 문서를 이어 붙인 packed 입력의 가변 길이 처리를 함께 해결해야 한다. 저자들은 단일 스텝의 cosine 유사도만으로는 이런 문제를 잡을 수 없다고 본다.

어떻게 동작하나

FLAMIA는 MXFP8 경로와 NVFP4/FP8 혼합 경로를 제공한다. 핵심은 token-block scaling으로, 32개 연속 토큰과 head dimension 전체가 하나의 2의 거듭제곱 지수를 공유한다. 혼합 경로에서 Q와 K는 토큰별로 head-dim 16개 원소마다 NVFP4 스케일을 쓰고, dS는 32×32 블록마다 E4M3 스케일을 쓴다. 8비트 MMA가 K-major와 transposed 레이아웃을 모두 지원하므로 하나의 E4M3 표현을 두 방향의 곱에 재사용할 수 있고, 여섯 개 논리 operand마다 양자화 표현을 중복 생성하는 비용을 줄인다. all-4-bit 경로도 지원하지만 기본값은 아니다. 4비트 MMA에는 MN-major 형태가 없어 방향별 복사가 필요하고, 여섯 복사본 양자화가 backward 시간의 큰 부분을 차지하며, 4비트 V는 대각 타일에서 per-key 스케일 비용을 추가하고, 평가된 backward 곱이 tensor-core-bound가 아니어서 이득이 제한된다.

무엇과 다른가

Causal 누수는 공유 스케일이 미래 토큰의 최댓값 변화에 반응하면서 발생한다. 미래 토큰이 스케일 그룹의 최댓값을 올리면 스케일이 바뀌고, 그 스케일로 복원되는 이전 토큰의 값도 달라질 수 있다. per-tensor 스케일은 그 영향이 시퀀스 전체로 번지는 극단적 사례다. FLAMIA는 어떤 operand 스케일도 32개 연속 위치를 넘지 않게 하고, Q와 K의 per-token NVFP4 스케일은 누수가 원천적으로 불가능하게 한다. 스케일 지수를 올림해 스케일된 최댓값이 E4M3 최대 코드값을 넘지 않게 하므로 클램핑으로 인한 코드 변화도 막는다. 4비트 V에는 대각 타일에서만 per-key 2의 거듭제곱 스케일을 적용해 P로 접어 넣고, dQ=dS·K의 K에도 K fold를 적용하며, FP16 dQ accumulator의 행별 스케일을 보이는 키만으로 결정하는 causal dQ scale을 선택적으로 제공한다.

어떻게 쓰나

수치 일관성도 별도 설계 대상이다. backward의 score 재계산은 forward와 같은 양자화된 Q·K 표현을 쓰고, tiling과 대각 처리, exponential 평가, row maximum 기준 스케일링을 forward와 맞춘다. D는 dP에 쓰는 것과 같은 양자화 dO로 계산한다. dS는 E4M3로 양자화하되 32×32 블록의 최대 절댓값에서 스케일을 잡아 작은 원소가 0으로 반올림되는 비율을 줄인다. head dimension 128에서는 정규화 Hadamard 변환도 기본으로 켜 채널 이상치를 분산한다.

전제와 한계

커널 수준에서는 양자화를 한 번의 launch로 융합해 Hadamard 회전, 블록 통계, 코드와 스케일 생성을 함께 처리하고, 출력과 반환 gradient를 BF16으로 유지해 별도 출력 스케일 보정 패스를 없앤다. MMA warp는 다음 타일의 QKᵀ를 두 개의 64-key 연산으로 나눠 현재 타일의 P·V 전후에 발행하고, NVFP4 경로는 확률을 log 영역에서 양자화해 나눗셈과 후속 스케일링 패스를 제거한다. B300에서는 완전히 unmasked된 타일에 tcgen05.ld.red를 사용하고, low-bit 2-CTA backward와 FP16 dQ partial의 bulk reduce-add, 데이터 의존 2의 거듭제곱 prescale을 결합한다. packed 문서는 문서 경계 안에서 토큰 블록과 통계를 유지하고 per-token 스케일 메타데이터를 쓰며, 유효하지 않은 query row에는 무한 log-sum-exp sentinel을 넣는다. longest-first 순서와 GPU 기반 CLC 스케줄링도 사용한다. 추론용으로는 Schraudolph와 Mitchell-P 방식의 비트 수준 exponential 근사도 평가한다.

성능 실험은 16개 query head, head dimension 128, MHA와 GQA 16/4, 동일 길이·편향·로그균일 문서를 섞은 9개 causal packed-document 워크로드에서 이뤄졌다. MXFP8 경로는 양자화를 포함한 forward+backward attention 호출에서 9개 causal 워크로드 모두 가장 낮은 지연을 기록해 B200에서 1.10~1.40배, B300에서 1.06~1.42배 빨랐다. 8192토큰 문서와 224개의 256토큰 문서가 섞인 워크로드에서는 B300 기준 training call 1.42배, backward 2.00배였고, cuDNN FP8은 약 1.68배, FA4 BF16은 1.40~1.42배, cuDNN BF16은 3.7~3.9배 느렸다. 커널만 측정하면 NVFP4 backward가 모든 워크로드에서 가장 빨라 B200 1.22~2.01배, B300 1.17~2.00배였고, MXFP8 forward는 B200 9개 모두, B300 8개에서 1.07~1.72배와 0.98~1.69배를 보였다. packed 입력은 T=131072에서 attention 호출당 peak memory가 2.7GB로, 정렬·패딩 버퍼의 16.8GB보다 작았다. 1B 모델 전체 학습에서 packed 입력은 16k·64k row 길이에서 각각 84.6k·96.8k tokens/s였고, 패딩 구성은 68.8k·84.3k였으며 128k에서는 메모리를 초과했다.

Ablation에서는 융합 양자화가 4×4096 토큰에서 MXFP8 전처리를 789µs에서 261µs로, NVFP4를 0.64ms에서 0.11ms로 줄였고 launch 수는 각각 17→1, 11→1이었다. NVFP4 forward에서 log-domain 확률 양자화가 14~19%, QKᵀ 분할이 4.7%, MXFP8의 score 스케일 fold가 4% 기여했다. B300에서 pure-E4M3 구성의 exponential emulation을 끄면 17~33%, 완전 unmasked 타일의 ld.red는 causal NVFP4 forward에 18~20% 이득이 있었다. 혼합 NVFP4/FP8 backward는 기하평균 1329 TFLOP/s로 all-4-bit 571, BF16 1015보다 높았다. all-4-bit에서는 여섯 복사본 양자화가 backward 시간의 45~52%를 차지했고, FP16 dQ partial과 데이터 의존 스케일은 backward를 23.8% 개선했다. causal dQ scale은 32768 토큰에서 backward 8~11%, training call 6~9%의 오버헤드를 냈다.

학습 검증은 1B Llama-3 스타일 모델을 100B FineWeb-Edu 토큰, 2k 문맥, 8개 B300에서 사전학습해 수행했다. MXFP8과 NVFP4/FP8은 최종 1000스텝 평균 손실과 5개 downstream 과제에서 baseline과 가까운 점수를 보이며 손실 증가가 미미했다. 반면 누수 실험은 더 경고적이다. enwik8에서 6층·폭 1024 byte-level GPT를 1024바이트 문맥으로 6000스텝 학습했을 때, 누수 있는 4비트 구성은 자체 커널로 평가하면 검증 손실 0.811~0.827로 BF16의 0.852보다 좋아 보였지만, 같은 가중치를 exact causal attention으로 평가하면 0.888~0.940로 나빠졌고 누수 없이 학습한 모델은 0.862였다. D와 dP에 서로 다른 dO 표현을 쓰면 500스텝 안에 발산했고, forward 4비트·backward 8비트 확률은 attention을 과도하게 집중시켰으며, forward 256 query·backward 128 query의 대각 처리는 단일 스텝 dQ cosine 0.945에도 gradient norm을 0.5에서 2×10^5로 키웠다. forward에만 exponential 근사를 쓰면 500스텝까지 발산했다. dS 스케일 ablation에서는 0으로 반올림되는 원소 비율이 0.4%에서 29%로 늘 때 BF16 대비 손실 격차가 0.0008에서 0.049로 커졌다.

개발자에게 이 논문의 실무적 메시지는 저정밀 attention을 “커널 한 스텝의 cosine”이나 “자체 커널로 잰 검증 손실”만으로 채택하면 안 된다는 것이다. 판단 기준은 양자화를 포함한 forward+backward 전체 지연, exact causal attention으로 재평가한 품질, forward와 backward의 표현·타일링·exponential 일관성, packed 문서에서 문서별 스케일이 섞이지 않는지다. 기본 선택은 combined training call에서 MXFP8이 안정적이고, NVFP4는 커널만 볼 때 backward가 강하지만 forward는 MXFP8보다 느리다. all-4-bit는 기본 경로가 아니며 더 느리고 손실도 약간 높다.

한계도 분명하다. 기본 스케일링은 32개 블록 내부의 잔여 의존성과 시퀀스 전역 backward accumulator 의존성을 남기며, 이를 제거하는 exact 옵션은 지원되는 batched causal 구성에만 적용된다. causal dQ scale은 32768 토큰에서 backward 8~11%, training call 6~9%의 비용을 추가한다. per-key V 스케일을 모든 타일에 적용하면 4096·32768 토큰에서 각각 2.33배·2.83배 느려지므로 대각 타일로 제한한다. 추론의 Mitchell-P 구성은 causal 워크로드에서 NVFP4 대비 1.25~1.77배 느리고 output cosine이 0.990~0.993에서 0.980~0.982로 떨어진다. 저자들은 평가된 워크로드와 구성 안에서 결론을 제시하며, 단일 스텝 유사도가 causal 정확성이나 학습 안정성을 보증하지 않는다고 명시한다.