본문으로 건너뛰기

[flashinfer] FlashInfer: BF16 어텐션 최적화 - SM100a의 P Barrier 개선 및 Tail-round Q-split 도입

PR 링크: flashinfer-ai/flashinfer#6258 상태: Merged | 변경: +2195 / -2735

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 Transformer 아키텍처의 핵심인 어텐션 메커니즘의 성능 최적화는 매우 중요한 과제가 되었습니다. 특히 BF16(bfloat16)과 같이 정밀도를 낮춘 데이터 타입을 활용하여 메모리 사용량을 줄이고 연산 속도를 높이려는 시도가 활발히 이루어지고 있습니다.

이번 PR은 FlashInfer 라이브러리에서 BF16 Packed-Varlen Attention의 성능을 더욱 향상시키기 위한 두 가지 주요 개선 사항을 다룹니다. 첫째, SM100a 아키텍처에 특화된 P barrier 동기화 방식을 최적화하여 불필요한 대기 시간을 줄였습니다. 둘째, Tail-round Q-split이라는 새로운 플래닝 기법을 도입하여, 연산량이 많은 일부 유닛을 더 작은 단위로 분할함으로써 전체적인 처리 시간을 단축했습니다. 이 글에서는 이러한 변경 사항들이 무엇이며, 왜 성능 향상으로 이어지는지 상세히 분석해보겠습니다.

코드 분석

이번 PR은 주로 flashinfer.experimental.minimax_h3_varlen_attention 모듈 내의 코드와 관련 설정을 변경했습니다. 핵심적인 변경 사항은 다음과 같습니다.

1. SM100a의 P Barrier 최적화

기존의 SM103a 아키텍처에서는 Softmax 연산 후 확률 타일을 발행하기 위해 두 번의 wait 명령을 사용했습니다. 하지만 SM100a 아키텍처에서는 이 과정을 한 번의 wait로 최적화하여 효율성을 높였습니다. 이는 특히 긴 시퀀스 길이를 처리하는 경우 성능 향상에 기여합니다.

Before:

- # SM103a 프로그램은 두 번의 wait 방식을 유지하며, 해당 아키텍처에서 더 나은 성능을 보였습니다.
+ # SM103a program keeps the two-wait form, which measured better there.

After:

- **SM100a single P barrier per PV group**: the softmax publishes the 128-key probability tile once (one `tcgen05.wait::st` and one arrive after the four 32-key stores) and the MMA warp runs each PV group as one `k(0, 128)` call behind one wait.
+ **SM100a single P barrier per PV group**: the softmax publishes the 128-key probability tile once (one `tcgen05.wait::st` and one arrive after the four 32-key stores) and the MMA warp runs each PV group as one `k(0, 128)` call behind one wait.

이 변경은 SM100a 아키텍처의 특성을 활용하여, Softmax 연산의 결과(확률 타일)를 기다리는 동안 발생하는 GPU 워프(warp)의 유휴 시간을 최소화합니다. 기존에는 4개의 32-키 스토어 이후에 wait가 발생했지만, 이제는 128-키 타일 전체를 한 번에 발행하고 단 한 번의 wait를 통해 후속 연산을 진행합니다. 이는 특히 center_10s_p2, center_15s_p2, center_6s_p1과 같이 긴 두 단계(two-stage) 행에서 1.010 ~ 1.029 배의 성능 향상을 가져왔습니다.

2. Tail-round Q-split 도입

두 번째 주요 개선 사항은 cached_bf16_segment_plan 함수에서 이루어진 'Tail-round Q-split'입니다. 이는 Persistent Grid의 Tail Round에서 발생하는 클러스터 유휴 시간을 줄이기 위한 기법입니다.

기존에는 Tail Round에서 연산할 유닛 수가 클러스터 수(num_clusters)로 나누어 떨어지지 않으면, 일부 클러스터는 다음 유닛을 기다리며 유휴 상태가 되었습니다. 이 PR에서는 이러한 유휴 시간을 활용하기 위해, 가장 비용이 많이 드는 두 단계(two-stage) 유닛을 두 개의 256-행 단일 단계(single-stage) 레코드로 분할하는 옵션을 도입했습니다. 이 분할은 전체적인 시뮬레이션된 Makespan(전체 작업 완료 시간)이 최소 1% (이후 0.5%로 조정됨) 개선될 때만 적용됩니다.

Before (Conceptual):

- ...the partial slot (-1 for an unsplit unit) -- in slot order: unit ``u`` of the table runs on cluster ``u % num_clusters`` as its ``u // num_clusters``-th unit (see :func:`assign_unit_slots`).  ``total_tiles`` counts scheduled units (split ranges included); ``combine_table`` holds ``COMBINE_WORDS`` int32
+ ...the partial slot (-1 for an unsplit unit) -- in slot order: unit ``u`` of the table runs on cluster ``u % num_clusters`` as its ``u // num_clusters``-th unit (see :func:`assign_unit_slots`).  ``total_tiles`` counts scheduled units (split ranges included); ``combine_table`` holds ``COMBINE_WORDS`` int32

After (Conceptual):

