[flashinfer] FlashInfer: DeepSeek-V4 NVFP4 sparse-MLA 디코딩을 위한 H64 타일 최적화

PR 링크: flashinfer-ai/flashinfer#6262 상태: Merged | 변경: +8834 / -814

들어가며

FlashInfer의 cake 백엔드에서 DeepSeek-V4 모델의 NVFP4 sparse-MLA 디코딩 성능을 개선하기 위한 최적화가 진행되었습니다. 기존에는 128행 Q 박스를 처리하는 범용적인 tile 멤버를 사용했으나, 헤드 수가 64개 이하인 경우 절반이 패딩으로 낭비되는 문제가 있었습니다. 본 PR은 이러한 워크로드에 특화된 H64 타일 멤버를 도입하여 SMEM(Shared Memory) 효율을 극대화하고 연산 대기 시간을 제거했습니다.

코드 분석

1. H64 타일 멤버 도입 (flashinfer/mla/cake_dsv4.py)

기존에는 모든 H <= 64 워크로드에 대해 128행 기반의 tile 멤버를 사용했습니다. 새로운 플래너는 이를 감지하여 전용 커널을 라우팅합니다.

# Before: 128-row tile member used for all cases
# After: H64 member selected for H <= 64
_nvfp4_plan = pick_tile_h64_if_eligible(H, ...)

2. SMEM 최적화 및 동기화 제거

기존 tile 멤버는 128행을 스테이징하면서 절반의 패딩으로 인해 q_ready 랑데뷰가 필요했습니다. H64 멤버는 64행만 스테이징하여 SMEM 사용량을 226,848 B에서 최대 192,032 B로 줄였으며, 데이터가 최종 슬롯에 직접 랜딩되도록 하여 동기화 오버헤드를 제거했습니다.

3. 디바이스 헬퍼의 콘텐츠 기반 공유 (csrc/cake_device_helpers/)

커널 생성 시 중복되는 헬퍼 함수들을 콘텐츠 해시 기반의 공유 헤더로 분리했습니다. 예를 들어, add_f32x2와 같은 연산은 다음과 같이 최적화된 인라인 어셈블리로 제공됩니다.

// csrc/cake_device_helpers/cake_add_f32x2_74578aa6314ab32b.cuh
__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) {
  float2 r;
  asm("add.rn.ftz.f32x2 %0, %1, %2;"
      : "=l"(*(unsigned long long*)&r)
      : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b));
  return r;
}

왜 이게 좋은가

이번 최적화는 단순히 코드 구조를 개선한 것을 넘어, 하드웨어 리소스(SMEM) 활용을 최적화하여 실질적인 성능 향상을 이끌어냈습니다.

• 성능 향상: B200 및 GB300 아키텍처에서 0.900.95x의 실행 시간(기존 대비 510% 속도 향상)을 기록했습니다.
• 리소스 효율: SMEM 사용량을 약 15% 이상 절감하여, 더 많은 워프가 동시에 실행될 수 있는 환경을 조성했습니다.
• 교훈: 범용적인 커널 템플릿보다는, 특정 데이터 크기(예: H=64)에 최적화된 커널을 동적으로 라우팅하는 것이 고성능 추론 엔진에서 매우 중요하다는 점을 보여줍니다.

결론

이번 변경은 FlashInfer의 cake 백엔드가 특정 아키텍처(SM100/SM103)의 특성을 얼마나 정교하게 활용할 수 있는지 보여주는 사례입니다. 특히 mbarrier를 활용한 클러스터 동기화와 인라인 어셈블리 최적화는 향후 다른 커널 최적화에도 중요한 레퍼런스가 될 것입니다.

참고 자료

• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-add
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#floating-point-instructions-fma

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