[flashinfer] FlashInfer Blackwell(SM103) 최적화: Sage 커널 및 Block-Sparse 성능 개선

PR 링크: flashinfer-ai/flashinfer#6014 상태: Merged | 변경: +3380 / -2164

들어가며

최근 FlashInfer 레포지토리에 NVIDIA Blackwell(SM103) 아키텍처를 타겟으로 한 대규모 커널 최적화 PR이 병합되었습니다. 이 작업은 특히 8-bit Sage 커널과 Block-Sparse 디코드 커널의 성능을 극대화하는 데 초점을 맞추고 있습니다. 기존 아키텍처 대비 Blackwell의 MUFU(Multi-Function Unit) 활용도를 높이고, 메모리 접근 패턴을 개선하여 특히 GB300 환경에서 괄목할 만한 성능 향상을 이끌어냈습니다.

코드 분석

1. Sage Scale 준비 로직의 분리 및 파이프라이닝

기존에는 Sage 커널이 호출될 때마다 스케일 값을 읽어오는 과정이 비효율적이었습니다. 이번 변경을 통해 스케일 값을 미리 하나의 버퍼로 정렬하여 읽어오는 전용 커널을 도입했습니다.

Before (기존 방식):

# Sage scale tensors are [heads, flat slots]
sage_kwargs = {
    "sage_q_scale_iter": q_scale.iterator,
    "sage_k_scale_iter": k_scale.iterator,
    "sage_v_scale_iter": v_scale.iterator,
}

After (최적화 방식):

# Attention reads every scale from the buffer the Sage prepare writes first
sage_kwargs = {"sage_scales_iter": sage_scales.iterator}
prepare_sage_scales(q_scale, k_scale, k_summary_scale, v_scale, v_mean, sage_scales, stream)

이 방식은 route prepare와 스케일 준비를 programmatic dependent launch로 연결하여, 어텐션 커널이 실행되기 전에 모든 데이터가 준비되도록 파이프라이닝을 최적화했습니다.

2. Softmax Rescale Threshold 도입

FP8 연산 시 Softmax의 수치적 안정성을 위해 softmax_rescale_threshold_log2 파라미터를 도입했습니다. 이를 통해 불필요한 리스케일링 연산을 줄이고 E4M3 포맷의 범위를 효율적으로 활용합니다.

# BlockSparseTSWrapper.plan(..., softmax_rescale_threshold_log2=8)
# 0 to 8, default 8. A row keeps its softmax anchor while a tile raises its maximum.

왜 이게 좋은가

이번 최적화의 핵심은 메모리 접근의 Coalescing과 커널 간 의존성 관리입니다.

1. 성능 수치: NVIDIA GB300(Blackwell Ultra) 환경에서 Sage FP8 기준, 기존 대비 약 20~25%의 커널 실행 시간 단축을 기록했습니다. 특히 VSA(Token-masked Block-sparse) 패턴에서 그 효과가 두드러집니다.
2. 일반적 교훈:
   • 데이터 레이아웃 최적화: 커널이 읽기 좋은 순서로 데이터를 미리 정렬(Pre-packing)하는 것은 고대역폭 메모리(HBM) 환경에서 필수적입니다.
   • 비동기 의존성 관리: programmatic dependent launch를 사용하여 커널 간의 유휴 시간을 최소화하는 것이 Blackwell과 같은 최신 아키텍처에서 성능을 뽑아내는 핵심입니다.

결론

이번 PR은 단순히 특정 커널의 속도를 높이는 것을 넘어, Blackwell 아키텍처의 특성을 반영한 새로운 커널 설계 패턴을 제시했습니다. 특히 SageAttention을 사용하는 대규모 언어 모델 추론 환경에서 큰 이점이 있을 것으로 기대됩니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async

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