- ...the partial slot (-1 for an unsplit unit) -- in slot order: unit ``u`` of the table runs on cluster ``u % num_clusters`` as its ``u // num_clusters``-th unit (see :func:`assign_unit_slots`).  ``total_tiles`` counts scheduled units (split ranges included); ``combine_table`` holds ``COMBINE_WORDS`` int32
+ ...the partial slot (-1 for an unsplit unit), ``q_half`` (0 for a whole 512-row unit, 1 / 2 for the first / second 256-row half of a tail-round Q-split unit, run as a single-stage unit) and two reserved words -- in slot order: unit ``u`` of the table runs on cluster ``u % num_clusters`` as its ``u // num_clusters``-th unit (see :func:`assign_unit_slots`), followed by ``num_clusters`` zero records (the kernel's one-unit-ahead record prefetch).  ``total_tiles`` counts scheduled units (K/V ranges and Q-split halves included; ``num_q_split_units`` counts the halved units); ``combine_table`` holds ``COMBINE_WORDS`` int32

cake_backend.py 파일에서는 BF16_Q_SPLIT_PLAN 상수가 True로 설정되었고, BF16_Q_SPLIT_MIN_GAIN은 0.01에서 0.005로 낮아져 더 적은 성능 향상으로도 분할이 시도될 수 있게 되었습니다. 이 변경은 README.md 파일에서도 상세히 설명되어 있으며, unit_table의 5번째 워드에 q_half 필드가 추가되어 분할된 유닛 정보를 나타냅니다.

이 Q-split 기법은 다음과 같은 특징을 가집니다:

  • 정확도 유지: 분할된 각 행은 원래 유닛의 K/V 범위와 출력 경로를 그대로 유지하므로, 최종 출력은 원래 계획과 비트 단위로 동일합니다. 즉, 정확도 손실이 없습니다.
  • K/V-split과의 상호 배제: Q-split 유닛은 K/V-split 유닛과 겹치지 않으므로, Combine 경로에는 영향을 주지 않습니다.
  • 계획 캐시 키 변경: cached_bf16_segment_plan의 캐시 키에 q_split 옵션이 추가되어, Q-split 적용 여부에 따라 다른 계획이 캐싱됩니다.

이 Tail-round Q-split은 B200 (SM100a) 및 GB300 (SM103a)에서 상당한 성능 향상을 보였습니다. 특히 B200에서는 center_8s_p4 (1.041 / 1.043 배), center_15s_p4 (1.030 / 1.032 배) 등에서, GB300에서는 center_6s_p4 (1.077 / 1.077 배) 등에서 높은 성능 향상이 측정되었습니다.

또한, PR 업데이트에서 BF16_Q_SPLIT_MIN_GAIN 상수가 0.01에서 0.005로 조정되면서, 더 작은 성능 개선(0.5% 이상)으로도 Q-split이 적용될 수 있게 되었습니다. 이로 인해 DiT(Diffusion Transformer) 관련 행에서 GB300은 1.1-1.8%, B200은 1.007% 이상의 성능 향상을 추가로 얻었습니다.

왜 이게 좋은가?

이번 PR의 변경 사항들은 다음과 같은 이유로 좋은 최적화라고 할 수 있습니다.

  1. 하드웨어 특성 활용: SM100a 아키텍처의 P barrier 동기화 방식을 최적화한 것은 특정 하드웨어의 성능 특성을 깊이 이해하고 이를 최대한 활용한 결과입니다. 이는 불필요한 동기화 오버헤드를 제거하여 GPU 워프의 효율성을 극대화합니다.
  2. 유휴 시간 활용: Tail-round Q-split은 GPU 연산에서 흔히 발생하는 유휴 시간을 단순히 버리는 것이 아니라, 이를 활용하여 더 작은 단위로 작업을 분할하고 전체 처리량을 늘리는 창의적인 방법입니다. 이는 특히 Tail Round와 같이 불균일한 작업 부하가 발생하는 경우에 효과적입니다.
  3. 성능 향상 수치: 측정 결과에 따르면, SM100a의 P barrier 개선은 최대 1.029배, Tail-round Q-split은 최대 1.077배 (이후 0.005 gain 적용 시 추가 향상)의 성능 향상을 가져왔습니다. 이는 실제 워크로드에서 상당한 속도 개선을 의미합니다.
  4. 정확도 보장: Q-split 기법은 연산 방식을 변경하지만, 최종 출력의 정확도는 원본과 동일하게 유지됩니다. 이는 성능 최적화 과정에서 가장 중요한 요소 중 하나인 정확성 손실을 방지합니다.
  5. 일반화 가능한 교훈:
    • 아키텍처별 최적화: 특정 GPU 아키텍처(SM100a vs SM103a)의 차이를 인지하고 각기 다른 최적화 전략을 적용하는 것은 중요합니다.
    • 동적 플래닝 및 분할: 작업 부하의 불균일성을 감지하고, 이를 해소하기 위해 동적으로 작업을 분할하는 기법은 다양한 병렬 컴퓨팅 환경에서 유용하게 적용될 수 있습니다.
    • 정밀도와 성능의 균형: BF16과 같은 저정밀도 데이터 타입을 활용하는 것은 성능 향상의 중요한 축이지만, 이와 동시에 정확도를 유지하는 것이 필수적입니다.

리뷰 피드백 반영

주요 리뷰어인 yyihuang 님은 여러 차례 테스트 실행을 요청하며 코드의 안정성과 정확성을 꼼꼼히 검증했습니다. 특히 tests/experimental/test_cake_minimax_h3_varlen_attention.py 테스트 스위트의 통과 여부를 지속적으로 확인한 것은, 새로운 Q-split 플래너가 기존의 다양한 시나리오(혼합 KV-split/Q-split 상호 배제, 기본 계획 경로, 서브 웨이브 그리드 성장, 반쪽 레코드 상속 등)에서도 올바르게 작동하는지 확인하는 데 중요한 역할을 했습니다. 이러한 테스트 중심의 검증 과정은 코드 변경의 신뢰도를 높이는 데 기여했습니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글