[flashinfer] FlashInfer: Blackwell W8A8 AlphaMoE Expert 계산 커널 퓨전으로 성능 비약적 향상

PR 링크: flashinfer-ai/flashinfer#4287 상태: Merged | 변경: +2396 / -0

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 Mixture-of-Experts (MoE) 아키텍처는 모델의 파라미터 수를 크게 늘리면서도 계산 비용을 효율적으로 유지하는 핵심 기술로 주목받고 있습니다. MoE 모델은 입력 토큰마다 소수의 '전문가(Expert)'를 선택하여 계산을 수행하며, 이 전문가들의 계산 효율성이 전체 모델의 성능을 좌우합니다. 특히, NVIDIA의 최신 Blackwell 아키텍처(SM100a/SM103a)와 8비트 가중치(W8) 및 8비트 활성화(A8) 양자화(W8A8)를 활용할 때, 이 전문가 계산의 최적화는 더욱 중요해집니다.

이번 FlashInfer PR(flashinfer-ai/flashinfer#4341)은 Blackwell GPU에 특화된 W8A8 AlphaMoE Expert 계산을 위한 고도로 최적화된 커널을 추가하여, 기존의 다단계 계산 파이프라인을 단일 커널로 퓨전(fuse)하는 혁신적인 개선을 이루어냈습니다. 이 최적화는 MoE 모델의 추론 성능을 비약적으로 향상시키는 데 기여합니다.

문제점: MoE Expert 계산의 병목

기존 MoE Expert 계산은 일반적으로 여러 개의 개별 CUDA 커널로 구성됩니다. PR 설명에 따르면, 이는 다음과 같은 5단계의 커널 시퀀스로 이루어져 있었습니다:

1. GEMM1 (General Matrix Multiply): Gate 및 Up Projection 계산
2. SwiGLU: 활성화 함수 적용
3. Intermediate FP8 Quantization: 중간 결과를 FP8로 재양자화
4. GEMM2: Down Projection 계산
5. Combine: 최종 결과 취합

각 단계가 별도의 커널로 실행될 때 발생하는 주요 병목 현상은 다음과 같습니다:

• 커널 런치 오버헤드: GPU에서 커널을 실행할 때마다 발생하는 CPU-GPU 동기화 및 런치 비용은 전체 실행 시간에 상당한 영향을 미칩니다.
• 글로벌 메모리 접근: 각 커널이 중간 결과를 글로벌 메모리에 쓰고 다음 커널이 이를 다시 읽는 과정에서 발생하는 불필요한 메모리 대역폭 소모는 성능 저하의 주된 원인입니다.
• 데이터 지역성 부족: 여러 커널에 걸쳐 데이터가 처리되면서 캐시 효율성이 떨어지고 데이터 지역성을 활용하기 어렵습니다.

이러한 문제점들을 해결하기 위해, PR은 이 5단계의 계산을 단일 커널로 퓨전하는 것을 목표로 했습니다.

핵심 최적화: 5개 커널을 하나로

이번 PR의 핵심은 alphamoe_fp8_block_scale_aligned_moe라는 새로운 CUDA 커널을 도입하여, 기존의 5단계 MoE Expert 계산을 단일 커널로 통합한 것입니다. 이 커널은 Blackwell 아키텍처의 특성을 최대한 활용하도록 저수준으로 최적화되었습니다.

csrc/alphamoe_sm100.cu 파일 분석

새로 추가된 csrc/alphamoe_sm100.cu 파일은 Blackwell GPU (SM100a/SM103a)를 위한 AlphaMoE 커널의 구현을 담고 있습니다. 이 파일은 TVM-FFI semantic wrapper를 통해 생성된 코드를 포함하며, 다음과 같은 주요 최적화 기법들을 사용합니다.

1. 커널 퓨전 (Kernel Fusion)

가장 중요한 변경사항은 여러 개의 논리적 연산을 하나의 물리적 CUDA 커널로 통합한 것입니다. 이는 커널 런치 오버헤드를 제거하고, 중간 결과를 글로벌 메모리에 저장할 필요 없이 레지스터나 공유 메모리(Shared Memory)에 유지할 수 있게 하여 메모리 대역폭 사용량을 극적으로 줄입니다.

PR 설명에 따르면, 이 커널은 gate/up projection, SwiGLU, intermediate FP8 requantization and down projection을 모두 한 번에 처리합니다.

2. W8A8 양자화 및 Blackwell Tensor Cores 활용

이 커널은 W8A8 (8비트 가중치, 8비트 활성화) 형식을 사용합니다. 이는 메모리 사용량을 줄이고, Blackwell 아키텍처의 Tensor Cores가 FP8 연산을 매우 효율적으로 수행할 수 있도록 합니다. 특히, tcgen05_mma_f8f6f4와 같은 저수준 Tensor Core 명령어를 직접 사용하여 최대 성능을 끌어냅니다.

// Before (Conceptual - multiple kernels)
// kernel_gemm1(activations, weights_up, ...);
// kernel_swiglu(intermediate_up, ...);
// kernel_quantize(swiglu_output, ...);
// kernel_gemm2(quantized_output, weights_down, ...);
// kernel_combine(final_output, ...);

// After (Fused kernel in alphamoe_sm100.cu)
__device__ __forceinline__ void tcgen05_mma_f8f6f4(
    int taddr, uint64_t a_desc, uint64_t b_desc,
    uint32_t i_desc, int enable_input_d) {
    asm volatile(
        

## 참고 자료
- https://pytorch.org/docs/stable/generated/torch.compile.html

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