[flashinfer] DeepSeek-V3 라우팅의 혁신: FlashInfer의 Cake 백엔드 가속 분석

PR 링크: flashinfer-ai/flashinfer#4587 상태: Merged | 변경: +8707 / -7

들어가며

최근 LLM 업계의 가장 뜨거운 감자는 단연 DeepSeek-V3입니다. DeepSeek-V3는 효율적인 MoE(Mixture of Experts) 구조를 통해 놀라운 성능을 보여주었지만, 그 이면에는 복잡한 라우팅 로직이 숨어 있습니다. 특히 Grouped Top-K 방식의 라우팅은 수천 개의 전문가 중 최적의 조합을 찾아야 하므로 연산 오버헤드가 상당합니다.

이번에 분석할 FlashInfer의 PR #574는 이 문제를 해결하기 위해 'Cake'라는 새로운 백엔드를 도입했습니다. 이 PR은 특히 NVIDIA의 최신 아키텍처인 Blackwell(B200, GB300) 환경에서 DeepSeek의 Fused Routing 성능을 비약적으로 향상시키는 것을 목표로 합니다. 실제 벤치마크 결과, GB300에서 기존 대비 최대 1.55배의 성능 향상을 기록했습니다.

시니어 엔지니어의 관점에서 이 PR이 왜 '미친 최적화'인지, 코드 수준에서 분석해 보겠습니다.

---

코드 분석: PTX와 Warp-level 최적화의 정수

이번 PR의 핵심은 cake_deepseek_fused_routing_kernels.cu 파일에 집약되어 있습니다. 이 파일은 사람이 직접 작성했다기보다, 최적화된 스케줄을 기반으로 생성된(Generated) 고성능 커널들로 가득 차 있습니다.

1. PTX 인라인 어셈블리를 통한 하드웨어 제어

기존의 일반적인 CUDA 커널이 C++ 추상화에 의존했다면, Cake 백엔드는 PTX(Parallel Thread Execution) 명령어를 직접 호출하여 하드웨어의 잠재력을 끝까지 끌어냅니다.

[Before: 일반적인 CUDA C++ 방식 (개념적)]

// 일반적인 방식은 표준 함수나 루프를 사용하여 reduction을 수행함
float max_val = -INFINITY;
for (int i = 0; i < 32; ++i) {
    max_val = max(max_val, shared_data[i]);
}

[After: Cake 백엔드의 PTX 최적화]

// csrc/fused_moe/cake_deepseek_fused_routing/cake_deepseek_fused_routing_kernels.cu

__device__ __forceinline__ int make_warp_uniform(int x) {
  int result;
  // shfl.sync를 통해 warp 내의 값을 즉시 브로드캐스팅
  asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x));
  return result;
}

// ... 중략 ...

unsigned int _warp_redux_u32_0;
// NVIDIA의 최신 하드웨어 가속 명령어인 redux.sync.max 사용
asm volatile("redux.sync.max.u32 %0, %1, 0xffffffff;"
             : "=r"(_warp_redux_u32_0)
             : "r"(first_encoded));

왜 좋은가?

• redux.sync.max.u32: Warp 내의 최댓값을 찾는 연산을 단 한 줄의 어셈블리 명령어로 처리합니다. 이는 기존의 __shfl_down_sync를 이용한 트리 구조의 reduction보다 훨씬 빠르며, 레지스터 압박(Register Pressure)을 줄여줍니다.
• make_warp_uniform: Warp 내의 모든 스레드가 동일한 값을 공유하도록 보장하여 제어 흐름의 분기(Divergence)를 최소화합니다.

2. Sigmoid 및 Tanh 근사 최적화

DeepSeek의 라우팅은 게이팅을 위해 Sigmoid 활성화 함수를 사용합니다. Cake 백엔드는 이를 위해 고정밀도 함수 대신 하드웨어 가속 근사 함수를 사용합니다.

[After: Tanh 근사를 이용한 Sigmoid 구현]

// csrc/fused_moe/cake_deepseek_fused_routing/cake_deepseek_fused_routing_kernels.cu

