본문으로 건너뛰기

[flashinfer] FlashInfer, RTX PRO 6000 Blackwell을 위한 Warp-Specialised Attention 커널로 성능 향상

PR 링크: flashinfer-ai/flashinfer#6060 상태: Merged | 변경: +6142 / -566

들어가며

최근 인공지능 모델, 특히 트랜스포머 아키텍처의 발전은 연산 집약적인 Attention 메커니즘의 효율적인 구현을 요구하고 있습니다. NVIDIA의 최신 GPU 아키텍처, 특히 RTX PRO 6000 Blackwell과 같은 고성능 하드웨어에서는 이러한 연산을 더욱 가속화할 수 있는 기회가 있습니다. 본 PR은 FlashInfer 라이브러리에서 cake_minimax_h3_sm120_nvfp4_varlen_attention_sm120a.cu 파일을 수정하여, 특정 조건에서 Warp-level 최적화를 적용한 새로운 Attention 커널을 도입함으로써 이러한 기회를 활용합니다. 이 변경은 특히 긴 시퀀스와 대규모 배치에서 Attention 연산의 성능을 향상시키는 것을 목표로 합니다.

코드 분석

이번 PR의 핵심은 csrc/cake_minimax_h3_sm120_nvfp4_varlen_attention_sm120a.cu 파일에 새로운 커널과 이를 선택하는 로직을 추가한 것입니다. 기존 커널은 그대로 유지하면서, 특정 하드웨어 및 입력 조건에 따라 더 효율적인 새로운 커널을 호출하도록 변경되었습니다.

1. 새로운 Warp-Specialised 커널 도입

PR은 두 개의 새로운 CUDA 커널을 도입했습니다. 이 커널들은 기존 커널과 동일한 FP32 연산을 동일한 순서로 수행하지만, Warp 수준에서의 최적화에 초점을 맞추고 있습니다. 구체적으로는 8개의 QK (Query-Key) 워프와 4개의 PV (Projection-Value) 워프를 사용하는 구조로, 총 384개의 스레드를 활용하여 하나의 CTA (Cooperative Thread Array)가 하나의 SM (Streaming Multiprocessor)을 차지하도록 설계되었습니다.

기존 커널의 __launch_bounds__(256, 2) 설정이 새로운 커널에서는 __launch_bounds__(256, LAUNCH_MIN_BLOCKS)로 변경되었으며, LAUNCH_MIN_BLOCKS는 2로 정의되어 있습니다. 이는 CTA 당 스레드 수를 256개로 고정하고, SM 당 CTA 수를 최대 2개로 제한하는 설정을 유지하면서도, 새로운 Warp-specialised 구조를 적용할 수 있도록 합니다.

Before:

-__global__ __launch_bounds__(256, 2) void
+__global__ __launch_bounds__(256, LAUNCH_MIN_BLOCKS) void

또한, #undef LAUNCH_MIN_BLOCKS와 같은 매크로 정의들이 각 커널의 끝부분에 추가되어, 해당 커널에서 사용된 매크로들이 정리되도록 했습니다. 이는 코드의 명확성을 높이고 잠재적인 매크로 충돌을 방지합니다.

2. Host-side Selector 로직 추가

새로운 커널의 성능 이점을 모든 상황에 적용하는 대신, PR은 호스트 측에서 특정 조건에 따라 새로운 커널을 선택하는 로직을 추가했습니다. 이는 minimax_h3_sm120_varlen_attention_nvfp4 및 _nodelta 엔트리 포인트에서 구현됩니다.

새로운 커널은 평균 세그먼트 길이 (tokens / segments)가 1536 토큰 이상이고, 사용 중인 GPU의 SM 수가 188개 이상일 때 (예: RTX PRO 6000 Blackwell) 선택됩니다. 이 조건을 만족하지 않는 경우에는 기존 커널이 그대로 사용됩니다. 이 조건부 로직은 다음과 같은 형태로 구현될 것으로 예상됩니다 (정확한 코드는 diff에 포함되지 않았지만, PR 설명에 기반합니다):

// pseudo-code
if (mean_segment_length >= 1536 && num_sms >= 188) {
    launch_new_kernel(...);
} else {
    launch_existing_kernel(...);
}

