[flashinfer] FlashInfer CAKE VSA 커널 최적화: 중복 제거와 비효율적인 CuTe DSL 제거로 얻은 코드 효율성

PR 링크: flashinfer-ai/flashinfer#6719 상태: Merged | 변경: +632 / -128250

들어가며FlashInfer는 고성능 딥러닝 추론을 위한 최적화된 커널 라이브러리로, 특히 Transformer 기반 모델의 Attention 메커니즘에 중점을 둡니다. 이번에 분석할 PR은 FlashInfer의 CAKE Variable Block-Sparse Attention (VSA) 커널에 대한 대규모 리팩토링 및 최적화 작업입니다.이 PR이 해결하고자 하는 핵심 문제는 다음과 같습니다.첫째, 코드 중복성: FP16과 BF16 커널 소스가 거의 동일한 내용으로 중복되어 있었고, SM90(Hopper) 아키텍처를 위한 런처 코드도 여러 파일에 걸쳐 중복되었습니다.둘째, 데드 코드: 호출되지 않는 호스트 코드와 오래된 설정 파일들이 코드베이스에 남아 있었습니다.셋째, 비효율적인 실험적 구현: CuTe DSL(Domain Specific Language)을 사용하여 SM90 커널을 구현한 실험적인 코드가 존재했지만, 기존 CUDA 구현 대비 성능 이점이 없고 오히려 호스트 오버헤드가 더 컸습니다.이러한 문제들은 코드 유지보수성을 저해하고, 빌드 복잡성을 증가시키며, 불필요한 코드 크기를 유발했습니다. 이번 PR은 이러한 문제들을 해결하여 코드베이스를 더욱 견고하고 효율적으로 만드는 것을 목표로 합니다.## 코드 분석### 1. ddae7c31c : FP16 block-128 커널을 BF16 소스로 통합 (CAKE_VSA_FP16 가드 사용)이 커밋의 핵심은 csrc/cake_vsa/cake_vsa_blk128_fp16_compact.cu 파일의 내용을 csrc/cake_vsa/cake_vsa_blk128_compact.cu 파일로 통합하고, #if CAKE_VSA_FP16 전처리기 가드를 사용하여 FP16과 BF16 구현을 하나의 소스 파일에서 관리하도록 변경한 것입니다. 이전에는 두 개의 거의 동일한 파일이 존재하여 유지보수 시 동일한 변경 사항을 두 번 적용해야 하는 비효율성이 있었습니다.#### Before (개념적):cpp// csrc/cake_vsa/cake_vsa_blk128_fp16_compact.cu__global__ __launch_bounds__(512, 1) voidkernel_flashinfer_blackwell_vsa_blk128_fp16_compact_sm100(...) { __half* q0_smem = reinterpret_cast<__half*>(smem_raw + 1024); // ... FP16 specific logic ...}``````cpp// csrc/cake_vsa/cake_vsa_blk128_compact.cu (BF16)__global__ __launch_bounds__(512, 1) voidkernel_flashinfer_blackwell_vsa_blk128_seed_sm100(...) { __nv_bfloat16* q0_smem = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); // ... BF16 specific logic ...}#### After (통합된 파일 csrc/cake_vsa/cake_vsa_blk128_compact.cu):```diff--- a/csrc/cake_vsa/cake_vsa_blk128_compact.cu+++ b/csrc/cake_vsa/cake_vsa_blk128_compact.cu@@ -490,7 +490,11 @@ device forceinline uint32_t make_warp_uniform(uint32_t val) { extern

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH__HALF.html
• https://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH__BFLOAT16.html
• https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__OCCUPANCY.html#group__CUDA__OCCUPANCY_1g50020104646549298282367258481358

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