[flashinfer] FlashInfer MoE 최적화: PDL 스케줄링 개선 및 GEMM2 균형 잡힌 스토어 구현

PR 링크: flashinfer-ai/flashinfer#3756 상태: Merged | 변경: +76 / -25

들어가며

FlashInfer는 대규모 언어 모델(LLM) 추론을 위한 고성능 커널 라이브러리입니다. 특히 Mixture-of-Experts(MoE) 모델에서 성능은 연산 효율성뿐만 아니라, 스케줄링 오버헤드와 메모리 접근 패턴에 크게 좌우됩니다. 이번 PR은 CuTe DSL 기반의 MoE 커널에서 두 가지 핵심적인 병목 현상을 해결합니다. 첫째는 PDL(Persistent Data Layout) 스케줄링의 시작 지연 문제이며, 둘째는 GEMM2 단계에서 발생하는 부분 타일(partial-tile) 처리 시의 워프 간 불균형 문제입니다.

코드 분석

1. PDL 스케줄링 최적화 (Startup Latency 개선)

기존에는 클러스터 동기화 이후에 첫 번째 타일을 스케줄링했기 때문에, 초기 워크로드 진입 시 유휴 시간이 발생했습니다. 이를 해결하기 위해 첫 번째 타일을 미리 방출(emit)하도록 변경했습니다.

Before:

# 기존에는 클러스터 대기 후 스케줄링 루프 시작
if cute.size(self.cluster_shape_mn) > 1:
    cute.arch.cluster_wait()
# ... 이후 스케줄링 루프 진입

After:

# 첫 타일을 미리 방출하여 클러스터 대기 중에도 스케줄러가 준비되도록 함
work_tile = tile_sched.initial_work_tile_info()
# ... (중략) ...
if warp_idx == self.sched_warp_id:
    if work_tile.is_valid_tile:
        # 첫 타일 정보 공유 메모리 기록 및 동기화
        self.sched_sync_barrier.arrive_and_wait()
        # ...
# 이후 클러스터 대기
if cute.size(self.cluster_shape_mn) > 1:
    cute.arch.cluster_wait()

이 변경을 통해 파이프라인의 첫 단계에서 발생하는 유휴 시간을 줄여 전체적인 처리량을 개선했습니다.

2. GEMM2 부분 타일 스토어 균형화

GEMM2의 에필로그(epilogue) 단계에서 부분 타일을 처리할 때, 특정 워프에 부하가 집중되는 문제가 있었습니다. 이를 워프 간 인터리빙(interleaving) 방식으로 분산시켰습니다.

After (Balanced Store):

if is_partial_tile:
    # 워프 인덱스를 재배치하여 부하 분산
    reduce_row = (epi_tidx % self.threads_per_warp) * len(self.epilog_warp_id) + (epi_tidx // self.threads_per_warp)
reduce_permuted_row = tile_m_start + reduce_row

리뷰어 liyuhannnnn의 제안에 따라 epilog_sync_barrier의 위치를 최적화하여, Accumulator 파이프라인 해제와 스토어 작업을 분리함으로써 MMA 워프가 더 빠르게 다음 타일로 넘어갈 수 있도록 개선했습니다.

왜 이게 좋은가

이번 최적화는 특히 토큰 수가 적은(64~256) 환경에서 최대 4.5% 이상의 성능 향상을 보여줍니다.

1. 파이프라인 효율성: 초기 타일을 미리 방출함으로써 cluster_wait 동안 발생하는 CPU/GPU 유휴 시간을 숨길 수 있습니다.
2. 부하 분산: reduce_row 계산 시 워프 인덱스를 인터리빙함으로써, 부분 타일 처리 시 특정 워프가 병목이 되는 현상을 방지했습니다. 이는 GPU의 SIMT 구조에서 워프 간 작업 균형을 맞추는 전형적인 고성능 커널 최적화 기법입니다.

이러한 변경은 커널의 'Tail Latency'를 줄이는 데 매우 효과적이며, 대규모 MoE 모델 추론 시 전체적인 처리량(Throughput)을 높이는 데 기여합니다.

참고 자료

• https://github.com/flashinfer-ai/flashinfer
• https://github.com/NVIDIA/cutlass

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