이러한 선택적 로직은 하드웨어 및 입력 데이터 특성에 따라 최적의 성능을 보장하면서도, 모든 경우에 새로운 커널을 적용했을 때 발생할 수 있는 잠재적인 성능 저하나 호환성 문제를 방지합니다. 또한, 새로운 커널은 기존 커널과 비트-정확도(bit-exact)를 보장하므로, 수치적 안정성 측면에서도 동일한 결과를 제공합니다.

3. 코드 정리 및 재사용성 향상

PR은 csrc/cake_minimax_h3_sm120_host.cuh 헤더 파일과 LAUNCH_MIN_BLOCKS 매크로와 같은 호스트-헬퍼 공유 및 생성기 디바이스 헬퍼를 활용합니다. 이는 코드 중복을 줄이고 유지보수성을 향상시키는 일반적인 소프트웨어 엔지니어링 모범 사례입니다.

왜 이게 좋은가?

성능 향상

RTX PRO 6000 Blackwell 서버 에디션에서 수행된 측정 결과에 따르면, 새로운 Warp-specialised 커널은 다양한 시퀀스 길이에서 상당한 성능 향상을 보여줍니다. 특히 seg4 (0, 8368, 16736, 25104, 33472) 형태에서는 1.0485배 (약 4.85%)의 연산 속도 향상을 기록했으며, single_4096 (0, 4096) 형태에서는 1.0564배 (약 5.64%)의 향상을 보였습니다. 이는 Attention 연산의 주요 병목 현상을 효과적으로 완화했음을 시사합니다.

Shape (cu_seqlens) existing kernel op ms new selector op ms operator speedup (min..max of 3 paired groups)
seg4 (0,8368,16736,25104,33472) 11.805 11.259 1.0485 (1.0480..1.0506)
single_4096 (0,4096) 0.792 0.750 1.0564 (1.0564..1.0574)

이러한 성능 향상은 특히 대규모 언어 모델(LLM)에서 Attention 연산이 차지하는 비중이 크기 때문에, 전체 모델의 추론 속도 향상에 직접적으로 기여할 수 있습니다. 또한, seg_empty와 같이 평균 세그먼트 길이가 짧은 경우에는 기존 커널이 그대로 사용되므로, 성능 저하 없이 일관된 성능을 유지합니다.

일반적 교훈

  1. 하드웨어 특화 최적화: 최신 GPU 아키텍처의 기능을 최대한 활용하기 위해 Warp 수준의 최적화를 적용하는 것은 성능 향상의 핵심입니다. 특히 Attention과 같이 병렬성이 높은 연산에서 이러한 접근 방식은 효과적입니다.
  2. 조건부 커널 선택: 모든 입력에 대해 단일 최적화된 커널을 사용하는 것보다, 입력 데이터의 특성(예: 시퀀스 길이, 배치 크기)과 하드웨어 사양(예: SM 수)에 따라 최적의 커널을 동적으로 선택하는 것이 더 나은 성능과 안정성을 제공할 수 있습니다.
  3. 비트-정확도 유지: 성능 향상을 추구하더라도 기존 커널과의 비트-정확도를 유지하는 것은 수치적 안정성과 결과의 신뢰성을 보장하는 데 필수적입니다. 이는 기존 모델과의 호환성을 유지하는 데도 중요합니다.
  4. 코드 재사용성 및 모듈화: 호스트 헬퍼 및 공유 디바이스 헬퍼를 활용하여 코드 중복을 줄이고 모듈성을 높이는 것은 장기적인 유지보수와 개발 효율성을 향상시킵니다.

결론

이번 PR은 FlashInfer 라이브러리의 성능을 한 단계 끌어올리는 중요한 개선 사항입니다. RTX PRO 6000 Blackwell과 같은 최신 GPU에서 Warp-specialised Attention 커널을 도입하고, 입력 조건에 따라 이를 선택적으로 사용하는 전략은 실제적인 성능 향상으로 이어졌습니다. 이는 LLM 개발자들이 더 빠르고 효율적인 모델을 구축하는 데 기여할 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글