기술 요약: FP8 어텐션에서의 P-Cast 정밀도 (P-Cast Precision in FP8 Attention)

1. 문제 정의

본 논문은 대규모 언어 모델(LLM)의 FP8 (E4M3) 어텐션 가속화에서 발생하는 결정적인 정밀도 병목 현상을 다룹다. 소프트맥스 확률 행렬 $P$를 $P \cdot V$ 행렬 곱셈 이전에 E4M3로 캐스팅(casting)하는 것은 상당한 처리량(throughput) 이득을 제공하지만, 해당 포맷의 3비트 가수부(mantissa)는 어텐션 싱크(Attention Sink) 현상 하에서 심각한 양자화 오차를 유발한다.

어텐션 싱크 현상에서 초기 토큰들은 불균형적으로 큰 어텐션 가중치를 받는다. 이를 전방 KV 블록 반복(forward KV block iteration) (토큰 $0 \to N$ 순서로 처리)과 결합하면, 싱크 토큰의 높은 로짓(logit) 점수가 실행 중인 소프트맥스 최댓값($m$)을 팽창시킨다. 이로 인해 후속되는 비싱크(non-sink) 확률 값($P$)들이 극도로 작아지게 된다. E4M3 양자화 하에서 이러한 작은 값들은 표현 가능한 범위를 벗어나 0으로 언더플로우(underflow)되는데, 이를 P-붕괴(P-collapse) 현상이라 부른다.

본 논문은 P-붕괴 발생 여부를 결정하는 두 가지 구현 선택지를 식별한다:

1. KV 블록 반복 순서: 온라인 소프트맥스가 블록을 전방($0 \to N$)으로 처리하는지 또는 역방향($N \to 0$)으로 처리하는지 여부.
2. 정적 스케일링 계수 ($S$): E4M3로 캐스팅하기 전(및 에필로그에서 다시 나누기 전) $P$에 적용되는 상수 승수.

2. 방법론 및 분석

이론적 프레임워크

저자들은 스케일 계수 $S$를 사용하는 전방 반복 시, 0으로 언더플로우되는 비싱크 $P$ 값의 비율에 대한 폐쇄형 수식(closed-form expression)을 도출하였다:
$$F(\Delta, S) = \Phi(\Delta + \delta_k - 6.93 - \ln S)$$
여기서:

• $\Delta$는 싱크 강도(sink strength, 싱크 토큰과 비싱크 토큰 간의 로짓 차이)이다.
• $\delta_k$는 싱크 블록 내 $k$개의 표준 가우시안 변수의 기대 최댓값이다 (예: $k=4$일 때 약 1.0).
• $6.93 \approx 10 \ln 2$는 E4M3의 round-to-zero 임계값($2^{-10}$)을 나타낸다.
• $\Phi$는 표준 정규 분포의 누적 분포 함수(CDF)이다.

이 분석은 P-붕괴가 $\Delta \approx 6\text{--}7$ 부근에서 활성화되는 임계 효과임을 밝혀내며, 이 지점에서 상당한 양의 비싱크 확률 질량이 암묵적으로 0으로 처리되어 출력값에 상당한 평균 제곱 오차(MSE)를 유발한다.

최적화 1: 역방향 KV 반복 (Reverse KV Iteration)

본 논문은 반복 순서를 역방향($N \to 0$)으로 바꾸면 싱크 블록의 처리가 마지막 단계로 지연됨을 증명한다. 모든 선행 반복 동안, 실행 중인 최댓값은 비싱크 토큰들에 의해서만 결정된다. 이론적 경계치에 따르면 실질적인 시퀀스 길이($N \le 10^6$)에 대해, 선행 싱크 최댓값은 충분히 낮게 유지되어 $S=256$과 결합될 경우 비싱크 $P$ 값이 언더플로우될 확률은 사실상 0($< 10^{-12}$)이다.

최적화 2: $S = 256$의 최적성

본 논문은 $S=256$ ($2^8$)이 다음 세 가지 조건을 동시에 만족하는 최적의 정적 스케일임을 구성적으로 설명한다:

