본문으로 건너뛰기

[flashinfer] [FlashInfer] SM120 FP8 Attention 최적화: Lazy Output Rescale를 통한 성능 향상

PR 링크: flashinfer-ai/flashinfer#7871 상태: Merged | 변경: +272 / -269

들어가며

LLM 추론 가속화 라이브러리인 FlashInfer는 최신 하드웨어 아키텍처인 NVIDIA Blackwell(SM120)을 지원하기 위해 지속적으로 커널을 최적화하고 있습니다. 이번에 분석할 PR은 MiniMax-H3 모델에서 사용되는 varlen_attention 커널, 그중에서도 FP8 형식을 사용하는 minimax_h3_sm120_varlen_attention_fp8의 성능 개선에 관한 것입니다.

이 PR의 핵심 아이디어는 'Lazy Output Rescale'입니다. Attention 메커니즘의 핵심인 Online Softmax 과정에서, 새로운 K-tile을 처리할 때마다 출력값(Accumulator)을 재조정(Rescale)해야 하는데, 만약 새로운 데이터가 기존의 최댓값(Row Maximum)을 갱신하지 않는다면 이 재조정 과정은 수학적으로 x * 1.0과 동일합니다. 이 PR은 이러한 불필요한 연산을 Warp 단위의 투표(Vote)를 통해 건너뜀으로써 성능을 향상시켰습니다.

코드 분석: Online Softmax의 최적화

1. 변수 명명 규칙 및 구조 개선

기존 코드에서는 _max_0, _max_1과 같이 추상적인 이름을 사용했으나, 변경된 코드에서는 max_k_0, max_k_1과 같이 K-tile 처리 과정에서의 최댓값임을 명확히 드러내도록 리팩토링되었습니다. 이는 코드 가독성을 높일 뿐만 아니라, 이후 추가된 Lazy Rescale 로직의 의도를 명확히 합니다.

Before:

// 기존의 추상적인 변수명
float _max_0 = fmaxf(s_acc[0], s_acc[1]);
float _max_1 = fmaxf(s_acc[2], s_acc[3]);
float _max_2 = fmaxf(_max_0, fmaxf(s_acc[4], s_acc[5]));
// ... 중략 ...
float _shfl_xor_sync_0 = __shfl_xor_sync(0xFFFFFFFF, _max_30, 1);
float _max_32 = fmaxf(_max_30, _shfl_xor_sync_0);

After:

// K-tile의 최댓값임을 명시하는 변수명으로 변경
float max_k_0 = fmaxf(s_acc[0], s_acc[1]);
float max_k_1 = fmaxf(s_acc[2], s_acc[3]);
float max_k_2 = fmaxf(max_k_0, fmaxf(s_acc[4], s_acc[5]));
// ... 중략 ...
// 이후 이 max_k 값들이 기존 max 값을 갱신하는지 체크하는 로직이 추가됨

2. Lazy Output Rescale 로직 (핵심 변경 사항)

PR 설명에 따르면, 핵심 로직은 각 Warp가 투표를 수행하는 부분에 있습니다. Online Softmax에서 출력 $O$는 다음과 같이 업데이트됩니다: $$O_{new} = O_{old} \times e^{m_{old} - m_{new}} + \text{new_contribution}$$ 여기서 $m$은 현재까지의 최댓값입니다. 만약 $m_{old} == m_{new}$라면, $e^0 = 1.0$이 되어 곱셈 연산이 무의미해집니다.

이 PR은 다음과 같은 논리를 적용했습니다:

  1. 각 Warp 내의 모든 스레드가 현재 K-tile에서 Row Maximum이 변경되었는지 확인합니다.
  2. __any_sync 또는 Warp 단위의 비교를 통해, Warp 내 어떤 행도 최댓값이 변하지 않았다면(rescale_factor == 1.0), 해당 Warp는 O에 대한 곱셈 연산을 통째로 건너뜁니다.
  3. x * 1.0 == x이므로, 이 최적화는 결과값에 비트 단위의 오차도 발생시키지 않습니다(Bitwise Equal).

왜 이게 좋은가?

1. 불필요한 산술 연산 제거 (Arithmetic Efficiency)

GPU 커널에서 exp 연산과 multiply 연산은 비용이 발생합니다. 특히 Attention 커널처럼 메모리 대역폭과 연산량이 모두 중요한 경우, 조건부로 연산을 생략하는 것은 처리량(Throughput) 향상에 직접적인 도움이 됩니다.

2. 성능 수치 (Speedup)

NVIDIA RTX PRO 6000 Blackwell 및 RTX 5090 환경에서 테스트된 결과, 긴 시퀀스 길이를 가진 데이터셋에서 유의미한 성능 향상이 관찰되었습니다.

  • center_15s (109,952 tokens): 두 GPU 모두에서 1.017x의 속도 향상을 기록했습니다.
  • center_10s (74,240 tokens): 1.018x (RTX PRO 6000) 속도 향상을 보였습니다.

약 1.5% ~ 1.8%의 성능 향상은 작아 보일 수 있지만, LLM 서빙 환경에서 Attention 커널이 차지하는 비중을 고려할 때 이는 매우 값진 최적화입니다. 특히 데이터 의존적인(Data-dependent) 최적화임에도 불구하고 최악의 경우에도 성능 저하가 거의 없다는 점이 강점입니다.

3. 일반적인 교훈

이 최적화는 "항상 수행되는 연산 중 항등원(Identity element) 연산이 발생하는 케이스가 있는가?"를 검토하는 것의 중요성을 보여줍니다. Online Softmax는 수렴성이 좋아 시퀀스가 길어질수록 최댓값이 변하지 않는 구간이 자주 발생하는데, 이 특성을 정확히 파고든 사례입니다.

결론

FlashInfer의 이번 PR은 Blackwell 아키텍처에서 FP8 Attention의 효율성을 한 단계 더 끌어올렸습니다. 수학적 동일성을 유지하면서도 하드웨어의 특성과 알고리즘의 동작 방식을 결합하여 실질적인 성능 이득을 얻어낸 훌륭한 엔지니어링 사례입니다. 특히 Warp 단위의 협력적 판단을 통해 분기 예측 실패의 리스크를 최소화하며 최적화를 달성했다는 점이 인상적입니다.

앞으로도 이러한 'Lazy' 기법은 다양한 커널 최적화에서 중요한 전략으로 활용될 것입니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글