[flashinfer] FlashInfer의 SM100/SM103 최적화: CAKE 기반 블록 희소 어텐션(VSA) 도입

PR 링크: flashinfer-ai/flashinfer#4593 상태: Merged | 변경: +35271 / -19

들어가며

최신 GPU 아키텍처인 Blackwell(SM100/SM103) 환경에서 어텐션 연산의 효율성은 모델 추론 속도를 결정짓는 핵심 요소입니다. 이번 PR에서는 CAKE IR을 활용하여 블록 희소 어텐션(Block-Sparse Attention, VSA) 커널을 새롭게 구현했습니다. 기존 구현 대비 특정 시나리오에서 최대 3.3배 이상의 성능 향상을 달성하며, 특히 대규모 시퀀스 처리 시 메모리 대역폭과 연산 효율을 극대화했습니다.

코드 분석

1. CAKE 기반 호스트 런처 (csrc/cake_vsa/cake_vsa_blk128_compact_host.cpp)

이 파일은 CAKE 소스 익스포터에 의해 생성된 호스트 런처로, Blackwell 아키텍처의 특성을 활용한 동적 공유 메모리(Dynamic SMEM) 관리를 수행합니다.

Before (기존 방식): 기존에는 표준적인 MaxSharedMemoryPerBlockOptin 제한 내에서만 커널을 실행했습니다.

After (개선된 방식): CUDA 13.4의 새로운 'Oversized Shared Memory' 모드를 활용하여, 표준 제한을 초과하는 공유 메모리 사용을 가능하게 했습니다.

#if CAKE_HAS_OVERSIZED_SMEM
  int oversized_max = 0;
  err = cuda_api::GetDeviceAttribute(&oversized_max, cuda_api::DeviceAttrType(150), device);
  TVM_FFI_CHECK(err == cuda_api::kSuccess && oversized_max >= smem_bytes, RuntimeError)
      << "dynamic smem " << smem_bytes << " B exceeds the standard opt-in ceiling";
  cache[device_id] = 2;
  return true;
#endif

이 코드는 cudaDevAttrOversizedSharedMemoryPerBlock(속성 150)을 쿼리하여, 하드웨어 가용 범위 내에서 더 큰 공유 메모리 블록을 할당함으로써 복잡한 희소 연산의 레지스터 압박을 완화합니다.

2. 빌드 시스템 설정 (.pre-commit-config.yaml)

생성된 코드는 수동 수정이 불가능하므로, clang-format 대상에서 제외하여 코드 무결성을 유지합니다.

-          (?x)^(3rdparty/.* flashinfer/jit/aot_config.py)$
+          (?x)^(3rdparty/.*|flashinfer/jit/aot_config.py|csrc/cake_vsa/.*)$

왜 이게 좋은가

성능 개선 수치

• Canonical blk64 (S=1024): 기존 0.25ms에서 0.07ms로 약 3.38배의 속도 향상.
• FastWan (E2E): 복잡한 시나리오에서도 약 1.2~1.3배의 성능 개선을 보이며, 기하 평균적으로 약 1.8배의 성능 향상을 기록했습니다.

기술적 교훈

1. Hardware-Specific Tuning: Blackwell 아키텍처의 Oversized Shared Memory와 같은 최신 기능을 활용하면, 기존의 공유 메모리 병목을 획기적으로 돌파할 수 있습니다.
2. Code Generation: CAKE와 같은 IR을 통해 커널을 생성하면, 수동 최적화가 어려운 복잡한 인덱싱 로직을 효율적으로 관리할 수 있습니다.
3. TMA(Tensor Memory Accelerator) 활용: 코드 내 TmaDeviceSlot 구현에서 볼 수 있듯이, 비동기 메모리 복사를 위한 TMA 디스크립터 관리 최적화가 필수적입니다.

리뷰어 피드백

리뷰 과정에서 hsr1234563은 향후 cute-dsl 버전으로의 전환을 언급했으나, yzh119는 CAKE IR이 생성하는 코드의 독립성과 효율성을 강조하며, 두 접근 방식의 상호 보완적 발전을 제안했습니다. 이는 FlashInfer가 성능 최적화를 위해 다양한 컴파일러 기술을 실험적으로 도입하고 있음을 보여줍니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html
• https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__DEVICE__DEPRECATED.html

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