1. 비트 정확 스케일링 (Bit-Exact Scaling): $2^k$ 형태의 멱수(power of two)로서, $S$에 의한 곱셈과 나눗셈은 IEEE 754 FP32에서 정확하게 수행되어 스케일링/언스케일링 과정에서 반올림 오차를 도입하지 않는다. 2의 거듭제곱이 아닌 스케일(예: $S=448$)은 체계적인 반올림 오차를 유발한다.
2. 최소 양자화 단계 ($d_p$): 저자들은 $d_p(S)$를 정규화된 최대 양자화 단계로 정의한다. 이들은 $d_p(S)$가 E4M3 수직선 상에서 "톱니 모양(sawtooth)" 함수를 형성하며, 2의 거듭제곱 값($S=2^k$)들이 최소값인 $2^{-4}$를 갖는 하단 포락선(lower envelope)에 위치함을 입증한다. $S=448$은 이 톱니 모양에서 14% 더 높은 곳에 위치하여 더 큰 최악의 경우 양자화 오차를 의미한다.
3. 최대 정규 범위 커버리지: 오버플로우를 일으키지 않는 2의 거듭제곱 스케일($S \le 448$) 중에서, $S=256$은 가장 크며, 이는 $P$ 도메인의 서브노멀(subnormal) 영역($2^{-6}/256 = 2^{-14}$)으로 가장 깊게 커버리지를 제공한다.

3. 주요 기여

1. P-붕괴의 정량화: 단순한 경험적 관찰을 넘어, P-cast 실패를 예측할 수 있는 폐쇄형 임계값 $\Delta_c = 6.93 + \ln S - \delta_k$를 도출하였다.
2. 충분성 증명: 역방향 반복과 $S=256$의 조합이 실질적인 시나리오에서 비싱크 값에 대한 제로 언더플로우를 보장함을 엄격하게 입증하였다.
3. $S=256$의 특성 규명: $S=256$이 비트 정확성, 최소 양자화 단계, 그리고 최대 정규 커버리지를 만족하는 유일한 스케일임을 확립하였다.
4. 커널 충실도 검증 (Kernel-Faithful Validation): (Q, K, V를 FP32로 유지하여 P-cast 효과를 격리함으로써) 생산 환경의 커널 시맨틱스를 시뮬레이션하여 이론적 주장을 검증하였다.

4. 실험 결과

$N=4096$ 및 다양한 싱크 강도($\Delta$)를 사용한 커널 충실도 시뮬레이션 결과:

• MSE 개선: 중간 정도의 싱크 강도($\Delta \approx 7$)에서, 기본 설정(전방 반복, $S=1$)은 최적화된 구성에 비해 3~10배 높은 MSE를 보인다.
• 스케일 비교: 전방 반복과 $S=256$을 사용할 경우, $S=448$(FlashInfer 및 TensorRT-LLM에서 사용)을 사용하는 전방 반복보다 약 10~15% 낮은 MSE를 기록하였으며, 이는 이론적인 $d_p$ 비율과 일치한다.
• 포화 (Saturation): 두 가지 최적화(역방향 + $S=256$)를 모두 적용했을 때, 결과는 P-붕괴 영역에서 전방 + $S=256$과 통계적으로 구분이 불가능하다. 두 방식 모두 고유한 E4M3 양자화 노이즈에 의해 결정되는 동일한 정밀도 하한선에 도달한다.
• 시퀀스 길이: 최적화된 구성이 기본 설정 대비 보이는 MSE 개선 효과는 시퀀스 길이가 길어질수록 커진다. 이는 긴 시퀀스일수록 위험에 처하는 비싱크 확률 질량이 증가하기 때문이다.

5. 의의 및 주장

본 논문은 새로운 기술을 제안하는 것이 아니라, 이미 FlashAttention-3/4(Reverse + $S=256$ 사용)에서 채택된 구현 선택지에 대한 정량적 설명을 제공하는 데 목적이 있다.

• 실질적 영향: 이 분석은 hpc-ops 커널을 $S=1$에서 $S=256$으로 업데이트하도록 유도하였다.
• 권장 사항: 현재 $S=1$ 또는 $S=448$을 사용하는 구현체에 $S=256$으로 전환할 것을 권장한다. 이는 "비용이 들지 않는(zero-cost)" 변경(비트 정확 스케일링은 무시할 만한 지연 시간만 추가함)이며, P-붕괴를 제거하고 P-cast 정밀도를 엄격하게 개선하거나 유지한다.
• 신중한 접근: 저자들은 P-cast MSE 개선이 3~10배로 유의미하지만, 최종적인 모델 퍼플렉시티(perplexity)나 태스크 정확도에 미치는 영향은 다른 FP8 노이즈 소스(QKV 양자화 등)에 의해 지배될 수 있으므로 보장되지 않는다고 명시하였다. 따라서 $S=256$으로의 전환을 모든 다운스트림 지표에 대한 보장된 품질 승리가 아닌, "결코 나빠지지 않는(strictly-no-worse)" 선택으로 규정한다.
• 범위: 본 분석은 E4M3 P-casting에 특화된 것이며, 전제 조건이 다른 다른 포맷(예: MXFP4)이나 퍼 채널(per-channel) 양자화 전략으로의 전이 가능성을 주장하지 않는다.