[flashinfer] NVFP4 MoE All-to-All 성능 최적화: Phased Dispatch 기법 분석

2026년 9월 8일수정: 2026년 9월 8일

PR 링크: flashinfer-ai/flashinfer#4067 상태: Merged | 변경: +226 / -33

들어가며

대규모 언어 모델(LLM)의 효율성을 높이기 위해 Mixture-of-Experts(MoE) 구조가 널리 사용되고 있습니다. MoE 아키텍처의 핵심 병목 중 하나는 서로 다른 GPU 간에 토큰 데이터를 주고받는 All-to-All Dispatch 과정입니다. 특히 NVIDIA의 새로운 데이터 포맷인 NVFP4를 사용할 경우, 활성화 함수(Activations), 스케일(Scales), 전문가 ID(Expert IDs), 라우팅 가중치(Routing Weights) 등 서로 크기가 다른 네 가지 페이로드를 동시에 전송해야 합니다.

기존의 FlashInfer 구현에서는 모든 CTA(Cooperative Thread Array) 스레드가 모든 페이로드에 대해 원격 목적지 포인터(Remote Destination Pointers)를 계산하는 범용적인 방식을 사용했습니다. 이는 구현이 간단하지만, 특정 스레드가 처리할 데이터가 없는 경우에도 포인터 계산 오버헤드를 감수해야 한다는 단점이 있었습니다.

이번 PR(#3846의 후속)은 NVFP4의 페이로드 구조가 예측 가능하다는 점에 착안하여, 각 페이로드별로 필요한 워커(Warp)만 할당하는 Phased Dispatch 기법을 도입했습니다. 이를 통해 특히 EP(Expert Parallelism) 규모가 큰 환경에서 비약적인 성능 향상을 이끌어냈습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. 워커 할당 전략의 변화

가장 먼저 눈에 띄는 변화는 vectorized_dispatch_impl 함수의 인터페이스 변경입니다. 기존에는 모든 스레드가 전체 blockDim.x를 기준으로 루프를 돌았지만, 이제는 특정 페이로드에 할당된 워커들만 작업을 수행하도록 worker_idx와 num_workers를 명시적으로 받습니다.

Before:

// 모든 스레드가 전체 blockDim.x를 기준으로 루프를 수행
template <int VEC_SIZE, int TOP_K>
__device__ void vectorized_dispatch_impl(..., int const* topk_send_indices) {
  // ... 포인터 계산 ...
  int const stride = blockDim.x * VEC_SIZE;
  for (int offset = threadIdx.x * VEC_SIZE; offset < bytes_per_token; offset += stride) {
    // 데이터 로드 및 저장
  }
}

After:

// 할당된 워커 서브셋만 포인터 설정 및 복사 비용을 지불
template <int VEC_SIZE, int TOP_K>
__device__ void vectorized_dispatch_impl(..., int worker_idx, int num_workers) {
  // ... 포인터 계산 ...
  int const stride = num_workers * VEC_SIZE;
  for (int offset = worker_idx * VEC_SIZE; offset < bytes_per_token; offset += stride) {
    // 데이터 로드 및 저장
  }
}

2. Phased Schedule 구현

moeA2ADispatchKernel 내부에서는 if constexpr을 사용하여 NVFP4 레이아웃일 때만 최적화된 경로를 타도록 설계되었습니다. 128개의 스레드로 구성된 CTA를 다음과 같이 쪼개어 사용합니다.

Warp 0, 1 (64 스레드): Packed Activations 처리
Warp 2 (32 스레드): Scales 처리
Warp 3 (32 스레드): Expert IDs 및 Routing Weights 처리 (각 16 스레드씩 분할)

Core Logic (After):

if constexpr (PHASE_H2048_NVFP4) {
  int warp_id = thread_idx / warpSize;
  int lane_id = thread_idx % warpSize;

  if (warp_id < kDispatchActivationWarps) {
    // Warp 0-1: Activations (H2048 기준 1024 bytes)
    vectorized_dispatch_impl<16, NUM_DESTINATIONS>(..., thread_idx, kNvfp4H2048ActivationWorkers);
  } else if (warp_id == kDispatchActivationWarps && lane_id < kNvfp4H2048ScaleWorkers) {
    // Warp 2: Scales
    vectorized_dispatch_impl<16, NUM_DESTINATIONS>(..., lane_id, kNvfp4H2048ScaleWorkers);
  } else if (warp_id == kDispatchActivationWarps + 1 && lane_id < 2 * kTopK22MetadataWorkers) {
    // Warp 3: Metadata (Expert IDs & Weights)
    bool second_tail = lane_id >= kTopK22MetadataWorkers;
    int payload_idx = second_tail ? 3 : 2;
    int worker_idx = second_tail ? lane_id - kTopK22MetadataWorkers : lane_id;
    vectorized_dispatch_impl<8, NUM_DESTINATIONS>(..., worker_idx, kTopK22MetadataWorkers);
  }
}

3. 유효성 검사 로직 추가

모든 케이스에 이 최적화를 적용할 수는 없으므로, isPhasedNvfp4PayloadLayout 함수를 통해 현재 레이아웃이 Phased Dispatch에 적합한지 검증합니다. 페이로드 개수가 4개인지, top_k가 특정 범위를 넘지 않는지, 데이터 정렬(Alignment)이 맞는지 등을 체크합니다.

왜 이게 좋은 최적화인가?

1. 불필요한 계산 제거 (Redundant Work Elimination)

기존 방식에서는 128개 스레드 모두가 4개의 페이로드 각각에 대해 목적지 주소를 계산했습니다. 하지만 실제로는 특정 페이로드의 크기가 작아 몇 개의 스레드만으로도 충분한 경우가 많습니다. Phased Dispatch는 각 스레드가 자신에게 할당된 페이로드에 대해서만 주소 계산과 복사를 수행하게 하여, GPU의 연산 자원과 메모리 대역폭 낭비를 줄였습니다.

2. 성능 향상 수치

GB200 환경에서 Nemotron Ultra(H2048, Top-K 22) 모델을 대상으로 테스트한 결과, 놀라운 성능 향상이 관찰되었습니다.

EP8 (2 Nodes): 토큰 수가 많을 때(8192) 지연시간이 308.292 us에서 96.494 us로 무려 68.70% 감소했습니다.
EP16 (4 Nodes): 동일 조건에서 지연시간이 52.63% 감소했습니다.
End-to-End: 전체 워크로드(TTFT) 기준 약 2.43%의 개선을 보였습니다.

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

PR Analysis 의 다른글

이전글 [sglang] SGLang의 새로운 스케줄링 정책: HRRN 도입을 통한 TTFT 최적화
현재글 : [flashinfer] NVFP4 MoE All-to-All 성능 최적화: Phased Dispatch 기법 분석
다음글 없음