[flashinfer] FlashInfer: Blackwell 아키텍처를 위한 결정론적 BGMV MoE 최적화

PR 링크: flashinfer-ai/flashinfer#4821 상태: Merged | 변경: +6937 / -2

들어가며

최신 대규모 언어 모델(LLM)의 핵심인 Mixture-of-Experts(MoE) 구조에서 BGMV(Batch GEMV) 연산은 추론 속도를 결정짓는 병목 구간입니다. 이번 PR은 NVIDIA의 차세대 아키텍처인 Blackwell(SM100) GPU의 성능을 극대화하기 위해, JIT(Just-In-Time) 컴파일 방식이 아닌, 특정 하드웨어에 최적화된 AOT(Ahead-Of-Time) 생성 커널을 도입하여 BGMV MoE 파이프라인을 획기적으로 개선했습니다.

코드 분석

1. Blackwell 전용 커널 생성 (csrc/blackwell_bgmv_moe/sm100a/)

기존의 범용적인 BGMV 구현과 달리, 이 PR은 특정 hidden size(2688, 3072)에 대해 Blackwell 아키텍처의 하드웨어 특성을 반영한 커널을 직접 생성하여 포함했습니다.

// csrc/blackwell_bgmv_moe/sm100a/blackwell_bgmv_moe_bf16_h2688_sm100a.cu
__global__ __launch_bounds__(128, 1) void kernel_flashinfer_bgmv_moe_shrink_bf16_h2688_r32_p4_s3(...) {
  // ... Blackwell 아키텍처의 Tensor Map 및 비동기 복사 최적화 적용
  asm volatile(
      "cp.async.cg.shared::cta.global [%0], [%1], 16;" ::"r"(
          x_smem_addr + (unsigned int)((tile % 3 * 4 * 1024 + pp_1 * 1024 + tid * 8) * 2)),
      "l"(reinterpret_cast<const __nv_bfloat16*>(x_raw) +
          (tokens[pp_1] * 2688 + (long long)k_base)));
}

이 코드는 cp.async 명령어를 사용하여 글로벌 메모리에서 공유 메모리로의 데이터 이동을 최적화하고, Blackwell의 SM100 아키텍처에 맞춘 레지스터 및 공유 메모리 사용량을 고정하여 성능을 극대화합니다.

2. CUDA Graph 기반의 결정론적 파이프라인 (prepare_bgmv_moe)

prepare_bgmv_moe 함수는 워크스페이스를 고정하고 전체 파이프라인을 하나의 CUDA Graph로 캡처하여 실행합니다. 이를 통해 커널 실행 간의 오버헤드를 제거하고, 결정론적(deterministic)인 연산 순서를 보장하여 출력의 일관성을 확보했습니다.

# 벤치마크 코드 내 사용 예시
plan = prepare_bgmv_moe(*inputs, backend="blackwell")
candidate_output = plan.run()

왜 이게 좋은가

이번 최적화는 Blackwell GPU 환경에서 기하 평균 기준 3.17배의 성능 향상을 보여주었습니다. 특히 토큰 수가 적을 때(예: 32개) 최대 4.14배의 속도 향상을 기록했습니다.

주요 교훈

1. 하드웨어 특화 커널의 힘: 범용 커널은 유연하지만, 특정 아키텍처(SM100)의 공유 메모리 구조와 비동기 복사 기능을 직접 활용하는 커널은 성능 격차를 크게 벌릴 수 있습니다.
2. CUDA Graph의 활용: 반복적인 MoE 연산에서 CUDA Graph를 사용하여 커널 런칭 오버헤드를 줄이는 것은 실시간 추론 성능에 필수적입니다.
3. 결정론적 연산: 원자적 연산(atomics)을 제거하고 고정된 입력 순서로 축소(reduction)를 수행함으로써, 성능 향상뿐만 아니라 결과의 재현성까지 확보했습니다.

참고 자료

• FlashInfer Documentation
• CUDA C++ Programming Guide - Asynchronous Copy

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-copy

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