프리트레인 LLM의 소프트맥스를 거칠게 근사해 B200 어텐션을 가속한다
Approximating Softmax in Pretrained LLMs: Model Sensitivity and Kernel Acceleration
무엇인가
NVIDIA Blackwell B200에서 텐서코어 처리량은 특수함수유닛(SFU)의 지수 함수 처리량보다 두 자릿수 이상 앞선다. SM당 클럭당 8192번의 행렬 연산이 가능한 반면 SFU가 내놓는 지수는 16개뿐이고, FlashAttention-4는 이미 일부 exp2 평가를 다항식 근사로 대체한 상태다. 이 논문은 여기서 질문을 바꾼다. 프리트레인된 트랜스포머가 정말 모든 원소에서 지수 함수를 정확히 평가해야 하는가. 저자들은 0.5B에서 72B까지 열 개의 frozen decoder-only 모델에서 추론 시 소프트맥스를 근사하며, 어떤 성질은 완화해도 되고 어떤 성질은 안 되는지를 실험으로 분리한다.
어떻게 동작하나
실험 설계의 핵심은 한 번에 한 속성만 바꾸고 나머지를 고정하는 것이다. 대상이 되는 속성은 어텐션 support(0이 아닌 가중치를 받는 키 위치 집합), 가중치 격자의 해상도, 그리고 고정된 구간 수를 어디에 배분하는지다. full-range grid는 각 행의 점수 범위 [0, C]를 K개 구간으로 나누되 e_a = C(1-ρ^a)/(1-ρ^K), ρ = R^(-1/(K-1))로 경계를 잡아 R>1이면 행 최댓값 근처가 더 촘촘해지게 만든다. support의 효과만 보기 위한 통제 실험에서는 Qwen2.5-1.5B의 각 타깃 레이어 아래는 모두 softmax로 두고, 같은 mean-threshold support 위에서 softmax 가중치와 균등 가중치만 바꿔 비교한다. 왜곡을 맞춘 비교도 있다. J0를 어떤 모델의 평균 row-wise JSD로 정의하고, 점수를 α<1로 눌러 평탄화(flattening)한 경우와 α>1로 날카롭게(sharpen)한 경우를 같은 JSD에 맞춘 뒤 NLL 차이를 본다.
무엇과 다른가
첫 번째 결과는 support와 해상도를 상당히 줄여도 된다는 것이다. Qwen2.5-1.5B에서 각 행의 가장 큰 절반만 남기면 perplexity가 0.131% 오르고, 4분의 1만 남기면 0.828%, 1%만 남기면 90.478%까지 치솟는다. 실패 구간은 작은 support 쪽에 있다. 반면 같은 support 위에서 가중치를 균등하게 주면 Qwen2.5-1.5B의 28개 타깃 레이어 전부에서 NLL이 증가했고, 모든 점별 95% 구간이 0을 배제했다. 즉 위치는 버려도 되지만 가중치의 지수적 성격 자체는 버릴 수 없다. K=32, R=4 격자는 열 개 모델의 모든 점추정에서 perplexity 증가가 1% 미만이었고, Gemma-2-2B만 구간이 1%를 걸쳤다.
어떻게 쓰나
두 번째 결과는 고정된 구간 수를 어디에 쓸지가 그 수 자체만큼 중요하다는 것이다. Qwen2.5-1.5B에서 K=21로 고정하고 R을 1에서 4로 바꾸면 가중치 없는 RMS 로그가중치 오차는 27.9% 늘어나는데, 확률 가중 오차는 29.2% 줄고 ΔNLL은 71.7% 줄어든다. R만 바꿨을 때 ΔNLL은 R=0.25의 0.0161에서 R=8의 0.00100까지 단조 감소한다. D_R(21)은 열 개 모델 점추정 모두 양수였고 아홉 개의 95% 구간이 0을 배제했다. 세 번째 결과는 스칼라 왜곡이 손실을 결정하지 않는다는 것이다. Qwen2.5-1.5B의 56개 개입에서 log ΔNLL과 log 평균 row-wise JSD는 기울기 1.05, R²=0.945로 맞지만, 같은 JSD로 보정한 평탄화와 날카롭게 하기는 모델별로 부호가 갈린다. 일곱 개 모델에서는 평탄화가 더 해롭고, Llama-3.2-1B(D_T=-0.00313)와 Qwen2.5-72B(-0.00119)에서는 덜 해로우며 Llama-3.1-70B는 미해결이었다. 모델 크기 순서를 따르지도 않는다.
전제와 한계
이 관찰이 제안 방법으로 이어진다. Rowmax-PoT는 각 행 최댓값에 고정한 옥타브 격자 위에서 지수 가중치를 유지하되, p̂_j ∝ 2^(-d_j), d_j = min{Kmax, floor(Δ_j/ln2 + 1/2)}로 로그 가중치를 정수로 반올림한다. Kmax=20에서 열 개 중 아홉 모델이 perplexity 증가 1% 미만이고 Llama-3.1-70B만 1.513%로 예외다. Rowmax-H15는 이를 FlashAttention-4에 특수화한 것으로, exp2를 옥타브당 두 값 {1, 1.5}×2^k로 대체한다. n = round_RNE(2x)로 반옥타브 단위 반올림한 뒤 ŵ(x) = (1 + 0.5(n mod 2))·2^(floor(n/2))를 계산하며, 커널 안에서는 FP32 덧셈 하나와 22비트 시프트로 끝난다. 기준점은 online-softmax 루프가 추적하는 running row maximum이라 별도로 캘리브레이션하는 스케일이 없다. sm_86 마이크로커널에 대한 정적 SASS 감사에서 원소당 3개 인스트럭션(깊이 3)으로, FA4의 큐빅 에뮬레이션 8개(깊이 8)보다 적다. 저자들은 이 구성이 Schraudolph(1999)의 근사를 반정수로 반올림한 지수 위에서 평가한 것과 이상적 연산자 수준에서 일치한다고 밝힌다.
커널 성능은 B200 한 장, batch 1, 16 heads, head dimension 128 조건에서 측정했다. 패치한 FP8 causal attention forward의 호스트 측 call latency는 4K에서 7.1%, 8K에서 12.4%, 16K에서 12.5% 빨라졌고, non-causal 8K에서는 25.8% 빨라졌다. causal 8K 기준 BF16은 5.4%, FP16은 5.3%로 텐서코어 비율이 두 배인 FP8의 절반 수준이다. 에너지는 causal 16K에서 forward당 board energy가 742.5mJ에서 680.3mJ로 8.4% 줄었고 평균 board power는 그대로였다. FA4 자체 설정과 비교한 별도 스윕에서는 기본 설정보다 12.3%, 테스트한 여덟 개 stock 에뮬레이션 설정 중 가장 빠른 것보다 9.3% 빨랐다. 다만 2K에서는 두 측정 모두 느려졌고, 1K는 세션 간에 부호가 바뀌었으며, 32K에서는 이득이 9.9%로 줄었다.
모델 정확도는 속도와 분리해서 측정했다. BF16 경로에서 세 패밀리 다섯 모델, 97개 paired 2K 토큰 블록, stock FA4를 베이스라인으로 perplexity가 0.091%에서 0.492% 올랐고 다섯 개 95% 구간 모두 0 위에 있었다. 가장 큰 상한은 Llama-3.1-70B의 0.583%다. FP8 경로는 Qwen2.5-1.5B 시뮬레이션만 수행했으며 ΔNLL 0.000716(perplexity +0.072%)이었다. 직접 FP8 실행은 BF16 활성값의 per-tensor 양자화 오차가 지배해 사용할 수 없었다고 밝힌다. 참고로 EXAQ도 행 최댓값에 기준점을 두고 있고, 동시 연구인 EFQ-Softmax는 {1, 1.5}×2^k 격자가 같지만 기준점(마이크로스케일링 블록 최댓값 대 어텐션 행 최댓값), 반올림 방식, 파라미터 탐색, 측정 범위가 다르다.
저자들이 명시한 한계는 분명하다. 결과는 frozen inference와 주로 WikiText-103 NLL에 관한 것이고, 커널 정확도는 8K·16K 타이밍 shape가 아니라 2K에서 BF16 다섯 모델로 측정했으며 FP8 정확도는 시뮬레이션이다. 타이밍은 B200 어텐션 forward만 다루고 full-model, decode, backward, B300/GB300은 포함하지 않는다. head dimension 64에서는 FP8 이득이 2.0~3.7%에 그쳤고, 성능 카운터가 없어 병목 설명은 추론에 머문다. 사전 등록한 "BF16 이득 3% 미만" 예측은 causal 8K에서 5.4%가 나오며 실패했다. 실무적으로 이 논문이 주는 교훈은 명확하다. 소프트맥스 근사는 커널이 얼마나 빨라지는지와 frozen 모델의 손실이 얼마나 늘어나는지를 따로 측정해야 하며, 같은 평균 왜곡이라도 모델에 따라 손실의 부호가 뒤집힐 수 있으므로 자기 모델에서 NLL을 직접 확인해야 한다.