본문으로 건너뛰기

[flashinfer] [FlashInfer] CUTLASS MoE 커널 최적화: 벡터화와 동적 스레드 할당으로 성능 한계 돌파하기

PR 링크: flashinfer-ai/flashinfer#3761 상태: Merged | 변경: +161 / -35

들어가며

최근 대규모 언어 모델(LLM)의 효율적인 추론을 위해 Mixture of Experts(MoE) 구조가 널리 채택되고 있습니다. MoE 모델은 수많은 파라미터를 가지면서도 실제 연산에는 일부 전문가(Expert)만 참여시키기 때문에 추론 속도가 빠르지만, 데이터의 레이아웃을 변경하거나 활성화 함수를 적용하는 과정에서 발생하는 Memory-bound 커널들이 전체 성능의 병목이 되곤 합니다.

이번 글에서는 FlashInfer 라이브러리의 CUTLASS 기반 MoE 백엔드에서 발견된 성능 저하 문제를 해결하고, Triton 구현체보다 느렸던 부분을 개선한 PR #3521의 핵심 최적화 기법을 분석합니다. 이 PR은 특히 Hopper(SM90/100) 아키텍처에서 BF16 데이터 타입의 성능을 끌어올리는 데 집중했습니다.

핵심 변경 사항 분석

1. 명시적 벡터화 (Explicit Vectorization)를 통한 메모리 대역폭 극대화

기존 코드에서는 expandInputRowsKernel과 doActivationKernel에서 데이터를 읽고 쓸 때 표준적인 배열 인덱싱 방식을 사용했습니다. 이론적으로는 컴파일러(NVCC)가 이를 최적화하여 벡터화된 로드(Vectorized Load)를 생성해야 하지만, 실제 SASS(Shader Assembly) 분석 결과 16바이트를 한 번에 읽는 LDG.E.128 대신 4바이트씩 네 번 읽는 LDG.E 명령어가 생성되고 있었습니다.

Before

// expandInputRowsKernel 내부
for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) {
  auto in_vec = source_row_ptr[elem_index];
  // ... 연산 ...
  dest_row_ptr[elem_index] = in_vec;
}

After

// cutlass::arch::global_load/store를 사용한 명시적 벡터화
for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) {
  DataElem in_vec;
  // 128-bit (16B) 로드를 명시적으로 강제
  cutlass::arch::global_load<DataElem, sizeof(DataElem)>(in_vec, source_row_ptr + elem_index, true);
  
  // ... 연산 ...

  auto out_vec = cutlass::platform::bit_cast<OutputElem>(in_vec);
  // 128-bit (16B) 스토어를 명시적으로 강제
  cutlass::arch::global_store<OutputElem, sizeof(OutputElem)>(out_vec, dest_row_ptr + elem_index, true);
}

왜 좋은가? CUDA 커널에서 메모리 대역폭을 최대로 활용하려면 트랜잭션 횟수를 줄여야 합니다. LDG.E.128은 한 번의 명령어로 128비트를 가져오므로, 여러 번의 작은 로드보다 오버헤드가 훨씬 적습니다. SASS 검증 결과, 이 변경을 통해 4x LDG.E가 1x LDG.E.128로 합쳐졌음을 확인했습니다.

2. 동적 스레드 할당을 통한 점유율(Occupancy) 개선

기존 doActivationKernel은 행(row)당 스레드 수를 256개로 고정(ACTIVATION_THREADS_PER_BLOCK = 256)하여 사용했습니다. 하지만 Qwen3-30B-A3B와 같이 중간 차원(Intermediate Size)이 작은 모델(예: 768)의 경우, 256개의 스레드는 과도하며 많은 워프(Warp)가 유휴 상태로 남게 됩니다.

Before

constexpr static int ACTIVATION_THREADS_PER_BLOCK = 256;
// 모든 워크로드에 대해 항상 256 스레드 사용

After

constexpr static int MAX_ACTIVATION_THREADS_PER_BLOCK = 256;

// doActivationKernelLauncher 내부에서 동적으로 결정 (코드 추정)
int threads = std::min(MAX_ACTIVATION_THREADS_PER_BLOCK, 
                       static_cast<int>(round_up(inter_size / ACTIVATION_ELEM_PER_THREAD, WARP_SIZE)));

또한, SM당 그리드 크기를 8~32 블록으로 제한하여 불필요한 워프 스케줄링 압박을 줄였습니다. 기존에는 고정적으로 8 블록을 사용했으나, 이제는 블록 크기에 따라 유연하게 조절됩니다.

왜 좋은가? 작은 워크로드에서 너무 많은 스레드를 할당하면 리소스 낭비가 발생하고, 이는 곧 낮은 점유율로 이어집니다. 워크로드 크기에 맞춰 스레드 수를 조절함으로써 GPU의 하드웨어 리소스를 더 효율적으로 분배할 수 있게 되었습니다.

3. 엄격한 메모리 정렬(Alignment) 검증

벡터화된 로드(LDG.E.128)는 데이터 주소가 16바이트 단위로 정렬되어 있어야 합니다. 이를 보장하기 위해 PR에서는 런타임 체크와 정적 어설션을 추가했습니다.

// expandInputRowsKernelLauncher 내부
TLLM_CHECK_WITH_INFO(((hidden_size * std::min(sizeof_bits<InputActivationsType>::value, 
                                              sizeof_bits<ExpandedActivationsType>::value)) % 128) == 0,
                     "hidden_size %ld rows are not a multiple of 16B; ...");

assert(reinterpret_cast<std::uintptr_t>(source_row_ptr) % sizeof(DataElem) == 0);

이러한 체크는 잘못된 메모리 접근으로 인한 런타임 에러를 방지하고, 최적화가 안전하게 적용될 수 있는 환경을 보장합니다.

성능 향상 결과

벤치마크 결과는 놀랍습니다.

  • SASS 검증: 16B 복사당 명령어 수가 4개에서 1개로 줄어든 것을 확인했습니다.
  • End-to-End 성능: B100(SM100) 환경에서 Qwen3 모델 기준 최대 1.08x의 성능 향상을 달성했습니다.
  • vLLM 통합 테스트: H200 장비에서 Qwen3-30B 모델 추론 시 TP4 설정에서 3.1%의 처리량(Throughput) 향상을 보였으며, 기존에 Triton 백엔드에 뒤처지던 성능 차이를 거의 극복했습니다.

일반적인 교훈: 시니어의 관점

  1. 컴파일러를 100% 믿지 마라: 고성능 컴퓨팅(HPC) 영역에서는 컴파일러가 최적의 어셈블리를 만들어줄 것이라 기대하기보다, SASS를 직접 확인하고 cutlass::arch::global_load와 같은 로우레벨 API를 사용하여 의도를 명확히 전달해야 합니다.
  2. 워크로드의 다양성을 고려하라: 대규모 모델뿐만 아니라 중간 규모나 작은 규모의 모델에서도 커널이 효율적으로 동작하도록 하드코딩된 매개변수(스레드 수 등)를 동적으로 최적화하는 것이 중요합니다.
  3. 정렬(Alignment)은 기본이다: 모든 고성능 메모리 최적화의 시작은 데이터 정렬입니다. 이를 무시하면 벡터화는 불가능합니다.

이 PR은 단순히 코드를 깔끔하게 만드는 것을 넘어, 하드웨어의 특성을 깊이 이해하고 이를 소프트웨어적으로 제어함으로써 실질적인 성능 이득을 얻어낸 훌륭한 사례입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글