본문으로 건너뛰기

[flashinfer] Blackwell 아키텍처를 위한 FlashInfer의 FP8/BF16 혼합 정밀도 Attention 최적화 분석

PR 링크: flashinfer-ai/flashinfer#5900 상태: Merged | 변경: +695 / -140

들어가며

최신 LLM 추론 라이브러리인 FlashInfer는 하드웨어의 한계를 밀어붙이는 최적화로 잘 알려져 있습니다. 이번 분석에서 다룰 PR은 NVIDIA의 차세대 아키텍처인 Blackwell(SM100/SM103)을 타겟으로, QK-BF16/PV-FP8 및 All-FP8 데이터 타입에 대한 Dense Context FMHA(Fused Multi-Head Attention) 커널의 성능을 비약적으로 향상시킨 사례입니다.

이 PR의 핵심 과제는 Blackwell의 새로운 하드웨어 유닛인 MUFU(Multi-Function Unit)와 TMA(Tensor Memory Accelerator)를 효율적으로 활용하여, 소프트맥스 계산과 MMA(Matrix Multiply-Accumulate) 연산 사이의 의존성 병목을 제거하는 것입니다.

코드 분석: 핵심 변경 사항

1. P-in-SMEM: 소프트맥스 결과의 SMEM 스테이징

기존 구조에서는 Softmax 결과인 $P$를 TMEM(Tensor Memory)에 직접 썼기 때문에, 다음 타일의 $QK$ 연산이 이전 타일의 $PV$ 연산이 끝날 때까지 기다려야 하는 의존성이 있었습니다. 이를 해결하기 위해 $P$를 SMEM(Shared Memory)에 스테이징하여 파이프라인을 분리했습니다.

[Before]

# 기존에는 P를 TMEM 스테이지에 직접 기록하여 의존성 발생
# Softmax(N) -> PV(N) -> QK(N+1) 순차적 대기

[After]

# flashinfer/attention/prims_ts/kernels/fmha_context/fmha_kernel.py

def _uses_smem_p(cfg: FmhaConfig, *, has_variable_window: bool) -> bool:
    """Stage fp8 P in SMEM for the dense query-paired D128 schedule..."""
    return (
        not cfg.single_qkv_instance
        and cfg.logical_head_dim_qk == 128
        and cfg.v_dtype.width == 8
        and cfg.o_dtype.width == 16
        and cfg.qk_mma_tiler[1] * cfg.v_dtype.width // 8 == 128
        and cfg.epi_tile[1] % 64 == 0
        and not has_variable_window
        and not cfg.causal_single_kv_tile
    )

이 변경을 통해 QK(N+1) 연산이 exp2(N)(소프트맥스 연산 중 일부)과 동시에 실행될 수 있게 되어, MMA 워프의 가동률이 극대화되었습니다.

2. Two-CTA UMMA 확장

Blackwell의 UMMA(User-Mode Memory Accelerator)를 활용하여 두 개의 CTA(Cooperative Thread Array)가 협력하는 방식을 FP8 데이터 타입까지 확장했습니다.

[Before]

# flashinfer/attention/prims_ts/context.py

and torch.finfo(geometry.qk_dtype).bits == 16

[After]

# flashinfer/attention/prims_ts/context.py

and torch.finfo(geometry.qk_dtype).bits in (8, 16)

기존에는 BF16(16-bit)에만 적용되던 Two-CTA 로직을 8-bit(FP8)까지 확장하여, Blackwell의 클러스터 스코프 배리어를 통한 데이터 공유 이점을 FP8 연산에서도 누릴 수 있게 되었습니다.

3. Lazy Correction 및 exp2 Offload

수치적 안정성을 유지하면서 연산량을 줄이기 위한 두 가지 기법이 도입되었습니다.

  • Lazy Correction: Row Max 값이 $2^8$ 이상 커지지 않는 한 $O$ (Output)의 리스케일링을 건너뜁니다.
  • exp2 Offload: B200에서 MUFU 유닛의 부하를 줄이기 위해 일부 exp2 연산을 FMA(Floating Point Multiply-Add) 파이프로 오프로드합니다.

왜 이게 좋은 최적화인가?

하드웨어 특화 파이프라이닝

Blackwell 아키텍처에서 소프트맥스 연산은 MUFU 유닛에 의해 바운드(Bound)되는 경향이 있습니다. 이 PR은 P-in-SMEM 기법을 통해 MUFU 연산(Softmax)과 MMA 연산(QK/PV)을 완벽하게 오버랩시켰습니다.

메모리 계층의 효율적 재분배

리뷰어 harrisonzhy의 지적에 따라, $D=128$ 환경에서 $P$를 SMEM에 넣기 위해 $O$ (Output) 스테이징을 64-wide 하프 타일로 나누어 저장하는 전략을 취했습니다. 이는 한정된 SMEM 자원을 쪼개어 파이프라인의 병목을 해결하는 고도의 자원 관리 기법입니다.

성능 수치 및 교훈

제공된 벤치마크 결과에 따르면, Wan2.2-T2V-A14B 모델의 Self-Attention 워크로드에서 B200 및 B300 모두 유의미한 성능 향상을 보였습니다. 특히 토큰 수가 많아질수록 최적화된 커널의 효율이 극대화되는 것을 확인할 수 있습니다.

일반적 교훈:

  1. 의존성 분리: 연산 유닛(MMA vs MUFU)이 다를 경우, 중간 결과를 SMEM 등에 스테이징하여 의존성을 끊으면 병렬성을 극대화할 수 있습니다.
  2. 정밀도 기반 게이팅: 모든 최적화가 모든 데이터 타입에 유리한 것은 아닙니다. 이 PR처럼 8-bit V 입력일 때만 특정 로직을 활성화하는 식의 정교한 게이팅이 필요합니다.

리뷰어 피드백 반영

리뷰 과정에서 harrisonzhy는 _uses_smem_p 조건이 $D=192$ 환경에서 문제를 일으킬 수 있음을 발견했습니다. 이에 따라 cfg.logical_head_dim_qk == 128 조건을 명시적으로 추가하여, 특정 레이아웃에서 발생할 수 있는 SMEM 용량 초과 및 파이프라인 빌드 오류를 방지했습니다. 또한 FP8 출력 시의 TMA 복사 단위(Granularity) 문제를 해결하기 위해 cfg.o_dtype.width == 16 조건을 추가하여 안정성을 확보했습니다.

결론

이번 FlashInfer의 업데이트는 단순히 새로운 하드웨어를 지원하는 것을 넘어, Blackwell의 독특한 메모리 및 연산 구조를 깊이 있게 이해하고 이를 소프트웨어 아키텍처에 녹여낸 훌륭한 사례입니다. FP8 추론이 대세가 되는 시점에서 이러한 저수준 최적화는 실제 서비스의 처리량(Throughput)에 직접적인 영향을 미칩니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글