float score_value = scores[row_base + (unsigned long long)expert];
float _tanh_approx_0;
// tanh.approx.f32를 사용하여 매우 빠르게 tanh 계산
asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_0) : "f"(0.5f * score_value));
sigmoid_value = 0.5f * _tanh_approx_0 + 0.5f;

왜 좋은가?

• 표준 expf() 기반의 Sigmoid는 연산 비용이 비쌉니다. 하지만 tanh.approx.f32는 GPU의 SFU(Special Function Unit)에서 직접 처리되어 처리량이 매우 높습니다. 딥러닝 추론에서는 이 정도의 근사 오차가 최종 정확도에 미치는 영향이 미미하면서도 속도는 수 배 빠릅니다.

3. Fused Grouped Top-K 로직

DeepSeek-V3의 핵심인 '그룹별 상위 전문가 선택' 로직이 하나의 커널 안에 완전히 융합(Fused)되어 있습니다.

// 1단계: 전문가별 점수 계산 및 Bias 적용
biased_value = sigmoid_value + bias[expert];

// 2단계: Warp 내에서 그룹별 최댓값(Group First/Second) 추출
// (위에서 언급한 redux.sync.max 활용)

// 3단계: Shared Memory를 통한 그룹 간 점수 취합 및 최종 Top-K 선택
if (lane_0 == 0) {
    if (group < n_group) {
        group_scores[group] = group_first + group_second;
    }
}
__syncthreads();

이 과정에서 글로벌 메모리 접근을 최소화하고, 모든 중간 결과물을 레지스터와 Shared Memory 내에서 처리합니다. 이는 메모리 대역폭이 병목인 MoE 라우팅에서 성능을 결정짓는 핵심 요소입니다.

---

왜 이게 좋은가: 성능과 안정성의 균형

1. 압도적인 성능 향상

PR 설명에 포함된 벤치마크 수치는 놀랍습니다.

• B200: 약 1.1x ~ 1.2x 성능 향상
• GB300: 약 1.3x ~ 1.55x 성능 향상
• 특히 특정 셰이프(T1 E384 k1 G1)에서는 기존 대비 9.38배라는 극적인 수치를 보여주기도 합니다. 이는 단순한 코드 수정을 넘어, 하드웨어 아키텍처에 최적화된 커널 스케줄링의 승리라고 볼 수 있습니다.

2. 리뷰어 피드백: 엔지니어링의 신중함

리뷰어 yzh119는 성능 향상을 반기면서도 중요한 지적을 남겼습니다.

"Could we avoid making Cake the default path in this PR? ... I would prefer to keep the existing backend as the default for now and put Cake behind an explicit opt-in/feature flag."

이는 아무리 성능이 좋아도 새로운 백엔드가 기존 프로덕션 환경에 미칠 수 있는 잠재적 위험(Regression)을 경계하는 시니어 엔지니어의 자세를 보여줍니다. 결과적으로 Cake 백엔드는 명시적인 옵션을 통해서만 활성화되도록 조정되었으며, 이는 시스템의 안정성을 유지하면서 혁신을 도입하는 모범 사례입니다.

---

결론

FlashInfer의 Cake 백엔드 도입은 "하드웨어를 깊게 이해하는 것이 곧 소프트웨어의 경쟁력"임을 다시 한번 증명했습니다. PTX를 통한 저수준 제어, SFU를 활용한 수학적 근사, 그리고 복잡한 MoE 로직의 커널 융합은 Blackwell 시대의 LLM 서빙 최적화가 나아가야 할 방향을 제시하고 있습니다.

DeepSeek-V3와 같은 거대 모델을 효율적으로 서빙하고자 하는 팀이라면, 이번 PR에 포함된 커널 구현 방식과 최적화 기법을 반드시 참고해야 할 것입니다.

참고 자료

• PTX ISA - Parallel Synchronization and Communication Instructions (redux)
• CUDA Math API - tanh.approx.f32
• FlashInfer GitHub Repository

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